fix(chat): harden custom agent validation recovery

Author: gtremper

docs/ai-chat/custom-agents.mdx

@@ -23,7 +23,7 @@ Inside the wrapper, pick one of two loop styles:
 
 Use `chat.withClientData({ schema })` to validate `payload.metadata`. Custom agents parse the initial payload and later message and action frames before passing them to `run`, `chat.messages`, or `chat.createSession`. Schema defaults and transforms are included in the value your code receives. Close frames are not validated.
 
-If validation fails in `run`, `chat.createSession()`, or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. If the initial payload is invalid, the task waits for the next valid frame instead of starting `run` with bad data. Without a schema, metadata is passed through unchanged.
+If validation fails in `run`, `chat.createSession()`, or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. If the initial payload is invalid, the task waits for the next valid frame instead of starting `run` with bad data. On a [head-start handover](/ai-chat/fast-starts#handover-with-custom-agents) boot with invalid client data, the SDK drains the warm handover signal first: a skip ends the run, and a real handover partial is discarded with a logged warning. Without a schema, metadata is passed through unchanged.
 
 `chat.messages.on()` is different because a subscribed frame can arrive while the current response is still streaming. Ending the turn at that point would cut off the response. The SDK skips the frame, logs the validation error, and calls `onClientDataValidationError` if you set it:
 

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

@@ -5539,10 +5539,13 @@ type ChatCustomAgentOptions<
    * `chat.messages.on()` cannot safely complete a turn that may still be
    * streaming, so subscribed frames are reported through this callback and
    * the task log instead.
+   *
+   * `payload.metadata` is typed `unknown`: this callback only fires when
+   * the metadata failed to parse, so it can be any shape the client sent.
    */
   onClientDataValidationError?: (event: {
     error: unknown;
-    payload: ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>;
+    payload: ChatTaskWirePayload<TUIMessage, unknown>;
   }) => Promise<void> | void;
   run: TaskOptions<
     TIdentifier,
@@ -5628,6 +5631,28 @@ function chatCustomAgent<
         );
       }
 
+      // A handover-prepare boot parks the warm handler's signal on
+      // `session.in`. Drain it with the handover facade BEFORE the message
+      // wait below — that facade consumes-and-discards non-message chunks
+      // and would swallow the signal (see `waitForHandover`). The warm
+      // partial cannot be spliced without valid clientData: mirror the
+      // normal flow for skip/crash (exit without a turn) and drop a real
+      // partial — the error chunk above already reported the failure.
+      if (payload.trigger === "handover-prepare") {
+        const signal = await waitForHandover({
+          payload,
+          timeout: "1h",
+          spanName: "waiting for handover signal (invalid clientData)",
+        });
+        if (!signal || signal.kind === "handover-skip") {
+          return;
+        }
+        logger.warn(
+          "chat.customAgent: dropping head-start handover partial — clientData failed validation",
+          { chatId: payload.chatId, isFinal: signal.isFinal }
+        );
+      }
+
       // The Session base payload is sticky across continuation runs. If it is
       // invalid, returning here would boot the same bad metadata again on the
       // next message. Stay attached and wait for a valid wire frame instead.
@@ -8561,7 +8586,10 @@ export interface ChatBuilder<
         options: ChatCustomAgentOptions<TId, undefined, TUIMessage>
       ) => Task<TId, ChatTaskWirePayload<TUIMessage, undefined>, unknown>
     : <TId extends string>(
-        options: ChatCustomAgentOptions<TId, TClientDataSchema, TUIMessage>
+        options: Omit<
+          ChatCustomAgentOptions<TId, TClientDataSchema, TUIMessage>,
+          "clientDataSchema"
+        >
       ) => Task<TId, ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>, unknown>;
 }
 

packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts

@@ -421,6 +421,98 @@ describe("chat.customAgent clientData validation", () => {
     }
   });
 
+  it("exits without a turn when a handover-prepare boot has invalid clientData and the warm handler skips", async () => {
+    const clientData: { userId: unknown } = { userId: 123 };
+    let runCalls = 0;
+
+    const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({
+      id: "custom-agent-client-data-handover-skip",
+      run: async () => {
+        runCalls++;
+        await chat.writeTurnComplete();
+      },
+    });
+
+    const harness = mockChatAgent(agent, {
+      chatId: "custom-agent-client-data-handover-skip-chat",
+      mode: "handover-prepare",
+      clientData,
+    });
+
+    try {
+      await waitFor(() =>
+        harness.allChunks.some((chunk) => (chunk as { type?: string }).type === "error")
+      );
+      expect(runCalls).toBe(0);
+
+      // The recovery path must drain the skip via the handover facade and
+      // end the run, mirroring the normal handover-skip exit.
+      await harness.sendHandoverSkip();
+
+      // The run has exited — a valid frame must NOT boot the loop. (Without
+      // the drain, the run would still be sitting in the message wait and
+      // would process it.) Fire-and-forget: no turn-complete will arrive.
+      clientData.userId = "user_123";
+      void harness.sendMessage(userMessage("late", "message-1")).catch(() => {});
+      await new Promise((resolve) => setTimeout(resolve, 100));
+      expect(runCalls).toBe(0);
+    } finally {
+      await harness.close();
+    }
+  });
+
+  it("drops the head-start partial and recovers on the next valid frame when a handover-prepare boot has invalid clientData", async () => {
+    const clientData: { userId: unknown } = { userId: 123 };
+    let runCalls = 0;
+    let receivedTrigger: string | undefined;
+    let receivedClientData: unknown;
+
+    const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({
+      id: "custom-agent-client-data-handover-drop",
+      run: async (payload) => {
+        runCalls++;
+        receivedTrigger = payload.trigger;
+        receivedClientData = payload.metadata;
+        await chat.writeTurnComplete();
+      },
+    });
+
+    const harness = mockChatAgent(agent, {
+      chatId: "custom-agent-client-data-handover-drop-chat",
+      mode: "handover-prepare",
+      clientData,
+    });
+
+    try {
+      await waitFor(() =>
+        harness.allChunks.some((chunk) => (chunk as { type?: string }).type === "error")
+      );
+      expect(runCalls).toBe(0);
+
+      // Resolves on the next turn-complete — the recovered message turn below.
+      const handover = harness.sendHandover({
+        partialAssistantMessage: [
+          { role: "assistant", content: [{ type: "text", text: "warm partial" }] },
+        ],
+      });
+      // Let the recovery drain consume the handover signal before the
+      // message frame goes out — a frame arriving mid-drain would be
+      // discarded by the handover facade (same as the pre-existing turn-0
+      // handover wait in chat.createSession).
+      await new Promise((resolve) => setTimeout(resolve, 50));
+
+      clientData.userId = "user_123";
+      await harness.sendMessage(userMessage("retry", "message-1"));
+      await handover;
+
+      expect(runCalls).toBe(1);
+      expect(receivedTrigger).toBe("submit-message");
+      expect(receivedClientData).toEqual({ userId: "user_123" });
+    } finally {
+      await harness.close();
+    }
+  });
+
   it("passes clientData through unchanged when no schema is configured", async () => {
     const clientData = { userId: "user_123", nested: { enabled: true } };
     let initialClientData: unknown;