diff --git a/src/agent/ag-ui/handler.test.ts b/src/agent/ag-ui/handler.test.ts index 8a6ebd7df8..a8e3b177ee 100644 --- a/src/agent/ag-ui/handler.test.ts +++ b/src/agent/ag-ui/handler.test.ts @@ -175,6 +175,131 @@ describe("agent/ag-ui-handler", () => { assertStringIncludes(body, '"delta":"hello from runtime"'); }); + it("bridges direct tool data events into the AG-UI stream", async () => { + const testAgent = createTestAgent(); + testAgent.agent.stream = async (input) => { + const publishDataEvent = input.context?.publishDataEvent; + if (typeof publishDataEvent === "function") { + await publishDataEvent({ + type: "test.report", + name: "test.report", + value: { status: "ready" }, + }); + } + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue( + encodeDataStreamEvent({ type: "message-start", messageId: "assistant-msg-1" }), + ); + controller.enqueue(encodeDataStreamEvent({ type: "text-start", id: "text-1" })); + controller.enqueue( + encodeDataStreamEvent({ type: "text-delta", id: "text-1", delta: "done" }), + ); + controller.enqueue(encodeDataStreamEvent({ type: "text-end", id: "text-1" })); + controller.close(); + }, + }); + + return { + toDataStreamResponse: () => + new Response(stream, { + headers: { "Content-Type": "text/event-stream" }, + }), + }; + }; + + const handler = createAgUiHandler({ agent: testAgent.agent }); + const response = await handler( + new Request("http://localhost/api/ag-ui", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + messages: [{ + id: "msg-1", + role: "user", + parts: [{ type: "text", text: "hello" }], + }], + }), + }), + ); + + const body = await response.text(); + assertStringIncludes(body, "event: Custom"); + assertStringIncludes(body, '"name":"test.report"'); + assertStringIncludes(body, '"status":"ready"'); + }); + + it("bridges injected-tools tool data events into the AG-UI stream exactly once", async () => { + const sessionManager = new RunResumeSessionManager<{ + result: unknown; + isError: boolean; + }>(); + const originalStream = AgentRuntime.prototype.stream; + + AgentRuntime.prototype.stream = async function ( + _messages, + context, + ): Promise> { + const publishDataEvent = context?.publishDataEvent; + if (typeof publishDataEvent === "function") { + await publishDataEvent({ + type: "test.report", + name: "test.report", + value: { status: "ready" }, + }); + } + + return new ReadableStream({ + start(controller) { + controller.enqueue( + encodeDataStreamEvent({ type: "message-start", messageId: "assistant-msg-1" }), + ); + controller.enqueue(encodeDataStreamEvent({ type: "text-start", id: "text-1" })); + controller.enqueue( + encodeDataStreamEvent({ type: "text-delta", id: "text-1", delta: "done" }), + ); + controller.enqueue(encodeDataStreamEvent({ type: "text-end", id: "text-1" })); + controller.close(); + }, + }); + }; + + try { + const handler = createAgUiHandler({ + agent: createTestAgent().agent, + sessionManager, + }); + + const response = await handler( + new Request("http://localhost/api/ag-ui", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + runId: "run_data_1", + threadId: crypto.randomUUID(), + messages: [{ + id: "msg-1", + role: "user", + parts: [{ type: "text", text: "hello" }], + }], + tools: [{ name: "client_confirm" }], + }), + }), + ); + + const body = await response.text(); + assertStringIncludes(body, "event: Custom"); + assertStringIncludes(body, '"name":"test.report"'); + assertStringIncludes(body, '"status":"ready"'); + // The injected path injects publishDataEvent and wraps the stream once, + // so the event must surface exactly once (no double-emit). + assertEquals(body.match(/"name":"test\.report"/g)?.length, 1); + } finally { + AgentRuntime.prototype.stream = originalStream; + } + }); + it("runs beforeStream before direct AG-UI streaming", async () => { const testAgent = createTestAgent(); const handler = createAgUiHandler({ diff --git a/src/agent/ag-ui/handler.ts b/src/agent/ag-ui/handler.ts index 51a1ea60fb..dcf3d0c97b 100644 --- a/src/agent/ag-ui/handler.ts +++ b/src/agent/ag-ui/handler.ts @@ -15,11 +15,16 @@ import { mapRuntimeEventToAgUi, } from "#veryfront/internal-agents/ag-ui-sse.ts"; import { streamDataStreamEvents } from "../streaming/data-stream.ts"; +import { + createToolExecutionDataEventBridgeStream, + type ToolExecutionDataEventPublisher, +} from "../streaming/tool-execution-data-event-bridge.ts"; import { type AgUiBeforeStream, applyBeforeStreamResult, extractLastUserText, } from "../service/before-stream.ts"; +import type { ToolExecutionDataEvent } from "#veryfront/tool/types.ts"; import { type AgUiRequest, normalizeAgUiMessages, @@ -49,6 +54,29 @@ function generateRunId(): string { return `run_${crypto.randomUUID().replaceAll("-", "")}`; } +function createToolDataEventBridge() { + const pendingEvents: ToolExecutionDataEvent[] = []; + let publishDataEvent: ToolExecutionDataEventPublisher = (event) => { + pendingEvents.push(event); + }; + + return { + publishDataEvent: (event: ToolExecutionDataEvent) => publishDataEvent(event), + wrapStream(baseStream: ReadableStream): ReadableStream { + return createToolExecutionDataEventBridgeStream({ + baseStream, + installPublisher(nextPublishDataEvent) { + publishDataEvent = nextPublishDataEvent; + while (pendingEvents.length > 0) { + const event = pendingEvents.shift(); + if (event) publishDataEvent(event); + } + }, + }); + }, + }; +} + function buildStreamContext( request: AgUiRequest, baseContext: Record, @@ -213,20 +241,25 @@ async function createAgUiDirectStreamResponse( await agent.clearMemory(); + const toolDataEvents = createToolDataEventBridge(); const result = await agent.stream({ messages, - context: finalContext, + context: { + ...finalContext, + publishDataEvent: toolDataEvents.publishDataEvent, + }, ...(request.model ? { model: request.model } : {}), ...(request.maxOutputTokens ? { maxOutputTokens: request.maxOutputTokens } : {}), }); const upstream = result.toDataStreamResponse(); + const upstreamBody = upstream.body ? toolDataEvents.wrapStream(upstream.body) : upstream.body; return await createAgUiStreamResponse({ agentId: agent.id, request, runId, threadId, - upstreamBody: upstream.body, + upstreamBody, upstreamStatus: upstream.status, upstreamStatusText: upstream.statusText, }); @@ -271,14 +304,19 @@ async function createAgUiInjectedToolsStreamResponse( }); let upstreamBody: ReadableStream; + const toolDataEvents = createToolDataEventBridge(); try { upstreamBody = await runtime.stream( messages, - finalContext, + { + ...finalContext, + publishDataEvent: toolDataEvents.publishDataEvent, + }, undefined, request.model, request.maxOutputTokens, ); + upstreamBody = toolDataEvents.wrapStream(upstreamBody); } catch (error) { sessionManager.failRun(runId); throw error; diff --git a/src/agent/ag-ui/runtime-chat-stream-encoder.ts b/src/agent/ag-ui/runtime-chat-stream-encoder.ts index 5fc48b0253..788a251df6 100644 --- a/src/agent/ag-ui/runtime-chat-stream-encoder.ts +++ b/src/agent/ag-ui/runtime-chat-stream-encoder.ts @@ -354,8 +354,16 @@ export function createAgUiRuntimeChatStreamEncoder( }); return events; } - default: + default: { + if (!event.type.startsWith("data-")) { + return events; + } + events.push({ + type: event.type as `data-${string}`, + data: event.data, + }); return events; + } } }, }; diff --git a/src/agent/react/use-chat/streaming/handler.ts b/src/agent/react/use-chat/streaming/handler.ts index 4889dc6f7a..204a9b2297 100644 --- a/src/agent/react/use-chat/streaming/handler.ts +++ b/src/agent/react/use-chat/streaming/handler.ts @@ -8,6 +8,7 @@ import type { ChatMessagePart, ChatToolPart } from "#veryfront/agent/react/use-c import { createAssistantMessage, generateClientId } from "#veryfront/agent/react/use-chat/utils.ts"; import { buildCurrentParts } from "#veryfront/agent/react/use-chat/streaming/parts-builder.ts"; import type { + OrderedMessagePart, OrderedReasoning, OrderedStep, OrderedToolCall, @@ -21,6 +22,7 @@ interface StreamingState { reasoningBlocks: Map; steps: Map; messageParts: ChatMessagePart[]; + dataParts: OrderedMessagePart[]; currentTextId: string; messageId: string; partOrderCounter: number; @@ -34,6 +36,7 @@ function createStreamingState(): StreamingState { reasoningBlocks: new Map(), steps: new Map(), messageParts: [], + dataParts: [], currentTextId: "", messageId: "", partOrderCounter: 0, @@ -50,7 +53,13 @@ export async function handleStreamingResponse( const state = createStreamingState(); const getBuildParts = (): ChatMessagePart[] => - buildCurrentParts(state.textBlocks, state.reasoningBlocks, state.toolCalls, state.steps); + buildCurrentParts( + state.textBlocks, + state.reasoningBlocks, + state.toolCalls, + state.steps, + state.dataParts, + ); let buffer = ""; @@ -101,7 +110,13 @@ export async function handleAgUiStreamingResponse( const state = createStreamingState(); const getBuildParts = (): ChatMessagePart[] => - buildCurrentParts(state.textBlocks, state.reasoningBlocks, state.toolCalls, state.steps); + buildCurrentParts( + state.textBlocks, + state.reasoningBlocks, + state.toolCalls, + state.steps, + state.dataParts, + ); const processDecodedEvents = (events: ChatStreamEvent[]) => { for (const event of events) { @@ -203,6 +218,15 @@ function processStreamEvent( return; default: + if (typeof parsed.type === "string" && parsed.type.startsWith("data-")) { + handleDataPart( + { type: parsed.type, data: parsed.data }, + state, + onUpdate, + getBuildParts, + ); + onData(parsed.data); + } return; } } @@ -289,7 +313,9 @@ function processChatStreamEvent( default: if (event.type.startsWith("data-")) { - onData((event as { data: unknown }).data); + const data = (event as { data: unknown }).data; + handleDataPart({ type: event.type, data }, state, onUpdate, getBuildParts); + onData(data); } return; } @@ -301,6 +327,7 @@ function handleStart(parsed: Record, state: StreamingState): vo state.toolCalls.clear(); state.reasoningBlocks.clear(); state.messageParts.length = 0; + state.dataParts.length = 0; } function handleTextStart(parsed: Record, state: StreamingState): void { @@ -351,6 +378,37 @@ function handleTextEnd(parsed: Record, state: StreamingState): } } +function handleDataPart( + parsed: { type: string; data?: unknown }, + state: StreamingState, + onUpdate: StreamingCallbacks["onUpdate"], + getBuildParts: () => ChatMessagePart[], +): void { + if (!isRenderableDataPartType(parsed.type)) { + return; + } + + if (!state.messageId) { + state.messageId = generateClientId("msg"); + } + + state.dataParts.push({ + order: state.partOrderCounter++, + part: { + type: parsed.type as `data-${string}`, + data: parsed.data, + }, + }); + + onUpdate?.(getBuildParts(), state.messageId); +} + +function isRenderableDataPartType(type: string): boolean { + return type !== "data-state-snapshot" && + type !== "data-state-delta" && + type !== "data-messages-snapshot"; +} + function handleToolInputStart( parsed: Record, state: StreamingState, diff --git a/src/agent/react/use-chat/streaming/parts-builder.ts b/src/agent/react/use-chat/streaming/parts-builder.ts index 7888d5eff1..d267b56104 100644 --- a/src/agent/react/use-chat/streaming/parts-builder.ts +++ b/src/agent/react/use-chat/streaming/parts-builder.ts @@ -1,5 +1,11 @@ import type { ChatMessagePart, ChatToolPart } from "../types.ts"; -import type { OrderedReasoning, OrderedStep, OrderedToolCall, TextBlock } from "./types.ts"; +import type { + OrderedMessagePart, + OrderedReasoning, + OrderedStep, + OrderedToolCall, + TextBlock, +} from "./types.ts"; interface OrderedPart { order: number; @@ -11,6 +17,7 @@ export function buildCurrentParts( reasoningBlocks: Map, toolCalls: Map, steps?: Map, + extraParts?: OrderedMessagePart[], ): ChatMessagePart[] { const orderedParts: OrderedPart[] = []; @@ -18,11 +25,21 @@ export function buildCurrentParts( addReasoningParts(orderedParts, reasoningBlocks); addToolParts(orderedParts, toolCalls); if (steps) addStepParts(orderedParts, steps); + if (extraParts) addExtraParts(orderedParts, extraParts); orderedParts.sort((a, b) => a.order - b.order); return orderedParts.map(({ part }) => part); } +function addExtraParts( + orderedParts: OrderedPart[], + extraParts: OrderedMessagePart[], +): void { + for (const { order, part } of extraParts) { + orderedParts.push({ order, part }); + } +} + function addTextParts( orderedParts: OrderedPart[], textBlocks: Map, diff --git a/src/agent/react/use-chat/streaming/types.ts b/src/agent/react/use-chat/streaming/types.ts index 190a88e22f..aa3be9adc7 100644 --- a/src/agent/react/use-chat/streaming/types.ts +++ b/src/agent/react/use-chat/streaming/types.ts @@ -39,3 +39,4 @@ export interface OrderedStep { export type OrderedToolCall = StreamingToolCall & { order: number }; export type OrderedReasoning = StreamingReasoning & { order: number }; +export type OrderedMessagePart = { order: number; part: ChatMessagePart }; diff --git a/src/agent/react/use-chat/types.ts b/src/agent/react/use-chat/types.ts index 0c3c032afc..29e6a13b4a 100644 --- a/src/agent/react/use-chat/types.ts +++ b/src/agent/react/use-chat/types.ts @@ -1,4 +1,5 @@ import type { + ChatDataPart, ChatDynamicToolPart, ChatMessage, ChatMessagePart, @@ -23,6 +24,7 @@ export type BrowserInferenceStatus = | "error"; export type { + ChatDataPart, ChatDynamicToolPart, ChatMessage, ChatMessagePart, diff --git a/src/agent/react/use-chat/use-chat.state.test.ts b/src/agent/react/use-chat/use-chat.state.test.ts index 61011f7c4c..325a2ea7ec 100644 --- a/src/agent/react/use-chat/use-chat.state.test.ts +++ b/src/agent/react/use-chat/use-chat.state.test.ts @@ -1,5 +1,5 @@ import "#veryfront/schemas/_test-setup.ts"; -import { assertEquals } from "#veryfront/testing/assert.ts"; +import { assertEquals, assertExists } from "#veryfront/testing/assert.ts"; import { describe, it } from "#veryfront/testing/bdd.ts"; import { handleAgUiStreamingResponse } from "#veryfront/agent/react/use-chat/streaming/index.ts"; import type { ChatMessage } from "./types.ts"; @@ -70,4 +70,91 @@ describe("use-chat internal state helpers", () => { assertEquals(resolveUseChatStreamHandler(undefined), handleAgUiStreamingResponse); assertEquals(resolveUseChatStreamHandler("ag-ui"), handleAgUiStreamingResponse); }); + + it("preserves AG-UI custom data events as assistant message parts", async () => { + const encoder = new TextEncoder(); + const body = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode([ + "event: Custom", + 'data: {"name":"dora.report","value":{"overall":"fail"}}', + "", + "event: TextMessageStart", + 'data: {"messageId":"msg-1","role":"assistant"}', + "", + "event: TextMessageContent", + 'data: {"messageId":"msg-1","delta":"Done"}', + "", + "event: TextMessageEnd", + 'data: {"messageId":"msg-1"}', + "", + "event: RunFinished", + 'data: {"threadId":"thread-1","runId":"run-1"}', + "", + "", + ].join("\n"))); + controller.close(); + }, + }); + const messages: ChatMessage[] = []; + + await handleAgUiStreamingResponse(body, { + onData: () => {}, + onMessage: (message) => messages.push(message), + }); + + const message = messages[0]; + assertExists(message); + assertEquals(message.parts, [ + { type: "data-dora.report", data: { overall: "fail" } }, + { type: "text", text: "Done", state: "done" }, + ]); + }); + + it("does not store AG-UI snapshot events as assistant message parts", async () => { + const encoder = new TextEncoder(); + const body = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode([ + "event: StateSnapshot", + 'data: {"snapshot":{"context":"private"}}', + "", + "event: MessagesSnapshot", + 'data: {"messages":[{"id":"u1","role":"user","content":"Question"}]}', + "", + "event: TextMessageStart", + 'data: {"messageId":"msg-1","role":"assistant"}', + "", + "event: TextMessageContent", + 'data: {"messageId":"msg-1","delta":"Done"}', + "", + "event: TextMessageEnd", + 'data: {"messageId":"msg-1"}', + "", + "event: RunFinished", + 'data: {"threadId":"thread-1","runId":"run-1"}', + "", + "", + ].join("\n"))); + controller.close(); + }, + }); + const messages: ChatMessage[] = []; + const dataEvents: unknown[] = []; + + await handleAgUiStreamingResponse(body, { + onData: (data) => dataEvents.push(data), + onMessage: (message) => messages.push(message), + }); + + const message = messages[0]; + assertExists(message); + assertEquals(message.parts, [ + { type: "text", text: "Done", state: "done" }, + ]); + assertEquals(dataEvents, [ + { context: "private" }, + [{ id: "u1", role: "user", content: "Question" }], + ]); + }); }); diff --git a/src/agent/runtime/chat-stream-handler.test.ts b/src/agent/runtime/chat-stream-handler.test.ts index 20148d02cd..5df37d10c0 100644 --- a/src/agent/runtime/chat-stream-handler.test.ts +++ b/src/agent/runtime/chat-stream-handler.test.ts @@ -4,6 +4,31 @@ import { describe, it } from "#veryfront/testing/bdd.ts"; import { createMockResult, createSSECollector } from "./chat-stream-handler.test-helpers.ts"; import { createStreamState, processStream } from "./chat-stream-handler.ts"; +function emptyAsyncIterable() { + return { + [Symbol.asyncIterator]() { + return { + async next() { + return { done: true as const, value: undefined }; + }, + }; + }, + }; +} + +function pendingAsyncIterable() { + return { + [Symbol.asyncIterator]() { + return { + async next(): Promise> { + await new Promise(() => {}); + return { done: true, value: undefined }; + }, + }; + }, + }; +} + describe("chat-stream-handler", () => { describe("createStreamState", () => { it("returns a clean initial state", () => { @@ -181,6 +206,149 @@ describe("chat-stream-handler", () => { ]); }); + it("finalizes streamed local tool input when the provider emits tool-input-end without a tool-call part", async () => { + const { events, controller, encoder } = createSSECollector(); + const state = createStreamState(); + + const result = createMockResult([ + { type: "tool-input-start", id: "tc-local-end", toolName: "retrieveDocumentEvidence" }, + { type: "tool-input-delta", id: "tc-local-end", delta: "{}" }, + { type: "tool-input-end", id: "tc-local-end" }, + { type: "finish", finishReason: "tool-calls", totalUsage: null }, + ]); + + await processStream(result, state, controller, encoder, "t", undefined); + + const toolCall = state.toolCalls.get("tc-local-end"); + assertEquals(toolCall?.id, "tc-local-end"); + assertEquals(toolCall?.name, "retrieveDocumentEvidence"); + assertEquals(toolCall?.arguments, "{}"); + assertEquals(toolCall?.inputAvailable, true); + assertEquals(events, [ + { + type: "tool-input-start", + toolCallId: "tc-local-end", + toolName: "retrieveDocumentEvidence", + }, + { + type: "tool-input-delta", + toolCallId: "tc-local-end", + inputTextDelta: "{}", + }, + { + type: "tool-input-available", + toolCallId: "tc-local-end", + toolName: "retrieveDocumentEvidence", + input: {}, + }, + ]); + }); + + it("finalizes parseable streamed local tool input when the provider finishes without a final tool-call part", async () => { + const { events, controller, encoder } = createSSECollector(); + const state = createStreamState(); + + const result = createMockResult([ + { type: "tool-input-start", id: "tc-finish", toolName: "retrieveDocumentEvidence" }, + { type: "tool-input-delta", id: "tc-finish", delta: '{"uploadId":"upload-1"}' }, + { type: "finish", finishReason: "tool-calls", totalUsage: null }, + ]); + + await processStream(result, state, controller, encoder, "t", undefined); + + const toolCall = state.toolCalls.get("tc-finish"); + assertEquals(toolCall?.inputAvailable, true); + assertEquals(events, [ + { + type: "tool-input-start", + toolCallId: "tc-finish", + toolName: "retrieveDocumentEvidence", + }, + { + type: "tool-input-delta", + toolCallId: "tc-finish", + inputTextDelta: '{"uploadId":"upload-1"}', + }, + { + type: "tool-input-available", + toolCallId: "tc-finish", + toolName: "retrieveDocumentEvidence", + input: { uploadId: "upload-1" }, + }, + ]); + }); + + it("does not emit duplicate input-available when tool-input-end is followed by tool-call", async () => { + const { events, controller, encoder } = createSSECollector(); + const state = createStreamState(); + + const result = createMockResult([ + { type: "tool-input-start", id: "tc-end-plus-call", toolName: "lookup" }, + { type: "tool-input-delta", id: "tc-end-plus-call", delta: '{"query":"DORA"}' }, + { type: "tool-input-end", id: "tc-end-plus-call" }, + { + type: "tool-call", + toolCallId: "tc-end-plus-call", + toolName: "lookup", + input: { query: "DORA" }, + }, + { type: "finish", finishReason: "tool-calls", totalUsage: null }, + ]); + + await processStream(result, state, controller, encoder, "t", undefined); + + assertEquals( + events.filter((event) => event.type === "tool-input-available"), + [ + { + type: "tool-input-available", + toolCallId: "tc-end-plus-call", + toolName: "lookup", + input: { query: "DORA" }, + }, + ], + ); + assertEquals(state.toolCalls.get("tc-end-plus-call")?.arguments, '{"query":"DORA"}'); + }); + + it("treats provider tool-input-available as the committed local tool call", async () => { + const { events, controller, encoder } = createSSECollector(); + const state = createStreamState(); + + const result = createMockResult([ + { + type: "tool-input-start", + id: "tc-provider-available", + toolName: "retrieveDocumentEvidence", + }, + { type: "tool-input-delta", id: "tc-provider-available", delta: "{}" }, + { + type: "tool-input-available", + toolCallId: "tc-provider-available", + toolName: "retrieveDocumentEvidence", + input: { uploadId: "upload-1" }, + }, + { type: "finish", finishReason: "tool-calls", totalUsage: null }, + ]); + + await processStream(result, state, controller, encoder, "t", undefined); + + const toolCall = state.toolCalls.get("tc-provider-available"); + assertEquals(toolCall?.inputAvailable, true); + assertEquals(toolCall?.arguments, '{"uploadId":"upload-1"}'); + assertEquals( + events.filter((event) => event.type === "tool-input-available"), + [ + { + type: "tool-input-available", + toolCallId: "tc-provider-available", + toolName: "retrieveDocumentEvidence", + input: { uploadId: "upload-1" }, + }, + ], + ); + }); + it("does not wait for provider stream cancellation after a committed local tool-call", async () => { const { controller, encoder } = createSSECollector(); const state = createStreamState(); @@ -232,6 +400,135 @@ describe("chat-stream-handler", () => { assertEquals(state.toolCalls.get("tc-local")?.inputAvailable, true); }); + it("allows a second local tool input to finish after a prior local tool was committed", async () => { + const { events, controller, encoder } = createSSECollector(); + const state = createStreamState(); + const result = { + fullStream: { + async *[Symbol.asyncIterator]() { + yield { type: "tool-input-start", id: "tc-a", toolName: "load-skill" }; + yield { type: "tool-input-delta", id: "tc-a", delta: '{"skillId":"dora"}' }; + yield { type: "tool-input-end", id: "tc-a" }; + yield { + type: "tool-input-start", + id: "tc-b", + toolName: "load-skill-reference", + }; + await new Promise((resolve) => setTimeout(resolve, 300)); + yield { + type: "tool-input-delta", + id: "tc-b", + delta: '{"skillId":"dora","reference":"references/article-17.md"}', + }; + yield { type: "tool-input-end", id: "tc-b" }; + yield { type: "finish", finishReason: "tool-calls", totalUsage: null }; + }, + }, + textStream: emptyAsyncIterable(), + }; + + await processStream(result, state, controller, encoder, "t", undefined); + + assertEquals(state.toolCalls.get("tc-a")?.inputAvailable, true); + assertEquals(state.toolCalls.get("tc-b")?.inputAvailable, true); + assertEquals( + events.filter((event) => event.type === "tool-input-available"), + [ + { + type: "tool-input-available", + toolCallId: "tc-a", + toolName: "load-skill", + input: { skillId: "dora" }, + }, + { + type: "tool-input-available", + toolCallId: "tc-b", + toolName: "load-skill-reference", + input: { skillId: "dora", reference: "references/article-17.md" }, + }, + ], + ); + }); + + it("does not cut off a slow active local tool input before the provider finishes it", async () => { + const { events, controller, encoder } = createSSECollector(); + const state = createStreamState(); + const result = { + fullStream: { + async *[Symbol.asyncIterator]() { + yield { + type: "tool-input-start", + id: "tc-slow", + toolName: "retrieveDocumentEvidence", + }; + yield { + type: "tool-input-delta", + id: "tc-slow", + delta: '{"uploadId":"upload-1",', + }; + await new Promise((resolve) => setTimeout(resolve, 2_100)); + yield { + type: "tool-input-delta", + id: "tc-slow", + delta: '"name":"sample-ict-services-agreement.docx"}', + }; + yield { type: "tool-input-end", id: "tc-slow" }; + yield { type: "finish", finishReason: "tool-calls", totalUsage: null }; + }, + }, + textStream: emptyAsyncIterable(), + }; + + await processStream(result, state, controller, encoder, "t", undefined); + + assertEquals(state.toolCalls.get("tc-slow")?.inputAvailable, true); + assertEquals( + events.filter((event) => event.type === "tool-input-available"), + [ + { + type: "tool-input-available", + toolCallId: "tc-slow", + toolName: "retrieveDocumentEvidence", + input: { + uploadId: "upload-1", + name: "sample-ict-services-agreement.docx", + }, + }, + ], + ); + }); + + it("times out an active local tool input instead of hanging the stream forever", async () => { + const { controller, encoder } = createSSECollector(); + const state = createStreamState(); + const result = { + fullStream: { + async *[Symbol.asyncIterator]() { + yield { type: "tool-input-start", id: "tc-a", toolName: "load-skill" }; + yield { type: "tool-input-delta", id: "tc-a", delta: '{"skillId":"dora"}' }; + yield { type: "tool-input-end", id: "tc-a" }; + yield { + type: "tool-input-start", + id: "tc-b", + toolName: "load-skill-reference", + }; + await new Promise(() => {}); + }, + }, + textStream: emptyAsyncIterable(), + }; + + const startedAt = Date.now(); + await processStream(result, state, controller, encoder, "t", { + localToolInputIdleTimeoutMs: 10, + }); + + assertEquals(state.finishReason, "tool-calls"); + assertEquals(state.toolCalls.get("tc-a")?.inputAvailable, true); + assertEquals(state.toolCalls.get("tc-b")?.inputAvailable, false); + assertEquals(Date.now() - startedAt >= 9, true); + }); + it("calls onChunk callback for each text delta", async () => { const { controller, encoder } = createSSECollector(); const state = createStreamState(); @@ -250,6 +547,43 @@ describe("chat-stream-handler", () => { assertEquals(chunks, ["a", "b"]); }); + it("times out an idle stream before any output starts", async () => { + const { events, controller, encoder } = createSSECollector(); + const state = createStreamState(); + const result = { + fullStream: pendingAsyncIterable(), + textStream: emptyAsyncIterable(), + }; + + await processStream(result, state, controller, encoder, "t", { + streamIdleTimeoutMs: 10, + }); + + assertEquals(state.finishReason, "stop"); + assertEquals(events, []); + }); + + it("times out an idle output stream after assistant output starts", async () => { + const { controller, encoder } = createSSECollector(); + const state = createStreamState(); + const result = { + fullStream: { + async *[Symbol.asyncIterator]() { + yield { type: "text-delta", text: "Ready." }; + await new Promise(() => {}); + }, + }, + textStream: emptyAsyncIterable(), + }; + + await processStream(result, state, controller, encoder, "t", { + streamIdleTimeoutMs: 10, + }); + + assertEquals(state.accumulatedText, "Ready."); + assertEquals(state.finishReason, "stop"); + }); + it("captures finish reason and usage", async () => { const { controller, encoder } = createSSECollector(); const state = createStreamState(); @@ -301,19 +635,30 @@ describe("chat-stream-handler", () => { throw new Error("error reading a body from connection"); }, }, - textStream: { - async *[Symbol.asyncIterator]() {}, - }, + textStream: emptyAsyncIterable(), }; await processStream(result, state, controller, encoder, "t", undefined); assertEquals(state.finishReason, "tool-calls"); assertEquals(state.toolCalls.size, 1); - assertEquals(events, []); + assertEquals(events, [ + { type: "tool-input-start", toolCallId: "tc-1", toolName: "gmail__get_email" }, + { + type: "tool-input-delta", + toolCallId: "tc-1", + inputTextDelta: '{"id":"msg-1"}', + }, + { + type: "tool-input-available", + toolCallId: "tc-1", + toolName: "gmail__get_email", + input: { id: "msg-1" }, + }, + ]); }); - it("buffers provisional tool-input-start and tool-input-delta until the tool call is committed", async () => { + it("commits buffered tool-input-start and tool-input-delta when the tool call finishes", async () => { const { events, controller, encoder } = createSSECollector(); const state = createStreamState(); @@ -330,8 +675,18 @@ describe("chat-stream-handler", () => { const tc = state.toolCalls.get("tc-1")!; assertEquals(tc.name, "search"); assertEquals(tc.arguments, '{"query":"test"}'); - assertEquals(tc.inputAvailable, false); - assertEquals(events, []); + assertEquals(tc.inputAvailable, true); + assertEquals(events, [ + { type: "tool-input-start", toolCallId: "tc-1", toolName: "search" }, + { type: "tool-input-delta", toolCallId: "tc-1", inputTextDelta: '{"query":' }, + { type: "tool-input-delta", toolCallId: "tc-1", inputTextDelta: '"test"}' }, + { + type: "tool-input-available", + toolCallId: "tc-1", + toolName: "search", + input: { query: "test" }, + }, + ]); }); it("replaces a transient empty-object placeholder when real streamed tool JSON begins", async () => { @@ -350,7 +705,18 @@ describe("chat-stream-handler", () => { const tc = state.toolCalls.get("tc-placeholder")!; assertEquals(tc.arguments, '{"skillId":"plan"}'); - assertEquals(events, []); + assertEquals(events, [ + { type: "tool-input-start", toolCallId: "tc-placeholder", toolName: "load_skill" }, + { type: "tool-input-delta", toolCallId: "tc-placeholder", inputTextDelta: "{}" }, + { type: "tool-input-delta", toolCallId: "tc-placeholder", inputTextDelta: '{"skillId":"' }, + { type: "tool-input-delta", toolCallId: "tc-placeholder", inputTextDelta: 'plan"}' }, + { + type: "tool-input-available", + toolCallId: "tc-placeholder", + toolName: "load_skill", + input: { skillId: "plan" }, + }, + ]); }); it("dedupes cumulative streamed tool argument buffers instead of corrupting the JSON payload", async () => { @@ -379,7 +745,28 @@ describe("chat-stream-handler", () => { tc.arguments, '{"path":"plans/report.md","content":"# Report\\n\\nExecutive summary"}', ); - assertEquals(events, []); + assertEquals(events, [ + { type: "tool-input-start", toolCallId: "tc-cumulative", toolName: "create_file" }, + { + type: "tool-input-delta", + toolCallId: "tc-cumulative", + inputTextDelta: '{"path":"plans/report.md","content":"# Report', + }, + { + type: "tool-input-delta", + toolCallId: "tc-cumulative", + inputTextDelta: '{"path":"plans/report.md","content":"# Report\\n\\nExecutive summary"}', + }, + { + type: "tool-input-available", + toolCallId: "tc-cumulative", + toolName: "create_file", + input: { + path: "plans/report.md", + content: "# Report\n\nExecutive summary", + }, + }, + ]); }); it("dedupes repeated placeholder-style cumulative tool deltas without swallowing parse errors", async () => { @@ -413,7 +800,29 @@ describe("chat-stream-handler", () => { tc.arguments, '{"path":"plans/report.md","content":"# Report\\n\\nExecutive summary"}', ); - assertEquals(events, []); + assertEquals(events, [ + { type: "tool-input-start", toolCallId: "tc-repeat-placeholder", toolName: "create_file" }, + { type: "tool-input-delta", toolCallId: "tc-repeat-placeholder", inputTextDelta: "{}" }, + { + type: "tool-input-delta", + toolCallId: "tc-repeat-placeholder", + inputTextDelta: '"path":"plans/report.md","content":"# Report', + }, + { + type: "tool-input-delta", + toolCallId: "tc-repeat-placeholder", + inputTextDelta: '"path":"plans/report.md","content":"# Report\\n\\nExecutive summary"}', + }, + { + type: "tool-input-available", + toolCallId: "tc-repeat-placeholder", + toolName: "create_file", + input: { + path: "plans/report.md", + content: "# Report\n\nExecutive summary", + }, + }, + ]); }); it("handles tool-call with full input object", async () => { diff --git a/src/agent/runtime/chat-stream-handler.ts b/src/agent/runtime/chat-stream-handler.ts index d03400e4bb..957d529152 100644 --- a/src/agent/runtime/chat-stream-handler.ts +++ b/src/agent/runtime/chat-stream-handler.ts @@ -14,6 +14,7 @@ import { mergeToolCallInput, mergeToolInputDelta, parseToolInputObject, + stripLeadingEmptyObjectPlaceholder, } from "../streaming/data-stream.ts"; import { isDynamicTool } from "./tool-helpers.ts"; import { serverLogger } from "#veryfront/utils"; @@ -24,6 +25,9 @@ import { stringifyToolError, throwIfAborted } from "./error-utils.ts"; const logger = serverLogger.component("agent"); const LOCAL_TOOL_COMMIT_GRACE_MS = 250; +const LOCAL_TOOL_INPUT_IDLE_MS = 15_000; +const STREAM_START_IDLE_MS = 60_000; +const STREAM_OUTPUT_IDLE_MS = 15_000; export interface StreamingToolCall { id: string; @@ -61,6 +65,8 @@ export interface ChatStreamCallbacks { completionTokens?: number; totalTokens?: number; }) => void; + localToolInputIdleTimeoutMs?: number; + streamIdleTimeoutMs?: number; } function isRecord(value: unknown): value is Record { @@ -75,6 +81,15 @@ function normalizeToolInputString(input: unknown): string { return JSON.stringify(input ?? null) ?? "null"; } +function tryParseToolInputObject(input: string): Record | null { + try { + const parsed = JSON.parse(stripLeadingEmptyObjectPlaceholder(input)); + return isRecord(parsed) ? parsed : null; + } catch { + return null; + } +} + function summarizeDebugValue(value: unknown): unknown { if (value instanceof Error) { return { @@ -270,6 +285,7 @@ export function processStream( let textOpen = false; let activeReasoningId: string | null = null; let shouldStopForCommittedLocalToolCall = false; + let hasActiveLocalToolInput = false; const normalizeReasoningId = (part: { id?: string }) => typeof part.id === "string" && part.id.length > 0 ? part.id : "reasoning"; @@ -329,6 +345,43 @@ export function processStream( activeReasoningId = null; }; + const commitParseablePendingToolInputs = () => { + for (const tc of state.toolCalls.values()) { + if (tc.inputAvailable === true || tc.providerExecuted === true) { + continue; + } + // A bare empty-object placeholder (`""` or `"{}"` after stripping + // transient prefixes) is provisional streamed input that never + // finalized into a real `tool-call`/`tool-input-end`. Committing it + // would mark `inputAvailable: true` and execute the tool with empty + // args. Leave it provisional so the runtime can recover by re-calling + // the model instead of executing a placeholder. + const stripped = stripLeadingEmptyObjectPlaceholder(tc.arguments); + if (stripped === "" || stripped === "{}") { + continue; + } + const parsedInput = tryParseToolInputObject(tc.arguments); + if (!parsedInput) { + continue; + } + tc.inputAvailable = true; + const dynamic = tc.dynamic ?? isDynamicTool(tc.name); + if (dynamic) { + tc.dynamic = true; + } + announceToolInputStart(tc); + sendSSE(controller, encoder, { + type: "tool-input-available", + toolCallId: tc.id, + toolName: tc.name, + input: parsedInput, + ...(tc.providerExecuted !== undefined ? { providerExecuted: tc.providerExecuted } : {}), + ...(dynamic ? { dynamic: true } : {}), + }); + shouldStopForCommittedLocalToolCall = true; + } + }; + const announceToolInputStart = (toolCall: StreamingToolCall) => { if (toolCall.inputAnnounced === true) { return; @@ -429,15 +482,39 @@ export function processStream( const streamIterator = result.fullStream[Symbol.asyncIterator](); while (true) { - const next = shouldStopForCommittedLocalToolCall + const shouldStopForIdleOutput = !hasActiveLocalToolInput && + !shouldStopForCommittedLocalToolCall && hasStreamOutput(state); + const shouldStopForIdleStart = !hasActiveLocalToolInput && + !shouldStopForCommittedLocalToolCall && !hasStreamOutput(state); + const next = hasActiveLocalToolInput + ? await readNextStreamPartWithTimeout( + streamIterator, + state, + callbacks?.localToolInputIdleTimeoutMs ?? LOCAL_TOOL_INPUT_IDLE_MS, + ) + : shouldStopForCommittedLocalToolCall ? await readNextStreamPartWithTimeout( streamIterator, state, LOCAL_TOOL_COMMIT_GRACE_MS, ) + : shouldStopForIdleOutput + ? await readNextStreamPartWithTimeout( + streamIterator, + state, + callbacks?.streamIdleTimeoutMs ?? STREAM_OUTPUT_IDLE_MS, + ) + : shouldStopForIdleStart + ? await readNextStreamPartWithTimeout( + streamIterator, + state, + callbacks?.streamIdleTimeoutMs ?? STREAM_START_IDLE_MS, + ) : await readNextStreamPart(streamIterator, state); if (next === "timeout") { - state.finishReason ??= "tool-calls"; + state.finishReason ??= shouldStopForIdleOutput || shouldStopForIdleStart + ? "stop" + : "tool-calls"; requestStreamIteratorReturn(streamIterator); break; } @@ -506,6 +583,8 @@ export function processStream( case "tool-input-start": { closeTextSegment(); closeReasoningSegment(); + shouldStopForCommittedLocalToolCall = false; + hasActiveLocalToolInput = true; const toolId = typedPart.id; state.toolCalls.set(toolId, { id: toolId, @@ -532,15 +611,86 @@ export function processStream( break; } + case "tool-input-end": { + closeTextSegment(); + closeReasoningSegment(); + const toolId = typedPart.id; + const tc = state.toolCalls.get(toolId); + if (!tc) break; + + tc.inputAvailable = true; + hasActiveLocalToolInput = false; + const dynamic = tc.dynamic ?? isDynamicTool(tc.name); + if (dynamic) { + tc.dynamic = true; + } + announceToolInputStart(tc); + sendSSE(controller, encoder, { + type: "tool-input-available", + toolCallId: toolId, + toolName: tc.name, + input: parseToolInputObject(tc.arguments), + ...(tc.providerExecuted !== undefined ? { providerExecuted: tc.providerExecuted } : {}), + ...(dynamic ? { dynamic: true } : {}), + }); + if (tc.providerExecuted !== true) { + shouldStopForCommittedLocalToolCall = true; + } + break; + } + + case "tool-input-available": { + closeTextSegment(); + closeReasoningSegment(); + const toolId = typedPart.toolCallId ?? typedPart.id; + if (!toolId) { + break; + } + hasActiveLocalToolInput = false; + const inputStr = normalizeToolInputString(typedPart.input); + const previous = state.toolCalls.get(toolId); + const previousArguments = previous?.arguments ?? ""; + const resolvedArguments = mergeToolCallInput(previousArguments, inputStr); + const wasInputAvailable = previous?.inputAvailable === true; + const dynamic = typedPart.dynamic ?? isDynamicTool(typedPart.toolName); + state.toolCalls.set(toolId, { + id: toolId, + name: typedPart.toolName, + arguments: resolvedArguments, + inputAvailable: true, + providerExecuted: typedPart.providerExecuted, + dynamic, + }); + + if (!wasInputAvailable) { + sendSSE(controller, encoder, { + type: "tool-input-available", + toolCallId: toolId, + toolName: typedPart.toolName, + input: parseToolInputObject(resolvedArguments), + ...(typedPart.providerExecuted !== undefined + ? { providerExecuted: typedPart.providerExecuted } + : {}), + ...(dynamic ? { dynamic: true } : {}), + }); + } + if (typedPart.providerExecuted !== true) { + shouldStopForCommittedLocalToolCall = true; + } + break; + } + case "tool-call": { closeTextSegment(); closeReasoningSegment(); // tool-call fires when the full tool call is available const toolId = typedPart.toolCallId; + hasActiveLocalToolInput = false; const inputStr = normalizeToolInputString(typedPart.input); const previous = state.toolCalls.get(toolId); const previousArguments = previous?.arguments ?? ""; const resolvedArguments = mergeToolCallInput(previousArguments, inputStr); + const wasInputAvailable = previous?.inputAvailable === true; const toolCall: StreamingToolCall = { id: toolId, name: typedPart.toolName, @@ -556,16 +706,18 @@ export function processStream( const dynamic = isDynamicTool(typedPart.toolName); const inputObj = parseToolInputObject(typedPart.input); announceToolInputStart(toolCall); - sendSSE(controller, encoder, { - type: "tool-input-available", - toolCallId: toolId, - toolName: typedPart.toolName, - input: inputObj, - ...(typedPart.providerExecuted !== undefined - ? { providerExecuted: typedPart.providerExecuted } - : {}), - ...(dynamic ? { dynamic: true } : {}), - }); + if (!wasInputAvailable) { + sendSSE(controller, encoder, { + type: "tool-input-available", + toolCallId: toolId, + toolName: typedPart.toolName, + input: inputObj, + ...(typedPart.providerExecuted !== undefined + ? { providerExecuted: typedPart.providerExecuted } + : {}), + ...(dynamic ? { dynamic: true } : {}), + }); + } if (typedPart.providerExecuted !== true) { shouldStopForCommittedLocalToolCall = true; } @@ -681,6 +833,9 @@ export function processStream( closeTextSegment(); closeReasoningSegment(); state.finishReason = typedPart.finishReason ?? null; + if (state.finishReason === "tool-calls") { + commitParseablePendingToolInputs(); + } if (typedPart.totalUsage) { const input = typedPart.totalUsage.inputTokens ?? 0; const output = typedPart.totalUsage.outputTokens ?? 0; diff --git a/src/agent/runtime/index.test.ts b/src/agent/runtime/index.test.ts index c8ce8a22ac..02fca628cd 100644 --- a/src/agent/runtime/index.test.ts +++ b/src/agent/runtime/index.test.ts @@ -6,6 +6,7 @@ import { collectFinalStreamToolResults, collectGeneratedToolResults, collectPersistedToolResults, + isRecoverablePlaceholderToolCall, isStreamedToolCallIncomplete, materializeStreamedToolCall, shouldContinueAfterStreamStep, @@ -134,6 +135,57 @@ describe("agent runtime streamed tool result collection", () => { assertEquals(shouldContinue, true); }); + it("stops after an unfinalized streamed tool call to avoid retry loops", () => { + const shouldContinue = shouldContinueAfterStreamStep({ + accumulatedText: "", + finishReason: "tool-calls", + toolCalls: new Map([ + [ + "toolu_incomplete_1", + { + id: "toolu_incomplete_1", + name: "load-skill-reference", + arguments: '{"skillId":"dora"', + inputAvailable: false, + }, + ], + ]), + toolResults: [], + }); + + assertEquals(shouldContinue, false); + }); + + it("stops instead of retrying when a step has both finalized and unfinalized tool calls", () => { + const shouldContinue = shouldContinueAfterStreamStep({ + accumulatedText: "", + finishReason: "tool-calls", + toolCalls: new Map([ + [ + "toolu_complete_1", + { + id: "toolu_complete_1", + name: "load-skill", + arguments: '{"skillId":"dora"}', + inputAvailable: true, + }, + ], + [ + "toolu_incomplete_1", + { + id: "toolu_incomplete_1", + name: "load-skill-reference", + arguments: '{"skillId":"dora"', + inputAvailable: false, + }, + ], + ]), + toolResults: [], + }); + + assertEquals(shouldContinue, false); + }); + it("continues finalized client-executed tool calls when the provider reports stop", () => { const shouldContinue = shouldContinueAfterStreamStep({ accumulatedText: "", @@ -298,6 +350,59 @@ describe("agent runtime streamed tool result collection", () => { ); }); + it("classifies a non-finalized empty-object placeholder as recoverable", () => { + assertEquals( + isRecoverablePlaceholderToolCall({ inputAvailable: false, arguments: "{}" }), + true, + ); + assertEquals( + isRecoverablePlaceholderToolCall({ inputAvailable: false, arguments: "" }), + true, + ); + assertEquals( + isRecoverablePlaceholderToolCall({ inputAvailable: undefined, arguments: "{}{}" }), + true, + ); + }); + + it("does not classify truncated partial JSON as a recoverable placeholder", () => { + assertEquals( + isRecoverablePlaceholderToolCall({ + inputAvailable: false, + arguments: '{"skillId":"dora"', + }), + false, + ); + }); + + it("does not classify a finalized tool call as a recoverable placeholder", () => { + assertEquals( + isRecoverablePlaceholderToolCall({ inputAvailable: true, arguments: "{}" }), + false, + ); + }); + + it("recovers a provisional empty-object placeholder by continuing the loop", () => { + const shouldContinue = shouldContinueAfterStreamStep({ + accumulatedText: "", + finishReason: "tool-calls", + toolCalls: new Map([ + [ + "toolu_placeholder_1", + { + id: "toolu_placeholder_1", + name: "review", + arguments: "{}", + inputAvailable: false, + }, + ], + ]), + toolResults: [], + }); + + assertEquals(shouldContinue, true); + }); + it("materializes a complete streamed tool call into a ready-to-execute part", () => { const materialized = materializeStreamedToolCall({ id: "toolu_complete", diff --git a/src/agent/runtime/index.ts b/src/agent/runtime/index.ts index 1a04a098b6..2c273c312b 100644 --- a/src/agent/runtime/index.ts +++ b/src/agent/runtime/index.ts @@ -45,6 +45,7 @@ import { type StreamingToolResult, } from "./chat-stream-handler.ts"; import { repairToolCall } from "./repair-tool-call.ts"; +import { stripLeadingEmptyObjectPlaceholder } from "../streaming/data-stream.ts"; import { MiddlewareChain } from "../middleware/chain.ts"; import { AGENT_DEFAULTS } from "./defaults.ts"; import { tryGetCacheKeyContext } from "#veryfront/cache/cache-key-builder.ts"; @@ -254,8 +255,34 @@ export function shouldContinueAfterStreamStep( return false; } + const streamedToolCalls = Array.from(state.toolCalls.values()); + const hasIncompleteToolCall = streamedToolCalls.some(isStreamedToolCallIncomplete); + const hasFinalizedClientToolCall = streamedToolCalls.some((toolCall) => + toolCall.inputAvailable === true && toolCall.providerExecuted !== true + ); + // A non-finalized call whose only accumulated arguments are a bare + // empty-object placeholder is provisional streamed input the model never + // committed. We can recover by re-calling the model, so it must not block + // continuation (unlike a truncated/dead partial-JSON call). + const hasIncompleteDeadToolCall = streamedToolCalls.some( + (toolCall) => + isStreamedToolCallIncomplete(toolCall) && + !isRecoverablePlaceholderToolCall(toolCall), + ); + const hasRecoverablePlaceholderToolCall = streamedToolCalls.some( + isRecoverablePlaceholderToolCall, + ); + if (state.finishReason === "tool-calls") { - return true; + if (hasIncompleteDeadToolCall) { + return false; + } + // Recover provisional placeholders by re-calling the model even when no + // client tool call finalized in this step. + if (hasRecoverablePlaceholderToolCall && !hasFinalizedClientToolCall) { + return true; + } + return !hasIncompleteToolCall && hasFinalizedClientToolCall; } if (state.finishReason !== "stop") { @@ -322,6 +349,30 @@ export function isStreamedToolCallIncomplete( return toolCall.inputAvailable !== true; } +/** + * A non-finalized streamed tool call is a "recoverable placeholder" when its + * accumulated `arguments` are empty or only the transient empty-object + * placeholder `"{}"` (after stripping leading placeholders). This happens when + * a provider emits `tool-input-start` + a `"{}"` `tool-input-delta` and then + * finishes the step WITHOUT ever sending the finalizing `tool-call` / + * `tool-input-end` event — the model never actually committed any arguments. + * + * Unlike an incomplete-dead tool call (which carries real truncated partial + * JSON and must stop the loop to avoid retry storms), a recoverable + * placeholder carries no committed intent, so the runtime can safely re-call + * the model to recover the real tool call. Such placeholders must NOT be + * executed and must NOT surface a stream-termination error. + */ +export function isRecoverablePlaceholderToolCall( + toolCall: Pick, +): boolean { + if (!isStreamedToolCallIncomplete(toolCall)) { + return false; + } + const stripped = stripLeadingEmptyObjectPlaceholder(toolCall.arguments); + return stripped === "" || stripped === "{}"; +} + /** * Classification of a streamed tool call when we reach end-of-stream and need * to persist it into the assistant message. Three distinct cases, each with @@ -1235,6 +1286,14 @@ export class AgentRuntime { const materialized = materializeStreamedToolCall(tc); streamParts.push(materialized.part); + if (materialized.kind === "incomplete" && isRecoverablePlaceholderToolCall(tc)) { + // Provisional empty-object placeholder that never finalized. The + // model never committed arguments; the loop will recover by + // re-calling the model. Persist the (empty) part for transparent + // history but surface no termination warning or error. + continue; + } + if (materialized.kind === "incomplete") { // Stream terminated before the provider emitted the finalizing // `tool-call` event for this block. The model never committed this @@ -1324,6 +1383,13 @@ export class AgentRuntime { for (const tc of streamedToolCalls) { throwIfAborted(abortSignal); + if (isRecoverablePlaceholderToolCall(tc)) { + // Provisional empty-object placeholder that never finalized. The + // model never committed arguments, so we neither execute it nor + // surface a stream-termination error — the loop continues and the + // next model call recovers the real tool call. + continue; + } if (isStreamedToolCallIncomplete(tc)) { // Stream ended before the provider finalized this tool call. We // cannot execute it — record a distinct stream-termination error diff --git a/src/agent/runtime/runtime-tool-types.ts b/src/agent/runtime/runtime-tool-types.ts index 0822fbce4d..8617f36586 100644 --- a/src/agent/runtime/runtime-tool-types.ts +++ b/src/agent/runtime/runtime-tool-types.ts @@ -74,6 +74,16 @@ export type RuntimeStreamPart = dynamic?: boolean; } | { type: "tool-input-delta"; id: string; delta: string } + | { type: "tool-input-end"; id: string } + | { + type: "tool-input-available"; + toolCallId?: string; + id?: string; + toolName: string; + input: unknown; + providerExecuted?: boolean; + dynamic?: boolean; + } | { type: "tool-call"; toolCallId: string; diff --git a/src/agent/streaming/tool-execution-data-event-bridge.test.ts b/src/agent/streaming/tool-execution-data-event-bridge.test.ts index 6ae3ebe45c..1173e433b6 100644 --- a/src/agent/streaming/tool-execution-data-event-bridge.test.ts +++ b/src/agent/streaming/tool-execution-data-event-bridge.test.ts @@ -4,6 +4,13 @@ import { describe, it } from "#veryfront/testing/bdd.ts"; import type { ToolExecutionDataEvent } from "../../tool/types.ts"; import { createToolExecutionDataEventBridgeStream } from "./tool-execution-data-event-bridge.ts"; +function requireStreamController( + controller: ReadableStreamDefaultController | null, +): ReadableStreamDefaultController { + if (!controller) throw new Error("Expected base stream controller"); + return controller; +} + describe("createToolExecutionDataEventBridgeStream", () => { it("emits published tool data events before forwarding upstream data stream chunks", async () => { const encoder = new TextEncoder(); @@ -38,8 +45,9 @@ describe("createToolExecutionDataEventBridgeStream", () => { }\n\n`, ); - baseController?.enqueue(encoder.encode('data: {"type":"message-finish"}\n\n')); - baseController?.close(); + const controller = requireStreamController(baseController); + controller.enqueue(encoder.encode('data: {"type":"message-finish"}\n\n')); + controller.close(); const forwardedChunk = await reader.read(); assertEquals(forwardedChunk.done, false); @@ -47,4 +55,42 @@ describe("createToolExecutionDataEventBridgeStream", () => { assertEquals(await reader.read(), { done: true, value: undefined }); }); + + it("emits named tool data events as data parts", async () => { + const encoder = new TextEncoder(); + const decoder = new TextDecoder(); + let publishDataEvent = (_event: ToolExecutionDataEvent) => {}; + let baseController: ReadableStreamDefaultController | null = null; + + const stream = createToolExecutionDataEventBridgeStream({ + baseStream: new ReadableStream({ + start(controller) { + baseController = controller; + }, + }), + installPublisher(nextPublishDataEvent) { + publishDataEvent = nextPublishDataEvent; + }, + }); + + const reader = stream.getReader(); + + publishDataEvent({ + type: "dora.report", + name: "dora.report", + value: { status: "ready" }, + }); + + const eventChunk = await reader.read(); + assertEquals(eventChunk.done, false); + assertEquals( + decoder.decode(eventChunk.value), + `data: ${JSON.stringify({ type: "data-dora.report", data: { status: "ready" } })}\n\n`, + ); + + const controller = requireStreamController(baseController); + controller.enqueue(encoder.encode('data: {"type":"message-finish"}\n\n')); + controller.close(); + await reader.cancel(); + }); }); diff --git a/src/agent/streaming/tool-execution-data-event-bridge.ts b/src/agent/streaming/tool-execution-data-event-bridge.ts index 6fc79c461a..7e6c2dd9d5 100644 --- a/src/agent/streaming/tool-execution-data-event-bridge.ts +++ b/src/agent/streaming/tool-execution-data-event-bridge.ts @@ -10,6 +10,13 @@ export type ToolExecutionDataEventBridgeStreamInput = { }; function serializeToolExecutionDataEvent(event: ToolExecutionDataEvent): Uint8Array { + if (typeof event.name === "string" && event.name.length > 0) { + const data = Object.hasOwn(event, "value") ? event.value : event.data; + return new TextEncoder().encode( + `data: ${JSON.stringify({ type: `data-${event.name}`, data })}\n\n`, + ); + } + return new TextEncoder().encode(`data: ${JSON.stringify({ type: "data", data: event })}\n\n`); } diff --git a/src/chat/protocol.ts b/src/chat/protocol.ts index 4bebe3851f..7e144587a3 100644 --- a/src/chat/protocol.ts +++ b/src/chat/protocol.ts @@ -66,6 +66,12 @@ export interface ChatStepPart { stepIndex: number; } +/** Public API contract for chat data part. */ +export interface ChatDataPart { + type: `data-${string}`; + data: unknown; +} + /** Public API contract for chat message part. */ export type ChatMessagePart = | ChatTextPart @@ -73,7 +79,8 @@ export type ChatMessagePart = | ChatToolPart | ChatToolResultPart | ChatDynamicToolPart - | ChatStepPart; + | ChatStepPart + | ChatDataPart; /** Message shape for chat. */ export interface ChatMessage {