Preserve tool data events in AG-UI streams

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<Uint8Array>({
+        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<ReadableStream<Uint8Array>> {
+      const publishDataEvent = context?.publishDataEvent;
+      if (typeof publishDataEvent === "function") {
+        await publishDataEvent({
+          type: "test.report",
+          name: "test.report",
+          value: { status: "ready" },
+        });
+      }
+
+      return new ReadableStream<Uint8Array>({
+        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({

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<Uint8Array>): ReadableStream<Uint8Array> {
+      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<string, unknown>,
@@ -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<Uint8Array>;
+  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;

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

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<string, OrderedReasoning>;
   steps: Map<number, OrderedStep>;
   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<string, unknown>, state: StreamingState): vo
   state.toolCalls.clear();
   state.reasoningBlocks.clear();
   state.messageParts.length = 0;
+  state.dataParts.length = 0;
 }
 
 function handleTextStart(parsed: Record<string, unknown>, state: StreamingState): void {
@@ -351,6 +378,37 @@ function handleTextEnd(parsed: Record<string, unknown>, 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<string, unknown>,
   state: StreamingState,