fix(sdk): keep mid-turn chat.inject messages in place in saved history

Fix `chat.inject` messages being saved out of order when called from a
tool mid-turn, causing prompt-cache misses. Background injections now
retain their correct position in conversation history without affecting
existing steering behavior or exposing injected content to clients.

Mono-RevId: e3f87cf4e10526bad0b1f2c53ebc97a7f99b19f4

Author: kathiekiwi

.changeset/chat-inject-mid-turn-order.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Fix `chat.inject` messages added from a tool mid-turn being saved before the assistant's response. They now stay after the tool result, where the model saw them, so the next turn keeps the same prompt prefix. If you roll `@trigger.dev/sdk` back to an earlier version, chats saved by this version lose these mid-turn `chat.inject` messages from the model's context, though the chat history shown in the UI keeps them.

docs/ai-chat/background-injection.mdx

@@ -70,6 +70,8 @@ export const myChat = chat.agent({
 
 If the background work finishes *during* a tool-call loop (not between turns), the messages are picked up at the next `prepareStep` boundary instead.
 
+An injection made from a tool mid-turn is recorded in the conversation history at the step where it happened, after the tool result. The next turn's prompt prefix is unchanged.
+
 ## Example: self-review
 
 A cheap model reviews the agent's response after each turn and injects coaching for the next one. Uses [Prompts](/ai/prompts) for the review prompt and `generateObject` for structured output.

docs/ai-chat/pending-messages.mdx

@@ -132,6 +132,8 @@ When messages are injected, the SDK automatically writes a `data-pending-message
 - Remove them from the pending overlay
 - Render them inline at the injection point in the assistant response
 
+Background injections appear as a different data part (`data-background-injected`) that `usePendingMessages` ignores. It is not a user message.
+
 A "pending message injected" span also appears in the run trace.
 
 ## Backend: chat.createSession

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

@@ -24,6 +24,13 @@ import type { InferUITools, ToolSet, UIDataTypes, UIMessage } from "ai";
  */
 export const PENDING_MESSAGE_INJECTED_TYPE = "data-pending-message-injected" as const;
 
+/**
+ * Message-part `type` value marking where `chat.inject` messages drained
+ * mid-turn reached the model. Separate from the pending-message type so
+ * clients that render injection markers do not render these.
+ */
+export const BACKGROUND_INJECTION_TYPE = "data-background-injected" as const;
+
 // Declared in `chat.ts` (a public subpath) so customer declaration emit can
 // name them — declaring them here breaks `declaration: true` consumers with
 // TS2742, since this module isn't reachable via the package exports map.

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

@@ -113,6 +113,7 @@ import {
   convertSteeredMessages,
   retainStepMessages,
   steeringMarkers,
+  withoutInjectionMarkers,
   type SteeringInjection,
 } from "./steeringContext.js";
 
@@ -134,6 +135,7 @@ import {
   type InferChatClientData,
   type InferChatUIMessage,
   type InferChatUIMessageFromTools,
+  BACKGROUND_INJECTION_TYPE,
   PENDING_MESSAGE_INJECTED_TYPE,
   upsertIncomingMessage,
   chatRunTags,
@@ -2901,13 +2903,6 @@ function spliceHandoverPartial(
  * @internal
  */
 const chatBackgroundQueueKey = locals.create<ModelMessage[]>("chat.backgroundQueue");
-/**
- * Background injections a step-boundary drain handed to the model this turn,
- * with the transcript message they followed. Reconciled into the model lane
- * and the persisted injections once the turn's response is in.
- */
-const chatPendingBackgroundKey =
-  locals.create<{ afterId: string; messages: ModelMessage[] }[]>("chat.pendingBackground");
 
 /**
  * System-role context injected mid-conversation, held for the instructions lane.
@@ -4450,6 +4445,34 @@ function modelFormOf(m: UIMessage, batch: UIMessage[], injected: ModelMessage[])
   return batch[0] === m ? injected : [];
 }
 
+/**
+ * Record a prepared injection and write its marker into the response, so the
+ * model lane rebuilds it at the step boundary where the model received it.
+ */
+function markInjection(
+  steps: CompactionStep[],
+  messageIds: string[],
+  messages: ModelMessage[],
+  markerMessages: Array<{ id: string; text: string }>,
+  kind?: "background"
+): void {
+  const id = generateMessageId();
+  const injections = locals.get(chatSteeringInjectionsKey) ?? new Map();
+  injections.set(id, { id, messageIds, messages, ...(kind ? { kind } : {}) });
+  locals.set(chatSteeringInjectionsKey, injections);
+  const response = managedResponse();
+  // prepareStep can run ahead of the UI consumer. Admit the marker only
+  // once the actual preceding finish-step chunk has entered this stream.
+  response.afterToolStep(
+    steps.at(-1)?.content.flatMap((part) => (part.toolCallId ? [part.toolCallId] : [])) ?? []
+  );
+  response.writeData({
+    type: kind === "background" ? BACKGROUND_INJECTION_TYPE : PENDING_MESSAGE_INJECTED_TYPE,
+    id,
+    data: { messageIds, messages: markerMessages },
+  });
+}
+
 /**
  * Drain the steering queue as a batch. Calls `shouldInject` once with all
  * pending messages. If it returns true, calls `prepareMessages` once to
@@ -4468,7 +4491,6 @@ async function drainSteeringQueue(
   steps: CompactionStep[],
   queueOverride?: SteeringQueueEntry[]
 ): Promise<DrainedSteering> {
-  if (steps.length === 0) managedResponse().beginGeneration();
   const queue = queueOverride ?? locals.get(chatSteeringQueueKey);
   if (!queue || queue.length === 0) return EMPTY_DRAIN;
 
@@ -4606,23 +4628,12 @@ async function drainSteeringQueue(
       }
 
       if (injected.length > 0) {
-        const id = generateMessageId();
-        const messageIds = claimedUIMessages.map((m) => m.id);
-        const injections = locals.get(chatSteeringInjectionsKey) ?? new Map();
-        injections.set(id, { id, messageIds, messages: injected });
-        locals.set(chatSteeringInjectionsKey, injections);
-        const response = managedResponse();
-        // prepareStep can run ahead of the UI consumer. Admit the marker only
-        // once the actual preceding finish-step chunk has entered this stream.
-        response.afterStep(steps.length);
-        response.writeData({
-          type: PENDING_MESSAGE_INJECTED_TYPE,
-          id,
-          data: {
-            messageIds,
-            messages: claimedUIMessages.map((m) => ({ id: m.id, text: textOfUIMessage(m) })),
-          },
-        });
+        markInjection(
+          steps,
+          claimedUIMessages.map((m) => m.id),
+          injected,
+          claimedUIMessages.map((m) => ({ id: m.id, text: textOfUIMessage(m) }))
+        );
       }
 
       // Fire onInjected callback
@@ -5439,13 +5450,7 @@ function toStreamTextOptions(
       if (bgQueue && bgQueue.length > 0) {
         const injected = bgQueue.splice(0); // drain
         resultMessages = [...(resultMessages ?? messages), ...injected];
-        const pendingBackground = locals.get(chatPendingBackgroundKey) ?? [];
-        pendingBackground.push({
-          afterId:
-            (locals.get(chatCurrentUIMessagesKey) as UIMessage[] | undefined)?.at(-1)?.id ?? "",
-          messages: injected,
-        });
-        locals.set(chatPendingBackgroundKey, pendingBackground);
+        markInjection(steps, [], injected, [], "background");
       }
 
       // 4. Parallel tool results in call order, the order the next turn's history
@@ -9164,23 +9169,6 @@ 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 background injections a step-boundary
-       * drain handed to the model this turn, and record them for persistence.
-       * Returns how many model messages were appended.
-       */
-      const reconcilePendingBackground = (): number => {
-        const pending = locals.get(chatPendingBackgroundKey);
-        if (!pending || pending.length === 0) return 0;
-        locals.set(chatPendingBackgroundKey, []);
-        let appended = 0;
-        for (const entry of pending) {
-          accumulatedMessages.push(...entry.messages);
-          laneInjections.push(entry);
-          appended += entry.messages.length;
-        }
-        return appended;
-      };
       /**
        * Give the model accumulator the steering messages a drain consumed,
        * in the form the model actually received. Appended, never reconverted
@@ -9211,6 +9199,15 @@ function chatAgent<
       // but in frontend-friendly UIMessage format (with parts, id, etc.).
       let accumulatedUIMessages: TUIMessage[] = [];
 
+      const cleanIncomingMessage = (msg: TUIMessage): TUIMessage => {
+        if (msg.role !== "assistant") return withoutInjectionMarkers(msg);
+        const cleaned = cleanupAbortedParts(msg);
+        // Markers are the agent's own; only a known assistant message keeps them.
+        return accumulatedUIMessages.some((m) => m.id === msg.id)
+          ? cleaned
+          : withoutInjectionMarkers(cleaned);
+      };
+
       // ── Snapshot + replay (gated on prior-state signals) ─────────────
       //
       // With `hydrateMessages` registered the customer owns history — the
@@ -10960,9 +10957,8 @@ function chatAgent<
             // Cleaning happens once here so `extractLastUserMessageText` and
             // every downstream consumer see the same message shape — and
             // `cleanupAbortedParts` no longer has to be re-applied below.
-            const cleanedIncomingMessages: TUIMessage[] = incomingMessages.map((msg) =>
-              msg.role === "assistant" ? cleanupAbortedParts(msg) : msg
-            );
+            const cleanedIncomingMessages: TUIMessage[] =
+              incomingMessages.map(cleanIncomingMessage);
             const lastUserMessage = extractLastUserMessageText(cleanedIncomingMessages);
 
             channelTurnRan =
@@ -11036,7 +11032,6 @@ function chatAgent<
                 locals.set(chatOutGateKey, createChatOutGate());
                 locals.set(chatCompactionStateKey, undefined);
                 locals.set(chatSteeringQueueKey, []);
-                locals.set(chatPendingBackgroundKey, []);
                 await locals.get(chatManagedResponseKey)?.close();
                 locals.set(chatManagedResponseKey, undefined);
                 locals.set(chatManagedResponseActiveKey, true);
@@ -12308,11 +12303,10 @@ function chatAgent<
                   // 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,
-                      response: capturedResponseMessage,
-                    }).reduce((n, e) => n + e.model.length, 0) + reconcilePendingBackground();
+                  const steerTailThisTurn = reconcilePendingSteer({
+                    turnNew: turnNewModelMessages,
+                    response: capturedResponseMessage,
+                  }).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
@@ -13346,7 +13340,8 @@ function chatAgent<
             // yet (a pre-run hook threw), so fold the wire message in for the
             // error event + snapshot — the cursor has already advanced past it,
             // so otherwise it survives in neither the snapshot nor the `.in` tail.
-            const erroredWireMessage = (currentWirePayload as { message?: TUIMessage }).message;
+            const rawWireMessage = (currentWirePayload as { message?: TUIMessage }).message;
+            const erroredWireMessage = rawWireMessage && cleanIncomingMessage(rawWireMessage);
             const erroredUIMessages =
               erroredWireMessage &&
               !accumulatedUIMessages.some((m) => m.id === erroredWireMessage.id)
@@ -13414,7 +13409,6 @@ function chatAgent<
             let erroredNewModelMessages: ModelMessage[] = [];
 
             const reconciledSteer = reconcilePendingSteer({ response: partialResponse });
-            const backgroundTailThisTurn = reconcilePendingBackground();
 
             if (!responseCommitted) {
               try {
@@ -13442,8 +13436,7 @@ function chatAgent<
                       accumulatedMessages,
                       erroredUIMessages[partialIdx]!,
                       partialResponse!,
-                      reconciledSteer.reduce((n, e) => n + e.model.length, 0) +
-                        backgroundTailThisTurn
+                      reconciledSteer.reduce((n, e) => n + e.model.length, 0)
                     );
                     if (!ok) {
                       logger.warn(

packages/trigger-sdk/src/v3/chat-react.ts

@@ -408,6 +408,41 @@ export type UsePendingMessagesReturn = {
   getInjectedMessages: (part: unknown) => InjectedMessage[];
 };
 
+/** @internal */
+export function isInjectionPart(part: unknown): boolean {
+  return (
+    typeof part === "object" &&
+    part !== null &&
+    (part as any).type === PENDING_MESSAGE_INJECTED_TYPE
+  );
+}
+
+/** @internal */
+export function injectedMessageIdsOf(part: unknown): string[] {
+  if (!isInjectionPart(part)) return [];
+  const ids = (part as any).data?.messageIds;
+  return Array.isArray(ids) ? ids : [];
+}
+
+/** @internal */
+export function injectedMessagesOf(part: unknown): InjectedMessage[] {
+  if (!isInjectionPart(part)) return [];
+  const msgs = (part as any).data?.messages;
+  return Array.isArray(msgs) ? msgs : [];
+}
+
+/** @internal */
+export function collectInjectedMessageIds(messages: ReadonlyArray<UIMessage>): string[] {
+  const ids: string[] = [];
+  for (const msg of messages) {
+    if (msg.role !== "assistant") continue;
+    for (const part of msg.parts ?? []) {
+      ids.push(...injectedMessageIdsOf(part));
+    }
+  }
+  return ids;
+}
+
 /**
  * React hook for managing pending messages (steering) during streaming.
  *
@@ -469,20 +504,10 @@ export function usePendingMessages<TUIMessage extends UIMessage = UIMessage>(
   useEffect(() => {
     if (status !== "streaming") return;
     let newlyInjected = false;
-    for (const msg of messages) {
-      if (msg.role !== "assistant") continue;
-      for (const part of msg.parts ?? []) {
-        if ((part as any).type === PENDING_MESSAGE_INJECTED_TYPE) {
-          const messageIds = (part as any).data?.messageIds;
-          if (Array.isArray(messageIds)) {
-            for (const id of messageIds) {
-              if (!injectedIdsRef.current.has(id)) {
-                injectedIdsRef.current.add(id);
-                newlyInjected = true;
-              }
-            }
-          }
-        }
+    for (const id of collectInjectedMessageIds(messages)) {
+      if (!injectedIdsRef.current.has(id)) {
+        injectedIdsRef.current.add(id);
+        newlyInjected = true;
       }
     }
     // Remove injected steering messages from the pending overlay immediately
@@ -579,31 +604,9 @@ export function usePendingMessages<TUIMessage extends UIMessage = UIMessage>(
     [transport, chatId, metadata]
   );
 
-  const isInjectionPoint = useCallback(
-    (part: unknown): boolean =>
-      typeof part === "object" &&
-      part !== null &&
-      (part as any).type === PENDING_MESSAGE_INJECTED_TYPE,
-    []
-  );
-
-  const getInjectedMessageIds = useCallback(
-    (part: unknown): string[] => {
-      if (!isInjectionPoint(part)) return [];
-      const ids = (part as any).data?.messageIds;
-      return Array.isArray(ids) ? ids : [];
-    },
-    [isInjectionPoint]
-  );
-
-  const getInjectedMessages = useCallback(
-    (part: unknown): InjectedMessage[] => {
-      if (!isInjectionPoint(part)) return [];
-      const msgs = (part as any).data?.messages;
-      return Array.isArray(msgs) ? msgs : [];
-    },
-    [isInjectionPoint]
-  );
+  const isInjectionPoint = isInjectionPart;
+  const getInjectedMessageIds = injectedMessageIdsOf;
+  const getInjectedMessages = injectedMessagesOf;
 
   const pending: PendingMessage[] = pendingMsgs.map((m) => ({
     id: m.id,