Custom agents

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",
  run: async (payload: ChatTaskWirePayload, { signal }) => {
    // One-time initialization — plain code, no hooks. Upsert, not create:
    // continuation runs boot with the row already in place.
    const clientData = payload.metadata as { userId: string };
    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) {
  if (turn.continuation && turn.number === 0) {
    const row = await db.chat.findUnique({ where: { id: turn.chatId } });
    const stored = (row?.messages ?? []) as UIMessage[];
    if (stored.length > 0) {
      // Keep any incoming message that isn't already persisted
      const incoming = turn.uiMessages.filter((m) => !stored.some((s) => s.id === m.id));
      await turn.setMessages([...stored, ...incoming]);
    }
  }

  // ... streamText + turn.complete as usual
}

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

  // Manual: pipe and capture separately
  const response = await chat.pipeAndCapture(result, { signal: turn.signal });

  if (response) {
    // Custom processing before accumulating
    await turn.addResponse(response);
  }

  // Custom persistence, analytics, etc.
  await db.chat.update({ ... });

  // Must call done() when not using complete()
  await turn.done();
}

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

      let response;
      try {
        response = await chat.pipeAndCapture(result, { signal: combinedSignal });
      } catch (error) {
        if (error instanceof Error && error.name === "AbortError") {
          if (runSignal.aborted) break;
          // Stop — fall through to accumulate partial
        } else {
          throw error;
        }
      }

      if (response) {
        const cleaned =
          stop.signal.aborted && !runSignal.aborted ? chat.cleanupAbortedParts(response) : response;
        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();
  },
});

const conversation = new chat.MessageAccumulator();

// Returns full accumulated ModelMessage[] for streamText
const messages = await conversation.addIncoming(
  payload.message ? [payload.message] : [],
  payload.trigger,
  turn
);

// After piping, add the response
const response = await chat.pipeAndCapture(result);
if (response) await conversation.addResponse(response);

// Access accumulated messages for persistence
conversation.uiMessages; // UIMessage[]
conversation.modelMessages; // ModelMessage[]