feat: add Extra High reasoning, account defaults, image editing, and proxy recovery

README.md

@@ -112,6 +112,8 @@ Models are bundled rather than discovered live. Registered context limits may di
 
 Run `/grok-cli-imagine <prompt>` 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 instructions>`. `--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.
 
 ## Commands
@@ -120,7 +122,8 @@ Run `/grok-cli-imagine <prompt>` 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 <prompt>` | Generate and preview an image. Supports `--aspect`, `--out`, and `--resolution 1k`. |
+| `/grok-cli-conv [status\|rotate]` | Show or rotate this session's Grok proxy conversation ID. |
+| `/grok-cli-imagine <prompt>` | 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
@@ -152,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`. |

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.
 

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,
     },

src/imagine/imageUrl.ts

@@ -0,0 +1,34 @@
+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) {
+  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]))
+      ? '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}`;
+}

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,
   };
 }

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)

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,
         );

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({

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',

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);

src/provider/proxyRetry.ts

@@ -0,0 +1,51 @@
+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<AssistantMessageEvent> {
+  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
+    ) {
+      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 };
+      return;
+    }
+    yield { type: 'done', reason: message.stopReason, message };
+    return;
+  }
+}

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) => {