From 0a81fd37a3b2c1712fa4260d66b45e4a4b7ce38d Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Sun, 20 Sep 2026 10:52:46 +0900 Subject: [PATCH 1/5] feat(models): expose Grok 4.6 Extra High reasoning Adapted from Ayagikei/pi-grok-cli commit cd0479319e7f65acbb4da48cdecf483aa3fc6bb3. Verify selection through Pi's thinking-level API. Co-authored-by: AyagiKei --- src/models/catalog.ts | 1 + tests/models/catalog.test.ts | 15 +++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/src/models/catalog.ts b/src/models/catalog.ts index 9be4c81..48360d4 100644 --- a/src/models/catalog.ts +++ b/src/models/catalog.ts @@ -88,6 +88,7 @@ const FALLBACK_MODELS: GrokCliModelConfig[] = [ cost: COST_46, contextWindow: 500_000, maxTokens: 30_000, + thinkingLevelMap: { xhigh: 'xhigh' }, }, { id: 'grok-4.20-0309-reasoning', diff --git a/tests/models/catalog.test.ts b/tests/models/catalog.test.ts index a93e738..5f912dc 100644 --- a/tests/models/catalog.test.ts +++ b/tests/models/catalog.test.ts @@ -1,3 +1,4 @@ +import { clampThinkingLevel, getSupportedThinkingLevels } from '@earendil-works/pi-ai'; import { afterEach, describe, expect, it } from 'vitest'; import { resolveModels, @@ -12,6 +13,20 @@ afterEach(() => { }); describe('model catalog', () => { + it('lets Pi select Extra High reasoning for Grok 4.6', () => { + delete process.env.PI_GROK_CLI_MODELS; + const config = resolveModels().find((model) => model.id === 'grok-4.6'); + if (!config) throw new Error('Grok 4.6 is missing'); + const model = { + ...config, + provider: 'grok-cli', + api: 'openai-responses' as const, + baseUrl: 'https://cli-chat-proxy.grok.com', + }; + expect(getSupportedThinkingLevels(model)).toContain('xhigh'); + expect(clampThinkingLevel(model, 'xhigh')).toBe('xhigh'); + }); + it('reports reasoning-effort support by normalized model name', () => { expect(supportsReasoningEffort('grok-4.3')).toBe(true); expect(supportsReasoningEffort('grok-4.5')).toBe(true); From e03d821086023b436692a159133ce606c4cacf73 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Sun, 20 Sep 2026 10:53:24 +0900 Subject: [PATCH 2/5] fix(accounts): remember selection for new sessions Persist explicit account activation as the vault default, matching the README, while retaining restored per-session selections. Adapted from Ayagikei/pi-grok-cli cd0479319e7f65acbb4da48cdecf483aa3fc6bb3. Co-authored-by: AyagiKei --- src/provider/accounts.ts | 1 + tests/provider/accounts.test.ts | 24 ++++++++++++++++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/src/provider/accounts.ts b/src/provider/accounts.ts index 5c22ee7..52533d4 100644 --- a/src/provider/accounts.ts +++ b/src/provider/accounts.ts @@ -278,6 +278,7 @@ function createAccountManager( if (!account.credential) { throw new Error(`Log in to “${account.label}” before making it active.`); } + vault.activeAccountId = account.id; return { id: account.id, slot: account.slot, label: account.label }; }); sessionSelection.select(_ctx, id); diff --git a/tests/provider/accounts.test.ts b/tests/provider/accounts.test.ts index 460c8a3..917d894 100644 --- a/tests/provider/accounts.test.ts +++ b/tests/provider/accounts.test.ts @@ -14,6 +14,7 @@ import { } from '../../src/provider/accounts.js'; import { getAccountVault, mutateAccountVault } from '../../src/provider/accountVault.js'; import { loadQuotaCache, saveQuotaUsage } from '../../src/provider/quotaCache.js'; +import { createSessionAccountSelection } from '../../src/provider/sessionAccountSelection.js'; import { deferred, oauthCredential, @@ -226,14 +227,33 @@ describe('vault account management', () => { await expect(accounts.activate(ctx, account.id)).rejects.toThrow('before making it active'); }); - it('activates an account only for the current Pi session', async () => { + it('activates an account for the current session and future sessions', async () => { const test = await selectedLoggedInAccount(); expect(test.appendEntry).toHaveBeenCalledWith('grok-cli-active-account-v1', { accountId: test.account.id, }); - expect((await getAccountVault()).activeAccountId).toBe('account-1'); + expect((await getAccountVault()).activeAccountId).toBe(test.account.id); expect(test.accounts.snapshot(ctx).accounts[1]).toMatchObject({ active: true }); + const freshSelection = createSessionAccountSelection({ appendEntry: vi.fn() }); + expect(freshSelection.accountId('new-session')).toBe(test.account.id); + freshSelection.restore({ + sessionManager: { + ...ctx.sessionManager, + getSessionId: () => 'existing-session', + getBranch: () => [ + { + type: 'custom', + id: 'entry-1', + parentId: null, + timestamp: new Date().toISOString(), + customType: 'grok-cli-active-account-v1', + data: { accountId: 'account-1' }, + }, + ], + }, + }); + expect(freshSelection.accountId('existing-session')).toBe('account-1'); }); it('selects another logged-in account when the active account logs out', async () => { From f4ee6865121f56f8374765eb22948a796dc757d7 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Sun, 20 Sep 2026 10:55:35 +0900 Subject: [PATCH 3/5] feat(imagine): edit local images through Grok Imagine Adapt Ayagikei/pi-grok-cli cd0479319e7f65acbb4da48cdecf483aa3fc6bb3. Resolve source paths from the session cwd and validate image signatures and dimensions before uploading. Cover command aliases, tool paths, malformed inputs, and the editing request. Co-authored-by: AyagiKei --- README.md | 4 +++- SECURITY.md | 2 +- src/imagine/generate.ts | 4 +++- src/imagine/imageUrl.ts | 19 +++++++++++++++ src/imagine/parseArgs.ts | 3 +++ src/imagine/register.ts | 5 ++-- src/imagine/tool.ts | 16 ++++++++++--- src/imagine/workflow.ts | 11 +++++++++ tests/imagine/generate.test.ts | 17 ++++++++++++++ tests/imagine/helpers.ts | 3 +++ tests/imagine/register.test.ts | 19 ++++++++++++++- tests/imagine/tool.test.ts | 43 ++++++++++++++++++++++++++++++++-- tests/provider/package.test.ts | 1 + 13 files changed, 136 insertions(+), 11 deletions(-) create mode 100644 src/imagine/imageUrl.ts diff --git a/README.md b/README.md index 4325748..290a848 100644 --- a/README.md +++ b/README.md @@ -112,6 +112,8 @@ Models are bundled rather than discovered live. Registered context limits may di 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. +To edit a local PNG, JPEG, or WebP image, use `/grok-cli-imagine --image "./source image.png" `. `--edit` is an alias. The `image_gen` tool accepts the same local path in its optional `image` argument. Relative paths use the session working directory. The source file is uploaded to Imagine; the edited image is saved separately unless you explicitly use `--out` to overwrite the source. + `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. ## Commands @@ -120,7 +122,7 @@ Run `/grok-cli-imagine ` to generate and preview a JPEG, or let any acti | --- | --- | | `/grok-cli-accounts [gui]` | Manage Grok accounts in the terminal, or add `gui` for the browser dashboard. | | `/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 ` | Generate or edit an image. Supports `--image`/`--edit`, `--aspect`, `--out`, and `--resolution 1k`. | | `/grok-cli-imagine:tool [on\|off\|status]` | Toggle, set, or report persistent model-callable `image_gen` availability. | ## Configuration diff --git a/SECURITY.md b/SECURITY.md index d188c9c..48bae06 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -23,7 +23,7 @@ The maintainer aims to acknowledge a complete report within 7 calendar days and - 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. -- 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. +- Grok Imagine sends the selected account's bearer token, prompt, options, and any local source image selected for editing 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. - Subscription tier, weekly allowance usage, 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. - 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. diff --git a/src/imagine/generate.ts b/src/imagine/generate.ts index fb6ddd4..458fa9e 100644 --- a/src/imagine/generate.ts +++ b/src/imagine/generate.ts @@ -98,13 +98,14 @@ export async function generateImage(options: { prompt: string; aspectRatio?: string; resolution?: string; + imageUrl?: string; baseUrl?: string; signal?: AbortSignal; fetchImpl?: typeof fetch; }) { const response = await requestWithRetry( options.fetchImpl ?? fetch, - `${(options.baseUrl ?? process.env.PI_GROK_CLI_IMAGINE_BASE_URL ?? 'https://api.x.ai/v1').replace(/\/+$/, '')}/images/generations`, + `${(options.baseUrl ?? process.env.PI_GROK_CLI_IMAGINE_BASE_URL ?? 'https://api.x.ai/v1').replace(/\/+$/, '')}/images/${options.imageUrl ? 'edits' : 'generations'}`, { method: 'POST', headers: { @@ -121,6 +122,7 @@ export async function generateImage(options: { aspect_ratio: normalizeAspectRatio(options.aspectRatio), resolution: options.resolution ?? '1k', response_format: 'b64_json', + ...(options.imageUrl ? { image: { url: options.imageUrl, type: 'image_url' } } : {}), }), signal: options.signal, }, diff --git a/src/imagine/imageUrl.ts b/src/imagine/imageUrl.ts new file mode 100644 index 0000000..23b0a56 --- /dev/null +++ b/src/imagine/imageUrl.ts @@ -0,0 +1,19 @@ +import { readFile } from 'node:fs/promises'; +import { getImageDimensions } from '@earendil-works/pi-tui'; + +export async function imageFileToDataUri(filePath: string, signal?: AbortSignal) { + const bytes = await readFile(filePath, { signal }); + const mime = bytes.subarray(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10])) + ? 'image/png' + : bytes.subarray(0, 3).equals(Buffer.from([255, 216, 255])) + ? 'image/jpeg' + : bytes.toString('ascii', 0, 4) === 'RIFF' && bytes.toString('ascii', 8, 12) === 'WEBP' + ? 'image/webp' + : undefined; + const data = bytes.toString('base64'); + const dimensions = mime ? getImageDimensions(data, mime) : null; + if (!mime || !dimensions || dimensions.widthPx < 1 || dimensions.heightPx < 1) { + throw new Error(`Unsupported image file: ${filePath}. Use a PNG, JPEG, or WebP image.`); + } + return `data:${mime};base64,${data}`; +} diff --git a/src/imagine/parseArgs.ts b/src/imagine/parseArgs.ts index b8125d0..dc4f3fa 100644 --- a/src/imagine/parseArgs.ts +++ b/src/imagine/parseArgs.ts @@ -19,6 +19,8 @@ export function parseImagineArgs(args: string) { ['--out', 'out'], ['-o', 'out'], ['--resolution', 'resolution'], + ['--image', 'image'], + ['--edit', 'image'], ]); for (let index = 0; index < tokens.length; index += 1) { @@ -42,6 +44,7 @@ export function parseImagineArgs(args: string) { prompt: prompt.join(' '), aspectRatio: normalizeAspectRatio(optionValues.get('aspect')), ...(optionValues.has('out') ? { outPath: optionValues.get('out') } : {}), + ...(optionValues.has('image') ? { imagePath: optionValues.get('image') } : {}), resolution, }; } diff --git a/src/imagine/register.ts b/src/imagine/register.ts index 5afd386..0df6275 100644 --- a/src/imagine/register.ts +++ b/src/imagine/register.ts @@ -63,17 +63,18 @@ export function registerImagineFeature( }); pi.registerCommand('grok-cli-imagine', { - description: 'Generate an image with Grok Imagine', + description: 'Generate or edit an image with Grok Imagine', handler: async (args, ctx) => { try { const parsed = parseImagineArgs(args); - ctx.ui.notify('Generating image…', 'info'); + ctx.ui.notify(parsed.imagePath ? 'Editing image…' : 'Generating image…', 'info'); const saved = await generateAndSaveImage( { ctx, prompt: parsed.prompt, aspectRatio: parsed.aspectRatio, resolution: parsed.resolution, + imagePath: parsed.imagePath, signal: ctx.signal, outPath: parsed.outPath ? isAbsolute(parsed.outPath) diff --git a/src/imagine/tool.ts b/src/imagine/tool.ts index 280c9b0..5570c83 100644 --- a/src/imagine/tool.ts +++ b/src/imagine/tool.ts @@ -11,7 +11,15 @@ import { } from './workflow.js'; const ImageGenParams = Type.Object({ - prompt: Type.String({ description: 'Text description of the image to generate.' }), + prompt: Type.String({ + description: 'Describe the image to generate, or the changes to apply to the source image.', + }), + image: Type.Optional( + Type.String({ + description: + 'Local PNG, JPEG, or WebP path to edit. Relative paths use the session working directory.', + }), + ), aspect_ratio: Type.Optional( Type.String({ description: @@ -47,7 +55,7 @@ export function registerImageGenTool( name: 'image_gen', label: 'Image Gen', description: - "Generate a new image from a text description using Imagine; returns the saved image's absolute path. For a request for one image, call this tool exactly once. Call it multiple times only when the user explicitly requests multiple images. Do not re-read or re-display the image unless the user asks.", + "Generate or edit an image with Grok Imagine; returns the saved image's absolute path. Pass image to edit an existing local file. For a request for one image, call this tool exactly once. Call it multiple times only when the user explicitly requests multiple images. Do not re-read or re-display the image unless the user asks.", promptGuidelines: [ 'For a request for one image, call image_gen exactly once. Call it multiple times only when the user explicitly requests multiple images.', 'Do not repeat the saved path unless the user asks for it; the image_gen result already displays a copyable path.', @@ -57,9 +65,11 @@ export function registerImageGenTool( try { const prompt = params.prompt.trim(); if (!prompt) throw new Error('Prompt is required'); + if (params.image !== undefined && !params.image.trim()) + throw new Error('Image path is required'); const aspectRatio = normalizeAspectRatio(params.aspect_ratio); const saved = await generateAndSaveImage( - { ctx, prompt, aspectRatio, signal }, + { ctx, prompt, aspectRatio, signal, imagePath: params.image?.trim() }, dependencies, resolveToken, ); diff --git a/src/imagine/workflow.ts b/src/imagine/workflow.ts index 63164c8..2842c9b 100644 --- a/src/imagine/workflow.ts +++ b/src/imagine/workflow.ts @@ -1,7 +1,9 @@ +import { resolve } from 'node:path'; import type { ExtensionContext } from '@earendil-works/pi-coding-agent'; import { convertToPng } from '@earendil-works/pi-coding-agent'; import { IMAGINE_AUTH_ERROR, resolveImagineToken } from './auth.js'; import { generateImage } from './generate.js'; +import { imageFileToDataUri } from './imageUrl.js'; import { saveImage, savePreviewImage } from './save.js'; export type ImagineDependencies = { @@ -31,6 +33,7 @@ export async function generateAndSaveImage( prompt: string; aspectRatio: string; resolution?: string; + imagePath?: string; signal?: AbortSignal; outPath?: string; }, @@ -45,6 +48,14 @@ export async function generateAndSaveImage( aspectRatio: options.aspectRatio, resolution: options.resolution, signal: options.signal, + ...(options.imagePath + ? { + imageUrl: await imageFileToDataUri( + resolve(options.ctx.cwd, options.imagePath), + options.signal, + ), + } + : {}), }); const persisted = options.ctx.sessionManager.getSessionFile() !== undefined; const saved = await dependencies.saveImage({ diff --git a/tests/imagine/generate.test.ts b/tests/imagine/generate.test.ts index 788251a..dfce74e 100644 --- a/tests/imagine/generate.test.ts +++ b/tests/imagine/generate.test.ts @@ -5,6 +5,23 @@ import { generateImage } from '../../src/imagine/generate.js'; afterEach(() => vi.useRealTimers()); describe('generateImage', () => { + it('sends a source image to the JSON editing endpoint', async () => { + const fetchImpl = vi.fn(async () => + Response.json({ data: [{ b64_json: '/9j/2Q==' }] }), + ); + await generateImage({ + token: 'secret', + prompt: 'Make it blue', + imageUrl: 'data:image/png;base64,source', + fetchImpl, + }); + expect(fetchImpl.mock.calls[0]?.[0]).toBe('https://api.x.ai/v1/images/edits'); + expect(JSON.parse(String(fetchImpl.mock.calls[0]?.[1]?.body))).toMatchObject({ + prompt: 'Make it blue', + image: { url: 'data:image/png;base64,source', type: 'image_url' }, + }); + }); + it('sends the captured Imagine request and returns JPEG base64', async () => { const fetchImpl = vi.fn(async () => Response.json({ data: [{ b64_json: '/9j/2Q==' }] }), diff --git a/tests/imagine/helpers.ts b/tests/imagine/helpers.ts index ace6dc2..89c1425 100644 --- a/tests/imagine/helpers.ts +++ b/tests/imagine/helpers.ts @@ -7,6 +7,9 @@ import type { savePreviewImage } from '../../src/imagine/save.js'; const tempDirs: string[] = []; +export const TEST_PNG_BASE64 = + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+aX1cAAAAASUVORK5CYII='; + afterEach(() => { for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true }); }); diff --git a/tests/imagine/register.test.ts b/tests/imagine/register.test.ts index 1623b03..34d2b19 100644 --- a/tests/imagine/register.test.ts +++ b/tests/imagine/register.test.ts @@ -5,7 +5,7 @@ import { describe, expect, it, vi } from 'vitest'; import { DEFAULT_CONFIG, loadConfig, saveConfig } from '../../src/config.js'; import { registerImagineFeature } from '../../src/imagine/register.js'; import { useEnvironmentToken, useTempHome } from '../stateTestHelpers.js'; -import { imagineDependencies } from './helpers.js'; +import { imagineDependencies, TEST_PNG_BASE64 } from './helpers.js'; const setupHome = useTempHome(); const setToken = useEnvironmentToken(); @@ -80,6 +80,23 @@ function setup( } describe('registerImagineFeature command', () => { + it.each(['--image', '--edit'])('edits a local image with %s', async (flag) => { + const extension = setup('token'); + writeFileSync(join(extension.home, 'source image.png'), Buffer.from(TEST_PNG_BASE64, 'base64')); + await extension.commands + .get('grok-cli-imagine') + ?.handler(`${flag} "source image.png" Make it blue`, { + ...extension.context, + cwd: extension.home, + }); + expect(extension.generate).toHaveBeenCalledWith( + expect.objectContaining({ + prompt: 'Make it blue', + imageUrl: `data:image/png;base64,${TEST_PNG_BASE64}`, + }), + ); + }); + it('registers the command, entry renderer, and image_gen tool', () => { const extension = setup('token'); expect(extension.commands.has('grok-cli-imagine')).toBe(true); diff --git a/tests/imagine/tool.test.ts b/tests/imagine/tool.test.ts index b4d66ea..2725a09 100644 --- a/tests/imagine/tool.test.ts +++ b/tests/imagine/tool.test.ts @@ -2,13 +2,13 @@ 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'; +import { imagineDependencies, TEST_PNG_BASE64 } from './helpers.js'; const setupHome = useTempHome(); const setToken = useEnvironmentToken(); function setup(token?: string, resolveToken?: () => Promise) { - setupHome(); + const home = setupHome(); setToken(token); let tool: Record | undefined; const dependencies = imagineDependencies(); @@ -22,6 +22,7 @@ function setup(token?: string, resolveToken?: () => Promise) resolveToken, ); const context = { + cwd: home, modelRegistry: { getApiKeyForProvider: vi.fn(async () => token) }, sessionManager: { getSessionDir: () => '/sessions', @@ -49,6 +50,41 @@ function setup(token?: string, resolveToken?: () => Promise) } describe('image_gen tool', () => { + it.each(['relative', 'absolute'])('edits a source image using its %s path', async (pathKind) => { + const test = setup('token'); + const path = join(test.context.cwd, 'source.bin'); + writeFileSync(path, Buffer.from(TEST_PNG_BASE64, 'base64')); + const result = await test.tool.execute( + 'edit', + { prompt: 'Make it blue', image: pathKind === 'absolute' ? path : 'source.bin' }, + undefined, + undefined, + test.context, + ); + expect(result.details.error).toBeUndefined(); + expect(test.generate).toHaveBeenCalledWith( + expect.objectContaining({ imageUrl: `data:image/png;base64,${TEST_PNG_BASE64}` }), + ); + }); + + it.each([ + 'not an image', + 'RIFF0000WAVEfmt ', + '\u0089PNG', + ])('rejects non-image or incomplete file content %j before generation', async (content) => { + const test = setup('token'); + writeFileSync(join(test.context.cwd, 'fake.png'), content); + const result = await test.tool.execute( + 'edit', + { prompt: 'Make it blue', image: 'fake.png' }, + undefined, + undefined, + test.context, + ); + expect(result.details.error).toMatch(/Unsupported image/); + expect(test.generate).not.toHaveBeenCalled(); + }); + it('returns path-only content and path details', async () => { const test = setup('token'); const signal = new AbortController().signal; @@ -174,3 +210,6 @@ describe('image_gen tool', () => { ).toContain('saved images/1.jpg (/missing.jpg)'); }); }); + +import { writeFileSync } from 'node:fs'; +import { join } from 'node:path'; diff --git a/tests/provider/package.test.ts b/tests/provider/package.test.ts index 7dd2eb0..7849206 100644 --- a/tests/provider/package.test.ts +++ b/tests/provider/package.test.ts @@ -55,6 +55,7 @@ describe('repository layout', () => { 'src/imagine/aspect.ts', 'src/imagine/auth.ts', 'src/imagine/generate.ts', + 'src/imagine/imageUrl.ts', 'src/imagine/parseArgs.ts', 'src/imagine/preview.ts', 'src/imagine/register.ts', From bd1865455ca143bfd17becb473c8c99698bb5b23 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Sun, 20 Sep 2026 10:59:38 +0900 Subject: [PATCH 4/5] feat(provider): recover proxy errors with bounded session retries Rotate and persist proxy conversation IDs on pre-stream HTTP 401, 502, or 520 failures, with at most two retries and a stable prompt-cache key. Preserve cancellation, partial streams, final error metadata, and account ownership. Adapted from Ayagikei/pi-grok-cli commits 08b47a5, ffd779e, and 1b87994, with Pi-native stream forwarding and retry guards. Co-authored-by: AyagiKei --- README.md | 3 + src/provider/proxyRetry.ts | 47 +++++++ src/provider/register.ts | 55 +++++--- src/provider/sessionConvId.ts | 63 +++++++++ tests/provider/package.test.ts | 2 + tests/provider/register.test.ts | 229 ++++++++++++++++++++++++++++++-- 6 files changed, 368 insertions(+), 31 deletions(-) create mode 100644 src/provider/proxyRetry.ts create mode 100644 src/provider/sessionConvId.ts diff --git a/README.md b/README.md index 290a848..b1fedaa 100644 --- a/README.md +++ b/README.md @@ -122,6 +122,7 @@ To edit a local PNG, JPEG, or WebP image, use `/grok-cli-imagine --image "./sour | --- | --- | | `/grok-cli-accounts [gui]` | Manage Grok accounts in the terminal, or add `gui` for the browser dashboard. | | `/grok-cli-usage` | Fetch current quota, update its cache, and show cached data if refresh fails. | +| `/grok-cli-conv [status\|rotate]` | Show or rotate this session's Grok proxy conversation ID. | | `/grok-cli-imagine ` | Generate or edit an image. Supports `--image`/`--edit`, `--aspect`, `--out`, and `--resolution 1k`. | | `/grok-cli-imagine:tool [on\|off\|status]` | Toggle, set, or report persistent model-callable `image_gen` availability. | @@ -154,6 +155,8 @@ See [Advanced configuration](./CONFIGURATION.md) for OAuth, callback, endpoint, ## Troubleshooting +For proxy HTTP 401, 502, or 520 errors before streaming starts, the extension rotates the conversation ID and retries up to twice. The prompt-cache key and selected account stay the same. Rotated IDs are saved in the Pi session. Recovery is best effort: an expired token still requires login. You can also rotate manually with `/grok-cli-conv rotate` before sending another request. + | Problem | What to do | | --- | --- | | grok-cli is missing from `/model` | Confirm the package appears in `pi list`, run `/login`, choose **Grok CLI**, then restart pi or run `/reload`. | diff --git a/src/provider/proxyRetry.ts b/src/provider/proxyRetry.ts new file mode 100644 index 0000000..2869b2e --- /dev/null +++ b/src/provider/proxyRetry.ts @@ -0,0 +1,47 @@ +import type { + AssistantMessage, + AssistantMessageEvent, + AssistantMessageEventStream, +} from '@earendil-works/pi-ai'; + +export async function* streamWithProxyRetry(options: { + start: () => AssistantMessageEventStream; + rotate?: () => void; + signal?: AbortSignal; + onMessage: (message: AssistantMessage) => void; +}): AsyncGenerator { + for (let attempt = 0; ; attempt += 1) { + const stream = options.start(); + let started = false; + for await (const event of stream) { + if (event.type === 'error') break; + if (event.type === 'done') { + options.onMessage(event.message); + yield event; + return; + } + started = true; + yield event; + } + + const message = await stream.result(); + if ( + !started && + !options.signal?.aborted && + message.stopReason === 'error' && + /^OpenAI API error \((401|502|520)\)/.test(message.errorMessage ?? '') && + attempt < 2 && + options.rotate + ) { + options.rotate(); + continue; + } + options.onMessage(message); + if (message.stopReason === 'error' || message.stopReason === 'aborted') { + yield { type: 'error', reason: message.stopReason, error: message }; + } else { + yield { type: 'done', reason: message.stopReason, message }; + } + return; + } +} diff --git a/src/provider/register.ts b/src/provider/register.ts index a725552..fd7a930 100644 --- a/src/provider/register.ts +++ b/src/provider/register.ts @@ -28,10 +28,12 @@ import { mutateAccountVault, } from './accountVault.js'; import { migrateSavedModelProviders } from './modelMigration.js'; +import { streamWithProxyRetry } from './proxyRetry.js'; import { removeQuotaUsage } from './quotaCache.js'; import { rememberRequestAccount } from './requestOwnership.js'; import { registerExhaustionRotation } from './rotation.js'; import { createSessionAccountSelection } from './sessionAccountSelection.js'; +import { registerSessionConvId } from './sessionConvId.js'; import { grokCliModelHeaders } from './stream.js'; import { registerUsageCommand } from './usage.js'; @@ -53,6 +55,7 @@ function accountCredential(credentials: OAuthCredentials): AccountCredential { export default function registerGrokCli(pi: ExtensionAPI) { const sessionSelection = createSessionAccountSelection(pi); + const convIds = registerSessionConvId(pi); let migrationWarning: string | undefined; let migrationError: string | undefined; let migrationErrorNotified = false; @@ -157,30 +160,40 @@ export default function registerGrokCli(pi: ExtensionAPI) { headers: grokCliModelHeaders(model.id), })), streamSimple(model, context, options?: SimpleStreamOptions) { - const accountId = sessionSelection.accountId(options?.sessionId); + const sessionId = options?.sessionId; + const accountId = sessionSelection.accountId(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; + return streamWithProxyRetry({ + start: () => + streamSimpleOpenAIResponses( + { + ...model, + baseUrl: route.baseUrl, + api: 'openai-responses', + } as Model<'openai-responses'>, + context, + { + ...options, + apiKey: route.token, + // Keep the retry budget here; SDK retries would reuse the failed conversation ID. + maxRetries: 0, + headers: { + ...options?.headers, + ...(sessionId ? { 'x-grok-conv-id': convIds.convId(sessionId) } : {}), + }, + }, + ), + rotate: sessionId + ? () => { + convIds.rotate(sessionId); + } + : undefined, + signal: options?.signal, + onMessage: (message) => rememberRequestAccount(message, route.accountId), + }); }); }, }); @@ -240,7 +253,7 @@ export default function registerGrokCli(pi: ExtensionAPI) { pi.on('before_provider_headers', (event, ctx) => { if (ctx.model?.provider !== 'grok-cli') return; - event.headers['x-grok-conv-id'] = ctx.sessionManager.getSessionId(); + event.headers['x-grok-conv-id'] = convIds.convId(ctx.sessionManager.getSessionId()); }); pi.on('before_provider_request', (event, ctx) => { diff --git a/src/provider/sessionConvId.ts b/src/provider/sessionConvId.ts new file mode 100644 index 0000000..a31174b --- /dev/null +++ b/src/provider/sessionConvId.ts @@ -0,0 +1,63 @@ +import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'; + +const SESSION_CONV_ENTRY = 'grok-cli-conv-id-v1'; + +function storedGeneration( + entry: ReturnType[number], +) { + if (entry.type !== 'custom' || entry.customType !== SESSION_CONV_ENTRY) return undefined; + if (!entry.data || typeof entry.data !== 'object' || Array.isArray(entry.data)) return undefined; + const generation = (entry.data as Record).generation; + return typeof generation === 'number' && Number.isSafeInteger(generation) && generation > 0 + ? generation + : undefined; +} + +export function registerSessionConvId(pi: ExtensionAPI) { + const generations = new Map(); + const convId = (sessionId: string) => { + const generation = generations.get(sessionId); + return generation ? `${sessionId}:${generation}` : sessionId; + }; + const rotate = (sessionId: string) => { + const generation = (generations.get(sessionId) ?? 0) + 1; + pi.appendEntry(SESSION_CONV_ENTRY, { generation }); + generations.set(sessionId, generation); + return convId(sessionId); + }; + const restore = (ctx: Pick) => { + const sessionId = ctx.sessionManager.getSessionId(); + const generation = ctx.sessionManager + .getBranch() + .reduceRight( + (generation, entry) => generation ?? storedGeneration(entry), + undefined, + ); + if (generation) generations.set(sessionId, generation); + else generations.delete(sessionId); + }; + + pi.on('session_start', (_event, ctx) => restore(ctx)); + pi.on('session_tree', (_event, ctx) => restore(ctx)); + pi.on('session_shutdown', (_event, ctx) => { + generations.delete(ctx.sessionManager.getSessionId()); + }); + pi.registerCommand('grok-cli-conv', { + description: 'Show or rotate the Grok CLI conversation ID', + handler: async (args, ctx) => { + const argument = args.trim().toLowerCase(); + if (argument && argument !== 'status' && argument !== 'rotate') { + ctx.ui.notify('Usage: /grok-cli-conv [status|rotate]', 'error'); + return; + } + const sessionId = ctx.sessionManager.getSessionId(); + ctx.ui.notify( + argument === 'rotate' + ? `Grok CLI conversation ID rotated to ${rotate(sessionId)}` + : `Grok CLI conversation ID: ${convId(sessionId)}`, + 'info', + ); + }, + }); + return { convId, rotate }; +} diff --git a/tests/provider/package.test.ts b/tests/provider/package.test.ts index 7849206..b42511e 100644 --- a/tests/provider/package.test.ts +++ b/tests/provider/package.test.ts @@ -72,11 +72,13 @@ describe('repository layout', () => { 'src/provider/billing.ts', 'src/provider/dashboard/server.ts', 'src/provider/modelMigration.ts', + 'src/provider/proxyRetry.ts', 'src/provider/quotaCache.ts', 'src/provider/register.ts', 'src/provider/requestOwnership.ts', 'src/provider/rotation.ts', 'src/provider/sessionAccountSelection.ts', + 'src/provider/sessionConvId.ts', 'src/provider/stream.ts', 'src/provider/usage.ts', 'src/shared/errors.ts', diff --git a/tests/provider/register.test.ts b/tests/provider/register.test.ts index 5107035..1adbb29 100644 --- a/tests/provider/register.test.ts +++ b/tests/provider/register.test.ts @@ -1,7 +1,12 @@ import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs'; +import { createServer } from 'node:http'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; -import type { OAuthLoginCallbacks } from '@earendil-works/pi-ai'; +import { + type AssistantMessage, + createAssistantMessageEventStream, + 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 { @@ -35,7 +40,7 @@ vi.mock('@earendil-works/pi-ai/compat', () => ({ })); interface CommandConfig { - handler: (args: string[], ctx: TestContext) => Promise; + handler: (args: string, ctx: TestContext) => Promise; } interface RegisteredTool { @@ -221,6 +226,210 @@ async function drain(stream: AsyncIterable | undefined) { } } +function proxyResponse(status?: number, started = false) { + const message: AssistantMessage = { + role: 'assistant', + content: [], + api: 'openai-responses', + provider: 'grok-cli', + model: 'grok-4.6', + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: status ? 'error' : 'stop', + ...(status ? { errorMessage: `OpenAI API error (${status}): proxy failure` } : {}), + timestamp: Date.now(), + }; + const stream = createAssistantMessageEventStream(); + if (started) stream.push({ type: 'start', partial: message }); + if (status) stream.push({ type: 'error', reason: 'error', error: message }); + else stream.push({ type: 'done', reason: 'stop', message }); + stream.end(message); + return stream; +} + +async function startProxyRequest( + extension: Awaited>, + options = {}, +) { + const provider = extension.providers.get('grok-cli'); + const model = provider?.models?.[0]; + if (!model || !provider.streamSimple) throw new Error('Grok CLI test model is missing.'); + return provider.streamSimple( + { + ...model, + provider: 'grok-cli', + api: 'openai-responses', + baseUrl: 'https://cli-chat-proxy.grok.com', + }, + { messages: [] }, + options, + ); +} + +describe('proxy conversation recovery', () => { + it('recovers through the real Pi HTTP adapter while preserving the prompt cache key', async () => { + await setAccount1Credential('one'); + writePiVaultMarker(); + const actual = await vi.importActual( + '@earendil-works/pi-ai/compat', + ); + mockProviderStream.mockImplementation(actual.streamSimpleOpenAIResponses); + const requests: { + conversationId: string | string[] | undefined; + cacheKey: unknown; + authorization: string | undefined; + }[] = []; + const server = createServer(async (request, response) => { + const chunks: Buffer[] = []; + for await (const chunk of request) chunks.push(Buffer.from(chunk)); + const payload = JSON.parse(Buffer.concat(chunks).toString()); + requests.push({ + conversationId: request.headers['x-grok-conv-id'], + cacheKey: payload.prompt_cache_key, + authorization: request.headers.authorization, + }); + if (requests.length < 3) { + response.writeHead(requests.length === 1 ? 401 : 520, { + 'Content-Type': 'application/json', + }); + response.end(JSON.stringify({ error: { message: 'proxy failure' } })); + return; + } + response.writeHead(200, { 'Content-Type': 'text/event-stream' }); + response.end( + `data: ${JSON.stringify({ type: 'response.completed', response: { id: 'response-1', status: 'completed', output: [], usage: { input_tokens: 1, output_tokens: 0, total_tokens: 1 } } })}\n\n`, + ); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (!address || typeof address === 'string') throw new Error('Missing test server address'); + vi.stubEnv('PI_GROK_CLI_BASE_URL', `http://127.0.0.1:${address.port}/v1`); + try { + const extension = await setupExtension(); + const stream = await startProxyRequest(extension, { sessionId: 'session-a', maxRetries: 5 }); + await drain(stream); + expect(await stream.result()).toMatchObject({ stopReason: 'stop', usage: { input: 1 } }); + expect(requests).toEqual( + ['session-a', 'session-a:1', 'session-a:2'].map((conversationId) => ({ + conversationId, + cacheKey: 'session-a', + authorization: 'Bearer one', + })), + ); + } finally { + vi.unstubAllEnvs(); + await new Promise((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())), + ); + } + }); + + it.each([ + 401, 502, 520, + ])('retries HTTP %i twice with new conversation IDs and the same cache key', async (status) => { + await setAccount1Credential('one'); + writePiVaultMarker(); + mockProviderStream + .mockImplementationOnce(() => proxyResponse(status)) + .mockImplementationOnce(() => proxyResponse(status)) + .mockImplementationOnce(() => proxyResponse()); + const extension = await setupExtension(); + const onPayload = vi.fn(); + const stream = await startProxyRequest(extension, { + sessionId: 'session-a', + headers: { custom: 'kept' }, + onPayload, + }); + const events = []; + for await (const event of stream) events.push(event); + + expect(events.map((event) => event.type)).toEqual(['done']); + expect((await stream.result()).stopReason).toBe('stop'); + expect(mockProviderStream.mock.calls.map((call) => call[2]?.headers)).toEqual([ + { custom: 'kept', 'x-grok-conv-id': 'session-a' }, + { custom: 'kept', 'x-grok-conv-id': 'session-a:1' }, + { custom: 'kept', 'x-grok-conv-id': 'session-a:2' }, + ]); + for (const call of mockProviderStream.mock.calls) { + expect(call[2]).toMatchObject({ apiKey: 'one', sessionId: 'session-a', onPayload }); + } + expect(extension.entries.filter((entry) => entry.customType === 'grok-cli-conv-id-v1')).toEqual( + [ + { customType: 'grok-cli-conv-id-v1', data: { generation: 1 } }, + { customType: 'grok-cli-conv-id-v1', data: { generation: 2 } }, + ], + ); + }); + + it.each([ + { status: 502, started: false, sessionId: 'session-a', abort: false, attempts: 3 }, + { status: 429, started: false, sessionId: 'session-a', abort: false, attempts: 1 }, + { status: 403, started: false, sessionId: 'session-a', abort: false, attempts: 1 }, + { status: 502, started: true, sessionId: 'session-a', abort: false, attempts: 1 }, + { status: 502, started: false, sessionId: undefined, abort: false, attempts: 1 }, + { status: 502, started: false, sessionId: 'session-a', abort: true, attempts: 1 }, + ])('delivers terminal failures without unsafe replay: $status, started=$started, abort=$abort, attempts=$attempts', async (scenario) => { + await setAccount1Credential('one'); + writePiVaultMarker(); + const controller = new AbortController(); + mockProviderStream.mockImplementation(() => { + if (scenario.abort) controller.abort(); + return proxyResponse(scenario.status, scenario.started); + }); + const extension = await setupExtension(); + const stream = await startProxyRequest(extension, { + sessionId: scenario.sessionId, + signal: controller.signal, + }); + const events = []; + for await (const event of stream) events.push(event); + expect(mockProviderStream).toHaveBeenCalledTimes(scenario.attempts); + expect(events.map((event) => event.type)).toEqual( + scenario.started ? ['start', 'error'] : ['error'], + ); + expect(await stream.result()).toMatchObject({ + role: 'assistant', + provider: 'grok-cli', + stopReason: 'error', + errorMessage: expect.stringContaining(`(${scenario.status})`), + }); + const { requestAccount } = await import('../../src/provider/requestOwnership.js'); + expect(requestAccount(await stream.result())).toBe('account-1'); + }); + + it('restores rotated IDs on reopen and follows the selected session branch', async () => { + const extension = await setupExtension(); + const command = extension.commands.get('grok-cli-conv'); + expect(command).toBeDefined(); + const ctx = sessionContext('session-a'); + ctx.model = { provider: 'grok-cli', id: 'grok-4.6' }; + await command?.handler('rotate', ctx); + await command?.handler('rotate', ctx); + const saved = extension.entries.filter((entry) => entry.customType === 'grok-cli-conv-id-v1'); + const reopened = await setupExtension(); + if (!ctx.sessionManager) throw new Error('Missing session manager'); + ctx.sessionManager.getBranch = () => saved.map((entry) => ({ type: 'custom', ...entry })); + await reopened.emit('session_start', {}, ctx); + const headers = { 'x-grok-conv-id': '' }; + await reopened.emit('before_provider_headers', { headers }, ctx); + expect(headers['x-grok-conv-id']).toBe('session-a:2'); + ctx.sessionManager.getBranch = () => []; + await reopened.emit('session_tree', {}, ctx); + await reopened.emit('before_provider_headers', { headers }, ctx); + expect(headers['x-grok-conv-id']).toBe('session-a'); + const other = sessionContext('session-b'); + other.model = ctx.model; + await reopened.emit('before_provider_headers', { headers }, other); + expect(headers['x-grok-conv-id']).toBe('session-b'); + }); +}); + function statusContext(notify: TestContext['ui']['notify']): TestContext { return { modelRegistry: { @@ -312,7 +521,7 @@ const billingFetchMock = ( async function runStatus(extension: Awaited>) { const notify = vi.fn(); - await extension.commands.get('grok-cli-usage')?.handler([], statusContext(notify)); + await extension.commands.get('grok-cli-usage')?.handler('', statusContext(notify)); return notify; } @@ -448,7 +657,7 @@ describe('Grok CLI status command', () => { const notify = vi.fn(); const getApiKeyForProvider = vi.fn(async () => 'provider-token'); - await extension.commands.get('grok-cli-usage')?.handler([], { + await extension.commands.get('grok-cli-usage')?.handler('', { ...statusContext(notify), modelRegistry: { ...statusContext(notify).modelRegistry, @@ -487,7 +696,7 @@ describe('Grok CLI status command', () => { const extension = await setupExtension(); const notify = vi.fn(); - const pending = extension.commands.get('grok-cli-usage')?.handler([], statusContext(notify)); + 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; @@ -603,7 +812,7 @@ describe('Grok CLI status command', () => { const extension = await setupExtension(); const notify = vi.fn(); - await extension.commands.get('grok-cli-usage')?.handler([], emptyStatusContext(notify)); + await extension.commands.get('grok-cli-usage')?.handler('', emptyStatusContext(notify)); expect(notify).toHaveBeenCalledOnce(); expect(notify).toHaveBeenCalledWith( @@ -617,7 +826,7 @@ describe('Grok CLI status command', () => { const extension = await setupExtension(); const notify = vi.fn(); - await extension.commands.get('grok-cli-usage')?.handler([], { + await extension.commands.get('grok-cli-usage')?.handler('', { modelRegistry: { getAll: () => Array.from({ length: 7 }, (_value, index) => ({ @@ -638,7 +847,7 @@ describe('Grok CLI status command', () => { const extension = await setupExtension(); const notify = vi.fn(); - await extension.commands.get('grok-cli-usage')?.handler([], { + await extension.commands.get('grok-cli-usage')?.handler('', { modelRegistry: { getAll: () => { throw new Error('registry unavailable'); @@ -655,7 +864,7 @@ describe('Grok CLI status command', () => { const extension = await setupExtension(); const notify = vi.fn(); - await extension.commands.get('grok-cli-usage')?.handler([], { + await extension.commands.get('grok-cli-usage')?.handler('', { modelRegistry: { getAll: () => { throw new XaiOAuthError('refresh failed', 'refresh_failed', true); @@ -831,7 +1040,7 @@ describe('Grok CLI provider registration', () => { 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); + await extension.commands.get('grok-cli-usage')?.handler('', context); expect(globalThis.fetch).toHaveBeenCalledWith( 'https://cli-chat-proxy.grok.com/v1/billing', From 84aed88d306523d738c97950ec89e9b9742c599c Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Sun, 20 Sep 2026 12:31:19 +0900 Subject: [PATCH 5/5] fix: bound image inputs and preserve proxy failures Why: - Oversized source images can exhaust memory, and session write errors can replace the original proxy failure. What: - Cap source images at 400 KiB before encoding, including files that grow during the read, and document the limit. - Stop recovery when conversation-ID persistence fails while preserving the proxy error and account ownership. - Cover the image-size boundary and failed rotation persistence through the registered tools and provider. Validation: - bun run check passed with 321 tests after both new regression cases failed before the fixes. - git diff --cached --check passed. LOC: src/ + 25 / - 6 = net +19 tests/ + 56 / - 3 = net +53 *.md + 1 / - 1 = net +0 Total + 82 / -10 = net +72 --- README.md | 2 +- src/imagine/imageUrl.ts | 19 +++++++++++++++++-- src/provider/proxyRetry.ts | 12 ++++++++---- tests/imagine/tool.test.ts | 27 +++++++++++++++++++++++++++ tests/provider/register.test.ts | 32 +++++++++++++++++++++++++++++--- 5 files changed, 82 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index b1fedaa..d4bc7da 100644 --- a/README.md +++ b/README.md @@ -112,7 +112,7 @@ Models are bundled rather than discovered live. Registered context limits may di 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. -To edit a local PNG, JPEG, or WebP image, use `/grok-cli-imagine --image "./source image.png" `. `--edit` is an alias. The `image_gen` tool accepts the same local path in its optional `image` argument. Relative paths use the session working directory. The source file is uploaded to Imagine; the edited image is saved separately unless you explicitly use `--out` to overwrite the source. +To edit a local PNG, JPEG, or WebP image, use `/grok-cli-imagine --image "./source image.png" `. `--edit` is an alias. The `image_gen` tool accepts the same local path in its optional `image` argument. Relative paths use the session working directory. Source files must be at most 400 KiB; resize or compress larger images before editing. The source file is uploaded to Imagine; the edited image is saved separately unless you explicitly use `--out` to overwrite the source. `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. diff --git a/src/imagine/imageUrl.ts b/src/imagine/imageUrl.ts index 23b0a56..dd4d5bd 100644 --- a/src/imagine/imageUrl.ts +++ b/src/imagine/imageUrl.ts @@ -1,8 +1,23 @@ -import { readFile } from 'node:fs/promises'; +import { createReadStream } from 'node:fs'; +import { stat } from 'node:fs/promises'; +import { buffer } from 'node:stream/consumers'; import { getImageDimensions } from '@earendil-works/pi-tui'; +const MAX_SOURCE_BYTES = 400 * 1024; + export async function imageFileToDataUri(filePath: string, signal?: AbortSignal) { - const bytes = await readFile(filePath, { signal }); + if ((await stat(filePath)).size > MAX_SOURCE_BYTES) { + throw new Error( + 'Source image exceeds the 400 KiB limit. Resize or compress it before editing.', + ); + } + // Read at most one byte beyond the limit, even if the file grows after stat. + const bytes = await buffer(createReadStream(filePath, { end: MAX_SOURCE_BYTES, signal })); + if (bytes.length > MAX_SOURCE_BYTES) { + throw new Error( + 'Source image exceeds the 400 KiB limit. Resize or compress it before editing.', + ); + } const mime = bytes.subarray(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10])) ? 'image/png' : bytes.subarray(0, 3).equals(Buffer.from([255, 216, 255])) diff --git a/src/provider/proxyRetry.ts b/src/provider/proxyRetry.ts index 2869b2e..3ecf04b 100644 --- a/src/provider/proxyRetry.ts +++ b/src/provider/proxyRetry.ts @@ -33,15 +33,19 @@ export async function* streamWithProxyRetry(options: { attempt < 2 && options.rotate ) { - options.rotate(); - continue; + try { + options.rotate(); + continue; + } catch { + // A failed session write must not replace the original proxy error. + } } options.onMessage(message); if (message.stopReason === 'error' || message.stopReason === 'aborted') { yield { type: 'error', reason: message.stopReason, error: message }; - } else { - yield { type: 'done', reason: message.stopReason, message }; + return; } + yield { type: 'done', reason: message.stopReason, message }; return; } } diff --git a/tests/imagine/tool.test.ts b/tests/imagine/tool.test.ts index 2725a09..6a3f108 100644 --- a/tests/imagine/tool.test.ts +++ b/tests/imagine/tool.test.ts @@ -50,6 +50,33 @@ function setup(token?: string, resolveToken?: () => Promise) } describe('image_gen tool', () => { + it.each([ + 400 * 1024 - 1, + 400 * 1024, + 400 * 1024 + 1, + ])('enforces the source-image size limit for a %i-byte file', async (size) => { + const test = setup('token'); + const bytes = Buffer.alloc(size); + Buffer.from(TEST_PNG_BASE64, 'base64').copy(bytes); + writeFileSync(join(test.context.cwd, 'source.png'), bytes); + const result = await test.tool.execute( + 'edit', + { prompt: 'Make it blue', image: 'source.png' }, + undefined, + undefined, + test.context, + ); + if (size > 400 * 1024) { + expect(result.details.error).toMatch(/400 KiB/); + expect(test.generate).not.toHaveBeenCalled(); + return; + } + expect(result.details.error).toBeUndefined(); + expect(test.generate).toHaveBeenCalledWith( + expect.objectContaining({ imageUrl: `data:image/png;base64,${bytes.toString('base64')}` }), + ); + }); + it.each(['relative', 'absolute'])('edits a source image using its %s path', async (pathKind) => { const test = setup('token'); const path = join(test.context.cwd, 'source.bin'); diff --git a/tests/provider/register.test.ts b/tests/provider/register.test.ts index 1adbb29..8030984 100644 --- a/tests/provider/register.test.ts +++ b/tests/provider/register.test.ts @@ -140,6 +140,9 @@ async function setupExtension(initialActiveTools = ['read', 'bash']) { const setModel = vi.fn(async (_model: { provider: string; id: string }) => true); const sendUserMessage = vi.fn(); const entries: { customType: string; data: unknown }[] = []; + const appendEntry = vi.fn((customType: string, data: unknown) => { + entries.push({ customType, data }); + }); const registerGrokCli = (await import('../../src/index.js')).default; registerGrokCli({ registerProvider(name: string, config: ProviderConfig) { @@ -153,9 +156,7 @@ async function setupExtension(initialActiveTools = ['read', 'bash']) { commands.set(name, config as CommandConfig); }, registerEntryRenderer() {}, - appendEntry(customType: string, data: unknown) { - entries.push({ customType, data }); - }, + appendEntry, registerTool(tool: RegisteredTool) { tools.set(tool.name, tool); }, @@ -182,6 +183,7 @@ async function setupExtension(initialActiveTools = ['read', 'bash']) { } as unknown as ExtensionAPI); return { commands, + appendEntry, providers, tools, handlers, @@ -273,6 +275,30 @@ async function startProxyRequest( } describe('proxy conversation recovery', () => { + it('preserves the proxy error and account ownership when rotation persistence fails', async () => { + await setAccount1Credential('one'); + writePiVaultMarker(); + mockProviderStream.mockImplementation(() => proxyResponse(502)); + const extension = await setupExtension(); + extension.appendEntry.mockImplementation(() => { + throw new Error('Session storage is full'); + }); + const stream = await startProxyRequest(extension, { sessionId: 'session-a' }); + const events = []; + for await (const event of stream) events.push(event); + expect(events.map((event) => event.type)).toEqual(['error']); + expect(mockProviderStream).toHaveBeenCalledTimes(1); + expect(await stream.result()).toMatchObject({ + stopReason: 'error', + errorMessage: 'OpenAI API error (502): proxy failure', + }); + const { requestAccount } = await import('../../src/provider/requestOwnership.js'); + expect(requestAccount(await stream.result())).toBe('account-1'); + const ctx = sessionContext('session-a'); + await extension.commands.get('grok-cli-conv')?.handler('status', ctx); + expect(ctx.ui.notify).toHaveBeenCalledWith('Grok CLI conversation ID: session-a', 'info'); + }); + it('recovers through the real Pi HTTP adapter while preserving the prompt cache key', async () => { await setAccount1Credential('one'); writePiVaultMarker();