diff --git a/CONFIGURATION.md b/CONFIGURATION.md index 4fdc7b9..d7d2587 100644 --- a/CONFIGURATION.md +++ b/CONFIGURATION.md @@ -15,3 +15,7 @@ Most users do not need these overrides. Only point base URL overrides at endpoin | `PI_GROK_CLI_IMAGINE_MODEL` | `grok-imagine-image-quality` | Override the Imagine image model. | OAuth logins store the effective main API base URL in each account's credentials. Changing or removing `PI_GROK_CLI_BASE_URL` or `GROK_CLI_BASE_URL` does not immediately redirect an already authenticated account; its stored URL remains active until token refresh or re-login. Log in again to apply the change immediately. + +Real account credentials are in `~/.pi/grok-cli/accounts.json`. The file uses mode `0600`, and updates use an interprocess lock and an atomic replacement. Pi's own credential store contains only the fixed non-secret marker that connects the `grok-cli` provider. + +When `GROK_CLI_OAUTH_TOKEN` is set, it overrides the saved account for requests. Account login is not available until you remove the variable and restart pi. If you use native `/logout` while this variable is set, restart pi so the extension can register the correct login method. diff --git a/README.md b/README.md index cd21290..508787f 100644 --- a/README.md +++ b/README.md @@ -16,8 +16,7 @@ Use your X Premium or SuperGrok subscription in [pi](https://pi.dev/) with a cle - **Multiple accounts:** Manage accounts and quotas from the terminal or browser dashboard, and automatically continue with another account when quota runs out. - **Usage tracking:** Check account limits, remaining credits, and reset times. - **Image generation:** Generate images directly or let Grok use the `image_gen` tool. -- **Vision routing:** Let text-only models understand images through a vision-capable Grok model. -- **Model-specific tools:** Use Cursor-compatible tool names only where Grok Build and Composer 2.5 need them. +- **Native image input:** Send images directly to supported Grok models, including Composer 2.5. > Requires pi 0.80.9 or newer and an xAI/Grok account with access to the selected model. Availability varies by account, plan, region, and xAI rollout. The Grok Build executable is not required. > @@ -53,7 +52,7 @@ Choose **Grok CLI**, then select one of these methods: /model grok-cli/grok-composer-2.5-fast ``` -Composer 2.5 is the recommended starting point for agentic coding. Choose an image-capable model when native image input matters. +Composer 2.5 is the recommended starting point for agentic coding and supports native image input. ### 4. Verify account usage @@ -73,12 +72,18 @@ Open the browser dashboard to view accounts and quotas in one place: Select **Add account** to optionally label the account and start browser login. -Prefer the terminal? Run `/grok-cli-accounts`, choose **＋ Add account**, and press Enter when the login command appears. Device-code login remains available there for headless environments. +Prefer the terminal? Run `/grok-cli-accounts`, choose **＋ Add account**, and complete the browser or device-code login. With at least two logged-in accounts, pi-grok-cli can continue an interrupted request with another account when the current account runs out of quota. Recently exhausted accounts are temporarily skipped, and rotation stops when no eligible account remains. Only add accounts you own or are authorized to access. +Pi shows one **Grok CLI** entry in `/login` and one `grok-cli` provider in `/model`. It does not add provider names such as `grok-cli-2`. + +Account selection is per session. The dashboard or `/grok-cli-accounts` sets the account for the current session, and the choice is stored in the session log so it is restored when the session is resumed or reopened. New sessions start with the most recently selected account and keep their own selection once changed. Requests always use the token captured at request start. + +Pi and pi-grok-cli keep separate login state. Native `/logout` removes Pi's non-secret connection marker. It does not delete accounts from the dashboard. Run `/login` and choose **Grok CLI** to reconnect the saved accounts. To delete a saved account, use `/grok-cli-accounts`. Account 1 is permanent, but you can log it out. + ## Manage installation ```bash @@ -90,47 +95,21 @@ pi remove npm:pi-grok-cli Models are bundled rather than discovered live. Registered context limits may differ from the limits xAI enforces. -| Model ID | Registered context | Reasoning | Input | Coding tools | -| --- | ---: | --- | --- | --- | -| `grok-composer-2.5-fast` | 200K | no | text; vision routing available | Cursor-compatible names | -| `grok-build` | 500K | yes | text + image | Cursor-compatible names | -| `grok-4.3` | 1M | yes | text + image | native pi names | -| `grok-4.5` | 500K | yes | text + image | native pi names | -| `grok-4.20-0309-reasoning` | 2M | yes | text + image | native pi names | -| `grok-4.20-0309-non-reasoning` | 2M | no | text + image | native pi names | -| `grok-4.20-multi-agent-0309` | 2M | yes | text + image | native pi names | - -## Coding tools and media - -### Coding tool compatibility - -Grok Build and Composer 2.5 use Cursor-compatible tool names. All other models and providers use pi's native tools. - -| Compatibility tools | Implementation | -| --- | --- | -| `Read`, `Write`, `Edit`, `Grep`, `LS`, `Shell` | Delegate execution and rendering to pi's native tools. | -| `StrReplace` | Preserves Cursor's literal replace-all behavior. | -| `Glob` | Preserves Cursor's newest-first path ordering. | -| `Delete` | Provides the file-deletion operation expected by these models. | -| `WebSearch` | Optionally replaces an active `web_search` tool when pi-web-access is installed. | - -pi applies `--tools` and `--exclude-tools` by exact name. For compatible models, use names such as `--tools Read,Grep` or `--exclude-tools Grep`. `--no-tools` disables all tools. +| Model ID | Registered context | Reasoning | Input | +| --- | ---: | --- | --- | +| `grok-composer-2.5-fast` | 200K | no | text + image | +| `grok-build` | 500K | yes | text + image | +| `grok-4.3` | 1M | yes | text + image | +| `grok-4.5` | 500K | yes | text + image | +| `grok-4.20-0309-reasoning` | 2M | yes | text + image | +| `grok-4.20-0309-non-reasoning` | 2M | no | text + image | +| `grok-4.20-multi-agent-0309` | 2M | yes | text + image | -Install pi-web-access 0.13.0 or newer to add optional `WebSearch` support: - -```bash -pi install npm:pi-web-access@^0.13.0 -``` - -Restart pi or run `/reload` after installing it in an existing session. Other models and providers retain the native `web_search` tool. - -### Vision routing for text-only models - -When a text-only model reads an image, pi-grok-cli describes it with an image-capable Grok model (`grok-build` by default) and passes back the description. Routing is enabled by default for up to four images, caches descriptions locally, and is bypassed for models that support images natively. +## Media ### Grok Imagine image generation -Run `/grok-cli-imagine ` to generate and preview a JPEG, or let any active model call the `image_gen` tool. Images use the current or most recently selected Grok account and are saved under the current session unless you request another path. +Run `/grok-cli-imagine ` to generate and preview a JPEG, or let any active model call the `image_gen` tool. Images use the current session's selected Grok account and are saved under the current session unless you request another path. `image_gen` is enabled by default across providers. Use `/grok-cli-imagine:tool [on|off|status]` to manage model access without disabling the direct command. @@ -142,9 +121,6 @@ Run `/grok-cli-imagine ` to generate and preview a JPEG, or let any acti | `/grok-cli-usage` | Fetch current quota, update its cache, and show cached data if refresh fails. | | `/grok-cli-imagine ` | Generate and preview an image. Supports `--aspect`, `--out`, and `--resolution 1k`. | | `/grok-cli-imagine:tool [on\|off\|status]` | Toggle, set, or report persistent model-callable `image_gen` availability. | -| `/grok-cli-vision` | Toggle vision routing on or off. | -| `/grok-cli-vision:status` | Show vision state, describer model, configuration path, and cache statistics. | -| `/grok-cli-vision:cache-clear` | Remove all cached image descriptions. | ## Configuration @@ -152,12 +128,17 @@ Extension-owned data is grouped under `~/.pi/grok-cli/`: ```text ~/.pi/grok-cli/ +├── accounts.json ├── config.json -├── vision-cache.json +├── config.v2.backup.json └── quota-cache.json ``` -`config.json` stores settings and account labels. The cache files store derived vision descriptions and quota responses. None contain OAuth tokens; pi stores credentials separately. Legacy settings are migrated automatically. +`accounts.json` stores account labels and OAuth credentials. It uses file mode `0600`. `config.json` stores non-account settings. `quota-cache.json` stores quota responses. `config.v2.backup.json` is an exact backup that is created only when the extension migrates an older multi-account configuration. Pi stores only a fixed, non-secret marker for the `grok-cli` provider. + +When you update from a release that used `grok-cli-N` providers, the extension copies known OAuth credentials to `accounts.json` and rewrites global and project model settings to `grok-cli`. Run `/login` once and choose **Grok CLI** to install Pi's marker. The old alias credentials can remain in Pi's `/logout` list. Remove those old alias entries after you confirm that the new accounts work. Do not remove the main `Grok CLI` entry unless you want to disconnect it. + +The extension does not replace saved session files because Pi can append to them from another process without a shared file lock. This prevents session data loss. When you resume an old session that selected `grok-cli-N`, Pi can show a model fallback warning. Select the required `grok-cli/` once in that session. ### Common environment variables @@ -179,11 +160,12 @@ See [Advanced configuration](./CONFIGURATION.md) for OAuth, callback, endpoint, | Authentication returns HTTP 401 or 403 | Run `/login` again and confirm the account can access the selected model. Replace an expired `GROK_CLI_OAUTH_TOKEN` if using an external token. | | A listed model is unavailable | Availability can differ by account or region, and the catalog is bundled rather than discovered live. Try another model or update the extension. | | Account rotation does not start | Confirm at least two accounts are logged in. Rotation responds only to Grok Build's final balance-exhausted error, not authentication failures, rate limits, or similar messages. | -| Images are not being described | Run `/grok-cli-vision:status`, confirm routing is on, and verify Grok authentication. Native image-capable models bypass routing by design. | +| A migrated account cannot send requests | Run `/login` once, choose **Grok CLI**, and restart pi if the model list does not update. | +| The dashboard says Pi is disconnected | Run `/login` and choose **Grok CLI**. Saved dashboard accounts remain available. | ## Security and data flow -Enabled tools run with your system permissions and can modify files or execute commands. Prompts, model context, tool data, and routed images are sent to the configured xAI endpoints. +Enabled tools run with your system permissions and can modify files or execute commands. Prompts, model context, tool data, and native image inputs are sent to the configured xAI endpoints. See [SECURITY.md](./SECURITY.md) for trust boundaries and private vulnerability reporting. Never include tokens, authorization codes, callback URLs, prompts, or private project data in public issues. diff --git a/SECURITY.md b/SECURITY.md index 5d9a7c7..e63aa00 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -18,15 +18,13 @@ The maintainer aims to acknowledge a complete report within 7 calendar days and ## Security boundaries -- pi-grok-cli and its tools run with the user's operating-system permissions. File and shell tools can access any path those permissions allow, including paths outside the workspace, and can execute commands when the model invokes them. +- pi-grok-cli runs with the user's operating-system permissions. Image generation can write to session storage, temporary storage, or a user-requested output path. - OAuth credentials returned by this extension are stored and refreshed by pi. pi-grok-cli does not read or modify Grok Build credentials. - Browser OAuth starts a temporary callback server on `127.0.0.1:56122` by default and falls back to an ephemeral port. `PI_GROK_CLI_CALLBACK_HOST` can bind it to another interface. The server validates the callback path and OAuth state and closes after login. Treat complete callback URLs and authorization codes as sensitive. - `/grok-cli-accounts gui` starts a temporary account-management server bound only to an OS-assigned `127.0.0.1` port. A random capability URL bootstraps a session cookie; subsequent mutations require same-origin and CSRF validation. The page receives account labels, status, and quota data. It never receives stored OAuth credentials, callback URLs, or environment-token values; a manually entered one-time authorization code is handled transiently during login and is never returned by `/api/state`. Treat the private dashboard URL as sensitive and do not share it while the server is running. - Prompts, conversation context, tool definitions, tool results, and native image inputs are sent to the configured Grok CLI proxy. -- Images handled by vision routing are sent to the configured describer model. The local vision cache stores descriptions and hashes, not raw images. - Grok Imagine sends the selected account's bearer token, generation prompt, and options to `https://api.x.ai/v1` or `PI_GROK_CLI_IMAGINE_BASE_URL`. Generated JPEGs and PNG previews are saved under session storage, a requested output path, or temporary storage. - Monthly and weekly quota totals and reset timestamps are cached per account in `~/.pi/grok-cli/quota-cache.json` with file mode `0600`. The cache does not contain OAuth tokens. -- Optional web search is delegated to `pi-web-access` and the providers configured there. - A configured main API base URL override is trusted with bearer tokens, prompts, conversation data, tool results, images, and billing queries. An Imagine base URL override is trusted with the bearer token and generation request described above. Expected provider behavior, account-specific model availability, subscription quotas, model-generated content, a model choosing an enabled local tool, and changes made after a user deliberately configures an untrusted endpoint are generally out of scope unless pi-grok-cli violates a documented boundary or exposes data beyond that configuration. diff --git a/bun.lock b/bun.lock index 97d1550..182eac7 100644 --- a/bun.lock +++ b/bun.lock @@ -4,10 +4,6 @@ "workspaces": { "": { "name": "pi-grok-cli", - "dependencies": { - "ignore": "7.0.5", - "minimatch": "10.2.5", - }, "devDependencies": { "@biomejs/biome": "2.4.16", "@earendil-works/pi-agent-core": "^0.80.9", @@ -28,11 +24,7 @@ "@earendil-works/pi-ai": ">=0.80.9", "@earendil-works/pi-coding-agent": ">=0.80.9", "@earendil-works/pi-tui": ">=0.80.9", - "pi-web-access": ">=0.13.0", }, - "optionalPeers": [ - "pi-web-access", - ], }, }, "packages": { diff --git a/knip.json b/knip.json index 3e3ab5c..5c61008 100644 --- a/knip.json +++ b/knip.json @@ -1,5 +1,5 @@ { "entry": ["src/shared/backup_worker.ts"], "project": ["**/*.ts", "!**/*.d.ts", "!.release-tools/**"], - "ignoreDependencies": ["ignore", "lint-staged", "minimatch"] + "ignoreDependencies": ["lint-staged"] } diff --git a/package.json b/package.json index c8f41ca..90f53c2 100644 --- a/package.json +++ b/package.json @@ -10,7 +10,6 @@ "grok-cli", "oauth", "xai-oauth", - "vision", "multimodal" ], "type": "module", @@ -56,13 +55,7 @@ "peerDependencies": { "@earendil-works/pi-ai": ">=0.80.9", "@earendil-works/pi-coding-agent": ">=0.80.9", - "@earendil-works/pi-tui": ">=0.80.9", - "pi-web-access": ">=0.13.0" - }, - "peerDependenciesMeta": { - "pi-web-access": { - "optional": true - } + "@earendil-works/pi-tui": ">=0.80.9" }, "devDependencies": { "@biomejs/biome": "2.4.16", @@ -84,9 +77,5 @@ "*": [ "biome check --write --no-errors-on-unmatched" ] - }, - "dependencies": { - "ignore": "7.0.5", - "minimatch": "10.2.5" } } diff --git a/src/config.ts b/src/config.ts index dfcc01c..afe040f 100644 --- a/src/config.ts +++ b/src/config.ts @@ -1,73 +1,25 @@ import { existsSync, readFileSync, unlinkSync } from 'node:fs'; -import { resolveModels } from './models/catalog.js'; import { getConfigPath, getLegacyConfigPath, getLegacyImagineConfigPath, - getLegacyVisionCachePath, - getLegacyVisionConfigPath, - getVisionCachePath, migrateStoredFile, writeFileAtomic, } from './storage.js'; -export { getConfigPath, getLegacyImagineConfigPath, getLegacyVisionConfigPath } from './storage.js'; - -export const CONFIG_VERSION = 2 as const; -const LEGACY_CONFIG_VERSION = 1; -export const DEFAULT_DESCRIBE_MODEL = 'grok-build'; -export const DEFAULT_MAX_IMAGES = 4; -export const DEFAULT_CACHE_MAX_ENTRIES = 100; +export { getConfigPath, getLegacyImagineConfigPath } from './storage.js'; +export const CONFIG_VERSION = 3 as const; export type ImagineConfig = { enabled: boolean }; - -export interface VisionConfig { - enabled: boolean; - model: string; - maxImages: number; - cacheEnabled: boolean; - cacheMaxEntries: number; -} - -export interface GrokCliAccount { - provider: string; - label: string; -} - -export interface AccountsConfig { - nextAccountNumber: number; - selectedProvider: string; - items: GrokCliAccount[]; -} - export interface GrokCliConfig { version: typeof CONFIG_VERSION; - accounts: AccountsConfig; imagine: ImagineConfig; - vision: VisionConfig; } export const DEFAULT_IMAGINE_CONFIG: ImagineConfig = { enabled: true }; - -export const DEFAULT_VISION_CONFIG: VisionConfig = { - enabled: true, - model: DEFAULT_DESCRIBE_MODEL, - maxImages: DEFAULT_MAX_IMAGES, - cacheEnabled: true, - cacheMaxEntries: DEFAULT_CACHE_MAX_ENTRIES, -}; - -export const DEFAULT_ACCOUNTS_CONFIG: AccountsConfig = { - nextAccountNumber: 2, - selectedProvider: 'grok-cli', - items: [{ provider: 'grok-cli', label: 'Account 1' }], -}; - export const DEFAULT_CONFIG: GrokCliConfig = { version: CONFIG_VERSION, - accounts: DEFAULT_ACCOUNTS_CONFIG, imagine: DEFAULT_IMAGINE_CONFIG, - vision: DEFAULT_VISION_CONFIG, }; export interface LoadedConfig { @@ -75,27 +27,8 @@ export interface LoadedConfig { warning?: string; } -type ParsedConfig = LoadedConfig & { valid: boolean; needsMigration: boolean }; - -type LegacyConfig = LoadedConfig & { - existingPaths: string[]; - recognizedPaths: string[]; -}; - -function defaultConfig(): GrokCliConfig { - return { - version: CONFIG_VERSION, - accounts: { - ...DEFAULT_ACCOUNTS_CONFIG, - items: DEFAULT_ACCOUNTS_CONFIG.items.map((account) => ({ ...account })), - }, - imagine: { ...DEFAULT_IMAGINE_CONFIG }, - vision: { ...DEFAULT_VISION_CONFIG }, - }; -} - function isObject(value: unknown): value is Record { - return Boolean(value) && typeof value === 'object' && !Array.isArray(value); + return typeof value === 'object' && value !== null && !Array.isArray(value); } function errorMessage(error: unknown) { @@ -114,95 +47,6 @@ export function hasTerminalControlCharacters(value: string) { }); } -const accountNumber = (provider: string) => { - if (provider === 'grok-cli') return 1; - const match = /^grok-cli-((?:[2-9]|[1-9]\d+))$/.exec(provider); - return match ? Number(match[1]) : undefined; -}; - -export function findAvailableAccountNumber( - providers: Iterable, - reservedProviders: Iterable = [], -) { - const unavailable = new Set([...providers, ...reservedProviders]); - const find = (number: number): number => - unavailable.has(`grok-cli-${number}`) ? find(number + 1) : number; - return find(2); -} - -function normalizeAccountsConfig(raw: unknown, warnings: string[]): AccountsConfig { - if (raw === undefined) return defaultConfig().accounts; - if (!isObject(raw) || !Array.isArray(raw.items)) { - warnings.push('accounts must be an object with an items array. Using defaults.'); - return defaultConfig().accounts; - } - - const invalid: unknown[] = []; - const providers = new Set(); - const labels = new Set(); - const baseIndex = raw.items.findIndex((value) => { - if (!isObject(value) || value.provider !== 'grok-cli' || typeof value.label !== 'string') { - return false; - } - const label = value.label.trim(); - return Boolean(label) && [...label].length <= 40 && !hasTerminalControlCharacters(label); - }); - const accountValues = - baseIndex >= 0 - ? [raw.items[baseIndex], ...raw.items.filter((_value, index) => index !== baseIndex)] - : [{ provider: 'grok-cli', label: 'Account 1' }, ...raw.items]; - const items = accountValues.flatMap((value) => { - if (!isObject(value) || typeof value.provider !== 'string' || typeof value.label !== 'string') { - invalid.push(value); - return []; - } - const label = value.label.trim(); - const normalizedLabel = label.toLocaleLowerCase(); - if ( - accountNumber(value.provider) === undefined || - !label || - [...label].length > 40 || - hasTerminalControlCharacters(label) || - providers.has(value.provider) || - labels.has(normalizedLabel) - ) { - invalid.push(value); - return []; - } - providers.add(value.provider); - labels.add(normalizedLabel); - return [{ provider: value.provider, label }]; - }); - - if (!providers.has('grok-cli')) { - items.unshift({ provider: 'grok-cli', label: 'Account 1' }); - providers.add('grok-cli'); - } else { - items.sort((left, right) => - left.provider === 'grok-cli' ? -1 : right.provider === 'grok-cli' ? 1 : 0, - ); - } - - if (invalid.length) - warnings.push('accounts contains invalid or duplicate entries. Ignoring them.'); - const selectedProvider = - typeof raw.selectedProvider === 'string' && providers.has(raw.selectedProvider) - ? raw.selectedProvider - : 'grok-cli'; - - return { - nextAccountNumber: findAvailableAccountNumber(providers), - selectedProvider, - items, - }; -} - -export function describableModels(): string[] { - return resolveModels() - .filter((model) => model.input.includes('image')) - .map((model) => model.id); -} - function normalizeImagineConfig(raw: unknown, warnings: string[]): ImagineConfig { if (raw === undefined) return { ...DEFAULT_IMAGINE_CONFIG }; if (!isObject(raw)) { @@ -216,143 +60,58 @@ function normalizeImagineConfig(raw: unknown, warnings: string[]): ImagineConfig return { ...DEFAULT_IMAGINE_CONFIG }; } -export function normalizeVisionConfig( - raw: Partial, - warnings: string[] = [], -): VisionConfig { - const config: VisionConfig = { ...DEFAULT_VISION_CONFIG }; - - if ('enabled' in raw) { - if (typeof raw.enabled === 'boolean') { - config.enabled = raw.enabled; - } else if (raw.enabled !== undefined) { - warnings.push('enabled must be true or false. Using enabled=true.'); - } - } - - if ('model' in raw) { - if (typeof raw.model === 'string' && describableModels().includes(raw.model)) { - config.model = raw.model; - } else if (raw.model !== undefined) { - warnings.push( - `Unknown model "${String(raw.model)}". Available: ${describableModels().join(', ')}. Using ${DEFAULT_DESCRIBE_MODEL}.`, - ); - } - } - - if ('maxImages' in raw) { - if ( - typeof raw.maxImages === 'number' && - Number.isFinite(raw.maxImages) && - raw.maxImages > 0 && - Number.isInteger(raw.maxImages) - ) { - config.maxImages = raw.maxImages; - } else if (raw.maxImages !== undefined) { - warnings.push(`maxImages must be a positive integer. Using ${DEFAULT_MAX_IMAGES}.`); - } - } - - if ('cacheEnabled' in raw) { - if (typeof raw.cacheEnabled === 'boolean') { - config.cacheEnabled = raw.cacheEnabled; - } else if (raw.cacheEnabled !== undefined) { - warnings.push('cacheEnabled must be true or false. Using cacheEnabled=true.'); - } - } - - if ('cacheMaxEntries' in raw) { - if ( - typeof raw.cacheMaxEntries === 'number' && - Number.isInteger(raw.cacheMaxEntries) && - raw.cacheMaxEntries > 0 - ) { - config.cacheMaxEntries = raw.cacheMaxEntries; - } else if (raw.cacheMaxEntries !== undefined) { - warnings.push( - `cacheMaxEntries must be a positive integer. Using ${DEFAULT_CACHE_MAX_ENTRIES}.`, - ); - } - } - - return config; -} - -function normalizeConfig( - raw: { accounts?: unknown; imagine?: unknown; vision?: unknown }, - warnings: string[], -): GrokCliConfig { - const vision = raw.vision; - if (vision !== undefined && !isObject(vision)) { - warnings.push('vision must be a JSON object. Using defaults.'); - } - return { - version: CONFIG_VERSION, - accounts: normalizeAccountsConfig(raw.accounts, warnings), - imagine: normalizeImagineConfig(raw.imagine, warnings), - vision: normalizeVisionConfig(isObject(vision) ? vision : {}, warnings), - }; -} - -function parseConfig(configPath: string): ParsedConfig { +function parseConfig(path: string): LoadedConfig & { supported: boolean; imagineDefined: boolean } { try { - const parsed: unknown = JSON.parse(readFileSync(configPath, 'utf8')); + const parsed: unknown = JSON.parse(readFileSync(path, 'utf8')); if (!isObject(parsed)) { return { - config: defaultConfig(), - needsMigration: false, - valid: false, - warning: `Config ${configPath} must be a JSON object. Using legacy settings or defaults.`, + config: structuredClone(DEFAULT_CONFIG), + supported: false, + imagineDefined: false, + warning: `Config ${path} must be a JSON object. Using legacy settings or defaults.`, }; } - if (parsed.version !== CONFIG_VERSION && parsed.version !== LEGACY_CONFIG_VERSION) { + if (parsed.version !== 1 && parsed.version !== 2 && parsed.version !== CONFIG_VERSION) { return { - config: defaultConfig(), - needsMigration: false, - valid: false, - warning: `Unsupported config version ${String(parsed.version)} in ${configPath}. Using legacy settings or defaults.`, + config: structuredClone(DEFAULT_CONFIG), + supported: false, + imagineDefined: false, + warning: `Unsupported config version ${String(parsed.version)} in ${path}. Using legacy settings or defaults.`, }; } const warnings: string[] = []; + const config = { + version: CONFIG_VERSION, + imagine: normalizeImagineConfig(parsed.imagine, warnings), + } satisfies GrokCliConfig; return { - config: normalizeConfig(parsed, warnings), - needsMigration: parsed.version === LEGACY_CONFIG_VERSION, - valid: true, - warning: warnings.length ? `Invalid ${configPath}: ${warnings.join(' ')}` : undefined, + config, + supported: true, + imagineDefined: Object.hasOwn(parsed, 'imagine'), + ...(warnings.length ? { warning: `Invalid ${path}: ${warnings.join(' ')}` } : {}), }; } catch (error) { return { - config: defaultConfig(), - needsMigration: false, - valid: false, - warning: `Could not read ${configPath}: ${errorMessage(error)}. Using legacy settings or defaults.`, + config: structuredClone(DEFAULT_CONFIG), + supported: false, + imagineDefined: false, + warning: `Could not read ${path}: ${errorMessage(error)}. Using legacy settings or defaults.`, }; } } -function parseLegacyImagine(configPath: string): { - config: ImagineConfig; - recognized: boolean; - warning?: string; -} { +function parseLegacyImagine(path: string) { try { - const parsed: unknown = JSON.parse(readFileSync(configPath, 'utf8')); + const parsed: unknown = JSON.parse(readFileSync(path, 'utf8')); if (!isObject(parsed)) { return { config: { ...DEFAULT_IMAGINE_CONFIG }, recognized: false, - warning: `Legacy config ${configPath} must be a JSON object.`, + warning: `Legacy config ${path} must be a JSON object.`, }; } - if ('enabled' in parsed) { - if (typeof parsed.enabled === 'boolean') { - return { config: { enabled: parsed.enabled }, recognized: true }; - } - return { - config: { ...DEFAULT_IMAGINE_CONFIG }, - recognized: false, - warning: `Invalid ${configPath}: enabled must be a boolean.`, - }; + if (typeof parsed.enabled === 'boolean') { + return { config: { enabled: parsed.enabled }, recognized: true }; } if (parsed.scope === 'grok-cli' || parsed.scope === 'all') { return { config: { enabled: true }, recognized: true }; @@ -360,191 +119,108 @@ function parseLegacyImagine(configPath: string): { return { config: { ...DEFAULT_IMAGINE_CONFIG }, recognized: false, - warning: `Invalid ${configPath}: expected enabled or a recognized scope.`, + warning: `Invalid ${path}: expected enabled or a recognized scope.`, }; } catch (error) { return { config: { ...DEFAULT_IMAGINE_CONFIG }, recognized: false, - warning: `Could not read ${configPath}: ${errorMessage(error)}.`, - }; - } -} - -function parseLegacyVision(configPath: string): { - config: VisionConfig; - recognized: boolean; - warning?: string; -} { - try { - const parsed: unknown = JSON.parse(readFileSync(configPath, 'utf8')); - if (!isObject(parsed)) { - return { - config: { ...DEFAULT_VISION_CONFIG }, - recognized: false, - warning: `Legacy config ${configPath} must be a JSON object.`, - }; - } - const warnings: string[] = []; - const config = normalizeVisionConfig(parsed, warnings); - return { - config, - recognized: warnings.length === 0, - warning: warnings.length ? `Invalid ${configPath}: ${warnings.join(' ')}` : undefined, - }; - } catch (error) { - return { - config: { ...DEFAULT_VISION_CONFIG }, - recognized: false, - warning: `Could not read ${configPath}: ${errorMessage(error)}.`, + warning: `Could not read ${path}: ${errorMessage(error)}.`, }; } } -function loadLegacyConfig(): LegacyConfig { - const imaginePath = getLegacyImagineConfigPath(); - const visionPath = getLegacyVisionConfigPath(); - const imagine = existsSync(imaginePath) ? parseLegacyImagine(imaginePath) : undefined; - const vision = existsSync(visionPath) ? parseLegacyVision(visionPath) : undefined; - return { - config: { - version: CONFIG_VERSION, - accounts: defaultConfig().accounts, - imagine: imagine?.config ?? { ...DEFAULT_IMAGINE_CONFIG }, - vision: vision?.config ?? { ...DEFAULT_VISION_CONFIG }, - }, - existingPaths: [imagine ? imaginePath : undefined, vision ? visionPath : undefined].filter( - (path): path is string => Boolean(path), - ), - recognizedPaths: [ - imagine?.recognized ? imaginePath : undefined, - vision?.recognized ? visionPath : undefined, - ].filter((path): path is string => Boolean(path)), - warning: combineWarnings([imagine?.warning, vision?.warning]), - }; -} - export function loadConfig(): LoadedConfig { - const configPath = existsSync(getConfigPath()) + const path = existsSync(getConfigPath()) ? getConfigPath() : existsSync(getLegacyConfigPath()) ? getLegacyConfigPath() : undefined; - if (!configPath) { - const legacy = loadLegacyConfig(); - return legacy.warning - ? { config: legacy.config, warning: legacy.warning } - : { config: legacy.config }; + if (path) { + const loaded = parseConfig(path); + if (loaded.supported) { + return loaded.warning + ? { config: loaded.config, warning: loaded.warning } + : { config: loaded.config }; + } + if (existsSync(getLegacyImagineConfigPath())) { + const legacy = parseLegacyImagine(getLegacyImagineConfigPath()); + return { + config: { version: CONFIG_VERSION, imagine: legacy.config }, + warning: combineWarnings([loaded.warning, legacy.warning]), + }; + } + return { config: loaded.config, ...(loaded.warning ? { warning: loaded.warning } : {}) }; } - const loaded = parseConfig(configPath); - if (loaded.valid) { - return loaded.warning - ? { config: loaded.config, warning: loaded.warning } - : { config: loaded.config }; + if (!existsSync(getLegacyImagineConfigPath())) { + return { config: structuredClone(DEFAULT_CONFIG) }; } - const legacy = loadLegacyConfig(); + const legacy = parseLegacyImagine(getLegacyImagineConfigPath()); return { - config: legacy.config, - warning: combineWarnings([loaded.warning, legacy.warning]), + config: { version: CONFIG_VERSION, imagine: legacy.config }, + ...(legacy.warning ? { warning: legacy.warning } : {}), }; } export function saveConfig(config: GrokCliConfig) { - writeFileAtomic(getConfigPath(), `${JSON.stringify(normalizeConfig(config, []), null, 2)}\n`); -} - -function removeLegacyConfigs(paths: string[]) { - return combineWarnings( - paths.flatMap((path) => { - try { - unlinkSync(path); - return []; - } catch (error) { - return [`Could not remove legacy config ${path}: ${errorMessage(error)}.`]; - } - }), + writeFileAtomic( + getConfigPath(), + `${JSON.stringify( + { version: CONFIG_VERSION, imagine: normalizeImagineConfig(config.imagine, []) }, + null, + 2, + )}\n`, ); } export function migrateLegacyConfig(): { warning?: string } { - const migratedLegacyConfig = existsSync(getLegacyConfigPath()) && !existsSync(getConfigPath()); - const storageWarning = combineWarnings([ - migrateStoredFile(getLegacyConfigPath(), getConfigPath(), true), - migrateStoredFile(getLegacyVisionCachePath(), getVisionCachePath()), - ]); + const storageWarning = migrateStoredFile(getLegacyConfigPath(), getConfigPath()); if (!existsSync(getConfigPath()) && existsSync(getLegacyConfigPath())) { - return { warning: storageWarning }; + return storageWarning ? { warning: storageWarning } : {}; } - const legacy = loadLegacyConfig(); + const legacyPath = getLegacyImagineConfigPath(); + const legacy = existsSync(legacyPath) ? parseLegacyImagine(legacyPath) : undefined; if (existsSync(getConfigPath())) { const loaded = parseConfig(getConfigPath()); - if (!loaded.valid) { - return { warning: combineWarnings([storageWarning, loaded.warning, legacy.warning]) }; + if (!loaded.supported) { + const warning = combineWarnings([storageWarning, loaded.warning, legacy?.warning]); + return warning ? { warning } : {}; } - if (loaded.needsMigration) { + if (legacy?.recognized && loaded.imagineDefined) { try { - saveConfig(loaded.config); - const verified = parseConfig(getConfigPath()); - if ( - !verified.valid || - verified.needsMigration || - JSON.stringify(verified.config) !== JSON.stringify(loaded.config) - ) { - return { - warning: combineWarnings([ - storageWarning, - verified.warning, - `Could not verify migrated config ${getConfigPath()}. Legacy files were preserved.`, - ]), - }; - } + unlinkSync(legacyPath); } catch (error) { return { warning: combineWarnings([ storageWarning, - `Could not migrate configuration ${getConfigPath()}: ${errorMessage(error)}. Legacy files were preserved.`, + legacy.warning, + `Could not remove legacy config ${legacyPath}: ${errorMessage(error)}.`, ]), }; } } - const cleanupWarning = removeLegacyConfigs([ - ...legacy.recognizedPaths, - ...(migratedLegacyConfig ? [getLegacyConfigPath()] : []), - ]); - const warning = combineWarnings([ - storageWarning, - loaded.warning, - legacy.warning, - cleanupWarning, - ]); + const warning = combineWarnings([storageWarning, loaded.warning, legacy?.warning]); return warning ? { warning } : {}; } - if (legacy.existingPaths.length === 0) { - return storageWarning ? { warning: storageWarning } : {}; - } - if (legacy.recognizedPaths.length !== legacy.existingPaths.length) { - return { warning: combineWarnings([storageWarning, legacy.warning]) }; + if (!legacy?.recognized) { + const warning = combineWarnings([storageWarning, legacy?.warning]); + return warning ? { warning } : {}; } try { - saveConfig(legacy.config); - const verified = parseConfig(getConfigPath()); - if (!verified.valid || JSON.stringify(verified.config) !== JSON.stringify(legacy.config)) { + saveConfig({ version: CONFIG_VERSION, imagine: legacy.config }); + if (JSON.stringify(loadConfig().config.imagine) !== JSON.stringify(legacy.config)) { return { - warning: combineWarnings([ - storageWarning, - verified.warning, - `Could not verify migrated config ${getConfigPath()}. Legacy files were preserved.`, - ]), + warning: `Could not verify migrated config ${getConfigPath()}. Legacy files were preserved.`, }; } + unlinkSync(legacyPath); + return storageWarning ? { warning: storageWarning } : {}; } catch (error) { return { warning: combineWarnings([ storageWarning, - `Could not migrate legacy configuration to ${getConfigPath()}: ${errorMessage(error)}. Legacy files were preserved.`, + `Could not migrate configuration to ${getConfigPath()}: ${errorMessage(error)}. Legacy files were preserved.`, ]), }; } - const warning = combineWarnings([storageWarning, removeLegacyConfigs(legacy.recognizedPaths)]); - return warning ? { warning } : {}; } diff --git a/src/imagine/register.ts b/src/imagine/register.ts index b554817..5afd386 100644 --- a/src/imagine/register.ts +++ b/src/imagine/register.ts @@ -9,6 +9,7 @@ import { DEFAULT_IMAGINE_DEPENDENCIES, generateAndSaveImage, type ImagineDependencies, + type ImagineTokenResolver, } from './workflow.js'; const ENTRY_TYPE = 'grok-cli-imagine'; @@ -29,6 +30,10 @@ function applyImageToolPreference(pi: ExtensionAPI, enabled: boolean) { pi.setActiveTools(nextTools); } +export function syncImageToolPreference(pi: ExtensionAPI) { + applyImageToolPreference(pi, loadConfig().config.imagine.enabled); +} + type ImagineEntry = { path: string; relativePath: string; @@ -39,6 +44,7 @@ type ImagineEntry = { export function registerImagineFeature( pi: ExtensionAPI, dependencies: ImagineDependencies = DEFAULT_IMAGINE_DEPENDENCIES, + resolveToken?: ImagineTokenResolver, ) { pi.registerEntryRenderer(ENTRY_TYPE, (entry, { expanded }, theme) => { if (!entry.data) return; @@ -76,6 +82,7 @@ export function registerImagineFeature( : undefined, }, dependencies, + resolveToken, ); pi.appendEntry(ENTRY_TYPE, { path: saved.absolutePath, @@ -128,10 +135,10 @@ export function registerImagineFeature( ); return; } - applyImageToolPreference(pi, enabled); + syncImageToolPreference(pi); ctx.ui.notify(`image_gen: ${enabled ? 'on' : 'off'}`, 'info'); }, }); - registerImageGenTool(pi, dependencies); + registerImageGenTool(pi, dependencies, resolveToken); } diff --git a/src/imagine/tool.ts b/src/imagine/tool.ts index cda4d68..280c9b0 100644 --- a/src/imagine/tool.ts +++ b/src/imagine/tool.ts @@ -7,6 +7,7 @@ import { DEFAULT_IMAGINE_DEPENDENCIES, generateAndSaveImage, type ImagineDependencies, + type ImagineTokenResolver, } from './workflow.js'; const ImageGenParams = Type.Object({ @@ -40,6 +41,7 @@ function errorResult(error: unknown) { export function registerImageGenTool( pi: ExtensionAPI, dependencies: ImagineDependencies = DEFAULT_IMAGINE_DEPENDENCIES, + resolveToken?: ImagineTokenResolver, ) { pi.registerTool({ name: 'image_gen', @@ -59,6 +61,7 @@ export function registerImageGenTool( const saved = await generateAndSaveImage( { ctx, prompt, aspectRatio, signal }, dependencies, + resolveToken, ); return { content: [ diff --git a/src/imagine/workflow.ts b/src/imagine/workflow.ts index ffdfe3d..63164c8 100644 --- a/src/imagine/workflow.ts +++ b/src/imagine/workflow.ts @@ -23,6 +23,8 @@ export type SavedImagineImage = Awaited> & { previewError?: string; }; +export type ImagineTokenResolver = (ctx: ExtensionContext) => Promise; + export async function generateAndSaveImage( options: { ctx: ExtensionContext; @@ -33,8 +35,9 @@ export async function generateAndSaveImage( outPath?: string; }, dependencies: ImagineDependencies, + resolveToken: ImagineTokenResolver = resolveImagineToken, ): Promise { - const token = await resolveImagineToken(options.ctx); + const token = await resolveToken(options.ctx); if (!token) throw new Error(IMAGINE_AUTH_ERROR); const generated = await dependencies.generateImage({ token, diff --git a/src/models/catalog.ts b/src/models/catalog.ts index 6d2c792..1484a37 100644 --- a/src/models/catalog.ts +++ b/src/models/catalog.ts @@ -39,7 +39,7 @@ const FALLBACK_MODELS: GrokCliModelConfig[] = [ id: 'grok-composer-2.5-fast', name: 'Composer 2.5 Fast (Grok CLI)', reasoning: false, - input: ['text'], + input: ['text', 'image'], cost: COST_COMPOSER_FAST, contextWindow: 200_000, maxTokens: 30_000, diff --git a/src/provider/accountMigration.ts b/src/provider/accountMigration.ts new file mode 100644 index 0000000..842acc3 --- /dev/null +++ b/src/provider/accountMigration.ts @@ -0,0 +1,280 @@ +import { randomUUID } from 'node:crypto'; +import { existsSync, readFileSync, unlinkSync } from 'node:fs'; +import type { Credential } from '@earendil-works/pi-ai'; +import { readStoredCredential } from '@earendil-works/pi-coding-agent'; +import { hasTerminalControlCharacters } from '../config.js'; +import { + acquireFileLock, + getAccountVaultPath, + getConfigBackupPath, + getConfigPath, + getLegacyConfigPath, + getQuotaCachePath, + writeFileAtomic, +} from '../storage.js'; +import { + ACCOUNT_1_ID, + ACCOUNT_VAULT_MARKER, + type AccountCredential, + getAccountVault, + mutateAccountVault, +} from './accountVault.js'; + +type ReleasedAccount = { provider: string; label: string; slot: number }; +type ReadCredential = (provider: string) => Credential | undefined; + +function isObject(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function releasedAccount(provider: unknown, label: unknown): ReleasedAccount | undefined { + if (typeof provider !== 'string' || typeof label !== 'string') return undefined; + const normalizedLabel = label.trim(); + if ( + !normalizedLabel || + [...normalizedLabel].length > 40 || + hasTerminalControlCharacters(normalizedLabel) + ) { + return undefined; + } + if (provider === 'grok-cli') return { provider, label: normalizedLabel, slot: 1 }; + const match = /^grok-cli-((?:[2-9]|[1-9]\d+))$/.exec(provider); + return match ? { provider, label: normalizedLabel, slot: Number(match[1]) } : undefined; +} + +function readReleasedConfig() { + const path = existsSync(getConfigPath()) + ? getConfigPath() + : existsSync(getLegacyConfigPath()) + ? getLegacyConfigPath() + : undefined; + if (!path) { + return { + accounts: [{ provider: 'grok-cli', label: 'Account 1', slot: 1 }], + imagine: { enabled: true }, + nextSlot: 2, + selectedProvider: 'grok-cli', + }; + } + const contents = readFileSync(path, 'utf8'); + const parsed: unknown = JSON.parse(contents); + if (!isObject(parsed)) throw new Error(`${path} must contain a JSON object`); + if (parsed.version === 3) return undefined; + if (parsed.version !== 1 && parsed.version !== 2) { + throw new Error(`${path} is not a supported version 1 or 2 configuration`); + } + const releasedAccounts: { + items: unknown[]; + selectedProvider?: unknown; + nextAccountNumber?: unknown; + } = + isObject(parsed.accounts) && Array.isArray(parsed.accounts.items) + ? { + items: parsed.accounts.items, + selectedProvider: parsed.accounts.selectedProvider, + nextAccountNumber: parsed.accounts.nextAccountNumber, + } + : { items: [] }; + const baseIndex = releasedAccounts.items.findIndex( + (account) => + isObject(account) && + releasedAccount(account.provider, account.label)?.provider === 'grok-cli', + ); + const accountValues = + baseIndex >= 0 + ? [ + releasedAccounts.items[baseIndex], + ...releasedAccounts.items.filter((_account, index) => index !== baseIndex), + ] + : [{ provider: 'grok-cli', label: 'Account 1' }, ...releasedAccounts.items]; + const providers = new Set(); + const labels = new Set(); + const accounts = accountValues.flatMap((account) => { + if (!isObject(account)) return []; + const value = releasedAccount(account.provider, account.label); + if (!value || providers.has(value.provider) || labels.has(value.label.toLocaleLowerCase())) { + return []; + } + providers.add(value.provider); + labels.add(value.label.toLocaleLowerCase()); + return [value]; + }); + const ordered = accounts.sort((left, right) => left.slot - right.slot); + return { + contents, + path, + accounts: ordered, + selectedProvider: + typeof releasedAccounts.selectedProvider === 'string' + ? releasedAccounts.selectedProvider + : 'grok-cli', + nextSlot: + typeof releasedAccounts.nextAccountNumber === 'number' && + Number.isInteger(releasedAccounts.nextAccountNumber) + ? Math.max( + releasedAccounts.nextAccountNumber, + Math.max(...ordered.map(({ slot }) => slot)) + 1, + ) + : Math.max(...ordered.map(({ slot }) => slot)) + 1, + imagine: + isObject(parsed.imagine) && typeof parsed.imagine.enabled === 'boolean' + ? { enabled: parsed.imagine.enabled } + : { enabled: true }, + }; +} + +function accountCredential(credential: Credential | undefined): AccountCredential | undefined { + if ( + credential?.type !== 'oauth' || + !credential.access || + credential.access === ACCOUNT_VAULT_MARKER || + !credential.refresh || + credential.refresh === ACCOUNT_VAULT_MARKER || + !Number.isFinite(credential.expires) + ) { + return undefined; + } + const { type: _type, ...oauth } = credential; + return structuredClone(oauth) as AccountCredential; +} + +async function migrateQuota(accountIds: Map) { + const release = await acquireFileLock(getQuotaCachePath()); + try { + if (!existsSync(getQuotaCachePath())) return; + let parsed: unknown; + try { + parsed = JSON.parse(readFileSync(getQuotaCachePath(), 'utf8')) as unknown; + } catch { + return; + } + if (!isObject(parsed) || parsed.version !== 1 || !isObject(parsed.accounts)) return; + const migratedIds = new Set(accountIds.values()); + const accounts = Object.fromEntries( + Object.entries(parsed.accounts).flatMap(([provider, value]) => { + const accountId = + accountIds.get(provider) ?? (migratedIds.has(provider) ? provider : undefined); + return accountId ? [[accountId, value]] : []; + }), + ); + writeFileAtomic( + getQuotaCachePath(), + `${JSON.stringify({ ...parsed, accounts }, null, 2)}\n`, + 0o600, + ); + } finally { + await release(); + } +} + +function preserveReleasedConfig(contents: string) { + if (existsSync(getConfigBackupPath())) return; + writeFileAtomic(getConfigBackupPath(), contents, 0o600); + if (readFileSync(getConfigBackupPath(), 'utf8') !== contents) { + throw new Error(`Could not verify ${getConfigBackupPath()}`); + } +} + +async function migrateReleasedAccountsUnlocked( + options: { readCredential?: ReadCredential } = {}, +): Promise<{ migrated: boolean }> { + const released = readReleasedConfig(); + const current = await getAccountVault(); + if ( + !existsSync(getConfigPath()) && + !existsSync(getLegacyConfigPath()) && + current.migration.legacyCredentialCopyComplete + ) { + return { migrated: false }; + } + if (!released && current.migration.legacyCredentialCopyComplete) return { migrated: false }; + if (!released) { + await mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + }); + return { migrated: true }; + } + const source = released; + const readCredential = options.readCredential ?? readStoredCredential; + + await mutateAccountVault((vault) => { + const accounts = source.accounts.map((account) => { + const existing = vault.accounts.find((candidate) => candidate.slot === account.slot); + const credential = accountCredential(readCredential(account.provider)); + return { + id: account.slot === 1 ? ACCOUNT_1_ID : (existing?.id ?? randomUUID()), + slot: account.slot, + label: account.label, + ...(credential ? { credential } : {}), + revision: existing?.revision ?? 0, + }; + }); + const selected = + accounts.find((account) => { + const sourceAccount = source.accounts.find( + (candidate) => candidate.provider === source.selectedProvider, + ); + return account.credential && account.slot === sourceAccount?.slot; + }) ?? accounts.find((account) => account.credential); + vault.accounts = accounts; + vault.nextSlot = source.nextSlot; + if (selected) vault.activeAccountId = selected.id; + else delete vault.activeAccountId; + vault.migration.markerInstallPending = accounts.some((account) => account.credential); + vault.migration.legacyCredentialCopyComplete = false; + }); + + const migrated = await getAccountVault(); + const accountIds = new Map( + migrated.accounts.map((account) => [ + account.slot === 1 ? 'grok-cli' : `grok-cli-${account.slot}`, + account.id, + ]), + ); + if ('contents' in source && typeof source.contents === 'string') { + preserveReleasedConfig(source.contents); + } + await migrateQuota(accountIds); + if ('contents' in source) { + writeFileAtomic( + getConfigPath(), + `${JSON.stringify({ version: 3, imagine: source.imagine }, null, 2)}\n`, + ); + if (source.path === getLegacyConfigPath()) unlinkSync(source.path); + } + await mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + }); + return { migrated: true }; +} + +export async function migrateReleasedAccounts( + options: { readCredential?: ReadCredential } = {}, +): Promise<{ migrated: boolean }> { + const release = await acquireFileLock(`${getAccountVaultPath()}.migration`); + try { + return await migrateReleasedAccountsUnlocked(options); + } finally { + await release(); + } +} + +export async function confirmMarkerInstallation( + readCredential: ReadCredential = readStoredCredential, +) { + if (!hasAccountVaultMarker(readCredential)) return false; + return mutateAccountVault((vault) => { + if (!vault.migration.markerInstallPending) return false; + vault.migration.markerInstallPending = false; + return true; + }); +} + +export function hasAccountVaultMarker(readCredential: ReadCredential = readStoredCredential) { + const credential = readCredential('grok-cli'); + return ( + credential?.type === 'oauth' && + credential.access === ACCOUNT_VAULT_MARKER && + credential.refresh === ACCOUNT_VAULT_MARKER + ); +} diff --git a/src/provider/accountRouting.ts b/src/provider/accountRouting.ts new file mode 100644 index 0000000..6d9eb5a --- /dev/null +++ b/src/provider/accountRouting.ts @@ -0,0 +1,100 @@ +import * as oauth from '../auth/oauth.js'; +import { getBaseUrl } from '../auth/oauth.js'; +import { confirmMarkerInstallation, hasAccountVaultMarker } from './accountMigration.js'; +import { + ACCOUNT_1_ID, + type AccountCredential, + getAccountVault, + mutateAccountVault, +} from './accountVault.js'; + +export interface AccountRoute { + accountId: string; + revision: number; + token: string; + baseUrl: string; + source: 'environment' | 'vault'; +} + +const refreshes = new Map>(); + +function route(accountId: string, revision: number, credential: AccountCredential): AccountRoute { + return { + accountId, + revision, + token: credential.access, + baseUrl: (credential.baseUrl ?? getBaseUrl()).replace(/\/+$/, ''), + source: 'vault', + }; +} + +async function refreshRoute( + accountId: string, + revision: number, + credential: AccountCredential, + fallbackToActive: boolean, +) { + const key = `${accountId}:${revision}`; + const pending = + refreshes.get(key) ?? + oauth + .refresh(credential as typeof credential & Record) + .then((refreshed) => refreshed as AccountCredential); + refreshes.set(key, pending); + try { + const refreshed = await pending; + const stored = await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === accountId); + if (!account?.credential || account.revision !== revision) return undefined; + account.credential = refreshed; + account.revision += 1; + return route(account.id, account.revision, account.credential); + }); + if (stored) return stored; + return resolveAccountRoute(fallbackToActive ? undefined : accountId); + } finally { + if (refreshes.get(key) === pending) refreshes.delete(key); + } +} + +export async function resolveAccountRoute(accountId?: string): Promise { + if (process.env.GROK_CLI_OAUTH_TOKEN) { + return { + accountId: ACCOUNT_1_ID, + revision: 0, + token: process.env.GROK_CLI_OAUTH_TOKEN, + baseUrl: getBaseUrl().replace(/\/+$/, ''), + source: 'environment', + }; + } + + let vault = await getAccountVault(); + if (vault.migration.markerInstallPending && (await confirmMarkerInstallation())) { + vault = await getAccountVault(); + } + if (!hasAccountVaultMarker()) { + throw new Error('Grok CLI account migration is ready. Please run /login and select Grok CLI.'); + } + const account = accountId + ? vault.accounts.find( + (candidate) => candidate.id === accountId && candidate.credential !== undefined, + ) + : (vault.accounts.find( + (candidate) => candidate.id === vault.activeAccountId && candidate.credential !== undefined, + ) ?? vault.accounts.find((candidate) => candidate.credential !== undefined)); + if (!account?.credential) { + throw new Error('Grok CLI login is required. Run /login and select Grok CLI.'); + } + if (!accountId && vault.activeAccountId !== account.id) { + await mutateAccountVault((latest) => { + const selected = latest.accounts.find( + (candidate) => candidate.id === account.id && candidate.credential !== undefined, + ); + if (selected) latest.activeAccountId = selected.id; + }); + } + if (Date.now() < account.credential.expires) { + return route(account.id, account.revision, account.credential); + } + return refreshRoute(account.id, account.revision, account.credential, accountId === undefined); +} diff --git a/src/provider/accountVault.ts b/src/provider/accountVault.ts new file mode 100644 index 0000000..a52bab4 --- /dev/null +++ b/src/provider/accountVault.ts @@ -0,0 +1,254 @@ +import { chmodSync, existsSync, mkdirSync, readFileSync } from 'node:fs'; +import { + acquireFileLock, + getAccountVaultPath, + getGrokCliDirectory, + writeFileAtomic, +} from '../storage.js'; + +export const ACCOUNT_1_ID = 'account-1'; +export const ACCOUNT_VAULT_MARKER = 'pi-grok-cli-account-vault-v1'; + +export interface AccountCredential { + access: string; + refresh: string; + expires: number; + tokenEndpoint?: string; + discovery?: { + authorization_endpoint: string; + token_endpoint: string; + device_authorization_endpoint?: string; + }; + idToken?: string; + tokenType?: string; + baseUrl?: string; +} + +export interface VaultAccount { + id: string; + slot: number; + label: string; + credential?: AccountCredential; + revision: number; +} + +export interface AccountVault { + version: 1; + migration: { + legacyCredentialCopyComplete: boolean; + markerInstallPending: boolean; + }; + nextSlot: number; + activeAccountId?: string; + accounts: VaultAccount[]; +} + +const defaultVault = (): AccountVault => ({ + version: 1, + migration: { + legacyCredentialCopyComplete: false, + markerInstallPending: false, + }, + nextSlot: 2, + accounts: [{ id: ACCOUNT_1_ID, slot: 1, label: 'Account 1', revision: 0 }], +}); + +function isObject(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function validOptionalString(value: unknown) { + return value === undefined || typeof value === 'string'; +} + +function parseCredential(value: unknown): AccountCredential | undefined { + if (!isObject(value)) return undefined; + if ( + typeof value.access !== 'string' || + !value.access || + value.access === ACCOUNT_VAULT_MARKER || + typeof value.refresh !== 'string' || + !value.refresh || + value.refresh === ACCOUNT_VAULT_MARKER || + typeof value.expires !== 'number' || + !Number.isFinite(value.expires) || + !validOptionalString(value.tokenEndpoint) || + !validOptionalString(value.idToken) || + !validOptionalString(value.tokenType) || + !validOptionalString(value.baseUrl) + ) { + return undefined; + } + if ( + value.discovery !== undefined && + (!isObject(value.discovery) || + typeof value.discovery.authorization_endpoint !== 'string' || + typeof value.discovery.token_endpoint !== 'string' || + !validOptionalString(value.discovery.device_authorization_endpoint)) + ) { + return undefined; + } + return { + access: value.access, + refresh: value.refresh, + expires: value.expires, + ...(typeof value.tokenEndpoint === 'string' ? { tokenEndpoint: value.tokenEndpoint } : {}), + ...(isObject(value.discovery) + ? { + discovery: { + authorization_endpoint: String(value.discovery.authorization_endpoint), + token_endpoint: String(value.discovery.token_endpoint), + ...(typeof value.discovery.device_authorization_endpoint === 'string' + ? { device_authorization_endpoint: value.discovery.device_authorization_endpoint } + : {}), + }, + } + : {}), + ...(typeof value.idToken === 'string' ? { idToken: value.idToken } : {}), + ...(typeof value.tokenType === 'string' ? { tokenType: value.tokenType } : {}), + ...(typeof value.baseUrl === 'string' ? { baseUrl: value.baseUrl } : {}), + }; +} + +function parseVault(value: unknown, path: string): AccountVault { + const invalid = (reason: string): never => { + throw new Error(`Invalid Grok CLI account vault at ${path}: ${reason}`); + }; + if (!isObject(value)) return invalid('root must be a JSON object'); + if (value.version !== 1) invalid('unsupported version'); + const raw = value; + if ( + !isObject(raw.migration) || + typeof raw.migration.legacyCredentialCopyComplete !== 'boolean' || + typeof raw.migration.markerInstallPending !== 'boolean' + ) { + invalid('invalid migration state'); + } + const migration = raw.migration as { + legacyCredentialCopyComplete: boolean; + markerInstallPending: boolean; + }; + if (!Number.isInteger(raw.nextSlot) || Number(raw.nextSlot) < 2) { + invalid('invalid next slot'); + } + if (!Array.isArray(raw.accounts) || !raw.accounts.length) return invalid('accounts are missing'); + + const ids = new Set(); + const slots = new Set(); + const labels = new Set(); + const accounts = raw.accounts.map((candidate, index): VaultAccount => { + if ( + !isObject(candidate) || + typeof candidate.id !== 'string' || + !candidate.id || + !Number.isInteger(candidate.slot) || + Number(candidate.slot) < 1 || + typeof candidate.label !== 'string' || + !candidate.label.trim() || + typeof candidate.revision !== 'number' || + !Number.isInteger(candidate.revision) || + Number(candidate.revision) < 0 + ) { + return invalid(`invalid account at index ${index}`); + } + const slot = Number(candidate.slot); + const label = candidate.label.trim(); + if (ids.has(candidate.id) || slots.has(slot) || labels.has(label.toLocaleLowerCase())) { + return invalid(`duplicate account at index ${index}`); + } + const credential = + candidate.credential === undefined ? undefined : parseCredential(candidate.credential); + if (candidate.credential !== undefined && !credential) { + return invalid(`invalid credential at index ${index}`); + } + ids.add(candidate.id); + slots.add(slot); + labels.add(label.toLocaleLowerCase()); + return { + id: candidate.id, + slot, + label, + ...(credential ? { credential } : {}), + revision: Number(candidate.revision), + }; + }); + + if (accounts[0]?.id !== ACCOUNT_1_ID || accounts[0].slot !== 1) { + invalid('Account 1 must be first and permanent'); + } + if (Number(raw.nextSlot) <= Math.max(...accounts.map((account) => account.slot))) { + invalid('next slot must be greater than every account slot'); + } + if ( + raw.activeAccountId !== undefined && + (typeof raw.activeAccountId !== 'string' || + !accounts.some( + (account) => account.id === raw.activeAccountId && account.credential !== undefined, + )) + ) { + invalid('active account must be logged in'); + } + + return { + version: 1, + migration: { + legacyCredentialCopyComplete: migration.legacyCredentialCopyComplete, + markerInstallPending: migration.markerInstallPending, + }, + nextSlot: Number(raw.nextSlot), + ...(typeof raw.activeAccountId === 'string' ? { activeAccountId: raw.activeAccountId } : {}), + accounts, + }; +} + +async function withVaultLock(action: (current: AccountVault) => T, write: boolean) { + const directory = getGrokCliDirectory(); + const path = getAccountVaultPath(); + mkdirSync(directory, { recursive: true, mode: 0o700 }); + chmodSync(directory, 0o700); + const release = await acquireFileLock(path); + try { + const current = existsSync(path) + ? parseVault(JSON.parse(readFileSync(path, 'utf8')) as unknown, path) + : defaultVault(); + const next = structuredClone(current); + const result = action(next); + if (!write) return result; + const validated = parseVault(next, path); + writeFileAtomic(path, `${JSON.stringify(validated, null, 2)}\n`, 0o600); + chmodSync(path, 0o600); + return result; + } catch (error) { + if (error instanceof SyntaxError) { + throw new Error(`Invalid Grok CLI account vault at ${path}: ${error.message}`); + } + throw error; + } finally { + await release(); + } +} + +export function getAccountVault() { + return withVaultLock((current) => structuredClone(current), false); +} + +export function getAccountVaultSync() { + const directory = getGrokCliDirectory(); + const path = getAccountVaultPath(); + mkdirSync(directory, { recursive: true, mode: 0o700 }); + chmodSync(directory, 0o700); + try { + return existsSync(path) + ? parseVault(JSON.parse(readFileSync(path, 'utf8')) as unknown, path) + : defaultVault(); + } catch (error) { + if (error instanceof SyntaxError) { + throw new Error(`Invalid Grok CLI account vault at ${path}: ${error.message}`); + } + throw error; + } +} + +export function mutateAccountVault(mutation: (vault: AccountVault) => T) { + return withVaultLock(mutation, true); +} diff --git a/src/provider/accounts.ts b/src/provider/accounts.ts index f5e984f..4e9b729 100644 --- a/src/provider/accounts.ts +++ b/src/provider/accounts.ts @@ -1,39 +1,48 @@ -import type { AuthInteraction } from '@earendil-works/pi-ai'; -import type { - ExtensionAPI, - ExtensionCommandContext, - ExtensionContext, - ModelRuntime, +import { spawn } from 'node:child_process'; +import { randomUUID } from 'node:crypto'; +import type { AuthInteraction, OAuthLoginCallbacks } from '@earendil-works/pi-ai'; +import { + type ExtensionAPI, + type ExtensionCommandContext, + type ExtensionContext, + readStoredCredential, } from '@earendil-works/pi-coding-agent'; -import { matchesKey, SelectList } from '@earendil-works/pi-tui'; +import * as oauth from '../auth/oauth.js'; +import { hasTerminalControlCharacters } from '../config.js'; +import { confirmMarkerInstallation } from './accountMigration.js'; +import { type AccountRoute, resolveAccountRoute } from './accountRouting.js'; import { - findAvailableAccountNumber, - type GrokCliAccount, - type GrokCliConfig, - hasTerminalControlCharacters, - loadConfig, - saveConfig, -} from '../config.js'; + ACCOUNT_1_ID, + ACCOUNT_VAULT_MARKER, + type AccountCredential, + getAccountVaultSync, + mutateAccountVault, + type VaultAccount, +} from './accountVault.js'; import { fetchBillingUsage } from './billing.js'; import { createAccountDashboard } from './dashboard/server.js'; import { - formatCachedQuota, isCachedQuotaFresh, loadQuotaCache, removeQuotaUsage, - saveQuotaUsage, + saveQuotaUsageWhen, } from './quotaCache.js'; +import { + createSessionAccountSelection, + type SessionAccountSelection, +} from './sessionAccountSelection.js'; export const GROK_CLI_PROVIDER = 'grok-cli'; export const DEFAULT_GROK_MODEL = 'grok-build'; - -type RegisterAccount = (account: GrokCliAccount) => void; +const noOp = () => {}; export type PlanTier = 'free' | 'supergrok-lite' | 'supergrok' | 'supergrok-heavy'; export interface AccountSnapshot { - provider: string; + id: string; + slot: number; label: string; + permanent: boolean; status: string; authenticated: boolean; active: boolean; @@ -55,116 +64,139 @@ export interface AccountSnapshot { } export interface AccountsSnapshot { + connected: boolean; accounts: AccountSnapshot[]; } export function isGrokCliProvider(provider: string | undefined): boolean { - return provider === GROK_CLI_PROVIDER || /^grok-cli-(?:[2-9]|[1-9]\d+)$/.test(provider ?? ''); -} - -function copyConfig(): GrokCliConfig { - const config = loadConfig().config; - return { - ...config, - accounts: { - ...config.accounts, - items: config.accounts.items.map((account) => ({ ...account })), - }, - imagine: { ...config.imagine }, - vision: { ...config.vision }, - }; + return provider === GROK_CLI_PROVIDER; } -function accountNumber(provider: string) { - if (provider === GROK_CLI_PROVIDER) return 1; - return Number(provider.slice('grok-cli-'.length)); +export function planTier(monthlyLimit: number): PlanTier { + if (monthlyLimit <= 0) return 'free'; + if (monthlyLimit <= 4000) return 'supergrok-lite'; + if (monthlyLimit <= 20000) return 'supergrok'; + return 'supergrok-heavy'; } -function defaultLabel(provider: string) { - return `Account ${accountNumber(provider)}`; +function defaultLabel(slot: number) { + return `Account ${slot}`; } -function labelError(config: GrokCliConfig, provider: string, label: string) { - if ([...label].length > 40) return 'Account labels must be 40 characters or fewer.'; - if (hasTerminalControlCharacters(label)) - return 'Account labels cannot contain control characters.'; +function normalizeLabel(accounts: VaultAccount[], id: string, slot: number, value: string) { + const label = value.trim() || defaultLabel(slot); + if ([...label].length > 40) throw new Error('Account labels must be 40 characters or fewer.'); + if (hasTerminalControlCharacters(label)) { + throw new Error('Account labels cannot contain control characters.'); + } if ( - config.accounts.items.some( + accounts.some( (account) => - account.provider !== provider && - account.label.toLocaleLowerCase() === label.toLocaleLowerCase(), + account.id !== id && account.label.toLocaleLowerCase() === label.toLocaleLowerCase(), ) ) { - return `An account named “${label}” already exists.`; + throw new Error(`An account named “${label}” already exists.`); } - return undefined; -} - -function normalizeLabel(config: GrokCliConfig, provider: string, value: string) { - const label = value.trim() || defaultLabel(provider); - const error = labelError(config, provider, label); - if (error) throw new Error(error); return label; } -async function promptLabel( - ctx: ExtensionCommandContext, - config: GrokCliConfig, - provider: string, - title: string, -) { - while (true) { - const input = await ctx.ui.input(title, defaultLabel(provider)); - if (input === undefined) return undefined; - const label = input.trim() || defaultLabel(provider); - const error = labelError(config, provider, label); - if (!error) return label; - ctx.ui.notify(error, 'error'); - } +function callbacks(interaction: AuthInteraction): OAuthLoginCallbacks { + return { + onAuth(info) { + interaction.notify({ type: 'auth_url', ...info }); + }, + onDeviceCode(info) { + interaction.notify({ type: 'device_code', ...info }); + }, + onPrompt(prompt) { + return interaction.prompt({ type: 'text', ...prompt }); + }, + onProgress(message) { + interaction.notify({ type: 'progress', message }); + }, + onManualCodeInput() { + return interaction.prompt({ + type: 'manual_code', + message: 'Paste the authorization code', + }); + }, + onSelect(prompt) { + return interaction.prompt({ type: 'select', ...prompt }); + }, + signal: interaction.signal, + }; } -function hasStoredAuth(ctx: ExtensionContext, provider: string) { - return ctx.modelRegistry.getProviderAuthStatus(provider).configured; +function openAuthorizationUrl(url: string, onError: () => void) { + const command: [string, string[]] = + process.platform === 'darwin' + ? ['open', [url]] + : process.platform === 'win32' + ? ['rundll32', ['url.dll,FileProtocolHandler', url]] + : ['xdg-open', [url]]; + const child = spawn(command[0], command[1], { detached: true, stdio: 'ignore' }); + child.once('error', onError); + child.unref(); } -function accountRuntime(ctx: ExtensionContext) { - const runtime = (ctx.modelRegistry as unknown as { runtime?: ModelRuntime }).runtime; - if (!runtime) throw new Error('Pi account authentication runtime is unavailable.'); - return runtime; +function terminalInteraction(ctx: ExtensionCommandContext): AuthInteraction { + return { + notify(event) { + if (event.type === 'auth_url') { + openAuthorizationUrl(event.url, () => { + ctx.ui.notify(`Open this URL to continue login: ${event.url}`, 'warning'); + }); + ctx.ui.notify(event.instructions ?? 'Complete the Grok CLI login in your browser.', 'info'); + return; + } + if (event.type === 'device_code') { + ctx.ui.notify(`Open ${event.verificationUri} and enter code ${event.userCode}.`, 'info'); + return; + } + ctx.ui.notify(event.message, 'info'); + }, + async prompt(prompt) { + if (prompt.type === 'select') { + const label = await ctx.ui.select( + prompt.message, + prompt.options.map((option) => option.label), + ); + return prompt.options.find((option) => option.label === label)?.id ?? ''; + } + return (await ctx.ui.input(prompt.message, prompt.placeholder)) ?? ''; + }, + }; } -function hasAccountAuth(ctx: ExtensionContext, provider: string) { - return ( - hasStoredAuth(ctx, provider) || - (provider === GROK_CLI_PROVIDER && Boolean(process.env.GROK_CLI_OAUTH_TOKEN)) - ); +function accountFrom(id: string) { + const account = getAccountVaultSync().accounts.find((candidate) => candidate.id === id); + if (!account) throw new Error(`Unknown Grok CLI account: ${id}`); + return account; } -function accountStatus(ctx: ExtensionContext, config: GrokCliConfig, provider: string) { - const environment = provider === GROK_CLI_PROVIDER && Boolean(process.env.GROK_CLI_OAUTH_TOKEN); - if (config.accounts.selectedProvider === provider && hasAccountAuth(ctx, provider)) { - return environment ? 'Active (environment)' : 'Active'; +export function accountRouteIsCurrent(route: AccountRoute) { + if (route.source === 'environment') { + return process.env.GROK_CLI_OAUTH_TOKEN === route.token; } - if (environment) return 'Authenticated (environment)'; - return hasStoredAuth(ctx, provider) ? 'Authenticated' : 'Login required'; + const account = getAccountVaultSync().accounts.find( + (candidate) => candidate.id === route.accountId, + ); + return account?.revision === route.revision && account.credential !== undefined; } -// Tier from the monthly credit cap: 0 free, up to 4,000 Lite, up to 20,000 SuperGrok, above Heavy. -export function planTier(monthlyLimit: number): PlanTier { - if (monthlyLimit <= 0) return 'free'; - if (monthlyLimit <= 4000) return 'supergrok-lite'; - if (monthlyLimit <= 20000) return 'supergrok'; - return 'supergrok-heavy'; +async function clearQuota(ctx: ExtensionContext, id: string) { + await removeQuotaUsage(id).catch(() => { + ctx.ui.notify('The account changed, but its cached quota could not be cleared.', 'warning'); + }); } -async function resolveAccountToken(ctx: ExtensionContext, provider: string) { - if (provider === GROK_CLI_PROVIDER && process.env.GROK_CLI_OAUTH_TOKEN) { - return process.env.GROK_CLI_OAUTH_TOKEN; - } - try { - return await ctx.modelRegistry.getApiKeyForProvider(provider); - } catch { - return undefined; +async function refreshQuota(route: AccountRoute, signal: AbortSignal) { + const usage = await fetchBillingUsage( + route.token, + AbortSignal.any([signal, AbortSignal.timeout(30_000)]), + ); + if (!(await saveQuotaUsageWhen(route.accountId, usage, () => accountRouteIsCurrent(route)))) { + throw new Error('The account changed while its quota was refreshing.'); } } @@ -180,569 +212,231 @@ async function refreshAccountBatches( await refreshAccountBatches(accounts.slice(3), refresh, signal); } -function showAccountSelector( - ctx: ExtensionCommandContext, - config: GrokCliConfig, - mode: 'main' | 'manage', - manager: AccountManager, -) { - return ctx.ui.custom((tui, theme, _keybindings, done) => { - const accountItems = new Map( - config.accounts.items.map((account) => [ - account.provider, - { - value: account.provider, - label: `${account.label} — ${accountStatus(ctx, config, account.provider)}`, - description: '', - }, - ]), - ); - const items = [ - ...accountItems.values(), - ...(mode === 'main' - ? [ - { value: 'action:add', label: '＋ Add account', description: '' }, - { value: 'action:manage', label: 'Manage accounts', description: '' }, - ] - : []), - ]; - const selectList = new SelectList( - items, - Math.min(items.length, 10), - { - selectedPrefix: (text) => theme.fg('accent', text), - selectedText: (text) => theme.fg('accent', text), - description: (text) => theme.fg('dim', text), - scrollInfo: (text) => theme.fg('dim', text), - noMatch: (text) => theme.fg('warning', text), - }, - { minPrimaryColumnWidth: 24, maxPrimaryColumnWidth: 40 }, - ); - const refreshingProviders = new Set(); - const failedProviders = new Set(); - const controller = new AbortController(); - let cache = loadQuotaCache(); - let status: string | undefined; - let refreshing = false; - let disposed = false; - - const updateDescriptions = () => { - for (const account of config.accounts.items) { - const item = accountItems.get(account.provider); - if (!item) continue; - const cached = cache.accounts[account.provider]; - item.description = !hasAccountAuth(ctx, account.provider) - ? 'Quota unavailable · login required' - : refreshingProviders.has(account.provider) - ? 'Refreshing…' - : failedProviders.has(account.provider) - ? cached - ? `${formatCachedQuota(cached)} · refresh failed` - : 'Refresh failed · press r to retry' - : cached - ? formatCachedQuota(cached) - : 'Quota not fetched · press r'; - } - selectList.invalidate(); - }; - - const close = (value: string | undefined) => { - if (disposed) return; - disposed = true; - controller.abort(); - done(value); - }; - - const startRefresh = async () => { - if (refreshing || disposed) return; - const accounts = config.accounts.items.filter((account) => - hasAccountAuth(ctx, account.provider), - ); - if (!accounts.length) { - status = 'No logged-in accounts to refresh'; - tui.requestRender(); - return; - } - refreshing = true; - failedProviders.clear(); - for (const account of accounts) refreshingProviders.add(account.provider); - status = `Refreshing quotas… 0/${accounts.length}`; - updateDescriptions(); - tui.requestRender(); - const result = await manager.refresh(ctx, controller.signal, (progress) => { - refreshingProviders.delete(progress.provider); - if (!progress.updated) failedProviders.add(progress.provider); - cache = loadQuotaCache(); - status = `Refreshing quotas… ${progress.completed}/${progress.total}`; - updateDescriptions(); - if (!disposed) tui.requestRender(); - }); - - refreshing = false; - if (disposed) return; - status = `Updated ${result.updated} accounts; ${result.failed.length} failed`; - tui.requestRender(); - }; - - selectList.onSelect = (item) => close(item.value); - selectList.onCancel = () => close(undefined); - updateDescriptions(); - - return { - render(width: number) { - return [ - theme.fg( - 'accent', - theme.bold(mode === 'main' ? 'Grok CLI accounts:' : 'Manage Grok CLI account:'), - ), - '', - ...selectList.render(width), - '', - ...(status ? [theme.fg('dim', status)] : []), - theme.fg( - 'dim', - `r refresh quotas · enter select · esc ${mode === 'main' ? 'close' : 'back'}`, - ), - ]; - }, - invalidate() { - selectList.invalidate(); - }, - handleInput(data: string) { - if (data === 'r' || data === 'R' || matchesKey(data, 'r')) { - void startRefresh(); - return; - } - selectList.handleInput(data); - if (!disposed) tui.requestRender(); - }, - dispose() { - if (disposed) return; - disposed = true; - controller.abort(); - }, - }; - }); -} - -function prefillLogin(ctx: ExtensionContext, provider: string) { - ctx.ui.setEditorText(`/login ${provider}`); +function fallbackActive(accounts: VaultAccount[], removedId: string) { + return accounts.find((account) => account.id !== removedId && account.credential !== undefined) + ?.id; } -async function addAccount( - ctx: ExtensionCommandContext, - config: GrokCliConfig, - manager: AccountManager, -) { - const provider = manager.nextProvider(); - const label = await promptLabel(ctx, config, provider, 'Label this Grok CLI account:'); - if (!label) return; - try { - prefillLogin(ctx, (await manager.add(ctx, label)).provider); - } catch (error) { - ctx.ui.notify( - `Could not add Grok CLI account: ${error instanceof Error ? error.message : String(error)}`, - 'error', - ); - } +function replaceVaultDefault(vault: ReturnType, removedId: string) { + if (vault.activeAccountId !== removedId) return; + const fallback = fallbackActive(vault.accounts, removedId); + if (fallback) vault.activeAccountId = fallback; + else delete vault.activeAccountId; } -async function renameAccount( - ctx: ExtensionCommandContext, - config: GrokCliConfig, - account: GrokCliAccount, - manager: AccountManager, -) { - const label = await promptLabel(ctx, config, account.provider, `Rename “${account.label}”:`); - if (!label) return; - try { - await manager.rename(ctx, account.provider, label); - } catch (error) { - ctx.ui.notify( - `Could not rename Grok CLI account: ${error instanceof Error ? error.message : String(error)}`, - 'error', - ); - } -} - -async function removeBaseAccount( - ctx: ExtensionCommandContext, - account: GrokCliAccount, - manager: AccountManager, -) { - if ( - !(await ctx.ui.confirm( - `Log out of “${account.label}”?`, - 'This removes its saved OAuth login from Pi.', - )) - ) { - return; - } - try { - const result = await manager.logout(ctx, account.provider); - if (result.warning) ctx.ui.notify(result.warning, 'warning'); - } catch (error) { - ctx.ui.notify( - `Could not log out of Grok CLI: ${error instanceof Error ? error.message : String(error)}`, - 'error', - ); - } -} - -async function fallbackBeforeRemoval( - pi: ExtensionAPI, - ctx: ExtensionContext, - config: GrokCliConfig, - removed: GrokCliAccount, -) { - const candidates = config.accounts.items.filter( - (account) => account.provider !== removed.provider && hasAccountAuth(ctx, account.provider), - ); - if (ctx.model?.provider !== removed.provider) { - return config.accounts.selectedProvider === removed.provider - ? (candidates[0]?.provider ?? GROK_CLI_PROVIDER) - : config.accounts.selectedProvider; - } - for (const account of candidates) { - const target = - ctx.modelRegistry.find(account.provider, ctx.model?.id ?? DEFAULT_GROK_MODEL) ?? - ctx.modelRegistry.find(account.provider, DEFAULT_GROK_MODEL); - if (target && (await pi.setModel(target))) return account.provider; - } - return candidates.length ? null : undefined; -} - -async function removeAlias( - ctx: ExtensionCommandContext, - account: GrokCliAccount, - manager: AccountManager, -) { - if ( - !(await ctx.ui.confirm( - `Remove “${account.label}”?`, - 'This removes its saved OAuth login from Pi and deletes the account slot.', - )) - ) { - return; - } - try { - const result = await manager.remove(ctx, account.provider); - if (result.warning) ctx.ui.notify(result.warning, 'warning'); - } catch (error) { - ctx.ui.notify( - `Could not remove Grok CLI account: ${error instanceof Error ? error.message : String(error)}`, - 'error', - ); - } -} - -async function manageAccount( - ctx: ExtensionCommandContext, - config: GrokCliConfig, - account: GrokCliAccount, - manager: AccountManager, -) { - const environment = - account.provider === GROK_CLI_PROVIDER && Boolean(process.env.GROK_CLI_OAUTH_TOKEN); - const loginLabel = hasStoredAuth(ctx, account.provider) ? 'Log in again' : 'Log in'; - const removeLabel = account.provider === GROK_CLI_PROVIDER ? 'Log out' : 'Log out and remove'; - const options = environment - ? ['Rename', 'Environment token instructions', 'Back'] - : ['Rename', loginLabel, removeLabel, 'Back']; - const action = await ctx.ui.select(`Manage “${account.label}”:`, options); - switch (action) { - case 'Rename': - await renameAccount(ctx, config, account, manager); - return; - case loginLabel: - prefillLogin(ctx, account.provider); - return; - case 'Environment token instructions': - ctx.ui.notify( - 'Unset GROK_CLI_OAUTH_TOKEN and restart Pi to remove the environment token.', - 'info', - ); - return; - case 'Log out': - await removeBaseAccount(ctx, account, manager); - return; - case 'Log out and remove': - await removeAlias(ctx, account, manager); - } -} - -export function resolveGrokProvider(ctx: Pick) { - const config = loadConfig().config; - if ( - isGrokCliProvider(ctx.model?.provider) && - config.accounts.items.some((account) => account.provider === ctx.model?.provider) - ) { - return ctx.model?.provider ?? GROK_CLI_PROVIDER; - } - return config.accounts.selectedProvider; +function logoutWarning(vault: ReturnType) { + return { + warning: vault.activeAccountId + ? undefined + : 'No logged-in account remains. Run /logout and select Grok CLI to disconnect Pi.', + }; } export async function resolveGrokToken( - ctx: Pick, + _ctx?: Pick, ): Promise { - const provider = resolveGrokProvider(ctx); - if (provider === GROK_CLI_PROVIDER && process.env.GROK_CLI_OAUTH_TOKEN) { - return process.env.GROK_CLI_OAUTH_TOKEN; - } try { - return await ctx.modelRegistry.getApiKeyForProvider(provider); + return (await resolveAccountRoute()).token; } catch { return undefined; } } function createAccountManager( - pi: ExtensionAPI, - registerAccount: RegisterAccount, - deferredRemovals: Set, + clearRecentExhaustion: (accountId: string) => void, + sessionSelection: SessionAccountSelection, ) { - let mutations = Promise.resolve(); - const accountGenerations = new Map(); - const mutate = (operation: () => Promise | T) => { - const result = mutations.then(operation, operation); - mutations = result.then( - () => undefined, - () => undefined, - ); - return result; - }; - const accountGeneration = (provider: string) => accountGenerations.get(provider) ?? 0; - const invalidateAccount = (provider: string) => - accountGenerations.set(provider, accountGeneration(provider) + 1); - const accountFrom = (config: GrokCliConfig, provider: string) => { - const account = config.accounts.items.find((candidate) => candidate.provider === provider); - if (!account) throw new Error(`Unknown Grok CLI account: ${provider}`); - return account; - }; - const clearQuota = async (provider: string) => { - try { - await removeQuotaUsage(provider); - return undefined; - } catch (error) { - return `Grok CLI quota cache cleanup failed: ${error instanceof Error ? error.message : String(error)}`; - } - }; - const refreshAccounts = async ( - ctx: ExtensionContext, - accounts: { account: GrokCliAccount; generation: number }[], - signal: AbortSignal, - onProgress?: (progress: { - provider: string; - completed: number; - total: number; - updated: boolean; - }) => void, - ) => { - let completed = 0; - let updated = 0; - const failed: string[] = []; - await refreshAccountBatches( - accounts, - async ({ account, generation }) => { - let succeeded = false; - try { - const token = await resolveAccountToken(ctx, account.provider); - if (!token) throw new Error('authentication unavailable'); - const usage = await fetchBillingUsage( - token, - AbortSignal.any([signal, AbortSignal.timeout(30_000)]), - ); - succeeded = await mutate(async () => { - if ( - accountGeneration(account.provider) !== generation || - !copyConfig().accounts.items.some( - (candidate) => candidate.provider === account.provider, - ) || - !hasAccountAuth(ctx, account.provider) - ) { - return false; - } - await saveQuotaUsage(account.provider, usage); - return true; - }); - if (succeeded) updated += 1; - } catch { - if (!signal.aborted) failed.push(account.provider); - } finally { - completed += 1; - onProgress?.({ - provider: account.provider, - completed, - total: accounts.length, - updated: succeeded, - }); - } - }, - signal, - ); - return { updated, failed }; - }; + function persistSessionFallback(ctx: ExtensionContext, removedId: string, selectedId?: string) { + if (selectedId !== removedId) return; + const fallback = sessionSelection.accountId(ctx.sessionManager.getSessionId()); + if (fallback) sessionSelection.select(ctx, fallback); + } return { - nextProvider() { - return `${GROK_CLI_PROVIDER}-${findAvailableAccountNumber( - copyConfig().accounts.items.map((account) => account.provider), - deferredRemovals, - )}`; + nextSlot() { + return getAccountVaultSync().nextSlot; }, add(_ctx: ExtensionContext, value: string) { - return mutate(() => { - const config = copyConfig(); - const provider = `${GROK_CLI_PROVIDER}-${findAvailableAccountNumber( - config.accounts.items.map((account) => account.provider), - deferredRemovals, - )}`; - const account = { provider, label: normalizeLabel(config, provider, value) }; - config.accounts.items.push(account); - config.accounts.nextAccountNumber = findAvailableAccountNumber( - config.accounts.items.map((candidate) => candidate.provider), - ); - saveConfig(config); - invalidateAccount(provider); - registerAccount(account); + return mutateAccountVault((vault) => { + const account = { + id: randomUUID(), + slot: vault.nextSlot, + label: normalizeLabel(vault.accounts, '', vault.nextSlot, value), + revision: 0, + }; + vault.nextSlot += 1; + vault.accounts.push(account); return account; }); }, - rename(_ctx: ExtensionContext, provider: string, value: string) { - return mutate(() => { - const config = copyConfig(); - const account = accountFrom(config, provider); - account.label = normalizeLabel(config, provider, value); - saveConfig(config); - registerAccount(account); - return account; + rename(_ctx: ExtensionContext, id: string, value: string) { + return mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === id); + if (!account) throw new Error(`Unknown Grok CLI account: ${id}`); + account.label = normalizeLabel(vault.accounts, id, account.slot, value); + return { id: account.id, slot: account.slot, label: account.label }; }); }, - activate(ctx: ExtensionContext, provider: string) { - return mutate(async () => { - const config = copyConfig(); - const account = accountFrom(config, provider); - if (!hasAccountAuth(ctx, provider)) { + async activate(_ctx: ExtensionContext, id: string) { + if (process.env.GROK_CLI_OAUTH_TOKEN) { + throw new Error('Saved accounts cannot be selected while the environment token is active.'); + } + const account = await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === id); + if (!account) throw new Error(`Unknown Grok CLI account: ${id}`); + if (!account.credential) { throw new Error(`Log in to “${account.label}” before making it active.`); } - const modelId = isGrokCliProvider(ctx.model?.provider) ? ctx.model?.id : DEFAULT_GROK_MODEL; - const model = - ctx.modelRegistry.find(provider, modelId ?? DEFAULT_GROK_MODEL) ?? - ctx.modelRegistry.find(provider, DEFAULT_GROK_MODEL); - if (!model) throw new Error(`Grok CLI model unavailable for “${account.label}”.`); - if (!(await pi.setModel(model))) { - throw new Error(`Could not switch to “${account.label}”; authentication is unavailable.`); - } - config.accounts.selectedProvider = provider; - saveConfig(config); - return account; + return { id: account.id, slot: account.slot, label: account.label }; }); + sessionSelection.select(_ctx, id); + return account; }, - async login(ctx: ExtensionContext, provider: string, interaction: AuthInteraction) { - const config = copyConfig(); - const account = accountFrom(config, provider); - if (provider === GROK_CLI_PROVIDER && process.env.GROK_CLI_OAUTH_TOKEN) { + async login(ctx: ExtensionContext, id: string, interaction: AuthInteraction) { + if (process.env.GROK_CLI_OAUTH_TOKEN) { throw new Error('Unset GROK_CLI_OAUTH_TOKEN before logging in from the dashboard.'); } - const result = await accountRuntime(ctx).login(account.provider, 'oauth', interaction); - invalidateAccount(provider); - return result; - }, - logout(ctx: ExtensionContext, provider: string) { - return mutate(async () => { - if (provider !== GROK_CLI_PROVIDER) { - throw new Error('Only the permanent base account can be logged out without removal.'); - } - if (process.env.GROK_CLI_OAUTH_TOKEN) { - throw new Error( - 'Unset GROK_CLI_OAUTH_TOKEN and restart Pi to remove the environment token.', - ); + const revision = accountFrom(id).revision; + const credential = (await oauth.login(callbacks(interaction))) as AccountCredential; + const account = await mutateAccountVault((vault) => { + const current = vault.accounts.find((candidate) => candidate.id === id); + if (!current) throw new Error('The account was removed while login was in progress.'); + if (current.revision !== revision) { + throw new Error('The account changed while login was in progress. Try again.'); } - const config = copyConfig(); - const account = accountFrom(config, provider); - await accountRuntime(ctx).logout(provider); - invalidateAccount(provider); - account.label = 'Account 1'; - saveConfig(config); - registerAccount(account); - return { warning: await clearQuota(provider) }; + current.credential = structuredClone(credential); + current.revision += 1; + vault.activeAccountId ??= current.id; + return { id: current.id, slot: current.slot, label: current.label }; }); + clearRecentExhaustion(id); + await clearQuota(ctx, id); + return account; }, - remove(ctx: ExtensionContext, provider: string) { - return mutate(async () => { - if (provider === GROK_CLI_PROVIDER) { - throw new Error('The permanent base account cannot be removed.'); - } - const config = copyConfig(); - const account = accountFrom(config, provider); - const fallback = await fallbackBeforeRemoval(pi, ctx, config, account); - if (fallback === null) { - throw new Error('Could not switch to another authenticated Grok CLI account.'); - } - await accountRuntime(ctx).logout(provider); - invalidateAccount(provider); - config.accounts.items = config.accounts.items.filter( - (candidate) => candidate.provider !== provider, + async logout(ctx: ExtensionContext, id: string) { + if (process.env.GROK_CLI_OAUTH_TOKEN) { + throw new Error( + 'Unset GROK_CLI_OAUTH_TOKEN and restart Pi to remove the environment token.', ); - config.accounts.selectedProvider = fallback ?? GROK_CLI_PROVIDER; - saveConfig(config); - const cacheWarning = await clearQuota(provider); - if (ctx.model?.provider === provider && fallback === undefined) { - deferredRemovals.add(provider); - return { - warning: - cacheWarning ?? - 'Grok CLI account removed. The current model now requires login and will disappear after you switch models.', - }; + } + const selectedAccountId = sessionSelection.accountId(ctx.sessionManager.getSessionId()); + const result = await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === id); + if (!account) throw new Error(`Unknown Grok CLI account: ${id}`); + delete account.credential; + account.revision += 1; + replaceVaultDefault(vault, id); + return logoutWarning(vault); + }); + persistSessionFallback(ctx, id, selectedAccountId); + await clearQuota(ctx, id); + return result; + }, + async remove(ctx: ExtensionContext, id: string, expectedRevision?: number) { + if (id === ACCOUNT_1_ID) throw new Error('The permanent Account 1 cannot be removed.'); + const selectedAccountId = sessionSelection.accountId(ctx.sessionManager.getSessionId()); + const result = await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === id); + if (!account) throw new Error(`Unknown Grok CLI account: ${id}`); + if ( + expectedRevision !== undefined && + (account.revision !== expectedRevision || account.credential !== undefined) + ) { + return undefined; } - pi.unregisterProvider(provider); - return { warning: cacheWarning }; + replaceVaultDefault(vault, id); + vault.accounts = vault.accounts.filter((account) => account.id !== id); + return logoutWarning(vault); }); + if (!result) return; + persistSessionFallback(ctx, id, selectedAccountId); + await clearQuota(ctx, id); + return result; }, async refresh( - ctx: ExtensionContext, + _ctx: ExtensionContext, signal: AbortSignal, onProgress?: (progress: { - provider: string; + id: string; completed: number; total: number; updated: boolean; }) => void, ) { - return refreshAccounts( - ctx, - copyConfig().accounts.items.flatMap((account) => - hasAccountAuth(ctx, account.provider) - ? [{ account, generation: accountGeneration(account.provider) }] - : [], - ), + const environment = Boolean(process.env.GROK_CLI_OAUTH_TOKEN); + const accounts = environment + ? getAccountVaultSync().accounts.filter((account) => account.id === ACCOUNT_1_ID) + : getAccountVaultSync().accounts.filter((account) => account.credential); + let completed = 0; + let updated = 0; + const failed: string[] = []; + await refreshAccountBatches( + accounts, + async (account) => { + let succeeded = false; + try { + const route = await resolveAccountRoute(environment ? undefined : account.id); + await refreshQuota(route, signal); + updated += 1; + succeeded = true; + } catch { + if (!signal.aborted) failed.push(account.id); + } finally { + completed += 1; + onProgress?.({ id: account.id, completed, total: accounts.length, updated: succeeded }); + } + }, signal, - onProgress, ); + return { updated, failed }; }, - async refreshOne(ctx: ExtensionContext, provider: string, signal: AbortSignal) { - const account = accountFrom(copyConfig(), provider); - if (!hasAccountAuth(ctx, provider)) { + async refreshOne(_ctx: ExtensionContext, id: string, signal: AbortSignal) { + const account = accountFrom(id); + if (!account.credential && !process.env.GROK_CLI_OAUTH_TOKEN) { throw new Error(`Log in to “${account.label}” before refreshing its quota.`); } - return refreshAccounts(ctx, [{ account, generation: accountGeneration(provider) }], signal); + try { + await refreshQuota(await resolveAccountRoute(id), signal); + return { updated: 1, failed: [] }; + } catch { + return { updated: 0, failed: [id] }; + } }, - snapshot(ctx: ExtensionContext): AccountsSnapshot { - const config = copyConfig(); + snapshot(_ctx: ExtensionContext): AccountsSnapshot { + const vault = getAccountVaultSync(); const cache = loadQuotaCache(); + const environment = Boolean(process.env.GROK_CLI_OAUTH_TOKEN); + const stored = readStoredCredential(GROK_CLI_PROVIDER); + const connected = + environment || + (stored?.type === 'oauth' && + stored.access === ACCOUNT_VAULT_MARKER && + stored.refresh === ACCOUNT_VAULT_MARKER); + const selectedAccountId = sessionSelection.accountId(_ctx.sessionManager.getSessionId()); return { - accounts: config.accounts.items.map((account) => { - const environment = - account.provider === GROK_CLI_PROVIDER && Boolean(process.env.GROK_CLI_OAUTH_TOKEN); - const authenticated = hasAccountAuth(ctx, account.provider); - const quota = cache.accounts[account.provider]; + connected, + accounts: vault.accounts.map((account, index) => { + const accountEnvironment = environment && account.id === ACCOUNT_1_ID; + const authenticated = accountEnvironment || account.credential !== undefined; + const active = accountEnvironment || (!environment && selectedAccountId === account.id); + const quota = cache.accounts[account.id]; + const slot = index + 1; return { - provider: account.provider, - label: account.label, - status: accountStatus(ctx, config, account.provider), + id: account.id, + slot, + label: + account.label === defaultLabel(account.slot) ? defaultLabel(slot) : account.label, + permanent: account.id === ACCOUNT_1_ID, + status: active + ? accountEnvironment + ? 'Active (environment)' + : 'Active' + : authenticated + ? 'Authenticated' + : 'Login required', authenticated, - active: config.accounts.selectedProvider === account.provider && authenticated, - environment, + active, + environment: accountEnvironment, ...(quota ? { plan: planTier(quota.monthly.monthlyLimit), @@ -758,37 +452,41 @@ function createAccountManager( }), }; }, - handleModelSelect(event: { - model: { provider: string }; - previousModel?: { provider: string }; - }) { - if (event.previousModel && deferredRemovals.delete(event.previousModel.provider)) { - pi.unregisterProvider(event.previousModel.provider); - } - if (!isGrokCliProvider(event.model.provider)) return; - return mutate(() => { - const config = copyConfig(); - if (!config.accounts.items.some((account) => account.provider === event.model.provider)) - return; - if (config.accounts.selectedProvider === event.model.provider) return; - config.accounts.selectedProvider = event.model.provider; - saveConfig(config); - }); - }, + handleModelSelect: noOp, }; } export type AccountManager = ReturnType; -export function registerAccountManagement(pi: ExtensionAPI, registerAccount: RegisterAccount) { - const deferredRemovals = new Set(); - const manager = createAccountManager(pi, registerAccount, deferredRemovals); +async function addAccount( + ctx: ExtensionCommandContext, + manager: AccountManager, + interaction: AuthInteraction, +) { + const slot = manager.nextSlot(); + const value = await ctx.ui.input('Label this Grok CLI account:', defaultLabel(slot)); + if (value === undefined) return; + const account = await manager.add(ctx, value); + try { + await manager.login(ctx, account.id, interaction); + } catch (error) { + await manager.remove(ctx, account.id); + throw error; + } +} + +export function registerAccountManagement( + pi: ExtensionAPI, + clearRecentExhaustion: (accountId: string) => void = () => {}, + sessionSelection = createSessionAccountSelection(pi), +) { + const manager = createAccountManager(clearRecentExhaustion, sessionSelection); const dashboard = createAccountDashboard(manager); pi.registerCommand('grok-cli-accounts', { - description: - 'Add, switch, rename, relogin, or remove Grok CLI accounts; pass gui to open the browser UI', + description: 'Add, select, rename, log in, log out, or remove Grok CLI accounts', handler: async (args, ctx) => { + await confirmMarkerInstallation(); const command = args.trim(); if (command === 'gui') { await dashboard.open(ctx); @@ -798,29 +496,62 @@ export function registerAccountManagement(pi: ExtensionAPI, registerAccount: Reg ctx.ui.notify('Usage: /grok-cli-accounts [gui]', 'warning'); return; } - const config = copyConfig(); - const choice = await showAccountSelector(ctx, config, 'main', manager); - if (choice === 'action:add') { - await addAccount(ctx, config, manager); + const snapshot = manager.snapshot(ctx); + const choice = await ctx.ui.select('Grok CLI accounts:', [ + ...snapshot.accounts.map( + (account) => `${account.slot}. ${account.label} — ${account.status}`, + ), + 'Add account', + 'Manage accounts', + ]); + if (!choice) return; + const interaction = terminalInteraction(ctx); + if (choice === 'Add account') { + try { + await addAccount(ctx, manager, interaction); + } catch (error) { + ctx.ui.notify(error instanceof Error ? error.message : String(error), 'error'); + } return; } - if (choice === 'action:manage') { - const latest = copyConfig(); - const selected = await showAccountSelector(ctx, latest, 'manage', manager); - const account = latest.accounts.items.find(({ provider }) => provider === selected); - if (account) { - await manageAccount(ctx, latest, account, manager); + if (choice === 'Manage accounts') { + const selected = await ctx.ui.select( + 'Manage Grok CLI account:', + snapshot.accounts.map((account) => `${account.slot}. ${account.label}`), + ); + const account = snapshot.accounts.find( + (candidate) => `${candidate.slot}. ${candidate.label}` === selected, + ); + if (!account) return; + const action = await ctx.ui.select(`Manage “${account.label}”:`, [ + 'Rename', + account.authenticated ? 'Log in again' : 'Log in', + 'Log out', + ...(account.permanent ? [] : ['Remove']), + 'Back', + ]); + try { + if (action === 'Rename') { + const label = await ctx.ui.input(`Rename “${account.label}”:`, account.label); + if (label !== undefined) await manager.rename(ctx, account.id, label); + } + if (action === 'Log in' || action === 'Log in again') { + await manager.login(ctx, account.id, interaction); + } + if (action === 'Log out') await manager.logout(ctx, account.id); + if (action === 'Remove') await manager.remove(ctx, account.id); + } catch (error) { + ctx.ui.notify(error instanceof Error ? error.message : String(error), 'error'); } return; } - const account = config.accounts.items.find(({ provider }) => provider === choice); + const account = snapshot.accounts.find( + (candidate) => `${candidate.slot}. ${candidate.label} — ${candidate.status}` === choice, + ); if (!account) return; - if (!hasAccountAuth(ctx, account.provider)) { - prefillLogin(ctx, account.provider); - return; - } try { - await manager.activate(ctx, account.provider); + if (!account.authenticated) await manager.login(ctx, account.id, interaction); + else await manager.activate(ctx, account.id); } catch (error) { ctx.ui.notify(error instanceof Error ? error.message : String(error), 'error'); } diff --git a/src/provider/dashboard/app.css b/src/provider/dashboard/app.css index abd94ed..73c7768 100644 --- a/src/provider/dashboard/app.css +++ b/src/provider/dashboard/app.css @@ -418,6 +418,7 @@ body:has(.link-pill.error) main { position: absolute; inset: -1px; border-radius: inherit; + corner-shape: inherit; padding: 1px; background: radial-gradient( 200px circle at var(--mx, 50%) var(--my, 50%), diff --git a/src/provider/dashboard/app.js b/src/provider/dashboard/app.js index 4a9db50..96bc3c5 100644 --- a/src/provider/dashboard/app.js +++ b/src/provider/dashboard/app.js @@ -21,7 +21,7 @@ let lastState = ''; let wasOffline = false; let entranceDone = false; let timer; -let pendingProviders = new Set(); +let pendingAccountIds = new Set(); let lastProgress = ''; let quotaRefreshInFlight = false; @@ -448,7 +448,7 @@ const animateGauge = (gauge, from, to) => { } }; -const quotaRow = (provider, label, usedLabel, metaText, remaining) => { +const quotaRow = (accountId, label, usedLabel, metaText, remaining) => { const row = element('div', 'quota'); const gauge = element( 'div', @@ -461,7 +461,7 @@ const quotaRow = (provider, label, usedLabel, metaText, remaining) => { gauge.setAttribute('aria-valuemax', '100'); gauge.setAttribute('aria-valuenow', String(Math.round(remaining))); gauge.append(element('span', 'gauge-value', `${Math.round(remaining)}%`)); - const key = `${provider}:${label}`; + const key = `${accountId}:${label}`; const previous = gaugeMemory.get(key); gaugeMemory.set(key, remaining); if (!entranceDone) animateGauge(gauge, 0, remaining); @@ -520,12 +520,12 @@ const actionButton = (label, action, kind = 'ghost') => { return button; }; -const startLogin = async (provider) => { +const startLogin = async (accountId) => { // Open the popup with the final URL instead of scripting a blank one: embedded // browsers (e.g. WKWebView) hand window.open('') an unusable about:blank view. try { - const ticket = await mutation(`/api/accounts/${provider}/login-ticket`, 'POST'); - if (!window.open(ticket.path, `grok-login-${provider}`)) { + const ticket = await mutation(`/api/accounts/${accountId}/login-ticket`, 'POST'); + if (!window.open(ticket.path, `grok-login-${accountId}`)) { showToast('Pop-up blocked. Allow pop-ups, then use Log in on the account card.', true); return; } @@ -553,7 +553,7 @@ const loginPanel = (account, isNew) => { cancel.dataset.action = 'Cancel login'; cancel.addEventListener('click', async () => { try { - await mutation(`/api/accounts/${account.provider}/login-cancel`, 'POST'); + await mutation(`/api/accounts/${account.id}/login-cancel`, 'POST'); await refreshState(true, true); } catch (error) { showToast(error.message, true); @@ -564,7 +564,7 @@ const loginPanel = (account, isNew) => { event.preventDefault(); if (!input.value.trim()) return; try { - await mutation(`/api/accounts/${account.provider}/login-code`, 'POST', { + await mutation(`/api/accounts/${account.id}/login-code`, 'POST', { code: input.value, }); input.value = ''; @@ -585,11 +585,11 @@ const loginPanel = (account, isNew) => { return panel; }; -const cardActions = (account) => { +const cardActions = (account, environmentMode) => { const actions = element('footer', 'card-actions'); const activate = async () => { try { - await mutation(`/api/accounts/${account.provider}/activate`, 'POST'); + await mutation(`/api/accounts/${account.id}/activate`, 'POST'); await refreshState(true, true); showToast(`Switched to ${account.label}.`); } catch (error) { @@ -605,7 +605,7 @@ const cardActions = (account) => { }); if (label === undefined) return; try { - const updated = await mutation(`/api/accounts/${account.provider}`, 'PATCH', { label }); + const updated = await mutation(`/api/accounts/${account.id}`, 'PATCH', { label }); await refreshState(true); showToast(`Renamed to ${updated.label}.`); } catch (error) { @@ -623,49 +623,45 @@ const cardActions = (account) => { }; const destructive = async () => { const confirmed = await modal({ - title: - account.provider === 'grok-cli' ? `Log out ${account.label}?` : `Remove ${account.label}?`, - message: - account.provider === 'grok-cli' - ? 'Removes the saved login. The account stays in the list — log in again to use it.' - : 'Removes this account and its saved login. You can add it again with Add account.', - confirm: account.provider === 'grok-cli' ? 'Log out' : 'Remove account', + title: account.permanent ? `Log out ${account.label}?` : `Remove ${account.label}?`, + message: account.permanent + ? 'Removes this account login from the vault. The account stays in the list. If no logged-in account remains, also run /logout in Pi.' + : 'Removes this account and its saved login. You can add it again with Add account.', + confirm: account.permanent ? 'Log out' : 'Remove account', danger: true, }); if (!confirmed) return; try { - await mutation( - account.provider === 'grok-cli' - ? '/api/accounts/grok-cli/logout' - : `/api/accounts/${account.provider}`, - account.provider === 'grok-cli' ? 'POST' : 'DELETE', + const result = await mutation( + account.permanent ? `/api/accounts/${account.id}/logout` : `/api/accounts/${account.id}`, + account.permanent ? 'POST' : 'DELETE', ); await refreshState(true, true); showToast( - account.provider === 'grok-cli' - ? `Logged out ${account.label}.` - : `Removed ${account.label}.`, + `${ + account.permanent ? `Logged out ${account.label}.` : `Removed ${account.label}.` + }${result.warning ? ` ${result.warning}` : ''}`, ); } catch (error) { showToast(error.message, true); } }; - if (!account.authenticated && !account.environment) { - actions.append(actionButton('Log in', () => startLogin(account.provider), 'primary')); - } else if (!account.active && account.authenticated) { + if (!account.authenticated && !account.environment && !environmentMode) { + actions.append(actionButton('Log in', () => startLogin(account.id), 'primary')); + } else if (!account.active && account.authenticated && !environmentMode) { actions.append(actionButton('Switch', activate, 'primary')); } - if (!account.environment && account.authenticated) { - actions.append(actionButton('Log in again', () => startLogin(account.provider))); + if (!account.environment && account.authenticated && !environmentMode) { + actions.append(actionButton('Log in again', () => startLogin(account.id))); } if (account.environment) { actions.append(actionButton('How to remove', tokenInstructions)); } actions.append(actionButton('Rename', rename)); - if (!account.environment) { + if (!account.environment && (!account.permanent || account.authenticated)) { const button = actionButton( - account.provider === 'grok-cli' ? 'Log out' : 'Remove', + account.permanent ? 'Log out' : 'Remove', destructive, 'danger push-right', ); @@ -674,10 +670,10 @@ const cardActions = (account) => { return actions; }; -const accountCard = (account, index, isNewPending, refreshing) => { +const accountCard = (account, index, isNewPending, refreshing, environmentMode) => { const card = element('article', `account-card${account.active ? ' active' : ''}`); - card.dataset.provider = account.provider; - card.style.viewTransitionName = `card-${account.provider}`; + card.dataset.accountId = account.id; + card.style.viewTransitionName = `card-${account.id}`; if (!entranceDone) { card.style.setProperty('--enter-delay', `${Math.min(index * 45, 220)}ms`); } @@ -687,7 +683,10 @@ const accountCard = (account, index, isNewPending, refreshing) => { titleRow.append(element('h2', '', account.label)); if (account.active) titleRow.append(element('span', 'active-badge', 'Active')); const metaRow = element('div', 'card-meta-row'); - metaRow.append(element('span', 'card-provider', account.provider), statusPill(account)); + metaRow.append( + element('span', 'card-provider', `grok-cli · Account ${account.slot}`), + statusPill(account), + ); if (account.plan) metaRow.append(element('span', 'plan-pill', PLAN_LABELS[account.plan])); head.append(titleRow, metaRow); @@ -703,7 +702,7 @@ const accountCard = (account, index, isNewPending, refreshing) => { const monthly = account.quota.monthly; body.append( quotaRow( - account.provider, + account.id, 'Monthly credits', isFree ? '' : `${Math.max(0, monthly.monthlyLimit - monthly.used).toLocaleString()} left`, isFree ? 'Not available' : `Resets ${dateLabel(monthly.billingPeriodEnd)}`, @@ -713,7 +712,7 @@ const accountCard = (account, index, isNewPending, refreshing) => { body.append( account.quota.weekly || isFree ? quotaRow( - account.provider, + account.id, 'Weekly credits', '', !isFree && account.quota.weekly @@ -750,14 +749,15 @@ const accountCard = (account, index, isNewPending, refreshing) => { body.append(element('p', 'card-error', errorText)); } - card.append(head, body, cardActions(account)); + card.append(head, body, cardActions(account, environmentMode)); return card; }; const render = (state) => { + const environmentMode = state.accounts.some((account) => account.environment); statsSummary.textContent = `${state.accounts.length} account${state.accounts.length === 1 ? '' : 's'}`; - linkState.className = 'link-pill ok'; - linkText.textContent = 'Synced'; + linkState.className = `link-pill ${state.connected ? 'ok' : 'error'}`; + linkText.textContent = state.connected ? 'Pi connected' : 'Pi disconnected'; refreshButton.disabled = state.refreshing; accountsRoot.setAttribute('aria-busy', String(state.refreshing)); accountsRoot.classList.toggle('refreshing', state.refreshing || quotaRefreshInFlight); @@ -766,25 +766,25 @@ const render = (state) => { const refocus = active instanceof HTMLElement && accountsRoot.contains(active) && active.dataset.action ? { - provider: active.closest('[data-provider]')?.dataset.provider, + accountId: active.closest('[data-account-id]')?.dataset.accountId, action: active.dataset.action, } : undefined; const codes = new Map( - [...accountsRoot.querySelectorAll('[data-provider] input[name="code"]')] - .map((input) => [input.closest('[data-provider]').dataset.provider, input.value]) + [...accountsRoot.querySelectorAll('[data-account-id] input[name="code"]')] + .map((input) => [input.closest('[data-account-id]').dataset.accountId, input.value]) .filter(([, value]) => value), ); if (entranceDone) accountsRoot.classList.add('settled'); const nextPending = new Set( state.accounts .filter((account) => account.login.state === 'pending') - .map((account) => account.provider), + .map((account) => account.id), ); // A login that left pending since the last render resolves audibly, not only visually. - for (const provider of pendingProviders) { - if (nextPending.has(provider)) continue; - const account = state.accounts.find((candidate) => candidate.provider === provider); + for (const accountId of pendingAccountIds) { + if (nextPending.has(accountId)) continue; + const account = state.accounts.find((candidate) => candidate.id === accountId); if (account?.login.state === 'success') showToast(`Logged in ${account.label}.`); if (account?.login.state === 'failed') showToast(account.login.error || 'Login failed.', true); if (account?.login.quotaError) showToast(account.login.quotaError, true); @@ -804,21 +804,22 @@ const render = (state) => { accountCard( account, index, - nextPending.has(account.provider) && !pendingProviders.has(account.provider), + nextPending.has(account.id) && !pendingAccountIds.has(account.id), state.refreshing || quotaRefreshInFlight, + environmentMode, ), ) : [element('p', 'grid-message', 'No accounts configured. Use Add account to connect one.')]; accountsRoot.replaceChildren(...children); - pendingProviders = nextPending; + pendingAccountIds = nextPending; entranceDone = true; - for (const [provider, value] of codes) { - const input = accountsRoot.querySelector(`[data-provider="${provider}"] input[name="code"]`); + for (const [accountId, value] of codes) { + const input = accountsRoot.querySelector(`[data-account-id="${accountId}"] input[name="code"]`); if (input) input.value = value; } - if (refocus?.provider) { + if (refocus?.accountId) { accountsRoot - .querySelector(`[data-provider="${refocus.provider}"] [data-action="${refocus.action}"]`) + .querySelector(`[data-account-id="${refocus.accountId}"] [data-action="${refocus.action}"]`) ?.focus(); } }; @@ -899,7 +900,7 @@ addAccount.addEventListener('click', async () => { try { const account = await mutation('/api/accounts', 'POST', { label }); await refreshState(true, true); - await startLogin(account.provider); + await startLogin(account.id); } catch (error) { showToast(error.message, true); } diff --git a/src/provider/dashboard/server.ts b/src/provider/dashboard/server.ts index e47970b..89a1adc 100644 --- a/src/provider/dashboard/server.ts +++ b/src/provider/dashboard/server.ts @@ -150,10 +150,8 @@ function escapeHtml(value: string) { return value.replace(/&/g, '&').replace(//g, '>'); } -function accountProvider(pathname: string, suffix = '') { - const match = new RegExp(`^/api/accounts/(grok-cli(?:-(?:[2-9]|[1-9][0-9]+))?)${suffix}$`).exec( - pathname, - ); +function accountId(pathname: string, suffix = '') { + const match = new RegExp(`^/api/accounts/([A-Za-z0-9-]+)${suffix}$`).exec(pathname); return match?.[1]; } @@ -240,8 +238,9 @@ export async function startAccountDashboard( let idleTimer: ReturnType | undefined; let closing: Promise | undefined; let refreshController: AbortController | undefined; - const loginTickets = new Map(); + const loginTickets = new Map(); const loginJobs = new Map(); + const newAccountRevisions = new Map(); const server = createServer((req, res) => { void handle(req, res).catch((error: unknown) => { @@ -302,8 +301,9 @@ export async function startAccountDashboard( const snapshot = manager.snapshot(ctx); return { refreshing: Boolean(refreshController), + connected: snapshot.connected, accounts: snapshot.accounts.map((account) => { - const job = loginJobs.get(account.provider); + const job = loginJobs.get(account.id); return { ...account, login: job @@ -319,10 +319,8 @@ export async function startAccountDashboard( }; }; - const startLogin = async (provider: string, res: ServerResponse) => { - const account = manager - .snapshot(ctx) - .accounts.find((candidate) => candidate.provider === provider); + const startLogin = async (id: string, res: ServerResponse) => { + const account = manager.snapshot(ctx).accounts.find((candidate) => candidate.id === id); if (!account) throw new HttpError(404, 'Account not found — it may have been removed.'); if (account.environment) { throw new HttpError( @@ -330,9 +328,10 @@ export async function startAccountDashboard( 'This account logs in with the GROK_CLI_OAUTH_TOKEN environment variable.', ); } - if (loginJobs.get(provider)?.state === 'pending') { + if (loginJobs.get(id)?.state === 'pending') { throw new HttpError(409, 'A login is already in progress for this account.'); } + const removeOnFailureRevision = newAccountRevisions.get(id); const controller = new AbortController(); let resolveManualCode = (_code: string) => {}; const manualCode = new Promise((resolve) => { @@ -344,7 +343,7 @@ export async function startAccountDashboard( progress: 'Waiting for xAI authorization', resolveManualCode, }; - loginJobs.set(provider, job); + loginJobs.set(id, job); let settledRedirect = false; let finishRedirect = () => {}; const redirected = new Promise((resolve) => { @@ -372,23 +371,28 @@ export async function startAccountDashboard( throw new Error('Interactive OAuth prompts are not supported in the dashboard.'); }, }; - void manager.login(ctx, provider, interaction).then( + void manager.login(ctx, id, interaction).then( async () => { + newAccountRevisions.delete(id); job.state = 'success'; job.progress = 'Login complete'; job.resolveManualCode(''); if (options.refreshAfterLogin === false) return; try { - const result = await manager.refreshOne(ctx, provider, controller.signal); + const result = await manager.refreshOne(ctx, id, controller.signal); if (result.failed.length) job.quotaError = 'Login succeeded, but quota refresh failed.'; } catch { job.quotaError = 'Login succeeded, but quota refresh failed.'; } }, - (error: unknown) => { + async (error: unknown) => { job.state = controller.signal.aborted ? 'cancelled' : 'failed'; job.error = controller.signal.aborted ? 'Login cancelled.' : publicError(error); job.resolveManualCode(''); + if (removeOnFailureRevision !== undefined && !controller.signal.aborted) { + newAccountRevisions.delete(id); + await manager.remove(ctx, id, removeOnFailureRevision).catch(() => undefined); + } if (!settledRedirect) { send( res, @@ -446,11 +450,13 @@ export async function startAccountDashboard( if (body.label !== undefined && typeof body.label !== 'string') { throw new HttpError(400, 'Account label must be text.'); } - json(res, 201, await manager.add(ctx, body.label ?? '')); + const account = await manager.add(ctx, body.label ?? ''); + newAccountRevisions.set(account.id, account.revision); + json(res, 201, account); return; } - const renameProvider = accountProvider(url.pathname); - if (req.method === 'PATCH' && renameProvider) { + const renameId = accountId(url.pathname); + if (req.method === 'PATCH' && renameId) { requireMutation(req); const body = objectBody( await readJson(req, options.bodyTimeoutMs ?? DEFAULT_BODY_TIMEOUT_MS), @@ -458,24 +464,25 @@ export async function startAccountDashboard( if (typeof body.label !== 'string') { throw new HttpError(400, 'Account label must be text.'); } - json(res, 200, await manager.rename(ctx, renameProvider, body.label)); + json(res, 200, await manager.rename(ctx, renameId, body.label)); return; } - const activateProvider = accountProvider(url.pathname, '/activate'); - if (req.method === 'POST' && activateProvider) { + const activateId = accountId(url.pathname, '/activate'); + if (req.method === 'POST' && activateId) { requireMutation(req); - json(res, 200, await manager.activate(ctx, activateProvider)); + json(res, 200, await manager.activate(ctx, activateId)); return; } - const logoutProvider = accountProvider(url.pathname, '/logout'); - if (req.method === 'POST' && logoutProvider) { + const logoutId = accountId(url.pathname, '/logout'); + if (req.method === 'POST' && logoutId) { requireMutation(req); - json(res, 200, await manager.logout(ctx, logoutProvider)); + json(res, 200, await manager.logout(ctx, logoutId)); return; } - if (req.method === 'DELETE' && renameProvider) { + if (req.method === 'DELETE' && renameId) { requireMutation(req); - json(res, 200, await manager.remove(ctx, renameProvider)); + newAccountRevisions.delete(renameId); + json(res, 200, await manager.remove(ctx, renameId)); return; } if (req.method === 'POST' && url.pathname === '/api/quotas/refresh') { @@ -489,12 +496,10 @@ export async function startAccountDashboard( } return; } - const ticketProvider = accountProvider(url.pathname, '/login-ticket'); - if (req.method === 'POST' && ticketProvider) { + const ticketId = accountId(url.pathname, '/login-ticket'); + if (req.method === 'POST' && ticketId) { requireMutation(req); - const account = manager - .snapshot(ctx) - .accounts.find((candidate) => candidate.provider === ticketProvider); + const account = manager.snapshot(ctx).accounts.find((candidate) => candidate.id === ticketId); if (!account) throw new HttpError(404, 'Account not found — it may have been removed.'); if (account.environment) { throw new HttpError( @@ -503,7 +508,7 @@ export async function startAccountDashboard( ); } const ticket = randomBytes(24).toString('base64url'); - loginTickets.set(ticket, { provider: ticketProvider, expiresAt: Date.now() + 60_000 }); + loginTickets.set(ticket, { accountId: ticketId, expiresAt: Date.now() + 60_000 }); json(res, 201, { path: `/oauth/${ticket}` }); return; } @@ -514,11 +519,11 @@ export async function startAccountDashboard( if (!loginTicket || loginTicket.expiresAt < Date.now()) { throw new HttpError(404, 'Login link expired — start again from the dashboard.'); } - await startLogin(loginTicket.provider, res); + await startLogin(loginTicket.accountId, res); return; } - const codeProvider = accountProvider(url.pathname, '/login-code'); - if (req.method === 'POST' && codeProvider) { + const codeId = accountId(url.pathname, '/login-code'); + if (req.method === 'POST' && codeId) { requireMutation(req); const body = objectBody( await readJson(req, options.bodyTimeoutMs ?? DEFAULT_BODY_TIMEOUT_MS), @@ -526,16 +531,16 @@ export async function startAccountDashboard( if (typeof body.code !== 'string' || !body.code.trim()) { throw new HttpError(400, 'Authorization code is required.'); } - const job = loginJobs.get(codeProvider); + const job = loginJobs.get(codeId); if (job?.state !== 'pending') throw new HttpError(409, 'No login is waiting for a code.'); job.resolveManualCode(body.code.trim()); json(res, 202, { accepted: true }); return; } - const cancelProvider = accountProvider(url.pathname, '/login-cancel'); - if (req.method === 'POST' && cancelProvider) { + const cancelId = accountId(url.pathname, '/login-cancel'); + if (req.method === 'POST' && cancelId) { requireMutation(req); - const job = loginJobs.get(cancelProvider); + const job = loginJobs.get(cancelId); if (job?.state !== 'pending') throw new HttpError(409, 'No login is running.'); job.controller.abort(); job.resolveManualCode(''); diff --git a/src/provider/modelMigration.ts b/src/provider/modelMigration.ts new file mode 100644 index 0000000..ac8f9f0 --- /dev/null +++ b/src/provider/modelMigration.ts @@ -0,0 +1,57 @@ +import { existsSync } from 'node:fs'; +import { join } from 'node:path'; +import { getAgentDir, SettingsManager } from '@earendil-works/pi-coding-agent'; + +const LEGACY_PROVIDER = /^grok-cli-(?:[2-9]|[1-9]\d+)$/; +const LEGACY_MODEL_PATTERN = /^grok-cli-(?:[2-9]|[1-9]\d+)(?=\/)/; + +async function migrateSettings(directory: string) { + if (!existsSync(join(directory, 'settings.json'))) return false; + const settings = SettingsManager.create(join(directory, '.pi-grok-cli-migration'), directory); + const current = settings.getGlobalSettings(); + const defaultProvider = + typeof current.defaultProvider === 'string' && LEGACY_PROVIDER.test(current.defaultProvider) + ? 'grok-cli' + : current.defaultProvider; + const enabledModels = Array.isArray(current.enabledModels) + ? [ + ...new Set( + current.enabledModels.map((pattern) => pattern.replace(LEGACY_MODEL_PATTERN, 'grok-cli')), + ), + ] + : current.enabledModels; + if ( + defaultProvider === current.defaultProvider && + JSON.stringify(enabledModels) === JSON.stringify(current.enabledModels) + ) { + return false; + } + if (defaultProvider !== current.defaultProvider && defaultProvider) { + settings.setDefaultProvider(defaultProvider); + } + if (JSON.stringify(enabledModels) !== JSON.stringify(current.enabledModels)) { + settings.setEnabledModels(enabledModels); + } + await settings.flush(); + const errors = settings.drainErrors(); + if (errors[0]) throw errors[0].error; + return true; +} + +export async function migrateSavedModelProviders(agentDir = getAgentDir(), cwd = process.cwd()) { + try { + const settingsChanged = ( + await Promise.all( + [...new Set([agentDir, join(cwd, '.pi')])].map((directory) => migrateSettings(directory)), + ) + ).some(Boolean); + return { migrated: settingsChanged }; + } catch (error) { + return { + migrated: false, + warning: `Could not migrate saved Grok model providers: ${ + error instanceof Error ? error.message : String(error) + }.`, + }; + } +} diff --git a/src/provider/quotaCache.ts b/src/provider/quotaCache.ts index 9b67af5..bccad27 100644 --- a/src/provider/quotaCache.ts +++ b/src/provider/quotaCache.ts @@ -1,5 +1,6 @@ -import { existsSync, readFileSync } from 'node:fs'; -import { getQuotaCachePath, writeFileAtomic } from '../storage.js'; +import { chmodSync, existsSync, mkdirSync, readFileSync } from 'node:fs'; +import { dirname } from 'node:path'; +import { acquireFileLock, getQuotaCachePath, writeFileAtomic } from '../storage.js'; import type { BillingUsage, MonthlyUsage, WeeklyUsage } from './billing.js'; export interface CachedQuota extends BillingUsage { @@ -80,51 +81,65 @@ export function loadQuotaCache(path = getQuotaCachePath()): QuotaCache { } } -const cacheUpdates = new Map>(); - async function updateQuotaCache( update: (cache: QuotaCache) => boolean, path = getQuotaCachePath(), ) { - const previous = cacheUpdates.get(path) ?? Promise.resolve(); - const operation = () => { + mkdirSync(dirname(path), { recursive: true, mode: 0o700 }); + chmodSync(dirname(path), 0o700); + const release = await acquireFileLock(path); + try { const cache = loadQuotaCache(path); - if (!update(cache)) return; + if (!update(cache)) return false; writeFileAtomic(path, `${JSON.stringify(cache, null, 2)}\n`, 0o600); - }; - const next = previous.then(operation, operation); - cacheUpdates.set(path, next); - try { - await next; + chmodSync(path, 0o600); + return true; } finally { - if (cacheUpdates.get(path) === next) cacheUpdates.delete(path); + await release(); } } +function storeQuota(cache: QuotaCache, accountId: string, usage: BillingUsage, updatedAt: string) { + cache.accounts[accountId] = { + updatedAt, + monthly: { ...usage.monthly }, + ...(usage.weekly ? { weekly: { ...usage.weekly } } : {}), + }; +} + export function saveQuotaUsage( - provider: string, + accountId: string, usage: BillingUsage, updatedAt = new Date().toISOString(), path = getQuotaCachePath(), ) { return updateQuotaCache((cache) => { - cache.accounts[provider] = { - updatedAt, - monthly: { ...usage.monthly }, - ...(usage.weekly ? { weekly: { ...usage.weekly } } : {}), - }; + storeQuota(cache, accountId, usage, updatedAt); return true; - }, path); + }, path).then(() => undefined); } -export function removeQuotaUsage(provider: string, path = getQuotaCachePath()) { +export function saveQuotaUsageWhen( + accountId: string, + usage: BillingUsage, + shouldSave: () => boolean, + path = getQuotaCachePath(), +) { return updateQuotaCache((cache) => { - if (!existsSync(path) || !cache.accounts[provider]) return false; - delete cache.accounts[provider]; + if (!shouldSave()) return false; + storeQuota(cache, accountId, usage, new Date().toISOString()); return true; }, path); } +export function removeQuotaUsage(accountId: string, path = getQuotaCachePath()) { + return updateQuotaCache((cache) => { + if (!existsSync(path) || !cache.accounts[accountId]) return false; + delete cache.accounts[accountId]; + return true; + }, path).then(() => undefined); +} + function formatAge(updatedAt: string, now: number) { const age = Math.max(0, now - new Date(updatedAt).getTime()); if (age < 60_000) return 'just now'; diff --git a/src/provider/register.ts b/src/provider/register.ts index e85fe13..a725552 100644 --- a/src/provider/register.ts +++ b/src/provider/register.ts @@ -1,143 +1,254 @@ -import type { ExtensionAPI, ProviderConfig } from '@earendil-works/pi-coding-agent'; +import { + lazyStream, + type Model, + type OAuthCredentials, + type SimpleStreamOptions, +} from '@earendil-works/pi-ai'; +import { streamSimpleOpenAIResponses } from '@earendil-works/pi-ai/compat'; +import { + type ExtensionAPI, + type ExtensionContext, + type ProviderConfig, + readStoredCredential, +} from '@earendil-works/pi-coding-agent'; import * as oauth from '../auth/oauth.js'; -import { getBaseUrl, type XaiOAuthCredentials } from '../auth/oauth.js'; -import { type GrokCliAccount, loadConfig, migrateLegacyConfig } from '../config.js'; -import { registerImagineFeature } from '../imagine/register.js'; +import { getBaseUrl } from '../auth/oauth.js'; +import { migrateLegacyConfig } from '../config.js'; +import { registerImagineFeature, syncImageToolPreference } from '../imagine/register.js'; import { resolveModels } from '../models/catalog.js'; import { sanitizePayload } from '../payload/sanitize.js'; -import { registerGrokTools } from '../tools/register.js'; -import { bindLivePiWebAccess, ensureWebSearchDelegate } from '../tools/webSearchDelegate.js'; -import { registerVisionFeature } from '../vision/register.js'; -import { isGrokCliProvider, registerAccountManagement } from './accounts.js'; +import { confirmMarkerInstallation, migrateReleasedAccounts } from './accountMigration.js'; +import { resolveAccountRoute } from './accountRouting.js'; +import { registerAccountManagement } from './accounts.js'; +import { + ACCOUNT_1_ID, + ACCOUNT_VAULT_MARKER, + type AccountCredential, + getAccountVault, + mutateAccountVault, +} from './accountVault.js'; +import { migrateSavedModelProviders } from './modelMigration.js'; import { removeQuotaUsage } from './quotaCache.js'; +import { rememberRequestAccount } from './requestOwnership.js'; import { registerExhaustionRotation } from './rotation.js'; +import { createSessionAccountSelection } from './sessionAccountSelection.js'; import { grokCliModelHeaders } from './stream.js'; -import { handoffGrokTools, restoreGrokTools, syncGrokTools } from './toolScope.js'; import { registerUsageCommand } from './usage.js'; -export default function registerGrokCli(pi: ExtensionAPI) { - const migration = migrateLegacyConfig(); - const baseUrl = getBaseUrl(); - const models = resolveModels(); - const exhaustionRotation = registerExhaustionRotation(pi); - - const registerAccount = (account: GrokCliAccount) => { - const oauthProvider = { - name: `Grok CLI — ${account.label}`, - usesCallbackServer: true, - - async login(callbacks) { - const credentials = await oauth.login(callbacks); - await removeQuotaUsage(account.provider).catch(() => undefined); - exhaustionRotation.clearRecentExhaustion(account.provider); - return credentials; - }, - - async refreshToken(credentials) { - return oauth.refresh(credentials); - }, - - getApiKey(credentials) { - return credentials.access; - }, - - modifyModels(registeredModels, credentials) { - const effectiveBaseUrl = String( - (credentials as XaiOAuthCredentials).baseUrl ?? getBaseUrl(), - ).replace(/\/+$/, ''); - - return registeredModels.map((model) => - model.provider === account.provider ? { ...model, baseUrl: effectiveBaseUrl } : model, - ); - }, - } satisfies NonNullable; - - pi.registerProvider(account.provider, { - name: `Grok CLI — ${account.label}`, - baseUrl, - ...(account.provider === 'grok-cli' ? { apiKey: '$GROK_CLI_OAUTH_TOKEN' } : {}), - api: 'openai-responses', - models: models.map((model) => ({ - id: model.id, - name: model.name, - reasoning: model.reasoning, - thinkingLevelMap: model.thinkingLevelMap, - input: model.input, - cost: model.cost, - contextWindow: model.contextWindow, - maxTokens: model.maxTokens, - // Carried as model.headers so the version-gate headers reach the server on - // every request even when the API-provider registry reverts to pi-ai's - // built-in openai-responses handler (see grokCliModelHeaders). - headers: grokCliModelHeaders(model.id), - })), - oauth: oauthProvider, - }); - }; +const marker = (): OAuthCredentials => ({ + access: ACCOUNT_VAULT_MARKER, + refresh: ACCOUNT_VAULT_MARKER, + expires: Number.MAX_SAFE_INTEGER, +}); - loadConfig().config.accounts.items.forEach(registerAccount); - const accountManagement = registerAccountManagement(pi, registerAccount); +function isMarker(credentials: OAuthCredentials) { + return ( + credentials.access === ACCOUNT_VAULT_MARKER && credentials.refresh === ACCOUNT_VAULT_MARKER + ); +} - const { webSearchRegistered } = registerGrokTools(pi); - registerImagineFeature(pi); +function accountCredential(credentials: OAuthCredentials): AccountCredential { + return structuredClone(credentials) as AccountCredential; +} - const syncTools = (model: { provider: string; id: string } | undefined, captureDelete = false) => - syncGrokTools(pi, model, { captureDelete, webSearchRegistered }); +export default function registerGrokCli(pi: ExtensionAPI) { + const sessionSelection = createSessionAccountSelection(pi); + let migrationWarning: string | undefined; + let migrationError: string | undefined; + let migrationErrorNotified = false; + let migrationNotice: string | undefined; + const migration = Promise.all([migrateReleasedAccounts(), migrateSavedModelProviders()]) + .then(async ([result, modelMigration]) => { + const legacyMigration = migrateLegacyConfig(); + migrationWarning = + [legacyMigration.warning, modelMigration.warning].filter(Boolean).join(' ') || undefined; + await confirmMarkerInstallation(); + const vault = await getAccountVault(); + if (result.migrated && vault.migration.markerInstallPending) { + migrationNotice = + 'Grok CLI accounts were copied to the secure account vault. Run /login and select Grok CLI once to finish setup. Old grok-cli-N login entries can then be removed with /logout.'; + } + }) + .catch((error: unknown) => { + migrationError = `Could not migrate Grok CLI accounts: ${ + error instanceof Error ? error.message : String(error) + }`; + }); + const exhaustionRotation = registerExhaustionRotation(pi, sessionSelection); + const storedCredential = readStoredCredential('grok-cli'); + const environmentApiKey = + process.env.GROK_CLI_OAUTH_TOKEN && storedCredential?.type !== 'oauth' + ? '$GROK_CLI_OAUTH_TOKEN' + : undefined; + + const oauthProvider = { + name: 'Grok CLI', + usesCallbackServer: true, + + async login(callbacks) { + await migration; + if (migrationError) throw new Error(migrationError); + if (process.env.GROK_CLI_OAUTH_TOKEN) { + throw new Error('Unset GROK_CLI_OAUTH_TOKEN before logging in to Grok CLI.'); + } + + const stored = readStoredCredential('grok-cli'); + const current = await getAccountVault(); + if (current.migration.markerInstallPending) { + if (stored?.type === 'oauth' && !isMarker(stored)) { + await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === ACCOUNT_1_ID); + if (!account) throw new Error('The permanent Grok CLI account is missing.'); + if (!account.credential) { + account.credential = accountCredential(stored); + account.revision += 1; + } + vault.activeAccountId ??= account.id; + }); + } + if ((await getAccountVault()).accounts.some((account) => account.credential)) + return marker(); + } + + if (current.accounts.some((account) => account.credential)) return marker(); + + const revision = current.accounts.find((account) => account.id === ACCOUNT_1_ID)?.revision; + if (revision === undefined) throw new Error('The permanent Grok CLI account is missing.'); + const credentials = await oauth.login(callbacks); + await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === ACCOUNT_1_ID); + if (!account) throw new Error('The permanent Grok CLI account is missing.'); + if (account.revision !== revision) { + throw new Error('The account changed while login was in progress. Try again.'); + } + account.credential = accountCredential(credentials); + account.revision += 1; + vault.activeAccountId ??= account.id; + }); + await removeQuotaUsage(ACCOUNT_1_ID).catch(() => undefined); + exhaustionRotation.clearRecentExhaustion(ACCOUNT_1_ID); + return marker(); + }, + + async refreshToken(credentials) { + if (isMarker(credentials) || process.env.GROK_CLI_OAUTH_TOKEN) return marker(); + return oauth.refresh(credentials); + }, + + getApiKey(credentials) { + return credentials.access; + }, + } satisfies NonNullable; + + pi.registerProvider('grok-cli', { + name: 'Grok CLI', + baseUrl: getBaseUrl(), + ...(environmentApiKey ? { apiKey: environmentApiKey } : { oauth: oauthProvider }), + api: 'openai-responses', + models: resolveModels().map((model) => ({ + id: model.id, + name: model.name, + reasoning: model.reasoning, + thinkingLevelMap: model.thinkingLevelMap, + input: model.input, + cost: model.cost, + contextWindow: model.contextWindow, + maxTokens: model.maxTokens, + headers: grokCliModelHeaders(model.id), + })), + streamSimple(model, context, options?: SimpleStreamOptions) { + const accountId = sessionSelection.accountId(options?.sessionId); + return lazyStream(model, async () => { + await migration; + if (migrationError) throw new Error(migrationError); + const route = await resolveAccountRoute(accountId); + const stream = streamSimpleOpenAIResponses( + { + ...model, + baseUrl: route.baseUrl, + api: 'openai-responses', + } as Model<'openai-responses'>, + context, + { + ...options, + apiKey: route.token, + }, + ); + void stream.result().then( + (message) => { + rememberRequestAccount(message, route.accountId); + }, + () => undefined, + ); + return stream; + }); + }, + }); - pi.on('model_select', (event) => { - accountManagement.handleModelSelect(event); - syncTools(event.model); + const accountManagement = registerAccountManagement( + pi, + exhaustionRotation.clearRecentExhaustion, + sessionSelection, + ); + const resolveSessionRoute = (ctx: Pick) => + resolveAccountRoute(sessionSelection.accountId(ctx.sessionManager.getSessionId())); + registerImagineFeature(pi, undefined, async (ctx) => (await resolveSessionRoute(ctx)).token); + + pi.on('model_select', (_event) => { + accountManagement.handleModelSelect(); + syncImageToolPreference(pi); }); - pi.on('before_agent_start', (_event, ctx) => { - syncTools(ctx.model); + pi.on('before_agent_start', () => { + syncImageToolPreference(pi); }); - pi.on('session_start', async (event, ctx) => { - if (migration.warning) { - ctx.ui.notify(`[pi-grok-cli] ${migration.warning}`, 'warning'); - delete migration.warning; + pi.on('session_start', async (_event, ctx) => { + await migration; + const accountId = sessionSelection.restore(ctx); + if (accountId && !process.env.GROK_CLI_OAUTH_TOKEN) sessionSelection.select(ctx, accountId); + if (migrationError && !migrationErrorNotified) { + ctx.ui.notify(`[pi-grok-cli] ${migrationError}`, 'warning'); + migrationErrorNotified = true; } - if (event.reason === 'new' || event.reason === 'resume' || event.reason === 'fork') { - restoreGrokTools(pi, ctx.sessionManager.getSessionFile()); + if (migrationWarning) { + ctx.ui.notify(`[pi-grok-cli] ${migrationWarning}`, 'warning'); + migrationWarning = undefined; } - syncTools(ctx.model, true); + if (migrationNotice) { + ctx.ui.notify(`[pi-grok-cli] ${migrationNotice}`, 'info'); + migrationNotice = undefined; + } + syncImageToolPreference(pi); if (process.env.GROK_CLI_OAUTH_TOKEN) { ctx.ui.notify( '[pi-grok-cli] Using GROK_CLI_OAUTH_TOKEN bypass — no auto-refresh, no model discovery', 'warning', ); } + }); - if (!webSearchRegistered) return; - - bindLivePiWebAccess(pi); - await ensureWebSearchDelegate(pi); - syncTools(ctx.model); + pi.on('session_tree', (_event, ctx) => { + const accountId = sessionSelection.restore(ctx); + if (accountId && !process.env.GROK_CLI_OAUTH_TOKEN) sessionSelection.select(ctx, accountId); }); - pi.on('session_shutdown', async (event) => { + pi.on('session_shutdown', async (_event, ctx) => { + sessionSelection.clear(ctx.sessionManager.getSessionId()); await accountManagement.closeDashboard(); - syncTools(undefined); - if (event.reason === 'new' || event.reason === 'resume' || event.reason === 'fork') { - handoffGrokTools(pi, event.targetSessionFile); - } }); pi.on('before_provider_headers', (event, ctx) => { - if (!isGrokCliProvider(ctx.model?.provider)) return; + if (ctx.model?.provider !== 'grok-cli') return; event.headers['x-grok-conv-id'] = ctx.sessionManager.getSessionId(); }); pi.on('before_provider_request', (event, ctx) => { - if (!isGrokCliProvider(ctx.model?.provider)) return; - + if (ctx.model?.provider !== 'grok-cli') return; const modelId = ctx.model?.id ?? ''; const sessionId = ctx.sessionManager?.getSessionId(); return sanitizePayload(event.payload as Record, modelId, sessionId, ctx.cwd); }); - registerUsageCommand(pi); - - registerVisionFeature(pi); + registerUsageCommand(pi, resolveSessionRoute); } diff --git a/src/provider/requestOwnership.ts b/src/provider/requestOwnership.ts new file mode 100644 index 0000000..e679432 --- /dev/null +++ b/src/provider/requestOwnership.ts @@ -0,0 +1,11 @@ +import type { AssistantMessage } from '@earendil-works/pi-ai'; + +const requestAccounts = new WeakMap(); + +export function rememberRequestAccount(message: AssistantMessage, accountId: string) { + requestAccounts.set(message, accountId); +} + +export function requestAccount(message: AssistantMessage) { + return requestAccounts.get(message); +} diff --git a/src/provider/rotation.ts b/src/provider/rotation.ts index ea2880a..3a0ac79 100644 --- a/src/provider/rotation.ts +++ b/src/provider/rotation.ts @@ -1,8 +1,12 @@ import type { AssistantMessage } from '@earendil-works/pi-ai'; -import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'; -import { loadConfig, saveConfig } from '../config.js'; -import { DEFAULT_GROK_MODEL, GROK_CLI_PROVIDER, isGrokCliProvider } from './accounts.js'; +import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; +import { getAccountVault } from './accountVault.js'; import { type CachedQuota, isCachedQuotaFresh, loadQuotaCache } from './quotaCache.js'; +import { requestAccount } from './requestOwnership.js'; +import { + createSessionAccountSelection, + type SessionAccountSelection, +} from './sessionAccountSelection.js'; export const EXHAUSTED_BALANCE_ERROR = 'OpenAI API error (402): 402 "Grok Build usage balance exhausted"'; @@ -11,8 +15,8 @@ export const ROTATION_CONTINUATION = const RECENT_EXHAUSTION_COOLDOWN_MS = 5 * 60_000; type FailedResponse = { + accountId: string; model: string; - provider: string; }; function isExhaustionMessage(message: unknown): message is AssistantMessage { @@ -20,23 +24,16 @@ function isExhaustionMessage(message: unknown): message is AssistantMessage { const candidate = message as Partial; return ( candidate.role === 'assistant' && - isGrokCliProvider(candidate.provider) && + candidate.provider === 'grok-cli' && candidate.stopReason === 'error' && candidate.errorMessage?.trim() === EXHAUSTED_BALANCE_ERROR ); } -function hasAccountAuth(ctx: ExtensionContext, provider: string) { - return ( - ctx.modelRegistry.getProviderAuthStatus(provider).configured || - (provider === GROK_CLI_PROVIDER && Boolean(process.env.GROK_CLI_OAUTH_TOKEN)) - ); -} - -function circularProviders(providers: string[], current: string) { - const index = providers.indexOf(current); - if (index < 0) return providers; - return [...providers.slice(index + 1), ...providers.slice(0, index)]; +function circularAccountIds(accountIds: string[], current: string) { + const index = accountIds.indexOf(current); + if (index < 0) return accountIds; + return [...accountIds.slice(index + 1), ...accountIds.slice(0, index)]; } function quotaScore(entry: CachedQuota | undefined, now: number) { @@ -51,43 +48,44 @@ function quotaScore(entry: CachedQuota | undefined, now: number) { return Math.min(monthly, Math.min(1, Math.max(0, 1 - entry.weekly.creditUsagePercent / 100))); } -function orderProvidersByQuota( - providers: string[], +function orderAccountsByQuota( + accountIds: string[], accounts: Record, now: number, ) { - const scored = providers.flatMap((provider, index) => { - const score = quotaScore(accounts[provider], now); - return score === undefined ? [] : [{ provider, index, score }]; + const scored = accountIds.flatMap((accountId, index) => { + const score = quotaScore(accounts[accountId], now); + return score === undefined ? [] : [{ accountId, index, score }]; }); const ranked = [...scored].sort( (left, right) => right.score - left.score || left.index - right.index, ); - return providers.map((provider, index) => { + return accountIds.map((accountId, index) => { const scoredIndex = scored.findIndex((candidate) => candidate.index === index); - return scoredIndex < 0 ? provider : ranked[scoredIndex].provider; + return scoredIndex < 0 ? accountId : ranked[scoredIndex].accountId; }); } -export function registerExhaustionRotation(pi: ExtensionAPI) { +export function registerExhaustionRotation( + pi: ExtensionAPI, + sessionSelection: SessionAccountSelection = createSessionAccountSelection(pi), +) { const exhausted = new Set(); - const unavailable = new Set(); const recentlyExhausted = new Map(); let pending: FailedResponse | undefined; let awaitingContinuation = false; const clearChain = () => { exhausted.clear(); - unavailable.clear(); pending = undefined; awaitingContinuation = false; }; - const isRecentlyExhausted = (provider: string, now: number) => { - const exhaustedAt = recentlyExhausted.get(provider); + const isRecentlyExhausted = (accountId: string, now: number) => { + const exhaustedAt = recentlyExhausted.get(accountId); if (exhaustedAt === undefined) return false; if (now - exhaustedAt < RECENT_EXHAUSTION_COOLDOWN_MS) return true; - recentlyExhausted.delete(provider); + recentlyExhausted.delete(accountId); return false; }; @@ -104,16 +102,16 @@ export function registerExhaustionRotation(pi: ExtensionAPI) { if (pending) clearChain(); }); - pi.on('message_end', (event) => { + pi.on('message_end', (event, ctx) => { + if (process.env.GROK_CLI_OAUTH_TOKEN) return; if (!isExhaustionMessage(event.message)) return; - const message = event.message; - const config = loadConfig().config; - if (!config.accounts.items.some((account) => account.provider === message.provider)) { - return; - } - pending = { provider: message.provider, model: message.model }; - exhausted.add(message.provider); - recentlyExhausted.set(message.provider, Date.now()); + const accountId = + requestAccount(event.message) ?? + sessionSelection.accountId(ctx.sessionManager.getSessionId()); + if (!accountId) return; + pending = { accountId, model: event.message.model }; + exhausted.add(accountId); + recentlyExhausted.set(accountId, Date.now()); awaitingContinuation = false; }); @@ -122,59 +120,53 @@ export function registerExhaustionRotation(pi: ExtensionAPI) { if (awaitingContinuation) clearChain(); return; } - const failed = pending; pending = undefined; - if (ctx.model?.provider !== failed.provider || ctx.model.id !== failed.model) { + if (ctx.model?.provider !== 'grok-cli' || ctx.model.id !== failed.model) { clearChain(); return; } - const config = loadConfig().config; - const authenticated = config.accounts.items.filter((account) => - hasAccountAuth(ctx, account.provider), + const now = Date.now(); + const quotas = loadQuotaCache().accounts; + const vault = await getAccountVault(); + const loggedIn = vault.accounts.filter((account) => account.credential); + const eligible = circularAccountIds( + vault.accounts.map((account) => account.id), + failed.accountId, + ).filter( + (accountId) => + accountId !== failed.accountId && + !exhausted.has(accountId) && + !isRecentlyExhausted(accountId, now) && + loggedIn.some((account) => account.id === accountId), ); - if (authenticated.length < 2) { + const selectedId = orderAccountsByQuota(eligible, quotas, now)[0]; + const selected = vault.accounts.find( + (account) => account.id === selectedId && account.credential, + ); + const outcome = { + loggedInIds: loggedIn.map((account) => account.id), + ...(selected + ? { + selected: { + id: selected.id, + failedLabel: + vault.accounts.find((account) => account.id === failed.accountId)?.label ?? + failed.accountId, + label: selected.label, + }, + } + : {}), + }; + if (outcome.loggedInIds.length < 2) { clearChain(); return; } - - const now = Date.now(); - const authenticatedProviders = new Set(authenticated.map((account) => account.provider)); - const eligible = circularProviders( - config.accounts.items.map((account) => account.provider), - failed.provider, - ).filter( - (provider) => - provider !== failed.provider && - !exhausted.has(provider) && - !isRecentlyExhausted(provider, now) && - !unavailable.has(provider) && - authenticatedProviders.has(provider), - ); - - for (const provider of orderProvidersByQuota(eligible, loadQuotaCache().accounts, now)) { - const model = - ctx.modelRegistry.find(provider, failed.model) ?? - ctx.modelRegistry.find(provider, DEFAULT_GROK_MODEL); - if (!model || !(await pi.setModel(model))) { - unavailable.add(provider); - continue; - } - - const refreshed = loadConfig().config; - const failedLabel = refreshed.accounts.items.find( - (account) => account.provider === failed.provider, - )?.label; - const selected = refreshed.accounts.items.find((account) => account.provider === provider); - if (!failedLabel || !selected) { - unavailable.add(provider); - continue; - } - refreshed.accounts.selectedProvider = provider; - saveConfig(refreshed); + if (outcome.selected) { + sessionSelection.select(ctx, outcome.selected.id); ctx.ui.notify( - `Grok CLI: “${failedLabel}” exhausted; switched to “${selected.label}” and continuing.`, + `Grok CLI: “${outcome.selected.failedLabel}” exhausted; switched to “${outcome.selected.label}” and continuing.`, 'info', ); awaitingContinuation = true; @@ -183,8 +175,8 @@ export function registerExhaustionRotation(pi: ExtensionAPI) { } if ( - authenticated.every( - (account) => exhausted.has(account.provider) || isRecentlyExhausted(account.provider, now), + outcome.loggedInIds.every( + (accountId) => exhausted.has(accountId) || isRecentlyExhausted(accountId, now), ) ) { ctx.ui.notify('Grok CLI: all logged-in accounts are exhausted.', 'warning'); @@ -199,8 +191,8 @@ export function registerExhaustionRotation(pi: ExtensionAPI) { }); return { - clearRecentExhaustion(provider: string) { - recentlyExhausted.delete(provider); + clearRecentExhaustion(accountId: string) { + recentlyExhausted.delete(accountId); }, }; } diff --git a/src/provider/sessionAccountSelection.ts b/src/provider/sessionAccountSelection.ts new file mode 100644 index 0000000..ca264b5 --- /dev/null +++ b/src/provider/sessionAccountSelection.ts @@ -0,0 +1,60 @@ +import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'; +import { ACCOUNT_1_ID, getAccountVaultSync } from './accountVault.js'; + +export const SESSION_ACCOUNT_ENTRY = 'grok-cli-active-account-v1'; + +function storedAccountId( + entry: ReturnType[number], +) { + if (entry.type !== 'custom' || entry.customType !== SESSION_ACCOUNT_ENTRY) return undefined; + if (!entry.data || typeof entry.data !== 'object' || Array.isArray(entry.data)) return undefined; + const accountId = (entry.data as Record).accountId; + return typeof accountId === 'string' && accountId ? accountId : undefined; +} + +export function createSessionAccountSelection(pi: Pick) { + const selections = new Map(); + const accountId = (sessionId: string | undefined) => { + if (process.env.GROK_CLI_OAUTH_TOKEN) return ACCOUNT_1_ID; + const vault = getAccountVaultSync(); + const selected = sessionId + ? vault.accounts.find( + (account) => account.id === selections.get(sessionId) && account.credential, + ) + : undefined; + if (selected) return selected.id; + if (sessionId) selections.delete(sessionId); + return ( + vault.accounts.find( + (account) => account.id === vault.activeAccountId && account.credential, + ) ?? vault.accounts.find((account) => account.credential) + )?.id; + }; + + return { + accountId, + restore(ctx: Pick) { + const sessionId = ctx.sessionManager.getSessionId(); + const restored = ctx.sessionManager + .getBranch() + .reduceRight( + (accountId, entry) => accountId ?? storedAccountId(entry), + undefined, + ); + if (restored) selections.set(sessionId, restored); + else selections.delete(sessionId); + return accountId(sessionId); + }, + select(ctx: Pick, accountId: string) { + const sessionId = ctx.sessionManager.getSessionId(); + if (selections.get(sessionId) === accountId) return; + selections.set(sessionId, accountId); + pi.appendEntry(SESSION_ACCOUNT_ENTRY, { accountId }); + }, + clear(sessionId: string) { + selections.delete(sessionId); + }, + }; +} + +export type SessionAccountSelection = ReturnType; diff --git a/src/provider/toolScope.ts b/src/provider/toolScope.ts deleted file mode 100644 index e6f1daa..0000000 --- a/src/provider/toolScope.ts +++ /dev/null @@ -1,220 +0,0 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { loadConfig } from '../config.js'; -import { isGrokCliProvider } from './accounts.js'; - -export const SHIM_MODEL_IDS = new Set(['grok-build', 'grok-composer-2.5-fast']); - -const TOOL_GROUPS = [ - { native: 'read', compatibility: ['Read'] }, - { native: 'write', compatibility: ['Write'] }, - { native: 'edit', compatibility: ['Edit', 'StrReplace'] }, - { native: 'grep', compatibility: ['Grep'] }, - { native: 'find', compatibility: ['Glob'] }, - { native: 'ls', compatibility: ['LS'] }, - { native: 'bash', compatibility: ['Shell'] }, -] as const; - -const WEB_SEARCH_GROUP = { native: 'web_search', compatibility: ['WebSearch'] } as const; -const LEGACY_DISCOVERY_TOOLS = ['Grep', 'Glob', 'LS'] as const; - -type ToolScopeState = { - deleteCaptured: boolean; - deleteEnabled: boolean; - initialized: boolean; - legacyDiscoveryCaptured: boolean; - legacyDiscoveryEnabled: string[]; - legacy: boolean; -}; - -type ToolScopeHandoff = { - activeTools: string[]; - state: ToolScopeState; -}; - -type ToolScopeHandoffStore = { - anonymous: ToolScopeHandoff[]; - bySessionFile: Map; -}; - -type ToolModel = { - provider: string; - id: string; -}; - -type SyncOptions = { - captureDelete?: boolean; - imagineEnabled?: boolean; - webSearchRegistered?: boolean; -}; - -type ToolScopeAPI = Pick & { - events?: ExtensionAPI['events']; - getAllTools: () => { name: string }[]; -}; - -const toolScopeStates = new WeakMap(); -const globalToolScope = globalThis as typeof globalThis & { - __piGrokCliToolScopeHandoffs?: ToolScopeHandoffStore; -}; -const toolScopeHandoffs: ToolScopeHandoffStore = globalToolScope.__piGrokCliToolScopeHandoffs ?? { - anonymous: [], - bySessionFile: new Map(), -}; -globalToolScope.__piGrokCliToolScopeHandoffs = toolScopeHandoffs; - -function isLegacyModel(model: ToolModel | undefined) { - return isGrokCliProvider(model?.provider) && SHIM_MODEL_IDS.has(model?.id ?? ''); -} - -function sameTools(currentTools: string[], nextTools: string[]) { - return ( - currentTools.length === nextTools.length && - currentTools.every((toolName, index) => toolName === nextTools[index]) - ); -} - -function stateKey(pi: ToolScopeAPI) { - return pi.events ?? pi; -} - -export function handoffGrokTools(pi: ToolScopeAPI, targetSessionFile?: string) { - const state = toolScopeStates.get(stateKey(pi)); - if (!state) return; - const handoff = { - activeTools: [...pi.getActiveTools()], - state: { ...state, legacyDiscoveryEnabled: [...state.legacyDiscoveryEnabled] }, - }; - if (targetSessionFile) { - toolScopeHandoffs.bySessionFile.set(targetSessionFile, handoff); - return; - } - toolScopeHandoffs.anonymous.push(handoff); -} - -export function restoreGrokTools(pi: ToolScopeAPI, sessionFile?: string) { - const handoff = sessionFile - ? toolScopeHandoffs.bySessionFile.get(sessionFile) - : toolScopeHandoffs.anonymous.shift(); - if (!handoff) return; - if (sessionFile) toolScopeHandoffs.bySessionFile.delete(sessionFile); - toolScopeStates.set(stateKey(pi), { - ...handoff.state, - legacyDiscoveryEnabled: [...handoff.state.legacyDiscoveryEnabled], - }); - const registeredTools = new Set(pi.getAllTools().map((tool) => tool.name)); - const activeTools = handoff.activeTools.filter((toolName) => registeredTools.has(toolName)); - if (!sameTools(pi.getActiveTools(), activeTools)) pi.setActiveTools(activeTools); -} - -export function syncGrokTools( - pi: ToolScopeAPI, - model: ToolModel | undefined, - options: SyncOptions = {}, -) { - const currentTools = pi.getActiveTools(); - const key = stateKey(pi); - const state = toolScopeStates.get(key) ?? { - deleteCaptured: false, - deleteEnabled: false, - initialized: false, - legacyDiscoveryCaptured: false, - legacyDiscoveryEnabled: [], - legacy: false, - }; - const legacy = isLegacyModel(model); - const imagineEnabled = options.imagineEnabled ?? loadConfig().config.imagine.enabled; - const groups = options.webSearchRegistered ? [...TOOL_GROUPS, WEB_SEARCH_GROUP] : TOOL_GROUPS; - const registeredTools = new Set(pi.getAllTools().map((tool) => tool.name)); - const registeredCompatibilityTools = groups.flatMap((group) => - group.compatibility.filter((name) => registeredTools.has(name)), - ); - const compatibilityToolsWereAutoActivated = - !state.initialized && - registeredCompatibilityTools.length > 0 && - registeredCompatibilityTools.every((name) => currentTools.includes(name)); - const compatibilityVocabularyAvailable = groups.every( - (group) => - !currentTools.includes(group.native) || - group.compatibility.some((name) => registeredTools.has(name)), - ); - - if (!state.legacyDiscoveryCaptured) { - state.legacyDiscoveryCaptured = true; - state.legacyDiscoveryEnabled = LEGACY_DISCOVERY_TOOLS.filter( - (name) => registeredTools.has(name) && currentTools.includes(name), - ); - } - if (state.legacy) { - state.legacyDiscoveryEnabled = LEGACY_DISCOVERY_TOOLS.filter( - (name) => registeredTools.has(name) && currentTools.includes(name), - ); - } - - if (options.captureDelete && !state.deleteCaptured) { - state.deleteCaptured = true; - state.deleteEnabled = currentTools.includes('Delete'); - } - if (state.legacy) state.deleteEnabled = currentTools.includes('Delete'); - - const handledGroups = new Set<(typeof groups)[number]>(); - const nextTools = currentTools.flatMap((toolName) => { - if (toolName === 'image_gen') return imagineEnabled ? [toolName] : []; - if (toolName === 'Delete') return legacy && state.deleteEnabled ? [toolName] : []; - - const group = groups.find( - (candidate) => - candidate.native === toolName || candidate.compatibility.some((name) => name === toolName), - ); - if (!group) return [toolName]; - if (handledGroups.has(group)) return []; - handledGroups.add(group); - - const nativeActive = currentTools.includes(group.native); - const activeCompatibility = currentTools.filter((name) => - group.compatibility.some((compatibilityName) => compatibilityName === name), - ); - if (!registeredTools.has(group.native)) return []; - if (!legacy || !compatibilityVocabularyAvailable) { - return nativeActive || - (!compatibilityToolsWereAutoActivated && activeCompatibility.length > 0) - ? [group.native] - : []; - } - if (nativeActive) { - const registeredCompatibility = group.compatibility.filter((name) => - registeredTools.has(name), - ); - return registeredCompatibility.length > 0 ? registeredCompatibility : [group.native]; - } - if (compatibilityToolsWereAutoActivated) { - return activeCompatibility.filter((name) => state.legacyDiscoveryEnabled.includes(name)); - } - return activeCompatibility.filter((name) => registeredTools.has(name)); - }); - - if (legacy && compatibilityVocabularyAvailable) { - nextTools.push( - ...state.legacyDiscoveryEnabled.filter( - (name) => registeredTools.has(name) && !nextTools.includes(name), - ), - ); - } - - if ( - legacy && - state.deleteEnabled && - registeredTools.has('Delete') && - !nextTools.includes('Delete') - ) { - nextTools.push('Delete'); - } - if (imagineEnabled && registeredTools.has('image_gen') && !nextTools.includes('image_gen')) { - nextTools.push('image_gen'); - } - - state.initialized = model !== undefined; - state.legacy = legacy; - toolScopeStates.set(key, state); - if (sameTools(currentTools, nextTools)) return; - pi.setActiveTools(nextTools); -} diff --git a/src/provider/usage.ts b/src/provider/usage.ts index f38d9c6..8a803c1 100644 --- a/src/provider/usage.ts +++ b/src/provider/usage.ts @@ -1,16 +1,20 @@ import type { Api, Model } from '@earendil-works/pi-ai'; -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; +import type { ExtensionAPI, ExtensionCommandContext } from '@earendil-works/pi-coding-agent'; import { XaiOAuthError } from '../shared/errors.js'; -import { isGrokCliProvider, resolveGrokProvider, resolveGrokToken } from './accounts.js'; +import { resolveAccountRoute } from './accountRouting.js'; +import { accountRouteIsCurrent, isGrokCliProvider } from './accounts.js'; import { fetchBillingUsage, formatQuota } from './billing.js'; -import { loadQuotaCache, saveQuotaUsage } from './quotaCache.js'; +import { loadQuotaCache, saveQuotaUsageWhen } from './quotaCache.js'; -export function registerUsageCommand(pi: Pick) { +export function registerUsageCommand( + pi: Pick, + resolveRoute: (ctx: ExtensionCommandContext) => ReturnType = () => + resolveAccountRoute(), +) { pi.registerCommand('grok-cli-usage', { description: 'Show Grok CLI provider status, quota, and token health', handler: async (_args, ctx) => { - const provider = resolveGrokProvider(ctx); - if (provider === 'grok-cli' && process.env.GROK_CLI_OAUTH_TOKEN) { + if (process.env.GROK_CLI_OAUTH_TOKEN) { ctx.ui.notify( '⚠️ Grok CLI: using GROK_CLI_OAUTH_TOKEN env bypass — no auto-refresh available', 'warning', @@ -27,17 +31,37 @@ export function registerUsageCommand(pi: Pick) return; } - const apiKey = await resolveGrokToken(ctx); - if (!apiKey) { - ctx.ui.notify(formatQuota(undefined).join('\n'), 'info'); + let routeError: unknown; + const route = await resolveRoute(ctx).catch((error: unknown) => { + routeError = error; + return undefined; + }); + if (!route) { + ctx.ui.notify( + [ + ...formatQuota(undefined), + ...(routeError === undefined + ? [] + : [ + ` Reason ${routeError instanceof Error ? routeError.message : String(routeError)}`, + ]), + ].join('\n'), + 'info', + ); return; } try { ctx.ui.notify('Fetching grok cli usage…', 'info'); - const usage = await fetchBillingUsage(apiKey, AbortSignal.timeout(30_000)); + const usage = await fetchBillingUsage(route.token, AbortSignal.timeout(30_000)); try { - await saveQuotaUsage(provider, usage); + if ( + !(await saveQuotaUsageWhen(route.accountId, usage, () => + accountRouteIsCurrent(route), + )) + ) { + throw new Error('The account changed while its usage was refreshing.'); + } } catch (error) { ctx.ui.notify( `Grok CLI quota cache update failed: ${error instanceof Error ? error.message : String(error)}`, @@ -50,7 +74,7 @@ export function registerUsageCommand(pi: Pick) `Grok CLI billing refresh failed: ${err instanceof Error ? err.message : String(err)}`, 'warning', ); - const cached = loadQuotaCache().accounts[provider]; + const cached = loadQuotaCache().accounts[route.accountId]; ctx.ui.notify( cached ? `Grok CLI cached usage from ${cached.updatedAt}:\n${formatQuota(cached).join('\n')}` diff --git a/src/storage.ts b/src/storage.ts index ac1e492..3234d87 100644 --- a/src/storage.ts +++ b/src/storage.ts @@ -1,10 +1,14 @@ import { randomUUID } from 'node:crypto'; import { + closeSync, existsSync, mkdirSync, + openSync, + readdirSync, readFileSync, renameSync, rmSync, + statSync, unlinkSync, writeFileSync, } from 'node:fs'; @@ -12,16 +16,126 @@ import { homedir } from 'node:os'; import { basename, dirname, join } from 'node:path'; const homePath = () => process.env.HOME || homedir(); +const LOCK_STALE_MS = 30_000; export const getGrokCliDirectory = () => join(homePath(), '.pi', 'grok-cli'); +export const getAccountVaultPath = () => join(getGrokCliDirectory(), 'accounts.json'); +export const getConfigBackupPath = () => join(getGrokCliDirectory(), 'config.v2.backup.json'); export const getConfigPath = () => join(getGrokCliDirectory(), 'config.json'); -export const getVisionCachePath = () => join(getGrokCliDirectory(), 'vision-cache.json'); export const getQuotaCachePath = () => join(getGrokCliDirectory(), 'quota-cache.json'); export const getLegacyConfigPath = () => join(homePath(), '.pi', 'grok-cli.json'); export const getLegacyImagineConfigPath = () => join(homePath(), '.pi', 'grok-cli-imagine.json'); -export const getLegacyVisionConfigPath = () => join(homePath(), '.pi', 'grok-cli-vision.json'); -export const getLegacyVisionCachePath = () => join(homePath(), '.pi', 'grok-cli-vision-cache.json'); + +function abandonedLock(path: string) { + try { + const owner = JSON.parse(readFileSync(path, 'utf8')) as unknown; + if (owner && typeof owner === 'object' && !Array.isArray(owner)) { + const pid = (owner as Record).pid; + if (typeof pid === 'number' && Number.isInteger(pid) && pid > 0) { + return !processIsRunning(pid); + } + } + } catch {} + try { + return Date.now() - statSync(path).mtimeMs > LOCK_STALE_MS; + } catch { + return false; + } +} + +function processIsRunning(pid: number) { + try { + process.kill(pid, 0); + return true; + } catch (error) { + return (error as NodeJS.ErrnoException).code !== 'ESRCH'; + } +} + +function recoveryPaths(lockPath: string) { + const prefix = `${basename(lockPath)}.recovery.`; + return readdirSync(dirname(lockPath)) + .filter((name) => name.startsWith(prefix)) + .map((name) => join(dirname(lockPath), name)); +} + +function recoveryInProgress(lockPath: string) { + return recoveryPaths(lockPath).some((path) => { + try { + const pid = Number(basename(path).split('.recovery.')[1]?.split('.')[0]); + if ( + Number.isInteger(pid) && + pid > 0 && + processIsRunning(pid) && + Date.now() - statSync(path).mtimeMs <= LOCK_STALE_MS + ) { + return true; + } + rmSync(path, { force: true }); + return false; + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error; + rmSync(path, { force: true }); + return false; + } + }); +} + +function releaseOwnedLock(lockPath: string, token: string) { + try { + const owner = JSON.parse(readFileSync(lockPath, 'utf8')) as Record; + if (owner.token === token) unlinkSync(lockPath); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error; + } +} + +function recoverAbandonedLock(lockPath: string) { + const recoveryPath = `${lockPath}.recovery.${process.pid}.${randomUUID()}`; + closeSync(openSync(recoveryPath, 'wx', 0o600)); + try { + if (abandonedLock(lockPath)) rmSync(lockPath, { force: true }); + } finally { + unlinkSync(recoveryPath); + } +} + +export async function acquireFileLock(path: string) { + const lockPath = `${path}.lock`; + const token = randomUUID(); + const startedAt = Date.now(); + mkdirSync(dirname(path), { recursive: true }); + while (true) { + if (Date.now() - startedAt >= LOCK_STALE_MS) { + throw new Error(`Timed out waiting for file lock: ${lockPath}`); + } + if (recoveryInProgress(lockPath)) { + await new Promise((resolve) => setTimeout(resolve, 25)); + continue; + } + try { + const file = openSync(lockPath, 'wx', 0o600); + try { + writeFileSync(file, JSON.stringify({ pid: process.pid, token })); + } finally { + closeSync(file); + } + if (recoveryInProgress(lockPath)) { + releaseOwnedLock(lockPath, token); + continue; + } + return async () => releaseOwnedLock(lockPath, token); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'EEXIST') throw error; + if (abandonedLock(lockPath)) { + recoverAbandonedLock(lockPath); + continue; + } + await new Promise((resolve) => setTimeout(resolve, 25)); + } + } +} export function writeFileAtomic(path: string, contents: string, mode?: number) { const tempPath = join(dirname(path), `.${basename(path)}.${process.pid}.${randomUUID()}.tmp`); diff --git a/src/tools/files.ts b/src/tools/files.ts deleted file mode 100644 index 960d5dc..0000000 --- a/src/tools/files.ts +++ /dev/null @@ -1,437 +0,0 @@ -import { existsSync, promises as fs, unlinkSync } from 'node:fs'; -import { resolve } from 'node:path'; -import { Type } from '@earendil-works/pi-ai'; -import { - createEditToolDefinition, - createLsToolDefinition, - createWriteToolDefinition, - type EditToolInput, - type ExtensionAPI, - type ExtensionContext, - generateUnifiedPatch, - type LsToolInput, - type WriteToolInput, - withFileMutationQueue, -} from '@earendil-works/pi-coding-agent'; -import { - booleanDetail, - fileError, - fileNotFound, - nativeRenderContext, - numberDetail, - recordFrom, - renderResultSummary, - stringFrom, - text, -} from './rendering.js'; - -type ReplacementEdit = { oldText: string; newText: string }; -type StrReplaceArgs = { path: string; old_str: string; new_str: string }; -type StrReplaceDetails = { - path: string; - replacements: number; - diff?: string; - diffTruncated?: boolean; -}; - -type ToolTheme = { - bold: (text: string) => string; - fg: (name: 'accent' | 'toolTitle', text: string) => string; -}; - -export const STR_REPLACE_DIFF_MAX_CHARS = 26_000; -const STR_REPLACE_COLLAPSED_DIFF_MAX_CHARS = 4_000; - -function parseJson(value: string): unknown { - try { - return JSON.parse(value); - } catch { - return undefined; - } -} - -function replacementEdit(value: unknown): ReplacementEdit | undefined { - const input = recordFrom(value); - if (!input) return undefined; - const oldText = - stringFrom(input.oldText) ?? stringFrom(input.old_string) ?? stringFrom(input.old_str); - const newText = - stringFrom(input.newText) ?? stringFrom(input.new_string) ?? stringFrom(input.new_str); - if (oldText === undefined || newText === undefined) return undefined; - return { oldText, newText }; -} - -function editList(value: unknown): ReplacementEdit[] | undefined { - const list = typeof value === 'string' ? parseJson(value) : value; - if (!Array.isArray(list)) return undefined; - const edits = list.map(replacementEdit); - if (edits.some((edit) => edit === undefined)) return undefined; - return edits.filter((edit): edit is ReplacementEdit => edit !== undefined); -} - -function normalizeWriteArgs(value: unknown): WriteToolInput { - const input = recordFrom(value); - return { - path: stringFrom(input?.path) ?? stringFrom(input?.file_path) ?? '', - content: stringFrom(input?.content) ?? stringFrom(input?.contents) ?? '', - }; -} - -function normalizeLsArgs(value: unknown): LsToolInput { - const input = recordFrom(value); - return { - path: stringFrom(input?.path), - limit: typeof input?.limit === 'number' ? input.limit : undefined, - }; -} - -function normalizeEditArgs(value: unknown): EditToolInput { - const input = recordFrom(value); - const topLevel = replacementEdit(input); - const nested = replacementEdit(input?.strReplace); - return { - path: stringFrom(input?.path) ?? stringFrom(input?.file_path) ?? '', - edits: - editList(input?.edits) ?? - editList(recordFrom(input?.multiStrReplace)?.edits) ?? - (topLevel ? [topLevel] : undefined) ?? - (nested ? [nested] : []), - }; -} - -function normalizeStrReplaceArgs(value: unknown): StrReplaceArgs { - const input = recordFrom(value); - return { - path: stringFrom(input?.path) ?? '', - old_str: - stringFrom(input?.old_str) ?? - stringFrom(input?.old_string) ?? - stringFrom(input?.oldText) ?? - stringFrom(recordFrom(input?.strReplace)?.oldText) ?? - '', - new_str: - stringFrom(input?.new_str) ?? - stringFrom(input?.new_string) ?? - stringFrom(input?.newText) ?? - stringFrom(recordFrom(input?.strReplace)?.newText) ?? - '', - }; -} - -function countOccurrences(content: string, search: string) { - let count = 0; - let offset = 0; - while (offset <= content.length - search.length) { - const match = content.indexOf(search, offset); - if (match === -1) return count; - count += 1; - offset = match + search.length; - } - return count; -} - -function boundedDiff(filePath: string, before: string, after: string) { - const fullDiff = generateUnifiedPatch(filePath, before, after); - if (fullDiff.length <= STR_REPLACE_DIFF_MAX_CHARS) { - return { diff: fullDiff, diffTruncated: false }; - } - const notice = '\n[Diff truncated at 26000 characters]'; - return { - diff: `${fullDiff.slice(0, STR_REPLACE_DIFF_MAX_CHARS - notice.length)}${notice}`, - diffTruncated: true, - }; -} - -function renderReplacementResult( - result: { content: { type: string; text?: string }[]; details: unknown }, - expanded: boolean, - isPartial: boolean, - theme: { - fg: (name: 'dim' | 'muted' | 'toolDiffAdded' | 'toolDiffRemoved', text: string) => string; - }, -) { - const replacements = numberDetail(result, 'replacements'); - const summary = - replacements === 0 - ? theme.fg('dim', 'No replacements') - : theme.fg('muted', `${replacements} replacement(s)`); - if (isPartial) return renderResultSummary(result, expanded, true, summary); - const details = recordFrom(result.details); - const diff = stringFrom(details?.diff); - if (!diff) return text(summary); - const limit = expanded ? STR_REPLACE_DIFF_MAX_CHARS : STR_REPLACE_COLLAPSED_DIFF_MAX_CHARS; - const preview = diff.length <= limit ? diff : `${diff.slice(0, limit)}\n[Diff preview truncated]`; - return text( - `${summary}\n${preview - .split('\n') - .map((line) => { - if (line.startsWith('+') && !line.startsWith('+++')) { - return theme.fg('toolDiffAdded', line); - } - if (line.startsWith('-') && !line.startsWith('---')) { - return theme.fg('toolDiffRemoved', line); - } - return line; - }) - .join('\n')}`, - ); -} - -function renderPathToolCall(toolName: string, filePath: string, theme: ToolTheme) { - return text(theme.fg('toolTitle', theme.bold(`${toolName} `)) + theme.fg('accent', filePath)); -} - -type FileDetails = { path: string; [key: string]: unknown }; - -function existingPathOrNotFound( - cwd: string, - requestedPath: string, - extraDetails: Omit, -) { - const target = resolve(cwd, requestedPath); - return existsSync(target) ? target : fileNotFound(target, extraDetails); -} - -export function registerFileTools(pi: ExtensionAPI) { - const nativeLs = createLsToolDefinition(process.cwd()); - pi.registerTool({ - ...nativeLs, - name: 'LS', - label: 'LS', - prepareArguments: normalizeLsArgs, - async execute( - toolCallId: string, - params: LsToolInput, - signal: AbortSignal | undefined, - onUpdate: Parameters[3], - ctx: ExtensionContext, - ) { - return createLsToolDefinition(ctx.cwd).execute( - toolCallId, - normalizeLsArgs(params), - signal, - onUpdate, - ctx, - ); - }, - renderCall( - args: unknown, - theme: Parameters>[1], - context: Parameters>[2], - ) { - const normalized = normalizeLsArgs(args); - if (!nativeLs.renderCall) return text(''); - return nativeLs.renderCall(normalized, theme, nativeRenderContext(context, normalized)); - }, - renderResult( - result: Parameters>[0], - options: Parameters>[1], - theme: Parameters>[2], - context: Parameters>[3], - ) { - const normalized = normalizeLsArgs(context.args); - if (!nativeLs.renderResult) return text(''); - return nativeLs.renderResult( - result, - options, - theme, - nativeRenderContext(context, normalized), - ); - }, - }); - - const nativeWrite = createWriteToolDefinition(process.cwd()); - pi.registerTool({ - ...nativeWrite, - name: 'Write', - label: 'Write', - prepareArguments: normalizeWriteArgs, - async execute( - toolCallId: string, - params: WriteToolInput, - signal: AbortSignal | undefined, - onUpdate: Parameters[3], - ctx: ExtensionContext, - ) { - return createWriteToolDefinition(ctx.cwd).execute( - toolCallId, - normalizeWriteArgs(params), - signal, - onUpdate, - ctx, - ); - }, - renderCall( - args: unknown, - theme: Parameters>[1], - context: Parameters>[2], - ) { - const normalized = normalizeWriteArgs(args); - if (!nativeWrite.renderCall) return text(''); - return nativeWrite.renderCall(normalized, theme, nativeRenderContext(context, normalized)); - }, - renderResult( - result: Parameters>[0], - options: Parameters>[1], - theme: Parameters>[2], - context: Parameters>[3], - ) { - const normalized = normalizeWriteArgs(context.args); - if (!nativeWrite.renderResult) return text(''); - return nativeWrite.renderResult( - result, - options, - theme, - nativeRenderContext(context, normalized), - ); - }, - }); - - const StrReplaceParams = Type.Object({ - path: Type.String({ description: 'Path to the file to modify' }), - old_str: Type.String({ description: 'String to search for (exact match)' }), - new_str: Type.String({ description: 'String to replace with' }), - }); - - pi.registerTool({ - name: 'StrReplace', - label: 'StrReplace', - description: - 'Replace all literal occurrences of a string in a file. old_str must not be empty.', - parameters: StrReplaceParams, - prepareArguments: normalizeStrReplaceArgs, - - async execute(_toolCallId, params, signal, _onUpdate, ctx) { - if (signal?.aborted) throw new Error('Operation aborted'); - if (params.old_str === '') throw new Error('old_str must not be empty'); - const filePath = resolve(ctx.cwd, params.path); - return withFileMutationQueue(filePath, async () => { - if (signal?.aborted) throw new Error('Operation aborted'); - const content = await fs.readFile(filePath, 'utf8'); - if (signal?.aborted) throw new Error('Operation aborted'); - const replacements = countOccurrences(content, params.old_str); - if (replacements === 0) { - return { - content: [ - { - type: 'text' as const, - text: `String not found in ${params.path}: "${params.old_str}"`, - }, - ], - details: { path: filePath, replacements: 0 }, - }; - } - - const nextContent = content.replaceAll(params.old_str, () => params.new_str); - const diff = boundedDiff(params.path, content, nextContent); - if (signal?.aborted) throw new Error('Operation aborted'); - await fs.writeFile(filePath, nextContent, 'utf8'); - if (signal?.aborted) throw new Error('Operation aborted'); - return { - content: [ - { - type: 'text' as const, - text: `Replaced ${replacements} occurrence(s) in ${params.path}`, - }, - ], - details: { path: filePath, replacements, ...diff } satisfies StrReplaceDetails, - }; - }); - }, - renderCall(args, theme) { - return renderPathToolCall('StrReplace', args.path, theme); - }, - renderResult(result, { expanded, isPartial }, theme) { - return renderReplacementResult(result, expanded, isPartial, theme); - }, - }); - - const nativeEdit = createEditToolDefinition(process.cwd()); - pi.registerTool({ - ...nativeEdit, - name: 'Edit', - label: 'Edit', - description: - 'Modify a file with exact text replacements. Every oldText must identify one unique, non-overlapping region in the original file. Use StrReplace for replace-all.', - prepareArguments: normalizeEditArgs, - async execute( - toolCallId: string, - params: EditToolInput, - signal: AbortSignal | undefined, - onUpdate: Parameters[3], - ctx: ExtensionContext, - ) { - return createEditToolDefinition(ctx.cwd).execute( - toolCallId, - normalizeEditArgs(params), - signal, - onUpdate, - ctx, - ); - }, - renderCall( - args: unknown, - theme: Parameters>[1], - context: Parameters>[2], - ) { - const normalized = normalizeEditArgs(args); - if (!nativeEdit.renderCall) return text(''); - return nativeEdit.renderCall(normalized, theme, nativeRenderContext(context, normalized)); - }, - renderResult( - result: Parameters>[0], - options: Parameters>[1], - theme: Parameters>[2], - context: Parameters>[3], - ) { - const normalized = normalizeEditArgs(context.args); - if (!nativeEdit.renderResult) return text(''); - return nativeEdit.renderResult( - result, - options, - theme, - nativeRenderContext(context, normalized), - ); - }, - }); - - const DeleteParams = Type.Object({ - path: Type.String({ description: 'Path to the file to delete' }), - }); - - pi.registerTool({ - name: 'Delete', - label: 'Delete', - description: 'Delete a file from the filesystem.', - parameters: DeleteParams, - - async execute(_toolCallId, params, _signal, _onUpdate, ctx) { - const filePath = resolve(ctx.cwd, params.path); - - try { - const resolved = existingPathOrNotFound(ctx.cwd, params.path, { deleted: false }); - if (typeof resolved !== 'string') return resolved; - unlinkSync(resolved); - return { - content: [{ type: 'text', text: `Successfully deleted ${params.path}` }], - details: { path: resolved, deleted: true }, - }; - } catch (error: unknown) { - return fileError(error, 'Delete', filePath, { deleted: false }); - } - }, - renderCall(args, theme) { - return renderPathToolCall('Delete', args.path, theme); - }, - renderResult(result, { expanded, isPartial }, theme) { - return renderResultSummary( - result, - expanded, - isPartial, - booleanDetail(result, 'deleted') - ? theme.fg('muted', 'Deleted') - : theme.fg('error', 'Not deleted'), - ); - }, - }); -} diff --git a/src/tools/glob.ts b/src/tools/glob.ts deleted file mode 100644 index 17a81f1..0000000 --- a/src/tools/glob.ts +++ /dev/null @@ -1,180 +0,0 @@ -import { lstat, readdir, readFile } from 'node:fs/promises'; -import { isAbsolute, join, relative, resolve, sep } from 'node:path'; -import { - DEFAULT_MAX_BYTES, - type FindToolDetails, - type FindToolInput, - formatSize, - truncateHead, -} from '@earendil-works/pi-coding-agent'; -import ignore from 'ignore'; -import { minimatch } from 'minimatch'; -import { normalizePath, recordFrom, stringFrom } from './rendering.js'; - -const DEFAULT_GLOB_LIMIT = 500; - -type IgnoreScope = { - root: string; - matcher: ReturnType; -}; - -type Candidate = { - absolutePath: string; - outputPath: string; - mtimeMs: number; -}; - -type GlobResult = { - content: [{ type: 'text'; text: string }]; - details: FindToolDetails | undefined; -}; - -function errorCode(error: unknown) { - if (!error || typeof error !== 'object' || !('code' in error)) return undefined; - return error.code; -} - -function throwIfAborted(signal?: AbortSignal) { - if (signal?.aborted) throw new Error('Operation aborted'); -} - -function compareCandidates(left: Candidate, right: Candidate) { - const timeDifference = right.mtimeMs - left.mtimeMs; - if (timeDifference !== 0) return timeDifference; - if (left.outputPath < right.outputPath) return -1; - if (left.outputPath > right.outputPath) return 1; - return 0; -} - -function outputPath(cwd: string, absolutePath: string) { - const workspaceRelative = relative(cwd, absolutePath); - const outsideWorkspace = - workspaceRelative === '..' || - workspaceRelative.startsWith(`..${sep}`) || - isAbsolute(workspaceRelative); - return outsideWorkspace ? absolutePath : normalizePath(workspaceRelative); -} - -function ignoredByScopes(absolutePath: string, directory: boolean, scopes: IgnoreScope[]) { - return scopes.reduce((ignored, scope) => { - const scopedPath = normalizePath(relative(scope.root, absolutePath)); - if (!scopedPath || scopedPath === '..' || scopedPath.startsWith('../')) return ignored; - const result = scope.matcher.test(directory ? `${scopedPath}/` : scopedPath); - if (result.unignored) return false; - if (result.ignored) return true; - return ignored; - }, false); -} - -async function scopesForDirectory(directory: string, scopes: IgnoreScope[], signal?: AbortSignal) { - throwIfAborted(signal); - const contents = await readFile(join(directory, '.gitignore'), 'utf8').catch((error: unknown) => { - if (errorCode(error) === 'ENOENT') return undefined; - throw error; - }); - if (contents === undefined) return scopes; - return [...scopes, { root: directory, matcher: ignore().add(contents) }]; -} - -export function normalizeGlobArgs(value: unknown): FindToolInput { - const input = recordFrom(value); - return { - pattern: stringFrom(input?.pattern) ?? stringFrom(input?.glob_pattern) ?? '', - path: stringFrom(input?.path) ?? stringFrom(input?.target_directory), - limit: typeof input?.limit === 'number' ? input.limit : undefined, - }; -} - -export async function executeGlob( - input: FindToolInput, - cwd: string, - signal?: AbortSignal, -): Promise { - const searchRoot = resolve(cwd, input.path ?? '.'); - const limit = - typeof input.limit === 'number' && Number.isFinite(input.limit) - ? Math.max(1, Math.floor(input.limit)) - : DEFAULT_GLOB_LIMIT; - const state = { totalMatches: 0, candidates: [] as Candidate[] }; - - throwIfAborted(signal); - const rootStats = await lstat(searchRoot).catch((error: unknown) => { - if (errorCode(error) === 'ENOENT') throw new Error(`Path not found: ${searchRoot}`); - throw error; - }); - if (!rootStats.isDirectory() || rootStats.isSymbolicLink()) { - throw new Error(`Not a directory: ${searchRoot}`); - } - - const visit = async (directory: string, inheritedScopes: IgnoreScope[]): Promise => { - throwIfAborted(signal); - const scopes = await scopesForDirectory(directory, inheritedScopes, signal); - throwIfAborted(signal); - const entries = await readdir(directory, { withFileTypes: true }); - - for (const entry of entries) { - const absolutePath = join(directory, entry.name); - throwIfAborted(signal); - const stats = await lstat(absolutePath).catch((error: unknown) => { - if (errorCode(error) === 'ENOENT') return undefined; - throw error; - }); - if (!stats || stats.isSymbolicLink()) continue; - - if (stats.isDirectory()) { - if (entry.name === '.git' || entry.name === 'node_modules') continue; - if (ignoredByScopes(absolutePath, true, scopes)) continue; - await visit(absolutePath, scopes); - continue; - } - if (!stats.isFile() || ignoredByScopes(absolutePath, false, scopes)) continue; - - const searchRelative = normalizePath(relative(searchRoot, absolutePath)); - if (!minimatch(searchRelative, input.pattern, { dot: true })) continue; - state.totalMatches += 1; - state.candidates.push({ - absolutePath, - outputPath: outputPath(cwd, absolutePath), - mtimeMs: stats.mtimeMs, - }); - state.candidates.sort(compareCandidates); - if (state.candidates.length > limit + 1) state.candidates.pop(); - } - }; - - await visit(searchRoot, []); - throwIfAborted(signal); - if (state.totalMatches === 0) { - return { content: [{ type: 'text', text: 'No files found' }], details: undefined }; - } - - const rawOutput = state.candidates - .sort(compareCandidates) - .slice(0, limit) - .map((candidate) => candidate.outputPath) - .join('\n'); - const truncation = truncateHead(rawOutput, { maxLines: Number.MAX_SAFE_INTEGER }); - const details: FindToolDetails = {}; - const notices: string[] = []; - if (state.totalMatches > limit) { - details.resultLimitReached = limit; - notices.push(`${limit} results limit reached. Use limit=${limit * 2} for more`); - } - if (truncation.truncated) { - details.truncation = truncation; - notices.push(`${formatSize(DEFAULT_MAX_BYTES)} limit reached`); - } - throwIfAborted(signal); - return { - content: [ - { - type: 'text', - text: - notices.length === 0 - ? truncation.content - : `${truncation.content}\n\n[${notices.join('. ')}]`, - }, - ], - details: Object.keys(details).length === 0 ? undefined : details, - }; -} diff --git a/src/tools/read.ts b/src/tools/read.ts deleted file mode 100644 index cbea77b..0000000 --- a/src/tools/read.ts +++ /dev/null @@ -1,70 +0,0 @@ -import { createReadToolDefinition, type ExtensionContext } from '@earendil-works/pi-coding-agent'; -import { nativeRenderContext, recordFrom, text } from './rendering.js'; - -type ReadInput = { - path: string; - offset?: number; - limit?: number; -}; - -/** - * Cursor/Grok-trained models call a capital-`Read` tool. Rather than - * reimplement file reading (which loses the native tool's image support, - * byte-aware truncation, and 1-indexed offset), this registers `Read` as a - * thin alias over pi's built-in `read` tool definition. - * - * The native `read` resolves paths against the `cwd` captured at construction - * time (it does not read `ctx.cwd`), so `execute` rebuilds the definition per - * call with the live `ctx.cwd`. Everything else — `parameters`, `description`, - * `renderCall`, `renderResult`, prompt guidance — is delegated to the native - * tool; only `name`/`label`/`prepareArguments` differ. - */ -export function createReadShim() { - const native = createReadToolDefinition(process.cwd()); - - const normalize = (args: unknown): ReadInput => { - const input = recordFrom(args); - if (!input) return { path: '' }; - const path = typeof input.path === 'string' ? input.path : input.file_path; - return { - path: typeof path === 'string' ? path : '', - offset: typeof input.offset === 'number' ? input.offset : undefined, - limit: typeof input.limit === 'number' ? input.limit : undefined, - }; - }; - - return { - ...native, - name: 'Read', - label: 'Read', - prepareArguments: normalize, - async execute( - toolCallId: string, - params: ReadInput, - signal: AbortSignal | undefined, - onUpdate: Parameters[3], - ctx: ExtensionContext, - ) { - return createReadToolDefinition(ctx.cwd).execute(toolCallId, params, signal, onUpdate, ctx); - }, - renderCall( - args: unknown, - theme: Parameters>[1], - context: Parameters>[2], - ) { - const normalized = normalize(args); - if (!native.renderCall) return text(''); - return native.renderCall(normalized, theme, nativeRenderContext(context, normalized)); - }, - renderResult( - result: Parameters>[0], - options: Parameters>[1], - theme: Parameters>[2], - context: Parameters>[3], - ) { - const normalized = normalize(context.args); - if (!native.renderResult) return text(''); - return native.renderResult(result, options, theme, nativeRenderContext(context, normalized)); - }, - }; -} diff --git a/src/tools/register.ts b/src/tools/register.ts deleted file mode 100644 index 55664f4..0000000 --- a/src/tools/register.ts +++ /dev/null @@ -1,30 +0,0 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { registerFileTools } from './files.js'; -import { createReadShim } from './read.js'; -import { registerSearchTools } from './search.js'; -import { registerShellTool } from './shell.js'; -import { registerWebSearchTool } from './webSearch.js'; -import { isPiWebAccessInstalled } from './webSearchDelegate.js'; - -/** Grok/Cursor shims always registered by this extension (excludes optional WebSearch). */ -export const GROK_SHIM_TOOL_NAMES = [ - 'Grep', - 'Glob', - 'LS', - 'Read', - 'Write', - 'StrReplace', - 'Edit', - 'Delete', - 'Shell', -] as const; - -export function registerGrokTools(pi: ExtensionAPI) { - const webSearchRegistered = isPiWebAccessInstalled(); - if (webSearchRegistered) registerWebSearchTool(pi); - registerSearchTools(pi); - registerFileTools(pi); - pi.registerTool(createReadShim()); - registerShellTool(pi); - return { webSearchRegistered }; -} diff --git a/src/tools/rendering.ts b/src/tools/rendering.ts deleted file mode 100644 index d5b18bd..0000000 --- a/src/tools/rendering.ts +++ /dev/null @@ -1,117 +0,0 @@ -import { type Component, Text } from '@earendil-works/pi-tui'; - -export type NativeRenderContext = { - args: TArgs; - toolCallId: string; - invalidate: () => void; - lastComponent: Component | undefined; - state: TState; - cwd: string; - executionStarted: boolean; - argsComplete: boolean; - isPartial: boolean; - expanded: boolean; - showImages: boolean; - isError: boolean; -}; - -export function nativeRenderContext( - context: NativeRenderContext, - args: TNativeArgs, -): NativeRenderContext { - return { ...context, args }; -} - -export function recordFrom(value: unknown): Record | undefined { - if (!value || typeof value !== 'object') return undefined; - return value as Record; -} - -export function stringFrom(value: unknown): string | undefined { - if (typeof value !== 'string') return undefined; - return value; -} - -export function normalizePath(filePath: string) { - return filePath.replaceAll('\\', '/'); -} - -export function text(value: string): Text { - return new Text(value, 0, 0); -} - -function firstText(result: { content: { type: string; text?: string }[] }) { - const first = result.content[0]; - if (first?.type !== 'text') return undefined; - return first.text; -} - -function renderResultText( - result: { content: { type: string; text?: string }[] }, - expanded: boolean, - summary: string, -) { - return text(expanded ? (firstText(result) ?? summary) : summary); -} - -export function renderRunning(isPartial: boolean): Text | undefined { - if (!isPartial) return undefined; - return text('Running...'); -} - -export function renderResultSummary( - result: { content: { type: string; text?: string }[] }, - expanded: boolean, - isPartial: boolean, - summary: string, -) { - return renderRunning(isPartial) ?? renderResultText(result, expanded, summary); -} - -function detailRecord(result: { details: unknown }): Record { - if (!result.details || typeof result.details !== 'object') return {}; - return result.details as Record; -} - -export function numberDetail(result: { details: unknown }, key: string): number { - const value = detailRecord(result)[key]; - if (typeof value !== 'number') return 0; - return value; -} - -export function booleanDetail(result: { details: unknown }, key: string): boolean { - return detailRecord(result)[key] === true; -} - -type FileDetails = { path: string; [key: string]: unknown }; - -type FileResult = { - content: [{ type: 'text'; text: string }]; - details: T; -}; - -export function fileNotFound( - filePath: string, - extraDetails: Omit, -): FileResult { - return { - content: [{ type: 'text', text: `File not found: ${filePath}` }], - details: { path: filePath, ...extraDetails } as T, - }; -} - -export function fileError( - error: unknown, - toolName: string, - filePath: string, - extraDetails: Omit, -): FileResult { - const message = - error && typeof error === 'object' && 'message' in error && typeof error.message === 'string' - ? error.message - : 'Unknown error'; - return { - content: [{ type: 'text', text: `${toolName} error: ${message}` }], - details: { path: filePath, ...extraDetails, failed: true, error: message } as unknown as T, - }; -} diff --git a/src/tools/search.ts b/src/tools/search.ts deleted file mode 100644 index 6b7f534..0000000 --- a/src/tools/search.ts +++ /dev/null @@ -1,125 +0,0 @@ -import { Type } from '@earendil-works/pi-ai'; -import { - createFindToolDefinition, - createGrepToolDefinition, - type ExtensionAPI, - type ExtensionContext, - type GrepToolInput, -} from '@earendil-works/pi-coding-agent'; -import { executeGlob, normalizeGlobArgs } from './glob.js'; -import { nativeRenderContext, recordFrom, stringFrom, text } from './rendering.js'; - -function normalizeGrepArgs(value: unknown): GrepToolInput { - const input = recordFrom(value); - return { - pattern: stringFrom(input?.pattern) ?? stringFrom(input?.query) ?? '', - path: stringFrom(input?.path), - glob: stringFrom(input?.glob) ?? stringFrom(input?.include) ?? stringFrom(input?.glob_filter), - ignoreCase: - typeof input?.ignoreCase === 'boolean' - ? input.ignoreCase - : typeof input?.ignore_case === 'boolean' - ? input.ignore_case - : undefined, - literal: typeof input?.literal === 'boolean' ? input.literal : undefined, - context: typeof input?.context === 'number' ? input.context : undefined, - limit: typeof input?.limit === 'number' ? input.limit : undefined, - }; -} - -export function registerSearchTools(pi: ExtensionAPI) { - const nativeGrep = createGrepToolDefinition(process.cwd()); - pi.registerTool({ - ...nativeGrep, - name: 'Grep', - label: 'Grep', - prepareArguments: normalizeGrepArgs, - async execute( - toolCallId: string, - params: GrepToolInput, - signal: AbortSignal | undefined, - onUpdate: Parameters[3], - ctx: ExtensionContext, - ) { - return createGrepToolDefinition(ctx.cwd).execute( - toolCallId, - normalizeGrepArgs(params), - signal, - onUpdate, - ctx, - ); - }, - renderCall( - args: unknown, - theme: Parameters>[1], - context: Parameters>[2], - ) { - const normalized = normalizeGrepArgs(args); - if (!nativeGrep.renderCall) return text(''); - return nativeGrep.renderCall(normalized, theme, nativeRenderContext(context, normalized)); - }, - renderResult( - result: Parameters>[0], - options: Parameters>[1], - theme: Parameters>[2], - context: Parameters>[3], - ) { - const normalized = normalizeGrepArgs(context.args); - if (!nativeGrep.renderResult) return text(''); - return nativeGrep.renderResult( - result, - options, - theme, - nativeRenderContext(context, normalized), - ); - }, - }); - - const GlobParams = Type.Object({ - pattern: Type.String({ - description: 'Glob pattern to match files (e.g. **/*.ts, src/**/*.json)', - }), - path: Type.Optional( - Type.String({ - description: 'Directory to search within. Defaults to current working directory.', - }), - ), - limit: Type.Optional(Type.Number({ description: 'Maximum number of files to return.' })), - }); - const nativeFind = createFindToolDefinition(process.cwd()); - pi.registerTool({ - name: 'Glob', - label: 'Glob', - description: - 'Find files matching a glob pattern, newest first. Returns workspace-relative or absolute paths.', - parameters: GlobParams, - prepareArguments: normalizeGlobArgs, - async execute(_toolCallId, params, signal, _onUpdate, ctx) { - return executeGlob(normalizeGlobArgs(params), ctx.cwd, signal); - }, - renderCall( - args: unknown, - theme: Parameters>[1], - context: Parameters>[2], - ) { - const normalized = normalizeGlobArgs(args); - if (!nativeFind.renderCall) return text(''); - return nativeFind.renderCall(normalized, theme, nativeRenderContext(context, normalized)); - }, - renderResult( - result: Parameters>[0], - options: Parameters>[1], - theme: Parameters>[2], - context: Parameters>[3], - ) { - const normalized = normalizeGlobArgs(context.args); - if (!nativeFind.renderResult) return text(''); - return nativeFind.renderResult( - result, - options, - theme, - nativeRenderContext(context, normalized), - ); - }, - }); -} diff --git a/src/tools/shell.ts b/src/tools/shell.ts deleted file mode 100644 index 3dfee95..0000000 --- a/src/tools/shell.ts +++ /dev/null @@ -1,87 +0,0 @@ -import { resolve } from 'node:path'; -import { Type } from '@earendil-works/pi-ai'; -import { - type BashToolInput, - createBashToolDefinition, - type ExtensionAPI, - type ExtensionContext, -} from '@earendil-works/pi-coding-agent'; -import { nativeRenderContext, recordFrom, stringFrom, text } from './rendering.js'; - -type ShellArgs = { - command: string; - working_directory?: string; - timeout: number; -}; - -function normalizeShellArgs(value: unknown): ShellArgs { - const input = recordFrom(value); - return { - command: stringFrom(input?.command) ?? stringFrom(input?.cmd) ?? '', - timeout: typeof input?.timeout === 'number' ? input.timeout : 120_000, - working_directory: stringFrom(input?.working_directory), - }; -} - -function nativeShellArgs(value: unknown): BashToolInput { - const input = normalizeShellArgs(value); - return { command: input.command, timeout: input.timeout / 1000 }; -} - -export function registerShellTool(pi: ExtensionAPI) { - const ShellParams = Type.Object({ - command: Type.String({ description: 'Shell command to execute' }), - working_directory: Type.Optional( - Type.String({ description: 'Working directory for the command' }), - ), - timeout: Type.Optional( - Type.Number({ description: 'Timeout in milliseconds (default: 120000)' }), - ), - }); - const nativeBash = createBashToolDefinition(process.cwd()); - - pi.registerTool({ - name: 'Shell', - label: 'Shell', - description: - 'Execute a shell command with native streaming, process-tree cancellation, and bounded output.', - parameters: ShellParams, - prepareArguments: normalizeShellArgs, - async execute( - toolCallId: string, - params: ShellArgs, - signal: AbortSignal | undefined, - onUpdate: Parameters[3], - ctx: ExtensionContext, - ) { - const normalized = normalizeShellArgs(params); - return createBashToolDefinition( - normalized.working_directory ? resolve(ctx.cwd, normalized.working_directory) : ctx.cwd, - ).execute(toolCallId, nativeShellArgs(normalized), signal, onUpdate, ctx); - }, - renderCall( - args: unknown, - theme: Parameters>[1], - context: Parameters>[2], - ) { - const normalized = nativeShellArgs(args); - if (!nativeBash.renderCall) return text(''); - return nativeBash.renderCall(normalized, theme, nativeRenderContext(context, normalized)); - }, - renderResult( - result: Parameters>[0], - options: Parameters>[1], - theme: Parameters>[2], - context: Parameters>[3], - ) { - const normalized = nativeShellArgs(context.args); - if (!nativeBash.renderResult) return text(''); - return nativeBash.renderResult( - result, - options, - theme, - nativeRenderContext(context, normalized), - ); - }, - }); -} diff --git a/src/tools/webSearch.ts b/src/tools/webSearch.ts deleted file mode 100644 index dee90a4..0000000 --- a/src/tools/webSearch.ts +++ /dev/null @@ -1,157 +0,0 @@ -import { StringEnum, Type } from '@earendil-works/pi-ai'; -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { Text } from '@earendil-works/pi-tui'; -import { renderRunning, text } from './rendering.js'; -import { - ensureWebSearchDelegate, - getWebSearchDelegate, - getWebSearchLoadError, -} from './webSearchDelegate.js'; - -const WEB_SEARCH_DESCRIPTION = - 'Search the web using Perplexity AI, Exa, or Gemini. Returns an AI-synthesized answer with source citations. For comprehensive research, prefer queries (plural) with 2-4 varied angles over a single query — each query gets its own synthesized answer, so varying phrasing and scope gives much broader coverage. When includeContent is true, full page content is fetched in the background. Searches auto-open the interactive browser curator and stream results live; set workflow to "none" to skip curation. Provider auto-selects: Exa (direct API with key, MCP fallback without), else Perplexity (needs key), else Gemini API (needs key), else Gemini Web (needs a supported Chromium-based browser login).'; - -const WebSearchParams = Type.Object({ - query: Type.Optional( - Type.String({ - description: - "Single search query. For research tasks, prefer 'queries' with multiple varied angles instead.", - }), - ), - queries: Type.Optional( - Type.Array(Type.String(), { - description: - 'Multiple queries searched in sequence, each returning its own synthesized answer. Prefer this for research — vary phrasing, scope, and angle across 2-4 queries to maximize coverage.', - }), - ), - numResults: Type.Optional( - Type.Number({ description: 'Results per query (default: 5, max: 20)' }), - ), - includeContent: Type.Optional(Type.Boolean({ description: 'Fetch full page content (async)' })), - recencyFilter: Type.Optional( - StringEnum(['day', 'week', 'month', 'year'], { description: 'Filter by recency' }), - ), - domainFilter: Type.Optional( - Type.Array(Type.String(), { - description: 'Limit to domains (prefix with - to exclude)', - }), - ), - provider: Type.Optional( - StringEnum(['auto', 'perplexity', 'gemini', 'exa'], { - description: 'Search provider (default: auto)', - }), - ), - workflow: Type.Optional( - StringEnum(['none', 'summary-review'], { - description: - 'Search workflow mode: none = no curator, summary-review = open curator with auto summary draft (default)', - }), - ), -}); - -function normalizeQueryList(raw: unknown[]): string[] { - return raw - .filter((q): q is string => typeof q === 'string') - .map((q) => q.trim()) - .filter((q) => q.length > 0); -} - -function queryListFromArgs(args: Record) { - const raw: unknown[] = Array.isArray(args.queries) - ? args.queries - : args.query !== undefined - ? [args.query] - : []; - return normalizeQueryList(raw); -} - -function missingDelegateMessage() { - const reason = - getWebSearchLoadError() ?? - 'pi-web-access web_search delegate not available. Install with: pi install npm:pi-web-access'; - return { - content: [ - { - type: 'text' as const, - text: `WebSearch requires pi-web-access: ${reason}`, - }, - ], - details: { error: reason }, - }; -} - -export function registerWebSearchTool(pi: ExtensionAPI) { - pi.registerTool({ - name: 'WebSearch', - label: 'Web Search', - description: WEB_SEARCH_DESCRIPTION, - promptSnippet: - 'Use for web research questions. Prefer {queries:[...]} with 2-4 varied angles over a single query for broader coverage.', - parameters: WebSearchParams, - - async execute(toolCallId, params, signal, onUpdate, ctx) { - await ensureWebSearchDelegate(pi); - const delegate = getWebSearchDelegate(); - if (!delegate) return missingDelegateMessage(); - const normalizedParams = { ...(params as Record) }; - if (Array.isArray(normalizedParams.queries)) { - normalizedParams.queries = normalizeQueryList(normalizedParams.queries); - } - if (typeof normalizedParams.query === 'string') { - const query = normalizedParams.query.trim(); - if (query) normalizedParams.query = query; - if (!query) delete normalizedParams.query; - } - return delegate(toolCallId, normalizedParams, signal, onUpdate, ctx); - }, - - renderCall(args, theme) { - const queryList = queryListFromArgs(args as Record); - if (queryList.length === 0) { - return text( - theme.fg('toolTitle', theme.bold('WebSearch ')) + theme.fg('error', '(no query)'), - ); - } - if (queryList.length === 1) { - const q = queryList[0]; - const display = q.length > 60 ? `${q.slice(0, 57)}...` : q; - return text( - theme.fg('toolTitle', theme.bold('WebSearch ')) + theme.fg('accent', `"${display}"`), - ); - } - const lines = [ - theme.fg('toolTitle', theme.bold('WebSearch ')) + - theme.fg('accent', `${queryList.length} queries`), - ]; - for (const q of queryList.slice(0, 5)) { - const display = q.length > 50 ? `${q.slice(0, 47)}...` : q; - lines.push(theme.fg('muted', ` "${display}"`)); - } - if (queryList.length > 5) { - lines.push(theme.fg('muted', ` ... and ${queryList.length - 5} more`)); - } - return new Text(lines.join('\n'), 0, 0); - }, - - renderResult(result, { expanded, isPartial }, theme) { - const running = renderRunning(isPartial); - if (running) return running; - - const details = result.details as { error?: string; totalResults?: number } | undefined; - if (details?.error) { - return text(theme.fg('error', `Error: ${details.error}`)); - } - - const summary = - typeof details?.totalResults === 'number' - ? theme.fg('success', `${details.totalResults} sources`) - : theme.fg('success', 'search complete'); - - if (!expanded) return text(summary); - - const textContent = result.content.find((c) => c.type === 'text')?.text ?? ''; - const preview = textContent.length > 800 ? `${textContent.slice(0, 800)}...` : textContent; - return new Text(`${summary}\n${theme.fg('dim', preview)}`, 0, 0); - }, - }); -} diff --git a/src/tools/webSearchDelegate.ts b/src/tools/webSearchDelegate.ts deleted file mode 100644 index 6342317..0000000 --- a/src/tools/webSearchDelegate.ts +++ /dev/null @@ -1,149 +0,0 @@ -import { existsSync } from 'node:fs'; -import { homedir } from 'node:os'; -import { join } from 'node:path'; -import { pathToFileURL } from 'node:url'; -import type { AgentToolResult, AgentToolUpdateCallback } from '@earendil-works/pi-agent-core'; -import type { ExtensionAPI, ToolDefinition } from '@earendil-works/pi-coding-agent'; -import { getAgentDir } from '@earendil-works/pi-coding-agent'; - -export const PI_WEB_SEARCH_TOOL = 'web_search'; -const MINIMUM_PI_WEB_ACCESS_VERSION = '0.13.0'; - -export type WebSearchExecute = ( - toolCallId: string, - params: Record, - signal: AbortSignal | undefined, - onUpdate: AgentToolUpdateCallback | undefined, - ctx: import('@earendil-works/pi-coding-agent').ExtensionContext, -) => Promise>; - -let webSearchExecute: WebSearchExecute | undefined; -let loadPromise: Promise | undefined; -let lastLoadError: string | undefined; -let boundLivePi: ExtensionAPI | undefined; -let bindGeneration = 0; -let importGeneration = 0; - -export function getWebSearchLoadError() { - return lastLoadError; -} - -function isCurrentBinding(pi: ExtensionAPI, generation: number) { - return bindGeneration === generation && (!boundLivePi || boundLivePi === pi); -} - -function resolvePiWebAccessEntry() { - for (const dir of [ - join(getAgentDir(), 'npm', 'node_modules', 'pi-web-access'), - join(homedir(), '.pi', 'agent', 'npm', 'node_modules', 'pi-web-access'), - ]) { - for (const file of ['index.ts', 'index.js']) { - const entry = join(dir, file); - if (existsSync(entry)) return entry; - } - } - return undefined; -} - -export function isPiWebAccessInstalled() { - return resolvePiWebAccessEntry() !== undefined; -} - -function createToolRegistrationAdapter( - pi: ExtensionAPI, - registerTool: (tool: ToolDefinition) => void, -) { - return new Proxy(pi, { - get(target, property, receiver) { - if (property === 'registerTool') return registerTool; - if (property === 'registerCommand' || property === 'registerShortcut' || property === 'on') { - return () => undefined; - } - return Reflect.get(target, property, receiver); - }, - }); -} - -async function captureWebSearchFromPublicFactory(pi: ExtensionAPI, generation: number) { - const entry = resolvePiWebAccessEntry(); - if (!entry) return; - - const module = (await import( - `${pathToFileURL(entry).href}?pi-grok-cli-load=${++importGeneration}` - )) as Record; - if (typeof module.default !== 'function') { - throw new Error( - `pi-web-access is incompatible. Install pi-web-access ${MINIMUM_PI_WEB_ACCESS_VERSION} or newer with a public default extension factory.`, - ); - } - - let webSearch: ToolDefinition | undefined; - await (module.default as (pi: ExtensionAPI) => void | Promise)( - createToolRegistrationAdapter(pi, (tool) => { - if (tool.name === PI_WEB_SEARCH_TOOL) webSearch = tool; - }), - ); - - if (!webSearch) { - if (!isCurrentBinding(pi, generation)) return; - lastLoadError = 'pi-web-access loaded but did not register web_search. Update pi-web-access.'; - return; - } - if (!isCurrentBinding(pi, generation)) return; - webSearchExecute = webSearch.execute.bind(webSearch) as WebSearchExecute; - lastLoadError = undefined; -} - -/** Remember the live session ExtensionAPI (bound after session_start). */ -export function bindLivePiWebAccess(pi: ExtensionAPI) { - bindGeneration += 1; - boundLivePi = pi; - webSearchExecute = undefined; - loadPromise = undefined; -} - -export async function ensureWebSearchDelegate( - pi?: ExtensionAPI, - isInstalled: () => boolean = isPiWebAccessInstalled, -) { - if (!isInstalled()) return; - - const livePi = pi ?? boundLivePi; - if (!livePi) return; - - const generation = bindGeneration; - if (!isCurrentBinding(livePi, generation) || webSearchExecute) return; - if (loadPromise) return loadPromise; - - loadPromise = (async () => { - if (isCurrentBinding(livePi, generation)) lastLoadError = undefined; - try { - await captureWebSearchFromPublicFactory(livePi, generation); - } catch (error) { - if (!isCurrentBinding(livePi, generation)) return; - lastLoadError = error instanceof Error ? error.message : String(error); - webSearchExecute = undefined; - } finally { - if (isCurrentBinding(livePi, generation)) loadPromise = undefined; - } - })(); - - return loadPromise; -} - -export function getWebSearchDelegate() { - return webSearchExecute; -} - -export function clearWebSearchDelegateForTests() { - bindGeneration += 1; - webSearchExecute = undefined; - loadPromise = undefined; - lastLoadError = undefined; - boundLivePi = undefined; -} - -export function setWebSearchDelegateForTests(execute: WebSearchExecute) { - webSearchExecute = execute; - lastLoadError = undefined; -} diff --git a/src/vision/cache.ts b/src/vision/cache.ts deleted file mode 100644 index 17930e2..0000000 --- a/src/vision/cache.ts +++ /dev/null @@ -1,116 +0,0 @@ -import { createHash } from 'node:crypto'; -import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'; -import { dirname } from 'node:path'; -import { getLegacyVisionCachePath, getVisionCachePath } from '../storage.js'; - -export const getCachePath = () => - existsSync(getVisionCachePath()) || !existsSync(getLegacyVisionCachePath()) - ? getVisionCachePath() - : getLegacyVisionCachePath(); - -export interface CacheEntry { - createdAt: string; - description: string; - imageHash: string; - mediaType: string; - model: string; - promptHash: string; -} - -export interface CacheFile { - version: 1; - entries: Record; -} - -export interface VisionImage { - data: string; - mimeType: string; -} - -function emptyCache(): CacheFile { - return { version: 1, entries: {} }; -} - -const cacheUpdates = new Map>(); - -export function loadCache(cachePath: string): CacheFile { - try { - const raw = JSON.parse(readFileSync(cachePath, 'utf-8')); - if ( - raw?.version === 1 && - raw.entries && - typeof raw.entries === 'object' && - !Array.isArray(raw.entries) - ) { - return raw as CacheFile; - } - } catch { - // Missing or invalid cache: start fresh. - } - return emptyCache(); -} - -export function saveCache(cache: CacheFile, cachePath: string) { - mkdirSync(dirname(cachePath), { recursive: true }); - writeFileSync(cachePath, JSON.stringify(cache, null, 2)); -} - -export async function updateCache(cachePath: string, update: (cache: CacheFile) => void) { - const previous = cacheUpdates.get(cachePath) ?? Promise.resolve(); - let release = () => {}; - const current = new Promise((resolve) => { - release = resolve; - }); - const next = previous.then(() => current); - cacheUpdates.set(cachePath, next); - - await previous; - try { - const cache = loadCache(cachePath); - update(cache); - saveCache(cache, cachePath); - } finally { - release(); - if (cacheUpdates.get(cachePath) === next) cacheUpdates.delete(cachePath); - } -} - -export function clearCache(cachePath: string) { - saveCache(emptyCache(), cachePath); -} - -export function cacheStats(cachePath: string): { entries: number; path: string } { - return { entries: Object.keys(loadCache(cachePath).entries).length, path: cachePath }; -} - -function hash(value: string | Buffer): string { - return createHash('sha256').update(value).digest('hex'); -} - -export function makeCacheKey(img: VisionImage, model: string, prompt: string): string { - const imageHash = hash(Buffer.from(img.data, 'base64')); - return hash(JSON.stringify({ imageHash, mediaType: img.mimeType, model, prompt })); -} - -export function makeCacheEntry( - img: VisionImage, - model: string, - prompt: string, - description: string, -): CacheEntry { - return { - createdAt: new Date().toISOString(), - description, - imageHash: hash(Buffer.from(img.data, 'base64')), - mediaType: img.mimeType || 'unknown', - model, - promptHash: hash(prompt), - }; -} - -export function pruneCache(cache: CacheFile, maxEntries: number) { - const entries = Object.entries(cache.entries).sort( - ([, a], [, b]) => Date.parse(b.createdAt) - Date.parse(a.createdAt), - ); - cache.entries = Object.fromEntries(entries.slice(0, maxEntries)); -} diff --git a/src/vision/describe.ts b/src/vision/describe.ts deleted file mode 100644 index f830111..0000000 --- a/src/vision/describe.ts +++ /dev/null @@ -1,358 +0,0 @@ -import type { ImageContent, TextContent } from '@earendil-works/pi-ai'; -import type { ExtensionContext, ToolResultEvent } from '@earendil-works/pi-coding-agent'; -import { getBaseUrl } from '../auth/oauth.js'; -import { loadConfig, type VisionConfig } from '../config.js'; -import { resolveGrokToken } from '../provider/accounts.js'; -import { grokCliModelHeaders } from '../provider/stream.js'; -import { - getCachePath, - loadCache, - makeCacheEntry, - makeCacheKey, - pruneCache, - updateCache, - type VisionImage, -} from './cache.js'; - -const REQUEST_TIMEOUT_MS = 30_000; -const MAX_ATTEMPTS = 3; -const RETRY_BASE_DELAY_MS = 500; -const DEFAULT_PROMPT = - 'Describe this image in detail. If it contains text, transcribe it exactly. ' + - 'If it shows code, reproduce it. If it shows a UI, describe layout and elements. ' + - 'Respond in the same language as any text in the image.'; - -function asObject(value: unknown): Record | undefined { - return typeof value === 'object' && value !== null - ? (value as Record) - : undefined; -} - -function createTimeoutSignal(parent?: AbortSignal): { signal: AbortSignal; cleanup: () => void } { - const controller = new AbortController(); - const timeout = setTimeout( - () => controller.abort(new Error('request timed out')), - REQUEST_TIMEOUT_MS, - ); - - const onAbort = () => controller.abort(parent?.reason ?? new Error('request cancelled')); - if (parent) { - if (parent.aborted) onAbort(); - else parent.addEventListener('abort', onAbort, { once: true }); - } - - return { - signal: controller.signal, - cleanup: () => { - clearTimeout(timeout); - parent?.removeEventListener('abort', onAbort); - }, - }; -} - -function sleep(ms: number, signal?: AbortSignal): Promise { - return new Promise((resolve, reject) => { - const timeout = setTimeout(resolve, ms); - if (!signal) return; - - const onAbort = () => { - clearTimeout(timeout); - reject(new Error('request cancelled')); - }; - if (signal.aborted) { - onAbort(); - return; - } - signal.addEventListener('abort', onAbort, { once: true }); - }); -} - -function isRetryableStatus(status: number): boolean { - return status === 408 || status === 409 || status === 425 || status === 429 || status >= 500; -} - -async function readErrorBody(res: Response): Promise { - const text = await res.text().catch(() => ''); - if (!text) return ''; - try { - const json = JSON.parse(text); - return json?.error?.message ?? json?.message ?? text; - } catch { - return text; - } -} - -function explainHttpError(status: number, body: string, model: string): string { - const detail = body ? `: ${body.slice(0, 500)}` : ''; - if (status === 401 || status === 403) { - return `Grok CLI rejected the API key (HTTP ${status}). Re-run /login grok-cli or set GROK_CLI_OAUTH_TOKEN${detail}`; - } - if (status === 400 || status === 404) { - return `Grok CLI rejected model "${model}" (HTTP ${status})${detail}`; - } - if (status === 429) { - return `Grok CLI rate limited the request (HTTP 429). Try again later${detail}`; - } - if (status >= 500) { - return `Grok CLI service error (HTTP ${status}). Retried automatically; try again later if it persists${detail}`; - } - return `Grok CLI request failed (HTTP ${status})${detail}`; -} - -function explainFetchError(err: unknown): string { - const message = err instanceof Error ? err.message : String(err); - if (err instanceof Error && (err.name === 'AbortError' || /timed out/i.test(message))) { - return `Grok CLI request timed out after ${REQUEST_TIMEOUT_MS / 1000}s`; - } - if (/cancelled|aborted/i.test(message)) { - return 'Grok CLI request was cancelled'; - } - return `Grok CLI network request failed: ${message}`; -} - -async function fetchWithRetry(url: string, init: RequestInit, model: string): Promise { - let lastError: string | undefined; - - for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt += 1) { - const { signal, cleanup } = createTimeoutSignal(init.signal ?? undefined); - try { - const res = await fetch(url, { ...init, signal }); - if (res.ok || !isRetryableStatus(res.status) || attempt === MAX_ATTEMPTS) { - return res; - } - lastError = explainHttpError(res.status, await readErrorBody(res), model); - } catch (err) { - lastError = explainFetchError(err); - if (attempt === MAX_ATTEMPTS || /cancelled|aborted/i.test(lastError)) { - throw new Error(`${lastError} after ${attempt} attempt${attempt === 1 ? '' : 's'}.`); - } - } finally { - cleanup(); - } - - await sleep(RETRY_BASE_DELAY_MS * 2 ** (attempt - 1), init.signal ?? undefined); - } - - throw new Error(lastError ?? 'Grok CLI request failed.'); -} - -function extractDescription(json: unknown): string { - const obj = asObject(json); - if (!obj) return ''; - if (typeof obj.output_text === 'string' && obj.output_text.trim()) return obj.output_text; - if (!Array.isArray(obj.output)) return ''; - - return obj.output - .map(asObject) - .filter((o): o is Record => o !== undefined && o.type === 'message') - .flatMap((o) => (Array.isArray(o.content) ? o.content : [])) - .map(asObject) - .filter( - (o): o is Record => - o !== undefined && o.type === 'output_text' && typeof o.text === 'string', - ) - .map((o) => o.text) - .join('\n'); -} - -/** Describe a single image via the Grok CLI Responses endpoint. */ -export async function describeImage( - img: VisionImage, - model: string, - prompt: string, - apiKey: string, - signal?: AbortSignal, -): Promise { - const url = `${getBaseUrl()}/responses`; - const headers = { - 'Content-Type': 'application/json', - Authorization: `Bearer ${apiKey}`, - ...grokCliModelHeaders(model), - 'x-grok-conv-id': `grok-cli-vision-${makeCacheKey(img, model, prompt).slice(0, 16)}`, - }; - - const res = await fetchWithRetry( - url, - { - method: 'POST', - headers, - body: JSON.stringify({ - model, - input: [ - { - role: 'user', - content: [ - { type: 'input_text', text: prompt }, - { - type: 'input_image', - image_url: `data:${img.mimeType};base64,${img.data}`, - detail: 'auto', - }, - ], - }, - ], - text: { format: { type: 'text' } }, - stream: false, - }), - signal, - }, - model, - ); - - if (!res.ok) { - const body = await readErrorBody(res); - throw new Error(explainHttpError(res.status, body, model)); - } - - const rawText = await res.text(); - let json: unknown; - try { - json = JSON.parse(rawText); - } catch (err) { - throw new Error( - `Grok CLI returned non-JSON response: ${err instanceof Error ? err.message : String(err)}`, - ); - } - const description = extractDescription(json); - if (!description.trim()) { - throw new Error('Grok CLI returned an empty description.'); - } - return description; -} - -function textContent(text: string): TextContent { - return { type: 'text', text }; -} - -function replaceImagesWithText( - content: (TextContent | ImageContent)[], - replacements: TextContent[], - skipped?: TextContent, -): TextContent[] { - let imageIndex = 0; - const parts = content.flatMap((part) => { - if (part.type === 'text') return [part]; - - const replacement = replacements[imageIndex]; - imageIndex += 1; - return replacement ? [replacement] : []; - }); - - return skipped ? [...parts, skipped] : parts; -} - -async function describeSingle( - img: ImageContent, - index: number, - config: VisionConfig, - cachePath: string, - apiKey: string, - ctx: ExtensionContext, -): Promise { - const visionImg: VisionImage = { data: img.data, mimeType: img.mimeType }; - const cacheKey = makeCacheKey(visionImg, config.model, DEFAULT_PROMPT); - - if (config.cacheEnabled) { - const hit = loadCache(cachePath).entries[cacheKey]; - if (hit) return hit.description; - } - - try { - const description = await describeImage( - visionImg, - config.model, - DEFAULT_PROMPT, - apiKey, - ctx.signal, - ); - if (config.cacheEnabled) { - await updateCache(cachePath, (cache) => { - cache.entries[cacheKey] = makeCacheEntry( - visionImg, - config.model, - DEFAULT_PROMPT, - description, - ); - pruneCache(cache, config.cacheMaxEntries); - }); - } - return description; - } catch (err) { - const msg = err instanceof Error ? err.message : String(err); - ctx.ui.notify(`[grok-cli-vision] description failed: ${msg}`, 'warning'); - return `[Image ${index + 1} — description unavailable: ${msg}]`; - } -} - -/** - * Tool-result handler that routes images from `read` results through a Grok CLI - * vision model when the active model is text-only. Returns text-only content to - * replace the image so the text-only model can reason over the description. - */ -export async function handleReadResult( - event: ToolResultEvent, - ctx: ExtensionContext, -): Promise<{ content: TextContent[] } | undefined> { - // Accept both the native `read` tool and the capital-`Read` Cursor shim, - // which delegates to the native definition. Both return image blocks. - if (event.toolName !== 'read' && event.toolName !== 'Read') return; - - // Route when the active model exists and does not declare image input. This - // matches pi's own read tool (getNonVisionImageNote), which treats any model - // whose input lacks "image" — including input:[] — as non-vision. - const modelInput = ctx.model?.input; - if (!modelInput || modelInput.includes('image')) return; - - const loaded = loadConfig(); - const config = loaded.config.vision; - if (config.enabled === false) return; - - const images = event.content.filter((c): c is ImageContent => c.type === 'image'); - if (images.length === 0) return; - - const selected = images.slice(0, config.maxImages); - const skipped = images.length - selected.length; - - const apiKey = await resolveGrokToken(ctx); - if (!apiKey) { - ctx.ui.notify( - '[grok-cli-vision] No API key — run /login grok-cli or set GROK_CLI_OAUTH_TOKEN', - 'warning', - ); - return { - content: replaceImagesWithText( - event.content, - selected.map(() => - textContent('[grok-cli-vision: image not described — not authenticated]'), - ), - skipped > 0 - ? textContent( - `[grok-cli-vision: ${skipped} additional image(s) omitted (maxImages=${config.maxImages}).]`, - ) - : undefined, - ), - }; - } - - if (loaded.warning) ctx.ui.notify(`[grok-cli-vision] ${loaded.warning}`, 'warning'); - - const cachePath = getCachePath(); - const descriptions = await Promise.all( - selected.map((img, index) => describeSingle(img, index, config, cachePath, apiKey, ctx)), - ); - - const parts = descriptions.map((description, index) => - textContent(`[Image ${index + 1} — described by ${config.model}]\n${description}`), - ); - return { - content: replaceImagesWithText( - event.content, - parts, - skipped > 0 - ? textContent( - `[grok-cli-vision: ${skipped} additional image(s) omitted (maxImages=${config.maxImages}).]`, - ) - : undefined, - ), - }; -} diff --git a/src/vision/register.ts b/src/vision/register.ts deleted file mode 100644 index 87dabcc..0000000 --- a/src/vision/register.ts +++ /dev/null @@ -1,55 +0,0 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { getConfigPath, loadConfig, saveConfig } from '../config.js'; -import { cacheStats, clearCache, getCachePath } from './cache.js'; -import { handleReadResult } from './describe.js'; - -export function registerVisionFeature(pi: ExtensionAPI) { - pi.on('tool_result', handleReadResult); - - pi.registerCommand('grok-cli-vision:status', { - description: 'Show grok-cli-vision status, describer model, and cache stats', - handler: async (_args, ctx) => { - const loaded = loadConfig(); - const config = loaded.config.vision; - const stats = cacheStats(getCachePath()); - ctx.ui.notify( - [ - `grok-cli-vision: ${config.enabled ? 'ON' : 'OFF'}`, - `describer: ${config.model}`, - `maxImages: ${config.maxImages}`, - `cache: ${config.cacheEnabled ? 'ON' : 'OFF'} (${stats.entries} entries, max ${config.cacheMaxEntries})`, - `config: ${getConfigPath()}`, - `cache file: ${stats.path}`, - loaded.warning ? `warning: ${loaded.warning}` : undefined, - ] - .filter(Boolean) - .join('\n'), - loaded.warning ? 'warning' : 'info', - ); - }, - }); - - pi.registerCommand('grok-cli-vision', { - description: 'Toggle grok-cli-vision image routing', - handler: async (_args, ctx) => { - const loaded = loadConfig(); - const enabled = !loaded.config.vision.enabled; - saveConfig({ - ...loaded.config, - vision: { ...loaded.config.vision, enabled }, - }); - ctx.ui.notify( - enabled ? `grok-cli-vision: ON (${loaded.config.vision.model})` : 'grok-cli-vision: OFF', - 'info', - ); - }, - }); - - pi.registerCommand('grok-cli-vision:cache-clear', { - description: 'Clear the grok-cli-vision response cache', - handler: async (_args, ctx) => { - clearCache(getCachePath()); - ctx.ui.notify('grok-cli-vision cache: cleared', 'info'); - }, - }); -} diff --git a/tests/config.test.ts b/tests/config.test.ts index 4eb2e09..182e4ea 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -1,11 +1,4 @@ -import { - chmodSync, - existsSync, - mkdirSync, - readdirSync, - readFileSync, - writeFileSync, -} from 'node:fs'; +import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'; import { join } from 'node:path'; import { describe, expect, it } from 'vitest'; import { @@ -15,27 +8,16 @@ import { migrateLegacyConfig, saveConfig, } from '../src/config.js'; -import { useTempHome } from './vision/helpers.js'; +import { useTempHome } from './stateTestHelpers.js'; const setupHome = useTempHome(); -const FEATURE_CONFIG = { - ...DEFAULT_CONFIG, - imagine: { enabled: false }, - vision: { ...DEFAULT_CONFIG.vision, maxImages: 2 }, -}; function paths(home: string) { return { - cache: join(home, '.pi', 'grok-cli', 'vision-cache.json'), config: join(home, '.pi', 'grok-cli', 'config.json'), data: join(home, '.pi', 'grok-cli'), - debug: join(home, '.pi', 'grok-cli-vision-debug.log'), imagine: join(home, '.pi', 'grok-cli-imagine.json'), - legacyCache: join(home, '.pi', 'grok-cli-vision-cache.json'), pi: join(home, '.pi'), - quota: join(home, '.pi', 'grok-cli-quota.json'), - tools: join(home, '.pi', 'grok-cli-tools.json'), - vision: join(home, '.pi', 'grok-cli-vision.json'), }; } @@ -43,271 +25,120 @@ function writeJson(path: string, value: unknown) { writeFileSync(path, JSON.stringify(value)); } -function writeConsolidatedConfig(home: string, value: unknown) { +function writeConfig(home: string, value: unknown) { mkdirSync(paths(home).data, { recursive: true }); writeJson(paths(home).config, value); } -function withoutDirectoryWrites(directory: string, action: () => T) { - chmodSync(directory, 0o500); - try { - return action(); - } finally { - chmodSync(directory, 0o700); - } -} - describe('Grok CLI configuration', () => { - it('uses defaults without creating a file when no configuration exists', () => { + it('uses version 3 defaults without creating a file', () => { setupHome(); expect(loadConfig()).toEqual({ config: DEFAULT_CONFIG }); - expect(existsSync(getConfigPath())).toBe(false); - expect(migrateLegacyConfig()).toEqual({}); + expect(DEFAULT_CONFIG).toEqual({ version: 3, imagine: { enabled: true } }); expect(existsSync(getConfigPath())).toBe(false); }); - it('atomically saves and loads the versioned configuration', () => { + it('saves only non-account settings in version 3', () => { const home = setupHome(); - saveConfig(FEATURE_CONFIG); - - expect(loadConfig()).toEqual({ config: FEATURE_CONFIG }); - expect(JSON.parse(readFileSync(paths(home).config, 'utf8')).version).toBe(2); - expect(readdirSync(paths(home).data).filter((name) => name.endsWith('.tmp'))).toEqual([]); - }); - - it('fills missing sections and normalizes invalid fields with warnings', () => { - const home = setupHome(); - writeConsolidatedConfig(home, { - version: 2, - vision: { enabled: 'yes', maxImages: -1, cacheEnabled: false }, - }); - - const loaded = loadConfig(); - - expect(loaded.config.imagine).toEqual(DEFAULT_CONFIG.imagine); - expect(loaded.config.vision).toEqual({ - ...DEFAULT_CONFIG.vision, - cacheEnabled: false, - }); - expect(loaded.warning).toMatch(/enabled must be true or false/); - expect(loaded.warning).toMatch(/maxImages must be a positive integer/); - }); - - it('falls back to legacy settings without overwriting an unsupported version', () => { - const home = setupHome(); - mkdirSync(paths(home).pi, { recursive: true }); - writeConsolidatedConfig(home, { version: 3, imagine: { enabled: false } }); - writeJson(paths(home).imagine, { enabled: false }); - - const migration = migrateLegacyConfig(); - expect(migration.warning).toMatch(/Unsupported.*version/); - expect(loadConfig().config.imagine.enabled).toBe(false); - expect(JSON.parse(readFileSync(paths(home).config, 'utf8')).version).toBe(3); - expect(existsSync(paths(home).imagine)).toBe(true); - }); + saveConfig({ version: 3, imagine: { enabled: false } }); - it('migrates both current legacy files and leaves non-config data untouched', () => { - const home = setupHome(); - mkdirSync(paths(home).pi, { recursive: true }); - writeJson(paths(home).imagine, { enabled: false }); - writeJson(paths(home).vision, { - enabled: false, - model: 'grok-build', - maxImages: 2, - cacheEnabled: false, - cacheMaxEntries: 25, + expect(loadConfig()).toEqual({ + config: { version: 3, imagine: { enabled: false } }, }); - for (const path of [paths(home).legacyCache, paths(home).tools, paths(home).quota]) { - writeJson(path, { keep: true }); - } - writeFileSync(paths(home).debug, 'keep'); - - expect(migrateLegacyConfig()).toEqual({}); - expect(loadConfig().config).toEqual({ - ...DEFAULT_CONFIG, + expect(JSON.parse(readFileSync(paths(home).config, 'utf8'))).toEqual({ + version: 3, imagine: { enabled: false }, - vision: { - enabled: false, - model: 'grok-build', - maxImages: 2, - cacheEnabled: false, - cacheMaxEntries: 25, - }, }); - expect(existsSync(paths(home).imagine)).toBe(false); - expect(existsSync(paths(home).vision)).toBe(false); - expect(existsSync(paths(home).legacyCache)).toBe(false); - expect(JSON.parse(readFileSync(paths(home).cache, 'utf8'))).toEqual({ keep: true }); - for (const path of [paths(home).tools, paths(home).quota, paths(home).debug]) { - expect(existsSync(path)).toBe(true); - } }); - it('migrates version 1 in place while preserving existing feature settings', () => { + it('reads released version 2 settings without changing the source file', () => { const home = setupHome(); - writeConsolidatedConfig(home, { ...FEATURE_CONFIG, version: 1 }); - - expect(loadConfig()).toEqual({ config: FEATURE_CONFIG }); - expect(JSON.parse(readFileSync(paths(home).config, 'utf8')).version).toBe(1); - - expect(migrateLegacyConfig()).toEqual({}); - expect(JSON.parse(readFileSync(paths(home).config, 'utf8'))).toEqual(FEATURE_CONFIG); - }); - - it('normalizes account metadata and reports discarded entries', () => { - const home = setupHome(); - writeConsolidatedConfig(home, { - ...DEFAULT_CONFIG, + const released = { + version: 2, accounts: { - nextAccountNumber: 2, - selectedProvider: 'missing', + nextAccountNumber: 3, + selectedProvider: 'grok-cli-2', items: [ - { provider: 'grok-cli', label: ' Personal ' }, + { provider: 'grok-cli', label: 'Personal' }, { provider: 'grok-cli-2', label: 'Work' }, - { provider: 'grok-cli-2', label: 'Duplicate provider' }, - { provider: 'grok-cli-3', label: 'work' }, - { provider: 'grok-cli-10', label: 'Account 10' }, - { provider: 'other', label: 'Other' }, - { provider: 'grok-cli-4', label: 'bad\nlabel' }, ], }, - }); - - const loaded = loadConfig(); + imagine: { enabled: false }, + }; + writeConfig(home, released); - expect(loaded.config.accounts).toEqual({ - nextAccountNumber: 3, - selectedProvider: 'grok-cli', - items: [ - { provider: 'grok-cli', label: 'Personal' }, - { provider: 'grok-cli-2', label: 'Work' }, - { provider: 'grok-cli-10', label: 'Account 10' }, - ], - }); - expect(loaded.warning).toContain('accounts'); + expect(loadConfig().config).toEqual({ version: 3, imagine: { enabled: false } }); + expect(migrateLegacyConfig()).toEqual({}); + expect(JSON.parse(readFileSync(paths(home).config, 'utf8'))).toEqual(released); }); - it('reserves the permanent base label when base metadata is missing', () => { + it('normalizes an invalid Imagine setting with a warning', () => { const home = setupHome(); - writeConsolidatedConfig(home, { - ...DEFAULT_CONFIG, - accounts: { - nextAccountNumber: 2, - selectedProvider: 'grok-cli-2', - items: [ - { provider: 'grok-cli-2', label: 'Account 1' }, - { provider: 'grok-cli-3', label: 'Work' }, - ], - }, - }); + writeConfig(home, { version: 3, imagine: { enabled: 'yes' } }); - expect(loadConfig().config.accounts).toEqual({ - nextAccountNumber: 2, - selectedProvider: 'grok-cli', - items: [ - { provider: 'grok-cli', label: 'Account 1' }, - { provider: 'grok-cli-3', label: 'Work' }, - ], - }); + const loaded = loadConfig(); + + expect(loaded.config).toEqual(DEFAULT_CONFIG); + expect(loaded.warning).toContain('imagine.enabled'); }); - it.each(['grok-cli', 'all'])('migrates released Imagine scope %s to enabled', (scope) => { + it('preserves an unsupported consolidated config', () => { const home = setupHome(); - mkdirSync(paths(home).pi, { recursive: true }); - writeJson(paths(home).imagine, { scope }); + writeConfig(home, { version: 4, imagine: { enabled: false } }); - expect(migrateLegacyConfig()).toEqual({}); - expect(loadConfig().config.imagine.enabled).toBe(true); - expect(existsSync(paths(home).imagine)).toBe(false); + expect(migrateLegacyConfig().warning).toContain('Unsupported config version 4'); + expect(JSON.parse(readFileSync(paths(home).config, 'utf8')).version).toBe(4); }); - it.each([ - ['Imagine', 'imagine'], - ['vision', 'vision'], - ] as const)('migrates a lone %s legacy file and defaults the other section', (_label, kind) => { + it.each(['grok-cli', 'all'])('migrates released Imagine scope %s', (scope) => { const home = setupHome(); mkdirSync(paths(home).pi, { recursive: true }); - if (kind === 'imagine') writeJson(paths(home).imagine, { enabled: false }); - if (kind === 'vision') writeJson(paths(home).vision, { maxImages: 2 }); + writeJson(paths(home).imagine, { scope }); expect(migrateLegacyConfig()).toEqual({}); - expect(loadConfig().config).toEqual({ - ...DEFAULT_CONFIG, - imagine: kind === 'imagine' ? { enabled: false } : DEFAULT_CONFIG.imagine, - vision: - kind === 'vision' ? { ...DEFAULT_CONFIG.vision, maxImages: 2 } : DEFAULT_CONFIG.vision, - }); + expect(loadConfig().config).toEqual(DEFAULT_CONFIG); + expect(existsSync(paths(home).imagine)).toBe(false); }); - it('preserves all legacy files and skips migration when one is malformed', () => { + it('migrates a lone Imagine setting and removes only its legacy file', () => { const home = setupHome(); mkdirSync(paths(home).pi, { recursive: true }); writeJson(paths(home).imagine, { enabled: false }); - writeFileSync(paths(home).vision, '{ nope'); - const migration = migrateLegacyConfig(); - - expect(migration.warning).toMatch(/Could not read/); - expect(existsSync(paths(home).config)).toBe(false); - expect(existsSync(paths(home).imagine)).toBe(true); - expect(existsSync(paths(home).vision)).toBe(true); - expect(loadConfig().config.imagine.enabled).toBe(false); - expect(loadConfig().config.vision).toEqual(DEFAULT_CONFIG.vision); + expect(migrateLegacyConfig()).toEqual({}); + expect(loadConfig().config).toEqual({ version: 3, imagine: { enabled: false } }); + expect(existsSync(paths(home).imagine)).toBe(false); }); - it('preserves legacy files and removes temporary output when migration cannot write', () => { + it('preserves a legacy Imagine setting when the consolidated config omits Imagine', () => { const home = setupHome(); - mkdirSync(paths(home).pi, { recursive: true }); + writeConfig(home, { version: 2 }); writeJson(paths(home).imagine, { enabled: false }); - const migration = withoutDirectoryWrites(paths(home).pi, migrateLegacyConfig); - - expect(migration.warning).toMatch(/Could not migrate/); - expect(existsSync(paths(home).config)).toBe(false); - expect(existsSync(paths(home).imagine)).toBe(true); - expect(readdirSync(paths(home).pi).filter((name) => name.endsWith('.tmp'))).toEqual([]); - }); - - it('keeps a valid consolidated file authoritative and cleans recognized legacy files', () => { - const home = setupHome(); - mkdirSync(paths(home).pi, { recursive: true }); - writeConsolidatedConfig(home, { - ...DEFAULT_CONFIG, - imagine: { enabled: false }, - }); - writeJson(paths(home).imagine, { enabled: true }); - writeJson(paths(home).vision, { maxImages: 2 }); expect(migrateLegacyConfig()).toEqual({}); - expect(loadConfig().config.imagine.enabled).toBe(false); - expect(loadConfig().config.vision.maxImages).toBe(DEFAULT_CONFIG.vision.maxImages); - expect(existsSync(paths(home).imagine)).toBe(false); - expect(existsSync(paths(home).vision)).toBe(false); - expect(migrateLegacyConfig()).toEqual({}); + expect(existsSync(paths(home).imagine)).toBe(true); }); - it('preserves malformed legacy files beside a valid consolidated config', () => { + it('preserves a malformed legacy file', () => { const home = setupHome(); mkdirSync(paths(home).pi, { recursive: true }); - writeConsolidatedConfig(home, DEFAULT_CONFIG); writeFileSync(paths(home).imagine, '{ nope'); - expect(migrateLegacyConfig().warning).toMatch(/Could not read/); + expect(migrateLegacyConfig().warning).toContain('Could not read'); expect(existsSync(paths(home).imagine)).toBe(true); - expect(loadConfig().config).toEqual(DEFAULT_CONFIG); + expect(existsSync(paths(home).config)).toBe(false); }); - it('preserves a recognized legacy file when cleanup fails and retries later', () => { + it('preserves a legacy file when the destination is not writable', () => { const home = setupHome(); mkdirSync(paths(home).pi, { recursive: true }); - writeConsolidatedConfig(home, DEFAULT_CONFIG); writeJson(paths(home).imagine, { enabled: false }); - const migration = withoutDirectoryWrites(paths(home).pi, migrateLegacyConfig); + writeFileSync(paths(home).data, 'not a directory'); + const result = migrateLegacyConfig(); - expect(migration.warning).toMatch(/Could not remove/); + expect(result.warning).toContain('Could not migrate'); expect(existsSync(paths(home).imagine)).toBe(true); - expect(migrateLegacyConfig()).toEqual({}); - expect(existsSync(paths(home).imagine)).toBe(false); }); }); diff --git a/tests/imagine/config.test.ts b/tests/imagine/config.test.ts index 66736b8..426ae34 100644 --- a/tests/imagine/config.test.ts +++ b/tests/imagine/config.test.ts @@ -2,7 +2,7 @@ import { mkdirSync, writeFileSync } from 'node:fs'; import { dirname } from 'node:path'; import { describe, expect, it } from 'vitest'; import { DEFAULT_CONFIG, getConfigPath, loadConfig, saveConfig } from '../../src/config.js'; -import { useTempHome } from '../vision/helpers.js'; +import { useTempHome } from '../stateTestHelpers.js'; const setupHome = useTempHome(); diff --git a/tests/imagine/helpers.ts b/tests/imagine/helpers.ts index 0563259..ace6dc2 100644 --- a/tests/imagine/helpers.ts +++ b/tests/imagine/helpers.ts @@ -1,7 +1,22 @@ +import { mkdtempSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; import type { convertToPng } from '@earendil-works/pi-coding-agent'; -import { vi } from 'vitest'; +import { afterEach, vi } from 'vitest'; import type { savePreviewImage } from '../../src/imagine/save.js'; +const tempDirs: string[] = []; + +afterEach(() => { + for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true }); +}); + +export function tempDir(prefix: string) { + const dir = mkdtempSync(join(tmpdir(), prefix)); + tempDirs.push(dir); + return dir; +} + export function imagineDependencies() { return { generateImage: vi.fn(async () => ({ b64: '/9j/2Q==', mimeType: 'image/jpeg' as const })), diff --git a/tests/imagine/preview.test.ts b/tests/imagine/preview.test.ts index 143192e..3bdbd62 100644 --- a/tests/imagine/preview.test.ts +++ b/tests/imagine/preview.test.ts @@ -3,7 +3,7 @@ import { join } from 'node:path'; import { resetCapabilitiesCache, setCapabilities } from '@earendil-works/pi-tui'; import { afterEach, describe, expect, it } from 'vitest'; import { imagePreview } from '../../src/imagine/preview.js'; -import { tempDir } from '../tools/toolTestHelpers.js'; +import { tempDir } from './helpers.js'; const JPEG = '/9j/2Q=='; const PNG = diff --git a/tests/imagine/register.test.ts b/tests/imagine/register.test.ts index 4e2d0f7..1623b03 100644 --- a/tests/imagine/register.test.ts +++ b/tests/imagine/register.test.ts @@ -4,13 +4,19 @@ import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; import { describe, expect, it, vi } from 'vitest'; import { DEFAULT_CONFIG, loadConfig, saveConfig } from '../../src/config.js'; import { registerImagineFeature } from '../../src/imagine/register.js'; -import { saveTestAccounts, useTempHome } from '../vision/helpers.js'; +import { useEnvironmentToken, useTempHome } from '../stateTestHelpers.js'; import { imagineDependencies } from './helpers.js'; const setupHome = useTempHome(); +const setToken = useEnvironmentToken(); -function setup(token?: string, initialActiveTools: readonly string[] = ['read']) { +function setup( + token?: string, + initialActiveTools: readonly string[] = ['read'], + resolveToken?: () => Promise, +) { const home = setupHome(); + setToken(token); const commands = new Map Promise }>(); const renderers = new Map(); const entries: { type: string; data: unknown }[] = []; @@ -41,6 +47,7 @@ function setup(token?: string, initialActiveTools: readonly string[] = ['read']) setActiveTools, } as unknown as ExtensionAPI, dependencies, + resolveToken, ); const notify = vi.fn(); const context = { @@ -119,18 +126,29 @@ describe('registerImagineFeature command', () => { ); }); - it('uses the last selected Grok alias while a non-Grok model is active', async () => { + it('uses the environment override while a non-Grok model is active', async () => { const extension = setup('work-token'); - saveTestAccounts(); await extension.commands.get('grok-cli-imagine')?.handler('cat', extension.context); - expect(extension.getApiKeyForProvider).toHaveBeenCalledWith('grok-cli-2'); + expect(extension.getApiKeyForProvider).not.toHaveBeenCalled(); expect(extension.generate).toHaveBeenCalledWith( expect.objectContaining({ token: 'work-token' }), ); }); + it('uses the supplied session token resolver', async () => { + const resolveToken = vi.fn(async () => 'session-token'); + const extension = setup(undefined, ['read'], resolveToken); + + await extension.commands.get('grok-cli-imagine')?.handler('cat', extension.context); + + expect(resolveToken).toHaveBeenCalledWith(extension.context); + expect(extension.generate).toHaveBeenCalledWith( + expect.objectContaining({ token: 'session-token' }), + ); + }); + it('keeps a successful JPEG when PNG preview conversion fails', async () => { const extension = setup('token'); extension.convert.mockResolvedValueOnce(null); @@ -238,22 +256,6 @@ describe('registerImagineFeature command', () => { expect(loadConfig().config.imagine).toEqual({ enabled: true }); }); - it('preserves vision settings when image_gen persistence changes', async () => { - const extension = setup('token', ['read', 'image_gen']); - saveConfig({ - ...DEFAULT_CONFIG, - vision: { ...DEFAULT_CONFIG.vision, enabled: false, maxImages: 2 }, - }); - - await extension.commands.get('grok-cli-imagine:tool')?.handler('off', extension.context); - - expect(loadConfig().config).toEqual({ - ...DEFAULT_CONFIG, - imagine: { enabled: false }, - vision: { ...DEFAULT_CONFIG.vision, enabled: false, maxImages: 2 }, - }); - }); - it('leaves active tools unchanged when persistence fails', async () => { const extension = setup('token', ['read', 'custom', 'image_gen']); writeFileSync(join(extension.home, '.pi'), 'not a directory'); diff --git a/tests/imagine/save.test.ts b/tests/imagine/save.test.ts index a3e6681..6bcdf14 100644 --- a/tests/imagine/save.test.ts +++ b/tests/imagine/save.test.ts @@ -2,7 +2,7 @@ import { promises as fs } from 'node:fs'; import { basename, join } from 'node:path'; import { describe, expect, it } from 'vitest'; import { saveImage, savePreviewImage } from '../../src/imagine/save.js'; -import { tempDir } from '../tools/toolTestHelpers.js'; +import { tempDir } from './helpers.js'; const JPEG = '/9j/2Q=='; diff --git a/tests/imagine/tool.test.ts b/tests/imagine/tool.test.ts index 988fd07..b4d66ea 100644 --- a/tests/imagine/tool.test.ts +++ b/tests/imagine/tool.test.ts @@ -1,9 +1,15 @@ import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; import { describe, expect, it, vi } from 'vitest'; import { registerImageGenTool } from '../../src/imagine/tool.js'; +import { useEnvironmentToken, useTempHome } from '../stateTestHelpers.js'; import { imagineDependencies } from './helpers.js'; -function setup(token?: string) { +const setupHome = useTempHome(); +const setToken = useEnvironmentToken(); + +function setup(token?: string, resolveToken?: () => Promise) { + setupHome(); + setToken(token); let tool: Record | undefined; const dependencies = imagineDependencies(); registerImageGenTool( @@ -13,6 +19,7 @@ function setup(token?: string) { }, } as ExtensionAPI, dependencies, + resolveToken, ); const context = { modelRegistry: { getApiKeyForProvider: vi.fn(async () => token) }, @@ -79,6 +86,16 @@ describe('image_gen tool', () => { expect(result.content.every((part: { type: string }) => part.type === 'text')).toBe(true); }); + it('uses the supplied session token resolver', async () => { + const resolveToken = vi.fn(async () => 'session-token'); + const test = setup(undefined, resolveToken); + + await test.tool.execute('call', { prompt: 'cat' }, undefined, undefined, test.context); + + expect(resolveToken).toHaveBeenCalledWith(test.context); + expect(test.generate).toHaveBeenCalledWith(expect.objectContaining({ token: 'session-token' })); + }); + it('requires exactly one call for singular requests while allowing explicit multiples', () => { const test = setup('token'); expect(test.tool.description).toContain( diff --git a/tests/models/catalog.test.ts b/tests/models/catalog.test.ts index 642dcad..bc788cc 100644 --- a/tests/models/catalog.test.ts +++ b/tests/models/catalog.test.ts @@ -41,7 +41,7 @@ describe('model catalog', () => { ]); expect(models.find((model) => model.id === 'grok-composer-2.5-fast')).toMatchObject({ contextWindow: 200_000, - input: ['text'], + input: ['text', 'image'], }); expect(models.find((model) => model.id === 'grok-build')).toMatchObject({ contextWindow: 500_000, diff --git a/tests/provider/accountMigration.test.ts b/tests/provider/accountMigration.test.ts new file mode 100644 index 0000000..112f6da --- /dev/null +++ b/tests/provider/accountMigration.test.ts @@ -0,0 +1,304 @@ +import { existsSync, readFileSync, statSync, writeFileSync } from 'node:fs'; +import { describe, expect, it, vi } from 'vitest'; +import { + confirmMarkerInstallation, + migrateReleasedAccounts, +} from '../../src/provider/accountMigration.js'; +import { ACCOUNT_VAULT_MARKER, getAccountVault } from '../../src/provider/accountVault.js'; +import { + acquireFileLock, + getConfigBackupPath, + getConfigPath, + getLegacyConfigPath, + getQuotaCachePath, +} from '../../src/storage.js'; +import { oauthCredential, useTempHome, writeTestJson } from '../stateTestHelpers.js'; + +const setupHome = useTempHome(); + +const releasedConfig = { + version: 2, + accounts: { + nextAccountNumber: 11, + selectedProvider: 'grok-cli-10', + items: [ + { provider: 'grok-cli', label: 'Personal' }, + { provider: 'grok-cli-2', label: 'Work' }, + { provider: 'grok-cli-10', label: 'Backup' }, + ], + }, + imagine: { enabled: false }, +}; +const releasedAccountLabels = [ + { slot: 1, label: 'Personal' }, + { slot: 2, label: 'Work' }, + { slot: 10, label: 'Backup' }, +]; + +describe('released account migration', () => { + it('moves v0.6 account state into the vault and keeps a downgrade backup', async () => { + setupHome(); + writeTestJson(getConfigPath(), releasedConfig); + writeTestJson(getQuotaCachePath(), { + version: 1, + accounts: { + 'grok-cli': { updatedAt: '2026-07-01T00:00:00.000Z', monthly: {} }, + 'grok-cli-10': { updatedAt: '2026-07-02T00:00:00.000Z', monthly: {} }, + }, + }); + const credentials = { + 'grok-cli': { ...oauthCredential('personal'), baseUrl: 'https://personal.example' }, + 'grok-cli-10': oauthCredential('backup'), + }; + const readCredential = vi.fn( + (provider: string) => credentials[provider as keyof typeof credentials], + ); + const original = readFileSync(getConfigPath(), 'utf8'); + + expect(await migrateReleasedAccounts({ readCredential })).toEqual({ migrated: true }); + + const vault = await getAccountVault(); + expect(vault.accounts.map(({ slot, label }) => ({ slot, label }))).toEqual( + releasedAccountLabels, + ); + expect(vault.accounts[0]?.credential).toMatchObject({ + access: 'personal', + baseUrl: 'https://personal.example', + }); + expect(vault.accounts[1]?.credential).toBeUndefined(); + expect(vault.activeAccountId).toBe(vault.accounts[2]?.id); + expect(vault.migration).toEqual({ + legacyCredentialCopyComplete: true, + markerInstallPending: true, + }); + expect(readFileSync(getConfigBackupPath(), 'utf8')).toBe(original); + expect(statSync(getConfigBackupPath()).mode & 0o777).toBe(0o600); + expect(JSON.parse(readFileSync(getConfigPath(), 'utf8'))).toEqual({ + version: 3, + imagine: { enabled: false }, + }); + expect(Object.keys(JSON.parse(readFileSync(getQuotaCachePath(), 'utf8')).accounts)).toEqual([ + vault.accounts[0]?.id, + vault.accounts[2]?.id, + ]); + }); + + it('migrates the previously supported version 1 account configuration', async () => { + setupHome(); + writeTestJson(getConfigPath(), { ...releasedConfig, version: 1 }); + + await migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli-10' ? oauthCredential('backup') : undefined, + }); + + const vault = await getAccountVault(); + expect(vault.accounts.map(({ slot, label }) => ({ slot, label }))).toEqual( + releasedAccountLabels, + ); + expect(vault.activeAccountId).toBe(vault.accounts[2].id); + }); + + it('migrates a released config with missing account metadata as Account 1', async () => { + setupHome(); + writeTestJson(getConfigPath(), { version: 2, imagine: { enabled: false } }); + const original = readFileSync(getConfigPath(), 'utf8'); + + await migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli' ? oauthCredential('personal') : undefined, + }); + + expect((await getAccountVault()).accounts).toMatchObject([ + { id: 'account-1', label: 'Account 1', credential: { access: 'personal' } }, + ]); + expect(readFileSync(getConfigBackupPath(), 'utf8')).toBe(original); + expect(JSON.parse(readFileSync(getConfigPath(), 'utf8'))).toEqual({ + version: 3, + imagine: { enabled: false }, + }); + }); + + it('normalizes tolerated legacy account entries before vault import', async () => { + setupHome(); + writeTestJson(getConfigPath(), { + version: 1, + accounts: { + nextAccountNumber: 5, + selectedProvider: 'grok-cli-2', + items: [ + { provider: 'grok-cli-2', label: 'Account 1' }, + { provider: 'grok-cli-3', label: 'Work' }, + { provider: 'grok-cli-4', label: 'Work' }, + ], + }, + imagine: { enabled: true }, + }); + + await migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli-3' ? oauthCredential('work') : undefined, + }); + + const vault = await getAccountVault(); + expect(vault.accounts.map(({ slot, label }) => ({ slot, label }))).toEqual([ + { slot: 1, label: 'Account 1' }, + { slot: 3, label: 'Work' }, + ]); + expect(vault.activeAccountId).toBe(vault.accounts[1].id); + }); + + it('completes migration when the derived quota cache is malformed', async () => { + setupHome(); + writeTestJson(getConfigPath(), releasedConfig); + writeFileSync(getQuotaCachePath(), '{malformed'); + + await expect( + migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli' ? oauthCredential('personal') : undefined, + }), + ).resolves.toEqual({ migrated: true }); + + expect(JSON.parse(readFileSync(getConfigPath(), 'utf8')).version).toBe(3); + expect((await getAccountVault()).accounts[0].credential?.access).toBe('personal'); + }); + + it('retries without changing account identities or duplicating records', async () => { + setupHome(); + writeTestJson(getConfigPath(), releasedConfig); + const readCredential = (provider: string) => + provider === 'grok-cli' ? oauthCredential('personal') : undefined; + + await migrateReleasedAccounts({ readCredential }); + const first = await getAccountVault(); + await migrateReleasedAccounts({ readCredential }); + + expect(await getAccountVault()).toEqual(first); + expect(readFileSync(getConfigBackupPath(), 'utf8')).toContain('"version": 2'); + }); + + it('serializes two upgrade migrations before either reads released state', async () => { + setupHome(); + writeTestJson(getConfigPath(), releasedConfig); + + await Promise.all([ + migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli' ? oauthCredential('first') : undefined, + }), + migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli' ? oauthCredential('stale-second') : undefined, + }), + ]); + + expect((await getAccountVault()).accounts[0].credential?.access).toBe('first'); + }); + + it('holds the quota-cache lock while it rewrites migrated account IDs', async () => { + setupHome(); + writeTestJson(getConfigPath(), releasedConfig); + writeTestJson(getQuotaCachePath(), { + version: 1, + accounts: { + 'grok-cli': { updatedAt: '2026-07-01T00:00:00.000Z', monthly: {} }, + }, + }); + const release = await acquireFileLock(getQuotaCachePath()); + const settled = vi.fn(); + const migration = migrateReleasedAccounts().finally(settled); + + try { + await vi.waitFor(async () => { + expect((await getAccountVault()).accounts).toHaveLength(3); + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(settled).not.toHaveBeenCalled(); + } finally { + await release(); + } + + await expect(migration).resolves.toEqual({ migrated: true }); + }); + + it('rejects an unsupported released configuration version', async () => { + setupHome(); + writeTestJson(getConfigPath(), { version: 4, imagine: { enabled: false } }); + + await expect(migrateReleasedAccounts()).rejects.toThrow('is not a supported version 1 or 2'); + expect(JSON.parse(readFileSync(getConfigPath(), 'utf8')).version).toBe(4); + expect(existsSync(getConfigBackupPath())).toBe(false); + }); + + it('imports all accounts from the released legacy config path', async () => { + setupHome(); + writeTestJson(getLegacyConfigPath(), releasedConfig); + + await migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli-10' ? oauthCredential('backup') : undefined, + }); + + expect((await getAccountVault()).accounts.map(({ slot, label }) => ({ slot, label }))).toEqual( + releasedAccountLabels, + ); + expect(existsSync(getLegacyConfigPath())).toBe(false); + expect(JSON.parse(readFileSync(getConfigPath(), 'utf8')).version).toBe(3); + }); + + it('migrates a normal Account 1 login when no extension config exists', async () => { + setupHome(); + + expect( + await migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli' ? oauthCredential('personal') : undefined, + }), + ).toEqual({ migrated: true }); + + const vault = await getAccountVault(); + expect(vault.accounts).toHaveLength(1); + expect(vault.accounts[0]?.credential?.access).toBe('personal'); + expect(vault.activeAccountId).toBe(vault.accounts[0]?.id); + expect(existsSync(getConfigBackupPath())).toBe(false); + }); + + it('rejects the vault marker as a real credential', async () => { + setupHome(); + + await migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli' + ? { + type: 'oauth', + access: ACCOUNT_VAULT_MARKER, + refresh: ACCOUNT_VAULT_MARKER, + expires: Number.MAX_SAFE_INTEGER, + } + : undefined, + }); + + expect((await getAccountVault()).accounts[0].credential).toBeUndefined(); + }); + + it('clears pending marker state only after the Pi marker is observed', async () => { + setupHome(); + await migrateReleasedAccounts({ + readCredential: (provider) => + provider === 'grok-cli' ? oauthCredential('personal') : undefined, + }); + + expect(await confirmMarkerInstallation(() => undefined)).toBe(false); + expect((await getAccountVault()).migration.markerInstallPending).toBe(true); + expect( + await confirmMarkerInstallation(() => ({ + type: 'oauth', + access: ACCOUNT_VAULT_MARKER, + refresh: ACCOUNT_VAULT_MARKER, + expires: Number.MAX_SAFE_INTEGER, + })), + ).toBe(true); + expect((await getAccountVault()).migration.markerInstallPending).toBe(false); + }); +}); diff --git a/tests/provider/accountRouting.test.ts b/tests/provider/accountRouting.test.ts new file mode 100644 index 0000000..3e83d35 --- /dev/null +++ b/tests/provider/accountRouting.test.ts @@ -0,0 +1,290 @@ +import { join } from 'node:path'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { mutateAccountVault } from '../../src/provider/accountVault.js'; +import { useTempHome, writeTestJson } from '../stateTestHelpers.js'; + +const { refresh } = vi.hoisted(() => ({ refresh: vi.fn() })); + +vi.mock('../../src/auth/oauth.js', async (importOriginal) => ({ + ...(await importOriginal()), + refresh, +})); + +const originalToken = process.env.GROK_CLI_OAUTH_TOKEN; +const setupHome = useTempHome(); + +function deferredRefresh() { + let resolve = (_value: { access: string; refresh: string; expires: number }) => {}; + const promise = new Promise<{ access: string; refresh: string; expires: number }>((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +async function startRefresh(accountId?: string) { + const deferred = deferredRefresh(); + refresh.mockReturnValue(deferred.promise); + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + const pending = resolveAccountRoute(accountId); + await vi.waitFor(() => expect(refresh).toHaveBeenCalledOnce()); + return { ...deferred, pending }; +} + +async function saveAccount1(options: { active?: boolean; pending?: boolean } = {}) { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'account-token', + refresh: 'account-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + if (options.active) vault.activeAccountId = 'account-1'; + vault.migration.markerInstallPending = options.pending ?? false; + }); +} + +beforeEach(() => { + setupHome(); + delete process.env.GROK_CLI_OAUTH_TOKEN; + refresh.mockReset(); + writeTestJson(join(process.env.HOME as string, '.pi', 'agent', 'auth.json'), { + 'grok-cli': { + type: 'oauth', + access: 'pi-grok-cli-account-vault-v1', + refresh: 'pi-grok-cli-account-vault-v1', + expires: Number.MAX_SAFE_INTEGER, + }, + }); +}); + +afterEach(() => { + if (originalToken === undefined) delete process.env.GROK_CLI_OAUTH_TOKEN; + else process.env.GROK_CLI_OAUTH_TOKEN = originalToken; +}); + +describe('account request routing', () => { + it('routes the environment token without reading a vault account', async () => { + process.env.GROK_CLI_OAUTH_TOKEN = 'environment-token'; + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + + await expect(resolveAccountRoute()).resolves.toMatchObject({ + source: 'environment', + token: 'environment-token', + }); + }); + + it('selects the first logged-in account when the saved active account is absent', async () => { + await saveAccount1(); + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + + await expect(resolveAccountRoute()).resolves.toMatchObject({ + accountId: 'account-1', + revision: 1, + source: 'vault', + token: 'account-token', + }); + }); + + it('blocks vault requests until Pi installs the marker', async () => { + writeTestJson(join(process.env.HOME as string, '.pi', 'agent', 'auth.json'), {}); + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'account-token', + refresh: 'account-refresh', + expires: Date.now() + 60_000, + }; + vault.activeAccountId = 'account-1'; + vault.migration.markerInstallPending = true; + }); + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + + await expect(resolveAccountRoute()).rejects.toThrow('run /login'); + }); + + it('blocks vault requests after native Pi logout removes the marker', async () => { + await saveAccount1({ active: true }); + writeTestJson(join(process.env.HOME as string, '.pi', 'agent', 'auth.json'), {}); + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + + await expect(resolveAccountRoute()).rejects.toThrow('run /login'); + }); + + it('unblocks the next request after Pi stores the marker', async () => { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'account-token', + refresh: 'account-refresh', + expires: Date.now() + 300_000, + }; + vault.activeAccountId = 'account-1'; + vault.migration.markerInstallPending = true; + }); + writeTestJson(join(process.env.HOME as string, '.pi', 'agent', 'auth.json'), { + 'grok-cli': { + type: 'oauth', + access: 'pi-grok-cli-account-vault-v1', + refresh: 'pi-grok-cli-account-vault-v1', + expires: Number.MAX_SAFE_INTEGER, + }, + }); + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + + await expect(resolveAccountRoute()).resolves.toMatchObject({ + accountId: 'account-1', + token: 'account-token', + }); + }); + + it('uses the stored OAuth expiry without applying the refresh skew again', async () => { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'still-valid-token', + refresh: 'account-refresh', + expires: Date.now() + 60_000, + }; + vault.accounts[0].revision = 1; + vault.activeAccountId = 'account-1'; + }); + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + + await expect(resolveAccountRoute()).resolves.toMatchObject({ token: 'still-valid-token' }); + expect(refresh).not.toHaveBeenCalled(); + }); + + it('shares one refresh and stores it only for the captured account revision', async () => { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'expired-token', + refresh: 'account-refresh', + expires: 1, + }; + vault.accounts[0].revision = 3; + vault.activeAccountId = 'account-1'; + }); + refresh.mockResolvedValue({ + access: 'fresh-token', + refresh: 'fresh-refresh', + expires: Date.now() + 300_000, + }); + const { resolveAccountRoute } = await import('../../src/provider/accountRouting.js'); + + const routes = await Promise.all([resolveAccountRoute(), resolveAccountRoute()]); + + expect(refresh).toHaveBeenCalledOnce(); + expect(routes.map((route) => route.token)).toEqual(['fresh-token', 'fresh-token']); + expect(routes.map((route) => route.revision)).toEqual([4, 4]); + }); + + it('does not overwrite a newer login with a stale refresh result', async () => { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'expired-token', + refresh: 'old-refresh', + expires: 1, + }; + vault.accounts[0].revision = 3; + vault.activeAccountId = 'account-1'; + }); + const request = await startRefresh(); + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'new-login', + refresh: 'new-login-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 4; + }); + + request.resolve({ + access: 'stale-refresh', + refresh: 'stale-refresh-token', + expires: Date.now() + 300_000, + }); + + await expect(request.pending).resolves.toMatchObject({ token: 'new-login', revision: 4 }); + }); + + it('keeps an explicitly requested account after a stale refresh result', async () => { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'active-token', + refresh: 'active-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + vault.accounts.push({ + id: 'work-id', + slot: 2, + label: 'Work', + credential: { access: 'expired-work', refresh: 'old-work-refresh', expires: 1 }, + revision: 3, + }); + vault.nextSlot = 3; + vault.activeAccountId = 'account-1'; + }); + const request = await startRefresh('work-id'); + await mutateAccountVault((vault) => { + const work = vault.accounts.find((account) => account.id === 'work-id'); + if (!work) throw new Error('missing test account'); + work.credential = { + access: 'new-work-login', + refresh: 'new-work-refresh', + expires: Date.now() + 300_000, + }; + work.revision = 4; + }); + + request.resolve({ + access: 'stale-work-refresh', + refresh: 'stale-work-refresh-token', + expires: Date.now() + 300_000, + }); + + await expect(request.pending).resolves.toMatchObject({ + accountId: 'work-id', + token: 'new-work-login', + revision: 4, + }); + }); + + it('falls back when the implicitly selected account logs out during refresh', async () => { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'expired-active', + refresh: 'expired-active-refresh', + expires: 1, + }; + vault.accounts[0].revision = 3; + vault.accounts.push({ + id: 'work-id', + slot: 2, + label: 'Work', + credential: { + access: 'work-token', + refresh: 'work-refresh', + expires: Date.now() + 300_000, + }, + revision: 1, + }); + vault.nextSlot = 3; + vault.activeAccountId = 'account-1'; + }); + const request = await startRefresh(); + await mutateAccountVault((vault) => { + delete vault.accounts[0].credential; + vault.accounts[0].revision = 4; + vault.activeAccountId = 'work-id'; + }); + + request.resolve({ + access: 'stale-active-refresh', + refresh: 'stale-active-refresh-token', + expires: Date.now() + 300_000, + }); + + await expect(request.pending).resolves.toMatchObject({ + accountId: 'work-id', + token: 'work-token', + }); + }); +}); diff --git a/tests/provider/accountVault.test.ts b/tests/provider/accountVault.test.ts new file mode 100644 index 0000000..8c18a2e --- /dev/null +++ b/tests/provider/accountVault.test.ts @@ -0,0 +1,126 @@ +import { chmodSync, existsSync, mkdirSync, readFileSync, statSync, writeFileSync } from 'node:fs'; +import { join } from 'node:path'; +import { describe, expect, it } from 'vitest'; +import { + ACCOUNT_1_ID, + getAccountVault, + getAccountVaultSync, + mutateAccountVault, +} from '../../src/provider/accountVault.js'; +import { acquireFileLock, getAccountVaultPath } from '../../src/storage.js'; +import { useTempHome } from '../stateTestHelpers.js'; + +const setupHome = useTempHome(); + +describe('Grok CLI account vault', () => { + it('creates a locked vault with permanent Account 1 on its first mutation', async () => { + const home = setupHome(); + + await mutateAccountVault((vault) => { + vault.accounts[0].label = 'Personal'; + }); + + expect(await getAccountVault()).toEqual({ + version: 1, + migration: { + legacyCredentialCopyComplete: false, + markerInstallPending: false, + }, + nextSlot: 2, + accounts: [ + { + id: ACCOUNT_1_ID, + slot: 1, + label: 'Personal', + revision: 0, + }, + ], + }); + expect(statSync(join(home, '.pi', 'grok-cli')).mode & 0o777).toBe(0o700); + expect(statSync(getAccountVaultPath()).mode & 0o777).toBe(0o600); + }); + + it('serializes concurrent mutations without losing accounts or slots', async () => { + setupHome(); + + await Promise.all( + Array.from({ length: 4 }, (_value, index) => + mutateAccountVault((vault) => { + const slot = vault.nextSlot; + vault.nextSlot += 1; + vault.accounts.push({ + id: `account-${index + 2}`, + slot, + label: `Account ${slot}`, + revision: 0, + }); + }), + ), + ); + + const vault = await getAccountVault(); + expect(vault.accounts.map((account) => account.slot)).toEqual([1, 2, 3, 4, 5]); + expect(vault.nextSlot).toBe(6); + }); + + it('rejects malformed or invalid state without changing the file', async () => { + const home = setupHome(); + const path = getAccountVaultPath(); + mkdirSync(join(home, '.pi', 'grok-cli'), { recursive: true }); + + for (const contents of [ + '{malformed', + JSON.stringify({ + version: 1, + migration: { legacyCredentialCopyComplete: false, markerInstallPending: false }, + nextSlot: 2, + activeAccountId: 'missing', + accounts: [{ id: ACCOUNT_1_ID, slot: 1, label: 'Account 1', revision: 0 }], + }), + ]) { + writeFileSync(path, contents, { mode: 0o600 }); + + await expect(getAccountVault()).rejects.toThrow(path); + expect(readFileSync(path, 'utf8')).toBe(contents); + } + }); + + it('distinguishes an invalid root from an unsupported vault version', async () => { + const home = setupHome(); + const path = getAccountVaultPath(); + mkdirSync(join(home, '.pi', 'grok-cli'), { recursive: true }); + + writeFileSync(path, '[]', { mode: 0o600 }); + await expect(getAccountVault()).rejects.toThrow('root must be a JSON object'); + + writeFileSync(path, '{"version":2}', { mode: 0o600 }); + await expect(getAccountVault()).rejects.toThrow('unsupported version'); + }); + + it('repairs permissive vault permissions after a successful mutation', async () => { + const home = setupHome(); + await mutateAccountVault(() => undefined); + chmodSync(join(home, '.pi', 'grok-cli'), 0o755); + chmodSync(getAccountVaultPath(), 0o644); + + await mutateAccountVault(() => undefined); + + expect(statSync(join(home, '.pi', 'grok-cli')).mode & 0o777).toBe(0o700); + expect(statSync(getAccountVaultPath()).mode & 0o777).toBe(0o600); + expect(existsSync(`${getAccountVaultPath()}.lock`)).toBe(false); + }); + + it('reads the atomic vault snapshot while another process holds the mutation lock', async () => { + setupHome(); + await mutateAccountVault((vault) => { + vault.accounts[0].label = 'Personal'; + }); + const release = await acquireFileLock(getAccountVaultPath()); + + try { + expect(getAccountVaultSync().accounts[0].label).toBe('Personal'); + } finally { + await release(); + } + }); +}); diff --git a/tests/provider/accounts.test.ts b/tests/provider/accounts.test.ts index bd4ca75..1fbb0fe 100644 --- a/tests/provider/accounts.test.ts +++ b/tests/provider/accounts.test.ts @@ -1,1111 +1,465 @@ -import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'; -import type { Component } from '@earendil-works/pi-tui'; -import { afterEach, describe, expect, it, vi } from 'vitest'; -import { DEFAULT_CONFIG, loadConfig, saveConfig } from '../../src/config.js'; +import { EventEmitter } from 'node:events'; +import type { AuthInteraction } from '@earendil-works/pi-ai'; +import type { + ExtensionAPI, + ExtensionCommandContext, + ExtensionContext, +} from '@earendil-works/pi-coding-agent'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { + GROK_CLI_PROVIDER, isGrokCliProvider, planTier, registerAccountManagement, - resolveGrokProvider, resolveGrokToken, } from '../../src/provider/accounts.js'; -import { loadQuotaCache, saveQuotaUsage } from '../../src/provider/quotaCache.js'; -import { oauthCredential, TEST_ACCOUNTS, useTempHome } from '../vision/helpers.js'; +import { getAccountVault, mutateAccountVault } from '../../src/provider/accountVault.js'; +import { loadQuotaCache } from '../../src/provider/quotaCache.js'; +import { + deferred, + oauthCredential, + setAccount1Credential, + useEnvironmentToken, + useTempHome, + writePiCredential, + writePiVaultMarker, +} from '../stateTestHelpers.js'; + +const { fetchBillingUsage, login, removeQuotaUsage, spawnProcess } = vi.hoisted(() => ({ + fetchBillingUsage: vi.fn(), + login: vi.fn(), + removeQuotaUsage: vi.fn(), + spawnProcess: vi.fn(), +})); + +vi.mock('node:child_process', () => ({ spawn: spawnProcess })); + +vi.mock('../../src/auth/oauth.js', async (importOriginal) => ({ + ...(await importOriginal()), + login, +})); + +vi.mock('../../src/provider/billing.js', () => ({ fetchBillingUsage })); + +vi.mock('../../src/provider/quotaCache.js', async (importOriginal) => ({ + ...(await importOriginal()), + removeQuotaUsage, +})); const setupHome = useTempHome(); -const originalFetch = globalThis.fetch; -const originalToken = process.env.GROK_CLI_OAUTH_TOKEN; - -afterEach(() => { - globalThis.fetch = originalFetch; - vi.restoreAllMocks(); - if (originalToken === undefined) delete process.env.GROK_CLI_OAUTH_TOKEN; - else process.env.GROK_CLI_OAUTH_TOKEN = originalToken; -}); - -function configureAccounts( - selectedProvider = 'grok-cli', - items = TEST_ACCOUNTS, - nextAccountNumber = 3, -) { - setupHome(); - saveConfig({ - ...DEFAULT_CONFIG, - accounts: { nextAccountNumber, selectedProvider, items }, - }); +const setEnvironmentToken = useEnvironmentToken(); + +const ctx = { + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], + }, + ui: { + notify: vi.fn(), + }, +} as unknown as ExtensionContext; + +const interaction = { + notify: vi.fn(), + prompt: vi.fn(), +} satisfies AuthInteraction; + +function manager(clearRecentExhaustion = (_accountId: string) => {}, appendEntry = vi.fn()) { + return registerAccountManagement( + { + appendEntry, + registerCommand: vi.fn(), + } as unknown as ExtensionAPI, + clearRecentExhaustion, + ).manager; } -const authenticatedAccounts = () => ({ - 'grok-cli': oauthCredential('personal'), - 'grok-cli-2': oauthCredential('work'), -}); - -const numberedAccounts = (count: number) => - Array.from({ length: count }, (_value, index) => ({ - provider: index === 0 ? 'grok-cli' : `grok-cli-${index + 1}`, - label: `Account ${index + 1}`, - })); - -function setupRefresh( - accounts: { provider: string }[], - action: (component: Component) => Promise, -) { - return setup({ - auth: Object.fromEntries( - accounts.map((account) => [account.provider, oauthCredential(account.provider)]), - ), - preserveHome: true, - customActions: [action], +async function addLoggedInAccount(accounts: ReturnType) { + const account = await accounts.add(ctx, 'Work'); + await mutateAccountVault((vault) => { + vault.accounts[1].credential = { + access: 'two', + refresh: 'two-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[1].revision = 1; }); + return account; } -async function runAccountsCommand(extension: ReturnType) { - await extension.commands.get('grok-cli-accounts')?.handler('', extension.context); +async function selectedLoggedInAccount() { + await setAccount1Credential('one'); + const appendEntry = vi.fn(); + const accounts = manager(undefined, appendEntry); + const account = await addLoggedInAccount(accounts); + await accounts.activate(ctx, account.id); + return { account, accounts, appendEntry }; } -async function runAccountsCommandWith(extension: ReturnType, args: string) { - await extension.commands.get('grok-cli-accounts')?.handler(args, extension.context); -} - -function setup( - options: { - auth?: Record>; - confirms?: boolean[]; - customActions?: ((component: Component) => Promise)[]; - inputs?: (string | undefined)[]; - model?: { provider: string; id: string }; - preserveHome?: boolean; - selections?: (string | undefined)[]; - setModel?: boolean[]; - } = {}, -) { - if (!options.preserveHome) setupHome(); - const commands = new Map Promise }>(); - const registerAccount = vi.fn(); - const unregisterProvider = vi.fn(); - const setModelResults = [...(options.setModel ?? [true])]; - const setModel = vi.fn(async () => setModelResults.shift() ?? true); - const pi = { - registerCommand(name: string, command: unknown) { - commands.set(name, command as { handler: (args: string, ctx: unknown) => Promise }); - }, - setModel, - unregisterProvider, - } as unknown as ExtensionAPI; - const accountManagement = registerAccountManagement(pi, registerAccount); - - const credentials = new Map(Object.entries(options.auth ?? {})); - const authStorage = { - get: (provider: string) => credentials.get(provider), - has: (provider: string) => credentials.has(provider), - set: (provider: string, credential: ReturnType) => - credentials.set(provider, credential), - }; - const selections = [...(options.selections ?? [])]; - const inputs = [...(options.inputs ?? [])]; - const confirms = [...(options.confirms ?? [])]; - const customActions = [...(options.customActions ?? [])]; - const notify = vi.fn(); - const setEditorText = vi.fn(); - const customRenders: string[][] = []; - const models = new Map(); - for (const provider of ['grok-cli', 'grok-cli-2', 'grok-cli-3']) { - for (const id of ['grok-build', 'grok-composer-2.5-fast']) { - models.set(`${provider}/${id}`, { provider, id }); - } - } - const context = { - model: options.model, - modelRegistry: { - runtime: { - login: vi.fn(), - logout: async (provider: string) => { - credentials.delete(provider); - }, - }, - find: (provider: string, id: string) => models.get(`${provider}/${id}`), - getProviderAuthStatus: (provider: string) => ({ configured: credentials.has(provider) }), - getApiKeyForProvider: async (provider: string) => { - const credential = authStorage.get(provider); - return credential?.type === 'oauth' ? credential.access : undefined; - }, - }, - ui: { - confirm: vi.fn(async () => confirms.shift() ?? false), - custom: vi.fn( - async ( - factory: ( - tui: { requestRender: () => void }, - theme: { bold: (text: string) => string; fg: (_color: string, text: string) => string }, - keybindings: unknown, - done: (value: string | undefined) => void, - ) => Component | Promise, - ) => { - let resolveResult = (_value: string | undefined) => {}; - const result = new Promise((resolve) => { - resolveResult = resolve; - }); - let component: Component; - const tui = { - requestRender() { - customRenders.push(component.render(160)); - }, - }; - component = await factory( - tui, - { bold: (text) => text, fg: (_color, text) => text }, - {}, - resolveResult, - ); - customRenders.push(component.render(160)); - const action = customActions.shift(); - if (action) { - await action(component); - return result; - } - const choice = selections.shift(); - if (choice === undefined) { - component.handleInput?.('\u001b'); - return result; - } - for (let index = 0; index < 20; index += 1) { - const selected = component.render(160).find((line) => line.startsWith('→ ')); - if (selected?.includes(choice)) { - component.handleInput?.('\r'); - return result; - } - component.handleInput?.('\u001b[B'); - } - throw new Error(`Could not select custom UI row: ${choice}`); - }, - ), - input: vi.fn(async () => inputs.shift()), - notify, - select: vi.fn(async () => selections.shift()), - setEditorText, +beforeEach(() => { + setupHome(); + setEnvironmentToken(undefined); + login.mockReset(); + removeQuotaUsage.mockReset(); + removeQuotaUsage.mockResolvedValue(undefined); + fetchBillingUsage.mockReset(); + fetchBillingUsage.mockResolvedValue({ + monthly: { + monthlyLimit: 2000, + used: 100, + billingPeriodEnd: '2026-08-01T00:00:00.000Z', }, - }; - - return { - authStorage, - accountManagement, - commands, - context, - customRenders, - notify, - registerAccount, - setEditorText, - setModel, - unregisterProvider, - }; -} - -describe('Grok CLI account helpers', () => { - it('recognizes only the base provider and valid numbered aliases', () => { - expect(isGrokCliProvider('grok-cli')).toBe(true); - expect(isGrokCliProvider('grok-cli-2')).toBe(true); - expect(isGrokCliProvider('grok-cli-10')).toBe(true); - expect(isGrokCliProvider('grok-cli-1')).toBe(false); - expect(isGrokCliProvider('grok-cli-work')).toBe(false); - }); - - it('maps monthly credit caps to plan tiers', () => { - expect(planTier(0)).toBe('free'); - expect(planTier(4000)).toBe('supergrok-lite'); - expect(planTier(20000)).toBe('supergrok'); - expect(planTier(20001)).toBe('supergrok-heavy'); }); - - it('resolves the current Grok alias or the persisted selection for other models', async () => { - configureAccounts('grok-cli-2'); - const getApiKeyForProvider = vi.fn(async (provider: string) => `${provider}-token`); - - expect( - resolveGrokProvider({ - model: { provider: 'grok-cli', id: 'grok-build' }, - } as unknown as Pick), - ).toBe('grok-cli'); - expect( - resolveGrokProvider({ - model: { provider: 'openai', id: 'gpt-5' }, - } as unknown as Pick), - ).toBe('grok-cli-2'); - expect( - await resolveGrokToken({ - model: { provider: 'openai', id: 'gpt-5' }, - modelRegistry: { getApiKeyForProvider }, - } as unknown as Pick), - ).toBe('grok-cli-2-token'); - expect(getApiKeyForProvider).toHaveBeenCalledWith('grok-cli-2'); + spawnProcess.mockReset(); + login.mockResolvedValue({ + access: 'new-access', + refresh: 'new-refresh', + expires: Date.now() + 300_000, }); }); -describe('/grok-cli-accounts', () => { - it('rejects unknown arguments without opening the TUI', async () => { - const extension = setup(); - - await runAccountsCommandWith(extension, 'wat'); - - expect(extension.notify).toHaveBeenCalledWith('Usage: /grok-cli-accounts [gui]', 'warning'); - expect(extension.context.ui.custom).not.toHaveBeenCalled(); - }); - - it('produces a credential-free account snapshot for alternate interfaces', async () => { - configureAccounts(); - await saveQuotaUsage('grok-cli-2', { - monthly: { - monthlyLimit: 2000, - used: 700, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - weekly: { - creditUsagePercent: 35, - billingPeriodEnd: '2026-07-20T00:00:00.000Z', - }, - }); - const extension = setup({ auth: authenticatedAccounts(), preserveHome: true }); - - const snapshot = extension.accountManagement.manager.snapshot( - extension.context as unknown as ExtensionContext, - ); - - expect(snapshot.accounts[1]).toMatchObject({ - provider: 'grok-cli-2', - label: 'Work', - authenticated: true, - active: false, - environment: false, - plan: 'supergrok-lite', - quota: { - monthly: { monthlyLimit: 2000, used: 700 }, - weekly: { creditUsagePercent: 35 }, - }, - }); - expect(JSON.stringify(snapshot)).not.toContain('work-token'); +describe('Grok provider identity', () => { + it('accepts only the single Grok CLI provider', () => { + expect(GROK_CLI_PROVIDER).toBe('grok-cli'); + expect(isGrokCliProvider('grok-cli')).toBe(true); + expect(isGrokCliProvider('grok-cli-2')).toBe(false); }); - it('reads authentication through the Pi 0.80.9 model-registry facade', () => { - configureAccounts(); - const extension = setup({ auth: authenticatedAccounts(), preserveHome: true }); - const modelRegistry = extension.context.modelRegistry as unknown as { - authStorage?: unknown; - getProviderAuthStatus?: (provider: string) => { configured: boolean }; - }; - delete modelRegistry.authStorage; - modelRegistry.getProviderAuthStatus = (provider) => ({ - configured: provider === 'grok-cli' || provider === 'grok-cli-2', - }); - - expect( - extension.accountManagement.manager.snapshot(extension.context as unknown as ExtensionContext) - .accounts, - ).toEqual([ - expect.objectContaining({ provider: 'grok-cli', authenticated: true }), - expect.objectContaining({ provider: 'grok-cli-2', authenticated: true }), + it('classifies plan tiers from monthly limits', () => { + expect([planTier(0), planTier(4000), planTier(20000), planTier(20001)]).toEqual([ + 'free', + 'supergrok-lite', + 'supergrok', + 'supergrok-heavy', ]); }); +}); - it('shares add and rename mutations with alternate interfaces', async () => { - const extension = setup(); +describe('vault account management', () => { + it('adds an account with an opaque stable ID and increasing display slot', async () => { + const accounts = manager(); - const account = await extension.accountManagement.manager.add( - extension.context as unknown as ExtensionContext, - ' Work ', - ); - await extension.accountManagement.manager.rename( - extension.context as unknown as ExtensionContext, - account.provider, - 'Client', - ); - - expect(account.provider).toBe('grok-cli-2'); - expect(loadConfig().config.accounts.items).toEqual([ - { provider: 'grok-cli', label: 'Account 1' }, - { provider: 'grok-cli-2', label: 'Client' }, - ]); - expect(extension.registerAccount).toHaveBeenNthCalledWith(1, { - provider: 'grok-cli-2', - label: 'Work', - }); - expect(extension.registerAccount).toHaveBeenNthCalledWith(2, { - provider: 'grok-cli-2', - label: 'Client', - }); - }); + const account = await accounts.add(ctx, 'Work'); - it('shows cached quota usage in both account selectors', async () => { - configureAccounts(); - await saveQuotaUsage( - 'grok-cli', - { - monthly: { - monthlyLimit: 2000, - used: 300, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - weekly: { - creditUsagePercent: 60, - billingPeriodEnd: '2026-07-20T00:00:00.000Z', - }, - }, - new Date().toISOString(), - ); - const extension = setup({ - auth: authenticatedAccounts(), - preserveHome: true, - selections: ['Manage accounts', 'Work — Authenticated', 'Back'], + expect(account).toMatchObject({ slot: 2, label: 'Work', revision: 0 }); + expect(account.id).toMatch(/^[0-9a-f-]{36}$/); + await expect(getAccountVault()).resolves.toMatchObject({ + nextSlot: 3, + accounts: [{ id: 'account-1' }, { id: account.id, slot: 2, label: 'Work' }], }); - - await runAccountsCommand(extension); - - expect(extension.context.ui.custom).toHaveBeenCalledTimes(2); - expect(extension.customRenders[0]?.join('\n')).toContain( - 'Monthly 300 / 2,000 used · Weekly 60% used', - ); - expect(extension.customRenders.at(-1)?.join('\n')).toContain('Quota not fetched · press r'); }); - it('refreshes every logged-in account with r without changing models and preserves selection', async () => { - configureAccounts(); - process.env.GROK_CLI_OAUTH_TOKEN = 'personal'; - const timeout = vi.spyOn(AbortSignal, 'timeout'); - const fetchMock = vi.fn(async (input, init) => { - const authorization = (init?.headers as Record).authorization; - const used = authorization === 'Bearer personal' ? 300 : 900; - if (String(input).includes('format=credits')) { - return Response.json({ - config: { - currentPeriod: { type: 'USAGE_PERIOD_TYPE_WEEKLY' }, - creditUsagePercent: used === 300 ? 60 : 25, - billingPeriodEnd: '2026-07-20T00:00:00.000Z', - }, - }); - } - return Response.json({ - config: { - monthlyLimit: { val: 2000 }, - used: { val: used }, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }); - }); - globalThis.fetch = fetchMock; - const extension = setup({ - auth: { 'grok-cli-2': oauthCredential('work') }, - preserveHome: true, - customActions: [ - async (component) => { - component.handleInput?.('\u001b[B'); - component.handleInput?.('r'); - component.handleInput?.('r'); - await vi.waitFor(() => { - expect(Object.keys(loadQuotaCache().accounts)).toHaveLength(2); - }); - expect(component.render(160).find((line) => line.startsWith('→ '))).toContain('Work'); - component.handleInput?.('\u001b'); - }, - ], - }); + it('logs in directly and preserves the account identity, slot, and label', async () => { + const accounts = manager(); + const account = await accounts.add(ctx, 'Work'); - await runAccountsCommand(extension); + await accounts.login(ctx, account.id, interaction); - expect(fetchMock).toHaveBeenCalledTimes(4); - expect(timeout).toHaveBeenCalledTimes(2); - expect(timeout).toHaveBeenCalledWith(30_000); - expect(loadQuotaCache().accounts['grok-cli']?.monthly.used).toBe(300); - expect(loadQuotaCache().accounts['grok-cli-2']?.monthly.used).toBe(900); - expect(extension.setModel).not.toHaveBeenCalled(); - expect(extension.customRenders.flat().join('\n')).toContain('Updated 2 accounts; 0 failed'); - }); - - it('keeps cached quota on a partial refresh failure and marks only that row failed', async () => { - configureAccounts(); - await saveQuotaUsage( - 'grok-cli-2', - { - monthly: { - monthlyLimit: 2000, - used: 700, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }, - new Date().toISOString(), - ); - globalThis.fetch = vi.fn(async (input, init) => { - const authorization = (init?.headers as Record).authorization; - if (authorization === 'Bearer work') return new Response('nope', { status: 500 }); - if (String(input).includes('format=credits')) - return new Response('no weekly', { status: 500 }); - return Response.json({ - config: { - monthlyLimit: { val: 2000 }, - used: { val: 300 }, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }); - }); - const extension = setup({ - auth: authenticatedAccounts(), - preserveHome: true, - customActions: [ - async (component) => { - component.handleInput?.('r'); - await vi.waitFor(() => { - expect(component.render(160).join('\n')).toContain('Updated 1 accounts; 1 failed'); - }); - expect(component.render(160).join('\n')).toContain( - 'Monthly 700 / 2,000 used · Weekly unavailable', - ); - expect(component.render(160).join('\n')).toContain('refresh failed'); - component.handleInput?.('\u001b'); + expect(login).toHaveBeenCalledOnce(); + await expect(getAccountVault()).resolves.toMatchObject({ + activeAccountId: account.id, + accounts: [ + { id: 'account-1', slot: 1, label: 'Account 1' }, + { + id: account.id, + slot: 2, + label: 'Work', + revision: 1, + credential: { access: 'new-access', refresh: 'new-refresh' }, }, ], }); - - await runAccountsCommand(extension); - - expect(loadQuotaCache().accounts['grok-cli']?.monthly.used).toBe(300); - expect(loadQuotaCache().accounts['grok-cli-2']?.monthly.used).toBe(700); }); - it('allows removal during refresh without restoring the removed quota cache', async () => { - configureAccounts(); - let releaseMonthly = () => {}; - globalThis.fetch = vi.fn(async (input) => { - if (String(input).includes('format=credits')) { - return new Response('no weekly', { status: 500 }); - } - await new Promise((resolve) => { - releaseMonthly = resolve; - }); - return Response.json({ - config: { - monthlyLimit: { val: 2000 }, - used: { val: 900 }, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }); - }); - const extension = setup({ auth: authenticatedAccounts(), preserveHome: true }); - const refresh = extension.accountManagement.manager.refreshOne( - extension.context as unknown as ExtensionContext, - 'grok-cli-2', - new AbortController().signal, - ); - await vi.waitFor(() => expect(globalThis.fetch).toHaveBeenCalledTimes(1)); - const removal = extension.accountManagement.manager.remove( - extension.context as unknown as ExtensionContext, - 'grok-cli-2', - ); + it('does not restore an account after logout completes during login', async () => { + const accounts = manager(); + const account = await accounts.add(ctx, 'Work'); + const authorization = deferred>(); + login.mockReturnValueOnce(authorization.promise); - await removal; - await extension.accountManagement.manager.add( - extension.context as unknown as ExtensionContext, - 'Replacement', - ); - expect(loadConfig().config.accounts.items).toEqual([ - { provider: 'grok-cli', label: 'Personal' }, - { provider: 'grok-cli-2', label: 'Replacement' }, - ]); - releaseMonthly(); - await refresh; + const pending = accounts.login(ctx, account.id, interaction); + await accounts.logout(ctx, account.id); + authorization.resolve(oauthCredential('late-login')); - expect(loadQuotaCache().accounts['grok-cli-2']).toBeUndefined(); + await expect(pending).rejects.toThrow('account changed while login was in progress'); + expect((await getAccountVault()).accounts[1]).toMatchObject({ revision: 1 }); + expect((await getAccountVault()).accounts[1].credential).toBeUndefined(); }); - it('allows logout during refresh without restoring the logged-out quota cache', async () => { - configureAccounts(); - let releaseMonthly = () => {}; - globalThis.fetch = vi.fn(async (input) => { - if (String(input).includes('format=credits')) return new Response(null, { status: 500 }); - await new Promise((resolve) => { - releaseMonthly = resolve; - }); - return Response.json({ - config: { - monthlyLimit: { val: 2000 }, - used: { val: 900 }, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }); - }); - const extension = setup({ auth: authenticatedAccounts(), preserveHome: true }); - const refresh = extension.accountManagement.manager.refreshOne( - extension.context as unknown as ExtensionContext, - 'grok-cli', - new AbortController().signal, - ); - await vi.waitFor(() => expect(globalThis.fetch).toHaveBeenCalledTimes(1)); - - await extension.accountManagement.manager.logout( - extension.context as unknown as ExtensionContext, - 'grok-cli', - ); - releaseMonthly(); - await refresh; - - expect(loadQuotaCache().accounts['grok-cli']).toBeUndefined(); - }); + it('keeps a successful login when quota-cache cleanup fails', async () => { + const accounts = manager(); + const account = await accounts.add(ctx, 'Work'); + removeQuotaUsage.mockRejectedValueOnce(new Error('cache unavailable')); - it('handles uppercase R with no authenticated accounts without fetching', async () => { - configureAccounts(); - const fetchMock = vi.fn(); - globalThis.fetch = fetchMock; - const extension = setup({ - preserveHome: true, - customActions: [ - async (component) => { - component.handleInput?.('R'); - expect(component.render(35).join('\n')).toContain('No logged-in accounts to refresh'); - expect(component.render(35).join('\n')).toContain('Personal — Login required'); - component.handleInput?.('\u001b'); - }, - ], + await expect(accounts.login(ctx, account.id, interaction)).resolves.toMatchObject({ + id: account.id, }); - await runAccountsCommand(extension); - - expect(fetchMock).not.toHaveBeenCalled(); + expect((await getAccountVault()).accounts[1].credential?.access).toBe('new-access'); }); - it('aborts outstanding quota requests when the selector closes', async () => { - const accounts = numberedAccounts(5); - configureAccounts('grok-cli', accounts, 6); - const aborted = vi.fn(); - globalThis.fetch = vi.fn( - (_input, init) => - new Promise((_resolve, reject) => { - init?.signal?.addEventListener('abort', () => { - aborted(); - reject(new DOMException('Aborted', 'AbortError')); - }); - }), - ); - const extension = setupRefresh(accounts, async (component) => { - component.handleInput?.('r'); - await vi.waitFor(() => expect(globalThis.fetch).toHaveBeenCalledTimes(3)); - component.handleInput?.('\u001b'); - }); + it('clears recent exhaustion after direct account login succeeds', async () => { + const clearRecentExhaustion = vi.fn(); + const accounts = manager(clearRecentExhaustion); + const account = await accounts.add(ctx, 'Work'); - await runAccountsCommand(extension); - await vi.waitFor(() => expect(aborted).toHaveBeenCalledTimes(3)); - await new Promise((resolve) => setTimeout(resolve, 0)); + await accounts.login(ctx, account.id, interaction); - expect(globalThis.fetch).toHaveBeenCalledTimes(3); - expect(loadQuotaCache().accounts).toEqual({}); + expect(clearRecentExhaustion).toHaveBeenCalledWith(account.id); }); - it('runs at most three account refreshes concurrently', async () => { - const accounts = numberedAccounts(5); - configureAccounts('grok-cli', accounts, 6); - const releases: (() => void)[] = []; - let active = 0; - let maximum = 0; - globalThis.fetch = vi.fn(async (input) => { - if (String(input).includes('format=credits')) - return new Response('no weekly', { status: 500 }); - active += 1; - maximum = Math.max(maximum, active); - await new Promise((resolve) => releases.push(resolve)); - active -= 1; - return Response.json({ - config: { - monthlyLimit: { val: 2000 }, - used: { val: 300 }, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }); - }); - const extension = setupRefresh(accounts, async (component) => { - component.handleInput?.('r'); - await vi.waitFor(() => expect(releases).toHaveLength(3)); - expect(maximum).toBe(3); - releases.splice(0).forEach((release) => { - release(); - }); - await vi.waitFor(() => expect(releases).toHaveLength(2)); - releases.splice(0).forEach((release) => { - release(); - }); - await vi.waitFor(() => expect(Object.keys(loadQuotaCache().accounts)).toHaveLength(5)); - component.handleInput?.('\u001b'); - }); - - await runAccountsCommand(extension); + it('rejects activation for a logged-out account', async () => { + const accounts = manager(); + const account = await accounts.add(ctx, 'Work'); - expect(maximum).toBe(3); + await expect(accounts.activate(ctx, account.id)).rejects.toThrow('before making it active'); }); - it('adds a labeled alias and pre-fills Pi native login', async () => { - const extension = setup({ selections: ['＋ Add account'], inputs: ['Work'] }); - - await runAccountsCommand(extension); + it('activates an account only for the current Pi session', async () => { + const test = await selectedLoggedInAccount(); - expect(loadConfig().config.accounts).toEqual({ - nextAccountNumber: 3, - selectedProvider: 'grok-cli', - items: [ - { provider: 'grok-cli', label: 'Account 1' }, - { provider: 'grok-cli-2', label: 'Work' }, - ], + expect(test.appendEntry).toHaveBeenCalledWith('grok-cli-active-account-v1', { + accountId: test.account.id, }); - expect(extension.registerAccount).toHaveBeenCalledWith({ - provider: 'grok-cli-2', - label: 'Work', - }); - expect(extension.setEditorText).toHaveBeenCalledWith('/login grok-cli-2'); - }); - - it('uses the lowest available alias number and its matching default label', async () => { - configureAccounts('grok-cli', [{ provider: 'grok-cli', label: 'Account 1' }], 4); - const extension = setup({ preserveHome: true, selections: ['＋ Add account'], inputs: [''] }); - - await runAccountsCommand(extension); - - expect(loadConfig().config.accounts.items.at(-1)).toEqual({ - provider: 'grok-cli-2', - label: 'Account 2', - }); - expect(loadConfig().config.accounts.nextAccountNumber).toBe(3); + expect((await getAccountVault()).activeAccountId).toBe('account-1'); + expect(test.accounts.snapshot(ctx).accounts[1]).toMatchObject({ active: true }); }); - it('switches a logged-in account while preserving the current Grok model', async () => { - configureAccounts(); - const extension = setup({ - auth: { - 'grok-cli': oauthCredential('personal'), - 'grok-cli-2': oauthCredential('work'), - }, - model: { provider: 'grok-cli', id: 'grok-composer-2.5-fast' }, - preserveHome: true, - selections: ['Work — Authenticated'], + it('selects another logged-in account when the active account logs out', async () => { + await setAccount1Credential('one'); + const accounts = manager(); + const account = await addLoggedInAccount(accounts); + await mutateAccountVault((vault) => { + vault.activeAccountId = account.id; }); - await runAccountsCommand(extension); + await accounts.logout(ctx, account.id); - expect(extension.setModel).toHaveBeenCalledWith({ - provider: 'grok-cli-2', - id: 'grok-composer-2.5-fast', + await expect(getAccountVault()).resolves.toMatchObject({ + activeAccountId: 'account-1', + accounts: [ + { id: 'account-1', credential: { access: 'one' } }, + { id: account.id, revision: 2 }, + ], }); - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli-2'); + expect((await getAccountVault()).accounts[1].credential).toBeUndefined(); }); - it('uses grok-build when switching from a non-Grok model', async () => { - configureAccounts(); - const extension = setup({ - auth: { 'grok-cli-2': oauthCredential('work') }, - model: { provider: 'openai', id: 'gpt-5' }, - preserveHome: true, - selections: ['Work — Authenticated'], - }); + it('persists another logged-in account when the selected session account logs out', async () => { + const test = await selectedLoggedInAccount(); + test.appendEntry.mockClear(); - await runAccountsCommand(extension); + await test.accounts.logout(ctx, test.account.id); - expect(extension.setModel).toHaveBeenCalledWith({ - provider: 'grok-cli-2', - id: 'grok-build', + expect(test.appendEntry).toHaveBeenCalledWith('grok-cli-active-account-v1', { + accountId: 'account-1', }); + expect(test.accounts.snapshot(ctx).accounts[0]).toMatchObject({ active: true }); }); - it('prefills login instead of switching an unauthenticated account', async () => { - configureAccounts(); - const extension = setup({ preserveHome: true, selections: ['Work — Login required'] }); + it('removes the final active selection when the final login logs out', async () => { + await setAccount1Credential('one'); - await runAccountsCommand(extension); + const result = await manager().logout(ctx, 'account-1'); - expect(extension.setEditorText).toHaveBeenCalledWith('/login grok-cli-2'); - expect(extension.setModel).not.toHaveBeenCalled(); + expect((await getAccountVault()).activeAccountId).toBeUndefined(); + expect(result.warning).toContain('/logout'); }); - it('does not persist a switch when Pi rejects the model change', async () => { - configureAccounts(); - const extension = setup({ - auth: { 'grok-cli-2': oauthCredential('work') }, - preserveHome: true, - selections: ['Work — Authenticated'], - setModel: [false], - }); + it('keeps Account 1 permanent and removes another account', async () => { + const accounts = manager(); + const extra = await accounts.add(ctx, 'Work'); - await runAccountsCommand(extension); + await expect(accounts.remove(ctx, 'account-1')).rejects.toThrow('cannot be removed'); + await accounts.remove(ctx, extra.id); - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli'); - expect(extension.notify).toHaveBeenCalledWith( - 'Could not switch to “Work”; authentication is unavailable.', - 'error', - ); + expect((await getAccountVault()).accounts.map((account) => account.id)).toEqual(['account-1']); }); - it('rejects duplicate, controlled, and overlong labels before adding an account', async () => { - configureAccounts(); - const extension = setup({ - preserveHome: true, - selections: ['＋ Add account'], - inputs: [' work ', 'bad\nlabel', 'bad\u009blabel', 'x'.repeat(41), 'Client'], - }); + it('shows contiguous account numbers after an earlier account is removed', async () => { + const accounts = manager(); + const removed = await accounts.add(ctx, ''); + const remaining = await accounts.add(ctx, ''); - await runAccountsCommand(extension); + await accounts.remove(ctx, removed.id); - expect(extension.notify.mock.calls.map(([message]) => message)).toEqual([ - 'An account named “work” already exists.', - 'Account labels cannot contain control characters.', - 'Account labels cannot contain control characters.', - 'Account labels must be 40 characters or fewer.', + expect(accounts.snapshot(ctx).accounts).toMatchObject([ + { id: 'account-1', slot: 1, label: 'Account 1' }, + { id: remaining.id, slot: 2, label: 'Account 2' }, ]); - expect(loadConfig().config.accounts.items.at(-1)).toEqual({ - provider: 'grok-cli-3', - label: 'Client', + expect((await getAccountVault()).accounts[1]).toMatchObject({ + id: remaining.id, + slot: 3, + label: 'Account 3', }); }); - it('renames an account and updates its provider display registration', async () => { - configureAccounts(); - await saveQuotaUsage('grok-cli-2', { - monthly: { - monthlyLimit: 2000, - used: 300, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }); - const extension = setup({ - selections: ['Manage accounts', 'Work — Login required', 'Rename'], - inputs: ['Client'], - preserveHome: true, - }); - - await runAccountsCommand(extension); - - expect(loadConfig().config.accounts.items[1]?.label).toBe('Client'); - expect(extension.registerAccount).toHaveBeenCalledWith({ - provider: 'grok-cli-2', - label: 'Client', - }); - expect(loadQuotaCache().accounts['grok-cli-2']?.monthly.used).toBe(300); - }); + it('hides logout state for a logged-out Account 1 snapshot', () => { + const snapshot = manager().snapshot(ctx).accounts[0]; - it('logs out and removes an inactive alias after confirmation', async () => { - configureAccounts(); - await saveQuotaUsage('grok-cli-2', { - monthly: { - monthlyLimit: 2000, - used: 300, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - }); - const extension = setup({ - auth: { - 'grok-cli': oauthCredential('personal'), - 'grok-cli-2': oauthCredential('work'), - }, - confirms: [true], - model: { provider: 'grok-cli', id: 'grok-build' }, - preserveHome: true, - selections: ['Manage accounts', 'Work — Authenticated', 'Log out and remove'], - }); - - await runAccountsCommand(extension); - - expect(extension.authStorage.has('grok-cli-2')).toBe(false); - expect(loadConfig().config.accounts.items).toEqual([ - { provider: 'grok-cli', label: 'Personal' }, - ]); - expect(extension.unregisterProvider).toHaveBeenCalledWith('grok-cli-2'); - expect(loadQuotaCache().accounts['grok-cli-2']).toBeUndefined(); - }); - - it('reuses grok-cli-2 for the next account after removing that alias', async () => { - configureAccounts(); - const extension = setup({ - auth: authenticatedAccounts(), - confirms: [true], - inputs: ['Replacement'], - model: { provider: 'grok-cli', id: 'grok-build' }, - preserveHome: true, - selections: [ - 'Manage accounts', - 'Work — Authenticated', - 'Log out and remove', - '＋ Add account', - ], - }); - - await runAccountsCommand(extension); - await runAccountsCommand(extension); - - expect(loadConfig().config.accounts.items).toEqual([ - { provider: 'grok-cli', label: 'Personal' }, - { provider: 'grok-cli-2', label: 'Replacement' }, - ]); - expect(extension.registerAccount).toHaveBeenCalledWith({ - provider: 'grok-cli-2', - label: 'Replacement', + expect(snapshot).toMatchObject({ + id: 'account-1', + permanent: true, + authenticated: false, + active: false, + status: 'Login required', }); - expect(extension.setEditorText).toHaveBeenCalledWith('/login grok-cli-2'); }); - it('does not reuse an active alias until its deferred unregister completes', async () => { - configureAccounts('grok-cli-2'); - const extension = setup({ - auth: { 'grok-cli-2': oauthCredential('work') }, - confirms: [true], - inputs: ['Temporary', 'Replacement'], - model: { provider: 'grok-cli-2', id: 'grok-build' }, - preserveHome: true, - selections: [ - 'Manage accounts', - 'Work — Active', - 'Log out and remove', - '＋ Add account', - '＋ Add account', - ], - }); - - await runAccountsCommand(extension); - await runAccountsCommand(extension); - - expect(loadConfig().config.accounts.items.at(-1)).toEqual({ - provider: 'grok-cli-3', - label: 'Temporary', - }); - - extension.accountManagement.handleModelSelect({ - model: { provider: 'openai' }, - previousModel: { provider: 'grok-cli-2' }, - }); - await runAccountsCommand(extension); + it('uses the environment token without exposing a vault credential', async () => { + process.env.GROK_CLI_OAUTH_TOKEN = 'environment-token'; - expect(extension.unregisterProvider).toHaveBeenCalledWith('grok-cli-2'); - expect(loadConfig().config.accounts.items.at(-1)).toEqual({ - provider: 'grok-cli-2', - label: 'Replacement', + expect(await resolveGrokToken()).toBe('environment-token'); + expect(manager().snapshot(ctx).accounts[0]).toMatchObject({ + environment: true, + authenticated: true, + active: true, }); - expect(extension.setEditorText.mock.calls.map(([value]) => value)).toEqual([ - '/login grok-cli-3', - '/login grok-cli-2', - ]); + await expect(manager().login(ctx, 'account-1', interaction)).rejects.toThrow( + 'Unset GROK_CLI_OAUTH_TOKEN', + ); }); - it('leaves an alias untouched when removal confirmation is cancelled', async () => { - configureAccounts(); - const extension = setup({ - auth: { 'grok-cli-2': oauthCredential('work') }, - confirms: [false], - preserveHome: true, - selections: ['Manage accounts', 'Work — Authenticated', 'Log out and remove'], + it('reports Pi connected only after Pi stores the vault marker', () => { + writePiCredential({ + access: 'released-access', + refresh: 'released-refresh', + expires: Date.now() + 300_000, }); + expect(manager().snapshot(ctx).connected).toBe(false); - await runAccountsCommand(extension); - - expect(extension.authStorage.has('grok-cli-2')).toBe(true); - expect(loadConfig().config.accounts.items).toHaveLength(2); - expect(extension.unregisterProvider).not.toHaveBeenCalled(); + writePiVaultMarker(); + expect(manager().snapshot(ctx).connected).toBe(true); }); - it('switches away before removing the active alias', async () => { - configureAccounts('grok-cli-2'); - const extension = setup({ - auth: { - 'grok-cli': oauthCredential('personal'), - 'grok-cli-2': oauthCredential('work'), + it('rejects saved-account activation while the environment token is active', async () => { + process.env.GROK_CLI_OAUTH_TOKEN = 'environment-token'; + const accounts = manager(); + const account = await accounts.add(ctx, 'Work'); + await mutateAccountVault((vault) => { + vault.accounts[1].credential = { + access: 'work', + refresh: 'work-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[1].revision = 1; + }); + + await expect(accounts.activate(ctx, account.id)).rejects.toThrow('environment token'); + }); + + it('reports the authorization URL when the browser opener fails', async () => { + const child = Object.assign(new EventEmitter(), { unref: vi.fn() }); + spawnProcess.mockReturnValue(child); + login.mockImplementationOnce(async (callbacks) => { + callbacks.onAuth({ url: 'https://accounts.x.ai/authorize?state=test' }); + child.emit('error', new Error('xdg-open missing')); + throw new Error('stop test login'); + }); + const commands = new Map< + string, + { handler: (args: string, ctx: ExtensionCommandContext) => unknown } + >(); + registerAccountManagement({ + registerCommand( + name: string, + command: { handler: (args: string, ctx: ExtensionCommandContext) => unknown }, + ) { + commands.set(name, command); }, - confirms: [true], - model: { provider: 'grok-cli-2', id: 'grok-composer-2.5-fast' }, - preserveHome: true, - selections: ['Manage accounts', 'Work — Active', 'Log out and remove'], - }); - - await runAccountsCommand(extension); - - expect(extension.setModel).toHaveBeenCalledWith({ - provider: 'grok-cli', - id: 'grok-composer-2.5-fast', - }); - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli'); - expect(extension.unregisterProvider).toHaveBeenCalledWith('grok-cli-2'); - }); - - it('does not switch a non-Grok model when removing the selected alias', async () => { - configureAccounts('grok-cli-2'); - const extension = setup({ - auth: authenticatedAccounts(), - confirms: [true], - model: { provider: 'openai', id: 'gpt-5' }, - preserveHome: true, - selections: ['Manage accounts', 'Work — Active', 'Log out and remove'], - }); - - await runAccountsCommand(extension); - - expect(extension.setModel).not.toHaveBeenCalled(); - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli'); - expect(extension.unregisterProvider).toHaveBeenCalledWith('grok-cli-2'); - }); - - it('keeps the selected account when removing another inactive alias', async () => { - configureAccounts( - 'grok-cli-2', - [...TEST_ACCOUNTS, { provider: 'grok-cli-3', label: 'Client' }], - 4, - ); - const extension = setup({ - auth: { - ...authenticatedAccounts(), - 'grok-cli-3': oauthCredential('client'), + } as unknown as ExtensionAPI); + const commandContext = { + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], }, - confirms: [true], - model: { provider: 'openai', id: 'gpt-5' }, - preserveHome: true, - selections: ['Manage accounts', 'Client — Authenticated', 'Log out and remove'], - }); - - await runAccountsCommand(extension); - - expect(extension.setModel).not.toHaveBeenCalled(); - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli-2'); - expect(extension.unregisterProvider).toHaveBeenCalledWith('grok-cli-3'); - }); - - it('defers unregistering an active alias when no authenticated fallback exists', async () => { - configureAccounts('grok-cli-2'); - const extension = setup({ - auth: { 'grok-cli-2': oauthCredential('work') }, - confirms: [true], - model: { provider: 'grok-cli-2', id: 'grok-build' }, - preserveHome: true, - selections: ['Manage accounts', 'Work — Active', 'Log out and remove'], - }); - - await runAccountsCommand(extension); + ui: { + input: vi.fn().mockResolvedValue('Work'), + notify: vi.fn(), + select: vi.fn().mockResolvedValue('Add account'), + }, + } as unknown as ExtensionCommandContext; - expect(extension.unregisterProvider).not.toHaveBeenCalled(); - expect(loadConfig().config.accounts.items).toEqual([ - { provider: 'grok-cli', label: 'Personal' }, - ]); + await commands.get('grok-cli-accounts')?.handler('', commandContext); - extension.accountManagement.handleModelSelect({ - model: { provider: 'openai' }, - previousModel: { provider: 'grok-cli-2' }, - }); - - expect(extension.unregisterProvider).toHaveBeenCalledWith('grok-cli-2'); + expect(commandContext.ui.notify).toHaveBeenCalledWith( + expect.stringContaining('https://accounts.x.ai/authorize?state=test'), + 'warning', + ); }); - it('prefills native Pi relogin from account management', async () => { - configureAccounts(); - const extension = setup({ - auth: { 'grok-cli-2': oauthCredential('work') }, - preserveHome: true, - selections: ['Manage accounts', 'Work — Authenticated', 'Log in again'], + it('does not save quota after the captured account logs out', async () => { + writePiVaultMarker(); + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'one', + refresh: 'one-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + vault.activeAccountId = 'account-1'; }); + let finishBilling = (_usage: Awaited>) => {}; + fetchBillingUsage.mockImplementationOnce( + () => + new Promise((resolve) => { + finishBilling = resolve; + }), + ); + const accounts = manager(); + const pending = accounts.refreshOne(ctx, 'account-1', new AbortController().signal); + await vi.waitFor(() => expect(fetchBillingUsage).toHaveBeenCalledOnce()); + await accounts.logout(ctx, 'account-1'); - await runAccountsCommand(extension); - - expect(extension.setEditorText).toHaveBeenCalledWith('/login grok-cli-2'); - }); - - it('keeps the base slot and resets its label when logging out', async () => { - configureAccounts('grok-cli', [{ provider: 'grok-cli', label: 'Personal' }], 2); - await saveQuotaUsage('grok-cli', { + finishBilling({ monthly: { monthlyLimit: 2000, - used: 300, + used: 100, billingPeriodEnd: '2026-08-01T00:00:00.000Z', }, }); - const extension = setup({ - auth: { 'grok-cli': oauthCredential('personal') }, - confirms: [true], - preserveHome: true, - selections: ['Manage accounts', 'Personal — Active', 'Log out'], - }); - - await runAccountsCommand(extension); - expect(extension.authStorage.has('grok-cli')).toBe(false); - expect(loadConfig().config.accounts.items).toEqual([ - { provider: 'grok-cli', label: 'Account 1' }, - ]); - expect(extension.unregisterProvider).not.toHaveBeenCalled(); - expect(loadQuotaCache().accounts['grok-cli']).toBeUndefined(); + await expect(pending).resolves.toEqual({ updated: 0, failed: ['account-1'] }); + expect(loadQuotaCache().accounts['account-1']).toBeUndefined(); }); - it('keeps the selected alias when logging out of the inactive base account', async () => { - configureAccounts('grok-cli-2'); - const extension = setup({ - auth: { - 'grok-cli': oauthCredential('personal'), - 'grok-cli-2': oauthCredential('work'), - }, - confirms: [true], - preserveHome: true, - selections: ['Manage accounts', 'Personal — Authenticated', 'Log out'], + it('refreshes at most three account quotas at one time', async () => { + writePiVaultMarker(); + await mutateAccountVault((vault) => { + vault.accounts = Array.from({ length: 7 }, (_value, index) => ({ + id: index === 0 ? 'account-1' : `account-${index + 1}`, + slot: index + 1, + label: `Account ${index + 1}`, + credential: { + access: `access-${index + 1}`, + refresh: `refresh-${index + 1}`, + expires: Date.now() + 300_000, + }, + revision: 1, + })); + vault.nextSlot = 8; + vault.activeAccountId = 'account-1'; }); - - await runAccountsCommand(extension); - - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli-2'); - expect(loadConfig().config.accounts.items[0]?.label).toBe('Account 1'); - }); - - it('explains that a base environment token cannot be logged out from Pi', async () => { - const original = process.env.GROK_CLI_OAUTH_TOKEN; - process.env.GROK_CLI_OAUTH_TOKEN = 'environment-token'; - try { - const extension = setup({ - selections: [ - 'Manage accounts', - 'Account 1 — Active (environment)', - 'Environment token instructions', - ], - }); - - await runAccountsCommand(extension); - - expect(extension.notify).toHaveBeenCalledWith( - 'Unset GROK_CLI_OAUTH_TOKEN and restart Pi to remove the environment token.', - 'info', - ); - expect(extension.authStorage.has('grok-cli')).toBe(false); - } finally { - if (original === undefined) delete process.env.GROK_CLI_OAUTH_TOKEN; - else process.env.GROK_CLI_OAUTH_TOKEN = original; - } - }); - - it('persists a Grok alias selected through Pi model controls', async () => { - configureAccounts(); - const extension = setup({ preserveHome: true }); - - await extension.accountManagement.handleModelSelect({ - model: { provider: 'grok-cli-2' }, - previousModel: { provider: 'grok-cli' }, + let active = 0; + let maximum = 0; + fetchBillingUsage.mockImplementation(async () => { + active += 1; + maximum = Math.max(maximum, active); + await new Promise((resolve) => setTimeout(resolve, 50)); + active -= 1; + return { + monthly: { + monthlyLimit: 2000, + used: 100, + billingPeriodEnd: '2026-08-01T00:00:00.000Z', + }, + }; }); - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli-2'); - }); - - it('applies model selections after an overlapping account mutation', async () => { - configureAccounts(); - const extension = setup({ auth: authenticatedAccounts(), preserveHome: true }); - let finishModelSwitch = () => {}; - extension.setModel.mockImplementationOnce( - () => - new Promise((resolve) => { - finishModelSwitch = () => resolve(true); - }), - ); - const activation = extension.accountManagement.manager.activate( - extension.context as unknown as ExtensionContext, - 'grok-cli-2', - ); - await vi.waitFor(() => expect(extension.setModel).toHaveBeenCalledTimes(1)); - - const selection = extension.accountManagement.handleModelSelect({ - model: { provider: 'grok-cli' }, - previousModel: { provider: 'grok-cli-2' }, + await expect(manager().refresh(ctx, new AbortController().signal)).resolves.toEqual({ + updated: 7, + failed: [], }); - finishModelSwitch(); - await Promise.all([activation, selection]); - - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli'); + expect(maximum).toBeLessThanOrEqual(3); }); }); diff --git a/tests/provider/dashboard.test.ts b/tests/provider/dashboard.test.ts index 0a2f387..d771c8f 100644 --- a/tests/provider/dashboard.test.ts +++ b/tests/provider/dashboard.test.ts @@ -1,16 +1,6 @@ import { once } from 'node:events'; import { createConnection } from 'node:net'; -import { - InMemoryCredentialStore, - type OAuthCredentials, - type OAuthLoginCallbacks, -} from '@earendil-works/pi-ai'; -import { - type ExtensionAPI, - type ExtensionContext, - ModelRegistry, - ModelRuntime, -} from '@earendil-works/pi-coding-agent'; +import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'; import { type HTMLDialogElement as BrowserDialog, type HTMLElement as BrowserElement, @@ -18,14 +8,21 @@ import { Window, } from 'happy-dom'; import { afterEach, describe, expect, it, vi } from 'vitest'; -import { DEFAULT_CONFIG, loadConfig, saveConfig } from '../../src/config.js'; import { registerAccountManagement } from '../../src/provider/accounts.js'; +import { getAccountVault, mutateAccountVault } from '../../src/provider/accountVault.js'; import { type AccountDashboardHandle, createAccountDashboard, startAccountDashboard, } from '../../src/provider/dashboard/server.js'; -import { oauthCredential, useTempHome } from '../vision/helpers.js'; +import { deferred, oauthCredential, useTempHome } from '../stateTestHelpers.js'; + +const { oauthLogin } = vi.hoisted(() => ({ oauthLogin: vi.fn() })); + +vi.mock('../../src/auth/oauth.js', async (importOriginal) => ({ + ...(await importOriginal()), + login: oauthLogin, +})); const setupHome = useTempHome(); const dashboards: AccountDashboardHandle[] = []; @@ -37,69 +34,29 @@ afterEach(async () => { async function setup() { setupHome(); - saveConfig(DEFAULT_CONFIG); - const credentials = new InMemoryCredentialStore(); - await credentials.modify('grok-cli', async () => oauthCredential('personal')); - const runtime = await ModelRuntime.create({ - credentials, - modelsPath: null, - allowModelNetwork: false, + oauthLogin.mockReset(); + oauthLogin.mockResolvedValue(oauthCredential('personal')); + await mutateAccountVault((vault) => { + vault.accounts[0].credential = oauthCredential('personal'); + vault.activeAccountId = 'account-1'; }); - const loginFlows = new Map< - string, - (callbacks: OAuthLoginCallbacks) => Promise - >(); - const registerAccount = (account: { provider: string; label: string }) => { - runtime.registerProvider(account.provider, { - name: account.label, - baseUrl: 'https://example.test', - api: 'openai-responses', - models: [ - { - id: 'grok-build', - name: 'Grok Build', - reasoning: true, - input: ['text'], - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, - contextWindow: 128_000, - maxTokens: 16_384, - }, - ], - oauth: { - name: account.label, - usesCallbackServer: true, - login: (callbacks) => - loginFlows.get(account.provider)?.(callbacks) ?? - Promise.reject(new Error('No test login flow registered.')), - refreshToken: async (credential) => credential, - getApiKey: (credential) => credential.access, - }, - }); - }; - registerAccount({ provider: 'grok-cli', label: 'Account 1' }); - const modelRegistry = new ModelRegistry(runtime); const pi = { + appendEntry: vi.fn(), registerCommand: vi.fn(), - setModel: vi.fn(async () => true), - unregisterProvider: vi.fn((provider: string) => runtime.unregisterProvider(provider)), } as unknown as ExtensionAPI; - const accountManagement = registerAccountManagement(pi, registerAccount); + const accountManagement = registerAccountManagement(pi); const ctx = { model: { provider: 'grok-cli', id: 'grok-build' }, - modelRegistry, + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], + }, ui: { notify: vi.fn() }, } as unknown as ExtensionContext; return { accountManagement, - credentials, ctx, - loginFlows, pi, - runtime, - async setCredential(provider: string, credential: ReturnType) { - await credentials.modify(provider, async () => credential); - await runtime.refresh({ allowNetwork: false }); - }, }; } @@ -126,7 +83,9 @@ const accountState = ( overrides: Record = {}, login: Record = { state: 'idle' }, ) => ({ - provider: 'grok-cli', + id: 'account-1', + slot: 1, + permanent: true, label: 'Account 1', status: 'Logged in', authenticated: true, @@ -314,6 +273,39 @@ async function openDashboard(options?: Parameters[ }; } +async function addDashboardAccount( + session: Awaited>, + label: string, +) { + const response = await fetch(`${session.dashboard.origin}/api/accounts`, { + method: 'POST', + headers: session.headers, + body: JSON.stringify({ label }), + }); + return { + response, + account: (await response.json()) as { id: string; label: string; slot: number }, + }; +} + +async function requestLoginTicket( + session: Awaited>, + accountId: string, +) { + const response = await fetch( + `${session.dashboard.origin}/api/accounts/${accountId}/login-ticket`, + { + method: 'POST', + headers: session.headers, + body: '{}', + }, + ); + return { + response, + path: ((await response.json()) as { path: string }).path, + }; +} + async function openIncompleteMutation(session: Awaited>) { const url = new URL(session.dashboard.origin); const socket = createConnection({ host: url.hostname, port: Number(url.port) }); @@ -337,16 +329,14 @@ async function openIncompleteMutation(session: Awaited) => boolean, ) { await vi.waitFor(async () => { const state = (await ( await fetch(`${dashboard.origin}/api/state`, { headers: { Cookie: cookie } }) ).json()) as { accounts: Record[] }; - expect(predicate(state.accounts.find((account) => account.provider === provider) ?? {})).toBe( - true, - ); + expect(predicate(state.accounts.find((account) => account.id === accountId) ?? {})).toBe(true); }); } @@ -390,7 +380,7 @@ describe('account dashboard loopback server', () => { expect(await state.json()).toMatchObject({ accounts: [ { - provider: 'grok-cli', + id: 'account-1', label: 'Account 1', authenticated: true, active: true, @@ -437,61 +427,66 @@ describe('account dashboard loopback server', () => { }); expect(added.status).toBe(201); - expect(await added.json()).toMatchObject({ provider: 'grok-cli-2', label: 'Work' }); - expect(loadConfig().config.accounts.items.at(-1)).toEqual({ - provider: 'grok-cli-2', + const account = (await added.json()) as { id: string; label: string; slot: number }; + expect(account).toMatchObject({ label: 'Work', slot: 2 }); + expect((await getAccountVault()).accounts.at(-1)).toMatchObject({ + id: account.id, label: 'Work', + slot: 2, }); }); - it('accepts account routes for two-digit provider aliases', async () => { + it('accepts account routes for opaque IDs after ten display slots', async () => { const session = await openDashboard(); + let lastId = ''; for (let account = 2; account <= 10; account += 1) { - await fetch(`${session.dashboard.origin}/api/accounts`, { + const response = await fetch(`${session.dashboard.origin}/api/accounts`, { method: 'POST', headers: session.headers, body: JSON.stringify({ label: `Account ${account}` }), }); + lastId = ((await response.json()) as { id: string }).id; } - const renamed = await fetch(`${session.dashboard.origin}/api/accounts/grok-cli-10`, { + const renamed = await fetch(`${session.dashboard.origin}/api/accounts/${lastId}`, { method: 'PATCH', headers: session.headers, body: JSON.stringify({ label: 'Account ten' }), }); expect(renamed.status).toBe(200); - expect(loadConfig().config.accounts.items.at(-1)).toEqual({ - provider: 'grok-cli-10', + expect((await getAccountVault()).accounts.at(-1)).toMatchObject({ + id: lastId, label: 'Account ten', }); }); it('renames, activates, logs out, and removes accounts through the shared manager', async () => { const session = await openDashboard(); - await fetch(`${session.dashboard.origin}/api/accounts`, { - method: 'POST', - headers: session.headers, - body: JSON.stringify({ label: 'Work' }), + const workId = (await addDashboardAccount(session, 'Work')).account.id; + await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === workId); + if (!account) throw new Error('Work account was not added.'); + account.credential = oauthCredential('work'); + account.revision = 1; }); - await session.extension.setCredential('grok-cli-2', oauthCredential('work')); - const renamed = await fetch(`${session.dashboard.origin}/api/accounts/grok-cli-2`, { + const renamed = await fetch(`${session.dashboard.origin}/api/accounts/${workId}`, { method: 'PATCH', headers: session.headers, body: JSON.stringify({ label: 'Client' }), }); - const activated = await fetch(`${session.dashboard.origin}/api/accounts/grok-cli-2/activate`, { + const activated = await fetch(`${session.dashboard.origin}/api/accounts/${workId}/activate`, { method: 'POST', headers: session.headers, body: '{}', }); - const removed = await fetch(`${session.dashboard.origin}/api/accounts/grok-cli-2`, { + const removed = await fetch(`${session.dashboard.origin}/api/accounts/${workId}`, { method: 'DELETE', headers: session.headers, body: '{}', }); - const loggedOut = await fetch(`${session.dashboard.origin}/api/accounts/grok-cli/logout`, { + const loggedOut = await fetch(`${session.dashboard.origin}/api/accounts/account-1/logout`, { method: 'POST', headers: session.headers, body: '{}', @@ -499,48 +494,36 @@ describe('account dashboard loopback server', () => { expect(renamed.status).toBe(200); expect(activated.status).toBe(200); - expect(session.extension.pi.setModel).toHaveBeenCalledWith( - expect.objectContaining({ provider: 'grok-cli-2', id: 'grok-build' }), - ); expect(removed.status).toBe(200); expect(loggedOut.status).toBe(200); - expect(loadConfig().config.accounts.items).toEqual([ - { provider: 'grok-cli', label: 'Account 1' }, + expect((await getAccountVault()).accounts).toMatchObject([ + { id: 'account-1', label: 'Account 1' }, ]); - expect(await session.extension.credentials.read('grok-cli')).toBeUndefined(); + expect((await getAccountVault()).accounts[0].credential).toBeUndefined(); }); it('redirects browser login without exposing credentials and accepts manual codes', async () => { const session = await openDashboard({ refreshAfterLogin: false }); - await fetch(`${session.dashboard.origin}/api/accounts`, { - method: 'POST', - headers: session.headers, - body: JSON.stringify({ label: 'Work' }), - }); - session.extension.loginFlows.set('grok-cli-2', async (callbacks) => { + const workId = (await addDashboardAccount(session, 'Work')).account.id; + oauthLogin.mockImplementationOnce(async (callbacks) => { callbacks.onAuth({ url: 'https://accounts.x.ai/authorize?state=browser-state' }); const code = await callbacks.onManualCodeInput?.(); if (code !== 'manual-code') throw new Error('manual code rejected'); return oauthCredential('dashboard-access'); }); - const ticket = await fetch(`${session.dashboard.origin}/api/accounts/grok-cli-2/login-ticket`, { - method: 'POST', - headers: session.headers, - body: '{}', - }); - const path = ((await ticket.json()) as { path: string }).path; - const redirect = await fetch(`${session.dashboard.origin}${path}`, { + const ticket = await requestLoginTicket(session, workId); + const redirect = await fetch(`${session.dashboard.origin}${ticket.path}`, { headers: { Cookie: session.cookie }, redirect: 'manual', }); - expect(ticket.status).toBe(201); + expect(ticket.response.status).toBe(201); expect(redirect.status).toBe(302); expect(redirect.headers.get('location')).toBe( 'https://accounts.x.ai/authorize?state=browser-state', ); - await fetch(`${session.dashboard.origin}/api/accounts/grok-cli-2/login-code`, { + await fetch(`${session.dashboard.origin}/api/accounts/${workId}/login-code`, { method: 'POST', headers: session.headers, body: JSON.stringify({ code: 'manual-code' }), @@ -548,7 +531,7 @@ describe('account dashboard loopback server', () => { await waitForAccount( session.dashboard, session.cookie, - 'grok-cli-2', + workId, (account) => account.authenticated === true, ); const state = await ( @@ -562,6 +545,92 @@ describe('account dashboard loopback server', () => { expect(state).not.toContain('browser-state'); }); + it('removes a new account when its first browser login fails', async () => { + const session = await openDashboard({ refreshAfterLogin: false }); + const workId = (await addDashboardAccount(session, 'Work')).account.id; + oauthLogin.mockRejectedValueOnce(new Error('login failed')); + const ticket = await requestLoginTicket(session, workId); + + expect( + ( + await fetch(`${session.dashboard.origin}${ticket.path}`, { + headers: { Cookie: session.cookie }, + }) + ).status, + ).toBe(502); + await vi.waitFor(async () => { + expect((await getAccountVault()).accounts.some((account) => account.id === workId)).toBe( + false, + ); + }); + }); + + it('keeps a new account authenticated by another process during its first login', async () => { + const session = await openDashboard({ refreshAfterLogin: false }); + const workId = (await addDashboardAccount(session, 'Work')).account.id; + const authorization = deferred>(); + oauthLogin.mockReturnValueOnce(authorization.promise); + const ticket = await requestLoginTicket(session, workId); + const login = fetch(`${session.dashboard.origin}${ticket.path}`, { + headers: { Cookie: session.cookie }, + }); + await vi.waitFor(() => expect(oauthLogin).toHaveBeenCalledOnce()); + await mutateAccountVault((vault) => { + const account = vault.accounts.find((candidate) => candidate.id === workId); + if (!account) throw new Error('Work account was not added.'); + account.credential = oauthCredential('other-process'); + account.revision += 1; + }); + authorization.resolve(oauthCredential('stale-login')); + + expect((await login).status).toBe(502); + await vi.waitFor(async () => { + expect( + (await getAccountVault()).accounts.find((account) => account.id === workId), + ).toMatchObject({ + revision: 1, + credential: { access: 'other-process' }, + }); + }); + }); + + it('keeps a new account when its first browser login is cancelled', async () => { + const session = await openDashboard({ refreshAfterLogin: false }); + const workId = (await addDashboardAccount(session, 'Work')).account.id; + const started = deferred(); + oauthLogin.mockImplementationOnce( + (callbacks) => + new Promise((_resolve, reject) => { + callbacks.signal?.addEventListener('abort', () => reject(new Error('Login cancelled')), { + once: true, + }); + started.resolve(); + }), + ); + const ticket = await requestLoginTicket(session, workId); + const login = fetch(`${session.dashboard.origin}${ticket.path}`, { + headers: { Cookie: session.cookie }, + }); + await started.promise; + + const cancelled = await fetch( + `${session.dashboard.origin}/api/accounts/${workId}/login-cancel`, + { + method: 'POST', + headers: session.headers, + body: '{}', + }, + ); + + expect(cancelled.status).toBe(202); + expect((await login).status).toBe(502); + await vi.waitFor(async () => { + expect((await getAccountVault()).accounts.some((account) => account.id === workId)).toBe( + true, + ); + }); + }); + it('rejects malformed or oversized mutations', async () => { const session = await openDashboard(); const malformed = await fetch(`${session.dashboard.origin}/api/accounts`, { @@ -713,23 +782,41 @@ describe('account dashboard loopback server', () => { const cases = [ { button: 'Switch', - initial: accountState({ active: false, label: 'Work', provider: 'grok-cli-2' }), - next: accountState({ label: 'Work', provider: 'grok-cli-2' }), + initial: accountState({ + id: 'work-id', + slot: 2, + permanent: false, + active: false, + label: 'Work', + }), + next: accountState({ id: 'work-id', slot: 2, permanent: false, label: 'Work' }), toast: 'Switched to Work.', }, { button: 'Rename', - initial: accountState({ label: 'Work', provider: 'grok-cli-2' }), - next: accountState({ label: 'Renamed', provider: 'grok-cli-2' }), + initial: accountState({ id: 'work-id', slot: 2, permanent: false, label: 'Work' }), + next: accountState({ + id: 'work-id', + slot: 2, + permanent: false, + label: 'Renamed', + }), response: { label: 'Renamed' }, value: 'Renamed', toast: 'Renamed to Renamed.', }, { button: 'Remove', - initial: accountState({ active: false, label: 'Work', provider: 'grok-cli-2' }), + initial: accountState({ + id: 'work-id', + slot: 2, + permanent: false, + active: false, + label: 'Work', + }), next: undefined, - toast: 'Removed Work.', + response: { warning: 'Credential cleanup continues in the background.' }, + toast: 'Removed Work. Credential cleanup continues in the background.', }, { button: 'Log out', @@ -777,6 +864,56 @@ describe('account dashboard loopback server', () => { } }); + it('hides log out when the permanent account is already logged out', async () => { + const session = await openDashboard(); + const browser = await browserDashboard(session, [ + { + refreshing: false, + accounts: [ + accountState({ + authenticated: false, + active: false, + status: 'Login required', + }), + ], + }, + ]); + const actions = [...browser.window.document.querySelectorAll('.card-actions button')].map( + (button) => button.textContent, + ); + + expect(actions).toContain('Log in'); + expect(actions).toContain('Rename'); + expect(actions).not.toContain('Log out'); + await browser.close(); + }); + + it('does not offer saved-account switching while an environment token is active', async () => { + const session = await openDashboard(); + const browser = await browserDashboard(session, [ + { + refreshing: false, + accounts: [ + accountState({ environment: true }), + accountState({ + id: 'work-id', + slot: 2, + permanent: false, + active: false, + label: 'Work', + }), + ], + }, + ]); + + expect( + [...browser.window.document.querySelectorAll('button')].map((button) => + button.textContent?.trim(), + ), + ).not.toContain('Switch'); + await browser.close(); + }); + it('throttles state-field drawing and resizes through ResizeObserver', async () => { const session = await openDashboard(); const browser = await browserDashboard( diff --git a/tests/provider/modelMigration.test.ts b/tests/provider/modelMigration.test.ts new file mode 100644 index 0000000..4bfbd84 --- /dev/null +++ b/tests/provider/modelMigration.test.ts @@ -0,0 +1,71 @@ +import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { afterEach, describe, expect, it } from 'vitest'; +import { migrateSavedModelProviders } from '../../src/provider/modelMigration.js'; +import { writeTestJson } from '../stateTestHelpers.js'; + +const tempDirs: string[] = []; + +function paths() { + const agentDir = mkdtempSync(join(tmpdir(), 'pi-grok-model-migration-')); + tempDirs.push(agentDir); + return { + agentDir, + cwd: join(agentDir, 'project'), + sessions: join(agentDir, 'sessions', 'project'), + }; +} + +afterEach(() => { + tempDirs.splice(0).forEach((path) => { + rmSync(path, { recursive: true, force: true }); + }); +}); + +describe('saved model provider migration', () => { + it('rewrites released settings without changing session data', async () => { + const target = paths(); + writeTestJson(join(target.agentDir, 'settings.json'), { + defaultProvider: 'grok-cli-2', + defaultModel: 'grok-build', + enabledModels: ['grok-cli-2/grok-build', 'grok-cli/grok-build', 'openai/gpt-5'], + }); + const sessionPath = join(target.sessions, 'session.jsonl'); + mkdirSync(target.sessions, { recursive: true }); + writeFileSync( + sessionPath, + `${JSON.stringify({ + type: 'model_change', + provider: 'grok-cli-10', + modelId: 'grok-build', + })}\n`, + ); + + await expect(migrateSavedModelProviders(target.agentDir, target.cwd)).resolves.toEqual({ + migrated: true, + }); + expect(JSON.parse(readFileSync(join(target.agentDir, 'settings.json'), 'utf8'))).toEqual({ + defaultProvider: 'grok-cli', + defaultModel: 'grok-build', + enabledModels: ['grok-cli/grok-build', 'openai/gpt-5'], + }); + expect(JSON.parse(readFileSync(sessionPath, 'utf8'))).toMatchObject({ + provider: 'grok-cli-10', + }); + await expect(migrateSavedModelProviders(target.agentDir, target.cwd)).resolves.toEqual({ + migrated: false, + }); + }); + + it('does not parse or replace malformed session data', async () => { + const target = paths(); + mkdirSync(target.sessions, { recursive: true }); + writeFileSync(join(target.sessions, 'session.jsonl'), '{ malformed\n'); + + const result = await migrateSavedModelProviders(target.agentDir, target.cwd); + + expect(result).toEqual({ migrated: false }); + expect(readFileSync(join(target.sessions, 'session.jsonl'), 'utf8')).toBe('{ malformed\n'); + }); +}); diff --git a/tests/provider/package.test.ts b/tests/provider/package.test.ts index 67cad75..7dd2eb0 100644 --- a/tests/provider/package.test.ts +++ b/tests/provider/package.test.ts @@ -33,15 +33,12 @@ describe('npm package manifest', () => { expect(existsSync(new URL('../../vitest.config.ts', import.meta.url))).toBe(true); }); - it('declares the Pi runtime version required by dashboard auth and web search', () => { + it('declares the Pi runtime version required by dashboard auth', () => { expect(packageJson.peerDependencies?.['@earendil-works/pi-ai']).toBe('>=0.80.9'); expect(packageJson.peerDependencies?.['@earendil-works/pi-coding-agent']).toBe('>=0.80.9'); expect(packageJson.peerDependencies?.['@earendil-works/pi-tui']).toBe('>=0.80.9'); - expect(packageJson.peerDependencies?.['pi-web-access']).toBe('>=0.13.0'); - expect(packageJson.dependencies?.jiti).toBeUndefined(); - expect(packageJson.dependencies?.typebox).toBeUndefined(); - expect(packageJson.dependencies?.ignore).toBe('7.0.5'); - expect(packageJson.dependencies?.minimatch).toBe('10.2.5'); + expect(packageJson.dependencies).toBeUndefined(); + expect(packageJson.devDependencies?.['@types/proper-lockfile']).toBeUndefined(); }); }); @@ -67,29 +64,22 @@ describe('repository layout', () => { 'src/index.ts', 'src/models/catalog.ts', 'src/payload/sanitize.ts', + 'src/provider/accountMigration.ts', + 'src/provider/accountRouting.ts', + 'src/provider/accountVault.ts', 'src/provider/accounts.ts', 'src/provider/billing.ts', 'src/provider/dashboard/server.ts', + 'src/provider/modelMigration.ts', 'src/provider/quotaCache.ts', 'src/provider/register.ts', + 'src/provider/requestOwnership.ts', 'src/provider/rotation.ts', + 'src/provider/sessionAccountSelection.ts', 'src/provider/stream.ts', - 'src/provider/toolScope.ts', 'src/provider/usage.ts', 'src/shared/errors.ts', 'src/storage.ts', - 'src/tools/files.ts', - 'src/tools/glob.ts', - 'src/tools/read.ts', - 'src/tools/register.ts', - 'src/tools/rendering.ts', - 'src/tools/search.ts', - 'src/tools/shell.ts', - 'src/tools/webSearch.ts', - 'src/tools/webSearchDelegate.ts', - 'src/vision/cache.ts', - 'src/vision/describe.ts', - 'src/vision/register.ts', ]); }); @@ -98,4 +88,13 @@ describe('repository layout', () => { expect(existsSync(new URL(`../../src/${file}`, import.meta.url))).toBe(false); } }); + + it('uses the exact Pi loader alias for OpenAI Responses streaming', () => { + const source = readFileSync(new URL('../../src/provider/register.ts', import.meta.url), 'utf8'); + + expect(source).toContain( + "import { streamSimpleOpenAIResponses } from '@earendil-works/pi-ai/compat';", + ); + expect(source).not.toContain('@earendil-works/pi-ai/api/openai-responses'); + }); }); diff --git a/tests/provider/quotaCache.test.ts b/tests/provider/quotaCache.test.ts index c0be262..8387db2 100644 --- a/tests/provider/quotaCache.test.ts +++ b/tests/provider/quotaCache.test.ts @@ -8,9 +8,10 @@ import { loadQuotaCache, removeQuotaUsage, saveQuotaUsage, + saveQuotaUsageWhen, } from '../../src/provider/quotaCache.js'; import { getQuotaCachePath } from '../../src/storage.js'; -import { useTempHome } from '../vision/helpers.js'; +import { useTempHome } from '../stateTestHelpers.js'; const setupHome = useTempHome(); const usage = (used: number, weekly = 60): BillingUsage => ({ @@ -87,6 +88,16 @@ describe('Grok CLI quota cache', () => { expect(Object.keys(loadQuotaCache().accounts)).toEqual(['grok-cli']); }); + it('conditionally saves usage to a path override', async () => { + const home = setupHome(); + const path = `${home}/quota-cache.json`; + + await saveQuotaUsageWhen('account-1', usage(300), () => true, path); + + expect(loadQuotaCache(path).accounts['account-1']?.monthly.used).toBe(300); + expect(existsSync(getQuotaCachePath())).toBe(false); + }); + it('runs a queued update after the previous update fails', async () => { setupHome(); const invalid = Object.defineProperty({}, 'monthly', { diff --git a/tests/provider/register.test.ts b/tests/provider/register.test.ts index df61071..2edbcc9 100644 --- a/tests/provider/register.test.ts +++ b/tests/provider/register.test.ts @@ -1,20 +1,28 @@ -import { existsSync, mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs'; +import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; -import type { Api, Model, OAuthCredentials, OAuthLoginCallbacks } from '@earendil-works/pi-ai'; +import type { OAuthLoginCallbacks } from '@earendil-works/pi-ai'; import type { ExtensionAPI, ProviderConfig } from '@earendil-works/pi-coding-agent'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -import { DEFAULT_CONFIG, loadConfig, saveConfig } from '../../src/config.js'; +import { + ACCOUNT_VAULT_MARKER, + getAccountVault, + mutateAccountVault, +} from '../../src/provider/accountVault.js'; import { loadQuotaCache, saveQuotaUsage } from '../../src/provider/quotaCache.js'; import { getQuotaCachePath } from '../../src/storage.js'; -import { GROK_SHIM_TOOL_NAMES } from '../../src/tools/register.js'; -import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; -import { plainTheme as theme } from '../tools/toolTestHelpers.js'; -import { saveTestAccounts } from '../vision/helpers.js'; - -const { mockOauthLogin, mockPiWebAccessInstalled } = vi.hoisted(() => ({ +import { + deferred, + oauthCredential, + saveTestAccounts, + setAccount1Credential, + writePiCredentials, + writePiVaultMarker, +} from '../stateTestHelpers.js'; + +const { mockOauthLogin, mockProviderStream } = vi.hoisted(() => ({ mockOauthLogin: vi.fn(), - mockPiWebAccessInstalled: vi.fn(() => true), + mockProviderStream: vi.fn(), })); vi.mock('../../src/auth/oauth.js', async (importOriginal) => { @@ -22,15 +30,9 @@ vi.mock('../../src/auth/oauth.js', async (importOriginal) => { return { ...actual, login: mockOauthLogin }; }); -vi.mock('../../src/tools/webSearchDelegate.js', async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - isPiWebAccessInstalled: () => mockPiWebAccessInstalled(), - bindLivePiWebAccess: vi.fn(), - ensureWebSearchDelegate: vi.fn(async () => undefined), - }; -}); +vi.mock('@earendil-works/pi-ai/compat', () => ({ + streamSimpleOpenAIResponses: mockProviderStream, +})); interface CommandConfig { handler: (args: string[], ctx: TestContext) => Promise; @@ -38,12 +40,6 @@ interface CommandConfig { interface RegisteredTool { name: string; - renderCall?: (...args: unknown[]) => Renderable; - renderResult?: (...args: unknown[]) => Renderable; -} - -interface Renderable { - render: (width: number) => string[]; } interface TestContext { @@ -55,6 +51,7 @@ interface TestContext { model?: { provider: string; id: string }; sessionManager?: { getSessionId: () => string; + getBranch?: () => unknown[]; }; ui: { notify: (message: string, level: string) => void; @@ -76,6 +73,27 @@ beforeEach(() => { refresh: 'new-refresh', expires: Date.now() + 60_000, }); + mockProviderStream.mockReset(); + mockProviderStream.mockImplementation(() => ({ + async *[Symbol.asyncIterator]() {}, + result: vi.fn(async () => ({ + role: 'assistant', + content: [], + api: 'openai-responses', + provider: 'grok-cli', + model: 'grok-build', + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: 'stop', + timestamp: Date.now(), + })), + })); process.env.TZ = 'America/New_York'; const dir = mkdtempSync(join(tmpdir(), 'pi-grok-cli-home-')); mkdirSync(join(dir, '.pi')); @@ -104,8 +122,7 @@ afterEach(() => { for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true }); }); -async function setupExtension(initialActiveTools = ['read', 'bash'], piWebAccessInstalled = true) { - vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(piWebAccessInstalled); +async function setupExtension(initialActiveTools = ['read', 'bash']) { const commands = new Map(); const providers = new Map(); const tools = new Map(); @@ -117,6 +134,7 @@ async function setupExtension(initialActiveTools = ['read', 'bash'], piWebAccess }); const setModel = vi.fn(async (_model: { provider: string; id: string }) => true); const sendUserMessage = vi.fn(); + const entries: { customType: string; data: unknown }[] = []; const registerGrokCli = (await import('../../src/index.js')).default; registerGrokCli({ registerProvider(name: string, config: ProviderConfig) { @@ -130,7 +148,9 @@ async function setupExtension(initialActiveTools = ['read', 'bash'], piWebAccess commands.set(name, config as CommandConfig); }, registerEntryRenderer() {}, - appendEntry() {}, + appendEntry(customType: string, data: unknown) { + entries.push({ customType, data }); + }, registerTool(tool: RegisteredTool) { tools.set(tool.name, tool); }, @@ -153,6 +173,7 @@ async function setupExtension(initialActiveTools = ['read', 'bash'], piWebAccess setActiveTools, setModel, sendUserMessage, + entries, } as unknown as ExtensionAPI); return { commands, @@ -162,16 +183,44 @@ async function setupExtension(initialActiveTools = ['read', 'bash'], piWebAccess setActiveTools, setModel, sendUserMessage, + entries, async emit(event: string, data: unknown, ctx: TestContext) { for (const handler of allHandlers.get(event) ?? []) await handler(data, ctx); }, getActiveTools: () => activeTools, - replaceActiveTools(nextTools: string[]) { - activeTools = [...nextTools]; + }; +} + +function sessionContext(sessionId: string, accountId?: string): TestContext { + return { + modelRegistry: { getAll: () => [] }, + sessionManager: { + getSessionId: () => sessionId, + getBranch: () => + accountId + ? [ + { + type: 'custom', + id: 'account-entry', + parentId: null, + timestamp: new Date().toISOString(), + customType: 'grok-cli-active-account-v1', + data: { accountId }, + }, + ] + : [], }, + ui: { notify: vi.fn() }, }; } +async function drain(stream: AsyncIterable | undefined) { + if (!stream) throw new Error('Grok CLI test stream is missing.'); + for await (const _event of stream) { + // The route setup completes before this empty test stream ends. + } +} + function statusContext(notify: TestContext['ui']['notify']): TestContext { return { modelRegistry: { @@ -180,6 +229,10 @@ function statusContext(notify: TestContext['ui']['notify']): TestContext { { provider: 'grok-cli', id: 'grok-composer-2.5-fast' }, ], }, + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], + }, ui: { notify }, }; } @@ -187,6 +240,10 @@ function statusContext(notify: TestContext['ui']['notify']): TestContext { function emptyStatusContext(notify: TestContext['ui']['notify']): TestContext { return { modelRegistry: { getAll: () => [] }, + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], + }, ui: { notify }, }; } @@ -195,34 +252,14 @@ function contextForModel(provider: string, id = `${provider}-model`): TestContex return { model: { provider, id }, modelRegistry: { getAll: () => [] }, + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], + }, ui: { notify: vi.fn() }, }; } -function renderText(component: Renderable): string { - return component - .render(120) - .map((line) => line.trimEnd()) - .join('\n'); -} - -function renderContext(args: Record = {}) { - return { - args, - toolCallId: 'tool-call-id', - invalidate: () => {}, - lastComponent: undefined, - state: {}, - cwd: '/project', - executionStarted: true, - argsComplete: true, - isPartial: false, - expanded: false, - showImages: true, - isError: false, - }; -} - function setupHome() { const dir = mkdtempSync(join(tmpdir(), 'pi-grok-cli-home-')); mkdirSync(join(dir, '.pi')); @@ -369,20 +406,20 @@ describe('Grok CLI status command', () => { ); }); - it('uses the registered provider token when no env token is set', async () => { + it('uses the active vault account when no environment token is set', async () => { delete process.env.GROK_CLI_OAUTH_TOKEN; setupHome(); - saveConfig({ - ...DEFAULT_CONFIG, - accounts: { - nextAccountNumber: 3, - selectedProvider: 'grok-cli-2', - items: [ - { provider: 'grok-cli', label: 'Personal' }, - { provider: 'grok-cli-2', label: 'Work' }, - ], - }, + await mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + vault.accounts[0].credential = { + access: 'provider-token', + refresh: 'provider-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + vault.activeAccountId = 'account-1'; }); + writePiVaultMarker(); const fetchMock = vi.fn(async () => billingResponse(4000, 100, '2026-07-01T00:00:00+00:00'), ); @@ -395,18 +432,52 @@ describe('Grok CLI status command', () => { ...statusContext(notify), modelRegistry: { ...statusContext(notify).modelRegistry, - getAll: () => [{ provider: 'grok-cli-2', id: 'grok-build' }], + getAll: () => [{ provider: 'grok-cli', id: 'grok-build' }], getApiKeyForProvider, }, }); - expect(getApiKeyForProvider).toHaveBeenCalledWith('grok-cli-2'); + expect(getApiKeyForProvider).not.toHaveBeenCalled(); expect(fetchMock.mock.calls[0]?.[1]?.headers).toMatchObject({ authorization: 'Bearer provider-token', }); expect(notify.mock.calls.at(-1)?.[0]).toContain('100 / 4,000 used 3%'); - expect(loadQuotaCache().accounts['grok-cli-2']?.monthly.used).toBe(100); - expect(loadQuotaCache().accounts['grok-cli']).toBeUndefined(); + expect(loadQuotaCache().accounts['account-1']?.monthly.used).toBe(100); + }); + + it('does not cache usage after the account changes during the request', async () => { + await mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + vault.accounts[0].credential = { + ...oauthCredential('one'), + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + vault.activeAccountId = 'account-1'; + }); + writePiVaultMarker(); + const monthly = deferred(); + globalThis.fetch = vi.fn(async (input) => { + const url = typeof input === 'string' ? input : input.toString(); + if (url.includes('format=credits')) { + return creditsResponse(10, '2026-07-14T00:19:56+00:00'); + } + return monthly.promise; + }); + const extension = await setupExtension(); + const notify = vi.fn(); + + const pending = extension.commands.get('grok-cli-usage')?.handler([], statusContext(notify)); + await vi.waitFor(() => expect(globalThis.fetch).toHaveBeenCalledOnce()); + await mutateAccountVault((vault) => { + delete vault.accounts[0].credential; + vault.accounts[0].revision += 1; + delete vault.activeAccountId; + }); + monthly.resolve(billingResponse(4000, 100, '2026-07-01T00:00:00+00:00')); + await pending; + + expect(loadQuotaCache().accounts['account-1']).toBeUndefined(); }); it('does not fetch billing when no token is available', async () => { @@ -422,6 +493,7 @@ describe('Grok CLI status command', () => { [ ' Usage:', ' no billing data available — run /login grok-cli or set GROK_CLI_OAUTH_TOKEN', + ' Reason Grok CLI account migration is ready. Please run /login and select Grok CLI.', ].join('\n'), ); }); @@ -436,7 +508,7 @@ describe('Grok CLI status command', () => { const extension = await setupExtension(); await runStatus(extension); - expect(loadQuotaCache().accounts['grok-cli']).toMatchObject({ + expect(loadQuotaCache().accounts['account-1']).toMatchObject({ monthly: { monthlyLimit: 4000, used: 1421 }, }); expect(timeout).toHaveBeenCalledWith(30_000); @@ -485,7 +557,7 @@ describe('Grok CLI status command', () => { process.env.GROK_CLI_OAUTH_TOKEN = 'env-token'; setupHome(); await saveQuotaUsage( - 'grok-cli', + 'account-1', { monthly: { monthlyLimit: 4000, @@ -580,149 +652,373 @@ describe('Grok CLI status command', () => { }); describe('Grok CLI provider registration', () => { - it('clears cached quota after a successful OAuth login', async () => { - await saveQuotaUsage('grok-cli', { - monthly: { - monthlyLimit: 2000, - used: 300, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, + it('registers one Grok CLI provider for all configured accounts', async () => { + saveTestAccounts(); + + const extension = await setupExtension(); + + expect([...extension.providers.keys()]).toEqual(['grok-cli']); + expect(extension.providers.get('grok-cli')?.name).toBe('Grok CLI'); + }); + + it('routes two live Pi sessions through their own selected accounts', async () => { + await mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + vault.accounts[0].credential = { + access: 'one', + refresh: 'one-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + vault.accounts.push({ + id: 'work-id', + slot: 2, + label: 'Work', + credential: { + access: 'two', + refresh: 'two-refresh', + expires: Date.now() + 300_000, + }, + revision: 1, + }); + vault.nextSlot = 3; + vault.activeAccountId = 'account-1'; }); + writePiVaultMarker(); + const first = await setupExtension(); + const second = await setupExtension(); + await first.emit( + 'session_start', + { type: 'session_start', reason: 'startup' }, + sessionContext('session-a', 'account-1'), + ); + await second.emit( + 'session_start', + { type: 'session_start', reason: 'startup' }, + sessionContext('session-b', 'work-id'), + ); + const firstProvider = first.providers.get('grok-cli'); + const secondProvider = second.providers.get('grok-cli'); + const firstModel = firstProvider?.models?.[0]; + const secondModel = secondProvider?.models?.[0]; + if (!firstModel || !secondModel) throw new Error('Grok CLI test model is missing.'); + + await Promise.all([ + drain( + firstProvider.streamSimple?.( + { + ...firstModel, + provider: 'grok-cli', + api: 'openai-responses', + baseUrl: 'https://cli-chat-proxy.grok.com', + }, + { messages: [] }, + { sessionId: 'session-a' }, + ), + ), + drain( + secondProvider.streamSimple?.( + { + ...secondModel, + provider: 'grok-cli', + api: 'openai-responses', + baseUrl: 'https://cli-chat-proxy.grok.com', + }, + { messages: [] }, + { sessionId: 'session-b' }, + ), + ), + ]); + + expect( + Object.fromEntries( + mockProviderStream.mock.calls.map((call) => [call[2]?.sessionId, call[2]?.apiKey]), + ), + ).toEqual({ 'session-a': 'one', 'session-b': 'two' }); + expect((await getAccountVault()).activeAccountId).toBe('account-1'); + }); + + it('observes a failed provider result without creating an unhandled rejection', async () => { + await setAccount1Credential('one'); + writePiVaultMarker(); + mockProviderStream.mockImplementationOnce(() => ({ + async *[Symbol.asyncIterator]() {}, + result: () => Promise.reject(new Error('provider failed')), + })); const extension = await setupExtension(); + const provider = extension.providers.get('grok-cli'); + const model = provider?.models?.[0]; + if (!model) throw new Error('Grok CLI test model is missing.'); - await extension.providers.get('grok-cli')?.oauth?.login({} as OAuthLoginCallbacks); + await drain( + provider.streamSimple?.( + { + ...model, + provider: 'grok-cli', + api: 'openai-responses', + baseUrl: 'https://cli-chat-proxy.grok.com', + }, + { messages: [] }, + { sessionId: 'session-a' }, + ), + ); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); - expect(mockOauthLogin).toHaveBeenCalledOnce(); - expect(loadQuotaCache().accounts['grok-cli']).toBeUndefined(); + it('stores the default account in a new Pi session', async () => { + await setAccount1Credential('one'); + const extension = await setupExtension(); + + await extension.emit( + 'session_start', + { type: 'session_start', reason: 'startup' }, + sessionContext('session-a'), + ); + + expect(extension.entries).toContainEqual({ + customType: 'grok-cli-active-account-v1', + data: { accountId: 'account-1' }, + }); }); - it('makes a recently exhausted account eligible after successful OAuth login', async () => { - saveTestAccounts('grok-cli'); + it('uses the session account for the usage command', async () => { + await mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + vault.accounts[0].credential = { + ...oauthCredential('one'), + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + vault.accounts.push({ + id: 'work-id', + slot: 2, + label: 'Work', + credential: { + ...oauthCredential('two'), + expires: Date.now() + 300_000, + }, + revision: 1, + }); + vault.nextSlot = 3; + vault.activeAccountId = 'account-1'; + }); + writePiVaultMarker(); + globalThis.fetch = vi.fn(async () => + billingResponse(4000, 100, '2026-07-01T00:00:00+00:00'), + ); const extension = await setupExtension(); - const context = { - model: { provider: 'grok-cli', id: 'grok-build' }, - modelRegistry: { - getProviderAuthStatus: (provider: string) => ({ - configured: provider.startsWith('grok-cli'), - }), - find: (provider: string, id: string) => ({ provider, id }), - getAll: () => [], + const context = sessionContext('session-b', 'work-id'); + context.modelRegistry.getAll = () => [{ provider: 'grok-cli', id: 'grok-build' }]; + await extension.emit('session_start', { type: 'session_start', reason: 'startup' }, context); + + await extension.commands.get('grok-cli-usage')?.handler([], context); + + expect(globalThis.fetch).toHaveBeenCalledWith( + 'https://cli-chat-proxy.grok.com/v1/billing', + expect.objectContaining({ + headers: expect.objectContaining({ authorization: 'Bearer two' }), + }), + ); + }); + + it('uses OAuth and a custom stream in stored-account mode', async () => { + delete process.env.GROK_CLI_OAUTH_TOKEN; + + const extension = await setupExtension(); + const provider = extension.providers.get('grok-cli'); + + expect(provider?.oauth).toBeDefined(); + expect(provider?.apiKey).toBeUndefined(); + expect(provider?.streamSimple).toBeTypeOf('function'); + }); + + it('uses only the environment API-key method when Pi has no stored OAuth credential', async () => { + process.env.GROK_CLI_OAUTH_TOKEN = 'environment-token'; + writePiCredentials({}); + + const provider = (await setupExtension()).providers.get('grok-cli'); + + expect(provider?.apiKey).toBe('$GROK_CLI_OAUTH_TOKEN'); + expect(provider?.oauth).toBeUndefined(); + expect(provider?.streamSimple).toBeTypeOf('function'); + }); + + it('keeps only OAuth registration when an environment token and stored OAuth coexist', async () => { + process.env.GROK_CLI_OAUTH_TOKEN = 'environment-token'; + writePiCredentials({ + 'grok-cli': { + type: 'oauth', + access: ACCOUNT_VAULT_MARKER, + refresh: ACCOUNT_VAULT_MARKER, + expires: Number.MAX_SAFE_INTEGER, }, - ui: { notify: vi.fn() }, - }; - const exhausted = (provider: string) => ({ - type: 'message_end', - message: { - role: 'assistant', - provider, - model: 'grok-build', - stopReason: 'error', - errorMessage: 'OpenAI API error (402): 402 "Grok Build usage balance exhausted"', + }); + + const provider = (await setupExtension()).providers.get('grok-cli'); + + expect(provider?.apiKey).toBeUndefined(); + expect(provider?.oauth).toBeDefined(); + }); + + it('stores a normal Grok login in Account 1 and returns only the Pi marker', async () => { + const extension = await setupExtension(); + + const result = await extension.providers + .get('grok-cli') + ?.oauth?.login({} as OAuthLoginCallbacks); + + expect(result).toEqual({ + access: ACCOUNT_VAULT_MARKER, + refresh: ACCOUNT_VAULT_MARKER, + expires: Number.MAX_SAFE_INTEGER, + }); + await expect(getAccountVault()).resolves.toMatchObject({ + activeAccountId: 'account-1', + accounts: [ + { + id: 'account-1', + credential: { + access: 'new-access', + refresh: 'new-refresh', + }, + revision: 1, + }, + ], + }); + }); + + it('does not restore Account 1 after it changes during OAuth login', async () => { + const authorization = deferred>(); + mockOauthLogin.mockReturnValueOnce(authorization.promise); + const extension = await setupExtension(); + + const pending = extension.providers.get('grok-cli')?.oauth?.login({} as OAuthLoginCallbacks); + await mutateAccountVault((vault) => { + vault.accounts[0].revision += 1; + }); + authorization.resolve(oauthCredential('late-login')); + + await expect(pending).rejects.toThrow('account changed while login was in progress'); + expect((await getAccountVault()).accounts[0]).toMatchObject({ revision: 1 }); + expect((await getAccountVault()).accounts[0].credential).toBeUndefined(); + }); + + it('reconnects an existing vault with no Pi credential or an existing marker', async () => { + await mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + vault.accounts[0].credential = { + access: 'saved-access', + refresh: 'saved-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision = 1; + vault.activeAccountId = 'account-1'; + }); + const extension = await setupExtension(); + const provider = extension.providers.get('grok-cli'); + + const disconnectedResult = await provider?.oauth?.login({} as OAuthLoginCallbacks); + + writePiCredentials({ + 'grok-cli': { + type: 'oauth', + access: ACCOUNT_VAULT_MARKER, + refresh: ACCOUNT_VAULT_MARKER, + expires: Number.MAX_SAFE_INTEGER, }, }); + const connectedResult = await provider?.oauth?.login({} as OAuthLoginCallbacks); - await extension.emit('message_end', exhausted('grok-cli'), context as TestContext); - await extension.emit('agent_settled', { type: 'agent_settled' }, context as TestContext); - context.model = { provider: 'grok-cli-2', id: 'grok-build' }; - await extension.emit( - 'input', - { type: 'input', source: 'extension', text: 'continue after rotation' }, - context as TestContext, + expect(disconnectedResult?.access).toBe(ACCOUNT_VAULT_MARKER); + expect(connectedResult?.access).toBe(ACCOUNT_VAULT_MARKER); + expect(mockOauthLogin).not.toHaveBeenCalled(); + }); + + it('installs the marker after copying the released Account 1 credential', async () => { + saveTestAccounts('grok-cli'); + writePiCredentials({ 'grok-cli': oauthCredential('released-account') }); + const extension = await setupExtension(); + + expect( + (await extension.providers.get('grok-cli')?.oauth?.login({} as OAuthLoginCallbacks))?.access, + ).toBe(ACCOUNT_VAULT_MARKER); + expect(mockOauthLogin).not.toHaveBeenCalled(); + const vault = await getAccountVault(); + expect(vault).toMatchObject({ + activeAccountId: 'account-1', + migration: { markerInstallPending: true }, + }); + expect(vault.accounts).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: 'account-1', + credential: expect.objectContaining({ access: 'released-account' }), + }), + ]), ); + }); + + it('does not replace a newer Account 1 login while installing the marker', async () => { + saveTestAccounts('grok-cli'); + writePiCredentials({ 'grok-cli': oauthCredential('released-account') }); + const extension = await setupExtension(); await extension.emit( - 'message_end', - { - type: 'message_end', - message: { - role: 'assistant', - provider: 'grok-cli-2', - model: 'grok-build', - stopReason: 'stop', - }, - }, - context as TestContext, + 'session_start', + { type: 'session_start', reason: 'startup' }, + sessionContext('session-a'), ); - await extension.emit('agent_settled', { type: 'agent_settled' }, context as TestContext); + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'new-dashboard-login', + refresh: 'new-dashboard-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts[0].revision += 1; + }); await extension.providers.get('grok-cli')?.oauth?.login({} as OAuthLoginCallbacks); - await extension.emit('message_end', exhausted('grok-cli-2'), context as TestContext); - await extension.emit('agent_settled', { type: 'agent_settled' }, context as TestContext); - expect(extension.setModel.mock.calls.map(([model]) => model.provider)).toEqual([ - 'grok-cli-2', - 'grok-cli', - ]); + expect((await getAccountVault()).accounts[0].credential?.access).toBe('new-dashboard-login'); + }); + + it('clears cached quota after a successful OAuth login', async () => { + await saveQuotaUsage('account-1', { + monthly: { + monthlyLimit: 2000, + used: 300, + billingPeriodEnd: '2026-08-01T00:00:00.000Z', + }, + }); + const extension = await setupExtension(); + + await extension.providers.get('grok-cli')?.oauth?.login({} as OAuthLoginCallbacks); + + expect(mockOauthLogin).toHaveBeenCalledOnce(); + expect(loadQuotaCache().accounts['account-1']).toBeUndefined(); }); it('registers provider metadata and OAuth helpers', async () => { const extension = await setupExtension(); const provider = extension.providers.get('grok-cli'); - expect(provider?.name).toBe('Grok CLI — Account 1'); + expect(provider?.name).toBe('Grok CLI'); expect(provider?.api).toBe('openai-responses'); - expect(provider?.apiKey).toBe('$GROK_CLI_OAUTH_TOKEN'); - expect(provider?.streamSimple).toBeUndefined(); + expect(provider?.apiKey).toBeUndefined(); + expect(provider?.streamSimple).toBeTypeOf('function'); expect(provider?.models?.map((model) => model.id)).toContain('grok-build'); expect(provider?.oauth?.usesCallbackServer).toBe(true); expect(provider?.oauth?.getApiKey({ access: 'access-token', refresh: '', expires: 0 })).toBe( 'access-token', ); - expect( - provider?.oauth?.modifyModels?.( - [ - { provider: 'grok-cli', id: 'grok-build', baseUrl: 'old' } as Model, - { provider: 'openai', id: 'gpt-4', baseUrl: 'keep' } as Model, - ], - { - access: 'access-token', - refresh: 'refresh-token', - expires: 123, - baseUrl: 'https://example.invalid/custom///', - } as OAuthCredentials, - ), - ).toEqual([ - { - provider: 'grok-cli', - id: 'grok-build', - baseUrl: 'https://example.invalid/custom', - }, - { provider: 'openai', id: 'gpt-4', baseUrl: 'keep' }, - ]); + expect(provider?.oauth?.modifyModels).toBeUndefined(); }); - it('registers configured aliases with independent OAuth keys and labeled names', async () => { + it('does not register provider aliases from the released account config', async () => { saveTestAccounts(); const extension = await setupExtension(); - const base = extension.providers.get('grok-cli'); - const alias = extension.providers.get('grok-cli-2'); - - expect(base?.name).toBe('Grok CLI — Personal'); - expect(base?.apiKey).toBe('$GROK_CLI_OAUTH_TOKEN'); - expect(alias?.name).toBe('Grok CLI — Work'); - expect(alias?.apiKey).toBeUndefined(); - expect(alias?.models?.map((model) => model.id)).toEqual(base?.models?.map((model) => model.id)); - expect( - alias?.oauth?.modifyModels?.( - [ - { provider: 'grok-cli', id: 'grok-build', baseUrl: 'base' } as Model, - { provider: 'grok-cli-2', id: 'grok-build', baseUrl: 'old' } as Model, - ], - { - access: 'access-token', - refresh: 'refresh-token', - expires: 123, - baseUrl: 'https://example.invalid/work///', - } as OAuthCredentials, - ), - ).toEqual([ - { provider: 'grok-cli', id: 'grok-build', baseUrl: 'base' }, - { - provider: 'grok-cli-2', - id: 'grok-build', - baseUrl: 'https://example.invalid/work', - }, - ]); + + expect([...extension.providers.keys()]).toEqual(['grok-cli']); }); it('adds conversation affinity headers only for Grok requests', async () => { @@ -753,7 +1049,7 @@ describe('Grok CLI provider registration', () => { sessionManager: { getSessionId: () => 'session-alias' }, }); - expect(aliasEvent.headers).toEqual({ 'x-grok-conv-id': 'session-alias' }); + expect(aliasEvent.headers).toEqual({}); }); it('sanitizes Grok provider requests with the current session id', async () => { @@ -790,11 +1086,7 @@ describe('Grok CLI provider registration', () => { }, ); - expect(aliasResult).toEqual({ - input: [], - instructions: 'alias instruction', - prompt_cache_key: 'session-alias', - }); + expect(aliasResult).toBeUndefined(); }); it('leaves non-Grok provider requests untouched', async () => { @@ -823,6 +1115,10 @@ describe('Grok CLI provider registration', () => { {}, { modelRegistry: { getAll: () => [] }, + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], + }, ui: { notify }, }, ); @@ -834,17 +1130,41 @@ describe('Grok CLI provider registration', () => { }); }); -describe('Grok CLI tool scoping', () => { +describe('Grok CLI feature registration', () => { it('migrates legacy configuration when the extension loads', async () => { const piDir = join(process.env.HOME as string, '.pi'); writeFileSync(join(piDir, 'grok-cli-imagine.json'), JSON.stringify({ enabled: false })); - await setupExtension(); + const extension = await setupExtension(); + await extension.emit( + 'session_start', + { type: 'session_start', reason: 'startup' }, + sessionContext('session-a'), + ); expect(existsSync(join(piDir, 'grok-cli', 'config.json'))).toBe(true); expect(existsSync(join(piDir, 'grok-cli-imagine.json'))).toBe(false); }); + it('imports the saved OAuth login before migrating a standalone Imagine setting', async () => { + const piDir = join(process.env.HOME as string, '.pi'); + writeFileSync(join(piDir, 'grok-cli-imagine.json'), JSON.stringify({ enabled: false })); + writePiCredentials({ 'grok-cli': oauthCredential('released-account') }); + const extension = await setupExtension(); + + await extension.emit( + 'session_start', + { type: 'session_start', reason: 'startup' }, + sessionContext('session-a'), + ); + + expect((await getAccountVault()).accounts[0].credential?.access).toBe('released-account'); + expect(JSON.parse(readFileSync(join(piDir, 'grok-cli', 'config.json'), 'utf8'))).toEqual({ + version: 3, + imagine: { enabled: false }, + }); + }); + it('reports a migration failure only once at session start', async () => { const piDir = join(process.env.HOME as string, '.pi'); writeFileSync(join(piDir, 'grok-cli-imagine.json'), '{ nope'); @@ -867,243 +1187,53 @@ describe('Grok CLI tool scoping', () => { ); }); - it('registers the Grok/Cursor-native tool shims', async () => { + it('keeps a failed account migration blocking after its warning is shown', async () => { + mkdirSync(join(process.env.HOME as string, '.pi', 'grok-cli')); + writeFileSync( + join(process.env.HOME as string, '.pi', 'grok-cli', 'config.json'), + JSON.stringify({ version: 99 }), + ); const extension = await setupExtension(); + const context = contextForModel('grok-cli', 'grok-build'); - expect([...extension.tools.keys()].sort()).toEqual( - [...GROK_SHIM_TOOL_NAMES, 'WebSearch', 'image_gen'].sort(), + await extension.handlers.get('session_start')?.( + { type: 'session_start', reason: 'startup' }, + context, ); - }); - - it('does not register WebSearch when pi-web-access is not installed', async () => { - const extension = await setupExtension(['read', 'bash'], false); - expect([...extension.tools.keys()].sort()).toEqual( - [...GROK_SHIM_TOOL_NAMES, 'image_gen'].sort(), - ); - expect(extension.tools.has('WebSearch')).toBe(false); + await expect( + extension.providers.get('grok-cli')?.oauth?.login({} as OAuthLoginCallbacks), + ).rejects.toThrow('Could not migrate Grok CLI accounts'); + expect(mockOauthLogin).not.toHaveBeenCalled(); }); - it('uses only enabled compatibility capabilities for exact legacy models', async () => { - const extension = await setupExtension(['read', 'custom_tool', 'web_search']); - - await extension.handlers.get('model_select')?.( - { model: { provider: 'grok-cli', id: 'grok-build' } }, - contextForModel('grok-cli', 'grok-build'), - ); + it('registers only the image generation tool', async () => { + const extension = await setupExtension(); - const next = extension.setActiveTools.mock.calls.at(-1)?.[0] as string[]; - expect(next).not.toContain('web_search'); - expect(next).toEqual(['Read', 'custom_tool', 'WebSearch', 'image_gen']); - expect(next).not.toEqual(expect.arrayContaining(['Write', 'Delete', 'Shell'])); + expect([...extension.tools.keys()]).toEqual(['image_gen']); }); - it('queues model selection persistence without blocking the event handler', async () => { - saveTestAccounts('grok-cli'); + it('does not change the active account when a Grok model is selected', async () => { + await setAccount1Credential('one'); const extension = await setupExtension(); const result = extension.handlers.get('model_select')?.( - { model: { provider: 'grok-cli-2', id: 'grok-build' } }, - contextForModel('grok-cli-2', 'grok-build'), + { model: { provider: 'grok-cli', id: 'grok-build' } }, + contextForModel('grok-cli', 'grok-build'), ); expect(result).toBeUndefined(); - await vi.waitFor(() => { - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli-2'); - }); + expect((await getAccountVault()).activeAccountId).toBe('account-1'); }); - it('translates compatibility names back to native tools for non-Grok models', async () => { - const extension = await setupExtension(['read', 'Grep', 'custom_tool', 'Shell']); + it('keeps native coding tools unchanged for Grok models', async () => { + const extension = await setupExtension(['read', 'write', 'edit', 'bash']); await extension.handlers.get('model_select')?.( - { model: { provider: 'openai', id: 'gpt-4' } }, - contextForModel('openai'), - ); - - expect(extension.setActiveTools).toHaveBeenLastCalledWith([ - 'read', - 'grep', - 'custom_tool', - 'bash', - 'image_gen', - ]); - }); - - it('syncs tool scope before each agent turn from the current context model', async () => { - const extension = await setupExtension(['read']); - - await extension.handlers.get('before_agent_start')?.( - {}, - contextForModel('grok-cli', 'grok-build'), - ); - - expect(extension.setActiveTools).toHaveBeenLastCalledWith(['Read', 'image_gen']); - }); - - it('does not update active tools when the selection is already correct', async () => { - const extension = await setupExtension(['Read', 'image_gen']); - - await extension.handlers.get('before_agent_start')?.( - {}, - contextForModel('grok-cli', 'grok-build'), - ); - - expect(extension.setActiveTools).not.toHaveBeenCalled(); - }); - - it('keeps modern and unknown Grok models on native tools', async () => { - for (const id of ['grok-4.5', 'future-model']) { - const extension = await setupExtension(['read', 'bash']); - - await extension.handlers.get('model_select')?.( - { model: { provider: 'grok-cli', id } }, - contextForModel('grok-cli', id), - ); - - expect(extension.getActiveTools()).toEqual(['read', 'bash', 'image_gen']); - } - }); - - it('captures Delete on first session start and restores native names on shutdown', async () => { - const extension = await setupExtension(['read', 'Delete']); - - await extension.handlers.get('session_start')?.( - { type: 'session_start', reason: 'startup' }, - contextForModel('grok-cli', 'grok-build'), - ); - expect(extension.getActiveTools()).toEqual(['Read', 'Delete', 'image_gen']); - - extension.replaceActiveTools(['Read', 'custom', 'image_gen']); - await extension.handlers.get('session_shutdown')?.( - { type: 'session_shutdown', reason: 'reload' }, - contextForModel('grok-cli', 'grok-build'), - ); - expect(extension.getActiveTools()).toEqual(['read', 'custom', 'image_gen']); - }); - - it('reconciles live tool changes before the next legacy prompt', async () => { - const extension = await setupExtension(['read']); - await extension.handlers.get('session_start')?.( - { type: 'session_start', reason: 'startup' }, - contextForModel('grok-cli', 'grok-build'), - ); - extension.replaceActiveTools(['Read', 'write', 'custom', 'image_gen']); - - await extension.handlers.get('before_agent_start')?.( - {}, - contextForModel('grok-cli', 'grok-build'), - ); - - expect(extension.getActiveTools()).toEqual(['Read', 'Write', 'custom', 'image_gen']); - }); - - it('leaves native web_search available when the optional adapter is unavailable', async () => { - const extension = await setupExtension(['web_search'], false); - - await extension.handlers.get('session_start')?.( - { type: 'session_start', reason: 'startup' }, + { model: { provider: 'grok-cli', id: 'grok-build' } }, contextForModel('grok-cli', 'grok-build'), ); - expect(extension.getActiveTools()).toEqual(['web_search', 'image_gen']); - }); -}); - -describe('Grok CLI tool rendering', () => { - it('adds renderers to every Grok tool shim', async () => { - const extension = await setupExtension(); - - for (const name of [...GROK_SHIM_TOOL_NAMES, 'WebSearch', 'image_gen']) { - expect(extension.tools.get(name)?.renderCall).toBeTypeOf('function'); - expect(extension.tools.get(name)?.renderResult).toBeTypeOf('function'); - } - }); - - it('delegates collapsed and expanded search output rendering', async () => { - const extension = await setupExtension(); - const grep = extension.tools.get('Grep'); - const result = { - content: [{ type: 'text', text: 'src/a.ts:1:match\nsrc/b.ts:2:match' }], - details: undefined, - }; - - const collapsed = renderText( - grep?.renderResult?.( - result, - { expanded: false, isPartial: false }, - theme, - renderContext({ pattern: 'match' }), - ) as Renderable, - ); - const expanded = renderText( - grep?.renderResult?.( - result, - { expanded: true, isPartial: false }, - theme, - renderContext({ pattern: 'match' }), - ) as Renderable, - ); - - expect(collapsed).toContain('src/a.ts:1:match'); - expect(collapsed).toContain('src/b.ts:2:match'); - expect(expanded).toContain('src/a.ts:1:match'); - }); - - it('uses native results for delegated tools and retained summaries for custom tools', async () => { - const extension = await setupExtension(); - - expect( - renderText( - extension.tools.get('Write')?.renderResult?.( - { - content: [{ type: 'text', text: 'long write output' }], - details: undefined, - }, - { expanded: false, isPartial: false }, - theme, - renderContext({ path: 'notes.txt', content: 'content' }), - ) as Renderable, - ), - ).toBe(''); - expect( - renderText( - extension.tools.get('StrReplace')?.renderResult?.( - { - content: [{ type: 'text', text: 'long replace output' }], - details: { replacements: 3 }, - }, - { expanded: false, isPartial: false }, - theme, - renderContext({ path: 'notes.txt', old_str: 'old', new_str: 'new' }), - ) as Renderable, - ), - ).toBe('3 replacement(s)'); - expect( - renderText( - extension.tools.get('Delete')?.renderResult?.( - { - content: [{ type: 'text', text: 'long delete output' }], - details: { deleted: true }, - }, - { expanded: false, isPartial: false }, - theme, - renderContext({ path: 'notes.txt' }), - ) as Renderable, - ), - ).toBe('Deleted'); - expect( - renderText( - extension.tools.get('Shell')?.renderResult?.( - { - content: [{ type: 'text', text: 'long shell output' }], - details: undefined, - }, - { expanded: false, isPartial: false }, - theme, - renderContext({ command: 'printf output' }), - ) as Renderable, - ), - ).toContain('long shell output'); + expect(extension.getActiveTools()).toEqual(['read', 'write', 'edit', 'bash', 'image_gen']); }); }); diff --git a/tests/provider/rotation.test.ts b/tests/provider/rotation.test.ts index 3a3304e..ac0bfa4 100644 --- a/tests/provider/rotation.test.ts +++ b/tests/provider/rotation.test.ts @@ -1,470 +1,225 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { afterEach, describe, expect, it, vi } from 'vitest'; -import { DEFAULT_CONFIG, loadConfig, saveConfig } from '../../src/config.js'; -import type { BillingUsage } from '../../src/provider/billing.js'; -import { saveQuotaUsage } from '../../src/provider/quotaCache.js'; +import { readFileSync } from 'node:fs'; +import type { AssistantMessage } from '@earendil-works/pi-ai'; +import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { getAccountVault, mutateAccountVault } from '../../src/provider/accountVault.js'; +import { rememberRequestAccount } from '../../src/provider/requestOwnership.js'; import { EXHAUSTED_BALANCE_ERROR, ROTATION_CONTINUATION, registerExhaustionRotation, } from '../../src/provider/rotation.js'; -import { useTempHome } from '../vision/helpers.js'; +import { createSessionAccountSelection } from '../../src/provider/sessionAccountSelection.js'; +import { acquireFileLock, getAccountVaultPath, writeFileAtomic } from '../../src/storage.js'; +import { useEnvironmentToken, useTempHome } from '../stateTestHelpers.js'; const setupHome = useTempHome(); -const THREE_ACCOUNTS = [ - { provider: 'grok-cli', label: 'Personal' }, - { provider: 'grok-cli-2', label: 'Work' }, - { provider: 'grok-cli-3', label: 'Client' }, -]; -const FOUR_ACCOUNTS = [...THREE_ACCOUNTS, { provider: 'grok-cli-4', label: 'Reserve' }]; -const NOW = Date.parse('2026-07-16T12:00:00.000Z'); -const originalFetch = globalThis.fetch; +const setEnvironmentToken = useEnvironmentToken(); -afterEach(() => { - vi.useRealTimers(); - globalThis.fetch = originalFetch; -}); - -function usage(monthlyRemaining: number, weeklyRemaining?: number): BillingUsage { +function exhaustedMessage(): AssistantMessage { return { - monthly: { - monthlyLimit: 100, - used: 100 - monthlyRemaining, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', + role: 'assistant', + api: 'openai-responses', + provider: 'grok-cli', + model: 'grok-build', + content: [], + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, - ...(weeklyRemaining === undefined - ? {} - : { - weekly: { - creditUsagePercent: 100 - weeklyRemaining, - billingPeriodEnd: '2026-07-20T00:00:00.000Z', - }, - }), + stopReason: 'error', + errorMessage: EXHAUSTED_BALANCE_ERROR, + timestamp: Date.now(), }; } -function setup( - options: { - accounts?: { provider: string; label: string }[]; - auth?: string[]; - current?: { provider: string; id: string }; - missingModels?: string[]; - setModel?: boolean[]; - } = {}, -) { - setupHome(); - const accounts = options.accounts ?? [ - { provider: 'grok-cli', label: 'Personal' }, - { provider: 'grok-cli-2', label: 'Work' }, - ]; - saveConfig({ - ...DEFAULT_CONFIG, - accounts: { - nextAccountNumber: accounts.length + 1, - selectedProvider: options.current?.provider ?? 'grok-cli', - items: accounts, - }, - }); - const handlers = new Map unknown)[]>(); - const authenticated = new Set(options.auth ?? accounts.map((account) => account.provider)); - const model = options.current ?? { provider: 'grok-cli', id: 'grok-build' }; - const context = { - model: { ...model }, - modelRegistry: { - getProviderAuthStatus: (provider: string) => ({ configured: authenticated.has(provider) }), - find: (provider: string, id: string) => - options.missingModels?.includes(`${provider}/${id}`) ? undefined : { provider, id }, - }, - ui: { notify: vi.fn() }, - }; - const setModelResults = [...(options.setModel ?? [])]; - const setModel = vi.fn(async (nextModel: { provider: string; id: string }) => { - const result = setModelResults.shift() ?? true; - if (result) context.model = nextModel; - return result; - }); +function extension() { + const handlers = new Map unknown)[]>(); const sendUserMessage = vi.fn(); + const setModel = vi.fn(); + const appendEntry = vi.fn(); const pi = { - on(event: string, handler: (event: unknown, ctx: unknown) => unknown) { + on(event: string, handler: (event: unknown, ctx: ExtensionContext) => unknown) { handlers.set(event, [...(handlers.get(event) ?? []), handler]); }, + appendEntry, sendUserMessage, setModel, } as unknown as ExtensionAPI; - const rotation = registerExhaustionRotation(pi); - + const selection = createSessionAccountSelection(pi); + const rotation = registerExhaustionRotation(pi, selection); + const notify = vi.fn(); + const ctx = { + model: { provider: 'grok-cli', id: 'grok-build' }, + sessionManager: { + getSessionId: () => 'session-a', + getBranch: () => [], + }, + ui: { notify }, + } as unknown as ExtensionContext; return { - context, - handlers, - notify: context.ui.notify, + rotation, + notify, + appendEntry, + selection, sendUserMessage, setModel, - rotation, - async emit(event: string, data: unknown = { type: event }) { - for (const handler of handlers.get(event) ?? []) await handler(data, context); + async emit(event: string, value: unknown = {}) { + for (const handler of handlers.get(event) ?? []) await handler(value, ctx); }, }; } -function assistant( - provider: string, - errorMessage = EXHAUSTED_BALANCE_ERROR, - model = 'grok-build', - stopReason = 'error', -) { - return { - type: 'message_end', - message: { - role: 'assistant', - provider, - model, - stopReason, - errorMessage, - }, - }; -} - -async function settleExhaustion(extension: ReturnType, provider: string) { - await extension.emit('message_end', assistant(provider)); - await extension.emit('agent_settled'); -} - -async function emitExtensionContinuation(extension: ReturnType) { - await extension.emit('input', { - type: 'input', - source: 'extension', - text: ROTATION_CONTINUATION, +async function addLoggedInAccounts() { + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'one', + refresh: 'one-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts.push( + { + id: 'account-2', + slot: 2, + label: 'Work', + credential: { + access: 'two', + refresh: 'two-refresh', + expires: Date.now() + 300_000, + }, + revision: 1, + }, + { + id: 'account-3', + slot: 3, + label: 'Client', + credential: { + access: 'three', + refresh: 'three-refresh', + expires: Date.now() + 300_000, + }, + revision: 1, + }, + ); + vault.nextSlot = 4; + vault.activeAccountId = 'account-1'; }); } -function switchedProviders(extension: ReturnType) { - return extension.setModel.mock.calls.map(([model]) => model.provider); +async function settleExhaustion(test: ReturnType, accountId: string) { + const message = exhaustedMessage(); + rememberRequestAccount(message, accountId); + await test.emit('message_end', { message }); + await test.emit('agent_settled'); } -async function setupAfterSuccessfulContinuation() { - vi.useFakeTimers(); - vi.setSystemTime(NOW); - const extension = setup(); - await settleExhaustion(extension, 'grok-cli'); - await emitExtensionContinuation(extension); - await extension.emit('message_end', assistant('grok-cli-2', '', 'grok-build', 'stop')); - await extension.emit('agent_settled'); - return extension; -} +beforeEach(() => { + setupHome(); + setEnvironmentToken(); + vi.useRealTimers(); +}); describe('Grok CLI exhaustion rotation', () => { - it('switches only after agent_settled, preserves the model, and continues once', async () => { - const extension = setup({ current: { provider: 'grok-cli', id: 'grok-composer-2.5-fast' } }); - - await extension.emit( - 'message_end', - assistant('grok-cli', EXHAUSTED_BALANCE_ERROR, 'grok-composer-2.5-fast'), - ); - - expect(extension.setModel).not.toHaveBeenCalled(); - expect(extension.sendUserMessage).not.toHaveBeenCalled(); - - await extension.emit('agent_settled'); - - expect(extension.setModel).toHaveBeenCalledOnce(); - expect(extension.setModel).toHaveBeenCalledWith({ - provider: 'grok-cli-2', - id: 'grok-composer-2.5-fast', + it('selects another logged-in account without changing the model provider', async () => { + await addLoggedInAccounts(); + const test = extension(); + await settleExhaustion(test, 'account-1'); + + expect(test.selection.accountId('session-a')).toBe('account-2'); + expect(test.appendEntry).toHaveBeenCalledWith('grok-cli-active-account-v1', { + accountId: 'account-2', }); - expect(loadConfig().config.accounts.selectedProvider).toBe('grok-cli-2'); - expect(extension.notify).toHaveBeenCalledWith( - 'Grok CLI: “Personal” exhausted; switched to “Work” and continuing.', + expect((await getAccountVault()).activeAccountId).toBe('account-1'); + expect(test.setModel).not.toHaveBeenCalled(); + expect(test.sendUserMessage).toHaveBeenCalledWith(ROTATION_CONTINUATION); + expect(test.notify).toHaveBeenCalledWith( + 'Grok CLI: “Account 1” exhausted; switched to “Work” and continuing.', 'info', ); - expect(extension.sendUserMessage).toHaveBeenCalledOnce(); - expect(extension.sendUserMessage).toHaveBeenCalledWith(ROTATION_CONTINUATION); - }); - - it('preserves configuration changes made while switching models', async () => { - const extension = setup(); - let releaseSwitch = () => {}; - extension.setModel.mockImplementation(async (model) => { - await new Promise((resolve) => { - releaseSwitch = resolve; - }); - extension.context.model = model; - return true; - }); - await extension.emit('message_end', assistant('grok-cli')); - const settling = extension.emit('agent_settled'); - await vi.waitFor(() => expect(extension.setModel).toHaveBeenCalledOnce()); - const concurrent = loadConfig().config; - concurrent.accounts.items.push({ provider: 'grok-cli-3', label: 'Added concurrently' }); - concurrent.vision.enabled = false; - saveConfig(concurrent); - - releaseSwitch(); - await settling; - - expect(loadConfig().config).toMatchObject({ - accounts: { - selectedProvider: 'grok-cli-2', - items: expect.arrayContaining([{ provider: 'grok-cli-3', label: 'Added concurrently' }]), - }, - vision: { enabled: false }, - }); - }); - - it.each([ - ['near match', 'OpenAI API error (402): 402 "Grok Build usage balance exhausted".'], - ['other 402', 'OpenAI API error (402): payment required'], - ['401', 'OpenAI API error (401): unauthorized'], - ['429', 'OpenAI API error (429): rate limited'], - ])('ignores %s errors', async (_name, errorMessage) => { - const extension = setup(); - - await extension.emit('message_end', assistant('grok-cli', errorMessage)); - await extension.emit('agent_settled'); - - expect(extension.setModel).not.toHaveBeenCalled(); - }); - - it('ignores non-Grok providers and non-error assistant messages', async () => { - const extension = setup(); - - await extension.emit('message_end', assistant('openai', EXHAUSTED_BALANCE_ERROR)); - await extension.emit( - 'message_end', - assistant('grok-cli', EXHAUSTED_BALANCE_ERROR, 'grok-build', 'stop'), - ); - await extension.emit('agent_settled'); - - expect(extension.setModel).not.toHaveBeenCalled(); - }); - - it('uses circular account order and skips login-required accounts', async () => { - const extension = setup({ - accounts: THREE_ACCOUNTS, - auth: ['grok-cli-2', 'grok-cli-3'], - current: { provider: 'grok-cli-3', id: 'grok-composer-2.5-fast' }, - }); - - await extension.emit( - 'message_end', - assistant('grok-cli-3', EXHAUSTED_BALANCE_ERROR, 'grok-composer-2.5-fast'), - ); - await extension.emit('agent_settled'); - - expect(extension.setModel).toHaveBeenCalledWith({ - provider: 'grok-cli-2', - id: 'grok-composer-2.5-fast', - }); }); - it('skips failed setModel candidates and falls back to grok-build when needed', async () => { - const extension = setup({ - accounts: THREE_ACCOUNTS, - current: { provider: 'grok-cli', id: 'grok-composer-2.5-fast' }, - missingModels: ['grok-cli-3/grok-composer-2.5-fast'], - setModel: [false, true], + it('attributes an in-flight failure to its captured account after selection changes', async () => { + await addLoggedInAccounts(); + await mutateAccountVault((vault) => { + vault.activeAccountId = 'account-3'; }); + const test = extension(); + await settleExhaustion(test, 'account-1'); - await extension.emit( - 'message_end', - assistant('grok-cli', EXHAUSTED_BALANCE_ERROR, 'grok-composer-2.5-fast'), - ); - await extension.emit('agent_settled'); - - expect(extension.setModel.mock.calls).toEqual([ - [{ provider: 'grok-cli-2', id: 'grok-composer-2.5-fast' }], - [{ provider: 'grok-cli-3', id: 'grok-build' }], - ]); - expect(extension.sendUserMessage).toHaveBeenCalledOnce(); + expect(test.selection.accountId('session-a')).toBe('account-2'); + expect((await getAccountVault()).activeAccountId).toBe('account-3'); }); - it('preserves attempted accounts across extension continuations and stops without wrapping', async () => { - const extension = setup({ - accounts: THREE_ACCOUNTS, - }); - - await settleExhaustion(extension, 'grok-cli'); - await emitExtensionContinuation(extension); - await settleExhaustion(extension, 'grok-cli-2'); - await emitExtensionContinuation(extension); - await settleExhaustion(extension, 'grok-cli-3'); - await extension.emit('agent_settled'); + it('keeps attempted accounts across extension continuations', async () => { + await addLoggedInAccounts(); + const test = extension(); + await settleExhaustion(test, 'account-1'); + await settleExhaustion(test, 'account-2'); - expect(switchedProviders(extension)).toEqual(['grok-cli-2', 'grok-cli-3']); - expect(extension.sendUserMessage).toHaveBeenCalledTimes(2); - expect(extension.notify).toHaveBeenCalledWith( - 'Grok CLI: all logged-in accounts are exhausted.', - 'warning', - ); + expect(test.selection.accountId('session-a')).toBe('account-3'); + expect((await getAccountVault()).activeAccountId).toBe('account-1'); + expect(test.sendUserMessage).toHaveBeenCalledTimes(2); }); - it('keeps recently exhausted accounts unavailable across new real user input', async () => { - vi.useFakeTimers(); - vi.setSystemTime(NOW); - const extension = setup(); - - await settleExhaustion(extension, 'grok-cli'); - await extension.emit('input', { type: 'input', source: 'interactive', text: 'try again' }); - await settleExhaustion(extension, 'grok-cli-2'); + it('reports when all logged-in accounts are exhausted', async () => { + await addLoggedInAccounts(); + const test = extension(); + for (const id of ['account-1', 'account-2', 'account-3']) { + await settleExhaustion(test, id); + } - expect(switchedProviders(extension)).toEqual(['grok-cli-2']); - expect(extension.notify).toHaveBeenCalledWith( + expect(test.notify).toHaveBeenLastCalledWith( 'Grok CLI: all logged-in accounts are exhausted.', 'warning', ); }); - it('keeps recently exhausted accounts unavailable after a successful continuation', async () => { - const extension = await setupAfterSuccessfulContinuation(); - await extension.emit('message_end', assistant('grok-cli-2')); - await extension.emit('agent_settled'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-2']); - }); - - it('makes an account eligible again exactly five minutes after exhaustion', async () => { - const extension = await setupAfterSuccessfulContinuation(); - vi.setSystemTime(NOW + 5 * 60_000); - await settleExhaustion(extension, 'grok-cli-2'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-2', 'grok-cli']); - }); - - it('clears recent exhaustion when a new session starts', async () => { - const extension = await setupAfterSuccessfulContinuation(); - await extension.emit('session_start'); - await settleExhaustion(extension, 'grok-cli-2'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-2', 'grok-cli']); - }); - - it('allows successful login to clear one account’s recent exhaustion', async () => { - const extension = await setupAfterSuccessfulContinuation(); - extension.rotation.clearRecentExhaustion('grok-cli'); - await settleExhaustion(extension, 'grok-cli-2'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-2', 'grok-cli']); - }); - - it('orders fresh cached quota by the tightest remaining window without fetching', async () => { - vi.useFakeTimers(); - vi.setSystemTime(NOW); - const extension = setup({ accounts: THREE_ACCOUNTS }); - const fetchMock = vi.fn(); - globalThis.fetch = fetchMock; - await saveQuotaUsage('grok-cli-2', usage(90, 10), new Date(NOW).toISOString()); - await saveQuotaUsage('grok-cli-3', usage(50), new Date(NOW).toISOString()); - - await settleExhaustion(extension, 'grok-cli'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-3']); - expect(fetchMock).not.toHaveBeenCalled(); - }); - - it('continues through missing models and failed switches in quota-ranked order', async () => { - vi.useFakeTimers(); - vi.setSystemTime(NOW); - const extension = setup({ - accounts: FOUR_ACCOUNTS, - current: { provider: 'grok-cli', id: 'grok-composer-2.5-fast' }, - missingModels: ['grok-cli-3/grok-composer-2.5-fast', 'grok-cli-3/grok-build'], - setModel: [false, true], - }); - await saveQuotaUsage('grok-cli-2', usage(50), new Date(NOW).toISOString()); - await saveQuotaUsage('grok-cli-3', usage(90), new Date(NOW).toISOString()); - await saveQuotaUsage('grok-cli-4', usage(70), new Date(NOW).toISOString()); - - await extension.emit( - 'message_end', - assistant('grok-cli', EXHAUSTED_BALANCE_ERROR, 'grok-composer-2.5-fast'), - ); - await extension.emit('agent_settled'); - - expect(extension.setModel.mock.calls).toEqual([ - [{ provider: 'grok-cli-4', id: 'grok-composer-2.5-fast' }], - [{ provider: 'grok-cli-2', id: 'grok-composer-2.5-fast' }], - ]); - expect(extension.sendUserMessage).toHaveBeenCalledOnce(); - }); - - it.each([ - ['missing', undefined], - ['stale', new Date(NOW - 30 * 60_000).toISOString()], - ['invalid', new Date(NOW).toISOString()], - ])('keeps a %s quota candidate in its circular slot', async (_name, updatedAt) => { - vi.useFakeTimers(); - vi.setSystemTime(NOW); - const extension = setup({ accounts: FOUR_ACCOUNTS, setModel: [false, true] }); - if (updatedAt) { - await saveQuotaUsage( - 'grok-cli-2', - updatedAt === new Date(NOW).toISOString() - ? { - monthly: { - monthlyLimit: 0, - used: 0, - billingPeriodEnd: '2026-08-01T00:00:00.000Z', - }, - } - : usage(50), - updatedAt, - ); - } - await saveQuotaUsage('grok-cli-3', usage(20), new Date(NOW).toISOString()); - await saveQuotaUsage('grok-cli-4', usage(80), new Date(NOW).toISOString()); - - await settleExhaustion(extension, 'grok-cli'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-2', 'grok-cli-4']); - }); - - it('preserves circular order when fresh quota scores tie', async () => { - vi.useFakeTimers(); - vi.setSystemTime(NOW); - const extension = setup({ accounts: THREE_ACCOUNTS }); - await saveQuotaUsage('grok-cli-2', usage(60, 40), new Date(NOW).toISOString()); - await saveQuotaUsage('grok-cli-3', usage(40), new Date(NOW).toISOString()); - - await settleExhaustion(extension, 'grok-cli'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-2']); - }); - - it('excludes a recently exhausted account even when it has the best cached quota', async () => { - vi.useFakeTimers(); - vi.setSystemTime(NOW); - const extension = setup({ - accounts: THREE_ACCOUNTS, - current: { provider: 'grok-cli-2', id: 'grok-build' }, - }); - await saveQuotaUsage('grok-cli', usage(10), new Date(NOW).toISOString()); - await saveQuotaUsage('grok-cli-2', usage(95), new Date(NOW).toISOString()); - await saveQuotaUsage('grok-cli-3', usage(20), new Date(NOW).toISOString()); - - await settleExhaustion(extension, 'grok-cli-2'); - await emitExtensionContinuation(extension); - await extension.emit('message_end', assistant('grok-cli-3', '', 'grok-build', 'stop')); - await extension.emit('agent_settled'); - extension.context.model = { provider: 'grok-cli', id: 'grok-build' }; - await settleExhaustion(extension, 'grok-cli'); - - expect(switchedProviders(extension)).toEqual(['grok-cli-3', 'grok-cli-3']); - }); - - it('cancels a pending rotation after a manual model change', async () => { - const extension = setup(); - - await extension.emit('message_end', assistant('grok-cli')); - await extension.emit('model_select', { - type: 'model_select', - model: { provider: 'openai', id: 'gpt-5' }, - }); - await extension.emit('model_select', { - type: 'model_select', - model: { provider: 'grok-cli', id: 'grok-build' }, - }); - await extension.emit('agent_settled'); - - expect(extension.setModel).not.toHaveBeenCalled(); - expect(extension.sendUserMessage).not.toHaveBeenCalled(); + it('lets a successful login clear one recent exhaustion record', async () => { + await addLoggedInAccounts(); + const test = extension(); + await settleExhaustion(test, 'account-1'); + test.rotation.clearRecentExhaustion('account-1'); + await settleExhaustion(test, 'account-2'); + + expect(test.selection.accountId('session-a')).toBe('account-3'); + expect((await getAccountVault()).activeAccountId).toBe('account-1'); + }); + + it('does not rotate saved accounts when an environment token is active', async () => { + setEnvironmentToken('environment-token'); + await addLoggedInAccounts(); + const test = extension(); + + await settleExhaustion(test, 'account-1'); + + expect((await getAccountVault()).activeAccountId).toBe('account-1'); + expect(test.sendUserMessage).not.toHaveBeenCalled(); + expect(test.notify).not.toHaveBeenCalled(); + }); + + it('selects another candidate when the first candidate is removed before commit', async () => { + await addLoggedInAccounts(); + const test = extension(); + const message = exhaustedMessage(); + rememberRequestAccount(message, 'account-1'); + await test.emit('message_end', { message }); + const release = await acquireFileLock(getAccountVaultPath()); + const pending = test.emit('agent_settled'); + await new Promise((resolve) => setTimeout(resolve, 20)); + const vault = JSON.parse(readFileSync(getAccountVaultPath(), 'utf8')) as Awaited< + ReturnType + >; + vault.accounts = vault.accounts.filter((account) => account.id !== 'account-2'); + writeFileAtomic(getAccountVaultPath(), `${JSON.stringify(vault, null, 2)}\n`, 0o600); + await release(); + + await expect(pending).resolves.toBeUndefined(); + expect(test.selection.accountId('session-a')).toBe('account-3'); + expect((await getAccountVault()).activeAccountId).toBe('account-1'); + expect(test.sendUserMessage).toHaveBeenCalledWith(ROTATION_CONTINUATION); }); }); diff --git a/tests/provider/sessionAccountSelection.test.ts b/tests/provider/sessionAccountSelection.test.ts new file mode 100644 index 0000000..ab3c56c --- /dev/null +++ b/tests/provider/sessionAccountSelection.test.ts @@ -0,0 +1,97 @@ +import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { mutateAccountVault } from '../../src/provider/accountVault.js'; +import { + createSessionAccountSelection, + SESSION_ACCOUNT_ENTRY, +} from '../../src/provider/sessionAccountSelection.js'; +import { useEnvironmentToken, useTempHome } from '../stateTestHelpers.js'; + +const setupHome = useTempHome(); +const setEnvironmentToken = useEnvironmentToken(); + +function context(sessionId: string, accountIds: string[] = []) { + return { + sessionManager: { + getSessionId: () => sessionId, + getBranch: () => + accountIds.map((accountId, index) => ({ + type: 'custom', + id: `entry-${index}`, + parentId: index ? `entry-${index - 1}` : null, + timestamp: new Date(index).toISOString(), + customType: SESSION_ACCOUNT_ENTRY, + data: { accountId }, + })), + }, + } as unknown as ExtensionContext; +} + +beforeEach(async () => { + setupHome(); + setEnvironmentToken(undefined); + await mutateAccountVault((vault) => { + vault.accounts[0].credential = { + access: 'one', + refresh: 'one-refresh', + expires: Date.now() + 300_000, + }; + vault.accounts.push({ + id: 'work-id', + slot: 2, + label: 'Work', + credential: { + access: 'two', + refresh: 'two-refresh', + expires: Date.now() + 300_000, + }, + revision: 1, + }); + vault.nextSlot = 3; + vault.activeAccountId = 'account-1'; + }); +}); + +describe('Pi session account selection', () => { + it('keeps account selections independent for two sessions', () => { + const appendEntry = vi.fn(); + const selection = createSessionAccountSelection({ appendEntry } as unknown as ExtensionAPI); + const first = context('session-a'); + const second = context('session-b'); + + selection.restore(first); + selection.restore(second); + selection.select(first, 'work-id'); + + expect(selection.accountId('session-a')).toBe('work-id'); + expect(selection.accountId('session-b')).toBe('account-1'); + expect(appendEntry).toHaveBeenCalledWith(SESSION_ACCOUNT_ENTRY, { accountId: 'work-id' }); + }); + + it('restores the last account selection on the active session branch', () => { + const selection = createSessionAccountSelection({ + appendEntry: vi.fn(), + } as unknown as ExtensionAPI); + + expect(selection.restore(context('session-a', ['account-1', 'work-id']))).toBe('work-id'); + expect(selection.accountId('session-a')).toBe('work-id'); + }); + + it('restores an account when the callback is detached', () => { + const selection = createSessionAccountSelection({ + appendEntry: vi.fn(), + } as unknown as ExtensionAPI); + const restore = selection.restore; + + expect(restore(context('session-a', ['work-id']))).toBe('work-id'); + }); + + it('returns another logged-in account when a stored selection is no longer valid', () => { + const selection = createSessionAccountSelection({ + appendEntry: vi.fn(), + } as unknown as ExtensionAPI); + + expect(selection.restore(context('session-a', ['removed-id']))).toBe('account-1'); + expect(selection.accountId('session-a')).toBe('account-1'); + }); +}); diff --git a/tests/provider/toolScope.integration.test.ts b/tests/provider/toolScope.integration.test.ts deleted file mode 100644 index 4c7086a..0000000 --- a/tests/provider/toolScope.integration.test.ts +++ /dev/null @@ -1,361 +0,0 @@ -import { mkdirSync } from 'node:fs'; -import { join } from 'node:path'; -import { fauxAssistantMessage, fauxProvider, InMemoryCredentialStore } from '@earendil-works/pi-ai'; -import { - AgentSessionRuntime, - type CreateAgentSessionRuntimeFactory, - createAgentSession, - createAgentSessionFromServices, - createAgentSessionServices, - DefaultResourceLoader, - type ExtensionFactory, - ModelRegistry, - ModelRuntime, - SessionManager, - SettingsManager, -} from '@earendil-works/pi-coding-agent'; -import { describe, expect, it } from 'vitest'; -import { handoffGrokTools, restoreGrokTools, syncGrokTools } from '../../src/provider/toolScope.js'; -import { registerGrokTools } from '../../src/tools/register.js'; -import { tempDir } from '../tools/toolTestHelpers.js'; - -const toolScopeExtension: ExtensionFactory = (pi) => { - const { webSearchRegistered } = registerGrokTools(pi); - - pi.on('session_start', (event, ctx) => { - if (event.reason === 'new' || event.reason === 'resume' || event.reason === 'fork') { - restoreGrokTools(pi, ctx.sessionManager.getSessionFile()); - } - syncGrokTools(pi, ctx.model, { - captureDelete: true, - imagineEnabled: false, - webSearchRegistered, - }); - }); - - pi.on('model_select', (event) => { - syncGrokTools(pi, event.model, { - imagineEnabled: false, - webSearchRegistered, - }); - }); - - pi.on('before_agent_start', (_event, ctx) => { - syncGrokTools(pi, ctx.model, { - imagineEnabled: false, - webSearchRegistered, - }); - }); - - pi.on('session_shutdown', (event) => { - syncGrokTools(pi, undefined, { - imagineEnabled: false, - webSearchRegistered, - }); - if (event.reason === 'new' || event.reason === 'resume' || event.reason === 'fork') { - handoffGrokTools(pi, event.targetSessionFile); - } - }); -}; - -async function toolScopeResources(cwd: string, agentDir: string) { - mkdirSync(agentDir); - const settingsManager = SettingsManager.create(cwd, agentDir); - const resourceLoader = new DefaultResourceLoader({ - cwd, - agentDir, - settingsManager, - extensionFactories: [{ name: 'tool-scope', factory: toolScopeExtension }], - noExtensions: true, - noSkills: true, - noPromptTemplates: true, - noThemes: true, - noContextFiles: true, - }); - await resourceLoader.reload(); - return { resourceLoader, settingsManager }; -} - -async function scopedSession(options: { - excludeTools?: string[]; - modelId?: string; - noTools?: 'all'; - tools?: string[]; -}) { - const cwd = tempDir('pi-grok-cli-tool-scope-integration-'); - const agentDir = join(cwd, '.pi-agent'); - const resources = await toolScopeResources(cwd, agentDir); - const model = { - ...fauxProvider().getModel(), - provider: 'grok-cli', - id: options.modelId ?? 'grok-build', - }; - const { session } = await createAgentSession({ - cwd, - agentDir, - model, - settingsManager: resources.settingsManager, - resourceLoader: resources.resourceLoader, - sessionManager: SessionManager.inMemory(cwd), - tools: options.tools, - excludeTools: options.excludeTools, - noTools: options.noTools, - }); - await session.bindExtensions({ shutdownHandler: () => {} }); - return session; -} - -async function scopedRuntime(persisted = false) { - const cwd = tempDir('pi-grok-cli-tool-scope-runtime-'); - const agentDir = join(cwd, '.pi-agent'); - const sessionDir = join(cwd, 'sessions'); - mkdirSync(agentDir); - const model = { - ...fauxProvider().getModel(), - provider: 'grok-cli', - id: 'grok-build', - }; - const createRuntime: CreateAgentSessionRuntimeFactory = async (options) => { - const services = await createAgentSessionServices({ - cwd: options.cwd, - agentDir: options.agentDir, - resourceLoaderOptions: { - extensionFactories: [{ name: 'tool-scope', factory: toolScopeExtension }], - noExtensions: true, - noSkills: true, - noPromptTemplates: true, - noThemes: true, - noContextFiles: true, - }, - }); - const created = await createAgentSessionFromServices({ - services, - sessionManager: options.sessionManager, - sessionStartEvent: options.sessionStartEvent, - model, - }); - return { ...created, services, diagnostics: services.diagnostics }; - }; - const created = await createRuntime({ - cwd, - agentDir, - sessionManager: persisted - ? SessionManager.create(cwd, sessionDir) - : SessionManager.inMemory(cwd), - }); - const runtime = new AgentSessionRuntime( - created.session, - created.services, - createRuntime, - created.diagnostics, - created.modelFallbackMessage, - ); - const bindExtensions = async (session: typeof created.session) => { - await session.bindExtensions({ shutdownHandler: () => {} }); - }; - await bindExtensions(runtime.session); - runtime.setRebindSession(bindExtensions); - return { cwd, runtime, sessionDir }; -} - -async function modelFacingTools(modelId: string) { - const cwd = tempDir('pi-grok-cli-tool-scope-prompt-'); - const agentDir = join(cwd, '.pi-agent'); - const faux = fauxProvider({ - provider: 'grok-cli', - models: [{ id: modelId }], - }); - const modelRuntime = await ModelRuntime.create({ - credentials: new InMemoryCredentialStore(), - modelsPath: null, - allowModelNetwork: false, - }); - const modelRegistry = new ModelRegistry(modelRuntime); - const model = faux.getModel(); - modelRuntime.registerProvider('grok-cli', { - api: faux.api, - apiKey: 'local-test-key', - baseUrl: model.baseUrl, - models: [ - { - id: model.id, - name: model.name, - reasoning: model.reasoning, - input: model.input, - cost: model.cost, - contextWindow: model.contextWindow, - maxTokens: model.maxTokens, - }, - ], - streamSimple: faux.provider.streamSimple.bind(faux.provider), - }); - const resources = await toolScopeResources(cwd, agentDir); - let names: string[] = []; - faux.setResponses([ - (context) => { - names = context.tools?.map((tool) => tool.name) ?? []; - return fauxAssistantMessage('done'); - }, - ]); - const { session } = await createAgentSession({ - cwd, - agentDir, - modelRuntime, - model: modelRegistry.find('grok-cli', modelId), - resourceLoader: resources.resourceLoader, - sessionManager: SessionManager.inMemory(cwd), - settingsManager: resources.settingsManager, - }); - try { - await session.bindExtensions({ shutdownHandler: () => {} }); - await session.prompt('Report available tools.'); - return names; - } finally { - session.dispose(); - modelRuntime.unregisterProvider('grok-cli'); - } -} - -describe('tool scope through Pi registry filtering', () => { - it.each([ - 'grok-build', - 'grok-composer-2.5-fast', - ])('sends the complete compatibility vocabulary to %s on a real agent turn', async (modelId) => { - expect(await modelFacingTools(modelId)).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Write', - 'Grep', - 'Glob', - 'LS', - 'Delete', - ]); - }); - - it('sends only native tools to a modern model on a real agent turn', async () => { - expect(await modelFacingTools('grok-4.5')).toEqual(['read', 'bash', 'edit', 'write']); - }); - - it('activates the complete legacy tool set on a normal startup', async () => { - const session = await scopedSession({}); - - expect(session.getActiveToolNames()).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Write', - 'Grep', - 'Glob', - 'LS', - 'Delete', - ]); - session.dispose(); - }); - - it('does not restore native capabilities excluded before extension startup', async () => { - const session = await scopedSession({ - excludeTools: ['write', 'edit', 'bash', 'Delete'], - }); - - expect(session.getActiveToolNames()).toEqual(['Read', 'Grep', 'Glob', 'LS']); - session.dispose(); - }); - - it('honors exact uppercase exclusions for legacy discovery defaults', async () => { - const session = await scopedSession({ - excludeTools: ['Grep', 'Glob', 'LS'], - }); - - expect(session.getActiveToolNames()).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Write', - 'Delete', - ]); - session.dispose(); - }); - - it('keeps all tools disabled with --no-tools', async () => { - const session = await scopedSession({ noTools: 'all' }); - - expect(session.getActiveToolNames()).toEqual([]); - session.dispose(); - }); - - it('keeps a native allowlisted capability without enabling filtered-out shims', async () => { - const session = await scopedSession({ tools: ['read'] }); - - expect(session.getActiveToolNames()).toEqual(['read']); - session.dispose(); - }); - - it('does not turn auto-enabled shims into native tools for a modern model', async () => { - const session = await scopedSession({ modelId: 'grok-4.5' }); - - expect(session.getActiveToolNames()).toEqual(['read', 'bash', 'edit', 'write']); - session.dispose(); - }); - - it('does not resurrect a disabled Delete tool across reload', async () => { - const session = await scopedSession({}); - session.setActiveToolsByName( - session.getActiveToolNames().filter((toolName) => toolName !== 'Delete'), - ); - - await session.reload(); - - expect(session.getActiveToolNames()).not.toContain('Delete'); - session.dispose(); - }); - - it('preserves live capability choices across a new-session runtime replacement', async () => { - const { runtime } = await scopedRuntime(); - runtime.session.setActiveToolsByName( - runtime.session - .getActiveToolNames() - .filter((toolName) => toolName !== 'Write' && toolName !== 'Delete'), - ); - - await runtime.newSession(); - - expect(runtime.session.getActiveToolNames()).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Grep', - 'Glob', - 'LS', - ]); - await runtime.dispose(); - }); - - it('preserves live capability choices across a resumed-session runtime replacement', async () => { - const { cwd, runtime, sessionDir } = await scopedRuntime(true); - const target = SessionManager.create(cwd, sessionDir); - target.appendCustomEntry('target', {}); - runtime.session.setActiveToolsByName( - runtime.session - .getActiveToolNames() - .filter((toolName) => toolName !== 'Write' && toolName !== 'Delete'), - ); - - await runtime.switchSession(target.getSessionFile() ?? ''); - - expect(runtime.session.getActiveToolNames()).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Grep', - 'Glob', - 'LS', - ]); - await runtime.dispose(); - }); -}); diff --git a/tests/provider/toolScope.test.ts b/tests/provider/toolScope.test.ts deleted file mode 100644 index b61ac83..0000000 --- a/tests/provider/toolScope.test.ts +++ /dev/null @@ -1,435 +0,0 @@ -import { describe, expect, it, vi } from 'vitest'; -import { handoffGrokTools, restoreGrokTools, syncGrokTools } from '../../src/provider/toolScope.js'; - -type SyncOptions = { - captureDelete?: boolean; - imagineEnabled?: boolean; - webSearchRegistered?: boolean; -}; - -const ALL_REGISTERED_TOOLS = [ - 'read', - 'write', - 'edit', - 'grep', - 'find', - 'ls', - 'bash', - 'web_search', - 'Read', - 'Write', - 'Edit', - 'StrReplace', - 'Grep', - 'Glob', - 'LS', - 'Shell', - 'WebSearch', - 'Delete', - 'image_gen', -]; - -const DEFAULT_AUTO_ACTIVE_TOOLS = [ - 'read', - 'bash', - 'edit', - 'write', - 'Read', - 'Write', - 'Edit', - 'StrReplace', - 'Grep', - 'Glob', - 'LS', - 'Shell', - 'WebSearch', - 'Delete', - 'image_gen', -]; -const AUTO_ACTIVE_SYNC_OPTIONS = { - captureDelete: true, - imagineEnabled: true, - webSearchRegistered: true, -}; - -function toolState(initialTools: string[], registeredTools = ALL_REGISTERED_TOOLS) { - let activeTools = [...initialTools]; - const setActiveTools = vi.fn((nextTools: string[]) => { - activeTools = [...nextTools]; - }); - const pi = { - getActiveTools: () => activeTools, - getAllTools: () => registeredTools.map((name) => ({ name })), - setActiveTools, - }; - - return { - pi, - setActiveTools, - tools: () => activeTools, - replaceTools(nextTools: string[]) { - activeTools = [...nextTools]; - }, - }; -} - -function sync( - state: ReturnType, - provider: string | undefined, - id: string | undefined, - options: SyncOptions = {}, -) { - syncGrokTools(state.pi, provider && id ? { provider, id } : undefined, { - imagineEnabled: false, - webSearchRegistered: false, - ...options, - }); -} - -describe('syncGrokTools', () => { - it.each([ - ['grok-cli', 'grok-build'], - ['grok-cli', 'grok-composer-2.5-fast'], - ['grok-cli-2', 'grok-build'], - ])('uses compatibility names for exact legacy model %s/%s', (provider, id) => { - const state = toolState(['read', 'write', 'edit', 'grep', 'find', 'ls', 'bash']); - - sync(state, provider, id, { captureDelete: true }); - - expect(state.tools()).toEqual([ - 'Read', - 'Write', - 'Edit', - 'StrReplace', - 'Grep', - 'Glob', - 'LS', - 'Shell', - ]); - }); - - it.each([ - ['grok-cli', 'grok-4.5'], - ['grok-cli', 'future-model'], - ['grok-cli', 'grok-build-preview'], - ['other-provider', 'grok-build'], - ])('keeps native names for %s/%s', (provider, id) => { - const state = toolState(['read', 'write', 'bash']); - - sync(state, provider, id, { captureDelete: true }); - - expect(state.tools()).toEqual(['read', 'write', 'bash']); - expect(state.setActiveTools).not.toHaveBeenCalled(); - }); - - it('translates only active capabilities and keeps image_gen independent', () => { - const state = toolState(['custom-before', 'read', 'custom-after']); - - sync(state, 'grok-cli', 'grok-build', { - captureDelete: true, - imagineEnabled: true, - }); - - expect(state.tools()).toEqual(['custom-before', 'Read', 'custom-after', 'image_gen']); - expect(state.tools()).not.toEqual( - expect.arrayContaining(['Write', 'Edit', 'StrReplace', 'Delete', 'Shell']), - ); - }); - - it('does not resurrect capabilities excluded from the initial tool selection', () => { - const state = toolState([]); - - sync(state, 'grok-cli', 'grok-build', { captureDelete: true }); - - expect(state.tools()).toEqual([]); - expect(state.setActiveTools).not.toHaveBeenCalled(); - }); - - it('keeps legacy discovery shims active after normal extension auto-activation', () => { - const state = toolState(DEFAULT_AUTO_ACTIVE_TOOLS); - - sync(state, 'grok-cli', 'grok-build', { - captureDelete: true, - imagineEnabled: true, - webSearchRegistered: true, - }); - - expect(state.tools()).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Write', - 'Grep', - 'Glob', - 'LS', - 'Delete', - 'image_gen', - ]); - }); - - it('restores captured legacy discovery defaults after starting on a modern model', () => { - const state = toolState(DEFAULT_AUTO_ACTIVE_TOOLS); - - sync(state, 'grok-cli', 'grok-4.5', AUTO_ACTIVE_SYNC_OPTIONS); - expect(state.tools()).toEqual(['read', 'bash', 'edit', 'write', 'image_gen']); - - sync(state, 'grok-cli', 'grok-composer-2.5-fast', { - imagineEnabled: true, - webSearchRegistered: true, - }); - expect(state.tools()).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Write', - 'image_gen', - 'Grep', - 'Glob', - 'LS', - 'Delete', - ]); - }); - - it('remembers disabled legacy discovery defaults across model switches', () => { - const state = toolState(DEFAULT_AUTO_ACTIVE_TOOLS); - sync(state, 'grok-cli', 'grok-build', { - captureDelete: true, - webSearchRegistered: true, - }); - state.replaceTools(state.tools().filter((name) => name !== 'Glob')); - - sync(state, 'grok-cli', 'grok-4.5', { webSearchRegistered: true }); - sync(state, 'grok-cli', 'grok-build', { webSearchRegistered: true }); - - expect(state.tools()).toEqual([ - 'Read', - 'Shell', - 'Edit', - 'StrReplace', - 'Write', - 'Grep', - 'LS', - 'Delete', - ]); - }); - - it('does not translate automatically activated shims into inactive native capabilities', () => { - const state = toolState(DEFAULT_AUTO_ACTIVE_TOOLS); - - sync(state, 'grok-cli', 'grok-4.5', AUTO_ACTIVE_SYNC_OPTIONS); - - expect(state.tools()).toEqual(['read', 'bash', 'edit', 'write', 'image_gen']); - }); - - it('honors native CLI exclusions even when Pi auto-activates compatibility tools', () => { - const state = toolState( - [ - 'read', - 'Read', - 'Write', - 'Edit', - 'StrReplace', - 'Grep', - 'Glob', - 'LS', - 'Shell', - 'WebSearch', - 'image_gen', - ], - ALL_REGISTERED_TOOLS.filter((name) => !['write', 'edit', 'bash', 'Delete'].includes(name)), - ); - - sync(state, 'grok-cli', 'grok-build', { - captureDelete: true, - imagineEnabled: true, - webSearchRegistered: true, - }); - - expect(state.tools()).toEqual(['Read', 'Grep', 'Glob', 'LS', 'image_gen']); - }); - - it('keeps an allowed native capability usable when Pi filters out its shim registration', () => { - const state = toolState(['read'], ['read']); - - sync(state, 'grok-cli', 'grok-build', { - captureDelete: true, - imagineEnabled: true, - webSearchRegistered: true, - }); - - expect(state.tools()).toEqual(['read']); - expect(state.setActiveTools).not.toHaveBeenCalled(); - }); - - it('keeps one native vocabulary when an active capability has no registered shim', () => { - const state = toolState( - [ - 'read', - 'bash', - 'edit', - 'write', - 'Write', - 'Edit', - 'StrReplace', - 'Grep', - 'Glob', - 'LS', - 'Shell', - 'WebSearch', - 'Delete', - 'image_gen', - ], - ALL_REGISTERED_TOOLS.filter((name) => name !== 'Read'), - ); - - sync(state, 'grok-cli', 'grok-build', { - captureDelete: true, - imagineEnabled: true, - webSearchRegistered: true, - }); - - expect(state.tools()).toEqual(['read', 'bash', 'edit', 'write', 'Delete', 'image_gen']); - }); - - it('preserves unmanaged tool order while entering and leaving legacy mode', () => { - const state = toolState(['custom-a', 'read', 'custom-b', 'bash', 'custom-c']); - - sync(state, 'grok-cli', 'grok-build', { captureDelete: true }); - expect(state.tools()).toEqual(['custom-a', 'Read', 'custom-b', 'Shell', 'custom-c']); - - sync(state, 'grok-cli', 'grok-4.5'); - expect(state.tools()).toEqual(['custom-a', 'read', 'custom-b', 'bash', 'custom-c']); - }); - - it('preserves individual compatibility choices while legacy mode remains active', () => { - const state = toolState(['read', 'write']); - sync(state, 'grok-cli', 'grok-build', { captureDelete: true }); - state.replaceTools(['Read']); - - sync(state, 'grok-cli', 'grok-build'); - expect(state.tools()).toEqual(['Read']); - - sync(state, 'grok-cli', 'grok-4.5'); - expect(state.tools()).toEqual(['read']); - }); - - it('translates a newly enabled native capability before the next legacy prompt', () => { - const state = toolState(['read']); - sync(state, 'grok-cli', 'grok-build', { captureDelete: true }); - state.replaceTools(['Read', 'write', 'custom']); - - sync(state, 'grok-cli', 'grok-build'); - - expect(state.tools()).toEqual(['Read', 'Write', 'custom']); - expect(state.tools()).not.toContain('write'); - }); - - it('switches directly between legacy models without changing capabilities', () => { - const state = toolState(['read', 'edit']); - sync(state, 'grok-cli', 'grok-build', { captureDelete: true }); - state.setActiveTools.mockClear(); - - sync(state, 'grok-cli', 'grok-composer-2.5-fast'); - - expect(state.tools()).toEqual(['Read', 'Edit', 'StrReplace']); - expect(state.setActiveTools).not.toHaveBeenCalled(); - }); - - it('restores live compatibility capabilities to native names on shutdown', () => { - const state = toolState(['read', 'write', 'bash']); - sync(state, 'grok-cli', 'grok-build', { captureDelete: true }); - state.replaceTools(['Read', 'Shell', 'custom']); - - sync(state, undefined, undefined); - - expect(state.tools()).toEqual(['read', 'bash', 'custom']); - expect(state.tools()).not.toEqual(expect.arrayContaining(['Write', 'write'])); - }); - - it('captures Delete once, preserves it across modes, and remembers a live disable', () => { - const state = toolState(['read', 'Delete']); - - sync(state, 'grok-cli', 'grok-4.5', { captureDelete: true }); - expect(state.tools()).toEqual(['read']); - - sync(state, 'grok-cli', 'grok-build'); - expect(state.tools()).toEqual(['Read', 'Delete']); - - state.replaceTools(['Read']); - sync(state, 'grok-cli', 'grok-build'); - sync(state, 'grok-cli', 'grok-4.5'); - sync(state, 'grok-cli', 'grok-build'); - expect(state.tools()).toEqual(['Read']); - }); - - it('never introduces Delete when the initial allowlist excluded it', () => { - const state = toolState(['read']); - - sync(state, 'grok-cli', 'grok-4.5', { captureDelete: true }); - sync(state, 'grok-cli', 'grok-build'); - - expect(state.tools()).toEqual(['Read']); - }); - - it('maps WebSearch only when its adapter is registered', () => { - const registered = toolState(['web_search']); - sync(registered, 'grok-cli', 'grok-build', { - captureDelete: true, - webSearchRegistered: true, - }); - expect(registered.tools()).toEqual(['WebSearch']); - sync(registered, 'grok-cli', 'grok-4.5', { webSearchRegistered: true }); - expect(registered.tools()).toEqual(['web_search']); - - const unavailable = toolState(['web_search']); - sync(unavailable, 'grok-cli', 'grok-build', { captureDelete: true }); - expect(unavailable.tools()).toEqual(['web_search']); - expect(unavailable.setActiveTools).not.toHaveBeenCalled(); - }); - - it('applies enabled image_gen across provider and model changes', () => { - const state = toolState(['read']); - - sync(state, 'openai', 'gpt-5', { captureDelete: true, imagineEnabled: true }); - expect(state.tools()).toEqual(['read', 'image_gen']); - sync(state, 'grok-cli', 'grok-build', { imagineEnabled: true }); - expect(state.tools()).toEqual(['Read', 'image_gen']); - sync(state, 'grok-cli', 'grok-4.5', { imagineEnabled: true }); - expect(state.tools()).toEqual(['read', 'image_gen']); - }); - - it('removes disabled image_gen without disturbing unrelated tools', () => { - const state = toolState(['custom-a', 'image_gen', 'custom-b', 'read']); - - sync(state, 'openai', 'gpt-5', { captureDelete: true }); - - expect(state.tools()).toEqual(['custom-a', 'custom-b', 'read']); - }); - - it('does not call setActiveTools when the ordered desired set is already active', () => { - const state = toolState(['custom', 'Read']); - - sync(state, 'grok-cli', 'grok-build', { captureDelete: true }); - - expect(state.setActiveTools).not.toHaveBeenCalled(); - }); - - it('restores only capabilities present in a replacement runtime registry', () => { - const source = toolState(['read', 'write', 'custom']); - sync(source, 'grok-cli', 'grok-build', { captureDelete: true }); - sync(source, undefined, undefined); - handoffGrokTools(source.pi, '/sessions/restricted-target.jsonl'); - const target = toolState( - ['read', 'Read', 'Delete', 'custom'], - ['read', 'Read', 'Delete', 'custom'], - ); - - restoreGrokTools(target.pi, '/sessions/restricted-target.jsonl'); - sync(target, 'grok-cli', 'grok-build', { captureDelete: true }); - - expect(target.tools()).toEqual(['Read', 'custom']); - expect(target.tools()).not.toEqual(expect.arrayContaining(['write', 'Write', 'Delete'])); - }); -}); diff --git a/tests/provider/usage.test.ts b/tests/provider/usage.test.ts new file mode 100644 index 0000000..1cb9803 --- /dev/null +++ b/tests/provider/usage.test.ts @@ -0,0 +1,36 @@ +import type { ExtensionAPI, ExtensionCommandContext } from '@earendil-works/pi-coding-agent'; +import { describe, expect, it, vi } from 'vitest'; +import { registerUsageCommand } from '../../src/provider/usage.js'; + +describe('Grok CLI usage command', () => { + it('shows the account-route error with the empty quota state', async () => { + const commands = new Map< + string, + { handler: (args: string, ctx: ExtensionCommandContext) => Promise } + >(); + registerUsageCommand( + { + registerCommand(name, command) { + commands.set(name, command); + }, + } as Pick, + async () => { + throw new Error('Account 1 needs login.'); + }, + ); + const notify = vi.fn(); + + await commands.get('grok-cli-usage')?.handler('', { + modelRegistry: { + getAll: () => [{ provider: 'grok-cli', id: 'grok-build' }], + }, + ui: { notify }, + } as unknown as ExtensionCommandContext); + + expect(notify).toHaveBeenCalledWith( + expect.stringContaining('no billing data available'), + 'info', + ); + expect(notify).toHaveBeenCalledWith(expect.stringContaining('Account 1 needs login.'), 'info'); + }); +}); diff --git a/tests/stateTestHelpers.ts b/tests/stateTestHelpers.ts new file mode 100644 index 0000000..8297eaa --- /dev/null +++ b/tests/stateTestHelpers.ts @@ -0,0 +1,103 @@ +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { afterEach } from 'vitest'; +import { ACCOUNT_VAULT_MARKER, mutateAccountVault } from '../src/provider/accountVault.js'; +import { getConfigPath, writeFileAtomic } from '../src/storage.js'; + +export const TEST_ACCOUNTS = [ + { provider: 'grok-cli', label: 'Personal' }, + { provider: 'grok-cli-2', label: 'Work' }, +]; + +export const oauthCredential = (access: string) => ({ + type: 'oauth' as const, + access, + refresh: `${access}-refresh`, + expires: Date.now() + 60_000, +}); + +export function deferred() { + let resolvePromise!: (value: T) => void; + const promise = new Promise((resolve) => { + resolvePromise = resolve; + }); + return { promise, resolve: (value: T) => resolvePromise(value) }; +} + +export function saveTestAccounts(selectedProvider = 'grok-cli-2') { + writeFileAtomic( + getConfigPath(), + `${JSON.stringify({ + version: 2, + accounts: { nextAccountNumber: 3, selectedProvider, items: TEST_ACCOUNTS }, + imagine: { enabled: true }, + })}\n`, + ); +} + +/** + * Point HOME at a fresh temp dir for the whole test file, restoring it on + * teardown. Returns a setup function that creates a new dir per call. + */ +export function useTempHome(): () => string { + const originalHome = process.env.HOME; + const dirs: string[] = []; + afterEach(() => { + if (originalHome === undefined) delete process.env.HOME; + else process.env.HOME = originalHome; + for (const dir of dirs.splice(0)) rmSync(dir, { recursive: true, force: true }); + }); + return () => { + const dir = mkdtempSync(join(tmpdir(), 'grok-cli-test-')); + dirs.push(dir); + process.env.HOME = dir; + return dir; + }; +} + +export function useEnvironmentToken() { + const originalToken = process.env.GROK_CLI_OAUTH_TOKEN; + afterEach(() => { + if (originalToken === undefined) delete process.env.GROK_CLI_OAUTH_TOKEN; + else process.env.GROK_CLI_OAUTH_TOKEN = originalToken; + }); + return (token?: string) => { + if (token === undefined) delete process.env.GROK_CLI_OAUTH_TOKEN; + else process.env.GROK_CLI_OAUTH_TOKEN = token; + }; +} + +export function setAccount1Credential(access: string) { + return mutateAccountVault((vault) => { + vault.migration.legacyCredentialCopyComplete = true; + vault.accounts[0].credential = oauthCredential(access); + vault.accounts[0].revision += 1; + vault.activeAccountId = 'account-1'; + }); +} + +export function writeTestJson(path: string, value: unknown) { + mkdirSync(join(path, '..'), { recursive: true }); + writeFileSync(path, `${JSON.stringify(value, null, 2)}\n`); +} + +export function writePiCredentials(credentials: Record) { + writeTestJson(join(process.env.HOME as string, '.pi', 'agent', 'auth.json'), credentials); +} + +export function writePiCredential(credential: { + access: string; + refresh: string; + expires: number; +}) { + writePiCredentials({ 'grok-cli': { type: 'oauth', ...credential } }); +} + +export function writePiVaultMarker() { + writePiCredential({ + access: ACCOUNT_VAULT_MARKER, + refresh: ACCOUNT_VAULT_MARKER, + expires: Number.MAX_SAFE_INTEGER, + }); +} diff --git a/tests/storage.test.ts b/tests/storage.test.ts index ba6e670..646c8e8 100644 --- a/tests/storage.test.ts +++ b/tests/storage.test.ts @@ -1,17 +1,23 @@ -import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'; +import { + existsSync, + mkdirSync, + readFileSync, + rmSync, + symlinkSync, + utimesSync, + writeFileSync, +} from 'node:fs'; import { join } from 'node:path'; -import { describe, expect, it } from 'vitest'; +import { describe, expect, it, vi } from 'vitest'; import { DEFAULT_CONFIG, loadConfig, migrateLegacyConfig } from '../src/config.js'; import { + acquireFileLock, getConfigPath, getGrokCliDirectory, getLegacyConfigPath, - getLegacyVisionCachePath, getQuotaCachePath, - getVisionCachePath, } from '../src/storage.js'; -import { getCachePath } from '../src/vision/cache.js'; -import { useTempHome } from './vision/helpers.js'; +import { useTempHome } from './stateTestHelpers.js'; const setupHome = useTempHome(); @@ -19,70 +25,156 @@ function writeJson(path: string, value: unknown) { writeFileSync(path, `${JSON.stringify(value, null, 2)}\n`); } +async function expectLockTimeout(lockPath: string) { + const pending = acquireFileLock(getQuotaCachePath()); + const rejection = expect(pending).rejects.toThrow(`Timed out waiting for file lock: ${lockPath}`); + await vi.advanceTimersByTimeAsync(30_000); + await rejection; +} + describe('Grok CLI storage', () => { it('groups extension-owned files under one directory without creating it on read', () => { const home = setupHome(); expect(getGrokCliDirectory()).toBe(join(home, '.pi', 'grok-cli')); expect(getConfigPath()).toBe(join(home, '.pi', 'grok-cli', 'config.json')); - expect(getVisionCachePath()).toBe(join(home, '.pi', 'grok-cli', 'vision-cache.json')); expect(getQuotaCachePath()).toBe(join(home, '.pi', 'grok-cli', 'quota-cache.json')); expect(loadConfig()).toEqual({ config: DEFAULT_CONFIG }); expect(existsSync(getGrokCliDirectory())).toBe(false); }); - it('migrates the consolidated config and vision cache after verified writes', () => { + it('migrates the consolidated config after a verified write', () => { const home = setupHome(); mkdirSync(join(home, '.pi'), { recursive: true }); writeJson(getLegacyConfigPath(), { ...DEFAULT_CONFIG, imagine: { enabled: false } }); - writeJson(getLegacyVisionCachePath(), { - version: 1, - entries: { cached: { description: 'keep' } }, - }); expect(migrateLegacyConfig()).toEqual({}); expect(loadConfig().config.imagine.enabled).toBe(false); - expect(JSON.parse(readFileSync(getConfigPath(), 'utf8')).version).toBe(2); - expect(JSON.parse(readFileSync(getVisionCachePath(), 'utf8'))).toEqual({ - version: 1, - entries: { cached: { description: 'keep' } }, - }); + expect(JSON.parse(readFileSync(getConfigPath(), 'utf8')).version).toBe(3); expect(existsSync(getLegacyConfigPath())).toBe(false); - expect(existsSync(getLegacyVisionCachePath())).toBe(false); - expect(getCachePath()).toBe(getVisionCachePath()); }); - it('keeps new files authoritative and preserves conflicting legacy files', () => { + it('keeps the new config authoritative and preserves a conflicting legacy file', () => { setupHome(); mkdirSync(getGrokCliDirectory(), { recursive: true }); writeJson(getConfigPath(), DEFAULT_CONFIG); writeJson(getLegacyConfigPath(), { ...DEFAULT_CONFIG, imagine: { enabled: false } }); - writeJson(getVisionCachePath(), { version: 1, entries: { current: {} } }); - writeJson(getLegacyVisionCachePath(), { version: 1, entries: { legacy: {} } }); const migration = migrateLegacyConfig(); expect(migration.warning).toContain(getLegacyConfigPath()); - expect(migration.warning).toContain(getLegacyVisionCachePath()); expect(loadConfig().config).toEqual(DEFAULT_CONFIG); expect(existsSync(getLegacyConfigPath())).toBe(true); - expect(existsSync(getLegacyVisionCachePath())).toBe(true); }); it('falls back to legacy files when the destination directory cannot be created', () => { const home = setupHome(); mkdirSync(join(home, '.pi'), { recursive: true }); writeJson(getLegacyConfigPath(), { ...DEFAULT_CONFIG, imagine: { enabled: false } }); - writeJson(getLegacyVisionCachePath(), { version: 1, entries: {} }); writeFileSync(getGrokCliDirectory(), 'not a directory'); const migration = migrateLegacyConfig(); expect(migration.warning).toMatch(/Could not migrate/); expect(loadConfig().config.imagine.enabled).toBe(false); - expect(getCachePath()).toBe(getLegacyVisionCachePath()); expect(existsSync(getLegacyConfigPath())).toBe(true); - expect(existsSync(getLegacyVisionCachePath())).toBe(true); + }); + + it('recovers an incomplete lock left by a stopped process', async () => { + setupHome(); + mkdirSync(getGrokCliDirectory(), { recursive: true }); + const lockPath = `${getQuotaCachePath()}.lock`; + writeFileSync(lockPath, '{'); + utimesSync(lockPath, new Date(0), new Date(0)); + vi.useFakeTimers(); + + try { + const pending = acquireFileLock(getQuotaCachePath()); + await vi.advanceTimersByTimeAsync(30_025); + const release = await pending; + await release(); + expect(existsSync(lockPath)).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it('recovers a stale lock owned by a stopped process', async () => { + setupHome(); + mkdirSync(getGrokCliDirectory(), { recursive: true }); + const lockPath = `${getQuotaCachePath()}.lock`; + writeJson(lockPath, { pid: 2_147_483_647, token: 'stale-owner' }); + utimesSync(lockPath, new Date(0), new Date(0)); + + const release = await acquireFileLock(getQuotaCachePath()); + await release(); + + expect(existsSync(lockPath)).toBe(false); + }); + + it('does not reclaim an old lock owned by a running process', async () => { + setupHome(); + mkdirSync(getGrokCliDirectory(), { recursive: true }); + const lockPath = `${getQuotaCachePath()}.lock`; + writeJson(lockPath, { pid: process.pid, token: 'live-owner' }); + utimesSync(lockPath, new Date(0), new Date(0)); + vi.useFakeTimers(); + + try { + await expectLockTimeout(lockPath); + expect(existsSync(lockPath)).toBe(true); + } finally { + vi.useRealTimers(); + rmSync(lockPath, { force: true }); + } + }); + + it('recovers a stale recovery barrier after its process ID is reused', async () => { + setupHome(); + mkdirSync(getGrokCliDirectory(), { recursive: true }); + const lockPath = `${getQuotaCachePath()}.lock`; + const recoveryPath = `${lockPath}.recovery.${process.pid}.stale`; + writeFileSync(recoveryPath, ''); + utimesSync(recoveryPath, new Date(0), new Date(0)); + vi.useFakeTimers(); + + try { + const release = await acquireFileLock(getQuotaCachePath()); + await release(); + expect(existsSync(recoveryPath)).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it('times out while recovery remains in progress', async () => { + setupHome(); + mkdirSync(getGrokCliDirectory(), { recursive: true }); + const lockPath = `${getQuotaCachePath()}.lock`; + const recoveryPath = `${lockPath}.recovery.${process.pid}.active`; + writeFileSync(recoveryPath, ''); + vi.useFakeTimers(); + + try { + await expectLockTimeout(lockPath); + } finally { + vi.useRealTimers(); + rmSync(recoveryPath, { force: true }); + } + }); + + it('continues when a recovery barrier disappears during polling', async () => { + setupHome(); + mkdirSync(getGrokCliDirectory(), { recursive: true }); + const lockPath = `${getQuotaCachePath()}.lock`; + const recoveryPath = `${lockPath}.recovery.${process.pid}.gone`; + symlinkSync(`${recoveryPath}.missing`, recoveryPath); + + const release = await acquireFileLock(getQuotaCachePath()); + await release(); + + expect(existsSync(recoveryPath)).toBe(false); + expect(existsSync(lockPath)).toBe(false); }); }); diff --git a/tests/tools/files.test.ts b/tests/tools/files.test.ts deleted file mode 100644 index 2d44baa..0000000 --- a/tests/tools/files.test.ts +++ /dev/null @@ -1,835 +0,0 @@ -import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'; -import { join } from 'node:path'; -import { describe, expect, it, vi } from 'vitest'; -import { registerFileTools } from '../../src/tools/files.js'; -import { - collectTools, - executePreparedTool, - executeTool, - executeToolWithOptions, - firstText, - prepareToolArguments, - renderToolCall, - renderToolResult, - type ToolResult, - tempDir, -} from './toolTestHelpers.js'; - -function expectStoryState(result: ToolResult, cwd: string, replacements: number, content: string) { - expect(result.details).toMatchObject({ - path: expectedPath(cwd, 'story.txt'), - replacements, - }); - expect(readFileSync(join(cwd, 'story.txt'), 'utf-8')).toBe(content); -} - -function expectedPath(cwd: string, ...parts: string[]) { - return join(cwd, ...parts); -} - -function strReplace(cwd: string, old_str: string, new_str: string) { - return executeTool( - collectTools(registerFileTools).get('StrReplace'), - { path: 'story.txt', old_str, new_str }, - cwd, - ); -} - -function strReplaceWithPreparedArgs(cwd: string, params: Record) { - return executePreparedTool( - collectTools(registerFileTools).get('StrReplace'), - { path: 'story.txt', ...params }, - cwd, - ); -} - -describe('file tools', () => { - it('lists directory contents including hidden files', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, '.hidden'), 'secret', 'utf-8'); - writeFileSync(join(cwd, 'visible.txt'), 'visible', 'utf-8'); - - const result = await executeTool(collectTools(registerFileTools).get('LS'), { path: '.' }, cwd); - - expect(firstText(result)).toContain('.hidden'); - expect(firstText(result)).toContain('visible.txt'); - expect(result.details).toBeUndefined(); - }); - - it('truncates oversized directory listings', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - Array.from({ length: 700 }, (_, index) => - writeFileSync(join(cwd, `${String(index).padStart(4, '0')}-${'x'.repeat(70)}`), '', 'utf-8'), - ); - - const result = await executeTool(collectTools(registerFileTools).get('LS'), { path: '.' }, cwd); - - expect(firstText(result).split('\n').slice(0, 500)).toHaveLength(500); - expect(firstText(result)).toContain('500 entries limit reached'); - expect(result.details).toMatchObject({ entryLimitReached: 500 }); - }); - - it('lists directory contents when Unix ls is not on PATH', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - const oldPath = process.env.PATH; - process.env.PATH = tempDir('pi-grok-cli-empty-bin-'); - vi.resetModules(); - writeFileSync(join(cwd, 'visible.txt'), 'visible', 'utf-8'); - - try { - const result = await executeTool( - collectTools((await import('../../src/tools/files.js')).registerFileTools).get('LS'), - { path: '.' }, - cwd, - ); - - expect(firstText(result)).toContain('visible.txt'); - expect(result.details).toBeUndefined(); - } finally { - process.env.PATH = oldPath; - vi.resetModules(); - } - }); - - it('reports filesystem errors for invalid file operations', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - mkdirSync(join(cwd, 'dir')); - writeFileSync(join(cwd, 'blocked'), 'not a directory', 'utf-8'); - const tools = collectTools(registerFileTools); - - const lsResult = executeTool(tools.get('LS'), { path: 'missing-dir' }, cwd); - const writeResult = executeTool( - tools.get('Write'), - { path: 'blocked/file.txt', content: 'content' }, - cwd, - ); - const replaceResult = executeTool( - tools.get('StrReplace'), - { path: 'dir', old_str: 'old', new_str: 'new' }, - cwd, - ); - const deleteResult = await executeTool(tools.get('Delete'), { path: 'dir' }, cwd); - - await expect(lsResult).rejects.toThrow(/Path not found/); - await expect(writeResult).rejects.toThrow(); - await expect(replaceResult).rejects.toThrow(); - expect(firstText(deleteResult).startsWith('Delete error:')).toBe(true); - expect(deleteResult.details).toEqual({ - path: join(cwd, 'dir'), - deleted: false, - failed: true, - error: expect.stringMatching( - /EISDIR: illegal operation on a directory|operation not permitted/, - ), - }); - }); - - it('allows file operations on paths outside the workspace', async () => { - const parent = tempDir('pi-grok-cli-files-parent-'); - const cwd = join(parent, 'workspace'); - mkdirSync(cwd); - const tools = collectTools(registerFileTools); - - const writeResult = await executeTool( - tools.get('Write'), - { path: '../outside.txt', content: 'hello' }, - cwd, - ); - expect(firstText(writeResult)).toBe('Successfully wrote 5 bytes to ../outside.txt'); - expect(existsSync(join(parent, 'outside.txt'))).toBe(true); - - const replaceResult = await executeTool( - tools.get('StrReplace'), - { path: '../outside.txt', old_str: 'hello', new_str: 'world' }, - cwd, - ); - expect(firstText(replaceResult)).toBe('Replaced 1 occurrence(s) in ../outside.txt'); - expect(readFileSync(join(parent, 'outside.txt'), 'utf-8')).toBe('world'); - - const lsResult = await executeTool(tools.get('LS'), { path: '..' }, cwd); - expect(firstText(lsResult)).toContain('outside.txt'); - }); - - it('writes Cursor-style contents arguments', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - - const result = await executePreparedTool( - collectTools(registerFileTools).get('Write'), - { path: 'nested/notes.txt', contents: 'alpha\nbeta' }, - cwd, - ); - - expect(firstText(result)).toBe('Successfully wrote 10 bytes to nested/notes.txt'); - expect(readFileSync(join(cwd, 'nested/notes.txt'), 'utf-8')).toBe('alpha\nbeta'); - expect(result.details).toBeUndefined(); - }); - - it('adopts the native Write result contract for multibyte content', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - const result = await executeTool( - collectTools(registerFileTools).get('Write'), - { path: 'emoji.txt', content: 'a🙂漢' }, - cwd, - ); - - expect(firstText(result)).toBe('Successfully wrote 4 bytes to emoji.txt'); - expect(result.details).toBeUndefined(); - }); - - it('replaces every exact string occurrence', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf-8'); - - const result = await strReplace(cwd, 'red', 'green'); - - expect(firstText(result)).toBe('Replaced 2 occurrence(s) in story.txt'); - expectStoryState(result, cwd, 2, 'green blue green'); - }); - - it('rejects empty replacement search strings without changing files', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf-8'); - - await expect(strReplace(cwd, '', 'green')).rejects.toThrow(/must not be empty/); - expect(readFileSync(join(cwd, 'story.txt'), 'utf8')).toBe('red blue red'); - }); - - it('treats replacement text as a literal string', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'abc', 'utf-8'); - - const result = await strReplace(cwd, 'a', '$&'); - - expect(firstText(result)).toBe('Replaced 1 occurrence(s) in story.txt'); - expectStoryState(result, cwd, 1, '$&bc'); - }); - - it('preserves BOM and CRLF bytes outside StrReplace matches', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'windows.txt'), '\uFEFFalpha\r\nbeta\r\n', 'utf8'); - - await executeTool( - collectTools(registerFileTools).get('StrReplace'), - { path: 'windows.txt', old_str: 'beta', new_str: 'gamma' }, - cwd, - ); - - expect(readFileSync(join(cwd, 'windows.txt'), 'utf8')).toBe('\uFEFFalpha\r\ngamma\r\n'); - }); - - it('replaces string occurrences with Grok and Cursor argument variants', async () => { - const oldStringCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(oldStringCwd, 'story.txt'), 'red blue red', 'utf-8'); - - const oldStringResult = await strReplaceWithPreparedArgs(oldStringCwd, { - old_string: 'red', - new_string: 'green', - }); - - expect(firstText(oldStringResult)).toBe('Replaced 2 occurrence(s) in story.txt'); - expectStoryState(oldStringResult, oldStringCwd, 2, 'green blue green'); - - const oldTextCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(oldTextCwd, 'story.txt'), 'red blue red', 'utf-8'); - - const oldTextResult = await strReplaceWithPreparedArgs(oldTextCwd, { - oldText: 'red', - newText: 'green', - }); - - expect(firstText(oldTextResult)).toBe('Replaced 2 occurrence(s) in story.txt'); - expectStoryState(oldTextResult, oldTextCwd, 2, 'green blue green'); - - const nestedCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(nestedCwd, 'story.txt'), 'red blue red', 'utf-8'); - - const nestedResult = await strReplaceWithPreparedArgs(nestedCwd, { - strReplace: { oldText: 'red', newText: 'green' }, - }); - - expect(firstText(nestedResult)).toBe('Replaced 2 occurrence(s) in story.txt'); - expectStoryState(nestedResult, nestedCwd, 2, 'green blue green'); - }); - - it('edits files with single, multiple, and stringified replacement inputs', async () => { - const singleCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(singleCwd, 'story.txt'), 'red blue', 'utf-8'); - - const singleResult = await executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { path: 'story.txt', oldText: 'red', newText: 'green' }, - singleCwd, - ); - - expect(firstText(singleResult)).toBe('Successfully replaced 1 block(s) in story.txt.'); - expect(singleResult.details).toMatchObject({ - diff: expect.stringContaining('green blue'), - patch: expect.stringContaining('story.txt'), - }); - expect(readFileSync(join(singleCwd, 'story.txt'), 'utf-8')).toBe('green blue'); - - const multipleCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(multipleCwd, 'story.txt'), 'red blue', 'utf-8'); - - const multipleResult = await executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { - path: 'story.txt', - edits: [ - { oldText: 'red', newText: 'green' }, - { oldText: 'blue', newText: 'yellow' }, - ], - }, - multipleCwd, - ); - - expect(firstText(multipleResult)).toBe('Successfully replaced 2 block(s) in story.txt.'); - expect(multipleResult.details).toMatchObject({ - diff: expect.stringContaining('green yellow'), - patch: expect.stringContaining('story.txt'), - }); - expect(readFileSync(join(multipleCwd, 'story.txt'), 'utf-8')).toBe('green yellow'); - - const stringifiedCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(stringifiedCwd, 'story.txt'), 'red blue', 'utf-8'); - - const stringifiedResult = await executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { - path: 'story.txt', - edits: JSON.stringify([{ oldText: 'red', newText: 'green' }]), - }, - stringifiedCwd, - ); - - expect(firstText(stringifiedResult)).toBe('Successfully replaced 1 block(s) in story.txt.'); - expect(stringifiedResult.details).toMatchObject({ - diff: expect.stringContaining('green blue'), - patch: expect.stringContaining('story.txt'), - }); - expect(readFileSync(join(stringifiedCwd, 'story.txt'), 'utf-8')).toBe('green blue'); - }); - - it('edits files with literal replacement text', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'abc', 'utf-8'); - - const result = await executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { path: 'story.txt', oldText: 'a', newText: '$&' }, - cwd, - ); - - expect(firstText(result)).toBe('Successfully replaced 1 block(s) in story.txt.'); - expect(result.details).toMatchObject({ - diff: expect.stringContaining('$&bc'), - patch: expect.stringContaining('story.txt'), - }); - expect(readFileSync(join(cwd, 'story.txt'), 'utf-8')).toBe('$&bc'); - }); - - it('rejects empty edit search strings without changing files', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf-8'); - - await expect( - executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { path: 'story.txt', oldText: '', newText: 'green' }, - cwd, - ), - ).rejects.toThrow(/oldText must not be empty/); - expect(readFileSync(join(cwd, 'story.txt'), 'utf-8')).toBe('red blue red'); - }); - - it('reports unsupported edit strategies without changing files', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf-8'); - - await expect( - executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { path: 'story.txt', applyPatch: { patchContent: 'patch' } }, - cwd, - ), - ).rejects.toThrow(/edits must contain at least one/); - expect(readFileSync(join(cwd, 'story.txt'), 'utf-8')).toBe('red blue red'); - }); - - it('rejects malformed edit lists without changing files', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf-8'); - const edit = collectTools(registerFileTools).get('Edit'); - - await expect( - executePreparedTool(edit, { path: 'story.txt', edits: '{ not json' }, cwd), - ).rejects.toThrow(/edits must contain at least one/); - await expect( - executePreparedTool(edit, { path: 'story.txt', edits: [{ oldText: 'red' }] }, cwd), - ).rejects.toThrow(/edits must contain at least one/); - expect(readFileSync(join(cwd, 'story.txt'), 'utf-8')).toBe('red blue red'); - }); - - it('reports no-match and filesystem failures from Edit', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf-8'); - mkdirSync(join(cwd, 'directory')); - const edit = collectTools(registerFileTools).get('Edit'); - - await expect( - executePreparedTool(edit, { path: 'story.txt', oldText: 'purple', newText: 'green' }, cwd), - ).rejects.toThrow(/Could not find the exact text/); - await expect( - executePreparedTool(edit, { path: 'directory', oldText: 'red', newText: 'green' }, cwd), - ).rejects.toThrow(/Could not edit file|EISDIR/); - expect(readFileSync(join(cwd, 'story.txt'), 'utf-8')).toBe('red blue red'); - }); - - it('leaves files unchanged when the replacement string is absent', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf-8'); - - const result = await strReplace(cwd, 'purple', 'green'); - - expect(firstText(result)).toBe('String not found in story.txt: "purple"'); - expectStoryState(result, cwd, 0, 'red blue red'); - }); - - it('deletes existing files and reports missing files', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'remove.txt'), 'delete me', 'utf-8'); - const tools = collectTools(registerFileTools); - - const deletedResult = await executeTool(tools.get('Delete'), { path: 'remove.txt' }, cwd); - - expect(firstText(deletedResult)).toBe('Successfully deleted remove.txt'); - expect(deletedResult.details).toEqual({ - path: expectedPath(cwd, 'remove.txt'), - deleted: true, - }); - expect(existsSync(join(cwd, 'remove.txt'))).toBe(false); - - const missingResult = await executeTool(tools.get('Delete'), { path: 'remove.txt' }, cwd); - - expect(firstText(missingResult)).toBe(`File not found: ${join(cwd, 'remove.txt')}`); - expect(missingResult.details).toEqual({ - path: join(cwd, 'remove.txt'), - deleted: false, - }); - }); - - it('renders file tool calls and result states', () => { - const tools = collectTools(registerFileTools); - - expect(renderToolCall(tools.get('LS'), { path: '.' })).toBe('ls .'); - expect(renderToolCall(tools.get('Write'), { path: 'notes.txt' })).toBe('write notes.txt'); - expect(renderToolCall(tools.get('StrReplace'), { path: 'notes.txt' })).toBe( - 'StrReplace notes.txt', - ); - expect(renderToolCall(tools.get('Edit'), { path: 'notes.txt' })).toContain('edit notes.txt'); - expect(renderToolCall(tools.get('Delete'), { path: 'notes.txt' })).toBe('Delete notes.txt'); - expect( - renderToolResult(tools.get('StrReplace'), { - content: [{ type: 'text', text: 'no replacement' }], - details: { replacements: 0 }, - }), - ).toBe('No replacements'); - expect( - renderToolResult(tools.get('Delete'), { - content: [{ type: 'text', text: 'not deleted' }], - details: { deleted: false }, - }), - ).toBe('Not deleted'); - expect( - renderToolResult(tools.get('Edit'), { - content: [{ type: 'text', text: 'edited' }], - details: undefined, - }), - ).toBe(''); - expect( - renderToolResult( - tools.get('LS'), - { - content: [{ type: 'text', text: 'full listing' }], - details: { path: '/tmp/project' }, - }, - { expanded: true, isPartial: false }, - ), - ).toContain('full listing'); - expect( - renderToolResult( - tools.get('Write'), - { - content: [{ type: 'text', text: 'writing' }], - details: { bytesWritten: 10 }, - }, - { expanded: false, isPartial: true }, - ), - ).toBe(''); - }); - - it('renders StrReplace removed and added diff lines with the theme diff colors', () => { - const theme = { - bold: (text: string) => text, - bg: (_name: string, text: string) => text, - fg: (name: string, text: string) => `<${name}>${text}`, - inverse: (text: string) => text, - }; - - const rendered = renderToolResult( - collectTools(registerFileTools).get('StrReplace'), - { - content: [{ type: 'text', text: 'replaced' }], - details: { - replacements: 1, - diff: ['--- notes.txt', '+++ notes.txt', '@@ -1 +1 @@', '-old', '+new'].join('\n'), - }, - }, - { expanded: true, isPartial: false }, - {}, - {}, - theme, - ); - - expect(rendered).toContain('-old'); - expect(rendered).toContain('+new'); - }); -}); - -describe('native file adapter contracts', () => { - it('normalizes Write file_path and contents aliases into native arguments', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - const write = collectTools(registerFileTools).get('Write'); - - expect( - prepareToolArguments(write, { file_path: 'nested/note.txt', contents: 'hello' }), - ).toEqual({ path: 'nested/note.txt', content: 'hello' }); - const result = await executePreparedTool( - write, - { file_path: 'nested/note.txt', contents: 'hello' }, - cwd, - ); - - expect(firstText(result)).toBe('Successfully wrote 5 bytes to nested/note.txt'); - expect(result.details).toBeUndefined(); - expect(readFileSync(join(cwd, 'nested', 'note.txt'), 'utf8')).toBe('hello'); - }); - - it('prevents an aborted native Write before changing the filesystem', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - const controller = new AbortController(); - controller.abort(); - - await expect( - executeToolWithOptions( - collectTools(registerFileTools).get('Write'), - { path: 'aborted.txt', content: 'must not exist' }, - cwd, - { signal: controller.signal }, - ), - ).rejects.toThrow(/aborted/i); - expect(existsSync(join(cwd, 'aborted.txt'))).toBe(false); - }); - - it('throws delegated Write filesystem failures', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'blocked'), 'not a directory', 'utf8'); - - await expect( - executeTool( - collectTools(registerFileTools).get('Write'), - { path: 'blocked/file.txt', content: 'content' }, - cwd, - ), - ).rejects.toThrow(); - }); - - it('serializes concurrent native Write and Edit mutations in request order', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'ordered.txt'), 'initial', 'utf8'); - const tools = collectTools(registerFileTools); - - await Promise.all([ - executeTool(tools.get('Write'), { path: 'ordered.txt', content: 'first' }, cwd), - executePreparedTool( - tools.get('Edit'), - { path: 'ordered.txt', oldText: 'first', newText: 'second' }, - cwd, - ), - ]); - - expect(readFileSync(join(cwd, 'ordered.txt'), 'utf8')).toBe('second'); - }); - - it('applies multiple Edit matches against the original content and returns native diffs', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'alpha beta', 'utf8'); - - const result = await executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { - path: 'story.txt', - edits: [ - { oldText: 'alpha', newText: 'beta' }, - { oldText: 'beta', newText: 'gamma' }, - ], - }, - cwd, - ); - - expect(readFileSync(join(cwd, 'story.txt'), 'utf8')).toBe('beta gamma'); - expect(firstText(result)).toBe('Successfully replaced 2 block(s) in story.txt.'); - expect(result.details).toMatchObject({ - diff: expect.stringContaining('beta gamma'), - patch: expect.stringContaining('story.txt'), - }); - }); - - it('rejects ambiguous and overlapping native Edit matches without writing', async () => { - const ambiguousCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(ambiguousCwd, 'story.txt'), 'red red', 'utf8'); - const edit = collectTools(registerFileTools).get('Edit'); - - await expect( - executePreparedTool( - edit, - { path: 'story.txt', oldText: 'red', newText: 'green' }, - ambiguousCwd, - ), - ).rejects.toThrow(/multiple locations|unique/i); - expect(readFileSync(join(ambiguousCwd, 'story.txt'), 'utf8')).toBe('red red'); - - const overlapCwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(overlapCwd, 'story.txt'), 'abcdef', 'utf8'); - await expect( - executePreparedTool( - edit, - { - path: 'story.txt', - edits: [ - { oldText: 'abcd', newText: 'one' }, - { oldText: 'cdef', newText: 'two' }, - ], - }, - overlapCwd, - ), - ).rejects.toThrow(/overlap/i); - expect(readFileSync(join(overlapCwd, 'story.txt'), 'utf8')).toBe('abcdef'); - }); - - it('preserves BOM and CRLF through native Edit', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'windows.txt'), '\uFEFFalpha\r\nbeta\r\n', 'utf8'); - - await executePreparedTool( - collectTools(registerFileTools).get('Edit'), - { file_path: 'windows.txt', old_string: 'beta', new_string: 'gamma' }, - cwd, - ); - - expect(readFileSync(join(cwd, 'windows.txt'), 'utf8')).toBe('\uFEFFalpha\r\ngamma\r\n'); - }); - - it('does not advertise unsupported Edit applyPatch input', () => { - const edit = collectTools(registerFileTools).get('Edit') as unknown as { - parameters: { properties: Record }; - }; - - expect(edit.parameters.properties).not.toHaveProperty('applyPatch'); - }); - - it('throws for empty StrReplace searches and operational failures', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf8'); - mkdirSync(join(cwd, 'directory')); - const replace = collectTools(registerFileTools).get('StrReplace'); - - await expect( - executeTool(replace, { path: 'story.txt', old_str: '', new_str: 'green' }, cwd), - ).rejects.toThrow(/must not be empty/i); - await expect( - executeTool(replace, { path: 'directory', old_str: 'red', new_str: 'green' }, cwd), - ).rejects.toThrow(); - expect(readFileSync(join(cwd, 'story.txt'), 'utf8')).toBe('red blue red'); - }); - - it('respects cancellation before StrReplace reads or writes', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf8'); - const controller = new AbortController(); - controller.abort(); - - await expect( - executeToolWithOptions( - collectTools(registerFileTools).get('StrReplace'), - { path: 'story.txt', old_str: 'red', new_str: 'green' }, - cwd, - { signal: controller.signal }, - ), - ).rejects.toThrow(/aborted/i); - expect(readFileSync(join(cwd, 'story.txt'), 'utf8')).toBe('red blue red'); - }); - - it('cancels StrReplace after asynchronous execution starts without writing', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'red blue red', 'utf8'); - const controller = new AbortController(); - const execution = executeToolWithOptions( - collectTools(registerFileTools).get('StrReplace'), - { path: 'story.txt', old_str: 'red', new_str: 'green' }, - cwd, - { signal: controller.signal }, - ); - - queueMicrotask(() => controller.abort()); - - await expect(execution).rejects.toThrow(/aborted/i); - expect(readFileSync(join(cwd, 'story.txt'), 'utf8')).toBe('red blue red'); - }); - - it('does not lose updates from concurrent StrReplace mutations', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'alpha beta alpha', 'utf8'); - const replace = collectTools(registerFileTools).get('StrReplace'); - - await Promise.all([ - executeTool(replace, { path: 'story.txt', old_str: 'alpha', new_str: 'one' }, cwd), - executeTool(replace, { path: 'story.txt', old_str: 'beta', new_str: 'two' }, cwd), - ]); - - expect(readFileSync(join(cwd, 'story.txt'), 'utf8')).toBe('one two one'); - }); - - it('stores separate bounded StrReplace hunks for non-contiguous replacements', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync( - join(cwd, 'story.txt'), - ['target', ...Array.from({ length: 30 }, (_, index) => `middle ${index}`), 'target'].join( - '\n', - ), - 'utf8', - ); - - const result = await executeTool( - collectTools(registerFileTools).get('StrReplace'), - { path: 'story.txt', old_str: 'target', new_str: 'replacement\nline' }, - cwd, - ); - const details = result.details as { diff?: string; diffTruncated?: boolean }; - - expect(details.diff?.match(/^@@/gm)).toHaveLength(2); - expect(details.diffTruncated).toBe(false); - expect(JSON.stringify(result.details).length).toBeLessThan(32_000); - expect( - renderToolResult( - collectTools(registerFileTools).get('StrReplace'), - result, - { expanded: true, isPartial: false }, - { path: 'story.txt', old_str: 'target', new_str: 'replacement\nline' }, - ).length, - ).toBeLessThanOrEqual(32_000); - }); - - it('summarizes very large StrReplace diffs instead of persisting unbounded output', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync( - join(cwd, 'large.txt'), - Array.from({ length: 1_000 }, (_, index) => `target ${index}`).join('\n'), - 'utf8', - ); - - const result = await executeTool( - collectTools(registerFileTools).get('StrReplace'), - { path: 'large.txt', old_str: 'target', new_str: 'replacement with more text' }, - cwd, - ); - const details = result.details as { diff?: string; diffTruncated?: boolean }; - - expect(details.diffTruncated).toBe(true); - expect(details.diff?.length ?? 0).toBeLessThanOrEqual(30_000); - expect(JSON.stringify(result.details).length).toBeLessThan(32_000); - const replace = collectTools(registerFileTools).get('StrReplace'); - const args = { - path: 'large.txt', - old_str: 'target', - new_str: 'replacement with more text', - }; - const collapsed = renderToolResult( - replace, - result, - { expanded: false, isPartial: false }, - args, - ); - const expanded = renderToolResult(replace, result, { expanded: true, isPartial: false }, args); - expect(collapsed).toContain('[Diff preview truncated]'); - expect(collapsed.length).toBeLessThanOrEqual(4_100); - expect(expanded).toContain('[Diff truncated at 26000 characters]'); - expect(expanded.length).toBeLessThanOrEqual(27_000); - }); - - it('keeps native Write previews bounded and renders delegated errors', () => { - const write = collectTools(registerFileTools).get('Write'); - const args = { - file_path: 'notes.txt', - contents: Array.from({ length: 12 }, (_, index) => `line ${index + 1}`).join('\n'), - }; - - const collapsed = renderToolCall(write, args); - const expanded = renderToolCall(write, args, { expanded: true }); - const partial = renderToolCall(write, args, { isPartial: true }); - expect(collapsed).toContain('line 10'); - expect(collapsed).not.toContain('line 11'); - expect(expanded).toContain('line 12'); - expect(partial).toContain('line 1'); - expect( - renderToolResult( - write, - { - content: [{ type: 'text', text: 'write failed visibly' }], - details: undefined, - }, - { expanded: false, isPartial: false }, - args, - { isError: true }, - ), - ).toContain('write failed visibly'); - }); - - it('renders native Edit diffs in collapsed and expanded result states', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - writeFileSync(join(cwd, 'story.txt'), 'alpha beta', 'utf8'); - const edit = collectTools(registerFileTools).get('Edit'); - const args = { path: 'story.txt', oldText: 'alpha', newText: 'gamma' }; - const result = await executePreparedTool(edit, args, cwd); - - expect(renderToolResult(edit, result, { expanded: false, isPartial: false }, args)).toContain( - 'gamma beta', - ); - expect(renderToolResult(edit, result, { expanded: true, isPartial: false }, args)).toContain( - 'gamma beta', - ); - }); - - it('adopts native LS defaults, suffixes, sorting, limits, and errors', async () => { - const cwd = tempDir('pi-grok-cli-files-'); - mkdirSync(join(cwd, 'Folder')); - writeFileSync(join(cwd, '.hidden'), '', 'utf8'); - writeFileSync(join(cwd, 'alpha'), '', 'utf8'); - writeFileSync(join(cwd, 'Beta'), '', 'utf8'); - const ls = collectTools(registerFileTools).get('LS'); - - expect(firstText(await executeTool(ls, {}, cwd)).split('\n')).toEqual([ - '.hidden', - 'alpha', - 'Beta', - 'Folder/', - ]); - const limited = await executeTool(ls, { limit: 2 }, cwd); - expect(firstText(limited)).toContain('2 entries limit reached'); - expect(limited.details).toMatchObject({ entryLimitReached: 2 }); - await expect(executeTool(ls, { path: 'missing' }, cwd)).rejects.toThrow(/Path not found/); - - const empty = tempDir('pi-grok-cli-files-'); - expect(firstText(await executeTool(ls, {}, empty))).toBe('(empty directory)'); - }); -}); diff --git a/tests/tools/glob.test.ts b/tests/tools/glob.test.ts deleted file mode 100644 index a72593b..0000000 --- a/tests/tools/glob.test.ts +++ /dev/null @@ -1,164 +0,0 @@ -import { mkdirSync, symlinkSync, utimesSync, writeFileSync } from 'node:fs'; -import { join } from 'node:path'; -import { describe, expect, it } from 'vitest'; -import { registerSearchTools } from '../../src/tools/search.js'; -import { - collectTools, - executePreparedTool, - executeTool, - executeToolWithOptions, - firstText, - tempDir, -} from './toolTestHelpers.js'; - -function globTool() { - return collectTools(registerSearchTools).get('Glob'); -} - -function writeAt(path: string, contents: string, timestamp: string) { - writeFileSync(path, contents, 'utf8'); - const time = new Date(timestamp); - utimesSync(path, time, time); -} - -describe('Glob compatibility contract', () => { - it('returns files only in globally newest-first order with a deterministic path tie-break', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - mkdirSync(join(cwd, 'src')); - mkdirSync(join(cwd, 'src', 'directory.ts')); - writeAt(join(cwd, 'src', 'older.ts'), 'old', '2026-01-01T00:00:00Z'); - writeAt(join(cwd, 'src', 'z-tied.ts'), 'z', '2026-01-02T00:00:00Z'); - writeAt(join(cwd, 'src', 'a-tied.ts'), 'a', '2026-01-02T00:00:00Z'); - - const result = await executeTool(globTool(), { pattern: '**/*.ts', path: 'src' }, cwd); - - expect(firstText(result).split('\n')).toEqual([ - 'src/a-tied.ts', - 'src/z-tied.ts', - 'src/older.ts', - ]); - }); - - it('supports both Cursor path and pattern alias pairs', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - mkdirSync(join(cwd, 'nested')); - writeFileSync(join(cwd, 'nested', 'match.ts'), '', 'utf8'); - - const nativeNames = await executePreparedTool( - globTool(), - { pattern: '*.ts', path: 'nested' }, - cwd, - ); - const cursorNames = await executePreparedTool( - globTool(), - { glob_pattern: '*.ts', target_directory: 'nested' }, - cwd, - ); - - expect(firstText(nativeNames)).toBe('nested/match.ts'); - expect(firstText(cursorNames)).toBe('nested/match.ts'); - }); - - it('includes hidden files while honoring root and nested gitignore rules', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - mkdirSync(join(cwd, 'src')); - mkdirSync(join(cwd, 'ignored')); - writeFileSync(join(cwd, '.gitignore'), 'ignored/\nroot-skip.ts\n', 'utf8'); - writeFileSync(join(cwd, 'src', '.gitignore'), 'nested-skip.ts\n', 'utf8'); - writeFileSync(join(cwd, '.hidden.ts'), '', 'utf8'); - writeFileSync(join(cwd, 'root-skip.ts'), '', 'utf8'); - writeFileSync(join(cwd, 'src', 'keep.ts'), '', 'utf8'); - writeFileSync(join(cwd, 'src', 'nested-skip.ts'), '', 'utf8'); - writeFileSync(join(cwd, 'ignored', 'ignored.ts'), '', 'utf8'); - - const result = await executeTool(globTool(), { pattern: '**/*.ts' }, cwd); - - expect(firstText(result).split('\n').sort()).toEqual(['.hidden.ts', 'src/keep.ts']); - }); - - it('traverses an explicitly selected ignored root but applies rules inside it', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - mkdirSync(join(cwd, 'ignored')); - writeFileSync(join(cwd, '.gitignore'), 'ignored/\n', 'utf8'); - writeFileSync(join(cwd, 'ignored', '.gitignore'), 'inside-skip.ts\n', 'utf8'); - writeFileSync(join(cwd, 'ignored', 'keep.ts'), '', 'utf8'); - writeFileSync(join(cwd, 'ignored', 'inside-skip.ts'), '', 'utf8'); - - const result = await executeTool(globTool(), { pattern: '*.ts', path: 'ignored' }, cwd); - - expect(firstText(result)).toBe('ignored/keep.ts'); - }); - - it('never traverses .git, node_modules, or symlinked directories', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - const outside = tempDir('pi-grok-cli-glob-outside-'); - mkdirSync(join(cwd, '.git')); - mkdirSync(join(cwd, 'node_modules')); - writeFileSync(join(cwd, '.git', 'config.ts'), '', 'utf8'); - writeFileSync(join(cwd, 'node_modules', 'package.ts'), '', 'utf8'); - writeFileSync(join(outside, 'escaped.ts'), '', 'utf8'); - symlinkSync(outside, join(cwd, 'linked')); - writeFileSync(join(cwd, 'visible.ts'), '', 'utf8'); - - const result = await executeTool(globTool(), { pattern: '**/*.ts' }, cwd); - - expect(firstText(result)).toBe('visible.ts'); - }); - - it('keeps outside-workspace results absolute and valid for follow-up reads', async () => { - const cwd = tempDir('pi-grok-cli-glob-workspace-'); - const outside = tempDir('pi-grok-cli-glob-outside-'); - writeFileSync(join(outside, 'external.ts'), '', 'utf8'); - - const result = await executeTool(globTool(), { pattern: '*.ts', path: outside }, cwd); - - expect(firstText(result)).toBe(join(outside, 'external.ts')); - }); - - it('applies the result limit after global sorting and reports explicit bounded truncation', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - for (let index = 0; index < 200; index++) { - writeAt( - join(cwd, `${String(index).padStart(3, '0')}.ts`), - '', - new Date(Date.UTC(2026, 0, 1, 0, 0, index)).toISOString(), - ); - } - - const result = await executeTool(globTool(), { pattern: '*.ts', limit: 2 }, cwd); - - expect(firstText(result).split('\n').slice(0, 2)).toEqual(['199.ts', '198.ts']); - expect(firstText(result)).toContain('2 results limit reached'); - expect(result.details).toMatchObject({ resultLimitReached: 2 }); - expect(JSON.stringify(result).length).toBeLessThan(2_000); - }); - - it('throws for missing roots and cancellation', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - const controller = new AbortController(); - controller.abort(); - - await expect( - executeTool(globTool(), { pattern: '*.ts', path: 'missing' }, cwd), - ).rejects.toThrow(/not found/i); - await expect( - executeToolWithOptions(globTool(), { pattern: '**/*' }, cwd, { - signal: controller.signal, - }), - ).rejects.toThrow(/aborted/i); - }); - - it('cancels after asynchronous traversal has started', async () => { - const cwd = tempDir('pi-grok-cli-glob-'); - mkdirSync(join(cwd, 'nested')); - writeFileSync(join(cwd, 'nested', 'match.ts'), '', 'utf8'); - const controller = new AbortController(); - const result = executeToolWithOptions(globTool(), { pattern: '**/*.ts' }, cwd, { - signal: controller.signal, - }); - - queueMicrotask(() => controller.abort()); - - await expect(result).rejects.toThrow(/aborted/i); - }); -}); diff --git a/tests/tools/nativeAdapter.integration.test.ts b/tests/tools/nativeAdapter.integration.test.ts deleted file mode 100644 index 472b6b3..0000000 --- a/tests/tools/nativeAdapter.integration.test.ts +++ /dev/null @@ -1,224 +0,0 @@ -import { mkdirSync, writeFileSync } from 'node:fs'; -import { join } from 'node:path'; -import { AgentHarness, type AgentTool, InMemorySessionRepo } from '@earendil-works/pi-agent-core'; -import { NodeExecutionEnv } from '@earendil-works/pi-agent-core/node'; -import { - createModels, - fauxAssistantMessage, - fauxProvider, - fauxToolCall, -} from '@earendil-works/pi-ai'; -import type { - ExtensionAPI, - ExtensionContext, - ToolDefinition, -} from '@earendil-works/pi-coding-agent'; -import { describe, expect, it } from 'vitest'; -import { registerFileTools } from '../../src/tools/files.js'; -import { createReadShim } from '../../src/tools/read.js'; -import { registerSearchTools } from '../../src/tools/search.js'; -import { registerShellTool } from '../../src/tools/shell.js'; -import { tempDir } from './toolTestHelpers.js'; - -function getDefinition(register: (pi: ExtensionAPI) => void, name: string) { - const definitions = new Map(); - register({ - registerTool(definition: ToolDefinition) { - definitions.set(definition.name, definition); - }, - } as unknown as ExtensionAPI); - const definition = definitions.get(name); - if (!definition) throw new Error(`${name} was not registered`); - return definition; -} - -function bridgeDefinition(definition: ToolDefinition, cwd: string): AgentTool { - return { - name: definition.name, - label: definition.label, - description: definition.description, - parameters: definition.parameters, - prepareArguments: definition.prepareArguments, - executionMode: definition.executionMode, - execute(toolCallId, params, signal, onUpdate) { - return definition.execute(toolCallId, params, signal, onUpdate, { cwd } as ExtensionContext); - }, - } as AgentTool; -} - -async function executeThroughPi( - definition: ToolDefinition, - args: Record, - cwd: string, -) { - const faux = fauxProvider(); - const models = createModels(); - models.setProvider(faux.provider); - faux.setResponses([ - fauxAssistantMessage(fauxToolCall(definition.name, args), { stopReason: 'toolUse' }), - fauxAssistantMessage('done'), - ]); - const environment = new NodeExecutionEnv({ cwd }); - const harness = new AgentHarness({ - env: environment, - session: await new InMemorySessionRepo().create(), - models, - model: faux.getModel(), - tools: [bridgeDefinition(definition, cwd)], - }); - const results: { - isError: boolean; - result: { content: { type: string; text?: string }[] }; - }[] = []; - harness.subscribe((event) => { - if (event.type === 'tool_execution_end') results.push(event); - }); - - await harness.prompt('Run the requested tool.'); - await environment.cleanup(); - models.deleteProvider(faux.provider.id); - - const result = results[0]; - if (!result) throw new Error('No tool result was emitted'); - return result; -} - -function resultText(result: Awaited>) { - return result.result.content.find((content) => content.type === 'text')?.text ?? ''; -} - -describe('native adapter error-state integration', () => { - it('marks a delegated Write filesystem failure as isError', async () => { - const cwd = tempDir('pi-grok-cli-native-integration-'); - writeFileSync(join(cwd, 'blocked'), 'not a directory', 'utf8'); - - const result = await executeThroughPi( - getDefinition(registerFileTools, 'Write'), - { path: 'blocked/file.txt', content: 'content' }, - cwd, - ); - - expect(result.isError).toBe(true); - expect(resultText(result)).toMatch(/EEXIST|not a directory/i); - }); - - it('marks a StrReplace operational filesystem failure as isError', async () => { - const cwd = tempDir('pi-grok-cli-native-integration-'); - mkdirSync(join(cwd, 'directory')); - - const result = await executeThroughPi( - getDefinition(registerFileTools, 'StrReplace'), - { path: 'directory', old_str: 'old', new_str: 'new' }, - cwd, - ); - - expect(result.isError).toBe(true); - expect(resultText(result)).toMatch(/directory|EISDIR/i); - }); - - it('marks a nonzero delegated Shell command as isError with native text', async () => { - const cwd = tempDir('pi-grok-cli-native-integration-'); - - const result = await executeThroughPi( - getDefinition(registerShellTool, 'Shell'), - { command: 'printf problem >&2; exit 7' }, - cwd, - ); - - expect(result.isError).toBe(true); - expect(resultText(result)).toMatch(/problem.*exited with code 7/is); - }); - - it('keeps successful delegated calls as isError false', async () => { - const cwd = tempDir('pi-grok-cli-native-integration-'); - - const result = await executeThroughPi( - getDefinition(registerFileTools, 'Write'), - { path: 'success.txt', content: 'content' }, - cwd, - ); - - expect(result.isError).toBe(false); - expect(resultText(result)).toContain('Successfully wrote'); - }); -}); - -describe('uppercase shim agent-loop integration', () => { - it('executes every local uppercase shim with Cursor aliases through Pi', async () => { - const cwd = tempDir('pi-grok-cli-native-integration-'); - mkdirSync(join(cwd, 'src')); - writeFileSync(join(cwd, 'src', 'story.txt'), 'alpha needle\n', 'utf8'); - writeFileSync(join(cwd, 'delete-me.txt'), 'delete', 'utf8'); - const cases: { - args: Record; - expected: RegExp; - name: string; - register: (pi: ExtensionAPI) => void; - }[] = [ - { - name: 'Read', - register: (pi) => pi.registerTool(createReadShim()), - args: { file_path: 'src/story.txt' }, - expected: /alpha needle/, - }, - { - name: 'Write', - register: registerFileTools, - args: { file_path: 'written.txt', contents: 'written' }, - expected: /Successfully wrote/, - }, - { - name: 'Edit', - register: registerFileTools, - args: { path: 'src/story.txt', old_string: 'alpha', new_string: 'beta' }, - expected: /Successfully replaced/, - }, - { - name: 'StrReplace', - register: registerFileTools, - args: { path: 'src/story.txt', old_string: 'needle', new_string: 'match' }, - expected: /Replaced 1 occurrence/, - }, - { - name: 'LS', - register: registerFileTools, - args: { path: 'src' }, - expected: /story\.txt/, - }, - { - name: 'Grep', - register: registerSearchTools, - args: { query: 'match', path: 'src', include: '*.txt' }, - expected: /story\.txt.*match/s, - }, - { - name: 'Glob', - register: registerSearchTools, - args: { glob_pattern: '*.txt', target_directory: 'src' }, - expected: /src\/story\.txt/, - }, - { - name: 'Shell', - register: registerShellTool, - args: { cmd: 'printf shell-ready' }, - expected: /shell-ready/, - }, - { - name: 'Delete', - register: registerFileTools, - args: { path: 'delete-me.txt' }, - expected: /Successfully deleted/, - }, - ]; - - for (const testCase of cases) { - const result = await executeThroughPi( - getDefinition(testCase.register, testCase.name), - testCase.args, - cwd, - ); - expect(result.isError, testCase.name).toBe(false); - expect(resultText(result), testCase.name).toMatch(testCase.expected); - } - }); -}); diff --git a/tests/tools/read.test.ts b/tests/tools/read.test.ts deleted file mode 100644 index d3d76fa..0000000 --- a/tests/tools/read.test.ts +++ /dev/null @@ -1,109 +0,0 @@ -import { copyFileSync, writeFileSync } from 'node:fs'; -import { join } from 'node:path'; -import { describe, expect, it } from 'vitest'; -import { createReadShim } from '../../src/tools/read.js'; -import { - executePreparedTool, - executeTool, - firstText, - prepareToolArguments, - renderToolCall, - renderToolResult, - type ToolResult, - tempDir, -} from './toolTestHelpers.js'; - -type LooseTool = Parameters[0]; - -const readShim = (() => createReadShim())() as unknown as LooseTool; - -describe('Read shim (native read alias)', () => { - it('registers under the capital-Read name Cursor-trained models call', () => { - const shim = createReadShim(); - expect(shim.name).toBe('Read'); - expect(shim.label).toBe('Read'); - // Delegates the native definition's execute, parameters, and description. - expect(typeof shim.execute).toBe('function'); - expect(shim.description).toMatch(/Read the contents of a file/i); - }); - - it('reads a text file via the native read tool with line numbers', async () => { - const cwd = tempDir('pi-grok-cli-read-'); - writeFileSync(join(cwd, 'notes.txt'), 'alpha\nbeta\ngamma', 'utf-8'); - - const result = (await executeTool(readShim, { path: 'notes.txt' }, cwd)) as ToolResult; - - // Native read renders each line (1-indexed). - expect(firstText(result)).toContain('alpha'); - expect(firstText(result)).toContain('beta'); - expect(firstText(result)).toContain('gamma'); - }); - - it('normalizes Cursor file_path onto the native path parameter', () => { - const prepared = prepareToolArguments(readShim, { file_path: 'src/app.ts', offset: 10 }); - expect(prepared).toEqual({ path: 'src/app.ts', offset: 10, limit: undefined }); - }); - - it('passes path through unchanged when already in native shape', () => { - const prepared = prepareToolArguments(readShim, { path: 'README.md', limit: 50 }); - expect(prepared).toEqual({ path: 'README.md', offset: undefined, limit: 50 }); - }); - - it('normalizes missing and invalid arguments to safe defaults', () => { - const shim = createReadShim(); - - expect(shim.prepareArguments(null)).toEqual({ path: '' }); - expect(shim.prepareArguments({ path: 42, offset: '10', limit: false })).toEqual({ - path: '', - offset: undefined, - limit: undefined, - }); - }); - - it('executes through the prepared Cursor-style arguments', async () => { - const cwd = tempDir('pi-grok-cli-read-'); - writeFileSync(join(cwd, 'story.txt'), 'once upon a time', 'utf-8'); - - const result = (await executePreparedTool( - readShim, - { file_path: 'story.txt' }, - cwd, - )) as ToolResult; - expect(firstText(result)).toContain('once upon a time'); - }); - - it('uses native path/range call rendering and expanded text results', async () => { - const cwd = tempDir('pi-grok-cli-read-'); - writeFileSync(join(cwd, 'notes.txt'), 'alpha\nbeta\ngamma', 'utf8'); - const args = { file_path: 'notes.txt', offset: 2, limit: 1 }; - const result = (await executePreparedTool(readShim, args, cwd)) as ToolResult; - - expect(renderToolCall(readShim, args)).toBe('read notes.txt:2-2'); - expect(renderToolResult(readShim, result, { expanded: false, isPartial: false }, args)).toBe( - '', - ); - expect( - renderToolResult(readShim, result, { expanded: true, isPartial: false }, args), - ).toContain('beta'); - }); - - it('keeps native image result content intact', async () => { - const cwd = tempDir('pi-grok-cli-read-'); - copyFileSync( - join( - process.cwd(), - 'node_modules', - '@earendil-works', - 'pi-coding-agent', - 'docs', - 'images', - 'exy.png', - ), - join(cwd, 'pixel.png'), - ); - - const result = (await executeTool(readShim, { path: 'pixel.png' }, cwd)) as ToolResult; - - expect(result.content.some((content) => content.type === 'image')).toBe(true); - }); -}); diff --git a/tests/tools/register.test.ts b/tests/tools/register.test.ts deleted file mode 100644 index 6668fbc..0000000 --- a/tests/tools/register.test.ts +++ /dev/null @@ -1,41 +0,0 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { afterEach, describe, expect, it, vi } from 'vitest'; -import { GROK_SHIM_TOOL_NAMES, registerGrokTools } from '../../src/tools/register.js'; -import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; - -afterEach(() => { - vi.restoreAllMocks(); -}); - -describe('Grok tool registration', () => { - it('registers shim tools with renderers', () => { - vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(false); - const toolNames: string[] = []; - - registerGrokTools({ - registerTool(tool: { name: string; renderCall?: unknown; renderResult?: unknown }) { - toolNames.push(tool.name); - expect(tool.renderCall).toBeTypeOf('function'); - expect(tool.renderResult).toBeTypeOf('function'); - }, - on() {}, - } as unknown as ExtensionAPI); - - expect(toolNames.sort()).toEqual([...GROK_SHIM_TOOL_NAMES].sort()); - expect(toolNames).not.toContain('WebSearch'); - }); - - it('registers WebSearch when pi-web-access is installed', () => { - vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(true); - const toolNames: string[] = []; - - registerGrokTools({ - registerTool(tool: { name: string }) { - toolNames.push(tool.name); - }, - on() {}, - } as unknown as ExtensionAPI); - - expect(toolNames).toContain('WebSearch'); - }); -}); diff --git a/tests/tools/rendering.test.ts b/tests/tools/rendering.test.ts deleted file mode 100644 index 1e14922..0000000 --- a/tests/tools/rendering.test.ts +++ /dev/null @@ -1,84 +0,0 @@ -import { describe, expect, it } from 'vitest'; -import { - booleanDetail, - fileError, - fileNotFound, - nativeRenderContext, - numberDetail, - renderResultSummary, - renderRunning, - text, -} from '../../src/tools/rendering.js'; -import { renderText } from './toolTestHelpers.js'; - -describe('tool rendering helpers', () => { - it('clones native render contexts with normalized args and preserves renderer state', () => { - const state = { startedAt: 1 }; - const lastComponent = text('existing'); - const context = { - args: { file_path: 'raw.txt' }, - toolCallId: 'call-1', - invalidate: () => {}, - lastComponent, - state, - cwd: '/project', - executionStarted: true, - argsComplete: true, - isPartial: false, - expanded: true, - showImages: true, - isError: false, - }; - - const normalized = nativeRenderContext(context, { path: 'raw.txt' }); - - expect(normalized).toEqual({ ...context, args: { path: 'raw.txt' } }); - expect(normalized).not.toBe(context); - expect(normalized.state).toBe(state); - expect(normalized.lastComponent).toBe(lastComponent); - expect(context.args).toEqual({ file_path: 'raw.txt' }); - }); - - it('renders summaries, expanded text, missing text fallback, and partial state', () => { - const result = { - content: [{ type: 'text', text: 'full output' }], - details: {}, - }; - - expect(renderText(text('plain'))).toBe('plain'); - expect(renderText(renderResultSummary(result, false, false, 'summary'))).toBe('summary'); - expect(renderText(renderResultSummary(result, true, false, 'summary'))).toBe('full output'); - expect(renderText(renderRunning(true) ?? text(''))).toBe('Running...'); - expect(renderRunning(false)).toBeUndefined(); - expect(renderText(renderResultSummary(result, false, true, 'summary'))).toBe('Running...'); - }); - - it('reads typed detail values with defaults for absent or invalid details', () => { - const result = { - content: [{ type: 'text', text: '' }], - details: { count: 2, deleted: true, invalid: null }, - }; - - expect(numberDetail(result, 'count')).toBe(2); - expect(numberDetail(result, 'deleted')).toBe(0); - expect(numberDetail({ details: undefined }, 'count')).toBe(0); - expect(booleanDetail(result, 'deleted')).toBe(true); - expect(booleanDetail(result, 'invalid')).toBe(false); - }); - - it('formats retained Delete errors with stable details', () => { - expect(fileNotFound('/tmp/missing.txt', { deleted: false })).toEqual({ - content: [{ type: 'text', text: 'File not found: /tmp/missing.txt' }], - details: { path: '/tmp/missing.txt', deleted: false }, - }); - expect(fileError({}, 'Delete', '/tmp/file.txt', { deleted: false })).toEqual({ - content: [{ type: 'text', text: 'Delete error: Unknown error' }], - details: { - path: '/tmp/file.txt', - deleted: false, - failed: true, - error: 'Unknown error', - }, - }); - }); -}); diff --git a/tests/tools/search.test.ts b/tests/tools/search.test.ts deleted file mode 100644 index 55cab42..0000000 --- a/tests/tools/search.test.ts +++ /dev/null @@ -1,433 +0,0 @@ -import { mkdirSync, rmSync, utimesSync, writeFileSync } from 'node:fs'; -import { join } from 'node:path'; -import { describe, expect, it, vi } from 'vitest'; -import { registerSearchTools } from '../../src/tools/search.js'; -import { - collectTools, - executePreparedTool, - executeTool, - executeToolWithOptions, - firstText, - prepareToolArguments, - renderToolCall, - renderToolResult, - type ToolResult, - tempDir, -} from './toolTestHelpers.js'; - -function setupProject() { - const dir = tempDir('pi-grok-cli-search-'); - mkdirSync(join(dir, 'src')); - writeFileSync(join(dir, 'src', 'alpha.ts'), 'needle\nhaystack\n', 'utf-8'); - writeFileSync(join(dir, 'src', 'beta.md'), 'needle in docs\n', 'utf-8'); - writeFileSync(join(dir, 'src', 'gamma.ts'), 'plain text\n', 'utf-8'); - return dir; -} - -function expectGrepResult(result: ToolResult) { - expect(firstText(result)).toContain('alpha.ts:1: needle'); - expect(firstText(result)).not.toContain('beta.md'); - expect(result.details).toBeUndefined(); -} - -function expectGlobResult(result: ToolResult) { - expect(firstText(result)).toContain('src/alpha.ts'); - expect(firstText(result)).toContain('src/gamma.ts'); - expect(firstText(result)).not.toContain('beta.md'); - expect(result.details).toBeUndefined(); -} - -async function withNoSearchBinaries( - run: (tools: ReturnType) => Promise, -) { - const oldPath = process.env.PATH; - process.env.PATH = tempDir('pi-grok-cli-empty-bin-'); - vi.resetModules(); - try { - await run(collectTools((await import('../../src/tools/search.js')).registerSearchTools)); - } finally { - process.env.PATH = oldPath; - vi.resetModules(); - } -} - -describe('search tools', () => { - it('greps matching file contents with include filters', async () => { - const cwd = setupProject(); - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'needle', path: 'src', include: '*.ts' }, - cwd, - ); - - expectGrepResult(result); - }); - - it('greps matching file contents with Cursor-style glob filters', async () => { - const cwd = setupProject(); - const result = await executePreparedTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'needle', path: 'src', glob_filter: '*.ts' }, - cwd, - ); - - expectGrepResult(result); - }); - - it('greps patterns that start with a dash', async () => { - const cwd = setupProject(); - writeFileSync(join(cwd, 'src', 'dash.ts'), '-export const value = 1\n', 'utf-8'); - - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: '-export', path: 'src/dash.ts' }, - cwd, - ); - - expect(firstText(result)).toBe('dash.ts:1: -export const value = 1'); - expect(result.details).toBeUndefined(); - }); - - it('includes file paths when grepping a single file', async () => { - const cwd = setupProject(); - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'needle', path: 'src/alpha.ts' }, - cwd, - ); - - expect(firstText(result)).toBe('alpha.ts:1: needle'); - expect(result.details).toBeUndefined(); - }); - - it('reports no grep matches as an empty result', async () => { - const cwd = setupProject(); - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'absent', path: 'src' }, - cwd, - ); - - expect(firstText(result)).toBe('No matches found'); - expect(result.details).toBeUndefined(); - }); - - it('reports grep command errors with empty match details', async () => { - const cwd = setupProject(); - await expect( - executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: '[', path: 'src' }, - cwd, - ), - ).rejects.toThrow(/regex parse error|Invalid regular expression/); - }); - - it('globs files under the requested path', async () => { - const cwd = setupProject(); - const result = await executeTool( - collectTools(registerSearchTools).get('Glob'), - { pattern: '**/*.ts', path: 'src' }, - cwd, - ); - - expectGlobResult(result); - }); - - it('globs files with Cursor-style glob pattern arguments', async () => { - const cwd = setupProject(); - const result = await executePreparedTool( - collectTools(registerSearchTools).get('Glob'), - { glob_pattern: '**/*.ts', path: 'src' }, - cwd, - ); - - expectGlobResult(result); - }); - - it('reports empty glob command results', async () => { - const cwd = setupProject(); - const result = await executeTool( - collectTools(registerSearchTools).get('Glob'), - { pattern: '**/*.json', path: 'src' }, - cwd, - ); - - expect(firstText(result)).toBe('No files found'); - expect(result.details).toBeUndefined(); - }); - - it('reports glob filesystem errors', async () => { - const cwd = setupProject(); - await expect( - executeTool( - collectTools(registerSearchTools).get('Glob'), - { pattern: '**/*.ts', path: 'missing' }, - cwd, - ), - ).rejects.toThrow(/Path not found/); - }); - - it('globs path-containing patterns without external search binaries', async () => { - const cwd = setupProject(); - await withNoSearchBinaries(async (fallbackTools) => { - const result = await executeTool(fallbackTools.get('Glob'), { pattern: 'src/**/*.ts' }, cwd); - - expectGlobResult(result); - }); - }); - - it('globs basename-only patterns relative to an explicit root', async () => { - const cwd = setupProject(); - await withNoSearchBinaries(async (fallbackTools) => { - const result = await executeTool( - fallbackTools.get('Glob'), - { pattern: '*.ts', path: 'src' }, - cwd, - ); - - expectGlobResult(result); - }); - }); - - it('globs files without ripgrep or Unix find on PATH', async () => { - const cwd = setupProject(); - writeFileSync(join(cwd, '.hidden.ts'), '', 'utf8'); - await withNoSearchBinaries(async (fallbackTools) => { - const result = await executeTool( - fallbackTools.get('Glob'), - { pattern: '**/*.ts', path: '.' }, - cwd, - ); - - expect(firstText(result)).toContain('.hidden.ts'); - expect(firstText(result)).toContain('src/alpha.ts'); - }); - }); - - it('greps literal metacharacters through the native engine', async () => { - const cwd = setupProject(); - writeFileSync(join(cwd, 'src', 'literal.ts'), 'needle.[\n', 'utf8'); - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'needle.[', path: 'src', glob: '*.ts', literal: true }, - cwd, - ); - - expect(firstText(result)).toBe('literal.ts:1: needle.['); - }); - - it('greps with path-containing native glob filters', async () => { - const cwd = setupProject(); - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'needle', path: 'src', glob: '**/*.ts' }, - cwd, - ); - - expect(firstText(result)).toBe('alpha.ts:1: needle'); - }); - - it('sorts glob results by modification time newest first', async () => { - const cwd = setupProject(); - const oldTime = new Date('2024-01-01T00:00:00.000Z'); - const newTime = new Date('2024-01-02T00:00:00.000Z'); - utimesSync(join(cwd, 'src', 'alpha.ts'), oldTime, oldTime); - utimesSync(join(cwd, 'src', 'gamma.ts'), newTime, newTime); - const result = await executeTool( - collectTools(registerSearchTools).get('Glob'), - { pattern: '**/*.ts', path: 'src' }, - cwd, - ); - - expect(firstText(result).split('\n')).toEqual(['src/gamma.ts', 'src/alpha.ts']); - }); - - it('omits glob candidates deleted before traversal', async () => { - const cwd = setupProject(); - const deleted = join(cwd, 'src', 'deleted.ts'); - writeFileSync(deleted, 'deleted\n', 'utf-8'); - rmSync(deleted); - - const result = await executeTool( - collectTools(registerSearchTools).get('Glob'), - { pattern: '**/*.ts', path: 'src' }, - cwd, - ); - - expect(firstText(result)).not.toContain('deleted.ts'); - }); - - it('breaks modification time ties with alphabetical ordering', async () => { - const cwd = setupProject(); - const sameTime = new Date('2024-06-01T00:00:00.000Z'); - utimesSync(join(cwd, 'src', 'gamma.ts'), sameTime, sameTime); - utimesSync(join(cwd, 'src', 'alpha.ts'), sameTime, sameTime); - - const result = await executeTool( - collectTools(registerSearchTools).get('Glob'), - { pattern: '**/*.ts', path: 'src' }, - cwd, - ); - - expect(firstText(result).split('\n')).toEqual(['src/alpha.ts', 'src/gamma.ts']); - }); - - it('renders grep calls and result states', () => { - const grep = collectTools(registerSearchTools).get('Grep'); - const result = { - content: [{ type: 'text', text: 'src/alpha.ts:1: needle' }], - details: undefined, - }; - - expect(renderToolCall(grep, { pattern: 'needle', path: 'src', include: '*.ts' })).toBe( - 'grep /needle/ in src (*.ts)', - ); - expect(renderToolResult(grep, result)).toContain('src/alpha.ts:1: needle'); - expect(renderToolResult(grep, result, { expanded: true, isPartial: false })).toContain( - 'src/alpha.ts:1: needle', - ); - expect( - renderToolResult(grep, { - content: [{ type: 'text', text: 'No matches found' }], - details: undefined, - }), - ).toContain('No matches found'); - expect(renderToolResult(grep, result, { expanded: false, isPartial: true })).toContain( - 'src/alpha.ts:1: needle', - ); - }); - - it('renders glob calls and result states', () => { - const glob = collectTools(registerSearchTools).get('Glob'); - const result = { - content: [{ type: 'text', text: 'src/alpha.ts\nsrc/gamma.ts' }], - details: undefined, - }; - - expect(renderToolCall(glob, { pattern: '**/*.ts', path: 'src' })).toBe('find **/*.ts in src'); - expect(renderToolResult(glob, result)).toContain('src/alpha.ts'); - expect( - renderToolResult(glob, { - content: [{ type: 'text', text: 'No files found' }], - details: undefined, - }), - ).toContain('No files found'); - expect(renderToolResult(glob, result, { expanded: false, isPartial: true })).toContain( - 'src/gamma.ts', - ); - }); -}); - -describe('native Grep adapter contracts', () => { - it('normalizes query and all observed glob aliases', () => { - const grep = collectTools(registerSearchTools).get('Grep'); - - expect( - prepareToolArguments(grep, { - query: 'needle', - path: 'src', - include: '*.ts', - ignore_case: true, - }), - ).toEqual({ - pattern: 'needle', - path: 'src', - glob: '*.ts', - ignoreCase: true, - literal: undefined, - context: undefined, - limit: undefined, - }); - expect(prepareToolArguments(grep, { pattern: 'x', glob_filter: '**/*.md' })).toMatchObject({ - pattern: 'x', - glob: '**/*.md', - }); - }); - - it('delegates literal, ignore-case, context, and limit behavior to native grep', async () => { - const cwd = tempDir('pi-grok-cli-search-'); - writeFileSync(join(cwd, 'sample.txt'), 'before\nNeedle.[\nafter\nneedle.[ again\n', 'utf8'); - - const result = await executePreparedTool( - collectTools(registerSearchTools).get('Grep'), - { - query: 'needle.[', - path: '.', - include: '*.txt', - ignoreCase: true, - literal: true, - context: 1, - limit: 1, - }, - cwd, - ); - - expect(firstText(result)).toContain('sample.txt:2: Needle.['); - expect(firstText(result)).toContain('sample.txt-1- before'); - expect(firstText(result)).toContain('1 matches limit reached'); - expect(result.details).toMatchObject({ matchLimitReached: 1 }); - }); - - it('uses native .gitignore handling', async () => { - const cwd = tempDir('pi-grok-cli-search-'); - mkdirSync(join(cwd, '.git')); - writeFileSync(join(cwd, '.gitignore'), 'ignored.txt\n', 'utf8'); - writeFileSync(join(cwd, 'ignored.txt'), 'needle\n', 'utf8'); - writeFileSync(join(cwd, 'visible.txt'), 'haystack\n', 'utf8'); - - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'needle', path: '.' }, - cwd, - ); - - expect(firstText(result)).toBe('No matches found'); - expect(result.details).toBeUndefined(); - }); - - it('reports native long-line and result truncation metadata', async () => { - const cwd = tempDir('pi-grok-cli-search-'); - writeFileSync(join(cwd, 'long.txt'), `needle ${'x'.repeat(700)}\n`, 'utf8'); - - const result = await executeTool( - collectTools(registerSearchTools).get('Grep'), - { pattern: 'needle', path: '.' }, - cwd, - ); - - expect(firstText(result)).toContain('Some lines truncated'); - expect(result.details).toMatchObject({ linesTruncated: true }); - }); - - it('throws native cancellation and invalid-regex errors', async () => { - const cwd = setupProject(); - const grep = collectTools(registerSearchTools).get('Grep'); - const controller = new AbortController(); - controller.abort(); - - await expect( - executeToolWithOptions(grep, { pattern: 'needle', path: 'src' }, cwd, { - signal: controller.signal, - }), - ).rejects.toThrow(/aborted/i); - await expect(executeTool(grep, { pattern: '[', path: 'src' }, cwd)).rejects.toThrow( - /regex parse error|invalid regular expression/i, - ); - }); - - it('uses native collapsed and expanded rendering with normalized aliases', async () => { - const cwd = setupProject(); - const grep = collectTools(registerSearchTools).get('Grep'); - const args = { query: 'needle', path: 'src', include: '*.ts' }; - const result = await executePreparedTool(grep, args, cwd); - - expect(renderToolCall(grep, args)).toBe('grep /needle/ in src (*.ts)'); - expect(renderToolResult(grep, result, { expanded: false, isPartial: false }, args)).toContain( - 'alpha.ts:1: needle', - ); - expect(renderToolResult(grep, result, { expanded: true, isPartial: false }, args)).toContain( - 'alpha.ts:1: needle', - ); - }); -}); diff --git a/tests/tools/shell.test.ts b/tests/tools/shell.test.ts deleted file mode 100644 index a8814ef..0000000 --- a/tests/tools/shell.test.ts +++ /dev/null @@ -1,295 +0,0 @@ -import { existsSync, writeFileSync } from 'node:fs'; -import { join } from 'node:path'; -import { describe, expect, it } from 'vitest'; -import { registerShellTool } from '../../src/tools/shell.js'; -import { - collectTools, - executePreparedTool, - executeTool, - executeToolWithOptions, - firstText, - prepareToolArguments, - renderToolCall, - renderToolResult, - tempDir, -} from './toolTestHelpers.js'; - -describe('shell tool', () => { - it('returns stdout and stderr in native stream order', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const result = await executeTool( - collectTools(registerShellTool).get('Shell'), - { command: 'printf stdout && printf stderr >&2' }, - cwd, - ); - - expect(firstText(result)).toBe('stdoutstderr'); - expect(result.details).toBeUndefined(); - }); - - it('runs commands in a resolved working directory', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - writeFileSync(join(cwd, 'target.txt'), 'from cwd', 'utf-8'); - const result = await executeTool( - collectTools(registerShellTool).get('Shell'), - { command: 'cat target.txt', working_directory: '.' }, - cwd, - ); - - expect(firstText(result)).toBe('from cwd'); - expect(result.details).toBeUndefined(); - }); - - it('returns a clear placeholder when commands produce no output', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const result = await executeTool( - collectTools(registerShellTool).get('Shell'), - { command: 'true' }, - cwd, - ); - - expect(firstText(result)).toBe('(no output)'); - expect(result.details).toBeUndefined(); - }); - - it('includes exit code, error message, and captured output on failure', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - await expect( - executeTool( - collectTools(registerShellTool).get('Shell'), - { command: 'printf before && printf problem >&2 && exit 7' }, - cwd, - ), - ).rejects.toThrow(/before.*problem.*Command exited with code 7/is); - }); - - it('truncates large successful and failed output', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const tools = collectTools(registerShellTool); - const largeOutput = 'node -e "for(let i=0;i<2500;i++) console.log(i)"'; - - const successResult = await executeTool(tools.get('Shell'), { command: largeOutput }, cwd); - const details = successResult.details as { - truncation?: { truncated: boolean }; - fullOutputPath?: string; - }; - - expect(details.truncation?.truncated).toBe(true); - expect(firstText(successResult)).toContain('Full output:'); - await expect( - executeTool( - tools.get('Shell'), - { command: 'node -e "for(let i=0;i<2500;i++) console.log(i); process.exit(9)"' }, - cwd, - ), - ).rejects.toThrow(/Command exited with code 9/); - }); - - it('truncates multibyte output by characters without hitting exec buffer limits', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const result = await executeTool( - collectTools(registerShellTool).get('Shell'), - { command: 'perl -e \'print "漢" x 50001\'' }, - cwd, - ); - - expect((result.details as { truncation?: { truncated: boolean } }).truncation?.truncated).toBe( - true, - ); - expect(firstText(result)).toContain('Full output:'); - expect(firstText(result).length).toBeLessThan(50_001); - }); - - it('kills commands that exceed the timeout', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const command = 'node -e "setTimeout(()=>{},10000)"'; - await expect( - executeTool(collectTools(registerShellTool).get('Shell'), { command, timeout: 100 }, cwd), - ).rejects.toThrow(/Command timed out after 0\.1 seconds/); - }); - - it('renders shell calls and result states', () => { - const shell = collectTools(registerShellTool).get('Shell'); - - expect( - renderToolCall(shell, { - command: 'pwd', - working_directory: 'src', - }), - ).toBe('$ pwd (timeout 120s)'); - expect(renderToolCall(shell, { command: 'pwd' })).toBe('$ pwd (timeout 120s)'); - expect( - renderToolResult(shell, { - content: [{ type: 'text', text: 'full output' }], - details: { exitCode: 0 }, - }), - ).toContain('full output'); - expect( - renderToolResult(shell, { - content: [{ type: 'text', text: 'spawn failed' }], - details: { exitCode: 'ENOENT' }, - }), - ).toContain('spawn failed'); - expect( - renderToolResult( - shell, - { - content: [{ type: 'text', text: 'full output' }], - details: { exitCode: 0 }, - }, - { expanded: true, isPartial: false }, - ), - ).toContain('full output'); - expect( - renderToolResult( - shell, - { - content: [{ type: 'text', text: 'still running' }], - details: { exitCode: 0 }, - }, - { expanded: false, isPartial: true }, - ), - ).toContain('still running'); - }); -}); - -describe('native Shell adapter contracts', () => { - it('normalizes cmd and preserves the 120-second compatibility default', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const shell = collectTools(registerShellTool).get('Shell'); - - expect(prepareToolArguments(shell, { cmd: 'printf ready' })).toEqual({ - command: 'printf ready', - timeout: 120_000, - working_directory: undefined, - }); - const result = await executePreparedTool(shell, { cmd: 'printf ready' }, cwd); - expect(firstText(result)).toBe('ready'); - }); - - it('converts millisecond timeouts to native seconds', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - - await expect( - executeTool( - collectTools(registerShellTool).get('Shell'), - { command: 'node -e "setTimeout(()=>{},10000)"', timeout: 100 }, - cwd, - ), - ).rejects.toThrow(/timed out after 0\.1 seconds/i); - }); - - it('streams native stdout and stderr updates', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const updates: string[] = []; - - const result = await executeToolWithOptions( - collectTools(registerShellTool).get('Shell'), - { - command: - "node -e \"process.stdout.write('first'); setTimeout(()=>process.stderr.write('second'),150)\"", - }, - cwd, - { onUpdate: (update) => updates.push(firstText(update)) }, - ); - - expect(updates.some((update) => update.includes('first'))).toBe(true); - expect(firstText(result)).toContain('first'); - expect(firstText(result)).toContain('second'); - }); - - it('throws native cancellation and nonzero-exit errors', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const shell = collectTools(registerShellTool).get('Shell'); - const controller = new AbortController(); - const execution = executeToolWithOptions( - shell, - { command: 'node -e "setTimeout(()=>{},10000)"' }, - cwd, - { signal: controller.signal }, - ); - setTimeout(() => controller.abort(), 50); - - await expect(execution).rejects.toThrow(/Command aborted/i); - await expect( - executeTool(shell, { command: 'printf before && printf problem >&2 && exit 7' }, cwd), - ).rejects.toThrow(/before.*problem.*exited with code 7/is); - }); - - it('kills descendant processes when native cancellation aborts Shell', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const marker = join(cwd, 'descendant-survived.txt'); - const child = `setTimeout(() => require('node:fs').writeFileSync(${JSON.stringify(marker)}, 'alive'), 500)`; - const parent = `require('node:child_process').spawn(process.execPath, ['-e', ${JSON.stringify(child)}], { stdio: 'ignore' }); setTimeout(() => {}, 10000)`; - const controller = new AbortController(); - const execution = executeToolWithOptions( - collectTools(registerShellTool).get('Shell'), - { command: `node -e ${JSON.stringify(parent)}` }, - cwd, - { signal: controller.signal }, - ); - setTimeout(() => controller.abort(), 100); - - await expect(execution).rejects.toThrow(/Command aborted/i); - await new Promise((resolve) => setTimeout(resolve, 600)); - expect(existsSync(marker)).toBe(false); - }); - - it('retains a full-output file when native truncation occurs', async () => { - const cwd = tempDir('pi-grok-cli-shell-'); - const result = await executeTool( - collectTools(registerShellTool).get('Shell'), - { command: 'node -e "for(let i=0;i<2500;i++) console.log(i)"' }, - cwd, - ); - const details = result.details as { - fullOutputPath?: string; - truncation?: { truncated: boolean }; - }; - - expect(details.truncation?.truncated).toBe(true); - expect(details.fullOutputPath).toBeTypeOf('string'); - expect(existsSync(details.fullOutputPath ?? '')).toBe(true); - expect(firstText(result)).toContain('Full output:'); - }); - - it('uses native Shell call and collapsed/expanded result rendering', () => { - const shell = collectTools(registerShellTool).get('Shell'); - const args = { cmd: 'pwd', timeout: 100 }; - const result = { - content: [{ type: 'text', text: 'one\ntwo\nthree\nfour\nfive\nsix' }], - details: {}, - }; - - expect(renderToolCall(shell, args)).toContain('$ pwd (timeout 0.1s)'); - const collapsed = renderToolResult(shell, result, { expanded: false, isPartial: false }, args); - const expanded = renderToolResult(shell, result, { expanded: true, isPartial: false }, args); - expect(collapsed).toContain('six'); - expect(collapsed).not.toContain('one'); - expect(expanded).toContain('one'); - expect(expanded).toContain('six'); - }); - - it('keeps native partial and error output visible', () => { - const shell = collectTools(registerShellTool).get('Shell'); - const args = { command: 'printf working' }; - - expect( - renderToolResult( - shell, - { content: [{ type: 'text', text: 'streaming output' }], details: undefined }, - { expanded: false, isPartial: true }, - args, - ), - ).toContain('streaming output'); - expect( - renderToolResult( - shell, - { content: [{ type: 'text', text: 'Command exited with code 7' }], details: undefined }, - { expanded: false, isPartial: false }, - args, - { isError: true }, - ), - ).toContain('Command exited with code 7'); - }); -}); diff --git a/tests/tools/toolTestHelpers.ts b/tests/tools/toolTestHelpers.ts deleted file mode 100644 index 808ada7..0000000 --- a/tests/tools/toolTestHelpers.ts +++ /dev/null @@ -1,224 +0,0 @@ -import { mkdtempSync, rmSync } from 'node:fs'; -import { tmpdir } from 'node:os'; -import { join } from 'node:path'; -import { stripVTControlCharacters } from 'node:util'; -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { afterEach } from 'vitest'; - -const tempDirs: string[] = []; - -afterEach(() => { - for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true }); -}); - -export type ToolResult = { - content: { type: string; text?: string }[]; - details: Record | undefined; -}; - -type ExtensionHandler = (event: unknown) => unknown; - -type Renderable = { render: (width: number) => string[] }; - -type ToolTheme = { - bold: (text: string) => string; - bg: (name: string, text: string) => string; - fg: (name: string, text: string) => string; - inverse: (text: string) => string; -}; - -type RenderContextOptions = { - argsComplete?: boolean; - cwd?: string; - executionStarted?: boolean; - expanded?: boolean; - isError?: boolean; - isPartial?: boolean; - showImages?: boolean; -}; - -type RegisteredTool = { - name: string; - prepareArguments?: (params: Record) => Record; - execute: ( - toolCallId: string, - params: Record, - signal: AbortSignal, - onUpdate: (result: ToolResult) => void, - ctx: { cwd: string }, - ) => Promise; - renderCall?: ( - args: Record, - theme: ToolTheme, - context: ToolRenderContext, - ) => Renderable; - renderResult?: ( - result: ToolResult, - state: { expanded: boolean; isPartial: boolean }, - theme: ToolTheme, - context: ToolRenderContext, - ) => Renderable; -}; - -type ToolRenderContext = { - args: Record; - toolCallId: string; - invalidate: () => void; - lastComponent: Renderable | undefined; - state: Record; - cwd: string; - executionStarted: boolean; - argsComplete: boolean; - isPartial: boolean; - expanded: boolean; - showImages: boolean; - isError: boolean; -}; - -export function collectTools(registerTools: (pi: ExtensionAPI) => void) { - const tools = new Map(); - const handlers = new Map(); - registerTools({ - registerTool(tool: RegisteredTool) { - tools.set(tool.name, tool); - }, - on(event: string, handler: ExtensionHandler) { - handlers.set(event, handler); - }, - } as unknown as ExtensionAPI); - return tools; -} - -export async function executeTool( - tool: RegisteredTool | undefined, - params: Record, - cwd: string, -) { - return executeToolWithOptions(tool, params, cwd); -} - -export async function executeToolWithOptions( - tool: RegisteredTool | undefined, - params: Record, - cwd: string, - options: { - signal?: AbortSignal; - onUpdate?: (result: ToolResult) => void; - } = {}, -) { - if (!tool) throw new Error('Tool was not registered'); - return tool.execute( - 'tool-call-id', - params, - options.signal ?? new AbortController().signal, - options.onUpdate ?? (() => {}), - { cwd }, - ); -} - -export function prepareToolArguments( - tool: RegisteredTool | undefined, - params: Record, -) { - if (!tool) throw new Error('Tool was not registered'); - return tool.prepareArguments?.(params) ?? params; -} - -export async function executePreparedTool( - tool: RegisteredTool | undefined, - params: Record, - cwd: string, -) { - if (!tool) throw new Error('Tool was not registered'); - return executeTool(tool, prepareToolArguments(tool, params), cwd); -} - -export function firstText(result: ToolResult) { - return result.content[0]?.text ?? ''; -} - -export function renderText(component: { render: (width: number) => string[] }) { - return component - .render(120) - .map((line) => stripVTControlCharacters(line).trimEnd()) - .join('\n'); -} - -export const plainTheme = { - bold: (text: string) => text, - bg: (_name: string, text: string) => text, - fg: (_name: string, text: string) => text, - inverse: (text: string) => text, -}; - -Object.assign(globalThis, { - [Symbol.for('@earendil-works/pi-coding-agent:theme')]: plainTheme, - [Symbol.for('@mariozechner/pi-coding-agent:theme')]: plainTheme, -}); - -export function renderToolCall( - tool: RegisteredTool | undefined, - args: Record, - options: RenderContextOptions = {}, -) { - if (!tool?.renderCall) throw new Error('Tool call renderer was not registered'); - return renderText( - tool.renderCall( - args, - plainTheme, - renderContext( - args, - { - expanded: options.expanded ?? false, - isPartial: options.isPartial ?? false, - }, - undefined, - options, - ), - ), - ); -} - -export function renderToolResult( - tool: RegisteredTool | undefined, - result: ToolResult, - state = { expanded: false, isPartial: false }, - args: Record = {}, - options: RenderContextOptions = {}, - theme: ToolTheme = plainTheme, -) { - if (!tool?.renderResult) { - throw new Error('Tool result renderer was not registered'); - } - return renderText( - tool.renderResult(result, state, theme, renderContext(args, state, result, options)), - ); -} - -function renderContext( - args: Record, - options = { expanded: false, isPartial: false }, - result?: ToolResult, - overrides: RenderContextOptions = {}, -): ToolRenderContext { - return { - args, - toolCallId: 'tool-call-id', - invalidate: () => {}, - lastComponent: undefined, - state: {}, - cwd: overrides.cwd ?? '/project', - executionStarted: overrides.executionStarted ?? result !== undefined, - argsComplete: overrides.argsComplete ?? true, - isPartial: options.isPartial, - expanded: options.expanded, - showImages: overrides.showImages ?? true, - isError: overrides.isError ?? false, - }; -} - -export function tempDir(prefix: string) { - const dir = mkdtempSync(join(tmpdir(), prefix)); - tempDirs.push(dir); - return dir; -} diff --git a/tests/tools/webSearch.test.ts b/tests/tools/webSearch.test.ts deleted file mode 100644 index 031b9ef..0000000 --- a/tests/tools/webSearch.test.ts +++ /dev/null @@ -1,220 +0,0 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { afterEach, describe, expect, it, vi } from 'vitest'; -import { registerGrokTools } from '../../src/tools/register.js'; -import { registerWebSearchTool } from '../../src/tools/webSearch.js'; -import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; -import { - clearWebSearchDelegateForTests, - setWebSearchDelegateForTests, -} from '../../src/tools/webSearchDelegate.js'; -import { - collectTools, - executeTool, - firstText, - renderToolCall, - renderToolResult, -} from './toolTestHelpers.js'; - -afterEach(() => { - clearWebSearchDelegateForTests(); -}); - -describe('WebSearch tool', () => { - it('registers WebSearch with renderers', () => { - const names: string[] = []; - registerWebSearchTool({ - registerTool(tool: { name: string; renderCall?: unknown; renderResult?: unknown }) { - names.push(tool.name); - expect(tool.renderCall).toBeTypeOf('function'); - expect(tool.renderResult).toBeTypeOf('function'); - }, - on() {}, - } as unknown as ExtensionAPI); - - expect(names).toContain('WebSearch'); - }); - - it('delegates execute to captured web_search', async () => { - setWebSearchDelegateForTests(async (_id, params) => ({ - content: [{ type: 'text', text: `delegated:${JSON.stringify(params)}` }], - details: { delegated: true }, - })); - - const tools = collectTools(registerWebSearchTool); - const result = await executeTool(tools.get('WebSearch'), { query: 'pi extensions' }, '/tmp'); - expect(firstText(result)).toBe('delegated:{"query":"pi extensions"}'); - expect(result.details).toEqual({ delegated: true }); - }); - - it('normalizes delegated queries', async () => { - setWebSearchDelegateForTests(async (_id, params) => ({ - content: [{ type: 'text', text: JSON.stringify(params) }], - details: {}, - })); - - const tools = collectTools(registerWebSearchTool); - const result = await executeTool( - tools.get('WebSearch'), - { - query: ' ', - queries: [' first query ', ' ', 'second query'], - numResults: 3, - }, - '/tmp', - ); - - expect(firstText(result)).toBe( - JSON.stringify({ - queries: ['first query', 'second query'], - numResults: 3, - }), - ); - }); - - it('reports missing pi-web-access when delegate was never captured', async () => { - vi.spyOn(webSearchDelegate, 'ensureWebSearchDelegate').mockResolvedValue(undefined); - vi.spyOn(webSearchDelegate, 'getWebSearchDelegate').mockReturnValue(undefined); - vi.spyOn(webSearchDelegate, 'getWebSearchLoadError').mockReturnValue( - 'pi-web-access is not installed. Run: pi install npm:pi-web-access', - ); - - const tools = collectTools(registerWebSearchTool); - const result = await executeTool(tools.get('WebSearch'), { query: 'test' }, '/tmp'); - expect(firstText(result)).toMatch(/pi-web-access|pi install npm:pi-web-access/i); - - vi.restoreAllMocks(); - }); - - it('registerGrokTools skips WebSearch when pi-web-access is not installed', () => { - vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(false); - const names: string[] = []; - registerGrokTools({ - registerTool(tool: { name: string }) { - names.push(tool.name); - }, - on() {}, - } as unknown as ExtensionAPI); - expect(names).not.toContain('WebSearch'); - vi.restoreAllMocks(); - }); - - it('renderCall shows WebSearch title', () => { - const tools = collectTools(registerWebSearchTool); - const line = renderToolCall(tools.get('WebSearch'), { query: 'hello world' }); - expect(line).toContain('WebSearch'); - expect(line).toContain('hello world'); - }); - - describe('renderCall', () => { - const tools = collectTools(registerWebSearchTool); - const ws = tools.get('WebSearch'); - - it('shows (no query) when query list is empty', () => { - expect(renderToolCall(ws, { queries: [] })).toContain('(no query)'); - }); - - it('shows (no query) with no query or queries param', () => { - expect(renderToolCall(ws, {})).toContain('(no query)'); - }); - - it('truncates a single long query', () => { - const long = 'a'.repeat(61); - const line = renderToolCall(ws, { query: long }); - expect(line).toContain('...'); - expect(line).not.toContain(long); - }); - - it('shows query count and list for multiple queries', () => { - const line = renderToolCall(ws, { queries: ['q1', 'q2', 'q3'] }); - expect(line).toContain('3 queries'); - expect(line).toContain('"q1"'); - expect(line).toContain('"q2"'); - expect(line).toContain('"q3"'); - }); - - it('shows truncation indicator for more than 5 queries', () => { - const queries = ['a', 'b', 'c', 'd', 'e', 'f', 'g']; - const line = renderToolCall(ws, { queries }); - expect(line).toContain('7 queries'); - expect(line).toContain('... and 2 more'); - }); - }); - - describe('renderResult', () => { - const tools = collectTools(registerWebSearchTool); - const ws = tools.get('WebSearch'); - - it('shows running state when partial', () => { - const result = { - content: [{ type: 'text', text: 'loading...' }], - details: {}, - }; - expect(renderToolResult(ws, result, { expanded: false, isPartial: true })).toBe('Running...'); - }); - - it('shows error message when details contain error', () => { - const result = { - content: [{ type: 'text', text: 'something broke' }], - details: { error: 'API key is invalid' }, - }; - const rendered = renderToolResult(ws, result); - expect(rendered).toContain('Error:'); - expect(rendered).toContain('API key is invalid'); - }); - - it('shows source count when totalResults is present', () => { - const result = { - content: [{ type: 'text', text: 'synthesized answer' }], - details: { totalResults: 12 }, - }; - expect(renderToolResult(ws, result)).toBe('12 sources'); - }); - - it('shows search complete when no totalResults', () => { - const result = { - content: [{ type: 'text', text: 'answer' }], - details: {}, - }; - expect(renderToolResult(ws, result)).toBe('search complete'); - }); - - it('shows expanded content with summary prefix', () => { - const result = { - content: [{ type: 'text', text: 'the full search answer' }], - details: { totalResults: 3 }, - }; - const rendered = renderToolResult(ws, result, { - expanded: true, - isPartial: false, - }); - expect(rendered).toContain('3 sources'); - expect(rendered).toContain('the full search answer'); - }); - - it('truncates long expanded text', () => { - const longText = 'x'.repeat(801); - const result = { - content: [{ type: 'text', text: longText }], - details: { totalResults: 1 }, - }; - const rendered = renderToolResult(ws, result, { - expanded: true, - isPartial: false, - }); - expect(rendered).toContain('...'); - expect(rendered).not.toContain(longText); - }); - }); - - it('does not register a provider-wide native web_search blocker', () => { - const events: string[] = []; - registerWebSearchTool({ - registerTool() {}, - on(event: string) { - events.push(event); - }, - } as unknown as ExtensionAPI); - - expect(events).not.toContain('tool_call'); - }); -}); diff --git a/tests/tools/webSearchDelegate.test.ts b/tests/tools/webSearchDelegate.test.ts deleted file mode 100644 index 6c16f68..0000000 --- a/tests/tools/webSearchDelegate.test.ts +++ /dev/null @@ -1,66 +0,0 @@ -import { afterEach, describe, expect, it, vi } from 'vitest'; -import { - bindLivePiWebAccess, - clearWebSearchDelegateForTests, - ensureWebSearchDelegate, - getWebSearchDelegate, - getWebSearchLoadError, - isPiWebAccessInstalled, - setWebSearchDelegateForTests, -} from '../../src/tools/webSearchDelegate.js'; - -afterEach(() => { - clearWebSearchDelegateForTests(); - vi.restoreAllMocks(); -}); - -describe('webSearchDelegate', () => { - it('returns delegate set for tests', async () => { - setWebSearchDelegateForTests(async () => ({ - content: [{ type: 'text', text: 'ok' }], - details: {}, - })); - - const delegate = getWebSearchDelegate(); - expect(delegate).toBeTypeOf('function'); - if (!delegate) throw new Error('expected delegate'); - const result = await delegate('id', {}, new AbortController().signal, undefined, { - cwd: '/tmp', - hasUI: false, - } as import('@earendil-works/pi-coding-agent').ExtensionContext); - const first = result.content[0]; - expect(first?.type === 'text' && first.text).toBe('ok'); - }); - - it('isPiWebAccessInstalled reflects agent install path', () => { - const installed = isPiWebAccessInstalled(); - expect(typeof installed).toBe('boolean'); - }); - - it('bindLivePiWebAccess resets delegate state', () => { - setWebSearchDelegateForTests(async () => ({ - content: [{ type: 'text', text: 'ok' }], - details: {}, - })); - expect(getWebSearchDelegate()).toBeTypeOf('function'); - - bindLivePiWebAccess({} as Parameters[0]); - - expect(getWebSearchDelegate()).toBeUndefined(); - }); - - it('ensureWebSearchDelegate returns undefined when pi-web-access is not installed', async () => { - const isInstalled = vi.fn(() => false); - const result = await ensureWebSearchDelegate(undefined, isInstalled); - expect(isInstalled).toHaveBeenCalledOnce(); - expect(result).toBeUndefined(); - expect(getWebSearchDelegate()).toBeUndefined(); - }); - - it('getWebSearchLoadError returns last error string', () => { - clearWebSearchDelegateForTests(); - // After clear, no error is set; the function should return undefined - const error = getWebSearchLoadError(); - expect(error).toBeUndefined(); - }); -}); diff --git a/tests/tools/webSearchDelegateRetry.test.ts b/tests/tools/webSearchDelegateRetry.test.ts deleted file mode 100644 index ea6bab8..0000000 --- a/tests/tools/webSearchDelegateRetry.test.ts +++ /dev/null @@ -1,258 +0,0 @@ -import { mkdirSync, mkdtempSync, writeFileSync } from 'node:fs'; -import { tmpdir } from 'node:os'; -import { join } from 'node:path'; -import { afterEach, describe, expect, it, vi } from 'vitest'; - -const testAgentDir = mkdtempSync(join(tmpdir(), 'pi-grok-cli-')); - -vi.mock('@earendil-works/pi-coding-agent', async (importOriginal) => ({ - ...(await importOriginal()), - getAgentDir: () => testAgentDir, -})); - -import { - bindLivePiWebAccess, - clearWebSearchDelegateForTests, - ensureWebSearchDelegate, - getWebSearchDelegate, - getWebSearchLoadError, -} from '../../src/tools/webSearchDelegate.js'; - -afterEach(() => { - clearWebSearchDelegateForTests(); - vi.unstubAllGlobals(); -}); - -describe('webSearchDelegate retry', () => { - it('loads the public extension factory and delegates the complete execution contract', async () => { - const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); - mkdirSync(extensionDir, { recursive: true }); - writeFileSync( - join(extensionDir, 'index.js'), - ` -export default function (pi) { - pi.registerTool({ - name: 'web_search', - execute: async (id, params, signal, onUpdate, ctx) => ({ - content: [{ type: 'text', text: JSON.stringify({ id, params, aborted: signal.aborted, hasUpdate: typeof onUpdate === 'function', cwd: ctx.cwd }) }], - details: {}, - }), - }) -} -`, - ); - const pi = { registerTool() {} } as unknown as Parameters[0]; - - await ensureWebSearchDelegate(pi); - - expect(getWebSearchLoadError()).toBeUndefined(); - const delegate = getWebSearchDelegate(); - expect(delegate).toBeTypeOf('function'); - if (!delegate) throw new Error('expected delegate'); - const result = await delegate( - 'call-1', - { query: 'xAI' }, - new AbortController().signal, - () => undefined, - { cwd: '/workspace' } as import('@earendil-works/pi-coding-agent').ExtensionContext, - ); - expect(result.content[0]).toEqual({ - type: 'text', - text: JSON.stringify({ - id: 'call-1', - params: { query: 'xAI' }, - aborted: false, - hasUpdate: true, - cwd: '/workspace', - }), - }); - }); - - it('isolates unrelated extension registrations while preserving API access', async () => { - const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); - mkdirSync(extensionDir, { recursive: true }); - writeFileSync( - join(extensionDir, 'index.js'), - ` -export default function (pi) { - pi.registerCommand('web-command', {}) - pi.registerShortcut('ctrl+w', {}) - pi.on('session_start', () => undefined) - pi.registerTool({ - name: 'web_search', - execute: async () => ({ content: [{ type: 'text', text: pi.marker }], details: {} }), - }) -} -`, - ); - const pi = { - marker: 'live API', - on: vi.fn(), - registerCommand: vi.fn(), - registerShortcut: vi.fn(), - registerTool: vi.fn(), - } as unknown as NonNullable[0]>; - - await ensureWebSearchDelegate(pi); - - expect(pi.on).not.toHaveBeenCalled(); - expect(pi.registerCommand).not.toHaveBeenCalled(); - expect(pi.registerShortcut).not.toHaveBeenCalled(); - expect(pi.registerTool).not.toHaveBeenCalled(); - const delegate = getWebSearchDelegate(); - expect(delegate).toBeTypeOf('function'); - if (!delegate) throw new Error('expected delegate'); - await expect( - delegate( - 'call-1', - {}, - undefined, - undefined, - {} as import('@earendil-works/pi-coding-agent').ExtensionContext, - ), - ).resolves.toMatchObject({ content: [{ type: 'text', text: 'live API' }] }); - }); - - it('retries after a failed delegate load', async () => { - const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); - mkdirSync(extensionDir, { recursive: true }); - writeFileSync( - join(extensionDir, 'index.js'), - ` -export default function (pi) { - globalThis.webSearchDelegateLoadAttempts = (globalThis.webSearchDelegateLoadAttempts ?? 0) + 1 - if (globalThis.webSearchDelegateLoadAttempts === 1) throw new Error('temporary load failure') - pi.registerTool({ - name: 'web_search', - execute: async () => ({ content: [{ type: 'text', text: 'ok' }], details: {} }), - }) -} -`, - ); - vi.stubGlobal('webSearchDelegateLoadAttempts', 0); - const pi = {} as Parameters[0]; - - await ensureWebSearchDelegate(pi); - expect(getWebSearchDelegate()).toBeUndefined(); - expect(getWebSearchLoadError()).toBe('temporary load failure'); - - await ensureWebSearchDelegate(pi); - expect(getWebSearchDelegate()).toBeTypeOf('function'); - expect(getWebSearchLoadError()).toBeUndefined(); - }); - - it('reports an incompatible package without a public default factory', async () => { - const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); - mkdirSync(extensionDir, { recursive: true }); - writeFileSync(join(extensionDir, 'index.js'), 'export const version = "old"\n'); - - await ensureWebSearchDelegate({} as Parameters[0]); - - expect(getWebSearchDelegate()).toBeUndefined(); - expect(getWebSearchLoadError()).toBe( - 'pi-web-access is incompatible. Install pi-web-access 0.13.0 or newer with a public default extension factory.', - ); - }); - - it('shares one public factory load across concurrent first calls', async () => { - const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); - mkdirSync(extensionDir, { recursive: true }); - writeFileSync( - join(extensionDir, 'index.js'), - ` -export default async function (pi) { - globalThis.concurrentDelegateLoads += 1 - await globalThis.finishConcurrentDelegateLoad - pi.registerTool({ name: 'web_search', execute: async () => ({ content: [], details: {} }) }) -} -`, - ); - let finish = () => {}; - vi.stubGlobal('concurrentDelegateLoads', 0); - vi.stubGlobal( - 'finishConcurrentDelegateLoad', - new Promise((resolve) => { - finish = resolve; - }), - ); - const pi = {} as Parameters[0]; - - const first = ensureWebSearchDelegate(pi); - const second = ensureWebSearchDelegate(pi); - const globals = globalThis as unknown as { concurrentDelegateLoads: number }; - await vi.waitFor(() => expect(globals.concurrentDelegateLoads).toBe(1)); - finish(); - await Promise.all([first, second]); - expect(globals.concurrentDelegateLoads).toBe(1); - }); - - it('reports extensions that do not register web_search', async () => { - const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); - mkdirSync(extensionDir, { recursive: true }); - writeFileSync( - join(extensionDir, 'index.js'), - ` -export default function (pi) { - pi.registerTool({ - name: 'fetch_content', - execute: async () => ({ content: [{ type: 'text', text: 'ok' }], details: {} }), - }) -} -`, - ); - const pi = {} as Parameters[0]; - - await ensureWebSearchDelegate(pi); - - expect(getWebSearchDelegate()).toBeUndefined(); - expect(getWebSearchLoadError()).toBe( - 'pi-web-access loaded but did not register web_search. Update pi-web-access.', - ); - }); - - it('does not let a stale load replace the delegate for a newer binding', async () => { - const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); - mkdirSync(extensionDir, { recursive: true }); - writeFileSync( - join(extensionDir, 'index.js'), - ` -export default async function (pi) { - const load = globalThis.webSearchDelegateLoads.shift() - load.started() - await load.wait - pi.registerTool({ - name: 'web_search', - execute: async () => ({ content: [{ type: 'text', text: load.name }], details: {} }), - }) -} -`, - ); - let startFirstLoad = () => {}; - const firstLoadStarted = new Promise((resolve) => { - startFirstLoad = resolve; - }); - let finishFirstLoad = () => {}; - const firstLoadWait = new Promise((resolve) => { - finishFirstLoad = resolve; - }); - vi.stubGlobal('webSearchDelegateLoads', [ - { name: 'first', started: startFirstLoad, wait: firstLoadWait }, - { name: 'second', started: () => {}, wait: Promise.resolve() }, - ]); - const firstPi = {} as Parameters[0]; - const secondPi = {} as Parameters[0]; - - bindLivePiWebAccess(firstPi); - const firstLoad = ensureWebSearchDelegate(); - await firstLoadStarted; - - bindLivePiWebAccess(secondPi); - await ensureWebSearchDelegate(); - const secondDelegate = getWebSearchDelegate(); - expect(secondDelegate).toBeTypeOf('function'); - - finishFirstLoad(); - await firstLoad; - expect(getWebSearchDelegate()).toBe(secondDelegate); - }); -}); diff --git a/tests/vision/cache.test.ts b/tests/vision/cache.test.ts deleted file mode 100644 index b573b37..0000000 --- a/tests/vision/cache.test.ts +++ /dev/null @@ -1,138 +0,0 @@ -import { mkdtempSync, rmSync } from 'node:fs'; -import { tmpdir } from 'node:os'; -import { join } from 'node:path'; -import { afterEach, describe, expect, it } from 'vitest'; -import { - cacheStats, - clearCache, - loadCache, - makeCacheEntry, - makeCacheKey, - pruneCache, - saveCache, - updateCache, - type VisionImage, -} from '../../src/vision/cache.js'; - -const tempDirs: string[] = []; - -afterEach(() => { - for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true }); -}); - -function cachePath() { - const dir = mkdtempSync(join(tmpdir(), 'grok-cli-vision-cache-')); - tempDirs.push(dir); - return join(dir, 'cache.json'); -} - -const image: VisionImage = { - data: Buffer.from('png-bytes').toString('base64'), - mimeType: 'image/png', -}; -const prompt = 'describe this'; - -describe('grok-cli-vision cache', () => { - it('starts empty for a missing cache file', () => { - const path = cachePath(); - expect(loadCache(path).entries).toEqual({}); - expect(cacheStats(path).entries).toBe(0); - }); - - it('produces stable keys for identical inputs', () => { - const keyA = makeCacheKey(image, 'grok-build', prompt); - const keyB = makeCacheKey(image, 'grok-build', prompt); - expect(keyA).toBe(keyB); - }); - - it('varies keys across model, prompt, and image bytes', () => { - expect(makeCacheKey(image, 'grok-4.3', prompt)).not.toBe( - makeCacheKey(image, 'grok-build', prompt), - ); - expect(makeCacheKey(image, 'grok-build', 'other prompt')).not.toBe( - makeCacheKey(image, 'grok-build', prompt), - ); - const otherImage: VisionImage = { - data: Buffer.from('other').toString('base64'), - mimeType: 'image/png', - }; - expect(makeCacheKey(otherImage, 'grok-build', prompt)).not.toBe( - makeCacheKey(image, 'grok-build', prompt), - ); - }); - - it('saves and reloads entries', () => { - const path = cachePath(); - const key = makeCacheKey(image, 'grok-build', prompt); - const cache = loadCache(path); - cache.entries[key] = makeCacheEntry(image, 'grok-build', prompt, 'a description'); - saveCache(cache, path); - - const reloaded = loadCache(path); - expect(reloaded.entries[key]?.description).toBe('a description'); - expect(reloaded.entries[key]?.model).toBe('grok-build'); - expect(cacheStats(path).entries).toBe(1); - }); - - it('serializes updates so concurrent writers preserve distinct entries', async () => { - const path = cachePath(); - const images: VisionImage[] = ['first', 'second'].map((value) => ({ - data: Buffer.from(value).toString('base64'), - mimeType: 'image/png', - })); - - await Promise.all( - images.map((img, i) => - updateCache(path, (cache) => { - cache.entries[makeCacheKey(img, 'grok-build', prompt)] = makeCacheEntry( - img, - 'grok-build', - prompt, - `d${i}`, - ); - }), - ), - ); - - expect(Object.values(loadCache(path).entries).map((entry) => entry.description)).toEqual([ - 'd0', - 'd1', - ]); - }); - - it('prunes to the most recent entries by createdAt', () => { - const path = cachePath(); - const cache = loadCache(path); - const images: VisionImage[] = Array.from({ length: 4 }, (_, i) => ({ - data: Buffer.from(`img${i}`).toString('base64'), - mimeType: 'image/png', - })); - images.forEach((img, i) => { - const key = makeCacheKey(img, 'grok-build', prompt); - const entry = makeCacheEntry(img, 'grok-build', prompt, `d${i}`); - entry.createdAt = new Date(2020, 0, i + 1).toISOString(); - cache.entries[key] = entry; - }); - - pruneCache(cache, 2); - expect(Object.keys(cache.entries)).toHaveLength(2); - const descriptions = Object.values(cache.entries).map((e) => e.description); - expect(descriptions).toEqual(['d3', 'd2']); - }); - - it('clearCache empties the cache file', () => { - const path = cachePath(); - const cache = loadCache(path); - cache.entries[makeCacheKey(image, 'grok-build', prompt)] = makeCacheEntry( - image, - 'grok-build', - prompt, - 'd', - ); - saveCache(cache, path); - - clearCache(path); - expect(loadCache(path).entries).toEqual({}); - expect(cacheStats(path).entries).toBe(0); - }); -}); diff --git a/tests/vision/config.test.ts b/tests/vision/config.test.ts deleted file mode 100644 index a8604d0..0000000 --- a/tests/vision/config.test.ts +++ /dev/null @@ -1,117 +0,0 @@ -import { mkdirSync, writeFileSync } from 'node:fs'; -import { dirname } from 'node:path'; -import { describe, expect, it } from 'vitest'; -import { - DEFAULT_CONFIG, - describableModels, - getConfigPath, - loadConfig, - normalizeVisionConfig, - saveConfig, -} from '../../src/config.js'; -import { useTempHome } from './helpers.js'; - -const setupHome = useTempHome(); - -function writeConfig(config: unknown) { - const configPath = getConfigPath(); - mkdirSync(dirname(configPath), { recursive: true }); - writeFileSync(configPath, JSON.stringify({ version: 1, vision: config })); -} - -describe('grok-cli-vision config', () => { - it('returns defaults when no config file exists', () => { - setupHome(); - const { config, warning } = loadConfig(); - expect(config.vision).toEqual(DEFAULT_CONFIG.vision); - expect(warning).toBeUndefined(); - }); - - it('loads a valid config', () => { - setupHome(); - writeConfig({ - enabled: false, - model: 'grok-build', - maxImages: 2, - cacheEnabled: false, - cacheMaxEntries: 25, - }); - - const { config, warning } = loadConfig(); - expect(config.vision).toEqual({ - enabled: false, - model: 'grok-build', - maxImages: 2, - cacheEnabled: false, - cacheMaxEntries: 25, - }); - expect(warning).toBeUndefined(); - }); - - it('falls back to defaults with a warning for invalid fields', () => { - setupHome(); - writeConfig({ - enabled: 'yes', - model: 'not-a-model', - maxImages: -3, - cacheEnabled: 'yes', - cacheMaxEntries: 0, - }); - - const { config, warning } = loadConfig(); - expect(config.vision).toEqual(DEFAULT_CONFIG.vision); - expect(warning).toMatch(/enabled must be true or false/); - expect(warning).toMatch(/Unknown model "not-a-model"/); - expect(warning).toMatch(/maxImages must be a positive integer/); - expect(warning).toMatch(/cacheEnabled must be true or false/); - expect(warning).toMatch(/cacheMaxEntries must be a positive integer/); - }); - - it('warns when the config file is not a JSON object', () => { - setupHome(); - const configPath = getConfigPath(); - mkdirSync(dirname(configPath), { recursive: true }); - writeFileSync(configPath, JSON.stringify([1, 2, 3])); - - const { config, warning } = loadConfig(); - expect(config).toEqual(DEFAULT_CONFIG); - expect(warning).toMatch(/must be a JSON object/); - }); - - it('warns on invalid JSON', () => { - setupHome(); - const configPath = getConfigPath(); - mkdirSync(dirname(configPath), { recursive: true }); - writeFileSync(configPath, '{ not json'); - - const { config, warning } = loadConfig(); - expect(config).toEqual(DEFAULT_CONFIG); - expect(warning).toMatch(/Could not read/); - }); - - it('saveConfig writes normalized values and round-trips through loadConfig', () => { - setupHome(); - saveConfig({ - ...DEFAULT_CONFIG, - vision: { ...DEFAULT_CONFIG.vision, enabled: false, maxImages: 3 }, - }); - - const { config } = loadConfig(); - expect(config.vision.enabled).toBe(false); - expect(config.vision.maxImages).toBe(3); - expect(config.vision.model).toBe('grok-build'); - }); - - it('normalizeConfig rejects a describer model that lacks image input', () => { - const warnings: string[] = []; - const config = normalizeVisionConfig({ model: 'grok-composer-2.5-fast' }, warnings); - expect(config.model).toBe('grok-build'); - expect(warnings.join(' ')).toMatch(/Unknown model "grok-composer-2.5-fast"/); - }); - - it('describableModels only lists image-capable models', () => { - const models = describableModels(); - expect(models).toContain('grok-build'); - expect(models).not.toContain('grok-composer-2.5-fast'); - }); -}); diff --git a/tests/vision/describe.test.ts b/tests/vision/describe.test.ts deleted file mode 100644 index 0639893..0000000 --- a/tests/vision/describe.test.ts +++ /dev/null @@ -1,429 +0,0 @@ -import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs'; -import { tmpdir } from 'node:os'; -import { dirname, join } from 'node:path'; -import type { ExtensionContext, ToolResultEvent } from '@earendil-works/pi-coding-agent'; -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -import { DEFAULT_CONFIG, getConfigPath } from '../../src/config.js'; -import { handleReadResult } from '../../src/vision/describe.js'; -import { saveTestAccounts } from './helpers.js'; - -const BASE_URL = 'https://cli-chat-proxy.grok.com/v1'; -const PNG = Buffer.from('fake-png-bytes').toString('base64'); - -const originalFetch = globalThis.fetch; -const originalHome = process.env.HOME; -const originalToken = process.env.GROK_CLI_OAUTH_TOKEN; -const tempDirs: string[] = []; -let fetchMock: ReturnType>; - -beforeEach(() => { - const dir = mkdtempSync(join(tmpdir(), 'grok-cli-vision-desc-')); - mkdirSync(join(dir, '.pi'), { recursive: true }); - tempDirs.push(dir); - process.env.HOME = dir; - delete process.env.GROK_CLI_OAUTH_TOKEN; - fetchMock = vi.fn(async () => - Response.json({ output_text: 'a screenshot of a button' }), - ); - globalThis.fetch = fetchMock; -}); - -afterEach(() => { - globalThis.fetch = originalFetch; - if (originalHome === undefined) delete process.env.HOME; - else process.env.HOME = originalHome; - if (originalToken === undefined) delete process.env.GROK_CLI_OAUTH_TOKEN; - else process.env.GROK_CLI_OAUTH_TOKEN = originalToken; - for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true }); - vi.useRealTimers(); -}); - -interface CtxOverrides { - modelInput?: ('text' | 'image')[]; - apiKey?: string; - signal?: AbortSignal; - withKey?: boolean; -} - -function buildCtx(overrides: CtxOverrides = {}): ExtensionContext { - const hasKey = overrides.withKey !== false; - return { - model: { input: overrides.modelInput ?? ['text'] }, - modelRegistry: hasKey - ? { getApiKeyForProvider: async () => overrides.apiKey ?? 'provider-token' } - : {}, - ui: { notify: vi.fn() }, - signal: overrides.signal, - } as unknown as ExtensionContext; -} - -function readEvent(content: unknown[], toolName = 'read'): ToolResultEvent { - return { - type: 'tool_result', - toolName, - toolCallId: 'call-1', - input: {}, - content, - isError: false, - details: undefined, - } as unknown as ToolResultEvent; -} - -function imageBlock(data = PNG): unknown { - return { type: 'image', data, mimeType: 'image/png' }; -} - -function lastBody(): Record { - const init = fetchMock.mock.calls.at(-1)?.[1] as RequestInit | undefined; - return JSON.parse(String(init?.body)) as Record; -} - -describe('handleReadResult — no-op cases', () => { - it('does nothing when the active model handles images natively', async () => { - const result = await handleReadResult( - readEvent([imageBlock()]), - buildCtx({ modelInput: ['text', 'image'] }), - ); - expect(result).toBeUndefined(); - expect(fetchMock).not.toHaveBeenCalled(); - }); - - it('routes when the model has no declared input (treated as non-vision, matching the read tool)', async () => { - const result = await handleReadResult(readEvent([imageBlock()]), buildCtx({ modelInput: [] })); - expect(result?.content[0]).toMatchObject({ type: 'text' }); - expect((result?.content[0] as { text: string }).text).toContain('described by grok-build'); - }); - - it('does nothing for a non-read tool', async () => { - const result = await handleReadResult(readEvent([imageBlock()], 'bash'), buildCtx()); - expect(result).toBeUndefined(); - expect(fetchMock).not.toHaveBeenCalled(); - }); - - it('does nothing when the read result has no images', async () => { - const result = await handleReadResult( - readEvent([{ type: 'text', text: 'just text' }]), - buildCtx(), - ); - expect(result).toBeUndefined(); - expect(fetchMock).not.toHaveBeenCalled(); - }); - - it('does nothing when routing is disabled via config', async () => { - const configPath = getConfigPath(); - mkdirSync(dirname(configPath), { recursive: true }); - writeFileSync( - configPath, - JSON.stringify({ - ...DEFAULT_CONFIG, - vision: { ...DEFAULT_CONFIG.vision, enabled: false }, - }), - ); - - const result = await handleReadResult(readEvent([imageBlock()]), buildCtx()); - expect(result).toBeUndefined(); - expect(fetchMock).not.toHaveBeenCalled(); - }); -}); - -describe('handleReadResult — image routing', () => { - it('describes an image via grok-build and replaces it with text', async () => { - const ctx = buildCtx(); - const result = await handleReadResult(readEvent([imageBlock()]), ctx); - - expect(fetchMock).toHaveBeenCalledOnce(); - expect(fetchMock.mock.calls[0]?.[0]).toBe(`${BASE_URL}/responses`); - - const body = lastBody(); - expect(body.model).toBe('grok-build'); - expect(body.stream).toBe(false); - const content = (body.input as Array<{ content: Array> }>)[0].content; - expect(content[1]).toEqual({ - type: 'input_image', - image_url: `data:image/png;base64,${PNG}`, - detail: 'auto', - }); - - const headers = (fetchMock.mock.calls[0]?.[1] as RequestInit).headers as Record; - expect(headers.Authorization).toBe('Bearer provider-token'); - expect(headers['x-grok-model-override']).toBe('grok-build'); - - expect(result?.content).toHaveLength(1); - expect(result?.content[0]).toEqual({ - type: 'text', - text: '[Image 1 — described by grok-build]\na screenshot of a button', - }); - }); - - it('reuses a cached description without calling the API on the second sight', async () => { - const ctx = buildCtx(); - await handleReadResult(readEvent([imageBlock()]), ctx); - await handleReadResult(readEvent([imageBlock()]), ctx); - - expect(fetchMock).toHaveBeenCalledOnce(); - }); - - it('prefers the GROK_CLI_OAUTH_TOKEN env var as the bearer', async () => { - process.env.GROK_CLI_OAUTH_TOKEN = 'env-token'; - await handleReadResult(readEvent([imageBlock()]), buildCtx({ apiKey: 'should-not-be-used' })); - - const headers = (fetchMock.mock.calls[0]?.[1] as RequestInit).headers as Record; - expect(headers.Authorization).toBe('Bearer env-token'); - }); - - it('uses the last selected Grok alias while a non-Grok model is active', async () => { - saveTestAccounts(); - const getApiKeyForProvider = vi.fn(async () => 'work-token'); - const ctx = { - ...buildCtx(), - model: { provider: 'openai', input: ['text'] }, - modelRegistry: { getApiKeyForProvider }, - } as unknown as ExtensionContext; - - await handleReadResult(readEvent([imageBlock()]), ctx); - - expect(getApiKeyForProvider).toHaveBeenCalledWith('grok-cli-2'); - const headers = (fetchMock.mock.calls[0]?.[1] as RequestInit).headers as Record; - expect(headers.Authorization).toBe('Bearer work-token'); - }); - - it('describes multiple images with preserved order and labels', async () => { - const result = await handleReadResult( - readEvent([ - imageBlock(Buffer.from('img-a').toString('base64')), - imageBlock(Buffer.from('img-b').toString('base64')), - ]), - buildCtx(), - ); - - expect(fetchMock).toHaveBeenCalledTimes(2); - expect(result?.content).toEqual([ - { type: 'text', text: '[Image 1 — described by grok-build]\na screenshot of a button' }, - { type: 'text', text: '[Image 2 — described by grok-build]\na screenshot of a button' }, - ]); - }); - - it('preserves text when replacing images in mixed read results', async () => { - const result = await handleReadResult( - readEvent([{ type: 'text', text: 'file metadata' }, imageBlock()]), - buildCtx(), - ); - - expect(result?.content).toEqual([ - { type: 'text', text: 'file metadata' }, - { type: 'text', text: '[Image 1 — described by grok-build]\na screenshot of a button' }, - ]); - }); - - it('stays silent while describing and writes the cache file', async () => { - const ctx = buildCtx(); - await handleReadResult(readEvent([imageBlock()]), ctx); - - expect(ctx.ui.notify).not.toHaveBeenCalledWith( - expect.stringMatching(/Describing image/), - 'info', - ); - const cachePath = join(process.env.HOME as string, '.pi', 'grok-cli', 'vision-cache.json'); - expect(existsSync(cachePath)).toBe(true); - const cache = JSON.parse(readFileSync(cachePath, 'utf-8')); - expect(Object.keys(cache.entries)).toHaveLength(1); - }); -}); - -describe('handleReadResult — failure handling', () => { - it('returns a text error and warns when the API rejects the request', async () => { - fetchMock = vi.fn(async () => new Response('bad model', { status: 400 })); - globalThis.fetch = fetchMock; - - const ctx = buildCtx(); - const result = await handleReadResult(readEvent([imageBlock()]), ctx); - - expect(fetchMock).toHaveBeenCalledOnce(); - expect(result?.content).toHaveLength(1); - expect((result?.content[0] as { text: string }).text).toMatch( - /Image 1 — description unavailable/, - ); - expect((result?.content[0] as { text: string }).text).toMatch(/HTTP 400/); - expect(ctx.ui.notify).toHaveBeenCalledWith( - expect.stringMatching(/description failed/), - 'warning', - ); - }); - - it('returns a text error when the network call fails', async () => { - fetchMock = vi.fn(async () => { - throw new Error('aborted'); - }); - globalThis.fetch = fetchMock; - - const ctx = buildCtx(); - const result = await handleReadResult(readEvent([imageBlock()]), ctx); - - expect(result?.content[0]).toMatchObject({ - type: 'text', - }); - expect((result?.content[0] as { text: string }).text).toMatch(/description unavailable/); - }); - - it('returns a not-authenticated note and never calls the API without a key', async () => { - const ctx = buildCtx({ withKey: false }); - const result = await handleReadResult(readEvent([imageBlock()]), ctx); - - expect(fetchMock).not.toHaveBeenCalled(); - expect(result?.content).toEqual([ - { type: 'text', text: '[grok-cli-vision: image not described — not authenticated]' }, - ]); - expect(ctx.ui.notify).toHaveBeenCalledWith(expect.any(String), 'warning'); - }); -}); - -describe('handleReadResult — response shapes and resilience', () => { - it('reads descriptions from the output[] message-item fallback shape', async () => { - fetchMock = vi.fn(async () => - Response.json({ - output: [ - { - type: 'message', - content: [{ type: 'output_text', text: 'array-shape description' }], - }, - ], - }), - ); - globalThis.fetch = fetchMock; - - const result = await handleReadResult(readEvent([imageBlock()]), buildCtx()); - expect(result?.content[0]).toMatchObject({ type: 'text' }); - expect((result?.content[0] as { text: string }).text).toContain('array-shape description'); - }); - - it('treats an empty description as a failure', async () => { - fetchMock = vi.fn(async () => Response.json({ output_text: ' ' })); - globalThis.fetch = fetchMock; - - const ctx = buildCtx(); - const result = await handleReadResult(readEvent([imageBlock()]), ctx); - expect((result?.content[0] as { text: string }).text).toMatch(/description unavailable/); - expect((result?.content[0] as { text: string }).text).toMatch(/empty description/); - }); - - it('retries and recovers after a transient 500', async () => { - const responses = [ - new Response('server error', { status: 500 }), - Response.json({ output_text: 'recovered' }), - ]; - let i = 0; - fetchMock = vi.fn(async () => responses[i++] ?? (responses.at(-1) as Response)); - globalThis.fetch = fetchMock; - - const result = await handleReadResult(readEvent([imageBlock()]), buildCtx()); - expect(fetchMock).toHaveBeenCalledTimes(2); - expect((result?.content[0] as { text: string }).text).toContain('recovered'); - }); - - it('reports an exhausted rate-limit retry sequence', async () => { - vi.useFakeTimers(); - fetchMock = vi.fn(async () => - Response.json({ error: { message: 'quota exhausted' } }, { status: 429 }), - ); - globalThis.fetch = fetchMock; - - const resultPromise = handleReadResult(readEvent([imageBlock()]), buildCtx()); - await vi.advanceTimersByTimeAsync(1_500); - const result = await resultPromise; - - expect(fetchMock).toHaveBeenCalledTimes(3); - expect((result?.content[0] as { text: string }).text).toContain( - 'Grok CLI rate limited the request (HTTP 429). Try again later: quota exhausted', - ); - }); - - it('reports an exhausted request-timeout retry sequence', async () => { - vi.useFakeTimers(); - fetchMock = vi.fn( - async (_input, init) => - new Promise((_resolve, reject) => { - const signal = init?.signal; - if (!signal) throw new Error('expected request signal'); - signal.addEventListener('abort', () => reject(signal.reason), { once: true }); - }), - ); - globalThis.fetch = fetchMock; - - const resultPromise = handleReadResult(readEvent([imageBlock()]), buildCtx()); - await vi.runAllTimersAsync(); - const result = await resultPromise; - - expect(fetchMock).toHaveBeenCalledTimes(3); - expect((result?.content[0] as { text: string }).text).toContain( - 'Grok CLI request timed out after 30s after 3 attempts.', - ); - expect(vi.getTimerCount()).toBe(0); - }); - - it('stops retrying when the request is cancelled', async () => { - const controller = new AbortController(); - fetchMock = vi.fn(async () => { - controller.abort(); - return new Response('server error', { status: 500 }); - }); - globalThis.fetch = fetchMock; - - const result = await handleReadResult( - readEvent([imageBlock()]), - buildCtx({ signal: controller.signal }), - ); - - expect(fetchMock).toHaveBeenCalledOnce(); - expect((result?.content[0] as { text: string }).text).toContain('request cancelled'); - }); - - it.each([ - [401, 'Grok CLI rejected the API key (HTTP 401)'], - [418, 'Grok CLI request failed (HTTP 418)'], - ])('explains non-retryable HTTP %i responses', async (status, message) => { - fetchMock = vi.fn(async () => new Response('', { status })); - globalThis.fetch = fetchMock; - - const result = await handleReadResult(readEvent([imageBlock()]), buildCtx()); - - expect(fetchMock).toHaveBeenCalledOnce(); - expect((result?.content[0] as { text: string }).text).toContain(message); - }); - - it('reports a successful non-JSON response', async () => { - fetchMock = vi.fn(async () => new Response('proxy error')); - globalThis.fetch = fetchMock; - - const result = await handleReadResult(readEvent([imageBlock()]), buildCtx()); - - expect((result?.content[0] as { text: string }).text).toContain( - 'Grok CLI returned non-JSON response:', - ); - }); - - it('caps described images at maxImages and notes the skipped remainder', async () => { - const configPath = getConfigPath(); - mkdirSync(dirname(configPath), { recursive: true }); - writeFileSync( - configPath, - JSON.stringify({ - ...DEFAULT_CONFIG, - vision: { ...DEFAULT_CONFIG.vision, maxImages: 1 }, - }), - ); - - const result = await handleReadResult( - readEvent([ - imageBlock(Buffer.from('first').toString('base64')), - imageBlock(Buffer.from('second').toString('base64')), - ]), - buildCtx(), - ); - - expect(fetchMock).toHaveBeenCalledOnce(); - expect(result?.content.at(-1)).toMatchObject({ type: 'text' }); - expect((result?.content.at(-1) as { text: string }).text).toMatch( - /1 additional image\(s\) omitted/, - ); - }); -}); diff --git a/tests/vision/helpers.ts b/tests/vision/helpers.ts deleted file mode 100644 index c7ae076..0000000 --- a/tests/vision/helpers.ts +++ /dev/null @@ -1,43 +0,0 @@ -import { mkdtempSync, rmSync } from 'node:fs'; -import { tmpdir } from 'node:os'; -import { join } from 'node:path'; -import { afterEach } from 'vitest'; -import { DEFAULT_CONFIG, saveConfig } from '../../src/config.js'; - -export const TEST_ACCOUNTS = [ - { provider: 'grok-cli', label: 'Personal' }, - { provider: 'grok-cli-2', label: 'Work' }, -]; - -export const oauthCredential = (access: string) => ({ - type: 'oauth' as const, - access, - refresh: `${access}-refresh`, - expires: Date.now() + 60_000, -}); - -export function saveTestAccounts(selectedProvider = 'grok-cli-2') { - saveConfig({ - ...DEFAULT_CONFIG, - accounts: { nextAccountNumber: 3, selectedProvider, items: TEST_ACCOUNTS }, - }); -} - -/** - * Point HOME at a fresh temp dir for the whole test file, restoring it on - * teardown. Returns a setup function that creates a new dir per call. - */ -export function useTempHome(): () => string { - const originalHome = process.env.HOME; - const dirs: string[] = []; - afterEach(() => { - if (originalHome === undefined) delete process.env.HOME; - else process.env.HOME = originalHome; - for (const dir of dirs.splice(0)) rmSync(dir, { recursive: true, force: true }); - }); - return () => { - const dir = mkdtempSync(join(tmpdir(), 'grok-cli-vision-')); - process.env.HOME = dir; - return dir; - }; -} diff --git a/tests/vision/register.test.ts b/tests/vision/register.test.ts deleted file mode 100644 index fe22d3f..0000000 --- a/tests/vision/register.test.ts +++ /dev/null @@ -1,117 +0,0 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { describe, expect, it, vi } from 'vitest'; -import { DEFAULT_CONFIG, loadConfig, saveConfig } from '../../src/config.js'; -import { useTempHome } from './helpers.js'; - -const setupHome = useTempHome(); - -interface CommandConfig { - handler: (args: string[], ctx: unknown) => Promise; -} - -type ToolResultHandler = (event: unknown, ctx: unknown) => unknown; - -async function setupExtension() { - const commands = new Map(); - const toolResultHandlers: ToolResultHandler[] = []; - const { registerVisionFeature } = await import('../../src/vision/register.js'); - registerVisionFeature({ - on(event: string, handler: ToolResultHandler) { - if (event === 'tool_result') toolResultHandlers.push(handler); - }, - registerCommand(name: string, config: unknown) { - commands.set(name, config as CommandConfig); - }, - } as unknown as ExtensionAPI); - return { commands, toolResultHandlers }; -} - -describe('registerVisionFeature', () => { - it('registers the tool_result handler and three commands', async () => { - setupHome(); - const { commands, toolResultHandlers } = await setupExtension(); - - expect(toolResultHandlers).toHaveLength(1); - expect([...commands.keys()].sort()).toEqual([ - 'grok-cli-vision', - 'grok-cli-vision:cache-clear', - 'grok-cli-vision:status', - ]); - }); - - it('status reports ON with the default describer and zero cache entries', async () => { - setupHome(); - const { commands } = await setupExtension(); - const notify = vi.fn(); - - await commands.get('grok-cli-vision:status')?.handler([], { - ui: { notify }, - }); - - const text = notify.mock.calls.at(-1)?.[0] as string; - expect(text).toMatch(/grok-cli-vision: ON/); - expect(text).toMatch(/describer: grok-build/); - expect(text).toMatch(/cache: ON \(0 entries/); - }); - - it('toggles enabled state and persists it to the config file', async () => { - setupHome(); - const { commands } = await setupExtension(); - const notify = vi.fn(); - saveConfig({ ...DEFAULT_CONFIG, imagine: { enabled: false } }); - - await commands.get('grok-cli-vision')?.handler([], { ui: { notify } }); - expect(loadConfig().config.vision.enabled).toBe(false); - expect(loadConfig().config.imagine.enabled).toBe(false); - expect(notify).toHaveBeenCalledWith('grok-cli-vision: OFF', 'info'); - - await commands.get('grok-cli-vision')?.handler([], { ui: { notify } }); - expect(loadConfig().config.vision.enabled).toBe(true); - expect(loadConfig().config.imagine.enabled).toBe(false); - expect(notify).toHaveBeenCalledWith('grok-cli-vision: ON (grok-build)', 'info'); - }); - - it('clears cached vision descriptions and notifies the user', async () => { - setupHome(); - const { commands } = await setupExtension(); - const { getCachePath, loadCache, saveCache } = await import('../../src/vision/cache.js'); - const notify = vi.fn(); - saveCache( - { - version: 1, - entries: { - cached: { - createdAt: '2026-07-12T00:00:00.000Z', - description: 'cached description', - imageHash: 'image', - mediaType: 'image/png', - model: 'grok-build', - promptHash: 'prompt', - }, - }, - }, - getCachePath(), - ); - - await commands.get('grok-cli-vision:cache-clear')?.handler([], { ui: { notify } }); - - expect(loadCache(getCachePath()).entries).toEqual({}); - expect(notify).toHaveBeenCalledWith('grok-cli-vision cache: cleared', 'info'); - }); - - it('delegates read tool results to the vision handler', async () => { - setupHome(); - const { toolResultHandlers } = await setupExtension(); - const result = await toolResultHandlers[0]?.( - { type: 'tool_result', toolName: 'bash', content: [{ type: 'text', text: 'x' }] }, - { - model: { input: ['text'] }, - modelRegistry: { getApiKeyForProvider: async () => 'token' }, - ui: { notify: vi.fn() }, - signal: undefined, - }, - ); - - expect(result).toBeUndefined(); - }); -});