From a81dc5eaca5ca2b7d25bc0701ae6910ce2442970 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 20 Dec 2025 22:00:53 +0100 Subject: [PATCH] Add configurable OAuth storage backend and respect --models in model selector - Add setOAuthStorage() and resetOAuthStorage() to pi-ai for custom storage backends - Configure coding-agent to use its own configurable OAuth path via getOAuthPath() - Model selector (/model command) now only shows models from --models scope when set - Rewrite OAuth documentation in pi-ai README with examples Fixes #255 --- packages/ai/README.md | 128 ++++++++++++++++-- packages/ai/src/utils/oauth/index.ts | 3 + packages/ai/src/utils/oauth/storage.ts | 119 +++++++++++----- packages/coding-agent/CHANGELOG.md | 4 + packages/coding-agent/DEVELOPMENT.md | 7 +- packages/coding-agent/src/core/oauth/index.ts | 14 +- packages/coding-agent/src/main.ts | 36 ++++- .../interactive/components/model-selector.ts | 59 +++++--- .../src/modes/interactive/interactive-mode.ts | 1 + 9 files changed, 300 insertions(+), 71 deletions(-) diff --git a/packages/ai/README.md b/packages/ai/README.md index 052d22e0..2a18ab3e 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -1115,29 +1115,137 @@ setApiKey('anthropic', 'sk-ant-...'); const key = getApiKey('openai'); ``` -## GitHub Copilot +## OAuth Providers -GitHub Copilot is available as a provider, requiring OAuth authentication via GitHub's device flow. +Several providers require OAuth authentication instead of static API keys. This library provides login flows, automatic token refresh, and credential storage for: -**Using with `@mariozechner/pi-coding-agent`**: Use `/login` and select "GitHub Copilot" to authenticate. All models are automatically enabled after login. Token stored in `~/.pi/agent/oauth.json`. +- **Anthropic** (Claude Pro/Max subscription) +- **GitHub Copilot** (Copilot subscription) +- **Google Gemini CLI** (Free Gemini 2.0/2.5 via Google Cloud Code Assist) +- **Antigravity** (Free Gemini 3, Claude, GPT-OSS via Google Cloud) -**Using standalone**: If you have a valid Copilot OAuth token (e.g., from the coding agent's `oauth.json`): +Credentials are stored in `~/.pi/agent/oauth.json` by default (with `chmod 600` permissions). Use `setOAuthStorage()` to configure a custom storage backend for different locations or environments (the coding-agent does this to respect its configurable config directory). + +### Using with @mariozechner/pi-coding-agent + +Use `/login` and select a provider to authenticate. Tokens are automatically refreshed when expired. + +### Programmatic OAuth + +For standalone usage, the library exposes low-level OAuth functions: ```typescript -import { getModel, complete } from '@mariozechner/pi-ai'; +import { + // Login functions (each implements provider-specific OAuth flow) + loginAnthropic, + loginGitHubCopilot, + loginGeminiCli, + loginAntigravity, + + // Token management + refreshToken, // Refresh token for any provider + getOAuthApiKey, // Get API key (auto-refreshes if expired) + + // Credential storage + loadOAuthCredentials, + saveOAuthCredentials, + removeOAuthCredentials, + hasOAuthCredentials, + listOAuthProviders, + getOAuthPath, + + // Types + type OAuthProvider, // 'anthropic' | 'github-copilot' | 'google-gemini-cli' | 'google-antigravity' + type OAuthCredentials, +} from '@mariozechner/pi-ai'; +``` + +### Login Flow Example + +Each provider has a different OAuth flow. Here's an example with GitHub Copilot: + +```typescript +import { loginGitHubCopilot, saveOAuthCredentials } from '@mariozechner/pi-ai'; + +const credentials = await loginGitHubCopilot({ + onAuth: (url, instructions) => { + // Display the URL and instructions to the user + console.log(`Open: ${url}`); + if (instructions) console.log(instructions); + }, + onPrompt: async (prompt) => { + // Prompt user for input (e.g., device code confirmation) + return await getUserInput(prompt.message); + }, + onProgress: (message) => { + // Optional: show progress updates + console.log(message); + } +}); + +// Save credentials for later use +saveOAuthCredentials('github-copilot', credentials); +``` + +### Using OAuth Tokens + +Call `getOAuthApiKey()` before **every** `complete()` or `stream()` call. This function checks token expiry and refreshes automatically when needed: + +```typescript +import { getModel, complete, getOAuthApiKey } from '@mariozechner/pi-ai'; const model = getModel('github-copilot', 'gpt-4o'); +// Always call getOAuthApiKey() right before the API call +// Do NOT cache the result - tokens expire and need refresh +const apiKey = await getOAuthApiKey('github-copilot'); +if (!apiKey) { + throw new Error('Not logged in to GitHub Copilot'); +} + const response = await complete(model, { messages: [{ role: 'user', content: 'Hello!' }] -}, { - apiKey: 'tid=...;exp=...;proxy-ep=...' // OAuth token from ~/.pi/agent/oauth.json -}); +}, { apiKey }); ``` -**Note**: OAuth tokens expire and need periodic refresh. The coding agent handles this automatically. +### Custom Storage Backend -If you get "The requested model is not supported" error, enable the model manually in VS Code: open Copilot Chat, click the model selector, select the model (warning icon), and click "Enable". +Override the default storage location with `setOAuthStorage()`: + +```typescript +import { setOAuthStorage, resetOAuthStorage } from '@mariozechner/pi-ai'; +import { readFileSync, writeFileSync } from 'fs'; + +// Custom file path +setOAuthStorage({ + load: () => { + try { + return JSON.parse(readFileSync('/custom/path/oauth.json', 'utf-8')); + } catch { + return {}; + } + }, + save: (storage) => { + writeFileSync('/custom/path/oauth.json', JSON.stringify(storage, null, 2)); + } +}); + +// In-memory storage (for testing or browser environments) +let memoryStorage = {}; +setOAuthStorage({ + load: () => memoryStorage, + save: (storage) => { memoryStorage = storage; } +}); + +// Reset to default (~/.pi/agent/oauth.json) +resetOAuthStorage(); +``` + +### Provider Notes + +**GitHub Copilot**: If you get "The requested model is not supported" error, enable the model manually in VS Code: open Copilot Chat, click the model selector, select the model (warning icon), and click "Enable". + +**Google Gemini CLI / Antigravity**: These use Google Cloud OAuth. The API key returned by `getOAuthApiKey()` is a JSON string containing both the token and project ID, which the library handles automatically. ## License diff --git a/packages/ai/src/utils/oauth/index.ts b/packages/ai/src/utils/oauth/index.ts index d8286be1..95d4599c 100644 --- a/packages/ai/src/utils/oauth/index.ts +++ b/packages/ai/src/utils/oauth/index.ts @@ -43,8 +43,11 @@ export { type OAuthCredentials, type OAuthProvider, type OAuthStorage, + type OAuthStorageBackend, removeOAuthCredentials, + resetOAuthStorage, saveOAuthCredentials, + setOAuthStorage, } from "./storage.js"; // ============================================================================ diff --git a/packages/ai/src/utils/oauth/storage.ts b/packages/ai/src/utils/oauth/storage.ts index 92860a01..3be2495a 100644 --- a/packages/ai/src/utils/oauth/storage.ts +++ b/packages/ai/src/utils/oauth/storage.ts @@ -1,5 +1,8 @@ /** - * OAuth credential storage for ~/.pi/agent/oauth.json + * OAuth credential storage with configurable backend. + * + * Default: ~/.pi/agent/oauth.json + * Override with setOAuthStorage() for custom storage locations or backends. */ import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; @@ -23,54 +26,104 @@ export interface OAuthStorage { export type OAuthProvider = "anthropic" | "github-copilot" | "google-gemini-cli" | "google-antigravity"; /** - * Get the path to the OAuth credentials file + * Storage backend interface. + * Implement this to use a custom storage location or backend. */ -export function getOAuthPath(): string { - return join(homedir(), ".pi", "agent", "oauth.json"); +export interface OAuthStorageBackend { + /** Load all OAuth credentials. Return empty object if none exist. */ + load(): OAuthStorage; + /** Save all OAuth credentials. */ + save(storage: OAuthStorage): void; } -/** - * Ensure the config directory exists - */ -function ensureConfigDir(): void { - const configDir = dirname(getOAuthPath()); - if (!existsSync(configDir)) { - mkdirSync(configDir, { recursive: true, mode: 0o700 }); - } -} +// ============================================================================ +// Default filesystem backend +// ============================================================================ -/** - * Load all OAuth credentials from ~/.pi/agent/oauth.json - */ -export function loadOAuthStorage(): OAuthStorage { - const filePath = getOAuthPath(); - if (!existsSync(filePath)) { +const DEFAULT_PATH = join(homedir(), ".pi", "agent", "oauth.json"); + +function defaultLoad(): OAuthStorage { + if (!existsSync(DEFAULT_PATH)) { return {}; } - try { - const content = readFileSync(filePath, "utf-8"); + const content = readFileSync(DEFAULT_PATH, "utf-8"); return JSON.parse(content); } catch { return {}; } } +function defaultSave(storage: OAuthStorage): void { + const configDir = dirname(DEFAULT_PATH); + if (!existsSync(configDir)) { + mkdirSync(configDir, { recursive: true, mode: 0o700 }); + } + writeFileSync(DEFAULT_PATH, JSON.stringify(storage, null, 2), "utf-8"); + chmodSync(DEFAULT_PATH, 0o600); +} + +// ============================================================================ +// Configurable backend +// ============================================================================ + +let currentBackend: OAuthStorageBackend = { + load: defaultLoad, + save: defaultSave, +}; + /** - * Save all OAuth credentials to ~/.pi/agent/oauth.json + * Configure the OAuth storage backend. + * + * @example + * // Custom file path + * setOAuthStorage({ + * load: () => JSON.parse(readFileSync('/custom/path/oauth.json', 'utf-8')), + * save: (storage) => writeFileSync('/custom/path/oauth.json', JSON.stringify(storage)) + * }); + * + * @example + * // In-memory storage (for testing) + * let memoryStorage = {}; + * setOAuthStorage({ + * load: () => memoryStorage, + * save: (storage) => { memoryStorage = storage; } + * }); */ -function saveOAuthStorage(storage: OAuthStorage): void { - ensureConfigDir(); - const filePath = getOAuthPath(); - writeFileSync(filePath, JSON.stringify(storage, null, 2), "utf-8"); - chmodSync(filePath, 0o600); +export function setOAuthStorage(backend: OAuthStorageBackend): void { + currentBackend = backend; +} + +/** + * Reset to default filesystem storage (~/.pi/agent/oauth.json) + */ +export function resetOAuthStorage(): void { + currentBackend = { load: defaultLoad, save: defaultSave }; +} + +/** + * Get the default OAuth path (for reference, may not be used if custom backend is set) + */ +export function getOAuthPath(): string { + return DEFAULT_PATH; +} + +// ============================================================================ +// Public API (uses current backend) +// ============================================================================ + +/** + * Load all OAuth credentials + */ +export function loadOAuthStorage(): OAuthStorage { + return currentBackend.load(); } /** * Load OAuth credentials for a specific provider */ export function loadOAuthCredentials(provider: string): OAuthCredentials | null { - const storage = loadOAuthStorage(); + const storage = currentBackend.load(); return storage[provider] || null; } @@ -78,18 +131,18 @@ export function loadOAuthCredentials(provider: string): OAuthCredentials | null * Save OAuth credentials for a specific provider */ export function saveOAuthCredentials(provider: string, creds: OAuthCredentials): void { - const storage = loadOAuthStorage(); + const storage = currentBackend.load(); storage[provider] = creds; - saveOAuthStorage(storage); + currentBackend.save(storage); } /** * Remove OAuth credentials for a specific provider */ export function removeOAuthCredentials(provider: string): void { - const storage = loadOAuthStorage(); + const storage = currentBackend.load(); delete storage[provider]; - saveOAuthStorage(storage); + currentBackend.save(storage); } /** @@ -103,6 +156,6 @@ export function hasOAuthCredentials(provider: string): boolean { * List all providers with OAuth credentials */ export function listOAuthProviders(): string[] { - const storage = loadOAuthStorage(); + const storage = currentBackend.load(); return Object.keys(storage); } diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index c11b5864..66145d5b 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -10,6 +10,10 @@ - **Google Antigravity OAuth provider**: Access Gemini 3, Claude (sonnet/opus thinking models), and GPT-OSS models for free via Google's Antigravity sandbox. Login with `/login` and select "Antigravity". Uses your Google account with rate limits. +### Changed + +- **Model selector respects --models scope**: The `/model` command now only shows models specified via `--models` flag when that flag is used, instead of showing all available models. This prevents accidentally selecting models from unintended providers. ([#255](https://github.com/badlogic/pi-mono/issues/255)) + ### Fixed - **Connection errors not retried**: Added "connection error" to the list of retryable errors so Anthropic connection drops trigger auto-retry instead of silently failing. ([#252](https://github.com/badlogic/pi-mono/issues/252)) diff --git a/packages/coding-agent/DEVELOPMENT.md b/packages/coding-agent/DEVELOPMENT.md index e979d2ee..c8467348 100644 --- a/packages/coding-agent/DEVELOPMENT.md +++ b/packages/coding-agent/DEVELOPMENT.md @@ -64,11 +64,8 @@ src/ │ ├── slash-commands.ts # loadSlashCommands() from ~/.pi/agent/commands/ │ ├── system-prompt.ts # buildSystemPrompt(), loadProjectContextFiles() │ │ -│ ├── oauth/ # OAuth authentication -│ │ ├── index.ts # OAuth exports -│ │ ├── anthropic.ts # Anthropic OAuth (Claude Pro/Max) -│ │ ├── github-copilot.ts # GitHub Copilot OAuth -│ │ └── storage.ts # Token persistence +│ ├── oauth/ # OAuth authentication (thin wrapper) +│ │ └── index.ts # Re-exports from @mariozechner/pi-ai with convenience wrappers │ │ │ ├── hooks/ # Hook system for extending behavior │ │ ├── index.ts # Hook exports diff --git a/packages/coding-agent/src/core/oauth/index.ts b/packages/coding-agent/src/core/oauth/index.ts index 8cfb6b64..3d841972 100644 --- a/packages/coding-agent/src/core/oauth/index.ts +++ b/packages/coding-agent/src/core/oauth/index.ts @@ -13,15 +13,25 @@ import { loginGitHubCopilot, type OAuthCredentials, type OAuthProvider, + type OAuthStorageBackend, refreshToken as refreshTokenFromAi, removeOAuthCredentials, + resetOAuthStorage, saveOAuthCredentials, + setOAuthStorage, } from "@mariozechner/pi-ai"; // Re-export types and functions -export type { OAuthCredentials, OAuthProvider }; +export type { OAuthCredentials, OAuthProvider, OAuthStorageBackend }; export { listOAuthProvidersFromAi as listOAuthProviders }; -export { getOAuthApiKey, loadOAuthCredentials, removeOAuthCredentials, saveOAuthCredentials }; +export { + getOAuthApiKey, + loadOAuthCredentials, + removeOAuthCredentials, + resetOAuthStorage, + saveOAuthCredentials, + setOAuthStorage, +}; // Types for OAuth flow export interface OAuthAuthInfo { diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index c94d0e8b..2400e4e4 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -3,13 +3,15 @@ */ import { Agent, type Attachment, ProviderTransport, type ThinkingLevel } from "@mariozechner/pi-agent-core"; -import { supportsXhigh } from "@mariozechner/pi-ai"; +import { setOAuthStorage, supportsXhigh } from "@mariozechner/pi-ai"; import chalk from "chalk"; +import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; +import { dirname } from "path"; import { type Args, parseArgs, printHelp } from "./cli/args.js"; import { processFileArguments } from "./cli/file-processor.js"; import { listModels } from "./cli/list-models.js"; import { selectSession } from "./cli/session-picker.js"; -import { getModelsPath, VERSION } from "./config.js"; +import { getModelsPath, getOAuthPath, VERSION } from "./config.js"; import { AgentSession } from "./core/agent-session.js"; import { discoverAndLoadCustomTools, type LoadedCustomTool } from "./core/custom-tools/index.js"; import { exportFromFile } from "./core/export-html.js"; @@ -27,6 +29,32 @@ import { initTheme, stopThemeWatcher } from "./modes/interactive/theme/theme.js" import { getChangelogPath, getNewEntries, parseChangelog } from "./utils/changelog.js"; import { ensureTool } from "./utils/tools-manager.js"; +/** Configure OAuth storage to use the coding-agent's configurable path */ +function configureOAuthStorage(): void { + const oauthPath = getOAuthPath(); + + setOAuthStorage({ + load: () => { + if (!existsSync(oauthPath)) { + return {}; + } + try { + return JSON.parse(readFileSync(oauthPath, "utf-8")); + } catch { + return {}; + } + }, + save: (storage) => { + const dir = dirname(oauthPath); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true, mode: 0o700 }); + } + writeFileSync(oauthPath, JSON.stringify(storage, null, 2), "utf-8"); + chmodSync(oauthPath, 0o600); + }, + }); +} + /** Check npm registry for new version (non-blocking) */ async function checkForNewVersion(currentVersion: string): Promise { try { @@ -142,6 +170,10 @@ async function prepareInitialMessage(parsed: Args): Promise<{ } export async function main(args: string[]) { + // Configure OAuth storage to use the coding-agent's configurable path + // This must happen before any OAuth operations + configureOAuthStorage(); + const parsed = parseArgs(args); if (parsed.version) { diff --git a/packages/coding-agent/src/modes/interactive/components/model-selector.ts b/packages/coding-agent/src/modes/interactive/components/model-selector.ts index 4c85bd8a..674d5a81 100644 --- a/packages/coding-agent/src/modes/interactive/components/model-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/model-selector.ts @@ -22,6 +22,11 @@ interface ModelItem { model: Model; } +interface ScopedModelItem { + model: Model; + thinkingLevel: string; +} + /** * Component that renders a model selector with search */ @@ -37,11 +42,13 @@ export class ModelSelectorComponent extends Container { private onCancelCallback: () => void; private errorMessage: string | null = null; private tui: TUI; + private scopedModels: ReadonlyArray; constructor( tui: TUI, currentModel: Model | null, settingsManager: SettingsManager, + scopedModels: ReadonlyArray, onSelect: (model: Model) => void, onCancel: () => void, ) { @@ -50,6 +57,7 @@ export class ModelSelectorComponent extends Container { this.tui = tui; this.currentModel = currentModel; this.settingsManager = settingsManager; + this.scopedModels = scopedModels; this.onSelectCallback = onSelect; this.onCancelCallback = onCancel; @@ -57,10 +65,12 @@ export class ModelSelectorComponent extends Container { this.addChild(new DynamicBorder()); this.addChild(new Spacer(1)); - // Add hint about API key filtering - this.addChild( - new Text(theme.fg("warning", "Only showing models with configured API keys (see README for details)"), 0, 0), - ); + // Add hint about model filtering + const hintText = + scopedModels.length > 0 + ? "Showing models from --models scope" + : "Only showing models with configured API keys (see README for details)"; + this.addChild(new Text(theme.fg("warning", hintText), 0, 0)); this.addChild(new Spacer(1)); // Create search input @@ -93,24 +103,35 @@ export class ModelSelectorComponent extends Container { } private async loadModels(): Promise { - // Load available models fresh (includes custom models from models.json) - const { models: availableModels, error } = await getAvailableModels(); + let models: ModelItem[]; - // If there's an error loading models.json, we'll show it via the "no models" path - // The error will be displayed to the user - if (error) { - this.allModels = []; - this.filteredModels = []; - this.errorMessage = error; - return; + // Use scoped models if provided via --models flag + if (this.scopedModels.length > 0) { + models = this.scopedModels.map((scoped) => ({ + provider: scoped.model.provider, + id: scoped.model.id, + model: scoped.model, + })); + } else { + // Load available models fresh (includes custom models from models.json) + const { models: availableModels, error } = await getAvailableModels(); + + // If there's an error loading models.json, we'll show it via the "no models" path + // The error will be displayed to the user + if (error) { + this.allModels = []; + this.filteredModels = []; + this.errorMessage = error; + return; + } + + models = availableModels.map((model) => ({ + provider: model.provider, + id: model.id, + model, + })); } - const models: ModelItem[] = availableModels.map((model) => ({ - provider: model.provider, - id: model.id, - model, - })); - // Sort: current model first, then by provider models.sort((a, b) => { const aIsCurrent = this.currentModel?.id === a.model.id && this.currentModel?.provider === a.provider; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index a7833795..9675adc8 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -1382,6 +1382,7 @@ export class InteractiveMode { this.ui, this.session.model, this.settingsManager, + this.session.scopedModels, async (model) => { try { await this.session.setModel(model);