From 34ed8613abd37bb1040d05d00ef146857055a0e0 Mon Sep 17 00:00:00 2001 From: Martin Date: Tue, 30 Jun 2026 22:34:53 +0200 Subject: [PATCH 1/3] feat(ai-proxy): Grok subscription provider and upstream routing --- src/ai-proxy/lib/model-meta.ts | 31 +- .../lib/providers/grok-subscription.ts | 59 ++- src/ai-proxy/lib/resolve-model.test.ts | 64 ++- src/ai-proxy/lib/resolve-model.ts | 92 +++- .../lib/rewrite-upstream-body.test.ts | 250 ++++++++- src/ai-proxy/lib/rewrite-upstream-body.ts | 473 +++++++++++++++++- src/ai-proxy/lib/server.ts | 9 +- src/ai-proxy/lib/thinking-folded.ts | 7 + src/utils/ai/grok/client.ts | 40 +- src/utils/ai/grok/models.test.ts | 13 + src/utils/ai/grok/models.ts | 16 +- src/utils/ai/grok/probe.test.ts | 8 +- 12 files changed, 999 insertions(+), 63 deletions(-) create mode 100644 src/utils/ai/grok/models.test.ts diff --git a/src/ai-proxy/lib/model-meta.ts b/src/ai-proxy/lib/model-meta.ts index 92c9e55f52..0feb08121a 100644 --- a/src/ai-proxy/lib/model-meta.ts +++ b/src/ai-proxy/lib/model-meta.ts @@ -8,6 +8,26 @@ import type { CopilotModelRecord } from "@app/utils/ai/github-copilot/types"; import type { GrokModelRecord } from "@app/utils/ai/grok"; import { GROK_STATIC_CATALOG, toProxyId } from "@app/utils/ai/grok"; +import { SafeJSON } from "@app/utils/json"; + +export function buildGrokModelDescription(meta: { + visibility: string; + speed: string; + thinking: string; + contextWindow?: number; + agentType?: string; + probeStatus?: string; +}): string { + return SafeJSON.stringify({ + visibility: meta.visibility, + speed: meta.speed, + thinking: meta.thinking, + contextWindow: meta.contextWindow, + agentType: meta.agentType, + probeStatus: meta.probeStatus, + }); +} + export function grokRecordToProxyMeta( account: AiProxyAccountConfig, record: GrokModelRecord, @@ -32,7 +52,16 @@ export function grokRecordToProxyMeta( billingPlane: "subscription", source: record.source, probeStatus: record.probeStatus, - description: record.description, + description: + record.description ?? + buildGrokModelDescription({ + visibility: record.visibility, + speed: record.speed, + thinking: record.thinking, + contextWindow: record.context_window, + agentType: record.agent_type, + probeStatus: record.probeStatus, + }), object: "model", created: 1_740_960_000, owned_by: providerKey(account), diff --git a/src/ai-proxy/lib/providers/grok-subscription.ts b/src/ai-proxy/lib/providers/grok-subscription.ts index d2626a05ec..3021ad0219 100644 --- a/src/ai-proxy/lib/providers/grok-subscription.ts +++ b/src/ai-proxy/lib/providers/grok-subscription.ts @@ -44,7 +44,7 @@ export class GrokSubscriptionProvider implements ProxyProvider { object: "model", created: model.created, owned_by: model.owned_by, - description: SafeDescription(model), + description: model.description, })); } @@ -75,14 +75,14 @@ export class GrokSubscriptionProvider implements ProxyProvider { private async forward(path: string, upstreamModel: string, bodyText: string, req: Request): Promise { const target = path.includes("responses") ? "responses" : "chat"; - const upstreamBody = prepareGrokUpstreamBody(bodyText, upstreamModel, target); + const prepared = prepareGrokUpstreamBody(bodyText, upstreamModel, target); const started = performance.now(); try { const upstream = await this.client.fetch(path, { method: "POST", - body: upstreamBody, - modelOverride: upstreamModel, + body: prepared.bodyText, + modelOverride: prepared.upstreamModel, signal: req.signal, headers: { Accept: req.headers.get("Accept") ?? "application/json", @@ -96,7 +96,9 @@ export class GrokSubscriptionProvider implements ProxyProvider { logger.warn( { account: this.account.name, - upstreamModel, + upstreamModel: prepared.upstreamModel, + requestedModel: upstreamModel, + imageRouted: prepared.imageRouted, path, status: upstream.status, elapsedMs, @@ -106,7 +108,15 @@ export class GrokSubscriptionProvider implements ProxyProvider { ); } else { logger.debug( - { account: this.account.name, upstreamModel, path, status: upstream.status, elapsedMs }, + { + account: this.account.name, + upstreamModel: prepared.upstreamModel, + requestedModel: upstreamModel, + imageRouted: prepared.imageRouted, + path, + status: upstream.status, + elapsedMs, + }, "ai-proxy: upstream request ok" ); } @@ -117,16 +127,29 @@ export class GrokSubscriptionProvider implements ProxyProvider { }); } catch (err) { if (err instanceof GrokAuthExpiredError) { + logger.warn( + { + account: this.account.name, + upstreamModel: prepared.upstreamModel, + requestedModel: upstreamModel, + imageRouted: prepared.imageRouted, + path, + elapsedMs: Math.round(performance.now() - started), + authPath: err.authPath, + }, + "ai-proxy: synthesizing 502 from GrokAuthExpiredError (upstream said 401/403 — see prior 'grok: upstream returned auth-status' log for body)" + ); + return new Response( SafeJSON.stringify({ error: { - message: err.message, - type: "auth_error", + message: `Upstream Grok auth expired or invalid — the ai-proxy host must refresh its Grok login. ${err.message}`, + type: "upstream_auth_error", code: "grok_auth_expired", }, }), { - status: 401, + status: 502, headers: { "Content-Type": "application/json" }, } ); @@ -136,21 +159,3 @@ export class GrokSubscriptionProvider implements ProxyProvider { } } } - -function SafeDescription(model: { - visibility: string; - speed: string; - thinking: string; - contextWindow?: number; - agentType?: string; - probeStatus?: string; -}): string { - return SafeJSON.stringify({ - visibility: model.visibility, - speed: model.speed, - thinking: model.thinking, - contextWindow: model.contextWindow, - agentType: model.agentType, - probeStatus: model.probeStatus, - }); -} diff --git a/src/ai-proxy/lib/resolve-model.test.ts b/src/ai-proxy/lib/resolve-model.test.ts index da3341538b..f8c5f3b381 100644 --- a/src/ai-proxy/lib/resolve-model.test.ts +++ b/src/ai-proxy/lib/resolve-model.test.ts @@ -1,5 +1,10 @@ import { describe, expect, it } from "bun:test"; import { parseProxyModelId, resolveModel } from "@app/ai-proxy/lib/resolve-model"; +import type { AiProxyAccountConfig } from "@app/ai-proxy/lib/types"; + +function grokAccount(name: string): AiProxyAccountConfig { + return { name, provider: "grok-subscription", providerSlug: "grok", enabled: true }; +} describe("resolve-model", () => { it("parses three-segment proxy ids", () => { @@ -14,24 +19,65 @@ describe("resolve-model", () => { expect(parseProxyModelId("genesiscz/grok/grok-build/extra").upstreamId).toBe("grok-build/extra"); }); - it("rejects bare upstream ids", () => { + it("rejects bare upstream ids in parseProxyModelId", () => { expect(() => parseProxyModelId("grok-build")).toThrow("must be //"); }); + it("rejects unknown bare ids when no account matches", () => { + expect(() => resolveModel("composer-2.5", [])).toThrow("No enabled account for model"); + expect(() => resolveModel("composer-2.5[fast=false]", [])).toThrow("No enabled account for model"); + }); + + it("resolves bare upstream ids", () => { + const accounts = [grokAccount("martin")]; + + const route = resolveModel("grok-build-0.1", accounts); + + expect(route.accountName).toBe("martin"); + expect(route.providerSlug).toBe("grok"); + expect(route.upstreamId).toBe("grok-build-0.1"); + }); + + it("resolves provider/upstream shorthand ids", () => { + const accounts = [grokAccount("martin")]; + + const route = resolveModel("grok/grok-build-0.1", accounts); + + expect(route.accountName).toBe("martin"); + expect(route.providerSlug).toBe("grok"); + expect(route.upstreamId).toBe("grok-build-0.1"); + }); + + it("resolves bare upstream ids when another enabled account uses an unimplemented provider", () => { + const accounts: AiProxyAccountConfig[] = [ + grokAccount("martin"), + { name: "work", provider: "xai-api-key", providerSlug: "xai", enabled: true }, + ]; + + const route = resolveModel("grok-build-0.1", accounts); + + expect(route.accountName).toBe("martin"); + }); + + it("rejects ambiguous bare upstream ids across multiple implemented grok accounts", () => { + const accounts = [grokAccount("martin"), grokAccount("work")]; + + expect(() => resolveModel("grok-build-0.1", accounts)).toThrow("Ambiguous model"); + }); + + it("rejects ambiguous provider/upstream ids across multiple accounts", () => { + const accounts = [grokAccount("martin"), grokAccount("work")]; + + expect(() => resolveModel("grok/grok-build-0.1", accounts)).toThrow("Ambiguous model"); + }); + it("rejects empty model-id segments", () => { expect(() => parseProxyModelId("genesiscz/grok/")).toThrow("must be //"); expect(() => parseProxyModelId("/grok/model")).toThrow("must be //"); }); it("resolves account by name and provider slug", () => { - const route = resolveModel("genesiscz/grok/grok-build", [ - { - name: "genesiscz", - provider: "grok-subscription", - providerSlug: "grok", - enabled: true, - }, - ]); + const route = resolveModel("genesiscz/grok/grok-build", [grokAccount("genesiscz")]); expect(route.account.provider).toBe("grok-subscription"); expect(route.upstreamId).toBe("grok-build"); diff --git a/src/ai-proxy/lib/resolve-model.ts b/src/ai-proxy/lib/resolve-model.ts index b089532df0..0d270db19e 100644 --- a/src/ai-proxy/lib/resolve-model.ts +++ b/src/ai-proxy/lib/resolve-model.ts @@ -1,3 +1,4 @@ +import { isProviderImplemented } from "@app/ai-proxy/lib/providers/registry"; import type { AiProxyAccountConfig } from "@app/ai-proxy/lib/types"; export interface ParsedModelId { @@ -6,6 +7,8 @@ export interface ParsedModelId { upstreamId: string; } +const FULL_MODEL_ID_HINT = "Use a full // id."; + export function parseProxyModelId(proxyModelId: string): ParsedModelId { const parts = proxyModelId.split("/"); @@ -28,15 +31,96 @@ export function parseProxyModelId(proxyModelId: string): ParsedModelId { }; } +function enabledImplementedAccounts(accounts: AiProxyAccountConfig[]) { + return accounts.filter((item) => item.enabled && isProviderImplemented(item.provider)); +} + +function resolveFromAccountMatches(matches: AiProxyAccountConfig[], upstreamId: string, requestedId: string) { + if (matches.length === 0) { + return undefined; + } + + if (matches.length > 1) { + const labels = matches.map((account) => `${account.name}/${account.providerSlug}`).join(", "); + + throw new Error( + `Ambiguous model '${requestedId}': multiple enabled accounts match (${labels}). ${FULL_MODEL_ID_HINT}` + ); + } + + const account = matches[0]; + + return { + accountName: account.name, + providerSlug: account.providerSlug, + upstreamId, + account, + }; +} + +function resolveBareUpstreamModel(upstreamId: string, accounts: AiProxyAccountConfig[]) { + const enabled = enabledImplementedAccounts(accounts); + const matches = enabled.filter((account) => account.providerSlug.length > 0); + + return resolveFromAccountMatches(matches, upstreamId, upstreamId); +} + +function resolveProviderUpstreamModel(providerSlug: string, upstreamId: string, accounts: AiProxyAccountConfig[]) { + const enabled = enabledImplementedAccounts(accounts); + const matches = enabled.filter((account) => account.providerSlug === providerSlug); + + return resolveFromAccountMatches(matches, upstreamId, `${providerSlug}/${upstreamId}`); +} + export function resolveModel(proxyModelId: string, accounts: AiProxyAccountConfig[]) { - const parsed = parseProxyModelId(proxyModelId); - const account = accounts.find( - (item) => item.enabled && item.name === parsed.accountName && item.providerSlug === parsed.providerSlug + const trimmed = proxyModelId.trim(); + + if (!trimmed) { + throw new Error(`Model id must be //, got: ${proxyModelId}`); + } + + const slashCount = (trimmed.match(/\//g) ?? []).length; + + if (slashCount === 0) { + const bareRoute = resolveBareUpstreamModel(trimmed, accounts); + + if (bareRoute) { + return bareRoute; + } + + throw new Error(`No enabled account for model '${proxyModelId}'. ${FULL_MODEL_ID_HINT}`); + } + + if (slashCount === 1) { + const [rawProviderSlug = "", rawUpstreamId = ""] = trimmed.split("/", 2); + const providerSlug = rawProviderSlug.trim(); + const upstreamId = rawUpstreamId.trim(); + + if (!providerSlug || !upstreamId) { + throw new Error( + `Model id must be / or //, got: ${proxyModelId}` + ); + } + + const providerRoute = resolveProviderUpstreamModel(providerSlug, upstreamId, accounts); + + if (providerRoute) { + return providerRoute; + } + + throw new Error( + `No enabled account for model '${proxyModelId}' (provider='${providerSlug}'). ${FULL_MODEL_ID_HINT}` + ); + } + + const parsed = parseProxyModelId(trimmed); + const account = enabledImplementedAccounts(accounts).find( + (item) => item.name === parsed.accountName && item.providerSlug === parsed.providerSlug ); if (!account) { throw new Error( - `No enabled account for model '${proxyModelId}' (account='${parsed.accountName}', provider='${parsed.providerSlug}')` + `No enabled account for model '${proxyModelId}' (account='${parsed.accountName}', provider='${parsed.providerSlug}'). ${FULL_MODEL_ID_HINT}` ); } diff --git a/src/ai-proxy/lib/rewrite-upstream-body.test.ts b/src/ai-proxy/lib/rewrite-upstream-body.test.ts index 3df9351f63..5deeed965b 100644 --- a/src/ai-proxy/lib/rewrite-upstream-body.test.ts +++ b/src/ai-proxy/lib/rewrite-upstream-body.test.ts @@ -1,5 +1,14 @@ import { describe, expect, it } from "bun:test"; -import { normalizeGrokTool, prepareGrokUpstreamBody } from "@app/ai-proxy/lib/rewrite-upstream-body"; +import { + GROK_IMAGE_FALLBACK_MODEL, + grokModelSupportsImages, + latestUserTurnHasImages, + normalizeGrokTool, + normalizeGrokToolForChat, + prepareGrokUpstreamBody, + requestHasImageContent, + resolveGrokUpstreamModelForImages, +} from "@app/ai-proxy/lib/rewrite-upstream-body"; import { SafeJSON } from "@app/utils/json"; describe("rewrite-upstream-body", () => { @@ -12,8 +21,37 @@ describe("rewrite-upstream-body", () => { "grok-composer-2.5-fast" ); - const parsed = SafeJSON.parse(rewritten) as { model: string }; + const parsed = SafeJSON.parse(rewritten.bodyText) as { model: string; enable_thinking?: boolean }; expect(parsed.model).toBe("grok-composer-2.5-fast"); + expect(parsed.enable_thinking).toBe(true); + expect(rewritten.imageRouted).toBe(false); + }); + + it("forces enable_thinking for grok-build reasoning models", () => { + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ + model: "martin/grok/grok-build-0.1", + messages: [{ role: "user", content: "hi" }], + }), + "grok-build-0.1" + ); + + const parsed = SafeJSON.parse(rewritten.bodyText) as { model: string; enable_thinking?: boolean }; + expect(parsed.model).toBe("grok-build-0.1"); + expect(parsed.enable_thinking).toBe(true); + }); + + it("does not force enable_thinking for non-reasoning grok models", () => { + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ + model: "grok-code-fast", + messages: [{ role: "user", content: "hi" }], + }), + "grok-code-fast" + ); + + const parsed = SafeJSON.parse(rewritten.bodyText) as { enable_thinking?: boolean }; + expect(parsed.enable_thinking).toBeUndefined(); }); it("flattens OpenAI nested function tools for Grok responses API", () => { @@ -30,6 +68,20 @@ describe("rewrite-upstream-body", () => { expect(normalized?.type).toBe("function"); }); + it("keeps nested function tools for Grok chat completions API", () => { + const normalized = normalizeGrokToolForChat({ + type: "function", + function: { + name: "Read", + description: "Read a file", + parameters: { type: "object", properties: {} }, + }, + }); + + expect(normalized?.type).toBe("function"); + expect((normalized?.function as { name: string }).name).toBe("Read"); + }); + it("converts custom Cursor tools to function tools", () => { const normalized = normalizeGrokTool({ type: "custom", @@ -41,7 +93,7 @@ describe("rewrite-upstream-body", () => { expect(normalized?.name).toBe("ApplyPatch"); }); - it("prepares agent body with normalized tools", () => { + it("prepares agent body with flattened tools for responses target", () => { const rewritten = prepareGrokUpstreamBody( SafeJSON.stringify({ model: "genesiscz/grok/grok-composer-2.5-fast", @@ -54,11 +106,199 @@ describe("rewrite-upstream-body", () => { { type: "custom", name: "ApplyPatch", format: { type: "grammar" } }, ], }), - "grok-composer-2.5-fast" + "grok-composer-2.5-fast", + "responses" ); - const parsed = SafeJSON.parse(rewritten) as { tools: Array<{ type: string; name: string }> }; + const parsed = SafeJSON.parse(rewritten.bodyText) as { tools: Array<{ type: string; name: string }> }; expect(parsed.tools).toHaveLength(2); expect(parsed.tools.every((tool) => tool.type === "function" && typeof tool.name === "string")).toBe(true); }); + + it("prepares chat body with nested function tools for chat target", () => { + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ + model: "genesiscz/grok/grok-composer-2.5-fast", + messages: [{ role: "user", content: "hi" }], + tools: [ + { + type: "function", + function: { name: "Read", description: "Read", parameters: { type: "object" } }, + }, + { type: "custom", name: "ApplyPatch", format: { type: "grammar" } }, + ], + }), + "grok-composer-2.5-fast", + "chat" + ); + + const parsed = SafeJSON.parse(rewritten.bodyText) as { + tools: Array<{ type: string; function: { name: string } }>; + max_tokens?: number; + max_output_tokens?: number; + }; + expect(parsed.tools).toHaveLength(2); + expect(parsed.tools.every((tool) => tool.type === "function" && typeof tool.function?.name === "string")).toBe( + true + ); + expect(parsed.max_output_tokens).toBeUndefined(); + }); + + it("detects image-capable grok models", () => { + expect(grokModelSupportsImages("grok-2-vision")).toBe(true); + expect(grokModelSupportsImages("grok-build")).toBe(true); + expect(grokModelSupportsImages("grok-composer-2.5-fast")).toBe(false); + }); + + it("routes composer requests with images to grok-build", () => { + const body = { + messages: [ + { + role: "user", + content: [ + { type: "text", text: "what is this" }, + { type: "image_url", image_url: { url: "data:image/png;base64,abc" } }, + ], + }, + ], + }; + + expect(requestHasImageContent(body)).toBe(true); + expect(latestUserTurnHasImages(body)).toBe(true); + expect(resolveGrokUpstreamModelForImages("grok-composer-2.5-fast", body)).toBe(GROK_IMAGE_FALLBACK_MODEL); + + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ model: "x", ...body }), + "grok-composer-2.5-fast", + "chat" + ); + const parsed = SafeJSON.parse(rewritten.bodyText) as { + model: string; + messages: Array<{ content: Array<{ type: string; image_url?: { url: string } }> }>; + }; + + expect(rewritten.upstreamModel).toBe(GROK_IMAGE_FALLBACK_MODEL); + expect(parsed.model).toBe(GROK_IMAGE_FALLBACK_MODEL); + expect(rewritten.imageRouted).toBe(true); + expect(parsed.messages[0]?.content?.[1]?.type).toBe("image_url"); + expect(parsed.messages[0]?.content?.[1]?.image_url?.url).toContain("data:image/png"); + }); + + it("follow-up turns stay on composer and replace historical images with text references", () => { + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ + model: "grok-composer-2.5-fast", + messages: [ + { + role: "user", + content: [ + { type: "text", text: "what is this" }, + { type: "image_url", image_url: { url: "data:image/png;base64,abc" } }, + ], + }, + { role: "assistant", content: "It looks red." }, + { role: "user", content: "thanks, now fix the bug" }, + ], + }), + "grok-composer-2.5-fast", + "chat" + ); + + const parsed = SafeJSON.parse(rewritten.bodyText) as { + model: string; + messages: Array<{ content: Array<{ type: string; text?: string }> | string }>; + }; + + expect(rewritten.upstreamModel).toBe("grok-composer-2.5-fast"); + expect(rewritten.imageRouted).toBe(false); + expect(parsed.model).toBe("grok-composer-2.5-fast"); + + const firstUserContent = parsed.messages[0]?.content; + expect(Array.isArray(firstUserContent)).toBe(true); + expect((firstUserContent as Array<{ type: string; text?: string }>)[1]?.type).toBe("text"); + expect((firstUserContent as Array<{ type: string; text?: string }>)[1]?.text).toContain("earlier turn"); + }); + + it("keeps images only in the latest user turn when routing to grok-build", () => { + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ + model: "grok-composer-2.5-fast", + messages: [ + { + role: "user", + content: [{ type: "image_url", image_url: { url: "data:image/png;base64,old" } }], + }, + { role: "assistant", content: "Red square." }, + { + role: "user", + content: [ + { type: "text", text: "and this one?" }, + { type: "image_url", image_url: { url: "data:image/png;base64,new" } }, + ], + }, + ], + }), + "grok-composer-2.5-fast", + "chat" + ); + + const parsed = SafeJSON.parse(rewritten.bodyText) as { + messages: Array<{ content: Array<{ type: string; text?: string; image_url?: { url: string } }> }>; + }; + + expect(rewritten.imageRouted).toBe(true); + expect(parsed.messages[0]?.content?.[0]?.type).toBe("text"); + expect(parsed.messages[2]?.content?.[1]?.type).toBe("image_url"); + expect(parsed.messages[2]?.content?.[1]?.image_url?.url).toContain("new"); + }); + + it("converts chat image_url to input_image for responses target", () => { + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ + model: "grok-composer-2.5-fast", + messages: [ + { + role: "user", + content: [ + { type: "text", text: "what is this" }, + { type: "image_url", image_url: { url: "data:image/png;base64,abc" } }, + ], + }, + ], + }), + "grok-composer-2.5-fast", + "responses" + ); + + const parsed = SafeJSON.parse(rewritten.bodyText) as { + input: Array<{ content: Array<{ type: string; image_url?: string }> }>; + }; + const imagePart = parsed.input[0]?.content?.[1]; + expect(imagePart?.type).toBe("input_image"); + expect(imagePart?.image_url).toBe("data:image/png;base64,abc"); + }); + + it("keeps image parts on grok-build without rerouting", () => { + const rewritten = prepareGrokUpstreamBody( + SafeJSON.stringify({ + model: "grok-build", + messages: [ + { + role: "user", + content: [{ type: "image_url", image_url: { url: "data:image/png;base64,abc" } }], + }, + ], + }), + "grok-build", + "chat" + ); + + const parsed = SafeJSON.parse(rewritten.bodyText) as { + model: string; + messages: Array<{ content: Array<{ type: string }> }>; + }; + expect(parsed.model).toBe("grok-build"); + expect(rewritten.imageRouted).toBe(false); + expect(parsed.messages[0]?.content?.[0]?.type).toBe("image_url"); + }); }); diff --git a/src/ai-proxy/lib/rewrite-upstream-body.ts b/src/ai-proxy/lib/rewrite-upstream-body.ts index f26114e32d..58b82c685d 100644 --- a/src/ai-proxy/lib/rewrite-upstream-body.ts +++ b/src/ai-proxy/lib/rewrite-upstream-body.ts @@ -1,10 +1,322 @@ import { ensureResponsesInput } from "@app/ai-proxy/lib/chat-to-responses-body"; +import { stripCursorThinkingBlocks } from "@app/ai-proxy/lib/thinking-folded"; import { logger } from "@app/logger"; +import { inferModelThinking } from "@app/utils/ai/grok/models"; import { SafeJSON } from "@app/utils/json"; import { isObject } from "@app/utils/object"; type JsonObject = Record; +const IMAGE_CONTENT_TYPES = new Set(["image", "image_url", "input_image"]); +export const GROK_IMAGE_FALLBACK_MODEL = "grok-build"; +const IMAGE_REFERENCE_TEXT = + "[Image attachment from an earlier turn — visual content omitted; refer to prior assistant messages for what was seen.]"; + +export function grokModelSupportsImages(modelId: string): boolean { + return /grok-build|vision/i.test(modelId); +} + +export function requestHasImageContent(body: JsonObject): boolean { + if (Array.isArray(body.messages)) { + for (const message of body.messages) { + if (!isObject(message)) { + continue; + } + + if (contentHasImage(message.content)) { + return true; + } + } + } + + if (Array.isArray(body.input) && inputHasImage(body.input)) { + return true; + } + + return false; +} + +function findLastUserMessageIndex(messages: unknown[]): number { + for (let index = messages.length - 1; index >= 0; index -= 1) { + const message = messages[index]; + + if (isObject(message) && message.role === "user") { + return index; + } + } + + return -1; +} + +function findLastUserInputIndex(input: unknown[]): number { + for (let index = input.length - 1; index >= 0; index -= 1) { + const item = input[index]; + + if (isObject(item) && item.role === "user") { + return index; + } + } + + return -1; +} + +export function latestUserTurnHasImages(body: JsonObject): boolean { + if (Array.isArray(body.messages)) { + const lastUserIndex = findLastUserMessageIndex(body.messages); + + if (lastUserIndex >= 0) { + const message = body.messages[lastUserIndex]; + + if (isObject(message) && contentHasImage(message.content)) { + return true; + } + } + } + + if (Array.isArray(body.input)) { + const lastUserIndex = findLastUserInputIndex(body.input); + + if (lastUserIndex >= 0) { + const item = body.input[lastUserIndex]; + + if (isObject(item) && (isImageContentPart(item) || contentHasImage(item.content))) { + return true; + } + } + + if (lastUserIndex === -1 && inputHasImage(body.input)) { + return true; + } + } + + return false; +} + +export function resolveGrokUpstreamModelForImages(upstreamModel: string, body: JsonObject): string { + if (!latestUserTurnHasImages(body)) { + return upstreamModel; + } + + if (grokModelSupportsImages(upstreamModel)) { + return upstreamModel; + } + + return GROK_IMAGE_FALLBACK_MODEL; +} + +function contentHasImage(content: unknown): boolean { + if (!Array.isArray(content)) { + return false; + } + + return content.some((part) => isImageContentPart(part)); +} + +function inputHasImage(input: unknown[]): boolean { + for (const item of input) { + if (!isObject(item)) { + continue; + } + + if (isImageContentPart(item)) { + return true; + } + + if (contentHasImage(item.content)) { + return true; + } + } + + return false; +} + +function isImageContentPart(part: unknown): boolean { + if (!isObject(part)) { + return false; + } + + if (typeof part.type === "string" && IMAGE_CONTENT_TYPES.has(part.type)) { + return true; + } + + if (part.image_url !== undefined || part.input_image !== undefined) { + return true; + } + + if (part.type === "image" && isObject(part.source)) { + return true; + } + + return false; +} + +function imageDataUrlFromPart(part: JsonObject): string | null { + if (part.type === "input_image") { + if (typeof part.image_url === "string") { + return part.image_url; + } + + if (isObject(part.image_url) && typeof part.image_url.url === "string") { + return part.image_url.url; + } + } + + if (part.type === "image_url") { + if (isObject(part.image_url) && typeof part.image_url.url === "string") { + return part.image_url.url; + } + + if (typeof part.image_url === "string") { + return part.image_url; + } + } + + if (part.type === "image" && isObject(part.source) && part.source.type === "base64") { + const mediaType = typeof part.source.media_type === "string" ? part.source.media_type : "image/png"; + const data = typeof part.source.data === "string" ? part.source.data : ""; + + if (data) { + return `data:${mediaType};base64,${data}`; + } + } + + return null; +} + +function normalizeImagePartForChat(part: JsonObject): JsonObject { + const dataUrl = imageDataUrlFromPart(part); + + if (dataUrl) { + return { + type: "image_url", + image_url: { url: dataUrl }, + }; + } + + return part; +} + +function normalizeImagePartForResponses(part: JsonObject): JsonObject { + const dataUrl = imageDataUrlFromPart(part); + + if (dataUrl) { + return { + type: "input_image", + image_url: dataUrl, + }; + } + + if (part.type === "text" && typeof part.text === "string") { + return { type: "input_text", text: part.text }; + } + + return part; +} + +function imageReferenceTextPart(target: "chat" | "responses"): JsonObject { + if (target === "responses") { + return { type: "input_text", text: IMAGE_REFERENCE_TEXT }; + } + + return { type: "text", text: IMAGE_REFERENCE_TEXT }; +} + +function processMessageContentImages(content: unknown, target: "chat" | "responses", replaceImages: boolean): unknown { + if (typeof content === "string" || content == null) { + return content; + } + + if (!Array.isArray(content)) { + return content; + } + + return content.map((part) => { + if (!isObject(part)) { + return part; + } + + if (isImageContentPart(part)) { + if (replaceImages) { + return imageReferenceTextPart(target); + } + + if (target === "chat") { + return normalizeImagePartForChat(part); + } + + return normalizeImagePartForResponses(part); + } + + if (target === "responses" && part.type === "text" && typeof part.text === "string") { + return { type: "input_text", text: part.text }; + } + + return part; + }); +} + +function applyGrokImageTurnPolicy( + body: JsonObject, + target: "chat" | "responses", + routeCurrentTurnToBuild: boolean +): void { + if (Array.isArray(body.messages)) { + const lastUserIndex = findLastUserMessageIndex(body.messages); + + for (let index = 0; index < body.messages.length; index += 1) { + const message = body.messages[index]; + + if (!isObject(message) || message.content === undefined) { + continue; + } + + const keepImages = routeCurrentTurnToBuild && index === lastUserIndex; + message.content = processMessageContentImages(message.content, target, !keepImages); + } + } + + if (Array.isArray(body.input)) { + const lastUserIndex = findLastUserInputIndex(body.input); + const keepUnscopedInputImages = routeCurrentTurnToBuild && lastUserIndex === -1; + + for (let index = 0; index < body.input.length; index += 1) { + const item = body.input[index]; + + if (!isObject(item)) { + continue; + } + + if (isImageContentPart(item)) { + const keepImages = routeCurrentTurnToBuild && (index === lastUserIndex || keepUnscopedInputImages); + + if (keepImages) { + body.input[index] = normalizeImagePartForResponses(item); + } else { + body.input[index] = { + type: "input_text", + text: IMAGE_REFERENCE_TEXT, + }; + } + + continue; + } + + if (item.content === undefined) { + continue; + } + + const keepImages = routeCurrentTurnToBuild && (index === lastUserIndex || keepUnscopedInputImages); + item.content = processMessageContentImages(item.content, "responses", !keepImages); + } + } +} + +export interface PreparedGrokUpstreamBody { + bodyText: string; + upstreamModel: string; + imageRouted: boolean; +} + export function normalizeGrokTool(tool: unknown): JsonObject | null { if (!isObject(tool)) { return null; @@ -49,6 +361,56 @@ export function normalizeGrokTool(tool: unknown): JsonObject | null { return null; } +export function normalizeGrokToolForChat(tool: unknown): JsonObject | null { + if (!isObject(tool)) { + return null; + } + + if (tool.type === "function" && isObject(tool.function)) { + const fn = tool.function; + if (typeof fn.name !== "string") { + return null; + } + + return { + type: "function", + function: { + name: fn.name, + description: fn.description ?? "", + parameters: fn.parameters ?? { type: "object", properties: {} }, + ...(fn.strict !== undefined ? { strict: fn.strict } : {}), + }, + }; + } + + if (tool.type === "function" && typeof tool.name === "string") { + return { + type: "function", + function: { + name: tool.name, + description: tool.description ?? "", + parameters: tool.parameters ?? { type: "object", properties: {} }, + ...(tool.strict !== undefined ? { strict: tool.strict } : {}), + }, + }; + } + + if (tool.type === "custom") { + const name = typeof tool.name === "string" ? tool.name : "custom_tool"; + + return { + type: "function", + function: { + name, + description: typeof tool.description === "string" ? tool.description : "Custom Cursor tool", + parameters: { type: "object", properties: {} }, + }, + }; + } + + return null; +} + export function normalizeGrokTools(tools: unknown): JsonObject[] | undefined { if (!Array.isArray(tools)) { return undefined; @@ -63,30 +425,117 @@ export function normalizeGrokTools(tools: unknown): JsonObject[] | undefined { return normalized; } +export function normalizeGrokToolsForChat(tools: unknown): JsonObject[] | undefined { + if (!Array.isArray(tools)) { + return undefined; + } + + const normalized = tools.map(normalizeGrokToolForChat).filter((tool): tool is JsonObject => tool !== null); + + if (normalized.length === 0) { + return undefined; + } + + return normalized; +} + +function ensureGrokThinkingEnabled(body: JsonObject, upstreamModel: string): void { + if (inferModelThinking(upstreamModel) !== "reasoning") { + return; + } + + if (body.enable_thinking === true) { + return; + } + + body.enable_thinking = true; +} + +function stripMirroredThinkingFromMessages(body: JsonObject): void { + if (!Array.isArray(body.messages)) { + return; + } + + for (const message of body.messages) { + if (!isObject(message) || message.role !== "assistant") { + continue; + } + + if (typeof message.content === "string") { + message.content = stripCursorThinkingBlocks(message.content); + } else if (Array.isArray(message.content)) { + for (const part of message.content) { + if (isObject(part) && part.type === "text" && typeof part.text === "string") { + part.text = stripCursorThinkingBlocks(part.text); + } + } + } + } +} + +function patchGrokAssistantReasoningForToolCalls(body: JsonObject): void { + if (!Array.isArray(body.messages)) { + return; + } + + for (const message of body.messages) { + if (!isObject(message) || message.role !== "assistant") { + continue; + } + + if (!Array.isArray(message.tool_calls) || message.tool_calls.length === 0) { + continue; + } + + if (typeof message.reasoning_content === "string" && message.reasoning_content.trim()) { + continue; + } + + message.reasoning_content = " "; + } +} + export function prepareGrokUpstreamBody( bodyText: string, upstreamModel: string, target: "chat" | "responses" = "chat" -): string { +): PreparedGrokUpstreamBody { try { const parsed = SafeJSON.parse(bodyText, { strict: true }); if (!isObject(parsed)) { - return bodyText; + return { bodyText, upstreamModel, imageRouted: false }; } const next: JsonObject = { ...parsed }; + const currentTurnHasImages = latestUserTurnHasImages(next); + const resolvedModel = resolveGrokUpstreamModelForImages(upstreamModel, next); + const imageRouted = resolvedModel !== upstreamModel; + + if (imageRouted) { + logger.debug( + { from: upstreamModel, to: resolvedModel, target }, + "ai-proxy: routing current Grok turn with images to vision-capable model" + ); + } if ("model" in next) { - next.model = upstreamModel; + next.model = resolvedModel; } - if ("max_tokens" in next && !("max_output_tokens" in next)) { - next.max_output_tokens = next.max_tokens; - delete next.max_tokens; + ensureGrokThinkingEnabled(next, resolvedModel); + patchGrokAssistantReasoningForToolCalls(next); + stripMirroredThinkingFromMessages(next); + + if (target === "responses") { + if ("max_tokens" in next && !("max_output_tokens" in next)) { + next.max_output_tokens = next.max_tokens; + delete next.max_tokens; + } } - const normalizedTools = normalizeGrokTools(next.tools); + const normalizedTools = + target === "responses" ? normalizeGrokTools(next.tools) : normalizeGrokToolsForChat(next.tools); if (normalizedTools) { next.tools = normalizedTools; } else if ("tools" in next) { @@ -101,12 +550,18 @@ export function prepareGrokUpstreamBody( delete next.n; } + applyGrokImageTurnPolicy(next, target, currentTurnHasImages); + const prepared = target === "responses" ? ensureResponsesInput(next) : next; - return SafeJSON.stringify(prepared); + return { + bodyText: SafeJSON.stringify(prepared), + upstreamModel: resolvedModel, + imageRouted, + }; } catch (err) { logger.debug({ err, upstreamModel, target }, "ai-proxy: prepareGrokUpstreamBody fallback"); - return rewriteBodyModel(bodyText, upstreamModel); + return { bodyText: rewriteBodyModel(bodyText, upstreamModel), upstreamModel, imageRouted: false }; } } diff --git a/src/ai-proxy/lib/server.ts b/src/ai-proxy/lib/server.ts index 2f75024eba..f985b9d7ac 100644 --- a/src/ai-proxy/lib/server.ts +++ b/src/ai-proxy/lib/server.ts @@ -22,7 +22,10 @@ const MAX_REQUEST_BODY_BYTES = 4 * 1024 * 1024; function mapProxyRequestError(err: unknown): { status: number; message: string } { if (err instanceof GrokAuthExpiredError || err instanceof CopilotAuthExpiredError) { - return { status: 401, message: err.message }; + return { + status: 502, + message: "Upstream provider auth expired or invalid — the ai-proxy host must refresh its login.", + }; } const message = err instanceof Error ? err.message : String(err); @@ -30,7 +33,8 @@ function mapProxyRequestError(err: unknown): { status: number; message: string } if ( message.startsWith("Model id must be") || message.startsWith("No enabled account for model") || - message.startsWith("Provider not loaded:") + message.startsWith("Provider not loaded:") || + message.startsWith("Ambiguous model") ) { return { status: 400, message }; } @@ -121,6 +125,7 @@ export function startAiProxyServer(runtime: AiProxyRuntime) { } const models = await buildProxyModelCatalog(config.accounts); + return new Response( SafeJSON.stringify({ object: "list", diff --git a/src/ai-proxy/lib/thinking-folded.ts b/src/ai-proxy/lib/thinking-folded.ts index bb64f315e9..3ad259cd0a 100644 --- a/src/ai-proxy/lib/thinking-folded.ts +++ b/src/ai-proxy/lib/thinking-folded.ts @@ -45,3 +45,10 @@ export function wrapReasoningForFoldedJson(reasoning: string, answer: string | n return body.trimEnd(); } + +const CURSOR_THINKING_BLOCK_RE = + /(?:<(?:think|thinking)\b[^>]*>[\s\S]*?(?:<\/(?:think|thinking)>|$)|]*>\s*]*>\s*(?:)?Thinking(?:<\/strong>)?\s*<\/summary>[\s\S]*?(?:<\/details>|$))\s*/gi; + +export function stripCursorThinkingBlocks(content: string): string { + return content.replace(CURSOR_THINKING_BLOCK_RE, "").replace(/^\r\n+/, ""); +} diff --git a/src/utils/ai/grok/client.ts b/src/utils/ai/grok/client.ts index 034e7daa15..5b941fb451 100644 --- a/src/utils/ai/grok/client.ts +++ b/src/utils/ai/grok/client.ts @@ -64,8 +64,25 @@ export class GrokSubscriptionClient { } } + private async ensureFreshTokenInMemory(): Promise { + if (!isTokenExpired(decodeJwtClaims(this.token))) { + return; + } + + const previousToken = this.token; + const reloaded = await this.reloadTokenFromDisk(); + + if (reloaded && reloaded !== previousToken) { + logger.debug("grok: in-memory token expired, reloaded auth.json from disk"); + } + + if (isTokenExpired(decodeJwtClaims(this.token))) { + throw new GrokAuthExpiredError(this.authPath); + } + } + async fetch(path: string, init?: RequestInit & { modelOverride?: string }): Promise { - this.assertTokenFresh(); + await this.ensureFreshTokenInMemory(); let response = await this.doFetch(path, init); @@ -80,6 +97,27 @@ export class GrokSubscriptionClient { } if (isAuthHttpStatus(response.status)) { + // Drain the upstream body so we can include it in the diagnostic log — + // otherwise we throw away the only clue about WHY upstream said 401. + // The body is small (xAI auth-fail bodies are <1KB); buffering is fine here. + let bodyExcerpt = ""; + try { + bodyExcerpt = (await response.text()).slice(0, 500); + } catch (err) { + bodyExcerpt = ``; + } + + logger.warn( + { + path, + upstreamStatus: response.status, + upstreamBodyExcerpt: bodyExcerpt, + modelOverride: init?.modelOverride, + authPath: this.authPath, + }, + "grok: upstream returned auth-status, throwing GrokAuthExpiredError" + ); + throw new GrokAuthExpiredError(this.authPath); } diff --git a/src/utils/ai/grok/models.test.ts b/src/utils/ai/grok/models.test.ts new file mode 100644 index 0000000000..2a96c0d853 --- /dev/null +++ b/src/utils/ai/grok/models.test.ts @@ -0,0 +1,13 @@ +import { describe, expect, it } from "bun:test"; +import { inferModelThinking } from "@app/utils/ai/grok/models"; + +describe("inferModelThinking", () => { + it("returns none for non-reasoning model ids instead of matching the broader reasoning regex", () => { + expect(inferModelThinking("grok-4-1-fast-non-reasoning")).toBe("none"); + }); + + it("still returns reasoning for ids that genuinely indicate reasoning", () => { + expect(inferModelThinking("grok-4-1-reasoning")).toBe("reasoning"); + expect(inferModelThinking("grok-build")).toBe("reasoning"); + }); +}); diff --git a/src/utils/ai/grok/models.ts b/src/utils/ai/grok/models.ts index a7090485fd..283d069c11 100644 --- a/src/utils/ai/grok/models.ts +++ b/src/utils/ai/grok/models.ts @@ -23,7 +23,7 @@ function seed( export const GROK_STATIC_CATALOG: GrokModelRecord[] = [ seed("grok-build", "high", "slow", "reasoning", "ok"), - seed("grok-composer-2.5-fast", "high", "fast", "none", "ok"), + seed("grok-composer-2.5-fast", "high", "fast", "reasoning", "ok"), seed("grok-build-0.1", "medium", "slow", "reasoning", "ok"), seed("grok-code-fast", "medium", "fast", "none", "ok"), seed("grok-code-fast-1", "medium", "fast", "none", "ok"), @@ -51,7 +51,7 @@ export const GROK_STATIC_CATALOG: GrokModelRecord[] = [ seed("grok-4.20-0309-non-reasoning", "medium", "fast", "none", "ok"), seed("grok-4.20-multi-agent-0309", "medium", "slow", "multi-agent", "ok"), seed("composer-2.5-fast", "low", "fast", "none", "fail"), - seed("grok-composer-2.5", "low", "fast", "none", "fail"), + seed("grok-composer-2.5", "low", "fast", "reasoning", "fail"), seed("grok-4.1-fast", "low", "fast", "none", "fail"), seed("grok-2", "low", "medium", "optional", "fail"), seed("grok-2-vision", "low", "medium", "optional", "fail"), @@ -82,19 +82,27 @@ export function inferModelSpeed(id: string): GrokModelSpeed { return "medium"; } +export function isGrokComposerModel(id: string): boolean { + return /grok-composer(?:-2\.5)?(?:-fast)?$/i.test(id) || /^composer-2\.5(?:-fast)?$/i.test(id); +} + export function inferModelThinking(id: string): GrokModelThinking { if (/multi-agent/i.test(id)) { return "multi-agent"; } - if (/reasoning|grok-build/i.test(id)) { + if (isGrokComposerModel(id)) { return "reasoning"; } - if (/non-reasoning|composer|code-fast|mini|fast/i.test(id)) { + if (/non-reasoning|code-fast|mini|fast/i.test(id)) { return "none"; } + if (/reasoning|grok-build/i.test(id)) { + return "reasoning"; + } + return "optional"; } diff --git a/src/utils/ai/grok/probe.test.ts b/src/utils/ai/grok/probe.test.ts index c72d41f886..847e024b44 100644 --- a/src/utils/ai/grok/probe.test.ts +++ b/src/utils/ai/grok/probe.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "bun:test"; -import { GROK_PROBE_CANDIDATES, GROK_STATIC_CATALOG, mergeModelCatalog, toProxyId } from "./models"; +import { GROK_PROBE_CANDIDATES, GROK_STATIC_CATALOG, inferModelThinking, mergeModelCatalog, toProxyId } from "./models"; describe("grok probe helpers", () => { it("includes researched static catalog ids", () => { @@ -7,6 +7,12 @@ describe("grok probe helpers", () => { expect(GROK_STATIC_CATALOG.length).toBeGreaterThanOrEqual(29); }); + it("treats grok-composer models as required reasoning, not optional", () => { + expect(inferModelThinking("grok-composer-2.5-fast")).toBe("reasoning"); + expect(inferModelThinking("grok-composer-2.5")).toBe("reasoning"); + expect(GROK_STATIC_CATALOG.find((model) => model.id === "grok-composer-2.5-fast")?.thinking).toBe("reasoning"); + }); + it("merges picker models over static catalog", () => { const merged = mergeModelCatalog( [ From f2ec00ffbd3d9d4abf2fff2dcbfad7f0da0ad3ed Mon Sep 17 00:00:00 2001 From: Martin Date: Tue, 30 Jun 2026 22:34:55 +0200 Subject: [PATCH 2/3] feat(ai-proxy): Cursor SSE translation and Grok chat enricher --- .../lib/chat-to-responses-body.test.ts | 2 +- .../lib/grok-chat-sse-enricher.test.ts | 53 +++ src/ai-proxy/lib/grok-chat-sse-enricher.ts | 342 ++++++++++++++++++ src/ai-proxy/lib/safe-stream-controller.ts | 37 ++ .../lib/translators/identity-pipeline.ts | 14 +- src/ai-proxy/lib/translators/index.test.ts | 34 ++ src/ai-proxy/lib/translators/index.ts | 22 +- .../responses-stream-translator.test.ts | 19 + .../responses-stream-translator.ts | 18 + .../lib/translators/responses-to-chat-sse.ts | 54 ++- 10 files changed, 583 insertions(+), 12 deletions(-) create mode 100644 src/ai-proxy/lib/grok-chat-sse-enricher.test.ts create mode 100644 src/ai-proxy/lib/grok-chat-sse-enricher.ts create mode 100644 src/ai-proxy/lib/safe-stream-controller.ts create mode 100644 src/ai-proxy/lib/translators/index.test.ts diff --git a/src/ai-proxy/lib/chat-to-responses-body.test.ts b/src/ai-proxy/lib/chat-to-responses-body.test.ts index 89fd962b34..5c12c1b024 100644 --- a/src/ai-proxy/lib/chat-to-responses-body.test.ts +++ b/src/ai-proxy/lib/chat-to-responses-body.test.ts @@ -61,7 +61,7 @@ describe("chat-to-responses-body", () => { "responses" ); - const parsed = SafeJSON.parse(rewritten) as { input?: unknown[]; messages?: unknown; model: string }; + const parsed = SafeJSON.parse(rewritten.bodyText) as { input?: unknown[]; messages?: unknown; model: string }; expect(parsed.model).toBe("grok-composer-2.5-fast"); expect(parsed.input?.length).toBe(1); expect(parsed.messages).toBeUndefined(); diff --git a/src/ai-proxy/lib/grok-chat-sse-enricher.test.ts b/src/ai-proxy/lib/grok-chat-sse-enricher.test.ts new file mode 100644 index 0000000000..4ed25c51c9 --- /dev/null +++ b/src/ai-proxy/lib/grok-chat-sse-enricher.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it } from "bun:test"; +import { enrichGrokChatResponse } from "@app/ai-proxy/lib/grok-chat-sse-enricher"; + +describe("grok-chat-sse-enricher", () => { + it("cursor mode rewrites model to proxy id and adds reasoning_items on first thinking delta", async () => { + const sse = [ + 'data: {"model":"grok-build-0.1","choices":[{"index":0,"delta":{"reasoning_content":"Hmm","role":"assistant"}}]}', + 'data: {"model":"grok-build-0.1","choices":[{"index":0,"delta":{"content":"Answer"}}]}', + "", + "data: [DONE]\r", + "", + ].join("\n"); + + const response = await enrichGrokChatResponse( + new Response(sse, { + headers: { "Content-Type": "text/event-stream" }, + }), + "martin/grok/grok-build-0.1", + "cursor" + ); + + const body = await response.text(); + + expect(body).toContain('"model":"martin/grok/grok-build-0.1"'); + expect(body).toContain('"reasoning_content":"Hmm"'); + expect(body).toContain('"reasoning_items"'); + expect(body).toContain('"type":"reasoning"'); + expect(body).toContain('"content":"Answer"'); + expect(body).not.toContain("
"); + }); + + it("folded mode moves reasoning into content only", async () => { + const sse = [ + 'data: {"model":"grok-composer-2.5-fast","choices":[{"delta":{"reasoning_content":"Hmm","role":"assistant"}}]}', + "", + "data: [DONE]", + "", + ].join("\n"); + + const response = await enrichGrokChatResponse( + new Response(sse, { + headers: { "Content-Type": "text/event-stream" }, + }), + "martin/grok/grok-composer-2.5-fast", + "folded" + ); + + const body = await response.text(); + + expect(body).toContain("
"); + expect(body).not.toContain("reasoning_content"); + }); +}); diff --git a/src/ai-proxy/lib/grok-chat-sse-enricher.ts b/src/ai-proxy/lib/grok-chat-sse-enricher.ts new file mode 100644 index 0000000000..4c1287c21f --- /dev/null +++ b/src/ai-proxy/lib/grok-chat-sse-enricher.ts @@ -0,0 +1,342 @@ +import { createFoldedStreamState, foldedAnswerPrefix, foldedReasoningPrefix } from "@app/ai-proxy/lib/thinking-folded"; +import { buildReasoningItem, serializeReasoningItems } from "@app/ai-proxy/lib/translators/reasoning"; +import type { ThinkingPresentationMode } from "@app/ai-proxy/lib/types"; +import { logger } from "@app/logger"; +import { SafeJSON } from "@app/utils/json"; +import { isObject } from "@app/utils/object"; +import { safeStreamControllerError } from "./safe-stream-controller"; + +type JsonObject = Record; + +interface StreamEnrichState { + foldedState: ReturnType; + roleSent: boolean; + reasoningItemId: string | null; +} + +function createStreamEnrichState(): StreamEnrichState { + return { + foldedState: createFoldedStreamState(), + roleSent: false, + reasoningItemId: null, + }; +} + +function rewriteChatPayloadModel(payload: unknown, responseModel: string): unknown { + if (!isObject(payload)) { + return payload; + } + + if (!("model" in payload)) { + return payload; + } + + return { + ...payload, + model: responseModel, + }; +} + +function isSseDonePayload(payload: string): boolean { + return payload.trim() === "[DONE]"; +} + +function deltaHasPayload(delta: JsonObject): boolean { + return ( + delta.content !== undefined || + delta.reasoning_content !== undefined || + delta.reasoning_items !== undefined || + delta.tool_calls !== undefined + ); +} + +function withAssistantRole(delta: JsonObject, state: StreamEnrichState): JsonObject { + if (typeof delta.role === "string") { + state.roleSent = true; + + return delta; + } + + if (state.roleSent || !deltaHasPayload(delta)) { + return delta; + } + + state.roleSent = true; + + return { + role: "assistant", + ...delta, + }; +} + +function enrichDeltaForCursor(delta: JsonObject, state: StreamEnrichState): JsonObject { + const next: JsonObject = { ...delta }; + + if (typeof delta.reasoning_content === "string" && delta.reasoning_content.length > 0) { + if (!state.reasoningItemId) { + state.reasoningItemId = `rs_${crypto.randomUUID()}`; + const reasoningItem = buildReasoningItem({ + id: state.reasoningItemId, + type: "reasoning", + content: [{ type: "reasoning_text", text: delta.reasoning_content }], + }); + + next.reasoning_items = serializeReasoningItems([reasoningItem]); + } + } + + return withAssistantRole(next, state); +} + +function enrichDeltaForFolded(delta: JsonObject, state: StreamEnrichState): JsonObject { + const next: JsonObject = { ...delta }; + + if (typeof delta.reasoning_content === "string") { + const prefix = foldedReasoningPrefix(state.foldedState); + next.content = `${prefix}${delta.reasoning_content}`; + delete next.reasoning_content; + } + + if (typeof delta.content === "string") { + const prefix = foldedAnswerPrefix(state.foldedState); + next.content = `${prefix}${delta.content}`; + } + + return next; +} + +function enrichChatPayload( + payload: JsonObject, + thinkingMode: ThinkingPresentationMode, + state: StreamEnrichState +): JsonObject { + const choices = payload.choices; + + if (!Array.isArray(choices) || choices.length === 0) { + return payload; + } + + let changed = false; + const enrichedChoices = choices.map((rawChoice) => { + if (!isObject(rawChoice)) { + return rawChoice; + } + + const delta = rawChoice.delta; + + if (!isObject(delta)) { + return rawChoice; + } + + let enrichedDelta = delta; + + if (thinkingMode === "cursor") { + enrichedDelta = enrichDeltaForCursor(delta, state); + } else if (thinkingMode === "folded") { + enrichedDelta = enrichDeltaForFolded(delta, state); + } + + if (enrichedDelta === delta) { + return rawChoice; + } + + changed = true; + + return { + ...rawChoice, + delta: enrichedDelta, + }; + }); + + if (!changed) { + return payload; + } + + return { + ...payload, + choices: enrichedChoices, + }; +} + +function transformChatPayload( + payload: unknown, + responseModel: string, + thinkingMode: ThinkingPresentationMode, + state: StreamEnrichState +): unknown { + let next = rewriteChatPayloadModel(payload, responseModel); + + if ((thinkingMode === "cursor" || thinkingMode === "folded") && isObject(next)) { + next = enrichChatPayload(next, thinkingMode, state); + } + + return next; +} + +function rewriteSseDataLine( + line: string, + responseModel: string, + thinkingMode: ThinkingPresentationMode, + state: StreamEnrichState +): string { + const prefix = "data:"; + const trimmed = line.trimStart(); + + if (!trimmed.startsWith(prefix)) { + return line; + } + + const payload = trimmed.slice(prefix.length).trim(); + + if (isSseDonePayload(payload)) { + return line; + } + + try { + const parsed = SafeJSON.parse(payload, { strict: true }); + const rewritten = transformChatPayload(parsed, responseModel, thinkingMode, state); + + if (rewritten === parsed) { + return line; + } + + const suffix = line.endsWith("\r\n") ? "\r\n" : line.endsWith("\n") ? "\n" : ""; + + return `data: ${SafeJSON.stringify(rewritten)}${suffix}`; + } catch (err) { + logger.debug({ err, responseModel }, "ai-proxy: enrichGrokChatSseLine fallback"); + return line; + } +} + +function enrichGrokChatSseStream( + stream: ReadableStream, + responseModel: string, + thinkingMode: ThinkingPresentationMode +): ReadableStream { + const decoder = new TextDecoder(); + const encoder = new TextEncoder(); + let buffer = ""; + const state = createStreamEnrichState(); + + return new ReadableStream({ + async start(controller) { + const reader = stream.getReader(); + let closed = false; + + try { + while (true) { + const { done, value } = await reader.read(); + + if (done) { + if (buffer.length > 0) { + controller.enqueue( + encoder.encode(rewriteSseDataLine(buffer, responseModel, thinkingMode, state)) + ); + } + + try { + controller.close(); + closed = true; + } catch (controllerErr) { + logger.warn( + { err: controllerErr, responseModel }, + "ai-proxy: enrichGrokChatSseStream controller.close() threw" + ); + } + break; + } + + buffer += decoder.decode(value, { stream: true }); + + let newlineIndex = buffer.indexOf("\n"); + + while (newlineIndex >= 0) { + const line = buffer.slice(0, newlineIndex + 1); + buffer = buffer.slice(newlineIndex + 1); + controller.enqueue( + encoder.encode(rewriteSseDataLine(line, responseModel, thinkingMode, state)) + ); + newlineIndex = buffer.indexOf("\n"); + } + } + } catch (err) { + logger.warn({ err, responseModel }, "ai-proxy: enrichGrokChatSseStream failed"); + + if (!safeStreamControllerError(controller, err, closed)) { + logger.debug( + { err, responseModel, closed }, + "ai-proxy: enrichGrokChatSseStream skipped controller.error (client abort or detached)" + ); + } + } finally { + try { + reader.releaseLock(); + } catch (lockErr) { + logger.debug( + { err: lockErr, responseModel }, + "ai-proxy: enrichGrokChatSseStream releaseLock failed" + ); + } + } + }, + }); +} + +export function enrichGrokChatCompletionJson( + bodyText: string, + responseModel: string, + thinkingMode: ThinkingPresentationMode +): string { + try { + const parsed = SafeJSON.parse(bodyText, { strict: true }); + + if (!isObject(parsed)) { + return bodyText; + } + + const state = createStreamEnrichState(); + const rewritten = transformChatPayload(parsed, responseModel, thinkingMode, state); + + return SafeJSON.stringify(rewritten); + } catch (err) { + logger.debug({ err, responseModel }, "ai-proxy: enrichGrokChatCompletionJson fallback"); + return bodyText; + } +} + +export async function enrichGrokChatResponse( + response: Response, + responseModel: string, + thinkingMode: ThinkingPresentationMode +): Promise { + if (!response.body) { + return response; + } + + const contentType = response.headers.get("content-type") ?? ""; + + if (contentType.includes("text/event-stream")) { + return new Response(enrichGrokChatSseStream(response.body, responseModel, thinkingMode), { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }); + } + + if (contentType.includes("application/json")) { + const bodyText = await response.text(); + const rewritten = enrichGrokChatCompletionJson(bodyText, responseModel, thinkingMode); + const headers = new Headers(response.headers); + headers.delete("content-length"); + headers.delete("etag"); + + return new Response(rewritten, { + status: response.status, + statusText: response.statusText, + headers, + }); + } + + return response; +} diff --git a/src/ai-proxy/lib/safe-stream-controller.ts b/src/ai-proxy/lib/safe-stream-controller.ts new file mode 100644 index 0000000000..297c535774 --- /dev/null +++ b/src/ai-proxy/lib/safe-stream-controller.ts @@ -0,0 +1,37 @@ +import { logger } from "@app/logger"; +import { isObject } from "@app/utils/object"; + +export function isStreamAbortError(err: unknown): boolean { + if (err instanceof DOMException && err.name === "AbortError") { + return true; + } + + if (err instanceof Error && err.name === "AbortError") { + return true; + } + + if (isObject(err) && err.name === "AbortError") { + return true; + } + + return false; +} + +export function safeStreamControllerError( + controller: ReadableStreamDefaultController, + err: unknown, + closed: boolean +): boolean { + if (closed || isStreamAbortError(err)) { + return false; + } + + try { + controller.error(err); + return true; + } catch (controllerErr) { + logger.debug({ err: controllerErr, originalErr: err }, "ai-proxy: controller.error() threw"); + + return false; + } +} diff --git a/src/ai-proxy/lib/translators/identity-pipeline.ts b/src/ai-proxy/lib/translators/identity-pipeline.ts index ca0b2e2870..2442181641 100644 --- a/src/ai-proxy/lib/translators/identity-pipeline.ts +++ b/src/ai-proxy/lib/translators/identity-pipeline.ts @@ -1,15 +1,21 @@ +import { enrichGrokChatResponse } from "@app/ai-proxy/lib/grok-chat-sse-enricher"; import type { ProxyProvider } from "@app/ai-proxy/lib/providers/types"; +import type { ThinkingPresentationMode } from "@app/ai-proxy/lib/types"; import { type PipelineResult, pipelineResult } from "@app/ai-proxy/lib/usage/pipeline-result"; export async function identityPipeline({ provider, upstreamModel, + proxyModel, + thinkingMode = "cursor", path, req, bodyText, }: { provider: ProxyProvider; upstreamModel: string; + proxyModel?: string; + thinkingMode?: ThinkingPresentationMode; path: "chat/completions" | "responses"; req: Request; bodyText: string; @@ -18,5 +24,11 @@ export async function identityPipeline({ return pipelineResult(await provider.responses(req, upstreamModel, bodyText)); } - return pipelineResult(await provider.chatCompletions(req, upstreamModel, bodyText)); + const upstream = await provider.chatCompletions(req, upstreamModel, bodyText); + + if (!proxyModel || provider.id !== "grok-subscription") { + return pipelineResult(upstream); + } + + return pipelineResult(await enrichGrokChatResponse(upstream, proxyModel, thinkingMode)); } diff --git a/src/ai-proxy/lib/translators/index.test.ts b/src/ai-proxy/lib/translators/index.test.ts new file mode 100644 index 0000000000..b96b88c054 --- /dev/null +++ b/src/ai-proxy/lib/translators/index.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, it } from "bun:test"; +import { shouldTranslateChatRequest } from "@app/ai-proxy/lib/translators/index"; + +describe("shouldTranslateChatRequest", () => { + it("skips responses translation for grok subscription passthrough", () => { + const req = new Request("http://127.0.0.1/v1/chat/completions", { + headers: { "User-Agent": "Cursor/1.0" }, + }); + + expect( + shouldTranslateChatRequest({ + mode: "auto", + req, + bodyText: '{"messages":[{"role":"user","content":"hi"}]}', + providerId: "grok-subscription", + }) + ).toBe(false); + }); + + it("still translates copilot subscription for Cursor", () => { + const req = new Request("http://127.0.0.1/v1/chat/completions", { + headers: { "User-Agent": "Cursor/1.0" }, + }); + + expect( + shouldTranslateChatRequest({ + mode: "auto", + req, + bodyText: '{"messages":[{"role":"user","content":"hi"}]}', + providerId: "github-copilot-subscription", + }) + ).toBe(true); + }); +}); diff --git a/src/ai-proxy/lib/translators/index.ts b/src/ai-proxy/lib/translators/index.ts index 6a3d687a6b..c89964dc21 100644 --- a/src/ai-proxy/lib/translators/index.ts +++ b/src/ai-proxy/lib/translators/index.ts @@ -5,11 +5,27 @@ import { responsesToChat } from "@app/ai-proxy/lib/translators/responses-to-chat import type { CursorTranslationMode, ThinkingPresentationMode } from "@app/ai-proxy/lib/types"; import type { PipelineResult } from "@app/ai-proxy/lib/usage/pipeline-result"; -export function shouldTranslateChatRequest(mode: CursorTranslationMode, req: Request, bodyText: string): boolean { +export function shouldTranslateChatRequest({ + mode, + req, + bodyText, + providerId, +}: { + mode: CursorTranslationMode; + req: Request; + bodyText: string; + providerId?: string; +}): boolean { if (mode === "off") { return false; } + // Grok subscription chat/completions already streams Cursor-native reasoning_content. + // Re-encoding via /responses drops role coalescing and breaks the thinking UI. + if (providerId === "grok-subscription") { + return false; + } + if (mode === "on") { return true; } @@ -34,7 +50,7 @@ export async function handleChatCompletions({ req: Request; bodyText: string; }): Promise { - if (shouldTranslateChatRequest(mode, req, bodyText)) { + if (shouldTranslateChatRequest({ mode, req, bodyText, providerId: provider.id })) { return responsesToChat({ provider, upstreamModel, @@ -48,6 +64,8 @@ export async function handleChatCompletions({ return identityPipeline({ provider, upstreamModel, + proxyModel, + thinkingMode, path: "chat/completions", req, bodyText, diff --git a/src/ai-proxy/lib/translators/responses-stream-translator.test.ts b/src/ai-proxy/lib/translators/responses-stream-translator.test.ts index 75c24cb8ea..841975f165 100644 --- a/src/ai-proxy/lib/translators/responses-stream-translator.test.ts +++ b/src/ai-proxy/lib/translators/responses-stream-translator.test.ts @@ -23,6 +23,25 @@ describe("responses-stream-translator", () => { expect(result?.delta?.reasoning_content).toBeUndefined(); }); + it("maps reasoning output_item.added to reasoning_items in cursor mode", () => { + const result = translateResponsesStreamEvent({ + event: { + type: "response.output_item.added", + output_index: 0, + item: { + id: "rs_1", + type: "reasoning", + summary: [], + status: "in_progress", + }, + }, + thinkingMode: "cursor", + }); + + expect(result?.delta?.reasoning_items?.[0]?.id).toBe("rs_1"); + expect(result?.delta?.reasoning_content).toBeUndefined(); + }); + it("maps Grok reasoning_text delta to reasoning_content in cursor mode", () => { const result = translateResponsesStreamEvent({ event: { type: "response.reasoning_text.delta", delta: "thinking" }, diff --git a/src/ai-proxy/lib/translators/responses-stream-translator.ts b/src/ai-proxy/lib/translators/responses-stream-translator.ts index 8b64c17942..349778a805 100644 --- a/src/ai-proxy/lib/translators/responses-stream-translator.ts +++ b/src/ai-proxy/lib/translators/responses-stream-translator.ts @@ -154,6 +154,24 @@ export function translateResponsesStreamEvent({ }, }; } + + if (item.type === "reasoning") { + if (useFoldedThinking) { + return null; + } + + if (!useCursorThinking) { + return null; + } + + const reasoningItem = buildReasoningItem(item); + + return { + delta: { + reasoning_items: serializeReasoningItems([reasoningItem]), + }, + }; + } } if (type === "response.function_call_arguments.delta" && typeof event.delta === "string") { diff --git a/src/ai-proxy/lib/translators/responses-to-chat-sse.ts b/src/ai-proxy/lib/translators/responses-to-chat-sse.ts index 1dc45b96c0..5f1c0cdf18 100644 --- a/src/ai-proxy/lib/translators/responses-to-chat-sse.ts +++ b/src/ai-proxy/lib/translators/responses-to-chat-sse.ts @@ -1,4 +1,5 @@ import type { ProxyProvider } from "@app/ai-proxy/lib/providers/types"; +import { safeStreamControllerError } from "@app/ai-proxy/lib/safe-stream-controller"; import { createFoldedStreamState } from "@app/ai-proxy/lib/thinking-folded"; import { type ChatStreamDelta, @@ -38,6 +39,25 @@ function chatChunk({ return SafeJSON.stringify(chunk); } +function withAssistantRole(delta: ChatStreamDelta, roleSent: { value: boolean }): ChatStreamDelta { + if (roleSent.value) { + return delta; + } + + const hasPayload = + delta.content !== undefined || + delta.reasoning_content !== undefined || + delta.reasoning_items !== undefined || + delta.tool_calls !== undefined; + + if (!hasPayload) { + return delta; + } + + roleSent.value = true; + return { role: "assistant", ...delta }; +} + export async function responsesToChatSse({ provider, upstreamModel, @@ -77,20 +97,24 @@ export async function responsesToChatSse({ if (!reader) { resolveBody(""); - controller.close(); + try { + controller.close(); + } catch (controllerErr) { + logger.debug( + { err: controllerErr, model: proxyModel }, + "ai-proxy: SSE early-return controller.close() threw" + ); + } return; } - const initialChunk = `data: ${chatChunk({ model: proxyModel, delta: { role: "assistant" } })}\n\n`; - outboundBuffer += initialChunk; - controller.enqueue(encoder.encode(initialChunk)); - let buffer = ""; let finishReason: string | null = null; let sentFinishReason = false; const decoder = new TextDecoder(); const foldedState = thinkingMode === "folded" ? createFoldedStreamState() : undefined; const toolCallIndexState = createToolCallIndexState(); + const roleSent = { value: false }; let streamSucceeded = false; try { @@ -133,9 +157,10 @@ export async function responsesToChatSse({ } if (translated.delta && Object.keys(translated.delta).length > 0) { + const delta = withAssistantRole(translated.delta, roleSent); const chunk = `data: ${chatChunk({ model: proxyModel, - delta: translated.delta, + delta, usage: translated.usage, })}\n\n`; outboundBuffer += chunk; @@ -174,12 +199,25 @@ export async function responsesToChatSse({ streamSucceeded = true; } catch (err) { logger.warn({ err, model: proxyModel }, "ai-proxy: SSE stream failed"); - controller.error(err); + + if (!safeStreamControllerError(controller, err, streamSucceeded)) { + logger.debug( + { err, model: proxyModel, streamSucceeded }, + "ai-proxy: SSE skipped controller.error (client abort or detached)" + ); + } } finally { resolveBody(outboundBuffer); if (streamSucceeded) { - controller.close(); + try { + controller.close(); + } catch (controllerErr) { + logger.warn( + { err: controllerErr, model: proxyModel }, + "ai-proxy: SSE controller.close() threw — controller already detached, likely client disconnected" + ); + } } } }, From 800152db9fa151cc2e5e6fea004974777f7aa3a3 Mon Sep 17 00:00:00 2001 From: Martin Date: Tue, 30 Jun 2026 22:34:56 +0200 Subject: [PATCH 3/3] feat(ai-proxy): thinking presentation defaults and config store --- src/ai-proxy/lib/config-store.test.ts | 4 ++-- src/ai-proxy/lib/config-store.ts | 2 +- src/ai-proxy/lib/types.ts | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ai-proxy/lib/config-store.test.ts b/src/ai-proxy/lib/config-store.test.ts index 7191827003..005205b4bd 100644 --- a/src/ai-proxy/lib/config-store.test.ts +++ b/src/ai-proxy/lib/config-store.test.ts @@ -37,9 +37,9 @@ describe("config-store migration", () => { expect(config.public?.tunnelName).toBeUndefined(); }); - it("defaults thinking presentation to raw", () => { + it("defaults thinking presentation to cursor", () => { const config = parseConfigJson(SafeJSON.stringify({})); - expect(config.translation.thinking).toBe("raw"); + expect(config.translation.thinking).toBe("cursor"); }); it("loadFresh reads disk without stale in-process cache", async () => { diff --git a/src/ai-proxy/lib/config-store.ts b/src/ai-proxy/lib/config-store.ts index dae757cbd1..ff93d03563 100644 --- a/src/ai-proxy/lib/config-store.ts +++ b/src/ai-proxy/lib/config-store.ts @@ -8,7 +8,7 @@ export function getDefaultConfig(): AiProxyConfig { return { listen: { host: "127.0.0.1", port: 8317 }, proxyApiKey: `aipx-${crypto.randomUUID().replace(/-/g, "")}`, - translation: { cursorAgent: "auto", thinking: "raw" }, + translation: { cursorAgent: "auto", thinking: "cursor" }, public: { mode: "none", basePath: "/ai" }, accounts: [], }; diff --git a/src/ai-proxy/lib/types.ts b/src/ai-proxy/lib/types.ts index 2e7aa29f74..2fd2adeec9 100644 --- a/src/ai-proxy/lib/types.ts +++ b/src/ai-proxy/lib/types.ts @@ -15,7 +15,7 @@ export interface AiProxyListenConfig { export interface AiProxyTranslationConfig { cursorAgent: CursorTranslationMode; - /** raw = inline content; cursor = reasoning_content; folded = HTML details in content */ + /** raw = passthrough; cursor = reasoning_content only (native thinking UI); folded =
in content */ thinking: ThinkingPresentationMode; }