fix(sdk): re-apply tool toModelOutput across chat.agent turns (#3790)

## Summary

`chat.agent` now takes a `tools` option. Until now tools only went to
`streamText` inside `run()`, so the SDK had no tools when it
re-converted the persisted `UIMessage` history at the start of each
turn. Any tool with a `toModelOutput` (raw image bytes into an image
content part, or a sub-agent transcript compressed to a summary) had its
transform applied on turn 1 and skipped from turn 2 onward, so the raw
output got JSON-stringified back into the prompt and the model lost the
transformed view.

Declaring `tools` on the config threads them into that conversion, so
`toModelOutput` runs on every turn. The resolved set is handed back,
typed, on the `run()` payload as `tools`:

```ts
const tools = { searchDocs, renderChart };

export const myChat = chat.agent({
  tools,
  run: async ({ messages, tools, signal }) =>
    streamText({ ...chat.toStreamTextOptions({ tools }), messages, abortSignal: signal }),
});
```

`tools` also accepts a per-turn function for tools that depend on the
user or a feature flag. Only `inputSchema` and `toModelOutput` are read
during conversion, never `execute`. Also exports
`InferChatUIMessageFromTools<typeof tools>` to derive the chat
`UIMessage` type from a tool set. No behavior change for agents that
don't declare `tools`.

Author: ericallam

.changeset/chat-agent-tools.md

@@ -0,0 +1,15 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Add a `tools` option to `chat.agent`. Declaring your tools here threads them into the SDK's internal `convertToModelMessages`, so each tool's `toModelOutput` is re-applied when prior-turn history is re-converted.
+
+```ts
+chat.agent({
+  tools: { readFile, search },
+  run: async ({ messages, tools, signal }) =>
+    streamText({ model, messages, tools, abortSignal: signal }),
+});
+```
+
+Also exports `InferChatUIMessageFromTools<typeof tools>` to derive the chat `UIMessage` type (typed tool parts) directly from a tool set.

packages/trigger-sdk/src/v3/ai-shared.ts

@@ -16,7 +16,7 @@
  */
 
 import type { Task, AnyTask } from "@trigger.dev/core/v3";
-import type { ModelMessage, UIMessage } from "ai";
+import type { InferUITools, ModelMessage, ToolSet, UIDataTypes, UIMessage } from "ai";
 
 /**
  * Message-part `type` value for the pending-message data part the agent
@@ -199,6 +199,26 @@ export type InferChatUIMessage<TTask extends AnyTask> = TTask extends Task<
   ? TUIM
   : UIMessage;
 
+/**
+ * Derive the chat `UIMessage` type for a given tool set. The tool-part types
+ * (`tool-${name}` with typed input/output) are inferred from the tools. Use
+ * this to declare the message type from your tools (e.g. to pass to
+ * `chat.withUIMessage<...>()` or to type the frontend) without hand-writing
+ * the `UIMessage<unknown, UIDataTypes, InferUITools<...>>` triple.
+ *
+ * @example
+ * ```ts
+ * import type { InferChatUIMessageFromTools } from "@trigger.dev/sdk/ai";
+ * const tools = { search, readFile };
+ * type ChatUiMessage = InferChatUIMessageFromTools<typeof tools>;
+ * ```
+ */
+export type InferChatUIMessageFromTools<TTools extends ToolSet> = UIMessage<
+  unknown,
+  UIDataTypes,
+  InferUITools<TTools>
+>;
+
 /**
  * Upsert an incoming wire message into the customer's DB-backed chain
  * inside a `hydrateMessages` hook. Returns `true` iff the chain was

packages/trigger-sdk/src/v3/ai.ts

@@ -102,7 +102,16 @@ const METADATA_KEY = "tool.execute.options";
  * stopped/aborted conversations with partial tool parts.
  */
 function toModelMessages(messages: UIMessage[]): Promise<ModelMessage[]> {
-  return convertToModelMessages(messages, { ignoreIncompleteToolCalls: true });
+  // Pass the resolved per-turn `tools` (if any) so the AI SDK can look up each
+  // tool's `toModelOutput` and re-apply it to prior-turn tool results. Without
+  // `tools` it falls back to JSON-stringifying the raw output (TRI-10149). The
+  // conditional spread keeps the options object byte-identical to the no-tools
+  // path when nothing was declared.
+  const tools = locals.get(chatResolvedToolsKey);
+  return convertToModelMessages(messages, {
+    ignoreIncompleteToolCalls: true,
+    ...(tools ? { tools } : {}),
+  });
 }
 
 export type ToolCallExecutionOptions = {
@@ -1425,7 +1434,10 @@ export type ChatTaskSignals = {
  * The full payload passed to a `chatAgent` run function.
  * Extends `ChatTaskPayload` (the wire payload) with abort signals.
  */
-export type ChatTaskRunPayload<TClientData = unknown> = ChatTaskPayload<TClientData> &
+export type ChatTaskRunPayload<
+  TClientData = unknown,
+  TTools extends ToolSet = ToolSet,
+> = ChatTaskPayload<TClientData> &
   ChatTaskSignals & {
     /**
      * Task run context — same object as the `ctx` passed to a standard `task({ run })` handler’s second argument.
@@ -1436,6 +1448,21 @@ export type ChatTaskRunPayload<TClientData = unknown> = ChatTaskPayload<TClientD
     previousTurnUsage?: LanguageModelUsage;
     /** Cumulative token usage across all completed turns so far. */
     totalUsage: LanguageModelUsage;
+    /**
+     * The resolved tool set for this turn, the same `tools` you declared on
+     * `chat.agent({ tools })` (or the result of the per-turn `tools` function).
+     * Pass straight to `streamText({ tools })` so you don't redeclare them:
+     *
+     * ```ts
+     * run: ({ messages, tools, signal }) =>
+     *   streamText({ model, messages, tools, abortSignal: signal })
+     * ```
+     *
+     * Declaring `tools` on the config is also what lets the SDK re-run each
+     * tool's `toModelOutput` when it re-converts prior-turn history (see the
+     * `tools` option on `chat.agent`). Empty object when no `tools` were declared.
+     */
+    tools: TTools;
   };
 
 // Input streams for bidirectional chat communication
@@ -2366,6 +2393,20 @@ const chatPrepareMessagesKey =
   locals.create<(event: PrepareMessagesEvent<unknown>) => ModelMessage[] | Promise<ModelMessage[]>>(
     "chat.prepareMessages"
   );
+/**
+ * @internal The raw `tools` option from `chat.agent({ tools })`, either a
+ * static `ToolSet` or a per-turn function. Set once at boot.
+ */
+const chatToolsOptionKey = locals.create<
+  ToolSet | ((event: ResolveToolsEvent<unknown>) => ToolSet | Promise<ToolSet>)
+>("chat.toolsOption");
+/**
+ * @internal The concrete `ToolSet` resolved for the current turn. Read by
+ * `toModelMessages` so `convertToModelMessages` can re-run `toModelOutput` on
+ * prior-turn tool results. Unset when no `tools` were declared (preserves the
+ * exact pre-feature conversion behavior).
+ */
+const chatResolvedToolsKey = locals.create<ToolSet>("chat.resolvedTools");
 
 /** @internal Flag set by `chat.requestUpgrade()` to exit the loop after the current turn. */
 const chatUpgradeRequestedKey = locals.create<boolean>("chat.upgradeRequested");
@@ -2626,6 +2667,25 @@ export type PrepareMessagesEvent<TClientData = unknown> = {
   clientData?: TClientData;
 };
 
+/**
+ * Event passed to the per-turn `tools` function form on `chat.agent`.
+ *
+ * Use this when the active tool set depends on per-turn context (the user, a
+ * feature flag, etc.). Return the `ToolSet` to use for converting this turn's
+ * history. Only `inputSchema` and `toModelOutput` are read during conversion,
+ * so a lightweight map (no `execute`) is fine.
+ */
+export type ResolveToolsEvent<TClientData = unknown> = {
+  /** The chat session ID. */
+  chatId: string;
+  /** The current turn number (0-indexed). */
+  turn: number;
+  /** Whether this run is continuing an existing chat. */
+  continuation: boolean;
+  /** Custom data from the frontend. */
+  clientData?: TClientData;
+};
+
 /**
  * Data shape for `data-compaction` stream chunks emitted during compaction.
  * Use to type the `data` field when rendering compaction parts in the frontend.
@@ -2800,6 +2860,41 @@ async function applyPrepareMessages(
   );
 }
 
+/**
+ * Resolve the `tools` option into a concrete `ToolSet` and cache it in locals so
+ * `toModelMessages` can pass it to `convertToModelMessages`. For the function
+ * form, invokes the user function with the given context (or the current turn
+ * context when no override is passed). Pass an `override` for the boot-time
+ * history conversion, which runs before the per-turn context exists and uses
+ * the run/continuation payload's `clientData`.
+ *
+ * Fails closed: a throwing resolver propagates rather than carrying a prior
+ * turn's set forward. The function form can gate capabilities by user or flag,
+ * so reusing stale tools would leak capabilities. No-op when no `tools` were
+ * declared.
+ * @internal
+ */
+async function resolveTurnTools(
+  override?: { chatId: string; turn: number; continuation: boolean; clientData: unknown }
+): Promise<void> {
+  const option = locals.get(chatToolsOptionKey);
+  if (!option) return;
+
+  if (typeof option !== "function") {
+    locals.set(chatResolvedToolsKey, option);
+    return;
+  }
+
+  const ctx = override ?? locals.get(chatTurnContextKey);
+  const resolved = await option({
+    chatId: ctx?.chatId ?? "",
+    turn: ctx?.turn ?? 0,
+    continuation: ctx?.continuation ?? false,
+    clientData: ctx?.clientData,
+  });
+  locals.set(chatResolvedToolsKey, resolved);
+}
+
 /**
  * Read the current compaction state. Returns the summary and base message count
  * if compaction has occurred in this turn, or `undefined` if not.
@@ -4250,6 +4345,7 @@ export type ChatAgentOptions<
   TClientDataSchema extends TaskSchema | undefined = undefined,
   TUIMessage extends UIMessage = UIMessage,
   TActionSchema extends TaskSchema | undefined = undefined,
+  TTools extends ToolSet = ToolSet,
 > = Omit<
   TaskOptions<
     TIdentifier,
@@ -4360,6 +4456,41 @@ export type ChatAgentOptions<
     >
   ) => Promise<unknown> | unknown;
 
+  /**
+   * The tools available to this agent.
+   *
+   * `chat.agent` doesn't call the model for you. Your tools still go to
+   * `streamText({ tools })` inside `run()`. Declaring them here additionally
+   * lets the SDK re-run each tool's
+   * [`toModelOutput`](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#tomodeloutput)
+   * when it re-converts persisted history on later turns. Without this, the
+   * AI SDK has no `tools` to look up `toModelOutput` against, so a tool's
+   * transformed result (e.g. raw image bytes → an image content part, or a
+   * sub-agent summary) silently degrades to its raw JSON output from turn 2
+   * onward.
+   *
+   * Only `inputSchema` and `toModelOutput` are read during conversion (never
+   * `execute`), so you may pass a lightweight map if you keep heavy execute
+   * deps out of this module.
+   *
+   * Pass either a static `ToolSet` or a function of per-turn context (for
+   * tools that depend on the user, a feature flag, etc.). The resolved set is
+   * available on the `run()` payload as `tools`.
+   *
+   * @example
+   * ```ts
+   * const tools = { read_file, search };
+   * chat.agent({
+   *   tools,
+   *   run: async ({ messages, tools, signal }) =>
+   *     streamText({ model, messages, tools, abortSignal: signal }),
+   * });
+   * ```
+   */
+  tools?:
+    | TTools
+    | ((event: ResolveToolsEvent<inferSchemaOut<TClientDataSchema>>) => TTools | Promise<TTools>);
+
   /**
    * The run function for the chat task.
    *
@@ -4370,7 +4501,9 @@ export type ChatAgentOptions<
    * **Auto-piping:** If this function returns a value with `.toUIMessageStream()`,
    * the stream is automatically piped to the frontend.
    */
-  run: (payload: ChatTaskRunPayload<inferSchemaOut<TClientDataSchema>>) => Promise<unknown>;
+  run: (
+    payload: ChatTaskRunPayload<inferSchemaOut<TClientDataSchema>, TTools>
+  ) => Promise<unknown>;
 
   /**
    * Called once at the start of every run boot — for the initial run, for
@@ -4951,8 +5084,9 @@ function chatAgent<
   TClientDataSchema extends TaskSchema | undefined = undefined,
   TUIMessage extends UIMessage = UIMessage,
   TActionSchema extends TaskSchema | undefined = undefined,
+  TTools extends ToolSet = ToolSet,
 >(
-  options: ChatAgentOptions<TIdentifier, TClientDataSchema, TUIMessage, TActionSchema>
+  options: ChatAgentOptions<TIdentifier, TClientDataSchema, TUIMessage, TActionSchema, TTools>
 ): Task<TIdentifier, ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>, unknown> {
   const {
     run: userRun,
@@ -4971,6 +5105,7 @@ function chatAgent<
     compaction,
     pendingMessages: pendingMessagesConfig,
     prepareMessages,
+    tools: toolsOption,
     onTurnComplete,
     maxTurns = 100,
     turnTimeout = "1h",
@@ -5049,6 +5184,25 @@ function chatAgent<
         locals.set(chatPrepareMessagesKey, prepareMessages);
       }
 
+      if (toolsOption) {
+        // Cast: the option's function form is typed against the parsed
+        // `clientData` (`ResolveToolsEvent<inferSchemaOut<...>>`), but the
+        // locals key uses the erased `ResolveToolsEvent<unknown>`. The runtime
+        // value is identical; this mirrors how `prepareMessages` is stored.
+        locals.set(
+          chatToolsOptionKey,
+          toolsOption as
+            | ToolSet
+            | ((event: ResolveToolsEvent<unknown>) => ToolSet | Promise<ToolSet>)
+        );
+        // Static tools are usable immediately. The function form is resolved
+        // just before the boot history conversion (with the payload's
+        // clientData) and again per-turn (see resolveTurnTools).
+        if (typeof toolsOption !== "function") {
+          locals.set(chatResolvedToolsKey, toolsOption);
+        }
+      }
+
       if (compaction) {
         locals.set(
           chatAgentCompactionKey,
@@ -5438,6 +5592,29 @@ function chatAgent<
         }
 
         if (accumulatedUIMessages.length > 0) {
+          // Resolve a function-form `tools` with the run/continuation payload's
+          // clientData so this conversion of the restored history applies each
+          // tool's toModelOutput (static tools were already seeded above). This
+          // only re-renders saved history, so it fails open: a resolver hiccup
+          // logs and converts without tools rather than blocking the resume.
+          // Per-turn resolveTurnTools still fails closed for live turns.
+          if (typeof toolsOption === "function") {
+            try {
+              await resolveTurnTools({
+                chatId: payload.chatId,
+                turn: 0,
+                continuation: payload.continuation ?? false,
+                clientData: parseClientData
+                  ? await parseClientData(payload.metadata)
+                  : payload.metadata,
+              });
+            } catch (error) {
+              logger.warn(
+                "chat.agent: tools() resolver threw at boot; restored history converted without toModelOutput",
+                { error: error instanceof Error ? error.message : String(error) }
+              );
+            }
+          }
           try {
             accumulatedMessages = await toModelMessages(accumulatedUIMessages);
           } catch (error) {
@@ -5958,6 +6135,11 @@ function chatAgent<
                     clientData,
                   });
 
+                  // Resolve the per-turn `tools` set now that turn context
+                  // (incl. parsed clientData) exists, so every toModelMessages
+                  // call this turn can re-apply tool `toModelOutput`.
+                  await resolveTurnTools();
+
                   // Per-turn stop controller (reset each turn)
                   const stopController = new AbortController();
                   currentStopController = stopController;
@@ -6613,6 +6795,7 @@ function chatAgent<
                         previousTurnUsage,
                         totalUsage: cumulativeUsage,
                         ctx,
+                        tools: locals.get(chatResolvedToolsKey) ?? {},
                         signal: combinedSignal,
                         cancelSignal,
                         stopSignal,
@@ -7512,11 +7695,11 @@ export interface ChatBuilder<
    * (backwards compatible).
    */
   agent: [TClientDataSchema] extends [undefined]
-  ? <TId extends string, TInfer extends TaskSchema | undefined = undefined, TAction extends TaskSchema | undefined = undefined>(
-    options: ChatAgentOptions<TId, TInfer, TUIMessage, TAction>
+  ? <TId extends string, TInfer extends TaskSchema | undefined = undefined, TAction extends TaskSchema | undefined = undefined, TTools extends ToolSet = ToolSet>(
+    options: ChatAgentOptions<TId, TInfer, TUIMessage, TAction, TTools>
   ) => Task<TId, ChatTaskWirePayload<TUIMessage, inferSchemaIn<TInfer>>, unknown>
-  : <TId extends string, TAction extends TaskSchema | undefined = undefined>(
-    options: Omit<ChatAgentOptions<TId, TClientDataSchema, TUIMessage, TAction>, "clientDataSchema">
+  : <TId extends string, TAction extends TaskSchema | undefined = undefined, TTools extends ToolSet = ToolSet>(
+    options: Omit<ChatAgentOptions<TId, TClientDataSchema, TUIMessage, TAction, TTools>, "clientDataSchema">
   ) => Task<TId, ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>, unknown>;
 
   /**
@@ -9145,7 +9328,11 @@ function chatLocal<T extends Record<string, unknown>>(options: { id: string }):
 // the browser graph. Re-exported here so `@trigger.dev/sdk/ai` consumers
 // still see them.
 import type { InferChatClientData, InferChatUIMessage } from "./ai-shared.js";
-export type { InferChatClientData, InferChatUIMessage } from "./ai-shared.js";
+export type {
+  InferChatClientData,
+  InferChatUIMessage,
+  InferChatUIMessageFromTools,
+} from "./ai-shared.js";
 
 /**
  * Options for {@link createChatStartSessionAction}.