github-featured/triggerdotdev--trigger.dev
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat(chat): add compaction option, pendingMessages steering, and usePendingMessages hook

Browse Source
This commit is contained in:
Eric Allam
2026-03-25 14:25:25 +00:00
parent f7442f1312
commit f05d638717
5 changed files with 939 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/trigger-sdk/src/v3/ai.ts
+473 -34
View File
@@ -600,6 +600,102 @@ export type ChatTaskCompactionOptions = {
/** @internal */
const chatTaskCompactionKey = locals.create<ChatTaskCompactionOptions>("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<boolean>;
/**
* 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<ModelMessage[]>;
/** Called when a message arrives during streaming (per-message). */
onReceived?: (event: PendingMessageReceivedEvent) => void | Promise<void>;
/** Called after a batch of messages is injected via `prepareStep`. */
onInjected?: (event: PendingMessagesInjectedEvent) => void | Promise<void>;
};
/**
* 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<PendingMessagesOptions>("chat.pendingMessages");
/** @internal */
const chatSteeringQueueKey = locals.create<SteeringQueueEntry[]>("chat.steeringQueue");
/** @internal — IDs of messages that were successfully injected via prepareStep */
const chatInjectedMessageIdsKey = locals.create<Set<string>>("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<ModelMessage[]> {
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<strin
const telemetry = prompt.toAISDKTelemetry(options?.telemetry);
Object.assign(result, telemetry);
// Auto-inject prepareStep when task-level compaction is configured.
// We build a custom prepareStep instead of using chatCompactionStep so we
// can pass the enriched SummarizeEvent to taskCompaction.summarize.
// Auto-inject prepareStep when compaction or pendingMessages is configured.
const taskCompaction = locals.get(chatTaskCompactionKey);
if (taskCompaction) {
const taskPendingMessages = locals.get(chatPendingMessagesKey);
if (taskCompaction || taskPendingMessages) {
result.prepareStep = async ({ messages, steps }: { messages: ModelMessage[]; steps: CompactionStep[] }) => {
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<void> {
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<void>;
/**
* 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 };
packages/trigger-sdk/src/v3/chat-react.ts
+252 -2
View File
@@ -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<TTask extends AnyTask = AnyTask>(
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<string, unknown>;
};
/** 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:
* <form onSubmit={(e) => {
* e.preventDefault();
* pending.send(input);
* setInput("");
* }}>
*
* // Render pending messages:
* {pending.pending.map(msg => (
* <div key={msg.id}>{msg.text} — {msg.injected ? "Injected" : "Pending"}</div>
* ))}
*
* // Render injection points inline in assistant messages:
* {msg.parts.map((part, i) =>
* pending.isInjectionPoint(part)
* ? <InjectionMarker key={i} ids={pending.getInjectedMessageIds(part)} />
* : <Part key={i} part={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<InternalMessage[]>([]);
const injectedIdsRef = useRef<Set<string>>(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<Set<string>>(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 };
}
packages/trigger-sdk/src/v3/chat.ts
+45
View File
@@ -384,6 +384,51 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
);
};
/**
* 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<string, unknown>,
): Promise<boolean> => {
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;
references/ai-chat/src/components/chat.tsx
+140 -30
View File
@@ -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<TtfbEntry[]>([]);
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<string | null>(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<void>((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<void>((resolve) => {
const check = setInterval(() => {
if (stateRef.current.status === target) { clearInterval(check); resolve(); }
}, 100);
}),
steerAfterDelay: (text: string, ms: number) => new Promise<void>((r) => setTimeout(() => { actionsRef.current.steer(text); r(); }, ms)),
queueAfterDelay: (text: string, ms: number) => new Promise<void>((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 (
<div className="flex h-full flex-col bg-white">
@@ -443,6 +509,25 @@ export function Chat({
return <ToolInvocation key={i} part={part} />;
}
if (pending.isInjectionPoint(part)) {
const injectedMsgs = pending.getInjectedMessages(part);
if (injectedMsgs.length === 0) return null;
return (
<div key={i} className="my-2 flex justify-end">
<div className="max-w-[60%]">
{injectedMsgs.map((m) => (
<div key={m.id} className="rounded-lg bg-purple-100 px-3 py-1.5 text-sm text-purple-800">
{m.text}
</div>
))}
<div className="mt-0.5 text-right text-[10px] text-purple-400">
injected mid-response
</div>
</div>
</div>
);
}
if (part.type.startsWith("data-")) {
return (
<div
@@ -472,18 +557,31 @@ export function Chat({
</div>
)}
{pendingMessage && (
<div className="flex justify-end">
{pending.pending.map((msg) => (
<div key={msg.id} className="flex justify-end">
<div className="max-w-[80%]">
<div className="rounded-lg bg-blue-600 px-4 py-2 text-sm text-white opacity-60">
{pendingMessage}
<div className={`rounded-lg px-4 py-2 text-sm text-white opacity-75 ${
msg.mode === "steering" ? "bg-purple-600" : "bg-gray-500"
}`}>
{msg.text}
</div>
<div className="mt-1 text-right text-[10px] text-gray-400">
Queued — will send when current response finishes
<div className="mt-1 flex items-center justify-end gap-2">
<span className="text-[10px] text-gray-400">
{msg.mode === "steering" ? "Steering — waiting for injection point" : "Queued for next turn"}
</span>
{msg.mode === "queued" && status === "streaming" && (
<button
type="button"
onClick={() => pending.promoteToSteering(msg.id)}
className="text-[10px] text-purple-500 hover:text-purple-700 underline"
>
Steer instead
</button>
)}
</div>
</div>
</div>
)}
))}
</div>
{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({
/>
<button
type="submit"
disabled={!input.trim() || !!pendingMessage}
disabled={!input.trim()}
className="rounded-lg bg-blue-600 px-4 py-2 text-sm font-medium text-white hover:bg-blue-700 disabled:opacity-50"
>
{status === "streaming" ? "Queue" : "Send"}
Send
</button>
{status === "streaming" && (
<button
type="button"
disabled={!input.trim()}
onClick={() => {
if (!input.trim()) return;
pending.queue(input);
setInput("");
}}
className="rounded-lg bg-gray-500 px-4 py-2 text-sm font-medium text-white hover:bg-gray-600 disabled:opacity-50"
>
Queue
</button>
)}
{status === "streaming" && (
<button
type="button"
references/ai-chat/src/trigger/chat.ts
+29
View File
@@ -280,6 +280,15 @@ export const aiChat = chat.task({
];
},
},
pendingMessages: {
// Inject user messages between tool-call steps so the agent can adjust
shouldInject: ({ steps }) => steps.length > 0,
prepare: ({ messages }) => messages.length === 1
? [{ role: "user" as const, content: (messages[0]!.parts?.[0] as any)?.text ?? "" }]
: [{ role: "user" as const, content: `The user sent ${messages.length} messages while you were working:\n\n${messages.map((m, i) => `${i + 1}. ${(m.parts?.[0] as any)?.text ?? ""}`).join("\n")}` }],
// onReceived/onInjected are optional — the SDK automatically writes
// a data-pending-message-injected chunk when injection happens.
},
prepareMessages: ({ messages, reason }) => {
// Add Anthropic cache breaks to the last message for prompt caching.
// Applied everywhere — run(), compaction rebuilds, compaction results.
@@ -585,6 +594,14 @@ export const aiChatRaw = task({
},
],
},
pendingMessages: {
// Inject with a prefix so the LLM knows these are mid-execution corrections
shouldInject: () => true,
prepare: ({ messages }) => [{
role: "user" as const,
content: [{ type: "text" as const, text: `[User sent ${messages.length} message(s) while you were working]:\n${messages.map(m => (m.parts?.[0] as any)?.text ?? "").join("\n")}` }],
}],
},
});
for (let turn = 0; turn < 100; turn++) {
@@ -623,6 +640,12 @@ export const aiChatRaw = task({
});
}
// Listen for steering messages during streaming
const steeringSub = chat.messages.on(async (msg) => {
const lastMsg = msg.messages?.[msg.messages.length - 1];
if (lastMsg) await conversation.steerAsync(lastMsg);
});
const result = streamText({
...chat.toStreamTextOptions({ registry }),
...(modelOverride ? { model: getModel(modelOverride) } : {}),
@@ -655,6 +678,8 @@ export const aiChatRaw = task({
} else {
throw error;
}
} finally {
steeringSub.off();
}
if (response) {
@@ -744,6 +769,10 @@ export const aiChatSession = task({
...uiMessages.slice(-4),
],
},
pendingMessages: {
// Always inject in the session variant
shouldInject: () => true,
},
});
for await (const turn of session) {