assistant-stream 0.3.13 → 0.3.14

package/src/core/tool/tool-types.ts

@@ -4,6 +4,26 @@ import type { AsyncIterableStream } from "../../utils";
 import type { StandardSchemaV1 } from "@standard-schema/spec";
 import type { ToolResponse } from "./ToolResponse";
 
+export type ToolModelContentPart =
+  | {
+      readonly type: "text";
+      readonly text: string;
+    }
+  | {
+      readonly type: "file";
+      readonly data: string;
+      readonly mediaType: string;
+      readonly filename?: string;
+    };
+
+export type ToolModelOutputFunction<TArgs, TResult> = (options: {
+  toolCallId: string;
+  input: TArgs;
+  output: TResult;
+}) =>
+  | readonly ToolModelContentPart[]
+  | Promise<readonly ToolModelContentPart[]>;
+
 /**
  * Interface for reading tool call arguments from a stream, which are
  * generated by a language learning model (LLM). Provides methods to
@@ -121,6 +141,7 @@ type BackendTool<
   parameters?: undefined;
   disabled?: undefined;
   execute?: undefined;
+  toModelOutput?: undefined;
   experimental_onSchemaValidationError?: undefined;
 };
 
@@ -134,6 +155,7 @@ type FrontendTool<
   parameters: StandardSchemaV1<TArgs> | JSONSchema7;
   disabled?: boolean;
   execute: ToolExecuteFunction<TArgs, TResult>;
+  toModelOutput?: ToolModelOutputFunction<TArgs, TResult>;
   experimental_onSchemaValidationError?: OnSchemaValidationErrorFunction<TResult>;
 };
 
@@ -147,6 +169,7 @@ type HumanTool<
   parameters: StandardSchemaV1<TArgs> | JSONSchema7;
   disabled?: boolean;
   execute?: undefined;
+  toModelOutput?: undefined;
   experimental_onSchemaValidationError?: undefined;
 };
 

package/src/core/tool/toolResultStream.test.ts

@@ -1,5 +1,10 @@
-import { describe, expect, it } from "vitest";
-import { unstable_runPendingTools } from "./toolResultStream";
+import { describe, expect, it, vi } from "vitest";
+import {
+  toolResultStream as unstable_toolResultStream,
+  unstable_runPendingTools,
+} from "./toolResultStream";
+import { ToolResponse } from "./ToolResponse";
+import type { AssistantStreamChunk } from "../AssistantStreamChunk";
 import type { AssistantMessage, ToolCallPart } from "../utils/types";
 import type { Tool } from "./tool-types";
 
@@ -397,4 +402,357 @@ describe("unstable_runPendingTools", () => {
       });
     });
   });
+
+  describe("toModelOutput", () => {
+    it("attaches modelContent from toModelOutput onto the resolved tool-call part", async () => {
+      const tool: Tool = {
+        parameters: { type: "object", properties: {} },
+        execute: async () => ({
+          mediaType: "application/pdf",
+          base64: "JVBERi0xLjQK",
+        }),
+        toModelOutput: ({ output }) => {
+          const o = output as { base64: string; mediaType: string };
+          return [
+            { type: "text", text: "PDF contents:" },
+            {
+              type: "file",
+              data: o.base64,
+              mediaType: o.mediaType,
+            },
+          ];
+        },
+      };
+
+      const message: AssistantMessage = {
+        role: "assistant",
+        status: { type: "requires-action", reason: "tool-calls" },
+        parts: [
+          {
+            type: "tool-call",
+            toolCallId: "tc-1",
+            toolName: "readPdf",
+            args: {},
+          } as ToolCallPart,
+        ],
+        content: [],
+        metadata: {
+          unstable_state: {},
+          unstable_data: [],
+          unstable_annotations: [],
+          steps: [],
+          custom: {},
+        },
+      };
+
+      const updated = await unstable_runPendingTools(
+        message,
+        { readPdf: tool },
+        new AbortController().signal,
+        async () => {},
+      );
+
+      expect(updated.parts[0]).toMatchObject({
+        type: "tool-call",
+        state: "result",
+        result: { mediaType: "application/pdf", base64: "JVBERi0xLjQK" },
+        modelContent: [
+          { type: "text", text: "PDF contents:" },
+          {
+            type: "file",
+            data: "JVBERi0xLjQK",
+            mediaType: "application/pdf",
+          },
+        ],
+      });
+    });
+
+    it("does not call toModelOutput when the ToolResponse already carries modelContent", async () => {
+      let called = false;
+      const tool: Tool = {
+        parameters: { type: "object", properties: {} },
+        execute: async () =>
+          new ToolResponse({
+            result: { ok: true },
+            modelContent: [{ type: "text", text: "preset" }],
+          }),
+        toModelOutput: () => {
+          called = true;
+          return [{ type: "text", text: "should not run" }];
+        },
+      };
+
+      const message: AssistantMessage = {
+        role: "assistant",
+        status: { type: "requires-action", reason: "tool-calls" },
+        parts: [
+          {
+            type: "tool-call",
+            toolCallId: "tc-1",
+            toolName: "preset",
+            args: {},
+          } as ToolCallPart,
+        ],
+        content: [],
+        metadata: {
+          unstable_state: {},
+          unstable_data: [],
+          unstable_annotations: [],
+          steps: [],
+          custom: {},
+        },
+      };
+
+      const updated = await unstable_runPendingTools(
+        message,
+        { preset: tool },
+        new AbortController().signal,
+        async () => {},
+      );
+
+      expect(called).toBe(false);
+      expect(updated.parts[0]).toMatchObject({
+        type: "tool-call",
+        state: "result",
+        modelContent: [{ type: "text", text: "preset" }],
+      });
+    });
+
+    it("falls back to the successful execute result when toModelOutput itself throws", async () => {
+      const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+      const tool: Tool = {
+        parameters: { type: "object", properties: {} },
+        execute: async () => ({ ok: true }),
+        toModelOutput: () => {
+          throw new Error("projection failed");
+        },
+      };
+
+      const message: AssistantMessage = {
+        role: "assistant",
+        status: { type: "requires-action", reason: "tool-calls" },
+        parts: [
+          {
+            type: "tool-call",
+            toolCallId: "tc-1",
+            toolName: "flaky",
+            args: {},
+          } as ToolCallPart,
+        ],
+        content: [],
+        metadata: {
+          unstable_state: {},
+          unstable_data: [],
+          unstable_annotations: [],
+          steps: [],
+          custom: {},
+        },
+      };
+
+      const updated = await unstable_runPendingTools(
+        message,
+        { flaky: tool },
+        new AbortController().signal,
+        async () => {},
+      );
+
+      expect(updated.parts[0]).toMatchObject({
+        type: "tool-call",
+        state: "result",
+        result: { ok: true },
+        isError: false,
+      });
+      expect(updated.parts[0]).not.toHaveProperty("modelContent");
+      expect(warn).toHaveBeenCalledWith(
+        expect.stringContaining(`tool "flaky" toModelOutput threw`),
+        expect.any(Error),
+      );
+      warn.mockRestore();
+    });
+
+    it("forwards modelContent through the streaming path (toolResultStream + ToolExecutionStream)", async () => {
+      const tool: Tool = {
+        parameters: { type: "object", properties: {} },
+        execute: async () => ({
+          mediaType: "application/pdf",
+          base64: "JVBERi0xLjQK",
+        }),
+        toModelOutput: ({ output }) => {
+          const o = output as { mediaType: string; base64: string };
+          return [
+            { type: "text", text: "PDF contents:" },
+            { type: "file", data: o.base64, mediaType: o.mediaType },
+          ];
+        },
+      };
+
+      const inputChunks: AssistantStreamChunk[] = [
+        {
+          type: "part-start",
+          path: [],
+          part: {
+            type: "tool-call",
+            toolCallId: "tc-stream-1",
+            toolName: "readPdf",
+          },
+        },
+        { type: "text-delta", path: [0], textDelta: "{}" },
+        { type: "tool-call-args-text-finish", path: [0] },
+        { type: "part-finish", path: [0] },
+      ];
+
+      const inputStream = new ReadableStream<AssistantStreamChunk>({
+        start(controller) {
+          for (const chunk of inputChunks) controller.enqueue(chunk);
+          controller.close();
+        },
+      });
+
+      const outputChunks: AssistantStreamChunk[] = [];
+      await inputStream
+        .pipeThrough(
+          unstable_toolResultStream(
+            { readPdf: tool },
+            new AbortController().signal,
+            async () => {},
+          ),
+        )
+        .pipeTo(
+          new WritableStream<AssistantStreamChunk>({
+            write(chunk) {
+              outputChunks.push(chunk);
+            },
+          }),
+        );
+
+      const resultChunk = outputChunks.find((c) => c.type === "result") as
+        | (AssistantStreamChunk & { type: "result" })
+        | undefined;
+      expect(resultChunk).toBeDefined();
+      expect(resultChunk?.result).toEqual({
+        mediaType: "application/pdf",
+        base64: "JVBERi0xLjQK",
+      });
+      expect(resultChunk?.isError).toBe(false);
+      expect(resultChunk?.modelContent).toEqual([
+        { type: "text", text: "PDF contents:" },
+        {
+          type: "file",
+          data: "JVBERi0xLjQK",
+          mediaType: "application/pdf",
+        },
+      ]);
+    });
+
+    it("falls back to the plain result when toModelOutput throws in the streaming path", async () => {
+      const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+      const tool: Tool = {
+        parameters: { type: "object", properties: {} },
+        execute: async () => ({ ok: true }),
+        toModelOutput: () => {
+          throw new Error("projection failed");
+        },
+      };
+
+      const inputChunks: AssistantStreamChunk[] = [
+        {
+          type: "part-start",
+          path: [],
+          part: {
+            type: "tool-call",
+            toolCallId: "tc-stream-err",
+            toolName: "flaky",
+          },
+        },
+        { type: "text-delta", path: [0], textDelta: "{}" },
+        { type: "tool-call-args-text-finish", path: [0] },
+        { type: "part-finish", path: [0] },
+      ];
+
+      const inputStream = new ReadableStream<AssistantStreamChunk>({
+        start(controller) {
+          for (const chunk of inputChunks) controller.enqueue(chunk);
+          controller.close();
+        },
+      });
+
+      const outputChunks: AssistantStreamChunk[] = [];
+      await inputStream
+        .pipeThrough(
+          unstable_toolResultStream(
+            { flaky: tool },
+            new AbortController().signal,
+            async () => {},
+          ),
+        )
+        .pipeTo(
+          new WritableStream<AssistantStreamChunk>({
+            write(chunk) {
+              outputChunks.push(chunk);
+            },
+          }),
+        );
+
+      const resultChunk = outputChunks.find((c) => c.type === "result") as
+        | (AssistantStreamChunk & { type: "result" })
+        | undefined;
+      expect(resultChunk).toBeDefined();
+      expect(resultChunk?.result).toEqual({ ok: true });
+      expect(resultChunk?.isError).toBe(false);
+      expect(resultChunk?.modelContent).toBeUndefined();
+      expect(warn).toHaveBeenCalledWith(
+        expect.stringContaining(`tool "flaky" toModelOutput threw`),
+        expect.any(Error),
+      );
+      warn.mockRestore();
+    });
+
+    it("does not call toModelOutput when the tool errors", async () => {
+      let called = false;
+      const tool: Tool = {
+        parameters: { type: "object", properties: {} },
+        execute: async () => {
+          throw new Error("boom");
+        },
+        toModelOutput: () => {
+          called = true;
+          return [{ type: "text", text: "should not run" }];
+        },
+      };
+
+      const message: AssistantMessage = {
+        role: "assistant",
+        status: { type: "requires-action", reason: "tool-calls" },
+        parts: [
+          {
+            type: "tool-call",
+            toolCallId: "tc-1",
+            toolName: "broken",
+            args: {},
+          } as ToolCallPart,
+        ],
+        content: [],
+        metadata: {
+          unstable_state: {},
+          unstable_data: [],
+          unstable_annotations: [],
+          steps: [],
+          custom: {},
+        },
+      };
+
+      try {
+        await unstable_runPendingTools(
+          message,
+          { broken: tool },
+          new AbortController().signal,
+          async () => {},
+        );
+      } catch {
+        // execute throws; toModelOutput must not have been consulted
+      }
+
+      expect(called).toBe(false);
+    });
+  });
 });

package/src/core/tool/toolResultStream.ts

@@ -87,7 +87,33 @@ function getToolResponse(
         abortSignal,
         human: (payload: unknown) => human(toolCall.toolCallId, payload),
       })) as unknown as ReadonlyJSONValue;
-      return ToolResponse.toResponse(result);
+      const response = ToolResponse.toResponse(result);
+      if (
+        tool.toModelOutput &&
+        !response.isError &&
+        response.modelContent === undefined
+      ) {
+        try {
+          const modelContent = await tool.toModelOutput({
+            toolCallId: toolCall.toolCallId,
+            input: toolCall.args,
+            output: response.result,
+          });
+          return new ToolResponse({
+            result: response.result,
+            artifact: response.artifact,
+            isError: response.isError,
+            messages: response.messages,
+            modelContent,
+          });
+        } catch (e) {
+          console.warn(
+            `[assistant-stream] tool "${toolCall.toolName}" toModelOutput threw; falling back to default projection.`,
+            e,
+          );
+        }
+      }
+      return response;
     })();
 
     return Promise.race([executePromise, abortPromise]);
@@ -169,6 +195,9 @@ export async function unstable_runPendingTools(
           ...(toolResponse.artifact !== undefined
             ? { artifact: toolResponse.artifact }
             : {}),
+          ...(toolResponse.modelContent !== undefined
+            ? { modelContent: toolResponse.modelContent }
+            : {}),
           result: toolResponse.result as ReadonlyJSONValue,
           isError: toolResponse.isError,
         };

package/src/core/utils/types.ts

@@ -2,6 +2,7 @@ import type {
   ReadonlyJSONObject,
   ReadonlyJSONValue,
 } from "../../utils/json/json-value";
+import type { ToolModelContentPart } from "../tool/tool-types";
 
 type TextStatus =
   | {
@@ -61,6 +62,7 @@ type ToolCallPartBase = {
   args: ReadonlyJSONObject;
   artifact?: ReadonlyJSONValue;
   result?: ReadonlyJSONValue;
+  modelContent?: readonly ToolModelContentPart[];
   isError?: boolean;
   parentId?: string;
 };
@@ -68,12 +70,14 @@ type ToolCallPartBase = {
 type ToolCallPartWithoutResult = ToolCallPartBase & {
   state: "partial-call" | "call";
   result?: undefined;
+  modelContent?: undefined;
 };
 
 type ToolCallPartWithResult = ToolCallPartBase & {
   state: "result";
   result: ReadonlyJSONValue;
   artifact?: ReadonlyJSONValue;
+  modelContent?: readonly ToolModelContentPart[];
   isError?: boolean;
 };
 

package/src/index.ts

@@ -35,7 +35,11 @@ export type {
   DataPart,
 } from "./core/utils/types";
 
-export type { Tool } from "./core/tool/tool-types";
+export type {
+  Tool,
+  ToolModelContentPart,
+  ToolModelOutputFunction,
+} from "./core/tool/tool-types";
 export { ToolResponse, type ToolResponseLike } from "./core/tool/ToolResponse";
 export { ToolExecutionStream } from "./core/tool/ToolExecutionStream";
 export type { ToolCallReader } from "./core/tool/tool-types";