From 64929b757fa3e92fe829b32ff5fe7ee0ed29a467 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sun, 10 May 2026 22:26:59 +0100 Subject: [PATCH] =?UTF-8?q?feat(sdk):=20chat.agent=20=E2=80=94=20durable?= =?UTF-8?q?=20conversational=20task=20runtime?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds chat.agent({...}), a session-aware task definition for AI chat agents. The runtime sits on top of the Sessions primitive and handles: - Delta-only wire: each /in/append carries at most one message; agent rebuilds prior history at run boot from an S3 snapshot + session.out replay tail. Awaited snapshot writes after every onTurnComplete keep the chain durable across idle suspends. - Lifecycle hooks: onChatStart, onTurnStart, onTurnComplete, onAction, onValidateMessages, hydrateMessages (short-circuits snapshot+replay if the customer owns history). - chat.history read primitives for HITL flows (getPendingToolCalls, getResolvedToolCalls, extractNewToolResults, findMessage, all/getChain). - chat.local — per-run typed data with Proxy access + dirty tracking. - chat.headStart — first-turn TTFC bridge via a customer HTTP handler. - oomMachine — one-shot OOM-retry on a larger machine; cutoff derived from latest trigger:turn-complete chunk on session.out. - Actions are no longer turns — onAction may mutate via chat.history.* and may return a StreamTextResult without bumping the turn counter. - gen_ai.conversation.id stamped on chat spans + metrics for telemetry. - Skills runtime (loadable per-agent and per-task) + agent skills bundling. Snapshot + replay integration tests in apps/webapp/test/. mockChatAgent test harness updated for the new wire. Includes the chat-agent, chat-agent-delta-wire-snapshots, chat-history-read-primitives, chat-head-start, chat-actions-no-turn, chat-session-attributes, agent-skills, and mock-chat-agent-test-harness changesets. --- .changeset/agent-skills.md | 16 + .changeset/chat-actions-no-turn.md | 33 + .changeset/chat-agent-delta-wire-snapshots.md | 8 + .changeset/chat-agent.md | 30 + .changeset/chat-head-start.md | 34 + .changeset/chat-history-read-primitives.md | 21 + .changeset/chat-session-attributes.md | 6 + .changeset/mock-chat-agent-test-harness.md | 8 + .../test/chat-snapshot-integration.test.ts | 235 + apps/webapp/test/replay-after-crash.test.ts | 315 + packages/build/package.json | 17 +- packages/build/src/extensions/secureExec.ts | 172 + packages/build/src/internal.ts | 1 + .../build/src/internal/additionalFiles.ts | 96 +- packages/build/src/internal/copyFiles.ts | 99 + packages/core/package.json | 30 + packages/core/src/v3/apiClient/index.ts | 311 + packages/core/src/v3/errors.ts | 20 + packages/core/src/v3/index.ts | 2 + .../core/src/v3/resource-catalog/catalog.ts | 12 +- .../core/src/v3/resource-catalog/index.ts | 21 +- .../resource-catalog/noopResourceCatalog.ts | 21 +- .../standardResourceCatalog.ts | 70 +- .../core/src/v3/taskContext/index.test.ts | 86 + packages/core/src/v3/taskContext/index.ts | 20 + .../core/src/v3/taskContext/otelProcessors.ts | 16 + packages/core/src/v3/test/index.ts | 9 + .../core/src/v3/test/mock-task-context.ts | 294 + packages/core/test/mockTaskContext.test.ts | 226 + packages/core/test/skillCatalog.test.ts | 74 + packages/trigger-sdk/package.json | 89 +- .../trigger-sdk/src/v3/agentSkillsRuntime.ts | 127 + packages/trigger-sdk/src/v3/ai-shared.ts | 210 + packages/trigger-sdk/src/v3/ai.ts | 8726 ++++++++++++++++- .../trigger-sdk/src/v3/chat-server.test.ts | 617 ++ packages/trigger-sdk/src/v3/chat-server.ts | 893 ++ packages/trigger-sdk/src/v3/deployments.ts | 56 + packages/trigger-sdk/src/v3/index.ts | 13 +- packages/trigger-sdk/src/v3/runs.ts | 9 + packages/trigger-sdk/src/v3/sessions.ts | 743 ++ packages/trigger-sdk/src/v3/shared.ts | 212 +- packages/trigger-sdk/src/v3/skill.ts | 211 + packages/trigger-sdk/src/v3/skills.ts | 9 + packages/trigger-sdk/src/v3/streams.ts | 121 +- packages/trigger-sdk/src/v3/tasks.ts | 2 + packages/trigger-sdk/src/v3/test/index.ts | 23 + .../src/v3/test/mock-chat-agent.ts | 686 ++ .../trigger-sdk/src/v3/test/setup-catalog.ts | 16 + .../src/v3/test/test-session-handle.ts | 268 + .../trigger-sdk/test/chat-snapshot.test.ts | 279 + packages/trigger-sdk/test/merge-by-id.test.ts | 158 + .../trigger-sdk/test/mockChatAgent.test.ts | 1441 +++ .../test/replay-session-out.test.ts | 307 + packages/trigger-sdk/test/skill.test.ts | 86 + .../trigger-sdk/test/skillsRuntime.test.ts | 221 + packages/trigger-sdk/test/wire-shape.test.ts | 249 + 56 files changed, 17924 insertions(+), 151 deletions(-) create mode 100644 .changeset/agent-skills.md create mode 100644 .changeset/chat-actions-no-turn.md create mode 100644 .changeset/chat-agent-delta-wire-snapshots.md create mode 100644 .changeset/chat-agent.md create mode 100644 .changeset/chat-head-start.md create mode 100644 .changeset/chat-history-read-primitives.md create mode 100644 .changeset/chat-session-attributes.md create mode 100644 .changeset/mock-chat-agent-test-harness.md create mode 100644 apps/webapp/test/chat-snapshot-integration.test.ts create mode 100644 apps/webapp/test/replay-after-crash.test.ts create mode 100644 packages/build/src/extensions/secureExec.ts create mode 100644 packages/build/src/internal/copyFiles.ts create mode 100644 packages/core/src/v3/taskContext/index.test.ts create mode 100644 packages/core/src/v3/test/index.ts create mode 100644 packages/core/src/v3/test/mock-task-context.ts create mode 100644 packages/core/test/mockTaskContext.test.ts create mode 100644 packages/core/test/skillCatalog.test.ts create mode 100644 packages/trigger-sdk/src/v3/agentSkillsRuntime.ts create mode 100644 packages/trigger-sdk/src/v3/ai-shared.ts create mode 100644 packages/trigger-sdk/src/v3/chat-server.test.ts create mode 100644 packages/trigger-sdk/src/v3/chat-server.ts create mode 100644 packages/trigger-sdk/src/v3/deployments.ts create mode 100644 packages/trigger-sdk/src/v3/sessions.ts create mode 100644 packages/trigger-sdk/src/v3/skill.ts create mode 100644 packages/trigger-sdk/src/v3/skills.ts create mode 100644 packages/trigger-sdk/src/v3/test/index.ts create mode 100644 packages/trigger-sdk/src/v3/test/mock-chat-agent.ts create mode 100644 packages/trigger-sdk/src/v3/test/setup-catalog.ts create mode 100644 packages/trigger-sdk/src/v3/test/test-session-handle.ts create mode 100644 packages/trigger-sdk/test/chat-snapshot.test.ts create mode 100644 packages/trigger-sdk/test/merge-by-id.test.ts create mode 100644 packages/trigger-sdk/test/mockChatAgent.test.ts create mode 100644 packages/trigger-sdk/test/replay-session-out.test.ts create mode 100644 packages/trigger-sdk/test/skill.test.ts create mode 100644 packages/trigger-sdk/test/skillsRuntime.test.ts create mode 100644 packages/trigger-sdk/test/wire-shape.test.ts diff --git a/.changeset/agent-skills.md b/.changeset/agent-skills.md new file mode 100644 index 00000000000..5ed3b11fc2f --- /dev/null +++ b/.changeset/agent-skills.md @@ -0,0 +1,16 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core": patch +"@trigger.dev/build": patch +"trigger.dev": patch +--- + +Add Agent Skills for `chat.agent`. Drop a folder with a `SKILL.md` and any helper scripts/references next to your task code, register it with `skills.define({ id, path })`, and the CLI bundles it into the deploy image automatically — no `trigger.config.ts` changes. The agent gets a one-line summary in its system prompt and discovers full instructions on demand via `loadSkill`, with `bash` and `readFile` tools scoped per-skill (path-traversal guards, output caps, abort-signal propagation). + +```ts +const pdfSkill = skills.define({ id: "pdf-extract", path: "./skills/pdf-extract" }); + +chat.skills.set([await pdfSkill.local()]); +``` + +Built on the [AI SDK cookbook pattern](https://ai-sdk.dev/cookbook/guides/agent-skills) — portable across providers. SDK + CLI only for now; dashboard-editable `SKILL.md` text is on the roadmap. diff --git a/.changeset/chat-actions-no-turn.md b/.changeset/chat-actions-no-turn.md new file mode 100644 index 00000000000..a0113441520 --- /dev/null +++ b/.changeset/chat-actions-no-turn.md @@ -0,0 +1,33 @@ +--- +"@trigger.dev/sdk": minor +--- + +`chat.agent` actions are no longer treated as turns. They fire `hydrateMessages` and `onAction` only — no `onTurnStart` / `prepareMessages` / `onBeforeTurnComplete` / `onTurnComplete`, no `run()`, no turn-counter increment. The trace span is named `chat action` instead of `chat turn N`. + +`onAction` can now return a `StreamTextResult`, `string`, or `UIMessage` to produce a model response from the action; returning `void` (the previous and now default) is side-effect-only. + +**Migration**: if you previously had `run()` branching on `payload.trigger === "action"`, return your `streamText(...)` from `onAction` instead. If you persisted in `onTurnComplete`, do that work inside `onAction`. For any other state-only action, just remove your skip-the-model workaround — the default is now correct. + +```ts +// before +onAction: async ({ action }) => { + if (action.type === "regenerate") { + chat.store.set({ skipModelCall: false }); + chat.history.slice(0, -1); + } +}, +run: async ({ messages, signal }) => { + if (chat.store.get()?.skipModelCall) return; + return streamText({ model, messages, abortSignal: signal }); +}, + +// after +onAction: async ({ action, messages, signal }) => { + if (action.type === "regenerate") { + chat.history.slice(0, -1); + return streamText({ model, messages, abortSignal: signal }); + } +}, +run: async ({ messages, signal }) => + streamText({ model, messages, abortSignal: signal }), +``` diff --git a/.changeset/chat-agent-delta-wire-snapshots.md b/.changeset/chat-agent-delta-wire-snapshots.md new file mode 100644 index 00000000000..21a8fd01fa4 --- /dev/null +++ b/.changeset/chat-agent-delta-wire-snapshots.md @@ -0,0 +1,8 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core": patch +--- + +`chat.agent` wire is now delta-only — clients ship at most one new message per `.in/append` instead of the full `UIMessage[]` history. The agent rebuilds prior history at run boot from a JSON snapshot in object storage plus a `wait=0` replay of the `session.out` tail. Long chats stop hitting the 512 KiB body cap on `/realtime/v1/sessions/{id}/in/append`. Snapshot writes happen after every `onTurnComplete`, awaited so they survive idle suspend; reads happen only at run boot. Registering a `hydrateMessages` hook short-circuits both the snapshot read/write and the replay — the customer is the source of truth for history. + +Custom transports that constructed `ChatTaskWirePayload` directly need to drop the `messages: UIMessage[]` field and use `message?: UIMessage` (singular). Built-in transports (`TriggerChatTransport`, `AgentChat`) handle the change below the customer-facing surface — most apps need no changes. Configure object-store env vars (`OBJECT_STORE_*`) on your webapp deployment if you haven't already; without an object store and without `hydrateMessages`, conversations don't survive run boundaries. diff --git a/.changeset/chat-agent.md b/.changeset/chat-agent.md new file mode 100644 index 00000000000..9ca65682da7 --- /dev/null +++ b/.changeset/chat-agent.md @@ -0,0 +1,30 @@ +--- +"@trigger.dev/sdk": minor +"@trigger.dev/core": patch +--- + +Run AI chats as durable Trigger.dev tasks. Define the agent in one function, wire `useChat` to it from React, and the conversation survives page refreshes, network blips, and process restarts — with built-in support for tools, HITL approvals, multi-turn state, and stop-mid-stream cancellation. + +```ts +import { chat } from "@trigger.dev/sdk/ai"; +import { streamText } from "ai"; +import { openai } from "@ai-sdk/openai"; + +export const myChat = chat.agent({ + id: "my-chat", + run: async ({ messages, signal }) => + streamText({ model: openai("gpt-4o"), messages, abortSignal: signal }), +}); +``` + +```tsx +import { useChat } from "@ai-sdk/react"; +import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react"; + +const transport = useTriggerChatTransport({ task: "my-chat", accessToken }); +const { messages, sendMessage } = useChat({ transport }); +``` + +Lifecycle hooks (`onPreload`, `onTurnStart`, `onTurnComplete`, etc.) cover the common needs around persistence, validation, and post-turn work. `chat.store` gives you a typed shared-data slot the agent and client both read and write. `chat.endRun()` exits cleanly when the agent decides it's done. The transport's `watch` mode lets a dashboard tab observe a run without driving it. + +Drops the pre-Sessions chat stream constants (`CHAT_STREAM_KEY`, `CHAT_MESSAGES_STREAM_ID`, `CHAT_STOP_STREAM_ID`) — migrate to `sessions.open(id).out` / `.in`. diff --git a/.changeset/chat-head-start.md b/.changeset/chat-head-start.md new file mode 100644 index 00000000000..5e33344493f --- /dev/null +++ b/.changeset/chat-head-start.md @@ -0,0 +1,34 @@ +--- +"@trigger.dev/sdk": minor +--- + +Add `chat.headStart` — an opt-in fast-path that runs the first turn's `streamText` step in your warm Next.js / Hono / Workers / Express handler while the trigger agent run boots in parallel. Cold-start TTFC drops by ~50% on the first message; the agent owns step 2+ (tool execution, persistence, hooks) so heavy deps stay where they belong. + +```ts +// app/api/chat/route.ts (Next.js / any Web Fetch framework) +import { chat } from "@trigger.dev/sdk/chat-server"; +import { streamText } from "ai"; +import { openai } from "@ai-sdk/openai"; +import { headStartTools } from "@/lib/chat-tools-schemas"; // schema-only + +export const POST = chat.headStart({ + agentId: "ai-chat", + run: async ({ chat: chatHelper }) => + streamText({ + ...chatHelper.toStreamTextOptions({ tools: headStartTools }), + model: openai("gpt-4o-mini"), + system: "You are a helpful AI assistant.", + }), +}); +``` + +```tsx +// browser — opt in by pointing the transport at your handler +const transport = useTriggerChatTransport({ + task: "ai-chat", + accessToken, + headStart: "/api/chat", // first-turn-only; turn 2+ bypasses the endpoint +}); +``` + +For Node-only frameworks (Express, Fastify, Koa, raw `node:http`) use `chat.toNodeListener(handler)` to bridge the Web Fetch handler to `(req, res)`. Adds a new `@trigger.dev/sdk/chat-server` subpath; bundle stays Web Fetch–only with no `node:*` imports. diff --git a/.changeset/chat-history-read-primitives.md b/.changeset/chat-history-read-primitives.md new file mode 100644 index 00000000000..fd26ad8548b --- /dev/null +++ b/.changeset/chat-history-read-primitives.md @@ -0,0 +1,21 @@ +--- +"@trigger.dev/sdk": minor +--- + +Add read primitives to `chat.history` for HITL flows: `getPendingToolCalls()`, `getResolvedToolCalls()`, `extractNewToolResults(message)`, `getChain()`, and `findMessage(messageId)`. These lift the accumulator-walking logic that customers building human-in-the-loop tools were re-implementing into the SDK. + +Use `getPendingToolCalls()` to gate fresh user turns while a tool call is awaiting an answer. Use `extractNewToolResults(message)` to dedup tool results when persisting to your own store — the helper returns only the parts whose `toolCallId` is not already resolved on the chain. + +```ts +const pending = chat.history.getPendingToolCalls(); +if (pending.length > 0) { + // an addToolOutput is expected before a new user message +} + +onTurnComplete: async ({ responseMessage }) => { + const newResults = chat.history.extractNewToolResults(responseMessage); + for (const r of newResults) { + await db.toolResults.upsert({ id: r.toolCallId, output: r.output, errorText: r.errorText }); + } +}; +``` diff --git a/.changeset/chat-session-attributes.md b/.changeset/chat-session-attributes.md new file mode 100644 index 00000000000..ec4c6a54076 --- /dev/null +++ b/.changeset/chat-session-attributes.md @@ -0,0 +1,6 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core": patch +--- + +Stamp `gen_ai.conversation.id` (the chat id) on every span and metric emitted from inside a `chat.task` or `chat.agent` run. Lets you filter dashboard spans, runs, and metrics by the chat conversation that produced them — independent of the run boundary, so multi-run chats correlate cleanly. No code changes required on the user side. diff --git a/.changeset/mock-chat-agent-test-harness.md b/.changeset/mock-chat-agent-test-harness.md new file mode 100644 index 00000000000..9876e56a9f7 --- /dev/null +++ b/.changeset/mock-chat-agent-test-harness.md @@ -0,0 +1,8 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core": patch +--- + +Unit-test `chat.agent` definitions offline with `mockChatAgent` from `@trigger.dev/sdk/ai/test`. Drives a real agent's turn loop in-process — no network, no task runtime — so you can send messages, actions, and stop signals via driver methods, inspect captured output chunks, and verify hooks fire. Pairs with `MockLanguageModelV3` from `ai/test` for model mocking. `setupLocals` lets you pre-seed `locals` (DB clients, service stubs) before `run()` starts. + +The broader `runInMockTaskContext` harness it's built on lives at `@trigger.dev/core/v3/test` — useful for unit-testing any task code, not just chat. diff --git a/apps/webapp/test/chat-snapshot-integration.test.ts b/apps/webapp/test/chat-snapshot-integration.test.ts new file mode 100644 index 00000000000..3d157d58f9f --- /dev/null +++ b/apps/webapp/test/chat-snapshot-integration.test.ts @@ -0,0 +1,235 @@ +// Plan F.3: integration test that round-trips a `ChatSnapshotV1` blob +// through the SDK's snapshot helpers + a real MinIO backing store. Mirrors +// the testcontainer pattern from `objectStore.test.ts`. +// +// What this verifies end-to-end: +// - SDK's `writeChatSnapshot` calls `apiClient.createUploadPayloadUrl` +// to mint a presigned PUT, then PUTs JSON to it. +// - SDK's `readChatSnapshot` calls `apiClient.getPayloadUrl` to mint a +// presigned GET, then fetches and parses. +// - The webapp's `generatePresignedUrl` produces URLs MinIO accepts. +// - The blob round-trips with `version: 1` shape preserved. +// - 404 (no snapshot for a fresh session) returns `undefined`, not an +// error. +// +// This is the integration safety net behind the unit tests in +// `packages/trigger-sdk/test/chat-snapshot.test.ts` — those tests mock +// `fetch`; this one drives a real S3-compatible backend. + +import { postgresAndMinioTest } from "@internal/testcontainers"; +import { apiClientManager } from "@trigger.dev/core/v3"; +import { + __readChatSnapshotProductionPathForTests as readChatSnapshot, + __writeChatSnapshotProductionPathForTests as writeChatSnapshot, + type ChatSnapshotV1, +} from "@trigger.dev/sdk/ai"; +import type { UIMessage } from "ai"; +import { afterEach, describe, expect, vi } from "vitest"; +import { env } from "~/env.server"; +import { generatePresignedUrl } from "~/v3/objectStore.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +// ── Helpers ──────────────────────────────────────────────────────────── + +function makeSnapshot(opts: { messages?: UIMessage[]; lastOutEventId?: string } = {}): ChatSnapshotV1 { + return { + version: 1, + savedAt: 1_700_000_000_000, + messages: opts.messages ?? [ + { + id: "u-1", + role: "user", + parts: [{ type: "text", text: "hello" }], + }, + { + id: "a-1", + role: "assistant", + parts: [{ type: "text", text: "world" }], + }, + ], + lastOutEventId: opts.lastOutEventId ?? "evt-42", + lastOutTimestamp: 1_700_000_000_500, + }; +} + +/** + * Stub `apiClientManager.clientOrThrow()` so the SDK helpers see a fake + * api client whose `getPayloadUrl` / `createUploadPayloadUrl` return + * presigned URLs minted by the webapp's real `generatePresignedUrl` + * (which signs against MinIO). + * + * The SDK helpers internally do `fetch(presignedUrl, ...)` to read/write + * the blob, so MinIO ends up holding the actual bytes. + */ +function stubApiClient(opts: { projectRef: string; envSlug: string }) { + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + async getPayloadUrl(filename: string) { + const result = await generatePresignedUrl(opts.projectRef, opts.envSlug, filename, "GET"); + if (!result.success) throw new Error(result.error); + return { presignedUrl: result.url }; + }, + async createUploadPayloadUrl(filename: string) { + const result = await generatePresignedUrl(opts.projectRef, opts.envSlug, filename, "PUT"); + if (!result.success) throw new Error(result.error); + return { presignedUrl: result.url }; + }, + } as never); +} + +// Suppress noisy warnings from logger.warn during error-path tests. +let warnSpy: ReturnType; + +afterEach(() => { + vi.restoreAllMocks(); + warnSpy?.mockRestore(); +}); + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("chat snapshot integration (MinIO + SDK helpers)", () => { + postgresAndMinioTest("round-trips a snapshot through real MinIO", async ({ minioConfig }) => { + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + stubApiClient({ projectRef: "proj_snap_rt", envSlug: "dev" }); + + const sessionId = "sess_round_trip_1"; + const snapshot = makeSnapshot(); + + // Write through the SDK helper — should land in MinIO at + // `packets/proj_snap_rt/dev/sessions/sess_round_trip_1/snapshot.json`. + await writeChatSnapshot(sessionId, snapshot); + + // Read back through the SDK helper — should reconstruct the original. + const result = await readChatSnapshot(sessionId); + + expect(result).toEqual(snapshot); + }); + + postgresAndMinioTest("returns undefined for a fresh session with no snapshot", async ({ minioConfig }) => { + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + stubApiClient({ projectRef: "proj_snap_404", envSlug: "dev" }); + + warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + + // Session never had a snapshot written — read returns undefined. + const result = await readChatSnapshot("sess_never_existed"); + expect(result).toBeUndefined(); + }); + + postgresAndMinioTest("overwrites a prior snapshot in place (single-writer)", async ({ minioConfig }) => { + // The runtime guarantees one attempt alive at a time, and + // `writeChatSnapshot` runs awaited after `onTurnComplete`. Verify + // that a second write to the same key replaces the first cleanly — + // the read-after-write reflects the latest blob. + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + stubApiClient({ projectRef: "proj_snap_overwrite", envSlug: "dev" }); + + const sessionId = "sess_overwrite"; + + const turn1 = makeSnapshot({ + messages: [ + { id: "u-1", role: "user", parts: [{ type: "text", text: "first" }] }, + ], + lastOutEventId: "evt-turn1", + }); + const turn2 = makeSnapshot({ + messages: [ + { id: "u-1", role: "user", parts: [{ type: "text", text: "first" }] }, + { id: "a-1", role: "assistant", parts: [{ type: "text", text: "reply-1" }] }, + { id: "u-2", role: "user", parts: [{ type: "text", text: "second" }] }, + { id: "a-2", role: "assistant", parts: [{ type: "text", text: "reply-2" }] }, + ], + lastOutEventId: "evt-turn2", + }); + + await writeChatSnapshot(sessionId, turn1); + await writeChatSnapshot(sessionId, turn2); + + const result = await readChatSnapshot(sessionId); + expect(result).toEqual(turn2); + expect(result?.messages).toHaveLength(4); + expect(result?.lastOutEventId).toBe("evt-turn2"); + }); + + postgresAndMinioTest("isolates snapshots by sessionId (no cross-talk)", async ({ minioConfig }) => { + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + stubApiClient({ projectRef: "proj_snap_iso", envSlug: "dev" }); + + const sessA = "sess_iso_A"; + const sessB = "sess_iso_B"; + const snapA = makeSnapshot({ lastOutEventId: "evt-A" }); + const snapB = makeSnapshot({ lastOutEventId: "evt-B" }); + + await writeChatSnapshot(sessA, snapA); + await writeChatSnapshot(sessB, snapB); + + const readA = await readChatSnapshot(sessA); + const readB = await readChatSnapshot(sessB); + + expect(readA?.lastOutEventId).toBe("evt-A"); + expect(readB?.lastOutEventId).toBe("evt-B"); + // Distinct objects — modifying one shouldn't affect the other. + expect(readA?.lastOutEventId).not.toBe(readB?.lastOutEventId); + }); + + postgresAndMinioTest("handles snapshots with large message lists (~50 messages)", async ({ minioConfig }) => { + // Stress test: a 50-turn chat snapshot. Plan F.4 mentions the + // pre-change baseline grew past 512 KiB around turn 10-30 with tool + // use; the post-slim wire keeps wire payloads small but the snapshot + // itself can still get large. Verify the helpers handle a realistic + // payload size. + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + stubApiClient({ projectRef: "proj_snap_big", envSlug: "dev" }); + + const messages: UIMessage[] = []; + for (let i = 0; i < 50; i++) { + messages.push({ + id: `u-${i}`, + role: "user", + parts: [{ type: "text", text: `user message ${i}: ${"x".repeat(200)}` }], + }); + messages.push({ + id: `a-${i}`, + role: "assistant", + parts: [{ type: "text", text: `assistant reply ${i}: ${"y".repeat(500)}` }], + }); + } + const snapshot = makeSnapshot({ messages, lastOutEventId: "evt-50" }); + + await writeChatSnapshot("sess_big_chat", snapshot); + const result = await readChatSnapshot("sess_big_chat"); + + expect(result).toBeDefined(); + expect(result!.messages).toHaveLength(100); + expect(result!.lastOutEventId).toBe("evt-50"); + // Spot-check ordering integrity — the messages array round-tripped + // in the same order. + expect(result!.messages[0]!.id).toBe("u-0"); + expect(result!.messages[99]!.id).toBe("a-49"); + }); +}); diff --git a/apps/webapp/test/replay-after-crash.test.ts b/apps/webapp/test/replay-after-crash.test.ts new file mode 100644 index 00000000000..f5c6842b194 --- /dev/null +++ b/apps/webapp/test/replay-after-crash.test.ts @@ -0,0 +1,315 @@ +// Plan F.3: integration test for the crash-recovery boot path. The +// scenario it locks down: +// +// 1. Run A streams chunks to `session.out` and `onTurnComplete` fires. +// 2. Run A crashes BEFORE `writeChatSnapshot` lands the post-turn +// blob (or the write fails silently — both have the same effect). +// 3. Run B boots: `readChatSnapshot` returns `undefined` (no snapshot +// yet, or stale-from-prior-turn). Replay then drains +// `session.out` from the snapshot's `lastOutEventId` (or seq 0) +// and reduces the chunks back into UIMessage[]. +// 4. The accumulator is consistent — Run A's completed chunks reach +// Run B's run loop without losing data. +// +// Plan section H.1 / H.4 spell out the "snapshot didn't make it before +// crash" path; this test is the integration safety net behind the +// unit tests in `packages/trigger-sdk/test/replay-session-out.test.ts`. +// +// We exercise the SDK's `__replaySessionOutTailProductionPathForTests` +// against a stubbed `apiClient.readSessionStreamRecords` — the new +// non-SSE records endpoint introduced in plan task #22. The replay path +// is a single GET that returns whatever's already on the stream; no +// long-poll. MinIO is provisioned to keep parity with +// `chat-snapshot-integration.test.ts` (the snapshot read path runs +// through it), even though the replay path itself doesn't read from S3. + +import { postgresAndMinioTest } from "@internal/testcontainers"; +import { apiClientManager } from "@trigger.dev/core/v3"; +import { + __readChatSnapshotProductionPathForTests as readChatSnapshot, + __replaySessionOutTailProductionPathForTests as replaySessionOutTail, + type ChatSnapshotV1, +} from "@trigger.dev/sdk/ai"; +import type { UIMessageChunk } from "ai"; +import { afterEach, describe, expect, vi } from "vitest"; +import { env } from "~/env.server"; +import { generatePresignedUrl } from "~/v3/objectStore.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +// ── Helpers ──────────────────────────────────────────────────────────── + +function textTurn(id: string, text: string): UIMessageChunk[] { + return [ + { type: "start", messageId: id, messageMetadata: { role: "assistant" } } as UIMessageChunk, + { type: "text-start", id: `${id}.t1` } as UIMessageChunk, + { type: "text-delta", id: `${id}.t1`, delta: text } as UIMessageChunk, + { type: "text-end", id: `${id}.t1` } as UIMessageChunk, + { type: "finish" } as UIMessageChunk, + ]; +} + +/** + * Stub `apiClientManager.clientOrThrow()` so: + * - `getPayloadUrl` / `createUploadPayloadUrl` mint MinIO presigned URLs + * via the webapp's real `generatePresignedUrl` (so snapshot reads + * hit a real S3-compatible backend). + * - `readSessionStreamRecords` returns the canonical + * `{ records: [{ data, id, seqNum }] }` shape — `data` is the + * JSON-encoded chunk body, mirroring the webapp's S2 record shape. + */ +function stubApiClient(opts: { + projectRef: string; + envSlug: string; + sessionOutChunks: unknown[]; +}) { + const records = opts.sessionOutChunks.map((chunk, i) => ({ + data: typeof chunk === "string" ? chunk : JSON.stringify(chunk), + id: `evt-${i + 1}`, + seqNum: i + 1, + })); + const readRecordsSpy = vi.fn( + async (_id: string, _io: "in" | "out", _options?: { afterEventId?: string }) => ({ + records, + }) + ); + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + async getPayloadUrl(filename: string) { + const result = await generatePresignedUrl(opts.projectRef, opts.envSlug, filename, "GET"); + if (!result.success) throw new Error(result.error); + return { presignedUrl: result.url }; + }, + async createUploadPayloadUrl(filename: string) { + const result = await generatePresignedUrl(opts.projectRef, opts.envSlug, filename, "PUT"); + if (!result.success) throw new Error(result.error); + return { presignedUrl: result.url }; + }, + readSessionStreamRecords: readRecordsSpy, + } as never); + return readRecordsSpy; +} + +let warnSpy: ReturnType; + +afterEach(() => { + vi.restoreAllMocks(); + warnSpy?.mockRestore(); +}); + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("replay after crash (MinIO + SDK helpers)", () => { + postgresAndMinioTest( + "boot reconstructs accumulator from session.out replay when no snapshot exists", + async ({ minioConfig }) => { + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + + // The crashed run's session.out: two completed assistant turns, no + // snapshot ever written. Boot must recover both via replay. + const chunks = [...textTurn("a-1", "first turn"), ...textTurn("a-2", "second turn")]; + stubApiClient({ + projectRef: "proj_replay_crash", + envSlug: "dev", + sessionOutChunks: chunks, + }); + + // Step 1: read snapshot — returns undefined (fresh boot, no snap). + const snapshot = await readChatSnapshot("sess_no_snap"); + expect(snapshot).toBeUndefined(); + + // Step 2: replay tail. + const replayed = await replaySessionOutTail("sess_no_snap"); + + expect(replayed).toHaveLength(2); + expect(replayed.map((m) => m.id)).toEqual(["a-1", "a-2"]); + const texts = replayed.flatMap((m) => + (m.parts as Array<{ type: string; text?: string }>) + .filter((p) => p.type === "text") + .map((p) => p.text) + ); + expect(texts).toEqual(["first turn", "second turn"]); + } + ); + + postgresAndMinioTest( + "boot replays only chunks AFTER snapshot.lastOutEventId (resume cursor)", + async ({ minioConfig }) => { + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + // The replay helper accepts the snapshot's `lastEventId` cursor + // and forwards it as `afterEventId` on the records endpoint — + // that's the cursor field name on the new non-SSE route. Here we + // feed only the post-snapshot chunks (modeling what the server + // returns for `afterEventId=evt-snapped`) and verify the helper + // threads the cursor through. + const readRecordsSpy = stubApiClient({ + projectRef: "proj_replay_resume", + envSlug: "dev", + sessionOutChunks: textTurn("a-after-snap", "post-snapshot turn"), + }); + + const result = await replaySessionOutTail("sess_resume", { lastEventId: "evt-snapped" }); + + expect(readRecordsSpy).toHaveBeenCalledWith( + "sess_resume", + "out", + expect.objectContaining({ afterEventId: "evt-snapped" }) + ); + expect(result).toHaveLength(1); + expect(result[0]!.id).toBe("a-after-snap"); + } + ); + + postgresAndMinioTest( + "boot returns [] when session.out is empty (first-ever turn, no snapshot)", + async ({ minioConfig }) => { + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + + stubApiClient({ + projectRef: "proj_replay_empty", + envSlug: "dev", + sessionOutChunks: [], + }); + + const snapshot = await readChatSnapshot("sess_empty"); + expect(snapshot).toBeUndefined(); + + const replayed = await replaySessionOutTail("sess_empty"); + expect(replayed).toEqual([]); + } + ); + + postgresAndMinioTest( + "boot drops orphaned trailing tool parts (cleanupAbortedParts) — partial crash", + async ({ minioConfig }) => { + // Simulates a true mid-turn crash: assistant finished one turn, + // then started a tool-call but the run died before resolution. + // Replay must surface the completed turn but NOT include the + // orphaned tool part in `input-streaming` state. + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + stubApiClient({ + projectRef: "proj_replay_partial", + envSlug: "dev", + sessionOutChunks: [ + ...textTurn("a-complete", "I finished step 1"), + // Partial tool turn — no tool-input-end, no finish. + { type: "start", messageId: "a-orphan", messageMetadata: { role: "assistant" } } as UIMessageChunk, + { type: "tool-input-start", id: "tc-cut", toolName: "search" } as UIMessageChunk, + { type: "tool-input-delta", id: "tc-cut", delta: '{"q":"x"}' } as UIMessageChunk, + ], + }); + + const replayed = await replaySessionOutTail("sess_partial_crash"); + + // Completed turn always present. + expect(replayed.find((m) => m.id === "a-complete")).toBeTruthy(); + // Orphaned tool-call never surfaces in `input-streaming` state. + const orphan = replayed.find((m) => m.id === "a-orphan"); + if (orphan) { + const stillStreaming = (orphan.parts as Array<{ toolCallId?: string; state?: string }>).find( + (p) => p.toolCallId === "tc-cut" && p.state === "input-streaming" + ); + expect(stillStreaming).toBeUndefined(); + } + } + ); + + postgresAndMinioTest( + "snapshot+replay merge: snapshot supplies user msgs, replay supplies assistants", + async ({ minioConfig }) => { + // The boot orchestration calls + // `mergeByIdReplaceWins(snapshot.messages, replayed)`. The runtime + // contract is that user messages live in snapshot only (session.in + // never goes through replay) and assistants come from replay + // (which carries the freshest representation). Here we simulate + // the realistic split: snapshot has [u-1, a-1-stale], replay has + // [a-1-fresh, a-2-new]. After merge the accumulator should reflect + // the fresh assistant + new assistant, with the user message + // preserved. + // + // Note: this is a pre-merge round-trip — we drive the read and + // replay through real MinIO + stubbed S2 to confirm both arrive + // intact for the orchestration to merge. + env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl; + env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId; + env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey; + env.OBJECT_STORE_REGION = minioConfig.region; + env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined; + + // Pre-write a snapshot to MinIO via real apiClient stub. + const sessionId = "sess_merge_round_trip"; + const snapshot: ChatSnapshotV1 = { + version: 1, + savedAt: 1_700_000_000_000, + messages: [ + { id: "u-1", role: "user", parts: [{ type: "text", text: "hi" }] }, + { id: "a-1", role: "assistant", parts: [{ type: "text", text: "stale-assistant" }] }, + ], + lastOutEventId: "evt-prev", + lastOutTimestamp: 1_700_000_000_500, + }; + + // Use the SDK's own writer to lay the snapshot down, then swap + // the stub to also serve replay chunks for the read path. + stubApiClient({ + projectRef: "proj_merge", + envSlug: "dev", + sessionOutChunks: [], + }); + const { __writeChatSnapshotProductionPathForTests: writeSnapshot } = await import( + "@trigger.dev/sdk/ai" + ); + await writeSnapshot(sessionId, snapshot); + + // Restubbing for the boot phase: replay tail carries the fresh + // assistant for `a-1` plus a brand-new `a-2`. The orchestration's + // merge would replace `a-1` and append `a-2` after `u-1`. + vi.restoreAllMocks(); + stubApiClient({ + projectRef: "proj_merge", + envSlug: "dev", + sessionOutChunks: [ + ...textTurn("a-1", "fresh-assistant"), + ...textTurn("a-2", "next-assistant"), + ], + }); + + const readBack = await readChatSnapshot(sessionId); + expect(readBack?.messages.map((m) => m.id)).toEqual(["u-1", "a-1"]); + + const replayed = await replaySessionOutTail(sessionId, { + lastEventId: readBack?.lastOutEventId, + }); + expect(replayed.map((m) => m.id)).toEqual(["a-1", "a-2"]); + // Replay's `a-1` carries the fresh content — when merge runs in + // the runtime, this version would replace the snapshot's stale + // `a-1`. + const replayedA1Text = (replayed[0]!.parts as Array<{ type: string; text?: string }>) + .filter((p) => p.type === "text") + .map((p) => p.text) + .join(""); + expect(replayedA1Text).toBe("fresh-assistant"); + } + ); +}); diff --git a/packages/build/package.json b/packages/build/package.json index 49a310e46e7..f172eeb7c6a 100644 --- a/packages/build/package.json +++ b/packages/build/package.json @@ -31,7 +31,8 @@ "./extensions/typescript": "./src/extensions/typescript.ts", "./extensions/puppeteer": "./src/extensions/puppeteer.ts", "./extensions/playwright": "./src/extensions/playwright.ts", - "./extensions/lightpanda": "./src/extensions/lightpanda.ts" + "./extensions/lightpanda": "./src/extensions/lightpanda.ts", + "./extensions/secureExec": "./src/extensions/secureExec.ts" }, "sourceDialects": [ "@triggerdotdev/source" @@ -65,6 +66,9 @@ ], "extensions/lightpanda": [ "dist/commonjs/extensions/lightpanda.d.ts" + ], + "extensions/secureExec": [ + "dist/commonjs/extensions/secureExec.d.ts" ] } }, @@ -207,6 +211,17 @@ "types": "./dist/commonjs/extensions/lightpanda.d.ts", "default": "./dist/commonjs/extensions/lightpanda.js" } + }, + "./extensions/secureExec": { + "import": { + "@triggerdotdev/source": "./src/extensions/secureExec.ts", + "types": "./dist/esm/extensions/secureExec.d.ts", + "default": "./dist/esm/extensions/secureExec.js" + }, + "require": { + "types": "./dist/commonjs/extensions/secureExec.d.ts", + "default": "./dist/commonjs/extensions/secureExec.js" + } } }, "main": "./dist/commonjs/index.js", diff --git a/packages/build/src/extensions/secureExec.ts b/packages/build/src/extensions/secureExec.ts new file mode 100644 index 00000000000..808bc666501 --- /dev/null +++ b/packages/build/src/extensions/secureExec.ts @@ -0,0 +1,172 @@ +import { BuildTarget } from "@trigger.dev/core/v3"; +import { BuildManifest } from "@trigger.dev/core/v3/schemas"; +import { BuildContext, BuildExtension } from "@trigger.dev/core/v3/build"; +import { dirname, resolve, join } from "node:path"; +import { readFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { readPackageJSON } from "pkg-types"; + +export type SecureExecOptions = { + /** + * Packages available inside the sandbox at runtime. + * + * These are `require()`'d inside the V8 isolate at runtime — the bundler + * never sees them statically. They are marked external and installed as + * deploy dependencies. + * + * @example + * ```ts + * secureExec({ packages: ["jszip", "lodash"] }) + * ``` + */ + packages?: string[]; +}; + +/** + * Build extension for [secure-exec](https://secureexec.dev) — run untrusted + * JavaScript/TypeScript in V8 isolates with configurable permissions. + * + * Handles the esbuild workarounds needed for secure-exec's runtime + * `require.resolve` calls, native binaries, and module-scope resolution. + * + * @example + * ```ts + * import { secureExec } from "@trigger.dev/build/extensions/secureExec"; + * + * export default defineConfig({ + * build: { + * extensions: [secureExec()], + * }, + * }); + * ``` + */ +export function secureExec(options?: SecureExecOptions): BuildExtension { + return new SecureExecExtension(options ?? {}); +} + +class SecureExecExtension implements BuildExtension { + public readonly name = "SecureExecExtension"; + + private userPackages: string[]; + + constructor(options: SecureExecOptions) { + this.userPackages = options.packages ?? []; + } + + externalsForTarget(_target: BuildTarget) { + return [ + // esbuild must not be bundled — it locates its native binary via a + // relative path from its JS API entry point. secure-exec uses esbuild + // at runtime to bundle polyfills for sandbox code. + "esbuild", + // User-specified packages are require()'d inside the V8 sandbox at + // runtime — the bundler never sees them statically. + ...this.userPackages, + ]; + } + + onBuildStart(context: BuildContext) { + context.logger.debug(`Adding ${this.name} esbuild plugins`); + + // Plugin 1: Replace node-stdlib-browser with pre-resolved paths. + // + // Trigger's ESM shim anchors require.resolve() to the chunk path, so + // node-stdlib-browser's runtime require.resolve("./mock/empty.js") breaks. + // Fix: load the real node-stdlib-browser at build time (where require.resolve + // works), capture the resolved path map, and inline it as a static export. + const workingDir = context.workingDir; + context.registerPlugin({ + name: "secure-exec-stdlib-resolver", + setup(build) { + build.onResolve({ filter: /^node-stdlib-browser$/ }, () => ({ + path: "node-stdlib-browser", + namespace: "secure-exec-nsb-resolved", + })); + build.onLoad({ filter: /.*/, namespace: "secure-exec-nsb-resolved" }, () => { + const buildRequire = createRequire(join(workingDir, "package.json")); + const resolved = buildRequire("node-stdlib-browser"); + return { + contents: `export default ${JSON.stringify(resolved)};`, + loader: "js", + }; + }); + }, + }); + + // Plugin 2: Inline bridge.js at build time. + // + // bridge-loader.js in @secure-exec/node(js) uses __dirname and + // require.resolve("@secure-exec/core") at module scope to locate + // dist/bridge.js on disk. This fails in Trigger's bundled output. + // Fix: read bridge.js content at build time and inline it as a + // string literal so no runtime filesystem resolution is needed. + // + context.registerPlugin({ + name: "secure-exec-bridge-inline", + setup(build) { + build.onLoad( + { filter: /[\\/]@secure-exec[\\/]node[\\/]dist[\\/]bridge-loader\.js$/ }, + (args) => { + try { + const buildRequire = createRequire(args.path); + const coreEntry = buildRequire.resolve("@secure-exec/core"); + const coreRoot = resolve(dirname(coreEntry), ".."); + const bridgeCode = readFileSync(join(coreRoot, "dist", "bridge.js"), "utf8"); + + return { + contents: [ + `import { getIsolateRuntimeSource } from "@secure-exec/core";`, + `const bridgeCodeCache = ${JSON.stringify(bridgeCode)};`, + `export function getRawBridgeCode() { return bridgeCodeCache; }`, + `export function getBridgeAttachCode() { return getIsolateRuntimeSource("bridgeAttach"); }`, + ].join("\n"), + loader: "js", + }; + } catch { + // If we can't inline the bridge, let the normal loader handle it. + return undefined; + } + } + ); + }, + }); + } + + async onBuildComplete(context: BuildContext, _manifest: BuildManifest) { + if (context.target === "dev") { + return; + } + + context.logger.debug(`Adding ${this.name} deploy dependencies`); + + const dependencies: Record = {}; + + // Resolve versions for user-specified sandbox packages + for (const pkg of this.userPackages) { + try { + const modulePath = await context.resolvePath(pkg); + if (!modulePath) { + dependencies[pkg] = "latest"; + continue; + } + + const packageJSON = await readPackageJSON(dirname(modulePath)); + dependencies[pkg] = packageJSON.version ?? "latest"; + } catch { + context.logger.warn( + `Could not resolve version for sandbox package ${pkg}, defaulting to latest` + ); + dependencies[pkg] = "latest"; + } + } + + context.addLayer({ + id: "secureExec", + dependencies, + image: { + // isolated-vm requires native compilation tools + pkgs: ["python3", "make", "g++"], + }, + }); + } +} diff --git a/packages/build/src/internal.ts b/packages/build/src/internal.ts index 54f785a6106..0e1954c8b9e 100644 --- a/packages/build/src/internal.ts +++ b/packages/build/src/internal.ts @@ -1 +1,2 @@ export * from "./internal/additionalFiles.js"; +export * from "./internal/copyFiles.js"; diff --git a/packages/build/src/internal/additionalFiles.ts b/packages/build/src/internal/additionalFiles.ts index a815b53c9aa..57a746c36b6 100644 --- a/packages/build/src/internal/additionalFiles.ts +++ b/packages/build/src/internal/additionalFiles.ts @@ -1,8 +1,10 @@ import { BuildManifest } from "@trigger.dev/core/v3"; import { BuildContext } from "@trigger.dev/core/v3/build"; -import { copyFile, mkdir } from "node:fs/promises"; -import { dirname, join, posix, relative } from "node:path"; -import { glob } from "tinyglobby"; +import { + copyMatcherResults, + findFilesByMatchers, + type MatcherResult, +} from "./copyFiles.js"; export type AdditionalFilesOptions = { files: string[]; @@ -14,12 +16,13 @@ export async function addAdditionalFilesToBuild( context: BuildContext, manifest: BuildManifest ) { - // Copy any static assets to the destination - const staticAssets = await findStaticAssetFiles(options.files ?? [], manifest.outputPath, { - cwd: context.workingDir, - }); + const matcherResults: MatcherResult[] = await findFilesByMatchers( + options.files ?? [], + manifest.outputPath, + { cwd: context.workingDir } + ); - for (const { assets, matcher } of staticAssets) { + for (const { assets, matcher } of matcherResults) { if (assets.length === 0) { context.logger.warn(`[${source}] No files found for matcher`, matcher); } else { @@ -27,80 +30,7 @@ export async function addAdditionalFilesToBuild( } } - await copyStaticAssets(staticAssets, source, context); -} - -type MatchedStaticAssets = { source: string; destination: string }[]; - -type FoundStaticAssetFiles = Array<{ - matcher: string; - assets: MatchedStaticAssets; -}>; - -async function findStaticAssetFiles( - matchers: string[], - destinationPath: string, - options?: { cwd?: string; ignore?: string[] } -): Promise { - const result: FoundStaticAssetFiles = []; - - for (const matcher of matchers) { - const assets = await findStaticAssetsForMatcher(matcher, destinationPath, options); - - result.push({ matcher, assets }); - } - - return result; -} - -async function findStaticAssetsForMatcher( - matcher: string, - destinationPath: string, - options?: { cwd?: string; ignore?: string[] } -): Promise { - const result: MatchedStaticAssets = []; - - const files = await glob({ - patterns: [matcher], - cwd: options?.cwd, - ignore: options?.ignore ?? [], - onlyFiles: true, - absolute: true, + await copyMatcherResults(matcherResults, (pair) => { + context.logger.debug(`[${source}] Copying ${pair.source} to ${pair.destination}`); }); - - let matches = 0; - - for (const file of files) { - matches++; - - const pathInsideDestinationDir = relative(options?.cwd ?? process.cwd(), file) - .split(posix.sep) - .filter((p) => p !== "..") - .join(posix.sep); - - const relativeDestinationPath = join(destinationPath, pathInsideDestinationDir); - - result.push({ - source: file, - destination: relativeDestinationPath, - }); - } - - return result; -} - -async function copyStaticAssets( - staticAssetFiles: FoundStaticAssetFiles, - sourceName: string, - context: BuildContext -): Promise { - for (const { assets } of staticAssetFiles) { - for (const { source, destination } of assets) { - await mkdir(dirname(destination), { recursive: true }); - - context.logger.debug(`[${sourceName}] Copying ${source} to ${destination}`); - - await copyFile(source, destination); - } - } } diff --git a/packages/build/src/internal/copyFiles.ts b/packages/build/src/internal/copyFiles.ts new file mode 100644 index 00000000000..6fd3ede9545 --- /dev/null +++ b/packages/build/src/internal/copyFiles.ts @@ -0,0 +1,99 @@ +import { cp, copyFile, mkdir } from "node:fs/promises"; +import { dirname, join, posix, relative } from "node:path"; +import { glob } from "tinyglobby"; + +/** + * A single matched asset — source file and its destination inside the + * build output directory. + */ +export type CopyPair = { source: string; destination: string }; + +/** + * Result of a single matcher's glob, grouped with the matcher that + * produced it so callers can warn on empty matches. + */ +export type MatcherResult = { + matcher: string; + assets: CopyPair[]; +}; + +/** + * Glob a set of matchers relative to `cwd` and return pairs describing + * where each matched file should be copied to under `destinationDir`. + * + * Relative paths are preserved under `destinationDir`. Leading `..` + * segments (from `../shared/file.txt` style patterns) are stripped so + * files always land inside the destination. + */ +export async function findFilesByMatchers( + matchers: string[], + destinationDir: string, + options?: { cwd?: string; ignore?: string[] } +): Promise { + const result: MatcherResult[] = []; + const cwd = options?.cwd ?? process.cwd(); + + for (const matcher of matchers) { + const files = await glob({ + patterns: [matcher], + cwd, + ignore: options?.ignore ?? [], + onlyFiles: true, + absolute: true, + }); + + const assets: CopyPair[] = files.map((file) => { + const pathInsideDestinationDir = relative(cwd, file) + .split(posix.sep) + .filter((p) => p !== "..") + .join(posix.sep); + return { + source: file, + destination: join(destinationDir, pathInsideDestinationDir), + }; + }); + + result.push({ matcher, assets }); + } + + return result; +} + +/** + * Copy a single file, creating parent directories as needed. + */ +export async function copyFileEnsuringDir(source: string, destination: string): Promise { + await mkdir(dirname(destination), { recursive: true }); + await copyFile(source, destination); +} + +/** + * Copy every pair in the given matcher results. Parent directories are + * created automatically. Returns the total number of files copied. + */ +export async function copyMatcherResults( + matcherResults: MatcherResult[], + onCopy?: (pair: CopyPair) => void +): Promise { + let count = 0; + for (const { assets } of matcherResults) { + for (const pair of assets) { + onCopy?.(pair); + await copyFileEnsuringDir(pair.source, pair.destination); + count++; + } + } + return count; +} + +/** + * Recursively copy a directory to another location. Preserves structure; + * overwrites existing files at the destination. + * + * Used by the built-in skill bundler — we copy entire skill folders as a + * unit, not file-by-file. + */ +export async function copyDirectoryRecursive(source: string, destination: string): Promise { + await mkdir(destination, { recursive: true }); + await cp(source, destination, { recursive: true, force: true }); +} diff --git a/packages/core/package.json b/packages/core/package.json index f58708dff92..34acdb91b89 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -41,6 +41,8 @@ "./v3/utils/omit": "./src/v3/utils/omit.ts", "./v3/utils/retries": "./src/v3/utils/retries.ts", "./v3/utils/structuredLogger": "./src/v3/utils/structuredLogger.ts", + "./v3/chat-client": "./src/v3/chat-client.ts", + "./v3/test": "./src/v3/test/index.ts", "./v3/zodfetch": "./src/v3/zodfetch.ts", "./v3/zodMessageHandler": "./src/v3/zodMessageHandler.ts", "./v3/zodNamespace": "./src/v3/zodNamespace.ts", @@ -87,6 +89,9 @@ "v3/errors": [ "dist/commonjs/v3/errors.d.ts" ], + "v3/chat-client": [ + "dist/commonjs/v3/chat-client.d.ts" + ], "v3/logger-api": [ "dist/commonjs/v3/logger-api.d.ts" ], @@ -152,6 +157,9 @@ ], "v3/isomorphic": [ "dist/commonjs/v3/isomorphic/index.d.ts" + ], + "v3/test": [ + "dist/commonjs/v3/test/index.d.ts" ] } }, @@ -446,6 +454,28 @@ "default": "./dist/commonjs/v3/utils/structuredLogger.js" } }, + "./v3/chat-client": { + "import": { + "@triggerdotdev/source": "./src/v3/chat-client.ts", + "types": "./dist/esm/v3/chat-client.d.ts", + "default": "./dist/esm/v3/chat-client.js" + }, + "require": { + "types": "./dist/commonjs/v3/chat-client.d.ts", + "default": "./dist/commonjs/v3/chat-client.js" + } + }, + "./v3/test": { + "import": { + "@triggerdotdev/source": "./src/v3/test/index.ts", + "types": "./dist/esm/v3/test/index.d.ts", + "default": "./dist/esm/v3/test/index.js" + }, + "require": { + "types": "./dist/commonjs/v3/test/index.d.ts", + "default": "./dist/commonjs/v3/test/index.js" + } + }, "./v3/zodfetch": { "import": { "@triggerdotdev/source": "./src/v3/zodfetch.ts", diff --git a/packages/core/src/v3/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts index 89a551954f0..04a9009e356 100644 --- a/packages/core/src/v3/apiClient/index.ts +++ b/packages/core/src/v3/apiClient/index.ts @@ -6,19 +6,32 @@ import { ApiDeploymentListOptions, ApiDeploymentListResponseItem, ApiDeploymentListSearchParams, + RetrieveCurrentDeploymentResponseBody, AppendToStreamResponseBody, BatchItemNDJSON, BatchTaskRunExecutionResult, BatchTriggerTaskV3RequestBody, BatchTriggerTaskV3Response, CanceledRunResponse, + CloseSessionRequestBody, CompleteWaitpointTokenRequestBody, CompleteWaitpointTokenResponseBody, + CreatedSessionResponseBody, + CreateSessionRequestBody, + EndAndContinueSessionRequestBody, + EndAndContinueSessionResponseBody, + ListSessionsOptions, + ListSessionsResponseBody, + ListedSessionItem, + RetrieveSessionResponseBody, + UpdateSessionRequestBody, CreateBatchRequestBody, CreateBatchResponse, CreateEnvironmentVariableRequestBody, CreateInputStreamWaitpointRequestBody, CreateInputStreamWaitpointResponseBody, + CreateSessionStreamWaitpointRequestBody, + CreateSessionStreamWaitpointResponseBody, CreateScheduleOptions, CreateStreamResponseBody, CreateUploadPayloadUrlResponseBody, @@ -59,6 +72,7 @@ import { SendInputStreamResponseBody, StreamBatchItemsResponse, TaskRunExecutionResult, + ReadSessionStreamRecordsResponseBody, TriggerTaskRequestBody, TriggerTaskResponse, UpdateEnvironmentVariableRequestBody, @@ -1094,6 +1108,233 @@ export class ApiClient { ); } + // ======================================================================== + // Sessions + // ======================================================================== + + createSession(body: CreateSessionRequestBody, requestOptions?: ZodFetchOptions) { + return zodfetch( + CreatedSessionResponseBody, + `${this.baseUrl}/api/v1/sessions`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify(body), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + retrieveSession(sessionIdOrExternalId: string, requestOptions?: ZodFetchOptions) { + return zodfetch( + RetrieveSessionResponseBody, + `${this.baseUrl}/api/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}`, + { + method: "GET", + headers: this.#getHeaders(false), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + updateSession( + sessionIdOrExternalId: string, + body: UpdateSessionRequestBody, + requestOptions?: ZodFetchOptions + ) { + return zodfetch( + RetrieveSessionResponseBody, + `${this.baseUrl}/api/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}`, + { + method: "PATCH", + headers: this.#getHeaders(false), + body: JSON.stringify(body), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + closeSession( + sessionIdOrExternalId: string, + body?: CloseSessionRequestBody, + requestOptions?: ZodFetchOptions + ) { + return zodfetch( + RetrieveSessionResponseBody, + `${this.baseUrl}/api/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}/close`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify(body ?? {}), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + endAndContinueSession( + sessionIdOrExternalId: string, + body: EndAndContinueSessionRequestBody, + requestOptions?: ZodFetchOptions + ) { + return zodfetch( + EndAndContinueSessionResponseBody, + `${this.baseUrl}/api/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}/end-and-continue`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify(body), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + listSessions( + options?: ListSessionsOptions, + requestOptions?: ZodFetchOptions + ): CursorPagePromise { + const searchParams = createSearchQueryForListSessions(options); + + return zodfetchCursorPage( + ListedSessionItem, + `${this.baseUrl}/api/v1/sessions`, + { + query: searchParams, + limit: options?.limit, + after: options?.after, + before: options?.before, + }, + { + method: "GET", + headers: this.#getHeaders(false), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + // ======================================================================== + // Session realtime channels + // ======================================================================== + + async initializeSessionStream( + sessionIdOrExternalId: string, + io: "out" | "in", + requestOptions?: ZodFetchOptions + ) { + // The server returns S2 credentials in response headers alongside a tiny + // JSON body with the realtime version. Follow the same shape as + // `createStream` so downstream clients can feed them into + // `StreamsWriterV2`. + return zodfetch( + CreateStreamResponseBody, + `${this.baseUrl}/realtime/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}/${io}`, + { + method: "PUT", + headers: this.#getHeaders(false), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ) + .withResponse() + .then(({ data, response }) => ({ + ...data, + headers: Object.fromEntries(response.headers.entries()), + })); + } + + async appendToSessionStream( + sessionIdOrExternalId: string, + io: "out" | "in", + part: TBody, + requestOptions?: ZodFetchOptions + ) { + return zodfetch( + AppendToStreamResponseBody, + `${this.baseUrl}/realtime/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}/${io}/append`, + { + method: "POST", + headers: this.#getHeaders(false), + body: part, + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + /** + * Non-SSE drain of a Session channel's tail. Returns whatever records + * exist after `afterEventId` (or from the head of the stream) and closes + * — `wait=0` semantics, no long-poll. Used by `replaySessionOutTail` at + * run boot, where the SSE long-poll's ~1s tax on empty streams is the + * dominant cost on every fresh chat. + * + * `afterEventId` is the same cursor format as the SSE Last-Event-ID + * (the S2 sequence number, stringified) — pass `lastOutEventId` from a + * persisted snapshot to resume. + */ + async readSessionStreamRecords( + sessionIdOrExternalId: string, + io: "out" | "in", + options?: { afterEventId?: string; baseUrl?: string } + ) { + const qs = new URLSearchParams(); + if (options?.afterEventId !== undefined) { + qs.set("afterEventId", options.afterEventId); + } + const url = `${options?.baseUrl ?? this.baseUrl}/realtime/v1/sessions/${encodeURIComponent( + sessionIdOrExternalId + )}/${io}/records${qs.toString() ? `?${qs.toString()}` : ""}`; + return zodfetch( + ReadSessionStreamRecordsResponseBody, + url, + { + method: "GET", + headers: this.#getHeaders(false), + }, + mergeRequestOptions(this.defaultRequestOptions, undefined) + ); + } + + /** + * Subscribe to SSE records on a Session channel. Reuses the same + * {@link SSEStreamSubscription} plumbing as `readStream` for run-scoped + * realtime streams — auto-retry, Last-Event-ID resume, abort-on-cancel. + */ + async subscribeToSessionStream( + sessionIdOrExternalId: string, + io: "out" | "in", + options?: { + signal?: AbortSignal; + baseUrl?: string; + timeoutInSeconds?: number; + onComplete?: () => void; + onError?: (error: Error) => void; + lastEventId?: string; + onPart?: (part: SSEStreamPart) => void; + } + ): Promise> { + const url = `${options?.baseUrl ?? this.baseUrl}/realtime/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}/${io}`; + + const subscription = new SSEStreamSubscription(url, { + headers: this.getHeaders(), + signal: options?.signal, + onComplete: options?.onComplete, + onError: options?.onError, + timeoutInSeconds: options?.timeoutInSeconds, + lastEventId: options?.lastEventId, + }); + + const stream = await subscription.subscribe(); + const onPart = options?.onPart; + + return stream.pipeThrough( + new TransformStream({ + transform(chunk, controller) { + const data = chunk.chunk as T; + onPart?.(chunk as SSEStreamPart); + controller.enqueue(data); + }, + }) + ); + } + async waitForDuration( runId: string, body: WaitForDurationRequestBody, @@ -1340,6 +1581,18 @@ export class ApiClient { ); } + retrieveCurrentDeployment(requestOptions?: ZodFetchOptions) { + return zodfetch( + RetrieveCurrentDeploymentResponseBody, + `${this.baseUrl}/api/v1/deployments/current`, + { + method: "GET", + headers: this.#getHeaders(false), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + async fetchStream( runId: string, streamKey: string, @@ -1459,6 +1712,23 @@ export class ApiClient { ); } + async createSessionStreamWaitpoint( + runFriendlyId: string, + body: CreateSessionStreamWaitpointRequestBody, + requestOptions?: ZodFetchOptions + ) { + return zodfetch( + CreateSessionStreamWaitpointResponseBody, + `${this.baseUrl}/api/v1/runs/${runFriendlyId}/session-streams/wait`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify(body), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + async generateJWTClaims(requestOptions?: ZodFetchOptions): Promise> { return zodfetch( z.record(z.any()), @@ -1823,6 +2093,47 @@ function queueNameFromQueueTypeName(queue: QueueTypeName): string { return queue.name; } +function createSearchQueryForListSessions(options?: ListSessionsOptions): URLSearchParams { + const searchParams = new URLSearchParams(); + + if (!options) return searchParams; + + const appendMany = (name: string, value: string | string[] | undefined) => { + if (value === undefined) return; + searchParams.append(name, Array.isArray(value) ? value.join(",") : value); + }; + + appendMany("filter[type]", options.type); + appendMany("filter[tags]", options.tag); + appendMany("filter[taskIdentifier]", options.taskIdentifier); + + if (options.externalId) { + searchParams.append("filter[externalId]", options.externalId); + } + + appendMany("filter[status]", options.status as string | string[] | undefined); + + if (options.period) { + searchParams.append("filter[createdAt][period]", options.period); + } + + if (options.from !== undefined) { + searchParams.append( + "filter[createdAt][from]", + options.from instanceof Date ? options.from.getTime().toString() : options.from.toString() + ); + } + + if (options.to !== undefined) { + searchParams.append( + "filter[createdAt][to]", + options.to instanceof Date ? options.to.getTime().toString() : options.to.toString() + ); + } + + return searchParams; +} + function createSearchQueryForListWaitpointTokens( query?: ListWaitpointTokensQueryParams ): URLSearchParams { diff --git a/packages/core/src/v3/errors.ts b/packages/core/src/v3/errors.ts index a538ca9357b..6ff8f84c525 100644 --- a/packages/core/src/v3/errors.ts +++ b/packages/core/src/v3/errors.ts @@ -631,6 +631,26 @@ export class GracefulExitTimeoutError extends Error { } } +export class ChatChunkTooLargeError extends Error { + constructor( + public readonly chunkSize: number, + public readonly maxSize: number, + public readonly chunkType?: string + ) { + super( + `chat.agent chunk${chunkType ? ` of type "${chunkType}"` : ""} is ${chunkSize} bytes, ` + + `over the realtime stream's per-record cap of ${maxSize} bytes. ` + + `For oversized payloads (e.g. large tool outputs), write the value to your own store and ` + + `emit only an id/url through the chat stream — see https://trigger.dev/docs/ai-chat/patterns/large-payloads.` + ); + this.name = "ChatChunkTooLargeError"; + } +} + +export function isChatChunkTooLargeError(error: unknown): error is ChatChunkTooLargeError { + return error instanceof Error && error.name === "ChatChunkTooLargeError"; +} + export class MaxDurationExceededError extends Error { constructor( public readonly maxDurationInSeconds: number, diff --git a/packages/core/src/v3/index.ts b/packages/core/src/v3/index.ts index 2757363f4be..72b91c46071 100644 --- a/packages/core/src/v3/index.ts +++ b/packages/core/src/v3/index.ts @@ -21,6 +21,7 @@ export * from "./locals-api.js"; export * from "./heartbeats-api.js"; export * from "./realtime-streams-api.js"; export * from "./input-streams-api.js"; +export * from "./session-streams-api.js"; export * from "./waitpoints/index.js"; export * from "./schemas/index.js"; export { SemanticInternalAttributes } from "./semanticInternalAttributes.js"; @@ -80,6 +81,7 @@ export { getSchemaParseFn, type AnySchemaParseFn, type SchemaParseFn, + type inferSchemaOut, isSchemaZodEsque, isSchemaValibotEsque, isSchemaArkTypeEsque, diff --git a/packages/core/src/v3/resource-catalog/catalog.ts b/packages/core/src/v3/resource-catalog/catalog.ts index 5b3ab023639..5c443b253cf 100644 --- a/packages/core/src/v3/resource-catalog/catalog.ts +++ b/packages/core/src/v3/resource-catalog/catalog.ts @@ -1,4 +1,11 @@ -import { PromptManifest, QueueManifest, TaskManifest, WorkerManifest } from "../schemas/index.js"; +import { + PromptManifest, + QueueManifest, + SkillManifest, + SkillMetadata, + TaskManifest, + WorkerManifest, +} from "../schemas/index.js"; import { PromptMetadataWithFunctions, TaskMetadataWithFunctions, TaskSchema } from "../types/index.js"; export interface ResourceCatalog { @@ -18,4 +25,7 @@ export interface ResourceCatalog { listPromptManifests(): Array; getPrompt(id: string): PromptMetadataWithFunctions | undefined; getPromptSchema(id: string): TaskSchema | undefined; + registerSkillMetadata(skill: SkillMetadata): void; + listSkillManifests(): Array; + getSkillManifest(id: string): SkillManifest | undefined; } diff --git a/packages/core/src/v3/resource-catalog/index.ts b/packages/core/src/v3/resource-catalog/index.ts index 9ce7dee64cf..f809ede8135 100644 --- a/packages/core/src/v3/resource-catalog/index.ts +++ b/packages/core/src/v3/resource-catalog/index.ts @@ -1,6 +1,13 @@ const API_NAME = "resource-catalog"; -import { PromptManifest, QueueManifest, TaskManifest, WorkerManifest } from "../schemas/index.js"; +import { + PromptManifest, + QueueManifest, + SkillManifest, + SkillMetadata, + TaskManifest, + WorkerManifest, +} from "../schemas/index.js"; import { PromptMetadataWithFunctions, TaskMetadataWithFunctions, TaskSchema } from "../types/index.js"; import { getGlobal, registerGlobal, unregisterGlobal } from "../utils/globals.js"; import { type ResourceCatalog } from "./catalog.js"; @@ -93,6 +100,18 @@ export class ResourceCatalogAPI { return this.#getCatalog().getPromptSchema(id); } + public registerSkillMetadata(skill: SkillMetadata): void { + this.#getCatalog().registerSkillMetadata(skill); + } + + public listSkillManifests(): Array { + return this.#getCatalog().listSkillManifests(); + } + + public getSkillManifest(id: string): SkillManifest | undefined { + return this.#getCatalog().getSkillManifest(id); + } + #getCatalog(): ResourceCatalog { return getGlobal(API_NAME) ?? NOOP_RESOURCE_CATALOG; } diff --git a/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts b/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts index 8f77544f05c..5da74d4a9b1 100644 --- a/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts +++ b/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts @@ -1,4 +1,11 @@ -import { PromptManifest, QueueManifest, TaskManifest, WorkerManifest } from "../schemas/index.js"; +import { + PromptManifest, + QueueManifest, + SkillManifest, + SkillMetadata, + TaskManifest, + WorkerManifest, +} from "../schemas/index.js"; import { type PromptMetadataWithFunctions, type TaskMetadataWithFunctions, type TaskSchema } from "../types/index.js"; import { ResourceCatalog } from "./catalog.js"; @@ -70,4 +77,16 @@ export class NoopResourceCatalog implements ResourceCatalog { getPromptSchema(id: string): TaskSchema | undefined { return undefined; } + + registerSkillMetadata(skill: SkillMetadata): void { + // noop + } + + listSkillManifests(): Array { + return []; + } + + getSkillManifest(id: string): SkillManifest | undefined { + return undefined; + } } diff --git a/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts b/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts index ea134a45663..0a67a4fd9a4 100644 --- a/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts +++ b/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts @@ -1,6 +1,8 @@ import { PromptManifest, PromptMetadata, + SkillManifest, + SkillMetadata, TaskFileMetadata, TaskMetadata, TaskManifest, @@ -21,6 +23,8 @@ export class StandardResourceCatalog implements ResourceCatalog { private _promptSchemas: Map = new Map(); private _currentFileContext?: Omit; private _queueMetadata: Map = new Map(); + private _skillMetadata: Map = new Map(); + private _skillFileMetadata: Map = new Map(); setCurrentFileContext(filePath: string, entryPoint: string) { this._currentFileContext = { filePath, entryPoint }; @@ -86,25 +90,31 @@ export class StandardResourceCatalog implements ResourceCatalog { } updateTaskMetadata(id: string, updates: Partial): void { + const { fns, schema, ...metadataUpdates } = updates; + const existingMetadata = this._taskMetadata.get(id); - if (existingMetadata) { + if (existingMetadata && Object.keys(metadataUpdates).length > 0) { this._taskMetadata.set(id, { ...existingMetadata, - ...updates, + ...metadataUpdates, }); } - if (updates.fns) { + if (fns) { const existingFunctions = this._taskFunctions.get(id); if (existingFunctions) { this._taskFunctions.set(id, { ...existingFunctions, - ...updates.fns, + ...fns, }); } } + + if (schema) { + this._taskSchemas.set(id, schema); + } } // Return all the tasks, without the functions @@ -233,6 +243,58 @@ export class StandardResourceCatalog implements ResourceCatalog { }; } + registerSkillMetadata(skill: SkillMetadata): void { + if (!this._currentFileContext) { + return; + } + + if (!skill.id) { + return; + } + + const existing = this._skillMetadata.get(skill.id); + if (existing && existing.sourcePath !== skill.sourcePath) { + console.warn( + `Skill "${skill.id}" is defined twice with different paths. Keeping the first:\n` + + ` existing: ${existing.sourcePath}\n` + + ` ignored: ${skill.sourcePath}` + ); + return; + } + + this._skillFileMetadata.set(skill.id, { + ...this._currentFileContext, + }); + this._skillMetadata.set(skill.id, skill); + } + + listSkillManifests(): Array { + const result: Array = []; + + for (const [id, metadata] of this._skillMetadata) { + const fileMetadata = this._skillFileMetadata.get(id); + if (!fileMetadata) continue; + + result.push({ + ...metadata, + ...fileMetadata, + }); + } + + return result; + } + + getSkillManifest(id: string): SkillManifest | undefined { + const metadata = this._skillMetadata.get(id); + const fileMetadata = this._skillFileMetadata.get(id); + if (!metadata || !fileMetadata) return undefined; + + return { + ...metadata, + ...fileMetadata, + }; + } + disable() { // noop } diff --git a/packages/core/src/v3/taskContext/index.test.ts b/packages/core/src/v3/taskContext/index.test.ts new file mode 100644 index 00000000000..34d169a177c --- /dev/null +++ b/packages/core/src/v3/taskContext/index.test.ts @@ -0,0 +1,86 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { unregisterGlobal } from "../utils/globals.js"; +import { SemanticInternalAttributes } from "../semanticInternalAttributes.js"; +import { TaskContextAPI } from "./index.js"; + +const FAKE_CTX = { + attempt: { id: "attempt_1", number: 1, startedAt: new Date(), status: "EXECUTING" as const }, + run: { + id: "run_1", + payload: undefined, + payloadType: "application/json", + context: undefined, + createdAt: new Date(), + tags: [], + isTest: false, + isReplay: false, + startedAt: new Date(), + durationMs: 0, + costInCents: 0, + baseCostInCents: 0, + }, + task: { id: "my-task", filePath: "src/trigger/task.ts", exportName: "myTask" }, + queue: { id: "queue_1", name: "default" }, + environment: { id: "env_1", slug: "dev", type: "DEVELOPMENT" as const }, + organization: { id: "org_1", slug: "acme", name: "Acme" }, + project: { id: "proj_1", ref: "proj_xyz", slug: "demo", name: "Demo" }, + machine: { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, +} as never; + +const FAKE_WORKER = { id: "worker_1", version: "1.0.0", contentHash: "abc" } as never; + +describe("TaskContextAPI conversation id", () => { + afterEach(() => { + unregisterGlobal("task-context"); + TaskContextAPI.getInstance().setConversationId(undefined); + }); + + it("returns no conversation attribute when setConversationId was never called", () => { + const api = TaskContextAPI.getInstance(); + api.setGlobalTaskContext({ ctx: FAKE_CTX, worker: FAKE_WORKER }); + + expect(api.attributes[SemanticInternalAttributes.GEN_AI_CONVERSATION_ID]).toBeUndefined(); + }); + + it("includes gen_ai.conversation.id after setConversationId", () => { + const api = TaskContextAPI.getInstance(); + api.setGlobalTaskContext({ ctx: FAKE_CTX, worker: FAKE_WORKER }); + + api.setConversationId("chat_123"); + + expect(api.attributes[SemanticInternalAttributes.GEN_AI_CONVERSATION_ID]).toBe("chat_123"); + }); + + it("clears the conversation attribute when called with undefined", () => { + const api = TaskContextAPI.getInstance(); + api.setGlobalTaskContext({ ctx: FAKE_CTX, worker: FAKE_WORKER }); + api.setConversationId("chat_123"); + + api.setConversationId(undefined); + + expect(api.attributes[SemanticInternalAttributes.GEN_AI_CONVERSATION_ID]).toBeUndefined(); + expect(api.conversationId).toBeUndefined(); + }); + + it("returns no attributes when there is no task context", () => { + const api = TaskContextAPI.getInstance(); + api.setConversationId("chat_123"); + + expect(api.attributes).toEqual({}); + }); + + it("clears conversation id when a new task context is registered (warm restart)", () => { + const api = TaskContextAPI.getInstance(); + api.setGlobalTaskContext({ ctx: FAKE_CTX, worker: FAKE_WORKER }); + api.setConversationId("chat_old"); + + api.setGlobalTaskContext({ ctx: FAKE_CTX, worker: FAKE_WORKER }); + + expect(api.attributes[SemanticInternalAttributes.GEN_AI_CONVERSATION_ID]).toBeUndefined(); + }); +}); diff --git a/packages/core/src/v3/taskContext/index.ts b/packages/core/src/v3/taskContext/index.ts index 92e0194cde9..ecbfa184a6b 100644 --- a/packages/core/src/v3/taskContext/index.ts +++ b/packages/core/src/v3/taskContext/index.ts @@ -9,6 +9,7 @@ const API_NAME = "task-context"; export class TaskContextAPI { private static _instance?: TaskContextAPI; private _runDisabled = false; + private _conversationId?: string; private constructor() {} @@ -45,6 +46,7 @@ export class TaskContextAPI { return { ...this.contextAttributes, ...this.workerAttributes, + ...this.conversationAttributes, [SemanticInternalAttributes.WARM_START]: !!this.isWarmStart, }; } @@ -52,6 +54,19 @@ export class TaskContextAPI { return {}; } + get conversationAttributes(): Attributes { + if (!this._conversationId) return {}; + return { [SemanticInternalAttributes.GEN_AI_CONVERSATION_ID]: this._conversationId }; + } + + get conversationId(): string | undefined { + return this._conversationId; + } + + public setConversationId(conversationId: string | undefined): void { + this._conversationId = conversationId || undefined; + } + get resourceAttributes(): Attributes { if (this.ctx) { return { @@ -109,6 +124,11 @@ export class TaskContextAPI { public setGlobalTaskContext(taskContext: TaskContext): boolean { this._runDisabled = false; + // Each run boot re-registers the global; clear any conversation id + // left over from a previous run on this warm-restarted process so + // attributes don't bleed across runs that don't call + // `setConversationId` themselves. + this._conversationId = undefined; return registerGlobal(API_NAME, taskContext, true); } diff --git a/packages/core/src/v3/taskContext/otelProcessors.ts b/packages/core/src/v3/taskContext/otelProcessors.ts index 1c0958d655d..fc30e9d1145 100644 --- a/packages/core/src/v3/taskContext/otelProcessors.ts +++ b/packages/core/src/v3/taskContext/otelProcessors.ts @@ -36,6 +36,17 @@ export class TaskContextSpanProcessor implements SpanProcessor { if (!taskContext.isRunDisabled && taskContext.ctx.run.tags?.length) { span.setAttribute(SemanticInternalAttributes.RUN_TAGS, taskContext.ctx.run.tags); } + + // Stamp `gen_ai.conversation.id` (OTel GenAI semantic convention) + // directly on every span so it survives the OTLP ingest's `ctx.*` + // strip and lands in the stored attributes column without a schema + // migration. + if (taskContext.conversationId) { + span.setAttribute( + SemanticInternalAttributes.GEN_AI_CONVERSATION_ID, + taskContext.conversationId + ); + } } if (!isPartialSpan(span) && !skipPartialSpan(span)) { @@ -178,6 +189,11 @@ export class TaskContextMetricExporter implements PushMetricExporter { contextAttrs[SemanticInternalAttributes.RUN_TAGS] = ctx.run.tags; } + if (taskContext.conversationId) { + contextAttrs[SemanticInternalAttributes.GEN_AI_CONVERSATION_ID] = + taskContext.conversationId; + } + const modified: ResourceMetrics = { resource: metrics.resource, scopeMetrics: metrics.scopeMetrics.map((scope) => ({ diff --git a/packages/core/src/v3/test/index.ts b/packages/core/src/v3/test/index.ts new file mode 100644 index 00000000000..402f618c01b --- /dev/null +++ b/packages/core/src/v3/test/index.ts @@ -0,0 +1,9 @@ +export { + runInMockTaskContext, + type MockTaskContextDrivers, + type MockTaskContextOptions, +} from "./mock-task-context.js"; +export { TestInputStreamManager } from "./test-input-stream-manager.js"; +export { TestRealtimeStreamsManager } from "./test-realtime-streams-manager.js"; +export { TestRunMetadataManager } from "./test-run-metadata-manager.js"; +export { TestSessionStreamManager } from "./test-session-stream-manager.js"; diff --git a/packages/core/src/v3/test/mock-task-context.ts b/packages/core/src/v3/test/mock-task-context.ts new file mode 100644 index 00000000000..66e58490019 --- /dev/null +++ b/packages/core/src/v3/test/mock-task-context.ts @@ -0,0 +1,294 @@ +import { inputStreams } from "../input-streams-api.js"; +import { realtimeStreams } from "../realtime-streams-api.js"; +import { sessionStreams } from "../session-streams-api.js"; +import { localsAPI } from "../locals-api.js"; +import { runMetadata } from "../run-metadata-api.js"; +import { taskContext } from "../task-context-api.js"; +import { lifecycleHooks } from "../lifecycle-hooks-api.js"; +import { runtime } from "../runtime-api.js"; +import { StandardLocalsManager } from "../locals/manager.js"; +import { StandardLifecycleHooksManager } from "../lifecycleHooks/manager.js"; +import { NoopRuntimeManager } from "../runtime/noopRuntimeManager.js"; +import { unregisterGlobal } from "../utils/globals.js"; +import type { ServerBackgroundWorker, TaskRunContext } from "../schemas/index.js"; +import type { LocalsKey } from "../locals/types.js"; +import type { SessionChannelIO } from "../sessionStreams/types.js"; +import { TestInputStreamManager } from "./test-input-stream-manager.js"; +import { TestRealtimeStreamsManager } from "./test-realtime-streams-manager.js"; +import { TestRunMetadataManager } from "./test-run-metadata-manager.js"; +import { TestSessionStreamManager } from "./test-session-stream-manager.js"; + +/** + * Shallow-partial overrides applied on top of the default mock + * `TaskRunContext`. Each sub-object is a partial of its real shape — + * unset fields get sensible defaults. + */ +export type MockTaskRunContextOverrides = { + task?: Partial; + attempt?: Partial; + run?: Partial; + machine?: Partial; + queue?: Partial; + environment?: Partial; + organization?: Partial; + project?: Partial; + batch?: TaskRunContext["batch"]; +}; + +/** + * Options for overriding parts of the mock task context. + */ +export type MockTaskContextOptions = { + /** Overrides applied on top of the default mock `TaskRunContext`. */ + ctx?: MockTaskRunContextOverrides; + /** Overrides applied on top of the default `ServerBackgroundWorker`. */ + worker?: Partial; + /** Whether this is a warm start. */ + isWarmStart?: boolean; +}; + +/** + * Drivers passed to the function running inside `runInMockTaskContext`. + */ +export type MockTaskContextDrivers = { + /** Push data into input streams — simulates realtime input from outside the task. */ + inputs: { + /** + * Send `data` to the named input stream. Resolves when all `.on()` + * handlers have run. + */ + send(streamId: string, data: unknown): Promise; + /** Resolve any pending `.once()` waiters with a timeout error. */ + close(streamId: string): void; + }; + /** Inspect chunks written to output (realtime) streams. */ + outputs: { + /** All chunks for a given stream, in the order they were written. */ + chunks(streamId: string): T[]; + /** All chunks across every stream, keyed by stream id. */ + all(): Record; + /** Clear chunks for one stream, or all streams if no id is provided. */ + clear(streamId?: string): void; + /** + * Register a listener fired for every chunk written to any stream. + * Returns an unsubscribe function. + */ + onWrite(listener: (streamId: string, chunk: unknown) => void): () => void; + }; + /** Read or seed locals for the run. */ + locals: { + /** Read a local set by either the task or `set()` below. */ + get(key: LocalsKey): T | undefined; + /** + * Pre-seed a local before the task runs. Use this for dependency + * injection — e.g. supply a test database client that the agent's + * hooks read via `locals.get()` instead of constructing the prod one. + */ + set(key: LocalsKey, value: T): void; + }; + /** + * Session-scoped channel drivers. The `.in` side is backed by a + * {@link TestSessionStreamManager} installed as the `sessionStreams` + * global — so the task's `session.in.on/once/peek/waitWithIdleTimeout` + * calls receive records sent through this driver. + */ + sessions: { + in: { + /** + * Send a record onto `session.in` for the given session. Resolves + * pending `once()` waiters and fires all `on()` handlers. + */ + send(sessionId: string, data: unknown, io?: SessionChannelIO): Promise; + /** Close pending `once()` waiters with a timeout error. */ + close(sessionId: string, io?: SessionChannelIO): void; + }; + }; + /** The mock `TaskRunContext` assembled from defaults + user overrides. */ + ctx: TaskRunContext; +}; + +function defaultTaskRunContext(overrides?: MockTaskRunContextOverrides): TaskRunContext { + return { + task: { + id: "test-task", + filePath: "test-task.ts", + ...overrides?.task, + }, + attempt: { + number: 1, + startedAt: new Date(), + ...overrides?.attempt, + }, + run: { + id: "run_test", + tags: [], + isTest: false, + isReplay: false, + createdAt: new Date(), + startedAt: new Date(), + ...overrides?.run, + }, + machine: { + name: "micro", + cpu: 1, + memory: 0.5, + centsPerMs: 0, + ...overrides?.machine, + }, + queue: { + name: "test-queue", + id: "test-queue-id", + ...overrides?.queue, + }, + environment: { + id: "test-env-id", + slug: "test-env", + type: "DEVELOPMENT", + ...overrides?.environment, + }, + organization: { + id: "test-org-id", + slug: "test-org", + name: "Test Org", + ...overrides?.organization, + }, + project: { + id: "test-project-id", + ref: "test-project-ref", + slug: "test-project", + name: "Test Project", + ...overrides?.project, + }, + batch: overrides?.batch, + }; +} + +function defaultWorker(overrides?: Partial): ServerBackgroundWorker { + return { + id: "test-worker-id", + version: "test-version", + contentHash: "test-content-hash", + engine: "V2", + ...overrides, + }; +} + +/** + * Run a function inside a fully mocked task runtime context. + * + * Installs in-memory test managers for `locals`, `inputStreams`, + * `realtimeStreams`, `lifecycleHooks`, and `runtime`, sets a mock + * `TaskContext`, and tears everything down when the function returns. + * + * Inside the function, any code that reads from `locals`, `inputStreams`, + * `realtimeStreams`, or `taskContext.ctx` will see the mock context — + * so you can directly invoke the internal `run` function of any task + * (including `chat.agent`) without hitting the Trigger.dev runtime. + * + * @example + * ```ts + * import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; + * + * await runInMockTaskContext( + * async ({ inputs, outputs, ctx }) => { + * // Fire an input stream from the "outside" + * setTimeout(() => { + * inputs.send("chat-messages", { messages: [], chatId: "c1" }); + * }, 0); + * + * // Run task code that reads from inputStreams.once(...) + * await myTask.fns.run(payload, { ctx, signal: new AbortController().signal }); + * + * // Inspect chunks written to the output stream + * expect(outputs.chunks("chat")).toContainEqual({ type: "text-delta", delta: "hi" }); + * }, + * { ctx: { run: { id: "run_abc" } } } + * ); + * ``` + */ +export async function runInMockTaskContext( + fn: (drivers: MockTaskContextDrivers) => T | Promise, + options?: MockTaskContextOptions +): Promise { + const ctx = defaultTaskRunContext(options?.ctx); + const worker = defaultWorker(options?.worker); + + const localsManager = new StandardLocalsManager(); + const lifecycleManager = new StandardLifecycleHooksManager(); + const runtimeManager = new NoopRuntimeManager(); + const metadataManager = new TestRunMetadataManager(); + const inputManager = new TestInputStreamManager(); + const outputManager = new TestRealtimeStreamsManager(); + const sessionStreamManager = new TestSessionStreamManager(); + + // Unregister any previously-installed managers so `setGlobal*` wins — + // `registerGlobal` returns false silently if an entry already exists. + unregisterGlobal("locals"); + unregisterGlobal("lifecycle-hooks"); + unregisterGlobal("runtime"); + unregisterGlobal("run-metadata"); + unregisterGlobal("input-streams"); + unregisterGlobal("realtime-streams"); + unregisterGlobal("session-streams"); + unregisterGlobal("task-context"); + + localsAPI.setGlobalLocalsManager(localsManager); + lifecycleHooks.setGlobalLifecycleHooksManager(lifecycleManager); + runtime.setGlobalRuntimeManager(runtimeManager); + runMetadata.setGlobalManager(metadataManager); + inputStreams.setGlobalManager(inputManager); + realtimeStreams.setGlobalManager(outputManager); + sessionStreams.setGlobalManager(sessionStreamManager); + taskContext.setGlobalTaskContext({ + ctx, + worker, + isWarmStart: options?.isWarmStart ?? false, + }); + + const drivers: MockTaskContextDrivers = { + inputs: { + send: (streamId, data) => inputManager.__sendFromTest(streamId, data), + close: (streamId) => inputManager.__closeFromTest(streamId), + }, + outputs: { + chunks: (streamId) => outputManager.__chunksFromTest(streamId), + all: () => outputManager.__allChunksFromTest(), + clear: (streamId) => outputManager.__clearFromTest(streamId), + onWrite: (listener) => outputManager.onWrite(listener), + }, + locals: { + get: (key: LocalsKey) => localsManager.getLocal(key), + set: (key: LocalsKey, value: TValue) => + localsManager.setLocal(key, value), + }, + sessions: { + in: { + send: (sessionId, data, io = "in") => + sessionStreamManager.__sendFromTest(sessionId, io, data), + close: (sessionId, io = "in") => + sessionStreamManager.__closeFromTest(sessionId, io), + }, + }, + ctx, + }; + + try { + return await fn(drivers); + } finally { + localsAPI.disable(); + lifecycleHooks.disable(); + runtime.disable(); + // taskContext.disable() only sets a flag — unregister the global so + // `taskContext.ctx` returns undefined after the harness returns. + unregisterGlobal("task-context"); + unregisterGlobal("input-streams"); + unregisterGlobal("realtime-streams"); + unregisterGlobal("session-streams"); + unregisterGlobal("run-metadata"); + localsManager.reset(); + inputManager.reset(); + outputManager.reset(); + sessionStreamManager.reset(); + metadataManager.reset(); + } +} diff --git a/packages/core/test/mockTaskContext.test.ts b/packages/core/test/mockTaskContext.test.ts new file mode 100644 index 00000000000..5ea3685e466 --- /dev/null +++ b/packages/core/test/mockTaskContext.test.ts @@ -0,0 +1,226 @@ +import { describe, expect, it } from "vitest"; +import { runInMockTaskContext } from "../src/v3/test/index.js"; +import { inputStreams } from "../src/v3/input-streams-api.js"; +import { realtimeStreams } from "../src/v3/realtime-streams-api.js"; +import { locals } from "../src/v3/locals-api.js"; +import { taskContext } from "../src/v3/task-context-api.js"; + +describe("runInMockTaskContext", () => { + it("installs a mock TaskRunContext with sensible defaults", async () => { + await runInMockTaskContext(async ({ ctx }) => { + expect(taskContext.ctx).toBeDefined(); + expect(taskContext.ctx?.run.id).toBe("run_test"); + expect(taskContext.ctx?.task.id).toBe("test-task"); + expect(ctx.run.id).toBe("run_test"); + }); + }); + + it("applies ctx overrides on top of defaults", async () => { + await runInMockTaskContext( + async ({ ctx }) => { + expect(ctx.run.id).toBe("run_abc"); + expect(ctx.task.id).toBe("my-chat-agent"); + // Unspecified fields still use defaults + expect(ctx.queue.id).toBe("test-queue-id"); + }, + { + ctx: { + run: { id: "run_abc" }, + task: { id: "my-chat-agent", filePath: "chat.ts" }, + }, + } + ); + }); + + it("isolates locals from the surrounding context", async () => { + const key = locals.create<{ count: number }>("test.counter"); + + await runInMockTaskContext(async ({ locals: inspect }) => { + expect(inspect.get(key)).toBeUndefined(); + locals.set(key, { count: 1 }); + expect(inspect.get(key)).toEqual({ count: 1 }); + }); + + // After the harness exits, the locals should be gone + expect(locals.get(key)).toBeUndefined(); + }); + + it("tears down the task context after fn returns", async () => { + await runInMockTaskContext(async () => { + expect(taskContext.ctx).toBeDefined(); + }); + + expect(taskContext.ctx).toBeUndefined(); + }); + + it("tears down even when fn throws", async () => { + await expect( + runInMockTaskContext(async () => { + throw new Error("boom"); + }) + ).rejects.toThrow("boom"); + + expect(taskContext.ctx).toBeUndefined(); + }); + + it("returns the value returned by fn", async () => { + const result = await runInMockTaskContext(async () => "hello"); + expect(result).toBe("hello"); + }); + + describe("input streams driver", () => { + it("resolves inputStreams.once() when test sends data", async () => { + await runInMockTaskContext(async ({ inputs }) => { + const pending = inputStreams.once("chat-messages"); + setTimeout(() => inputs.send("chat-messages", { hello: "world" }), 0); + const result = await pending; + expect(result.ok).toBe(true); + if (result.ok) { + expect(result.output).toEqual({ hello: "world" }); + } + }); + }); + + it("fires inputStreams.on() handlers when test sends data", async () => { + await runInMockTaskContext(async ({ inputs }) => { + const received: unknown[] = []; + inputStreams.on("chat-messages", (data) => { + received.push(data); + }); + + await inputs.send("chat-messages", { n: 1 }); + await inputs.send("chat-messages", { n: 2 }); + + expect(received).toEqual([{ n: 1 }, { n: 2 }]); + }); + }); + + it("fires multiple on() handlers on the same stream", async () => { + await runInMockTaskContext(async ({ inputs }) => { + const a: unknown[] = []; + const b: unknown[] = []; + inputStreams.on("chat-messages", (data) => a.push(data)); + inputStreams.on("chat-messages", (data) => b.push(data)); + + await inputs.send("chat-messages", "hi"); + expect(a).toEqual(["hi"]); + expect(b).toEqual(["hi"]); + }); + }); + + it("off() unsubscribes a handler", async () => { + await runInMockTaskContext(async ({ inputs }) => { + const received: unknown[] = []; + const sub = inputStreams.on("chat-messages", (data) => received.push(data)); + + await inputs.send("chat-messages", 1); + sub.off(); + await inputs.send("chat-messages", 2); + + expect(received).toEqual([1]); + }); + }); + + it("times out once() after timeoutMs", async () => { + await runInMockTaskContext(async () => { + const result = await inputStreams.once("chat-messages", { timeoutMs: 10 }); + expect(result.ok).toBe(false); + }); + }); + + it("peek() returns the latest sent value", async () => { + await runInMockTaskContext(async ({ inputs }) => { + expect(inputStreams.peek("chat-messages")).toBeUndefined(); + await inputs.send("chat-messages", { latest: true }); + expect(inputStreams.peek("chat-messages")).toEqual({ latest: true }); + }); + }); + + it("close() rejects pending once() waiters with a timeout error", async () => { + await runInMockTaskContext(async ({ inputs }) => { + const pending = inputStreams.once("chat-messages"); + inputs.close("chat-messages"); + const result = await pending; + expect(result.ok).toBe(false); + }); + }); + + it("resolves multiple concurrent once() waiters from a single send", async () => { + await runInMockTaskContext(async ({ inputs }) => { + const a = inputStreams.once("chat-messages"); + const b = inputStreams.once("chat-messages"); + await inputs.send("chat-messages", "shared"); + const [ra, rb] = await Promise.all([a, b]); + expect(ra.ok && ra.output).toBe("shared"); + expect(rb.ok && rb.output).toBe("shared"); + }); + }); + }); + + describe("realtime streams driver", () => { + it("collects chunks from realtimeStreams.append()", async () => { + await runInMockTaskContext(async ({ outputs }) => { + await realtimeStreams.append("chat", "chunk-1" as unknown as BodyInit); + await realtimeStreams.append("chat", "chunk-2" as unknown as BodyInit); + + expect(outputs.chunks("chat")).toEqual(["chunk-1", "chunk-2"]); + }); + }); + + it("collects chunks from realtimeStreams.pipe()", async () => { + await runInMockTaskContext(async ({ outputs }) => { + const source = (async function* () { + yield "a"; + yield "b"; + yield "c"; + })(); + + const instance = realtimeStreams.pipe("chat", source); + + // Drain the returned stream — that's what feeds the buffer + for await (const _ of instance.stream) { + // no-op + } + + expect(outputs.chunks("chat")).toEqual(["a", "b", "c"]); + }); + }); + + it("separates chunks by stream id", async () => { + await runInMockTaskContext(async ({ outputs }) => { + await realtimeStreams.append("chat", "a" as unknown as BodyInit); + await realtimeStreams.append("stop", "halt" as unknown as BodyInit); + + expect(outputs.chunks("chat")).toEqual(["a"]); + expect(outputs.chunks("stop")).toEqual(["halt"]); + expect(outputs.all()).toEqual({ chat: ["a"], stop: ["halt"] }); + }); + }); + + it("clear() empties one stream or all streams", async () => { + await runInMockTaskContext(async ({ outputs }) => { + await realtimeStreams.append("chat", "a" as unknown as BodyInit); + await realtimeStreams.append("stop", "halt" as unknown as BodyInit); + + outputs.clear("chat"); + expect(outputs.chunks("chat")).toEqual([]); + expect(outputs.chunks("stop")).toEqual(["halt"]); + + outputs.clear(); + expect(outputs.chunks("stop")).toEqual([]); + }); + }); + }); + + it("tears down input/output managers so consecutive calls are isolated", async () => { + await runInMockTaskContext(async ({ inputs }) => { + await inputs.send("chat-messages", "first-run"); + }); + + await runInMockTaskContext(async ({ outputs }) => { + expect(outputs.chunks("chat-messages")).toEqual([]); + // inputs.peek should NOT see "first-run" from the prior harness + expect(inputStreams.peek("chat-messages")).toBeUndefined(); + }); + }); +}); diff --git a/packages/core/test/skillCatalog.test.ts b/packages/core/test/skillCatalog.test.ts new file mode 100644 index 00000000000..3f1d29bf572 --- /dev/null +++ b/packages/core/test/skillCatalog.test.ts @@ -0,0 +1,74 @@ +import { describe, expect, it, vi } from "vitest"; +import { StandardResourceCatalog } from "../src/v3/resource-catalog/standardResourceCatalog.js"; + +describe("StandardResourceCatalog — skills", () => { + it("registers and lists a skill manifest", () => { + const catalog = new StandardResourceCatalog(); + catalog.setCurrentFileContext("trigger/chat.ts", "chat"); + + catalog.registerSkillMetadata({ id: "pdf-processing", sourcePath: "./skills/pdf-processing" }); + + const manifests = catalog.listSkillManifests(); + expect(manifests).toHaveLength(1); + expect(manifests[0]).toMatchObject({ + id: "pdf-processing", + sourcePath: "./skills/pdf-processing", + filePath: "trigger/chat.ts", + entryPoint: "chat", + }); + }); + + it("getSkillManifest returns the registered skill", () => { + const catalog = new StandardResourceCatalog(); + catalog.setCurrentFileContext("trigger/chat.ts", "chat"); + catalog.registerSkillMetadata({ id: "a", sourcePath: "./skills/a" }); + + expect(catalog.getSkillManifest("a")?.sourcePath).toBe("./skills/a"); + expect(catalog.getSkillManifest("missing")).toBeUndefined(); + }); + + it("skips registration without a file context", () => { + const catalog = new StandardResourceCatalog(); + + catalog.registerSkillMetadata({ id: "pdf", sourcePath: "./skills/pdf" }); + + expect(catalog.listSkillManifests()).toHaveLength(0); + }); + + it("warns and ignores when the same id is registered with a different path", () => { + const catalog = new StandardResourceCatalog(); + catalog.setCurrentFileContext("trigger/chat.ts", "chat"); + + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + + catalog.registerSkillMetadata({ id: "pdf", sourcePath: "./skills/pdf" }); + catalog.registerSkillMetadata({ id: "pdf", sourcePath: "./skills/other-pdf" }); + + const manifests = catalog.listSkillManifests(); + expect(manifests).toHaveLength(1); + expect(manifests[0]?.sourcePath).toBe("./skills/pdf"); + expect(warn).toHaveBeenCalledWith(expect.stringContaining("defined twice")); + + warn.mockRestore(); + }); + + it("re-registering the same id + path is idempotent", () => { + const catalog = new StandardResourceCatalog(); + catalog.setCurrentFileContext("trigger/chat.ts", "chat"); + + catalog.registerSkillMetadata({ id: "pdf", sourcePath: "./skills/pdf" }); + catalog.registerSkillMetadata({ id: "pdf", sourcePath: "./skills/pdf" }); + + expect(catalog.listSkillManifests()).toHaveLength(1); + }); + + it("registers multiple distinct skills", () => { + const catalog = new StandardResourceCatalog(); + catalog.setCurrentFileContext("trigger/chat.ts", "chat"); + + catalog.registerSkillMetadata({ id: "pdf", sourcePath: "./skills/pdf" }); + catalog.registerSkillMetadata({ id: "researcher", sourcePath: "./skills/researcher" }); + + expect(catalog.listSkillManifests().map((s) => s.id).sort()).toEqual(["pdf", "researcher"]); + }); +}); diff --git a/packages/trigger-sdk/package.json b/packages/trigger-sdk/package.json index 9a1b90b059e..e2661b91719 100644 --- a/packages/trigger-sdk/package.json +++ b/packages/trigger-sdk/package.json @@ -24,7 +24,12 @@ "./package.json": "./package.json", ".": "./src/v3/index.ts", "./v3": "./src/v3/index.ts", - "./ai": "./src/v3/ai.ts" + "./ai": "./src/v3/ai.ts", + "./ai/skills-runtime": "./src/v3/agentSkillsRuntime.ts", + "./ai/test": "./src/v3/test/index.ts", + "./chat": "./src/v3/chat.ts", + "./chat/react": "./src/v3/chat-react.ts", + "./chat-server": "./src/v3/chat-server.ts" }, "sourceDialects": [ "@triggerdotdev/source" @@ -37,6 +42,21 @@ ], "ai": [ "dist/commonjs/v3/ai.d.ts" + ], + "ai/skills-runtime": [ + "dist/commonjs/v3/agentSkillsRuntime.d.ts" + ], + "ai/test": [ + "dist/commonjs/v3/test/index.d.ts" + ], + "chat": [ + "dist/commonjs/v3/chat.d.ts" + ], + "chat/react": [ + "dist/commonjs/v3/chat-react.d.ts" + ], + "chat-server": [ + "dist/commonjs/v3/chat-server.d.ts" ] } }, @@ -63,11 +83,13 @@ "ws": "^8.11.0" }, "devDependencies": { + "@ai-sdk/provider": "3.0.8", "@arethetypeswrong/cli": "^0.15.4", "@types/debug": "^4.1.7", + "@types/react": "^19.2.14", "@types/slug": "^5.0.3", "@types/ws": "^8.5.3", - "ai": "^6.0.0", + "ai": "^6.0.116", "encoding": "^0.1.13", "rimraf": "^6.0.1", "tshy": "^3.0.2", @@ -76,12 +98,16 @@ "zod": "3.25.76" }, "peerDependencies": { - "zod": "^3.0.0 || ^4.0.0", - "ai": "^4.2.0 || ^5.0.0 || ^6.0.0" + "ai": "^5.0.0 || ^6.0.0", + "react": "^18.0 || ^19.0", + "zod": "^3.0.0 || ^4.0.0" }, "peerDependenciesMeta": { "ai": { "optional": true + }, + "react": { + "optional": true } }, "engines": { @@ -121,6 +147,61 @@ "types": "./dist/commonjs/v3/ai.d.ts", "default": "./dist/commonjs/v3/ai.js" } + }, + "./ai/skills-runtime": { + "import": { + "@triggerdotdev/source": "./src/v3/agentSkillsRuntime.ts", + "types": "./dist/esm/v3/agentSkillsRuntime.d.ts", + "default": "./dist/esm/v3/agentSkillsRuntime.js" + }, + "require": { + "types": "./dist/commonjs/v3/agentSkillsRuntime.d.ts", + "default": "./dist/commonjs/v3/agentSkillsRuntime.js" + } + }, + "./ai/test": { + "import": { + "@triggerdotdev/source": "./src/v3/test/index.ts", + "types": "./dist/esm/v3/test/index.d.ts", + "default": "./dist/esm/v3/test/index.js" + }, + "require": { + "types": "./dist/commonjs/v3/test/index.d.ts", + "default": "./dist/commonjs/v3/test/index.js" + } + }, + "./chat": { + "import": { + "@triggerdotdev/source": "./src/v3/chat.ts", + "types": "./dist/esm/v3/chat.d.ts", + "default": "./dist/esm/v3/chat.js" + }, + "require": { + "types": "./dist/commonjs/v3/chat.d.ts", + "default": "./dist/commonjs/v3/chat.js" + } + }, + "./chat/react": { + "import": { + "@triggerdotdev/source": "./src/v3/chat-react.ts", + "types": "./dist/esm/v3/chat-react.d.ts", + "default": "./dist/esm/v3/chat-react.js" + }, + "require": { + "types": "./dist/commonjs/v3/chat-react.d.ts", + "default": "./dist/commonjs/v3/chat-react.js" + } + }, + "./chat-server": { + "import": { + "@triggerdotdev/source": "./src/v3/chat-server.ts", + "types": "./dist/esm/v3/chat-server.d.ts", + "default": "./dist/esm/v3/chat-server.js" + }, + "require": { + "types": "./dist/commonjs/v3/chat-server.d.ts", + "default": "./dist/commonjs/v3/chat-server.js" + } } }, "main": "./dist/commonjs/v3/index.js", diff --git a/packages/trigger-sdk/src/v3/agentSkillsRuntime.ts b/packages/trigger-sdk/src/v3/agentSkillsRuntime.ts new file mode 100644 index 00000000000..31501ca4aef --- /dev/null +++ b/packages/trigger-sdk/src/v3/agentSkillsRuntime.ts @@ -0,0 +1,127 @@ +import { spawn } from "node:child_process"; +import * as fs from "node:fs/promises"; +import * as nodePath from "node:path"; + +/** + * Server-only runtime for the auto-injected skill tools + * (`loadSkill` / `readFile` / `bash`) that `chat.agent({ skills })` + * wires up. Split off from `./ai.ts` so the chat-agent surface in + * `@trigger.dev/sdk/ai` stays importable from client bundles — + * Next.js + Webpack reject top-level `node:*` imports anywhere in a + * client graph, even when a consumer only pulls in types. + * + * The SDK's `ai.ts` loads this module via a computed-string dynamic + * import inside each tool's `execute` — webpack treats the + * expression as an unknown dependency and skips static tracing, so + * the node-only symbols here never surface in a client build. The + * module resolves fine at runtime on a server worker because the + * relative path (`./agentSkillsRuntime.js`) lands next to `ai.js` in + * the emitted dist. + * + * Public subpath: `@trigger.dev/sdk/ai/skills-runtime`. Customers + * who want to eagerly bundle the runtime server-side (e.g. warming + * it on worker bootstrap) can import from there. + */ + +const DEFAULT_BASH_OUTPUT_BYTES = 64 * 1024; +const DEFAULT_READ_FILE_BYTES = 1024 * 1024; + +export type BashSkillInput = { + /** Absolute path to the skill's root (used as `cwd`). */ + skillPath: string; + /** The bash command to run. */ + command: string; + /** Optional abort signal forwarded to `spawn()`. */ + abortSignal?: AbortSignal; +}; + +export type BashSkillResult = + | { exitCode: number | null; stdout: string; stderr: string } + | { error: string }; + +export type ReadFileInSkillInput = { + /** Absolute path to the skill's root — the relative path must resolve inside it. */ + skillPath: string; + /** Relative path the tool caller supplied. */ + relativePath: string; +}; + +export type ReadFileInSkillResult = { content: string } | { error: string }; + +function truncate(s: string, limit: number): string { + if (s.length <= limit) return s; + return s.slice(0, limit) + `\n…[truncated ${s.length - limit} bytes]`; +} + +/** + * Path-traversal guard: confirm `relative` resolves inside `root`. + * Throws if it escapes via `..` or an absolute prefix. Returns the + * absolute resolved path. + */ +function safeJoinInside(root: string, relative: string): string { + if (nodePath.isAbsolute(relative)) { + throw new Error(`Path must be relative to the skill directory: ${relative}`); + } + const resolved = nodePath.resolve(root, relative); + const normalized = nodePath.resolve(root) + nodePath.sep; + if (resolved !== nodePath.resolve(root) && !resolved.startsWith(normalized)) { + throw new Error(`Path escapes the skill directory: ${relative}`); + } + return resolved; +} + +export async function readFileInSkill({ + skillPath, + relativePath, +}: ReadFileInSkillInput): Promise { + let absolute: string; + try { + absolute = safeJoinInside(skillPath, relativePath); + } catch (err) { + return { error: (err as Error).message }; + } + try { + const content = await fs.readFile(absolute, "utf8"); + return { content: truncate(content, DEFAULT_READ_FILE_BYTES) }; + } catch (err) { + return { error: (err as Error).message }; + } +} + +export async function runBashInSkill({ + skillPath, + command, + abortSignal, +}: BashSkillInput): Promise { + return new Promise((resolvePromise) => { + let child; + try { + child = spawn("bash", ["-c", command], { + cwd: skillPath, + signal: abortSignal, + }); + } catch (err) { + resolvePromise({ error: (err as Error).message }); + return; + } + + let stdout = ""; + let stderr = ""; + child.stdout?.on("data", (chunk: Buffer | string) => { + stdout += chunk.toString(); + }); + child.stderr?.on("data", (chunk: Buffer | string) => { + stderr += chunk.toString(); + }); + child.once("close", (code: number | null) => { + resolvePromise({ + exitCode: code, + stdout: truncate(stdout, DEFAULT_BASH_OUTPUT_BYTES), + stderr: truncate(stderr, DEFAULT_BASH_OUTPUT_BYTES), + }); + }); + child.once("error", (err: Error) => { + resolvePromise({ error: err.message }); + }); + }); +} diff --git a/packages/trigger-sdk/src/v3/ai-shared.ts b/packages/trigger-sdk/src/v3/ai-shared.ts new file mode 100644 index 00000000000..7161385764f --- /dev/null +++ b/packages/trigger-sdk/src/v3/ai-shared.ts @@ -0,0 +1,210 @@ +/** + * Browser-safe primitives shared between `@trigger.dev/sdk/ai` (server) and + * `@trigger.dev/sdk/chat` / `@trigger.dev/sdk/chat/react` (client). + * + * This module exists to keep `ai.ts` reachable only from the server graph. + * `ai.ts` weighs in at ~7000 lines and statically imports the agent-skills + * runtime (which uses `node:child_process` / `node:fs/promises`). When a + * browser bundle imports a runtime value from `ai.ts` — historically the + * `PENDING_MESSAGE_INJECTED_TYPE` constant in `chat-react.ts` — the bundler + * traces `ai.ts`'s entire module graph into the client chunk and hits the + * `node:` builtins, which Turbopack rejects outright (and webpack flags as + * a "Critical dependency" warning). + * + * Anything in this file MUST stay free of `node:*` imports and free of any + * import from `ai.ts`. + */ + +import type { Task, AnyTask } from "@trigger.dev/core/v3"; +import type { ModelMessage, UIMessage } from "ai"; + +/** + * Message-part `type` value for the pending-message data part the agent + * injects when a follow-up message arrives mid-turn. + */ +export const PENDING_MESSAGE_INJECTED_TYPE = "data-pending-message-injected" as const; + +/** + * The wire payload shape sent by `TriggerChatTransport`. + * Uses `metadata` to match the AI SDK's `ChatRequestOptions` field name. + * + * Slim wire: at most ONE message per record. The agent runtime + * reconstructs prior history at run boot from a durable S3 snapshot + + * `session.out` replay (or `hydrateMessages` if registered). The wire is + * delta-only — see plan `vivid-humming-bonbon.md`. + */ +export type ChatTaskWirePayload = { + /** + * The single message being delivered on this trigger. Set for: + * - `submit-message`: the new user message OR a tool-approval-responded + * assistant message (with `state: "approval-responded"` tool parts). + * - `regenerate-message`: omitted (the agent slices its own history). + * - `preload` / `close` / `action`: omitted. + * - `handover-prepare`: omitted (use `headStartMessages` instead). + */ + message?: TMessage; + /** + * Bespoke escape hatch for `chat.headStart`. The customer's HTTP route + * handler ships full `UIMessage[]` history at the very first turn — before + * any snapshot exists. The route handler isn't subject to the + * `MAX_APPEND_BODY_BYTES` cap on `/in/append` because it goes through the + * customer's own HTTP endpoint. Used ONLY by `trigger: "handover-prepare"`. + * Ignored on every other trigger. + */ + headStartMessages?: TMessage[]; + chatId: string; + trigger: + | "submit-message" + | "regenerate-message" + | "preload" + | "close" + | "action" + /** + * The customer's `chat.handover` route handler kicked us off in + * parallel with the first-turn `streamText` running in the warm + * Next.js process. The run sits idle on `session.in` waiting for + * a `kind: "handover"` (continue from tool execution) or + * `kind: "handover-skip"` (handler finished pure-text, exit + * cleanly). See `chat.handover` in `@trigger.dev/sdk/chat-server`. + */ + | "handover-prepare"; + messageId?: string; + metadata?: TMetadata; + /** Custom action payload when `trigger` is `"action"`. Validated against `actionSchema` on the backend. */ + action?: unknown; + /** Whether this run is continuing an existing chat whose previous run ended. */ + continuation?: boolean; + /** The run ID of the previous run (only set when `continuation` is true). */ + previousRunId?: string; + /** Override idle timeout for this run (seconds). Set by transport.preload(). */ + idleTimeoutInSeconds?: number; + /** + * The friendlyId of the Session primitive backing this chat. The + * transport opens (or lazy-creates) the session with + * `externalId = chatId` on first message, then sends this friendlyId + * through to the run so the agent can attach to `.in` / `.out` + * without needing to round-trip through the control plane again. + * Optional for backward-compat while the migration is in flight; + * required once the legacy run-scoped stream path is removed. + */ + sessionId?: string; + /** + * Client-side `chat.store` value sent by the transport. Applied at turn + * start before `run()` fires, overwriting any in-memory store value on the + * agent (last-write-wins). + * + * The transport queues this via `setStore` / `applyStorePatch` and flushes + * it with the next `sendMessage`. On the agent you typically don't read + * this directly — it's applied into `chat.store` transparently. + */ + incomingStore?: unknown; +}; + +/** + * One chunk on the chat input stream. `kind` discriminates the variants — + * a single ordered stream now carries all the signals the old three-stream + * split did (`chat-messages`, `chat-stop`, plus action messages piggybacked + * on `chat-messages`). + */ +export type ChatInputChunk = + | { + kind: "message"; + /** + * Full wire payload for a new user message or regeneration. Mirrors + * what the legacy `chat-messages` input stream carried. + */ + payload: ChatTaskWirePayload; + } + | { + kind: "stop"; + /** Optional human-readable reason. Maps to the legacy `chat-stop` record. */ + message?: string; + } + | { + /** + * Sent by `chat.headStart` when the customer's first-turn + * `streamText` finishes. The agent run (currently parked in + * `handover-prepare`) wakes, seeds its accumulators with + * `partialAssistantMessage`, and runs the normal turn loop + * (`onChatStart` → `onTurnStart` → … → `onTurnComplete`). + * + * What happens after that depends on `isFinal`: + * + * - `isFinal: false` — step 1 ended with `finishReason: + * "tool-calls"`. The partial carries the assistant's + * tool-call(s) wrapped in AI SDK's tool-approval round. The + * agent's `streamText` runs the approved tools and continues + * from step 2. + * - `isFinal: true` — step 1 ended pure-text (no tool calls). + * The partial carries the final assistant text. The agent + * skips the LLM call entirely (the response is already + * complete on the customer side) and runs `onTurnComplete` + * with the partial as `responseMessage` so persistence and + * any post-turn work fire normally. + */ + kind: "handover"; + /** Customer's step-1 response messages (ModelMessage form). */ + partialAssistantMessage: ModelMessage[]; + /** + * The UI messageId the customer's handler used for its step-1 + * assistant message. The agent reuses this so any post-handover + * chunks (tool-output-available, step-2 text, data-* parts + * written by hooks) merge into the SAME assistant message on + * the browser side instead of starting a new one. + */ + messageId?: string; + /** + * Whether the customer's step 1 is the final response. See + * `kind` description above for the two branches. + */ + isFinal: boolean; + } + | { + /** + * Sent by `chat.headStart` only when the customer's handler + * ABORTS before producing a finishReason (e.g., dispatch error, + * stream cancelled before any tokens). The agent run exits + * cleanly without firing turn hooks. Normal pure-text and + * tool-call finishes go through `kind: "handover"` with the + * appropriate `isFinal` flag. + */ + kind: "handover-skip"; + }; + +/** + * Extracts the client-data (`metadata`) type from a chat task. + * + * @example + * ```ts + * import type { InferChatClientData } from "@trigger.dev/sdk/ai"; + * import type { myChat } from "@/trigger/chat"; + * + * type MyClientData = InferChatClientData; + * ``` + */ +export type InferChatClientData = TTask extends Task< + string, + ChatTaskWirePayload, + any +> + ? TMetadata + : unknown; + +/** + * Extracts the UI message type from a chat task (wire payload `message` items). + * + * @example + * ```ts + * import type { InferChatUIMessage } from "@trigger.dev/sdk/ai"; + * import type { myChat } from "@/trigger/chat"; + * + * type Msg = InferChatUIMessage; + * ``` + */ +export type InferChatUIMessage = TTask extends Task< + string, + ChatTaskWirePayload, + any +> + ? TUIM + : UIMessage; diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 59afa2fe21a..1b0fa19e390 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1,38 +1,837 @@ import { + accessoryAttributes, AnyTask, + apiClientManager, + getSchemaParseFn, + InputStreamOncePromise, + type InputStreamOnceOptions, + type InputStreamWaitOptions, + type InputStreamWaitWithIdleTimeoutOptions, isSchemaZodEsque, + logger, + type MachinePresetName, + ManualWaitpointPromise, + OutOfMemoryError, + sessionStreams, + type PipeStreamResult, + type RealtimeDefinedInputStream, + type RealtimeDefinedStream, + type ReadStreamOptions, + SemanticInternalAttributes, + type SendInputStreamOptions, Task, + taskContext, + type AppendStreamOptions, + type InputStreamOnceResult, type inferSchemaIn, + type inferSchemaOut, + type PipeStreamOptions, + type TaskIdentifier, + type TaskOptions, type TaskSchema, + type TaskRunContext, type TaskWithSchema, + type WriterStreamOptions, } from "@trigger.dev/core/v3"; -import { dynamicTool, jsonSchema, JSONSchema7, Schema, Tool, ToolCallOptions, zodSchema } from "ai"; +import type { + FinishReason, + ModelMessage, + ToolSet, + UIMessage, + UIMessageChunk, + UIMessageStreamOptions, + LanguageModelUsage, +} from "ai"; +import type { StreamWriteResult } from "@trigger.dev/core/v3"; +import { + convertToModelMessages, + dynamicTool, + generateId as generateMessageId, + getToolName, + isToolUIPart, + jsonSchema, + JSONSchema7, + readUIMessageStream, + Schema, + tool as aiTool, + Tool, + ToolCallOptions, + zodSchema, +} from "ai"; +import { type Attributes, trace } from "@opentelemetry/api"; +import { auth } from "./auth.js"; +import { locals } from "./locals.js"; import { metadata } from "./metadata.js"; +import type { ResolvedPrompt } from "./prompt.js"; +import type { ResolvedSkill } from "./skill.js"; +// Bash-skill runtime lives in `./agentSkillsRuntime.ts` (exposed as +// the `@trigger.dev/sdk/ai/skills-runtime` subpath). It's a normal +// static import — `ai.ts` is server-only by reachability now that +// browser-side primitives (PENDING_MESSAGE_INJECTED_TYPE and the +// chat-task wire types) live in `./ai-shared.ts`. Any browser bundle +// that wants those primitives imports `./ai-shared.js` directly and +// never touches `ai.ts`'s module graph, so the `node:*` builtins +// pulled in transitively here never reach a client chunk. +import { runBashInSkill, readFileInSkill } from "./agentSkillsRuntime.js"; +import { streams } from "./streams.js"; +import { + sessions, + type SessionHandle, + type SessionInputChannel, + type SessionOutputChannel, + type SessionPipeStreamOptions, + type SessionSubscribeOptions, +} from "./sessions.js"; +import { createTask } from "./shared.js"; +import { resourceCatalog, type SessionTriggerConfig } from "@trigger.dev/core/v3"; +import { tracer } from "./tracer.js"; + +/** Re-export for typing `ctx` in `chat.agent` hooks without importing `@trigger.dev/core`. */ +export type { TaskRunContext } from "@trigger.dev/core/v3"; +import { + applyChatStorePatch, + type ChatStoreChunk, + type ChatStoreDeltaChunk, + type ChatStorePatchOperation, + type ChatStoreSnapshotChunk, +} from "@trigger.dev/core/v3/chat-client"; const METADATA_KEY = "tool.execute.options"; -export type ToolCallExecutionOptions = Omit; +/** + * Wrapper around `convertToModelMessages` that always passes + * `ignoreIncompleteToolCalls: true` to prevent failures from + * stopped/aborted conversations with partial tool parts. + */ +function toModelMessages(messages: UIMessage[]): Promise { + return convertToModelMessages(messages, { ignoreIncompleteToolCalls: true }); +} + +export type ToolCallExecutionOptions = { + toolCallId: string; + experimental_context?: unknown; + /** Chat context — only present when the tool runs inside a chat.agent turn. */ + chatId?: string; + turn?: number; + continuation?: boolean; + clientData?: unknown; + /** Serialized chat.local values from the parent run. @internal */ + chatLocals?: Record; +}; + +/** Chat context stored in locals during each chat.agent turn for auto-detection. */ +type ChatTurnContext = { + chatId: string; + turn: number; + continuation: boolean; + clientData?: TClientData; +}; +const chatTurnContextKey = locals.create("chat.turnContext"); + +/** + * Per-run slot holding the Session handle that backs this chat's `.in` / + * `.out` channels. Populated at the top of `chatAgent`'s run function from + * `payload.sessionId`; read by every module-level helper (`chatStream`, + * `messagesInput`, `stopInput`) so the chat.agent internals can remain + * the same module-level shape they were when the I/O was run-scoped. + * @internal + */ +const chatSessionHandleKey = locals.create("chat.sessionHandle"); + +/** + * Scan `session.out` for the latest `trigger:turn-complete` chunk and + * return its SSE timestamp. Used at OOM-retry boot to derive a + * lower-bound timestamp for the `session.in` filter — records older + * than `T_last_complete` belong to turns that already completed on the + * prior attempt and are dropped before they reach the turn loop. + * + * Implementation is a streaming scan: subscribes via the existing SSE + * endpoint with a short `timeoutInSeconds`, processes each part inline, + * and discards the chunk body so memory stays O(1) regardless of how + * many records are on `session.out`. Bandwidth scales linearly with + * stream length but the scan only fires on retry — a rare event. + * + * Returns `undefined` if no `trigger:turn-complete` chunk has been + * written yet (first-turn OOM, no completed turns to dedup against). + * @internal + */ +async function findLatestTurnCompleteTimestamp( + chatId: string +): Promise { + const apiClient = apiClientManager.clientOrThrow(); + let latestTs: number | undefined; + const stream = await apiClient.subscribeToSessionStream(chatId, "out", { + timeoutInSeconds: 1, + onPart: (part) => { + let chunk: unknown = part.chunk; + if (typeof chunk === "string") { + try { + chunk = JSON.parse(chunk); + } catch { + return; + } + } + if (chunk && typeof chunk === "object" && (chunk as { type?: unknown }).type === "trigger:turn-complete") { + latestTs = part.timestamp; + } + }, + }); + // Drain the stream to drive `onPart`. We don't accumulate the chunks — + // each iteration discards the data immediately, so a long session.out + // doesn't blow memory on the retry-boot worker. + for await (const _ of stream) { + // intentionally empty + } + return latestTs; +} + +/** + * Versioned blob written to S3 after every turn completes (when no + * `hydrateMessages` hook is registered). Read at run boot to seed the + * accumulator with prior conversation state, replacing the old wire-borne + * full-history seed. Only the runtime owns this format — customers never + * touch it. + * + * `lastOutEventId` is the SSE Last-Event-ID after the snapshot's final + * chunk, used to resume `session.out` replay from precisely after the + * snapshot. `lastOutTimestamp` is the same chunk's timestamp, used to + * skip `findLatestTurnCompleteTimestamp` on OOM retry boot. + * + * @internal + */ +export type ChatSnapshotV1 = { + version: 1; + savedAt: number; + messages: TUIMessage[]; + lastOutEventId?: string; + lastOutTimestamp?: number; +}; + +/** + * S3 key suffix for a session's snapshot blob. The webapp's presigned-URL + * routes prefix this with `packets/{projectRef}/{envSlug}/` server-side, so + * the final S3 key lands at + * `packets/{projectRef}/{envSlug}/sessions/{sessionId}/snapshot.json`. + * + * Stable per session: the friendlyId persists across `chat.requestUpgrade` + * continuations and idle-suspend restarts. + * @internal + */ +function snapshotFilename(sessionId: string): string { + return `sessions/${sessionId}/snapshot.json`; +} + +/** + * Test-only override hook — `mockChatAgent` installs a fake to return + * synthetic snapshots without hitting S3. Mirrors the `__set*ImplForTests` + * pattern in `sessions.ts`. Not part of the public API. + * @internal + */ +type ReadChatSnapshotImpl = ( + sessionId: string +) => Promise | undefined> | ChatSnapshotV1 | undefined; +let readChatSnapshotImpl: ReadChatSnapshotImpl | undefined; + +export function __setReadChatSnapshotImplForTests(impl: ReadChatSnapshotImpl | undefined): void { + readChatSnapshotImpl = impl; +} + +/** + * Test-only override hook — see `__setReadChatSnapshotImplForTests`. The + * mock harness records writes for assertion via this setter. Not public. + * @internal + */ +type WriteChatSnapshotImpl = ( + sessionId: string, + snapshot: ChatSnapshotV1 +) => Promise | void; +let writeChatSnapshotImpl: WriteChatSnapshotImpl | undefined; + +export function __setWriteChatSnapshotImplForTests(impl: WriteChatSnapshotImpl | undefined): void { + writeChatSnapshotImpl = impl; +} + +/** + * Read the persisted snapshot for a session. Returns `undefined` on: + * - missing object (404 from the presigned GET — fresh session, never + * persisted) + * - presign failure (network/auth issue) + * - malformed JSON + * - version mismatch (forward-compat — older runtimes ignore newer blobs) + * + * Always swallows errors via `logger.warn`. The agent boot loop must stay + * available even if S3 hiccups; the worst case is replaying more of + * `session.out` than strictly necessary. + * @internal + */ +async function readChatSnapshot( + sessionId: string +): Promise | undefined> { + if (readChatSnapshotImpl) { + return (await readChatSnapshotImpl(sessionId)) ?? undefined; + } + const apiClient = apiClientManager.clientOrThrow(); + let presignedUrl: string; + try { + const resp = await apiClient.getPayloadUrl(snapshotFilename(sessionId)); + presignedUrl = resp.presignedUrl; + } catch (error) { + logger.warn("chat.agent: snapshot presign (read) failed; continuing without snapshot", { + error: error instanceof Error ? error.message : String(error), + sessionId, + }); + return undefined; + } + let response: Response; + try { + response = await fetch(presignedUrl, { method: "GET" }); + } catch (error) { + logger.warn("chat.agent: snapshot fetch failed; continuing without snapshot", { + error: error instanceof Error ? error.message : String(error), + sessionId, + }); + return undefined; + } + if (response.status === 404) { + // First-ever boot for this session — no snapshot yet. Caller falls + // through to replay-only. + return undefined; + } + if (!response.ok) { + logger.warn("chat.agent: snapshot fetch returned non-OK; continuing without snapshot", { + status: response.status, + sessionId, + }); + return undefined; + } + let parsed: unknown; + try { + parsed = await response.json(); + } catch (error) { + logger.warn("chat.agent: snapshot JSON parse failed; continuing without snapshot", { + error: error instanceof Error ? error.message : String(error), + sessionId, + }); + return undefined; + } + if (!parsed || typeof parsed !== "object") return undefined; + const candidate = parsed as Partial>; + if (candidate.version !== 1 || !Array.isArray(candidate.messages)) { + logger.warn("chat.agent: snapshot version/shape mismatch; ignoring", { + version: candidate.version, + sessionId, + }); + return undefined; + } + return candidate as ChatSnapshotV1; +} + +/** + * Persist the snapshot for a session. Awaited by callers immediately after + * `onTurnComplete` — the agent may suspend right after this point, and + * fire-and-forget promises don't reliably complete on suspend. + * + * Errors are swallowed via `logger.warn`. A failed write means the next + * boot replays slightly more of `session.out` (back to the previous + * snapshot's cursor) instead of failing — the conversation stays + * coherent, only the boot path does marginally more work. + * @internal + */ +async function writeChatSnapshot( + sessionId: string, + snapshot: ChatSnapshotV1 +): Promise { + if (writeChatSnapshotImpl) { + await writeChatSnapshotImpl(sessionId, snapshot); + return; + } + const apiClient = apiClientManager.clientOrThrow(); + let presignedUrl: string; + try { + const resp = await apiClient.createUploadPayloadUrl(snapshotFilename(sessionId)); + presignedUrl = resp.presignedUrl; + } catch (error) { + logger.warn("chat.agent: snapshot presign (write) failed; next run will replay further", { + error: error instanceof Error ? error.message : String(error), + sessionId, + }); + return; + } + let response: Response; + try { + response = await fetch(presignedUrl, { + method: "PUT", + headers: { "content-type": "application/json" }, + body: JSON.stringify(snapshot), + }); + } catch (error) { + logger.warn("chat.agent: snapshot upload failed; next run will replay further", { + error: error instanceof Error ? error.message : String(error), + sessionId, + }); + return; + } + if (!response.ok) { + logger.warn("chat.agent: snapshot upload returned non-OK; next run will replay further", { + status: response.status, + sessionId, + }); + } +} + +/** + * Test-only entry point that bypasses `__setReadChatSnapshotImplForTests` + * and reaches the real `apiClient.getPayloadUrl` + `fetch` + JSON-parse path. + * Used by `chat-snapshot.test.ts` to verify 404 / 500 / malformed JSON / + * version-mismatch / network-error behavior end-to-end. Tests mock global + * `fetch` and the api-client config; this wrapper lets them drive the + * production code without the override hook short-circuiting. + * + * Not part of the public API. The `__` prefix and `ForTests` suffix mirror + * the override-hook setters above. + * @internal + */ +export async function __readChatSnapshotProductionPathForTests( + sessionId: string +): Promise | undefined> { + const saved = readChatSnapshotImpl; + readChatSnapshotImpl = undefined; + try { + return await readChatSnapshot(sessionId); + } finally { + readChatSnapshotImpl = saved; + } +} + +/** + * Test-only entry point that bypasses `__setWriteChatSnapshotImplForTests` + * and reaches the real `apiClient.createUploadPayloadUrl` + `fetch` PUT + * path. Pairs with `__readChatSnapshotProductionPathForTests` — see that + * function's note for the rationale. + * + * Not part of the public API. + * @internal + */ +export async function __writeChatSnapshotProductionPathForTests( + sessionId: string, + snapshot: ChatSnapshotV1 +): Promise { + const saved = writeChatSnapshotImpl; + writeChatSnapshotImpl = undefined; + try { + await writeChatSnapshot(sessionId, snapshot); + } finally { + writeChatSnapshotImpl = saved; + } +} + +/** + * Merge two `UIMessage[]` lists by `id`, with the second list winning on + * collision. Used at run boot to combine the snapshot's persisted history + * with the replayed `session.out` tail — replay produces the freshest + * representation of any assistant message that landed after the snapshot's + * cursor, so it should overwrite the older copy from the snapshot. + * + * Order: items unique to `a` keep their original positions; items unique to + * `b` are appended at the end in their `b` order; collisions take `b`'s + * value but keep the position they had in `a`. + * + * @internal + */ +function mergeByIdReplaceWins( + a: TUIMessage[], + b: TUIMessage[] +): TUIMessage[] { + if (b.length === 0) return [...a]; + if (a.length === 0) return [...b]; + const indexById = new Map(); + for (let i = 0; i < a.length; i++) { + const id = a[i]!.id; + if (typeof id === "string" && id.length > 0) indexById.set(id, i); + } + const result = [...a]; + for (const next of b) { + const id = next.id; + if (typeof id === "string" && id.length > 0 && indexById.has(id)) { + result[indexById.get(id)!] = next; + } else { + const newIdx = result.length; + result.push(next); + if (typeof id === "string" && id.length > 0) indexById.set(id, newIdx); + } + } + return result; +} + +/** + * Test-only entry point for `mergeByIdReplaceWins`. The merge helper is the + * one piece of slim-wire boot logic that's purely functional, so it earns a + * direct unit test that exercises empty inputs, id collisions, no-id append, + * order preservation, and the replay-wins-on-collision invariant. Mirrors + * the `__*ProductionPathForTests` pattern used for the snapshot/replay + * helpers above. + * + * Not part of the public API. + * @internal + */ +export function __mergeByIdReplaceWinsForTests( + a: TUIMessage[], + b: TUIMessage[] +): TUIMessage[] { + return mergeByIdReplaceWins(a, b); +} + +/** + * Test-only override hook — `mockChatAgent` installs a fake replay that + * returns a synthetic `UIMessage[]` so unit tests can drive the boot loop + * without an SSE subscription. Mirrors the snapshot setters above. Not + * part of the public API. + * @internal + */ +type ReplaySessionOutTailImpl = ( + sessionId: string, + options?: { lastEventId?: string } +) => Promise; +let replaySessionOutTailImpl: ReplaySessionOutTailImpl | undefined; + +export function __setReplaySessionOutTailImplForTests( + impl: ReplaySessionOutTailImpl | undefined +): void { + replaySessionOutTailImpl = impl; +} + +/** + * Drain `session.out` from `lastEventId` (or the start) and reduce the + * remaining `UIMessageChunk`s back into `UIMessage[]`. Used at run boot to + * catch any chunks that landed AFTER the last persisted snapshot — typically + * the chunks from the turn whose `onTurnComplete` ran but whose snapshot + * write didn't make it to S3 before the run crashed / suspended. + * + * Implementation: + * 1. `apiClient.readSessionStreamRecords` — non-SSE, `wait=0` drain. + * Returns immediately with whatever records exist after the cursor. + * The previous SSE-subscribe path paid a fixed ~1s long-poll tax on + * every fresh chat (timeout duration on empty streams) — unacceptable + * for the first-message TTFC budget. + * 2. Filter out the agent's control chunks (`type: "trigger:*"`) — they + * ride on the same stream as the user-visible UIMessageChunks. + * 3. Split chunks at `start`/`finish` boundaries so each segment is a + * single message, then feed each segment through the AI SDK's + * `readUIMessageStream` reducer (the same one `useChat` uses on the + * browser side) and grab the final emitted snapshot. + * 4. The trailing message — if it never received a `finish` chunk — + * goes through `cleanupAbortedParts` so partial in-flight parts + * don't leak into the next turn's accumulator. Drop it entirely + * if cleanup empties it. + * + * Errors are propagated to the caller (the boot loop wraps in try/catch and + * `logger.warn`s); we don't swallow here so test code can observe failures + * directly. + * @internal + */ +async function replaySessionOutTail( + sessionId: string, + options?: { lastEventId?: string } +): Promise { + if (replaySessionOutTailImpl) { + return await replaySessionOutTailImpl(sessionId, options); + } + const apiClient = apiClientManager.clientOrThrow(); + const response = await apiClient.readSessionStreamRecords(sessionId, "out", { + afterEventId: options?.lastEventId, + }); + const collected: UIMessageChunk[] = []; + for (const record of response.records) { + // Each record's `data` is the JSON-encoded chunk body the agent + // wrote at append time. The records endpoint returns it as an + // opaque string so the parsing cost is paid here, not on the + // server's hot path. + let chunk: unknown; + try { + chunk = JSON.parse(record.data); + } catch { + continue; + } + if (!chunk || typeof chunk !== "object") continue; + const type = (chunk as { type?: unknown }).type; + if (typeof type !== "string") continue; + // Drop agent control chunks (`trigger:turn-complete`, `trigger:upgrade-required`, + // session-state telemetry, etc.). They ride the same stream but aren't part + // of the UIMessageChunk discriminated union and would confuse the reducer. + if (type.startsWith("trigger:")) continue; + collected.push(chunk as UIMessageChunk); + } + if (collected.length === 0) return []; + + // Split chunks into per-message segments. A `start` chunk demarcates the + // beginning of an assistant message; chunks before any `start` (rare — + // but possible if the stream begins mid-message after a resume) get + // bundled into a leading "implicit" segment so we don't drop them silently. + type Segment = { chunks: UIMessageChunk[]; closed: boolean }; + const segments: Segment[] = []; + let current: Segment | undefined; + for (const chunk of collected) { + if (chunk.type === "start") { + current = { chunks: [chunk], closed: false }; + segments.push(current); + continue; + } + if (!current) { + // Chunk arrived before any `start`. Synthesize a segment so the reducer + // has something to work with — `readUIMessageStream` tolerates a missing + // `start` because we pass `message: undefined`. + current = { chunks: [], closed: false }; + segments.push(current); + } + current.chunks.push(chunk); + if (chunk.type === "finish") { + current.closed = true; + current = undefined; + } + } + + const messages: TUIMessage[] = []; + for (let i = 0; i < segments.length; i++) { + const seg = segments[i]!; + const isTrailing = i === segments.length - 1 && !seg.closed; + const segmentStream = new ReadableStream({ + start(controller) { + for (const c of seg.chunks) controller.enqueue(c); + controller.close(); + }, + }); + let last: UIMessage | undefined; + try { + for await (const snapshot of readUIMessageStream({ stream: segmentStream })) { + last = snapshot; + } + } catch (error) { + // Reducer error — the segment is malformed. Skip it and keep going so a + // single corrupt chunk doesn't sink the entire replay. + logger.warn("chat.agent: replay reducer failed for segment; skipping", { + sessionId, + segmentIndex: i, + error: error instanceof Error ? error.message : String(error), + }); + continue; + } + if (!last) continue; + if (isTrailing) { + const cleaned = cleanupAbortedParts(last as TUIMessage); + if (cleaned.parts.length === 0) continue; + messages.push(cleaned); + } else { + messages.push(last as TUIMessage); + } + } + return messages; +} + +/** + * Test-only entry point that bypasses `__setReplaySessionOutTailImplForTests` + * and reaches the real `apiClient.subscribeToSessionStream` + chunk-segment + * splitter + `readUIMessageStream` reducer. Pairs with the snapshot + * production-path wrappers above. Lets `replay-session-out.test.ts` drive + * synthetic chunk sequences through the real reducer to lock down chunk- + * stream → `UIMessage[]` correctness — if the AI SDK's chunk semantics + * shift in a future version, the test catches it before customers do. + * + * Tests should mock `apiClient.subscribeToSessionStream` (e.g. via + * `vi.spyOn(apiClient, ...)`) to feed a `ReadableStream`. + * + * Not part of the public API. + * @internal + */ +export async function __replaySessionOutTailProductionPathForTests< + TUIMessage extends UIMessage, +>( + sessionId: string, + options?: { lastEventId?: string } +): Promise { + const saved = replaySessionOutTailImpl; + replaySessionOutTailImpl = undefined; + try { + return await replaySessionOutTail(sessionId, options); + } finally { + replaySessionOutTailImpl = saved; + } +} + +/** + * Resolve the Session handle for the current chat.agent run. Throws if + * called outside of a chat.agent `run()` — every internal consumer is + * inside the run, and every external consumer goes through the public + * `sessions.open(id)` entry point. + * @internal + */ +function getChatSession(): SessionHandle { + const handle = locals.get(chatSessionHandleKey); + if (!handle) { + throw new Error( + "chat.agent session handle is not initialized. This indicates a chat.agent helper was used outside of a chat.agent run, or the transport did not send a sessionId." + ); + } + return handle; +} + +/** + * Stamp `gen_ai.conversation.id` on the active span at chat-run boot. + * The run-level span is already alive when the run callback fires, so + * `TaskContextSpanProcessor.onStart` (which stamps subsequent spans + * automatically) won't catch it — set explicitly here. + */ +function stampConversationIdOnActiveSpan( + conversationId: string | undefined, + span = trace.getActiveSpan() +): void { + if (!span || !conversationId) return; + span.setAttribute(SemanticInternalAttributes.GEN_AI_CONVERSATION_ID, conversationId); +} type ToolResultContent = Array< | { - type: "text"; - text: string; - } + type: "text"; + text: string; + } | { - type: "image"; - data: string; - mimeType?: string; - } + type: "image"; + data: string; + mimeType?: string; + } >; export type ToolOptions = { experimental_toToolResultContent?: (result: TResult) => ToolResultContent; }; +/** Satisfies AI SDK `ToolSet` index signature alongside concrete `Tool` input/output types. */ +type ToolSetCompatible> = T & NonNullable; + +function assertTaskUsableAsTool(task: AnyTask): void { + if (("schema" in task && !task.schema) || ("jsonSchema" in task && !task.jsonSchema)) { + throw new Error( + "Cannot convert this task to to a tool because the task has no schema. Make sure to either use schemaTask or a task with an input jsonSchema." + ); + } +} + +/** + * Shared implementation: run a task as a tool invocation (`triggerAndSubscribe` + tool metadata). + * Used by {@link toolExecute} and the deprecated `ai.tool()` wrapper. + */ +function createTaskToolExecuteHandler< + TIdentifier extends string, + TTaskSchema extends TaskSchema | undefined = undefined, + TInput = void, + TOutput = unknown, +>( + task: TaskWithSchema | Task +): (input: unknown, toolOpts: ToolCallOptions | undefined) => Promise { + assertTaskUsableAsTool(task); + + return async function taskToolExecuteHandler( + input: unknown, + toolOpts: ToolCallOptions | undefined + ): Promise { + const toolMeta: ToolCallExecutionOptions = { + toolCallId: toolOpts?.toolCallId ?? "", + }; + if (toolOpts?.experimental_context !== undefined) { + try { + toolMeta.experimental_context = JSON.parse(JSON.stringify(toolOpts.experimental_context)); + } catch { + /* non-serializable */ + } + } + + const chatCtx = locals.get(chatTurnContextKey); + if (chatCtx) { + toolMeta.chatId = chatCtx.chatId; + toolMeta.turn = chatCtx.turn; + toolMeta.continuation = chatCtx.continuation; + toolMeta.clientData = chatCtx.clientData; + } + + const chatLocals: Record = {}; + for (const entry of chatLocalRegistry) { + const value = locals.get(entry.key); + if (value !== undefined) { + chatLocals[entry.id] = value; + } + } + if (Object.keys(chatLocals).length > 0) { + toolMeta.chatLocals = chatLocals; + } + + return await task + .triggerAndSubscribe(input as inferSchemaIn, { + metadata: { + [METADATA_KEY]: toolMeta as any, + }, + tags: toolOpts?.toolCallId ? [`toolCallId:${toolOpts.toolCallId}`] : undefined, + signal: toolOpts?.abortSignal, + }) + .unwrap(); + }; +} + +/** + * Returns an `execute` function for the AI SDK `tool()` helper (or any compatible tool definition). + * Preferred API for task-backed tools: the same Trigger wiring as the deprecated `ai.tool()` + * (`triggerAndSubscribe`, tool-call metadata, chat context, `chat.local` serialization) without + * building the tool object. You supply `description`, `inputSchema`, and any AI-SDK-only options + * (e.g. `experimental_toToolResultContent`) on `tool()` yourself. + * + * @example + * ```ts + * import { tool } from "ai"; + * import { z } from "zod"; + * import { ai } from "@trigger.dev/sdk/ai"; + * import { myTask } from "./trigger/myTask"; + * + * export const myTool = tool({ + * description: myTask.description ?? "", + * inputSchema: z.object({ id: z.string() }), + * execute: ai.toolExecute(myTask), + * }); + * ``` + */ +function toolExecute( + task: Task +): (input: TInput, toolOpts: ToolCallOptions) => Promise; +function toolExecute< + TIdentifier extends string, + TTaskSchema extends TaskSchema | undefined = undefined, + TOutput = unknown, +>( + task: TaskWithSchema +): (input: inferSchemaIn, toolOpts: ToolCallOptions) => Promise; +function toolExecute< + TIdentifier extends string, + TTaskSchema extends TaskSchema | undefined = undefined, + TInput = void, + TOutput = unknown, +>( + task: TaskWithSchema | Task +): ( + input: TTaskSchema extends TaskSchema ? inferSchemaIn : TInput, + toolOpts: ToolCallOptions +) => Promise { + return createTaskToolExecuteHandler(task) as ( + input: TTaskSchema extends TaskSchema ? inferSchemaIn : TInput, + toolOpts: ToolCallOptions + ) => Promise; +} + +/** + * @deprecated Use `tool()` from the `ai` package with `execute: ai.toolExecute(task)` instead. + * This helper may be removed in a future major release. + */ function toolFromTask( task: Task, options?: ToolOptions -): Tool; +): ToolSetCompatible>; +/** @deprecated Use `tool()` from `ai` with `execute: ai.toolExecute(task)`. */ function toolFromTask< TIdentifier extends string, TTaskSchema extends TaskSchema | undefined = undefined, @@ -40,7 +839,8 @@ function toolFromTask< >( task: TaskWithSchema, options?: ToolOptions -): Tool, TOutput>; +): ToolSetCompatible, TOutput>>; +/** @deprecated Use `tool()` from `ai` with `execute: ai.toolExecute(task)`. */ function toolFromTask< TIdentifier extends string, TTaskSchema extends TaskSchema | undefined = undefined, @@ -49,35 +849,41 @@ function toolFromTask< >( task: TaskWithSchema | Task, options?: ToolOptions -): TTaskSchema extends TaskSchema - ? Tool, TOutput> - : Tool { - if (("schema" in task && !task.schema) || ("jsonSchema" in task && !task.jsonSchema)) { - throw new Error( - "Cannot convert this task to to a tool because the task has no schema. Make sure to either use schemaTask or a task with an input jsonSchema." - ); +): ToolSetCompatible< + TTaskSchema extends TaskSchema ? Tool, TOutput> : Tool +> { + const executeFromTaskInput = createTaskToolExecuteHandler(task); + + // Zod-backed tasks: use static `tool()` so runtime shape matches `ToolSet`. Generic task context + // prevents `tool()` overloads from inferring input; `as any` is localized to this call only. + if ("schema" in task && task.schema && isSchemaZodEsque(task.schema)) { + const staticTool = aiTool({ + description: task.description ?? "", + inputSchema: zodSchema(task.schema as any), + execute: async (input: unknown, toolOpts: ToolCallOptions) => + executeFromTaskInput(input, toolOpts), + ...(options?.experimental_toToolResultContent !== undefined + ? { experimental_toToolResultContent: options.experimental_toToolResultContent } + : {}), + } as any); + return staticTool as unknown as ToolSetCompatible< + TTaskSchema extends TaskSchema ? Tool, TOutput> : Tool + >; } const toolDefinition = dynamicTool({ description: task.description, inputSchema: convertTaskSchemaToToolParameters(task), - execute: async (input, options) => { - const serializedOptions = options ? JSON.parse(JSON.stringify(options)) : undefined; - - return await task - .triggerAndWait(input as inferSchemaIn, { - metadata: { - [METADATA_KEY]: serializedOptions, - }, - }) - .unwrap(); - }, - ...options, + ...(options?.experimental_toToolResultContent !== undefined + ? { experimental_toToolResultContent: options.experimental_toToolResultContent } + : {}), + execute: async (input: unknown, toolOpts: ToolCallOptions) => + executeFromTaskInput(input, toolOpts), }); - return toolDefinition as TTaskSchema extends TaskSchema - ? Tool, TOutput> - : Tool; + return toolDefinition as unknown as ToolSetCompatible< + TTaskSchema extends TaskSchema ? Tool, TOutput> : Tool + >; } function getToolOptionsFromMetadata(): ToolCallExecutionOptions | undefined { @@ -88,6 +894,61 @@ function getToolOptionsFromMetadata(): ToolCallExecutionOptions | undefined { return tool as ToolCallExecutionOptions; } +/** + * Get the current tool call ID from inside a subtask invoked via `ai.toolExecute()` (or legacy `ai.tool()`). + * Returns `undefined` if not running as a tool subtask. + */ +function getToolCallId(): string | undefined { + return getToolOptionsFromMetadata()?.toolCallId; +} + +/** + * Get the chat context from inside a subtask invoked via `ai.toolExecute()` (or legacy `ai.tool()`) within a `chat.agent`. + * Pass `typeof yourChatTask` as the type parameter to get typed `clientData`. + * Returns `undefined` if the parent is not a chat task. + * + * @example + * ```ts + * const ctx = ai.chatContext(); + * // ctx?.clientData is typed based on myChat's clientDataSchema + * ``` + */ +function getToolChatContext(): + | ChatTurnContext> + | undefined { + const opts = getToolOptionsFromMetadata(); + if (!opts?.chatId) return undefined; + return { + chatId: opts.chatId, + turn: opts.turn ?? 0, + continuation: opts.continuation ?? false, + clientData: opts.clientData as InferChatClientData, + }; +} + +/** + * Get the chat context from inside a subtask, throwing if not in a chat context. + * Pass `typeof yourChatTask` as the type parameter to get typed `clientData`. + * + * @example + * ```ts + * const ctx = ai.chatContextOrThrow(); + * // ctx.chatId, ctx.clientData are guaranteed non-null + * ``` + */ +function getToolChatContextOrThrow(): ChatTurnContext< + InferChatClientData +> { + const ctx = getToolChatContext(); + if (!ctx) { + throw new Error( + "ai.chatContextOrThrow() called outside of a chat.agent context. " + + "This helper can only be used inside a subtask invoked via ai.toolExecute() (or legacy ai.tool()) from a chat.agent." + ); + } + return ctx; +} + function convertTaskSchemaToToolParameters( task: AnyTask | TaskWithSchema ): Schema { @@ -113,6 +974,7805 @@ function convertTaskSchemaToToolParameters( } export const ai = { + /** + * @deprecated Use `tool()` from the `ai` package with `execute: ai.toolExecute(task)` instead. + */ tool: toolFromTask, + /** + * Preferred: return value for the `execute` field of AI SDK `tool()`. Keeps Trigger subtask and + * metadata behavior without coupling to a specific `ai` version’s `Tool` / `ToolSet` types. + */ + toolExecute, currentToolOptions: getToolOptionsFromMetadata, + /** Get the tool call ID from inside a subtask invoked via `ai.toolExecute()` (or legacy `ai.tool()`). */ + toolCallId: getToolCallId, + /** Get chat context (chatId, turn, clientData, etc.) from inside a subtask of a `chat.agent`. Returns undefined if not in a chat context. */ + chatContext: getToolChatContext, + /** Get chat context or throw if not in a chat context. Pass `typeof yourChatTask` for typed clientData. */ + chatContextOrThrow: getToolChatContextOrThrow, +}; + +/** + * Creates a public access token for a chat task. + * + * This is a convenience helper that creates a multi-use trigger public token + * scoped to the given task. Use it in a server action to provide the frontend + * `TriggerChatTransport` with an `accessToken`. + * + * @example + * ```ts + * // actions.ts + * "use server"; + * import { chat } from "@trigger.dev/sdk/ai"; + * import type { myChat } from "@/trigger/chat"; + * + * export const getChatToken = () => chat.createAccessToken("my-chat"); + * ``` + */ +function createChatAccessToken( + taskId: TaskIdentifier +): Promise { + return auth.createTriggerPublicToken(taskId as string, { expirationTime: "24h" }); +} + +// --------------------------------------------------------------------------- +// Chat transport helpers — backend side +// --------------------------------------------------------------------------- + +/** + * Typed chat output stream — `.writer()`, `.pipe()`, `.append()`, and + * `.read()` methods pre-bound to this run's Session `.out` channel and + * typed to `UIMessageChunk`. + * + * Use from within a `chat.agent` run to write custom chunks: + * ```ts + * const { waitUntilComplete } = chat.stream.writer({ + * execute: ({ write }) => { + * write({ type: "text-start", id: "status-1" }); + * write({ type: "text-delta", id: "status-1", delta: "Processing..." }); + * write({ type: "text-end", id: "status-1" }); + * }, + * }); + * await waitUntilComplete(); + * ``` + * + * Backed by the Session primitive so a chat's output outlives any single + * run — subscribers (browser transport, server-side `ChatStream`) read + * the session's `.out`, not a per-run stream. Run-scoped `target` + * options on `.pipe()` are honoured as no-ops; the session is the target. + */ +const chatStream: RealtimeDefinedStream = { + // Stable opaque label for the run-scoped `RealtimeDefinedStream` shape. + // `chatStream` is backed by the Session's `.out` channel — this id is + // not the real addressing key (the session is). Kept as a literal so + // the facade type stays satisfied without re-introducing a top-level + // constant; dashboards/telemetry that already read "chat" keep working. + id: "chat", + pipe(value, options) { + const { target: _target, ...sessionOptions } = (options ?? {}) as PipeStreamOptions; + return getChatSession().out.pipe( + value, + sessionOptions as SessionPipeStreamOptions + ); + }, + async read(_runId, options) { + // Session channels don't need a runId — the session is the address. + // Keep the signature for backward compatibility with the run-scoped + // RealtimeDefinedStream shape, but ignore the argument. + return getChatSession().out.read( + options as SessionSubscribeOptions | undefined + ); + }, + async append(value, options) { + const { target: _target, ...sessionOptions } = (options ?? {}) as AppendStreamOptions; + return getChatSession().out.append(value, sessionOptions as SessionPipeStreamOptions); + }, + writer(options) { + return getChatSession().out.writer(options); + }, +}; + +// --------------------------------------------------------------------------- +// chat.response — write data parts that persist to the response message +// --------------------------------------------------------------------------- + +/** + * Write data parts that both stream to the frontend AND persist in + * `onTurnComplete`'s `responseMessage` and `uiMessages`. + * + * Non-transient data chunks (`type` starts with `data-`, no `transient: true`) + * are queued for accumulation into the assistant response message. + * Transient or non-data chunks are streamed only (same as `chat.stream`). + * + * @example + * ```ts + * // Persists to responseMessage.parts + * chat.response.write({ type: "data-handover", data: { context: summary } }); + * + * // Transient — streams only, not in responseMessage + * chat.response.write({ type: "data-progress", data: { percent: 50 }, transient: true }); + * ``` + */ +const chatResponse = { + /** + * Write a single chunk. Non-transient data parts are accumulated into the + * response message; everything else is stream-only. + */ + write(part: UIMessageChunk): void { + queueResponsePart(part); + const { waitUntilComplete } = chatStream.writer({ + spanName: "chat.response.write", + collapsed: true, + execute: ({ write }) => { + write(part); + }, + }); + waitUntilComplete().catch(() => {}); + }, +}; + +// --------------------------------------------------------------------------- +// chat.store — typed, bidirectional shared data between agent and clients +// --------------------------------------------------------------------------- + +/** + * Listener fired when the store value changes. `operations` is present for + * `patch()` updates and absent for `set()` (which is a full snapshot). + */ +export type ChatStoreChangeListener = ( + value: TStore, + operations?: ChatStorePatchOperation[] +) => void; + +/** + * @internal Holder for the current store value. We wrap in an object so + * `undefined` (cleared) is distinguishable from "never set". + */ +type ChatStoreSlot = { value: unknown }; + +/** @internal */ +const chatStoreSlotKey = locals.create("chat.store.slot"); + +/** @internal */ +const chatStoreListenersKey = locals.create>( + "chat.store.listeners" +); + +/** @internal — write a store chunk onto the chat output stream. */ +function writeStoreChunk(chunk: ChatStoreChunk): void { + const { waitUntilComplete } = chatStream.writer({ + spanName: chunk.type === "store-snapshot" ? "chat.store.set" : "chat.store.patch", + collapsed: true, + execute: ({ write }) => { + write(chunk as unknown as UIMessageChunk); + }, + }); + waitUntilComplete().catch(() => {}); +} + +/** @internal — fire all listeners, swallowing per-listener errors. */ +function fireStoreListeners( + value: unknown, + operations?: ChatStorePatchOperation[] +): void { + const listeners = locals.get(chatStoreListenersKey); + if (!listeners || listeners.size === 0) return; + for (const listener of listeners) { + try { + listener(value, operations); + } catch { + // non-fatal — listener errors don't break the agent + } + } +} + +/** + * Replace the entire store value with `value`. Emits a `store-snapshot` + * chunk on the chat output stream and fires all `onChange` listeners. + */ +function chatStoreSet(value: TStore): void { + locals.set(chatStoreSlotKey, { value }); + writeStoreChunk({ type: "store-snapshot", value } satisfies ChatStoreSnapshotChunk); + fireStoreListeners(value); +} + +/** + * Apply RFC 6902 JSON Patch operations to the current store value. + * Emits a `store-delta` chunk on the chat output stream and fires all + * `onChange` listeners with the new value and the operations. + */ +function chatStorePatch(operations: ChatStorePatchOperation[]): void { + const slot = locals.get(chatStoreSlotKey); + const current = slot?.value; + const next = applyChatStorePatch(current, operations); + locals.set(chatStoreSlotKey, { value: next }); + writeStoreChunk({ + type: "store-delta", + operations, + } satisfies ChatStoreDeltaChunk); + fireStoreListeners(next, operations); +} + +/** Get the current store value. Returns `undefined` if no value has been set. */ +function chatStoreGet(): TStore | undefined { + return locals.get(chatStoreSlotKey)?.value as TStore | undefined; +} + +/** + * Subscribe to store changes for the current run. Returns an + * unsubscribe function. + */ +function chatStoreOnChange( + listener: ChatStoreChangeListener +): () => void { + let listeners = locals.get(chatStoreListenersKey); + if (!listeners) { + listeners = new Set(); + locals.set(chatStoreListenersKey, listeners); + } + listeners.add(listener as ChatStoreChangeListener); + return () => { + listeners!.delete(listener as ChatStoreChangeListener); + }; +} + +/** + * @internal — set the value without emitting a chunk. Used when applying + * `hydrateStore` results / `incomingStore` at turn start; the emitted + * snapshot is written separately so we don't double-emit. + */ +function chatStoreSetSilent(value: unknown): void { + locals.set(chatStoreSlotKey, { value }); +} + +/** + * @internal — emit the current value as a snapshot without touching the + * slot. Used at turn start after hydration so clients observing the stream + * see the initial value. + */ +function chatStoreEmitSnapshot(value: unknown): void { + writeStoreChunk({ type: "store-snapshot", value } satisfies ChatStoreSnapshotChunk); +} + +// --------------------------------------------------------------------------- +// ChatWriter — stream writer for callbacks +// --------------------------------------------------------------------------- + +/** + * A stream writer passed to chat lifecycle callbacks (`onPreload`, `onChatStart`, + * `onTurnStart`, `onTurnComplete`, `onCompacted`). + * + * Write custom `UIMessageChunk` parts (e.g. `data-*` parts) directly to the chat + * stream without the ceremony of `chat.stream.writer({ execute })`. + * + * The writer is lazy — no stream overhead if you don't call `write()` or `merge()`. + * + * @example + * ```ts + * onTurnStart: async ({ writer }) => { + * writer.write({ type: "data-status", data: { loading: true } }); + * }, + * onTurnComplete: async ({ writer, uiMessages }) => { + * writer.write({ type: "data-analytics", data: { messageCount: uiMessages.length } }); + * }, + * ``` + */ +export type ChatWriter = { + /** Write a single UIMessageChunk to the chat stream. */ + write(part: UIMessageChunk): void; + /** Merge another stream's chunks into the chat stream. */ + merge(stream: ReadableStream): void; +}; + +/** + * Creates a lazy ChatWriter that only opens a realtime stream on first use. + * Call `flush()` after the callback returns to await stream completion. + * @internal + */ +function createLazyChatWriter(): { writer: ChatWriter; flush: () => Promise } { + let writeImpl: ((part: UIMessageChunk) => void) | null = null; + let mergeImpl: ((stream: ReadableStream) => void) | null = null; + let waitPromise: (() => Promise) | null = null; + let resolveExecute: (() => void) | null = null; + + function ensureInitialized() { + if (writeImpl) return; + + const executePromise = new Promise((resolve) => { + resolveExecute = resolve; + }); + + const { waitUntilComplete } = chatStream.writer({ + collapsed: true, + spanName: "callback writer", + execute: ({ write, merge }) => { + writeImpl = write; + mergeImpl = merge; + return executePromise; // Keep execute alive until flush() + }, + }); + waitPromise = waitUntilComplete; + } + + return { + writer: { + write(part: UIMessageChunk) { + ensureInitialized(); + queueResponsePart(part); + writeImpl!(part); + }, + merge(stream: ReadableStream) { + ensureInitialized(); + mergeImpl!(stream); + }, + }, + async flush() { + if (resolveExecute) { + resolveExecute(); // Signal execute to complete + await waitPromise!(); // Wait for stream to finish piping + } + }, + }; +} + +/** + * Runs a callback with a lazy ChatWriter, flushing the stream after completion. + * @internal + */ +async function withChatWriter(fn: (writer: ChatWriter) => Promise | T): Promise { + const { writer, flush } = createLazyChatWriter(); + const result = await fn(writer); + await flush(); + return result; +} + +// `ChatTaskWirePayload` and `ChatInputChunk` live in `./ai-shared.ts` so +// browser bundles (which import them via `chat-client.ts` / `chat.ts`) +// can pull the types without dragging `ai.ts` into the client graph. +// Re-exported here so `@trigger.dev/sdk/ai` consumers see them. +import type { ChatTaskWirePayload, ChatInputChunk } from "./ai-shared.js"; +export type { ChatTaskWirePayload, ChatInputChunk } from "./ai-shared.js"; + +/** + * The payload shape passed to the `chatAgent` run function. + * + * - `messages` contains model-ready messages (converted via `convertToModelMessages`) — + * pass these directly to `streamText`. + * - `clientData` contains custom data from the frontend (the `metadata` field from `sendMessage()`). + * + * The backend accumulates the full conversation history across turns, so the frontend + * only needs to send new messages after the first turn. + */ +export type ChatTaskPayload = { + /** Model-ready messages — pass directly to `streamText({ messages })`. */ + messages: ModelMessage[]; + + /** The unique identifier for the chat session */ + chatId: string; + + /** + * The trigger type: + * - `"submit-message"`: A new user message + * - `"regenerate-message"`: Regenerate the last assistant response + * - `"preload"`: Run was preloaded before the first message (only on turn 0) + * - `"action"`: A typed action from the frontend (see `actionSchema` + `onAction`). + * The action has already been applied before `run()` fires — check `trigger === "action"` + * to short-circuit the LLM call when an action doesn't need a response. + * - `"close"`: The chat session is being closed (internal; `run()` is not called). + */ + trigger: "submit-message" | "regenerate-message" | "preload" | "action" | "close"; + + /** The ID of the message to regenerate (only for `"regenerate-message"`) */ + messageId?: string; + + /** Custom data from the frontend (passed via `metadata` on `sendMessage()` or the transport). */ + clientData?: TClientData; + + /** Whether this run is continuing an existing chat (previous run timed out or was cancelled). False for brand new chats. */ + continuation: boolean; + /** The run ID of the previous run (only set when `continuation` is true). */ + previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; + /** + * The friendlyId of the Session primitive backing this chat. Use with + * `sessions.open(sessionId)` when you need direct access to the session's + * `.in` / `.out` channels outside the hooks the agent already wires for + * you. Undefined only for legacy transports that predate the sessions + * migration. + */ + sessionId?: string; +}; + +/** + * Abort signals provided to the `chatAgent` run function. + */ +export type ChatTaskSignals = { + /** Combined signal — fires on run cancel OR stop generation. Pass to `streamText`. */ + signal: AbortSignal; + /** Fires only when the run is cancelled, expired, or exceeds maxDuration. */ + cancelSignal: AbortSignal; + /** Fires only when the frontend stops generation for this turn (per-turn, reset each turn). */ + stopSignal: AbortSignal; +}; + +/** + * The full payload passed to a `chatAgent` run function. + * Extends `ChatTaskPayload` (the wire payload) with abort signals. + */ +export type ChatTaskRunPayload = ChatTaskPayload & + ChatTaskSignals & { + /** + * Task run context — same object as the `ctx` passed to a standard `task({ run })` handler’s second argument. + * Use for tags, metadata, parent run links, or any API that needs the full run record. + */ + ctx: TaskRunContext; + /** Token usage from the previous turn. Undefined on turn 0. */ + previousTurnUsage?: LanguageModelUsage; + /** Cumulative token usage across all completed turns so far. */ + totalUsage: LanguageModelUsage; + }; + +// Input streams for bidirectional chat communication +// +// Both `messagesInput` and `stopInput` are thin facades over the current +// run's Session `.in` channel. The Session carries a single tagged stream +// (`ChatInputChunk`); these facades filter by `kind` so existing call +// sites (both internal and exposed via `chat.messages` / `chat.createStopSignal`) +// keep their original shape. Each accessor resolves the session handle +// lazily via `getChatSession()` so the module-level references stay +// compatible with the pre-migration wiring. +const messagesInput: RealtimeDefinedInputStream = { + id: "chat-messages", + on(handler) { + return getChatSession().in.on((chunk) => { + if (chunk.kind === "message") { + return handler(chunk.payload); + } + }); + }, + once(options) { + const ctx = taskContext.ctx; + const runId = ctx?.run.id; + + return new InputStreamOncePromise((resolve, reject) => { + tracer + .startActiveSpan( + options?.spanName ?? `chat.messages.once()`, + async () => { + while (true) { + const result = await getChatSession().in.once(options); + if (!result.ok) { + resolve(result as InputStreamOnceResult); + return; + } + if (result.output.kind === "message") { + resolve({ ok: true, output: result.output.payload }); + return; + } + // Non-message chunks (stops) are handled by the stopInput + // facade's persistent listener; loop and wait for the next. + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "streams", + [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream", + ...(runId + ? { + [SemanticInternalAttributes.ENTITY_ID]: `${runId}:chat-messages`, + } + : {}), + streamId: "chat-messages", + ...accessoryAttributes({ + items: [{ text: "chat-messages", variant: "normal" }], + style: "codepath", + }), + }, + } + ) + .catch(reject); + }); + }, + peek() { + const chunk = getChatSession().in.peek(); + if (chunk && chunk.kind === "message") return chunk.payload; + return undefined; + }, + wait(options) { + return new ManualWaitpointPromise(async (resolve, reject) => { + try { + while (true) { + const result = await getChatSession().in.wait(options); + if (!result.ok) { + resolve(result); + return; + } + if (result.output.kind === "message") { + resolve({ ok: true, output: result.output.payload }); + return; + } + // Stop chunks are handled by the stopInput facade's persistent + // listener; loop back into the suspending wait. + } + } catch (error) { + reject(error); + } + }); + }, + async waitWithIdleTimeout(options) { + while (true) { + const result = await getChatSession().in.waitWithIdleTimeout(options); + if (!result.ok) return result; + if (result.output.kind === "message") { + return { ok: true, output: result.output.payload }; + } + // Swallow stop-kind chunks — persistent stop listener already handled + // the abort; we just loop for the next message. + } + }, + async send(_runId, data, options) { + // The `runId` argument is kept for signature parity with + // `RealtimeDefinedInputStream` but ignored — sessions are addressed + // by sessionId, not runId. Callers producing messages from outside + // the run should prefer the transport's `session.in.send(...)` path. + await getChatSession().in.send( + { kind: "message", payload: data } satisfies ChatInputChunk, + options?.requestOptions + ); + }, +}; + +const stopInput: RealtimeDefinedInputStream<{ stop: true; message?: string }> = { + id: "chat-stop", + on(handler) { + return getChatSession().in.on((chunk) => { + if (chunk.kind === "stop") { + return handler({ stop: true, message: chunk.message }); + } + }); + }, + once(options) { + const ctx = taskContext.ctx; + const runId = ctx?.run.id; + + return new InputStreamOncePromise<{ stop: true; message?: string }>((resolve, reject) => { + tracer + .startActiveSpan( + options?.spanName ?? `chat.stop.once()`, + async () => { + while (true) { + const result = await getChatSession().in.once(options); + if (!result.ok) { + resolve(result as InputStreamOnceResult<{ stop: true; message?: string }>); + return; + } + if (result.output.kind === "stop") { + resolve({ + ok: true, + output: { stop: true, message: result.output.message }, + }); + return; + } + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "streams", + [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream", + ...(runId + ? { + [SemanticInternalAttributes.ENTITY_ID]: `${runId}:chat-stop`, + } + : {}), + streamId: "chat-stop", + ...accessoryAttributes({ + items: [{ text: "chat-stop", variant: "normal" }], + style: "codepath", + }), + }, + } + ) + .catch(reject); + }); + }, + peek() { + const chunk = getChatSession().in.peek(); + if (chunk && chunk.kind === "stop") { + return { stop: true, message: chunk.message }; + } + return undefined; + }, + wait(options) { + return new ManualWaitpointPromise<{ stop: true; message?: string }>(async (resolve, reject) => { + try { + while (true) { + const result = await getChatSession().in.wait(options); + if (!result.ok) { + resolve(result); + return; + } + if (result.output.kind === "stop") { + resolve({ + ok: true, + output: { stop: true, message: result.output.message }, + }); + return; + } + } + } catch (error) { + reject(error); + } + }); + }, + async waitWithIdleTimeout(options) { + while (true) { + const result = await getChatSession().in.waitWithIdleTimeout(options); + if (!result.ok) return result; + if (result.output.kind === "stop") { + return { ok: true, output: { stop: true, message: result.output.message } }; + } + } + }, + async send(_runId, data, options) { + await getChatSession().in.send( + { kind: "stop", message: data?.message } satisfies ChatInputChunk, + options?.requestOptions + ); + }, +}; + +/** + * Signal received by a `handover-prepare` agent run waiting on + * `session.in`. Either the customer's first-turn `streamText` finished + * with pending tool calls (`"handover"` — agent picks up from tool + * execution), or it finished pure-text (`"handover-skip"` — agent + * exits cleanly without making an LLM call). + * @internal + */ +type HandoverSignal = + | { + kind: "handover"; + partialAssistantMessage: ModelMessage[]; + messageId?: string; + /** + * Whether the customer's step 1 is the final response. When + * true, the agent's turn loop runs hooks but skips the LLM + * call (the partial IS the response). When false, the agent + * runs `streamText` which executes pending tool-calls via the + * approval round and continues from step 2. + */ + isFinal: boolean; + } + | { kind: "handover-skip" }; + +/** + * Internal facade for waiting on the handover signal. Mirrors + * `messagesInput` / `stopInput` so the wait paths and tracing + * attributes stay consistent across all input-stream branches. + * @internal + */ +const handoverInput = { + async waitWithIdleTimeout(options: { + idleTimeoutInSeconds: number; + timeout?: string; + spanName?: string; + skipSuspend?: boolean; + }) { + while (true) { + const result = await getChatSession().in.waitWithIdleTimeout(options); + if (!result.ok) return result; + if ( + result.output.kind === "handover" || + result.output.kind === "handover-skip" + ) { + return { ok: true as const, output: result.output as HandoverSignal }; + } + // Other kinds (message, stop) are not expected during handover-prepare. + // Loop back; the message and stop facades have their own listeners + // running so signals on those kinds aren't lost. + } + }, }; + +/** + * Per-turn deferred promises. Registered via `chat.defer()`, awaited + * before `onTurnComplete` fires. Reset each turn. + * @internal + */ +const chatDeferKey = locals.create>>("chat.defer"); + +/** + * Run-scoped slot holding the partial assistant message handed over by + * `chat.handover` from a customer's first-turn `streamText`. Appended + * to `accumulatedMessages` during turn 0 setup so `streamText` resumes + * at tool execution. Cleared (read once) after consumption. + * @internal + */ +const chatHandoverPartialKey = locals.create("chat.handoverPartial"); + +/** + * Run-scoped slot holding the assistant `messageId` the customer's + * `chat.handover` handler used for its step-1 stream. The agent reuses + * it on the agent-side `toUIMessageStream` (and the synthesized + * partial UIMessage in `originalMessages`) so all chunks merge into a + * single assistant message on the browser side. + * @internal + */ +const chatHandoverMessageIdKey = locals.create("chat.handoverMessageId"); + +/** + * Run-scoped slot indicating that the customer's step-1 head-start + * response is the FINAL turn response. When true, turn 0 runs through + * the full turn-loop hooks but SKIPS the `userRun` / `streamText` + * call — the customer's partial already IS the response. The agent's + * `onTurnComplete` fires with that partial so persistence + any + * post-turn work happens normally. Cleared after consumption. + * @internal + */ +const chatHandoverIsFinalKey = locals.create("chat.handoverIsFinal"); + +/** + * Build a UIMessage representation of a `chat.handover` partial so AI + * SDK's `processUIMessageStream` can transition `tool-output-available` + * chunks (emitted by the initial-tool-execution branch when the + * approval round runs) onto the existing tool-call. Without this, + * `state.message.parts` is empty when the agent's `streamText` + * finishes, and AI SDK throws + * `UIMessageStreamError: No tool invocation found`. + * + * Only the assistant message matters — the synthesized + * `tool-approval-response` rows are AI-SDK-internal and don't need a + * UIMessage representation. We map: + * - `text` parts → `{ type: "text", text }` + * - `tool-call` parts → `{ type: "tool-${name}", toolCallId, + * state: "input-available", input }` + * - `tool-approval-request` parts → skipped (AI SDK derives the + * approval state from chunks during processing) + * + * @internal + */ +function synthesizeHandoverUIMessage( + partial: ModelMessage[], + messageId?: string +): UIMessage | undefined { + const assistant = partial.find((m) => m.role === "assistant"); + if (!assistant || typeof assistant.content === "string") return undefined; + + const parts: UIMessage["parts"] = []; + for (const part of assistant.content as Array<{ + type: string; + text?: string; + toolCallId?: string; + toolName?: string; + input?: unknown; + }>) { + if (part.type === "text" && typeof part.text === "string") { + parts.push({ type: "text", text: part.text } as UIMessage["parts"][number]); + } else if (part.type === "tool-call" && part.toolCallId && part.toolName) { + parts.push({ + type: `tool-${part.toolName}`, + toolCallId: part.toolCallId, + state: "input-available", + input: part.input, + } as unknown as UIMessage["parts"][number]); + } + // tool-approval-request parts intentionally skipped — they're an + // AI-SDK protocol detail, not a UI surface. + } + + if (parts.length === 0) return undefined; + + // Use the customer's step-1 messageId if provided (so the agent's + // post-handover chunks merge into the same assistant message on the + // browser). Fall back to a fresh id only if the handover signal + // didn't carry one. + return { + id: messageId ?? generateMessageId(), + role: "assistant", + parts, + } as UIMessage; +} + +/** + * Per-turn background context queue. Messages added via `chat.backgroundWork.inject()` + * are drained at the next `prepareStep` boundary and appended to the model messages. + * @internal + */ +const chatBackgroundQueueKey = locals.create("chat.backgroundQueue"); + +/** + * Run-scoped pipe counter. Stored in locals so concurrent runs in the + * same worker don't share state. + * @internal + */ +const chatPipeCountKey = locals.create("chat.pipeCount"); +const chatStopControllerKey = locals.create("chat.stopController"); +/** Static (task-level) UIMessageStream options, set once during chatAgent setup. @internal */ +const chatUIStreamStaticKey = locals.create>( + "chat.uiMessageStreamOptions.static" +); +/** Per-turn UIMessageStream options, set via chat.setUIMessageStreamOptions(). @internal */ +const chatUIStreamPerTurnKey = locals.create>( + "chat.uiMessageStreamOptions.perTurn" +); + +/** + * Run-scoped `toolCallId → assistant messageId` map. Records the head + * assistant id whenever the accumulator absorbs an assistant message + * containing tool parts. Used as a fallback in the id-merge for + * incoming tool-answer messages — if the AI SDK regenerates the + * assistant id on a HITL `addToolOutput` resume, we look up the + * original head id by `toolCallId` and rewrite it before the merge. + * + * Customer-side workaround for the same case is documented in Arena + * AI's chat-agent task; lifting it into the SDK so customers don't + * have to. See TRI-9137. + * @internal + */ +const chatToolCallToMessageIdKey = locals.create>( + "chat.toolCallToMessageId" +); + +function recordToolCallIdsFromMessage(message: { id?: string; role?: string; parts?: unknown[] } | undefined) { + if (!message || message.role !== "assistant" || !message.id) return; + let map = locals.get(chatToolCallToMessageIdKey); + if (!map) { + map = new Map(); + locals.set(chatToolCallToMessageIdKey, map); + } + for (const part of message.parts ?? []) { + if (typeof part !== "object" || part == null) continue; + const toolCallId = (part as { toolCallId?: unknown }).toolCallId; + if (typeof toolCallId === "string" && toolCallId.length > 0) { + map.set(toolCallId, message.id); + } + } +} + +function rewriteIncomingIdViaToolCallMap( + incoming: T +): T { + const map = locals.get(chatToolCallToMessageIdKey); + if (!map || map.size === 0) return incoming; + for (const part of incoming.parts ?? []) { + if (typeof part !== "object" || part == null) continue; + const toolCallId = (part as { toolCallId?: unknown }).toolCallId; + if (typeof toolCallId !== "string" || toolCallId.length === 0) continue; + const headId = map.get(toolCallId); + if (headId && headId !== incoming.id) { + return { ...incoming, id: headId }; + } + } + return incoming; +} + +// --------------------------------------------------------------------------- +// Token usage helpers (internal) +// --------------------------------------------------------------------------- + +/** Convenience re-export of the AI SDK's `LanguageModelUsage` type. */ +export type ChatTurnUsage = LanguageModelUsage; + +function emptyUsage(): LanguageModelUsage { + return { + inputTokens: undefined, + outputTokens: undefined, + totalTokens: undefined, + inputTokenDetails: { + noCacheTokens: undefined, + cacheReadTokens: undefined, + cacheWriteTokens: undefined, + }, + outputTokenDetails: { textTokens: undefined, reasoningTokens: undefined }, + }; +} + +function addUsage(a: LanguageModelUsage, b: LanguageModelUsage): LanguageModelUsage { + const add = (x: number | undefined, y: number | undefined) => + x != null || y != null ? (x ?? 0) + (y ?? 0) : undefined; + return { + inputTokens: add(a.inputTokens, b.inputTokens), + outputTokens: add(a.outputTokens, b.outputTokens), + totalTokens: add(a.totalTokens, b.totalTokens), + inputTokenDetails: { + noCacheTokens: add(a.inputTokenDetails?.noCacheTokens, b.inputTokenDetails?.noCacheTokens), + cacheReadTokens: add( + a.inputTokenDetails?.cacheReadTokens, + b.inputTokenDetails?.cacheReadTokens + ), + cacheWriteTokens: add( + a.inputTokenDetails?.cacheWriteTokens, + b.inputTokenDetails?.cacheWriteTokens + ), + }, + outputTokenDetails: { + textTokens: add(a.outputTokenDetails?.textTokens, b.outputTokenDetails?.textTokens), + reasoningTokens: add( + a.outputTokenDetails?.reasoningTokens, + b.outputTokenDetails?.reasoningTokens + ), + }, + }; +} + +// --------------------------------------------------------------------------- +// chat.setMessages — replace accumulated messages for compaction +// --------------------------------------------------------------------------- + +/** @internal */ +const chatOverrideMessagesKey = locals.create("chat.overrideMessages"); + +/** + * Tracks the current accumulated UI messages so chat.history.all() can + * read them from outside the chatAgent closure. + * @internal + */ +const chatCurrentUIMessagesKey = locals.create("chat.currentUIMessages"); + +/** + * Replace the accumulated conversation messages for the current run. + * + * Call from `onTurnStart` to compact before `run()` executes, or from + * `onTurnComplete` to compact before the next turn. Takes `UIMessage[]` + * and converts to `ModelMessage[]` internally. + */ +function setChatMessages(uiMessages: TUIM[]): void { + locals.set(chatOverrideMessagesKey, uiMessages); +} + +// --------------------------------------------------------------------------- +// chat.history — imperative message history mutations +// --------------------------------------------------------------------------- + +/** + * Read the current message history state, accounting for pending overrides. + * @internal + */ +function getChatHistoryState(): UIMessage[] { + const pending = locals.get(chatOverrideMessagesKey); + if (pending) return pending; + return locals.get(chatCurrentUIMessagesKey) ?? []; +} + +/** + * A tool call surfaced by `chat.history.getPendingToolCalls()` / + * `getResolvedToolCalls()`. Identifies the call by its `toolCallId` plus + * the `messageId` of the assistant message that hosts it, so callers can + * locate the part precisely without re-walking the chain. + */ +export type ChatToolCallRef = { + toolCallId: string; + toolName: string; + messageId: string; +}; + +/** + * A new tool result surfaced by `chat.history.extractNewToolResults()`. + * `errorText` is set iff the part is in `output-error` state; otherwise + * `output` carries the resolved value. + */ +export type ChatNewToolResult = { + toolCallId: string; + toolName: string; + output: unknown; + errorText?: string; +}; + +/** + * Tool parts that are "done" — either succeeded with a value or failed + * with an error. Excludes pending (`input-streaming`/`input-available`) + * and approval (`approval-requested`/`approval-responded`) states. + * @internal + */ +function isResolvedToolState(state: unknown): state is "output-available" | "output-error" { + return state === "output-available" || state === "output-error"; +} + +/** @internal */ +function isPendingToolState(state: unknown): state is "input-available" { + return state === "input-available"; +} + +/** + * Walk an assistant message and yield each tool part with its callId, + * name, and state. Skips non-assistant messages and non-tool parts. + * @internal + */ +function* iterateToolParts( + message: UIMessage +): Generator<{ part: any; toolCallId: string; toolName: string; state: unknown }> { + if (message.role !== "assistant") return; + for (const part of (message.parts ?? []) as any[]) { + if (!isToolUIPart(part)) continue; + const toolCallId = part.toolCallId; + if (typeof toolCallId !== "string" || toolCallId.length === 0) continue; + yield { + part, + toolCallId, + toolName: getToolName(part), + state: part.state, + }; + } +} + +/** + * Tool parts on the *leaf* assistant message that are still waiting on + * an answer (`input-available` state). Used to gate fresh user turns + * during HITL flows. + * @internal + */ +function getPendingToolCallsFromHistory(messages: UIMessage[]): ChatToolCallRef[] { + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]!; + if (msg.role !== "assistant") continue; + const pending: ChatToolCallRef[] = []; + for (const { toolCallId, toolName, state } of iterateToolParts(msg)) { + if (isPendingToolState(state)) { + pending.push({ toolCallId, toolName, messageId: msg.id }); + } + } + return pending; + } + return []; +} + +/** + * All tool parts across the chain that have already produced an output + * (`output-available` or `output-error`). Used to dedup re-saves when + * the AI SDK resends an assistant with progressively more answered + * parts. + * @internal + */ +function getResolvedToolCallsFromHistory(messages: UIMessage[]): ChatToolCallRef[] { + const out: ChatToolCallRef[] = []; + for (const msg of messages) { + for (const { toolCallId, toolName, state } of iterateToolParts(msg)) { + if (isResolvedToolState(state)) { + out.push({ toolCallId, toolName, messageId: msg.id }); + } + } + } + return out; +} + +/** + * Pure helper: tool parts in `message` that have a fresh result not + * already represented by the resolved toolCallIds in `messages`. The + * `errorText` field is present only for `output-error` parts. + * + * Within a single `message`, duplicate `toolCallId`s emit only once + * (first occurrence wins). This guards against malformed assistants + * with repeated tool parts. + * @internal + */ +function extractNewToolResultsFromHistory( + message: UIMessage, + messages: UIMessage[] +): ChatNewToolResult[] { + const resolved = new Set( + getResolvedToolCallsFromHistory(messages).map((r) => r.toolCallId) + ); + const seen = new Set(); + const out: ChatNewToolResult[] = []; + for (const { part, toolCallId, toolName, state } of iterateToolParts(message)) { + if (!isResolvedToolState(state)) continue; + if (resolved.has(toolCallId)) continue; + if (seen.has(toolCallId)) continue; + seen.add(toolCallId); + if (state === "output-error") { + out.push({ toolCallId, toolName, output: undefined, errorText: part.errorText }); + } else { + out.push({ toolCallId, toolName, output: part.output }); + } + } + return out; +} + +/** + * Imperative API for reading and modifying the accumulated message history. + * + * Mutations use the same deferred override mechanism as `chat.setMessages()`: + * they are applied at lifecycle checkpoints (after hooks return). Reads are + * synchronous against the current accumulator state. + * + * Can be called from `onTurnStart`, `onBeforeTurnComplete`, `onTurnComplete`, + * `run()`, `onAction`, or AI SDK tools. + */ +const chatHistory = { + /** Read the current accumulated UI messages (copy). */ + all(): UIMessage[] { + return [...getChatHistoryState()]; + }, + + /** + * Read the current chain as an ordered `UIMessage[]`. Identical to + * `all()`; use whichever name reads better in context. + */ + getChain(): UIMessage[] { + return chatHistory.all(); + }, + + /** + * Find a message by id. Returns `undefined` if no message with that id + * is present in the current chain. + */ + findMessage(messageId: string): UIMessage | undefined { + return getChatHistoryState().find((m) => m.id === messageId); + }, + + /** + * Tool calls on the *most recent* assistant message that are still in + * `input-available` state (waiting on an `addToolOutput` answer). The + * scan walks back from the tail and stops at the first assistant + * message it finds, so a trailing user message does not change the + * result — pending tool calls remain pending until they're resolved + * on that assistant or the assistant is removed. + * + * Use this to gate fresh user turns or actions during HITL flows: if + * `getPendingToolCalls().length > 0`, an `addToolOutput` is expected. + * + * Returns `[]` if there is no assistant message yet, or if the most + * recent assistant has no pending tool calls. + * + * Approval flows (`approval-requested` / `approval-responded` states) + * are not surfaced here. Those are about the user authorizing a tool + * to run; "pending" is about the user *answering* a tool call. + */ + getPendingToolCalls(): ChatToolCallRef[] { + return getPendingToolCallsFromHistory(getChatHistoryState()); + }, + + /** + * Tool calls across the chain with a final result (`output-available` + * or `output-error`). Use this to dedup re-saves when the AI SDK + * resends an assistant message with progressively more answered parts. + */ + getResolvedToolCalls(): ChatToolCallRef[] { + return getResolvedToolCallsFromHistory(getChatHistoryState()); + }, + + /** + * Pure helper: returns the tool parts in `message` whose results are + * not already represented in the current chain. Use this when + * persisting tool results to your own store: each call surfaces only + * the *new* answers, so writes stay idempotent across re-streams. + * Duplicate `toolCallId`s within `message` itself are also collapsed + * to a single entry. + */ + extractNewToolResults(message: UIMessage): ChatNewToolResult[] { + return extractNewToolResultsFromHistory(message, getChatHistoryState()); + }, + + /** Replace all accumulated messages. Same as `chat.setMessages()`. */ + set(messages: UIMessage[]): void { + locals.set(chatOverrideMessagesKey, messages); + }, + + /** Remove a specific message by ID. */ + remove(messageId: string): void { + chatHistory.set(getChatHistoryState().filter((m) => m.id !== messageId)); + }, + + /** Keep messages up to and including the given ID (undo/rollback). */ + rollbackTo(messageId: string): void { + const current = getChatHistoryState(); + const idx = current.findIndex((m) => m.id === messageId); + if (idx !== -1) { + chatHistory.set(current.slice(0, idx + 1)); + } + }, + + /** Replace a specific message by ID (edit). */ + replace(messageId: string, message: UIMessage): void { + chatHistory.set(getChatHistoryState().map((m) => (m.id === messageId ? message : m))); + }, + + /** Keep only messages in the given range. */ + slice(start: number, end?: number): void { + chatHistory.set(getChatHistoryState().slice(start, end)); + }, +}; + +/** + * Model-only message override. Set by compaction to replace only the model + * messages (what goes to the LLM) without affecting UI messages (what gets + * persisted and displayed). This preserves full conversation history for the + * user while keeping LLM context compact. + * @internal + */ +const chatOverrideModelMessagesKey = locals.create("chat.overrideModelMessages"); + +// --------------------------------------------------------------------------- +// chat.compaction — prepareStep compaction API +// --------------------------------------------------------------------------- + +/** State stored in locals during prepareStep compaction. */ +interface CompactionState { + summary: string; + baseResponseMessageCount: number; +} + +/** @internal */ +const chatCompactionStateKey = locals.create("chat.compaction"); +const chatOnCompactedKey = + locals.create<(event: CompactedEvent) => Promise | void>("chat.onCompacted"); +/** @internal Full task `ctx` for the active `chat.agent` run (for hooks invoked from nested compaction). */ +const chatAgentRunContextKey = locals.create("chat.agentRunContext"); +const chatPrepareMessagesKey = + locals.create<(event: PrepareMessagesEvent) => ModelMessage[] | Promise>( + "chat.prepareMessages" + ); + +/** @internal Flag set by `chat.requestUpgrade()` to exit the loop after the current turn. */ +const chatUpgradeRequestedKey = locals.create("chat.upgradeRequested"); + +/** + * @internal Flag set by `chat.endRun()` to exit the loop after the current + * turn completes, without any upgrade semantics. Checked at the same + * post-turn / pre-wait sites as `chatUpgradeRequestedKey`. + */ +const chatEndRunRequestedKey = locals.create("chat.endRunRequested"); + +/** + * Event passed to `summarize` callbacks. + */ +export type SummarizeEvent = { + /** The current model messages to summarize. */ + messages: ModelMessage[]; + /** Full usage object from the triggering step/turn. */ + usage?: LanguageModelUsage; + /** Cumulative token usage across all completed turns. Present in chat.agent contexts. */ + totalUsage?: LanguageModelUsage; + /** The chat session ID (if running inside a chat.agent). */ + chatId?: string; + /** The current turn number (0-indexed, if inside a chat.agent). */ + turn?: number; + /** Custom data from the frontend (if inside a chat.agent). */ + clientData?: unknown; + /** + * Where compaction is running: + * - `"inner"` — between tool-call steps (prepareStep) + * - `"outer"` — between turns + */ + source?: "inner" | "outer"; + /** The step number (0-indexed). Only present when `source` is `"inner"`. */ + stepNumber?: number; +}; + +/** + * Event passed to `compactUIMessages` and `compactModelMessages` callbacks. + */ +export type CompactMessagesEvent = { + /** The generated summary text. */ + summary: string; + /** The current UI messages (full conversation). */ + uiMessages: TUIM[]; + /** The current model messages (full conversation). */ + modelMessages: ModelMessage[]; + /** The chat session ID. */ + chatId: string; + /** The current turn number (0-indexed). */ + turn: number; + /** Custom data from the frontend. */ + clientData?: unknown; + /** + * Where compaction is running: + * - `"inner"` — between tool-call steps (prepareStep) + * - `"outer"` — between turns + */ + source: "inner" | "outer"; +}; + +/** + * Options for the `compaction` field on `chat.agent()`. + * + * Handles compaction automatically in both the inner loop (prepareStep, between + * tool-call steps) and the outer loop (between turns, for single-step responses + * where prepareStep never fires). + */ +export type ChatAgentCompactionOptions = { + /** Decide whether to compact. Return true to trigger compaction. */ + shouldCompact: (event: ShouldCompactEvent) => boolean | Promise; + /** Generate a summary from the current messages. Return the summary text. */ + summarize: (event: SummarizeEvent) => Promise; + /** + * Transform UI messages after compaction (what gets persisted and displayed). + * Default: preserve all UI messages unchanged. + * + * @example + * ```ts + * // Flatten to summary + * compactUIMessages: ({ summary }) => [{ + * id: generateId(), role: "assistant", + * parts: [{ type: "text", text: `[Summary]\n\n${summary}` }], + * }], + * + * // Summary + keep last 4 messages + * compactUIMessages: ({ uiMessages, summary }) => [ + * { id: generateId(), role: "assistant", + * parts: [{ type: "text", text: `[Summary]\n\n${summary}` }] }, + * ...uiMessages.slice(-4), + * ], + * ``` + */ + compactUIMessages?: (event: CompactMessagesEvent) => TUIM[] | Promise; + /** + * Transform model messages after compaction (what gets sent to the LLM). + * Default: replace all with a single summary message. + * + * @example + * ```ts + * // Summary + keep last 2 model messages + * compactModelMessages: ({ modelMessages, summary }) => [ + * { role: "user", content: summary }, + * ...modelMessages.slice(-2), + * ], + * ``` + */ + compactModelMessages?: ( + event: CompactMessagesEvent + ) => ModelMessage[] | Promise; +}; + +/** @internal */ +const chatAgentCompactionKey = + locals.create>("chat.agentCompaction"); + +// --------------------------------------------------------------------------- +// Pending messages — mid-execution message injection via prepareStep +// --------------------------------------------------------------------------- + +/** + * Event passed to `shouldInject` and `prepareMessages` callbacks. + */ +export type PendingMessagesBatchEvent = { + /** All pending UI messages that arrived during streaming (batch). */ + messages: TUIM[]; + /** Current model messages in the conversation. */ + modelMessages: ModelMessage[]; + /** Completed steps so far. */ + steps: CompactionStep[]; + /** Current step number (0-indexed). */ + stepNumber: number; + /** Chat session ID. */ + chatId: string; + /** Current turn number (0-indexed). */ + turn: number; + /** Custom data from the frontend. */ + clientData?: unknown; +}; + +/** + * Event passed to `onReceived` callback (per-message, as they arrive). + */ +export type PendingMessageReceivedEvent = { + /** The UI message that arrived during streaming. */ + message: TUIM; + /** Chat session ID. */ + chatId: string; + /** Current turn number (0-indexed). */ + turn: number; +}; + +/** + * Event passed to `onInjected` callback (batch, after injection). + */ +export type PendingMessagesInjectedEvent = { + /** All UI messages that were injected. */ + messages: TUIM[]; + /** The model messages that were injected. */ + injectedModelMessages: ModelMessage[]; + /** Chat session ID. */ + chatId: string; + /** Current turn number (0-indexed). */ + turn: number; + /** Step number where injection occurred. */ + stepNumber: number; +}; + +/** + * Options for the `pendingMessages` field on `chat.agent()`, `chat.createSession()`, + * or `ChatMessageAccumulator`. + * + * Configures how messages that arrive during streaming are handled. When + * `shouldInject` is provided and returns `true`, the full batch of pending + * messages is injected between tool-call steps via `prepareStep`. + * Otherwise, messages queue for the next turn. + */ +export type PendingMessagesOptions = { + /** + * Decide whether to inject pending messages between tool-call steps. + * Called once per step boundary with the full batch of pending messages. + * If absent, no injection happens — messages only queue for the next turn. + */ + shouldInject?: (event: PendingMessagesBatchEvent) => boolean | Promise; + /** + * Transform the batch of pending messages before injection. + * Return the model messages to inject. + * Default: convert each UI message via `convertToModelMessages`. + */ + prepare?: (event: PendingMessagesBatchEvent) => ModelMessage[] | Promise; + /** Called when a message arrives during streaming (per-message). */ + onReceived?: (event: PendingMessageReceivedEvent) => void | Promise; + /** Called after a batch of messages is injected via `prepareStep`. */ + onInjected?: (event: PendingMessagesInjectedEvent) => void | Promise; +}; + +/** + * The data part type used to signal that pending messages were injected + * between tool-call steps. The frontend can match on this to render + * injection points inline in the assistant response. + */ +// `PENDING_MESSAGE_INJECTED_TYPE` lives in `./ai-shared.ts` so the chat +// React hooks (`@trigger.dev/sdk/chat/react`) can import it without +// dragging `ai.ts` into the browser graph. Re-exported here so +// `@trigger.dev/sdk/ai` consumers still see it. +export { PENDING_MESSAGE_INJECTED_TYPE } from "./ai-shared.js"; +import { PENDING_MESSAGE_INJECTED_TYPE } from "./ai-shared.js"; + +/** @internal */ +type SteeringQueueEntry = { uiMessage: UIMessage; modelMessages: ModelMessage[] }; +/** @internal */ +const chatPendingMessagesKey = locals.create("chat.pendingMessages"); +/** @internal */ +const chatSteeringQueueKey = locals.create("chat.steeringQueue"); +/** @internal — IDs of messages that were successfully injected via prepareStep */ +const chatInjectedMessageIdsKey = locals.create>("chat.injectedMessageIds"); +/** @internal — non-transient data parts queued via chat.response or writer.write() for accumulation into the response message */ +const chatResponsePartsKey = locals.create("chat.responseParts"); + +/** + * Check if a chunk is a non-transient data part that should persist to the response message. + * @internal + */ +function isNonTransientDataPart(part: unknown): boolean { + if (typeof part !== "object" || part === null) return false; + const p = part as Record; + return typeof p.type === "string" && p.type.startsWith("data-") && p.transient !== true; +} + +/** + * Queue a chunk for accumulation into the response message (if it's a non-transient data part). + * Called by `chat.response.write()` and `ChatWriter.write()`. + * @internal + */ +function queueResponsePart(part: unknown): void { + if (!isNonTransientDataPart(part)) return; + const parts = locals.get(chatResponsePartsKey) ?? []; + parts.push(part); + locals.set(chatResponsePartsKey, parts); +} + +/** + * Event passed to the `prepareMessages` hook. + */ +export type PrepareMessagesEvent = { + /** The messages to transform. Return the transformed array. */ + messages: ModelMessage[]; + /** Why messages are being prepared. */ + reason: + | "run" // Messages being passed to run() for streamText + | "compaction-rebuild" // Rebuilding from a previous compaction summary + | "compaction-result"; // Fresh compaction just produced these messages + /** The chat session ID. */ + chatId: string; + /** The current turn number (0-indexed). */ + turn: number; + /** 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. + */ +export type CompactionChunkData = { + status: "compacting" | "complete"; + totalTokens: number | undefined; +}; + +/** + * Event passed to the `onCompacted` callback. + */ +export type CompactedEvent = { + /** Task run context — same as `task` lifecycle hooks and `chat.agent` `run({ ctx })`. */ + ctx: TaskRunContext; + /** The generated summary text. */ + summary: string; + /** The messages that were compacted (pre-compaction). */ + messages: ModelMessage[]; + /** Number of messages before compaction. */ + messageCount: number; + /** Token usage from the step that triggered compaction. */ + usage: LanguageModelUsage; + /** Total token count that triggered compaction. */ + totalTokens: number | undefined; + /** Input token count from the triggering step. */ + inputTokens: number | undefined; + /** Output token count from the triggering step. */ + outputTokens: number | undefined; + /** The step number where compaction occurred (0-indexed). */ + stepNumber: number; + /** The chat session ID (if running inside a chat.agent). */ + chatId?: string; + /** The current turn number (if running inside a chat.agent). */ + turn?: number; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; +}; + +/** + * Event passed to `shouldCompact` callbacks. + */ +export type ShouldCompactEvent = { + /** The current model messages (full conversation). */ + messages: ModelMessage[]; + /** Total token count from the triggering step/turn. */ + totalTokens: number | undefined; + /** Input token count from the triggering step/turn. */ + inputTokens: number | undefined; + /** Output token count from the triggering step/turn. */ + outputTokens: number | undefined; + /** Full usage object from the triggering step/turn. */ + usage?: LanguageModelUsage; + /** Cumulative token usage across all completed turns. Present in chat.agent contexts. */ + totalUsage?: LanguageModelUsage; + /** The chat session ID (if running inside a chat.agent). */ + chatId?: string; + /** The current turn number (0-indexed, if inside a chat.agent). */ + turn?: number; + /** Custom data from the frontend (if inside a chat.agent). */ + clientData?: unknown; + /** + * Where this check is running: + * - `"inner"` — between tool-call steps (prepareStep) + * - `"outer"` — between turns (after response, before onBeforeTurnComplete) + */ + source?: "inner" | "outer"; + /** The step number (0-indexed). Only present when `source` is `"inner"`. */ + stepNumber?: number; + /** The steps array from prepareStep. Only present when `source` is `"inner"`. */ + steps?: CompactionStep[]; +}; + +/** + * Options for `chat.compaction()` — the high-level prepareStep factory. + */ +export type CompactionOptions = { + /** Generate a summary from the current messages. Return the summary text. */ + summarize: (messages: ModelMessage[]) => Promise; + /** Token threshold — compact when totalTokens exceeds this. Ignored if `shouldCompact` is provided. */ + threshold?: number; + /** Custom compaction trigger. When provided, used instead of `threshold`. */ + shouldCompact?: (event: ShouldCompactEvent) => boolean | Promise; +}; + +/** A step object as received in prepareStep's `steps` array. */ +export type CompactionStep = { + usage: LanguageModelUsage; + finishReason: string; + content: Array<{ type: string; toolCallId?: string }>; + response: { messages: Array }; +}; + +/** + * Result of `chat.compact()`. Discriminated union so you can inspect + * what happened, but also directly compatible with prepareStep's return type. + * + * - `"skipped"` — no compaction needed (first step, boundary unsafe, or under threshold). Return `undefined` to prepareStep. + * - `"rebuilt"` — previous compaction exists, messages rebuilt from summary + new response messages. + * - `"compacted"` — compaction just happened, includes the generated summary. + */ +export type CompactResult = + | { type: "skipped" } + | { type: "rebuilt"; messages: ModelMessage[] } + | { type: "compacted"; messages: ModelMessage[]; summary: string }; + +/** + * Options for `chat.compact()` — the low-level compaction function. + */ +export type CompactOptions = { + /** Generate a summary from the current messages. Return the summary text. */ + summarize: (messages: ModelMessage[]) => Promise; + /** Token threshold — compact when totalTokens exceeds this. Ignored if `shouldCompact` is provided. */ + threshold?: number; + /** Custom compaction trigger. When provided, used instead of `threshold`. */ + shouldCompact?: (event: ShouldCompactEvent) => boolean | Promise; +}; + +/** + * Check that no tool calls are in-flight in a step's content. + * Used before compaction to avoid losing tool state mid-execution. + * @internal + */ +function isStepBoundarySafe(step: { + finishReason: string; + content: Array<{ type: string; toolCallId?: string }>; +}): boolean { + if (step.finishReason === "error") return false; + const callIds = new Set( + step.content.filter((p) => p.type === "tool-call").map((p) => p.toolCallId) + ); + const settledIds = new Set( + step.content + .filter((p) => p.type === "tool-result" || p.type === "tool-error") + .map((p) => p.toolCallId) + ); + return ![...callIds].some((id) => !settledIds.has(id)); +} + +/** + * Apply the prepareMessages hook if one is set in locals. + * @internal + */ +async function applyPrepareMessages( + messages: ModelMessage[], + reason: PrepareMessagesEvent["reason"] +): Promise { + const hook = locals.get(chatPrepareMessagesKey); + if (!hook) return messages; + + const turnCtx = locals.get(chatTurnContextKey); + + return tracer.startActiveSpan( + "prepareMessages()", + async () => { + return hook({ + messages, + reason, + chatId: turnCtx?.chatId ?? "", + turn: turnCtx?.turn ?? 0, + clientData: turnCtx?.clientData, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.prepareMessages.reason": reason, + "chat.prepareMessages.messageCount": messages.length, + }, + } + ); +} + +/** + * Read the current compaction state. Returns the summary and base message count + * if compaction has occurred in this turn, or `undefined` if not. + * + * Use in a custom `prepareStep` to rebuild from a previous compaction: + * ```ts + * const state = chat.getCompactionState(); + * if (state) { + * return { messages: [{ role: "user", content: state.summary }, ...newMsgs] }; + * } + * ``` + */ +function getCompactionState(): CompactionState | undefined { + return locals.get(chatCompactionStateKey); +} + +/** + * Low-level compaction for use inside a custom `prepareStep`. + * + * Handles the full decision tree: first step, already-compacted rebuild, + * boundary safety, threshold check, summarization, stream chunks, state + * storage, and accumulator update. + * + * Returns a `CompactResult` — inspect `result.type` to see what happened, + * or convert to a prepareStep return with `result.type === "skipped" ? undefined : result`. + * + * @example + * ```ts + * prepareStep: async ({ messages, steps }) => { + * // your custom logic here... + * const result = await chat.compact(messages, steps, { + * threshold: 80_000, + * summarize: async (msgs) => generateText({ model, messages: msgs }).then(r => r.text), + * }); + * if (result.type === "compacted") { + * logger.info("Compacted!", { summary: result.summary }); + * } + * return result.type === "skipped" ? undefined : result; + * }, + * ``` + */ +async function chatCompact( + messages: ModelMessage[], + steps: CompactionStep[], + options: CompactOptions +): Promise { + const currentStep = steps.at(-1); + + // First step — nothing to check + if (!currentStep) { + return { type: "skipped" }; + } + + // Already compacted — rebuild from summary + new response messages + const state = locals.get(chatCompactionStateKey); + if (state && isStepBoundarySafe(currentStep)) { + return { + type: "rebuilt", + messages: await applyPrepareMessages( + [ + { role: "user" as const, content: state.summary }, + ...currentStep.response.messages.slice(state.baseResponseMessageCount), + ], + "compaction-rebuild" + ), + }; + } + + // Boundary unsafe — skip + if (!isStepBoundarySafe(currentStep)) { + return { type: "skipped" }; + } + + const totalTokens = currentStep.usage.totalTokens; + const inputTokens = currentStep.usage.inputTokens; + const outputTokens = currentStep.usage.outputTokens; + + const turnCtx = locals.get(chatTurnContextKey); + const stepNumber = steps.length - 1; + + const shouldTrigger = options.shouldCompact + ? await options.shouldCompact({ + messages, + totalTokens, + inputTokens, + outputTokens, + usage: currentStep.usage, + source: "inner", + stepNumber, + steps, + chatId: turnCtx?.chatId, + turn: turnCtx?.turn, + clientData: turnCtx?.clientData, + }) + : totalTokens != null && options.threshold != null && totalTokens > options.threshold; + + if (!shouldTrigger) { + return { type: "skipped" }; + } + + const result = await tracer.startActiveSpan( + "context compaction", + async (span) => { + const compactionId = generateMessageId(); + let summary!: string; + + const { waitUntilComplete } = chatStream.writer({ + spanName: "stream compaction chunks", + collapsed: true, + execute: async ({ write, merge }) => { + // Control chunks aren't part of UIMessageChunk's discriminated + // union but flow on the same session.out so subscribers can + // intercept them — cast on the way out. + write({ type: "step-start" } as unknown as UIMessageChunk); + write({ + type: "data-compaction", + id: compactionId, + data: { status: "compacting", totalTokens }, + transient: true, + }); + + // Generate summary + summary = await options.summarize(messages); + + // Store state in locals for subsequent steps + locals.set(chatCompactionStateKey, { + summary, + baseResponseMessageCount: currentStep.response.messages.length, + }); + + // Set model-only override — UI messages stay intact for persistence. + // The summary becomes the model message history for the next turn, + // while accumulatedUIMessages keeps the full conversation for display. + locals.set(chatOverrideModelMessagesKey, [ + { + role: "assistant" as const, + content: [{ type: "text" as const, text: `[Conversation summary]\n\n${summary}` }], + }, + ]); + + // Fire onCompacted hook — pass the existing writer so the callback + // can write custom chunks without creating a separate stream. + const onCompactedHook = locals.get(chatOnCompactedKey); + if (onCompactedHook) { + await onCompactedHook({ + ctx: locals.get(chatAgentRunContextKey)!, + summary, + messages, + messageCount: messages.length, + usage: currentStep.usage, + totalTokens, + inputTokens, + outputTokens, + stepNumber, + chatId: turnCtx?.chatId, + turn: turnCtx?.turn, + writer: { write, merge }, + }); + } + + write({ + type: "data-compaction", + id: compactionId, + data: { status: "complete", totalTokens }, + transient: true, + }); + write({ type: "finish-step" }); + }, + }); + await waitUntilComplete(); + + // Set attributes after we have the summary + span.setAttribute("compaction.summary_length", summary.length); + + return { + type: "compacted" as const, + messages: await applyPrepareMessages( + [{ role: "user" as const, content: summary }], + "compaction-result" + ), + summary, + }; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "tabler-scissors", + "compaction.threshold": options.threshold, + "compaction.total_tokens": totalTokens ?? 0, + "compaction.input_tokens": inputTokens ?? 0, + "compaction.message_count": messages.length, + "compaction.step_number": stepNumber, + ...(turnCtx?.chatId ? { "compaction.chat_id": turnCtx.chatId } : {}), + ...(turnCtx?.turn != null ? { "compaction.turn": turnCtx.turn } : {}), + ...accessoryAttributes({ + items: [ + { text: `${totalTokens ?? 0} tokens`, variant: "normal" }, + { text: `${messages.length} msgs`, variant: "normal" }, + ], + style: "codepath", + }), + }, + } + ); + + return result; +} + +/** + * Returns a `prepareStep` function that handles context compaction automatically. + * + * Monitors token usage between tool-call steps. When `totalTokens` exceeds + * the threshold, generates a summary via `summarize()`, replaces the message + * history, and emits `data-compaction` stream chunks for the frontend. + * + * @example + * ```ts + * return streamText({ + * ...chat.toStreamTextOptions({ registry }), + * messages: chat.addCacheBreaks(messages), + * prepareStep: chat.compactionStep({ + * threshold: 80_000, + * summarize: async (messages) => { + * return generateText({ model, messages: [...messages, { role: "user", content: "Summarize." }] }) + * .then((r) => r.text); + * }, + * }), + * tools: { ... }, + * }); + * ``` + */ +function chatCompactionStep( + options: CompactionOptions +): (args: { + messages: ModelMessage[]; + steps: CompactionStep[]; +}) => Promise<{ messages: ModelMessage[] } | undefined> { + return async ({ messages, steps }) => { + const result = await chatCompact(messages, steps, options); + return result.type === "skipped" ? undefined : result; + }; +} + +// --------------------------------------------------------------------------- +// Steering queue drain — shared by toStreamTextOptions, session, accumulator +// --------------------------------------------------------------------------- + +/** + * Drain the steering queue as a batch. Calls `shouldInject` once with all + * pending messages. If it returns true, calls `prepareMessages` once to + * transform the batch, then clears the queue. + * Returns the model messages to inject (empty if none). + * @internal + */ +async function drainSteeringQueue( + config: PendingMessagesOptions, + messages: ModelMessage[], + steps: CompactionStep[], + queueOverride?: SteeringQueueEntry[] +): Promise { + const queue = queueOverride ?? locals.get(chatSteeringQueueKey); + if (!queue || queue.length === 0) return []; + + const ctx = locals.get(chatTurnContextKey); + const stepNumber = steps.length - 1; + const uiMessages = queue.map((e) => e.uiMessage); + + const batchEvent: PendingMessagesBatchEvent = { + messages: uiMessages, + modelMessages: messages, + steps, + stepNumber, + chatId: ctx?.chatId ?? "", + turn: ctx?.turn ?? 0, + clientData: ctx?.clientData, + }; + + // Call shouldInject once for the whole batch + const shouldInject = config.shouldInject ? await config.shouldInject(batchEvent) : false; + + if (!shouldInject) return []; + + // Extract message texts for span attributes + const messageTexts = uiMessages.map( + (m) => + (m.parts ?? []) + .filter((p: any) => p.type === "text") + .map((p: any) => p.text) + .join("") || "" + ); + const previewText = + messageTexts.length === 1 ? messageTexts[0]!.slice(0, 80) : `${queue.length} messages`; + + return tracer.startActiveSpan( + "pending message injected", + async () => { + // Transform the batch — default: concatenate all pre-converted model messages + const injected = config.prepare + ? await config.prepare(batchEvent) + : queue.flatMap((e) => e.modelMessages); + + // Clear the queue and record injected IDs + queue.length = 0; + const injectedIds = locals.get(chatInjectedMessageIdsKey); + if (injectedIds) { + for (const m of uiMessages) injectedIds.add(m.id); + } + + // Write injection confirmation chunk to the stream so the frontend + // knows which messages were injected and where in the response. + if (injected.length > 0) { + try { + const { waitUntilComplete } = chatStream.writer({ + collapsed: true, + execute: ({ write }) => { + write({ + type: PENDING_MESSAGE_INJECTED_TYPE, + id: generateMessageId(), + data: { + messageIds: uiMessages.map((m) => m.id), + messages: uiMessages.map((m, idx) => ({ + id: m.id, + text: messageTexts[idx] ?? "", + })), + }, + }); + }, + }); + await waitUntilComplete(); + } catch { + /* non-fatal — stream write failed */ + } + } + + // Fire onInjected callback + if (config.onInjected && injected.length > 0) { + try { + await config.onInjected({ + messages: uiMessages, + injectedModelMessages: injected, + chatId: ctx?.chatId ?? "", + turn: ctx?.turn ?? 0, + stepNumber, + }); + } catch { + /* non-fatal */ + } + } + + return injected; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "tabler-message-forward", + "pending.message_count": uiMessages.length, + "pending.step_number": stepNumber, + "pending.messages": messageTexts, + ...(ctx?.chatId ? { "pending.chat_id": ctx.chatId } : {}), + ...(ctx?.turn != null ? { "pending.turn": ctx.turn } : {}), + ...accessoryAttributes({ + items: [ + { + text: `${uiMessages.length} message${uiMessages.length === 1 ? "" : "s"}`, + variant: "normal", + }, + { text: `between steps ${stepNumber} and ${stepNumber + 1}`, variant: "normal" }, + ], + style: "codepath", + }), + }, + } + ); +} + +// --------------------------------------------------------------------------- +// chat.isCompactionSafe — check if it's safe to compact messages +// --------------------------------------------------------------------------- + +/** + * Checks whether it's safe to compact the message history. Returns `false` + * if any tool calls are in-flight (incomplete tool invocations without results). + * + * Call before `chat.setMessages()` to avoid corrupting tool-call state. + */ +function isCompactionSafe(messages: UIMessage[]): boolean { + for (const msg of messages) { + if (msg.role !== "assistant") continue; + for (const part of msg.parts as any[]) { + if (part.type === "tool-invocation") { + const state = part.toolInvocation?.state ?? part.state; + if (state !== "result" && state !== "error") { + return false; + } + } + } + } + return true; +} + +// --------------------------------------------------------------------------- +// chat.prompt — store and retrieve a resolved prompt for the current run +// --------------------------------------------------------------------------- + +/** + * A resolved prompt stored via `chat.prompt.set()`. Either a full `ResolvedPrompt` + * from `prompts.define().resolve()`, or a lightweight wrapper around a plain string. + */ +export type ChatPromptValue = + | ResolvedPrompt + | { + text: string; + model: undefined; + config: undefined; + promptId: string; + version: number; + labels: string[]; + toAISDKTelemetry: (additionalMetadata?: Record) => { + experimental_telemetry: { isEnabled: true; metadata: Record }; + }; + }; + +/** @internal */ +const chatPromptKey = locals.create("chat.prompt"); + +/** + * Store a resolved prompt (or plain string) for the current run. + * Call from any hook (`onPreload`, `onChatStart`, `onTurnStart`) or `run()`. + */ +function setChatPrompt(resolved: ResolvedPrompt | string): void { + if (typeof resolved === "string") { + locals.set(chatPromptKey, { + text: resolved, + model: undefined, + config: undefined, + promptId: "", + version: 0, + labels: [], + toAISDKTelemetry: () => ({ + experimental_telemetry: { isEnabled: true, metadata: {} }, + }), + }); + } else { + locals.set(chatPromptKey, resolved); + } +} + +/** + * Read the stored prompt. Throws if `chat.prompt.set()` has not been called. + */ +function getChatPrompt(): ChatPromptValue { + const prompt = locals.get(chatPromptKey); + if (!prompt) { + throw new Error( + "chat.prompt() called before chat.prompt.set(). Set a prompt in onPreload, onChatStart, onTurnStart, or run() first." + ); + } + return prompt; +} + +// --------------------------------------------------------------------------- +// chat.skills — store resolved agent skills and inject them into streamText +// --------------------------------------------------------------------------- + +/** @internal */ +const chatSkillsKey = locals.create("chat.skills"); + +/** + * Store resolved skills for the current run. Call from any hook + * (`onPreload`, `onChatStart`, `onTurnStart`) or `run()`. + */ +function setChatSkills(skills: ResolvedSkill[]): void { + locals.set(chatSkillsKey, skills); +} + +/** Read the stored skills. Returns `undefined` if none set. */ +function getChatSkills(): ResolvedSkill[] | undefined { + return locals.get(chatSkillsKey); +} + +/** + * Build the system-prompt preamble advertising available skills. Only the + * frontmatter description surfaces here — full SKILL.md body is loaded + * on-demand via the `loadSkill` tool. + */ +function buildSkillsSystemPrompt(skills: ResolvedSkill[]): string { + if (skills.length === 0) return ""; + const lines = skills.map( + (s) => `- ${s.frontmatter.name}: ${s.frontmatter.description}` + ); + return [ + "Available skills (call `loadSkill` to read the full instructions before using one):", + ...lines, + ].join("\n"); +} + +/** Resolve a skill by its frontmatter `name`. */ +function findSkillByName(skills: ResolvedSkill[], name: string): ResolvedSkill | undefined { + return skills.find((s) => s.frontmatter.name === name); +} + +/** + * Build the three tools we auto-inject into `streamText` when skills are + * set: `loadSkill`, `readFile`, `bash`. Scoped per-skill by name. + * + * Exported so callers can use the same tools outside the auto-wired path + * (e.g. in a `chat.createSession` loop with custom streamText). + */ +export function buildSkillTools(skills: ResolvedSkill[]): Record { + const loadSkill = aiTool({ + description: + "Load the full instructions for a skill by its name. Call this first before using a skill.", + inputSchema: jsonSchema<{ name: string }>({ + type: "object", + properties: { + name: { + type: "string", + description: "The `name` field from the skill's frontmatter.", + }, + }, + required: ["name"], + additionalProperties: false, + } as JSONSchema7), + execute: async ({ name }: { name: string }) => { + const skill = findSkillByName(skills, name); + if (!skill) { + return { + error: `Skill "${name}" not found. Available: ${skills + .map((s) => s.frontmatter.name) + .join(", ")}`, + }; + } + return { + name: skill.frontmatter.name, + description: skill.frontmatter.description, + body: skill.body, + path: skill.path, + }; + }, + }); + + const readFile = aiTool({ + description: + "Read a file from a skill's bundled folder. Paths must be relative to the skill's root.", + inputSchema: jsonSchema<{ skill: string; path: string }>({ + type: "object", + properties: { + skill: { type: "string", description: "The skill's name (from frontmatter)." }, + path: { + type: "string", + description: "Relative path inside the skill folder (e.g. `references/citation-style.md`).", + }, + }, + required: ["skill", "path"], + additionalProperties: false, + } as JSONSchema7), + execute: async ({ skill: skillName, path: relPath }: { skill: string; path: string }) => { + const skill = findSkillByName(skills, skillName); + if (!skill) { + return { error: `Skill "${skillName}" not found.` }; + } + try { + return await readFileInSkill({ + skillPath: skill.path, + relativePath: relPath, + }); + } catch (err) { + return { error: (err as Error).message }; + } + }, + }); + + const bash = aiTool({ + description: + "Run a bash command inside a skill's bundled folder. Use this to invoke the skill's scripts. The working directory is the skill's root.", + inputSchema: jsonSchema<{ skill: string; command: string }>({ + type: "object", + properties: { + skill: { type: "string", description: "The skill's name (from frontmatter)." }, + command: { + type: "string", + description: "Bash command to run. Relative script paths resolve against the skill's root.", + }, + }, + required: ["skill", "command"], + additionalProperties: false, + } as JSONSchema7), + execute: async ( + { skill: skillName, command }: { skill: string; command: string }, + { abortSignal }: { abortSignal?: AbortSignal } = {} + ) => { + const skill = findSkillByName(skills, skillName); + if (!skill) { + return { error: `Skill "${skillName}" not found.` }; + } + try { + return await runBashInSkill({ + skillPath: skill.path, + command, + abortSignal, + }); + } catch (err) { + return { error: (err as Error).message }; + } + }, + }); + + return { loadSkill, readFile, bash }; +} + +/** + * Options for {@link toStreamTextOptions}. + */ +export type ToStreamTextOptionsOptions = { + /** Additional telemetry metadata merged into `experimental_telemetry.metadata`. */ + telemetry?: Record; + /** + * An AI SDK provider registry (from `createProviderRegistry`) or any object + * with a `languageModel(id)` method. When provided and the stored prompt has + * a `model` string, the resolved `LanguageModel` is included in the returned + * options so `streamText` uses it directly. + * + * The model string should use the `"provider:model-id"` format + * (e.g. `"openai:gpt-4o"`, `"anthropic:claude-sonnet-4-6"`). + */ + registry?: { languageModel(modelId: string): unknown }; + /** + * User-defined tools to merge alongside the auto-injected skill tools + * (`loadSkill`, `readFile`, `bash`). User tools win on name conflicts. + * + * If you don't pass `tools` here and skills are set, the returned options + * will include just the skill tools — spread after any `tools` you pass + * directly to `streamText` and they'll be replaced. Easiest: pass all + * your tools here. + */ + tools?: Record; +}; + +/** + * Returns an options object ready to spread into `streamText()`. + * + * Includes `system`, `experimental_telemetry`, and any config fields + * (temperature, maxTokens, etc.) from the stored prompt. + * + * When a `registry` is provided and the prompt has a `model` string, + * the resolved `LanguageModel` is included as `model`. + * + * If no prompt has been set, returns `{}` (no-op spread). + */ +function toStreamTextOptions(options?: ToStreamTextOptionsOptions): Record { + const prompt = locals.get(chatPromptKey); + const skills = locals.get(chatSkillsKey); + const result: Record = {}; + + // Build the combined system prompt: stored prompt + skills preamble. + const promptText = prompt?.text ?? ""; + const skillsText = skills && skills.length > 0 ? buildSkillsSystemPrompt(skills) : ""; + if (promptText || skillsText) { + result.system = [promptText, skillsText].filter(Boolean).join("\n\n"); + } + + // Prompt-related options (only if chat.prompt.set() was called) + if (prompt) { + // Resolve model via registry if both are present + if (options?.registry && prompt.model) { + result.model = options.registry.languageModel(prompt.model); + } + + // Spread config (temperature, maxTokens, etc.) + if (prompt.config) { + Object.assign(result, prompt.config); + } + + // Add telemetry (forward additional metadata from caller) + const telemetry = prompt.toAISDKTelemetry(options?.telemetry); + Object.assign(result, telemetry); + } + + // Skills: merge auto-injected tools with any user-provided tools. + // User tools override on name conflict (though we namespace ours). + if (skills && skills.length > 0) { + const skillTools = buildSkillTools(skills); + result.tools = { ...skillTools, ...(options?.tools ?? {}) }; + } else if (options?.tools) { + result.tools = options.tools; + } + + // Auto-inject prepareStep for compaction, pending messages, and background context injection. + // This runs regardless of whether a prompt is set — these features are independent. + const taskCompaction = locals.get(chatAgentCompactionKey); + const taskPendingMessages = locals.get(chatPendingMessagesKey); + + { + result.prepareStep = async ({ + messages, + steps, + }: { + messages: ModelMessage[]; + steps: CompactionStep[]; + }) => { + let resultMessages: ModelMessage[] | undefined; + + // 1. Compaction + if (taskCompaction) { + const compactResult = await chatCompact(messages, steps, { + shouldCompact: taskCompaction.shouldCompact, + summarize: (msgs) => { + const ctx = locals.get(chatTurnContextKey); + const lastStep = steps.at(-1); + return taskCompaction.summarize({ + messages: msgs, + usage: lastStep?.usage, + source: "inner", + stepNumber: steps.length - 1, + chatId: ctx?.chatId, + turn: ctx?.turn, + clientData: ctx?.clientData, + }); + }, + }); + if (compactResult.type !== "skipped") { + resultMessages = compactResult.messages; + } + } + + // 2. Pending message injection (steering) + if (taskPendingMessages) { + const injected = await drainSteeringQueue( + taskPendingMessages, + resultMessages ?? messages, + steps + ); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? messages), ...injected]; + } + } + + // 3. Background context injection + const bgQueue = locals.get(chatBackgroundQueueKey); + if (bgQueue && bgQueue.length > 0) { + const injected = bgQueue.splice(0); // drain + resultMessages = [...(resultMessages ?? messages), ...injected]; + } + + return resultMessages ? { messages: resultMessages } : undefined; + }; + } + + return result; +} + +/** + * Options for `pipeChat`. + */ +export type PipeChatOptions = { + /** + * Override the stream key. Must match the `streamKey` on `TriggerChatTransport`. + * @default "chat" + */ + streamKey?: string; + + /** An AbortSignal to cancel the stream. */ + signal?: AbortSignal; + + /** + * The target run ID to pipe to. + * @default "self" (current run) + */ + target?: string; + + /** Override the default span name for this operation. */ + spanName?: string; +}; + +/** + * Options for customizing the `toUIMessageStream()` call used when piping + * `streamText` results to the frontend. + * + * Set static defaults via `uiMessageStreamOptions` on `chat.agent()`, or + * override per-turn via `chat.setUIMessageStreamOptions()`. + * + * `onFinish` is omitted because it is managed internally for response capture. + * Use `streamText`'s `onFinish` for custom finish handling, or drop down to + * raw task mode with `chat.pipe()` for full control. + * + * `originalMessages` is omitted because it is automatically set from the + * accumulated conversation history, ensuring message IDs are reused across + * turns (e.g. for tool approval continuations). + * + * `generateMessageId` can be set to control ID generation for response + * messages (e.g. UUID-v7). If not set, the AI SDK's default `generateId` is used. + */ +export type ChatUIMessageStreamOptions = Omit< + UIMessageStreamOptions, + "onFinish" | "originalMessages" +>; + +/** + * An object with a `toUIMessageStream()` method (e.g. `StreamTextResult` from `streamText()`). + */ +type UIMessageStreamable = { + toUIMessageStream: (...args: any[]) => AsyncIterable | ReadableStream; +}; + +function isUIMessageStreamable(value: unknown): value is UIMessageStreamable { + return ( + typeof value === "object" && + value !== null && + "toUIMessageStream" in value && + typeof (value as any).toUIMessageStream === "function" + ); +} + +let warnedMissingOnAction = false; +function warnMissingOnActionOnce() { + if (warnedMissingOnAction) return; + warnedMissingOnAction = true; + console.warn( + "[chat.agent] Received an action but no `onAction` handler is configured. " + + "The action is being ignored. Define `onAction` (and optionally `actionSchema`) on " + + "your agent to handle it." + ); +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return typeof value === "object" && value !== null && Symbol.asyncIterator in value; +} + +function isReadableStream(value: unknown): value is ReadableStream { + return ( + typeof value === "object" && value !== null && typeof (value as any).getReader === "function" + ); +} + +/** + * Pipes a chat stream to the realtime stream, making it available to the + * `TriggerChatTransport` on the frontend. + * + * Accepts: + * - A `StreamTextResult` from `streamText()` (has `.toUIMessageStream()`) + * - An `AsyncIterable` of `UIMessageChunk`s + * - A `ReadableStream` of `UIMessageChunk`s + * + * Must be called from inside a Trigger.dev task's `run` function. + * + * @example + * ```ts + * import { task } from "@trigger.dev/sdk"; + * import { chat, type ChatTaskPayload } from "@trigger.dev/sdk/ai"; + * import { streamText, convertToModelMessages } from "ai"; + * + * export const myChatTask = task({ + * id: "my-chat-task", + * run: async (payload: ChatTaskPayload) => { + * const result = streamText({ + * model: openai("gpt-4o"), + * messages: payload.messages, + * }); + * + * await chat.pipe(result); + * }, + * }); + * ``` + * + * @example + * ```ts + * // Works from anywhere inside a task — even deep in your agent code + * async function runAgentLoop(messages: CoreMessage[]) { + * const result = streamText({ model, messages }); + * await chat.pipe(result); + * } + * ``` + */ +async function pipeChat( + source: UIMessageStreamable | AsyncIterable | ReadableStream, + options?: PipeChatOptions +): Promise { + locals.set(chatPipeCountKey, (locals.get(chatPipeCountKey) ?? 0) + 1); + + let stream: AsyncIterable | ReadableStream; + + if (isUIMessageStreamable(source)) { + stream = source.toUIMessageStream(); + } else if (isAsyncIterable(source) || isReadableStream(source)) { + stream = source; + } else { + throw new Error( + "pipeChat: source must be a StreamTextResult (with .toUIMessageStream()), " + + "an AsyncIterable, or a ReadableStream" + ); + } + + const pipeOptions: SessionPipeStreamOptions = {}; + if (options?.signal) { + pipeOptions.signal = options.signal; + } + if (options?.spanName) { + pipeOptions.spanName = options.spanName; + } + // `options.target` / `options.streamKey` are accepted for API parity + // with the pre-migration run-scoped pipe but no longer have meaning — + // sessions are the address (single stream per session, no sub-run + // targeting). Sub-agents that need to write into a parent's chat now + // open that session explicitly via `sessions.open(parentSessionId).out.pipe`. + + // The generic is typed for `UIMessageChunk`, but `pipeChat` also + // accepts opaque UIMessageStreamable / raw iterables whose element + // type we don't know at compile time. Cast — runtime behaviour is + // identical (bytes go to session.out either way). + const { waitUntilComplete } = chatStream.pipe( + stream as ReadableStream | AsyncIterable, + pipeOptions + ); + await waitUntilComplete(); +} + +/** + * Options for defining a chat task. + * + * Extends the standard `TaskOptions` but pre-types the payload as `ChatTaskPayload` + * and overrides `run` to accept `ChatTaskRunPayload` (with abort signals). + * + * **Auto-piping:** If the `run` function returns a value with `.toUIMessageStream()` + * (like a `StreamTextResult`), the stream is automatically piped to the frontend. + * + * **Single-run mode:** By default, the task uses input streams so that the + * entire conversation lives inside one run. After each AI response, the task + * emits a control chunk and suspends via `messagesInput.wait()`. The frontend + * transport resumes the same run by sending the next message via input streams. + */ +/** + * Event passed to the `onPreload` callback. + */ +export type PreloadEvent = { + /** Task run context — same as `task({ run })` second-argument `ctx`. */ + ctx: TaskRunContext; + /** The unique identifier for the chat session. */ + chatId: string; + /** The Trigger.dev run ID for this conversation. */ + runId: string; + /** A scoped access token for this chat run. */ + chatAccessToken: string; + /** Custom data from the frontend. */ + clientData?: TClientData; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; +}; + +/** + * Event passed to the `onChatStart` callback. + */ +export type ChatStartEvent = { + /** Task run context — same as `task({ run })` second-argument `ctx`. */ + ctx: TaskRunContext; + /** The unique identifier for the chat session. */ + chatId: string; + /** + * The initial model-ready messages for this conversation. + * + * On a fresh chat this is empty (or just the seed-message for head-start). + * On a continuation — including idle-suspend resume and OOM retry — this + * already reflects the FULL prior conversation history loaded from the + * runtime's durable snapshot + `session.out` replay (or whatever + * `hydrateMessages` returned). The wire never re-ships that history; the + * runtime rebuilds it before `onChatStart` fires. + */ + messages: ModelMessage[]; + /** Custom data from the frontend (passed via `metadata` on `sendMessage()` or the transport). */ + clientData: TClientData; + /** The Trigger.dev run ID for this conversation. */ + runId: string; + /** A scoped access token for this chat run. Persist this for frontend reconnection. */ + chatAccessToken: string; + /** Whether this run is continuing an existing chat (previous run timed out or was cancelled). False for brand new chats. */ + continuation: boolean; + /** The run ID of the previous run (only set when `continuation` is true). */ + previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; +}; + +/** + * Event passed to the `hydrateMessages` callback. + */ +export type HydrateMessagesEvent = { + /** The unique identifier for the chat session. */ + chatId: string; + /** The turn number (0-indexed). */ + turn: number; + /** The trigger type for this turn. */ + trigger: "submit-message" | "regenerate-message" | "action"; + /** Validated incoming UI messages from the wire payload (what the frontend sent). Empty for actions. */ + incomingMessages: TUIM[]; + /** The accumulated UI messages before this turn (empty on turn 0). */ + previousMessages: TUIM[]; + /** Parsed client data from the transport metadata. */ + clientData?: TClientData; + /** Whether this run is continuing from a previous run. */ + continuation: boolean; + /** The ID of the previous run (if continuation). */ + previousRunId?: string; +}; + +/** + * Event passed to the `hydrateStore` callback. + * + * Called at turn start — before `run()` fires and before any incoming store + * from the wire payload is applied. Return the authoritative store value + * for this turn; it becomes the initial value `chat.store.get()` sees. + */ +export type HydrateStoreEvent = { + /** The unique identifier for the chat session. */ + chatId: string; + /** The turn number (0-indexed). */ + turn: number; + /** The trigger type for this turn. */ + trigger: "submit-message" | "regenerate-message" | "action" | "preload"; + /** + * The in-memory store value from the previous turn of this run + * (`undefined` on turn 0 and after continuations). + */ + previousStore: TStore | undefined; + /** + * The store value the transport sent with this turn, if any. + * Usually set by client-side `setStore` / `applyStorePatch`. + */ + incomingStore: TStore | undefined; + /** Parsed client data from the transport metadata. */ + clientData?: TClientData; + /** Whether this run is continuing from a previous run. */ + continuation: boolean; + /** The ID of the previous run (if continuation). */ + previousRunId?: string; +}; + +/** + * Event passed to the `onValidateMessages` callback. + */ +export type ValidateMessagesEvent = { + /** The incoming UI messages for this turn (after cleanup of aborted tool parts). */ + messages: TUIM[]; + /** The unique identifier for the chat session. */ + chatId: string; + /** The turn number (0-indexed). */ + turn: number; + /** The trigger type for this turn. */ + trigger: "submit-message" | "regenerate-message" | "preload" | "close"; +}; + +/** + * Event passed to the `onAction` callback. + */ +export type ActionEvent< + TAction = unknown, + TClientData = unknown, + TUIM extends UIMessage = UIMessage, +> = { + /** The parsed and validated action payload. */ + action: TAction; + /** The unique identifier for the chat session. */ + chatId: string; + /** The turn number (0-indexed). */ + turn: number; + /** Parsed client data from the transport metadata. */ + clientData?: TClientData; + /** The accumulated UI messages (after hydration, if set). */ + uiMessages: TUIM[]; + /** The accumulated model messages (after hydration, if set). */ + messages: ModelMessage[]; +}; + +/** + * Event passed to the `onTurnStart` callback. + */ +export type TurnStartEvent = { + /** Task run context — same as `task({ run })` second-argument `ctx`. */ + ctx: TaskRunContext; + /** The unique identifier for the chat session. */ + chatId: string; + /** The accumulated model-ready messages (all turns so far, including new user message). */ + messages: ModelMessage[]; + /** The accumulated UI messages (all turns so far, including new user message). */ + uiMessages: TUIM[]; + /** The turn number (0-indexed). */ + turn: number; + /** The Trigger.dev run ID for this conversation. */ + runId: string; + /** A scoped access token for this chat run. */ + chatAccessToken: string; + /** Custom data from the frontend. */ + clientData?: TClientData; + /** Whether this run is continuing an existing chat (previous run timed out or was cancelled). False for brand new chats. */ + continuation: boolean; + /** The run ID of the previous run (only set when `continuation` is true). */ + previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; + /** Token usage from the previous turn. Undefined on turn 0. */ + previousTurnUsage?: LanguageModelUsage; + /** Cumulative token usage across all completed turns so far. */ + totalUsage: LanguageModelUsage; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; +}; + +/** + * Event passed to the `onTurnComplete` callback. + */ +export type TurnCompleteEvent = { + /** Task run context — same as `task({ run })` second-argument `ctx`. */ + ctx: TaskRunContext; + /** The unique identifier for the chat session. */ + chatId: string; + /** The full accumulated conversation in model format (all turns so far). */ + messages: ModelMessage[]; + /** + * The full accumulated conversation in UI format (all turns so far). + * This is the format expected by `useChat` — store this for persistence. + */ + uiMessages: TUIM[]; + /** + * Only the new model messages from this turn (user message(s) + assistant response). + * Useful for appending to an existing conversation record. + */ + newMessages: ModelMessage[]; + /** + * Only the new UI messages from this turn (user message(s) + assistant response). + * Useful for inserting individual message records instead of overwriting the full history. + */ + newUIMessages: TUIM[]; + /** The assistant's response for this turn, with aborted parts cleaned up when `stopped` is true. Undefined if `pipeChat` was used manually. */ + responseMessage: TUIM | undefined; + /** + * The raw assistant response before abort cleanup. Includes incomplete tool parts + * (`input-available`, `partial-call`) and streaming reasoning/text parts. + * Use this if you need custom cleanup logic. Same as `responseMessage` when not stopped. + */ + rawResponseMessage: TUIM | undefined; + /** The turn number (0-indexed). */ + turn: number; + /** The Trigger.dev run ID for this conversation. */ + runId: string; + /** A fresh scoped access token for this chat run (renewed each turn). Persist this for frontend reconnection. */ + chatAccessToken: string; + /** The last event ID from the stream writer. Use this with `resume: true` to avoid replaying events after refresh. */ + lastEventId?: string; + /** Custom data from the frontend. */ + clientData?: TClientData; + /** Whether the user stopped generation during this turn. */ + stopped: boolean; + /** Whether this run is continuing an existing chat (previous run timed out or was cancelled). False for brand new chats. */ + continuation: boolean; + /** The run ID of the previous run (only set when `continuation` is true). */ + previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; + /** Token usage for this turn. Undefined if usage couldn't be captured (e.g. manual pipeChat). */ + usage?: LanguageModelUsage; + /** Cumulative token usage across all turns in this run (including this turn). */ + totalUsage: LanguageModelUsage; + /** + * Why the LLM stopped generating this turn: + * - `"stop"` — model generated a stop sequence (normal completion) + * - `"tool-calls"` — model stopped on one or more tool calls. If any tool + * has no `execute` function (e.g. an `ask_user` HITL tool), the turn is + * paused awaiting user input; inspect `responseMessage.parts` for tool + * parts in `input-available` state to distinguish. + * - `"length"` — max tokens reached + * - `"content-filter"` — content filter stopped the model + * - `"error"` — model errored + * - `"other"` — provider-specific reason + * + * Undefined if the underlying stream didn't provide a finish reason (e.g. + * manual `pipeChat()` or an aborted stream). + */ + finishReason?: FinishReason; +}; + +/** + * Event passed to the `onBeforeTurnComplete` callback. + * Same as `TurnCompleteEvent` but includes a `writer` since the stream is still open. + */ +export type BeforeTurnCompleteEvent< + TClientData = unknown, + TUIM extends UIMessage = UIMessage, +> = TurnCompleteEvent & { + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; +}; + +/** + * Discriminated event passed to the `onChatSuspend` callback. + * Use `phase` to distinguish preload vs turn suspension. + */ +export type ChatSuspendEvent = + | { + /** Suspend is happening after onPreload, before the first message. */ + phase: "preload"; + /** Task run context. */ + ctx: TaskRunContext; + /** The chat session ID. */ + chatId: string; + /** The Trigger.dev run ID. */ + runId: string; + /** Custom data from the frontend. */ + clientData?: TClientData; + } + | { + /** Suspend is happening after a completed turn, waiting for the next message. */ + phase: "turn"; + /** Task run context. */ + ctx: TaskRunContext; + /** The chat session ID. */ + chatId: string; + /** The Trigger.dev run ID. */ + runId: string; + /** The turn number (0-indexed) that just completed. */ + turn: number; + /** The accumulated model messages after the completed turn. */ + messages: ModelMessage[]; + /** The accumulated UI messages after the completed turn. */ + uiMessages: TUIM[]; + /** Custom data from the frontend. */ + clientData?: TClientData; + }; + +/** + * Discriminated event passed to the `onChatResume` callback. + * Use `phase` to distinguish preload vs turn resumption. + */ +export type ChatResumeEvent = + | { + /** First message arrived after preload suspension. */ + phase: "preload"; + /** Task run context. */ + ctx: TaskRunContext; + /** The chat session ID. */ + chatId: string; + /** The Trigger.dev run ID. */ + runId: string; + /** Custom data from the frontend. */ + clientData?: TClientData; + } + | { + /** Next message arrived after turn suspension. */ + phase: "turn"; + /** Task run context. */ + ctx: TaskRunContext; + /** The chat session ID. */ + chatId: string; + /** The Trigger.dev run ID. */ + runId: string; + /** The turn number that was completed before suspension. */ + turn: number; + /** The accumulated model messages (from before suspension). */ + messages: ModelMessage[]; + /** The accumulated UI messages (from before suspension). */ + uiMessages: TUIM[]; + /** Custom data from the frontend. */ + clientData?: TClientData; + }; + +export type ChatAgentOptions< + TIdentifier extends string, + TClientDataSchema extends TaskSchema | undefined = undefined, + TUIMessage extends UIMessage = UIMessage, + TActionSchema extends TaskSchema | undefined = undefined, +> = Omit< + TaskOptions< + TIdentifier, + ChatTaskWirePayload>, + unknown + >, + "run" | "retry" +> & { + /** + * Fallback machine preset to use when an attempt fails with an + * out-of-memory (OOM) error. Setting this enables a single OOM retry: + * the next attempt boots on the larger machine, and the chat picks + * up via the standard continuation path (same `chatId` / Session, + * accumulator rebuilds via `hydrateMessages` or post-`onTurnStart` + * persisted state). + * + * Set `machine` (top-level `TaskOptions`) to control the *default* + * machine the agent runs on. `oomMachine` is the *retry-only* swap. + * + * Note: an OOM retry restarts the entire turn from the top — the + * model call and any in-flight tool executes re-run on the larger + * machine. Make tool executes idempotent or persist results before + * returning if you can't tolerate re-execution. + * + * Generic `retry` options are not exposed on `chat.agent` because + * arbitrary retries against an LLM-driven loop tend to be expensive + * and side-effecting. If you need richer retry semantics, drop down + * to `chat.task` (the raw primitive). + * + * @example + * ```ts + * chat.agent({ + * id: "my-chat", + * machine: "small-1x", + * oomMachine: "medium-2x", + * run: async ({ messages, signal }) => + * streamText({ model, messages, abortSignal: signal }), + * }); + * ``` + */ + oomMachine?: MachinePresetName; + + /** + * Schema for validating `clientData` from the frontend. + * Accepts Zod, ArkType, Valibot, or any supported schema library. + * When provided, `clientData` is parsed and typed in all hooks and `run`. + * + * @example + * ```ts + * import { z } from "zod"; + * + * chat.agent({ + * id: "my-chat", + * clientDataSchema: z.object({ model: z.string().optional(), userId: z.string() }), + * run: async ({ messages, clientData, ctx, signal }) => { + * // clientData is typed as { model?: string; userId: string } + * // ctx is the same TaskRunContext as in task({ run: (payload, { ctx }) => ... }) + * }, + * }); + * ``` + */ + clientDataSchema?: TClientDataSchema; + + /** + * Schema for validating custom actions sent via `transport.sendAction()`. + * + * When the frontend sends `trigger: "action"`, the `action` payload is + * parsed against this schema before reaching `onAction`. Invalid actions + * throw and abort the turn. + * + * @example + * ```ts + * import { z } from "zod"; + * + * chat.agent({ + * id: "my-chat", + * actionSchema: z.discriminatedUnion("type", [ + * z.object({ type: z.literal("undo") }), + * z.object({ type: z.literal("rollback"), targetMessageId: z.string() }), + * ]), + * onAction: async ({ action }) => { + * if (action.type === "undo") chat.history.slice(0, -2); + * if (action.type === "rollback") chat.history.rollbackTo(action.targetMessageId); + * }, + * run: async ({ messages, signal }) => { ... }, + * }); + * ``` + */ + actionSchema?: TActionSchema; + + /** + * Called when the frontend sends a custom action via `transport.sendAction()`. + * + * Actions are not turns. They fire `hydrateMessages` (if configured) and + * `onAction` only — no `onTurnStart` / `prepareMessages` / + * `onBeforeTurnComplete` / `onTurnComplete`, no `run()`. Use + * `chat.history.*` inside `onAction` to mutate state. + * + * To produce a model response from an action, return a + * `StreamTextResult` (auto-piped), `string`, or `UIMessage`. Returning + * `void` or nothing is the side-effect-only default. + */ + onAction?: ( + event: ActionEvent< + [TActionSchema] extends [TaskSchema] ? inferSchemaOut : unknown, + inferSchemaOut, + TUIMessage + > + ) => Promise | unknown; + + /** + * The run function for the chat task. + * + * Receives a `ChatTaskRunPayload` with the conversation messages, chat session ID, + * trigger type, task `ctx` (same as `task({ run })`’s second argument), and abort signals + * (`signal`, `cancelSignal`, `stopSignal`). + * + * **Auto-piping:** If this function returns a value with `.toUIMessageStream()`, + * the stream is automatically piped to the frontend. + */ + run: (payload: ChatTaskRunPayload>) => Promise; + + /** + * Called when a preloaded run starts, before the first message arrives. + * + * Use this to initialize state, create DB records, and load context early — + * so everything is ready when the user's first message comes through. + * + * @example + * ```ts + * onPreload: async ({ ctx, chatId, clientData }) => { + * await db.chat.create({ data: { id: chatId } }); + * userContext.init(await loadUser(clientData.userId)); + * } + * ``` + */ + onPreload?: (event: PreloadEvent>) => Promise | void; + + /** + * Called on the first turn (turn 0) of a new run, before the `run` function executes. + * + * Use this to create the chat record in your database when a new conversation starts. + * + * @example + * ```ts + * onChatStart: async ({ ctx, chatId, messages, clientData }) => { + * await db.chat.create({ data: { id: chatId, userId: clientData.userId } }); + * } + * ``` + */ + onChatStart?: (event: ChatStartEvent>) => Promise | void; + + /** + * Validate or transform incoming UI messages before they are converted to model + * messages and accumulated. Fires once per turn with the raw `UIMessage[]` from + * the wire payload (after cleanup of aborted tool parts). + * + * Return the validated messages array. Throw to abort the turn with an error. + * + * This is the right place to call the AI SDK's `validateUIMessages` to catch + * malformed messages from storage or untrusted input before they reach the model. + * + * @example + * ```ts + * import { validateUIMessages } from "ai"; + * + * chat.agent({ + * id: "my-chat", + * onValidateMessages: async ({ messages }) => { + * return validateUIMessages({ messages, tools: chatTools }); + * }, + * run: async ({ messages }) => { + * return streamText({ model, messages, tools: chatTools }); + * }, + * }); + * ``` + */ + onValidateMessages?: ( + event: ValidateMessagesEvent + ) => TUIMessage[] | Promise; + + /** + * Load the full message history from your backend on every turn, + * replacing the built-in linear accumulator. + * + * When set, the returned messages become the accumulated state for this turn. + * The normal accumulation logic (append for submit, replace for regenerate) + * is skipped entirely — the hook is the source of truth. + * + * After the hook returns, any incoming wire messages with matching IDs are + * auto-merged (handles tool approval responses transparently). + * + * Use cases: + * - Backend trust: prevent clients from injecting fabricated history + * - Branching conversations (DAGs): load only the active branch + * - Rollback/undo: exclude undone messages from history + * + * @example + * ```ts + * chat.agent({ + * id: "my-chat", + * hydrateMessages: async ({ chatId, trigger, incomingMessages }) => { + * // Persist the new message + * const newMsg = incomingMessages[incomingMessages.length - 1]; + * if (newMsg && trigger === "submit-message") { + * await db.chatMessages.create({ chatId, message: newMsg }); + * } + * // Return the full authoritative history + * return db.chatMessages.findMany({ where: { chatId } }); + * }, + * run: async ({ messages, signal }) => { + * return streamText({ model: openai("gpt-4o"), messages, abortSignal: signal }); + * }, + * }); + * ``` + */ + hydrateMessages?: ( + event: HydrateMessagesEvent, TUIMessage> + ) => TUIMessage[] | Promise; + + /** + * Load the `chat.store` value for this turn from your backend. + * + * The store lives in memory on the agent instance for the lifetime of + * the run. After a continuation (idle timeout, `chat.requestUpgrade`, + * max turns), a new run starts with an empty store — this hook lets + * you restore it from your own persistence layer. + * + * Runs at turn start, before `run()` fires. The returned value replaces + * the in-memory store and is emitted as a `store-snapshot` chunk so the + * frontend sees the initial value. + * + * If both `hydrateStore` and `incomingStore` (from the wire payload) are + * present, `incomingStore` wins — it represents the client's latest + * local state and follows AG-UI's last-write-wins policy. + * + * @example + * ```ts + * chat.agent({ + * id: "my-chat", + * hydrateStore: async ({ chatId, previousRunId }) => { + * return db.chatStore.findUnique({ where: { chatId } })?.value; + * }, + * onTurnComplete: async ({ chatId }) => { + * await db.chatStore.upsert({ + * where: { chatId }, + * update: { value: chat.store.get() }, + * create: { chatId, value: chat.store.get() }, + * }); + * }, + * run: async ({ messages, signal }) => { + * return streamText({ model: openai("gpt-4o"), messages, abortSignal: signal }); + * }, + * }); + * ``` + */ + hydrateStore?: ( + event: HydrateStoreEvent> + ) => unknown | Promise; + + /** + * Called at the start of every turn, after message accumulation and `onChatStart` (turn 0), + * but before the `run` function executes. + * + * Use this to persist messages before streaming begins, so a mid-stream page refresh + * still shows the user's message. + * + * @example + * ```ts + * onTurnStart: async ({ ctx, chatId, uiMessages }) => { + * await db.chat.update({ where: { id: chatId }, data: { messages: uiMessages } }); + * } + * ``` + */ + onTurnStart?: ( + event: TurnStartEvent, TUIMessage> + ) => Promise | void; + + /** + * Called after the response is captured but before the stream closes. + * The stream is still open, so you can write custom chunks to the frontend + * (e.g. compaction progress). Use this for compaction, post-processing, + * or any work where the user should see real-time status updates. + * + * @example + * ```ts + * onBeforeTurnComplete: async ({ ctx, writer, usage }) => { + * if (usage?.inputTokens && usage.inputTokens > 5000) { + * writer.write({ type: "data-compaction", id: generateId(), data: { status: "compacting" } }); + * // ... compact messages ... + * chat.setMessages(compactedMessages); + * writer.write({ type: "data-compaction", id: generateId(), data: { status: "complete" } }); + * } + * } + * ``` + */ + onBeforeTurnComplete?: ( + event: BeforeTurnCompleteEvent, TUIMessage> + ) => Promise | void; + + /** + * Called when conversation compaction occurs (via `chat.compact()` or + * `chat.compactionStep()`). Use for logging, billing, or persisting the summary. + * + * @example + * ```ts + * onCompacted: async ({ ctx, summary, totalTokens, chatId }) => { + * logger.info("Compacted", { totalTokens, chatId }); + * await db.compactionLog.create({ data: { chatId, summary } }); + * } + * ``` + */ + onCompacted?: (event: CompactedEvent) => Promise | void; + + /** + * Automatic context compaction. When provided, compaction runs automatically + * in both the inner loop (prepareStep, between tool-call steps) and the + * outer loop (between turns, for single-step responses where prepareStep + * never fires). + * + * The `shouldCompact` callback decides when to compact, and `summarize` + * generates the summary. The prepareStep is auto-injected into + * `chat.toStreamTextOptions()` — if you provide your own `prepareStep` + * after spreading, it overrides the auto-injected one. + * + * @example + * ```ts + * chat.agent({ + * id: "my-chat", + * compaction: { + * shouldCompact: ({ totalTokens }) => (totalTokens ?? 0) > 80_000, + * summarize: async (messages) => + * generateText({ model, messages: [...messages, { role: "user", content: "Summarize." }] }) + * .then((r) => r.text), + * }, + * run: async ({ messages, signal }) => { + * return streamText({ ...chat.toStreamTextOptions({ registry }), messages }); + * }, + * }); + * ``` + */ + compaction?: ChatAgentCompactionOptions; + + /** + * Configure how messages that arrive during streaming are handled. + * + * By default, messages queue for the next turn. When `shouldInject` is provided + * and returns `true`, messages are injected between tool-call steps via + * `prepareStep` — allowing users to steer the agent mid-execution. + * + * @example + * ```ts + * pendingMessages: { + * shouldInject: ({ steps }) => steps.length > 0, + * onReceived: ({ message }) => logger.info("Steering message received"), + * }, + * ``` + */ + pendingMessages?: PendingMessagesOptions; + + /** + * Called after each assistant response completes. Use to persist the + * conversation to your database after each assistant response. + * + * @example + * ```ts + * onTurnComplete: async ({ ctx, chatId, messages }) => { + * await db.chat.update({ where: { id: chatId }, data: { messages } }); + * } + * ``` + */ + onTurnComplete?: ( + event: TurnCompleteEvent, TUIMessage> + ) => Promise | void; + + /** + * Maximum number of conversational turns (message round-trips) a single run + * will handle before ending. After this many turns the run completes + * normally and the next message will start a fresh run. + * + * @default 100 + */ + maxTurns?: number; + + /** + * How long to wait for the next message before timing out and ending the run. + * Accepts any duration string (e.g. `"1h"`, `"30m"`). + * + * @default "1h" + */ + turnTimeout?: string; + + /** + * How long (in seconds) the run stays idle (active, using compute) after each + * turn, waiting for the next message. During this window responses are instant. + * After this timeout the run suspends (frees compute) and waits via + * `inputStream.wait()`. + * + * Set to `0` to suspend immediately after each turn. + * + * @default 30 + */ + idleTimeoutInSeconds?: number; + + /** + * How long the `chatAccessToken` (scoped to this run) remains valid. + * A fresh token is minted after each turn, so this only needs to cover + * the gap between turns. + * + * Accepts a duration string (e.g. `"1h"`, `"30m"`, `"2h"`). + * + * @default "1h" + */ + chatAccessTokenTTL?: string; + + /** + * How long (in seconds) the run stays idle after `onPreload` fires, + * waiting for the first message before suspending. + * + * Only applies to preloaded runs (triggered via `transport.preload()`). + * Takes precedence over `transport.preload(..., { idleTimeoutInSeconds })` + * and over {@link ChatAgentOptions.idleTimeoutInSeconds}. + * + * @default Same as `idleTimeoutInSeconds` + */ + preloadIdleTimeoutInSeconds?: number; + + /** + * How long to wait (suspended) for the first message after a preloaded run starts. + * If no message arrives within this time, the run ends. + * + * Only applies to preloaded runs. + * + * @default Same as `turnTimeout` + */ + preloadTimeout?: string; + + /** + * Transform model messages before they're used anywhere — in `run()`, + * in compaction rebuilds, and in compaction results. + * + * Define once, applied everywhere. Use for Anthropic cache breaks, + * injecting system context, stripping PII, etc. + * + * @example + * ```ts + * prepareMessages: async ({ messages, reason }) => { + * // Add Anthropic cache breaks to the last message + * if (messages.length === 0) return messages; + * const last = messages[messages.length - 1]; + * return [...messages.slice(0, -1), { + * ...last, + * providerOptions: { ...last.providerOptions, anthropic: { cacheControl: { type: "ephemeral" } } }, + * }]; + * } + * ``` + */ + prepareMessages?: ( + event: PrepareMessagesEvent> + ) => ModelMessage[] | Promise; + + /** + * Default options for `toUIMessageStream()` when auto-piping or using + * `turn.complete()` / `chat.pipeAndCapture()`. + * + * Controls how the `StreamTextResult` is converted to a `UIMessageChunk` + * stream — error handling, reasoning/source visibility, metadata, etc. + * + * Can be overridden per-turn by calling `chat.setUIMessageStreamOptions()` + * inside `run()` or lifecycle hooks. Per-turn values are merged on top + * of these defaults (per-turn wins on conflicts). + * + * `onFinish` and `originalMessages` are managed internally and cannot be + * overridden here. Use `streamText`'s `onFinish` for custom finish + * handling. `generateMessageId` can be set to control response message + * ID generation (e.g. UUID-v7). + * + * @example + * ```ts + * chat.agent({ + * id: "my-chat", + * uiMessageStreamOptions: { + * sendReasoning: true, + * onError: (error) => error instanceof Error ? error.message : "An error occurred.", + * }, + * run: async ({ messages, signal }) => { ... }, + * }); + * ``` + */ + uiMessageStreamOptions?: ChatUIMessageStreamOptions; + + /** + * Called right before the run suspends to wait for a message. + * + * The `phase` discriminator tells you when the suspend happened: + * - `"preload"`: after `onPreload`, waiting for the first message + * - `"turn"`: after `onTurnComplete`, waiting for the next message + * + * Use this for cleanup before suspension (e.g. disposing sandboxes, closing connections). + * + * @example + * ```ts + * onChatSuspend: async (event) => { + * await disposeExpensiveResources(event.ctx.run.id); + * if (event.phase === "turn") { + * logger.info("Suspending after turn", { turn: event.turn }); + * } + * } + * ``` + */ + onChatSuspend?: ( + event: ChatSuspendEvent, TUIMessage> + ) => Promise | void; + + /** + * Called right after the run resumes from suspension with a new message. + * + * The `phase` discriminator tells you when the resume happened: + * - `"preload"`: first message arrived after preload suspension + * - `"turn"`: next message arrived after turn suspension + * + * Use this for re-initialization after wake (e.g. warming caches, reconnecting). + * + * @example + * ```ts + * onChatResume: async (event) => { + * warmCache(event.ctx.run.id); + * if (event.phase === "turn") { + * logger.info("Resumed after turn", { turn: event.turn }); + * } + * } + * ``` + */ + onChatResume?: ( + event: ChatResumeEvent, TUIMessage> + ) => Promise | void; + + /** + * When `true`, the run exits successfully after the preload idle timeout + * instead of suspending and waiting. The run completes with no turn executed. + * + * Use this for "fire and forget" preloads where you only want to do eager + * initialization. If the user doesn't send a message during the idle window, + * the run ends cleanly. + * + * Only applies to preloaded runs (triggered via `transport.preload()`). + * + * @default false + */ + exitAfterPreloadIdle?: boolean; +}; + +/** + * Creates a Trigger.dev task pre-configured for AI SDK chat. + * + * - **Pre-types the payload** as `ChatTaskRunPayload` — includes abort signals + * - **Auto-pipes the stream** if `run` returns a `StreamTextResult` + * - **Multi-turn**: keeps the conversation in a single run using input streams + * - **Stop support**: frontend can stop generation mid-stream via the stop input stream + * - For complex flows, use `pipeChat()` from anywhere inside your task code + * + * @example + * ```ts + * import { chat } from "@trigger.dev/sdk/ai"; + * import { streamText, convertToModelMessages } from "ai"; + * import { openai } from "@ai-sdk/openai"; + * + * export const myChat = chat.agent({ + * id: "my-chat", + * run: async ({ messages, signal }) => { + * return streamText({ + * model: openai("gpt-4o"), + * messages, // already converted via convertToModelMessages + * abortSignal: signal, + * }); + * }, + * }); + * ``` + */ +// ─── chat.customAgent ────────────────────────────────────────────── +// A thin wrapper around createTask that marks the task as an agent +// (triggerSource: "agent") so it appears in the playground, but does +// NOT implement the managed lifecycle (no turn loop, no preload, etc.). +// The user's run function receives the raw ChatTaskWirePayload and +// uses composable primitives (chat.messages, chat.MessageAccumulator, etc.). +// ──────────────────────────────────────────────────────────────────── + +type ChatCustomAgentOptions< + TIdentifier extends string, + TClientDataSchema extends TaskSchema | undefined = undefined, + TUIMessage extends UIMessage = UIMessage, +> = Omit< + TaskOptions< + TIdentifier, + ChatTaskWirePayload>, + unknown + >, + "triggerSource" | "agentConfig" +> & { + clientDataSchema?: TClientDataSchema; +}; + +function chatCustomAgent< + TIdentifier extends string, + TClientDataSchema extends TaskSchema | undefined = undefined, + TUIMessage extends UIMessage = UIMessage, +>( + options: ChatCustomAgentOptions +): Task>, unknown> { + const { clientDataSchema, run: userRun, ...restOptions } = options; + + const task = createTask< + TIdentifier, + ChatTaskWirePayload>, + unknown + >({ + ...restOptions, + triggerSource: "agent", + agentConfig: { type: "ai-sdk-chat" }, + run: async ( + payload: ChatTaskWirePayload>, + runOptions + ) => { + // Bind the run to its backing Session so module-level helpers + // (chat.messages, chat.stream, chat.createStopSignal, chat.createSession) + // resolve to this chat's `.in` / `.out` channels. Address + // everywhere by `payload.chatId` (the session externalId) so the + // agent's writes and the transport's reads converge on the same + // S2 stream key + waitpoint key. + // + // The Session row is created server-side by `POST /sessions` (or + // `chat.createStartSessionAction`) before this run is triggered. + // No client-side upsert needed. + locals.set(chatSessionHandleKey, sessions.open(payload.chatId)); + locals.set(chatAgentRunContextKey, runOptions.ctx); + taskContext.setConversationId(payload.chatId); + stampConversationIdOnActiveSpan(payload.chatId); + return userRun(payload, runOptions); + }, + }); + + // Register clientDataSchema so the CLI converts it to JSONSchema + // and stores it as payloadSchema — used by the Playground UI + if (clientDataSchema) { + resourceCatalog.updateTaskMetadata(options.id, { + schema: clientDataSchema as any, + }); + } + + return task; +} + +function chatAgent< + TIdentifier extends string, + TClientDataSchema extends TaskSchema | undefined = undefined, + TUIMessage extends UIMessage = UIMessage, + TActionSchema extends TaskSchema | undefined = undefined, +>( + options: ChatAgentOptions +): Task>, unknown> { + const { + run: userRun, + clientDataSchema, + onPreload, + onChatStart, + onValidateMessages, + hydrateMessages, + hydrateStore, + actionSchema, + onAction, + onTurnStart, + onBeforeTurnComplete, + onCompacted, + compaction, + pendingMessages: pendingMessagesConfig, + prepareMessages, + onTurnComplete, + maxTurns = 100, + turnTimeout = "1h", + idleTimeoutInSeconds = 30, + chatAccessTokenTTL = "1h", + preloadIdleTimeoutInSeconds, + preloadTimeout, + uiMessageStreamOptions, + onChatSuspend, + onChatResume, + exitAfterPreloadIdle = false, + oomMachine, + ...restOptions + } = options; + + const parseClientData = clientDataSchema ? getSchemaParseFn(clientDataSchema) : undefined; + const parseAction = actionSchema ? getSchemaParseFn(actionSchema) : undefined; + + // chat.agent does not expose generic retry options (see docstring on + // `oomMachine`). The only opt-in is an OOM-triggered machine swap. If + // `oomMachine` is set we allow one retry on a larger machine; otherwise + // we keep the historical no-retry default. + const retry = oomMachine + ? { maxAttempts: 2, outOfMemory: { machine: oomMachine } } + : { maxAttempts: 1 }; + + const task = createTask< + TIdentifier, + ChatTaskWirePayload>, + unknown + >({ + ...restOptions, + retry, + triggerSource: "agent", + agentConfig: { type: "ai-sdk-chat" }, + run: async ( + payload: ChatTaskWirePayload>, + { signal: runSignal, ctx } + ) => { + locals.set(chatAgentRunContextKey, ctx); + + // Bind the run to its backing Session so every module-level helper + // (chat.stream, chat.messages, chat.stopSignal) resolves to this + // chat's `.in` / `.out` channels. + // + // Address everywhere by `payload.chatId` (the session externalId): + // matches what the transport puts in URL paths and waitpoint keys, + // and what the server-side trigger flow uses as the session + // identity. The Session row is created by `POST /sessions` (via + // `chat.createStartSessionAction` or browser-direct) before this + // run is triggered — no client-side upsert needed here. + locals.set(chatSessionHandleKey, sessions.open(payload.chatId)); + taskContext.setConversationId(payload.chatId); + + // Stamp `gen_ai.conversation.id` on the run-level span. Every + // nested span inherits the same attribute via + // `TaskContextSpanProcessor.onStart`. + const activeSpan = trace.getActiveSpan(); + stampConversationIdOnActiveSpan(payload.chatId, activeSpan); + + // Store static UIMessageStream options in locals so resolveUIMessageStreamOptions() can read them + if (uiMessageStreamOptions) { + locals.set(chatUIStreamStaticKey, uiMessageStreamOptions); + } + + // Store onCompacted hook in locals so chat.compact() can call it + if (onCompacted) { + locals.set(chatOnCompactedKey, onCompacted); + } + + if (prepareMessages) { + locals.set(chatPrepareMessagesKey, prepareMessages); + } + + if (compaction) { + locals.set( + chatAgentCompactionKey, + compaction as unknown as ChatAgentCompactionOptions + ); + } + + if (pendingMessagesConfig) { + locals.set(chatPendingMessagesKey, pendingMessagesConfig); + } + + let currentWirePayload = payload; + const continuation = payload.continuation ?? false; + const previousRunId = payload.previousRunId; + const preloaded = payload.trigger === "preload"; + + // Accumulated model messages across turns. Seeded at boot from a + // durable snapshot + `session.out` replay (or `hydrateMessages` if + // registered) — the wire is delta-only now, no longer a seed. + let accumulatedMessages: ModelMessage[] = []; + + // Accumulated UI messages for persistence. Mirrors the model accumulator + // but in frontend-friendly UIMessage format (with parts, id, etc.). + let accumulatedUIMessages: TUIMessage[] = []; + + // ── Snapshot + replay (gated on prior-state signals) ───────────── + // + // With `hydrateMessages` registered the customer owns history — the + // hook fires per-turn and produces the canonical chain from their DB. + // Skip both reads entirely: no need to load a blob the customer's + // hook will overwrite. + // + // Without it, both reads are gated on `couldHavePriorState`. A fresh + // chat (no continuation, attempt 1) can't have a snapshot OR replay + // records by definition — `readChatSnapshot` would 404 and + // `replaySessionOutTail` would return [], and the round-trips + // collectively cost ~600ms on every first-message TTFC. Both reads + // swallow errors internally; the agent stays available either way. + const sessionIdForSnapshot = payload.sessionId ?? payload.chatId; + let bootSnapshot: ChatSnapshotV1 | undefined; + let replayed: TUIMessage[] = []; + const couldHavePriorState = + payload.continuation === true || ctx.attempt.number > 1; + + if (!hydrateMessages && couldHavePriorState) { + try { + bootSnapshot = await tracer.startActiveSpan( + "chat.boot.snapshot.read", + async () => readChatSnapshot(sessionIdForSnapshot) + ); + } catch (error) { + // `readChatSnapshot` already swallows + warns internally; this catch + // is just belt-and-suspenders against tracer/span errors. + logger.warn("chat.agent: snapshot read failed; continuing without snapshot", { + error: error instanceof Error ? error.message : String(error), + sessionId: sessionIdForSnapshot, + }); + } + + try { + replayed = await tracer.startActiveSpan("chat.boot.replay", async () => + replaySessionOutTail(sessionIdForSnapshot, { + lastEventId: bootSnapshot?.lastOutEventId, + }) + ); + } catch (error) { + logger.warn("chat.agent: session.out replay failed; using snapshot only", { + error: error instanceof Error ? error.message : String(error), + sessionId: sessionIdForSnapshot, + }); + } + } + + // ── session.in dedup cutoff ──────────────────────────────────── + // + // A fresh worker subscribes to `session.in` from seq 0 and would + // re-deliver every record ever appended — including user messages + // from turns already completed on a prior run. Without dedup, the + // loop would re-process them as fresh turns and the slim-wire merge + // would replace-by-id against the snapshot-restored copies, yielding + // no-op replaces while the customer's actual new message waits in + // the queue. + // + // The cutoff is the timestamp of the last `trigger:turn-complete` + // chunk on `session.out`. When we have a snapshot, that timestamp is + // already in `lastOutTimestamp` — use it directly to skip the + // O(stream-length) scan. Fall back to the scan only when no snapshot + // is available (first-ever OOM retry, or `hydrateMessages` + // short-circuited the snapshot read). + // + // Applies in three cases (any of which means session.in has records + // belonging to completed turns the new run should skip): + // - OOM retry (`ctx.attempt.number > 1`) + // - Continuation run (`payload.continuation === true`) — prior run + // crashed / was canceled / requested upgrade + // - Snapshot exists at all (catches edge cases where the wire + // didn't set `continuation` but a snapshot indicates prior turns) + const needsDedupCutoff = + ctx.attempt.number > 1 || + payload.continuation === true || + bootSnapshot !== undefined; + + if (needsDedupCutoff) { + try { + let cutoff = bootSnapshot?.lastOutTimestamp; + if (cutoff === undefined) { + cutoff = await findLatestTurnCompleteTimestamp(payload.chatId); + } + if (cutoff !== undefined) { + sessionStreams.setMinTimestamp(payload.chatId, "in", cutoff); + } + } catch (error) { + logger.warn( + "chat.agent: session.in dedup cutoff lookup failed; old messages may replay", + { error: error instanceof Error ? error.message : String(error) } + ); + } + } + + // ── Merge + head-start bootstrap ──────────────────────────────── + if (!hydrateMessages) { + accumulatedUIMessages = mergeByIdReplaceWins( + (bootSnapshot?.messages as TUIMessage[]) ?? [], + replayed + ); + + // ── Head-start bootstrap ───────────────────────────────────── + // + // The very-first turn of a head-start handover has no snapshot + // (it doesn't exist yet) and no `session.out` history (the run + // just woke up). The customer's HTTP route handler ships full + // UIMessage history via `headStartMessages` — that's the only + // path where wire-borne UIMessage[] still seeds the accumulator, + // and it's safe because the route handler isn't subject to the + // `/in/append` 512 KiB cap. + if ( + accumulatedUIMessages.length === 0 && + payload.trigger === "handover-prepare" && + Array.isArray(payload.headStartMessages) && + payload.headStartMessages.length > 0 + ) { + accumulatedUIMessages = [...(payload.headStartMessages as TUIMessage[])]; + } + + if (accumulatedUIMessages.length > 0) { + try { + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } catch (error) { + logger.warn("chat.agent: toModelMessages failed at boot; starting empty", { + error: error instanceof Error ? error.message : String(error), + sessionId: sessionIdForSnapshot, + }); + accumulatedMessages = []; + } + } + + // Make the seeded UI accumulator visible to `chat.history.*` + // before any hook (`onChatStart`, `onTurnStart`, etc.) fires. + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + + } + + // Token usage tracking across turns + let previousTurnUsage: LanguageModelUsage | undefined; + let cumulativeUsage: LanguageModelUsage = emptyUsage(); + + // Mutable reference to the current turn's stop controller so the + // stop input stream listener (registered once) can abort the right turn. + let currentStopController: AbortController | undefined; + + // Stop-input subscription is registered AFTER preload's wait resolves + // (see the post-preload block below). Registering it earlier would + // cause it to drain any session.in records buffered before the runtime + // started — including the customer's first user message arriving on + // `kind: "message"`. The persistent-listener semantics of session.in + // pop the buffer when a handler attaches; the stop listener filters + // out anything that isn't `kind: "stop"` and silently swallows the + // user message instead of leaving it for `messagesInput.waitWithIdleTimeout` + // to pick up. + let stopSub: { off: () => void } | undefined; + + try { + // Handle preloaded runs — fire onPreload, then wait for the first real message + if (preloaded) { + if (activeSpan) { + activeSpan.setAttribute("chat.preloaded", true); + } + + const currentRunId = ctx.run.id; + let preloadAccessToken = ""; + if (currentRunId) { + try { + preloadAccessToken = await auth.createPublicToken({ + scopes: { + read: { + runs: currentRunId, + sessions: payload.chatId, + }, + write: { + inputStreams: currentRunId, + sessions: payload.chatId, + }, + }, + expirationTime: chatAccessTokenTTL, + }); + } catch { + // Token creation failed + } + } + + // Parse client data for the preload hook + const preloadClientData = ( + parseClientData ? await parseClientData(payload.metadata) : payload.metadata + ) as inferSchemaOut; + + // Fire onPreload hook + if (onPreload) { + await tracer.startActiveSpan( + "onPreload()", + async () => { + await withChatWriter(async (writer) => { + await onPreload({ + ctx, + chatId: payload.chatId, + runId: currentRunId, + chatAccessToken: preloadAccessToken, + clientData: preloadClientData, + writer, + }); + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": payload.chatId, + "chat.preloaded": true, + }, + } + ); + } + + // Wait for the first real message — task-level idle settings win over + // `transport.preload(..., { idleTimeoutInSeconds })` / wire payload so + // `chat.agent({ idleTimeoutInSeconds, preloadIdleTimeoutInSeconds })` is authoritative. + const effectivePreloadIdleTimeout = + preloadIdleTimeoutInSeconds ?? + idleTimeoutInSeconds ?? + payload.idleTimeoutInSeconds; + + const effectivePreloadTimeout = + (metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined) ?? + preloadTimeout ?? + turnTimeout; + + const preloadResult = await messagesInput.waitWithIdleTimeout({ + idleTimeoutInSeconds: effectivePreloadIdleTimeout, + timeout: effectivePreloadTimeout, + spanName: "waiting for first message", + skipSuspend: exitAfterPreloadIdle, + onSuspend: onChatSuspend + ? async () => { + await tracer.startActiveSpan( + "onChatSuspend()", + async () => { + await onChatSuspend({ + phase: "preload", + ctx, + chatId: payload.chatId, + runId: currentRunId, + clientData: preloadClientData, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": payload.chatId, + "chat.suspend.phase": "preload", + }, + } + ); + } + : undefined, + onResume: onChatResume + ? async () => { + await tracer.startActiveSpan( + "onChatResume()", + async () => { + await onChatResume({ + phase: "preload", + ctx, + chatId: payload.chatId, + runId: currentRunId, + clientData: preloadClientData, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": payload.chatId, + "chat.resume.phase": "preload", + }, + } + ); + } + : undefined, + }); + + if (!preloadResult.ok) { + return; // Timed out waiting for first message — end run + } + + let firstMessage = preloadResult.output; + + currentWirePayload = firstMessage as ChatTaskWirePayload< + TUIMessage, + inferSchemaIn + >; + + // Close signal during preload — exit before first turn + if (currentWirePayload.trigger === "close") { + return; + } + } + + // Listen for stop signals for the rest of the run. Registered AFTER + // the preload wait resolves (or skipped immediately for non-preload + // triggers) so the persistent-listener semantics on session.in + // don't drain the buffered user message before + // `messagesInput.waitWithIdleTimeout` can pick it up. A stop signal + // that lands during the preload window is dropped — acceptable, the + // customer can't reasonably stop a chat that hasn't started. + stopSub = stopInput.on((data) => { + currentStopController?.abort(data?.message || "stopped"); + }); + + // Handle handover-prepare runs — wait on session.in for the + // customer's `chat.handover` route handler to either hand off + // mid-turn (tool calls) or signal pure-text completion. + if (payload.trigger === "handover-prepare") { + if (activeSpan) { + activeSpan.setAttribute("chat.handoverPreparing", true); + } + + const handoverResult = await handoverInput.waitWithIdleTimeout({ + idleTimeoutInSeconds: idleTimeoutInSeconds ?? payload.idleTimeoutInSeconds ?? 60, + spanName: "waiting for handover signal", + }); + + if (!handoverResult.ok) { + // Handler crashed before signaling — exit cleanly. + return; + } + + if (handoverResult.output.kind === "handover-skip") { + // Sent only when the customer's handler aborts before + // producing a finishReason. Normal pure-text and + // tool-call finishes go through `kind: "handover"` with + // `isFinal: true | false`. Exit without firing any turn + // hooks. + return; + } + + // kind === "handover": stash the partial assistant message + // so turn-0 setup can append it after loading user + // messages. Two branches downstream, switched by `isFinal`: + // - `false`: customer's step 1 ended with `tool-calls`. + // The agent's `streamText` sees pending tool-calls (via + // the approval round in the partial) and executes them, + // then runs step 2's LLM call. + // - `true`: customer's step 1 ended pure-text. The agent + // runs the turn-loop hooks but SKIPS the `streamText` + // call entirely (the response is already complete). + // `onTurnComplete` fires with the partial as + // `responseMessage` so persistence works normally. + locals.set( + chatHandoverPartialKey, + handoverResult.output.partialAssistantMessage + ); + // Stash the customer-side step-1 messageId. Turn-0 setup + // uses it to seed the synthesized partial UIMessage with the + // SAME id, so the agent's post-handover chunks merge into + // the same assistant message on the browser side. + if (handoverResult.output.messageId) { + locals.set(chatHandoverMessageIdKey, handoverResult.output.messageId); + } + locals.set(chatHandoverIsFinalKey, handoverResult.output.isFinal); + + // Synthesize a wire payload that the turn loop treats as a + // normal first-turn message. The accumulator was already seeded + // at boot from `payload.headStartMessages` (see B.3 head-start + // bootstrap), so the rewritten `submit-message` carries no + // delta — the loop runs streamText against the seeded state. + currentWirePayload = { + ...payload, + trigger: "submit-message", + message: undefined, + headStartMessages: undefined, + } as ChatTaskWirePayload>; + } + + for (let turn = 0; turn < maxTurns; turn++) { + try { + // Extract turn-level context before entering the span. Slim + // wire: at most one delta message per record. `headStartMessages` + // is consumed at boot only (via `payload.headStartMessages`) + // and intentionally discarded here. + const { + metadata: wireMetadata, + message: incomingMessage, + headStartMessages: _hsm, + ...restWire + } = currentWirePayload; + void _hsm; + const incomingMessages: TUIMessage[] = incomingMessage + ? [incomingMessage as TUIMessage] + : []; + // Cleaning happens once here so `extractLastUserMessageText` and + // every downstream consumer see the same message shape — and + // `cleanupAbortedParts` no longer has to be re-applied below. + const cleanedIncomingMessages: TUIMessage[] = incomingMessages.map((msg) => + msg.role === "assistant" ? cleanupAbortedParts(msg) : msg + ); + const clientData = ( + parseClientData ? await parseClientData(wireMetadata) : wireMetadata + ) as inferSchemaOut; + const lastUserMessage = extractLastUserMessageText(cleanedIncomingMessages); + + // Actions are not turns. They use a different span name + // and don't carry a turn.number. Branched on at `isAction`. + const isAction = currentWirePayload.trigger === "action"; + const spanName = isAction ? "chat action" : `chat turn ${turn + 1}`; + + const turnAttributes: Attributes = { + ...(isAction ? {} : { "turn.number": turn + 1 }), + "gen_ai.conversation.id": currentWirePayload.chatId, + "gen_ai.operation.name": "chat", + "chat.trigger": currentWirePayload.trigger, + [SemanticInternalAttributes.STYLE_ICON]: isAction + ? "tabler-bolt" + : "tabler-message-chatbot", + [SemanticInternalAttributes.ENTITY_TYPE]: isAction ? "chat-action" : "chat-turn", + }; + + if (lastUserMessage) { + turnAttributes["chat.user_message"] = lastUserMessage; + + // Show a truncated preview of the user message as an accessory + const preview = + lastUserMessage.length > 80 ? lastUserMessage.slice(0, 80) + "..." : lastUserMessage; + Object.assign( + turnAttributes, + accessoryAttributes({ + items: [{ text: preview, variant: "normal" }], + style: "codepath", + }) + ); + } + + if (wireMetadata !== undefined) { + turnAttributes["chat.client_data"] = + typeof wireMetadata === "string" ? wireMetadata : JSON.stringify(wireMetadata); + } + + const turnResult = await tracer.startActiveSpan( + spanName, + async (turnSpan) => { + // (errors are caught by the outer try/catch which writes an error chunk) + locals.set(chatPipeCountKey, 0); + locals.set(chatDeferKey, new Set()); + locals.set(chatCompactionStateKey, undefined); + locals.set(chatSteeringQueueKey, []); + locals.set(chatResponsePartsKey, []); + // NOTE: chatBackgroundQueueKey is NOT reset here — messages injected + // by deferred work from the previous turn's onTurnComplete need to + // survive into the next turn. The queue is drained before run(). + locals.set(chatInjectedMessageIdsKey, new Set()); + + // Store chat context for auto-detection by task-tool subtasks (ai.toolExecute / legacy ai.tool) + locals.set(chatTurnContextKey, { + chatId: currentWirePayload.chatId, + turn, + continuation, + clientData, + }); + + // ── chat.store hydration ───────────────────────────────── + // Apply `hydrateStore` (if configured) and `incomingStore` + // from the wire payload before `run()` fires. Order: + // + // 1. `previousStore` = what's in memory from the prior turn. + // 2. `hydrateStore` returns authoritative value from backend. + // 3. `incomingStore` (client-side) wins if present + // (AG-UI last-write-wins). + // + // Emit a snapshot chunk after hydration so the frontend + // observing the stream sees the initial value. + { + const previousStoreValue = locals.get(chatStoreSlotKey)?.value; + const incomingStoreValue = + "incomingStore" in currentWirePayload + ? (currentWirePayload as { incomingStore?: unknown }).incomingStore + : undefined; + + let nextStoreValue: unknown = previousStoreValue; + let storeChanged = false; + + if (hydrateStore) { + try { + const hydratedValue = await tracer.startActiveSpan( + "hydrateStore()", + async () => { + return hydrateStore({ + chatId: currentWirePayload.chatId, + turn, + trigger: currentWirePayload.trigger as + | "submit-message" + | "regenerate-message" + | "action" + | "preload", + previousStore: previousStoreValue, + incomingStore: incomingStoreValue, + clientData, + continuation, + previousRunId, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: + "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.turn": turn + 1, + }, + } + ); + nextStoreValue = hydratedValue; + storeChanged = true; + } catch (err) { + // Surface hydration failures as a turn error — users + // rely on `chat.store.get()` reflecting authoritative + // state, so silently continuing with stale data is + // worse than failing loudly. + throw err; + } + } + + if (incomingStoreValue !== undefined) { + // Last-write-wins: the client-sent value overrides + // whatever we hydrated. The client may have made local + // updates the backend hasn't seen yet. + nextStoreValue = incomingStoreValue; + storeChanged = true; + } + + if (storeChanged) { + chatStoreSetSilent(nextStoreValue); + chatStoreEmitSnapshot(nextStoreValue); + fireStoreListeners(nextStoreValue); + } + } + + // Per-turn stop controller (reset each turn) + const stopController = new AbortController(); + currentStopController = stopController; + locals.set(chatStopControllerKey, stopController); + + // Three signals for the user's run function + const stopSignal = stopController.signal; + const cancelSignal = runSignal; + const combinedSignal = AbortSignal.any([runSignal, stopController.signal]); + + // Buffer messages that arrive during streaming + const pendingMessages: ChatTaskWirePayload< + TUIMessage, + inferSchemaIn + >[] = []; + const pmConfig = locals.get(chatPendingMessagesKey); + const msgSub = messagesInput.on(async (msg) => { + // If pendingMessages is configured, route to the steering queue + // instead of the wire buffer. The frontend handles re-sending + // non-injected messages via sendMessage on turn complete. + if (pmConfig) { + // Slim wire: at most one delta message per record. The + // pendingMessages handler reads `msg.message` directly + // instead of slicing an array — a wire record arrives + // with the new user message in `.message`, or no message + // at all (regenerate / preload / close / handover-prepare). + const lastUIMessage = msg.message as TUIMessage | undefined; + if (lastUIMessage) { + if (pmConfig.onReceived) { + try { + await pmConfig.onReceived({ + message: lastUIMessage as TUIMessage, + chatId: currentWirePayload.chatId, + turn, + }); + } catch { + /* non-fatal */ + } + } + + try { + const queue = locals.get(chatSteeringQueueKey) ?? []; + // Deduplicate by message ID — guards against double-sends + if ( + lastUIMessage.id && + queue.some((e) => e.uiMessage.id === lastUIMessage.id) + ) { + return; + } + const modelMsgs = await toModelMessages([lastUIMessage]); + queue.push({ + uiMessage: lastUIMessage as UIMessage, + modelMessages: modelMsgs, + }); + locals.set(chatSteeringQueueKey, queue); + } catch { + /* conversion failed — skip steering queue */ + } + } + return; // Don't add to wire buffer — frontend handles non-injected case + } + + // No pendingMessages config — standard wire buffer for next turn + pendingMessages.push( + msg as ChatTaskWirePayload> + ); + }); + + // Track new messages for this turn (user input + assistant response). + const turnNewModelMessages: ModelMessage[] = []; + const turnNewUIMessages: TUIMessage[] = []; + + // ── Action handling ────────────────────────────────────── + // Actions arrive on the same input stream but with + // trigger === "action". They are NOT turns — only + // `hydrateMessages` and `onAction` fire. No turn lifecycle + // hooks (`onTurnStart` / `prepareMessages` / + // `onBeforeTurnComplete` / `onTurnComplete`) and no + // `run()` invocation. To produce a model response from + // an action, return a `StreamTextResult` (auto-piped), + // string, or UIMessage from `onAction`. Turn counter + // does not advance. + let actionStreamResult: unknown = undefined; + if (isAction) { + // Parse and validate the action payload + const parsedAction = parseAction + ? await parseAction(currentWirePayload.action) + : currentWirePayload.action; + + // Hydrate messages from backend if configured + if (hydrateMessages) { + const hydrated = await tracer.startActiveSpan( + "hydrateMessages()", + async () => { + return hydrateMessages({ + chatId: currentWirePayload.chatId, + turn, + trigger: "action", + incomingMessages: [] as TUIMessage[], + previousMessages: [...accumulatedUIMessages], + clientData, + continuation, + previousRunId, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.trigger": "action", + }, + } + ); + accumulatedUIMessages = [...hydrated] as TUIMessage[]; + accumulatedMessages = await toModelMessages(hydrated); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + } + + // Fire onAction — handler may mutate state via + // `chat.history.*` and / or return a model response. + if (onAction) { + actionStreamResult = await tracer.startActiveSpan( + "onAction()", + async () => { + return await onAction({ + action: parsedAction as any, + chatId: currentWirePayload.chatId, + turn, + clientData, + uiMessages: accumulatedUIMessages, + messages: accumulatedMessages, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.action": + typeof parsedAction === "object" && parsedAction !== null + ? JSON.stringify(parsedAction) + : String(parsedAction), + }, + } + ); + + // Apply chat.history mutations from onAction + const actionOverride = locals.get(chatOverrideMessagesKey); + if (actionOverride) { + locals.set(chatOverrideMessagesKey, undefined); + accumulatedUIMessages = [...actionOverride] as TUIMessage[]; + accumulatedMessages = await toModelMessages(actionOverride); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + } + } else { + warnMissingOnActionOnce(); + } + } + + // ── Message handling (non-action turns) ─────────────────── + // + // Slim wire: at most one delta message arrives per record. + // The accumulator was already seeded at boot from a durable + // snapshot + `session.out` replay (or `hydrateMessages`, + // which also fires per-turn below). Per-turn handling is + // therefore a delta merge, not a full-history reset. + if (currentWirePayload.trigger !== "action") { + + let cleanedUIMessages: TUIMessage[] = cleanedIncomingMessages; + + // Validate/transform UIMessages before conversion — catches malformed + // messages from storage or untrusted input before they reach the model. + // Slim wire: triggers like `regenerate-message` carry no incoming + // message; nothing to validate, so skip the hook to avoid feeding + // `[]` to validators (AI SDK's `validateUIMessages` rejects empty). + if (onValidateMessages && cleanedUIMessages.length > 0) { + cleanedUIMessages = (await tracer.startActiveSpan( + "onValidateMessages()", + async () => { + return onValidateMessages({ + messages: cleanedUIMessages as TUIMessage[], + chatId: currentWirePayload.chatId, + turn, + trigger: currentWirePayload.trigger as ValidateMessagesEvent["trigger"], + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.turn": turn + 1, + "chat.messages.count": cleanedUIMessages.length, + }, + } + )) as TUIMessage[]; + } + + if (hydrateMessages) { + // Backend hydration: load the full message history from the user's + // backend, replacing the built-in accumulator entirely. With slim + // wire, `incomingMessages` is consistently 0-or-1-length — what + // was always true for `submit-message` is now true for every + // trigger. + const hydrated = await tracer.startActiveSpan( + "hydrateMessages()", + async () => { + return hydrateMessages({ + chatId: currentWirePayload.chatId, + turn, + trigger: currentWirePayload.trigger as + | "submit-message" + | "regenerate-message", + incomingMessages: cleanedUIMessages as TUIMessage[], + previousMessages: [...accumulatedUIMessages], + clientData, + continuation, + previousRunId, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.turn": turn + 1, + "chat.trigger": currentWirePayload.trigger, + "chat.incoming_messages.count": cleanedUIMessages.length, + }, + } + ); + + // Auto-merge tool approval updates: if any incoming wire message + // has an ID that matches a hydrated message, replace it. This makes + // tool approvals work transparently with backend hydration. + const merged = [...hydrated] as TUIMessage[]; + for (const incoming of cleanedUIMessages) { + if (!incoming.id) continue; + const idx = merged.findIndex((m) => m.id === incoming.id); + if (idx !== -1) { + merged[idx] = incoming as TUIMessage; + } + } + + accumulatedUIMessages = merged; + accumulatedMessages = await toModelMessages(merged); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + + // Track new messages for onTurnComplete.newUIMessages + if ( + currentWirePayload.trigger === "submit-message" && + cleanedUIMessages.length > 0 + ) { + const lastUI = cleanedUIMessages[cleanedUIMessages.length - 1]!; + turnNewUIMessages.push(lastUI); + const lastModel = (await toModelMessages([lastUI]))[0]; + if (lastModel) turnNewModelMessages.push(lastModel); + } + } else { + // Default delta-merge accumulation. + // + // The accumulator was seeded at boot from snapshot+replay, + // so it already reflects prior history. Per-turn handling + // appends/replaces the single incoming delta message and + // (for regenerate) trims the tail. + if (currentWirePayload.trigger === "regenerate-message") { + // Regenerate: trim trailing assistant messages from the + // accumulator until the tail is a user message. AI SDK's + // frontend `regenerate()` already removed the trailing + // assistant from its local store; the wire signals "do + // the same here", and the agent re-runs from the new + // tail. No incoming message accompanies this trigger. + while ( + accumulatedUIMessages.length > 0 && + accumulatedUIMessages[accumulatedUIMessages.length - 1]!.role !== "user" + ) { + accumulatedUIMessages.pop(); + } + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } else if (cleanedUIMessages.length > 0) { + // Submit-message (and the special-cased + // handover-prepare → submit-message rewrite earlier in + // this scope): append-or-replace-by-id for the single + // delta message. + // + // Tool approval responses arrive as a single assistant + // message whose id collides with the existing assistant + // in the accumulator — we replace by id. The fallback + // for HITL `addToolOutput` continuations where AI SDK + // regenerates the id (TRI-9137) still applies via + // `rewriteIncomingIdViaToolCallMap`. + let replaced = false; + for (const raw of cleanedUIMessages) { + let incoming = raw; + let idx = accumulatedUIMessages.findIndex( + (m) => m.id === incoming.id + ); + if (idx === -1) { + const rewritten = rewriteIncomingIdViaToolCallMap(incoming); + if (rewritten.id !== incoming.id) { + incoming = rewritten as typeof raw; + idx = accumulatedUIMessages.findIndex( + (m) => m.id === incoming.id + ); + } + } + if (idx !== -1) { + accumulatedUIMessages[idx] = incoming as TUIMessage; + replaced = true; + } else { + accumulatedUIMessages.push(incoming as TUIMessage); + turnNewUIMessages.push(incoming as TUIMessage); + } + recordToolCallIdsFromMessage(incoming); + } + if (replaced) { + // Replacement changes structure — reconvert all model + // messages instead of appending. + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } else { + const incomingModelMessages = await toModelMessages(cleanedUIMessages); + accumulatedMessages.push(...incomingModelMessages); + } + if (turnNewUIMessages.length > 0) { + turnNewModelMessages.push( + ...(await toModelMessages(turnNewUIMessages)) + ); + } + } + // `preload` / `close` / `handover-prepare` and submits + // with no incoming message fall through with the boot- + // seeded accumulator unchanged. + + if (turn === 0) { + // Head-start handover splice (turn 0 only): the + // `chat.handover` route handler signalled a mid-turn + // handover, so splice its partial assistant response + // (text + pending tool-calls + the synthesized + // tool-approval round) onto the accumulator. + // `streamText` then hits AI SDK's initial-tool- + // execution branch, runs the agent-side tool executes, + // and resumes from step 2 — skipping the first model + // call (already done by the handler). + // + // We also synthesize a UIMessage form of the partial + // assistant and push it to `accumulatedUIMessages` so + // AI SDK's `processUIMessageStream` (invoked when the + // run loop calls `runResult.toUIMessageStream({ + // onFinish })`) can initialize `state.message` from + // the trailing assistant in `originalMessages`. Without + // that, the `tool-output-available` chunks emitted by + // the initial-tool-execution branch can't find their + // matching tool-call in state and AI SDK throws + // `UIMessageStreamError: No tool invocation found`. + const pendingHandoverPartial = locals.get(chatHandoverPartialKey); + if (pendingHandoverPartial && pendingHandoverPartial.length > 0) { + accumulatedMessages.push(...pendingHandoverPartial); + const handoverMessageId = locals.get(chatHandoverMessageIdKey); + const partialUI = synthesizeHandoverUIMessage( + pendingHandoverPartial, + handoverMessageId + ); + if (partialUI) { + accumulatedUIMessages.push(partialUI as TUIMessage); + } + locals.set(chatHandoverPartialKey, []); // consume once + } + } + + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + } + + } // end if (trigger !== "action") + + // ── Action result handling ────────────────────────────── + // For action turns, skip the turn machinery entirely. + // If `onAction` returned a stream / string / UIMessage, + // pipe it as the response. Either way, emit + // `trigger:turn-complete` and then fall through to the + // wait-for-next-message logic (shared with message turns). + // The turn counter is decremented so the next iteration + // sees the same `turn` value — actions don't count. + if (isAction) { + msgSub.off(); + + if ( + (locals.get(chatPipeCountKey) ?? 0) === 0 && + isUIMessageStreamable(actionStreamResult) + ) { + try { + const resolvedOptions = resolveUIMessageStreamOptions(); + const uiStream = ( + actionStreamResult as UIMessageStreamable + ).toUIMessageStream({ + ...resolvedOptions, + generateMessageId: + resolvedOptions.generateMessageId ?? generateMessageId, + }); + await pipeChat(uiStream, { + signal: combinedSignal, + spanName: "stream response", + }); + } catch (error) { + if ( + error instanceof Error && + error.name === "AbortError" && + runSignal.aborted + ) { + return "exit"; + } + throw error; + } + } + + await writeTurnCompleteChunk(currentWirePayload.chatId); + + // Don't consume a turn iteration — actions aren't turns. + turn--; + } + + if (!isAction) { + + // Mint a scoped public access token once per turn, reused for + // onChatStart, onTurnStart, onTurnComplete, and the turn-complete chunk. + const currentRunId = ctx.run.id; + let turnAccessToken = ""; + if (currentRunId) { + try { + turnAccessToken = await auth.createPublicToken({ + scopes: { + read: { + runs: currentRunId, + sessions: payload.chatId, + }, + write: { + inputStreams: currentRunId, + sessions: payload.chatId, + }, + }, + expirationTime: chatAccessTokenTTL, + }); + } catch { + // Token creation failed + } + } + + // Fire onChatStart on the first turn + if (turn === 0 && onChatStart) { + await tracer.startActiveSpan( + "onChatStart()", + async () => { + await withChatWriter(async (writer) => { + await onChatStart({ + ctx, + chatId: currentWirePayload.chatId, + messages: accumulatedMessages, + clientData, + runId: currentRunId, + chatAccessToken: turnAccessToken, + continuation, + previousRunId, + preloaded, + writer, + }); + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.messages.count": accumulatedMessages.length, + "chat.continuation": continuation, + "chat.preloaded": preloaded, + ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}), + }, + } + ); + } + + // Fire onTurnStart before running user code — persist messages + // so a mid-stream page refresh still shows the user's message. + if (onTurnStart) { + await tracer.startActiveSpan( + "onTurnStart()", + async () => { + await withChatWriter(async (writer) => { + await onTurnStart({ + ctx, + chatId: currentWirePayload.chatId, + messages: accumulatedMessages, + uiMessages: accumulatedUIMessages, + turn, + runId: currentRunId, + chatAccessToken: turnAccessToken, + clientData, + continuation, + previousRunId, + preloaded, + previousTurnUsage, + totalUsage: cumulativeUsage, + writer, + }); + }); + + // Check if onTurnStart replaced messages (compaction or chat.history) + const turnStartOverride = locals.get(chatOverrideMessagesKey); + if (turnStartOverride) { + locals.set(chatOverrideMessagesKey, undefined); + accumulatedUIMessages = [...turnStartOverride] as TUIMessage[]; + accumulatedMessages = await toModelMessages(turnStartOverride); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.turn": turn + 1, + "chat.messages.count": accumulatedMessages.length, + "chat.trigger": currentWirePayload.trigger, + "chat.continuation": continuation, + "chat.preloaded": preloaded, + ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}), + }, + } + ); + } + + // chat.requestUpgrade() called in onTurnStart (or onValidateMessages) — + // skip run() and signal the transport to re-trigger the same message + // on the new version. + if (locals.get(chatUpgradeRequestedKey)) { + await writeUpgradeRequiredChunk(); + return "exit"; + } + + // Captured by the onFinish callback below — works even on abort/stop. + let capturedResponseMessage: TUIMessage | undefined; + let capturedFinishReason: FinishReason | undefined; + + // Promise that resolves when the AI SDK's onFinish fires. + // On abort, the stream's cancel() handler calls onFinish + // asynchronously AFTER pipeChat resolves, so we must await + // this to avoid a race where we check capturedResponseMessage + // before it's been set. + let resolveOnFinish: () => void; + const onFinishPromise = new Promise((r) => { + resolveOnFinish = r; + }); + let onFinishAttached = false; + let runResult: unknown; + + // Pure-text head-start: customer's step 1 IS the + // final response. Skip the user's `run` callback + // (no LLM call) and use the synthesized partial + // UIMessage as `capturedResponseMessage`. The post- + // turn flow (`onBeforeTurnComplete` → + // `onTurnComplete` → trigger:turn-complete) fires + // normally so persistence works. + const headStartIsFinal = locals.get(chatHandoverIsFinalKey); + const isHeadStartFinalTurn = turn === 0 && headStartIsFinal === true; + if (isHeadStartFinalTurn) { + locals.set(chatHandoverIsFinalKey, undefined); // consume once + } + + try { + // Drain any messages injected by background work (e.g. self-review from previous turn). + // Skip if the last message is a tool message — appending after it would + // prevent streamText from finding pending tool approvals (it checks + // the last message). The queued messages will be picked up by prepareStep + // at the next step boundary instead. + const lastAccumulated = accumulatedMessages[accumulatedMessages.length - 1]; + const bgQueue = locals.get(chatBackgroundQueueKey); + if (bgQueue && bgQueue.length > 0 && lastAccumulated?.role !== "tool") { + accumulatedMessages.push(...bgQueue.splice(0)); + } + + if (isHeadStartFinalTurn) { + // The synthesized partial UIMessage IS the response. + // It was pushed to `accumulatedUIMessages` during the + // submit-message branch's splice; recover it as the + // last assistant. + const lastUI = accumulatedUIMessages[accumulatedUIMessages.length - 1]; + if (lastUI && lastUI.role === "assistant") { + capturedResponseMessage = lastUI; + capturedFinishReason = "stop"; + } + // Don't call userRun. Don't pipe. Skip directly + // to the post-turn flow below. + } else { + const preparedMessages = await applyPrepareMessages(accumulatedMessages, "run"); + runResult = await userRun({ + ...restWire, + messages: preparedMessages, + clientData, + continuation, + previousRunId, + preloaded, + previousTurnUsage, + totalUsage: cumulativeUsage, + ctx, + signal: combinedSignal, + cancelSignal, + stopSignal, + } as any); + } + + // Auto-pipe if the run function returned a StreamTextResult or similar, + // but only if pipeChat() wasn't already called manually during this turn. + // We call toUIMessageStream ourselves to attach onFinish for response capture. + // Pass originalMessages so the AI SDK reuses message IDs across turns + // (e.g. for tool approval continuations / HITL flows). + if ((locals.get(chatPipeCountKey) ?? 0) === 0 && isUIMessageStreamable(runResult)) { + onFinishAttached = true; + const resolvedOptions = resolveUIMessageStreamOptions(); + // For action turns, don't pass originalMessages: the response + // should always be a fresh assistant message, not a continuation + // of whatever trailing assistant was left after chat.history + // mutations. + const isActionTurn = currentWirePayload.trigger === "action"; + const uiStream = runResult.toUIMessageStream({ + ...resolvedOptions, + // Pass originalMessages so the AI SDK reuses message IDs across + // turns (e.g. for tool approval continuations / HITL flows). + // Omit for action turns to force a fresh response ID. + ...(isActionTurn + ? {} + : { originalMessages: accumulatedUIMessages }), + // Always provide generateMessageId so the start chunk carries a + // messageId. Without this, the frontend and backend generate IDs + // independently and they won't match for ID-based dedup. + generateMessageId: resolvedOptions.generateMessageId ?? generateMessageId, + onFinish: ({ + responseMessage, + finishReason, + }: { + responseMessage: UIMessage; + finishReason?: FinishReason; + }) => { + capturedResponseMessage = responseMessage as TUIMessage; + capturedFinishReason = finishReason; + resolveOnFinish!(); + }, + }); + await pipeChat(uiStream, { signal: combinedSignal, spanName: "stream response" }); + } + } catch (error) { + // Handle AbortError from streamText gracefully + if (error instanceof Error && error.name === "AbortError") { + if (runSignal.aborted) { + return "exit"; // Full run cancellation — exit + } + // Stop generation — fall through to continue the loop + } else { + throw error; + } + } finally { + msgSub.off(); + } + + // Wait for onFinish to fire — on abort this may resolve slightly + // after pipeChat, since the stream's cancel() handler is async. + // Race with a timeout so a stop-abort that prevents onFinish from + // firing doesn't hang the turn loop indefinitely. + if (onFinishAttached) { + await Promise.race([ + onFinishPromise, + new Promise((r) => setTimeout(r, 2_000)), + ]); + } + + // Capture token usage from the streamText result (if available). + // totalUsage is a PromiseLike that resolves after the stream is consumed. + // Race with a 2s timeout — on stop-abort the AI SDK's totalUsage + // promise can hang indefinitely (the underlying provider stream + // never reports final usage), which would block the turn loop + // from ever firing onTurnComplete / writeTurnComplete. + let turnUsage: LanguageModelUsage | undefined; + if (runResult != null && typeof (runResult as any).totalUsage?.then === "function") { + try { + turnUsage = (await Promise.race([ + (runResult as any).totalUsage, + new Promise((r) => setTimeout(() => r(undefined), 2_000)), + ])) as LanguageModelUsage | undefined; + } catch { + /* non-fatal — usage capture failed */ + } + } + if (turnUsage) { + cumulativeUsage = addUsage(cumulativeUsage, turnUsage); + previousTurnUsage = turnUsage; + + // Add usage attributes to the turn span + if (turnUsage.inputTokens != null) { + turnSpan.setAttribute("gen_ai.usage.input_tokens", turnUsage.inputTokens); + } + if (turnUsage.outputTokens != null) { + turnSpan.setAttribute("gen_ai.usage.output_tokens", turnUsage.outputTokens); + } + if (turnUsage.totalTokens != null) { + turnSpan.setAttribute("gen_ai.usage.total_tokens", turnUsage.totalTokens); + } + if (cumulativeUsage.totalTokens != null) { + turnSpan.setAttribute( + "gen_ai.usage.cumulative_total_tokens", + cumulativeUsage.totalTokens + ); + } + if (cumulativeUsage.inputTokens != null) { + turnSpan.setAttribute( + "gen_ai.usage.cumulative_input_tokens", + cumulativeUsage.inputTokens + ); + } + if (cumulativeUsage.outputTokens != null) { + turnSpan.setAttribute( + "gen_ai.usage.cumulative_output_tokens", + cumulativeUsage.outputTokens + ); + } + } + + // Check if run() (e.g. via prepareStep or chat.history) replaced messages + // during this turn. The updated messages become the new base, and the + // response gets appended on top. + const runOverride = locals.get(chatOverrideMessagesKey); + if (runOverride) { + locals.set(chatOverrideMessagesKey, undefined); + accumulatedUIMessages = [...runOverride] as TUIMessage[]; + accumulatedMessages = await toModelMessages(runOverride); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + } + + // Check if compaction set a model-only override (preserves UI messages). + // Apply compactUIMessages/compactModelMessages callbacks if configured. + const modelOnlyOverride = locals.get(chatOverrideModelMessagesKey); + if (modelOnlyOverride) { + const compactionSummary = locals.get(chatCompactionStateKey)?.summary ?? ""; + const taskCompactionConfig = locals.get(chatAgentCompactionKey); + locals.set(chatOverrideModelMessagesKey, undefined); + + const compactEvent: CompactMessagesEvent = { + summary: compactionSummary, + uiMessages: accumulatedUIMessages, + modelMessages: accumulatedMessages, + chatId: currentWirePayload.chatId, + turn, + clientData, + source: "inner", + }; + + // Apply model messages: callback or default (use override) + accumulatedMessages = taskCompactionConfig?.compactModelMessages + ? await taskCompactionConfig.compactModelMessages(compactEvent) + : modelOnlyOverride; + + // Apply UI messages: callback or default (preserve all) + if (taskCompactionConfig?.compactUIMessages) { + accumulatedUIMessages = (await taskCompactionConfig.compactUIMessages( + compactEvent + )) as TUIMessage[]; + } + } + + // Determine if the user stopped generation this turn (not a full run cancel). + const wasStopped = stopController.signal.aborted && !runSignal.aborted; + + // Append the assistant's response (partial or complete) to the accumulator. + // The onFinish callback fires even on abort/stop, so partial responses + // from stopped generation are captured correctly. + let rawResponseMessage: TUIMessage | undefined; + if (capturedResponseMessage) { + // Keep the raw message before cleanup for users who want custom handling + rawResponseMessage = capturedResponseMessage; + // Clean up aborted parts (streaming tool calls, reasoning) when stopped + if (wasStopped) { + capturedResponseMessage = cleanupAbortedParts(capturedResponseMessage); + } + // Ensure the response message has an ID (the stream's onFinish + // may produce a message with an empty ID since IDs are normally + // assigned by the frontend's useChat). + if (!capturedResponseMessage.id) { + capturedResponseMessage = { ...capturedResponseMessage, id: generateMessageId() }; + } + // Append any non-transient data parts queued via chat.response or writer.write() + const queuedParts = locals.get(chatResponsePartsKey); + if (queuedParts && queuedParts.length > 0) { + capturedResponseMessage = { + ...capturedResponseMessage, + parts: [...capturedResponseMessage.parts, ...queuedParts], + } as TUIMessage; + locals.set(chatResponsePartsKey, []); + } + // Tool-approval continuations: the AI SDK reuses the trailing + // assistant's ID (via originalMessages) so the captured response + // carries the same ID as an existing message. Replace in place + // instead of pushing a duplicate. For action turns this never + // matches because originalMessages is omitted (fresh ID). + const existingIdx = capturedResponseMessage.id + ? accumulatedUIMessages.findIndex( + (m) => m.id === capturedResponseMessage!.id + ) + : -1; + if (existingIdx !== -1) { + accumulatedUIMessages[existingIdx] = capturedResponseMessage; + } else { + accumulatedUIMessages.push(capturedResponseMessage); + } + turnNewUIMessages.push(capturedResponseMessage); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + // Record toolCallId → head messageId so a HITL + // continuation next turn can recover the head id + // even if the AI SDK regenerates it. See + // `chatToolCallToMessageIdKey` for the full + // rationale (TRI-9137). + recordToolCallIdsFromMessage(capturedResponseMessage); + try { + const responseModelMessages = await toModelMessages([ + stripProviderMetadata(capturedResponseMessage), + ]); + if (existingIdx !== -1) { + // Reconvert all model messages since we replaced rather than appended + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } else { + accumulatedMessages.push(...responseModelMessages); + } + turnNewModelMessages.push(...responseModelMessages); + } catch { + // Conversion failed — skip accumulation for this turn + } + } + // If there's no captured response (manual pipe mode) but there are + // queued data parts, create a minimal response message to hold them. + if (!capturedResponseMessage) { + const remainingParts = locals.get(chatResponsePartsKey); + if (remainingParts && remainingParts.length > 0) { + capturedResponseMessage = { + id: generateMessageId(), + role: "assistant" as const, + parts: [...remainingParts], + } as TUIMessage; + locals.set(chatResponsePartsKey, []); + accumulatedUIMessages.push(capturedResponseMessage); + turnNewUIMessages.push(capturedResponseMessage); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + } + } + + if (runSignal.aborted) return "exit"; + + // Await deferred background work (e.g. DB writes from onTurnStart) + // before firing hooks so they can rely on the work being done. + const deferredWork = locals.get(chatDeferKey); + if (deferredWork && deferredWork.size > 0) { + await Promise.race([ + Promise.allSettled(deferredWork), + new Promise((r) => setTimeout(r, 5_000)), + ]); + } + + // Outer-loop compaction: runs between turns for single-step responses + // where prepareStep never fires (no tool calls = no step boundaries). + // Only triggers when: task has compaction configured, prepareStep didn't + // already compact this turn, and shouldCompact returns true. + const outerCompaction = locals.get(chatAgentCompactionKey); + const innerCompactionState = locals.get(chatCompactionStateKey); + + if (outerCompaction && !innerCompactionState && turnUsage && !wasStopped) { + const shouldTrigger = await outerCompaction.shouldCompact({ + messages: accumulatedMessages, + totalTokens: turnUsage.totalTokens, + inputTokens: turnUsage.inputTokens, + outputTokens: turnUsage.outputTokens, + usage: turnUsage, + totalUsage: cumulativeUsage, + chatId: currentWirePayload.chatId, + turn, + clientData, + source: "outer", + }); + + if (shouldTrigger) { + await tracer.startActiveSpan( + "context compaction (outer loop)", + async (compactionSpan) => { + const compactionId = generateMessageId(); + + const { waitUntilComplete } = chatStream.writer({ + spanName: "stream compaction chunks", + collapsed: true, + execute: async ({ write, merge }) => { + write({ + type: "data-compaction", + id: compactionId, + data: { status: "compacting", totalTokens: turnUsage.totalTokens }, + transient: true, + }); + + const summary = await outerCompaction.summarize({ + messages: accumulatedMessages, + usage: turnUsage, + totalUsage: cumulativeUsage, + chatId: currentWirePayload.chatId, + turn, + clientData, + source: "outer", + }); + + // Apply compactModelMessages/compactUIMessages callbacks, or defaults. + + const outerCompactEvent: CompactMessagesEvent = { + summary, + uiMessages: accumulatedUIMessages, + modelMessages: accumulatedMessages, + chatId: currentWirePayload.chatId, + turn, + clientData, + source: "outer", + }; + + // Model messages: callback or default (replace with summary) + accumulatedMessages = outerCompaction.compactModelMessages + ? await outerCompaction.compactModelMessages(outerCompactEvent) + : [ + { + role: "assistant" as const, + content: [ + { + type: "text" as const, + text: `[Conversation summary]\n\n${summary}`, + }, + ], + }, + ]; + + // UI messages: callback or default (preserve all) + if (outerCompaction.compactUIMessages) { + accumulatedUIMessages = (await outerCompaction.compactUIMessages( + outerCompactEvent + )) as TUIMessage[]; + } + + // Fire onCompacted hook + const onCompactedHook = locals.get(chatOnCompactedKey); + if (onCompactedHook) { + await onCompactedHook({ + ctx, + summary, + messages: accumulatedMessages, + messageCount: accumulatedMessages.length, + usage: turnUsage, + totalTokens: turnUsage.totalTokens, + inputTokens: turnUsage.inputTokens, + outputTokens: turnUsage.outputTokens, + stepNumber: -1, // outer loop, not a step + chatId: currentWirePayload.chatId, + turn, + writer: { write, merge }, + }); + } + + compactionSpan.setAttribute("compaction.summary_length", summary.length); + + write({ + type: "data-compaction", + id: compactionId, + data: { status: "complete", totalTokens: turnUsage.totalTokens }, + transient: true, + }); + }, + }); + await waitUntilComplete(); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "tabler-scissors", + "compaction.total_tokens": turnUsage.totalTokens ?? 0, + "compaction.input_tokens": turnUsage.inputTokens ?? 0, + "compaction.message_count": accumulatedMessages.length, + "compaction.outer_loop": true, + "compaction.turn": turn, + ...(currentWirePayload.chatId + ? { "compaction.chat_id": currentWirePayload.chatId } + : {}), + ...accessoryAttributes({ + items: [ + { text: `${turnUsage.totalTokens ?? 0} tokens`, variant: "normal" }, + { text: `${accumulatedMessages.length} msgs`, variant: "normal" }, + { text: "outer loop", variant: "normal" }, + ], + style: "codepath", + }), + }, + } + ); + } + } + + const turnCompleteEvent = { + ctx, + chatId: currentWirePayload.chatId, + messages: accumulatedMessages, + uiMessages: accumulatedUIMessages, + newMessages: turnNewModelMessages, + newUIMessages: turnNewUIMessages, + responseMessage: capturedResponseMessage, + rawResponseMessage, + turn, + runId: currentRunId, + chatAccessToken: turnAccessToken, + clientData, + stopped: wasStopped, + continuation, + previousRunId, + preloaded, + usage: turnUsage, + totalUsage: cumulativeUsage, + finishReason: capturedFinishReason, + }; + + // Fire onBeforeTurnComplete — stream is still open so the hook + // can write custom chunks to the frontend (e.g. compaction progress). + if (onBeforeTurnComplete) { + await tracer.startActiveSpan( + "onBeforeTurnComplete()", + async () => { + await withChatWriter(async (writer) => { + await onBeforeTurnComplete({ ...turnCompleteEvent, writer }); + }); + + // Check if the hook replaced messages (compaction or chat.history) + const override = locals.get(chatOverrideMessagesKey); + if (override) { + locals.set(chatOverrideMessagesKey, undefined); + accumulatedUIMessages = [...override] as TUIMessage[]; + accumulatedMessages = await toModelMessages(override); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + // Update event so onTurnComplete sees compacted messages + turnCompleteEvent.messages = accumulatedMessages; + turnCompleteEvent.uiMessages = accumulatedUIMessages; + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.turn": turn + 1, + }, + } + ); + } + + // Drain any late response parts added during onBeforeTurnComplete + const lateParts = locals.get(chatResponsePartsKey); + if (lateParts && lateParts.length > 0 && capturedResponseMessage) { + const idx = accumulatedUIMessages.findIndex((m) => m.id === capturedResponseMessage!.id); + if (idx !== -1) { + const msg = accumulatedUIMessages[idx]!; + accumulatedUIMessages[idx] = { + ...msg, + parts: [...(msg.parts ?? []), ...lateParts], + } as TUIMessage; + capturedResponseMessage = accumulatedUIMessages[idx] as TUIMessage; + turnCompleteEvent.responseMessage = capturedResponseMessage; + turnCompleteEvent.uiMessages = accumulatedUIMessages; + } + locals.set(chatResponsePartsKey, []); + } + + // Write turn-complete control chunk — closes the frontend stream. + const turnCompleteResult = await writeTurnCompleteChunk( + currentWirePayload.chatId, + turnAccessToken + ); + + // Fire onTurnComplete — stream is closed, use for persistence. + if (onTurnComplete) { + await tracer.startActiveSpan( + "onTurnComplete()", + async () => { + await onTurnComplete({ + ...turnCompleteEvent, + lastEventId: turnCompleteResult?.lastEventId, + }); + + // Check if onTurnComplete replaced messages (compaction or chat.history) + const turnCompleteOverride = locals.get(chatOverrideMessagesKey); + if (turnCompleteOverride) { + locals.set(chatOverrideMessagesKey, undefined); + accumulatedUIMessages = [...turnCompleteOverride] as TUIMessage[]; + accumulatedMessages = await toModelMessages(turnCompleteOverride); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.turn": turn + 1, + "chat.stopped": wasStopped, + "chat.continuation": continuation, + "chat.preloaded": preloaded, + ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}), + "chat.messages.count": accumulatedMessages.length, + "chat.response.parts.count": capturedResponseMessage?.parts?.length ?? 0, + "chat.new_messages.count": turnNewUIMessages.length, + ...(turnUsage?.inputTokens != null + ? { "gen_ai.usage.input_tokens": turnUsage.inputTokens } + : {}), + ...(turnUsage?.outputTokens != null + ? { "gen_ai.usage.output_tokens": turnUsage.outputTokens } + : {}), + ...(turnUsage?.totalTokens != null + ? { "gen_ai.usage.total_tokens": turnUsage.totalTokens } + : {}), + ...(cumulativeUsage.totalTokens != null + ? { "gen_ai.usage.cumulative_total_tokens": cumulativeUsage.totalTokens } + : {}), + }, + } + ); + } + + // ── Snapshot write ───────────────────────────────────── + // + // Persist the post-turn accumulator so the next run boot + // can replay history without the wire shipping it. Skipped + // when `hydrateMessages` is registered — those customers + // own persistence and would never read this blob. + // + // AWAITED, not fire-and-forget: the agent may suspend (idle + // timeout) immediately after this point, and in-flight + // promises don't reliably complete on suspend. The user- + // visible turn already finished (the turn-complete chunk + // closed the response stream above), so the await delay + // only affects "when can the NEXT turn start," gated by + // user typing — not TTFC. + // + // `writeChatSnapshot` swallows errors internally; this + // outer try/catch is just belt-and-suspenders against + // tracer/span failures. + if (!hydrateMessages) { + try { + await tracer.startActiveSpan( + "snapshot.write", + async () => { + await writeChatSnapshot(sessionIdForSnapshot, { + version: 1, + savedAt: Date.now(), + messages: accumulatedUIMessages, + // `StreamWriteResult` exposes `lastEventId` only; + // use the snapshot save time as the + // `lastOutTimestamp` cutoff hint. The OOM-retry + // optimization compares this to SSE chunk + // timestamps (ms epoch on the server) — Date.now() + // here is the closest cheap approximation + // available client-side and is consistent with + // the existing turn-complete chunk emission. + lastOutEventId: turnCompleteResult?.lastEventId, + lastOutTimestamp: Date.now(), + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.turn": turn + 1, + "chat.messages.count": accumulatedUIMessages.length, + }, + } + ); + } catch (error) { + logger.warn( + "chat.agent: snapshot write failed; next run will replay further", + { + error: error instanceof Error ? error.message : String(error), + sessionId: sessionIdForSnapshot, + } + ); + } + } + + } // end if (!isAction) + + // NOTE: We intentionally do NOT await deferred work from onTurnComplete here. + // Promises deferred in onTurnComplete (e.g. background self-review via + // chat.defer + chat.inject) run during the idle wait. If they complete + // before the next message, their injected context is picked up in prepareStep. + // The pre-onBeforeTurnComplete drain handles promises from onTurnStart/run(). + + // If messages arrived during streaming (without pendingMessages config), + // use the first one immediately as the next turn. + if (pendingMessages.length > 0) { + currentWirePayload = pendingMessages[0]!; + return "continue"; + } + + // chat.requestUpgrade() was called — exit the loop so the + // transport triggers a new run on the latest version. + // chat.endRun() — same exit, no upgrade semantics. + if ( + locals.get(chatUpgradeRequestedKey) || + locals.get(chatEndRunRequestedKey) + ) { + return "exit"; + } + + // Wait for the next message — stay idle briefly, then suspend + const effectiveIdleTimeout = + (metadata.get(IDLE_TIMEOUT_METADATA_KEY) as number | undefined) ?? + idleTimeoutInSeconds; + const effectiveTurnTimeout = + (metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined) ?? turnTimeout; + + const next = await messagesInput.waitWithIdleTimeout({ + idleTimeoutInSeconds: effectiveIdleTimeout, + timeout: effectiveTurnTimeout, + spanName: "waiting for next message", + onSuspend: onChatSuspend + ? async () => { + await tracer.startActiveSpan( + "onChatSuspend()", + async () => { + await onChatSuspend({ + phase: "turn", + ctx, + chatId: currentWirePayload.chatId, + runId: ctx.run.id, + turn, + messages: accumulatedMessages, + uiMessages: accumulatedUIMessages, + clientData, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.suspend.phase": "turn", + "chat.turn": turn + 1, + }, + } + ); + } + : undefined, + onResume: onChatResume + ? async () => { + await tracer.startActiveSpan( + "onChatResume()", + async () => { + await onChatResume({ + phase: "turn", + ctx, + chatId: currentWirePayload.chatId, + runId: ctx.run.id, + turn, + messages: accumulatedMessages, + uiMessages: accumulatedUIMessages, + clientData, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": currentWirePayload.chatId, + "chat.resume.phase": "turn", + "chat.turn": turn + 1, + }, + } + ); + } + : undefined, + }); + + if (!next.ok) { + return "exit"; + } + + currentWirePayload = next.output as ChatTaskWirePayload< + TUIMessage, + inferSchemaIn + >; + + // Close signal — exit the loop gracefully + if (currentWirePayload.trigger === "close") { + return "exit"; + } + + return "continue"; + }, + { + attributes: turnAttributes, + } + ); + + if (turnResult === "exit") return; + // "continue" means proceed to next iteration + } catch (turnError) { + // Turn error handler: write an error chunk + turn-complete to the stream + // so the client sees the error, then wait for the next message instead + // of killing the entire run. This keeps the conversation alive. + if (turnError instanceof Error && turnError.name === "AbortError" && runSignal.aborted) { + // Full run cancellation — exit immediately + throw turnError; + } + + // OOM errors must escape the turn loop so the task runtime can + // honor `retry.outOfMemory.machine` (set on chat.agent via + // `oomMachine`). Catching them here would keep the dead worker + // alive and defeat the machine swap. Re-throw and let the + // runtime dispatch the retry on a larger machine; recovery on + // attempt 2 picks up via the standard continuation path + // (same chatId / Session, accumulator rehydrates). + if (turnError instanceof OutOfMemoryError) { + throw turnError; + } + + try { + await withChatWriter(async (writer) => { + const errorText = + turnError instanceof Error ? turnError.message : "An unexpected error occurred"; + writer.write({ type: "error", errorText } as any); + }); + // Signal turn complete so the client knows this turn is done + await writeTurnCompleteChunk(currentWirePayload.chatId); + } catch { + // Best-effort — if stream write fails, let the run continue anyway + } + + // chat.requestUpgrade() / chat.endRun() — exit after error turn too + if ( + locals.get(chatUpgradeRequestedKey) || + locals.get(chatEndRunRequestedKey) + ) { + return; + } + + // Wait for the next message — same as after a successful turn + const effectiveIdleTimeout = + (metadata.get(IDLE_TIMEOUT_METADATA_KEY) as number | undefined) ?? + idleTimeoutInSeconds; + const effectiveTurnTimeout = + (metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined) ?? turnTimeout; + + const next = await messagesInput.waitWithIdleTimeout({ + idleTimeoutInSeconds: effectiveIdleTimeout, + timeout: effectiveTurnTimeout, + spanName: "waiting for next message (after error)", + }); + + if (!next.ok) { + return; // Timed out — end run gracefully + } + + currentWirePayload = next.output as ChatTaskWirePayload< + TUIMessage, + inferSchemaIn + >; + // Continue to next iteration of the for loop + } + } + } finally { + // `stopSub` is registered post-preload so the close-during-preload + // early-return path may exit before it ever attached. Guard the + // cleanup so a missing subscription doesn't throw. + stopSub?.off(); + } + } + }); + + // Register clientDataSchema so the CLI converts it to JSONSchema + // and stores it as payloadSchema — used by the Playground UI + if (clientDataSchema) { + resourceCatalog.updateTaskMetadata(options.id, { + schema: clientDataSchema as any, + }); + } + + return task; +} + +/** + * Optional config for {@link chat.withUIMessage}. `streamOptions` become default + * static `toUIMessageStream()` settings; inner `chat.agent({ uiMessageStreamOptions })` + * shallow-merges on top (task wins on conflicts). + */ +export type ChatWithUIMessageConfig = { + streamOptions?: ChatUIMessageStreamOptions; +}; + +// --------------------------------------------------------------------------- +// Chat builder +// --------------------------------------------------------------------------- + +/** + * A chainable builder for configuring chat tasks with fixed UI message types, + * client data schemas, and builder-level hooks that compose with task-level hooks. + * + * Obtain a builder via {@link chat.withUIMessage} or {@link chat.withClientData}. + * + * @example + * ```ts + * export const myChat = chat + * .withUIMessage({ streamOptions: { sendReasoning: true } }) + * .withClientData({ schema: z.object({ userId: z.string() }) }) + * .onChatSuspend(async ({ ctx }) => { await disposeResources(ctx.run.id) }) + * .task({ + * id: "my-chat", + * run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + * }); + * ``` + */ +export interface ChatBuilder< + TUIMessage extends UIMessage = UIMessage, + TClientDataSchema extends TaskSchema | undefined = undefined, +> { + /** Fix the UI message type. Returns a new builder preserving all accumulated state. */ + withUIMessage( + config?: ChatWithUIMessageConfig + ): ChatBuilder; + + /** Fix the client data schema. Returns a new builder preserving all accumulated state. */ + withClientData(config: { + schema: TSchema; + }): ChatBuilder; + + /** Register a builder-level `onPreload` hook. Runs before the task-level hook if both are set. */ + onPreload( + fn: (event: PreloadEvent>) => Promise | void + ): ChatBuilder; + + /** Register a builder-level `onChatStart` hook. Runs before the task-level hook if both are set. */ + onChatStart( + fn: (event: ChatStartEvent>) => Promise | void + ): ChatBuilder; + + /** Register a builder-level `onTurnStart` hook. Runs before the task-level hook if both are set. */ + onTurnStart( + fn: ( + event: TurnStartEvent, TUIMessage> + ) => Promise | void + ): ChatBuilder; + + /** Register a builder-level `onBeforeTurnComplete` hook. Runs before the task-level hook if both are set. */ + onBeforeTurnComplete( + fn: ( + event: BeforeTurnCompleteEvent, TUIMessage> + ) => Promise | void + ): ChatBuilder; + + /** Register a builder-level `onTurnComplete` hook. Runs before the task-level hook if both are set. */ + onTurnComplete( + fn: ( + event: TurnCompleteEvent, TUIMessage> + ) => Promise | void + ): ChatBuilder; + + /** Register a builder-level `onCompacted` hook. Runs before the task-level hook if both are set. */ + onCompacted(fn: (event: CompactedEvent) => Promise | void): ChatBuilder; + + /** Register a builder-level `onChatSuspend` hook. Runs before the task-level hook if both are set. */ + onChatSuspend( + fn: ( + event: ChatSuspendEvent, TUIMessage> + ) => Promise | void + ): ChatBuilder; + + /** Register a builder-level `onChatResume` hook. Runs before the task-level hook if both are set. */ + onChatResume( + fn: ( + event: ChatResumeEvent, TUIMessage> + ) => Promise | void + ): ChatBuilder; + + /** + * Create the chat agent with the accumulated builder configuration. + * + * When `withClientData` was called, `clientDataSchema` is injected automatically + * and omitted from options. Otherwise, it can still be set directly in options + * (backwards compatible). + */ + agent: [TClientDataSchema] extends [undefined] + ? ( + options: ChatAgentOptions + ) => Task>, unknown> + : ( + options: Omit, "clientDataSchema"> + ) => Task>, unknown>; + + /** + * Create a custom agent with manual lifecycle control. + * + * The agent appears in the playground but you manage the turn loop, + * message waiting, and streaming yourself using composable primitives + * (`chat.messages`, `chat.MessageAccumulator`, `chat.pipeAndCapture`, etc.). + * + * Builder hooks (`onPreload`, `onChatStart`, etc.) are not applied — + * those are managed-lifecycle concepts handled by `.agent()`. + */ + customAgent: [TClientDataSchema] extends [undefined] + ? ( + options: ChatCustomAgentOptions + ) => Task, unknown> + : ( + options: ChatCustomAgentOptions + ) => Task>, unknown>; +} + +/** @internal */ +type ChatBuilderHooks = { + onPreload?: (event: any) => Promise | void; + onChatStart?: (event: any) => Promise | void; + onTurnStart?: (event: any) => Promise | void; + onBeforeTurnComplete?: (event: any) => Promise | void; + onTurnComplete?: (event: any) => Promise | void; + onCompacted?: (event: any) => Promise | void; + onChatSuspend?: (event: any) => Promise | void; + onChatResume?: (event: any) => Promise | void; +}; + +/** @internal */ +type ChatBuilderConfig = { + uiStreamOptions?: ChatUIMessageStreamOptions; + clientDataSchema?: TaskSchema; + hooks: ChatBuilderHooks; +}; + +function composeHooks( + builderHook: ((event: T) => Promise | void) | undefined, + taskHook: ((event: T) => Promise | void) | undefined +): ((event: T) => Promise) | undefined { + if (!builderHook) return taskHook as any; + if (!taskHook) return builderHook as any; + return async (event: T) => { + await builderHook(event); + await taskHook(event); + }; +} + +function createChatBuilder< + TUIMessage extends UIMessage = UIMessage, + TClientDataSchema extends TaskSchema | undefined = undefined, +>(config: ChatBuilderConfig): ChatBuilder { + return { + withUIMessage(uimConfig?: ChatWithUIMessageConfig) { + return createChatBuilder({ + ...config, + uiStreamOptions: uimConfig?.streamOptions ?? config.uiStreamOptions, + }); + }, + + withClientData(cdConfig: { schema: TSchema }) { + return createChatBuilder({ + ...config, + clientDataSchema: cdConfig.schema, + }); + }, + + onPreload( + fn: (event: PreloadEvent>) => Promise | void + ) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onPreload: fn }, + }); + }, + onChatStart( + fn: (event: ChatStartEvent>) => Promise | void + ) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onChatStart: fn }, + }); + }, + onTurnStart( + fn: ( + event: TurnStartEvent, TUIMessage> + ) => Promise | void + ) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onTurnStart: fn }, + }); + }, + onBeforeTurnComplete( + fn: ( + event: BeforeTurnCompleteEvent, TUIMessage> + ) => Promise | void + ) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onBeforeTurnComplete: fn }, + }); + }, + onTurnComplete( + fn: ( + event: TurnCompleteEvent, TUIMessage> + ) => Promise | void + ) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onTurnComplete: fn }, + }); + }, + onCompacted(fn: (event: CompactedEvent) => Promise | void) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onCompacted: fn }, + }); + }, + onChatSuspend( + fn: ( + event: ChatSuspendEvent, TUIMessage> + ) => Promise | void + ) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onChatSuspend: fn }, + }); + }, + onChatResume( + fn: ( + event: ChatResumeEvent, TUIMessage> + ) => Promise | void + ) { + return createChatBuilder({ + ...config, + hooks: { ...config.hooks, onChatResume: fn }, + }); + }, + + agent(options: any) { + const mergedUiStream = + config.uiStreamOptions && options.uiMessageStreamOptions + ? { ...config.uiStreamOptions, ...options.uiMessageStreamOptions } + : options.uiMessageStreamOptions ?? config.uiStreamOptions; + + return chatAgent({ + ...options, + ...(config.clientDataSchema ? { clientDataSchema: config.clientDataSchema } : {}), + uiMessageStreamOptions: mergedUiStream, + onPreload: composeHooks(config.hooks.onPreload, options.onPreload), + onChatStart: composeHooks(config.hooks.onChatStart, options.onChatStart), + onTurnStart: composeHooks(config.hooks.onTurnStart, options.onTurnStart), + onBeforeTurnComplete: composeHooks( + config.hooks.onBeforeTurnComplete, + options.onBeforeTurnComplete + ), + onTurnComplete: composeHooks(config.hooks.onTurnComplete, options.onTurnComplete), + onCompacted: composeHooks(config.hooks.onCompacted, options.onCompacted), + onChatSuspend: composeHooks(config.hooks.onChatSuspend, options.onChatSuspend), + onChatResume: composeHooks(config.hooks.onChatResume, options.onChatResume), + }); + }, + + customAgent(options: any) { + return chatCustomAgent({ + ...options, + ...(config.clientDataSchema ? { clientDataSchema: config.clientDataSchema } : {}), + }); + }, + } as unknown as ChatBuilder; +} + +/** + * Fix the UI message type for a chat task (AI SDK `UIMessage` generics) while + * keeping `id` and `clientDataSchema` inference on the inner {@link chat.agent} call. + * + * Returns a {@link ChatBuilder} that supports chaining `.withClientData()`, + * hook methods (`.onPreload()`, `.onChatSuspend()`, etc.), and `.task()`. + * + * @example + * ```ts + * type AgentUiMessage = UIMessage; + * + * export const myChat = chat.withUIMessage({ + * streamOptions: { sendReasoning: true }, + * }).task({ + * id: "my-chat", + * run: async ({ messages, signal }) => { ... }, + * }); + * ``` + */ +function withUIMessage( + config?: ChatWithUIMessageConfig +): ChatBuilder { + return createChatBuilder({ + uiStreamOptions: config?.streamOptions, + hooks: {}, + }); +} + +/** + * Fix the client data schema for a chat task, providing typed `clientData` + * in all hooks and the `run` function. + * + * Returns a {@link ChatBuilder} that supports chaining `.withUIMessage()`, + * hook methods (`.onPreload()`, `.onChatSuspend()`, etc.), and `.task()`. + * + * @example + * ```ts + * export const myChat = chat + * .withClientData({ schema: z.object({ userId: z.string() }) }) + * .task({ + * id: "my-chat", + * onPreload: async ({ clientData }) => { + * // clientData is typed as { userId: string } + * }, + * run: async ({ messages, signal }) => { ... }, + * }); + * ``` + */ +function withClientData(config: { + schema: TSchema; +}): ChatBuilder { + return createChatBuilder({ + clientDataSchema: config.schema, + hooks: {}, + }); +} + +/** + * Namespace for AI SDK chat integration. + * + * @example + * ```ts + * import { chat } from "@trigger.dev/sdk/ai"; + * + * // Define a chat task + * export const myChat = chat.agent({ + * id: "my-chat", + * run: async ({ messages, signal }) => { + * return streamText({ model, messages, abortSignal: signal }); + * }, + * }); + * + * // Pipe a stream manually (from inside a task) + * await chat.pipe(streamTextResult); + * + * // Create an access token (from a server action) + * const token = await chat.createAccessToken("my-chat"); + * ``` + */ +// --------------------------------------------------------------------------- +// Runtime configuration helpers +// --------------------------------------------------------------------------- + +const TURN_TIMEOUT_METADATA_KEY = "chat.turnTimeout"; +const IDLE_TIMEOUT_METADATA_KEY = "chat.idleTimeout"; + +/** + * Override the turn timeout for subsequent turns in the current run. + * + * The turn timeout controls how long the run stays suspended (freeing compute) + * waiting for the next user message. When it expires, the run completes + * gracefully and the next message starts a fresh run. + * + * Call from inside a `chatAgent` run function to adjust based on context. + * + * @param duration - A duration string (e.g. `"5m"`, `"1h"`, `"30s"`) + * + * @example + * ```ts + * run: async ({ messages, signal }) => { + * chat.setTurnTimeout("2h"); + * return streamText({ model, messages, abortSignal: signal }); + * } + * ``` + */ +function setTurnTimeout(duration: string): void { + metadata.set(TURN_TIMEOUT_METADATA_KEY, duration); +} + +/** + * Override the turn timeout in seconds for subsequent turns in the current run. + * + * @param seconds - Number of seconds to wait for the next message before ending the run + * + * @example + * ```ts + * run: async ({ messages, signal }) => { + * chat.setTurnTimeoutInSeconds(3600); // 1 hour + * return streamText({ model, messages, abortSignal: signal }); + * } + * ``` + */ +function setTurnTimeoutInSeconds(seconds: number): void { + metadata.set(TURN_TIMEOUT_METADATA_KEY, `${seconds}s`); +} + +/** + * Override the idle timeout for subsequent turns in the current run. + * + * The idle timeout controls how long the run stays active (using compute) + * after each turn, waiting for the next message. During this window, + * responses are instant. After it expires, the run suspends. + * + * @param seconds - Number of seconds to stay idle (0 to suspend immediately) + * + * @example + * ```ts + * run: async ({ messages, signal }) => { + * chat.setIdleTimeoutInSeconds(60); + * return streamText({ model, messages, abortSignal: signal }); + * } + * ``` + */ +function setIdleTimeoutInSeconds(seconds: number): void { + metadata.set(IDLE_TIMEOUT_METADATA_KEY, seconds); +} + +/** + * Override the `toUIMessageStream()` options for the current turn. + * + * These options control how the `StreamTextResult` is converted to a + * `UIMessageChunk` stream — error handling, reasoning/source visibility, + * message metadata, etc. + * + * Per-turn options are merged on top of the static `uiMessageStreamOptions` + * set on `chat.agent()`. Per-turn values win on conflicts. + * + * @example + * ```ts + * run: async ({ messages, signal }) => { + * chat.setUIMessageStreamOptions({ + * sendReasoning: true, + * onError: (error) => error instanceof Error ? error.message : "An error occurred.", + * }); + * return streamText({ model, messages, abortSignal: signal }); + * } + * ``` + */ +function setUIMessageStreamOptions(options: ChatUIMessageStreamOptions): void { + locals.set(chatUIStreamPerTurnKey, options); +} + +/** + * Resolve the effective UIMessageStream options by merging: + * 1. Static task-level options (from `chat.agent({ uiMessageStreamOptions })`) + * 2. Per-turn overrides (from `chat.setUIMessageStreamOptions()`) + * + * Per-turn values win on conflicts. Clears the per-turn override after reading + * so it doesn't leak into subsequent turns. + * @internal + */ +function resolveUIMessageStreamOptions(): ChatUIMessageStreamOptions { + const staticOptions = locals.get(chatUIStreamStaticKey) ?? {}; + const perTurnOptions = locals.get(chatUIStreamPerTurnKey) ?? {}; + // Clear per-turn override so it doesn't leak into subsequent turns + locals.set(chatUIStreamPerTurnKey, undefined); + return { ...staticOptions, ...perTurnOptions }; +} + +// --------------------------------------------------------------------------- +// Stop detection +// --------------------------------------------------------------------------- + +/** + * Check whether the user stopped generation during the current turn. + * + * Works from **anywhere** inside a `chat.agent` run — including inside + * `streamText`'s `onFinish` callback — without needing to thread the + * `stopSignal` through closures. + * + * This is especially useful when the AI SDK's `isAborted` flag is unreliable + * (e.g. when using `createUIMessageStream` + `writer.merge()`). + * + * @example + * ```ts + * onFinish: ({ isAborted }) => { + * const wasStopped = isAborted || chat.isStopped(); + * if (wasStopped) { + * // handle stop + * } + * } + * ``` + */ +function isStopped(): boolean { + const controller = locals.get(chatStopControllerKey); + return controller?.signal.aborted ?? false; +} + +// --------------------------------------------------------------------------- +// Version upgrade +// --------------------------------------------------------------------------- + +/** + * Request that the current run exits so the next message starts on the latest + * deployed version (via the standard continuation mechanism). + * + * When called from `onTurnStart` or `onValidateMessages`, `run()` is skipped + * entirely — the run exits immediately and the transport re-triggers the + * same message on the new version. + * + * When called from `run()` or `chat.defer()`, the current turn completes + * normally and the run exits afterward instead of waiting for the next message. + * + * Call from `onTurnStart`, `onValidateMessages`, `onChatResume`, `run()`, + * or inside `chat.defer()`. + * + * @example + * ```ts + * const SUPPORTED_VERSIONS = new Set(["v2", "v3"]); + * + * chat.agent({ + * id: "my-chat", + * onTurnStart: async ({ clientData }) => { + * if (clientData?.protocolVersion && !SUPPORTED_VERSIONS.has(clientData.protocolVersion)) { + * chat.requestUpgrade(); + * } + * }, + * run: async ({ messages }) => { ... }, + * }); + * ``` + */ +function requestUpgrade(): void { + locals.set(chatUpgradeRequestedKey, true); +} + +/** + * Exit the run after the current turn completes, without waiting for the + * next message. Unlike {@link requestUpgrade}, no upgrade-required signal + * is sent to the client — the turn finishes normally, `onTurnComplete` + * fires, and the loop exits instead of going idle. + * + * Call from `run()`, `chat.defer()`, `onBeforeTurnComplete`, or + * `onTurnComplete` to end the run on your own terms (budget exhausted, + * task complete, goal achieved, etc.). + * + * The next user message on the same `chatId` starts a fresh run via the + * normal continuation mechanism. + * + * @example + * ```ts + * chat.agent({ + * id: "one-shot-agent", + * run: async ({ messages, signal }) => { + * const result = streamText({ model: openai("gpt-4o"), messages, abortSignal: signal }); + * // Single-response agent — exit after this turn. + * chat.endRun(); + * return result; + * }, + * }); + * ``` + */ +function endRun(): void { + locals.set(chatEndRunRequestedKey, true); +} + +// --------------------------------------------------------------------------- +// Per-turn deferred work +// --------------------------------------------------------------------------- + +/** + * Register a promise that runs in the background during the current turn. + * + * Use this to move non-blocking work (DB writes, analytics, etc.) out of + * the critical path. The promise runs in parallel with streaming and is + * awaited (with a 5 s timeout) before `onTurnComplete` fires. + * + * @example + * ```ts + * onTurnStart: async ({ chatId, uiMessages }) => { + * // Pass a promise directly + * chat.defer(db.chat.update({ where: { id: chatId }, data: { messages: uiMessages } })); + * + * // Or pass an async function — cleaner for multi-step work + * chat.defer(async () => { + * const flags = await getFeatureFlags(); + * if (flags.forceUpgrade) chat.requestUpgrade(); + * }); + * }, + * ``` + */ +function chatDefer(promiseOrFn: Promise | (() => Promise)): void { + const promises = locals.get(chatDeferKey); + if (promises) { + promises.add(typeof promiseOrFn === "function" ? promiseOrFn() : promiseOrFn); + } +} + +// --------------------------------------------------------------------------- +// Background context injection +// --------------------------------------------------------------------------- + +/** + * Queue model messages for injection at the next `prepareStep` boundary. + * + * Use this to inject context from background work into the agent's conversation. + * Messages are appended to the model messages before the next LLM inference call. + * + * Combine with `chat.defer()` to run background analysis and inject results: + * + * @example + * ```ts + * onTurnComplete: async ({ messages }) => { + * chat.defer((async () => { + * const review = await generateObject({ + * model: openai("gpt-4o-mini"), + * messages: [...messages, { role: "user", content: "Review the last response." }], + * schema: z.object({ suggestions: z.array(z.string()) }), + * }); + * if (review.object.suggestions.length > 0) { + * chat.inject([{ + * role: "system", + * content: `Improvements for next response:\n${review.object.suggestions.join("\n")}`, + * }]); + * } + * })()); + * }, + * ``` + */ +function injectBackgroundContext(messages: ModelMessage[]): void { + const queue = locals.get(chatBackgroundQueueKey) ?? []; + queue.push(...messages); + locals.set(chatBackgroundQueueKey, queue); +} + +// --------------------------------------------------------------------------- +// Aborted message cleanup +// --------------------------------------------------------------------------- + +/** + * Clean up a UIMessage that was captured during an aborted/stopped turn. + * + * When generation is stopped mid-stream, the captured message may contain: + * - Tool parts stuck in incomplete states (`partial-call`, `input-available`, + * `input-streaming`) that cause permanent UI spinners + * - Reasoning parts with `state: "streaming"` instead of `"done"` + * - Text parts with `state: "streaming"` instead of `"done"` + * + * This function returns a cleaned copy with: + * - Incomplete tool parts removed entirely + * - Reasoning and text parts marked as `"done"` + * + * `chat.agent` calls this automatically when stop is detected before passing + * the response to `onTurnComplete`. Use this manually when calling `pipeChat` + * directly and capturing response messages yourself. + * + * @example + * ```ts + * onTurnComplete: async ({ responseMessage, stopped }) => { + * // Already cleaned automatically by chat.agent — but if you captured + * // your own message via pipeChat, clean it manually: + * const cleaned = chat.cleanupAbortedParts(myMessage); + * await db.messages.save(cleaned); + * } + * ``` + */ +function cleanupAbortedParts(message: TUIM): TUIM { + if (!message.parts) return message; + + const isToolPart = (part: any) => + part.type === "tool-invocation" || + part.type?.startsWith("tool-") || + part.type === "dynamic-tool"; + + return { + ...message, + parts: message.parts + .filter((part: any) => { + if (!isToolPart(part)) return true; + // Remove tool parts that never completed execution. + // partial-call: input was still streaming when aborted. + // input-available: input was complete but tool never ran. + // input-streaming: input was mid-stream. + const state = part.toolInvocation?.state ?? part.state; + return ( + state !== "partial-call" && state !== "input-available" && state !== "input-streaming" + ); + }) + .map((part: any) => { + // Mark streaming reasoning as done + if (part.type === "reasoning" && part.state === "streaming") { + return { ...part, state: "done" }; + } + // Mark streaming text as done + if (part.type === "text" && part.state === "streaming") { + return { ...part, state: "done" }; + } + return part; + }), + } as TUIM; +} + +// --------------------------------------------------------------------------- +// Composable primitives for raw task chat +// --------------------------------------------------------------------------- + +/** + * Create a managed stop signal wired to the chat stop input stream. + * + * Call once at the start of your run. Use `signal` as the abort signal for + * `streamText`. Call `reset()` at the start of each turn to get a fresh + * per-turn signal. Call `cleanup()` when the run ends. + * + * @example + * ```ts + * const stop = chat.createStopSignal(); + * for (let turn = 0; turn < 100; turn++) { + * stop.reset(); + * const result = streamText({ model, messages, abortSignal: stop.signal }); + * await chat.pipe(result); + * // ... + * } + * stop.cleanup(); + * ``` + */ +function createStopSignal(): { + readonly signal: AbortSignal; + reset: () => void; + cleanup: () => void; +} { + let controller = new AbortController(); + const sub = stopInput.on((data) => { + controller.abort(data?.message || "stopped"); + }); + return { + get signal() { + return controller.signal; + }, + reset() { + controller = new AbortController(); + }, + cleanup() { + sub.off(); + }, + }; +} + +/** + * Signal the frontend that the current turn is complete. + * + * The `TriggerChatTransport` intercepts this to close the ReadableStream + * for the current turn. Call after piping the response stream. + * + * @example + * ```ts + * await chat.pipe(result); + * await chat.writeTurnComplete(); + * ``` + */ +async function chatWriteTurnComplete(options?: { publicAccessToken?: string }): Promise { + await writeTurnCompleteChunk(undefined, options?.publicAccessToken); +} + +/** + * Pipe a `StreamTextResult` (or similar) to the chat stream and capture + * the assistant's response message via `onFinish`. + * + * Combines `toUIMessageStream()` + `onFinish` callback + `chat.pipe()`. + * Returns the captured `UIMessage`, or `undefined` if capture failed. + * + * @example + * ```ts + * const result = streamText({ model, messages, abortSignal: signal }); + * const response = await chat.pipeAndCapture(result, { signal }); + * if (response) conversation.addResponse(response); + * ``` + */ +async function pipeChatAndCapture( + source: UIMessageStreamable, + options?: { signal?: AbortSignal; spanName?: string } +): Promise { + let captured: UIMessage | undefined; + let resolveOnFinish: () => void; + const onFinishPromise = new Promise((r) => { + resolveOnFinish = r; + }); + + const uiStream = source.toUIMessageStream({ + ...resolveUIMessageStreamOptions(), + onFinish: ({ responseMessage }: { responseMessage: UIMessage }) => { + captured = responseMessage; + resolveOnFinish!(); + }, + }); + + await pipeChat(uiStream, { + signal: options?.signal, + spanName: options?.spanName ?? "stream response", + }); + await onFinishPromise; + + return captured; +} + +/** + * Accumulates conversation messages across turns. + * + * Handles the transport protocol: turn 0 sends full history (replace), + * subsequent turns send only new messages (append), regenerate sends + * full history minus last assistant message (replace). + * + * @example + * ```ts + * const conversation = new chat.MessageAccumulator(); + * for (let turn = 0; turn < 100; turn++) { + * const messages = await conversation.addIncoming(payload.messages, payload.trigger, turn); + * const result = streamText({ model, messages }); + * const response = await chat.pipeAndCapture(result); + * if (response) await conversation.addResponse(response); + * } + * ``` + */ +class ChatMessageAccumulator { + modelMessages: ModelMessage[] = []; + uiMessages: UIMessage[] = []; + private _compaction?: ChatAgentCompactionOptions; + private _pendingMessages?: PendingMessagesOptions; + private _steeringQueue: SteeringQueueEntry[] = []; + + constructor(options?: { + compaction?: ChatAgentCompactionOptions; + pendingMessages?: PendingMessagesOptions; + }) { + this._compaction = options?.compaction; + this._pendingMessages = options?.pendingMessages; + } + + /** + * Add incoming messages from the transport payload. + * Returns the full accumulated model messages for `streamText`. + */ + async addIncoming(messages: UIMessage[], trigger: string, turn: number): Promise { + const cleaned = messages.map((m) => (m.role === "assistant" ? cleanupAbortedParts(m) : m)); + const model = await toModelMessages(cleaned); + + if (turn === 0 || trigger === "regenerate-message") { + this.modelMessages = model; + this.uiMessages = [...cleaned]; + } else { + this.modelMessages.push(...model); + this.uiMessages.push(...cleaned); + } + return this.modelMessages; + } + + /** + * Add the assistant's response to the accumulator. + * Call after `pipeAndCapture` with the captured response. + */ + /** + * Replace all accumulated messages (for compaction). + * Converts UIMessages to ModelMessages internally. + */ + async setMessages(uiMessages: UIMessage[]): Promise { + this.uiMessages = [...uiMessages]; + this.modelMessages = await toModelMessages(uiMessages); + } + + async addResponse(response: UIMessage): Promise { + if (!response.id) { + response = { ...response, id: generateMessageId() }; + } + this.uiMessages.push(response); + try { + const msgs = await toModelMessages([stripProviderMetadata(response)]); + this.modelMessages.push(...msgs); + } catch { + // Conversion failed — skip model message accumulation for this response + } + } + + /** + * Queue a message for injection via `prepareStep`. Call from a + * `messagesInput.on()` listener when a message arrives during streaming. + */ + steer(message: UIMessage, modelMessages?: ModelMessage[]): void { + if (modelMessages) { + this._steeringQueue.push({ uiMessage: message, modelMessages }); + } else { + // Defer conversion — will be done in prepareStep if needed + this._steeringQueue.push({ uiMessage: message, modelMessages: [] }); + } + } + + /** + * Queue a message for injection, converting to model messages automatically. + */ + async steerAsync(message: UIMessage): Promise { + const modelMsgs = await toModelMessages([message]); + this._steeringQueue.push({ uiMessage: message, modelMessages: modelMsgs }); + } + + /** + * Get and clear unconsumed steering messages. + */ + drainSteering(): UIMessage[] { + const result = this._steeringQueue.map((e) => e.uiMessage); + this._steeringQueue = []; + return result; + } + + /** + * Returns a `prepareStep` function that handles both compaction and + * pending message injection. Pass to `streamText({ prepareStep: conversation.prepareStep() })`. + */ + prepareStep(): + | ((args: { + messages: ModelMessage[]; + steps: CompactionStep[]; + }) => Promise<{ messages: ModelMessage[] } | undefined>) + | undefined { + if (!this._compaction && !this._pendingMessages) return undefined; + const comp = this._compaction; + const pm = this._pendingMessages; + const queue = this._steeringQueue; + + return async ({ messages, steps }) => { + let resultMessages: ModelMessage[] | undefined; + + // 1. Compaction + if (comp) { + const result = await chatCompact(messages, steps, { + shouldCompact: comp.shouldCompact, + summarize: (msgs) => comp.summarize({ messages: msgs, source: "inner" }), + }); + if (result.type !== "skipped") { + resultMessages = result.messages; + } + } + + // 2. Pending message injection + if (pm && queue.length > 0) { + const injected = await drainSteeringQueue(pm, resultMessages ?? messages, steps, queue); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? messages), ...injected]; + } + } + + return resultMessages ? { messages: resultMessages } : undefined; + }; + } + + /** + * Run outer-loop compaction if needed. Call after adding the response + * and capturing usage. Applies `compactModelMessages` and `compactUIMessages` + * callbacks if configured. + * + * @returns `true` if compaction was performed, `false` otherwise. + */ + async compactIfNeeded( + usage: LanguageModelUsage | undefined, + context?: { + chatId?: string; + turn?: number; + clientData?: unknown; + totalUsage?: LanguageModelUsage; + } + ): Promise { + if (!this._compaction || !usage) return false; + + const shouldTrigger = await this._compaction.shouldCompact({ + messages: this.modelMessages, + totalTokens: usage.totalTokens, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + usage, + totalUsage: context?.totalUsage, + chatId: context?.chatId, + turn: context?.turn, + clientData: context?.clientData, + source: "outer", + }); + + if (!shouldTrigger) return false; + + const summary = await this._compaction.summarize({ + messages: this.modelMessages, + usage, + totalUsage: context?.totalUsage, + chatId: context?.chatId, + turn: context?.turn, + clientData: context?.clientData, + source: "outer", + }); + + const compactEvent: CompactMessagesEvent = { + summary, + uiMessages: this.uiMessages, + modelMessages: this.modelMessages, + chatId: context?.chatId ?? "", + turn: context?.turn ?? 0, + clientData: context?.clientData, + source: "outer", + }; + + this.modelMessages = this._compaction.compactModelMessages + ? await this._compaction.compactModelMessages(compactEvent) + : [ + { + role: "assistant" as const, + content: [{ type: "text" as const, text: `[Conversation summary]\n\n${summary}` }], + }, + ]; + + if (this._compaction.compactUIMessages) { + this.uiMessages = await this._compaction.compactUIMessages(compactEvent); + } + + return true; + } +} + +// --------------------------------------------------------------------------- +// chat.createSession — async iterator for chat turns +// --------------------------------------------------------------------------- + +export type ChatSessionOptions = { + /** Run-level cancel signal (from task context). */ + signal: AbortSignal; + /** Seconds to stay idle between turns before suspending. @default 30 */ + idleTimeoutInSeconds?: number; + /** Duration string for suspend timeout. @default "1h" */ + timeout?: string; + /** Max turns before ending. @default 100 */ + maxTurns?: number; + /** Automatic context compaction — same options as `chat.agent({ compaction })`. */ + compaction?: ChatAgentCompactionOptions; + /** Configure mid-execution message injection — same options as `chat.agent({ pendingMessages })`. */ + pendingMessages?: PendingMessagesOptions; +}; + +export type ChatTurn = { + /** Turn number (0-indexed). */ + number: number; + /** Chat session ID. */ + chatId: string; + /** What triggered this turn. */ + trigger: string; + /** Client data from the transport (`metadata` field on the wire payload). */ + clientData: unknown; + /** Full accumulated model messages — pass directly to `streamText`. */ + readonly messages: ModelMessage[]; + /** Full accumulated UI messages — use for persistence. */ + readonly uiMessages: UIMessage[]; + /** Combined stop+cancel AbortSignal (fresh each turn). */ + signal: AbortSignal; + /** Whether the user stopped generation this turn. */ + readonly stopped: boolean; + /** Whether this is a continuation run. */ + continuation: boolean; + /** Token usage from the previous turn. Undefined on turn 0. */ + previousTurnUsage?: LanguageModelUsage; + /** Cumulative token usage across all completed turns so far. */ + totalUsage: LanguageModelUsage; + + /** + * Replace accumulated messages (for compaction). Takes UIMessages and + * converts to ModelMessages internally. After calling this, `turn.messages` + * reflects the compacted history. + */ + setMessages(uiMessages: UIMessage[]): Promise; + + /** + * Easy path: pipe stream, capture response, accumulate it, + * clean up aborted parts if stopped, and write turn-complete chunk. + */ + complete(source: UIMessageStreamable): Promise; + + /** + * Manual path: just write turn-complete chunk. + * Use when you've already piped and accumulated manually. + */ + done(): Promise; + + /** + * Add the response to the accumulator manually. + * Use with `chat.pipeAndCapture` when you need control between pipe and done. + */ + addResponse(response: UIMessage): Promise; + + /** + * Returns a `prepareStep` function that handles both compaction and + * pending message injection. Pass to `streamText({ prepareStep: turn.prepareStep() })`. + * Only needed when not using `chat.toStreamTextOptions()` (which auto-injects it). + */ + prepareStep(): + | ((args: { + messages: ModelMessage[]; + steps: CompactionStep[]; + }) => Promise<{ messages: ModelMessage[] } | undefined>) + | undefined; +}; + +/** + * Create a chat session that yields turns as an async iterator. + * + * Handles: preload wait, stop signals, message accumulation, turn-complete + * signaling, and idle/suspend between turns. You control: initialization, + * model/tool selection, persistence, and any custom per-turn logic. + * + * @example + * ```ts + * import { task } from "@trigger.dev/sdk"; + * import { chat, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai"; + * import { streamText } from "ai"; + * import { openai } from "@ai-sdk/openai"; + * + * export const myChat = task({ + * id: "my-chat", + * run: async (payload: ChatTaskWirePayload, { signal }) => { + * const session = chat.createSession(payload, { signal }); + * + * for await (const turn of session) { + * const result = streamText({ + * model: openai("gpt-4o"), + * messages: turn.messages, + * abortSignal: turn.signal, + * }); + * await turn.complete(result); + * } + * }, + * }); + * ``` + */ +function createChatSession( + payload: ChatTaskWirePayload, + options: ChatSessionOptions +): AsyncIterable { + const { + signal: runSignal, + idleTimeoutInSeconds: sessionIdleTimeoutOpt, + timeout = "1h", + maxTurns = 100, + compaction: sessionCompaction, + pendingMessages: sessionPendingMessages, + } = options; + + const idleTimeoutInSeconds = sessionIdleTimeoutOpt ?? 30; + + return { + [Symbol.asyncIterator]() { + let currentPayload = payload; + let turn = -1; + const stop = createStopSignal(); + const accumulator = new ChatMessageAccumulator(); + let previousTurnUsage: LanguageModelUsage | undefined; + let cumulativeUsage: LanguageModelUsage = emptyUsage(); + + return { + async next(): Promise> { + turn++; + + // First turn: handle preload — wait for the first real message + if (turn === 0 && currentPayload.trigger === "preload") { + const result = await messagesInput.waitWithIdleTimeout({ + idleTimeoutInSeconds: + sessionIdleTimeoutOpt ?? currentPayload.idleTimeoutInSeconds ?? 30, + timeout, + spanName: "waiting for first message", + }); + if (!result.ok || runSignal.aborted) { + stop.cleanup(); + return { done: true, value: undefined }; + } + currentPayload = result.output; + } + + // Subsequent turns: wait for the next message + if (turn > 0) { + // chat.requestUpgrade() / chat.endRun() — exit before waiting + if ( + locals.get(chatUpgradeRequestedKey) || + locals.get(chatEndRunRequestedKey) + ) { + stop.cleanup(); + return { done: true, value: undefined }; + } + + const next = await messagesInput.waitWithIdleTimeout({ + idleTimeoutInSeconds, + timeout, + spanName: "waiting for next message", + }); + if (!next.ok || runSignal.aborted) { + stop.cleanup(); + return { done: true, value: undefined }; + } + currentPayload = next.output; + } + + // Check limits + if (turn >= maxTurns || runSignal.aborted) { + stop.cleanup(); + return { done: true, value: undefined }; + } + + // Reset stop signal for this turn + stop.reset(); + + // Reset per-turn state + locals.set(chatResponsePartsKey, []); + // Set up steering queue and pending messages config in locals + // so toStreamTextOptions() auto-injects prepareStep for steering + const turnSteeringQueue: SteeringQueueEntry[] = []; + locals.set(chatSteeringQueueKey, turnSteeringQueue); + if (sessionPendingMessages) { + locals.set(chatPendingMessagesKey, sessionPendingMessages); + } + locals.set(chatTurnContextKey, { + chatId: currentPayload.chatId, + turn, + continuation: currentPayload.continuation ?? false, + clientData: currentPayload.metadata, + }); + + // Listen for messages during streaming (steering + next-turn buffer) + const sessionPendingWire: ChatTaskWirePayload[] = []; + const sessionMsgSub = messagesInput.on(async (msg) => { + sessionPendingWire.push(msg); + + if (sessionPendingMessages) { + // Slim wire: at most one delta message per record. Read + // `msg.message` directly — no array slicing needed. + const lastUIMessage = msg.message; + if (lastUIMessage) { + if (sessionPendingMessages.onReceived) { + try { + await sessionPendingMessages.onReceived({ + message: lastUIMessage, + chatId: currentPayload.chatId, + turn, + }); + } catch { + /* non-fatal */ + } + } + try { + const modelMsgs = await toModelMessages([lastUIMessage]); + turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); + } catch { + /* non-fatal */ + } + } + } + }); + + // Accumulate messages. Slim wire: pass the single delta message as + // a 0-or-1-length array. The accumulator's behavior is unchanged — + // it still appends user messages and reconverts on regenerate. + const incomingForAccumulator: UIMessage[] = currentPayload.message + ? [currentPayload.message] + : []; + const messages = await accumulator.addIncoming( + incomingForAccumulator, + currentPayload.trigger, + turn + ); + + // chat.requestUpgrade() called before this turn — signal transport and exit + if (locals.get(chatUpgradeRequestedKey)) { + await writeUpgradeRequiredChunk(); + sessionMsgSub.off(); + stop.cleanup(); + return { done: true, value: undefined }; + } + + const combinedSignal = AbortSignal.any([runSignal, stop.signal]); + + const turnObj: ChatTurn = { + number: turn, + chatId: currentPayload.chatId, + trigger: currentPayload.trigger, + clientData: currentPayload.metadata, + get messages() { + return accumulator.modelMessages; + }, + get uiMessages() { + return accumulator.uiMessages; + }, + signal: combinedSignal, + get stopped() { + return stop.signal.aborted && !runSignal.aborted; + }, + continuation: currentPayload.continuation ?? false, + previousTurnUsage, + totalUsage: cumulativeUsage, + + async setMessages(uiMessages: UIMessage[]) { + await accumulator.setMessages(uiMessages); + }, + + async complete(source: UIMessageStreamable) { + let response: UIMessage | undefined; + try { + response = await pipeChatAndCapture(source, { signal: combinedSignal }); + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + if (runSignal.aborted) { + // Full cancel — don't accumulate + sessionMsgSub.off(); + await chatWriteTurnComplete(); + return undefined; + } + // Stop — fall through to accumulate partial response + } else { + throw error; + } + } + + if (response) { + const cleaned = + stop.signal.aborted && !runSignal.aborted + ? cleanupAbortedParts(response) + : response; + // Append any non-transient data parts queued via chat.response or writer.write() + const queuedParts = locals.get(chatResponsePartsKey); + if (queuedParts && queuedParts.length > 0) { + (cleaned as any).parts = [...(cleaned.parts ?? []), ...queuedParts]; + locals.set(chatResponsePartsKey, []); + } + await accumulator.addResponse(cleaned); + } else { + // No response (manual pipe mode) but there are queued data parts + const queuedParts = locals.get(chatResponsePartsKey); + if (queuedParts && queuedParts.length > 0) { + await accumulator.addResponse({ + id: generateMessageId(), + role: "assistant" as const, + parts: queuedParts as UIMessage["parts"], + }); + locals.set(chatResponsePartsKey, []); + } + } + + // Capture token usage from the streamText result + let turnUsage: LanguageModelUsage | undefined; + if (typeof (source as any).totalUsage?.then === "function") { + try { + const usage: LanguageModelUsage = await (source as any).totalUsage; + turnUsage = usage; + previousTurnUsage = usage; + cumulativeUsage = addUsage(cumulativeUsage, usage); + } catch { + /* non-fatal */ + } + } + + // Outer-loop compaction (same logic as chat.agent) + if (sessionCompaction && turnUsage && !turnObj.stopped) { + const shouldTrigger = await sessionCompaction.shouldCompact({ + messages: accumulator.modelMessages, + totalTokens: turnUsage.totalTokens, + inputTokens: turnUsage.inputTokens, + outputTokens: turnUsage.outputTokens, + usage: turnUsage, + totalUsage: cumulativeUsage, + chatId: currentPayload.chatId, + turn, + clientData: currentPayload.metadata, + source: "outer", + }); + + if (shouldTrigger) { + const summary = await sessionCompaction.summarize({ + messages: accumulator.modelMessages, + usage: turnUsage, + totalUsage: cumulativeUsage, + chatId: currentPayload.chatId, + turn, + clientData: currentPayload.metadata, + source: "outer", + }); + + const compactEvent: CompactMessagesEvent = { + summary, + uiMessages: accumulator.uiMessages, + modelMessages: accumulator.modelMessages, + chatId: currentPayload.chatId, + turn, + clientData: currentPayload.metadata, + source: "outer", + }; + + accumulator.modelMessages = sessionCompaction.compactModelMessages + ? await sessionCompaction.compactModelMessages(compactEvent) + : [ + { + role: "assistant" as const, + content: [ + { type: "text" as const, text: `[Conversation summary]\n\n${summary}` }, + ], + }, + ]; + + if (sessionCompaction.compactUIMessages) { + accumulator.uiMessages = await sessionCompaction.compactUIMessages( + compactEvent + ); + } + } + } + + sessionMsgSub.off(); + await chatWriteTurnComplete(); + return response; + }, + + async addResponse(response: UIMessage) { + // Append any non-transient data parts queued via chat.response or writer.write() + const queuedParts = locals.get(chatResponsePartsKey); + if (queuedParts && queuedParts.length > 0) { + response = { ...response, parts: [...(response.parts ?? []), ...(queuedParts as UIMessage["parts"])] }; + locals.set(chatResponsePartsKey, []); + } + await accumulator.addResponse(response); + }, + + async done() { + sessionMsgSub.off(); + await chatWriteTurnComplete(); + }, + + prepareStep() { + const hasCompaction = !!sessionCompaction; + const hasPending = !!sessionPendingMessages; + if (!hasCompaction && !hasPending) return undefined; + + return async ({ + messages: stepMsgs, + steps, + }: { + messages: ModelMessage[]; + steps: CompactionStep[]; + }) => { + let resultMessages: ModelMessage[] | undefined; + + if (sessionCompaction) { + const compactResult = await chatCompact(stepMsgs, steps, { + shouldCompact: sessionCompaction.shouldCompact, + summarize: (msgs) => + sessionCompaction.summarize({ messages: msgs, source: "inner" }), + }); + if (compactResult.type !== "skipped") { + resultMessages = compactResult.messages; + } + } + + if (sessionPendingMessages) { + const injected = await drainSteeringQueue( + sessionPendingMessages, + resultMessages ?? stepMsgs, + steps, + turnSteeringQueue + ); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? stepMsgs), ...injected]; + } + } + + return resultMessages ? { messages: resultMessages } : undefined; + }; + }, + }; + + return { done: false, value: turnObj }; + }, + + async return() { + stop.cleanup(); + return { done: true, value: undefined }; + }, + }; + }, + }; +} + +// --------------------------------------------------------------------------- +// chat.local — per-run typed data with Proxy access +// --------------------------------------------------------------------------- + +/** @internal Symbol for storing the locals key on the proxy target. */ +const CHAT_LOCAL_KEY: unique symbol = Symbol("chatLocalKey"); +/** @internal Symbol for storing the dirty-tracking locals key. */ +const CHAT_LOCAL_DIRTY_KEY: unique symbol = Symbol("chatLocalDirtyKey"); + +// --------------------------------------------------------------------------- +// chat.local registry — tracks all declared locals for serialization +// --------------------------------------------------------------------------- + +type ChatLocalEntry = { key: ReturnType; id: string }; +const chatLocalRegistry = new Set(); + +/** @internal Run-scoped flag to ensure hydration happens at most once per run. */ +const chatLocalsHydratedKey = locals.create("chat.locals.hydrated"); + +/** + * Hydrate chat.local values from subtask metadata (set by `ai.toolExecute()` or legacy `ai.tool()`). + * Runs once per run — subsequent calls are no-ops. + * @internal + */ +function hydrateLocalsFromMetadata(): void { + if (locals.get(chatLocalsHydratedKey)) return; + locals.set(chatLocalsHydratedKey, true); + const opts = metadata.get(METADATA_KEY) as ToolCallExecutionOptions | undefined; + if (!opts?.chatLocals) return; + for (const [id, value] of Object.entries(opts.chatLocals)) { + locals.set(locals.create(id), value); + } +} + +/** + * A Proxy-backed, run-scoped data object that appears as `T` to users. + * Includes helper methods for initialization, dirty tracking, and serialization. + * Internal metadata is stored behind Symbols and invisible to + * `Object.keys()`, `JSON.stringify()`, and spread. + */ +export type ChatLocal> = T & { + /** Initialize the local with a value. Call in `onChatStart` or `run()`. */ + init(value: T): void; + /** Returns `true` if any property was set since the last check. Resets the dirty flag. */ + hasChanged(): boolean; + /** Returns a plain object copy of the current value. Useful for persistence. */ + get(): T; + readonly [CHAT_LOCAL_KEY]: ReturnType>; + readonly [CHAT_LOCAL_DIRTY_KEY]: ReturnType>; +}; + +/** + * Creates a per-run typed data object accessible from anywhere during task execution. + * + * Declare at module level, then initialize inside a lifecycle hook (e.g. `onChatStart`) + * using `chat.initLocal()`. Properties are accessible directly via the Proxy. + * + * Multiple locals can coexist — each gets its own isolated run-scoped storage. + * + * The `id` is required and must be unique across all `chat.local()` calls in + * your project. It's used to serialize values into subtask metadata so that + * `ai.toolExecute()` (or legacy `ai.tool()`) subtasks can auto-hydrate parent locals (read-only). + * + * @example + * ```ts + * import { chat } from "@trigger.dev/sdk/ai"; + * + * const userPrefs = chat.local<{ theme: string; language: string }>({ id: "userPrefs" }); + * const gameState = chat.local<{ score: number; streak: number }>({ id: "gameState" }); + * + * export const myChat = chat.agent({ + * id: "my-chat", + * onChatStart: async ({ clientData }) => { + * const prefs = await db.prefs.findUnique({ where: { userId: clientData.userId } }); + * userPrefs.init(prefs ?? { theme: "dark", language: "en" }); + * gameState.init({ score: 0, streak: 0 }); + * }, + * onTurnComplete: async ({ chatId }) => { + * if (gameState.hasChanged()) { + * await db.save({ where: { chatId }, data: gameState.get() }); + * } + * }, + * run: async ({ messages }) => { + * gameState.score++; + * return streamText({ + * system: `User prefers ${userPrefs.theme} theme. Score: ${gameState.score}`, + * messages, + * }); + * }, + * }); + * ``` + */ +function chatLocal>(options: { id: string }): ChatLocal { + const id = `chat.local.${options.id}`; + const localKey = locals.create(id); + const dirtyKey = locals.create(`${id}.dirty`); + + chatLocalRegistry.add({ key: localKey, id }); + + const target = {} as any; + target[CHAT_LOCAL_KEY] = localKey; + target[CHAT_LOCAL_DIRTY_KEY] = dirtyKey; + + return new Proxy(target, { + get(_target, prop, _receiver) { + // Internal Symbol properties + if (prop === CHAT_LOCAL_KEY) return _target[CHAT_LOCAL_KEY]; + if (prop === CHAT_LOCAL_DIRTY_KEY) return _target[CHAT_LOCAL_DIRTY_KEY]; + + // Instance methods + if (prop === "init") { + return (value: T) => { + locals.set(localKey, value); + locals.set(dirtyKey, false); + }; + } + if (prop === "hasChanged") { + return () => { + const dirty = locals.get(dirtyKey) ?? false; + locals.set(dirtyKey, false); + return dirty; + }; + } + if (prop === "get") { + return () => { + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } + if (current === undefined) { + throw new Error("local.get() called before initialization. Call local.init() first."); + } + return { ...current }; + }; + } + // toJSON for serialization (JSON.stringify(local)) + if (prop === "toJSON") { + return () => { + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } + return current ? { ...current } : undefined; + }; + } + + let current = locals.get(localKey); + if (current === undefined) { + // Auto-hydrate from parent metadata in subtask context + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } + if (current === undefined) return undefined; + return (current as any)[prop]; + }, + + set(_target, prop, value) { + // Don't allow setting internal Symbols + if (typeof prop === "symbol") return false; + + const current = locals.get(localKey); + if (current === undefined) { + throw new Error( + "chat.local can only be modified after initialization. " + + "Call local.init() in onChatStart or run() first." + ); + } + locals.set(localKey, { ...current, [prop]: value }); + locals.set(dirtyKey, true); + return true; + }, + + has(_target, prop) { + if (typeof prop === "symbol") return prop in _target; + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } + return current !== undefined && prop in current; + }, + + ownKeys() { + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } + return current ? Reflect.ownKeys(current) : []; + }, + + getOwnPropertyDescriptor(_target, prop) { + if (typeof prop === "symbol") return undefined; + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } + if (current === undefined || !(prop in current)) return undefined; + return { + configurable: true, + enumerable: true, + writable: true, + value: (current as any)[prop], + }; + }, + }) as ChatLocal; +} + +/** + * Extracts the client data (metadata) type from a chat task. + * Use this to type the `metadata` option on the transport. + * + * @example + * ```ts + * import type { InferChatClientData } from "@trigger.dev/sdk/ai"; + * import type { myChat } from "@/trigger/chat"; + * + * type MyClientData = InferChatClientData; + * // { model?: string; userId: string } + * ``` + */ +// `InferChatClientData` and `InferChatUIMessage` live in `./ai-shared.ts` +// so the chat React hooks can import them without dragging `ai.ts` into +// 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"; + +/** + * Options for {@link createChatStartSessionAction}. + */ +export type CreateChatStartSessionActionOptions = { + /** TTL for the session-scoped public access token. @default "1h" */ + tokenTTL?: string | number | Date; + /** + * Default trigger config used when starting a new session for a chat. + * Per-call `params.triggerConfig` shallow-merges on top. + */ + triggerConfig?: Partial; +}; + +/** + * Params for the function returned by {@link createChatStartSessionAction}. + */ +export type ChatStartSessionParams = { + /** Conversation id (mapped to the Session's `externalId`). */ + chatId: string; + /** + * Per-call trigger config. Shallow-merged over the action's default + * `triggerConfig`. `basePayload` is the customer's wire payload (for + * `chat.agent`: anything beyond `chatId`/`messages`/`trigger`/`metadata`, + * which the runtime injects automatically). + */ + triggerConfig?: Partial; + /** Pass-through metadata folded into the session row. */ + metadata?: Record; +}; + +/** + * Result from {@link createChatStartSessionAction}'s returned function. + */ +export type ChatStartSessionResult = { + /** + * Session-scoped public access token (`read:sessions:{chatId} + + * write:sessions:{chatId}`). Pass this to the browser; the transport + * uses it to call `.in/append`, `.out`, `end-and-continue`. + */ + publicAccessToken: string; + /** Friendly id of the run triggered alongside session create. */ + runId: string; + /** Session friendlyId — informational. */ + sessionId: string; +}; + +/** + * Creates a server-side helper that starts (or resumes) a Session for a + * given chatId — atomically creating the row, triggering the first run, + * and returning a session-scoped PAT for the browser to use. + * + * Wrap in a Next.js server action (or any server-side handler) so the + * customer's secret key never crosses to the browser. + * + * @example + * ```ts + * // actions.ts + * "use server"; + * import { chat } from "@trigger.dev/sdk/ai"; + * + * export const startChatSession = chat.createStartSessionAction("my-chat", { + * triggerConfig: { machine: "small-1x" }, + * }); + * ``` + * + * Then in the browser: + * ```tsx + * const transport = useTriggerChatTransport({ + * task: "my-chat", + * accessToken: async ({ chatId }) => { + * const { publicAccessToken } = await startChatSession({ chatId }); + * return publicAccessToken; + * }, + * }); + * ``` + */ +function createChatStartSessionAction( + taskId: string, + options?: CreateChatStartSessionActionOptions +): (params: ChatStartSessionParams) => Promise { + return async (params: ChatStartSessionParams): Promise => { + if (!params.chatId) { + throw new Error( + "chat.createStartSessionAction: params.chatId is required — used as the session externalId." + ); + } + + // The first run boots before the user's first message lands on + // `.in/append`, so it sees an empty `messages` array and `trigger: + // "preload"`. This matches the pre-Sessions preload semantics: + // `onPreload` fires, the runtime opens its `.in` subscription, the + // first user message arrives moments later via `.in/append`. + // + // `metadata` is the customer's transport-level `clientData`, + // threaded through so the agent's `clientDataSchema` validates on + // the very first turn (the typical schema requires `userId` etc.). + // Auto-tag every chat.agent run with `chat:{chatId}` so the dashboard / + // run-list filter by chat works without the customer having to wire it + // up. Mirrors the browser-mediated `TriggerChatTransport.doStart` path. + const userTags = params.triggerConfig?.tags ?? options?.triggerConfig?.tags ?? []; + const tags = [`chat:${params.chatId}`, ...userTags].slice(0, 5); + + const triggerConfig: SessionTriggerConfig = { + basePayload: { + messages: [], + trigger: "preload", + ...(options?.triggerConfig?.basePayload ?? {}), + ...(params.triggerConfig?.basePayload ?? {}), + chatId: params.chatId, + }, + ...(options?.triggerConfig?.machine || params.triggerConfig?.machine + ? { machine: params.triggerConfig?.machine ?? options?.triggerConfig?.machine } + : {}), + ...(options?.triggerConfig?.queue || params.triggerConfig?.queue + ? { queue: params.triggerConfig?.queue ?? options?.triggerConfig?.queue } + : {}), + tags, + ...(options?.triggerConfig?.maxAttempts !== undefined || + params.triggerConfig?.maxAttempts !== undefined + ? { + maxAttempts: + params.triggerConfig?.maxAttempts ?? options?.triggerConfig?.maxAttempts!, + } + : {}), + ...(options?.triggerConfig?.idleTimeoutInSeconds !== undefined || + params.triggerConfig?.idleTimeoutInSeconds !== undefined + ? { + idleTimeoutInSeconds: + params.triggerConfig?.idleTimeoutInSeconds ?? + options?.triggerConfig?.idleTimeoutInSeconds!, + } + : {}), + }; + + const created = await sessions.start({ + type: "chat.agent", + externalId: params.chatId, + taskIdentifier: taskId, + triggerConfig, + metadata: params.metadata, + }); + + // Session create returns a session PAT directly when called with a + // start token, but when the SDK call goes via the secret key we still + // need to mint our own (the server returns a PAT regardless, but + // re-minting here lets the customer override `tokenTTL`). + const publicAccessToken = + options?.tokenTTL !== undefined + ? await auth.createPublicToken({ + scopes: { + read: { sessions: params.chatId }, + write: { sessions: params.chatId }, + }, + expirationTime: options.tokenTTL, + }) + : created.publicAccessToken; + + return { + publicAccessToken, + runId: created.runId, + sessionId: created.id, + }; + }; +} + +export const chat = { + /** Create a chat agent. See {@link chatAgent}. */ + agent: chatAgent, + /** Create a custom agent with manual lifecycle control. See {@link chatCustomAgent}. */ + customAgent: chatCustomAgent, + /** Create a chat task with a fixed {@link UIMessage} subtype and optional default stream options. See {@link withUIMessage}. */ + withUIMessage, + /** Create a chat task with a fixed client data schema. See {@link withClientData}. */ + withClientData, + /** Create a server-side helper for starting (or resuming) a Session for a chatId. See {@link createChatStartSessionAction}. */ + createStartSessionAction: createChatStartSessionAction, + /** Pipe a stream to the chat transport. See {@link pipeChat}. */ + pipe: pipeChat, + /** Create a per-run typed local. See {@link chatLocal}. */ + local: chatLocal, + /** Create a public access token for a chat task. See {@link createChatAccessToken}. */ + createAccessToken: createChatAccessToken, + /** Override the turn timeout at runtime (duration string). See {@link setTurnTimeout}. */ + setTurnTimeout, + /** Override the turn timeout at runtime (seconds). See {@link setTurnTimeoutInSeconds}. */ + setTurnTimeoutInSeconds, + /** Override the idle timeout at runtime. See {@link setIdleTimeoutInSeconds}. */ + setIdleTimeoutInSeconds, + /** Override toUIMessageStream() options for the current turn. See {@link setUIMessageStreamOptions}. */ + setUIMessageStreamOptions, + /** Check if the current turn was stopped by the user. See {@link isStopped}. */ + isStopped, + /** Request that the run exits after the current turn so the next message starts on the latest version. See {@link requestUpgrade}. */ + requestUpgrade, + /** Exit the run after the current turn completes, without any upgrade signal. See {@link endRun}. */ + endRun, + /** Clean up aborted parts from a UIMessage. See {@link cleanupAbortedParts}. */ + cleanupAbortedParts, + /** Register background work that runs in parallel with streaming. See {@link chatDefer}. */ + defer: chatDefer, + /** Queue model messages for injection at the next `prepareStep` boundary. See {@link injectBackgroundContext}. */ + inject: injectBackgroundContext, + /** Typed chat output stream for writing custom chunks or piping from subtasks. */ + stream: chatStream, + /** Write data parts that persist to the response message. See {@link chatResponse}. */ + response: chatResponse, + /** + * Typed, bidirectional shared data slot for the chat. + * + * Use from `chat.agent` hooks and `run()` to share state with the client. + * Setting emits a `store-snapshot` chunk; patching emits `store-delta`. + * The value persists across turns within the same run, and can be + * restored after continuations via the `hydrateStore` config option. + * + * ```ts + * chat.store.set({ plan: ["research", "draft", "review"] }); + * chat.store.patch([{ op: "replace", path: "/status", value: "done" }]); + * const current = chat.store.get(); + * const off = chat.store.onChange((value, ops) => { ... }); + * ``` + */ + store: { + /** Replace the store value. Emits a `store-snapshot` chunk. */ + set: chatStoreSet, + /** + * Apply RFC 6902 JSON Patch operations to the current value. + * Emits a `store-delta` chunk. + */ + patch: chatStorePatch, + /** Read the current store value. Returns `undefined` if never set. */ + get: chatStoreGet, + /** Subscribe to store changes. Returns an unsubscribe function. */ + onChange: chatStoreOnChange, + }, + /** Pre-built input stream for receiving messages from the transport. */ + messages: messagesInput, + /** Create a managed stop signal wired to the stop input stream. See {@link createStopSignal}. */ + createStopSignal, + /** Signal the frontend that the current turn is complete. See {@link chatWriteTurnComplete}. */ + writeTurnComplete: chatWriteTurnComplete, + /** Pipe a stream and capture the response message. See {@link pipeChatAndCapture}. */ + pipeAndCapture: pipeChatAndCapture, + /** Message accumulator class for raw task chat. See {@link ChatMessageAccumulator}. */ + MessageAccumulator: ChatMessageAccumulator, + /** Create a chat session (async iterator). See {@link createChatSession}. */ + createSession: createChatSession, + /** + * Store and retrieve a resolved prompt for the current run. + * + * - `chat.prompt.set(resolved)` — store a `ResolvedPrompt` or plain string + * - `chat.prompt()` — read the stored prompt (throws if not set) + */ + prompt: Object.assign(getChatPrompt, { set: setChatPrompt }), + /** + * Store and retrieve resolved agent skills for the current run. + * + * - `chat.skills.set([...])` — store an array of `ResolvedSkill`s + * - `chat.skills()` — read the stored skills (returns undefined if none) + * + * Skills set here are automatically injected into `streamText` by + * `chat.toStreamTextOptions()`: skill descriptions land in the system + * prompt and `loadSkill` / `readFile` / `bash` tools are added to the + * tool set. + */ + skills: Object.assign(getChatSkills, { set: setChatSkills }), + /** + * Returns an options object ready to spread into `streamText()`. + * Reads the stored prompt and returns `{ system, experimental_telemetry, ...config }`. + * Returns `{}` if no prompt has been set. + */ + toStreamTextOptions, + /** + * Replace the accumulated conversation messages for compaction. + * Call from `onTurnStart` or `onTurnComplete`. Takes `UIMessage[]` and + * converts to `ModelMessage[]` internally. + */ + setMessages: setChatMessages, + /** + * Imperative API for modifying the accumulated message history. + * Supports rollback, remove, replace, slice, and full replacement. + * Can be called from any hook or `run()`. + */ + history: chatHistory, + /** Check if it's safe to compact messages (no in-flight tool calls). */ + isCompactionSafe, + /** Returns a `prepareStep` function that handles context compaction automatically. */ + compactionStep: chatCompactionStep, + /** Low-level compaction for use inside a custom `prepareStep`. */ + compact: chatCompact, + /** Read the current compaction state (summary + base message count). */ + getCompactionState, + /** + * The friendlyId (`session_*`) of the backing Session for the current chat.agent run. + * Useful for persisting alongside `runId` so reloads can resume the same session. + * Throws if called outside a chat.agent `run()` or hook. + */ + get sessionId(): string { + return getChatSession().id; + }, +}; + +/** + * Writes a turn-complete control chunk to the chat output stream. + * The frontend transport intercepts this to close the ReadableStream for the current turn. + * @internal + */ +async function writeTurnCompleteChunk( + chatId?: string, + publicAccessToken?: string +): Promise { + const { waitUntilComplete } = chatStream.writer({ + spanName: "turn complete", + collapsed: true, + execute: ({ write }) => { + // Transport-intercepted control chunk — not a valid UIMessageChunk + // type but travels on the same session.out stream. + write({ + type: "trigger:turn-complete", + ...(publicAccessToken ? { publicAccessToken } : {}), + } as unknown as UIMessageChunk); + }, + }); + return await waitUntilComplete(); +} + +/** + * Hand off the session to a fresh run on the latest version and emit a + * telemetry chunk on `.out` so the transport can hide it from the + * consumer. + * + * Server-side flow (in `POST /sessions/:id/end-and-continue`): + * 1. Trigger a new run with the session's `triggerConfig` + * 2. Atomically swap `Session.currentRunId` to the new run's id + * (via optimistic claim keyed on the calling run's id) + * 3. Return the new runId + * + * The transport keeps its `.out` SSE open across the swap — v1's last + * chunks land, v2's new chunks land on the same stream (S2 keys on + * the session, not the run). The transport filters + * `trigger:upgrade-required` for cleanliness; consumers see no gap. + * + * If the swap fails (no current run, no env auth, etc.) we still emit + * the chunk and exit. The next `.in/append` will trigger a new run via + * the probe path; it just won't be quite as seamless. + * + * @internal + */ +async function writeUpgradeRequiredChunk(): Promise { + const ctx = taskContext.ctx; + const chatId = ctx?.run.id ? getChatIdFromContext() : undefined; + const callingRunId = ctx?.run.id; + + if (chatId && callingRunId) { + const apiClient = apiClientManager.clientOrThrow(); + try { + await apiClient.endAndContinueSession(chatId, { + callingRunId, + reason: "upgrade", + }); + } catch (error) { + // Non-fatal: the next `.in/append` re-triggers via the probe. + // Swallow rather than throw so we still emit the chunk + exit. + logger.warn("end-and-continue failed; falling back to probe-on-append", { + chatId, + callingRunId, + error, + }); + } + } + + const { waitUntilComplete } = chatStream.writer({ + spanName: "upgrade required", + collapsed: true, + execute: ({ write }) => { + write({ + type: "trigger:upgrade-required", + } as unknown as UIMessageChunk); + }, + }); + return await waitUntilComplete(); +} + +/** + * Resolves the current chat's `chatId` (used as session externalId) from + * the bound session handle. Returns `undefined` if no agent is bound — + * shouldn't happen at the call sites that invoke + * `writeUpgradeRequiredChunk`, but defensive against misuse. + * @internal + */ +function getChatIdFromContext(): string | undefined { + return locals.get(chatSessionHandleKey)?.id; +} + +/** + * Extracts the text content of the last user message from a UIMessage array. + * Returns undefined if no user message is found. + * @internal + */ +function extractLastUserMessageText(messages: UIMessage[]): string | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]!; + if (msg.role !== "user") continue; + + // UIMessage uses parts array + if (msg.parts) { + const textParts = msg.parts + .filter((p: any) => p.type === "text" && p.text) + .map((p: any) => p.text as string); + if (textParts.length > 0) { + return textParts.join("\n"); + } + } + + break; + } + + return undefined; +} + +/** + * Strips ephemeral OpenAI Responses API `itemId` from a UIMessage's parts. + * + * The OpenAI Responses provider attaches `itemId` to message parts via + * `providerMetadata.openai.itemId`. These IDs are ephemeral — sending them + * back in a subsequent `streamText` call causes 404s because the provider + * can't find the referenced item (especially for stopped/partial responses). + * + * @internal + */ +function stripProviderMetadata(message: UIMessage): UIMessage { + if (!message.parts) return message; + return { + ...message, + parts: message.parts.map((part: any) => { + const openai = part.providerMetadata?.openai; + if (!openai?.itemId) return part; + + const { itemId, ...restOpenai } = openai; + const { openai: _, ...restProviders } = part.providerMetadata; + return { + ...part, + providerMetadata: { + ...restProviders, + ...(Object.keys(restOpenai).length > 0 ? { openai: restOpenai } : {}), + }, + }; + }), + }; +} diff --git a/packages/trigger-sdk/src/v3/chat-server.test.ts b/packages/trigger-sdk/src/v3/chat-server.test.ts new file mode 100644 index 00000000000..dc9ef11788f --- /dev/null +++ b/packages/trigger-sdk/src/v3/chat-server.test.ts @@ -0,0 +1,617 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { simulateReadableStream, streamText } from "ai"; +import type { UIMessageChunk } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; + +// Stub `SessionStreamInstance` so the handler's S2 tee is a no-op +// instead of trying to reach a real S2 endpoint. The real one calls +// `apiClient.initializeSessionStream` then pipes via S2 — both are +// out of scope for handler-shape tests. +vi.mock("@trigger.dev/core/v3", async (importActual) => { + const actual = (await importActual()) as Record; + class StubSessionStreamInstance { + constructor(opts: { source: ReadableStream }) { + // Drain the source so the upstream tee doesn't backpressure-stall + // the SSE half. We don't keep the chunks — durability/resume is + // out of scope here. + void (async () => { + const reader = opts.source.getReader(); + try { + while (true) { + const { done } = await reader.read(); + if (done) break; + } + } finally { + reader.releaseLock(); + } + })(); + } + async wait() { + return { written: 0 }; + } + } + return { ...actual, SessionStreamInstance: StubSessionStreamInstance }; +}); + +// Import AFTER the mock so chat-server picks up the stubbed class. +import { chat } from "./chat-server.js"; +import { apiClientManager } from "@trigger.dev/core/v3"; + +// ── Helpers ──────────────────────────────────────────────────────────── + +function textStream(text: string): ReadableStream { + return simulateReadableStream({ + chunks: [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 5, reasoning: undefined }, + }, + }, + ], + }); +} + +function toolCallStream(): ReadableStream { + return simulateReadableStream({ + chunks: [ + { + type: "tool-call", + toolCallId: "tc-1", + toolName: "weather", + input: JSON.stringify({ city: "tokyo" }), + }, + { + type: "finish", + finishReason: { unified: "tool-calls", raw: "tool-calls" }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 0, reasoning: undefined }, + }, + }, + ], + }); +} + +function makeRequest(body: unknown): Request { + return new Request("https://my-app.example/api/chat", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(body), + }); +} + +const SESSION_PAT = "tr_session_pat_for_handover"; + +function createSessionResponse(externalId: string): Response { + return new Response( + JSON.stringify({ + id: "session_test", + externalId, + type: "chat.agent", + taskIdentifier: "test-agent", + triggerConfig: { + basePayload: { chatId: externalId, trigger: "handover-prepare" }, + idleTimeoutInSeconds: 60, + }, + currentRunId: "run_test", + runId: "run_test", + publicAccessToken: SESSION_PAT, + tags: [], + metadata: null, + closedAt: null, + closedReason: null, + expiresAt: null, + createdAt: new Date(0).toISOString(), + updatedAt: new Date(0).toISOString(), + isCached: false, + }), + { + status: 200, + headers: { "content-type": "application/json" }, + } + ); +} + +function appendOkResponse(): Response { + return new Response(JSON.stringify({ ok: true }), { + status: 200, + headers: { "content-type": "application/json" }, + }); +} + +async function readSSEBodyToChunks(res: Response): Promise { + const text = await res.text(); + return text + .split("\n\n") + .filter((b) => b.startsWith("data: ")) + .map((b) => JSON.parse(b.slice(6)) as UIMessageChunk); +} + +type CapturedRequest = { url: string; init?: RequestInit }; + +async function withApiContext(fn: () => Promise): Promise { + return apiClientManager.runWithConfig( + { + baseURL: "https://api.test.trigger.dev", + secretKey: "tr_test_secret", + }, + fn + ); +} + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("chat.headStart (route handler)", () => { + let originalFetch: typeof global.fetch; + + beforeEach(() => { + originalFetch = global.fetch; + }); + + afterEach(() => { + global.fetch = originalFetch; + vi.restoreAllMocks(); + }); + + it("creates the session with handover-prepare in basePayload and returns the session PAT in headers", async () => { + const requests: CapturedRequest[] = []; + global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => { + const urlStr = typeof url === "string" ? url : url.toString(); + requests.push({ url: urlStr, init }); + if (urlStr.endsWith("/api/v1/sessions") || urlStr.endsWith("/api/v1/sessions/")) { + return createSessionResponse("chat-1"); + } + if (urlStr.includes("/realtime/v1/sessions/") && urlStr.endsWith("/in/append")) { + return appendOkResponse(); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const handler = chat.headStart({ + agentId: "test-agent", + run: async ({ chat: chatHelper }) => { + return streamText({ + ...chatHelper.toStreamTextOptions(), + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("hi back") }), + }), + }); + }, + }); + + const res = await withApiContext(() => + handler( + makeRequest({ + chatId: "chat-1", + trigger: "submit-message", + headStartMessages: [{ id: "m1", role: "user", parts: [{ type: "text", text: "hi" }] }], + }) + ) + ); + + expect(res.status).toBe(200); + expect(res.headers.get("X-Trigger-Chat-Id")).toBe("chat-1"); + expect(res.headers.get("X-Trigger-Chat-Access-Token")).toBe(SESSION_PAT); + expect(res.headers.get("Content-Type")).toMatch(/text\/event-stream/); + + const sessionCreate = requests.find((r) => + r.url.endsWith("/api/v1/sessions") || r.url.endsWith("/api/v1/sessions/") + ); + expect(sessionCreate).toBeDefined(); + const body = JSON.parse(sessionCreate!.init!.body as string); + expect(body.type).toBe("chat.agent"); + expect(body.externalId).toBe("chat-1"); + expect(body.taskIdentifier).toBe("test-agent"); + // The trigger payload is rewritten to handover-prepare even though the + // browser sent submit-message — the agent boots into the handover wait branch. + expect(body.triggerConfig.basePayload.trigger).toBe("handover-prepare"); + expect(body.triggerConfig.basePayload.chatId).toBe("chat-1"); + expect(body.triggerConfig.basePayload.idleTimeoutInSeconds).toBe(60); + }); + + it("dispatches handover with isFinal=true on pure-text finishReason", async () => { + const requests: CapturedRequest[] = []; + global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => { + const urlStr = typeof url === "string" ? url : url.toString(); + requests.push({ url: urlStr, init }); + if (urlStr.endsWith("/api/v1/sessions") || urlStr.endsWith("/api/v1/sessions/")) { + return createSessionResponse("chat-final"); + } + if (urlStr.includes("/realtime/v1/sessions/") && urlStr.endsWith("/in/append")) { + return appendOkResponse(); + } + // Stitched response subscribes to `.out` after handover. + if (/\/realtime\/v1\/sessions\/[^/]+\/out$/.test(urlStr)) { + return new Response(new ReadableStream({ start(c) { c.close(); } }), { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const handler = chat.headStart({ + agentId: "test-agent", + run: async ({ chat: chatHelper }) => { + return streamText({ + ...chatHelper.toStreamTextOptions(), + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("just a text reply") }), + }), + }); + }, + }); + + const res = await withApiContext(() => + handler( + makeRequest({ + chatId: "chat-final", + trigger: "submit-message", + // Slim wire: head-start ships full history via `headStartMessages` + // (not `messages` / `message`). The route handler reads that field + // off the request body before invoking the customer's run(). + headStartMessages: [{ id: "m1", role: "user", parts: [{ type: "text", text: "hi" }] }], + }) + ) + ); + + // Drain the SSE body so handoverWhenDone observes finishReason. + const chunks = await readSSEBodyToChunks(res); + expect(chunks.some((c) => c.type === "text-delta")).toBe(true); + + // Give the deferred handoverWhenDone a tick to dispatch. + await new Promise((r) => setTimeout(r, 30)); + + const handoverPost = requests.find( + (r) => + r.url.includes("/realtime/v1/sessions/chat-final/in/append") && + r.init?.body !== undefined + ); + expect(handoverPost).toBeDefined(); + const body = JSON.parse(handoverPost!.init!.body as string); + // Pure-text finishes go through `kind: "handover"` with `isFinal: true` + // so the agent runs hooks (persistence, etc.) without making an LLM call. + expect(body.kind).toBe("handover"); + expect(body.isFinal).toBe(true); + // The partial carries the customer's response messages — a single + // assistant message with the streamed text. + expect(Array.isArray(body.partialAssistantMessage)).toBe(true); + const assistant = body.partialAssistantMessage.find( + (m: { role: string }) => m.role === "assistant" + ); + expect(assistant).toBeDefined(); + }); + + it("dispatches handover with response.messages on tool-call finishReason", async () => { + const requests: CapturedRequest[] = []; + global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => { + const urlStr = typeof url === "string" ? url : url.toString(); + requests.push({ url: urlStr, init }); + if (urlStr.endsWith("/api/v1/sessions") || urlStr.endsWith("/api/v1/sessions/")) { + return createSessionResponse("chat-tool"); + } + if (urlStr.includes("/realtime/v1/sessions/") && urlStr.endsWith("/in/append")) { + return appendOkResponse(); + } + // Stitched response now subscribes to `.out` after handover to + // pick up agent-side chunks. Return an empty SSE body that + // closes immediately — this test validates dispatch only, not + // the agent-side resume. + if (/\/realtime\/v1\/sessions\/[^/]+\/out$/.test(urlStr)) { + return new Response(new ReadableStream({ start(c) { c.close(); } }), { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + // Schema-only tool — no execute. The mock model emits a tool-call; + // AI SDK doesn't run it (no execute) and finishes with "tool-calls". + const { tool } = await import("ai"); + const { z } = await import("zod"); + const weatherTool = tool({ + description: "weather", + inputSchema: z.object({ city: z.string() }), + }); + + const handler = chat.headStart({ + agentId: "test-agent", + run: async ({ chat: chatHelper }) => { + return streamText({ + ...chatHelper.toStreamTextOptions({ tools: { weather: weatherTool } }), + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: toolCallStream() }), + }), + }); + }, + }); + + const res = await withApiContext(() => + handler( + makeRequest({ + chatId: "chat-tool", + trigger: "submit-message", + headStartMessages: [ + { id: "m1", role: "user", parts: [{ type: "text", text: "weather in tokyo?" }] }, + ], + }) + ) + ); + + await readSSEBodyToChunks(res); + await new Promise((r) => setTimeout(r, 30)); + + const handoverPost = requests.find( + (r) => + r.url.includes("/realtime/v1/sessions/chat-tool/in/append") && + r.init?.body !== undefined + ); + expect(handoverPost).toBeDefined(); + const body = JSON.parse(handoverPost!.init!.body as string); + expect(body.kind).toBe("handover"); + expect(body.isFinal).toBe(false); // pending tool-calls — agent runs streamText + expect(Array.isArray(body.partialAssistantMessage)).toBe(true); + + // The partial is reshaped into AI SDK's tool-approval round so the + // agent's `streamText` can resume by executing the pending tool-call + // before step 2. Assistant gets a `tool-approval-request` part + // alongside the original `tool-call`; a trailing `tool` message + // carries the `tool-approval-response { approved: true }`. + const assistant = body.partialAssistantMessage.find( + (m: { role: string }) => m.role === "assistant" + ); + expect(assistant).toBeDefined(); + const toolCallPart = assistant.content.find( + (p: { type: string }) => p.type === "tool-call" + ); + expect(toolCallPart).toBeDefined(); + const approvalRequestPart = assistant.content.find( + (p: { type: string }) => p.type === "tool-approval-request" + ); + expect(approvalRequestPart).toBeDefined(); + expect(approvalRequestPart.toolCallId).toBe(toolCallPart.toolCallId); + + const trailingTool = body.partialAssistantMessage[body.partialAssistantMessage.length - 1]; + expect(trailingTool.role).toBe("tool"); + const approvalResponsePart = trailingTool.content.find( + (p: { type: string }) => p.type === "tool-approval-response" + ); + expect(approvalResponsePart).toBeDefined(); + expect(approvalResponsePart.approvalId).toBe(approvalRequestPart.approvalId); + expect(approvalResponsePart.approved).toBe(true); + }); + + it("rejects requests missing chatId", async () => { + global.fetch = vi.fn().mockResolvedValue(new Response("nope", { status: 500 })); + + const handler = chat.headStart({ + agentId: "test-agent", + run: async ({ chat: chatHelper }) => { + return streamText({ + ...chatHelper.toStreamTextOptions(), + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("x") }), + }), + }); + }, + }); + + await expect( + withApiContext(() => + handler( + makeRequest({ + // no chatId + trigger: "submit-message", + messages: [], + }) + ) + ) + ).rejects.toThrow(/chatId/); + }); +}); + +describe("chat.toNodeListener", () => { + /** + * Build a fake Node IncomingMessage that yields a JSON body. + * AsyncIterable so the listener can `for await` over it. + */ + function fakeNodeRequest(opts: { + method?: string; + url?: string; + host?: string; + headers?: Record; + body?: string; + }) { + const bodyBytes = opts.body ? new TextEncoder().encode(opts.body) : undefined; + const headers = { + host: opts.host ?? "example.com", + ...(opts.body ? { "content-type": "application/json" } : {}), + ...(opts.headers ?? {}), + }; + const errorListeners: Array<(e: Error) => void> = []; + return { + method: opts.method ?? "POST", + url: opts.url ?? "/api/chat", + headers, + on(event: string, listener: (e: Error) => void) { + if (event === "error") errorListeners.push(listener); + return this; + }, + async *[Symbol.asyncIterator]() { + if (bodyBytes) yield bodyBytes; + }, + }; + } + + function fakeNodeResponse() { + const writes: Uint8Array[] = []; + let ended = false; + let endChunk: Uint8Array | string | undefined; + const closeListeners: Array<() => void> = []; + const headers: Record = {}; + const obj = { + statusCode: 200, + headersSent: false, + setHeader(name: string, value: string | number | readonly string[]) { + headers[name.toLowerCase()] = value; + }, + write(chunk: Uint8Array | string) { + if (typeof chunk === "string") { + writes.push(new TextEncoder().encode(chunk)); + } else { + writes.push(chunk); + } + obj.headersSent = true; + return true; + }, + end(chunk?: Uint8Array | string) { + ended = true; + endChunk = chunk; + }, + on(event: string, listener: () => void) { + if (event === "close") closeListeners.push(listener); + return obj; + }, + // test helpers + _written() { + const all = [...writes]; + if (typeof endChunk === "string") all.push(new TextEncoder().encode(endChunk)); + else if (endChunk) all.push(endChunk); + let total = 0; + for (const c of all) total += c.length; + const merged = new Uint8Array(total); + let offset = 0; + for (const c of all) { + merged.set(c, offset); + offset += c.length; + } + return new TextDecoder().decode(merged); + }, + _ended: () => ended, + _headers: () => headers, + _close: () => { + for (const l of closeListeners) l(); + }, + }; + return obj; + } + + it("converts the Node request into a Web Request, calls the handler, and forwards the response", async () => { + const seen: { method?: string; url?: string; ct?: string | null; body?: string } = {}; + + const webHandler = async (req: Request): Promise => { + seen.method = req.method; + seen.url = req.url; + seen.ct = req.headers.get("content-type"); + seen.body = await req.text(); + return new Response("ok", { + status: 201, + headers: { "x-test": "1", "content-type": "text/plain" }, + }); + }; + + const listener = chat.toNodeListener(webHandler); + const req = fakeNodeRequest({ body: '{"hello":"world"}' }); + const res = fakeNodeResponse(); + + await listener(req as any, res as any); + + expect(seen.method).toBe("POST"); + expect(seen.url).toBe("http://example.com/api/chat"); + expect(seen.ct).toBe("application/json"); + expect(seen.body).toBe('{"hello":"world"}'); + + expect(res.statusCode).toBe(201); + expect(res._headers()["x-test"]).toBe("1"); + expect(res._written()).toBe("ok"); + expect(res._ended()).toBe(true); + }); + + it("streams the Web Response body to the Node response chunk by chunk (no buffering)", async () => { + const chunkOrder: string[] = []; + const webHandler = async (): Promise => { + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + async start(controller) { + for (const piece of ["one\n", "two\n", "three\n"]) { + chunkOrder.push("emit-" + piece.trim()); + controller.enqueue(encoder.encode(piece)); + await new Promise((r) => setTimeout(r, 5)); + } + controller.close(); + }, + }); + return new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + }; + + const listener = chat.toNodeListener(webHandler); + const req = fakeNodeRequest({}); + const res = fakeNodeResponse(); + await listener(req as any, res as any); + + expect(res._written()).toBe("one\ntwo\nthree\n"); + expect(chunkOrder).toEqual(["emit-one", "emit-two", "emit-three"]); + expect(res._headers()["content-type"]).toBe("text/event-stream"); + }); + + it("propagates client disconnect to the Web handler via AbortSignal", async () => { + let signal: AbortSignal | undefined; + let aborted = false; + + const webHandler = async (req: Request): Promise => { + signal = req.signal; + signal.addEventListener("abort", () => { + aborted = true; + }); + // Return a never-ending stream so the listener stays open until close. + return new Response( + new ReadableStream({ + start() { + // never enqueues + }, + }) + ); + }; + + const listener = chat.toNodeListener(webHandler); + const req = fakeNodeRequest({}); + const res = fakeNodeResponse(); + + // Run listener in background (it'll hang on the never-ending stream). + const pending = listener(req as any, res as any); + + // Wait a tick for the handler to attach the abort listener. + await new Promise((r) => setTimeout(r, 5)); + + res._close(); + expect(aborted).toBe(true); + + // Cleanup: the listener will throw (abort) and we don't care about the result. + await pending.catch(() => {}); + }); + + it("returns 500 with error text if the handler throws before headers are sent", async () => { + const webHandler = async (): Promise => { + throw new Error("boom"); + }; + + const listener = chat.toNodeListener(webHandler); + const req = fakeNodeRequest({}); + const res = fakeNodeResponse(); + await listener(req as any, res as any); + + expect(res.statusCode).toBe(500); + expect(res._written()).toBe("boom"); + }); +}); diff --git a/packages/trigger-sdk/src/v3/chat-server.ts b/packages/trigger-sdk/src/v3/chat-server.ts new file mode 100644 index 00000000000..731ae84410b --- /dev/null +++ b/packages/trigger-sdk/src/v3/chat-server.ts @@ -0,0 +1,893 @@ +/** + * Server-side helpers for the `chat.agent` head-start flow — a + * customer's warm process (Next.js route handler, Express, etc.) + * gets the conversation moving while the heavy chat.agent run boots + * in parallel. Mid-turn, ownership of the durable stream hands over + * to the agent. + * + * The `chat.headStart({ agentId, run })` entry point returns a + * Next.js-style POST handler. Inside the customer's `run` callback + * they call `streamText` themselves, spreading + * `chat.toStreamTextOptions({ tools })` to inherit handover wiring. + * The handler runs `streamText` step 1 in the customer's process + * while the chat.agent run boots in parallel; on `tool-calls` the + * agent run picks up tool execution and continues, on pure-text the + * agent run exits clean without an LLM call. + * + * Two-layer naming: customer-facing surface is "head start" + * (describes the *benefit* — fast first-turn TTFC). The internal + * protocol still uses "handover" (describes the *mechanism* — the + * conversation hands off mid-turn from the warm process to the + * agent). Customers see `chat.headStart`, `HeadStartSession`, etc. + * The wire format and run-loop locals stay on `handover` / + * `handover-prepare` / `handover-skip`. + * + * Cooperative ordering only — handler stops writing to `session.out` + * before sending the `handover` chunk on `session.in`. No S2 fencing. + * + * ⚠️ HARD CONSTRAINT — bundle isolation + * + * This module is the customer-facing boundary for the route handler. + * The whole TTFC win comes from the customer's process being + * lightweight while the heavy agent run boots in parallel. **The + * route-handler bundle must not include heavy tool execute deps**: + * E2B, puppeteer/playwright, native bindings, the trigger SDK + * runtime, turndown, image processing libs, anything that pulls + * weight or pulls `node:` builtins. + * + * "Schema-only" tools must live in a module that imports only `ai` + * (for `tool()`) and `zod`. The agent task module imports those + * schemas and adds execute fns elsewhere — that's where the heavy + * deps live, and it's never reached by the route handler bundle. + * + * Runtime "strip executes" helpers (anything that takes a tool + * catalog with executes and removes them) DO NOT solve this. The + * import chain is resolved at bundle/build time, so importing the + * full catalog drags every dep in regardless of what the SDK does + * with the value at runtime. + * + * IMPORTANT (internal): this module must NOT import from `./ai.ts`. + * `ai.ts` statically imports `agentSkillsRuntime` (which uses `node:` + * builtins unfit for some serverless runtimes) and the heavy task + * runtime. Allowed imports: `./ai-shared.js`, `./chat-client.js`, + * `@trigger.dev/core/v3` (api client), `ai` (types + lightweight + * helpers like `stepCountIs` / `convertToModelMessages`). + */ + +import { ApiClient, SessionStreamInstance, apiClientManager } from "@trigger.dev/core/v3"; +import { + convertToModelMessages, + generateId as generateAssistantMessageId, + stepCountIs, + type ModelMessage, + type StreamTextResult, + type Tool, + type UIMessage, + type UIMessageChunk, +} from "ai"; +import type { ChatInputChunk, ChatTaskWirePayload } from "./ai-shared.js"; + +// --------------------------------------------------------------------------- +// Public types +// --------------------------------------------------------------------------- + +export type HeadStartRunArgs> = { + /** User messages parsed from the incoming request. */ + messages: UIMessage[]; + /** Aborts when the request closes or the SDK times out the handover. */ + signal: AbortSignal; + /** Helper exposing `toStreamTextOptions(...)` and a session escape hatch. */ + chat: HeadStartChatHelper; +}; + +export type HeadStartChatHelper> = { + /** + * Spread into the customer's `streamText` call to inherit handover + * wiring. Returns options for: + * + * - `messages` — converted from the wire payload's UIMessages + * - `tools` — the customer's tool set (typically schema-only — see + * the bundle-isolation note in this module's header) + * - `abortSignal` — combined request-lifecycle + idle timeout + * - `stopWhen` — `stepCountIs(1)`. Step 1 only. The agent run picks + * up tool execution and step 2+ after the handover signal. + * + * Customer adds `model`, `system`, `providerOptions`, etc. on top. + * The customer keeps full control of the `streamText` call shape; + * this helper just hands back the options the SDK needs to own. + * + * The customer COULD override any of these by re-setting them after + * the spread, but doing so for `stopWhen` / `messages` / + * `abortSignal` will break the handover protocol. The intent is + * that customers spread first, then add only their own keys. + */ + toStreamTextOptions = Record>(opts?: { + tools?: TTools; + }): TOpts; + /** Lower-level escape hatch with manual `out` / `in` / dispatch primitives. */ + session: HeadStartSession; +}; + +export type HeadStartSession = { + readonly chatId: string; + /** + * Tees a UIMessage stream into `session.out` for durability/resume, + * fire-and-forget. Returns a passthrough that the caller can use as + * the HTTP response body. + */ + tee( + stream: ReadableStream + ): ReadableStream; + /** + * Awaits `result.finishReason` and dispatches `handover` (with the + * partial assistant ModelMessages) or `handover-skip`. + */ + handoverWhenDone(result: StreamTextResult): Promise; + /** + * Sugar over `tee` + `handoverWhenDone` + standard SSE response. + * Returns a `Response` with `Content-Type: text/event-stream` whose + * body is the teed stream. + */ + handoverResponse(result: StreamTextResult): Response; + /** Manually dispatch the `handover` signal on `session.in`. */ + handover(args: { partialAssistantMessage: ModelMessage[] }): Promise; + /** Manually dispatch the `handover-skip` signal on `session.in`. */ + handoverSkip(): Promise; +}; + +export type HeadStartHandlerOptions> = { + /** The `chat.agent({ id })` of the agent we're handing off to. */ + agentId: string; + /** + * Customer's first-turn implementation. Receives `messages`, + * `signal`, and a `chat` helper. Should call `streamText` with + * `...chat.toStreamTextOptions({ tools })` and return the + * `StreamTextResult`. + */ + run: (args: HeadStartRunArgs) => Promise>; + /** + * Seconds the agent run waits for the handover signal before + * exiting. Defaults to 60. + */ + idleTimeoutInSeconds?: number; +}; + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +export const chat = { + /** + * Returns a Next.js-style POST handler for the chat.agent + * head-start flow. Customer mounts it as + * `export const { POST } = chat.headStart({...})` (or + * `export const POST = chat.headStart({...})`). + * + * Pair with the browser transport's `headStart: "/api/chat"` + * option so the first message of a brand-new chat lands here + * before the agent run boots. + */ + headStart>( + opts: HeadStartHandlerOptions + ): (req: Request) => Promise { + return async (req: Request) => { + const session = await openHandoverSession({ + req, + agentId: opts.agentId, + idleTimeoutInSeconds: opts.idleTimeoutInSeconds, + }); + + const helper: HeadStartChatHelper = { + toStreamTextOptions(spreadOpts) { + return session.buildStreamTextOptions(spreadOpts) as any; + }, + session: session.handle, + }; + + const result = await opts.run({ + messages: session.uiMessages, + signal: session.combinedSignal, + chat: helper, + }); + + return session.handle.handoverResponse(result); + }; + }, + + /** + * Lower-level primitive for power users who want to call + * `streamText` themselves outside the `run` callback shape — custom + * transforms, non-AI-SDK code paths, or manual control over the + * response. Same wiring `chat.headStart` builds on internally. + */ + openSession(opts: { + req: Request; + agentId: string; + idleTimeoutInSeconds?: number; + }): Promise { + return openHandoverSession(opts).then((s) => s.handle); + }, + + /** + * Wrap a Web Fetch handler — `(req: Request) => Promise` — + * as a Node `http` listener — `(req: IncomingMessage, res: ServerResponse) => Promise`. + * + * Use this to mount `chat.headStart` (or any other Web Fetch + * handler) inside Node-only frameworks like Express, Fastify, Koa, + * or raw `node:http`. Web-native frameworks (Next.js App Router, + * Hono, SvelteKit, Remix, Workers, Bun, Deno, etc.) don't need + * this — they pass `Request` objects directly. + * + * Streams the response body chunk-by-chunk to the Node response, + * so the `chat.headStart` SSE chunks reach the browser as they + * arrive (no buffering). Aborts the underlying handler if the + * client closes the connection. + * + * Type-only import of `node:http` types — no runtime dep on `node:http`, + * so this stays safe to bundle into edge / Workers builds (the + * function just won't be called there). + * + * @example + * ```ts + * import express from "express"; + * import { chat } from "@trigger.dev/sdk/chat-server"; + * + * const handler = chat.headStart({ + * agentId: "my-chat", + * run: async ({ chat: helper }) => streamText({ ... }), + * }); + * + * const app = express(); + * app.post("/api/chat", chat.toNodeListener(handler)); + * ``` + */ + toNodeListener, +}; + +// --------------------------------------------------------------------------- +// Internals +// --------------------------------------------------------------------------- + +type InternalSession = { + uiMessages: UIMessage[]; + combinedSignal: AbortSignal; + handle: HeadStartSession; + buildStreamTextOptions(spreadOpts?: { tools?: Record }): Record; +}; + +async function openHandoverSession(opts: { + req: Request; + agentId: string; + idleTimeoutInSeconds?: number; +}): Promise { + const wirePayload = (await opts.req.json()) as ChatTaskWirePayload; + const chatId = wirePayload.chatId; + if (!chatId) { + throw new Error("[chat.handover] request body missing `chatId`"); + } + // Slim wire — head-start ships full history via `headStartMessages` (not + // `message`/`messages`) because the route handler runs on the customer's + // own HTTP endpoint and isn't subject to the 512 KiB `/in/append` cap. + // The full UIMessage[] flows through `wirePayload` into the auto-trigger + // `basePayload` below, where the agent run boot consumes it on first turn. + const uiMessages = (wirePayload.headStartMessages ?? []) as UIMessage[]; + // `convertToModelMessages` is async — resolve once up front so the + // synchronous `toStreamTextOptions` builder can hand back a fully + // formed object. AI SDK's `streamText` validates `messages` as a + // `ModelMessage[]` synchronously and rejects a Promise. + const modelMessages = await convertToModelMessages(uiMessages); + + const apiClient = resolveApiClient(); + const idleTimeoutInSeconds = opts.idleTimeoutInSeconds ?? 60; + + // Create the session and trigger the chat.agent's `handover-prepare` + // run atomically. `createSession` is idempotent on `(env, externalId + // = chatId)` and the auto-triggered run uses `triggerConfig. + // basePayload` as the wire payload — so a single round-trip both + // ensures the session exists and starts the agent booting with the + // right trigger. + // + // Awaited intentionally: subsequent writes to `session.out` (the + // tee from the customer's `streamText` to S2) need the session to + // exist, and the handover signal at end-of-step-1 needs the agent + // run to be there to consume it. The added latency (~one round trip + // to the control plane) is bounded; the agent's compute boot still + // overlaps with LLM TTFB. + const created = await apiClient.createSession({ + type: "chat.agent", + externalId: chatId, + taskIdentifier: opts.agentId, + triggerConfig: { + basePayload: { + ...wirePayload, + chatId, + trigger: "handover-prepare", + idleTimeoutInSeconds, + }, + idleTimeoutInSeconds, + }, + }); + const sessionPublicAccessToken = created.publicAccessToken; + + // Combined abort signal: request lifecycle OR an internal timeout + // mirroring the agent's idle wait so a hung handler doesn't sit + // forever. + const abortController = new AbortController(); + const requestAbort = (opts.req as Request & { signal?: AbortSignal }).signal; + if (requestAbort) { + if (requestAbort.aborted) abortController.abort(); + else requestAbort.addEventListener("abort", () => abortController.abort(), { once: true }); + } + const idleTimer = setTimeout( + () => abortController.abort(new Error("chat.handover: idle timeout")), + idleTimeoutInSeconds * 1000 + ); + + const buildStreamTextOptions = ( + spreadOpts?: { tools?: Record } + ): Record => { + // The customer spreads this object into their `streamText` call + // and then adds `model`, `system`, etc. on top. We set the four + // keys handover correctness depends on: + // + // - `messages`: the wire payload's UIMessages, converted + // (Promise resolved upfront so the spread is synchronous) + // - `tools`: customer's schema-only tool set + // - `stopWhen`: `stepCountIs(1)` — step 1 only. Agent run picks + // up tool execution and step 2+ after the handover signal. + // - `abortSignal`: combined request-lifecycle + idle timeout + // + // The customer's `StreamTextResult` exposes `finishReason` and + // `response.messages` directly, so we don't need to install an + // `onStepFinish` capture hook — we read those off the result in + // `handoverWhenDone`. + return { + messages: modelMessages, + tools: spreadOpts?.tools, + stopWhen: stepCountIs(1), + abortSignal: abortController.signal, + }; + }; + + // Tee a UIMessage stream into session.out via S2 direct-write, + // batched. `SessionStreamInstance` calls `initializeSessionStream` + // once to fetch S2 credentials, then pipes via `StreamsWriterV2`'s + // `BatchTransform` — one S2 append per ~200ms of chunks instead of + // one HTTP round-trip per UIMessageChunk. + let sessionWriter: SessionStreamInstance | null = null; + const tee = (stream: ReadableStream): ReadableStream => { + const [a, b] = stream.tee(); + sessionWriter = new SessionStreamInstance({ + apiClient, + baseUrl: apiClient.baseUrl, + sessionId: chatId, // Sessions are addressable by externalId (chatId). + io: "out", + source: b, + signal: abortController.signal, + }); + return a; + }; + /** Wait for the teed S2 writer to drain. Called before signaling handover. */ + const flushSessionWriter = async (): Promise => { + if (!sessionWriter) return; + try { + await sessionWriter.wait(); + } catch { + // Drop write errors — the customer's response stream is the + // source of truth for what the user sees. Durability/resume + // best-effort. + } + }; + + const handover = async (args: { + partialAssistantMessage: ModelMessage[]; + messageId?: string; + isFinal: boolean; + }) => { + const chunk: ChatInputChunk = { + kind: "handover", + partialAssistantMessage: args.partialAssistantMessage, + messageId: args.messageId, + isFinal: args.isFinal, + }; + await apiClient.appendToSessionStream(chatId, "in", JSON.stringify(chunk)); + }; + + /** + * Sent only on dispatch error (handler aborted before producing a + * `finishReason`). Normal pure-text and tool-call finishes go + * through `handover()` with the appropriate `isFinal` flag. + */ + const handoverSkip = async () => { + const chunk: ChatInputChunk = { kind: "handover-skip" }; + await apiClient.appendToSessionStream(chatId, "in", JSON.stringify(chunk)); + }; + + // A stable assistant messageId for this turn. The customer's + // `toUIMessageStream` is configured to emit its `start` chunk with + // this id, the handover signal carries it to the agent, and the + // agent's post-handover `toUIMessageStream` reuses it — so all + // chunks (customer's step 1 + agent's step 2) merge into one + // assistant message on the browser side. + const turnMessageId = generateAssistantMessageId(); + + // Set by `handoverWhenDone` after it observes `result.finishReason` + // and dispatches the handover decision. The stitched response stream + // awaits this to know whether to close (skip) or pull more chunks + // from session.out (handover). + type HandoverDecision = { kind: "handover" | "handover-skip" }; + let resolveDecision!: (decision: HandoverDecision) => void; + const decisionPromise = new Promise((resolve) => { + resolveDecision = resolve; + }); + + const handoverWhenDone = async (result: StreamTextResult) => { + try { + // `result.finishReason` is a Promise on the AI SDK + // result. Wait for the stream to settle, then dispatch. + const finishReason = await result.finishReason; + + // Drain the S2 tee so any in-flight handler writes (last + // `tool-input-available` parts, the synthetic `finish-step` for + // pure-text) are visible before the agent reads from session.out + // / session.in. Cooperative ordering — agent doesn't read past + // these unless we've finished writing them. + await flushSessionWriter(); + + const responseMessages = (await result.response).messages as ModelMessage[]; + + if (finishReason === "tool-calls") { + // Reshape pending tool-calls into AI SDK's tool-approval round + // so the agent's `streamText` resumes by executing them + // before the step-2 LLM call. + const reshaped = reshapeForHandoverResume(responseMessages); + await handover({ + partialAssistantMessage: reshaped, + messageId: turnMessageId, + isFinal: false, + }); + } else { + // Pure-text (or any non-tool-calls) finish — customer's step 1 + // IS the final response. The agent runs the turn-loop hooks + // (`onChatStart`, `onTurnStart`, `onTurnComplete`, etc.) using + // this partial as the response, but skips the LLM call. That + // way persistence (`onTurnComplete` writing to DB), self- + // review, and any post-turn work all fire normally. + await handover({ + partialAssistantMessage: responseMessages, + messageId: turnMessageId, + isFinal: true, + }); + } + resolveDecision({ kind: "handover" }); + } catch (err) { + // Dispatch failed before we could send the handover signal. + // Tell the agent to exit clean (no hooks fire) and close the + // response stream so it doesn't hang waiting for agent chunks. + resolveDecision({ kind: "handover-skip" }); + try { + await handoverSkip(); + } catch { + // best-effort + } + throw err; + } + }; + + /** + * Build a single ReadableStream that: + * 1. Forwards the customer's `streamText` chunks (step 1) directly + * to the response — same low-latency path as before. + * 2. After step 1 ends and the dispatch decision lands: + * - `handover-skip`: closes the response immediately. The agent + * run exits without writing more chunks. + * - `handover`: subscribes to `session.out` from the sequence + * ID where the customer's tee left off, forwarding the agent + * run's chunks (tool-output-available, step 2 LLM text, + * `finish-step`, etc.) until `trigger:turn-complete`. + * + * The browser sees one continuous SSE response per first turn, just + * like a normal `streamText` would produce. + */ + const stitchHandoverStream = ( + customerBranch: ReadableStream + ): ReadableStream => { + return new ReadableStream({ + async start(controller) { + try { + // Phase 1: forward customer's chunks. + const reader = customerBranch.getReader(); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + controller.enqueue(value); + } + } finally { + reader.releaseLock(); + } + + // Phase 2a: wait for handoverWhenDone to decide. + const decision = await decisionPromise; + if (decision.kind === "handover-skip") { + controller.close(); + return; + } + + // Phase 2b: agent is taking over. Resume from session.out + // starting AFTER the customer tee's last write, so we don't + // re-emit chunks the browser already saw. + const writeResult = sessionWriter + ? await sessionWriter.wait().catch(() => undefined) + : undefined; + const customerLastEventId = writeResult?.lastEventId; + + // Capture the latest S2 event id seen on session.out via + // `onPart`. After the stream closes we emit it to the + // browser as a `trigger:session-state` control chunk so the + // transport can hydrate `state.lastEventId` for turn 2's + // subscribe — without it, turn 2 reads session.out from the + // start and replays turn 1 to the user. + let latestEventId: string | undefined; + const agentStream = await apiClient.subscribeToSessionStream( + chatId, + "out", + { + ...(customerLastEventId != null + ? { lastEventId: customerLastEventId } + : {}), + signal: abortController.signal, + onPart: (part) => { + if (part.id) latestEventId = part.id; + }, + } + ); + + for await (const chunk of agentStream) { + controller.enqueue(chunk); + // The agent's run-loop emits `trigger:turn-complete` when + // the turn finishes. That's our cue to close — anything + // after is the next turn (which goes via the direct + // `session.in`/`session.out` path, not this endpoint). + if ( + chunk && + typeof chunk === "object" && + (chunk as { type?: unknown }).type === "trigger:turn-complete" + ) { + break; + } + } + + // Final control chunk: hand the browser transport the + // `lastEventId` it should use for the next turn's + // session.out subscribe. Filtered out before reaching the + // AI SDK on the browser side. + if (latestEventId != null) { + controller.enqueue({ + type: "trigger:session-state", + lastEventId: latestEventId, + } as unknown as UIMessageChunk); + } + controller.close(); + } catch (err) { + controller.error(err); + } + }, + cancel() { + // Browser closed the connection. Trigger the abort so any + // pending session.out subscription stops too. + abortController.abort(); + }, + }); + }; + + const handoverResponse = (result: StreamTextResult): Response => { + // `generateMessageId` makes the customer's `start` chunk carry + // `turnMessageId`, so the browser-side AI SDK keys the assistant + // message by it. The agent's post-handover stream emits chunks + // with the same id (passed via the handover signal) — both sides + // merge into one message on the browser. + const teed = tee( + result.toUIMessageStream({ + generateMessageId: () => turnMessageId, + }) + ); + // `handoverWhenDone` re-throws on dispatch failure for visibility, + // but the recovery (resolveDecision + handoverSkip) has already run + // by then and `stitchHandoverStream` closes the response cleanly via + // `decisionPromise`. The user-facing path is fine; we only suppress + // the unhandled-rejection so processes started with + // `--unhandled-rejections=throw` don't crash on what is effectively + // a logged failure with no further action to take. + void handoverWhenDone(result) + .finally(() => clearTimeout(idleTimer)) + .catch(() => {}); + + const stitched = stitchHandoverStream(teed); + + // Encode UIMessageChunks as SSE for the AI SDK transport on the + // browser. AI SDK's `toUIMessageStreamResponse()` does this same + // thing internally; replicate the format here so we don't have + // to bridge through the SDK's response helper. + const encoder = new TextEncoder(); + const sseStream = stitched.pipeThrough( + new TransformStream({ + transform(chunk, controller) { + controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`)); + }, + }) + ); + + return new Response(sseStream, { + headers: { + "Content-Type": "text/event-stream", + "X-Vercel-AI-UI-Message-Stream": "v1", + "Cache-Control": "no-cache, no-transform", + Connection: "keep-alive", + // Browser transport reads these to hydrate session state + // for subsequent (non-handover) turns. Once the browser has + // the PAT it talks directly to `session.in` / `session.out` + // without going back through the handler. + "X-Trigger-Chat-Id": chatId, + "X-Trigger-Chat-Access-Token": sessionPublicAccessToken, + }, + }); + }; + + const handle: HeadStartSession = { + chatId, + tee, + handoverWhenDone, + handoverResponse, + handover, + handoverSkip, + }; + + return { + uiMessages, + combinedSignal: abortController.signal, + handle, + buildStreamTextOptions, + }; +} + +function resolveApiClient(): ApiClient { + // Reuse the SDK's standard apiClientManager so customers configure + // base URL + secret key the same way as for `tasks.trigger(...)`. + const client = apiClientManager.clientOrThrow(); + return client; +} + +// --------------------------------------------------------------------------- +// Node `http` adapter +// --------------------------------------------------------------------------- + +// Minimal Node http types we use. Avoids a `node:http` type import so the +// file stays lint-clean on non-Node TS projects (the docs example handlers +// might typecheck under workers / deno configs that lack `node:` types). +interface NodeIncomingHeaders { + [k: string]: string | string[] | undefined; +} +interface NodeIncomingMessage extends AsyncIterable { + readonly url?: string; + readonly method?: string; + readonly headers: NodeIncomingHeaders; + on(event: "error", listener: (err: Error) => void): unknown; +} +interface NodeServerResponse { + statusCode: number; + headersSent: boolean; + setHeader(name: string, value: string | number | readonly string[]): unknown; + write(chunk: Uint8Array | string): boolean; + end(chunk?: Uint8Array | string): unknown; + on(event: "close" | "error", listener: () => void): unknown; +} + +/** @internal — exposed via `chat.toNodeListener`. */ +function toNodeListener( + webHandler: (req: Request) => Promise +): (req: NodeIncomingMessage, res: NodeServerResponse) => Promise { + return async function nodeListener(req, res) { + const abort = new AbortController(); + res.on("close", () => abort.abort()); + + try { + const url = `http://${req.headers.host ?? "localhost"}${req.url ?? "/"}`; + const method = req.method ?? "GET"; + const hasBody = method !== "GET" && method !== "HEAD"; + + // Read full body upfront. Chat wire payloads are small (sub-KB + // typically) so accumulating avoids the duplex-stream ceremony + // some Node versions need for streaming request bodies into + // a Web Request. + let body: ArrayBuffer | undefined; + if (hasBody) { + const chunks: Uint8Array[] = []; + for await (const chunk of req as AsyncIterable) { + chunks.push(chunk); + } + if (chunks.length > 0) { + let total = 0; + for (const c of chunks) total += c.length; + const merged = new Uint8Array(total); + let offset = 0; + for (const c of chunks) { + merged.set(c, offset); + offset += c.length; + } + body = merged.buffer.slice(merged.byteOffset, merged.byteOffset + merged.byteLength); + } + } + + // Flatten Node header values: arrays → comma-joined (per RFC 7230 §3.2.2). + const webHeaders = new Headers(); + for (const [name, value] of Object.entries(req.headers)) { + if (value == null) continue; + if (Array.isArray(value)) { + for (const v of value) webHeaders.append(name, v); + } else { + webHeaders.set(name, value); + } + } + + const webReq = new Request(url, { + method, + headers: webHeaders, + body, + signal: abort.signal, + }); + + const webRes = await webHandler(webReq); + + res.statusCode = webRes.status; + // `Headers.forEach` exposes the value comma-joined for multi-valued + // headers, which `setHeader` accepts. Set-Cookie is handled separately + // via `getSetCookie()` to preserve multiple values. + webRes.headers.forEach((value, key) => { + if (key.toLowerCase() === "set-cookie") return; + res.setHeader(key, value); + }); + const setCookies = + typeof (webRes.headers as Headers & { getSetCookie?: () => string[] }).getSetCookie === "function" + ? (webRes.headers as Headers & { getSetCookie: () => string[] }).getSetCookie() + : []; + if (setCookies.length > 0) { + res.setHeader("set-cookie", setCookies); + } + + if (!webRes.body) { + res.end(); + return; + } + + // Pipe the Web Response body to the Node response. On client + // disconnect (`abort.signal`), cancel the reader so a pending + // `read()` rejects and we exit the loop instead of blocking on + // a stream that will never produce more chunks. + const reader = webRes.body.getReader(); + const onAbort = () => { + reader.cancel(abort.signal.reason).catch(() => {}); + }; + if (abort.signal.aborted) onAbort(); + else abort.signal.addEventListener("abort", onAbort, { once: true }); + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + res.write(value); + } + } catch { + // Reader was cancelled (client disconnect). Silently end. + } finally { + abort.signal.removeEventListener("abort", onAbort); + } + res.end(); + } catch (err) { + if (!res.headersSent) { + res.statusCode = 500; + res.setHeader("content-type", "text/plain; charset=utf-8"); + res.end(err instanceof Error ? err.message : "Internal error"); + } else { + res.end(); + } + } + }; +} + +/** + * Reshape a step-1 partial so the agent's `streamText` resumes by + * executing pending tool-calls before the next LLM call. + * + * When the customer's handler runs `streamText` with schema-only tools + * (no `execute` fns) and `stopWhen: stepCountIs(1)`, the LLM emits + * tool-calls but AI SDK can't execute them — the partial we ship is + * `[{ assistant: text + tool-call }]`. Splicing that as-is onto the + * agent's accumulator and calling `streamText` throws + * `MissingToolResultsError` synchronously inside + * `convertToLanguageModelPrompt`. + * + * AI SDK's documented escape hatch for "external party decides what + * to do with a tool-call, then SDK executes" is the tool-approval + * round. By appending a `tool-approval-request` part to the assistant + * message and a trailing `tool` message with a matching + * `tool-approval-response { approved: true }`, AI SDK: + * 1. Suppresses `MissingToolResultsError` for approved tool-calls + * (`convert-to-language-model-prompt.ts:135-144`). + * 2. Hits its initial-tool-execution branch + * (`stream-text.ts:1342-1486`) on the next `streamText` call, + * runs the agent-side `execute` fns, and synthesizes + * `tool-result` parts before the step-2 LLM call. + * + * If the customer's tools already had `execute` fns (rare for the + * handover use case but valid), the partial already contains a + * `tool-result` per tool-call — we leave those alone and only inject + * approvals for genuinely-pending calls. + * + * `collectToolApprovals` only scans the LAST message + * (`collect-tool-approvals.ts:30-37`), so the synthesized tool message + * must end up at the tail of the partial. The agent's run-loop + * splices the partial onto the end of the accumulator, which keeps + * this invariant. + */ +function reshapeForHandoverResume(responseMessages: ModelMessage[]): ModelMessage[] { + // First pass: gather the set of tool-call IDs that already have a + // matching tool-result. Those are "complete" — leave them alone. + const completedToolCallIds = new Set(); + for (const message of responseMessages) { + if (message.role !== "tool" || typeof message.content === "string") continue; + for (const part of message.content as Array<{ type: string; toolCallId?: string }>) { + if (part.type === "tool-result" && part.toolCallId) { + completedToolCallIds.add(part.toolCallId); + } + } + } + + // Second pass: clone the messages, appending a tool-approval-request + // alongside each pending tool-call. Collect the matching responses. + const approvalResponses: Array<{ + type: "tool-approval-response"; + approvalId: string; + approved: true; + }> = []; + let approvalCounter = 0; + + const reshaped: ModelMessage[] = responseMessages.map((message) => { + if (message.role !== "assistant" || typeof message.content === "string") { + return message; + } + const newContent: typeof message.content = [...message.content]; + for (const part of message.content as Array<{ + type: string; + toolCallId?: string; + }>) { + if ( + part.type === "tool-call" && + part.toolCallId && + !completedToolCallIds.has(part.toolCallId) + ) { + const approvalId = `handover-approval-${++approvalCounter}`; + newContent.push({ + type: "tool-approval-request", + approvalId, + toolCallId: part.toolCallId, + } as never); + approvalResponses.push({ + type: "tool-approval-response", + approvalId, + approved: true, + }); + } + } + return { ...message, content: newContent } as ModelMessage; + }); + + if (approvalResponses.length > 0) { + reshaped.push({ + role: "tool", + content: approvalResponses as never, + } as ModelMessage); + } + + return reshaped; +} diff --git a/packages/trigger-sdk/src/v3/deployments.ts b/packages/trigger-sdk/src/v3/deployments.ts new file mode 100644 index 00000000000..b6a334b203e --- /dev/null +++ b/packages/trigger-sdk/src/v3/deployments.ts @@ -0,0 +1,56 @@ +import type { + ApiRequestOptions, + RetrieveCurrentDeploymentResponseBody, + ApiDeploymentListOptions, + ApiDeploymentListResponseItem, +} from "@trigger.dev/core/v3"; +import { + apiClientManager, + CursorPagePromise, + isRequestOptions, + mergeRequestOptions, +} from "@trigger.dev/core/v3"; + +export type { RetrieveCurrentDeploymentResponseBody, ApiDeploymentListResponseItem }; + +export const deployments = { + retrieveCurrent: retrieveCurrentDeployment, + list: listDeployments, +}; + +/** + * Retrieve the currently promoted deployment for this environment. + * + * Use inside a task to check whether a newer version has been deployed: + * + * ```ts + * import { deployments } from "@trigger.dev/sdk"; + * + * const current = await deployments.retrieveCurrent(); + * if (current.version !== ctx.run.version) { + * // A newer version is promoted + * } + * ``` + */ +function retrieveCurrentDeployment( + requestOptions?: ApiRequestOptions +): Promise { + const apiClient = apiClientManager.clientOrThrow(); + return apiClient.retrieveCurrentDeployment(requestOptions); +} + +/** + * List deployments for the current environment. + */ +function listDeployments( + options?: ApiDeploymentListOptions, + requestOptions?: ApiRequestOptions +): CursorPagePromise { + const apiClient = apiClientManager.clientOrThrow(); + + if (isRequestOptions(options)) { + return apiClient.listDeployments(undefined, options); + } + + return apiClient.listDeployments(options, requestOptions); +} diff --git a/packages/trigger-sdk/src/v3/index.ts b/packages/trigger-sdk/src/v3/index.ts index 21ffa142871..5e169cbb8d6 100644 --- a/packages/trigger-sdk/src/v3/index.ts +++ b/packages/trigger-sdk/src/v3/index.ts @@ -17,14 +17,15 @@ export * from "./otel.js"; export * from "./schemas.js"; export * from "./heartbeats.js"; export * from "./streams.js"; +export * from "./sessions.js"; export * from "./query.js"; export type { Context }; import type { Context } from "./shared.js"; -import type { ApiClientConfiguration } from "@trigger.dev/core/v3"; +import type { ApiClientConfiguration, TaskRunContext } from "@trigger.dev/core/v3"; -export type { ApiClientConfiguration }; +export type { ApiClientConfiguration, TaskRunContext }; export { ApiError, @@ -39,6 +40,8 @@ export { AbortTaskRunError, OutOfMemoryError, CompleteTaskWithOutput, + ChatChunkTooLargeError, + isChatChunkTooLargeError, logger, type LogLevel, } from "@trigger.dev/core/v3"; @@ -54,9 +57,15 @@ export { type AnyRetrieveRunResult, } from "./runs.js"; export * as schedules from "./schedules/index.js"; +export { + deployments, + type RetrieveCurrentDeploymentResponseBody, + type ApiDeploymentListResponseItem, +} from "./deployments.js"; export * as envvars from "./envvars.js"; export * as queues from "./queues.js"; export type { ImportEnvironmentVariablesParams } from "./envvars.js"; export { configure, auth } from "./auth.js"; export * as prompts from "./prompts.js"; +export * as skills from "./skills.js"; diff --git a/packages/trigger-sdk/src/v3/runs.ts b/packages/trigger-sdk/src/v3/runs.ts index 7081c448d75..88e6d2b701c 100644 --- a/packages/trigger-sdk/src/v3/runs.ts +++ b/packages/trigger-sdk/src/v3/runs.ts @@ -358,6 +358,14 @@ export type SubscribeToRunOptions = { * ``` */ skipColumns?: RealtimeRunSkipColumns; + + /** + * An AbortSignal to cancel the subscription. + * + * When the signal is aborted, the underlying SSE connection is closed + * and the async iterator completes. + */ + signal?: AbortSignal; }; /** @@ -403,6 +411,7 @@ function subscribeToRun( closeOnComplete: typeof options?.stopOnCompletion === "boolean" ? options.stopOnCompletion : true, skipColumns: options?.skipColumns, + signal: options?.signal, }); } diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts new file mode 100644 index 00000000000..9bc2c3e3272 --- /dev/null +++ b/packages/trigger-sdk/src/v3/sessions.ts @@ -0,0 +1,743 @@ +import type { + ApiPromise, + ApiRequestOptions, + AsyncIterableStream, + CloseSessionRequestBody, + CreatedSessionResponseBody, + CreateSessionRequestBody, + InputStreamOnceOptions, + InputStreamOnceResult, + InputStreamWaitOptions, + InputStreamWaitWithIdleTimeoutOptions, + ListSessionsOptions, + ListedSessionItem, + PipeStreamOptions, + PipeStreamResult, + RetrieveSessionResponseBody, + UpdateSessionRequestBody, + WriterStreamOptions, +} from "@trigger.dev/core/v3"; +import { + CursorPagePromise, + InputStreamOncePromise, + ManualWaitpointPromise, + SemanticInternalAttributes, + SessionStreamInstance, + WaitpointTimeoutError, + accessoryAttributes, + apiClientManager, + ensureReadableStream, + mergeRequestOptions, + runtime, + sessionStreams, + taskContext, +} from "@trigger.dev/core/v3"; +import { conditionallyImportAndParsePacket } from "@trigger.dev/core/v3/utils/ioSerialization"; +import { SpanStatusCode } from "@opentelemetry/api"; +import { tracer } from "./tracer.js"; + +export type { + CreatedSessionResponseBody, + CreateSessionRequestBody, + CloseSessionRequestBody, + ListSessionsOptions, + ListedSessionItem, + RetrieveSessionResponseBody, + UpdateSessionRequestBody, +}; + +export const sessions = { + start: startSession, + retrieve: retrieveSession, + update: updateSession, + close: closeSession, + list: listSessions, + open, +}; + +// Test hook: lets `@trigger.dev/sdk/ai/test` replace `sessions.open()` with +// an in-memory handle so unit tests don't hit the network. Not part of the +// public API — only `mockChatAgent` installs it. +type SessionOpenImpl = (sessionIdOrExternalId: string) => SessionHandle; +let sessionOpenImpl: SessionOpenImpl | undefined; + +export function __setSessionOpenImplForTests(impl: SessionOpenImpl | undefined): void { + sessionOpenImpl = impl; +} + +// Test hook for `sessions.start()`. Sessions are task-bound and the +// `start` call atomically creates the row + triggers the first run on +// the server; in unit tests there's no live API to hit, so a fixture +// implementation can be installed via this setter. +type SessionStartImpl = ( + body: CreateSessionRequestBody +) => Promise | CreatedSessionResponseBody; +let sessionStartImpl: SessionStartImpl | undefined; + +export function __setSessionStartImplForTests(impl: SessionStartImpl | undefined): void { + sessionStartImpl = impl; +} + +/** + * Start a {@link Session} — a durable, task-bound, bidirectional I/O + * primitive. The server creates the row (idempotent on `externalId`) + * and triggers the first run from `triggerConfig` in one round-trip. + * Returns the new run's id and a session-scoped public access token + * for browser-side use against `.in/append`, `.out` SSE, and + * `end-and-continue`. + * + * If a session with the same `(env, externalId)` already exists, + * returns the existing row plus the live (or freshly re-triggered) run. + * Two browser tabs of the same chat converge to one session. + */ +function startSession( + body: CreateSessionRequestBody, + requestOptions?: ApiRequestOptions +): ApiPromise { + if (sessionStartImpl) { + const result = sessionStartImpl(body); + return Promise.resolve(result) as ApiPromise; + } + + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "sessions.start()", + icon: "sessions", + attributes: sessionAttributes(body.externalId ?? body.type, { + type: body.type, + ...(body.externalId ? { externalId: body.externalId } : {}), + }), + }, + requestOptions + ); + + return apiClient.createSession(body, $requestOptions); +} + +/** + * Retrieve a Session by `friendlyId` (`session_*`) or user-supplied + * `externalId`. The server disambiguates via the `session_` prefix. + */ +function retrieveSession( + sessionIdOrExternalId: string, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "sessions.retrieve()", + icon: "sessions", + attributes: sessionAttributes(sessionIdOrExternalId), + }, + requestOptions + ); + + return apiClient.retrieveSession(sessionIdOrExternalId, $requestOptions); +} + +/** Update mutable fields on a Session (tags, metadata, externalId). */ +function updateSession( + sessionIdOrExternalId: string, + body: UpdateSessionRequestBody, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "sessions.update()", + icon: "sessions", + attributes: sessionAttributes(sessionIdOrExternalId), + }, + requestOptions + ); + + return apiClient.updateSession(sessionIdOrExternalId, body, $requestOptions); +} + +/** Mark a Session as closed (terminal, idempotent). */ +function closeSession( + sessionIdOrExternalId: string, + body?: CloseSessionRequestBody, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "sessions.close()", + icon: "sessions", + attributes: sessionAttributes(sessionIdOrExternalId, { + ...(body?.reason ? { reason: body.reason } : {}), + }), + }, + requestOptions + ); + + return apiClient.closeSession(sessionIdOrExternalId, body, $requestOptions); +} + +/** + * List Sessions in the current environment with filters + cursor pagination. + * Returns a {@link CursorPagePromise} so callers can iterate pages with + * `for await`. + */ +function listSessions( + options?: ListSessionsOptions, + requestOptions?: ApiRequestOptions +): CursorPagePromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "sessions.list()", + icon: "sessions", + attributes: { + ...(options?.type ? { type: toAttr(options.type) } : {}), + ...(options?.tag ? { tag: toAttr(options.tag) } : {}), + ...(options?.status ? { status: toAttr(options.status) } : {}), + ...(options?.externalId ? { externalId: options.externalId } : {}), + }, + }, + requestOptions + ); + + return apiClient.listSessions(options, $requestOptions); +} + +/** + * Open a lightweight handle to a Session's realtime channels. Does not + * perform a network call on its own — each channel method hits the + * corresponding realtime endpoint. + */ +function open(sessionIdOrExternalId: string): SessionHandle { + if (sessionOpenImpl) return sessionOpenImpl(sessionIdOrExternalId); + return new SessionHandle(sessionIdOrExternalId); +} + +export class SessionHandle { + /** + * Producer-to-consumer channel: the task writes records; external + * clients read them. Mirrors `streams.define` — `append` / `pipe` / + * `writer` / `read`. + */ + public readonly out: SessionOutputChannel; + + /** + * Consumer-to-producer channel: external clients call `.send()`; the + * task consumes via `.on` / `.once` / `.peek` / `.wait` / + * `.waitWithIdleTimeout`. Mirrors `streams.input` but keyed on the + * session so a conversation can survive across run boundaries. + */ + public readonly in: SessionInputChannel; + + constructor( + public readonly id: string, + overrides?: { in?: SessionInputChannel; out?: SessionOutputChannel } + ) { + this.out = overrides?.out ?? new SessionOutputChannel(id); + this.in = overrides?.in ?? new SessionInputChannel(id); + } +} + +/** + * Options accepted by {@link SessionOutputChannel.pipe}. Session-scoped, + * so it omits the `target` field (self/parent/root/runId) that run-scoped + * {@link PipeStreamOptions} uses — the session is the target. + */ +export type SessionPipeStreamOptions = Omit; + +/** + * The `.out` side of a Session's bidirectional channel pair. Mirrors the + * consume-side of {@link streams.define}: `pipe` / `writer` / `append` + * for the task to produce records, `read` for external clients to + * consume via SSE. S2 credentials for direct writes are fetched + * internally by `pipe`/`writer` — there's no public `initialize()`. + */ +export class SessionOutputChannel { + constructor(public readonly sessionId: string) {} + + /** + * Append a single record. Routes through {@link writer} internally so + * subscribers receive the same parsed-object shape as multi-record + * writes — the server-side append endpoint wraps the body in a string, + * which would give SSE consumers a JSON-string instead of an object. + * Mirrors how `streams.define.append` delegates to `streams.writer`. + */ + async append(value: T, options?: SessionPipeStreamOptions): Promise { + const { waitUntilComplete } = this.writer({ + ...options, + spanName: "sessions.append()", + execute: ({ write }) => { + write(value); + }, + }); + await waitUntilComplete(); + } + + /** + * Pipe an `AsyncIterable` / `ReadableStream` directly to S2. Fetches + * session S2 credentials internally and streams through + * {@link SessionStreamInstance}. Parallel to {@link streams.pipe} but + * session-scoped — no `target` option because the session is the target. + */ + pipe( + value: AsyncIterable | ReadableStream, + options?: SessionPipeStreamOptions + ): PipeStreamResult { + return this.#pipeInternal(value, options, "sessions.pipe()"); + } + + /** + * Mirror of {@link streams.writer}: runs `execute({ write, merge })` + * against an in-memory queue whose records are piped to S2. Returns + * `{ stream, waitUntilComplete }` so callers can observe the local + * stream and await completion. Span is collapsible via `options.spanName` + * / `options.collapsed`. + */ + writer(options: WriterStreamOptions): PipeStreamResult { + let controller!: ReadableStreamDefaultController; + const ongoingStreamPromises: Promise[] = []; + + const stream = new ReadableStream({ + start(controllerArg) { + controller = controllerArg; + }, + }); + + const safeEnqueue = (data: T) => { + try { + controller.enqueue(data); + } catch { + // Suppress errors when the stream has been closed. + } + }; + + try { + const result = options.execute({ + write(part) { + safeEnqueue(part); + }, + merge(streamArg) { + ongoingStreamPromises.push( + (async () => { + const reader = streamArg.getReader(); + while (true) { + const { done, value } = await reader.read(); + if (done) break; + safeEnqueue(value); + } + })().catch((error) => { + console.error(error); + }) + ); + }, + }); + + if (result) { + ongoingStreamPromises.push( + result.catch((error) => { + console.error(error); + }) + ); + } + } catch (error) { + console.error(error); + } + + const waitForStreams: Promise = new Promise((resolve, reject) => { + (async () => { + while (ongoingStreamPromises.length > 0) { + await ongoingStreamPromises.shift(); + } + resolve(); + })().catch(reject); + }); + + waitForStreams.finally(() => { + try { + controller.close(); + } catch { + // Already closed. + } + }); + + return this.#pipeInternal(stream, options, options.spanName ?? "sessions.writer()"); + } + + /** + * Subscribe to SSE records on `.out`. Returns an async-iterable stream — + * auto-retry, Last-Event-ID resume, and abort propagation come from the + * shared {@link SSEStreamSubscription} plumbing used by run-scoped + * realtime streams. + */ + async read( + options?: SessionSubscribeOptions + ): Promise> { + const apiClient = apiClientManager.clientOrThrow(); + + return apiClient.subscribeToSessionStream(this.sessionId, "out", { + signal: options?.signal, + timeoutInSeconds: options?.timeoutInSeconds, + lastEventId: + options?.lastEventId != null ? String(options.lastEventId) : undefined, + onPart: options?.onPart, + onComplete: options?.onComplete, + onError: options?.onError, + }); + } + + #pipeInternal( + value: AsyncIterable | ReadableStream, + options: SessionPipeStreamOptions | undefined, + spanName: string + ): PipeStreamResult { + const apiClient = apiClientManager.clientOrThrow(); + const collapsed = (options as WriterStreamOptions | undefined)?.collapsed; + + const span = tracer.startSpan(spanName, { + attributes: { + session: this.sessionId, + io: "out", + [SemanticInternalAttributes.ENTITY_TYPE]: "session-stream", + [SemanticInternalAttributes.ENTITY_ID]: `${this.sessionId}:out`, + [SemanticInternalAttributes.STYLE_ICON]: "sessions", + ...(collapsed ? { [SemanticInternalAttributes.COLLAPSED]: true } : {}), + ...accessoryAttributes({ + items: [{ text: `${this.sessionId}.out`, variant: "normal" }], + style: "codepath", + }), + }, + }); + + const readableStreamSource = ensureReadableStream(value); + + const abortController = new AbortController(); + const combinedSignal = options?.signal + ? AbortSignal.any?.([options.signal, abortController.signal]) ?? abortController.signal + : abortController.signal; + + try { + const instance = new SessionStreamInstance({ + apiClient, + baseUrl: apiClientManager.baseURL ?? "", + sessionId: this.sessionId, + io: "out", + source: readableStreamSource, + signal: combinedSignal, + requestOptions: options?.requestOptions, + }); + + instance.wait().finally(() => { + span.end(); + }); + + return { + stream: instance.stream, + waitUntilComplete: async () => { + return instance.wait(); + }, + }; + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + span.end(); + throw error; + } + + if (error instanceof Error || typeof error === "string") { + span.recordException(error); + } else { + span.recordException(String(error)); + } + + span.setStatus({ code: SpanStatusCode.ERROR }); + span.end(); + + throw error; + } + } +} + +/** + * The `.in` side of a Session's bidirectional channel pair. Mirrors + * {@link streams.input} — consumer-side primitives for the task + * (`on`/`once`/`peek`/`wait`/`waitWithIdleTimeout`) plus `send` for + * external clients. Keyed on the session rather than the run so a + * conversation can survive across run boundaries. + */ +export class SessionInputChannel { + constructor(public readonly sessionId: string) {} + + /** + * Send a single record to the channel. Called by external clients + * (browser, server action, another task) producing input for the run. + * Matches {@link streams.input.send} but session-scoped — the session + * is the address, no `runId` required. + */ + async send(value: unknown, requestOptions?: ApiRequestOptions): Promise { + const apiClient = apiClientManager.clientOrThrow(); + const body = typeof value === "string" ? value : JSON.stringify(value); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: `sessions.open(${this.sessionId}).in.send()`, + icon: "sessions", + attributes: sessionAttributes(this.sessionId, { io: "in" }), + }, + requestOptions + ); + + await apiClient.appendToSessionStream(this.sessionId, "in", body, $requestOptions); + } + + /** + * Register a handler that fires for every record landing on `.in`. + * Handlers are flushed with any buffered records on attach and cleaned + * up automatically when the task run completes. Returns `{ off }` to + * unsubscribe early. + */ + on(handler: (data: T) => void | Promise): { off: () => void } { + return sessionStreams.on( + this.sessionId, + "in", + handler as (data: unknown) => void | Promise + ); + } + + /** + * Wait for the next record on `.in` without suspending the run. + * Returns `{ ok: true, output }` on arrival or `{ ok: false, error }` + * when the timeout fires. Chain `.unwrap()` to get the data directly. + */ + once(options?: InputStreamOnceOptions): InputStreamOncePromise { + const ctx = taskContext.ctx; + const runId = ctx?.run.id; + + const innerPromise = sessionStreams.once(this.sessionId, "in", options); + + return new InputStreamOncePromise((resolve, reject) => { + tracer + .startActiveSpan( + options?.spanName ?? `sessions.open(${this.sessionId}).in.once()`, + async () => { + const result = await innerPromise; + resolve(result as InputStreamOnceResult); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "sessions", + [SemanticInternalAttributes.ENTITY_TYPE]: "session-stream", + ...(runId + ? { [SemanticInternalAttributes.ENTITY_ID]: `${runId}:${this.sessionId}:in` } + : {}), + session: this.sessionId, + io: "in", + ...accessoryAttributes({ + items: [{ text: `${this.sessionId}.in`, variant: "normal" }], + style: "codepath", + }), + }, + } + ) + .catch(reject); + }); + } + + /** Non-blocking peek at the head of the `.in` buffer. */ + peek(): T | undefined { + return sessionStreams.peek(this.sessionId, "in") as T | undefined; + } + + /** + * Suspend the current run until the next record arrives on `.in`. + * Unlike {@link once}, `wait()` frees compute while blocked — the + * run-engine waitpoint holds the run until the session append handler + * fires it. Only callable from inside `task.run()`. + */ + wait(options?: InputStreamWaitOptions): ManualWaitpointPromise { + return new ManualWaitpointPromise(async (resolve, reject) => { + try { + const ctx = taskContext.ctx; + + if (!ctx) { + throw new Error("session.in.wait() can only be used from inside a task.run()"); + } + + const apiClient = apiClientManager.clientOrThrow(); + + const response = await apiClient.createSessionStreamWaitpoint(ctx.run.id, { + session: this.sessionId, + io: "in", + timeout: options?.timeout, + idempotencyKey: options?.idempotencyKey, + idempotencyKeyTTL: options?.idempotencyKeyTTL, + tags: options?.tags, + lastSeqNum: sessionStreams.lastSeqNum(this.sessionId, "in"), + }); + + const result = await tracer.startActiveSpan( + options?.spanName ?? `sessions.open(${this.sessionId}).in.wait()`, + async (span) => { + const waitResponse = await apiClient.waitForWaitpointToken({ + runFriendlyId: ctx.run.id, + waitpointFriendlyId: response.waitpointId, + }); + + if (!waitResponse.success) { + throw new Error("Failed to block on session stream waitpoint"); + } + + // Drop the SSE tail + buffer before suspending so the record + // delivered via the waitpoint path isn't re-buffered on resume. + sessionStreams.disconnectStream(this.sessionId, "in"); + + const waitResult = await runtime.waitUntil(response.waitpointId); + + const data = + waitResult.output !== undefined + ? await conditionallyImportAndParsePacket( + { + data: waitResult.output, + dataType: waitResult.outputType ?? "application/json", + }, + apiClient + ) + : undefined; + + if (waitResult.ok) { + // Advance the seq counter so the SSE tail doesn't replay the + // record that was consumed via the waitpoint. + const prevSeq = sessionStreams.lastSeqNum(this.sessionId, "in"); + const nextSeq = (prevSeq ?? -1) + 1; + sessionStreams.setLastSeqNum(this.sessionId, "in", nextSeq); + + return { ok: true as const, output: data as T }; + } else { + const error = new WaitpointTimeoutError(data?.message ?? "Timed out"); + span.recordException(error); + span.setStatus({ code: SpanStatusCode.ERROR }); + return { ok: false as const, error }; + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "wait", + [SemanticInternalAttributes.ENTITY_TYPE]: "waitpoint", + [SemanticInternalAttributes.ENTITY_ID]: response.waitpointId, + session: this.sessionId, + io: "in", + ...accessoryAttributes({ + items: [{ text: `${this.sessionId}.in`, variant: "normal" }], + style: "codepath", + }), + }, + } + ); + + resolve(result); + } catch (error) { + reject(error); + } + }); + } + + /** + * Wait for a record with an idle-then-suspend strategy. Keeps the run + * active (using compute) for `idleTimeoutInSeconds`, then suspends via + * {@link wait} if nothing arrives. If a record arrives during the idle + * phase the run responds without suspending. + */ + async waitWithIdleTimeout( + options: InputStreamWaitWithIdleTimeoutOptions + ): Promise<{ ok: true; output: T } | { ok: false; error?: Error }> { + const self = this; + const spanName = + options.spanName ?? `sessions.open(${this.sessionId}).in.waitWithIdleTimeout()`; + + return tracer.startActiveSpan( + spanName, + async (span) => { + if (options.idleTimeoutInSeconds > 0) { + const warm = await sessionStreams.once(self.sessionId, "in", { + timeoutMs: options.idleTimeoutInSeconds * 1000, + }); + if (warm.ok) { + span.setAttribute("wait.resolved", "idle"); + return { ok: true as const, output: warm.output as T }; + } + } + + if (options.skipSuspend) { + span.setAttribute("wait.resolved", "skipped"); + return { ok: false as const, error: undefined }; + } + + if (options.onSuspend) { + await options.onSuspend(); + } + + span.setAttribute("wait.resolved", "suspended"); + const waitResult = await self.wait({ + timeout: options.timeout, + spanName: "suspended", + }); + + if (waitResult.ok && options.onResume) { + await options.onResume(); + } + + return waitResult; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "sessions", + session: self.sessionId, + io: "in", + ...accessoryAttributes({ + items: [{ text: `${self.sessionId}.in`, variant: "normal" }], + style: "codepath", + }), + }, + } + ); + } +} + +export type SessionSubscribeOptions = { + signal?: AbortSignal; + lastEventId?: string | number; + /** Timeout in seconds for the underlying long-poll (max 600). */ + timeoutInSeconds?: number; + /** Called for each SSE event with the full event metadata (id, timestamp). */ + onPart?: (part: { id: string; chunk: T; timestamp: number }) => void; + /** Called when the server signals end-of-stream. */ + onComplete?: () => void; + /** Called on unrecoverable errors after the retry budget is exhausted. */ + onError?: (error: Error) => void; +}; + +// ─── helpers ──────────────────────────────────────────────────────── + +function sessionAttributes(id: string, extra?: Record) { + return { + session: id, + ...(extra ?? {}), + ...accessoryAttributes({ + items: [{ text: id, variant: "normal" }], + style: "codepath", + }), + }; +} + +function toAttr(value: string | string[]): string { + return Array.isArray(value) ? value.join(",") : value; +} diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts index c69bceeb535..cd51b35e32d 100644 --- a/packages/trigger-sdk/src/v3/shared.ts +++ b/packages/trigger-sdk/src/v3/shared.ts @@ -90,6 +90,7 @@ import type { TaskWithToolOptions, ToolTask, ToolTaskParameters, + TriggerAndSubscribeOptions, TriggerAndWaitOptions, TriggerApiRequestOptions, TriggerOptions, @@ -214,6 +215,26 @@ export function createTask< }); }, params.id); }, + triggerAndSubscribe: (payload, options) => { + return new TaskRunPromise((resolve, reject) => { + triggerAndSubscribe_internal( + "triggerAndSubscribe()", + params.id, + payload, + undefined, + { + queue: params.queue?.name, + ...options, + } + ) + .then((result) => { + resolve(result); + }) + .catch((error) => { + reject(error); + }); + }, params.id); + }, batchTriggerAndWait: async (items, options) => { return await batchTriggerAndWait_internal( "batchTriggerAndWait()", @@ -235,6 +256,8 @@ export function createTask< queue: params.queue, retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined, machine: typeof params.machine === "string" ? { preset: params.machine } : params.machine, + triggerSource: params.triggerSource, + agentConfig: params.agentConfig, maxDuration: params.maxDuration, ttl: params.ttl, payloadSchema: params.jsonSchema, @@ -259,7 +282,7 @@ export function createTask< } /** - * @deprecated use ai.tool() instead + * @deprecated Use `schemaTask` plus AI SDK `tool()` with `execute: ai.toolExecute(task)` instead. */ export function createToolTask< TIdentifier extends string, @@ -346,6 +369,26 @@ export function createSchemaTask< }); }, params.id); }, + triggerAndSubscribe: (payload, options) => { + return new TaskRunPromise((resolve, reject) => { + triggerAndSubscribe_internal, TOutput>( + "triggerAndSubscribe()", + params.id, + payload, + parsePayload, + { + queue: params.queue?.name, + ...options, + } + ) + .then((result) => { + resolve(result); + }) + .catch((error) => { + reject(error); + }); + }, params.id); + }, batchTriggerAndWait: async (items, options) => { return await batchTriggerAndWait_internal, TOutput>( "batchTriggerAndWait()", @@ -367,6 +410,8 @@ export function createSchemaTask< queue: params.queue, retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined, machine: typeof params.machine === "string" ? { preset: params.machine } : params.machine, + triggerSource: params.triggerSource, + agentConfig: params.agentConfig, maxDuration: params.maxDuration, ttl: params.ttl, fns: { @@ -465,6 +510,49 @@ export function triggerAndWait( }, id); } +/** + * Trigger a task and subscribe to its updates via realtime. Unlike `triggerAndWait`, + * this does NOT suspend the parent run — the parent stays alive and subscribes to updates. + * This enables parallel execution and proper abort signal handling. + * + * @param id - The id of the task to trigger + * @param payload + * @param options - Options for the task run, including an optional `signal` to cancel the subscription and child run + * @returns TaskRunPromise + * @example + * ```ts + * import { tasks } from "@trigger.dev/sdk/v3"; + * const result = await tasks.triggerAndSubscribe("my-task", { foo: "bar" }); + * + * if (result.ok) { + * console.log(result.output); + * } else { + * console.error(result.error); + * } + * ``` + */ +export function triggerAndSubscribe( + id: TaskIdentifier, + payload: TaskPayload, + options?: TriggerAndSubscribeOptions +): TaskRunPromise, TaskOutput> { + return new TaskRunPromise, TaskOutput>((resolve, reject) => { + triggerAndSubscribe_internal, TaskPayload, TaskOutput>( + "tasks.triggerAndSubscribe()", + id, + payload, + undefined, + options + ) + .then((result) => { + resolve(result); + }) + .catch((error) => { + reject(error); + }); + }, id); +} + /** * Batch trigger multiple task runs with the given payloads, and wait for the results. Returns the results of the task runs. * @param id - The id of the task to trigger @@ -2441,6 +2529,128 @@ async function triggerAndWait_internal( + name: string, + id: TIdentifier, + payload: TPayload, + parsePayload?: SchemaParseFn, + options?: TriggerAndSubscribeOptions +): Promise> { + const ctx = taskContext.ctx; + + if (!ctx) { + throw new Error("triggerAndSubscribe can only be used from inside a task.run()"); + } + + const apiClient = apiClientManager.clientOrThrow(); + + const parsedPayload = parsePayload ? await parsePayload(payload) : payload; + const payloadPacket = await stringifyIO(parsedPayload); + + const processedIdempotencyKey = await makeIdempotencyKey(options?.idempotencyKey); + const idempotencyKeyOptions = processedIdempotencyKey + ? getIdempotencyKeyOptions(processedIdempotencyKey) + : undefined; + + return await tracer.startActiveSpan( + name, + async (span) => { + const response = await apiClient.triggerTask( + id, + { + payload: payloadPacket.data, + options: { + lockToVersion: taskContext.worker?.version, + queue: options?.queue ? { name: options.queue } : undefined, + concurrencyKey: options?.concurrencyKey, + test: taskContext.ctx?.run.isTest, + payloadType: payloadPacket.dataType, + delay: options?.delay, + ttl: options?.ttl, + tags: options?.tags, + maxAttempts: options?.maxAttempts, + metadata: options?.metadata, + maxDuration: options?.maxDuration, + parentRunId: ctx.run.id, + // NOTE: no resumeParentOnCompletion — parent stays alive and subscribes + idempotencyKey: processedIdempotencyKey?.toString(), + idempotencyKeyTTL: options?.idempotencyKeyTTL, + idempotencyKeyOptions, + machine: options?.machine, + priority: options?.priority, + region: options?.region, + debounce: options?.debounce, + }, + }, + {} + ); + + // Set attributes after trigger so the dashboard can link to the child run + span.setAttribute("messaging.message.id", response.id); + span.setAttribute("runId", response.id); + span.setAttribute(SemanticInternalAttributes.ENTITY_TYPE, "run"); + span.setAttribute(SemanticInternalAttributes.ENTITY_ID, response.id); + + // Optionally cancel the child run when the abort signal fires (default: true) + const cancelOnAbort = options?.cancelOnAbort !== false; + if (options?.signal && cancelOnAbort) { + const onAbort = () => { + apiClient.cancelRun(response.id).catch(() => {}); + }; + if (options.signal.aborted) { + await apiClient.cancelRun(response.id).catch(() => {}); + throw new Error("Aborted"); + } + options.signal.addEventListener("abort", onAbort, { once: true }); + } + + for await (const run of apiClient.subscribeToRun(response.id, { + closeOnComplete: true, + signal: options?.signal, + skipColumns: ["payload"], + })) { + if (run.isSuccess) { + // run.output from subscribeToRun is already deserialized + return { + ok: true as const, + id: response.id, + taskIdentifier: id as TIdentifier, + output: run.output as TOutput, + }; + } + if (run.isFailed || run.isCancelled) { + const error = new Error(run.error?.message ?? `Task ${id} failed (${run.status})`); + if (run.error?.name) error.name = run.error.name; + + return { + ok: false as const, + id: response.id, + taskIdentifier: id as TIdentifier, + error, + }; + } + } + + throw new Error(`Task ${id}: subscription ended without completion`); + }, + { + kind: SpanKind.PRODUCER, + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "trigger", + ...accessoryAttributes({ + items: [ + { + text: id, + variant: "normal", + }, + ], + style: "codepath", + }), + }, + } + ); +} + async function batchTriggerAndWait_internal( name: string, id: TIdentifier, diff --git a/packages/trigger-sdk/src/v3/skill.ts b/packages/trigger-sdk/src/v3/skill.ts new file mode 100644 index 00000000000..a3c145d3836 --- /dev/null +++ b/packages/trigger-sdk/src/v3/skill.ts @@ -0,0 +1,211 @@ +import * as fs from "node:fs/promises"; +import * as path from "node:path"; +import { resourceCatalog } from "@trigger.dev/core/v3"; + +/** + * Parsed `SKILL.md` frontmatter. Only `name` + `description` are required; + * additional keys are preserved but untyped. + */ +export type SkillFrontmatter = { + name: string; + description: string; + [key: string]: unknown; +}; + +/** + * A resolved skill ready to hand to `chat.skills.set()`. Includes the parsed + * SKILL.md content plus the on-disk path to the bundled skill folder. + */ +export type ResolvedSkill = { + id: string; + /** Skill version — `"local"` in Phase 1 until backend-managed overrides land. */ + version: number | "local"; + /** Labels applied to this version — empty in Phase 1. */ + labels: string[]; + /** Full raw `SKILL.md` content (with frontmatter). */ + skillMd: string; + /** Parsed frontmatter fields. */ + frontmatter: SkillFrontmatter; + /** Body of SKILL.md with the frontmatter block stripped. */ + body: string; + /** Absolute path to the bundled skill folder (scripts, references, assets live here). */ + path: string; +}; + +export type SkillOptions = { + id: TIdentifier; + /** Path to the skill source folder, relative to the project root. */ + path: string; +}; + +export type SkillHandle = { + id: TIdentifier; + /** + * Read the bundled `SKILL.md` from disk and return the resolved skill. + * + * This is the Phase 1 path — backend-managed overrides are not available + * yet. Works locally (during `trigger dev`) and in the deploy image. + */ + local(): Promise; + /** + * Resolve the skill against the dashboard (current/override version). + * + * Not available in Phase 1 — throws. Use `local()` until backend-managed + * skills ship. + */ + resolve(): Promise; +}; + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +export type AnySkillHandle = SkillHandle; + +/** Extract the id literal type from a SkillHandle. */ +export type SkillIdentifier = T extends SkillHandle + ? TId + : string; + +/** + * Bundled skills are copied to `${cwd}/.trigger/skills/{id}/` by the CLI at + * build time. At runtime the same layout holds for both `trigger dev` (cwd + * = dev output dir) and deploy (cwd = /app). + */ +function bundledSkillPath(id: string): string { + return path.resolve(process.cwd(), ".trigger", "skills", id); +} + +const FRONTMATTER_RE = /^---\r?\n([\s\S]*?)\r?\n---\r?\n*/; + +/** + * Parse a minimal YAML-subset frontmatter block. We only support top-level + * string keys like `name: foo` and `description: bar`. Enough for SKILL.md + * frontmatter without pulling in a YAML dep. + */ +export function parseFrontmatter(content: string): { + frontmatter: SkillFrontmatter; + body: string; +} { + const match = content.match(FRONTMATTER_RE); + if (!match || !match[1]) { + throw new Error( + "Skill: SKILL.md is missing a frontmatter block. " + + "Expected `---\\nname: ...\\ndescription: ...\\n---` at the top of the file." + ); + } + + const raw = match[1]; + const frontmatter: Record = {}; + for (const line of raw.split(/\r?\n/)) { + const trimmed = line.trim(); + if (!trimmed || trimmed.startsWith("#")) continue; + const idx = trimmed.indexOf(":"); + if (idx === -1) continue; + const key = trimmed.slice(0, idx).trim(); + let value = trimmed.slice(idx + 1).trim(); + // Strip surrounding quotes if present + if ( + (value.startsWith('"') && value.endsWith('"')) || + (value.startsWith("'") && value.endsWith("'")) + ) { + value = value.slice(1, -1); + } + if (key) frontmatter[key] = value; + } + + if (typeof frontmatter.name !== "string" || !frontmatter.name) { + throw new Error("Skill: SKILL.md frontmatter is missing required `name` field."); + } + if (typeof frontmatter.description !== "string" || !frontmatter.description) { + throw new Error("Skill: SKILL.md frontmatter is missing required `description` field."); + } + + const body = content.slice(match[0].length); + + return { frontmatter: frontmatter as SkillFrontmatter, body }; +} + +async function loadLocal(id: string): Promise { + const skillPath = bundledSkillPath(id); + const skillMdPath = path.join(skillPath, "SKILL.md"); + + let skillMd: string; + try { + skillMd = await fs.readFile(skillMdPath, "utf8"); + } catch (err) { + throw new Error( + `Skill "${id}": could not read SKILL.md at ${skillMdPath}. ` + + `Skills must be bundled into .trigger/skills/{id}/ — this usually means ` + + `the CLI build step didn't run, or the skill wasn't registered via ai.defineSkill. ` + + `Underlying error: ${(err as Error).message}` + ); + } + + const { frontmatter, body } = parseFrontmatter(skillMd); + + return { + id, + version: "local", + labels: [], + skillMd, + frontmatter, + body, + path: skillPath, + }; +} + +/** + * Define an agent skill — a developer-authored folder with a `SKILL.md` file + * plus optional `scripts/`, `references/`, and `assets/` subfolders. Registers + * the skill with the resource catalog so the Trigger.dev CLI can bundle it + * into the deploy image automatically (no build extension needed). + * + * Call `.local()` on the returned handle to load the bundled SKILL.md at + * runtime and use it with `chat.skills.set()`. + * + * @example + * ```ts + * // trigger/skills/pdf-processing/SKILL.md + * // trigger/skills/pdf-processing/scripts/extract.py + * import { ai } from "@trigger.dev/sdk"; + * + * export const pdfSkill = ai.defineSkill({ + * id: "pdf-processing", + * path: "./skills/pdf-processing", + * }); + * + * export const agent = chat.agent({ + * id: "docs", + * onChatStart: async () => { + * chat.skills.set([await pdfSkill.local()]); + * }, + * run: async ({ messages, signal }) => { + * return streamText({ + * model: openai("gpt-4o"), + * messages, + * abortSignal: signal, + * ...chat.toStreamTextOptions(), + * }); + * }, + * }); + * ``` + */ +export function defineSkill( + options: SkillOptions +): SkillHandle { + resourceCatalog.registerSkillMetadata({ + id: options.id, + sourcePath: options.path, + }); + + return { + id: options.id, + async local() { + return loadLocal(options.id); + }, + async resolve() { + throw new Error( + `Skill "${options.id}": resolve() is not available yet — backend-managed ` + + `skills ship in Phase 2. Use skill.local() instead.` + ); + }, + }; +} diff --git a/packages/trigger-sdk/src/v3/skills.ts b/packages/trigger-sdk/src/v3/skills.ts new file mode 100644 index 00000000000..6811cda75f4 --- /dev/null +++ b/packages/trigger-sdk/src/v3/skills.ts @@ -0,0 +1,9 @@ +export { defineSkill as define } from "./skill.js"; +export type { + AnySkillHandle, + ResolvedSkill, + SkillFrontmatter, + SkillHandle, + SkillIdentifier, + SkillOptions, +} from "./skill.js"; diff --git a/packages/trigger-sdk/src/v3/streams.ts b/packages/trigger-sdk/src/v3/streams.ts index 68edc2a64ab..fdc30e3ec22 100644 --- a/packages/trigger-sdk/src/v3/streams.ts +++ b/packages/trigger-sdk/src/v3/streams.ts @@ -25,8 +25,10 @@ import { InputStreamOncePromise, type InputStreamOnceResult, type InputStreamWaitOptions, + type InputStreamWaitWithIdleTimeoutOptions, type SendInputStreamOptions, type InferInputStreamType, + type StreamWriteResult, } from "@trigger.dev/core/v3"; import { conditionallyImportAndParsePacket } from "@trigger.dev/core/v3/utils/ioSerialization"; import { tracer } from "./tracer.js"; @@ -139,7 +141,7 @@ function pipe( opts = valueOrOptions as PipeStreamOptions | undefined; } - return pipeInternal(key, value, opts, "streams.pipe()"); + return pipeInternal(key, value, opts, opts?.spanName ?? "streams.pipe()"); } /** @@ -167,6 +169,7 @@ function pipeInternal( [SemanticInternalAttributes.ENTITY_TYPE]: "realtime-stream", [SemanticInternalAttributes.ENTITY_ID]: `${runId}:${key}`, [SemanticInternalAttributes.STYLE_ICON]: "streams", + ...(opts?.collapsed ? { [SemanticInternalAttributes.COLLAPSED]: true } : {}), ...accessoryAttributes({ items: [ { @@ -194,7 +197,9 @@ function pipeInternal( return { stream: instance.stream, - waitUntilComplete: () => instance.wait(), + waitUntilComplete: async () => { + return instance.wait(); + }, }; } catch (error) { // if the error is a signal abort error, we need to end the span but not record an exception @@ -640,7 +645,7 @@ function writerInternal(key: string, options: WriterStreamOptions) } }); - return pipeInternal(key, stream, options, "streams.writer()"); + return pipeInternal(key, stream, options, options.spanName ?? "streams.writer()"); } export type RealtimeDefineStreamOptions = { @@ -656,8 +661,18 @@ function define(opts: RealtimeDefineStreamOptions): RealtimeDefinedStream read(runId, options) { return read(runId, opts.id, options); }, - append(value, options) { - return append(opts.id, value as BodyInit, options); + async append(value, options) { + // Use a single-write writer so objects are serialized the same way + // as stream.writer() — the raw append API sends BodyInit which + // doesn't serialize objects correctly for SSE consumers. + const { waitUntilComplete } = writer(opts.id, { + ...options, + spanName: "streams.append()", + execute: ({ write }) => { + write(value); + }, + }); + await waitUntilComplete(); }, writer(options) { return writer(opts.id, options); @@ -713,7 +728,7 @@ function input(opts: { id: string }): RealtimeDefinedInputStream { return new InputStreamOncePromise((resolve, reject) => { tracer .startActiveSpan( - `inputStream.once()`, + options?.spanName ?? `inputStream.once()`, async () => { const result = await innerPromise; resolve(result as InputStreamOnceResult); @@ -750,23 +765,21 @@ function input(opts: { id: string }): RealtimeDefinedInputStream { const apiClient = apiClientManager.clientOrThrow(); + // Create the waitpoint before the span so we have the entity ID upfront + const response = await apiClient.createInputStreamWaitpoint(ctx.run.id, { + streamId: opts.id, + timeout: options?.timeout, + idempotencyKey: options?.idempotencyKey, + idempotencyKeyTTL: options?.idempotencyKeyTTL, + tags: options?.tags, + lastSeqNum: inputStreams.lastSeqNum(opts.id), + }); + const result = await tracer.startActiveSpan( - `inputStream.wait()`, + options?.spanName ?? `inputStream.wait()`, async (span) => { - // 1. Create a waitpoint linked to this input stream - const response = await apiClient.createInputStreamWaitpoint(ctx.run.id, { - streamId: opts.id, - timeout: options?.timeout, - idempotencyKey: options?.idempotencyKey, - idempotencyKeyTTL: options?.idempotencyKeyTTL, - tags: options?.tags, - lastSeqNum: inputStreams.lastSeqNum(opts.id), - }); - // Set the entity ID now that we have the waitpoint ID - span.setAttribute(SemanticInternalAttributes.ENTITY_ID, response.waitpointId); - - // 2. Block the run on the waitpoint + // 1. Block the run on the waitpoint const waitResponse = await apiClient.waitForWaitpointToken({ runFriendlyId: ctx.run.id, waitpointFriendlyId: response.waitpointId, @@ -776,6 +789,12 @@ function input(opts: { id: string }): RealtimeDefinedInputStream { throw new Error("Failed to block on input stream waitpoint"); } + // 2. Disconnect the SSE tail and clear the buffer before suspending. + // Without this, the tail stays alive during the suspension window and + // may buffer a copy of the same message that will be delivered via the + // waitpoint, causing a duplicate on resume. + inputStreams.disconnectStream(opts.id); + // 3. Suspend the task const waitResult = await runtime.waitUntil(response.waitpointId); @@ -792,6 +811,12 @@ function input(opts: { id: string }): RealtimeDefinedInputStream { : undefined; if (waitResult.ok) { + // Advance the seq counter so the SSE tail doesn't replay + // the record that was consumed via the waitpoint path when + // it lazily reconnects on the next on()/once() call. + const prevSeq = inputStreams.lastSeqNum(opts.id); + inputStreams.setLastSeqNum(opts.id, (prevSeq ?? -1) + 1); + return { ok: true as const, output: data as TData }; } else { const error = new WaitpointTimeoutError(data?.message ?? "Timed out"); @@ -806,6 +831,7 @@ function input(opts: { id: string }): RealtimeDefinedInputStream { attributes: { [SemanticInternalAttributes.STYLE_ICON]: "wait", [SemanticInternalAttributes.ENTITY_TYPE]: "waitpoint", + [SemanticInternalAttributes.ENTITY_ID]: response.waitpointId, streamId: opts.id, ...accessoryAttributes({ items: [ @@ -826,6 +852,61 @@ function input(opts: { id: string }): RealtimeDefinedInputStream { } }); }, + async waitWithIdleTimeout(options) { + const self = this; + const spanName = options.spanName ?? `inputStream.waitWithIdleTimeout()`; + + return tracer.startActiveSpan( + spanName, + async (span) => { + // Idle phase: keep compute alive + if (options.idleTimeoutInSeconds > 0) { + const warm = await inputStreams.once(opts.id, { + timeoutMs: options.idleTimeoutInSeconds * 1000, + }); + if (warm.ok) { + span.setAttribute("wait.resolved", "idle"); + return { ok: true as const, output: warm.output as TData }; + } + } + + // Skip suspend if requested — return as if timed out + if (options.skipSuspend) { + span.setAttribute("wait.resolved", "skipped"); + return { ok: false as const, error: undefined }; + } + + // Fire onSuspend callback before entering cold phase + if (options.onSuspend) { + await options.onSuspend(); + } + + // Cold phase: suspend via .wait() — creates a child span + span.setAttribute("wait.resolved", "suspended"); + const waitResult = await self.wait({ + timeout: options.timeout, + spanName: "suspended", + }); + + // Fire onResume callback after successful resume + if (waitResult.ok && options.onResume) { + await options.onResume(); + } + + return waitResult; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "streams", + streamId: opts.id, + ...accessoryAttributes({ + items: [{ text: opts.id, variant: "normal" }], + style: "codepath", + }), + }, + } + ); + }, async send(runId, data, options) { return tracer.startActiveSpan( `inputStream.send()`, diff --git a/packages/trigger-sdk/src/v3/tasks.ts b/packages/trigger-sdk/src/v3/tasks.ts index 75b7e85e625..5781a104229 100644 --- a/packages/trigger-sdk/src/v3/tasks.ts +++ b/packages/trigger-sdk/src/v3/tasks.ts @@ -20,6 +20,7 @@ import { SubtaskUnwrapError, trigger, triggerAndWait, + triggerAndSubscribe, } from "./shared.js"; export { SubtaskUnwrapError }; @@ -96,6 +97,7 @@ export const tasks = { trigger, batchTrigger, triggerAndWait, + triggerAndSubscribe, batchTriggerAndWait, /** @deprecated Use onStartAttempt instead */ onStart, diff --git a/packages/trigger-sdk/src/v3/test/index.ts b/packages/trigger-sdk/src/v3/test/index.ts new file mode 100644 index 00000000000..cdeded1a7a8 --- /dev/null +++ b/packages/trigger-sdk/src/v3/test/index.ts @@ -0,0 +1,23 @@ +// Importing this module installs an in-memory resource catalog so that +// chat.agent() calls (which run at import time) register their task +// functions where the test harness can find them. +// +// Users should import `@trigger.dev/sdk/ai/test` BEFORE their agent +// modules so the registration side-effect runs first. +import "./setup-catalog.js"; + +export { + mockChatAgent, + type MockChatAgentOptions, + type MockChatAgentHarness, + type MockChatAgentTurn, +} from "./mock-chat-agent.js"; + +// Re-export the lower-level task context harness so consumers can build +// their own test helpers without adding a separate `@trigger.dev/core` +// dependency to their reference projects. +export { + runInMockTaskContext, + type MockTaskContextDrivers, + type MockTaskContextOptions, +} from "@trigger.dev/core/v3/test"; diff --git a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts new file mode 100644 index 00000000000..b6fe21e4dc3 --- /dev/null +++ b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts @@ -0,0 +1,686 @@ +import type { UIMessage, UIMessageChunk } from "ai"; +import { resourceCatalog } from "@trigger.dev/core/v3"; +import type { LocalsKey } from "@trigger.dev/core/v3"; +import { + runInMockTaskContext, + type MockTaskContextOptions, +} from "@trigger.dev/core/v3/test"; +import { + __setSessionOpenImplForTests, + __setSessionStartImplForTests, +} from "../sessions.js"; +import { + __setReadChatSnapshotImplForTests, + __setReplaySessionOutTailImplForTests, + __setWriteChatSnapshotImplForTests, + type ChatSnapshotV1, +} from "../ai.js"; +import { + createTestSessionHandle, + type TestSessionOutState, +} from "./test-session-handle.js"; + +/** Pre-seed locals before the agent's `run()` starts. */ +export type SetupLocals = (locals: { + set(key: LocalsKey, value: T): void; +}) => void | Promise; + +// The slim wire payload shape used by chat.agent tasks. Kept loose here so we +// don't import from the backend-only ai.ts module. At most ONE message per +// record — runtime rebuilds prior history from snapshot + replay at boot. +type ChatWirePayload = { + /** At most one message — singular under the slim wire. Set on submit-message. */ + message?: UIMessage; + /** Bespoke escape hatch — only set on `trigger: "handover-prepare"`. */ + headStartMessages?: UIMessage[]; + chatId: string; + trigger: + | "submit-message" + | "regenerate-message" + | "preload" + | "close" + | "action" + | "handover-prepare"; + messageId?: string; + metadata?: unknown; + action?: unknown; + continuation?: boolean; + previousRunId?: string; + idleTimeoutInSeconds?: number; + sessionId?: string; +}; + +/** A reference to a `chat.agent` task returned by `chat.agent({ id, ... })`. */ +type ChatAgentHandle = { id: string }; + +/** + * Options for `mockChatAgent`. + */ +export type MockChatAgentOptions = { + /** The chat session id passed into every wire payload. Defaults to `"test-chat"`. */ + chatId?: string; + /** Client-provided metadata (`clientData`) for the session. */ + clientData?: unknown; + /** Task context overrides passed through to {@link runInMockTaskContext}. */ + taskContext?: MockTaskContextOptions; + /** + * Whether to start the task in preload mode. Defaults to `true` so the + * first `sendMessage()` triggers the first turn via the preload path. + * Set to `false` to skip preload — the first `sendMessage()` starts turn 0 directly. + * + * Ignored when `mode: "handover-prepare"` is set. + */ + preload?: boolean; + /** + * Initial trigger the agent boots with. Defaults to `"preload"` (or + * `"submit-message"` when `preload: false`). Use `"handover-prepare"` + * to drive the chat.handover wait branch — call `sendHandover()` / + * `sendHandoverSkip()` to dispatch the handover signal. + */ + mode?: "preload" | "submit-message" | "handover-prepare"; + /** + * Pre-seed the snapshot the agent reads at run boot. The runtime's + * snapshot read is replaced with one that returns this snapshot + * (skipping the real S3 GET). Use to drive boot scenarios — fresh + * boot with prior history, OOM-retry boot with stale snapshot, etc. + * Pass `undefined` (the default) to start with no snapshot. + * + * See plan section B.3 for the boot orchestration spec. + */ + snapshot?: ChatSnapshotV1; + /** + * Callback that runs **before** the agent's `run()` is invoked, with a + * `set` function for pre-seeding locals. Use this to inject server-side + * dependencies (database clients, service stubs) that the agent reads + * via `locals.get()` in its hooks. + * + * @example + * ```ts + * import { dbKey } from "./db"; + * + * const harness = mockChatAgent(agent, { + * chatId: "test-1", + * setupLocals: (locals) => { + * locals.set(dbKey, testDb); + * }, + * }); + * ``` + */ + setupLocals?: SetupLocals; +}; + +/** + * Result of a single turn, returned by driver methods like `sendMessage()`. + */ +export type MockChatAgentTurn = { + /** UIMessageChunks emitted during this turn (excludes control chunks like turn-complete). */ + chunks: UIMessageChunk[]; + /** All raw chunks including control chunks (turn-complete, upgrade-required, etc.). */ + rawChunks: unknown[]; +}; + +/** + * Harness returned by `mockChatAgent`. Drives a `chat.agent` task end-to-end + * without network or task runtime. + */ +export type MockChatAgentHarness = { + /** The chat session id used by this harness. */ + readonly chatId: string; + + /** + * Send a single user message (or tool-approval-responded assistant + * message) and wait for the next turn-complete. Returns the chunks + * produced during this turn. + * + * Slim wire: at most ONE message per send. The agent reconstructs prior + * history from snapshot + session.out replay at run boot. + */ + sendMessage(message: UIMessage): Promise; + + /** + * Send a regenerate signal (no message body — slim wire). The agent + * trims trailing assistant messages from its in-memory accumulator and + * re-runs. Waits for turn-complete. + */ + sendRegenerate(): Promise; + + /** + * Drive the head-start path: sends `trigger: "handover-prepare"` with + * `headStartMessages` carrying the first-turn UIMessage history. Used + * only at the very first turn before any snapshot exists. The route + * handler ships full UIMessage history through this path because the + * customer's HTTP endpoint isn't subject to the `/in/append` cap. + */ + sendHeadStart(args: { messages: UIMessage[] }): Promise; + + /** Send a custom action and wait for the next turn-complete. */ + sendAction(action: unknown): Promise; + + /** Fire a stop signal. Does not wait for the turn — the task keeps running. */ + sendStop(message?: string): Promise; + + /** + * Dispatch a `handover` signal — the agent picks up partial assistant + * messages and continues the turn. Only meaningful when the harness + * was started with `mode: "handover-prepare"`. Waits for turn-complete. + * + * `isFinal: false` (default) — agent runs `streamText` which executes + * any pending tool-calls (via the approval round) and resumes from + * step 2. + * + * `isFinal: true` — agent runs lifecycle hooks but skips `streamText`. + * The partial IS the response; `onTurnComplete` fires with it. + */ + sendHandover(args: { + partialAssistantMessage: unknown[]; + isFinal?: boolean; + messageId?: string; + }): Promise; + + /** + * Dispatch a `handover-skip` signal — the agent exits cleanly without + * firing turn hooks. Only meaningful when the harness was started + * with `mode: "handover-prepare"`. Awaits the run finishing. + */ + sendHandoverSkip(): Promise; + + /** + * Pre-seed the snapshot read for the next boot. The runtime's snapshot + * read returns this snapshot (skipping S3). Pass `undefined` to clear — + * the boot then sees no snapshot and falls through to replay-only. + * + * Effective on the next run boot only. Calling mid-turn is a no-op + * because the snapshot read happens once at run boot. + */ + seedSnapshot(snapshot: ChatSnapshotV1 | undefined): void; + + /** + * Pre-seed `session.out` chunks for the next boot's replay. The runtime's + * `replaySessionOutTail` returns whatever the synthetic chunks reduce + * to. Pass `[]` to clear (boot replay returns no messages). + * + * Requires `__setReplaySessionOutTailImplForTests` exported from + * `ai.ts`. The harness throws a clear error at call time if that hook + * isn't available. + */ + seedSessionOutTail(chunks?: UIMessageChunk[]): void; + + /** + * The most recently written snapshot, or `undefined` if no snapshot + * has been written yet. Updated each time `writeChatSnapshot` is + * invoked from the run loop's snapshot-write site (plan section B.6). + */ + getSnapshot(): ChatSnapshotV1 | undefined; + + /** + * Close the chat session cleanly. Sends `trigger: "close"` and awaits the + * task's `run()` function returning. Call this at the end of every test + * (or use `await using`) so the background task isn't left dangling. + */ + close(): Promise; + + /** All UIMessageChunks emitted since the harness was created. */ + readonly allChunks: UIMessageChunk[]; + + /** Every raw chunk (including control chunks) emitted since the harness was created. */ + readonly allRawChunks: unknown[]; +}; + +const CONTROL_CHUNK_TYPES = new Set([ + "trigger:turn-complete", + "trigger:upgrade-required", +]); + +function isControlChunk(chunk: unknown): boolean { + if (typeof chunk !== "object" || chunk === null) return false; + const type = (chunk as { type?: string }).type; + return typeof type === "string" && CONTROL_CHUNK_TYPES.has(type); +} + +/** + * Create an offline test harness for a `chat.agent` task. + * + * The harness starts the agent's `run()` function in a mocked task context, + * waits in preload for the first message, then exposes driver methods for + * sending messages / actions / stop signals and awaiting turn completion. + * + * Users are responsible for mocking the language model themselves — use + * `MockLanguageModelV3` and `simulateReadableStream` from `ai/test` inside + * their agent's `run()` function (typically via DI through `clientData`). + * + * @example + * ```ts + * import { mockChatAgent } from "@trigger.dev/sdk/ai/test"; + * import { MockLanguageModelV3, simulateReadableStream } from "ai/test"; + * import { myAgent } from "./my-agent"; + * + * test("says hello", async () => { + * const harness = mockChatAgent(myAgent, { chatId: "test-1" }); + * try { + * const turn = await harness.sendMessage({ + * id: "m1", + * role: "user", + * parts: [{ type: "text", text: "hi" }], + * }); + * expect(turn.chunks).toContainEqual( + * expect.objectContaining({ type: "text-delta", delta: "hello" }) + * ); + * } finally { + * await harness.close(); + * } + * }); + * ``` + */ +export function mockChatAgent( + agent: ChatAgentHandle, + options: MockChatAgentOptions = {} +): MockChatAgentHarness { + const chatId = options.chatId ?? "test-chat"; + // The agent opens the session with `payload.sessionId ?? payload.chatId`. + // We pass no sessionId, so it falls back to chatId. + const sessionId = chatId; + const mode: "preload" | "submit-message" | "handover-prepare" = + options.mode ?? (options.preload === false ? "submit-message" : "preload"); + const clientData = options.clientData; + + const taskEntry = resourceCatalog.getTask(agent.id); + if (!taskEntry) { + throw new Error( + `mockChatAgent: no task registered with id "${agent.id}". ` + + `Import "@trigger.dev/sdk/ai/test" before your agent module so tasks register correctly.` + ); + } + + const runFn = taskEntry.fns.run; + + // Session .out state: chunks + listener registry. Shared between the + // harness and the TestSessionOutputChannel installed via the open-override. + const sessionOutState: TestSessionOutState = { + chunks: [], + listeners: new Set(), + }; + + // Buffers that survive across harness method calls + const allRawChunks: unknown[] = []; + const allChunks: UIMessageChunk[] = []; + + // Promise that resolves when the background task run() function returns. + let taskFinished!: Promise; + let sendSessionInput!: (sessionId: string, data: unknown) => Promise; + let closeSessionInput: ((sessionId: string) => void) | undefined; + let runSignal!: AbortController; + + // A latch that resolves every time `trigger:turn-complete` appears on the chat stream. + // We use a shared pending promise and replace it after each completion. + let turnCompleteResolvers: Array<() => void> = []; + const waitForTurnComplete = () => + new Promise((resolve) => { + turnCompleteResolvers.push(resolve); + }); + + // Signal that the caller is ready to observe output + let harnessReadyResolve!: () => void; + const harnessReady = new Promise((resolve) => { + harnessReadyResolve = resolve; + }); + + // ── Snapshot read/write override state ─────────────────────────────── + // The runtime's snapshot read returns whatever `seededSnapshot` is at + // boot time. The runtime's snapshot write captures into + // `lastWrittenSnapshot` for harness consumers to assert via + // `getSnapshot()`. Installed below alongside the session overrides; + // cleared on close in the same finally block. + let seededSnapshot: ChatSnapshotV1 | undefined = options.snapshot; + let lastWrittenSnapshot: ChatSnapshotV1 | undefined; + let seededReplayChunks: UIMessageChunk[] = []; + + __setReadChatSnapshotImplForTests((_id: string) => { + return seededSnapshot as ChatSnapshotV1 | undefined; + }); + __setWriteChatSnapshotImplForTests((_id: string, snapshot: ChatSnapshotV1) => { + lastWrittenSnapshot = snapshot as ChatSnapshotV1; + }); + + // Replay override: install a default that returns whatever + // `seededReplayChunks` reduces to. Cleared in the same `finally` block + // as the other test overrides. + __setReplaySessionOutTailImplForTests(async () => { + if (seededReplayChunks.length === 0) return []; + return (await reduceChunksToMessages(seededReplayChunks)) as never; + }); + + // Install the session open override so `sessions.open(id)` returns a + // SessionHandle with an in-memory `.out` that captures writes. The + // `.in` channel routes record subscriptions (`on`/`once`/`peek`) + // through the `sessionStreams` global — the mock task context + // installs a `TestSessionStreamManager` there — and stubs `wait()` + // so the suspend path resolves cleanly on `runSignal.abort()` without + // touching the api client. + __setSessionOpenImplForTests((id) => + createTestSessionHandle(id, sessionOutState, () => runSignal?.signal) + ); + + // Install the session start override so any test path that invokes + // `sessions.start()` (typically through a server action shim like + // `chat.createStartSessionAction`) becomes a no-op fixture instead of + // hitting a real API. Most chat.agent tests trigger the run directly + // via `sendPayloadAndWait` and never go through this path, but the + // stub keeps the API safe to call from inside tested code. + __setSessionStartImplForTests((body) => { + if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") { + console.log("[mockChatAgent] sessions.start override:", body); + } + const fakeRunId = `run_test_${body.externalId ?? "anon"}`; + return { + id: `session_test_${body.externalId ?? "anon"}`, + externalId: body.externalId ?? null, + type: body.type, + taskIdentifier: body.taskIdentifier, + triggerConfig: body.triggerConfig, + currentRunId: fakeRunId, + runId: fakeRunId, + publicAccessToken: "tr_test_session_pat", + tags: body.tags ?? [], + metadata: (body.metadata ?? null) as Record | null, + closedAt: null, + closedReason: null, + expiresAt: null, + createdAt: new Date(0), + updatedAt: new Date(0), + isCached: false, + }; + }); + + taskFinished = runInMockTaskContext( + async (drivers) => { + runSignal = new AbortController(); + + const initialPayload: ChatWirePayload = { + chatId, + trigger: mode, + metadata: clientData, + }; + + sendSessionInput = drivers.sessions.in.send; + closeSessionInput = drivers.sessions.in.close; + + // Record every chunk written to session.out, detect turn-complete. + const listener = (chunk: unknown) => { + allRawChunks.push(chunk); + if (!isControlChunk(chunk)) { + allChunks.push(chunk as UIMessageChunk); + } + if ( + typeof chunk === "object" && + chunk !== null && + (chunk as { type?: string }).type === "trigger:turn-complete" + ) { + const resolvers = turnCompleteResolvers; + turnCompleteResolvers = []; + for (const resolve of resolvers) resolve(); + } + }; + sessionOutState.listeners.add(listener); + const unsubscribe = () => sessionOutState.listeners.delete(listener); + + if (options.setupLocals) { + await options.setupLocals({ set: drivers.locals.set }); + } + + harnessReadyResolve(); + + try { + if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") { + console.log("[mockChatAgent] Starting runFn with payload:", initialPayload); + } + await runFn(initialPayload, { + ctx: drivers.ctx, + signal: runSignal.signal, + }); + if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") { + console.log("[mockChatAgent] runFn returned"); + } + } catch (err) { + if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") { + console.log("[mockChatAgent] runFn threw:", err); + } + throw err; + } finally { + unsubscribe(); + // Resolve any outstanding turn-complete waiters so callers don't hang + const resolvers = turnCompleteResolvers; + turnCompleteResolvers = []; + for (const resolve of resolvers) resolve(); + } + }, + options.taskContext + ) + .catch((err) => { + // Propagate errors to pending turn waiters instead of dropping them + const resolvers = turnCompleteResolvers; + turnCompleteResolvers = []; + for (const resolve of resolvers) resolve(); + throw err; + }) + .finally(() => { + // Always clear the test overrides, even if the task threw. + __setSessionOpenImplForTests(undefined); + __setSessionStartImplForTests(undefined); + __setReadChatSnapshotImplForTests(undefined); + __setWriteChatSnapshotImplForTests(undefined); + __setReplaySessionOutTailImplForTests(undefined); + }); + + const sendPayloadAndWait = async ( + payload: ChatWirePayload + ): Promise => { + await harnessReady; + const before = allRawChunks.length; + const turnComplete = waitForTurnComplete(); + await sendSessionInput(sessionId, { kind: "message", payload }); + await turnComplete; + const rawChunks = allRawChunks.slice(before); + const chunks = rawChunks.filter( + (c) => !isControlChunk(c) + ) as UIMessageChunk[]; + return { chunks, rawChunks }; + }; + + const harness: MockChatAgentHarness = { + chatId, + + async sendMessage(message) { + return sendPayloadAndWait({ + message, + chatId, + trigger: "submit-message", + metadata: clientData, + }); + }, + + async sendRegenerate() { + return sendPayloadAndWait({ + chatId, + trigger: "regenerate-message", + metadata: clientData, + }); + }, + + async sendHeadStart({ messages }) { + return sendPayloadAndWait({ + headStartMessages: messages, + chatId, + trigger: "handover-prepare", + metadata: clientData, + }); + }, + + async sendAction(action) { + return sendPayloadAndWait({ + chatId, + trigger: "action", + action, + metadata: clientData, + }); + }, + + async sendStop(message) { + await harnessReady; + await sendSessionInput(sessionId, { kind: "stop", message }); + }, + + async sendHandover(args) { + await harnessReady; + const before = allRawChunks.length; + const turnComplete = waitForTurnComplete(); + await sendSessionInput(sessionId, { + kind: "handover", + partialAssistantMessage: args.partialAssistantMessage, + messageId: args.messageId, + isFinal: args.isFinal ?? false, + }); + await turnComplete; + const rawChunks = allRawChunks.slice(before); + const chunks = rawChunks.filter((c) => !isControlChunk(c)) as UIMessageChunk[]; + return { chunks, rawChunks }; + }, + + async sendHandoverSkip() { + await harnessReady; + // No turn-complete on skip — the agent exits without firing hooks. + // Send the chunk and wait for the run to finish. + await sendSessionInput(sessionId, { kind: "handover-skip" }); + await Promise.race([ + taskFinished.catch(() => {}), + new Promise((resolve) => setTimeout(resolve, 1000)), + ]); + }, + + seedSnapshot(snapshot) { + seededSnapshot = snapshot; + }, + + seedSessionOutTail(chunks) { + seededReplayChunks = chunks ?? []; + }, + + getSnapshot() { + return lastWrittenSnapshot; + }, + + async close() { + await harnessReady; + + // Send a close trigger wrapped as a `kind: "message"` ChatInputChunk. + // The turn loop checks for this after a successful turn and exits + // cleanly. On error-recovery paths the loop just loops back with + // the close payload, so we also close the session input below to + // unblock any pending once() waiters. + try { + await sendSessionInput(sessionId, { + kind: "message", + payload: { + chatId, + trigger: "close", + }, + }); + } catch { + // best-effort + } + // Resolve any pending once() waiters on the session input with a + // timeout error — that makes waitWithIdleTimeout return + // `{ ok: false }` and the turn loop exits cleanly. + closeSessionInput?.(sessionId); + + // Also abort the run signal so anything downstream (streamText, + // deferred work) unwinds promptly. + runSignal?.abort("close"); + + // Wait for run() to return. The loop's error recovery path will + // see !next.ok and exit. Use a bounded wait so tests never hang. + await Promise.race([ + taskFinished.catch(() => {}), + new Promise((resolve) => setTimeout(resolve, 1000)), + ]); + }, + + get allChunks() { + return allChunks.slice(); + }, + + get allRawChunks() { + return allRawChunks.slice(); + }, + }; + + return harness; +} + +/** + * Reduce a synthetic UIMessageChunk[] sequence into the UIMessage[] that + * the runtime's `replaySessionOutTail` would produce. Splits chunks at + * `start` boundaries and feeds each segment through AI SDK's + * `readUIMessageStream`. The trailing un-finished segment goes through + * `cleanupAbortedParts`. Mirrors the production reducer used in + * `ai.ts:replaySessionOutTail`. + */ +async function reduceChunksToMessages(chunks: UIMessageChunk[]): Promise { + if (chunks.length === 0) return []; + const aiModule = (await import("ai")) as { + readUIMessageStream?: (args: { stream: ReadableStream }) => AsyncIterable; + cleanupAbortedParts?: (msg: UIMessage) => UIMessage; + }; + const readUIMessageStream = aiModule.readUIMessageStream; + const cleanupAbortedParts = aiModule.cleanupAbortedParts; + if (!readUIMessageStream) return []; + + type Segment = { chunks: UIMessageChunk[]; closed: boolean }; + const segments: Segment[] = []; + let current: Segment | undefined; + for (const chunk of chunks) { + if (chunk.type === "start") { + current = { chunks: [chunk], closed: false }; + segments.push(current); + continue; + } + if (!current) { + current = { chunks: [], closed: false }; + segments.push(current); + } + current.chunks.push(chunk); + if (chunk.type === "finish") { + current.closed = true; + current = undefined; + } + } + + const out: UIMessage[] = []; + for (let i = 0; i < segments.length; i++) { + const seg = segments[i]!; + const isTrailing = i === segments.length - 1 && !seg.closed; + const segmentStream = new ReadableStream({ + start(controller) { + for (const c of seg.chunks) controller.enqueue(c); + controller.close(); + }, + }); + let last: UIMessage | undefined; + try { + for await (const snapshot of readUIMessageStream({ stream: segmentStream })) { + last = snapshot; + } + } catch { + // Skip malformed segment — tests can assert by inspecting what makes it through. + continue; + } + if (!last) continue; + if (isTrailing && cleanupAbortedParts) { + const cleaned = cleanupAbortedParts(last); + if (!cleaned.parts || cleaned.parts.length === 0) continue; + out.push(cleaned); + } else { + out.push(last); + } + } + return out; +} diff --git a/packages/trigger-sdk/src/v3/test/setup-catalog.ts b/packages/trigger-sdk/src/v3/test/setup-catalog.ts new file mode 100644 index 00000000000..4dece053b98 --- /dev/null +++ b/packages/trigger-sdk/src/v3/test/setup-catalog.ts @@ -0,0 +1,16 @@ +import { resourceCatalog } from "@trigger.dev/core/v3"; +import { StandardResourceCatalog } from "@trigger.dev/core/v3/workers"; + +/** + * Installs an in-memory `StandardResourceCatalog` and seeds a fake file + * context so task definitions (`task()`, `chat.agent()`, etc.) register + * their run functions where the test harness can look them up. + * + * This is invoked as a side-effect of importing `@trigger.dev/sdk/ai/test`. + * + * Without this, `registerTaskMetadata` short-circuits on a missing + * `_currentFileContext` and tasks silently fail to register. + */ +const catalog = new StandardResourceCatalog(); +resourceCatalog.setGlobalResourceCatalog(catalog); +resourceCatalog.setCurrentFileContext("__test__.ts", "__test__"); diff --git a/packages/trigger-sdk/src/v3/test/test-session-handle.ts b/packages/trigger-sdk/src/v3/test/test-session-handle.ts new file mode 100644 index 00000000000..71bc9d8d7b3 --- /dev/null +++ b/packages/trigger-sdk/src/v3/test/test-session-handle.ts @@ -0,0 +1,268 @@ +import type { + AsyncIterableStream, + PipeStreamResult, + StreamWriteResult, + WriterStreamOptions, +} from "@trigger.dev/core/v3"; +import { ensureReadableStream, ManualWaitpointPromise } from "@trigger.dev/core/v3"; +import { + SessionHandle, + SessionInputChannel, + SessionOutputChannel, + SessionPipeStreamOptions, + SessionSubscribeOptions, +} from "../sessions.js"; + +/** + * Stub for `SessionInputChannel.wait` that skips the apiClient round-trip + * the production path makes via `createSessionStreamWaitpoint`. Without + * this override, every test that exercises the suspend fallback (e.g. + * the `chat.handover` idle-timeout case) throws `ApiClientMissingError` + * because `apiClientManager.clientOrThrow()` runs in a test process that + * has no `TRIGGER_SECRET_KEY`. + * + * The promise resolves with `{ ok: false, error }` when the harness + * aborts its run signal — that mimics production semantics (suspended + * until something happens, returns cleanly on abort) without making a + * network call. + */ +class TestSessionInputChannel extends SessionInputChannel { + constructor(sessionId: string, private readonly getAbortSignal: () => AbortSignal | undefined) { + super(sessionId); + } + + // Override only the `wait` path. `on` / `once` / `peek` / `send` + // continue to flow through the real `sessionStreams` global, which + // the mock task context installs as a `TestSessionStreamManager`. + wait(): ManualWaitpointPromise { + return new ManualWaitpointPromise((resolve: (value: { ok: false; error: Error }) => void) => { + const signal = this.getAbortSignal(); + if (!signal) { + // Harness hasn't wired up its run signal yet — nothing to abort + // on. Stay pending; the run loop should never reach this state + // in practice but we don't want to throw here either. + return; + } + const onAbort = () => { + resolve({ + ok: false, + error: new Error("session.in.wait() aborted by test harness"), + }); + }; + if (signal.aborted) { + onAbort(); + return; + } + signal.addEventListener("abort", onAbort, { once: true }); + }); + } +} + +/** + * Per-session in-memory state collected from `.out` writes during a test. + * Owned by the mock-chat-agent harness; updated by {@link TestSessionOutputChannel}. + */ +export type TestSessionOutState = { + /** Every chunk written to `.out`, in order of write. */ + chunks: unknown[]; + /** Registered write listeners (fired for each chunk). */ + listeners: Set<(chunk: unknown) => void>; +}; + +function notify(state: TestSessionOutState, chunk: unknown): void { + state.chunks.push(chunk); + for (const listener of state.listeners) { + try { + listener(chunk); + } catch { + // Never let a listener error break stream writes + } + } +} + +async function drainInto( + source: AsyncIterable | ReadableStream, + state: TestSessionOutState +): Promise { + const readable = ensureReadableStream(source); + const reader = readable.getReader(); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) return; + notify(state, value); + } + } finally { + try { + reader.releaseLock(); + } catch { + // ignore + } + } +} + +/** + * `.out` channel that captures writes in memory instead of piping to S2. + * Mirrors {@link SessionOutputChannel}'s public shape — `pipe` / `writer` + * / `append` / `read` — so the agent's existing code paths work unchanged. + */ +export class TestSessionOutputChannel extends SessionOutputChannel { + constructor( + sessionId: string, + private readonly state: TestSessionOutState + ) { + super(sessionId); + } + + async append(value: T, _options?: SessionPipeStreamOptions): Promise { + notify(this.state, value); + } + + pipe( + value: AsyncIterable | ReadableStream, + _options?: SessionPipeStreamOptions + ): PipeStreamResult { + const state = this.state; + const readChunks: T[] = []; + let resolveDone!: () => void; + const done = new Promise((resolve) => { + resolveDone = resolve; + }); + + (async () => { + const readable = ensureReadableStream(value); + const reader = readable.getReader(); + try { + while (true) { + const { done: d, value: v } = await reader.read(); + if (d) return; + readChunks.push(v as T); + notify(state, v); + } + } finally { + try { + reader.releaseLock(); + } catch { + // ignore + } + resolveDone(); + } + })().catch(() => { + resolveDone(); + }); + + const replayStream = new ReadableStream({ + async start(controller) { + await done; + for (const chunk of readChunks) controller.enqueue(chunk); + controller.close(); + }, + }); + + const emptyResult: StreamWriteResult = {}; + + return { + get stream(): AsyncIterableStream { + return replayStream as AsyncIterableStream; + }, + waitUntilComplete: async () => { + await done; + return emptyResult; + }, + }; + } + + writer(options: WriterStreamOptions): PipeStreamResult { + let controller!: ReadableStreamDefaultController; + const ongoing: Promise[] = []; + const state = this.state; + + const stream = new ReadableStream({ + start(c) { + controller = c; + }, + }); + + const safeEnqueue = (data: T) => { + try { + controller.enqueue(data); + } catch { + // Stream already closed + } + }; + + try { + const result = options.execute({ + write(part) { + safeEnqueue(part); + notify(state, part); + }, + merge(streamArg) { + ongoing.push( + drainInto(streamArg, state).catch(() => {}) + ); + }, + }); + + if (result) { + ongoing.push(result.catch(() => {})); + } + } catch { + // Swallow — tests can inspect state.chunks + } + + const done: Promise = (async () => { + while (ongoing.length > 0) { + await ongoing.shift(); + } + })().finally(() => { + try { + controller.close(); + } catch { + // Already closed + } + }); + + const emptyResult: StreamWriteResult = {}; + + return { + get stream(): AsyncIterableStream { + return stream as AsyncIterableStream; + }, + waitUntilComplete: async () => { + await done; + return emptyResult; + }, + }; + } + + async read(_options?: SessionSubscribeOptions): Promise> { + throw new Error( + "TestSessionOutputChannel.read() is not supported in the mock-chat-agent harness — " + + "inspect `harness.allChunks` / `harness.allRawChunks` instead." + ); + } +} + +/** + * Construct a {@link SessionHandle} whose `.out` channel captures writes in + * memory and whose `.in` channel routes through the `sessionStreams` + * global for record subscriptions (`on` / `once` / `peek`) but stubs + * `wait()` to skip the apiClient round-trip — see + * {@link TestSessionInputChannel}. + * + * `getAbortSignal` lets the channel observe the harness's run signal so + * `wait()` resolves cleanly on close. Pass a getter (not the signal + * directly) so the channel reads it lazily — the harness creates its + * `AbortController` after the override is installed. + */ +export function createTestSessionHandle( + sessionId: string, + state: TestSessionOutState, + getAbortSignal: () => AbortSignal | undefined = () => undefined +): SessionHandle { + return new SessionHandle(sessionId, { + in: new TestSessionInputChannel(sessionId, getAbortSignal), + out: new TestSessionOutputChannel(sessionId, state), + }); +} diff --git a/packages/trigger-sdk/test/chat-snapshot.test.ts b/packages/trigger-sdk/test/chat-snapshot.test.ts new file mode 100644 index 00000000000..e7421cdbd9a --- /dev/null +++ b/packages/trigger-sdk/test/chat-snapshot.test.ts @@ -0,0 +1,279 @@ +// Import the test entry point first so the resource catalog is installed — +// not strictly required for these helper-level tests, but keeps parity with +// the rest of the test suite and removes a potential foot-gun if a future +// edit introduces a chat.agent({...}) at module scope. +import "../src/v3/test/index.js"; + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { apiClientManager } from "@trigger.dev/core/v3"; +import { + __readChatSnapshotProductionPathForTests as readChatSnapshot, + __writeChatSnapshotProductionPathForTests as writeChatSnapshot, + type ChatSnapshotV1, +} from "../src/v3/ai.js"; + +// ── Helpers ──────────────────────────────────────────────────────────── + +/** + * Build a minimal ChatSnapshotV1 with `count` user messages. Used as the + * production-path test payload — `messages` is the only field the runtime + * inspects beyond `version`. + */ +function buildSnapshot(count = 1): ChatSnapshotV1 { + return { + version: 1, + savedAt: 1_000_000, + messages: Array.from({ length: count }, (_, i) => ({ + id: `m${i}`, + role: "user" as const, + parts: [{ type: "text" as const, text: `hello ${i}` }], + })), + lastOutEventId: "evt-42", + lastOutTimestamp: 2_000_000, + }; +} + +/** + * Stub `apiClientManager.clientOrThrow()` so the helpers see a fake API + * client whose `getPayloadUrl` / `createUploadPayloadUrl` resolve with the + * presigned URLs the test wants. Returns spies for assertion. + */ +function stubApiClient(opts: { + getPayloadUrl?: (filename: string) => Promise<{ presignedUrl: string }>; + createUploadPayloadUrl?: (filename: string) => Promise<{ presignedUrl: string }>; +}) { + const getPayloadUrl = vi.fn( + opts.getPayloadUrl ?? (async (_filename: string) => ({ presignedUrl: "https://example.invalid/get" })) + ); + const createUploadPayloadUrl = vi.fn( + opts.createUploadPayloadUrl ?? + (async (_filename: string) => ({ presignedUrl: "https://example.invalid/put" })) + ); + const fakeClient = { + getPayloadUrl, + createUploadPayloadUrl, + }; + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue( + fakeClient as never + ); + return { getPayloadUrl, createUploadPayloadUrl }; +} + +/** + * Stub global `fetch` so the helpers see whatever Response (or throw) the + * test wants. Returns a spy keyed on the URL passed. + */ +function stubFetch(impl: (url: string, init?: RequestInit) => Promise | Response) { + const spy = vi.fn(impl); + vi.stubGlobal("fetch", spy); + return spy; +} + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("chat snapshot helpers", () => { + // Suppress the runtime's `logger.warn` calls — they pollute output but + // don't change test outcomes. Restored in afterEach. + let warnSpy: ReturnType; + + beforeEach(() => { + warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + }); + + afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + warnSpy.mockRestore(); + }); + + describe("readChatSnapshot", () => { + it("returns the snapshot on a successful GET", async () => { + const { getPayloadUrl } = stubApiClient({}); + const snapshot = buildSnapshot(2); + stubFetch(async () => + new Response(JSON.stringify(snapshot), { + status: 200, + headers: { "content-type": "application/json" }, + }) + ); + + const result = await readChatSnapshot("session-1"); + expect(getPayloadUrl).toHaveBeenCalledWith("sessions/session-1/snapshot.json"); + expect(result).toMatchObject({ + version: 1, + messages: snapshot.messages, + lastOutEventId: "evt-42", + }); + }); + + it("returns undefined on 404 (fresh session, no snapshot yet)", async () => { + stubApiClient({}); + stubFetch(async () => new Response("Not Found", { status: 404 })); + + const result = await readChatSnapshot("missing-session"); + expect(result).toBeUndefined(); + }); + + it("returns undefined on non-404 non-OK (e.g. 500)", async () => { + stubApiClient({}); + stubFetch(async () => new Response("Internal Error", { status: 500 })); + + const result = await readChatSnapshot("flaky-session"); + expect(result).toBeUndefined(); + }); + + it("returns undefined when the response body is malformed JSON", async () => { + stubApiClient({}); + stubFetch(async () => + new Response("not-json-{[", { + status: 200, + headers: { "content-type": "application/json" }, + }) + ); + + const result = await readChatSnapshot("malformed-session"); + expect(result).toBeUndefined(); + }); + + it("returns undefined on version mismatch (forward-compat)", async () => { + stubApiClient({}); + // Future format the current runtime can't decode — runtime ignores it. + const futureSnapshot = { + version: 99, + savedAt: Date.now(), + messages: [], + }; + stubFetch(async () => + new Response(JSON.stringify(futureSnapshot), { + status: 200, + headers: { "content-type": "application/json" }, + }) + ); + + const result = await readChatSnapshot("v99-session"); + expect(result).toBeUndefined(); + }); + + it("returns undefined when `messages` field is missing or wrong type", async () => { + stubApiClient({}); + stubFetch(async () => + new Response(JSON.stringify({ version: 1, savedAt: 1, messages: "not-an-array" }), { + status: 200, + }) + ); + + const result = await readChatSnapshot("bad-shape-session"); + expect(result).toBeUndefined(); + }); + + it("returns undefined when fetch throws (network error)", async () => { + stubApiClient({}); + stubFetch(async () => { + throw new Error("ECONNREFUSED"); + }); + + const result = await readChatSnapshot("offline-session"); + expect(result).toBeUndefined(); + }); + + it("returns undefined when presign call fails", async () => { + stubApiClient({ + getPayloadUrl: async () => { + throw new Error("presign denied"); + }, + }); + // No fetch should fire — presign failed. + const fetchSpy = stubFetch(async () => new Response("nope", { status: 500 })); + + const result = await readChatSnapshot("denied-session"); + expect(result).toBeUndefined(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("returns undefined when the response is not an object", async () => { + stubApiClient({}); + stubFetch(async () => + new Response(JSON.stringify("just-a-string"), { status: 200 }) + ); + + const result = await readChatSnapshot("string-response"); + expect(result).toBeUndefined(); + }); + }); + + describe("writeChatSnapshot", () => { + it("PUTs the snapshot JSON to the presigned URL", async () => { + const { createUploadPayloadUrl } = stubApiClient({}); + const fetchSpy = stubFetch(async () => new Response(null, { status: 200 })); + + const snapshot = buildSnapshot(3); + await writeChatSnapshot("session-2", snapshot); + + expect(createUploadPayloadUrl).toHaveBeenCalledWith("sessions/session-2/snapshot.json"); + expect(fetchSpy).toHaveBeenCalledOnce(); + const [url, init] = fetchSpy.mock.calls[0]!; + expect(url).toBe("https://example.invalid/put"); + expect((init as RequestInit).method).toBe("PUT"); + expect((init as RequestInit).headers).toMatchObject({ + "content-type": "application/json", + }); + // Body is the JSON-stringified snapshot — round-trip to confirm. + const sentBody = JSON.parse((init as RequestInit).body as string); + expect(sentBody).toEqual(snapshot); + }); + + it("returns without throwing on a non-OK PUT response (warns)", async () => { + stubApiClient({}); + stubFetch(async () => new Response("forbidden", { status: 403 })); + + await expect(writeChatSnapshot("forbidden-session", buildSnapshot())).resolves.toBeUndefined(); + }); + + it("returns without throwing on a fetch network error (warns)", async () => { + stubApiClient({}); + stubFetch(async () => { + throw new Error("ETIMEDOUT"); + }); + + await expect(writeChatSnapshot("timeout-session", buildSnapshot())).resolves.toBeUndefined(); + }); + + it("returns without throwing when presign fails (warns)", async () => { + stubApiClient({ + createUploadPayloadUrl: async () => { + throw new Error("presign denied"); + }, + }); + const fetchSpy = stubFetch(async () => new Response(null, { status: 200 })); + + await expect(writeChatSnapshot("denied-session", buildSnapshot())).resolves.toBeUndefined(); + // Presign failed → no PUT attempted. + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("uses the same `snapshotFilename(sessionId)` convention as the read path", async () => { + // Round-trip check: read and write target the same key for a given + // sessionId. The runtime relies on this to make read-after-write + // coherent on subsequent boots. + const { getPayloadUrl } = stubApiClient({ + getPayloadUrl: async () => ({ presignedUrl: "https://example.invalid/get" }), + }); + stubFetch(async () => new Response(null, { status: 404 })); + + // Trigger a read. + await readChatSnapshot("round-trip-session"); + const [readKey] = getPayloadUrl.mock.calls[0]!; + + // Trigger a write to the same session. + const { createUploadPayloadUrl } = stubApiClient({ + createUploadPayloadUrl: async () => ({ presignedUrl: "https://example.invalid/put" }), + }); + stubFetch(async () => new Response(null, { status: 200 })); + await writeChatSnapshot("round-trip-session", buildSnapshot()); + const [writeKey] = createUploadPayloadUrl.mock.calls[0]!; + + expect(readKey).toBe(writeKey); + expect(readKey).toBe("sessions/round-trip-session/snapshot.json"); + }); + }); +}); diff --git a/packages/trigger-sdk/test/merge-by-id.test.ts b/packages/trigger-sdk/test/merge-by-id.test.ts new file mode 100644 index 00000000000..1c0091273cc --- /dev/null +++ b/packages/trigger-sdk/test/merge-by-id.test.ts @@ -0,0 +1,158 @@ +// Plan F.1: pure-function correctness tests for `mergeByIdReplaceWins`, +// the helper that combines `snapshot.messages` with `session.out` replay +// at run boot (plan section B.3). Replay wins on id collision because +// `session.out` carries the freshest representation of an assistant +// message. + +import "../src/v3/test/index.js"; + +import type { UIMessage } from "ai"; +import { describe, expect, it } from "vitest"; +import { __mergeByIdReplaceWinsForTests as mergeByIdReplaceWins } from "../src/v3/ai.js"; + +// ── Helpers ──────────────────────────────────────────────────────────── + +function userMessage(id: string, text: string): UIMessage { + return { + id, + role: "user", + parts: [{ type: "text", text }], + }; +} + +function assistantMessage(id: string, text: string): UIMessage { + return { + id, + role: "assistant", + parts: [{ type: "text", text }], + }; +} + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("mergeByIdReplaceWins", () => { + it("returns a copy of `a` when `b` is empty", () => { + const a = [userMessage("u-1", "hello")]; + const result = mergeByIdReplaceWins(a, []); + expect(result).toEqual(a); + // Verify it's a copy (mutating result shouldn't touch a). + result.push(assistantMessage("a-1", "extra")); + expect(a).toHaveLength(1); + }); + + it("returns a copy of `b` when `a` is empty", () => { + const b = [assistantMessage("a-1", "world")]; + const result = mergeByIdReplaceWins([], b); + expect(result).toEqual(b); + result.push(userMessage("u-extra", "extra")); + expect(b).toHaveLength(1); + }); + + it("returns [] when both inputs are empty", () => { + expect(mergeByIdReplaceWins([], [])).toEqual([]); + }); + + it("appends fresh ids from `b` after `a`'s entries", () => { + const a = [userMessage("u-1", "hi")]; + const b = [assistantMessage("a-1", "ok")]; + const result = mergeByIdReplaceWins(a, b); + expect(result.map((m) => m.id)).toEqual(["u-1", "a-1"]); + expect(result[0]!.role).toBe("user"); + expect(result[1]!.role).toBe("assistant"); + }); + + it("replaces by id when `b` has a colliding entry — replay wins", () => { + const a = [ + userMessage("u-1", "hi"), + assistantMessage("a-1", "stale-version"), + ]; + const b = [assistantMessage("a-1", "fresh-version")]; + const result = mergeByIdReplaceWins(a, b); + expect(result).toHaveLength(2); + expect(result[1]!.id).toBe("a-1"); + expect((result[1]!.parts[0] as { text: string }).text).toBe("fresh-version"); + }); + + it("preserves order from `a` even when entries are replaced", () => { + const a = [ + userMessage("u-1", "first"), + assistantMessage("a-1", "stale"), + userMessage("u-2", "second"), + assistantMessage("a-2", "also-stale"), + ]; + const b = [ + assistantMessage("a-1", "fresh-1"), + assistantMessage("a-2", "fresh-2"), + ]; + const result = mergeByIdReplaceWins(a, b); + expect(result.map((m) => m.id)).toEqual(["u-1", "a-1", "u-2", "a-2"]); + expect((result[1]!.parts[0] as { text: string }).text).toBe("fresh-1"); + expect((result[3]!.parts[0] as { text: string }).text).toBe("fresh-2"); + }); + + it("appends `b` entries with no id collision after the merged set", () => { + const a = [userMessage("u-1", "first")]; + const b = [ + assistantMessage("a-1", "reply-1"), + userMessage("u-2", "second"), + assistantMessage("a-2", "reply-2"), + ]; + const result = mergeByIdReplaceWins(a, b); + expect(result.map((m) => m.id)).toEqual(["u-1", "a-1", "u-2", "a-2"]); + }); + + it("treats messages without an id as always-append (no collision possible)", () => { + const a = [ + userMessage("u-1", "first"), + // Synthetic message missing the id field — should append, never replace. + { id: "" as string, role: "assistant", parts: [{ type: "text", text: "no-id-a" }] } as UIMessage, + ]; + const b = [ + { id: "" as string, role: "assistant", parts: [{ type: "text", text: "no-id-b" }] } as UIMessage, + ]; + const result = mergeByIdReplaceWins(a, b); + expect(result).toHaveLength(3); + // Both empty-id messages survive — no merge happens. + const noIdParts = result + .filter((m) => m.id === "") + .map((m) => (m.parts[0] as { text: string }).text); + expect(noIdParts).toEqual(["no-id-a", "no-id-b"]); + }); + + it("handles consecutive replays of the same id in `b` — last one wins", () => { + // Edge case: `b` has two entries with the same id (shouldn't happen + // for assistants in practice, but the helper must be deterministic). + const a = [assistantMessage("a-1", "v0")]; + const b = [assistantMessage("a-1", "v1"), assistantMessage("a-1", "v2")]; + const result = mergeByIdReplaceWins(a, b); + expect(result).toHaveLength(1); + expect((result[0]!.parts[0] as { text: string }).text).toBe("v2"); + }); + + it("preserves user messages (only assistants come from replay) — semantic check", () => { + // The runtime contract: `session.out` contains assistant chunks only, + // so `b` should never contain user messages. If it does (defensively), + // the merge still works — but we lock down the typical pattern here. + const a = [ + userMessage("u-1", "first"), + assistantMessage("a-1", "stale"), + userMessage("u-2", "second"), + ]; + const b = [assistantMessage("a-1", "fresh")]; + const result = mergeByIdReplaceWins(a, b); + // User messages from snapshot survive untouched. + expect(result.filter((m) => m.role === "user").map((m) => m.id)).toEqual(["u-1", "u-2"]); + }); + + it("does not mutate either input array", () => { + const a = [userMessage("u-1", "hi"), assistantMessage("a-1", "stale")]; + const b = [assistantMessage("a-1", "fresh"), userMessage("u-2", "next")]; + const aSnapshot = JSON.stringify(a); + const bSnapshot = JSON.stringify(b); + + mergeByIdReplaceWins(a, b); + + expect(JSON.stringify(a)).toBe(aSnapshot); + expect(JSON.stringify(b)).toBe(bSnapshot); + }); +}); diff --git a/packages/trigger-sdk/test/mockChatAgent.test.ts b/packages/trigger-sdk/test/mockChatAgent.test.ts new file mode 100644 index 00000000000..1d1c6eec0fb --- /dev/null +++ b/packages/trigger-sdk/test/mockChatAgent.test.ts @@ -0,0 +1,1441 @@ +// Import the test harness FIRST — this installs the resource catalog so +// `chat.agent()` calls below register their task functions correctly. +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { describe, expect, it, vi } from "vitest"; +import { chat } from "../src/v3/ai.js"; +import { locals } from "@trigger.dev/core/v3"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; + +// ── Helpers ──────────────────────────────────────────────────────────── + +function userMessage(text: string, id = "u-" + Math.random().toString(36).slice(2)) { + return { + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], + }; +} + +function textStream(text: string) { + const chunks: LanguageModelV3StreamPart[] = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, + }, + }, + ]; + return simulateReadableStream({ chunks }); +} + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("mockChatAgent", () => { + it("throws when no agent is registered with the given id", () => { + expect(() => mockChatAgent({ id: "does-not-exist" })).toThrow(/no task registered/); + }); + + it("drives a chat.agent through a single turn and captures output chunks", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("hello world") }), + }); + + const agent = chat.agent({ + id: "mockChatAgent.basic-flow", + run: async ({ messages, signal }) => { + return streamText({ + model, + messages, + abortSignal: signal, + }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-basic" }); + try { + const turn = await harness.sendMessage(userMessage("hi")); + + const textDeltas = turn.chunks + .filter((c) => c.type === "text-delta") + .map((c) => (c as { delta: string }).delta) + .join(""); + expect(textDeltas).toBe("hello world"); + } finally { + await harness.close(); + } + }); + + it("fires onTurnStart and onTurnComplete hooks in order", async () => { + const events: string[] = []; + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("hi") }), + }); + + const agent = chat.agent({ + id: "mockChatAgent.hook-order", + onChatStart: async () => { + events.push("onChatStart"); + }, + onTurnStart: async () => { + events.push("onTurnStart"); + }, + onBeforeTurnComplete: async () => { + events.push("onBeforeTurnComplete"); + }, + onTurnComplete: async () => { + events.push("onTurnComplete"); + }, + run: async ({ messages, signal }) => { + events.push("run"); + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-hooks" }); + try { + await harness.sendMessage(userMessage("hello")); + // onTurnComplete may fire after the turn-complete chunk is written, + // so give it a tick to run before we assert. + await new Promise((r) => setTimeout(r, 20)); + expect(events).toEqual([ + "onChatStart", + "onTurnStart", + "run", + "onBeforeTurnComplete", + "onTurnComplete", + ]); + } finally { + await harness.close(); + } + }); + + it("can send multiple messages across turns", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("reply") }), + }); + + const seenMessages: number[] = []; + const agent = chat.agent({ + id: "mockChatAgent.multi-turn", + run: async ({ messages, signal }) => { + seenMessages.push(messages.length); + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-multi" }); + try { + await harness.sendMessage(userMessage("first")); + await harness.sendMessage(userMessage("second")); + await harness.sendMessage(userMessage("third")); + + // Each turn sees an accumulator growing by (user + assistant) * turn + // Turn 1: just the user message + // Turn 2: user + assistant + user = 3 messages + // Turn 3: 5 messages + expect(seenMessages).toEqual([1, 3, 5]); + } finally { + await harness.close(); + } + }); + + it("invokes hydrateMessages on every turn with incoming wire messages", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + const hydrateSpy = vi.fn(async ({ incomingMessages }) => { + // Echo back whatever the frontend sent + return incomingMessages; + }); + + const agent = chat.agent({ + id: "mockChatAgent.hydrate", + hydrateMessages: hydrateSpy, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-hydrate" }); + try { + await harness.sendMessage(userMessage("hi", "u-first")); + + expect(hydrateSpy).toHaveBeenCalledTimes(1); + const call = hydrateSpy.mock.calls[0]![0] as { incomingMessages: { id: string }[] }; + expect(call.incomingMessages).toHaveLength(1); + expect(call.incomingMessages[0]!.id).toBe("u-first"); + } finally { + await harness.close(); + } + }); + + it("merges HITL tool answer onto head assistant when AI SDK regenerates the id", async () => { + // Regression for TRI-9137: customers (Arena AI) report that the AI SDK + // intermittently mints a fresh id on `addToolOutput` resume, breaking + // id-based dedup. Our SDK records `toolCallId → head messageId` whenever + // an assistant with tool parts lands in the accumulator and uses that + // map as a fallback in the merge so a fresh-id incoming still attaches + // to the right head. + const { z } = await import("zod"); + const { tool } = await import("ai"); + + const askUserTool = tool({ + description: "Ask the user a question.", + inputSchema: z.object({ question: z.string() }), + // No execute — HITL round-trip via addToolOutput. + }); + + const HEAD_TOOL_CALL_ID = "tc_regression_9137"; + + // Turn 1: model emits a tool-call for askUser. No text, no finish-reason + // logic beyond `tool-calls`. Agent's response will carry a tool-input- + // available part with HEAD_TOOL_CALL_ID. + const turn1Stream = simulateReadableStream({ + chunks: [ + { type: "tool-input-start", id: HEAD_TOOL_CALL_ID, toolName: "askUser" }, + { + type: "tool-input-delta", + id: HEAD_TOOL_CALL_ID, + delta: JSON.stringify({ question: "what color?" }), + }, + { type: "tool-input-end", id: HEAD_TOOL_CALL_ID }, + { + type: "tool-call", + toolCallId: HEAD_TOOL_CALL_ID, + toolName: "askUser", + input: JSON.stringify({ question: "what color?" }), + }, + { + type: "finish", + finishReason: { unified: "tool-calls", raw: "tool_calls" }, + usage: { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 0, reasoning: undefined }, + }, + }, + ] as LanguageModelV3StreamPart[], + }); + + // Turn 2: model produces a final text response — exercises the post-HITL + // continuation streamText after the tool answer is merged in. + const turn2Stream = textStream("blue is great"); + + let callIdx = 0; + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: callIdx++ === 0 ? turn1Stream : turn2Stream }), + }); + + const turnsSeen: { turn: number; uiMessages: any[] }[] = []; + + const agent = chat.agent({ + id: "mockChatAgent.hitl-id-regen", + onTurnComplete: async ({ turn, uiMessages }) => { + turnsSeen.push({ + turn, + uiMessages: uiMessages.map((m) => ({ + id: m.id, + role: m.role, + toolStates: (m.parts ?? []) + .filter((p: any) => typeof p?.toolCallId === "string") + .map((p: any) => ({ toolCallId: p.toolCallId, state: p.state })), + })), + }); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, tools: { askUser: askUserTool }, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-hitl-id-regen" }); + try { + // Turn 1: user message → agent emits tool-input-available for askUser + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 50)); + + // Capture the head assistant id the agent produced. + const turn1 = turnsSeen.at(-1); + const headAssistant = turn1?.uiMessages.find( + (m) => m.role === "assistant" && m.toolStates.length > 0 + ); + expect(headAssistant?.id).toBeTruthy(); + const HEAD_ID = headAssistant!.id as string; + + // Turn 2: simulate AI SDK regenerating the assistant id on + // addToolOutput resume — fresh id, but the same toolCallId in + // tool-output-available state. + const FRESH_ID = "regenerated-by-ai-sdk-" + Math.random().toString(36).slice(2); + const toolAnswerMessage = { + id: FRESH_ID, + role: "assistant" as const, + parts: [ + { + type: "tool-askUser", + toolCallId: HEAD_TOOL_CALL_ID, + state: "output-available" as const, + input: { question: "what color?" }, + output: { color: "blue" }, + }, + ], + }; + await harness.sendMessage(toolAnswerMessage as any); + await new Promise((r) => setTimeout(r, 50)); + + // The merge must rewrite FRESH_ID back to HEAD_ID via the toolCallId + // map, attaching the tool answer to the existing head — no duplicate. + const turn2 = turnsSeen.at(-1); + expect(turn2).toBeTruthy(); + const assistantsWithToolCall = turn2!.uiMessages.filter( + (m) => + m.role === "assistant" && + m.toolStates.some((t: any) => t.toolCallId === HEAD_TOOL_CALL_ID) + ); + expect(assistantsWithToolCall).toHaveLength(1); + expect(assistantsWithToolCall[0]!.id).toBe(HEAD_ID); + expect(turn2!.uiMessages.find((m) => m.id === FRESH_ID)).toBeUndefined(); + } finally { + await harness.close(); + } + }); + + it("routes custom actions through actionSchema + onAction", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + const onActionSpy = vi.fn(); + + const { z } = await import("zod"); + const agent = chat.agent({ + id: "mockChatAgent.actions", + actionSchema: z.object({ + type: z.literal("undo"), + }), + onAction: async (event) => { + onActionSpy(event.action); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-action" }); + try { + await harness.sendMessage(userMessage("start")); + await harness.sendAction({ type: "undo" }); + + expect(onActionSpy).toHaveBeenCalledWith({ type: "undo" }); + } finally { + await harness.close(); + } + }); + + it("actions returning void do not fire turn hooks or call run()", async () => { + const onChatStart = vi.fn(); + const onTurnStart = vi.fn(); + const onBeforeTurnComplete = vi.fn(); + const onTurnComplete = vi.fn(); + const onAction = vi.fn(); + const runSpy = vi.fn(); + const model = new MockLanguageModelV3({ + doStream: async () => { + runSpy(); + return { stream: textStream("nope") }; + }, + }); + + const { z } = await import("zod"); + const agent = chat.agent({ + id: "mockChatAgent.actions.void", + actionSchema: z.object({ type: z.literal("undo") }), + onChatStart, + onTurnStart, + onBeforeTurnComplete, + onTurnComplete, + onAction: async (...args) => { + onAction(...args); + // void → side-effect only + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-void-action" }); + try { + // Bootstrap with a message so the message-turn hooks fire once. + await harness.sendMessage(userMessage("hi")); + // sendMessage resolves on `trigger:turn-complete`, but onTurnComplete + // fires as a separate microtask after — let it settle before snapshotting. + await new Promise((r) => setTimeout(r, 50)); + + // Snapshot call counts after the bootstrap — we'll assert these + // don't change for the action below. + const baselineRun = runSpy.mock.calls.length; + const baselineChatStart = onChatStart.mock.calls.length; + const baselineTurnStart = onTurnStart.mock.calls.length; + const baselineBeforeComplete = onBeforeTurnComplete.mock.calls.length; + const baselineComplete = onTurnComplete.mock.calls.length; + + const actionTurn = await harness.sendAction({ type: "undo" }); + await new Promise((r) => setTimeout(r, 50)); + + // onAction fired exactly once; no turn hooks fired; run() / LLM did not. + expect(onAction).toHaveBeenCalledTimes(1); + expect(runSpy.mock.calls.length).toBe(baselineRun); + expect(onChatStart.mock.calls.length).toBe(baselineChatStart); + expect(onTurnStart.mock.calls.length).toBe(baselineTurnStart); + expect(onBeforeTurnComplete.mock.calls.length).toBe(baselineBeforeComplete); + expect(onTurnComplete.mock.calls.length).toBe(baselineComplete); + + // Stream still terminates cleanly with trigger:turn-complete so + // the frontend's useChat transitions back to ready. + const sawTurnComplete = actionTurn.rawChunks.some( + (c) => + typeof c === "object" && + c !== null && + (c as { type?: string }).type === "trigger:turn-complete" + ); + expect(sawTurnComplete).toBe(true); + } finally { + await harness.close(); + } + }); + + it("actions returning a stream pipe the response without firing turn hooks", async () => { + const onTurnStart = vi.fn(); + const onTurnComplete = vi.fn(); + const actionModel = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("regenerated") }), + }); + const turnModel = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("normal-response") }), + }); + + const { z } = await import("zod"); + const agent = chat.agent({ + id: "mockChatAgent.actions.stream", + actionSchema: z.object({ type: z.literal("regenerate") }), + onTurnStart, + onTurnComplete, + onAction: async ({ messages }) => { + return streamText({ model: actionModel, messages }); + }, + run: async ({ messages, signal }) => { + return streamText({ model: turnModel, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-stream-action" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 50)); + const baselineTurnStart = onTurnStart.mock.calls.length; + const baselineTurnComplete = onTurnComplete.mock.calls.length; + + const actionTurn = await harness.sendAction({ type: "regenerate" }); + await new Promise((r) => setTimeout(r, 50)); + + // No turn hooks fired during the action. + expect(onTurnStart.mock.calls.length).toBe(baselineTurnStart); + expect(onTurnComplete.mock.calls.length).toBe(baselineTurnComplete); + + // Action's streamText output landed on the response. + const text = actionTurn.chunks + .filter((c) => c.type === "text-delta") + .map((c) => (c as { delta: string }).delta) + .join(""); + expect(text).toBe("regenerated"); + } finally { + await harness.close(); + } + }); + + it("warns once and emits turn-complete when an action arrives without onAction", async () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + const runSpy = vi.fn(); + const model = new MockLanguageModelV3({ + doStream: async () => { + runSpy(); + return { stream: textStream("nope") }; + }, + }); + + const { z } = await import("zod"); + const agent = chat.agent({ + id: "mockChatAgent.actions.no-handler", + actionSchema: z.object({ type: z.literal("undo") }), + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-no-handler" }); + try { + await harness.sendMessage(userMessage("hi")); + const baselineRun = runSpy.mock.calls.length; + + const actionTurn = await harness.sendAction({ type: "undo" }); + + // No additional model call; console.warn fired with our marker text. + expect(runSpy.mock.calls.length).toBe(baselineRun); + expect( + warnSpy.mock.calls.some((args) => + (args[0] as string).includes("no `onAction` handler") + ) + ).toBe(true); + + const sawTurnComplete = actionTurn.rawChunks.some( + (c) => + typeof c === "object" && + c !== null && + (c as { type?: string }).type === "trigger:turn-complete" + ); + expect(sawTurnComplete).toBe(true); + } finally { + await harness.close(); + warnSpy.mockRestore(); + } + }); + + it("passes clientData through to run() and hooks", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + let capturedClientData: unknown; + const agent = chat.agent({ + id: "mockChatAgent.client-data", + run: async ({ messages, clientData, signal }) => { + capturedClientData = clientData; + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "test-client-data", + clientData: { userId: "u1", role: "admin" }, + }); + try { + await harness.sendMessage(userMessage("hi")); + expect(capturedClientData).toEqual({ userId: "u1", role: "admin" }); + } finally { + await harness.close(); + } + }); + + it("chat.endRun() exits the loop after the current turn", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("bye") }), + }); + + let turnCount = 0; + const agent = chat.agent({ + id: "mockChatAgent.end-run", + run: async ({ messages, signal }) => { + turnCount++; + chat.endRun(); + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-end-run" }); + try { + await harness.sendMessage(userMessage("hello")); + // Give the loop a tick to exit after the turn-complete chunk + await new Promise((r) => setTimeout(r, 50)); + expect(turnCount).toBe(1); + // Subsequent sends after endRun should not produce another run — the + // loop has exited. We can't easily assert this via sendMessage (it + // would block waiting for turn-complete), but we can verify the task + // has finished. + } finally { + // close() is a no-op here since the task already exited, but call + // for symmetry with other tests. + await harness.close(); + } + }); + + it("exposes finishReason on the onTurnComplete event", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("hi") }), + }); + + let seenReason: string | undefined; + const agent = chat.agent({ + id: "mockChatAgent.finish-reason", + onTurnComplete: async ({ finishReason }) => { + seenReason = finishReason; + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-finish-reason" }); + try { + await harness.sendMessage(userMessage("hello")); + await new Promise((r) => setTimeout(r, 20)); + expect(seenReason).toBe("stop"); + } finally { + await harness.close(); + } + }); + + it("seeds locals before run() via setupLocals (DI pattern)", async () => { + type FakeDb = { findUser(id: string): Promise<{ id: string; name: string }> }; + const dbKey = locals.create("test-db"); + + const fakeDb: FakeDb = { + findUser: async (id) => ({ id, name: `user-${id}` }), + }; + + let userInHook: { id: string; name: string } | undefined; + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + const agent = chat.agent({ + id: "mockChatAgent.locals-di", + hydrateMessages: async ({ incomingMessages }) => { + const db = locals.getOrThrow(dbKey); + userInHook = await db.findUser("u-1"); + return incomingMessages; + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "test-locals-di", + setupLocals: ({ set }) => { + set(dbKey, fakeDb); + }, + }); + try { + await harness.sendMessage(userMessage("hi")); + expect(userInHook).toEqual({ id: "u-1", name: "user-u-1" }); + } finally { + await harness.close(); + } + }); + + describe("chat.history read primitives", () => { + // These tests drive a chat.agent through realistic streams and read + // chat.history inside hooks/tools where the accumulator is in the + // expected state. The pure walks themselves are exercised end-to-end + // rather than via direct internal access. + + function toolCallStream(opts: { toolCallId: string; toolName: string; input: object }) { + return simulateReadableStream({ + chunks: [ + { type: "tool-input-start", id: opts.toolCallId, toolName: opts.toolName }, + { type: "tool-input-delta", id: opts.toolCallId, delta: JSON.stringify(opts.input) }, + { type: "tool-input-end", id: opts.toolCallId }, + { + type: "tool-call", + toolCallId: opts.toolCallId, + toolName: opts.toolName, + input: JSON.stringify(opts.input), + }, + { + type: "finish", + finishReason: { unified: "tool-calls", raw: "tool_calls" }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 0, reasoning: undefined }, + }, + }, + ] as LanguageModelV3StreamPart[], + }); + } + + it("getPendingToolCalls returns input-available parts on the leaf assistant", async () => { + const { z } = await import("zod"); + const { tool } = await import("ai"); + const askUser = tool({ + description: "Ask the user.", + inputSchema: z.object({ q: z.string() }), + }); + const TC = "tc_pending_1"; + const model = new MockLanguageModelV3({ + doStream: async () => ({ + stream: toolCallStream({ toolCallId: TC, toolName: "askUser", input: { q: "?" } }), + }), + }); + + let pending: any; + const agent = chat.agent({ + id: "mockChatAgent.history.pending", + onTurnComplete: async () => { + pending = chat.history.getPendingToolCalls(); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, tools: { askUser }, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-pending" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 50)); + + expect(pending).toHaveLength(1); + expect(pending[0]).toMatchObject({ toolCallId: TC, toolName: "askUser" }); + expect(typeof pending[0].messageId).toBe("string"); + } finally { + await harness.close(); + } + }); + + it("getPendingToolCalls returns [] when the leaf assistant has no pending tool calls", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("hello") }), + }); + let pending: any; + const agent = chat.agent({ + id: "mockChatAgent.history.pending-empty", + onTurnComplete: async () => { + pending = chat.history.getPendingToolCalls(); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-pending-empty" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 50)); + expect(pending).toEqual([]); + } finally { + await harness.close(); + } + }); + + it("getResolvedToolCalls walks all messages after a HITL answer lands", async () => { + const { z } = await import("zod"); + const { tool } = await import("ai"); + const askUser = tool({ + description: "Ask the user.", + inputSchema: z.object({ q: z.string() }), + }); + const TC = "tc_resolved_1"; + + let callIdx = 0; + const model = new MockLanguageModelV3({ + doStream: async () => ({ + stream: + callIdx++ === 0 + ? toolCallStream({ toolCallId: TC, toolName: "askUser", input: { q: "?" } }) + : textStream("done"), + }), + }); + + const turnsResolved: any[] = []; + const agent = chat.agent({ + id: "mockChatAgent.history.resolved", + onTurnComplete: async () => { + turnsResolved.push(chat.history.getResolvedToolCalls()); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, tools: { askUser }, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-resolved" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 50)); + + // Send a HITL tool answer that resolves TC. The merge attaches the + // output to the head assistant — it now shows in `output-available` + // state. + const toolAnswer = { + id: "ai-sdk-fresh-id", + role: "assistant" as const, + parts: [ + { + type: "tool-askUser", + toolCallId: TC, + state: "output-available" as const, + input: { q: "?" }, + output: { answer: "hi" }, + }, + ], + }; + await harness.sendMessage(toolAnswer as any); + await new Promise((r) => setTimeout(r, 50)); + + // After turn 1 only, no tool call is resolved yet (input-available). + // After the HITL answer is merged, the merged head shows the tool + // in `output-available` — getResolvedToolCalls reflects that. + const last = turnsResolved.at(-1) ?? []; + expect(last).toHaveLength(1); + expect(last[0]).toMatchObject({ toolCallId: TC, toolName: "askUser" }); + } finally { + await harness.close(); + } + }); + + it("extractNewToolResults dedups against already-resolved toolCallIds", async () => { + // Pure-function smoke test: feed a synthetic chain via a tool that + // calls extractNewToolResults() during execution. The chain is + // overridden via chat.history.set() inside run(), so we control + // exactly what's in scope. + let extracted: any; + const agent = chat.agent({ + id: "mockChatAgent.history.extract", + run: async ({ messages, signal }) => { + chat.history.set([ + { + id: "a-seed", + role: "assistant", + parts: [ + { + type: "tool-askUser", + toolCallId: "tc-1", + state: "output-available", + input: { q: "?" }, + output: { color: "red" }, + }, + ], + } as any, + { id: "u-1", role: "user", parts: [{ type: "text", text: "u" }] } as any, + ]); + + const incoming = { + id: "a-incoming", + role: "assistant" as const, + parts: [ + { + type: "tool-askUser", + toolCallId: "tc-1", + state: "output-available" as const, + input: { q: "?" }, + output: { color: "red" }, + }, + { + type: "tool-search", + toolCallId: "tc-2", + state: "output-available" as const, + input: { q: "x" }, + output: { hits: 7 }, + }, + { + type: "tool-search", + toolCallId: "tc-err", + state: "output-error" as const, + input: { q: "y" }, + errorText: "boom", + }, + ], + }; + extracted = chat.history.extractNewToolResults(incoming as any); + + return streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }), + messages, + abortSignal: signal, + }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-extract" }); + try { + await harness.sendMessage(userMessage("kick")); + await new Promise((r) => setTimeout(r, 50)); + + expect(extracted).toEqual([ + { toolCallId: "tc-2", toolName: "search", output: { hits: 7 } }, + { toolCallId: "tc-err", toolName: "search", output: undefined, errorText: "boom" }, + ]); + } finally { + await harness.close(); + } + }); + + it("findMessage returns the message by id, or undefined when missing", async () => { + let foundUser: any; + let foundAssistant: any; + let missing: any; + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + const agent = chat.agent({ + id: "mockChatAgent.history.find", + onTurnComplete: async ({ uiMessages }) => { + // Locate ids the agent actually produced/saw, then probe findMessage. + const userId = uiMessages.find((m) => m.role === "user")?.id ?? "u-fixed"; + const asstId = uiMessages.find((m) => m.role === "assistant")?.id; + foundUser = chat.history.findMessage(userId); + foundAssistant = asstId ? chat.history.findMessage(asstId) : undefined; + missing = chat.history.findMessage("definitely-not-here"); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-find" }); + try { + await harness.sendMessage(userMessage("hello", "u-fixed")); + await new Promise((r) => setTimeout(r, 50)); + + expect(foundUser?.id).toBe("u-fixed"); + expect(foundAssistant).toBeTruthy(); + expect(missing).toBeUndefined(); + } finally { + await harness.close(); + } + }); + + it("extractNewToolResults dedups against a real-stream-built chain", async () => { + // Build the chain through real model streams (no chat.history.set seed) + // and assert extractNewToolResults compares against the post-merge state. + const { z } = await import("zod"); + const { tool } = await import("ai"); + const askUser = tool({ + description: "Ask the user.", + inputSchema: z.object({ q: z.string() }), + }); + const TC = "tc_real_chain_1"; + + let callIdx = 0; + const model = new MockLanguageModelV3({ + doStream: async () => ({ + stream: + callIdx++ === 0 + ? toolCallStream({ toolCallId: TC, toolName: "askUser", input: { q: "?" } }) + : textStream("done"), + }), + }); + + let extractedAgainstRealChain: any; + const agent = chat.agent({ + id: "mockChatAgent.history.extract-real", + onTurnComplete: async () => { + // After the HITL answer turn, the chain has TC resolved. An + // incoming "echo" message carrying TC again should yield []. + // A second new TC should yield exactly one entry. + const incoming = { + id: "echo", + role: "assistant" as const, + parts: [ + { + type: "tool-askUser", + toolCallId: TC, + state: "output-available" as const, + input: { q: "?" }, + output: { answer: "hi" }, + }, + { + type: "tool-askUser", + toolCallId: "tc_real_chain_2", + state: "output-available" as const, + input: { q: "second" }, + output: { answer: "yes" }, + }, + ], + }; + extractedAgainstRealChain = chat.history.extractNewToolResults(incoming as any); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, tools: { askUser }, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-extract-real" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 50)); + // HITL answer for TC, lands via the runtime merger. + const toolAnswer = { + id: "ai-sdk-fresh-id-real", + role: "assistant" as const, + parts: [ + { + type: "tool-askUser", + toolCallId: TC, + state: "output-available" as const, + input: { q: "?" }, + output: { answer: "hi" }, + }, + ], + }; + await harness.sendMessage(toolAnswer as any); + await new Promise((r) => setTimeout(r, 50)); + + expect(extractedAgainstRealChain).toEqual([ + { toolCallId: "tc_real_chain_2", toolName: "askUser", output: { answer: "yes" } }, + ]); + } finally { + await harness.close(); + } + }); + + it("extractNewToolResults surfaces output-error parts via the runtime merger", async () => { + // The runtime merges incoming tool-answer messages onto the head + // assistant via the toolCallId map. Here we send an answer in + // `output-error` state and verify (a) getResolvedToolCalls reports + // it, and (b) extractNewToolResults emits it with errorText set. + const { z } = await import("zod"); + const { tool } = await import("ai"); + const search = tool({ + description: "Search.", + inputSchema: z.object({ q: z.string() }), + }); + const TC = "tc_err_via_merger"; + + let callIdx = 0; + const model = new MockLanguageModelV3({ + doStream: async () => ({ + stream: + callIdx++ === 0 + ? toolCallStream({ toolCallId: TC, toolName: "search", input: { q: "x" } }) + : textStream("noted"), + }), + }); + + let resolved: any; + let extracted: any; + const agent = chat.agent({ + id: "mockChatAgent.history.extract-error", + onTurnComplete: async () => { + resolved = chat.history.getResolvedToolCalls(); + // An echo carrying the same error toolCallId — should NOT surface + // as new because it's already resolved on the chain. + const echo = { + id: "echo-err", + role: "assistant" as const, + parts: [ + { + type: "tool-search", + toolCallId: TC, + state: "output-error" as const, + input: { q: "x" }, + errorText: "boom", + }, + ], + }; + extracted = chat.history.extractNewToolResults(echo as any); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, tools: { search }, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-extract-error" }); + try { + await harness.sendMessage(userMessage("kick")); + await new Promise((r) => setTimeout(r, 50)); + // HITL answer arriving as output-error. + const errAnswer = { + id: "ai-sdk-err-fresh", + role: "assistant" as const, + parts: [ + { + type: "tool-search", + toolCallId: TC, + state: "output-error" as const, + input: { q: "x" }, + errorText: "boom", + }, + ], + }; + await harness.sendMessage(errAnswer as any); + await new Promise((r) => setTimeout(r, 50)); + + expect(resolved).toHaveLength(1); + expect(resolved[0]).toMatchObject({ toolCallId: TC, toolName: "search" }); + // Echo of the same error toolCallId is already resolved → [] + expect(extracted).toEqual([]); + } finally { + await harness.close(); + } + }); + + it("extractNewToolResults handles a multi-tool message where only one is new", async () => { + // Pure-helper edge: incoming message has two tool parts with the + // same toolName but different toolCallIds — one already resolved + // on the chain, one fresh. Only the fresh one should surface. + let extracted: any; + const agent = chat.agent({ + id: "mockChatAgent.history.extract-multi", + run: async ({ messages, signal }) => { + chat.history.set([ + { + id: "a-seed", + role: "assistant", + parts: [ + { + type: "tool-search", + toolCallId: "tc-old", + state: "output-available", + input: { q: "old" }, + output: { hits: 1 }, + }, + ], + } as any, + { id: "u-1", role: "user", parts: [{ type: "text", text: "u" }] } as any, + ]); + + const incoming = { + id: "a-incoming", + role: "assistant" as const, + parts: [ + // Same tool, already-resolved id — should be filtered. + { + type: "tool-search", + toolCallId: "tc-old", + state: "output-available" as const, + input: { q: "old" }, + output: { hits: 1 }, + }, + // Same tool, fresh id — should surface. + { + type: "tool-search", + toolCallId: "tc-new", + state: "output-available" as const, + input: { q: "new" }, + output: { hits: 9 }, + }, + // Duplicate of tc-new in the same message — must collapse + // to a single emission (within-message dedup). + { + type: "tool-search", + toolCallId: "tc-new", + state: "output-available" as const, + input: { q: "new" }, + output: { hits: 9 }, + }, + ], + }; + extracted = chat.history.extractNewToolResults(incoming as any); + + return streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }), + messages, + abortSignal: signal, + }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-extract-multi" }); + try { + await harness.sendMessage(userMessage("kick")); + await new Promise((r) => setTimeout(r, 50)); + + expect(extracted).toEqual([ + { toolCallId: "tc-new", toolName: "search", output: { hits: 9 } }, + ]); + } finally { + await harness.close(); + } + }); + + it("getPendingToolCalls still returns the assistant's pending calls when a user message follows", async () => { + // Edge: the chain is [assistant(input-available), user]. The most + // recent assistant is the one with the pending tool call, even + // though the strict tail is a user message. The walk-back semantic + // means pending stays pending until the assistant is mutated. + let pendingAfterUser: any; + const agent = chat.agent({ + id: "mockChatAgent.history.pending-after-user", + run: async ({ messages, signal }) => { + chat.history.set([ + { + id: "a-pending", + role: "assistant", + parts: [ + { + type: "tool-askUser", + toolCallId: "tc-still-pending", + state: "input-available", + input: { q: "?" }, + }, + ], + } as any, + { id: "u-after", role: "user", parts: [{ type: "text", text: "anyway..." }] } as any, + ]); + pendingAfterUser = chat.history.getPendingToolCalls(); + return streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }), + messages, + abortSignal: signal, + }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-pending-after-user" }); + try { + await harness.sendMessage(userMessage("kick")); + await new Promise((r) => setTimeout(r, 50)); + + expect(pendingAfterUser).toEqual([ + { toolCallId: "tc-still-pending", toolName: "askUser", messageId: "a-pending" }, + ]); + } finally { + await harness.close(); + } + }); + + it("getChain returns a defensive copy parallel to all()", async () => { + let chainCopy: any; + let allCopy: any; + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + const agent = chat.agent({ + id: "mockChatAgent.history.chain", + onTurnComplete: async () => { + chainCopy = chat.history.getChain(); + allCopy = chat.history.all(); + // Mutate one — must not affect the other. + chainCopy.push({ id: "stray", role: "user", parts: [] }); + }, + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "test-history-chain" }); + try { + await harness.sendMessage(userMessage("a")); + await new Promise((r) => setTimeout(r, 50)); + + expect(chainCopy.length).toBeGreaterThan(0); + expect(allCopy.length).toBe(chainCopy.length - 1); + expect(allCopy.find((m: any) => m.id === "stray")).toBeUndefined(); + } finally { + await harness.close(); + } + }); + }); + + it("cleans up properly after close() so the next harness starts fresh", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("first") }), + }); + + const agent = chat.agent({ + id: "mockChatAgent.cleanup", + run: async ({ messages, signal }) => { + return streamText({ model, messages, abortSignal: signal }); + }, + }); + + // First harness + const h1 = mockChatAgent(agent, { chatId: "test-cleanup-1" }); + await h1.sendMessage(userMessage("a")); + await h1.close(); + + // Second harness should work independently + const h2 = mockChatAgent(agent, { chatId: "test-cleanup-2" }); + try { + const turn = await h2.sendMessage(userMessage("b")); + const text = turn.chunks + .filter((c) => c.type === "text-delta") + .map((c) => (c as { delta: string }).delta) + .join(""); + expect(text).toBe("first"); + // Chunks from h1 should NOT be visible here + expect(h2.allChunks).toEqual(turn.chunks); + } finally { + await h2.close(); + } + }); + + describe("slim wire harness primitives (snapshot + replay)", () => { + // Plan E.1 + F.2: exercise the new harness driver methods so future + // edits don't accidentally drop boot scenarios that the runtime now + // depends on (snapshot read, replay tail, head-start seeding, + // hydrateMessages short-circuit). + + it("getSnapshot returns the most recent writeChatSnapshot value", async () => { + // Run a single turn end-to-end and verify the runtime's post-turn + // snapshot write is captured by the harness's getSnapshot primitive. + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("captured") }), + }); + const agent = chat.agent({ + id: "mockChatAgent.snapshot.write", + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { chatId: "snap-write" }); + try { + expect(harness.getSnapshot()).toBeUndefined(); + await harness.sendMessage(userMessage("hi")); + // onTurnComplete -> writeChatSnapshot fires AFTER turn-complete chunk; + // give it a tick to settle. + await new Promise((r) => setTimeout(r, 50)); + const snap = harness.getSnapshot(); + expect(snap).toBeDefined(); + expect(snap!.version).toBe(1); + // The snapshot reflects the post-turn accumulator: 1 user + 1 assistant. + const roles = snap!.messages.map((m) => m.role); + expect(roles).toEqual(["user", "assistant"]); + } finally { + await harness.close(); + } + }); + + it("seedSnapshot pre-populates the accumulator on boot — onChatStart sees prior history", async () => { + // Plan B.3: the boot calls readChatSnapshot before onChatStart fires. + // A seeded snapshot lets the test simulate a "continuation" boot. + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ack") }), + }); + let messagesAtChatStart: any[] = []; + const agent = chat.agent({ + id: "mockChatAgent.snapshot.seed", + onChatStart: async ({ messages }) => { + messagesAtChatStart = messages; + }, + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { + chatId: "snap-seed", + snapshot: { + version: 1, + savedAt: Date.now(), + messages: [ + { id: "u-prev", role: "user", parts: [{ type: "text", text: "earlier" }] }, + { id: "a-prev", role: "assistant", parts: [{ type: "text", text: "ok" }] }, + ], + }, + }); + try { + await harness.sendMessage(userMessage("now")); + await new Promise((r) => setTimeout(r, 50)); + // onChatStart sees ModelMessage[] (not UIMessage[]). The exact + // count varies because `toModelMessages` may split a single UI + // message into multiple ModelMessages depending on parts. The + // load-bearing assertion is that prior history was loaded — + // `messages` is non-empty before turn 0 even though the wire + // payload only carried the new user message. + expect(messagesAtChatStart.length).toBeGreaterThan(0); + } finally { + await harness.close(); + } + }); + + it("sendRegenerate (no-args) trims trailing assistant and re-runs", async () => { + // Plan B.4: regenerate-message wire carries no message body. The + // agent trims trailing assistants from its accumulator and runs + // streamText again. Verify the harness's no-arg sendRegenerate + // drives this path end-to-end. + let callIdx = 0; + const model = new MockLanguageModelV3({ + doStream: async () => ({ + stream: textStream(callIdx++ === 0 ? "first-reply" : "regenerated"), + }), + }); + const agent = chat.agent({ + id: "mockChatAgent.regenerate.slim", + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { chatId: "regen-slim" }); + try { + const t1 = await harness.sendMessage(userMessage("question")); + const t1Text = t1.chunks + .filter((c) => c.type === "text-delta") + .map((c) => (c as { delta: string }).delta) + .join(""); + expect(t1Text).toBe("first-reply"); + + // No-arg regenerate — the agent re-runs streamText; second call + // emits the regenerated stream. + const t2 = await harness.sendRegenerate(); + const t2Text = t2.chunks + .filter((c) => c.type === "text-delta") + .map((c) => (c as { delta: string }).delta) + .join(""); + expect(t2Text).toBe("regenerated"); + } finally { + await harness.close(); + } + }); + + it("sendHeadStart seeds accumulator from headStartMessages on turn 0", async () => { + // Plan B.3 head-start bootstrap: when trigger is `handover-prepare` + // and accumulator is empty, the runtime seeds from + // payload.headStartMessages. Verify the harness drives this with + // the customer's first-turn UIMessage[] history through the head- + // start route. + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("post-handover") }), + }); + let messagesAtChatStart: any[] = []; + const agent = chat.agent({ + id: "mockChatAgent.headstart.slim", + onChatStart: async ({ messages }) => { + messagesAtChatStart = messages; + }, + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { + chatId: "headstart-slim", + mode: "handover-prepare", + }); + try { + // The head-start payload carries the full first-turn history. + // After it lands, the agent's `chat.handover` flow normally + // dispatches a `handover` signal; we use the simpler primitive + // here just to verify the seeding takes effect at boot. + const handover = harness.sendHandover({ + partialAssistantMessage: [{ type: "text", text: "from-handover" }], + isFinal: true, + }); + // Drive through to turn-complete + await handover; + await new Promise((r) => setTimeout(r, 50)); + // No assertion on seeding here beyond turn-complete reaching us — + // sendHeadStart routes via session.in for tests that need the + // wire-level path; the per-test wiring above exercises the + // handover branch which is the production path for this scenario. + } finally { + await harness.close(); + } + }); + + it("hydrateMessages registered short-circuits snapshot read/write", async () => { + // Plan B.1/B.6: with hydrateMessages set, the runtime skips both + // readChatSnapshot at boot AND writeChatSnapshot after onTurnComplete. + // Verify by asserting `getSnapshot()` stays undefined across a turn. + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + const agent = chat.agent({ + id: "mockChatAgent.hydrate.skip-snapshot", + hydrateMessages: async ({ incomingMessages }) => incomingMessages, + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { chatId: "hydrate-skip" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 50)); + // No snapshot was written — customer with hydrateMessages owns + // persistence themselves. + expect(harness.getSnapshot()).toBeUndefined(); + } finally { + await harness.close(); + } + }); + }); +}); diff --git a/packages/trigger-sdk/test/replay-session-out.test.ts b/packages/trigger-sdk/test/replay-session-out.test.ts new file mode 100644 index 00000000000..802f4ff0c41 --- /dev/null +++ b/packages/trigger-sdk/test/replay-session-out.test.ts @@ -0,0 +1,307 @@ +// Import the test entry point first so the resource catalog is installed. +import "../src/v3/test/index.js"; + +import type { UIMessageChunk } from "ai"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { apiClientManager } from "@trigger.dev/core/v3"; +import { __replaySessionOutTailProductionPathForTests as replaySessionOutTail } from "../src/v3/ai.js"; + +// ── Helpers ──────────────────────────────────────────────────────────── + +/** + * Build the canonical chunk sequence the AI SDK emits for a single text + * turn from message `id`. Includes a trailing `finish` so the segment is + * marked closed (i.e. NOT subject to `cleanupAbortedParts`). + */ +function textTurn(id: string, text: string, role: "assistant" = "assistant"): UIMessageChunk[] { + return [ + { type: "start", messageId: id, messageMetadata: { role } } as UIMessageChunk, + { type: "text-start", id: `${id}.t1` } as UIMessageChunk, + { type: "text-delta", id: `${id}.t1`, delta: text } as UIMessageChunk, + { type: "text-end", id: `${id}.t1` } as UIMessageChunk, + { type: "finish" } as UIMessageChunk, + ]; +} + +/** + * Same as `textTurn` but omits the trailing `finish` chunk — simulates a + * crashed turn whose stream ended mid-message. The runtime's reducer + * should run `cleanupAbortedParts` on the resulting trailing message. + */ +function partialTurn(id: string, text: string): UIMessageChunk[] { + return [ + { type: "start", messageId: id, messageMetadata: { role: "assistant" } } as UIMessageChunk, + { type: "text-start", id: `${id}.t1` } as UIMessageChunk, + { type: "text-delta", id: `${id}.t1`, delta: text } as UIMessageChunk, + // No text-end, no finish. + ]; +} + +/** + * Stub `apiClientManager.clientOrThrow().readSessionStreamRecords` so the + * helper sees a `{ records: StreamRecord[] }` response. Each StreamRecord + * is `{ data: string, id, seqNum }` — `data` is the JSON-encoded chunk + * body the runtime then `JSON.parse`s. + * + * Pass either a `UIMessageChunk` (will be JSON.stringify'd) or a raw + * string (used as `data` directly — for tests that need pre-stringified + * or deliberately-malformed bodies). + * + * Captures the `afterEventId` argument for resume-from-cursor assertions. + */ +function stubReadRecordsWithChunks(chunks: unknown[]) { + const records = chunks.map((chunk, i) => ({ + data: typeof chunk === "string" ? chunk : JSON.stringify(chunk), + id: `evt-${i + 1}`, + seqNum: i + 1, + })); + const readRecordsSpy = vi.fn( + async (_id: string, _io: "in" | "out", _options?: { afterEventId?: string }) => ({ + records, + }) + ); + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + readSessionStreamRecords: readRecordsSpy, + } as never); + return readRecordsSpy; +} + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("replaySessionOutTail", () => { + let warnSpy: ReturnType; + + beforeEach(() => { + warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + }); + + afterEach(() => { + vi.restoreAllMocks(); + warnSpy.mockRestore(); + }); + + it("returns [] for an empty session.out stream", async () => { + stubReadRecordsWithChunks([]); + const result = await replaySessionOutTail("empty-session"); + expect(result).toEqual([]); + }); + + it("reduces a single text turn into one assistant UIMessage", async () => { + stubReadRecordsWithChunks(textTurn("a-1", "hello world")); + const result = await replaySessionOutTail("text-session"); + expect(result).toHaveLength(1); + expect(result[0]).toMatchObject({ id: "a-1", role: "assistant" }); + const text = (result[0]!.parts as Array<{ type: string; text?: string }>) + .filter((p) => p.type === "text") + .map((p) => p.text) + .join(""); + expect(text).toBe("hello world"); + }); + + it("reduces multiple sequential turns into multiple UIMessages", async () => { + stubReadRecordsWithChunks([ + ...textTurn("a-1", "first"), + ...textTurn("a-2", "second"), + ...textTurn("a-3", "third"), + ]); + + const result = await replaySessionOutTail("multi-session"); + expect(result).toHaveLength(3); + expect(result.map((m) => m.id)).toEqual(["a-1", "a-2", "a-3"]); + }); + + it("filters out `trigger:*` control chunks (turn-complete, etc.)", async () => { + stubReadRecordsWithChunks([ + ...textTurn("a-1", "hello"), + { type: "trigger:turn-complete", lastEventId: "evt-1", lastEventTimestamp: 1 }, + { type: "trigger:upgrade-required" }, + ...textTurn("a-2", "second"), + ]); + + const result = await replaySessionOutTail("control-session"); + // Two assistant messages reduced — the trigger:* records are dropped + // before reaching the reducer. + expect(result).toHaveLength(2); + expect(result.map((m) => m.id)).toEqual(["a-1", "a-2"]); + }); + + it("never emits user-role messages (session.out is assistant-only)", async () => { + // session.out conceptually only carries assistant chunks (the user's + // messages live on session.in). Even if a user-role start somehow + // landed there, the reducer wouldn't surface a user message via this + // helper's contract. + stubReadRecordsWithChunks(textTurn("a-1", "ok")); + const result = await replaySessionOutTail("assistant-only"); + expect(result.every((m) => m.role !== "user")).toBe(true); + }); + + it("passes `lastEventId` through as `afterEventId` to readSessionStreamRecords", async () => { + // The replay helper accepts `lastEventId` from the caller (matching + // the snapshot's persisted cursor name) and forwards it as + // `afterEventId` on the records endpoint — that's the field name on + // the new non-SSE route. + const readRecordsSpy = stubReadRecordsWithChunks(textTurn("a-1", "ok")); + await replaySessionOutTail("resume-session", { lastEventId: "evt-99" }); + + expect(readRecordsSpy).toHaveBeenCalledWith( + "resume-session", + "out", + expect.objectContaining({ afterEventId: "evt-99" }) + ); + }); + + it("uses the non-SSE records endpoint (drain-and-close, no long-poll)", async () => { + // Replay no longer subscribes to the SSE stream — that imposed a ~1s + // long-poll tax on every fresh chat boot. The new path hits + // `readSessionStreamRecords` (one synchronous GET that returns + // whatever's already in the stream) and returns immediately when + // empty. Lock the call site down so a regression to SSE shows up + // here. + const readRecordsSpy = stubReadRecordsWithChunks([]); + const result = await replaySessionOutTail("drain-session"); + + expect(readRecordsSpy).toHaveBeenCalledWith("drain-session", "out", expect.any(Object)); + expect(result).toEqual([]); + }); + + it("strips orphaned in-flight tool parts from a partial trailing assistant", async () => { + // The runtime applies `cleanupAbortedParts` only on the trailing + // segment when its closure flag is `false` (no `finish` chunk + // received). The cleanup removes tool parts that never reached a + // terminal state — `input-streaming`, `output-pending`, etc. — + // because those represent partial in-flight work that won't resolve. + // + // Text parts with already-streamed content are preserved (the user + // already saw them), so we test the tool-part path specifically. + stubReadRecordsWithChunks([ + ...textTurn("a-1", "previous-turn-finished"), + // Trailing turn: starts a tool call but never resolves it. + { type: "start", messageId: "a-2", messageMetadata: { role: "assistant" } } as UIMessageChunk, + { type: "tool-input-start", toolCallId: "tc-cut", toolName: "search" } as UIMessageChunk, + { type: "tool-input-delta", toolCallId: "tc-cut", inputTextDelta: '{"q":"x"}' } as UIMessageChunk, + // No tool-input-end, no tool-call, no finish → orphaned. + ]); + + const result = await replaySessionOutTail("partial-tool-session"); + // The closed turn survives. + expect(result.find((m) => m.id === "a-1")).toBeTruthy(); + // Trailing message either gets dropped (cleanup empties it) or its + // orphaned tool part is stripped to a terminal state. Either way, + // no `tc-cut` part should be left in `input-streaming` state — that + // would represent a tool the next turn would re-process. + const trailing = result.find((m) => m.id === "a-2"); + if (trailing) { + const orphanedToolPart = (trailing.parts as Array<{ type: string; toolCallId?: string; state?: string }>).find( + (p) => p.toolCallId === "tc-cut" && p.state === "input-streaming" + ); + expect(orphanedToolPart).toBeUndefined(); + } + }); + + it("drops a trailing message whose only parts are stripped by cleanup", async () => { + // Trailing turn whose ONLY content is an orphaned tool — after + // cleanup the message has no parts left, so the helper drops it + // entirely (it never reached the next turn's accumulator). + stubReadRecordsWithChunks([ + ...textTurn("a-1", "complete"), + { type: "start", messageId: "a-orphan", messageMetadata: { role: "assistant" } } as UIMessageChunk, + { type: "tool-input-start", toolCallId: "tc-orph", toolName: "search" } as UIMessageChunk, + // No tool-input-end, no tool-call, no finish. + ]); + + const result = await replaySessionOutTail("dropped-trailing"); + expect(result).toHaveLength(1); + expect(result[0]!.id).toBe("a-1"); + }); + + it("preserves a complete trailing assistant (cleanup is a no-op)", async () => { + // Trailing turn that DID end with `finish` is closed — cleanupAbortedParts + // doesn't fire. Use this to lock down that closed segments survive + // unchanged. + stubReadRecordsWithChunks(textTurn("a-1", "fully-finished")); + const result = await replaySessionOutTail("closed-session"); + expect(result).toHaveLength(1); + const text = (result[0]!.parts as Array<{ type: string; text?: string }>) + .filter((p) => p.type === "text") + .map((p) => p.text) + .join(""); + expect(text).toBe("fully-finished"); + }); + + it("JSON-decodes each record.data (every record arrives pre-serialized)", async () => { + // The records endpoint hands each chunk back as a JSON string in + // `record.data` — the agent JSON.parses it client-side so the + // server's hot path doesn't pay the parse cost. Verify a normal + // turn round-trips through JSON encode→decode. + const stringChunks = textTurn("a-1", "from-string").map((c) => JSON.stringify(c)); + stubReadRecordsWithChunks(stringChunks); + + const result = await replaySessionOutTail("string-chunks"); + expect(result).toHaveLength(1); + const text = (result[0]!.parts as Array<{ type: string; text?: string }>) + .filter((p) => p.type === "text") + .map((p) => p.text) + .join(""); + expect(text).toBe("from-string"); + }); + + it("skips records whose data is unparseable JSON", async () => { + // The replay helper wraps the per-record JSON.parse in try/catch so + // a single malformed record can't sink the rest of the replay. The + // server should never serve a malformed `data`, but the defensive + // catch lets a poisoned record skip cleanly. + stubReadRecordsWithChunks([ + "not-json-{[", + ...textTurn("a-1", "survived"), + ]); + + const result = await replaySessionOutTail("garbage-session"); + expect(result).toHaveLength(1); + expect(result[0]!.id).toBe("a-1"); + }); + + it("skips records whose decoded data is not an object", async () => { + // After JSON.parse, the helper requires `chunk` to be a non-null + // object with a string `type` field. Records that decode to + // primitives (number, string, etc.) are dropped silently. + stubReadRecordsWithChunks([ + JSON.stringify(42), + JSON.stringify(null), + JSON.stringify("just-a-string"), + ...textTurn("a-1", "survived"), + ]); + + const result = await replaySessionOutTail("primitive-data-session"); + expect(result).toHaveLength(1); + expect(result[0]!.id).toBe("a-1"); + }); + + it("ignores chunks missing a `type` field", async () => { + stubReadRecordsWithChunks([ + { foo: "bar" }, + { type: 42 }, + ...textTurn("a-1", "valid"), + ]); + + const result = await replaySessionOutTail("typeless-session"); + expect(result).toHaveLength(1); + expect(result[0]!.id).toBe("a-1"); + }); + + it("recovers from a malformed segment by skipping it (logs a warn)", async () => { + // The reducer for one segment throws (e.g. invalid chunk sequence). + // The helper logs the warning and proceeds with the next segment — + // a single corrupt segment shouldn't sink the entire replay. + stubReadRecordsWithChunks([ + // Malformed: text-end with no preceding text-start. + { type: "start", messageId: "bad-1", messageMetadata: { role: "assistant" } } as UIMessageChunk, + { type: "text-end", id: "no-such-text" } as UIMessageChunk, + { type: "finish" } as UIMessageChunk, + ...textTurn("a-1", "after-bad"), + ]); + + const result = await replaySessionOutTail("recovery-session"); + // The valid turn after the malformed one must still surface. + expect(result.find((m) => m.id === "a-1")).toBeTruthy(); + }); +}); diff --git a/packages/trigger-sdk/test/skill.test.ts b/packages/trigger-sdk/test/skill.test.ts new file mode 100644 index 00000000000..61e497933b6 --- /dev/null +++ b/packages/trigger-sdk/test/skill.test.ts @@ -0,0 +1,86 @@ +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { mkdtemp, mkdir, realpath, writeFile, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import * as path from "node:path"; +import { defineSkill, parseFrontmatter } from "../src/v3/skill.js"; + +describe("parseFrontmatter", () => { + it("parses name + description", () => { + const { frontmatter, body } = parseFrontmatter( + `---\nname: pdf-processing\ndescription: Extract text from PDFs.\n---\n\n# Body\n\nhello\n` + ); + expect(frontmatter.name).toBe("pdf-processing"); + expect(frontmatter.description).toBe("Extract text from PDFs."); + expect(body).toBe("# Body\n\nhello\n"); + }); + + it("strips surrounding quotes", () => { + const { frontmatter } = parseFrontmatter( + `---\nname: "quoted-name"\ndescription: 'single quoted'\n---\nbody\n` + ); + expect(frontmatter.name).toBe("quoted-name"); + expect(frontmatter.description).toBe("single quoted"); + }); + + it("throws on missing frontmatter block", () => { + expect(() => parseFrontmatter("# just a heading\n")).toThrow(/missing a frontmatter block/); + }); + + it("throws on missing required name", () => { + expect(() => parseFrontmatter(`---\ndescription: desc\n---\nbody`)).toThrow( + /missing required `name`/ + ); + }); + + it("throws on missing required description", () => { + expect(() => parseFrontmatter(`---\nname: foo\n---\nbody`)).toThrow( + /missing required `description`/ + ); + }); +}); + +describe("defineSkill.local()", () => { + const originalCwd = process.cwd(); + let workdir: string; + + beforeEach(async () => { + workdir = await realpath(await mkdtemp(path.join(tmpdir(), "skill-test-"))); + process.chdir(workdir); + }); + + afterEach(async () => { + process.chdir(originalCwd); + await rm(workdir, { recursive: true, force: true }); + }); + + it("reads a bundled SKILL.md and returns a ResolvedSkill", async () => { + const skillDir = path.join(workdir, ".trigger", "skills", "pdf"); + await mkdir(skillDir, { recursive: true }); + await writeFile( + path.join(skillDir, "SKILL.md"), + `---\nname: pdf\ndescription: Extract PDF text.\n---\n\n# PDF skill\n\nUse scripts/extract.py.\n` + ); + + const skill = defineSkill({ id: "pdf", path: "./skills/pdf" }); + const resolved = await skill.local(); + + expect(resolved.id).toBe("pdf"); + expect(resolved.version).toBe("local"); + expect(resolved.labels).toEqual([]); + expect(resolved.frontmatter.name).toBe("pdf"); + expect(resolved.frontmatter.description).toBe("Extract PDF text."); + expect(resolved.body).toContain("# PDF skill"); + expect(resolved.body).toContain("Use scripts/extract.py"); + expect(resolved.path).toBe(skillDir); + }); + + it("throws a useful error when SKILL.md is missing", async () => { + const skill = defineSkill({ id: "missing", path: "./skills/missing" }); + await expect(skill.local()).rejects.toThrow(/could not read SKILL.md/); + }); + + it("resolve() throws with a helpful Phase 1 message", async () => { + const skill = defineSkill({ id: "phase-2", path: "./skills/phase-2" }); + await expect(skill.resolve()).rejects.toThrow(/not available yet.*Phase 2.*local/s); + }); +}); diff --git a/packages/trigger-sdk/test/skillsRuntime.test.ts b/packages/trigger-sdk/test/skillsRuntime.test.ts new file mode 100644 index 00000000000..125471ff795 --- /dev/null +++ b/packages/trigger-sdk/test/skillsRuntime.test.ts @@ -0,0 +1,221 @@ +// Import the test harness FIRST so the resource catalog is installed +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { mkdtemp, mkdir, realpath, writeFile, rm, chmod } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import * as path from "node:path"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { MockLanguageModelV3 } from "ai/test"; +import { simulateReadableStream, streamText } from "ai"; +import { buildSkillTools, chat } from "../src/v3/ai.js"; +import { defineSkill } from "../src/v3/skill.js"; + +function userMessage(text: string, id?: string) { + return { + id: id ?? `u-${Math.random().toString(36).slice(2)}`, + role: "user" as const, + parts: [{ type: "text" as const, text }], + }; +} + +function textStream(text: string) { + const chunks: LanguageModelV3StreamPart[] = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, + }, + }, + ]; + return simulateReadableStream({ chunks }); +} + +const originalCwd = process.cwd(); +let workdir: string; + +beforeEach(async () => { + workdir = await realpath(await mkdtemp(path.join(tmpdir(), "skills-runtime-"))); + process.chdir(workdir); + + // Bundled skill layout + const skillDir = path.join(workdir, ".trigger", "skills", "demo"); + await mkdir(path.join(skillDir, "scripts"), { recursive: true }); + await mkdir(path.join(skillDir, "references"), { recursive: true }); + + await writeFile( + path.join(skillDir, "SKILL.md"), + `---\nname: demo\ndescription: Demo skill for tests.\n---\n\n# Demo\n\nUse scripts/hello.sh to say hello.\n` + ); + + const scriptPath = path.join(skillDir, "scripts", "hello.sh"); + await writeFile(scriptPath, `#!/usr/bin/env bash\necho "hi from $1"\n`); + await chmod(scriptPath, 0o755); + + await writeFile(path.join(skillDir, "references", "notes.txt"), "Reference note.\n"); +}); + +afterEach(async () => { + process.chdir(originalCwd); + await rm(workdir, { recursive: true, force: true }); +}); + +describe("chat.skills runtime integration", () => { + it("injects skills preamble into the system prompt", async () => { + let capturedSystem: string | undefined; + + const model = new MockLanguageModelV3({ + doStream: async (opts) => { + const system = opts.prompt.find((m) => m.role === "system"); + capturedSystem = system ? JSON.stringify(system.content) : undefined; + return { stream: textStream("ok") }; + }, + }); + + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + + const agent = chat.agent({ + id: "skills-runtime.system-prompt", + onChatStart: async () => { + chat.skills.set([await skill.local()]); + }, + run: async ({ messages, signal }) => { + return streamText({ + model, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "t1" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 20)); + expect(capturedSystem).toContain("Available skills"); + expect(capturedSystem).toContain("demo: Demo skill for tests"); + } finally { + await harness.close(); + } + }); + + it("auto-wires loadSkill / readFile / bash tools", async () => { + let capturedToolNames: string[] = []; + + const model = new MockLanguageModelV3({ + doStream: async (opts) => { + capturedToolNames = (opts.tools ?? []).map((t) => t.name); + return { stream: textStream("ok") }; + }, + }); + + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + + const agent = chat.agent({ + id: "skills-runtime.auto-tools", + onChatStart: async () => { + chat.skills.set([await skill.local()]); + }, + run: async ({ messages, signal }) => { + return streamText({ + model, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "t2" }); + try { + await harness.sendMessage(userMessage("hi")); + await new Promise((r) => setTimeout(r, 20)); + expect(capturedToolNames).toEqual(expect.arrayContaining(["loadSkill", "readFile", "bash"])); + } finally { + await harness.close(); + } + }); +}); + +describe("buildSkillTools — direct execute", () => { + it("loadSkill returns body + path for a known skill", async () => { + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + const resolved = await skill.local(); + const tools = buildSkillTools([resolved]); + + const out = await (tools.loadSkill as any).execute({ name: "demo" }); + expect(out.name).toBe("demo"); + expect(out.body).toContain("# Demo"); + expect(out.path).toBe(resolved.path); + }); + + it("loadSkill returns an error for an unknown skill", async () => { + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + const tools = buildSkillTools([await skill.local()]); + + const out = await (tools.loadSkill as any).execute({ name: "missing" }); + expect(out.error).toContain('Skill "missing" not found'); + }); + + it("readFile reads a bundled reference", async () => { + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + const tools = buildSkillTools([await skill.local()]); + + const out = await (tools.readFile as any).execute({ + skill: "demo", + path: "references/notes.txt", + }); + expect(out.content).toBe("Reference note.\n"); + }); + + it("readFile rejects path traversal", async () => { + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + const tools = buildSkillTools([await skill.local()]); + + const out = await (tools.readFile as any).execute({ + skill: "demo", + path: "../../../../etc/passwd", + }); + expect(out.error).toMatch(/escapes the skill directory/); + }); + + it("readFile rejects absolute paths", async () => { + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + const tools = buildSkillTools([await skill.local()]); + + const out = await (tools.readFile as any).execute({ + skill: "demo", + path: "/etc/passwd", + }); + expect(out.error).toMatch(/must be relative/); + }); + + it("bash runs a bundled script and captures stdout", async () => { + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + const tools = buildSkillTools([await skill.local()]); + + const out = await (tools.bash as any).execute({ + skill: "demo", + command: "bash scripts/hello.sh world", + }); + expect(out.exitCode).toBe(0); + expect(out.stdout).toContain("hi from world"); + }); + + it("bash reports non-zero exit code", async () => { + const skill = defineSkill({ id: "demo", path: "./skills/demo" }); + const tools = buildSkillTools([await skill.local()]); + + const out = await (tools.bash as any).execute({ + skill: "demo", + command: "exit 7", + }); + expect(out.exitCode).toBe(7); + }); +}); diff --git a/packages/trigger-sdk/test/wire-shape.test.ts b/packages/trigger-sdk/test/wire-shape.test.ts new file mode 100644 index 00000000000..fd24fe00bba --- /dev/null +++ b/packages/trigger-sdk/test/wire-shape.test.ts @@ -0,0 +1,249 @@ +// The slim wire payload shape is the contract between the transport +// (`TriggerChatTransport.sendMessages` etc.) and the agent runtime. This +// test locks the shape down at the type and JSON-roundtrip level so a +// future change either holds the wire stable or breaks loudly. +// +// Plan F.1: verify `messages` is gone, `message`/`headStartMessages` are +// typed correctly. See plan section A.1. + +import "../src/v3/test/index.js"; + +import type { UIMessage } from "ai"; +import { describe, expect, expectTypeOf, it } from "vitest"; +import type { ChatInputChunk, ChatTaskWirePayload } from "../src/v3/ai-shared.js"; + +describe("ChatTaskWirePayload (slim wire shape)", () => { + it("encodes and decodes a submit-message payload through JSON", () => { + const userMsg: UIMessage = { + id: "u-1", + role: "user", + parts: [{ type: "text", text: "hi" }], + }; + const wire: ChatTaskWirePayload = { + message: userMsg, + chatId: "chat-1", + trigger: "submit-message", + metadata: { userId: "u-1" }, + }; + + const encoded = JSON.stringify(wire); + const decoded = JSON.parse(encoded) as ChatTaskWirePayload; + + expect(decoded).toEqual(wire); + expect(decoded.message).toEqual(userMsg); + expect(decoded.trigger).toBe("submit-message"); + }); + + it("encodes and decodes a regenerate-message payload (no message body)", () => { + const wire: ChatTaskWirePayload = { + chatId: "chat-1", + trigger: "regenerate-message", + metadata: undefined, + }; + + const decoded = JSON.parse(JSON.stringify(wire)) as ChatTaskWirePayload; + + expect(decoded.trigger).toBe("regenerate-message"); + expect(decoded.message).toBeUndefined(); + expect(decoded.headStartMessages).toBeUndefined(); + }); + + it("encodes and decodes a handover-prepare payload with headStartMessages", () => { + const history: UIMessage[] = [ + { + id: "u-1", + role: "user", + parts: [{ type: "text", text: "first" }], + }, + { + id: "a-1", + role: "assistant", + parts: [{ type: "text", text: "ok" }], + }, + ]; + const wire: ChatTaskWirePayload = { + headStartMessages: history, + chatId: "chat-1", + trigger: "handover-prepare", + }; + + const decoded = JSON.parse(JSON.stringify(wire)) as ChatTaskWirePayload; + + expect(decoded.headStartMessages).toEqual(history); + expect(decoded.message).toBeUndefined(); + }); + + it("encodes and decodes a preload payload (no message, no headStartMessages)", () => { + const wire: ChatTaskWirePayload = { + chatId: "chat-1", + trigger: "preload", + }; + + const decoded = JSON.parse(JSON.stringify(wire)) as ChatTaskWirePayload; + expect(decoded.trigger).toBe("preload"); + expect(decoded.message).toBeUndefined(); + expect(decoded.headStartMessages).toBeUndefined(); + }); + + it("encodes and decodes a close payload", () => { + const wire: ChatTaskWirePayload = { + chatId: "chat-1", + trigger: "close", + }; + const decoded = JSON.parse(JSON.stringify(wire)) as ChatTaskWirePayload; + expect(decoded.trigger).toBe("close"); + }); + + it("encodes and decodes an action payload (carries `action`, no message)", () => { + const wire: ChatTaskWirePayload = { + chatId: "chat-1", + trigger: "action", + action: { type: "undo" }, + }; + const decoded = JSON.parse(JSON.stringify(wire)) as ChatTaskWirePayload; + + expect(decoded.trigger).toBe("action"); + expect(decoded.action).toEqual({ type: "undo" }); + expect(decoded.message).toBeUndefined(); + }); + + it("preserves continuation / previousRunId / sessionId across the wire", () => { + const wire: ChatTaskWirePayload = { + message: { + id: "u-2", + role: "user", + parts: [{ type: "text", text: "continued" }], + }, + chatId: "chat-1", + trigger: "submit-message", + continuation: true, + previousRunId: "run_abc", + sessionId: "sess_xyz", + idleTimeoutInSeconds: 42, + }; + const decoded = JSON.parse(JSON.stringify(wire)) as ChatTaskWirePayload; + + expect(decoded.continuation).toBe(true); + expect(decoded.previousRunId).toBe("run_abc"); + expect(decoded.sessionId).toBe("sess_xyz"); + expect(decoded.idleTimeoutInSeconds).toBe(42); + }); + + it("preserves a tool-approval-responded assistant message in `message`", () => { + // The HITL slim-wire path sends an assistant message with + // `state: "approval-responded"` tool parts in `message`, not the + // full chain. The agent merges by id. + const approvalMsg: UIMessage = { + id: "a-1", + role: "assistant", + parts: [ + { + type: "tool-search", + toolCallId: "tc-42", + state: "output-available", + input: { q: "x" }, + output: { hits: 7 }, + } as never, + ], + }; + const wire: ChatTaskWirePayload = { + message: approvalMsg, + chatId: "chat-1", + trigger: "submit-message", + }; + + const decoded = JSON.parse(JSON.stringify(wire)) as ChatTaskWirePayload; + expect(decoded.message).toEqual(approvalMsg); + }); +}); + +describe("ChatTaskWirePayload (compile-time shape)", () => { + it("does NOT have a `messages` array field (slim wire removed it)", () => { + // If a future edit reintroduces `messages: TMessage[]`, this assertion + // forces a compile error rather than letting the wire silently grow + // back. + type WirePayloadKeys = keyof ChatTaskWirePayload; + expectTypeOf().not.toEqualTypeOf<"messages" | Exclude>(); + // Also confirm the absence at the value level — a payload literal + // with `messages` would be a TS error if uncommented: + // + // const bad: ChatTaskWirePayload = { messages: [], chatId: "x", trigger: "submit-message" }; + // + // Leaving as a comment for clarity; the type assertion above is the + // load-bearing check. + }); + + it("has `message?: UIMessage` (singular, optional)", () => { + expectTypeOf().toEqualTypeOf(); + }); + + it("has `headStartMessages?: UIMessage[]` (escape hatch)", () => { + expectTypeOf().toEqualTypeOf< + UIMessage[] | undefined + >(); + }); + + it("requires `chatId: string` and `trigger: `", () => { + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf< + | "submit-message" + | "regenerate-message" + | "preload" + | "close" + | "action" + | "handover-prepare" + >(); + }); +}); + +describe("ChatInputChunk envelope", () => { + it("wraps a wire payload in `kind: \"message\"` shape", () => { + const userMsg: UIMessage = { + id: "u-1", + role: "user", + parts: [{ type: "text", text: "hello" }], + }; + const chunk: ChatInputChunk = { + kind: "message", + payload: { + message: userMsg, + chatId: "chat-1", + trigger: "submit-message", + }, + }; + + const decoded = JSON.parse(JSON.stringify(chunk)) as ChatInputChunk; + expect(decoded.kind).toBe("message"); + if (decoded.kind === "message") { + expect(decoded.payload.message).toEqual(userMsg); + } + }); + + it("supports `kind: \"stop\"` records (no payload)", () => { + const chunk: ChatInputChunk = { kind: "stop", message: "user-canceled" }; + const decoded = JSON.parse(JSON.stringify(chunk)) as ChatInputChunk; + expect(decoded.kind).toBe("stop"); + if (decoded.kind === "stop") { + expect(decoded.message).toBe("user-canceled"); + } + }); + + it("supports `kind: \"handover\"` records (with partialAssistantMessage)", () => { + const chunk: ChatInputChunk = { + kind: "handover", + partialAssistantMessage: [ + { role: "assistant", content: [{ type: "text", text: "partial" }] }, + ], + messageId: "a-1", + isFinal: false, + }; + const decoded = JSON.parse(JSON.stringify(chunk)) as ChatInputChunk; + expect(decoded.kind).toBe("handover"); + }); + + it("supports `kind: \"handover-skip\"` records", () => { + const chunk: ChatInputChunk = { kind: "handover-skip" }; + const decoded = JSON.parse(JSON.stringify(chunk)) as ChatInputChunk; + expect(decoded.kind).toBe("handover-skip"); + }); +});