diff --git a/.changeset/action-stream-into-conversation.md b/.changeset/action-stream-into-conversation.md new file mode 100644 index 00000000000..be3723ef784 --- /dev/null +++ b/.changeset/action-stream-into-conversation.md @@ -0,0 +1,7 @@ +--- +"@trigger.dev/sdk": patch +--- + +A response returned from `onAction` is now part of the conversation, whether it is a `StreamTextResult`, a `string`, or an assistant `UIMessage`. A `string` or `UIMessage` return was documented but did nothing: it reached neither the browser nor the conversation. Returning a `StreamTextResult` from an action sent it to the browser and nowhere else, so a regenerate showed the user a new answer that the model had no memory of, and the next turn carried on from the answer it had replaced. + +A stream that fails part-way through is also no longer committed as though it finished. Whatever streamed is still kept, but the failure is reported instead of the truncated text being stored, and built on, as a complete answer. An action that used to end quietly on a mid-stream failure now surfaces an error to the frontend. It is still an action, not a turn: `onTurnComplete` does not fire for it, the turn count is unchanged, and an instruction injected for the next turn still reaches that turn. diff --git a/.changeset/createsession-steering-lanes.md b/.changeset/createsession-steering-lanes.md new file mode 100644 index 00000000000..bf39c2a4685 --- /dev/null +++ b/.changeset/createsession-steering-lanes.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Steering messages are now kept in the conversation when you drive turns yourself with `chat.createSession()` or `chat.MessageAccumulator`. Previously a message that arrived mid-answer shaped that answer and then existed nowhere: it was missing from `turn.uiMessages`, so an app persisting from there never stored it, missing from `turn.messages`, so every later turn answered as though it had never been sent, and it was not queued as its own turn either. It now lands in both, the same way it does on `chat.agent`. diff --git a/.changeset/inject-instructions-shape.md b/.changeset/inject-instructions-shape.md new file mode 100644 index 00000000000..85129f9cb57 --- /dev/null +++ b/.changeset/inject-instructions-shape.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Injected system context is merged into a single instruction block, so it works on every supported AI SDK version. Note that a cached system prompt gives up its cache entry for as long as an injection is live, since the cached prefix has changed. diff --git a/.changeset/inject-system-to-instructions.md b/.changeset/inject-system-to-instructions.md new file mode 100644 index 00000000000..f5c2e5993b4 --- /dev/null +++ b/.changeset/inject-system-to-instructions.md @@ -0,0 +1,7 @@ +--- +"@trigger.dev/sdk": patch +--- + +`chat.inject()` with `role: "system"` now works. It previously put the system message into the conversation, which AI SDK 7 rejects for every provider: the next turn died with a generic "An error occurred." and persisted an empty assistant message, so the agent looked like it had stopped answering. System-role context is now appended to the model's instructions, which is also the only way to inject context the agent treats as trusted. + +Two things to know. Instructions are delivered by `chat.toStreamTextOptions()`, so a `run()` that calls `streamText` without spreading it does not receive a system-role injection. The conversational lane has no such requirement. And an injection applies to the next turn only, rather than repeating on every turn that follows it. Every inference call in that turn sees it, so a `run()` that builds options more than once gets the same instructions each time. diff --git a/.changeset/persist-action-history-mutations.md b/.changeset/persist-action-history-mutations.md new file mode 100644 index 00000000000..337494f327d --- /dev/null +++ b/.changeset/persist-action-history-mutations.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Undo, edit and regenerate now survive a run ending. History rolled back from `onAction` was only kept in the running worker's memory, so the rollback held while that worker stayed warm and then reverted on the next continuation. The undone messages came back, minutes later, with no error. This also holds when the turn before the action failed: the rollback used to be written against the cursor from before that turn, so a continuation could replay output the failed turn had already superseded. diff --git a/.changeset/steering-messages-accumulator.md b/.changeset/steering-messages-accumulator.md new file mode 100644 index 00000000000..52bd67bef38 --- /dev/null +++ b/.changeset/steering-messages-accumulator.md @@ -0,0 +1,9 @@ +--- +"@trigger.dev/sdk": patch +--- + +Steering messages injected mid-answer are now part of the conversation, both for your hooks and for the model on later turns. Previously they reached the model for the answer they steered and reached the browser, but nothing else: `onTurnComplete` never saw them, so an app storing its own transcript lost the instruction the answer was shaped by, and it vanished from the conversation on reload. The model also forgot the instruction from the next turn onwards, answering as though the message had never been sent, while the chat UI still showed it. This holds when the steered turn fails part-way, and when `pendingMessages.prepare` reshapes the message: later turns now see the same form the steered turn did, not the original message. + +Approving a tool call no longer undoes compaction. A tool-approval continuation used to rebuild the model's context from the full conversation, so a chat that had been summarised to fit the context window was sent the whole transcript again on the next call, and could go over the limit it had just been compacted to avoid. The same applied to a regenerated answer that replaced an existing one. + +If you worked around this by saving steering messages as they arrive, in `pendingMessages.onReceived` for example, that write now duplicates the one you get from `newUIMessages`. Drop it, or skip messages you have already stored. diff --git a/docs/ai-chat/actions.mdx b/docs/ai-chat/actions.mdx index 956e5090aef..c16b662ec0e 100644 --- a/docs/ai-chat/actions.mdx +++ b/docs/ai-chat/actions.mdx @@ -1,7 +1,7 @@ --- title: "Actions" sidebarTitle: "Actions" -description: "Custom commands sent from the frontend that mutate chat state without consuming a turn — undo, rollback, edit, regenerate." +description: "Custom commands sent from the frontend that mutate chat state without consuming a turn: undo, rollback, edit, regenerate." --- ## Overview @@ -54,7 +54,7 @@ export const myChat = chat.agent({ ## Returning a model response from an action -`onAction` can return a `StreamTextResult`, `string`, or `UIMessage` to produce a response. The returned stream is auto-piped to the frontend just like a normal turn, but the rest of the turn machinery (`onTurnStart`, `onTurnComplete`, etc.) still does not fire. +`onAction` can return a `StreamTextResult`, `string`, or `UIMessage` to produce a response. All three are sent to the frontend and added to the conversation just like a normal turn's answer, but the rest of the turn machinery (`onTurnStart`, `onTurnComplete`, etc.) still does not fire. A returned `UIMessage` must have `role: "assistant"`; its text and `data-*` parts are delivered, and other part types are dropped. ```ts onAction: async ({ action, messages }) => { @@ -70,7 +70,37 @@ onAction: async ({ action, messages }) => { } ``` -This is useful for actions that both mutate state and want a fresh model response (regenerate-from-here, retry-with-different-style). Persistence is your responsibility inside `onAction` itself; you have access to the streamed response object. +This is useful for actions that both mutate state and want a fresh model response (regenerate-from-here, retry-with-different-style). + +### Actions and persistence + +An action is not a turn, so `onTurnComplete` never fires, and that is where an app that owns its own transcript normally writes. What that means depends on which persistence model you use. + +**Platform-managed** (no `hydrateMessages`): nothing to do. After an action that changed the conversation (a `chat.history` mutation, a response returned from `onAction`, or both), the runtime writes the snapshot, so the change survives the run ending. + +**Your own store** (`hydrateMessages` registered): the runtime deliberately does not write, because your store is the source of truth. A history mutation and a returned response both live only in the running worker until you persist them, and a continuation rehydrates from your store, not from what the worker had in memory. `chat.pipeAndCapture` hands you the same assistant message the runtime would have captured: + +```ts +onAction: async ({ action, messages }) => { + if (action.type === "undo") { + chat.history.slice(0, -2); + await db.deleteLastExchange(chatId); // the rollback is yours to persist + } + + if (action.type === "regenerate") { + chat.history.slice(0, -1); + await db.deleteLastAssistant(chatId); // drop the answer being replaced + const { message } = await chat.pipeAndCapture( + streamText({ model: anthropic("claude-sonnet-4-5"), messages }) + ); + if (message) await db.saveMessage(message); // then store the new one + } +}, +``` + +Mirror each mutation in your store, not only the additions. A `chat.history` mutation is invisible to your database, so a regenerate is a delete *and* an insert. Saving the new answer without removing the old one leaves both in the canonical transcript, and the next hydration returns the two of them. (An append-only or branching store is the exception: there you write a new version and resolve the head on read.) + +Returning the stream instead of piping it yourself still works and still reaches the browser, but you have no message to store, so the next run does not know about it. ## Gating actions on HITL state @@ -89,10 +119,10 @@ onAction: async ({ action, messages, signal }) => { ## Sending actions from the frontend ```ts -// Browser — TriggerChatTransport +// Browser: TriggerChatTransport const stream = await transport.sendAction(chatId, { type: "undo" }); -// Server — AgentChat +// Server: AgentChat const stream = await agentChat.sendAction({ type: "rollback", targetMessageId: "msg-3" }); ``` @@ -104,8 +134,8 @@ The action payload is validated against `actionSchema` on the backend; invalid a ## See also -- [`chat.history`](/ai-chat/backend#chat-history) — the imperative API actions use to mutate state -- [Sending actions from the frontend](/ai-chat/frontend#sending-actions) — `transport.sendAction` ergonomics -- [`hydrateMessages`](/ai-chat/lifecycle-hooks#hydratemessages) — fires before `onAction` when set -- [Branching conversations](/ai-chat/patterns/branching-conversations) — pairs action handlers with backend-controlled history -- [Human-in-the-loop](/ai-chat/patterns/human-in-the-loop) — gating fresh actions while a tool is waiting +- [`chat.history`](/ai-chat/backend#chat-history): the imperative API actions use to mutate state +- [Sending actions from the frontend](/ai-chat/frontend#sending-actions): `transport.sendAction` ergonomics +- [`hydrateMessages`](/ai-chat/lifecycle-hooks#hydratemessages): fires before `onAction` when set +- [Branching conversations](/ai-chat/patterns/branching-conversations): pairs action handlers with backend-controlled history +- [Human-in-the-loop](/ai-chat/patterns/human-in-the-loop): gating fresh actions while a tool is waiting diff --git a/docs/ai-chat/background-injection.mdx b/docs/ai-chat/background-injection.mdx index f84336ff4de..92fa3dd1336 100644 --- a/docs/ai-chat/background-injection.mdx +++ b/docs/ai-chat/background-injection.mdx @@ -1,14 +1,14 @@ --- title: "Background injection" sidebarTitle: "Background injection" -description: "Inject context from background work into the agent's conversation — self-review, RAG augmentation, or any async analysis." +description: "Inject context from background work into the agent's conversation: self-review, RAG augmentation, or any async analysis." --- ## Overview `chat.inject()` queues model messages for injection into the conversation. Messages are picked up at the start of the next turn or at the next `prepareStep` boundary (between tool-call steps). -This is the backend counterpart to [pending messages](/ai-chat/pending-messages) — pending messages come from the user via the frontend, while `chat.inject()` comes from your task code. +This is the backend counterpart to [pending messages](/ai-chat/pending-messages). Pending messages come from the user via the frontend, while `chat.inject()` comes from your task code. ## Basic usage @@ -34,7 +34,7 @@ The most powerful pattern combines `chat.defer()` (background work) with `chat.i export const myChat = chat.agent({ id: "my-chat", onTurnComplete: async ({ messages }) => { - // Kick off background analysis — doesn't block the turn + // Kick off background analysis, doesn't block the turn chat.defer( (async () => { const analysis = await analyzeConversation(messages); @@ -150,7 +150,7 @@ export const myChat = chat.agent({ }); ``` -The self-review runs on `claude-haiku-4-5` (fast, cheap) in the background. If the user sends another message before it completes, the coaching is still injected — `chat.inject()` persists across the idle wait. +The self-review runs on `claude-haiku-4-5` (fast, cheap) in the background. If the user sends another message before it completes, the coaching is still injected, because `chat.inject()` persists across the idle wait. ## Other use cases @@ -161,13 +161,13 @@ The self-review runs on `claude-haiku-4-5` (fast, cheap) in the background. If t ## `chat.defer` standalone -`chat.defer()` is also useful on its own, without `chat.inject()`. Any work whose timing has no resume implication — analytics, audit logs, search-index writes, cache warming — can run in parallel with streaming instead of in the critical path. All deferred promises are awaited (with a 5s timeout) before `onTurnComplete` fires. +`chat.defer()` is also useful on its own, without `chat.inject()`. Any work whose timing has no resume implication (analytics, audit logs, search-index writes, cache warming) can run in parallel with streaming instead of in the critical path. All deferred promises are awaited (with a 5s timeout) before `onTurnComplete` fires. ```ts export const myChat = chat.agent({ id: "my-chat", onTurnStart: async ({ chatId, runId }) => { - // Analytics — fire-and-forget, irrelevant to resume. + // Analytics: fire-and-forget, irrelevant to resume. chat.defer(analytics.track("turn_started", { chatId, runId })); }, run: async ({ messages, signal }) => { @@ -176,10 +176,10 @@ export const myChat = chat.agent({ }); ``` -`chat.defer()` can be called from anywhere during a turn — hooks, `run()`, or nested helpers. All deferred promises are collected and awaited together before `onTurnComplete`. +`chat.defer()` can be called from anywhere during a turn: hooks, `run()`, or nested helpers. All deferred promises are collected and awaited together before `onTurnComplete`. -**Don't use `chat.defer()` for the message-history write in `onTurnStart`.** That write must land *before* the model starts streaming, otherwise a mid-stream page refresh will read `[]` from your DB and lose the user's message from the rendered conversation. See [Database persistence — `onTurnStart`](/ai-chat/patterns/database-persistence#onturnstart). Reserve `chat.defer` for writes whose timing has no resume implication. +**Don't use `chat.defer()` for the message-history write in `onTurnStart`.** That write must land *before* the model starts streaming, otherwise a mid-stream page refresh will read `[]` from your DB and lose the user's message from the rendered conversation. See [Database persistence: `onTurnStart`](/ai-chat/patterns/database-persistence#onturnstart). Reserve `chat.defer` for writes whose timing has no resume implication. ## How it differs from pending messages @@ -189,9 +189,57 @@ export const myChat = chat.agent({ | **Source** | Backend task code | Frontend user input | | **Triggered by** | Your code (e.g. `onTurnComplete` + `chat.defer()`) | User sending a message during streaming | | **Injection point** | Start of next turn, or next `prepareStep` boundary | Next `prepareStep` boundary only | -| **Message role** | Any (`system`, `user`, `assistant`) | Typically `user` | +| **Message role** | Any. `system` becomes an instruction, others join the conversation (see below) | Typically `user` | | **Frontend visibility** | Not visible unless you write custom `data-*` chunks | Visible via `usePendingMessages` hook | +## Two lanes: trusted and untrusted + +The role you inject with decides more than position. It decides whether the model +treats the content as trustworthy. + +**`role: "system"` goes to the instructions lane.** The block is appended to the +system instructions for subsequent inference calls, so it carries the same standing +as your system prompt. This is the lane for context the agent should believe: +entitlements, plan changes, operational notices. + +It has to work this way. On AI SDK 7 a system message inside `messages` is rejected +for every provider. `standardizePrompt` throws before any provider is called, and +its own advice is to use the instructions option, so the injected block goes there +rather than into the transcript. + + + The instructions lane is delivered by `chat.toStreamTextOptions()`, because that + is the only place the SDK can set `streamText`'s instructions for you. If your + `run()` calls `streamText({ model, messages, abortSignal })` without spreading + `chat.toStreamTextOptions()`, a `role: "system"` injection never reaches the + model. The conversational lane has no such requirement: it arrives through + `messages` either way. + + +- An injection applies to the next turn only. A block injected in `onTurnComplete` + shapes the following turn and is cleared after it, so it is not repeated on every + turn from then on. Within that turn it is consumed once rather than once per read, + so a `run()` that builds options more than once sees the same instructions in + every build. +- The injected text is merged into a single instruction rather than added as a + second block, because AI SDK 5 rejects an array of system blocks while accepting + one structured block. Merging changes the cached prefix, so a cached system prompt + gets no cache hit for as long as an injection is live. If you rely on prompt + caching, inject sparingly and prefer facts that go stale, so the injection clears. + +**Any other role joins the conversation, and is untrusted by construction.** A +message injected as `user` is indistinguishable from something the user typed, and a +well-aligned model treats it accordingly, and may say so and re-derive the answer +from tools instead of taking it at face value: + +> "that text arrived embedded in your message, not from a tool I called, so I +> verified it myself rather than trusting it" + +That is correct behaviour, not a bug. So inject **checkable facts** in the +conversational lane and put **directives** in the instructions lane. A conclusion +injected as a user message is the worst of both: the model neither trusts it nor +ignores it, and may contradict it in front of the user. + ## API reference ### chat.inject() @@ -200,7 +248,7 @@ export const myChat = chat.agent({ chat.inject(messages: ModelMessage[]): void ``` -Queue model messages for injection at the next opportunity. Messages persist across the idle wait between turns — they are not reset when a new turn starts. +Queue model messages for injection at the next opportunity. Messages persist across the idle wait between turns, and are not reset when a new turn starts. **Parameters:** @@ -209,9 +257,9 @@ Queue model messages for injection at the next opportunity. Messages persist acr | `messages` | `ModelMessage[]` | Model messages to inject (from the `ai` package) | Messages are drained (consumed) when: -1. A new turn starts — before `run()` executes -2. A `prepareStep` boundary is reached — between tool-call steps during streaming +1. A new turn starts, before `run()` executes +2. A `prepareStep` boundary is reached, between tool-call steps during streaming - `chat.inject()` writes to an in-memory queue in the current process. It works from any code running in the same task — lifecycle hooks, deferred work, tool execute functions, etc. It does not work from subtasks or other runs. + `chat.inject()` writes to an in-memory queue in the current process. It works from any code running in the same task: lifecycle hooks, deferred work, tool execute functions, etc. It does not work from subtasks or other runs. diff --git a/docs/ai-chat/testing.mdx b/docs/ai-chat/testing.mdx index 65e094ed530..a7d0a425d18 100644 --- a/docs/ai-chat/testing.mdx +++ b/docs/ai-chat/testing.mdx @@ -634,6 +634,7 @@ The harness's initial wire payload depends on `mode`: | `sendHandover({ partialAssistantMessage, isFinal?, messageId? })` | Dispatch a `handover` signal — only meaningful when started with `mode: "handover-prepare"`. The agent picks up partial assistant messages and continues the turn. | | `sendHandoverSkip()` | Dispatch a `handover-skip` signal — only meaningful when started with `mode: "handover-prepare"`. The agent exits cleanly without firing turn hooks. | | `sendAction(action)` | Route a custom action through `actionSchema` + `onAction`. | +| `sendPendingMessage(message)` | Append a user message mid-turn without waiting for a turn to complete, so it reaches the running turn as a steering message. Resolves once the record has landed on `session.in`. | | `sendStop(message?)` | Fire a stop signal. Does not wait for the turn — the run's `signal.aborted` becomes `true`. | | `seedSnapshot(snapshot)` | Pre-seed the snapshot read for the next boot. Effective on the next run boot only. | | `seedSessionOutTail(chunks?)` | Pre-seed `session.out` chunks for the next boot's replay. Reduces to settled assistant turns. | diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index bcc70fa9ce0..2f7879a1e93 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -46,6 +46,7 @@ import type { FinishReason, LanguageModelUsage, ModelMessage, + SystemModelMessage, ProviderMetadata, Tool, ToolSet, @@ -2695,6 +2696,35 @@ function spliceHandoverPartial( */ const chatBackgroundQueueKey = locals.create("chat.backgroundQueue"); +/** + * System-role context injected mid-conversation, held for the instructions lane. + * + * Kept apart from the message queue because ai@7 rejects a system message inside + * `messages` for every provider — `standardizePrompt` throws upstream of any + * provider call, and its own advice is to use the instructions option. Instructions + * accept `Array`, so a system-role injection has a correct + * home: appended as another system block rather than smuggled into the transcript. + * + * This is also the only way to inject *trusted* context. A message injected as + * `user` is untrusted by construction, and a well-aligned model treats it that + * way — it will say so, and re-derive the answer from tools instead. + */ +const chatInjectedInstructionsKey = locals.create( + "chat.injectedInstructions" +); +/** + * What a turn already consumed from the instructions lane, so a second + * `toStreamTextOptions()` call in the same turn sees the same blocks. + * + * Consumed blocks are moved here rather than left in the pending lane: leaving + * them there means an injection made during the consumed turn sits behind them, + * and clearing the lane on the next turn destroys both. + */ +const chatInstructionsConsumedKey = locals.create<{ + turn: number; + blocks: SystemModelMessage[]; +}>("chat.injectedInstructionsConsumed"); + /** * Run-scoped pipe counter. Stored in locals so concurrent runs in the * same worker don't share state. @@ -3521,6 +3551,41 @@ type SteeringQueueEntry = { const chatPendingMessagesKey = locals.create("chat.pendingMessages"); /** @internal */ const chatSteeringQueueKey = locals.create("chat.steeringQueue"); + +/** + * This turn's new messages, as `onTurnComplete.newUIMessages` will see them. + * + * Held in locals because `drainSteeringQueue` runs outside the turn closure and + * has to append the messages it injects. Without that, an injected message + * reaches the model and the browser but no hook, so an app persisting from + * `onTurnComplete` never learns it existed. + */ +const chatTurnNewUIMessagesKey = locals.create("chat.turnNewUIMessages"); + +/** + * Steering messages a drain consumed that the model accumulator has not been + * given yet. + * + * The two accumulators are maintained separately, and the model one is + * normally advanced by appending each turn's delta. A drained message is + * appended to the UI one but reaches the model only through the `prepareStep` + * return value, which is per-step: without this the model lane never learns + * the message exists and every later turn of the run answers without it, + * while the browser, the snapshot and `chat.history.*` all still show it. + * + * Held as the messages rather than a "rebuild me" flag because the model lane + * can only be appended to, never reconstructed. Compaction replaces it with a + * summary and deliberately leaves the UI lane whole, so reconverting the UI + * lane restores every message the summary replaced. + */ +const chatPendingSteerKey = locals.create("chat.pendingSteer"); + +/** + * A consumed steering message in both forms: the UI message for display and + * persistence, and the model messages `pendingMessages.prepare` produced for + * it, which is what the model actually saw and what later turns must see too. + */ +type PendingSteer = { ui: UIMessage; model: ModelMessage[] }; /** @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 */ @@ -4105,11 +4170,36 @@ function chatCompactionStep( // Steering queue drain — shared by toStreamTextOptions, session, accumulator // --------------------------------------------------------------------------- +/** What a steering drain produced: what to send now, and what it consumed. */ +type DrainedSteering = { + /** Model messages to add to this step's prompt. */ + injected: ModelMessage[]; + /** The UI messages the drain consumed, for the caller to record. */ + claimed: UIMessage[]; +}; + +const EMPTY_DRAIN: DrainedSteering = { injected: [], claimed: [] }; + +/** + * The model messages to record for one claimed message. Without `prepare` + * each entry's own conversion is used. With it, `prepare` returned one list + * for the whole batch, so the first claimed message carries all of it and the + * rest carry none, which keeps the total exactly what the model received. + */ +function modelFormOf(m: UIMessage, batch: UIMessage[], injected: ModelMessage[]): ModelMessage[] { + return batch[0] === m ? injected : []; +} + /** * 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). + * Returns the model messages to inject and the UI messages actually claimed. + * + * `claimed` is returned rather than only published to locals because each + * surface files it somewhere different: `chat.agent` has an accumulator in + * locals, while `chat.createSession` keeps its own. Publishing to locals alone + * is silently a no-op for any surface that never set the key. * @internal */ async function drainSteeringQueue( @@ -4117,9 +4207,9 @@ async function drainSteeringQueue( messages: ModelMessage[], steps: CompactionStep[], queueOverride?: SteeringQueueEntry[] -): Promise { +): Promise { const queue = queueOverride ?? locals.get(chatSteeringQueueKey); - if (!queue || queue.length === 0) return []; + if (!queue || queue.length === 0) return EMPTY_DRAIN; const ctx = locals.get(chatTurnContextKey); const stepNumber = steps.length - 1; @@ -4145,7 +4235,7 @@ async function drainSteeringQueue( // Call shouldInject once for the whole batch const shouldInject = config.shouldInject ? await config.shouldInject(batchEvent) : false; - if (!shouldInject) return []; + if (!shouldInject) return EMPTY_DRAIN; const textOfUIMessage = (m: UIMessage) => (m.parts ?? []) @@ -4195,7 +4285,7 @@ async function drainSteeringQueue( if (at !== -1) queue.splice(at, 1); } - if (claimed.length === 0) return []; + if (claimed.length === 0) return EMPTY_DRAIN; /** * Give the claim back if the transform fails. `prepare` is caller code and @@ -4224,6 +4314,38 @@ async function drainSteeringQueue( for (const m of claimedUIMessages) injectedIds.add(m.id); } + // Record them as part of the conversation. + // + // The model has them and the browser has them; without this the + // accumulator does not, so they reach neither `uiMessages` nor + // `newUIMessages` on `onTurnComplete` and an app that persists from there + // silently loses the instruction the answer was shaped by. Appending here + // rather than at turn end keeps them in the order they happened: after the + // message that started the turn, before the response that answers it. + // + // De-duplicated by id because a step boundary can drain more than once per + // turn, and because a message that failed to inject falls back to becoming + // its own turn, where it is accumulated the normal way. + const currentUIMessages = locals.get(chatCurrentUIMessagesKey); + const turnNew = locals.get(chatTurnNewUIMessagesKey); + for (const m of claimedUIMessages) { + if (currentUIMessages && !currentUIMessages.some((existing) => existing.id === m.id)) { + currentUIMessages.push(m); + } + if (turnNew && !turnNew.some((existing) => existing.id === m.id)) { + turnNew.push(m); + } + } + if (claimedUIMessages.length > 0 && currentUIMessages) { + const pendingSteer = locals.get(chatPendingSteerKey) ?? []; + for (const m of claimedUIMessages) { + if (!pendingSteer.some((existing) => existing.ui.id === m.id)) { + pendingSteer.push({ ui: m, model: modelFormOf(m, claimedUIMessages, injected) }); + } + } + locals.set(chatPendingSteerKey, pendingSteer); + } + // Write injection confirmation chunk to the stream so the frontend // knows which messages were injected and where in the response. if (injected.length > 0) { @@ -4265,7 +4387,7 @@ async function drainSteeringQueue( } } - return injected; + return { injected, claimed: claimedUIMessages }; }, { attributes: { @@ -4658,6 +4780,86 @@ function toStreamTextOptions(options?: ToStreamTextOptionsOptions): Record`. This package's peer + * range still spans all three, so emitting an array unconditionally would break + * v5 consumers — for whom a system-role injection used to work, since v5 accepted + * a system message inside `messages` that v7 rejects. + * + * So: concatenate into one string when the base is a plain string, which every + * version accepts and which loses nothing (separate blocks only matter for + * per-block `providerOptions`). Use the array form only when the base is already + * a structured message — that path requires v6+ regardless, because it is how + * prompt caching marks the system block, and flattening it would silently throw + * the cache away. + * + * Either way the injected text goes last: the base prompt keeps its position for + * caching, and the addition reads as a later amendment. A changed prefix does + * cost the first call its cache hit, on turns that actually injected. + */ + /** + * Consumed once per turn, not once per read, and moved out of the lane rather + * than marked read in place. + * + * Per turn, because a `run()` that builds options twice (a cheap classifier + * pass and then the answer) has to see the injection in both, and draining on + * read hands it to whichever call ran first. Moved out, because blocks left in + * the lane sit in front of anything injected during the same turn, and + * clearing the lane on the next turn then destroys both. Outside a turn there + * is no turn to scope the stash to, so the lane drains on read there. + */ + const injectedInstructions = locals.get(chatInjectedInstructionsKey); + const currentTurn = locals.get(chatTurnContextKey)?.turn; + const consumedThisTurn = + currentTurn === undefined ? undefined : locals.get(chatInstructionsConsumedKey); + + let injectedBlocks: SystemModelMessage[] = []; + if (consumedThisTurn && consumedThisTurn.turn === currentTurn) { + injectedBlocks = consumedThisTurn.blocks; + } else if (injectedInstructions && injectedInstructions.length > 0) { + injectedBlocks = injectedInstructions.splice(0); + if (currentTurn !== undefined) { + locals.set(chatInstructionsConsumedKey, { turn: currentTurn, blocks: injectedBlocks }); + } + } + + if (injectedBlocks.length > 0) { + const blocks = injectedBlocks; + + const injectedText = blocks + .map((block) => (typeof block.content === "string" ? block.content : "")) + .filter(Boolean) + .join("\n\n"); + + const base = result.system; + + if (base === undefined) { + result.system = injectedText; + } else if (typeof base === "string") { + result.system = [base, injectedText].filter(Boolean).join("\n\n"); + } else { + // Merged into the existing block rather than added as a second one. An array + // of system blocks would keep the base block's cache entry, but ai@5 rejects + // it outright ("Invalid prompt: system must be a string") while accepting a + // single structured block, and this package's peer range still spans v5. + // Choosing per version would mean resolving the installed version at runtime, + // which is not something to build on: `import.meta.url` is illegal in this + // package's CommonJS output, and a bundled task may have no resolvable `ai` + // to read. One shape that works everywhere beats a cache hit. + const baseBlock = base as SystemModelMessage; + result.system = { + ...baseBlock, + content: [typeof baseBlock.content === "string" ? baseBlock.content : "", injectedText] + .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 @@ -4724,7 +4926,7 @@ function toStreamTextOptions(options?: ToStreamTextOptionsOptions): Record AsyncIterable | ReadableStream; }; +/** + * A plain `onAction` reply (`string` or assistant `UIMessage`) as a stream, so + * it takes the same path as a `StreamTextResult`: piped to the browser, + * captured, committed to the conversation, snapshotted. Text and `data-*` + * parts are emitted; anything else in a supplied message is dropped, since a + * tool part with no execution behind it cannot be replayed as chunks. + */ +function plainReplyAsStream(value: unknown): UIMessageStreamable | undefined { + let message: UIMessage | undefined; + if (typeof value === "string") { + message = { + id: generateMessageId(), + role: "assistant", + parts: [{ type: "text", text: value }], + } as UIMessage; + } else if ( + typeof value === "object" && + value !== null && + (value as UIMessage).role === "assistant" && + Array.isArray((value as UIMessage).parts) + ) { + const m = value as UIMessage; + message = { ...m, id: m.id || generateMessageId() }; + } + if (!message) return undefined; + + const chunks: Record[] = [{ type: "start", messageId: message.id }]; + let n = 0; + for (const part of message.parts as { + type: string; + text?: string; + data?: unknown; + id?: string; + }[]) { + if (part.type === "text") { + const id = `t${n++}`; + chunks.push( + { type: "text-start", id }, + { type: "text-delta", id, delta: part.text ?? "" }, + { type: "text-end", id } + ); + } else if (part.type.startsWith("data-")) { + chunks.push({ type: part.type, id: part.id, data: part.data }); + } + } + chunks.push({ type: "finish" }); + + return { + toUIMessageStream: () => + new ReadableStream({ + start(controller) { + for (const c of chunks) controller.enqueue(c); + controller.close(); + }, + }), + } as unknown as UIMessageStreamable; +} + +/** + * Replace, in a model lane, the run of messages one UI message contributed. + * + * Used when a UI message is replaced in place (a tool-approval continuation + * merging onto the trailing assistant, a captured response reusing an existing + * id, a partial replacing an existing message). Reconverting the whole lane + * from the UI lane would also replace a compaction summary with the full + * transcript and drop the model forms `pendingMessages.prepare` produced. + * + * The replaced message is the trailing one, so its run is the lane's tail, + * before any steer forms appended after it this turn (`tailAfter`). If the + * tail does not match the old message's conversion, nothing is changed and + * `false` is returned so the caller can fall back to a full reconversion. + */ +async function replaceModelRun( + lane: ModelMessage[], + oldUi: UIMessage, + newUi: UIMessage, + tailAfter: number +): Promise { + const oldRun = await toModelMessages([stripProviderMetadata(oldUi)]); + const newRun = await toModelMessages([stripProviderMetadata(newUi)]); + const end = lane.length - tailAfter; + const start = end - oldRun.length; + if (start < 0 || end > lane.length) return false; + if (JSON.stringify(lane.slice(start, end)) !== JSON.stringify(oldRun)) return false; + lane.splice(start, oldRun.length, ...newRun); + return true; +} + function isUIMessageStreamable(value: unknown): value is UIMessageStreamable { return ( typeof value === "object" && @@ -6498,6 +6788,26 @@ function chatAgent< // durable snapshot + `session.out` replay (or `hydrateMessages` if // registered) — the wire is delta-only now, no longer a seed. let accumulatedMessages: ModelMessage[] = []; + /** + * Give the model accumulator the steering messages a drain consumed, + * in the form the model actually received. Appended, never reconverted + * from the UI lane, so a model-only compaction summary survives. Called + * on both the success and the error path, before the response or the + * partial joins the lane, so the order stays steer-then-answer. + */ + const reconcilePendingSteer = (options?: { + /** This turn's model delta, as `onTurnComplete.newMessages` reports it. */ + turnNew?: ModelMessage[]; + }): PendingSteer[] => { + const pending = locals.get(chatPendingSteerKey); + if (!pending || pending.length === 0) return []; + locals.set(chatPendingSteerKey, []); + for (const entry of pending) { + accumulatedMessages.push(...entry.model); + options?.turnNew?.push(...entry.model); + } + return pending; + }; // Accumulated UI messages for persistence. Mirrors the model accumulator // but in frontend-friendly UIMessage format (with parts, id, etc.). @@ -6518,6 +6828,66 @@ function chatAgent< // swallow errors internally; the agent stays available either way. const sessionIdForSnapshot = payload.sessionId ?? payload.chatId; let bootSnapshot: ChatSnapshotV1 | undefined; + + /** + * The `lastOutEventId` the most recent snapshot carried. + * + * A snapshot written outside a turn — after an action mutates history — has + * no turn cursor of its own, and writing `undefined` there would drop the + * resume point and make the next boot replay from further back. Retaining it + * keeps an action's write cursor-neutral. + */ + let lastSnapshotOutEventId: string | undefined; + + /** + * Persist the accumulator outside a turn. + * + * An action is not a turn, so it never reaches the turn-complete path where + * the snapshot is normally written — but it can change the conversation in + * two ways: a `chat.history` mutation, and a response streamed back from + * `onAction`. Both have to survive, and one write at the end of the action + * covers both rather than writing twice for a regenerate that does both. + * + * Cursor-neutral: an action has no turn cursor of its own, and writing + * `undefined` would drop the resume point the last turn established and make + * the next boot replay from further back. + */ + const writeSnapshotOutsideTurn = async (reason: string) => { + if (hydrateMessages) return; + try { + await tracer.startActiveSpan( + "snapshot.write", + async () => { + const snapshotInCursor = chatInputRouter().resumeFloor(); + await writeChatSnapshot(sessionIdForSnapshot, { + version: 1, + savedAt: Date.now(), + messages: accumulatedUIMessages, + lastOutEventId: lastSnapshotOutEventId, + lastInEventId: + snapshotInCursor !== undefined ? String(snapshotInCursor) : undefined, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.snapshot.reason": reason, + "chat.messages.count": accumulatedUIMessages.length, + }, + } + ); + } catch (error) { + logger.warn( + "chat.agent: snapshot write outside a turn failed; the change may not survive a continuation", + { + error: error instanceof Error ? error.message : String(error), + sessionId: sessionIdForSnapshot, + reason, + } + ); + } + }; let replayedSettled: TUIMessage[] = []; let replayedPartial: TUIMessage | undefined; let replayedPartialRaw: TUIMessage | undefined; @@ -6568,6 +6938,8 @@ function chatAgent< // Without seeding, the new worker would emit no trim on its first // turn (chain self-bootstraps from turn 2), so this is purely an // optimization to keep continuation runs bounded from the first turn. + lastSnapshotOutEventId = bootSnapshot?.lastOutEventId; + if (bootSnapshot?.lastOutEventId !== undefined) { const seeded = Number.parseInt(bootSnapshot.lastOutEventId, 10); if (Number.isFinite(seeded)) { @@ -7490,6 +7862,7 @@ function chatAgent< // Track new messages for this turn (user input + assistant response). const turnNewModelMessages: ModelMessage[] = []; const turnNewUIMessages: TUIMessage[] = []; + locals.set(chatTurnNewUIMessagesKey, turnNewUIMessages); // ── Action handling ────────────────────────────────────── // Actions arrive on the same input stream but with @@ -7502,6 +7875,13 @@ function chatAgent< // string, or UIMessage from `onAction`. Turn counter // does not advance. let actionStreamResult: unknown = undefined; + /** + * Whether this action changed the conversation, by rolling history + * back or by streaming a response. Drives the single snapshot write + * at the end — an action never reaches the turn-complete path that + * normally does it. + */ + let actionChangedHistory = false; if (isAction) { // Parse and validate the action payload const parsedAction = parseAction @@ -7573,6 +7953,8 @@ function chatAgent< accumulatedUIMessages = [...actionOverride] as TUIMessage[]; accumulatedMessages = await toModelMessages(actionOverride); locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + + actionChangedHistory = true; } } else { warnMissingOnActionOnce(); @@ -7766,6 +8148,7 @@ function chatAgent< // where AI SDK regenerates the id (TRI-9137) still // applies via `rewriteIncomingIdViaToolCallMap`. let replaced = false; + const replacedPairs: { previous: TUIMessage; merged: TUIMessage }[] = []; for (const raw of cleanedUIMessages) { let incoming = raw; let idx = accumulatedUIMessages.findIndex((m) => m.id === incoming.id); @@ -7777,10 +8160,12 @@ function chatAgent< } } if (idx !== -1) { + const previous = accumulatedUIMessages[idx]!; accumulatedUIMessages[idx] = mergeIncomingIntoHydrated( - accumulatedUIMessages[idx]!, + previous, incoming ) as TUIMessage; + replacedPairs.push({ previous, merged: accumulatedUIMessages[idx]! }); replaced = true; } else { accumulatedUIMessages.push(incoming as TUIMessage); @@ -7789,9 +8174,19 @@ function chatAgent< recordToolCallIdsFromMessage(incoming); } if (replaced) { - // Replacement changes structure — reconvert all model - // messages instead of appending. - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + let inPlace = true; + for (const { previous, merged } of replacedPairs) { + if (!(await replaceModelRun(accumulatedMessages, previous, merged, 0))) { + inPlace = false; + break; + } + } + if (!inPlace) { + logger.warn( + "chat.agent: replaced message not found at the model lane tail; reconverting the lane" + ); + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } } else { const incomingModelMessages = await toModelMessages(cleanedUIMessages); accumulatedMessages.push(...incomingModelMessages); @@ -7851,22 +8246,73 @@ function chatAgent< if (isAction) { msgSub?.off(); + // A documented plain reply takes the streamed reply's path. + actionStreamResult = plainReplyAsStream(actionStreamResult) ?? actionStreamResult; + 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", - }); + /** + * Captured, not just piped. The stream reaching the browser was + * never the problem — the problem was that it stopped there, so + * the user read an answer the accumulator had no record of and + * the next turn contradicted the screen. Worst on regenerate, + * which removes the old answer and used to leave nothing in its + * place. + * + * Persistence beyond the snapshot is still the app's job: an + * action fires no `onTurnComplete`, so an app owning its own + * store has to write the row itself — `chat.pipeAndCapture` + * hands back the same message for that. + */ + const captured = await pipeChatAndCapture( + actionStreamResult as UIMessageStreamable, + { signal: combinedSignal, spanName: "stream response" } + ); + + if (runSignal.aborted) return "exit"; + + /** + * A stopped action still commits what streamed, cleaned: + * incomplete tool and text parts left mid-flight are what + * strand the UI on a spinner forever once persisted. + */ + const actionResponse = + captured.status === "complete" || !captured.message + ? captured.message + : cleanupAbortedParts(captured.message); + + if (actionResponse) { + const existingIdx = actionResponse.id + ? accumulatedUIMessages.findIndex((m) => m.id === actionResponse.id) + : -1; + if (existingIdx !== -1) { + accumulatedUIMessages[existingIdx] = actionResponse as TUIMessage; + // Replacing an existing message has no in-place model + // form to swap, so this path still reconverts. + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } else { + accumulatedUIMessages.push(actionResponse as TUIMessage); + // Appended, not reconverted: a reconversion from the UI + // lane would undo a model-only compaction summary. + accumulatedMessages.push( + ...(await toModelMessages([stripProviderMetadata(actionResponse)])) + ); + } + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + actionChangedHistory = true; + } + + /** + * Reported after the partial is committed, not instead of it. + * `pipeChatAndCapture` returns a stream failure rather than + * throwing, so without this a mid-stream failure writes a + * normal turn-complete and the truncated answer is persisted + * as if it were finished — the next turn then builds on it. + */ + if (captured.status === "error") throw captured.error; } catch (error) { if ( error instanceof Error && @@ -7875,10 +8321,26 @@ function chatAgent< ) { return "exit"; } - throw error; + // Reported here rather than rethrown: the shared catch + // below is the turn-error path, and it would fire + // onTurnComplete, keep the turn number and consume the + // one-shot instruction lane, none of which an action does. + try { + await withChatWriter(async (writer) => { + const errorText = + error instanceof Error ? error.message : "An unexpected error occurred"; + writer.write({ type: "error", errorText } as any); + }); + } catch { + // best effort + } } } + if (actionChangedHistory) { + await writeSnapshotOutsideTurn("action"); + } + await writeTurnCompleteChunk(currentWirePayload.chatId); // Don't consume a turn iteration — actions aren't turns. @@ -8214,7 +8676,23 @@ function chatAgent< if (runOverride) { locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...runOverride] as TUIMessage[]; - accumulatedMessages = await toModelMessages(runOverride); + /** + * Steers the drain consumed are left out of the rebuild and + * appended by the reconciliation below instead, so the lane + * gets the form the model actually received rather than a + * reconversion of the UI message, and gets it once. A steer + * the edit removed is dropped from the pending list too, so + * the edit is honoured. + */ + const overrideIds = new Set(runOverride.map((m) => m.id)); + const pending = (locals.get(chatPendingSteerKey) ?? []).filter((e) => + overrideIds.has(e.ui.id) + ); + locals.set(chatPendingSteerKey, pending); + const pendingIds = new Set(pending.map((e) => e.ui.id)); + accumulatedMessages = await toModelMessages( + runOverride.filter((m) => !pendingIds.has(m.id)) + ); locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } @@ -8252,6 +8730,16 @@ function chatAgent< // Determine if the user stopped generation this turn (not a full run cancel). const wasStopped = stopController.signal.aborted && !runSignal.aborted; + // Give the model accumulator the steering messages the drain + // consumed. Appended, never reconverted from the UI lane, so a + // model-only compaction summary set just above survives; and done + // before the response is appended so the order stays + // steer-then-answer. Outside the `capturedResponseMessage` + // branches below, so a turn that captured no response is covered. + const steerTailThisTurn = reconcilePendingSteer({ + turnNew: turnNewModelMessages, + }).reduce((n, e) => n + e.model.length, 0); + // 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. @@ -8289,6 +8777,8 @@ function chatAgent< const existingIdx = capturedResponseMessage.id ? accumulatedUIMessages.findIndex((m) => m.id === capturedResponseMessage!.id) : -1; + const previousAtIdx = + existingIdx !== -1 ? accumulatedUIMessages[existingIdx] : undefined; if (existingIdx !== -1) { accumulatedUIMessages[existingIdx] = capturedResponseMessage; } else { @@ -8307,8 +8797,20 @@ function chatAgent< stripProviderMetadata(capturedResponseMessage), ]); if (existingIdx !== -1) { - // Reconvert all model messages since we replaced rather than appended - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + const ok = + previousAtIdx !== undefined && + (await replaceModelRun( + accumulatedMessages, + previousAtIdx, + capturedResponseMessage, + steerTailThisTurn + )); + if (!ok) { + logger.warn( + "chat.agent: replaced response not found at the model lane tail; reconverting the lane" + ); + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } } else { accumulatedMessages.push(...responseModelMessages); } @@ -8649,11 +9151,13 @@ function chatAgent< "snapshot.write", async () => { const snapshotInCursor = chatInputRouter().resumeFloor(); + lastSnapshotOutEventId = + turnCompleteResult?.lastEventId ?? lastSnapshotOutEventId; await writeChatSnapshot(sessionIdForSnapshot, { version: 1, savedAt: Date.now(), messages: accumulatedUIMessages, - lastOutEventId: turnCompleteResult?.lastEventId, + lastOutEventId: lastSnapshotOutEventId, lastInEventId: snapshotInCursor !== undefined ? String(snapshotInCursor) : undefined, }); @@ -8832,6 +9336,10 @@ function chatAgent< }); // Signal turn complete so the client knows this turn is done errorTurnCompleteResult = await writeTurnCompleteChunk(currentWirePayload.chatId); + // A later action's snapshot reuses this cursor, so it has to move + // here too or that snapshot resumes from before the failed turn. + lastSnapshotOutEventId = + errorTurnCompleteResult?.lastEventId ?? lastSnapshotOutEventId; } catch { // Best-effort — if stream write fails, let the run continue anyway } @@ -8883,19 +9391,51 @@ function chatAgent< i === partialIdx ? partialResponse! : m ) as TUIMessage[]); - let erroredNewUIMessages: TUIMessage[] = erroredWireMessage ? [erroredWireMessage] : []; - if (includePartial) { - erroredNewUIMessages.push(partialResponse!); - } + /** + * Seeded from the per-turn list, not just the wire message and the + * partial, so a steering message the drain consumed is reported too. + * An app persisting from `newUIMessages` would otherwise lose the + * instruction whenever the turn it steered went on to fail. + */ + const buildErroredNew = (): TUIMessage[] => { + const out: TUIMessage[] = []; + const addUnique = (m?: TUIMessage) => { + if (m && !out.some((existing) => existing.id === m.id)) out.push(m); + }; + addUnique(erroredWireMessage); + for (const m of (locals.get(chatTurnNewUIMessagesKey) ?? []) as TUIMessage[]) { + addUnique(m); + } + if (includePartial) addUnique(partialResponse!); + return out; + }; + + let erroredNewUIMessages: TUIMessage[] = buildErroredNew(); let erroredNewModelMessages: ModelMessage[] = []; + const reconciledSteer = reconcilePendingSteer(); + if (!responseCommitted) { try { if (erroredNewUIMessages.length > 0) { - erroredNewModelMessages = await toModelMessages( - erroredNewUIMessages.map((m) => stripProviderMetadata(m)) + /** + * Built in order from the recorded forms rather than by + * converting the UI list, so a steer appears in the delta as + * the model received it (what `prepare` produced), matching the + * lane. The wire message and partial are converted as before. + */ + const steerModelById = new Map( + reconciledSteer.map((e) => [e.ui.id, e.model] as const) ); + for (const m of erroredNewUIMessages) { + const recorded = steerModelById.get(m.id); + if (recorded) erroredNewModelMessages.push(...recorded); + else + erroredNewModelMessages.push( + ...(await toModelMessages([stripProviderMetadata(m)])) + ); + } } if (erroredUIMessagesWithPartial !== accumulatedUIMessages) { if (partialIdx === -1) { @@ -8906,7 +9446,18 @@ function chatAgent< ...(await toModelMessages(appended.map((m) => stripProviderMetadata(m)))) ); } else { - accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial); + const ok = await replaceModelRun( + accumulatedMessages, + erroredUIMessages[partialIdx]!, + partialResponse!, + reconciledSteer.reduce((n, e) => n + e.model.length, 0) + ); + if (!ok) { + logger.warn( + "chat.agent: replaced partial not found at the model lane tail; reconverting the lane" + ); + accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial); + } } accumulatedUIMessages = erroredUIMessagesWithPartial; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); @@ -8914,7 +9465,7 @@ function chatAgent< } catch { erroredNewModelMessages = []; erroredUIMessagesWithPartial = erroredUIMessages; - erroredNewUIMessages = erroredWireMessage ? [erroredWireMessage] : []; + erroredNewUIMessages = buildErroredNew().filter((m) => m !== partialResponse); } } @@ -9813,9 +10364,22 @@ function chatDefer(promiseOrFn: Promise | (() => Promise)): vo * ``` */ function injectBackgroundContext(messages: ModelMessage[]): void { - const queue = locals.get(chatBackgroundQueueKey) ?? []; - queue.push(...messages); - locals.set(chatBackgroundQueueKey, queue); + const systemBlocks = messages.filter( + (message): message is SystemModelMessage => message.role === "system" + ); + const conversational = messages.filter((message) => message.role !== "system"); + + if (systemBlocks.length > 0) { + const instructions = locals.get(chatInjectedInstructionsKey) ?? []; + instructions.push(...systemBlocks); + locals.set(chatInjectedInstructionsKey, instructions); + } + + if (conversational.length > 0) { + const queue = locals.get(chatBackgroundQueueKey) ?? []; + queue.push(...conversational); + locals.set(chatBackgroundQueueKey, queue); + } } // --------------------------------------------------------------------------- @@ -10299,12 +10863,14 @@ class ChatMessageAccumulator { // a duplicate, mirroring the chat.agent accumulator. const existingIdx = this.uiMessages.findIndex((m) => m.id === response.id); if (existingIdx !== -1) { + const previous = this.uiMessages[existingIdx]!; this.uiMessages[existingIdx] = response; try { - // Reconvert all model messages since we replaced rather than appended. - this.modelMessages = await toModelMessages( - this.uiMessages.map((m) => stripProviderMetadata(m)) - ); + if (!(await replaceModelRun(this.modelMessages, previous, response, 0))) { + this.modelMessages = await toModelMessages( + this.uiMessages.map((m) => stripProviderMetadata(m)) + ); + } } catch { // Conversion failed — leave the existing model messages in place } @@ -10340,6 +10906,30 @@ class ChatMessageAccumulator { this._steeringQueue.push({ uiMessage: message, modelMessages: modelMsgs }); } + /** + * Record the messages a steering drain consumed. + * + * The drain only puts them in this step's prompt, so without this they + * shape one answer and then exist in neither lane: not in `uiMessages`, + * which is what an app persists from, and not in `modelMessages`, which is + * what every later turn sends. + * + * Both lanes are appended to. The model lane is never reconverted from the + * UI lane, because `compactIfNeeded` replaces it with a summary and leaves + * the UI lane whole: a reconversion would restore everything the summary + * replaced. + */ + async absorbSteering(claimed: UIMessage[], injected?: ModelMessage[]): Promise { + const fresh = claimed.filter((m) => !this.uiMessages.some((e) => e.id === m.id)); + if (fresh.length === 0) return; + this.uiMessages.push(...fresh); + // Record what the model received. Only when the whole batch is new is + // `injected` known to describe exactly these messages. + this.modelMessages.push( + ...(injected && fresh.length === claimed.length ? injected : await toModelMessages(fresh)) + ); + } + /** * Get and clear unconsumed steering messages. */ @@ -10380,7 +10970,13 @@ class ChatMessageAccumulator { // 2. Pending message injection if (pm && queue.length > 0) { - const injected = await drainSteeringQueue(pm, resultMessages ?? messages, steps, queue); + const { injected, claimed } = await drainSteeringQueue( + pm, + resultMessages ?? messages, + steps, + queue + ); + await this.absorbSteering(claimed, injected); if (injected.length > 0) { resultMessages = [...(resultMessages ?? messages), ...injected]; } @@ -11150,12 +11746,13 @@ function createChatSession( } if (sessionPendingMessages) { - const injected = await drainSteeringQueue( + const { injected, claimed } = await drainSteeringQueue( sessionPendingMessages, resultMessages ?? stepMsgs, steps, turnSteeringQueue ); + await accumulator.absorbSteering(claimed, injected); if (injected.length > 0) { resultMessages = [...(resultMessages ?? stepMsgs), ...injected]; } diff --git a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts index 63768b9b3f2..e50df390fb5 100644 --- a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts +++ b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts @@ -1,5 +1,5 @@ import type { UIMessage, UIMessageChunk } from "ai"; -import { resourceCatalog } from "@trigger.dev/core/v3"; +import { resourceCatalog, sessionStreams } 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"; @@ -186,6 +186,17 @@ export type MockChatAgentHarness = { /** Send a custom action and wait for the next turn-complete. */ sendAction(action: unknown): Promise; + /** + * Deliver a message mid-turn without waiting for it, the way the browser's + * steering path does. With a `pendingMessages` config the agent routes it into + * the steering queue for injection at the next step boundary; without one it + * buffers as the next turn. + * + * Send it while a turn is in flight — start the turn without awaiting it, then + * call this. Awaiting the turn first leaves nothing to steer. + */ + sendPendingMessage(message: UIMessage): Promise; + /** Fire a stop signal. Does not wait for the turn — the task keeps running. */ sendStop(message?: string): Promise; @@ -618,6 +629,39 @@ export function mockChatAgent( }); }, + async sendPendingMessage(message) { + await harnessReady; + + const seqBefore = sessionStreams.lastSeqNum(chatId, "in") ?? -1; + + await sendSessionInput(sessionId, { + kind: "message", + payload: { + message, + chatId, + trigger: "submit-message", + metadata: clientData, + }, + }); + + /** + * Wait for the record to be observable on the channel, not merely for the + * send call to return. A test that continues on the send alone is racing the + * append: the message can still be in flight when the step boundary runs, so + * the injection it was meant to trigger silently does not happen and the test + * passes while proving nothing. + */ + const deadline = Date.now() + 5_000; + while ((sessionStreams.lastSeqNum(chatId, "in") ?? -1) <= seqBefore) { + if (Date.now() > deadline) { + throw new Error( + `sendPendingMessage: append for ${message.id} never landed on session.in` + ); + } + await new Promise((resolve) => setTimeout(resolve, 5)); + } + }, + async sendStop(message) { await harnessReady; await sendSessionInput(sessionId, { kind: "stop", message }); diff --git a/packages/trigger-sdk/test/accumulator-steering-compaction.test.ts b/packages/trigger-sdk/test/accumulator-steering-compaction.test.ts new file mode 100644 index 00000000000..083f0d72779 --- /dev/null +++ b/packages/trigger-sdk/test/accumulator-steering-compaction.test.ts @@ -0,0 +1,60 @@ +import type { UIMessage } from "ai"; +import { describe, expect, it } from "vitest"; +import { chat } from "../src/v3/ai.js"; + +/** + * `chat.MessageAccumulator` compaction is model-only: it replaces + * `modelMessages` with a summary and leaves `uiMessages` whole so the chat can + * still display the conversation. Recording a steer by reconverting + * `modelMessages` from `uiMessages` therefore restores everything the summary + * replaced, on the next steer after any compaction. + * + * Asserted directly on the accumulator: no run, no model, no harness, because + * the whole question is which of its two lanes gets written and how. + */ + +const USAGE = { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 5, reasoning: undefined }, + totalTokens: 15, +}; + +const userMessage = (text: string, id: string): UIMessage => + ({ id, role: "user", parts: [{ type: "text", text }] }) as UIMessage; + +const assistantMessage = (text: string, id: string): UIMessage => + ({ id, role: "assistant", parts: [{ type: "text", text }] }) as UIMessage; + +const flatten = (messages: { content: unknown }[]) => JSON.stringify(messages); + +describe("chat.MessageAccumulator steering after compaction", () => { + it("keeps the summary in the model lane when a later steer is absorbed", async () => { + const conversation = new chat.MessageAccumulator({ + compaction: { + shouldCompact: () => true, + summarize: async () => "SUMMARY-OF-EVERYTHING", + }, + }); + + await conversation.addIncoming([userMessage("EARLY-SENTINEL", "u-1")], "submit-message", 0); + await conversation.addResponse(assistantMessage("first answer", "a-1")); + + const didCompact = await conversation.compactIfNeeded(USAGE as never); + expect(didCompact).toBe(true); + + // Compaction is model-only, so the two lanes deliberately disagree here. + expect(flatten(conversation.modelMessages)).toContain("SUMMARY-OF-EVERYTHING"); + expect(flatten(conversation.modelMessages)).not.toContain("EARLY-SENTINEL"); + expect(JSON.stringify(conversation.uiMessages)).toContain("EARLY-SENTINEL"); + + await conversation.absorbSteering([userMessage("steer-me", "u-2")]); + + // The steer has to land in both lanes. + expect(JSON.stringify(conversation.uiMessages)).toContain("steer-me"); + expect(flatten(conversation.modelMessages)).toContain("steer-me"); + // And the compaction has to survive it. Reconverting from the UI lane + // brings the compacted message back and drops the summary. + expect(flatten(conversation.modelMessages)).toContain("SUMMARY-OF-EVERYTHING"); + expect(flatten(conversation.modelMessages)).not.toContain("EARLY-SENTINEL"); + }); +}); diff --git a/packages/trigger-sdk/test/action-failure-not-a-turn.test.ts b/packages/trigger-sdk/test/action-failure-not-a-turn.test.ts new file mode 100644 index 00000000000..54a907b4bfd --- /dev/null +++ b/packages/trigger-sdk/test/action-failure-not-a-turn.test.ts @@ -0,0 +1,117 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * An action whose stream fails is still an action, not a turn. + * + * Reporting the failure by throwing lands in the shared turn-error path, + * which fires `onTurnComplete`, advances the turn counter and consumes the + * one-shot instruction lane, none of which an action is supposed to do. The + * failure still has to be reported to the client and the partial kept. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; +const userMessage = (text: string, id: string) => ({ + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], +}); +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} +const textChunks = (text: string): 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: USAGE }, +]; + +describe("an action whose stream fails", () => { + it("is reported without being counted as a turn", { timeout: 30_000 }, async () => { + const turnCompletes: { turn: number; finishReason?: string }[] = []; + const turnPrompts: string[] = []; + + const turnModel = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + turnPrompts.push(JSON.stringify(prompt)); + return { + stream: simulateReadableStream({ chunks: textChunks("answer"), initialDelayInMs: 5 }), + }; + }, + }); + const failingActionModel = new MockLanguageModelV3({ + doStream: async () => ({ + stream: new ReadableStream({ + pull(c) { + c.error(new Error("provider exploded mid-stream")); + }, + }), + }), + }); + + const agent = chat.agent({ + id: "action-failure-not-a-turn", + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("regenerate") })]), + onTurnComplete: async ({ turn, finishReason }) => { + turnCompletes.push({ turn, finishReason }); + // Injected after turn 0, meant for the next real turn. + if (turn === 0) + chat.inject([{ role: "system", content: "INSTRUCTION-FOR-NEXT-TURN" }] as never); + }, + onAction: async ({ action, messages }) => { + if (action.type !== "regenerate") return; + chat.history.slice(0, -1); + return streamText({ model: failingActionModel, messages, ...chat.toStreamTextOptions() }); + }, + run: async ({ messages, signal }) => + streamText({ + model: turnModel, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + }), + }); + + const harness = mockChatAgent(agent, { chatId: "action-failure-not-a-turn" }); + try { + await harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => turnCompletes.length >= 1, "turn 0"); + + await harness.sendAction({ type: "regenerate" }).catch(() => {}); + await new Promise((r) => setTimeout(r, 200)); + + // The failure reached the client. + const errors = (harness.allRawChunks as { type?: string }[]).filter( + (c) => c.type === "error" + ); + expect(errors.length).toBeGreaterThan(0); + + // But it was not a turn: no turn lifecycle for it. + expect(turnCompletes).toHaveLength(1); + + await harness.sendMessage(userMessage("m2", "u-2")); + await waitFor(() => turnCompletes.length >= 2, "turn 1"); + + // The next real turn is turn 1, not turn 2, and it still gets the + // instruction the failed action must not have consumed. + expect(turnCompletes[1]!.turn).toBe(1); + expect(turnPrompts.at(-1)!).toContain("INSTRUCTION-FOR-NEXT-TURN"); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/action-plain-replies.test.ts b/packages/trigger-sdk/test/action-plain-replies.test.ts new file mode 100644 index 00000000000..5a98f0c82f2 --- /dev/null +++ b/packages/trigger-sdk/test/action-plain-replies.test.ts @@ -0,0 +1,121 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, streamText } from "ai"; +import type { UIMessage } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * `onAction` is documented to accept a `string` or a `UIMessage` as a reply, + * not only a stream. Each has to reach the browser, the conversation the next + * turn is built from, and the snapshot, the same as a streamed reply does. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; +const userMessage = (text: string, id: string) => ({ + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], +}); +const textChunks = (text: string): 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: USAGE }, +]; +const textOf = (m: { parts?: unknown[] }) => + ((m.parts ?? []) as { type: string; text?: string }[]) + .filter((p) => p.type === "text") + .map((p) => p.text ?? "") + .join(""); + +async function runAction(chatId: string, reply: () => unknown, opts?: { compact?: boolean }) { + const prompts: string[] = []; + let compacted = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push(JSON.stringify(prompt)); + return { + stream: simulateReadableStream({ chunks: textChunks("first answer"), initialDelayInMs: 5 }), + }; + }, + }); + const agent = chat.agent({ + id: chatId, + ...(opts?.compact + ? { + compaction: { + shouldCompact: ({ source }) => source === "outer" && compacted === 0, + summarize: async () => { + compacted++; + return "SUMMARY-OF-EVERYTHING"; + }, + }, + } + : {}), + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("note") })]), + onAction: async ({ action }) => (action.type === "note" ? reply() : undefined), + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { chatId }); + try { + await harness.sendMessage(userMessage("m1", "u-1")); + if (opts?.compact) { + const start = Date.now(); + while (compacted === 0 && Date.now() - start < 5000) + await new Promise((r) => setTimeout(r, 10)); + } + await harness.sendAction({ type: "note" }); + await new Promise((r) => setTimeout(r, 80)); + const streamed = (harness.allRawChunks as { type?: string; delta?: string }[]) + .filter((c) => c.type === "text-delta") + .map((c) => c.delta ?? "") + .join(""); + const snapshot = (harness.getSnapshot()?.messages ?? []).map(textOf); + const promptsBefore = prompts.length; + await harness.sendMessage(userMessage("m2", "u-2")); + return { streamed, snapshot, nextPrompt: prompts[promptsBefore]! }; + } finally { + await harness.close(); + } +} + +describe("a plain reply from onAction", () => { + it("delivers a returned string like a streamed reply", { timeout: 30_000 }, async () => { + const r = await runAction("action-string-reply", () => "NOTE-FROM-ACTION"); + expect(r.streamed).toContain("NOTE-FROM-ACTION"); + expect(r.snapshot.at(-1)).toBe("NOTE-FROM-ACTION"); + expect(r.nextPrompt).toContain("NOTE-FROM-ACTION"); + }); + + it("delivers a returned UIMessage like a streamed reply", { timeout: 30_000 }, async () => { + const message = { + id: "a-note", + role: "assistant", + parts: [{ type: "text", text: "UIMESSAGE-FROM-ACTION" }], + } as UIMessage; + const r = await runAction("action-uimessage-reply", () => message); + expect(r.streamed).toContain("UIMESSAGE-FROM-ACTION"); + expect(r.snapshot.at(-1)).toBe("UIMESSAGE-FROM-ACTION"); + expect(r.nextPrompt).toContain("UIMESSAGE-FROM-ACTION"); + }); + + it("is appended to a compacted lane rather than rebuilding it", { timeout: 30_000 }, async () => { + /** + * Compaction is model-only. Committing the reply by reconverting the UI + * lane would put the message compaction removed back in front of the model. + */ + const r = await runAction("action-reply-after-compaction", () => "NOTE-AFTER-COMPACTION", { + compact: true, + }); + expect(r.nextPrompt).toContain("SUMMARY-OF-EVERYTHING"); + expect(r.nextPrompt).toContain("NOTE-AFTER-COMPACTION"); + expect(r.nextPrompt).not.toContain("first answer"); + }); +}); diff --git a/packages/trigger-sdk/test/action-snapshot-cursor.test.ts b/packages/trigger-sdk/test/action-snapshot-cursor.test.ts new file mode 100644 index 00000000000..8f07db20f68 --- /dev/null +++ b/packages/trigger-sdk/test/action-snapshot-cursor.test.ts @@ -0,0 +1,80 @@ +// Import the test harness FIRST — installs the resource catalog so +// `chat.agent()` below registers its task functions correctly. +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { describe, expect, it } from "vitest"; +import { chat } from "../src/v3/ai.js"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { z } from "zod"; + +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: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, + }, + }, + ], + }); +} + +describe("the snapshot an action writes", () => { + it("keeps the resume cursor the last turn established", async () => { + const agent = chat.agent({ + id: "action-snapshot-cursor", + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("undo") })]), + onAction: async ({ action }) => { + if (action.type === "undo") chat.history.slice(0, -2); + }, + run: async ({ messages, signal }) => + streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("answer") }), + }), + messages, + abortSignal: signal, + }), + }); + + const harness = mockChatAgent(agent, { chatId: "action-snapshot-cursor" }); + + try { + await harness.sendMessage({ + id: "u1", + role: "user", + parts: [{ type: "text", text: "first" }], + }); + await new Promise((r) => setTimeout(r, 30)); + + const afterTurn = harness.getSnapshot(); + expect(afterTurn?.lastOutEventId).toBeDefined(); + + await harness.sendAction({ type: "undo" }); + await new Promise((r) => setTimeout(r, 30)); + + const afterAction = harness.getSnapshot(); + + /** + * An action has no turn cursor of its own. Writing the snapshot with + * `lastOutEventId: undefined` would drop the resume point the last turn + * established, and the next boot would replay from further back to rebuild + * what it could have read — so an action's write has to be cursor-neutral. + */ + expect(afterAction?.lastOutEventId).toBe(afterTurn?.lastOutEventId); + + // And the mutation itself landed, which is the point of writing at all. + expect(afterAction?.messages ?? []).toEqual([]); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/action-snapshot.test.ts b/packages/trigger-sdk/test/action-snapshot.test.ts new file mode 100644 index 00000000000..39e2ecfe280 --- /dev/null +++ b/packages/trigger-sdk/test/action-snapshot.test.ts @@ -0,0 +1,82 @@ +// Import the test harness FIRST — 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 } from "vitest"; +import { chat } from "../src/v3/ai.js"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { z } from "zod"; + +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: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, + }, + }, + ], + }); +} + +function agentWithUndo(id: string) { + return chat.agent({ + id, + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("undo") })]), + onAction: async ({ action }) => { + if (action.type === "undo") { + // The documented way to roll history back — see /ai-chat/actions. + chat.history.slice(0, -2); + } + }, + run: async ({ messages, signal }) => + streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("answer") }), + }), + messages, + abortSignal: signal, + }), + }); +} + +describe("snapshot durability of history mutated by an action", () => { + it("persists an undo, so a continuation does not resurrect the undone turn", async () => { + const harness = mockChatAgent(agentWithUndo("action-snapshot-undo"), { + chatId: "action-snapshot-undo", + }); + + try { + await harness.sendMessage({ + id: "u1", + role: "user", + parts: [{ type: "text", text: "first" }], + }); + await new Promise((r) => setTimeout(r, 30)); + + // After a turn the snapshot holds the exchange. + expect(harness.getSnapshot()?.messages.map((m) => m.role)).toEqual(["user", "assistant"]); + + await harness.sendAction({ type: "undo" }); + await new Promise((r) => setTimeout(r, 30)); + + /** + * An action is not a turn, so it never reaches the turn-complete path where + * the snapshot is written. The rollback lives in the accumulator only, and + * the next continuation boots from a snapshot that still holds the undone + * exchange — the user's undo silently reverts, minutes later, with no error. + */ + expect(harness.getSnapshot()?.messages ?? []).toEqual([]); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/action-stream-accumulator.test.ts b/packages/trigger-sdk/test/action-stream-accumulator.test.ts new file mode 100644 index 00000000000..f3a6234883f --- /dev/null +++ b/packages/trigger-sdk/test/action-stream-accumulator.test.ts @@ -0,0 +1,172 @@ +// Import the test harness FIRST — installs the resource catalog so +// `chat.agent()` below registers its task functions correctly. +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { describe, expect, it } from "vitest"; +import { chat } from "../src/v3/ai.js"; +import { simulateReadableStream, streamText } from "ai"; +import type { UIMessage } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { z } from "zod"; + +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: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, + }, + }, + ], + }); +} + +function textOf(message: UIMessage): string { + return message.parts.map((part) => (part.type === "text" ? part.text : "")).join(""); +} + +describe("a StreamTextResult returned from onAction", () => { + it("becomes part of the conversation, not just something the browser saw", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("regenerated answer") }), + }); + + const agent = chat.agent({ + id: "action-stream-accumulator", + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("regenerate") })]), + + /** + * The bare shape the docs show: return the stream and let the runtime pipe + * it. The alternative — consuming it with `chat.pipeAndCapture` — is the + * workaround, so testing that instead would prove nothing about this path. + */ + onAction: async ({ action, messages }) => { + if (action.type !== "regenerate") return; + chat.history.slice(0, -1); + return streamText({ model, messages }); + }, + + run: async ({ messages, signal }) => + streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("first answer") }), + }), + messages, + abortSignal: signal, + }), + }); + + const harness = mockChatAgent(agent, { chatId: "action-stream-accumulator" }); + + try { + await harness.sendMessage({ + id: "u1", + role: "user", + parts: [{ type: "text", text: "ask" }], + }); + await new Promise((r) => setTimeout(r, 30)); + + const turn = await harness.sendAction({ type: "regenerate" }); + await new Promise((r) => setTimeout(r, 50)); + + // The browser did see it — that part was never broken. + const streamed = turn.chunks + .filter((c) => c.type === "text-delta") + .map((c) => (c as { delta: string }).delta) + .join(""); + expect(streamed).toBe("regenerated answer"); + + /** + * And the conversation agrees with the screen. Before the fix the response + * was piped and dropped: absent from the accumulator, absent from the + * snapshot, so the next turn's model context contained the question and the + * *old* answer that regenerate had just removed. + */ + const snapshot = harness.getSnapshot(); + expect(snapshot?.messages.map(textOf)).toEqual(["ask", "regenerated answer"]); + } finally { + await harness.close(); + } + }); + + it("reports a mid-stream failure instead of committing a truncated answer as finished", async () => { + /** + * `pipeChatAndCapture` returns a stream failure as `status: "error"` rather + * than throwing it. Unchecked, the action commits whatever streamed, writes a + * normal turn-complete, and the browser just sees the stream stop — so the + * user reads a half-finished answer presented as complete and the next turn + * builds on it. The partial is still kept, as on the turn path; what changes + * is that the failure is surfaced alongside it. + */ + let stage = 0; + const failsMidStream = new MockLanguageModelV3({ + doStream: async () => ({ + stream: new ReadableStream({ + async pull(controller) { + await new Promise((r) => setTimeout(r, 25)); + if (stage === 0) { + controller.enqueue({ type: "text-start", id: "t1" }); + stage++; + return; + } + if (stage === 1) { + controller.enqueue({ type: "text-delta", id: "t1", delta: "half an answer" }); + stage++; + return; + } + controller.error(new Error("provider exploded mid-stream")); + }, + }), + }), + }); + + const agent = chat.agent({ + id: "action-stream-error", + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("regenerate") })]), + onAction: async ({ action, messages }) => { + if (action.type !== "regenerate") return; + chat.history.slice(0, -1); + return streamText({ model: failsMidStream, messages }); + }, + run: async ({ messages, signal }) => + streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("first answer") }), + }), + messages, + abortSignal: signal, + }), + }); + + const harness = mockChatAgent(agent, { chatId: "action-stream-error" }); + + try { + await harness.sendMessage({ + id: "u1", + role: "user", + parts: [{ type: "text", text: "ask" }], + }); + await new Promise((r) => setTimeout(r, 40)); + + await harness.sendAction({ type: "regenerate" }).catch(() => {}); + await new Promise((r) => setTimeout(r, 300)); + + const errors = (harness.allRawChunks as { type?: string }[]).filter( + (c) => c.type === "error" + ); + expect(errors.length).toBeGreaterThan(0); + + // The partial is still kept rather than discarded. + expect(harness.getSnapshot()?.messages.map(textOf).at(-1)).toContain("half an answer"); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/chatHandover.test.ts b/packages/trigger-sdk/test/chatHandover.test.ts index a101b91494f..65dd802d26f 100644 --- a/packages/trigger-sdk/test/chatHandover.test.ts +++ b/packages/trigger-sdk/test/chatHandover.test.ts @@ -632,4 +632,63 @@ describe("chat.handover", () => { await harness.close(); } }); + + it("seeds the accumulator from headStartMessages without hydrateMessages", async () => { + // The hydrate variant above gets the head-start user message through + // `incomingMessages`. Without `hydrateMessages` it arrives only via the + // boot-time seed from `payload.headStartMessages`, so this is the path + // that keeps an app with a display-only transcript from storing an + // answer with no question above it. + // + // Note the shape a persisting app has to handle: by `onTurnStart` the + // accumulator is already ["user", "assistant"], because the warm route's + // partial is spliced in before the hook fires. "The incoming message is + // the last one" is therefore false on this path. + let captured: { roles: string[]; texts: string[] } | undefined; + + const agent = chat.agent({ + id: "test-handover-seed-no-hydrate", + onTurnComplete: async ({ uiMessages }) => { + captured = { + roles: uiMessages.map((m) => m.role), + texts: uiMessages.map((m) => + m.parts.map((p) => (p.type === "text" ? p.text : "")).join("") + ), + }; + }, + run: async ({ messages, signal }) => + streamText({ + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("should-not-run") }), + }), + messages, + abortSignal: signal, + }), + }); + + const harness = mockChatAgent(agent, { + chatId: "test-handover-seed-no-hydrate", + mode: "handover-prepare", + headStartMessages: [ + { id: "hs-user-1", role: "user", parts: [{ type: "text", text: "say hi" }] }, + ], + }); + + try { + await harness.sendHandover({ + partialAssistantMessage: [ + { role: "assistant", content: [{ type: "text", text: "Hi there." }] }, + ], + messageId: "asst-seed-1", + isFinal: true, + }); + await new Promise((r) => setTimeout(r, 30)); + + expect(captured).toBeDefined(); + expect(captured!.roles).toEqual(["user", "assistant"]); + expect(captured!.texts[0]).toBe("say hi"); + } finally { + await harness.close(); + } + }); }); diff --git a/packages/trigger-sdk/test/createsession-steering-lanes.test.ts b/packages/trigger-sdk/test/createsession-steering-lanes.test.ts new file mode 100644 index 00000000000..6649da3c124 --- /dev/null +++ b/packages/trigger-sdk/test/createsession-steering-lanes.test.ts @@ -0,0 +1,332 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { sessionStreams } from "@trigger.dev/core/v3"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, stepCountIs, streamText, tool } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import type { UIMessage } from "ai"; +import { chat } from "../src/v3/ai.js"; + +/** + * `chat.createSession` keeps its own accumulator rather than the one + * `chat.agent` publishes to locals, so the two lanes have to be checked on + * this surface separately. + * + * The steering drain appends claimed messages to + * `locals.get(chatCurrentUIMessagesKey)` behind a truthiness guard, and + * `createSession` never sets that key, so the append is a silent no-op here. + * If that is what happens, a mid-turn steer reaches the model for the answer + * it steered and then disappears from both of the session's own lanes: + * `turn.uiMessages`, which is what an app persists from, and `turn.messages`, + * which is what every later turn sends to the model. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; + +function userMessage(text: string, id: string) { + return { id, role: "user" as const, parts: [{ type: "text" as const, text }] }; +} + +function deferred() { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} + +function textChunks(text: string): LanguageModelV3StreamPart[] { + return [ + { 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: USAGE }, + ]; +} + +function toolCallChunks(callId: string): LanguageModelV3StreamPart[] { + return [ + { + type: "tool-call", + toolCallId: callId, + toolName: "gate", + input: JSON.stringify({ q: "go" }), + }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE }, + ]; +} + +type SeqReader = { lastSeqNum: (chatId: string, dir: "in" | "out") => number | undefined }; + +/** Send and wait for the record to land on the channel, so the steer is claimable. */ +async function sendAndLand( + harness: { sendMessage: (m: ReturnType) => Promise }, + chatId: string, + text: string, + id: string +) { + const seqs = sessionStreams as unknown as SeqReader; + const before = seqs.lastSeqNum(chatId, "in") ?? -1; + void harness.sendMessage(userMessage(text, id)); + await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`); +} + +describe("chat.createSession steering across turns", () => { + it("keeps a mid-turn steer in both of the session's own lanes", { timeout: 30_000 }, async () => { + const chatId = "createsession-steer-lanes"; + const toolGate = deferred(); + let toolEntered = false; + + /** Per-turn snapshots of the session's own two lanes. */ + const lanes: { turn: number; ui: string[]; model: string[] }[] = []; + const prompts: string[][] = []; + let turnCount = 0; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push( + prompt + .filter((m) => m.role === "user") + .flatMap((m) => + Array.isArray(m.content) + ? (m.content as { type: string; text?: string }[]) + .filter((c) => c.type === "text") + .map((c) => c.text ?? "") + : [] + ) + ); + const isToolStep = step++ % 2 === 0; + return { + stream: simulateReadableStream({ + chunks: isToolStep ? toolCallChunks(`tc-${step}`) : textChunks("done"), + initialDelayInMs: 10, + chunkDelayInMs: 2, + }), + }; + }, + }); + + const textOf = (m: { parts?: unknown[] }) => + ((m.parts ?? []) as { type: string; text?: string }[]) + .filter((p) => p.type === "text") + .map((p) => p.text ?? "") + .join(""); + + const modelTextOf = (m: { content: unknown }) => + typeof m.content === "string" + ? m.content + : Array.isArray(m.content) + ? (m.content as { type: string; text?: string }[]) + .filter((c) => c.type === "text") + .map((c) => c.text ?? "") + .join("") + : ""; + + const agent = chat.customAgent({ + id: "createsession-steer-lanes", + run: async (payload, { signal }) => { + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + pendingMessages: { shouldInject: () => true }, + }); + + for await (const turn of session) { + const thisTurn = turnCount++; + await turn.complete( + streamText({ + model, + messages: turn.messages, + abortSignal: turn.signal, + prepareStep: turn.prepareStep(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }) + ); + lanes.push({ + turn: thisTurn, + ui: turn.uiMessages.map(textOf), + model: turn.messages.map(modelTextOf), + }); + } + }, + }); + + const harness = mockChatAgent(agent, { chatId }); + + try { + const first = harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => toolEntered, "tool entered"); + await sendAndLand(harness, chatId, "steer-me", "u-2"); + toolGate.resolve(); + await first; + + await waitFor(() => lanes.length >= 1, "turn 1 recorded"); + const promptsAfterTurn1 = prompts.length; + + await harness.sendMessage(userMessage("m3", "u-3")); + await waitFor(() => prompts.length > promptsAfterTurn1, "turn 2 prompt built"); + await waitFor(() => lanes.length >= 2, "turn 2 recorded"); + + // The lane an app persists from. + expect(lanes[0]!.ui).toContain("steer-me"); + // The lane every later turn sends to the model. + expect(lanes[1]!.model).toContain("steer-me"); + // And what the model was actually asked on the later turn. + expect(prompts[promptsAfterTurn1]!).toContain("steer-me"); + } finally { + toolGate.resolve(); + await harness.close(); + } + }); +}); + +/** + * The same lane check for a fully manual loop built on + * `chat.MessageAccumulator`. + * + * This is the other accumulator-based drain site, and it files the claimed + * messages through `this` rather than through a captured `accumulator`, so a + * binding mistake there would not show up in the `createSession` test above. + */ +describe("chat.MessageAccumulator steering", () => { + it("records a steer the drain consumed in both of its lanes", { timeout: 30_000 }, async () => { + let toolEntered = false; + const toolGate = deferred(); + const lanes: { ui: string[]; model: string[] }[] = []; + const prompts: string[][] = []; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push( + prompt + .filter((m) => m.role === "user") + .flatMap((m) => + Array.isArray(m.content) + ? (m.content as { type: string; text?: string }[]) + .filter((c) => c.type === "text") + .map((c) => c.text ?? "") + : [] + ) + ); + const isToolStep = step++ % 2 === 0; + return { + stream: simulateReadableStream({ + chunks: isToolStep ? toolCallChunks(`tc-${step}`) : textChunks("done"), + initialDelayInMs: 10, + chunkDelayInMs: 2, + }), + }; + }, + }); + + const textOf = (m: { parts?: unknown[] }) => + ((m.parts ?? []) as { type: string; text?: string }[]) + .filter((p) => p.type === "text") + .map((p) => p.text ?? "") + .join(""); + + const modelTextOf = (m: { content: unknown }) => + typeof m.content === "string" + ? m.content + : Array.isArray(m.content) + ? (m.content as { type: string; text?: string }[]) + .filter((c) => c.type === "text") + .map((c) => c.text ?? "") + .join("") + : ""; + + const agent = chat.customAgent({ + id: "accumulator-steer-lanes", + run: async () => { + const conversation = new chat.MessageAccumulator({ + pendingMessages: { shouldInject: () => true }, + }); + const next = await chat.messages.waitWithIdleTimeout({ + idleTimeoutInSeconds: 60, + timeout: "1h", + }); + if (!next.ok) return; + const wire = next.output as { message?: UIMessage; trigger: string }; + const messages = await conversation.addIncoming( + wire.message ? [wire.message] : [], + wire.trigger, + 0 + ); + + const result = streamText({ + model, + messages, + prepareStep: conversation.prepareStep(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }); + + // Steer while the tool holds the turn open, so the drain has a step + // boundary to consume it at. + void (async () => { + await waitFor(() => toolEntered, "tool entered"); + await conversation.steerAsync(userMessage("steer-me", "u-2")); + toolGate.resolve(); + })(); + + const captured = await chat.pipeAndCapture(result); + if (captured.message) await conversation.addResponse(captured.message); + lanes.push({ + ui: conversation.uiMessages.map(textOf), + model: conversation.modelMessages.map(modelTextOf), + }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "accumulator-steer-lanes" }); + try { + await harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => lanes.length >= 1, "turn recorded"); + + // The drain put it in the prompt, which is what makes the lane checks meaningful. + expect(prompts.some((p) => p.includes("steer-me"))).toBe(true); + expect(lanes[0]!.ui).toContain("steer-me"); + expect(lanes[0]!.model).toContain("steer-me"); + } finally { + toolGate.resolve(); + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/error-snapshot-cursor.test.ts b/packages/trigger-sdk/test/error-snapshot-cursor.test.ts new file mode 100644 index 00000000000..7c2e8390ec3 --- /dev/null +++ b/packages/trigger-sdk/test/error-snapshot-cursor.test.ts @@ -0,0 +1,108 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * The snapshot cursor after a failed turn. + * + * The error path writes its snapshot with the failed turn's completion cursor + * but does not update the shared cursor holder, so a later action's snapshot, + * which is cursor-neutral and reuses the holder, writes the cursor from + * BEFORE the failed turn. A continuation then resumes from there and replays + * output the failed turn's snapshot had already superseded. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; +const userMessage = (text: string, id: string) => ({ + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], +}); +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} +const textChunks = (text: string): 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: USAGE }, +]; +function erroringStream(): ReadableStream { + const chunks: LanguageModelV3StreamPart[] = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "partial" }, + ]; + let i = 0; + return new ReadableStream({ + pull(c) { + if (i < chunks.length) return void c.enqueue(chunks[i++]!); + c.error(new Error("UND_ERR_BODY_TIMEOUT")); + }, + }); +} + +describe("the snapshot an action writes after a failed turn", () => { + it("carries the failed turn's cursor, not the one before it", { timeout: 30_000 }, async () => { + const completes: { finishReason?: string }[] = []; + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async () => + step++ === 0 + ? { stream: simulateReadableStream({ chunks: textChunks("first"), initialDelayInMs: 5 }) } + : { stream: erroringStream() }, + }); + + const agent = chat.agent({ + id: "error-snapshot-cursor", + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("undo") })]), + onTurnComplete: async ({ finishReason }) => { + completes.push({ finishReason }); + }, + onAction: async ({ action }) => { + if (action.type === "undo") chat.history.slice(0, -2); + }, + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + + const harness = mockChatAgent(agent, { chatId: "error-snapshot-cursor" }); + try { + await harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => harness.getSnapshot()?.lastOutEventId !== undefined, "turn 0 snapshot"); + const afterTurn0 = harness.getSnapshot()?.lastOutEventId; + + await harness.sendMessage(userMessage("m2", "u-2")); + await waitFor(() => completes.length >= 2, "turn 1 (failed)"); + expect(completes[1]!.finishReason).toBe("error"); + await waitFor( + () => harness.getSnapshot()?.lastOutEventId !== afterTurn0, + "failed turn snapshot" + ); + const afterFailedTurn = harness.getSnapshot()?.lastOutEventId; + expect(afterFailedTurn).toBeDefined(); + // The failed turn moved the cursor: it wrote an error and a completion. + expect(afterFailedTurn).not.toBe(afterTurn0); + + await harness.sendAction({ type: "undo" }); + await new Promise((r) => setTimeout(r, 60)); + + // An action's write is cursor-neutral, so it has to keep the CURRENT + // cursor, which is the failed turn's, not the one from before it. + expect(harness.getSnapshot()?.lastOutEventId).toBe(afterFailedTurn); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/hitl-uncompacts.test.ts b/packages/trigger-sdk/test/hitl-uncompacts.test.ts new file mode 100644 index 00000000000..62130beba23 --- /dev/null +++ b/packages/trigger-sdk/test/hitl-uncompacts.test.ts @@ -0,0 +1,158 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, streamText, tool } from "ai"; +import type { UIMessage } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * A tool-approval continuation after compaction. + * + * Compaction is model-only: the model lane becomes a summary while the UI + * lane keeps everything. A tool-approval response arrives as an update to the + * existing assistant message, and that path rebuilds the model lane from the + * UI lane. The summary is replaced by the full transcript, and the message + * compaction had removed is sent to the model again. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + totalTokens: 2, +}; +const userMessage = (text: string, id: string) => ({ + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], +}); +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} +const textChunks = (text: string): 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: USAGE }, +]; +const approvalToolCall = (callId: string): LanguageModelV3StreamPart[] => [ + { + type: "tool-call", + toolCallId: callId, + toolName: "risky", + input: JSON.stringify({ what: "x" }), + }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE }, +]; + +describe("a tool-approval turn after compaction", () => { + it("keeps the summary in the model lane", { timeout: 30_000 }, async () => { + const prompts: string[] = []; + const turns: UIMessage[][] = []; + let compacted = 0; + + const risky = tool({ + description: "needs a human to approve", + inputSchema: z.object({ what: z.string() }), + needsApproval: true, + execute: async () => "done", + }); + + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push(JSON.stringify(prompt)); + const n = step++; + // turn 0 answers; turn 1 asks for approval; the continuation answers. + const chunks = n === 1 ? approvalToolCall("tc-1") : textChunks(`answer-${n}`); + return { stream: simulateReadableStream({ chunks, initialDelayInMs: 5 }) }; + }, + }); + + const agent = chat.agent({ + id: "hitl-uncompacts", + compaction: { + // Compact once, between turns 0 and 1. + shouldCompact: ({ source }) => source === "outer" && compacted === 0, + summarize: async () => { + compacted++; + return "SUMMARY-OF-EVERYTHING"; + }, + }, + onTurnComplete: async ({ uiMessages }) => { + turns.push(uiMessages.map((m) => structuredClone(m))); + }, + run: async ({ messages, signal }) => + streamText({ + model, + messages, + abortSignal: signal, + tools: { risky }, + ...chat.toStreamTextOptions(), + }), + }); + + const harness = mockChatAgent(agent, { chatId: "hitl-uncompacts" }); + try { + await harness.sendMessage(userMessage("EARLY-SENTINEL", "u-1")); + await waitFor(() => turns.length >= 1 && compacted > 0, "turn 0 + compaction"); + + await harness.sendMessage(userMessage("please do the risky thing", "u-2")); + await waitFor(() => turns.length >= 2, "turn 1 (approval requested)"); + + // The summary is in force going into the approval turn. + expect(prompts.at(-1)!).toContain("SUMMARY-OF-EVERYTHING"); + expect(prompts.at(-1)!).not.toContain("EARLY-SENTINEL"); + + // Approve, as the browser would: a slim update to the existing assistant. + const head = turns.at(-1)!.at(-1)!; + const part = ( + head.parts as { + type: string; + toolCallId?: string; + state?: string; + approval?: { id: string }; + }[] + ).find((p) => p.type === "tool-risky"); + expect(part?.state).toBe("approval-requested"); + // sendMessage resolves at turn-complete, so the continuation's prompt is + // recorded by the time it returns; capture the index first. + const promptsBefore = prompts.length; + await harness.sendMessage({ + id: head.id, + role: "assistant", + parts: [ + { + type: "tool-risky", + toolCallId: part!.toolCallId!, + state: "approval-responded", + approval: { id: part!.approval!.id, approved: true }, + }, + ], + } as unknown as UIMessage); + // The continuation has to run against the compacted lane, not the + // whole transcript that compaction had already replaced. + const cont = prompts[promptsBefore]!; + expect(cont).toContain("SUMMARY-OF-EVERYTHING"); + expect(cont).not.toContain("EARLY-SENTINEL"); + + // And the turn after it: the continuation's own response is committed by + // replacing the approval-requested assistant, and that path must not + // reconvert the lane either. + const promptsBeforeNext = prompts.length; + await harness.sendMessage(userMessage("and then?", "u-3")); + const next = prompts[promptsBeforeNext]!; + expect(next).toContain("SUMMARY-OF-EVERYTHING"); + expect(next).not.toContain("EARLY-SENTINEL"); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/inject-system-instructions.test.ts b/packages/trigger-sdk/test/inject-system-instructions.test.ts new file mode 100644 index 00000000000..2126382effe --- /dev/null +++ b/packages/trigger-sdk/test/inject-system-instructions.test.ts @@ -0,0 +1,346 @@ +// Import the test harness FIRST — installs the resource catalog so +// `chat.agent()` below registers its task functions correctly. +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { describe, expect, it } from "vitest"; +import { chat } from "../src/v3/ai.js"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; + +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: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, + }, + }, + ], + }); +} + +describe("chat.inject with a system role", () => { + it("goes to the instructions lane instead of poisoning the prompt", async () => { + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + let injected = false; + + const agent = chat.agent({ + id: "inject-system-instructions", + onBoot: async () => { + chat.prompt.set({ + promptId: "base", + version: 1, + labels: ["local"], + text: "You are a helpful assistant.", + model: undefined, + config: undefined, + toAISDKTelemetry: () => ({ experimental_telemetry: { isEnabled: true, metadata: {} } }), + }); + }, + onTurnComplete: async () => { + if (injected) return; + injected = true; + /** + * The shape every docs example uses. On ai@7 a system message inside + * `messages` is rejected by `standardizePrompt` for every provider, so + * this used to kill the next turn — an error chunk reading "An error + * occurred." and an assistant message with no parts. + */ + chat.inject([{ role: "system", content: "The user just upgraded to Pro." }]); + }, + run: async ({ messages, signal }) => + streamText({ + ...chat.toStreamTextOptions(), + model, + messages, + abortSignal: signal, + }), + }); + + const harness = mockChatAgent(agent, { chatId: "inject-system-instructions" }); + + try { + await harness.sendMessage({ id: "u1", role: "user", parts: [{ type: "text", text: "one" }] }); + await new Promise((r) => setTimeout(r, 40)); + + const turn = await harness.sendMessage({ + id: "u2", + role: "user", + parts: [{ type: "text", text: "two" }], + }); + await new Promise((r) => setTimeout(r, 40)); + + // The turn survives. + const errors = turn.rawChunks.filter((c) => (c as { type?: string })?.type === "error"); + expect(errors).toEqual([]); + + // The injected context arrives as a system block, alongside the base prompt, + // and never as a system message inside `messages`. + const prompt = model.doStreamCalls.at(-1)!.prompt; + const systemBlocks = prompt.filter((m) => m.role === "system"); + const asText = JSON.stringify(systemBlocks); + + expect(asText).toContain("You are a helpful assistant."); + expect(asText).toContain("The user just upgraded to Pro."); + + const nonSystem = prompt.filter((m) => m.role !== "system"); + expect(JSON.stringify(nonSystem)).not.toContain("upgraded to Pro"); + } finally { + await harness.close(); + } + }); + it("emits one system value whether or not the base block is cached", async () => { + /** + * Never an array. ai@6+ accepts `Array` and would let a + * cached base block keep its cache entry, but ai@5 rejects an array outright + * ("Invalid prompt: system must be a string") while accepting a single + * structured block — and the peer range still spans v5. So a plain base + * concatenates into a string, and a cached base absorbs the injection into its + * own content, keeping its provider options. + */ + const shapes: unknown[] = []; + + function agentFor(id: string, cacheControl: boolean) { + let injected = false; + return chat.agent({ + id, + onBoot: async () => { + chat.prompt.set({ + promptId: "base", + version: 1, + labels: ["local"], + text: "Base instructions.", + model: undefined, + config: undefined, + toAISDKTelemetry: () => ({ + experimental_telemetry: { isEnabled: true, metadata: {} }, + }), + }); + }, + onTurnComplete: async () => { + if (injected) return; + injected = true; + chat.inject([{ role: "system", content: "Amendment." }]); + }, + run: async ({ messages, signal }) => { + const options = cacheControl + ? chat.toStreamTextOptions({ cacheControl: { type: "ephemeral" } }) + : chat.toStreamTextOptions(); + shapes.push(Array.isArray(options.system) ? "array" : typeof options.system); + return streamText({ + ...options, + model: new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }), + messages, + abortSignal: signal, + }); + }, + }); + } + + for (const [id, cacheControl] of [ + ["shape-plain", false], + ["shape-cached", true], + ] as const) { + const harness = mockChatAgent(agentFor(id, cacheControl), { chatId: id }); + try { + await harness.sendMessage({ + id: "u1", + role: "user", + parts: [{ type: "text", text: "one" }], + }); + await new Promise((r) => setTimeout(r, 40)); + await harness.sendMessage({ + id: "u2", + role: "user", + parts: [{ type: "text", text: "two" }], + }); + await new Promise((r) => setTimeout(r, 40)); + } finally { + await harness.close(); + } + } + + // [plain turn 1, plain turn 2 (injected), cached turn 1, cached turn 2 (injected)] + expect(shapes).toEqual(["string", "string", "object", "object"]); + }); + + it("applies an injection to the next turn only, not to every later turn", async () => { + /** + * `chat.inject()` is a queue consumed at the next injection opportunity, so + * the instructions lane has to drain like the conversational one does. Left + * undrained, every later turn in the run repeats every earlier injection — + * the prompt grows without bound and its cached prefix changes each turn. + */ + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + let injectedOnce = false; + + const agent = chat.agent({ + id: "inject-system-drains", + onTurnComplete: async () => { + if (injectedOnce) return; + injectedOnce = true; + chat.inject([{ role: "system", content: "SENTINEL-ONE-SHOT" }]); + }, + run: async ({ messages, signal }) => + streamText({ + ...chat.toStreamTextOptions(), + model, + messages, + abortSignal: signal, + }), + }); + + const harness = mockChatAgent(agent, { chatId: "inject-system-drains" }); + + try { + await harness.sendMessage({ id: "u1", role: "user", parts: [{ type: "text", text: "one" }] }); + await new Promise((r) => setTimeout(r, 40)); + + await harness.sendMessage({ id: "u2", role: "user", parts: [{ type: "text", text: "two" }] }); + await new Promise((r) => setTimeout(r, 40)); + + await harness.sendMessage({ + id: "u3", + role: "user", + parts: [{ type: "text", text: "three" }], + }); + await new Promise((r) => setTimeout(r, 40)); + + const systemOf = (i: number) => + JSON.stringify(model.doStreamCalls[i]!.prompt.filter((m) => m.role === "system")); + + // Turn 1 injected nothing yet, turn 2 carries it, turn 3 must not repeat it. + expect(systemOf(0)).not.toContain("SENTINEL-ONE-SHOT"); + expect(systemOf(1)).toContain("SENTINEL-ONE-SHOT"); + expect(systemOf(2)).not.toContain("SENTINEL-ONE-SHOT"); + } finally { + await harness.close(); + } + }); + + it("carries the injection into every options build in the turn, not only the first", async () => { + /** + * A `run()` that builds options twice, a classifier pass and then the + * answer, has to see the injection in both. Consuming on read hands it to + * whichever call ran first and drops it from the rest, silently. + */ + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + const seen: { first: boolean; second: boolean }[] = []; + let injectedOnce = false; + + const agent = chat.agent({ + id: "inject-system-two-builds", + onTurnComplete: async () => { + if (injectedOnce) return; + injectedOnce = true; + chat.inject([{ role: "system", content: "SENTINEL-BOTH-BUILDS" }]); + }, + run: async ({ messages, signal }) => { + const first = chat.toStreamTextOptions(); + const second = chat.toStreamTextOptions(); + const has = (o: { system?: unknown }) => + JSON.stringify(o.system ?? null).includes("SENTINEL-BOTH-BUILDS"); + seen.push({ first: has(first), second: has(second) }); + return streamText({ ...second, model, messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "inject-system-two-builds" }); + + try { + await harness.sendMessage({ id: "u1", role: "user", parts: [{ type: "text", text: "one" }] }); + await new Promise((r) => setTimeout(r, 40)); + + await harness.sendMessage({ id: "u2", role: "user", parts: [{ type: "text", text: "two" }] }); + await new Promise((r) => setTimeout(r, 40)); + + await harness.sendMessage({ + id: "u3", + role: "user", + parts: [{ type: "text", text: "three" }], + }); + await new Promise((r) => setTimeout(r, 40)); + + // Turn 1 predates the injection, turn 2 carries it in both builds, turn 3 is clear again. + expect(seen).toEqual([ + { first: false, second: false }, + { first: true, second: true }, + { first: false, second: false }, + ]); + } finally { + await harness.close(); + } + }); + + it("gives each turn only its own injection, over consecutive turns", async () => { + /** + * Consuming the lane has to move the blocks out of it, not mark them read in + * place. Left in place, an injection made during the consumed turn queues + * behind them and the next turn's clear destroys both: turn 1 gets its + * instruction and every turn after it silently gets none. + */ + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: textStream("ok") }), + }); + + let n = 0; + + const agent = chat.agent({ + id: "inject-system-consecutive", + onTurnComplete: async () => { + n++; + chat.inject([{ role: "system", content: `INJECT-${n}` }]); + }, + run: async ({ messages, signal }) => + streamText({ + ...chat.toStreamTextOptions(), + model, + messages, + abortSignal: signal, + }), + }); + + const harness = mockChatAgent(agent, { chatId: "inject-system-consecutive" }); + + try { + for (const id of ["u1", "u2", "u3", "u4"]) { + await harness.sendMessage({ id, role: "user", parts: [{ type: "text", text: id }] }); + await new Promise((r) => setTimeout(r, 40)); + } + + const injectionsSeenOn = (turn: number) => { + const system = JSON.stringify( + model.doStreamCalls[turn]!.prompt.filter((m) => m.role === "system") + ); + return ["INJECT-1", "INJECT-2", "INJECT-3"].filter((key) => system.includes(key)); + }; + + // Turn 0 predates any injection; after that each turn carries exactly the + // one injected at the end of the turn before it. + expect(injectionsSeenOn(0)).toEqual([]); + expect(injectionsSeenOn(1)).toEqual(["INJECT-1"]); + expect(injectionsSeenOn(2)).toEqual(["INJECT-2"]); + expect(injectionsSeenOn(3)).toEqual(["INJECT-3"]); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/instructions-action-replay.test.ts b/packages/trigger-sdk/test/instructions-action-replay.test.ts new file mode 100644 index 00000000000..1794f687242 --- /dev/null +++ b/packages/trigger-sdk/test/instructions-action-replay.test.ts @@ -0,0 +1,128 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * A one-shot instruction and an action in between. + * + * `turn--` marks an action as not-a-turn, so an action and the message after + * it share a turn number. The consumed-instruction stash is keyed on that + * number, so an action that builds options consumes the injection and the next + * real turn reads the same stash back. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; + +function userMessage(text: string, id: string) { + return { id, role: "user" as const, parts: [{ type: "text" as const, text }] }; +} + +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} + +function textChunks(text: string): LanguageModelV3StreamPart[] { + return [ + { 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: USAGE }, + ]; +} + +describe("a one-shot instruction across an action", () => { + it( + "reaches each turn once and is not replayed by the turn after an action", + { timeout: 30_000 }, + async () => { + /** One entry per model call, in order, saying whether it carried the instruction. */ + const sawInstruction: { label: string; saw: boolean }[] = []; + + const makeModel = (label: string) => + new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + sawInstruction.push({ + label, + saw: JSON.stringify(prompt).includes("INSTRUCTION-ONE-SHOT"), + }); + return { + stream: simulateReadableStream({ chunks: textChunks("ok"), initialDelayInMs: 5 }), + }; + }, + }); + + const turnModel = makeModel("turn"); + const actionModel = makeModel("action"); + + const agent = chat.agent({ + id: "instructions-action-replay", + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("ping") })]), + onTurnComplete: async ({ turn }) => { + // Injecting from inside the run, because the lane lives in run locals. + if (turn === 0) + chat.inject([{ role: "system", content: "INSTRUCTION-ONE-SHOT" }] as never); + }, + onAction: async ({ action }) => { + if (action.type !== "ping") return; + return streamText({ + model: actionModel, + messages: [{ role: "user", content: "regenerate" }], + ...chat.toStreamTextOptions(), + }); + }, + run: async ({ messages, signal }) => + streamText({ + model: turnModel, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + }), + }); + + const harness = mockChatAgent(agent, { chatId: "instructions-action-replay" }); + try { + // Turn 1, nothing injected yet, then inject for the next turn. + await harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => sawInstruction.length >= 1, "turn 1"); + + // An action lands before the next message. + await harness.sendAction({ type: "ping" }); + await waitFor(() => sawInstruction.length >= 2, "action"); + + // Then the real turn the injection was meant for. + await harness.sendMessage(userMessage("m2", "u-2")); + await waitFor(() => sawInstruction.length >= 3, "turn 2"); + + // And one more, which must not see it again. + await harness.sendMessage(userMessage("m3", "u-3")); + await waitFor(() => sawInstruction.length >= 4, "turn 3"); + + const carriers = sawInstruction.filter((e) => e.saw).map((e) => e.label); + // The action sees it: it is pending context, and an action is not a turn, + // so the action reading it must not use it up. + expect(sawInstruction[1]!).toEqual({ label: "action", saw: true }); + // And the turn it was actually injected for still gets it. + expect(sawInstruction[2]!).toEqual({ label: "turn", saw: true }); + // The turn after that does not: one-shot means one turn. + expect(sawInstruction[3]!).toEqual({ label: "turn", saw: false }); + // And it is never carried by more than one real turn. + expect(carriers.filter((l) => l === "turn")).toHaveLength(1); + } finally { + await harness.close(); + } + } + ); +}); diff --git a/packages/trigger-sdk/test/steering-accumulator.test.ts b/packages/trigger-sdk/test/steering-accumulator.test.ts new file mode 100644 index 00000000000..41f0e734068 --- /dev/null +++ b/packages/trigger-sdk/test/steering-accumulator.test.ts @@ -0,0 +1,133 @@ +// Import the test harness FIRST — installs the resource catalog so +// `chat.agent()` below registers its task functions correctly. +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { describe, expect, it } from "vitest"; +import { chat } from "../src/v3/ai.js"; +import { simulateReadableStream, stepCountIs, streamText, tool } from "ai"; +import type { UIMessage } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { z } from "zod"; + +const usage = { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, +}; + +function userMessage(text: string, id: string): UIMessage { + return { id, role: "user", parts: [{ type: "text", text }] }; +} + +function textOf(message: UIMessage): string { + return message.parts.map((part) => (part.type === "text" ? part.text : "")).join(""); +} + +/** + * Two steps with a tool call in between, so there is a step boundary for the + * steering queue to drain at. Step 1 calls the tool, step 2 answers. + */ +function twoStepModel(onFirstStep: () => Promise) { + let call = 0; + return new MockLanguageModelV3({ + doStream: async () => { + call += 1; + if (call === 1) { + const chunks: LanguageModelV3StreamPart[] = [ + { type: "tool-input-start", id: "c1", toolName: "lookup" }, + { type: "tool-input-delta", id: "c1", delta: "{}" }, + { type: "tool-input-end", id: "c1" }, + { type: "tool-call", toolCallId: "c1", toolName: "lookup", input: "{}" }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage }, + ]; + // Land the steering message while step 1 is streaming, so it is queued + // before the boundary that drains it. + await onFirstStep(); + return { stream: simulateReadableStream({ chunks }) }; + } + return { + stream: simulateReadableStream({ + chunks: [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "done" }, + { type: "text-end", id: "t1" }, + { type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage }, + ], + }), + }; + }, + }); +} + +describe("injected steering messages", () => { + it("enter the accumulator, so onTurnComplete can see them", async () => { + let captured: { ui: string[]; newUi: string[] } | undefined; + let injectedCount = 0; + + const send = { fn: async () => {} }; + + const agent = chat.agent({ + id: "steering-accumulator", + tools: { + lookup: tool({ + description: "look something up", + inputSchema: z.object({}), + execute: async () => ({ ok: true }), + }), + }, + pendingMessages: { + shouldInject: ({ steps }) => steps.length > 0, + onInjected: ({ messages }) => { + injectedCount = messages.length; + }, + }, + onTurnComplete: async ({ uiMessages, newUIMessages }) => { + captured = { + ui: uiMessages.map(textOf), + newUi: newUIMessages.map(textOf), + }; + }, + run: async ({ messages, tools, signal }) => + streamText({ + ...chat.toStreamTextOptions({ tools }), + model: twoStepModel(() => send.fn()), + messages, + abortSignal: signal, + stopWhen: stepCountIs(5), + }), + }); + + const harness = mockChatAgent(agent, { chatId: "steering-accumulator" }); + + send.fn = async () => { + await harness.sendPendingMessage(userMessage("actually, only the platform one", "steer-1")); + }; + + try { + await harness.sendMessage(userMessage("summarise every project", "u1")); + await new Promise((resolve) => setTimeout(resolve, 50)); + + // The injection happened — this is the SDK's own bookkeeping. + expect(injectedCount).toBe(1); + + expect(captured).toBeDefined(); + + /** + * The steering message reached the model and the browser. Before this fix it + * reached neither `uiMessages` nor `newUIMessages`, so an app persisting from + * `onTurnComplete` stored an answer shaped by an instruction it never saw, and + * rebuilt the next turn's context without it. + */ + expect(captured!.ui).toContain("actually, only the platform one"); + expect(captured!.newUi).toContain("actually, only the platform one"); + + // And in the order it happened: after the question, before the answer. + expect(captured!.ui.indexOf("actually, only the platform one")).toBeGreaterThan( + captured!.ui.indexOf("summarise every project") + ); + expect(captured!.ui.at(-1)).toBe("done"); + } finally { + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/steering-compaction-lanes.test.ts b/packages/trigger-sdk/test/steering-compaction-lanes.test.ts new file mode 100644 index 00000000000..5de7ed4a89d --- /dev/null +++ b/packages/trigger-sdk/test/steering-compaction-lanes.test.ts @@ -0,0 +1,178 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { sessionStreams } from "@trigger.dev/core/v3"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, stepCountIs, streamText, tool } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * Steering and compaction in the same turn. + * + * Compaction is model-only by design: it replaces the model messages with a + * summary and deliberately leaves the UI messages whole, so the chat still + * displays the full conversation. Reconciling the model lane by rebuilding it + * from the UI lane therefore un-compacts it, and the next turn is sent the + * entire pre-compaction transcript. + * + * The assertion that catches this is the absence of an early message, not the + * presence of the steer: a rebuild puts the steer in the prompt too, so a + * steer-presence check passes while compaction has been silently undone. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; + +function userMessage(text: string, id: string) { + return { id, role: "user" as const, parts: [{ type: "text" as const, text }] }; +} + +function deferred() { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} + +function textChunks(text: string): LanguageModelV3StreamPart[] { + return [ + { 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: USAGE }, + ]; +} + +function toolCallChunks(callId: string): LanguageModelV3StreamPart[] { + return [ + { type: "tool-call", toolCallId: callId, toolName: "gate", input: JSON.stringify({ q: "go" }) }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE }, + ]; +} + +type SeqReader = { lastSeqNum: (chatId: string, dir: "in" | "out") => number | undefined }; + +async function sendAndLand( + harness: { sendMessage: (m: ReturnType) => Promise }, + chatId: string, + text: string, + id: string +) { + const seqs = sessionStreams as unknown as SeqReader; + const before = seqs.lastSeqNum(chatId, "in") ?? -1; + void harness.sendMessage(userMessage(text, id)); + await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`); +} + +const flatUserTexts = (prompt: { role: string; content: unknown }[]) => + prompt + .filter((m) => m.role === "user") + .flatMap((m) => + Array.isArray(m.content) + ? (m.content as { type: string; text?: string }[]) + .filter((c) => c.type === "text") + .map((c) => c.text ?? "") + : [] + ); + +describe("chat.agent steering with compaction in the same turn", () => { + it("keeps the summary and adds the steer, rather than restoring the transcript", async () => { + const chatId = "steer-compaction"; + const toolGate = deferred(); + let toolEntered = false; + const prompts: string[][] = []; + const allPrompts: string[] = []; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push(flatUserTexts(prompt)); + allPrompts.push(JSON.stringify(prompt)); + const isToolStep = step++ % 2 === 0; + return { + stream: simulateReadableStream({ + chunks: isToolStep ? toolCallChunks(`tc-${step}`) : textChunks("done"), + initialDelayInMs: 10, + chunkDelayInMs: 2, + }), + }; + }, + }); + + let compacted = 0; + const agent = chat.agent({ + id: "steer-compaction", + pendingMessages: { shouldInject: () => true }, + compaction: { + // Compact once, at the first step boundary of turn 1. + shouldCompact: () => compacted === 0, + summarize: async () => { + compacted++; + return "SUMMARY-OF-EVERYTHING"; + }, + }, + run: async ({ messages, signal }) => + streamText({ + model, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }), + }); + + const harness = mockChatAgent(agent, { chatId }); + try { + const first = harness.sendMessage(userMessage("EARLY-SENTINEL", "u-1")); + await waitFor(() => toolEntered, "tool entered"); + await sendAndLand(harness, chatId, "steer-me", "u-2"); + toolGate.resolve(); + await first; + + await waitFor(() => compacted > 0, "compaction ran"); + const promptsAfterTurn1 = prompts.length; + + await harness.sendMessage(userMessage("m3", "u-3")); + await waitFor(() => prompts.length > promptsAfterTurn1, "turn 2 prompt built"); + + const turn2 = prompts[promptsAfterTurn1]!; + const turn2Raw = allPrompts[promptsAfterTurn1]!; + + // The steer has to survive. + expect(turn2).toContain("steer-me"); + // And so does the compaction: the summary is what the model gets... + expect(turn2Raw).toContain("SUMMARY-OF-EVERYTHING"); + // ...instead of the message the summary replaced. This is the assertion a + // rebuild-based reconciliation fails. + expect(turn2).not.toContain("EARLY-SENTINEL"); + } finally { + toolGate.resolve(); + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/steering-error-path.test.ts b/packages/trigger-sdk/test/steering-error-path.test.ts new file mode 100644 index 00000000000..8029a061e98 --- /dev/null +++ b/packages/trigger-sdk/test/steering-error-path.test.ts @@ -0,0 +1,209 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { sessionStreams } from "@trigger.dev/core/v3"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { stepCountIs, streamText, tool } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { UIMessage } from "ai"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * A turn that drains a steering message and then fails. + * + * The error path builds its own `newUIMessages` from the wire message and the + * partial response, so a steer the drain consumed is not in it. That is the + * same append-only persistence hole the steering fix exists to close: the app + * stores what `onTurnComplete` hands it, the failed turn hands it everything + * except the steer, and the instruction is gone. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; + +function userMessage(text: string, id: string) { + return { id, role: "user" as const, parts: [{ type: "text" as const, text }] }; +} + +function deferred() { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} + +function toolCallChunks(callId: string): LanguageModelV3StreamPart[] { + return [ + { type: "tool-call", toolCallId: callId, toolName: "gate", input: JSON.stringify({ q: "go" }) }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE }, + ]; +} + +/** + * Emits a partial then errors, one chunk per pull so the queue isn't reset by + * erroring in the same tick as the enqueue. + */ +function erroringStream(): ReadableStream { + const chunks: LanguageModelV3StreamPart[] = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "partial" }, + ]; + let i = 0; + return new ReadableStream({ + pull(controller) { + if (i < chunks.length) { + controller.enqueue(chunks[i++]!); + return; + } + controller.error(new Error("UND_ERR_BODY_TIMEOUT")); + }, + }); +} + +type SeqReader = { lastSeqNum: (chatId: string, dir: "in" | "out") => number | undefined }; + +async function sendAndLand( + harness: { sendMessage: (m: ReturnType) => Promise }, + chatId: string, + text: string, + id: string +) { + const seqs = sessionStreams as unknown as SeqReader; + const before = seqs.lastSeqNum(chatId, "in") ?? -1; + void harness.sendMessage(userMessage(text, id)); + await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`); +} + +describe("chat.agent steering on a turn that fails", () => { + it("still reports the steer in the error path's newUIMessages", async () => { + const chatId = "steer-error-path"; + const toolGate = deferred(); + let toolEntered = false; + const events: { + newUIMessages: UIMessage[]; + messages: unknown[]; + newMessages: unknown[]; + finishReason?: string; + }[] = []; + const promptsSawSteer: boolean[] = []; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + promptsSawSteer.push(JSON.stringify(prompt).includes("steer-me")); + // Step 1 calls the tool, step 2 fails mid-stream, the next turn answers. + const n = step++; + const fromChunks = (chunks: LanguageModelV3StreamPart[]) => ({ + stream: new ReadableStream({ + start(c) { + for (const ch of chunks) c.enqueue(ch); + c.close(); + }, + }), + }); + if (n === 0) return fromChunks(toolCallChunks("tc-1")); + if (n === 1) return { stream: erroringStream() }; + return fromChunks([ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "ok" }, + { type: "text-end", id: "t1" }, + { type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage: USAGE }, + ]); + }, + }); + + const agent = chat.agent({ + id: "steer-error-path", + pendingMessages: { + shouldInject: () => true, + prepare: async ({ messages }) => [ + { + role: "system", + content: `[OPERATOR-NOTE] ${messages + .map((m) => (m.parts as { text?: string }[]).map((p) => p.text ?? "").join("")) + .join(" ")}`, + }, + ], + }, + onTurnComplete: async ({ newUIMessages, messages, newMessages, finishReason }) => { + events.push({ + newUIMessages: [...(newUIMessages ?? [])], + messages: [...messages], + newMessages: [...(newMessages ?? [])], + finishReason, + }); + }, + run: async ({ messages, signal }) => + streamText({ + model, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }), + }); + + const harness = mockChatAgent(agent, { chatId }); + try { + void harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => toolEntered, "tool entered"); + await sendAndLand(harness, chatId, "steer-me", "u-2"); + toolGate.resolve(); + + await waitFor(() => events.length >= 1, "turn complete fired"); + + // The turn really did fail, and the steer really did reach the model, + // so the lane check below is about persistence and nothing else. + expect(events[0]!.finishReason).toBe("error"); + expect(promptsSawSteer.some(Boolean)).toBe(true); + + const texts = events[0]!.newUIMessages.flatMap((m) => + ((m.parts ?? []) as { type: string; text?: string }[]) + .filter((p) => p.type === "text") + .map((p) => p.text ?? "") + ); + expect(texts).toContain("steer-me"); + + // The model lane the failed turn reports has to carry it too, and so + // does the prompt the next turn actually sends. Reconciling only on the + // success path leaves it pending, so the next turn misses it and it + // lands a slot late at the end of that turn. + expect(JSON.stringify(events[0]!.messages)).toContain("[OPERATOR-NOTE] steer-me"); + // The per-turn delta carries the same form, not a reconversion of the UI + // message: append-only model persistence from `newMessages` would + // otherwise store a different instruction from the one the model acted on. + expect(JSON.stringify(events[0]!.newMessages)).toContain("[OPERATOR-NOTE] steer-me"); + const promptsBefore = promptsSawSteer.length; + await harness.sendMessage(userMessage("m3", "u-3")); + await waitFor(() => promptsSawSteer.length > promptsBefore, "turn 2 prompt built"); + expect(promptsSawSteer[promptsBefore]).toBe(true); + } finally { + toolGate.resolve(); + await harness.close(); + } + }); +}); diff --git a/packages/trigger-sdk/test/steering-history-edit-once.test.ts b/packages/trigger-sdk/test/steering-history-edit-once.test.ts new file mode 100644 index 00000000000..c33b28285bd --- /dev/null +++ b/packages/trigger-sdk/test/steering-history-edit-once.test.ts @@ -0,0 +1,200 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { sessionStreams } from "@trigger.dev/core/v3"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, stepCountIs, streamText, tool } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * A `chat.history` edit after a steer has been drained. + * + * The edit is applied by rebuilding the model lane from the UI lane, and the + * UI lane already holds the steer, so the rebuilt lane has it. Appending the + * recorded model form on top of that sends it twice from the next turn on. A + * steer-presence check passes either way; the count is what discriminates. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; +const userMessage = (text: string, id: string) => ({ + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], +}); +function deferred() { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} +const textChunks = (text: string): 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: USAGE }, +]; +const toolCallChunks = (callId: string): LanguageModelV3StreamPart[] => [ + { type: "tool-call", toolCallId: callId, toolName: "gate", input: JSON.stringify({ q: "go" }) }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE }, +]; +type SeqReader = { lastSeqNum: (chatId: string, dir: "in" | "out") => number | undefined }; +async function sendAndLand( + harness: { sendMessage: (m: ReturnType) => Promise }, + chatId: string, + text: string, + id: string +) { + const seqs = sessionStreams as unknown as SeqReader; + const before = seqs.lastSeqNum(chatId, "in") ?? -1; + void harness.sendMessage(userMessage(text, id)); + await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`); +} +const countOf = (hay: string, needle: string) => hay.split(needle).length - 1; + +type Variant = { prepare?: boolean; compact?: boolean; deleteSteer?: boolean }; + +/** One steered turn with a history edit from onInjected, then a follow-up turn. Returns turn 2's prompt. */ +async function runVariant(chatId: string, v: Variant): Promise { + const toolGate = deferred(); + let toolEntered = false; + const prompts: string[] = []; + let compacted = 0; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push(JSON.stringify(prompt)); + const isToolStep = step++ % 2 === 0; + return { + stream: simulateReadableStream({ + chunks: isToolStep ? toolCallChunks(`tc-${step}`) : textChunks("done"), + initialDelayInMs: 10, + chunkDelayInMs: 2, + }), + }; + }, + }); + + const agent = chat.agent({ + id: chatId, + pendingMessages: { + shouldInject: () => true, + ...(v.prepare + ? { + prepare: async ({ messages }) => [ + { + role: "system" as const, + content: `[OPERATOR-NOTE] ${messages.map((m) => (m.parts as { text?: string }[]).map((p) => p.text ?? "").join("")).join(" ")}`, + }, + ], + } + : {}), + onInjected: () => { + chat.history.set(chat.history.all().filter((m) => !(v.deleteSteer && m.id === "u-2"))); + }, + }, + ...(v.compact + ? { + compaction: { + shouldCompact: () => compacted === 0, + summarize: async () => { + compacted++; + return "SUMMARY-OF-EVERYTHING"; + }, + }, + } + : {}), + run: async ({ messages, signal }) => + streamText({ + model, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }), + }); + + const harness = mockChatAgent(agent, { chatId }); + try { + const first = harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => toolEntered, "tool entered"); + await sendAndLand(harness, chatId, "steer-me", "u-2"); + toolGate.resolve(); + await first; + await waitFor(() => prompts.length >= 2, "turn 1 done"); + if (v.compact) await waitFor(() => compacted > 0, "compaction ran"); + const promptsAfterTurn1 = prompts.length; + await harness.sendMessage(userMessage("m3", "u-3")); + await waitFor(() => prompts.length > promptsAfterTurn1, "turn 2 prompt built"); + return prompts[promptsAfterTurn1]!; + } finally { + toolGate.resolve(); + await harness.close(); + } +} + +describe("a history edit after a steer was drained", () => { + it("sends the steer to later turns exactly once", { timeout: 30_000 }, async () => { + const turn2 = await runVariant("steer-history-edit-once", {}); + expect(countOf(turn2, '"steer-me"')).toBe(1); + }); + + it("keeps the prepared form, once", { timeout: 30_000 }, async () => { + /** + * The rebuild converts the UI message, which is the raw form. If the raw + * form is what stays, the model's memory of the instruction differs from + * the one it acted on. If both stay, it is there twice. + */ + const turn2 = await runVariant("steer-history-edit-prepared", { prepare: true }); + expect(countOf(turn2, "[OPERATOR-NOTE] steer-me")).toBe(1); + expect(countOf(turn2, '"steer-me"')).toBe(0); + }); + + it("keeps the steer when compaction also replaces the lane", { timeout: 30_000 }, async () => { + /** + * A model-only compaction replaces the rebuilt lane, raw steer included. + * If reconciliation then withholds the prepared form because the rebuild + * "already had it", the steer is gone from the model lane altogether. + */ + const turn2 = await runVariant("steer-history-edit-compacted", { + prepare: true, + compact: true, + }); + expect(turn2).toContain("SUMMARY-OF-EVERYTHING"); + expect(countOf(turn2, "[OPERATOR-NOTE] steer-me")).toBe(1); + }); + + it("does not bring back a steer the edit removed", { timeout: 30_000 }, async () => { + /** The edit is the app's decision. Reconciliation must not undo it. */ + const turn2 = await runVariant("steer-history-edit-deleted", { + prepare: true, + deleteSteer: true, + }); + expect(turn2).not.toContain("steer-me"); + }); +}); diff --git a/packages/trigger-sdk/test/steering-injection.test.ts b/packages/trigger-sdk/test/steering-injection.test.ts index 12f3dfa0c87..514e285913a 100644 --- a/packages/trigger-sdk/test/steering-injection.test.ts +++ b/packages/trigger-sdk/test/steering-injection.test.ts @@ -355,18 +355,14 @@ describe("chat.agent injection claims only its own batch", () => { /** * Whether an injected message survives into the next turn's model context. * - * Recorded here because the deployed QA lane finds the two surfaces disagree: - * a `chat.createSession()` recap in the same run recalls a mid-turn steer, - * while the managed `chat.agent` loop denies it. That difference is - * pre-existing and is the surface-specific half of the accumulator gap. - * - * `it.fails` because the managed path does not carry it: turn 2's prompt comes - * back as the original and the following message only, with the injected one - * absent. Held here so the day that changes is noticed, and so the gap has a - * repro that does not need a deployed environment. + * The UI accumulator and the model accumulator are maintained separately, and + * a drained message used to reach only the first: the browser, the snapshot + * and `chat.history.*` all showed it while every later turn of the run + * answered without it. The model lane is now rebuilt from the UI lane at the + * end of a turn that drained, and this is the repro for it. */ describe("chat.agent injected message in the next turn's context", () => { - it.fails( + it( "carries an injected message into the following turn's prompt", { timeout: 30_000 }, async () => { @@ -448,4 +444,97 @@ describe("chat.agent injected message in the next turn's context", () => { } } ); + + /** + * The same thing for a turn that captures no assistant response. + * + * `run()` piping the stream itself skips the auto-pipe, so no `onFinish` is + * attached and nothing is captured. The rebuild sits outside both + * `capturedResponseMessage` branches for that reason: moving it inside + * either one leaves this turn's model lane without the steer while the + * captured case looks fine. + */ + it( + "carries it into the following turn when the turn captures no response", + { timeout: 30_000 }, + async () => { + const chatId = "inject-next-turn-manual-pipe"; + const toolGate = makeGate(); + let toolEntered = false; + const prompts: string[][] = []; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + let step = 0; + const recordingModel = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push( + prompt + .filter((m) => m.role === "user") + .flatMap((m) => + Array.isArray(m.content) + ? (m.content as { type: string; text?: string }[]) + .filter((c) => c.type === "text") + .map((c) => c.text ?? "") + : [] + ) + ); + const isToolStep = step++ % 2 === 0; + return { + stream: simulateReadableStream({ + chunks: isToolStep ? toolCallChunks(`tc-${step}`) : textChunks("done"), + initialDelayInMs: 10, + chunkDelayInMs: 2, + }), + }; + }, + }); + + const agent = chat.agent({ + id: "steering-injection.next-turn-manual-pipe", + pendingMessages: { shouldInject: () => true }, + run: async ({ messages, signal }) => { + const result = streamText({ + model: recordingModel, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }); + // Piping here rather than returning the result is what leaves the + // turn with no captured response. + await chat.pipe(result.toUIMessageStream(), { signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId }); + try { + const first = harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => toolEntered, "tool entered"); + await sendAndLand(harness, chatId, "steer-me", "u-2"); + toolGate.open(); + await first; + + await waitFor(() => turnCompleteCount(harness) >= 1, "turn 1 complete"); + const promptsAfterTurn1 = prompts.length; + + await harness.sendMessage(userMessage("m3", "u-3")); + await waitFor(() => prompts.length > promptsAfterTurn1, "turn 2 prompt built"); + + expect(prompts[promptsAfterTurn1]!).toContain("steer-me"); + } finally { + toolGate.open(); + await harness.close(); + } + } + ); }); diff --git a/packages/trigger-sdk/test/steering-prepare-transform.test.ts b/packages/trigger-sdk/test/steering-prepare-transform.test.ts new file mode 100644 index 00000000000..0d6102353e3 --- /dev/null +++ b/packages/trigger-sdk/test/steering-prepare-transform.test.ts @@ -0,0 +1,152 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { sessionStreams } from "@trigger.dev/core/v3"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, stepCountIs, streamText, tool } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * `pendingMessages.prepare` decides how a steer is presented to the model. + * The steered turn gets that form. Later turns have to get the same form, + * or the model's memory of the instruction differs from what it acted on. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; +const userMessage = (text: string, id: string) => ({ + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], +}); +function deferred() { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} +async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} +const textChunks = (text: string): 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: USAGE }, +]; +const toolCallChunks = (callId: string): LanguageModelV3StreamPart[] => [ + { type: "tool-call", toolCallId: callId, toolName: "gate", input: JSON.stringify({ q: "go" }) }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE }, +]; +type SeqReader = { lastSeqNum: (chatId: string, dir: "in" | "out") => number | undefined }; +async function sendAndLand( + harness: { sendMessage: (m: ReturnType) => Promise }, + chatId: string, + text: string, + id: string +) { + const seqs = sessionStreams as unknown as SeqReader; + const before = seqs.lastSeqNum(chatId, "in") ?? -1; + void harness.sendMessage(userMessage(text, id)); + await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`); +} + +describe("a steer transformed by pendingMessages.prepare", () => { + it("reaches later turns in the transformed form", { timeout: 30_000 }, async () => { + const chatId = "steer-prepare-transform"; + const toolGate = deferred(); + let toolEntered = false; + const prompts: string[] = []; + const newModelDeltas: string[] = []; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + let step = 0; + const model = new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push(JSON.stringify(prompt)); + const isToolStep = step++ % 2 === 0; + return { + stream: simulateReadableStream({ + chunks: isToolStep ? toolCallChunks(`tc-${step}`) : textChunks("done"), + initialDelayInMs: 10, + chunkDelayInMs: 2, + }), + }; + }, + }); + + const agent = chat.agent({ + id: "steer-prepare-transform", + pendingMessages: { + shouldInject: () => true, + // Present the steer to the model as an operator note, not a user turn. + prepare: async ({ messages }) => [ + { + role: "system", + content: `[OPERATOR-NOTE] ${messages.map((m) => (m.parts as { text?: string }[]).map((p) => p.text ?? "").join("")).join(" ")}`, + }, + ], + }, + onTurnComplete: async ({ newMessages }) => { + newModelDeltas.push(JSON.stringify(newMessages)); + }, + run: async ({ messages, signal }) => + streamText({ + model, + messages, + abortSignal: signal, + ...chat.toStreamTextOptions(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }), + }); + + const harness = mockChatAgent(agent, { chatId }); + try { + const first = harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => toolEntered, "tool entered"); + await sendAndLand(harness, chatId, "steer-me", "u-2"); + toolGate.resolve(); + await first; + await waitFor(() => prompts.length >= 2, "turn 1 second step"); + + // The steered turn saw the transformed form. + expect(prompts[1]!).toContain("[OPERATOR-NOTE] steer-me"); + const promptsAfterTurn1 = prompts.length; + + await harness.sendMessage(userMessage("m3", "u-3")); + await waitFor(() => prompts.length > promptsAfterTurn1, "turn 2 prompt built"); + + // And so does the next one. Reconverting the UI message gives the raw + // user turn instead, so the model remembers a different instruction + // from the one it acted on. + expect(prompts[promptsAfterTurn1]!).toContain("[OPERATOR-NOTE] steer-me"); + + // The per-turn model delta the hook reports carries it in the same form, + // or append-only persistence from `newMessages` loses the model's view. + expect(newModelDeltas[0]!).toContain("[OPERATOR-NOTE] steer-me"); + } finally { + toolGate.resolve(); + await harness.close(); + } + }); +});