diff --git a/src/adapters/zed-openai-stream.ts b/src/adapters/zed-openai-stream.ts new file mode 100644 index 00000000000..02f98d0dc93 --- /dev/null +++ b/src/adapters/zed-openai-stream.ts @@ -0,0 +1,340 @@ +import type { AdapterEvent, OcxMessagePhase, OcxUsage } from "../types"; +import type { TranslatorBudget } from "../lib/translator-budget"; +import { decodeServerSentEvents } from "../lib/sse-decoder"; + +function isRecord(value: unknown): value is Record { + return !!value && typeof value === "object" && !Array.isArray(value); +} + +function toolKey(tool: Record): string | undefined { + if (typeof tool.type !== "string" || !tool.type) return undefined; + const name = typeof tool.name === "string" ? tool.name : ""; + return tool.type + "\0" + name; +} + +/** + * Zed's OpenAI Responses schema accepts message and call items, not Codex's private + * additional_tools wrapper. Drop every wrapper. Promote only wrappers at or after the + * replayed prefix: earlier ones are history, not the current turn's tool catalog. + * A non-array tools field is left untouched so a malformed catalog is not replaced + * by a partial promotion. + */ +export function liftZedAdditionalTools(providerRequest: Record, replayPrefixLength: number): void { + if (!Array.isArray(providerRequest.input)) return; + if (providerRequest.tools !== undefined && !Array.isArray(providerRequest.tools)) return; + + const currentTurnStart = Number.isFinite(replayPrefixLength) + ? Math.min(providerRequest.input.length, Math.max(0, Math.trunc(replayPrefixLength))) + : 0; + const input: unknown[] = []; + const promoted: unknown[] = []; + let changed = false; + for (const [index, item] of providerRequest.input.entries()) { + if (isRecord(item) && item.type === "additional_tools" && Array.isArray(item.tools)) { + changed = true; + if (index >= currentTurnStart) promoted.push(...item.tools); + continue; + } + input.push(item); + } + if (!changed) return; + + const seen = new Set(); + const tools: unknown[] = []; + const existing = Array.isArray(providerRequest.tools) ? providerRequest.tools : []; + for (const tool of [...existing, ...promoted]) { + if (!isRecord(tool)) { + tools.push(tool); + continue; + } + const key = toolKey(tool); + if (key && seen.has(key)) continue; + if (key) seen.add(key); + tools.push(tool); + } + providerRequest.input = input; + providerRequest.tools = tools; +} + +interface CallState { + key: string; + callId: string; + name: string; + args: string; + pendingFull?: string; + started: boolean; + ended: boolean; +} + +function phaseOf(item: Record): OcxMessagePhase | undefined { + return item.phase === "commentary" || item.phase === "final_answer" ? item.phase : undefined; +} + +function messageText(item: Record): string { + if (!Array.isArray(item.content)) return ""; + return item.content + .filter((part): part is Record => isRecord(part) && part.type === "output_text" && typeof part.text === "string") + .map(part => String(part.text)) + .join(""); +} + +function usageFromResponse(response: unknown): OcxUsage | undefined { + if (!isRecord(response) || !isRecord(response.usage)) return undefined; + const usage = response.usage; + if (typeof usage.input_tokens !== "number" || typeof usage.output_tokens !== "number") return undefined; + const inputDetails = isRecord(usage.input_tokens_details) ? usage.input_tokens_details : undefined; + const outputDetails = isRecord(usage.output_tokens_details) ? usage.output_tokens_details : undefined; + return { + inputTokens: usage.input_tokens, + outputTokens: usage.output_tokens, + ...(typeof usage.total_tokens === "number" ? { totalTokens: usage.total_tokens } : {}), + ...(typeof inputDetails?.cached_tokens === "number" + ? { cachedInputTokens: inputDetails.cached_tokens, cacheReadInputTokens: inputDetails.cached_tokens } + : {}), + ...(typeof outputDetails?.reasoning_tokens === "number" + ? { reasoningOutputTokens: outputDetails.reasoning_tokens } + : {}), + }; +} + +function failureMessage(payload: Record): string { + const source = isRecord(payload.response) ? payload.response : payload; + const err = source.error ?? payload.error; + if (typeof err === "string" && err) return err; + if (isRecord(err) && typeof err.message === "string" && err.message) return err.message; + if (typeof source.message === "string" && source.message) return source.message; + return "Zed OpenAI request failed"; +} + +function incompleteReason(payload: Record): string { + const source = isRecord(payload.response) ? payload.response : payload; + const details = isRecord(source.incomplete_details) ? source.incomplete_details : undefined; + if (typeof details?.reason === "string" && details.reason) return details.reason; + return failureMessage(payload); +} + +/** + * Translate Zed's OpenAI Responses SSE into adapter events. + * + * The shared Responses passthrough parser only keeps message text, because its other caller + * is compaction. Using it here drops function and custom calls, and Codex then stores the + * first sentence as final_answer and ends the turn. A later synthetic response.completed + * with empty output (Zed's stream-ended frame) must not wipe usage taken from the real + * completion, and completed output must not repeat text or calls that already streamed. + */ +export async function* translateZedOpenAiStream( + body: ReadableStream | null, + budget: TranslatorBudget, +): AsyncGenerator { + if (!body) { + yield { type: "error", message: "Zed OpenAI response had no body" }; + return; + } + + const calls = new Map(); + const phases = new Map(); + let streamedText = ""; + let sawThinking = false; + let sawRawReasoning = false; + let usage: OcxUsage | undefined; + let terminal = false; + + function* closeOpenCallsExcept(keep: CallState): Generator { + for (const other of calls.values()) { + if (other !== keep && other.started && !other.ended) { + other.ended = true; + yield { type: "tool_call_end" }; + } + } + } + + function* startIfNeeded(state: CallState): Generator { + if (state.started || state.ended || !state.name) return; + yield* closeOpenCallsExcept(state); + state.started = true; + yield { type: "tool_call_start", id: state.callId, name: state.name }; + if (state.args) yield { type: "tool_call_delta", arguments: state.args }; + } + + function* appendArgs(state: CallState, fragment: string): Generator { + if (!fragment || state.ended) return; + if (!state.started) { + if (!state.name) { + state.args += fragment; + return; + } + yield* startIfNeeded(state); + } + state.args += fragment; + yield { type: "tool_call_delta", arguments: fragment }; + } + + function rememberFull(state: CallState, full: unknown): void { + if (typeof full === "string" && full.length > 0) state.pendingFull = full; + } + + function* finishCall(state: CallState): Generator { + if (state.ended) return; + const full = state.pendingFull; + if (typeof full === "string" && full.length > state.args.length && full.startsWith(state.args)) { + yield* appendArgs(state, full.slice(state.args.length)); + } + yield* startIfNeeded(state); + if (!state.started) return; + state.ended = true; + yield { type: "tool_call_end" }; + } + + function ensureCall(key: string, callId: string, name: string): CallState { + let state = calls.get(key); + if (!state) { + state = { key, callId, name, args: "", started: false, ended: false }; + calls.set(key, state); + return state; + } + if (callId) state.callId = callId; + if (name) state.name = name; + return state; + } + + function* emitNewText(text: string, phase?: OcxMessagePhase): Generator { + if (!text) return; + if (!streamedText) { + streamedText = text; + yield { type: "text_delta", text, ...(phase ? { phase } : {}) }; + return; + } + if (text === streamedText || streamedText.endsWith(text)) return; + if (text.startsWith(streamedText)) { + const tail = text.slice(streamedText.length); + streamedText = text; + if (tail) yield { type: "text_delta", text: tail }; + return; + } + streamedText += text; + yield { type: "text_delta", text, ...(phase ? { phase } : {}) }; + } + + function* observeToolItem(item: Record, done: boolean): Generator { + if (item.type !== "function_call" && item.type !== "custom_tool_call") return; + const key = typeof item.id === "string" ? item.id : typeof item.call_id === "string" ? item.call_id : undefined; + if (!key) return; + const callId = typeof item.call_id === "string" ? item.call_id : key; + const name = typeof item.name === "string" ? item.name : ""; + const state = ensureCall(key, callId, name); + rememberFull(state, item.type === "custom_tool_call" ? item.input : item.arguments); + if (done) yield* finishCall(state); + else yield* startIfNeeded(state); + } + + function* emitReasoningItem(item: Record): Generator { + if (!sawThinking && Array.isArray(item.summary)) { + const text = item.summary + .filter((part): part is Record => isRecord(part) && typeof part.text === "string") + .map(part => String(part.text)) + .join(""); + if (text) { + sawThinking = true; + yield { type: "thinking_delta", thinking: text }; + } + } + if (!sawRawReasoning && Array.isArray(item.content)) { + const text = item.content + .filter((part): part is Record => isRecord(part) && part.type === "reasoning_text" && typeof part.text === "string") + .map(part => String(part.text)) + .join(""); + if (text) { + sawRawReasoning = true; + yield { type: "reasoning_raw_delta", text }; + } + } + } + + function* consumeItem(item: Record, done: boolean): Generator { + if (item.type === "message") { + const phase = phaseOf(item); + if (typeof item.id === "string" && phase) phases.set(item.id, phase); + if (done) yield* emitNewText(messageText(item), phase); + return; + } + if (item.type === "function_call" || item.type === "custom_tool_call") { + yield* observeToolItem(item, done); + return; + } + if (done && item.type === "reasoning") yield* emitReasoningItem(item); + } + + function* onPayload(payload: Record): Generator { + if (payload.type === "error" || payload.type === "response.failed") { + terminal = true; + yield { type: "error", message: failureMessage(payload) }; + return; + } + if (payload.type === "response.incomplete") { + terminal = true; + for (const state of calls.values()) yield* finishCall(state); + yield { type: "incomplete", reason: incompleteReason(payload) }; + return; + } + if (payload.type === "response.output_item.added" || payload.type === "response.output_item.done") { + if (isRecord(payload.item)) yield* consumeItem(payload.item, payload.type === "response.output_item.done"); + return; + } + if (payload.type === "response.output_text.delta" && typeof payload.delta === "string" && payload.delta) { + const itemId = typeof payload.item_id === "string" ? payload.item_id : undefined; + const phase = itemId ? phases.get(itemId) : undefined; + streamedText += payload.delta; + yield { type: "text_delta", text: payload.delta, ...(phase ? { phase } : {}) }; + return; + } + if (payload.type === "response.reasoning_summary_text.delta" && typeof payload.delta === "string" && payload.delta) { + sawThinking = true; + yield { type: "thinking_delta", thinking: payload.delta }; + return; + } + if (payload.type === "response.reasoning_text.delta" && typeof payload.delta === "string" && payload.delta) { + sawRawReasoning = true; + yield { type: "reasoning_raw_delta", text: payload.delta }; + return; + } + if ( + (payload.type === "response.function_call_arguments.delta" || payload.type === "response.custom_tool_call_input.delta") + && typeof payload.delta === "string" + ) { + const key = typeof payload.item_id === "string" ? payload.item_id : undefined; + if (!key) return; + yield* appendArgs(ensureCall(key, key, ""), payload.delta); + return; + } + if (payload.type === "response.function_call_arguments.done" || payload.type === "response.custom_tool_call_input.done") { + const key = typeof payload.item_id === "string" ? payload.item_id : undefined; + if (!key) return; + const state = ensureCall(key, key, ""); + rememberFull(state, payload.type === "response.custom_tool_call_input.done" ? payload.input : payload.arguments); + yield* finishCall(state); + return; + } + if (payload.type === "response.completed") { + const response = isRecord(payload.response) ? payload.response : undefined; + if (response && Array.isArray(response.output)) { + for (const item of response.output) { + if (isRecord(item)) yield* consumeItem(item, true); + } + } + const nextUsage = usageFromResponse(response); + if (nextUsage) usage = nextUsage; + } + } + + for await (const event of decodeServerSentEvents(body, { translatorBudget: budget })) { + let payload: unknown; + try { payload = JSON.parse(event.data); } catch { continue; } + if (!isRecord(payload)) continue; + yield* onPayload(payload); + if (terminal) return; + } + + for (const state of calls.values()) yield* finishCall(state); + yield { type: "done", ...(usage ? { usage } : {}) }; +} + diff --git a/src/adapters/zed.ts b/src/adapters/zed.ts index 0004090896c..322e504f76a 100644 --- a/src/adapters/zed.ts +++ b/src/adapters/zed.ts @@ -5,6 +5,7 @@ import type { AnthropicCacheRetention } from "./anthropic/cache-retention"; import { createGoogleAdapter } from "./google"; import { createOpenAIChatAdapter } from "./openai-chat"; import { createResponsesPassthroughAdapter } from "./openai-responses"; +import { liftZedAdditionalTools, translateZedOpenAiStream } from "./zed-openai-stream"; import { createAdapterPhysicalSend } from "./physical-send"; import type { AdapterEvent, OcxParsedRequest, OcxProviderConfig } from "../types"; import { createTranslatorBudget, releaseTranslatedEvent, retainTranslatedEvent, type TranslatorBudget } from "../lib/translator-budget"; @@ -314,6 +315,12 @@ export function createZedAdapter(provider: OcxProviderConfig, cacheRetention?: A : item, ); } + if (zedProvider === "open_ai") { + // Codex carries this turn's client tools in a private additional_tools input item. + // Zed rejects that variant, so promote the current turn and drop every wrapper + // before the built-in tool filter runs. + liftZedAdditionalTools(providerRequest, parsed._replayPrefixLen ?? 0); + } if (zedProvider === "open_ai" && Array.isArray(providerRequest.tools)) { // Zed's Responses decoder supports function/custom tools, not Codex's built-in web_search. const supported = (tool: unknown) => isRecord(tool) && (tool.type === "function" || tool.type === "custom"); @@ -410,7 +417,11 @@ export function createZedAdapter(provider: OcxProviderConfig, cacheRetention?: A headers: { "Content-Type": "text/event-stream" }, }); try { - yield* delegate.adapter.parseStream(translated, budget); + // The Responses passthrough parser only keeps compaction text. Zed is not a + // passthrough route, so function and custom calls have to be translated or + // Codex ends the turn on the first sentence. + if (delegate.provider === "open_ai") yield* translateZedOpenAiStream(translated.body, budget); + else yield* delegate.adapter.parseStream(translated, budget); } finally { // Some delegates only unlock their reader when a collector stops early. try { await translated.body?.cancel(); } catch { /* already closed/errored */ } diff --git a/tests/providers/zed-hosted-provider.test.ts b/tests/providers/zed-hosted-provider.test.ts index 7442777fba2..cb1c889edc9 100644 --- a/tests/providers/zed-hosted-provider.test.ts +++ b/tests/providers/zed-hosted-provider.test.ts @@ -362,6 +362,120 @@ describe("Zed OpenAI tool selection after filtering", () => { }); }); +describe("Zed OpenAI tool relay", () => { + const provider: OcxProviderConfig = { + adapter: "zed", + baseUrl: "https://cloud.zed.dev", + authMode: "oauth", + apiKey: "access-token", + }; + + function catalogFetch(): void { + const responses = [ + jsonResponse({ default_organization_id: "org-1" }), + jsonResponse({ token: "llm-token-1" }), + jsonResponse({ models: [{ id: "gpt-6.1-sol", provider: "open_ai" }] }), + ]; + globalThis.fetch = (async () => { + const response = responses.shift(); + if (!response) throw new Error("unexpected Zed catalog fetch"); + return response; + }) as typeof globalThis.fetch; + } + + test("promotes the current turn's additional_tools and drops replayed wrappers", async () => { + catalogFetch(); + const already = { type: "function", name: "already", parameters: { type: "object", properties: {} } }; + const parsed = parseRequest({ + model: "gpt-6.1-sol", + input: [ + { type: "additional_tools", tools: [{ type: "function", name: "old_tool", parameters: { type: "object", properties: {} } }] }, + { type: "message", role: "user", content: [{ type: "input_text", text: "earlier" }] }, + { type: "additional_tools", tools: [ + { type: "function", name: "shell", parameters: { type: "object", properties: {} } }, + { type: "web_search" }, + already, + ] }, + { type: "message", role: "user", content: [{ type: "input_text", text: "now" }] }, + ], + tools: [already], + stream: true, + }); + parsed._zedAuthContext = { userId: "user-1" }; + parsed._replayPrefixLen = 2; + const adapter = withTestTranslatorBudget(createZedAdapter(provider)); + const request = await adapter.buildRequest(parsed); + const body = JSON.parse(request.body).provider_request; + expect(body.input.map((item: { type: string }) => item.type)).toEqual(["message", "message"]); + expect(JSON.stringify(body.input)).not.toContain("old_tool"); + expect(JSON.stringify(body.input)).not.toContain("additional_tools"); + expect(body.tools).toEqual([ + already, + { type: "function", name: "shell", parameters: { type: "object", properties: {} } }, + ]); + }); + + async function relay(frames: unknown[]): Promise { + catalogFetch(); + const parsed = parseRequest({ model: "gpt-6.1-sol", input: "hello", stream: true }); + parsed._zedAuthContext = { userId: "user-1" }; + const adapter = withTestTranslatorBudget(createZedAdapter(provider)); + await adapter.buildRequest(parsed); + const ndjson = frames.map(frame => JSON.stringify(frame)).join("\n") + "\n"; + const events = []; + for await (const event of adapter.parseStream(new Response(ndjson))) events.push(event); + return events; + } + + test("relays commentary text, function-call deltas, and usage across the synthetic completion", async () => { + const events = await relay([ + { event: { type: "response.output_item.added", item: { id: "msg_1", type: "message", role: "assistant", phase: "commentary", content: [] } } }, + { event: { type: "response.output_text.delta", item_id: "msg_1", delta: "Looking." } }, + { event: { type: "response.output_item.added", item: { id: "fc_1", type: "function_call", call_id: "call_1", name: "shell", arguments: "" } } }, + { event: { type: "response.function_call_arguments.delta", item_id: "fc_1", delta: "{\"cmd\":" } }, + { event: { type: "response.function_call_arguments.delta", item_id: "fc_1", delta: "\"ls\"}" } }, + { event: { type: "response.function_call_arguments.done", item_id: "fc_1", arguments: "{\"cmd\":\"ls\"}" } }, + { event: { type: "response.completed", response: { + output: [ + { type: "message", id: "msg_1", phase: "commentary", content: [{ type: "output_text", text: "Looking." }] }, + { type: "function_call", id: "fc_1", call_id: "call_1", name: "shell", arguments: "{\"cmd\":\"ls\"}" }, + ], + usage: { input_tokens: 11, output_tokens: 6 }, + } } }, + { status: "stream_ended" }, + ]); + expect(events).toEqual([ + { type: "text_delta", text: "Looking.", phase: "commentary" }, + { type: "tool_call_start", id: "call_1", name: "shell" }, + { type: "tool_call_delta", arguments: "{\"cmd\":" }, + { type: "tool_call_delta", arguments: "\"ls\"}" }, + { type: "tool_call_end" }, + { type: "done", usage: { inputTokens: 11, outputTokens: 6 } }, + ]); + }); + + test("relays a custom tool that arrives only in the completed output without repeating text", async () => { + const events = await relay([ + { event: { type: "response.output_text.delta", delta: "hi" } }, + { event: { type: "response.completed", response: { + output: [ + { type: "message", id: "msg_1", phase: "commentary", content: [{ type: "output_text", text: "hi" }] }, + { type: "custom_tool_call", id: "ctc_1", call_id: "call_c", name: "apply_patch", input: "patch body" }, + ], + usage: { input_tokens: 3, output_tokens: 2 }, + } } }, + { status: "stream_ended" }, + ]); + expect(events).toEqual([ + { type: "text_delta", text: "hi" }, + { type: "tool_call_start", id: "call_c", name: "apply_patch" }, + { type: "tool_call_delta", arguments: "patch body" }, + { type: "tool_call_end" }, + { type: "done", usage: { inputTokens: 3, outputTokens: 2 } }, + ]); + }); +}); + describe("Zed inference send accounting", () => { async function fixture(limit: number, full = false) { const catalog = [jsonResponse({ default_organization_id: "org-1" }),