diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 285c62cef..89b758e34 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -600,6 +600,102 @@ export type ChatTaskCompactionOptions = { /** @internal */ const chatTaskCompactionKey = locals.create("chat.taskCompaction"); +// --------------------------------------------------------------------------- +// Pending messages — mid-execution message injection via prepareStep +// --------------------------------------------------------------------------- + +/** + * Event passed to `shouldInject` and `prepareMessages` callbacks. + */ +export type PendingMessagesBatchEvent = { + /** All pending UI messages that arrived during streaming (batch). */ + messages: UIMessage[]; + /** Current model messages in the conversation. */ + modelMessages: ModelMessage[]; + /** Completed steps so far. */ + steps: CompactionStep[]; + /** Current step number (0-indexed). */ + stepNumber: number; + /** Chat session ID. */ + chatId: string; + /** Current turn number (0-indexed). */ + turn: number; + /** Custom data from the frontend. */ + clientData?: unknown; +}; + +/** + * Event passed to `onReceived` callback (per-message, as they arrive). + */ +export type PendingMessageReceivedEvent = { + /** The UI message that arrived during streaming. */ + message: UIMessage; + /** Chat session ID. */ + chatId: string; + /** Current turn number (0-indexed). */ + turn: number; +}; + +/** + * Event passed to `onInjected` callback (batch, after injection). + */ +export type PendingMessagesInjectedEvent = { + /** All UI messages that were injected. */ + messages: UIMessage[]; + /** The model messages that were injected. */ + injectedModelMessages: ModelMessage[]; + /** Chat session ID. */ + chatId: string; + /** Current turn number (0-indexed). */ + turn: number; + /** Step number where injection occurred. */ + stepNumber: number; +}; + +/** + * Options for the `pendingMessages` field on `chat.task()`, `chat.createSession()`, + * or `ChatMessageAccumulator`. + * + * Configures how messages that arrive during streaming are handled. When + * `shouldInject` is provided and returns `true`, the full batch of pending + * messages is injected between tool-call steps via `prepareStep`. + * Otherwise, messages queue for the next turn. + */ +export type PendingMessagesOptions = { + /** + * Decide whether to inject pending messages between tool-call steps. + * Called once per step boundary with the full batch of pending messages. + * If absent, no injection happens — messages only queue for the next turn. + */ + shouldInject?: (event: PendingMessagesBatchEvent) => boolean | Promise; + /** + * Transform the batch of pending messages before injection. + * Return the model messages to inject. + * Default: convert each UI message via `convertToModelMessages`. + */ + prepare?: (event: PendingMessagesBatchEvent) => ModelMessage[] | Promise; + /** Called when a message arrives during streaming (per-message). */ + onReceived?: (event: PendingMessageReceivedEvent) => void | Promise; + /** Called after a batch of messages is injected via `prepareStep`. */ + onInjected?: (event: PendingMessagesInjectedEvent) => void | Promise; +}; + +/** + * The data part type used to signal that pending messages were injected + * between tool-call steps. The frontend can match on this to render + * injection points inline in the assistant response. + */ +export const PENDING_MESSAGE_INJECTED_TYPE = "data-pending-message-injected" as const; + +/** @internal */ +type SteeringQueueEntry = { uiMessage: UIMessage; modelMessages: ModelMessage[] }; +/** @internal */ +const chatPendingMessagesKey = locals.create("chat.pendingMessages"); +/** @internal */ +const chatSteeringQueueKey = locals.create("chat.steeringQueue"); +/** @internal — IDs of messages that were successfully injected via prepareStep */ +const chatInjectedMessageIdsKey = locals.create>("chat.injectedMessageIds"); + /** * Event passed to the `prepareMessages` hook. */ @@ -1009,6 +1105,129 @@ function chatCompactionStep(options: CompactionOptions): (args: { messages: Mode }; } +// --------------------------------------------------------------------------- +// Steering queue drain — shared by toStreamTextOptions, session, accumulator +// --------------------------------------------------------------------------- + +/** + * Drain the steering queue as a batch. Calls `shouldInject` once with all + * pending messages. If it returns true, calls `prepareMessages` once to + * transform the batch, then clears the queue. + * Returns the model messages to inject (empty if none). + * @internal + */ +async function drainSteeringQueue( + config: PendingMessagesOptions, + messages: ModelMessage[], + steps: CompactionStep[], + queueOverride?: SteeringQueueEntry[], +): Promise { + const queue = queueOverride ?? locals.get(chatSteeringQueueKey); + if (!queue || queue.length === 0) return []; + + const ctx = locals.get(chatTurnContextKey); + const stepNumber = steps.length - 1; + const uiMessages = queue.map((e) => e.uiMessage); + + const batchEvent: PendingMessagesBatchEvent = { + messages: uiMessages, + modelMessages: messages, + steps, + stepNumber, + chatId: ctx?.chatId ?? "", + turn: ctx?.turn ?? 0, + clientData: ctx?.clientData, + }; + + // Call shouldInject once for the whole batch + const shouldInject = config.shouldInject + ? await config.shouldInject(batchEvent) + : false; + + if (!shouldInject) return []; + + // Extract message texts for span attributes + const messageTexts = uiMessages.map((m) => + (m.parts ?? []).filter((p: any) => p.type === "text").map((p: any) => p.text).join("") || "" + ); + const previewText = messageTexts.length === 1 + ? messageTexts[0]!.slice(0, 80) + : `${queue.length} messages`; + + return tracer.startActiveSpan( + "pending message injected", + async () => { + // Transform the batch — default: concatenate all pre-converted model messages + const injected = config.prepare + ? await config.prepare(batchEvent) + : queue.flatMap((e) => e.modelMessages); + + // Clear the queue and record injected IDs + queue.length = 0; + const injectedIds = locals.get(chatInjectedMessageIdsKey); + if (injectedIds) { + for (const m of uiMessages) injectedIds.add(m.id); + } + + // Write injection confirmation chunk to the stream so the frontend + // knows which messages were injected and where in the response. + if (injected.length > 0) { + try { + const { waitUntilComplete } = streams.writer(CHAT_STREAM_KEY, { + collapsed: true, + execute: ({ write }) => { + write({ + type: PENDING_MESSAGE_INJECTED_TYPE, + id: generateMessageId(), + data: { + messageIds: uiMessages.map((m) => m.id), + messages: uiMessages.map((m, idx) => ({ + id: m.id, + text: messageTexts[idx] ?? "", + })), + }, + }); + }, + }); + await waitUntilComplete(); + } catch { /* non-fatal — stream write failed */ } + } + + // Fire onInjected callback + if (config.onInjected && injected.length > 0) { + try { + await config.onInjected({ + messages: uiMessages, + injectedModelMessages: injected, + chatId: ctx?.chatId ?? "", + turn: ctx?.turn ?? 0, + stepNumber, + }); + } catch { /* non-fatal */ } + } + + return injected; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "tabler-message-forward", + "pending.message_count": uiMessages.length, + "pending.step_number": stepNumber, + "pending.messages": messageTexts, + ...(ctx?.chatId ? { "pending.chat_id": ctx.chatId } : {}), + ...(ctx?.turn != null ? { "pending.turn": ctx.turn } : {}), + ...accessoryAttributes({ + items: [ + { text: `${uiMessages.length} message${uiMessages.length === 1 ? "" : "s"}`, variant: "normal" }, + { text: `between steps ${stepNumber} and ${stepNumber + 1}`, variant: "normal" }, + ], + style: "codepath", + }), + }, + } + ); +} + // --------------------------------------------------------------------------- // chat.isCompactionSafe — check if it's safe to compact messages // --------------------------------------------------------------------------- @@ -1143,29 +1362,50 @@ function toStreamTextOptions(options?: ToStreamTextOptionsOptions): Record { - const compactResult = await chatCompact(messages, steps, { - shouldCompact: taskCompaction.shouldCompact, - summarize: (msgs) => { - const ctx = locals.get(chatTurnContextKey); - const lastStep = steps.at(-1); - return taskCompaction.summarize({ - messages: msgs, - usage: lastStep?.usage, - source: "inner", - stepNumber: steps.length - 1, - chatId: ctx?.chatId, - turn: ctx?.turn, - clientData: ctx?.clientData, - }); - }, - }); - return compactResult.type === "skipped" ? undefined : compactResult; + let resultMessages: ModelMessage[] | undefined; + + // 1. Compaction + if (taskCompaction) { + const compactResult = await chatCompact(messages, steps, { + shouldCompact: taskCompaction.shouldCompact, + summarize: (msgs) => { + const ctx = locals.get(chatTurnContextKey); + const lastStep = steps.at(-1); + return taskCompaction.summarize({ + messages: msgs, + usage: lastStep?.usage, + source: "inner", + stepNumber: steps.length - 1, + chatId: ctx?.chatId, + turn: ctx?.turn, + clientData: ctx?.clientData, + }); + }, + }); + if (compactResult.type !== "skipped") { + resultMessages = compactResult.messages; + } + } + + // 2. Pending message injection (steering) + if (taskPendingMessages) { + const injected = await drainSteeringQueue( + taskPendingMessages, + resultMessages ?? messages, + steps, + ); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? messages), ...injected]; + } + } + + return resultMessages ? { messages: resultMessages } : undefined; }; } @@ -1590,7 +1830,23 @@ export type ChatTaskOptions< compaction?: ChatTaskCompactionOptions; /** - * Called after the stream closes for this turn. Use this to persist the + * Configure how messages that arrive during streaming are handled. + * + * By default, messages queue for the next turn. When `shouldInject` is provided + * and returns `true`, messages are injected between tool-call steps via + * `prepareStep` — allowing users to steer the agent mid-execution. + * + * @example + * ```ts + * pendingMessages: { + * shouldInject: ({ steps }) => steps.length > 0, + * onReceived: ({ message }) => logger.info("Steering message received"), + * }, + * ``` + */ + pendingMessages?: PendingMessagesOptions; + + /** * conversation to your database after each assistant response. * * @example @@ -1756,6 +2012,7 @@ function chatTask< onBeforeTurnComplete, onCompacted, compaction, + pendingMessages: pendingMessagesConfig, prepareMessages, onTurnComplete, maxTurns = 100, @@ -1799,6 +2056,10 @@ function chatTask< locals.set(chatTaskCompactionKey, compaction); } + if (pendingMessagesConfig) { + locals.set(chatPendingMessagesKey, pendingMessagesConfig); + } + let currentWirePayload = payload; const continuation = payload.continuation ?? false; const previousRunId = payload.previousRunId; @@ -1948,6 +2209,8 @@ function chatTask< locals.set(chatPipeCountKey, 0); locals.set(chatDeferKey, new Set()); locals.set(chatCompactionStateKey, undefined); + locals.set(chatSteeringQueueKey, []); + locals.set(chatInjectedMessageIdsKey, new Set()); // Store chat context for auto-detection by ai.tool subtasks locals.set(chatTurnContextKey, { @@ -1969,7 +2232,39 @@ function chatTask< // Buffer messages that arrive during streaming const pendingMessages: ChatTaskWirePayload[] = []; - const msgSub = messagesInput.on((msg) => { + const pmConfig = locals.get(chatPendingMessagesKey); + const msgSub = messagesInput.on(async (msg) => { + // If pendingMessages is configured, route to the steering queue + // instead of the wire buffer. The frontend handles re-sending + // non-injected messages via sendMessage on turn complete. + if (pmConfig) { + const lastUIMessage = msg.messages?.[msg.messages.length - 1]; + if (lastUIMessage) { + if (pmConfig.onReceived) { + try { + await pmConfig.onReceived({ + message: lastUIMessage, + chatId: currentWirePayload.chatId, + turn, + }); + } catch { /* non-fatal */ } + } + + try { + const queue = locals.get(chatSteeringQueueKey) ?? []; + // Deduplicate by message ID — guards against double-sends + if (lastUIMessage.id && queue.some((e) => e.uiMessage.id === lastUIMessage.id)) { + return; + } + const modelMsgs = await toModelMessages([lastUIMessage]); + queue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); + locals.set(chatSteeringQueueKey, queue); + } catch { /* conversion failed — skip steering queue */ } + } + return; // Don't add to wire buffer — frontend handles non-injected case + } + + // No pendingMessages config — standard wire buffer for next turn pendingMessages.push(msg); }); @@ -2508,7 +2803,8 @@ function chatTask< ); } - // If messages arrived during streaming, use the first one immediately + // If messages arrived during streaming (without pendingMessages config), + // use the first one immediately as the next turn. if (pendingMessages.length > 0) { currentWirePayload = pendingMessages[0]!; return "continue"; @@ -2909,9 +3205,12 @@ class ChatMessageAccumulator { modelMessages: ModelMessage[] = []; uiMessages: UIMessage[] = []; private _compaction?: ChatTaskCompactionOptions; + private _pendingMessages?: PendingMessagesOptions; + private _steeringQueue: SteeringQueueEntry[] = []; - constructor(options?: { compaction?: ChatTaskCompactionOptions }) { + constructor(options?: { compaction?: ChatTaskCompactionOptions; pendingMessages?: PendingMessagesOptions }) { this._compaction = options?.compaction; + this._pendingMessages = options?.pendingMessages; } /** @@ -2965,19 +3264,73 @@ class ChatMessageAccumulator { } /** - * Returns a `prepareStep` function for inner-loop compaction. - * Only available when `compaction` was provided to the constructor. - * Pass the result to `streamText({ prepareStep: conversation.prepareStep() })`. + * Queue a message for injection via `prepareStep`. Call from a + * `messagesInput.on()` listener when a message arrives during streaming. + */ + steer(message: UIMessage, modelMessages?: ModelMessage[]): void { + if (modelMessages) { + this._steeringQueue.push({ uiMessage: message, modelMessages }); + } else { + // Defer conversion — will be done in prepareStep if needed + this._steeringQueue.push({ uiMessage: message, modelMessages: [] }); + } + } + + /** + * Queue a message for injection, converting to model messages automatically. + */ + async steerAsync(message: UIMessage): Promise { + const modelMsgs = await toModelMessages([message]); + this._steeringQueue.push({ uiMessage: message, modelMessages: modelMsgs }); + } + + /** + * Get and clear unconsumed steering messages. + */ + drainSteering(): UIMessage[] { + const result = this._steeringQueue.map((e) => e.uiMessage); + this._steeringQueue = []; + return result; + } + + /** + * Returns a `prepareStep` function that handles both compaction and + * pending message injection. Pass to `streamText({ prepareStep: conversation.prepareStep() })`. */ prepareStep(): ((args: { messages: ModelMessage[]; steps: CompactionStep[] }) => Promise<{ messages: ModelMessage[] } | undefined>) | undefined { - if (!this._compaction) return undefined; + if (!this._compaction && !this._pendingMessages) return undefined; const comp = this._compaction; + const pm = this._pendingMessages; + const queue = this._steeringQueue; + return async ({ messages, steps }) => { - const result = await chatCompact(messages, steps, { - shouldCompact: comp.shouldCompact, - summarize: (msgs) => comp.summarize({ messages: msgs, source: "inner" }), - }); - return result.type === "skipped" ? undefined : result; + let resultMessages: ModelMessage[] | undefined; + + // 1. Compaction + if (comp) { + const result = await chatCompact(messages, steps, { + shouldCompact: comp.shouldCompact, + summarize: (msgs) => comp.summarize({ messages: msgs, source: "inner" }), + }); + if (result.type !== "skipped") { + resultMessages = result.messages; + } + } + + // 2. Pending message injection + if (pm && queue.length > 0) { + const injected = await drainSteeringQueue( + pm, + resultMessages ?? messages, + steps, + queue, + ); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? messages), ...injected]; + } + } + + return resultMessages ? { messages: resultMessages } : undefined; }; } @@ -3058,6 +3411,8 @@ export type ChatSessionOptions = { maxTurns?: number; /** Automatic context compaction — same options as `chat.task({ compaction })`. */ compaction?: ChatTaskCompactionOptions; + /** Configure mid-execution message injection — same options as `chat.task({ pendingMessages })`. */ + pendingMessages?: PendingMessagesOptions; }; export type ChatTurn = { @@ -3108,6 +3463,13 @@ export type ChatTurn = { * Use with `chat.pipeAndCapture` when you need control between pipe and done. */ addResponse(response: UIMessage): Promise; + + /** + * Returns a `prepareStep` function that handles both compaction and + * pending message injection. Pass to `streamText({ prepareStep: turn.prepareStep() })`. + * Only needed when not using `chat.toStreamTextOptions()` (which auto-injects it). + */ + prepareStep(): ((args: { messages: ModelMessage[]; steps: CompactionStep[] }) => Promise<{ messages: ModelMessage[] } | undefined>) | undefined; }; /** @@ -3151,6 +3513,7 @@ function createChatSession( timeout = "1h", maxTurns = 100, compaction: sessionCompaction, + pendingMessages: sessionPendingMessages, } = options; return { @@ -3203,6 +3566,45 @@ function createChatSession( // Reset stop signal for this turn stop.reset(); + // Set up steering queue and pending messages config in locals + // so toStreamTextOptions() auto-injects prepareStep for steering + const turnSteeringQueue: SteeringQueueEntry[] = []; + locals.set(chatSteeringQueueKey, turnSteeringQueue); + if (sessionPendingMessages) { + locals.set(chatPendingMessagesKey, sessionPendingMessages); + } + locals.set(chatTurnContextKey, { + chatId: currentPayload.chatId, + turn, + continuation: currentPayload.continuation ?? false, + clientData: currentPayload.metadata, + }); + + // Listen for messages during streaming (steering + next-turn buffer) + const sessionPendingWire: ChatTaskWirePayload[] = []; + const sessionMsgSub = messagesInput.on(async (msg) => { + sessionPendingWire.push(msg); + + if (sessionPendingMessages) { + const lastUIMessage = msg.messages?.[msg.messages.length - 1]; + if (lastUIMessage) { + if (sessionPendingMessages.onReceived) { + try { + await sessionPendingMessages.onReceived({ + message: lastUIMessage, + chatId: currentPayload.chatId, + turn, + }); + } catch { /* non-fatal */ } + } + try { + const modelMsgs = await toModelMessages([lastUIMessage]); + turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); + } catch { /* non-fatal */ } + } + } + }); + // Accumulate messages const messages = await accumulator.addIncoming( currentPayload.messages, @@ -3237,6 +3639,7 @@ function createChatSession( if (error instanceof Error && error.name === "AbortError") { if (runSignal.aborted) { // Full cancel — don't accumulate + sessionMsgSub.off(); await chatWriteTurnComplete(); return undefined; } @@ -3310,6 +3713,7 @@ function createChatSession( } } + sessionMsgSub.off(); await chatWriteTurnComplete(); return response; }, @@ -3319,8 +3723,43 @@ function createChatSession( }, async done() { + sessionMsgSub.off(); await chatWriteTurnComplete(); }, + + prepareStep() { + const hasCompaction = !!sessionCompaction; + const hasPending = !!sessionPendingMessages; + if (!hasCompaction && !hasPending) return undefined; + + return async ({ messages: stepMsgs, steps }: { messages: ModelMessage[]; steps: CompactionStep[] }) => { + let resultMessages: ModelMessage[] | undefined; + + if (sessionCompaction) { + const compactResult = await chatCompact(stepMsgs, steps, { + shouldCompact: sessionCompaction.shouldCompact, + summarize: (msgs) => sessionCompaction.summarize({ messages: msgs, source: "inner" }), + }); + if (compactResult.type !== "skipped") { + resultMessages = compactResult.messages; + } + } + + if (sessionPendingMessages) { + const injected = await drainSteeringQueue( + sessionPendingMessages, + resultMessages ?? stepMsgs, + steps, + turnSteeringQueue, + ); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? stepMsgs), ...injected]; + } + } + + return resultMessages ? { messages: resultMessages } : undefined; + }; + }, }; return { done: false, value: turnObj }; diff --git a/packages/trigger-sdk/src/v3/chat-react.ts b/packages/trigger-sdk/src/v3/chat-react.ts index 612f0c184..ac8d2a388 100644 --- a/packages/trigger-sdk/src/v3/chat-react.ts +++ b/packages/trigger-sdk/src/v3/chat-react.ts @@ -23,13 +23,14 @@ * ``` */ -import { useEffect, useRef } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import { TriggerChatTransport, type TriggerChatTransportOptions, } from "./chat.js"; import type { AnyTask, TaskIdentifier } from "@trigger.dev/core/v3"; -import type { InferChatClientData } from "./ai.js"; +import { PENDING_MESSAGE_INJECTED_TYPE, type InferChatClientData } from "./ai.js"; +import type { UIMessage, ChatRequestOptions } from "ai"; /** * Options for `useTriggerChatTransport`, with a type-safe `task` field. @@ -93,3 +94,252 @@ export function useTriggerChatTransport( return ref.current; } + +// --------------------------------------------------------------------------- +// usePendingMessages — manage steering messages during streaming +// --------------------------------------------------------------------------- + +/** A pending message tracked by `usePendingMessages`. */ +export type PendingMessage = { + id: string; + text: string; + /** How this message is being handled. */ + mode: "steering" | "queued"; + /** Whether the backend confirmed this message was injected mid-response. */ + injected: boolean; +}; + +/** Options for `usePendingMessages`. */ +export type UsePendingMessagesOptions = { + /** The chat transport instance. */ + transport: TriggerChatTransport; + /** The chat session ID. */ + chatId: string; + /** The current useChat status. */ + status: string; + /** The current messages from useChat. */ + messages: UIMessage[]; + /** The setMessages function from useChat. */ + setMessages: (fn: UIMessage[] | ((prev: UIMessage[]) => UIMessage[])) => void; + /** The sendMessage function from useChat. */ + sendMessage: (message: { text: string }, options?: ChatRequestOptions) => void; + /** Metadata to include when sending (e.g. `{ model }` for model selection). */ + metadata?: Record; +}; + +/** A message embedded in an injection point data part. */ +export type InjectedMessage = { + id: string; + text: string; +}; + +/** Return value of `usePendingMessages`. */ +export type UsePendingMessagesReturn = { + /** Current pending messages with their mode and injection status. */ + pending: PendingMessage[]; + /** Send a steering message during streaming, or a normal message when ready. */ + steer: (text: string) => void; + /** Queue a message for the next turn (sent after current response finishes). */ + queue: (text: string) => void; + /** Promote a queued message to a steering message (sends via input stream immediately). */ + promoteToSteering: (id: string) => void; + /** Check if an assistant message part is an injection point. */ + isInjectionPoint: (part: unknown) => boolean; + /** Get the injected message IDs from an injection point part. */ + getInjectedMessageIds: (part: unknown) => string[]; + /** Get the injected messages (id + text) from an injection point part. Self-contained — works after turn complete. */ + getInjectedMessages: (part: unknown) => InjectedMessage[]; +}; + +/** + * React hook for managing pending messages (steering) during streaming. + * + * Handles: + * - Sending messages via input stream during streaming (bypassing useChat) + * - Tracking which messages were injected mid-response vs queued for next turn + * - Inserting injected messages into the conversation on turn complete + * - Auto-sending non-injected messages as the next turn + * + * @example + * ```tsx + * const pending = usePendingMessages({ + * transport, chatId, status, messages, setMessages, sendMessage, + * metadata: { model }, + * }); + * + * // In the form: + *
{ + * e.preventDefault(); + * pending.send(input); + * setInput(""); + * }}> + * + * // Render pending messages: + * {pending.pending.map(msg => ( + *
{msg.text} — {msg.injected ? "Injected" : "Pending"}
+ * ))} + * + * // Render injection points inline in assistant messages: + * {msg.parts.map((part, i) => + * pending.isInjectionPoint(part) + * ? + * : + * )} + * ``` + */ +export function usePendingMessages(options: UsePendingMessagesOptions): UsePendingMessagesReturn { + const { transport, chatId, status, messages, setMessages, sendMessage, metadata } = options; + + // Internal state: track messages with their mode + type InternalMessage = UIMessage & { _mode: "steering" | "queued" }; + const [pendingMsgs, setPendingMsgs] = useState([]); + const injectedIdsRef = useRef>(new Set()); + const prevStatusRef = useRef(status); + + // Watch for injection confirmation chunks in streaming messages + 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; + } + } + } + } + } + } + // Remove injected steering messages from the pending overlay immediately + if (newlyInjected) { + setPendingMsgs((prev) => prev.filter((m) => !injectedIdsRef.current.has(m.id))); + } + }, [status, messages]); + + // Handle turn completion + useEffect(() => { + const turnCompleted = prevStatusRef.current === "streaming" && status === "ready"; + prevStatusRef.current = status; + if (!turnCompleted) return; + + // Auto-send non-injected messages as the next turn. + // This includes queued messages AND steering messages that weren't + // injected (arrived too late, no prepareStep boundary, etc.). + // Note: steering messages were also sent via sendPendingMessage to + // the backend's wire buffer, so the backend may already have them. + // Calling sendMessage here ensures useChat subscribes to the response. + const toSend = pendingMsgs.filter( + (m) => !injectedIdsRef.current.has(m.id) + ); + + // Clean up + setPendingMsgs([]); + injectedIdsRef.current.clear(); + promotedIdsRef.current.clear(); + + // Auto-send as next turn + if (toSend.length > 0) { + const text = toSend.map((m) => (m.parts?.[0] as any)?.text ?? "").join("\n"); + sendMessage({ text }, metadata ? { metadata } : undefined); + } + }, [status, pendingMsgs, sendMessage, metadata, messages]); + + // Send a steering message (injected mid-response via prepareStep) + const steer = useCallback( + (text: string) => { + if (status === "streaming") { + const msg: InternalMessage = { + id: crypto.randomUUID(), + role: "user", + parts: [{ type: "text", text }], + _mode: "steering", + }; + transport.sendPendingMessage(chatId, msg, metadata); + setPendingMsgs((prev) => [...prev, msg]); + } else { + // Not streaming — just send normally + sendMessage({ text }, metadata ? { metadata } : undefined); + } + }, + [status, transport, chatId, sendMessage, metadata] + ); + + // Queue a message for the next turn (no injection attempt) + const queue = useCallback( + (text: string) => { + if (status === "streaming") { + const msg: InternalMessage = { + id: crypto.randomUUID(), + role: "user", + parts: [{ type: "text", text }], + _mode: "queued", + }; + setPendingMsgs((prev) => [...prev, msg]); + } else { + sendMessage({ text }, metadata ? { metadata } : undefined); + } + }, + [status, sendMessage, metadata] + ); + + // Promote a queued message to steering (send via input stream immediately) + const promotedIdsRef = useRef>(new Set()); + const promoteToSteering = useCallback( + (id: string) => { + // Guard against double-click — ref check is synchronous + if (promotedIdsRef.current.has(id)) { + console.log("[usePendingMessages] promote blocked — already promoted:", id); + return; + } + console.log("[usePendingMessages] promoting:", id); + promotedIdsRef.current.add(id); + + setPendingMsgs((prev) => { + const msg = prev.find((m) => m.id === id); + if (!msg || msg._mode !== "queued") return prev; + transport.sendPendingMessage(chatId, msg, metadata); + return prev.map((m) => m.id === id ? { ...m, _mode: "steering" as const } : m); + }); + }, + [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 pending: PendingMessage[] = pendingMsgs.map((m) => ({ + id: m.id, + text: (m.parts?.[0] as any)?.text ?? "", + mode: m._mode, + injected: injectedIdsRef.current.has(m.id), + })); + + return { pending, steer, queue, promoteToSteering, isInjectionPoint, getInjectedMessageIds, getInjectedMessages }; +} diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 820132abe..1a5037984 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -384,6 +384,51 @@ export class TriggerChatTransport implements ChatTransport { ); }; + /** + * Send a message to the running task via input stream without disrupting + * the current streaming response. Use this to send steering/pending messages + * while the agent is actively streaming. + * + * Unlike `sendMessage()` from useChat, this does NOT: + * - Add the message to useChat's local message state + * - Cancel the active stream subscription + * - Start a new response stream + * + * The message is delivered to the task's `messagesInput.on()` listener + * and can be injected between tool-call steps via the `pendingMessages` + * configuration. + * + * @returns `true` if the message was sent, `false` if there's no active session. + */ + sendPendingMessage = async ( + chatId: string, + message: UIMessage, + metadata?: Record, + ): Promise => { + const session = this.sessions.get(chatId); + if (!session?.runId) return false; + + const mergedMetadata = + this.defaultMetadata || metadata + ? { ...(this.defaultMetadata ?? {}), ...(metadata ?? {}) } + : undefined; + + const payload = { + messages: [message], + chatId, + trigger: "submit-message" as const, + metadata: mergedMetadata, + }; + + try { + const apiClient = new ApiClient(this.baseURL, session.publicAccessToken); + await apiClient.sendInputStream(session.runId, CHAT_MESSAGES_STREAM_ID, payload); + return true; + } catch { + return false; + } + }; + reconnectToStream = async ( options: { chatId: string; diff --git a/references/ai-chat/src/components/chat.tsx b/references/ai-chat/src/components/chat.tsx index ec837c74f..cedb282ab 100644 --- a/references/ai-chat/src/components/chat.tsx +++ b/references/ai-chat/src/components/chat.tsx @@ -4,6 +4,7 @@ import type { UIMessage } from "ai"; import { useChat } from "@ai-sdk/react"; import type { TriggerChatTransport } from "@trigger.dev/sdk/chat"; import type { CompactionChunkData } from "@trigger.dev/sdk/ai"; +import { usePendingMessages } from "@trigger.dev/sdk/chat/react"; import { useEffect, useRef, useState } from "react"; import { Streamdown } from "streamdown"; import { MODEL_OPTIONS } from "@/lib/models"; @@ -271,7 +272,7 @@ export function Chat({ const turnCounter = useRef(0); const [ttfbHistory, setTtfbHistory] = useState([]); - const { messages, sendMessage, stop, status, error } = useChat({ + const { messages, setMessages, sendMessage, stop, status, error } = useChat({ id: chatId, messages: initialMessages, transport, @@ -305,31 +306,96 @@ export function Chat({ } }, [status, messages]); - // Pending message to send after the current turn completes - const [pendingMessage, setPendingMessage] = useState(null); + // Pending messages — handles steering messages during streaming + const pending = usePendingMessages({ + transport, + chatId, + status, + messages, + setMessages, + sendMessage, + metadata: { model }, + }); - // Handle turn completion: persist messages and auto-send pending message + // Expose test helpers for automated testing via Chrome DevTools. + // All actions go through refs so closures always call the latest version. + const stateRef = useRef({ status, messages, pending: pending.pending }); + stateRef.current = { status, messages, pending: pending.pending }; + + const actionsRef = useRef({ + steer: pending.steer, + queue: pending.queue, + promote: pending.promoteToSteering, + send: (text: string) => { + turnCounter.current++; + sendTimestamp.current = Date.now(); + sendMessage({ text }, { metadata: { model } }); + }, + stop, + }); + actionsRef.current = { + steer: pending.steer, + queue: pending.queue, + promote: pending.promoteToSteering, + send: (text: string) => { + turnCounter.current++; + sendTimestamp.current = Date.now(); + sendMessage({ text }, { metadata: { model } }); + }, + stop, + }; + + useEffect(() => { + (window as any).__chat = { + get status() { return stateRef.current.status; }, + get messages() { return stateRef.current.messages; }, + get pending() { return stateRef.current.pending; }, + get runId() { return transport.getSession(chatId)?.runId ?? session?.runId ?? null; }, + chatId, + steer: (text: string) => actionsRef.current.steer(text), + queue: (text: string) => actionsRef.current.queue(text), + promote: (id: string) => actionsRef.current.promote(id), + send: (text: string) => actionsRef.current.send(text), + stop: () => actionsRef.current.stop(), + // Wait for a tool call to appear, then steer + steerOnToolCall: (text: string) => new Promise((resolve) => { + const check = setInterval(() => { + const { messages: msgs } = stateRef.current; + const lastMsg = msgs[msgs.length - 1]; + const hasTool = lastMsg?.role === "assistant" && lastMsg.parts?.some( + (p: any) => p.type?.startsWith("tool-") || p.type === "dynamic-tool" + ); + if (hasTool) { + clearInterval(check); + console.log("[__chat] steerOnToolCall: tool detected, steering now. status:", stateRef.current.status); + actionsRef.current.steer(text); + resolve(); + } + }, 200); + }), + // Wait for status to become a value + waitForStatus: (target: string) => new Promise((resolve) => { + const check = setInterval(() => { + if (stateRef.current.status === target) { clearInterval(check); resolve(); } + }, 100); + }), + steerAfterDelay: (text: string, ms: number) => new Promise((r) => setTimeout(() => { actionsRef.current.steer(text); r(); }, ms)), + queueAfterDelay: (text: string, ms: number) => new Promise((r) => setTimeout(() => { actionsRef.current.queue(text); r(); }, ms)), + }; + return () => { delete (window as any).__chat; }; + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [chatId]); + + // Persist messages when a turn completes const prevStatus = useRef(status); useEffect(() => { const turnCompleted = prevStatus.current === "streaming" && status === "ready"; prevStatus.current = status; - if (!turnCompleted) return; - - // Persist messages when a turn completes if (messages.length > 0) { onMessagesChange?.(chatId, messages); } - - // Auto-send the pending message - if (pendingMessage) { - const text = pendingMessage; - setPendingMessage(null); - turnCounter.current++; - sendTimestamp.current = Date.now(); - sendMessage({ text }, { metadata: { model } }); - } - }, [status, messages, chatId, onMessagesChange, sendMessage, pendingMessage, model]); + }, [status, messages, chatId, onMessagesChange]); return (
@@ -443,6 +509,25 @@ export function Chat({ return ; } + if (pending.isInjectionPoint(part)) { + const injectedMsgs = pending.getInjectedMessages(part); + if (injectedMsgs.length === 0) return null; + return ( +
+
+ {injectedMsgs.map((m) => ( +
+ {m.text} +
+ ))} +
+ injected mid-response +
+
+
+ ); + } + if (part.type.startsWith("data-")) { return (
)} - {pendingMessage && ( -
+ {pending.pending.map((msg) => ( +
-
- {pendingMessage} +
+ {msg.text}
-
- Queued — will send when current response finishes +
+ + {msg.mode === "steering" ? "Steering — waiting for injection point" : "Queued for next turn"} + + {msg.mode === "queued" && status === "streaming" && ( + + )}
- )} + ))}
{error && ( @@ -507,13 +605,11 @@ export function Chat({ onSubmit={(e) => { e.preventDefault(); if (!input.trim()) return; - if (status === "streaming") { - setPendingMessage(input); - } else { + if (status !== "streaming") { turnCounter.current++; sendTimestamp.current = Date.now(); - sendMessage({ text: input }, { metadata: { model } }); } + pending.steer(input); setInput(""); }} className="shrink-0 border-t border-gray-200 bg-white p-4" @@ -528,11 +624,25 @@ export function Chat({ /> + {status === "streaming" && ( + + )} {status === "streaming" && (