diff --git a/src/main/services/provider-tool-turns/chat-format.ts b/src/main/services/provider-tool-turns/chat-format.ts index fef93611..2dcf05f4 100644 --- a/src/main/services/provider-tool-turns/chat-format.ts +++ b/src/main/services/provider-tool-turns/chat-format.ts @@ -1,6 +1,6 @@ -/** OpenAI Chat Completions encoding/decoding (xAI, Groq, Mistral, Cerebras, OpenRouter). */ -import { isRecord, type JsonRecord } from "../provider-http"; -import { parseArguments, type AgentTurnMessage, type ToolCall, type ToolDefinition } from "../provider-tool-turn"; +/** OpenAI Chat Completions encoding/decoding (xAI, Groq, Mistral, Cerebras, OpenRouter, Kimi family). */ +import { isRecord, type JsonRecord, type ProviderEventStreamMessage } from "../provider-http"; +import { parseArguments, type AgentTurnMessage, type AgentTurnResult, type ToolCall, type ToolDefinition } from "../provider-tool-turn"; export function chatToolDefs(tools: ToolDefinition[]): unknown[] { return tools.map((tool) => ({ type: "function", function: { name: tool.name, description: tool.description, parameters: tool.parameters } })); @@ -44,3 +44,52 @@ export function parseChatChoice(data: JsonRecord): { text: string; toolCalls: To } return { text: text.trim(), toolCalls }; } + +/** + * Accumulate a Chat Completions SSE stream into the final turn result. + * + * Streaming here is load-bearing, not cosmetic: reasoning models (the Kimi + * family) can think well past the 60s header deadline before a buffered + * response would start, so the turn must consume deltas as they arrive and + * let the stream idle deadline govern instead. Tool calls arrive as indexed + * fragments — the opening chunk of a call carries id and name, later chunks + * append argument text. Usage-only chunks have no choices; reasoning_content + * deltas keep the stream alive but are not part of the result. + */ +export async function accumulateChatStream(messages: AsyncIterable): Promise { + const raw: string[] = []; + const textParts: string[] = []; + const fragments = new Map(); + for await (const message of messages) { + raw.push(message.rawData); + if (message.data.error !== undefined) throw new Error(streamErrorText(message.data.error)); + const choices = Array.isArray(message.data.choices) ? message.data.choices : []; + const first = choices.length > 0 && isRecord(choices[0]) ? choices[0] : null; + const delta = first && isRecord(first.delta) ? first.delta : null; + if (!delta) continue; + if (typeof delta.content === "string") textParts.push(delta.content); + const calls = Array.isArray(delta.tool_calls) ? delta.tool_calls : []; + for (const call of calls) { + if (!isRecord(call)) continue; + const index = typeof call.index === "number" ? call.index : fragments.size; + const fragment = fragments.get(index) ?? { id: "", name: "", args: "" }; + if (typeof call.id === "string" && call.id) fragment.id = call.id; + const fn = isRecord(call.function) ? call.function : null; + if (fn && typeof fn.name === "string" && fn.name) fragment.name = fn.name; + if (fn && typeof fn.arguments === "string") fragment.args += fn.arguments; + fragments.set(index, fragment); + } + } + if (raw.length === 0) throw new Error("Chat completions stream ended without any events."); + const toolCalls: ToolCall[] = []; + for (const [, fragment] of [...fragments.entries()].sort(([a], [b]) => a - b)) { + if (fragment.id && fragment.name) toolCalls.push({ id: fragment.id, name: fragment.name, arguments: parseArguments(fragment.args) }); + } + return { text: textParts.join("").trim(), toolCalls, rawResponseJson: JSON.stringify(raw) }; +} + +function streamErrorText(error: unknown): string { + if (typeof error === "string") return error; + if (isRecord(error) && typeof error.message === "string") return error.message; + return "Provider stream failed."; +} diff --git a/src/main/services/provider-tool-turns/turns.ts b/src/main/services/provider-tool-turns/turns.ts index f17c1c3c..296628b8 100644 --- a/src/main/services/provider-tool-turns/turns.ts +++ b/src/main/services/provider-tool-turns/turns.ts @@ -8,7 +8,7 @@ import { postEventStream, postJson, type JsonRecord } from "../provider-http"; import type { AgentTurnMessage, AgentTurnResult, ProviderToolTurnInput, TurnCredential } from "../provider-tool-turn"; import { accumulateResponsesStream, parseResponsesOutput, responsesToolDefs, toResponsesInput } from "./responses-format"; import { anthropicToolDefs, parseAnthropicContent, toAnthropicMessages } from "./anthropic-format"; -import { chatToolDefs, parseChatChoice, toChatMessages } from "./chat-format"; +import { accumulateChatStream, chatToolDefs, parseChatChoice, toChatMessages } from "./chat-format"; import { accumulateGeminiStream, antigravityRequestBody, parseGeminiResponse, vertexRequestBody } from "./gemini-format"; import { bedrockToolConfig, parseBedrockOutput, toBedrockMessages } from "./bedrock-format"; import { signAwsRequest } from "./aws-sigv4"; @@ -51,21 +51,29 @@ export async function runOpenAiResponsesTurn(input: ProviderToolTurnInput): Prom return { ...parsed, rawResponseJson: JSON.stringify(response.data) }; } -// ---- OpenAI Chat Completions (xAI, Groq, Mistral, Cerebras, OpenRouter) ---- +// ---- OpenAI Chat Completions, streamed (xAI, Groq, Mistral, Cerebras, OpenRouter, Kimi, Kimi CN) ---- + +/** + * Streaming is required, not an optimization: a non-streaming completion + * sends its response headers only after the WHOLE generation finishes, so a + * reasoning model that thinks past the 60s header deadline gets aborted with + * "Provider request timed out." even though it is working. Streamed turns + * emit chunks continuously and only a truly dead connection times out. + */ +async function runChatCompletionsStream(endpoint: string, headers: Record, input: ProviderToolTurnInput): Promise { + const body: JsonRecord = { + model: input.model, + messages: toChatMessages(input.instructions, input.messages), + tools: chatToolDefs(input.tools), + tool_choice: "auto", + stream: true, + }; + const response = await postEventStream(endpoint, headers, body); + return accumulateChatStream(response.messages); +} export function makeChatCompletionsTurn(endpoint: string, label: string): (input: ProviderToolTurnInput) => Promise { - return async (input) => { - const body: JsonRecord = { - model: input.model, - messages: toChatMessages(input.instructions, input.messages), - tools: chatToolDefs(input.tools), - tool_choice: "auto", - stream: false, - }; - const response = await postJson(endpoint, { authorization: `Bearer ${requireApiKey(input.credential, label)}` }, body); - const parsed = parseChatChoice(response.data); - return { ...parsed, rawResponseJson: JSON.stringify(response.data) }; - }; + return async (input) => runChatCompletionsStream(endpoint, { authorization: `Bearer ${requireApiKey(input.credential, label)}` }, input); } // ---- Anthropic Messages (api key) ---- @@ -109,20 +117,11 @@ export async function runClaudeProTurn(input: ProviderToolTurnInput): Promise { const access = requireAccessToken(input.credential, "Kimi Code"); - const body: JsonRecord = { - model: input.model, - messages: toChatMessages(input.instructions, input.messages), - tools: chatToolDefs(input.tools), - tool_choice: "auto", - stream: false, - }; - const response = await postJson(`${KIMI_CODE_BASE_URL}/chat/completions`, { authorization: `Bearer ${access}` }, body); - const parsed = parseChatChoice(response.data); - return { ...parsed, rawResponseJson: JSON.stringify(response.data) }; + return runChatCompletionsStream(`${KIMI_CODE_BASE_URL}/chat/completions`, { authorization: `Bearer ${access}` }, input); } // ---- GPT account / OpenAI Codex (OAuth, SSE) ---- diff --git a/tests/unit/provider-tool-turns.test.ts b/tests/unit/provider-tool-turns.test.ts index b87f4ae4..e24e63b4 100644 --- a/tests/unit/provider-tool-turns.test.ts +++ b/tests/unit/provider-tool-turns.test.ts @@ -5,7 +5,7 @@ import { describe, expect, test } from "vitest"; import type { JsonRecord, ProviderEventStreamMessage } from "../../src/main/services/provider-http"; import type { AgentTurnMessage } from "../../src/main/services/provider-tool-turn"; import { accumulateResponsesStream, parseResponsesOutput, responsesToolDefs, toResponsesInput } from "../../src/main/services/provider-tool-turns/responses-format"; -import { parseChatChoice, toChatMessages } from "../../src/main/services/provider-tool-turns/chat-format"; +import { accumulateChatStream, parseChatChoice, toChatMessages } from "../../src/main/services/provider-tool-turns/chat-format"; import { parseAnthropicContent, toAnthropicMessages } from "../../src/main/services/provider-tool-turns/anthropic-format"; import { accumulateGeminiStream, antigravityRequestBody } from "../../src/main/services/provider-tool-turns/gemini-format"; import { bedrockToolConfig, parseBedrockOutput, toBedrockMessages } from "../../src/main/services/provider-tool-turns/bedrock-format"; @@ -90,6 +90,33 @@ describe("openai chat completions encoding", () => { const data = { choices: [{ message: { content: "hi", tool_calls: [{ id: "c9", function: { name: "f", arguments: '{"a":1}' } }] } }] }; expect(parseChatChoice(data)).toEqual({ text: "hi", toolCalls: [{ id: "c9", name: "f", arguments: { a: 1 } }] }); }); + + test("accumulates streamed content deltas and indexed tool-call fragments", async () => { + // The Kimi coding endpoint streams reasoning_content first (which must be + // ignored), then content deltas, then tool calls split across chunks: the + // opening fragment carries index/id/name, later fragments append arguments. + const result = await accumulateChatStream(stream([ + { id: "cmpl-1", choices: [{ delta: { role: "assistant", reasoning_content: "thinking..." } }] }, + { choices: [{ delta: { content: "Staffing " } }] }, + { choices: [{ delta: { content: "the team." } }] }, + { choices: [{ delta: { tool_calls: [{ index: 0, id: "c1", type: "function", function: { name: "hire_agent", arguments: "" } }] } }] }, + { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '{"name":' } }, { index: 1, id: "c2", type: "function", function: { name: "create_task", arguments: '{"title":"x"}' } }] } }] }, + { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '"Sam"}' } }] }, finish_reason: "tool_calls" }] }, + { choices: [], usage: { prompt_tokens: 10, completion_tokens: 5 } }, + ])); + expect(result.text).toBe("Staffing the team."); + expect(result.toolCalls).toEqual([ + { id: "c1", name: "hire_agent", arguments: { name: "Sam" } }, + { id: "c2", name: "create_task", arguments: { title: "x" } }, + ]); + }); + + test("fails the turn on a streamed error event and on an empty stream", async () => { + await expect(accumulateChatStream(stream([ + { error: { message: "The usage limit has been reached" } }, + ]))).rejects.toThrow("The usage limit has been reached"); + await expect(accumulateChatStream(stream([]))).rejects.toThrow("Chat completions stream ended without any events."); + }); }); describe("anthropic encoding", () => {