Custom agents

import { chat } from "@trigger.dev/sdk/ai";
import { streamText, stepCountIs } from "ai";
import { anthropic } from "@ai-sdk/anthropic";
import { z } from "zod";

export const myChat = chat
  .withClientData({ schema: z.object({ userId: z.string() }) })
  .customAgent({
    id: "my-chat",
    run: async (payload, { signal }) => {
      // One-time initialization — plain code, no hooks. Upsert, not create:
      // continuation runs boot with the row already in place.
      const clientData = payload.metadata!;
      await db.chat.upsert({
        where: { id: payload.chatId },
        create: { id: payload.chatId, userId: clientData.userId },
        update: {},
      });

      const session = chat.createSession(payload, {
        signal,
        idleTimeoutInSeconds: 60,
        timeout: "1h",
      });

      for await (const turn of session) {
        // Persist the incoming user message BEFORE streaming — this is your
        // onTurnStart equivalent. Without it, a page reload mid-stream
        // restores the assistant text (replayed from the session) but loses
        // the user message that prompted it.
        await db.chat.update({
          where: { id: turn.chatId },
          data: { messages: turn.uiMessages },
        });

        const result = streamText({
          model: anthropic("claude-sonnet-4-5"),
          messages: turn.messages,
          abortSignal: turn.signal,
          stopWhen: stepCountIs(15),
        });

        // Pipe, capture, accumulate, and signal turn-complete — all in one call
        await turn.complete(result);

        // Persist the full exchange after the turn — your onTurnComplete equivalent
        await db.chat.update({
          where: { id: turn.chatId },
          data: { messages: turn.uiMessages },
        });
      }
    },
  });

for await (const turn of session) {
  const result = streamText({
    model: anthropic("claude-sonnet-4-5"),
    messages: turn.messages,
    abortSignal: turn.signal, // fires on a user stop OR a run cancel
    stopWhen: stepCountIs(15),
  });

  await turn.complete(result);

  if (turn.stopped) {
    // user stopped this turn — the partial response is already accumulated
  }
}

import { chat, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai";
import { streamText, stepCountIs } from "ai";
import { anthropic } from "@ai-sdk/anthropic";

export const myChat = chat.customAgent({
  id: "my-chat-raw",
  run: async (payload: ChatTaskWirePayload, { signal: runSignal }) => {
    let currentPayload = payload;

    // Handle preload — wait for the first real message
    if (currentPayload.trigger === "preload") {
      const result = await chat.messages.waitWithIdleTimeout({
        idleTimeoutInSeconds: 60,
        timeout: "1h",
        spanName: "waiting for first message",
      });
      if (!result.ok) return;
      currentPayload = result.output;
    }

    const stop = chat.createStopSignal();
    const conversation = new chat.MessageAccumulator();

    // Continuation runs (cancel, crash, upgrade) start with an empty
    // accumulator — fetch stored history so turn 0 can seed it.
    let continuationSeed: UIMessage[] = [];
    if (currentPayload.continuation) {
      const row = await db.chat.findUnique({ where: { id: currentPayload.chatId } });
      continuationSeed = (row?.messages ?? []) as UIMessage[];
    }

    for (let turn = 0; turn < 100; turn++) {
      stop.reset();

      // The wire payload carries at most one new message per turn. Turn 0
      // REPLACES the accumulator, so seed stored history through
      // addIncoming together with the incoming message — a setMessages
      // call before the loop would be wiped here.
      const incoming = currentPayload.message ? [currentPayload.message] : [];
      const turnInput =
        turn === 0 && continuationSeed.length > 0
          ? [...continuationSeed.filter((s) => !incoming.some((m) => m.id === s.id)), ...incoming]
          : incoming;
      const messages = await conversation.addIncoming(turnInput, currentPayload.trigger, turn);

      // Persist the incoming user message before streaming so a
      // mid-stream reload doesn't lose it.
      await db.chat.update({
        where: { id: currentPayload.chatId },
        data: { messages: conversation.uiMessages },
      });

      const combinedSignal = AbortSignal.any([runSignal, stop.signal]);

      const result = streamText({
        model: anthropic("claude-sonnet-4-5"),
        messages,
        abortSignal: combinedSignal,
        stopWhen: stepCountIs(15),
      });

      const { message, status, error } = await chat.pipeAndCapture(result, {
        signal: combinedSignal,
      });
      // pipeAndCapture never throws: a user stop returns status "aborted" with
      // the partial message, and a failure returns status "error" with `error`.
      if (status === "error") throw error;

      if (message) {
        const cleaned =
          stop.signal.aborted && !runSignal.aborted ? chat.cleanupAbortedParts(message) : message;
        await conversation.addResponse(cleaned);
      }

      if (runSignal.aborted) break;

      // Persist, analytics, etc.
      await db.chat.update({
        where: { id: currentPayload.chatId },
        data: { messages: conversation.uiMessages },
      });

      await chat.writeTurnComplete();

      // Wait for the next message
      const next = await chat.messages.waitWithIdleTimeout({
        idleTimeoutInSeconds: 60,
        timeout: "1h",
        spanName: "waiting for next message",
      });
      if (!next.ok) break;
      currentPayload = next.output;
    }

    stop.cleanup();
  },
});