From 843777fbd472114b3812733c64be9df1fc9d4568 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 20:26:17 -0700 Subject: [PATCH 01/20] feat(chat): add custom agent mailbox helpers --- .changeset/tidy-mailboxes-wait.md | 6 + docs/ai-chat/custom-agents.mdx | 42 ++- docs/ai-chat/reference.mdx | 2 +- .../core/src/v3/apiClient/runStream.test.ts | 2 + packages/core/src/v3/apiClient/runStream.ts | 4 + packages/core/src/v3/sessionStreams/index.ts | 49 +++- .../src/v3/sessionStreams/manager.test.ts | 124 ++++++++- .../core/src/v3/sessionStreams/manager.ts | 223 ++++++++++------ .../core/src/v3/sessionStreams/noopManager.ts | 40 ++- packages/core/src/v3/sessionStreams/types.ts | 44 ++++ .../core/src/v3/test/mock-task-context.ts | 11 +- .../v3/test/test-session-stream-manager.ts | 200 +++++++++++---- packages/trigger-sdk/src/v3/ai.ts | 57 ++++- .../test/chat-messages-mailbox.test.ts | 242 ++++++++++++++++++ .../trigger-sdk/test/mockChatAgent.test.ts | 9 +- 15 files changed, 913 insertions(+), 142 deletions(-) create mode 100644 .changeset/tidy-mailboxes-wait.md create mode 100644 packages/trigger-sdk/test/chat-messages-mailbox.test.ts diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md new file mode 100644 index 00000000000..99013758d35 --- /dev/null +++ b/.changeset/tidy-mailboxes-wait.md @@ -0,0 +1,6 @@ +--- +"@trigger.dev/core": patch +"@trigger.dev/sdk": patch +--- + +Custom agent loops can now inspect pending chat input without consuming it and consume one mailbox record at a time with `chat.messages.hasPending()` and `chat.messages.next()`. Mailbox records include stable identifiers for tracing and redelivery. diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 197bff6b5e1..e7ed3f1e5b2 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -213,7 +213,7 @@ For full control, skip `createSession` and compose the primitives directly: | Primitive | Description | | ------------------------------- | -------------------------------------------------------------------------------------------- | -| `chat.messages` | Input stream for incoming messages — use `.waitWithIdleTimeout()` to wait for the next turn | +| `chat.messages` | Mailbox for incoming messages — inspect buffered input, consume one record, or suspend until the next turn | | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | | `chat.pipeAndCapture(result)` | Pipe a stream and capture the response; returns `{ message, status, error }` | | `chat.writeTurnComplete()` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | @@ -221,6 +221,46 @@ For full control, skip `createSession` and compose the primitives directly: | `chat.pipe(stream)` | Pipe a stream to the frontend (no response capture) | | `chat.cleanupAbortedParts(msg)` | Clean up incomplete parts from a stopped response | +### `chat.messages` mailbox + +`chat.messages` exposes the incoming message mailbox for hand-rolled loops: + +| Method | Behavior | +| --- | --- | +| `peek()` | Return the buffer head when it is a message, without consuming it; otherwise return `undefined` | +| `hasPending()` | Resolve `true` when any message is buffered; does not consume it | +| `next({ timeoutInSeconds? })` | Consume exactly one message record in channel order, or resolve `undefined` when the optional timeout elapses | +| `on(handler)` | Consume messages as they arrive and invoke the handler | +| `waitWithIdleTimeout(options)` | Wait warm, then suspend the run until the next message arrives | + +`hasPending()` checks the local, already-delivered buffer. It does not query the +remote Session channel or start a subscription. Use `waitWithIdleTimeout()` when +the loop needs to idle until future input arrives. + +`next()` returns a readonly record envelope: + +```ts +const record = await chat.messages.next({ timeoutInSeconds: 5 }); +if (record) { + console.log(record.id, record.seqNum); + currentPayload = record.payload; +} +``` + +- `id` is the append's stable idempotency key. +- `seqNum` is the monotonic sequence on this Session's `.in` channel. +- `payload` is the existing `ChatTaskWirePayload` delivered by the other mailbox methods. + +Both identifiers remain the same if the record is delivered again after a +reconnect. Each `next()` call commits only the record it returns, so a loop that +owns its own turn sequencing never advances past input it has not taken. By +contrast, `on()` commits a record as soon as it dispatches the handler; avoid +mixing `on()` and `next()` when a single loop owns mailbox consumption. + +The Session `.in` channel also carries control records such as handovers. If one +comes before a message, `next()` leaves it for its own consumer and waits until +that record has been handled. + A complete loop: ```ts trigger/my-chat-raw.ts diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index da08f9a0473..65e98a05aa4 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -506,7 +506,7 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`. | `chat.pipeAndCapture(source, options?)` | Pipe and capture the response; returns `{ message, status, error }` | | `chat.writeTurnComplete(options?)` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | -| `chat.messages` | Input stream for incoming messages — use `.waitWithIdleTimeout()` | +| `chat.messages` | Incoming message mailbox; supports non-consuming `.peek()` / `.hasPending()`, single-record `.next()`, `.on()`, and suspend-aware `.waitWithIdleTimeout()` | | `chat.local({ id })` | Create a per-run typed local (see [`chat.local`](/ai-chat/chat-local)) | | `chat.createStartSessionAction(taskId, options?)` | Returns a server action that creates a chat Session + triggers the first run + returns a session-scoped PAT. Idempotent on `(env, externalId)`. | | `chat.waitForHandover(options)` | Wait for a [`chat.headStart`](/ai-chat/fast-starts#handover-with-custom-agents) handover signal in a custom loop. Returns the signal or `null`. `chat.MessageAccumulator` wraps this as `consumeHandover()` / `applyHandover()` | diff --git a/packages/core/src/v3/apiClient/runStream.test.ts b/packages/core/src/v3/apiClient/runStream.test.ts index 3a266f2a918..ee3f3df22a6 100644 --- a/packages/core/src/v3/apiClient/runStream.test.ts +++ b/packages/core/src/v3/apiClient/runStream.test.ts @@ -492,6 +492,7 @@ describe("SSEStreamSubscription v2 batch parsing — record kinds", () => { }); type ParsedPart = { + recordId?: string; id: string; chunk: unknown; headers?: ReadonlyArray; @@ -548,6 +549,7 @@ describe("SSEStreamSubscription v2 batch parsing — record kinds", () => { const parts = await sub.subscribe().then(drain); expect(parts).toHaveLength(1); + expect(parts[0]!.recordId).toBe("p1"); expect(parts[0]!.id).toBe("5"); expect(parts[0]!.chunk).toEqual({ type: "text-delta", delta: "hi" }); expect(parts[0]!.headers).toEqual([]); diff --git a/packages/core/src/v3/apiClient/runStream.ts b/packages/core/src/v3/apiClient/runStream.ts index b0d43ef3f99..b01fc6e9643 100644 --- a/packages/core/src/v3/apiClient/runStream.ts +++ b/packages/core/src/v3/apiClient/runStream.ts @@ -170,6 +170,9 @@ export interface StreamSubscriptionFactory { } export type SSEStreamPart = { + /** Stable logical record id from the S2 data envelope (`X-Part-Id` on append). */ + recordId?: string; + /** S2 sequence number in decimal-string form. */ id: string; chunk: TChunk; timestamp: number; @@ -502,6 +505,7 @@ export class SSEStreamSubscription implements StreamSubscription { chunkController.enqueue({ type: "part", part: { + recordId: parsedBody?.id, id: record.seq_num.toString(), chunk: parsedBody?.data, timestamp: record.timestamp, diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 21e2e8d2450..073bb2ac514 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -1,6 +1,12 @@ import { getGlobal, registerGlobal } from "../utils/globals.js"; import { NoopSessionStreamManager } from "./noopManager.js"; -import type { InputStreamOncePromise, SessionChannelIO, SessionStreamManager } from "./types.js"; +import type { + InputStreamOncePromise, + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "./types.js"; import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; const API_NAME = "session-streams"; @@ -43,10 +49,51 @@ export class SessionStreamsAPI implements SessionStreamManager { return this.#getManager().once(sessionId, io, options); } + public onceRecord( + sessionId: string, + io: SessionChannelIO, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { + const manager = this.#getManager(); + if (!manager.onceRecord) { + throw new Error("The configured Session stream manager does not support record metadata"); + } + return manager.onceRecord(sessionId, io, options); + } + + public onceRecordWhere( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { + const manager = this.#getManager(); + if (!manager.onceRecordWhere) { + throw new Error("The configured Session stream manager does not support selective records"); + } + return manager.onceRecordWhere(sessionId, io, predicate, options); + } + public peek(sessionId: string, io: SessionChannelIO): unknown | undefined { return this.#getManager().peek(sessionId, io); } + public peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined { + return this.#getManager().peekRecord?.(sessionId, io); + } + + public peekRecordWhere( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate + ): SessionStreamRecord | undefined { + const manager = this.#getManager(); + if (!manager.peekRecordWhere) { + throw new Error("The configured Session stream manager does not support selective records"); + } + return manager.peekRecordWhere(sessionId, io, predicate); + } + public lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.#getManager().lastSeqNum(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.test.ts b/packages/core/src/v3/sessionStreams/manager.test.ts index 9b489616f74..29693e60fce 100644 --- a/packages/core/src/v3/sessionStreams/manager.test.ts +++ b/packages/core/src/v3/sessionStreams/manager.test.ts @@ -11,7 +11,7 @@ import type { SSEStreamPart } from "../apiClient/runStream.js"; // an empty stream synchronously triggers a tight reconnect loop, so the // mock parks indefinitely instead. function singleShotApiClient( - records: Array<{ id: string; chunk: unknown; timestamp: number }> + records: Array<{ id: string; recordId?: string; chunk: unknown; timestamp: number }> ): ApiClient { let delivered = false; return { @@ -160,3 +160,125 @@ describe("StandardSessionStreamManager — minTimestamp filter", () => { manager.disconnect(); }); }); + +describe("StandardSessionStreamManager — record metadata", () => { + const sessionId = "session-records"; + const io = "in" as const; + const records = [ + { + id: "41", + recordId: "part-stable-1", + chunk: { kind: "message", payload: { id: "u1" } }, + timestamp: 1000, + }, + { + id: "42", + recordId: "part-stable-2", + chunk: { kind: "message", payload: { id: "u2" } }, + timestamp: 2000, + }, + ]; + + it("consumes one record at a time with stable id and sequence metadata", async () => { + const manager = new StandardSessionStreamManager( + singleShotApiClient(records), + "http://localhost" + ); + + const first = await manager.onceRecord(sessionId, io); + expect(first).toEqual({ + ok: true, + output: { + id: "part-stable-1", + seqNum: 41, + data: { kind: "message", payload: { id: "u1" } }, + }, + }); + expect(manager.peekRecord(sessionId, io)).toEqual({ + id: "part-stable-2", + seqNum: 42, + data: { kind: "message", payload: { id: "u2" } }, + }); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(41); + + const second = await manager.onceRecord(sessionId, io); + expect(second.ok && second.output.id).toBe("part-stable-2"); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(42); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); + + it("returns the same envelope when a record is redelivered", async () => { + const firstDelivery = new StandardSessionStreamManager( + singleShotApiClient([records[0]!]), + "http://localhost" + ); + const redelivery = new StandardSessionStreamManager( + singleShotApiClient([records[0]!]), + "http://localhost" + ); + + const first = await firstDelivery.onceRecord(sessionId, io); + const replayed = await redelivery.onceRecord(sessionId, io); + + expect(first).toEqual(replayed); + + firstDelivery.disconnectStream(sessionId, io); + firstDelivery.disconnect(); + redelivery.disconnectStream(sessionId, io); + redelivery.disconnect(); + }); + + it("does not consume a matching record past an earlier unmatched record", async () => { + const manager = new StandardSessionStreamManager( + singleShotApiClient([ + { + id: "50", + recordId: "handover-1", + chunk: { kind: "handover" }, + timestamp: 1000, + }, + { + id: "51", + recordId: "message-1", + chunk: { kind: "message", payload: { id: "u1" } }, + timestamp: 2000, + }, + ]), + "http://localhost" + ); + + const pendingMessage = manager.onceRecordWhere( + sessionId, + io, + (record) => (record.data as { kind?: string }).kind === "message", + { timeoutMs: 200 } + ); + + expect(manager.peekRecordWhere(sessionId, io, (record) => record.id === "message-1")).toEqual({ + id: "message-1", + seqNum: 51, + data: { kind: "message", payload: { id: "u1" } }, + }); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBeUndefined(); + + const handover = await manager.onceRecord(sessionId, io); + expect(handover).toEqual({ + ok: true, + output: { id: "handover-1", seqNum: 50, data: { kind: "handover" } }, + }); + await expect(pendingMessage).resolves.toEqual({ + ok: true, + output: { + id: "message-1", + seqNum: 51, + data: { kind: "message", payload: { id: "u1" } }, + }, + }); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(51); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); +}); diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index fb87b211643..184d65d2443 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -3,7 +3,12 @@ import type { InputStreamOnceResult } from "../inputStreams/types.js"; import { InputStreamOncePromise, InputStreamTimeoutError } from "../inputStreams/types.js"; import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; import { computeReconnectDelayMs } from "../utils/reconnectBackoff.js"; -import type { SessionChannelIO, SessionStreamManager } from "./types.js"; +import type { + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "./types.js"; import { controlSubtype } from "./wireProtocol.js"; // A handler that synchronously returns `true` CONSUMES the record: it is @@ -13,8 +18,9 @@ import { controlSubtype } from "./wireProtocol.js"; type SessionStreamHandler = (data: unknown) => void | boolean | Promise; type OnceWaiter = { - resolve: (result: InputStreamOnceResult) => void; + resolve: (result: InputStreamOnceResult) => void; reject: (error: Error) => void; + predicate?: SessionStreamRecordPredicate; timeoutHandle?: ReturnType; // The abort signal and its handler are tracked on the waiter so any // resolution path (dispatch / timeout / explicit removal) can detach @@ -44,19 +50,7 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { export class StandardSessionStreamManager implements SessionStreamManager { private handlers = new Map>(); private onceWaiters = new Map(); - private buffer = new Map(); - // Parallel to `buffer`: the SSE seq_num of each buffered record. Same - // length and order as `buffer[key]`. Used so that when `once()` shifts - // a buffered record into a waiter, the cursor (`lastDispatchedSeqNums`) - // can advance to that record's seq. Kept as a separate map so the - // existing `peek()` shape (returns `unknown`) stays unchanged. - // - // Entries are `number | undefined` so the array stays length-locked - // with `buffer` even if a record arrives without a parseable seq — - // shifting `undefined` is just a no-op for the cursor advance, but - // the slot still gets consumed. Drifting lengths would map seq_nums - // to the wrong records on subsequent shifts. - private bufferSeqNums = new Map>(); + private buffer = new Map(); private tails = new Map(); // Per-stream lower-bound timestamp filter. When set, records whose // SSE timestamp is <= the bound are dropped before dispatch — used by @@ -123,28 +117,21 @@ export class StandardSessionStreamManager implements SessionStreamManager { // duplicating turns. const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - const seqList = this.bufferSeqNums.get(key) ?? []; - const keptRecords: unknown[] = []; - // Kept in lock-step with `keptRecords` — drifting lengths would map - // seq_nums to the wrong records on subsequent shifts. - const keptSeqNums: Array = []; - for (let i = 0; i < buffered.length; i++) { - const consumed = this.#invokeHandler(handler, buffered[i]); + const keptRecords: SessionStreamRecord[] = []; + for (const record of buffered) { + const consumed = this.#invokeHandler(handler, record.data); if (consumed) { - const s = seqList[i]; - if (s !== undefined) this.#advanceLastDispatched(key, s); + this.#advanceLastDispatched(key, record.seqNum); } else { - keptRecords.push(buffered[i]); - keptSeqNums.push(seqList[i]); + keptRecords.push(record); } } if (keptRecords.length > 0) { this.buffer.set(key, keptRecords); - this.bufferSeqNums.set(key, keptSeqNums); } else { this.buffer.delete(key); - this.bufferSeqNums.delete(key); } + this.#drainOnceWaitersFromBuffer(key); } return { @@ -162,6 +149,37 @@ export class StandardSessionStreamManager implements SessionStreamManager { io: SessionChannelIO, options?: InputStreamOnceOptions ): InputStreamOncePromise { + const recordPromise = this.onceRecord(sessionId, io, options); + return new InputStreamOncePromise((resolve, reject) => { + recordPromise.then((result) => { + resolve(result.ok ? { ok: true, output: result.output.data } : result); + }, reject); + }); + } + + onceRecord( + sessionId: string, + io: SessionChannelIO, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return this.#onceRecord(sessionId, io, undefined, options); + } + + onceRecordWhere( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return this.#onceRecord(sessionId, io, predicate, options); + } + + #onceRecord( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { const key = keyFor(sessionId, io); this.explicitlyDisconnected.delete(key); @@ -169,23 +187,30 @@ export class StandardSessionStreamManager implements SessionStreamManager { const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - const data = buffered.shift()!; - const seqList = this.bufferSeqNums.get(key); - const shiftedSeqNum = seqList?.shift(); - if (buffered.length === 0) { - this.buffer.delete(key); - this.bufferSeqNums.delete(key); - } - if (shiftedSeqNum !== undefined) { - this.#advanceLastDispatched(key, shiftedSeqNum); + const record = buffered[0]!; + if (!predicate || predicate(record)) { + buffered.shift(); + if (buffered.length === 0) { + this.buffer.delete(key); + } + this.#advanceLastDispatched(key, record.seqNum); + this.#drainOnceWaitersFromBuffer(key); + return new InputStreamOncePromise((resolve) => { + resolve({ ok: true, output: record }); + }); } - return new InputStreamOncePromise((resolve) => { - resolve({ ok: true, output: data }); - }); } - return new InputStreamOncePromise((resolve, reject) => { - const waiter: OnceWaiter = { resolve, reject }; + return new InputStreamOncePromise((resolve, reject) => { + const waiter: OnceWaiter = { resolve, reject, predicate }; + + if (predicate && options?.timeoutMs === 0) { + resolve({ + ok: false, + error: new InputStreamTimeoutError(key, 0), + }); + return; + } if (options?.signal) { if (options.signal.aborted) { @@ -222,9 +247,19 @@ export class StandardSessionStreamManager implements SessionStreamManager { } peek(sessionId: string, io: SessionChannelIO): unknown | undefined { - const buffered = this.buffer.get(keyFor(sessionId, io)); - if (buffered && buffered.length > 0) return buffered[0]; - return undefined; + return this.peekRecord(sessionId, io)?.data; + } + + peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined { + return this.buffer.get(keyFor(sessionId, io))?.[0]; + } + + peekRecordWhere( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate + ): SessionStreamRecord | undefined { + return this.buffer.get(keyFor(sessionId, io))?.find(predicate); } lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { @@ -248,6 +283,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { } #advanceLastDispatched(key: string, seqNum: number): void { + if (!Number.isFinite(seqNum)) return; const current = this.lastDispatchedSeqNums.get(key); if (current === undefined || seqNum > current) { this.lastDispatchedSeqNums.set(key, seqNum); @@ -267,16 +303,12 @@ export class StandardSessionStreamManager implements SessionStreamManager { const key = keyFor(sessionId, io); const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - buffered.shift(); - const seqList = this.bufferSeqNums.get(key); - const shiftedSeqNum = seqList?.shift(); + const record = buffered.shift()!; if (buffered.length === 0) { this.buffer.delete(key); - this.bufferSeqNums.delete(key); - } - if (shiftedSeqNum !== undefined) { - this.#advanceLastDispatched(key, shiftedSeqNum); } + this.#advanceLastDispatched(key, record.seqNum); + this.#drainOnceWaitersFromBuffer(key); return true; } return false; @@ -296,7 +328,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.tails.delete(key); } this.buffer.delete(key); - this.bufferSeqNums.delete(key); // Reset the backoff counter so a future re-attach starts fresh — // an explicit disconnect is a deliberate teardown, not evidence of // a broken backend. @@ -350,7 +381,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { } this.onceWaiters.clear(); this.buffer.clear(); - this.bufferSeqNums.clear(); } #ensureTailConnected(sessionId: string, io: SessionChannelIO): void { @@ -454,7 +484,11 @@ export class StandardSessionStreamManager implements SessionStreamManager { // keep as string } } - this.#dispatch(key, data, Number.isFinite(seqNum) ? seqNum : undefined); + this.#dispatch(key, { + id: part.recordId ?? part.id, + seqNum, + data, + }); }, onComplete: () => { if (this.debug) { @@ -479,27 +513,21 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } - #dispatch(key: string, data: unknown, seqNum: number | undefined): void { + #dispatch(key: string, record: SessionStreamRecord): void { // Any record flowing through = healthy connection; reset the backoff // counter so the next disconnect starts fresh. this.reconnectAttempts.delete(key); - const waiters = this.onceWaiters.get(key); - if (waiters && waiters.length > 0) { - const waiter = waiters.shift()!; - if (waiters.length === 0) this.onceWaiters.delete(key); - if (waiter.timeoutHandle) clearTimeout(waiter.timeoutHandle); - if (waiter.signal && waiter.abortHandler) { - waiter.signal.removeEventListener("abort", waiter.abortHandler); - } + const existingBuffer = this.buffer.get(key); + const waiter = + existingBuffer && existingBuffer.length > 0 ? undefined : this.#takeOnceWaiter(key, record); + if (waiter) { // Record was consumed directly by a waiter — advance the // committed-consume cursor immediately. Buffered-then-shifted // records advance the cursor in `once()` / `shiftBuffer()`. - if (seqNum !== undefined) { - this.#advanceLastDispatched(key, seqNum); - } - waiter.resolve({ ok: true, output: data }); - this.#invokeHandlers(key, data); + this.#advanceLastDispatched(key, record.seqNum); + waiter.resolve({ ok: true, output: record }); + this.#invokeHandlers(key, record.data); return; } @@ -511,11 +539,9 @@ export class StandardSessionStreamManager implements SessionStreamManager { // second turn. Records no handler consumed (e.g. a message arriving // while only the stop facade is attached during preload) are buffered // so a subsequent `once()` can still pick them up. - const consumed = this.#invokeHandlers(key, data); + const consumed = this.#invokeHandlers(key, record.data); if (consumed) { - if (seqNum !== undefined) { - this.#advanceLastDispatched(key, seqNum); - } + this.#advanceLastDispatched(key, record.seqNum); return; } @@ -524,17 +550,48 @@ export class StandardSessionStreamManager implements SessionStreamManager { buffered = []; this.buffer.set(key, buffered); } - buffered.push(data); - let bufferedSeqs = this.bufferSeqNums.get(key); - if (!bufferedSeqs) { - bufferedSeqs = []; - this.bufferSeqNums.set(key, bufferedSeqs); + buffered.push(record); + this.#drainOnceWaitersFromBuffer(key); + } + + #takeOnceWaiter(key: string, record: SessionStreamRecord): OnceWaiter | undefined { + const waiters = this.onceWaiters.get(key); + if (!waiters) return undefined; + + const index = waiters.findIndex((waiter) => { + if (!waiter.predicate) return true; + try { + return waiter.predicate(record); + } catch (error) { + if (this.debug) { + console.error("[SessionStreamManager] Record predicate error:", error); + } + return false; + } + }); + if (index === -1) return undefined; + + const [waiter] = waiters.splice(index, 1); + if (waiters.length === 0) this.onceWaiters.delete(key); + if (waiter!.timeoutHandle) clearTimeout(waiter!.timeoutHandle); + if (waiter!.signal && waiter!.abortHandler) { + waiter!.signal.removeEventListener("abort", waiter!.abortHandler); + } + return waiter; + } + + #drainOnceWaitersFromBuffer(key: string): void { + const buffered = this.buffer.get(key); + while (buffered && buffered.length > 0) { + const record = buffered[0]!; + const waiter = this.#takeOnceWaiter(key, record); + if (!waiter) return; + + buffered.shift(); + if (buffered.length === 0) this.buffer.delete(key); + this.#advanceLastDispatched(key, record.seqNum); + waiter.resolve({ ok: true, output: record }); } - // Always push, even when `seqNum` is undefined (e.g. NaN from a - // malformed `part.id`). Skipping the push here would drift the two - // arrays apart and misattribute seq_nums to records on the next - // shift. - bufferedSeqs.push(seqNum); } /** Returns true when any handler consumed the record. All handlers are invoked regardless. */ diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index f2d355d24ef..1a68c36e7e7 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -1,6 +1,11 @@ import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; import { InputStreamOncePromise } from "../inputStreams/types.js"; -import type { SessionChannelIO, SessionStreamManager } from "./types.js"; +import type { + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "./types.js"; export class NoopSessionStreamManager implements SessionStreamManager { on( @@ -21,10 +26,43 @@ export class NoopSessionStreamManager implements SessionStreamManager { }); } + onceRecord( + _sessionId: string, + _io: SessionChannelIO, + _options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return new InputStreamOncePromise(() => { + // Never resolves in noop mode. + }); + } + + onceRecordWhere( + _sessionId: string, + _io: SessionChannelIO, + _predicate: SessionStreamRecordPredicate, + _options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return new InputStreamOncePromise(() => { + // Never resolves in noop mode. + }); + } + peek(_sessionId: string, _io: SessionChannelIO): unknown | undefined { return undefined; } + peekRecord(_sessionId: string, _io: SessionChannelIO): SessionStreamRecord | undefined { + return undefined; + } + + peekRecordWhere( + _sessionId: string, + _io: SessionChannelIO, + _predicate: SessionStreamRecordPredicate + ): SessionStreamRecord | undefined { + return undefined; + } + lastSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { return undefined; } diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index ae24259b3fc..d80f8d6cec3 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -12,6 +12,21 @@ export type { InputStreamOnceResult }; export type SessionChannelIO = "out" | "in"; +/** + * One durable Session channel record. + * + * `id` is the append's stable idempotency key. `seqNum` is the record's + * monotonic S2 sequence within the Session channel. Both stay stable when + * the same record is delivered again after a reconnect. + */ +export type SessionStreamRecord = Readonly<{ + id: string; + seqNum: number; + data: T; +}>; + +export type SessionStreamRecordPredicate = (record: SessionStreamRecord) => boolean; + /** * Manager for Session channel reads: a session-scoped parallel to * {@link InputStreamManager} keyed on `(sessionId, io)` instead of @@ -42,9 +57,38 @@ export interface SessionStreamManager { options?: InputStreamOnceOptions ): InputStreamOncePromise; + /** Wait for and consume the next record, including its durable metadata. */ + onceRecord?( + sessionId: string, + io: SessionChannelIO, + options?: InputStreamOnceOptions + ): InputStreamOncePromise; + + /** + * Wait for and consume the next record accepted by `predicate`. + * Earlier unmatched records stay buffered and block consumption so the + * committed cursor never advances past them. + */ + onceRecordWhere?( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate, + options?: InputStreamOnceOptions + ): InputStreamOncePromise; + /** Non-blocking peek at the head of the channel buffer. */ peek(sessionId: string, io: SessionChannelIO): unknown | undefined; + /** Non-blocking peek at the head record, including its durable metadata. */ + peekRecord?(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined; + + /** Non-blocking peek at the first buffered record accepted by `predicate`. */ + peekRecordWhere?( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate + ): SessionStreamRecord | undefined; + /** Last S2 sequence number seen on the given channel. */ lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined; diff --git a/packages/core/src/v3/test/mock-task-context.ts b/packages/core/src/v3/test/mock-task-context.ts index 5fbe1957613..cd65ac24d46 100644 --- a/packages/core/src/v3/test/mock-task-context.ts +++ b/packages/core/src/v3/test/mock-task-context.ts @@ -113,7 +113,12 @@ export type MockTaskContextDrivers = { * Send a record onto `session.in` for the given session. Resolves * pending `once()` waiters and fires all `on()` handlers. */ - send(sessionId: string, data: unknown, io?: SessionChannelIO): Promise; + send( + sessionId: string, + data: unknown, + io?: SessionChannelIO, + metadata?: { id?: string; seqNum?: number } + ): Promise; /** Close pending `once()` waiters with a timeout error. */ close(sessionId: string, io?: SessionChannelIO): void; }; @@ -277,9 +282,9 @@ export async function runInMockTaskContext( }, sessions: { in: { - send: (sessionId, data, io = "in") => + send: (sessionId, data, io = "in", metadata) => sessionStreamManager instanceof TestSessionStreamManager - ? sessionStreamManager.__sendFromTest(sessionId, io, data) + ? sessionStreamManager.__sendFromTest(sessionId, io, data, metadata) : Promise.reject( new Error("drivers.sessions.in.send requires the default TestSessionStreamManager") ), diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index 0e08441d4c3..c4865fcfe53 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -1,10 +1,16 @@ import type { InputStreamOnceResult } from "../inputStreams/types.js"; import { InputStreamOncePromise, InputStreamTimeoutError } from "../inputStreams/types.js"; import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; -import type { SessionChannelIO, SessionStreamManager } from "../sessionStreams/types.js"; +import type { + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "../sessionStreams/types.js"; type OnceWaiter = { - resolve: (value: InputStreamOnceResult) => void; + resolve: (value: InputStreamOnceResult) => void; + predicate?: SessionStreamRecordPredicate; timer?: ReturnType; signal?: AbortSignal; abortHandler?: () => void; @@ -31,7 +37,7 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { export class TestSessionStreamManager implements SessionStreamManager { private handlers = new Map>(); private onceWaiters = new Map(); - private buffer = new Map(); + private buffer = new Map(); private seqNums = new Map(); private dispatchedSeqNums = new Map(); @@ -55,21 +61,26 @@ export class TestSessionStreamManager implements SessionStreamManager { // messages into every newly attached per-turn handler. const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - const kept: unknown[] = []; - for (const data of buffered) { + const kept: SessionStreamRecord[] = []; + for (const record of buffered) { let consumed = false; try { - consumed = handler(data) === true; + consumed = handler(record.data) === true; } catch { // Never let a handler error break test state } - if (!consumed) kept.push(data); + if (consumed) { + this.#advanceLastDispatched(key, record.seqNum); + } else { + kept.push(record); + } } if (kept.length > 0) { this.buffer.set(key, kept); } else { this.buffer.delete(key); } + this.#drainOnceWaitersFromBuffer(key); } return { @@ -84,9 +95,40 @@ export class TestSessionStreamManager implements SessionStreamManager { io: SessionChannelIO, options?: InputStreamOnceOptions ): InputStreamOncePromise { + const recordPromise = this.onceRecord(sessionId, io, options); + return new InputStreamOncePromise((resolve, reject) => { + recordPromise.then((result) => { + resolve(result.ok ? { ok: true, output: result.output.data } : result); + }, reject); + }); + } + + onceRecord( + sessionId: string, + io: SessionChannelIO, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return this.#onceRecord(sessionId, io, undefined, options); + } + + onceRecordWhere( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return this.#onceRecord(sessionId, io, predicate, options); + } + + #onceRecord( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined, + options?: InputStreamOnceOptions + ): InputStreamOncePromise { const key = keyFor(sessionId, io); - return new InputStreamOncePromise((resolve) => { + return new InputStreamOncePromise((resolve) => { if (options?.signal?.aborted) { resolve({ ok: false, @@ -97,13 +139,18 @@ export class TestSessionStreamManager implements SessionStreamManager { const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - const next = buffered.shift(); - if (buffered.length === 0) this.buffer.delete(key); - resolve({ ok: true, output: next }); - return; + const next = buffered[0]!; + if (!predicate || predicate(next)) { + buffered.shift(); + if (buffered.length === 0) this.buffer.delete(key); + this.#advanceLastDispatched(key, next.seqNum); + this.#drainOnceWaitersFromBuffer(key); + resolve({ ok: true, output: next }); + return; + } } - const waiter: OnceWaiter = { resolve, signal: options?.signal }; + const waiter: OnceWaiter = { resolve, predicate, signal: options?.signal }; if (options?.timeoutMs !== undefined) { waiter.timer = setTimeout(() => { @@ -138,9 +185,19 @@ export class TestSessionStreamManager implements SessionStreamManager { } peek(sessionId: string, io: SessionChannelIO): unknown | undefined { - const buffered = this.buffer.get(keyFor(sessionId, io)); - if (buffered && buffered.length > 0) return buffered[0]; - return undefined; + return this.peekRecord(sessionId, io)?.data; + } + + peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined { + return this.buffer.get(keyFor(sessionId, io))?.[0]; + } + + peekRecordWhere( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate + ): SessionStreamRecord | undefined { + return this.buffer.get(keyFor(sessionId, io))?.find(predicate); } lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { @@ -152,15 +209,14 @@ export class TestSessionStreamManager implements SessionStreamManager { } lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - // `__sendFromTest` carries no seq numbers, so this only reflects - // explicit `setLastDispatchedSeqNum` calls (e.g. the waitpoint - // delivery path). Full cursor behaviour is exercised via the real - // manager. return this.dispatchedSeqNums.get(keyFor(sessionId, io)); } setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void { - const key = keyFor(sessionId, io); + this.#advanceLastDispatched(keyFor(sessionId, io), seqNum); + } + + #advanceLastDispatched(key: string, seqNum: number): void { const current = this.dispatchedSeqNums.get(key); if (current === undefined || seqNum > current) { this.dispatchedSeqNums.set(key, seqNum); @@ -180,8 +236,10 @@ export class TestSessionStreamManager implements SessionStreamManager { const key = keyFor(sessionId, io); const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - buffered.shift(); + const record = buffered.shift()!; if (buffered.length === 0) this.buffer.delete(key); + this.#advanceLastDispatched(key, record.seqNum); + this.#drainOnceWaitersFromBuffer(key); return true; } return false; @@ -235,39 +293,53 @@ export class TestSessionStreamManager implements SessionStreamManager { * resolves. Consumption is decided on the synchronous return value, * exactly like production. */ - async __sendFromTest(sessionId: string, io: SessionChannelIO, data: unknown): Promise { + async __sendFromTest( + sessionId: string, + io: SessionChannelIO, + data: unknown, + metadata?: { id?: string; seqNum?: number } + ): Promise { const key = keyFor(sessionId, io); + const seqNum = metadata?.seqNum ?? (this.seqNums.get(key) ?? -1) + 1; + const record: SessionStreamRecord = { + id: metadata?.id ?? `test-record-${seqNum}`, + seqNum, + data, + }; + const lastSeqNum = this.seqNums.get(key); + if (lastSeqNum === undefined || seqNum > lastSeqNum) { + this.seqNums.set(key, seqNum); + } - const waiters = this.onceWaiters.get(key); - if (waiters && waiters.length > 0) { - const w = waiters.shift()!; - if (waiters.length === 0) this.onceWaiters.delete(key); - if (w.timer) clearTimeout(w.timer); - if (w.signal && w.abortHandler) { - w.signal.removeEventListener("abort", w.abortHandler); - } - w.resolve({ ok: true, output: data }); - await this.#invokeHandlers(key, data); + const existingBuffer = this.buffer.get(key); + const waiter = + existingBuffer && existingBuffer.length > 0 ? undefined : this.#takeOnceWaiter(key, record); + if (waiter) { + this.#advanceLastDispatched(key, record.seqNum); + waiter.resolve({ ok: true, output: record }); + await this.#invokeHandlers(key, record.data); return; } - const consumed = await this.#invokeHandlers(key, data); - if (consumed) return; + const consumed = await this.#invokeHandlers(key, record.data); + if (consumed) { + this.#advanceLastDispatched(key, record.seqNum); + return; + } // Re-check waiters: handler invocation above is awaited (unlike the // synchronous production dispatch), and the runtime commonly registers // its next `once()` during that window — e.g. the turn loop reaching // `waitWithIdleTimeout` while a handler settles. Without this second // look the record would be buffered while the fresh waiter hangs. - const lateWaiters = this.onceWaiters.get(key); - if (lateWaiters && lateWaiters.length > 0) { - const w = lateWaiters.shift()!; - if (lateWaiters.length === 0) this.onceWaiters.delete(key); - if (w.timer) clearTimeout(w.timer); - if (w.signal && w.abortHandler) { - w.signal.removeEventListener("abort", w.abortHandler); - } - w.resolve({ ok: true, output: data }); + const bufferedAfterHandlers = this.buffer.get(key); + const lateWaiter = + bufferedAfterHandlers && bufferedAfterHandlers.length > 0 + ? undefined + : this.#takeOnceWaiter(key, record); + if (lateWaiter) { + this.#advanceLastDispatched(key, record.seqNum); + lateWaiter.resolve({ ok: true, output: record }); return; } @@ -276,7 +348,45 @@ export class TestSessionStreamManager implements SessionStreamManager { buffered = []; this.buffer.set(key, buffered); } - buffered.push(data); + buffered.push(record); + this.#drainOnceWaitersFromBuffer(key); + } + + #takeOnceWaiter(key: string, record: SessionStreamRecord): OnceWaiter | undefined { + const waiters = this.onceWaiters.get(key); + if (!waiters) return undefined; + + const index = waiters.findIndex((waiter) => { + if (!waiter.predicate) return true; + try { + return waiter.predicate(record); + } catch { + return false; + } + }); + if (index === -1) return undefined; + + const [waiter] = waiters.splice(index, 1); + if (waiters.length === 0) this.onceWaiters.delete(key); + if (waiter!.timer) clearTimeout(waiter!.timer); + if (waiter!.signal && waiter!.abortHandler) { + waiter!.signal.removeEventListener("abort", waiter!.abortHandler); + } + return waiter; + } + + #drainOnceWaitersFromBuffer(key: string): void { + const buffered = this.buffer.get(key); + while (buffered && buffered.length > 0) { + const record = buffered[0]!; + const waiter = this.#takeOnceWaiter(key, record); + if (!waiter) return; + + buffered.shift(); + if (buffered.length === 0) this.buffer.delete(key); + this.#advanceLastDispatched(key, record.seqNum); + waiter.resolve({ ok: true, output: record }); + } } /** diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 844d506079b..52fcba19584 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1543,7 +1543,31 @@ export type ChatTaskRunPayload< // keep their original shape. Each accessor resolves the session handle // lazily via `getChatSession()` so the module-level references stay // compatible with the pre-migration wiring. -const messagesInput: RealtimeDefinedInputStream = { +/** + * One message record delivered through {@link chat.messages}. + * + * `id` is the append's stable idempotency key and `seqNum` is its monotonic + * sequence on the Session `.in` channel. Both remain stable if the record is + * delivered again after a reconnect. + */ +export type ChatMessageRecord = Readonly<{ + id: string; + seqNum: number; + payload: ChatTaskWirePayload; +}>; + +export type ChatMessages = RealtimeDefinedInputStream & { + /** Whether a delivered message is waiting in the local buffer. Does not consume it. */ + hasPending(): Promise; + /** Consume one message record, or return `undefined` when the optional timeout elapses. */ + next(options?: { timeoutInSeconds?: number }): Promise; +}; + +function isChatMessageRecord(record: { data: unknown }): boolean { + return (record.data as ChatInputChunk | undefined)?.kind === "message"; +} + +const messagesInput: ChatMessages = { id: "chat-messages", on(handler) { return getChatSession().in.on((chunk) => { @@ -1607,6 +1631,37 @@ const messagesInput: RealtimeDefinedInputStream = { if (chunk && chunk.kind === "message") return chunk.payload; return undefined; }, + async hasPending() { + const session = getChatSession(); + return sessionStreams.peekRecordWhere(session.id, "in", isChatMessageRecord) !== undefined; + }, + async next(options) { + const timeoutInSeconds = options?.timeoutInSeconds; + if ( + timeoutInSeconds !== undefined && + (!Number.isFinite(timeoutInSeconds) || timeoutInSeconds < 0) + ) { + throw new TypeError( + "chat.messages.next() timeoutInSeconds must be a finite non-negative number" + ); + } + + const session = getChatSession(); + const result = await sessionStreams.onceRecordWhere( + session.id, + "in", + isChatMessageRecord, + timeoutInSeconds === undefined ? undefined : { timeoutMs: timeoutInSeconds * 1000 } + ); + if (!result.ok) return undefined; + + const chunk = result.output.data as Extract; + return { + id: result.output.id, + seqNum: result.output.seqNum, + payload: chunk.payload, + }; + }, wait(options) { return new ManualWaitpointPromise(async (resolve, reject) => { try { diff --git a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts new file mode 100644 index 00000000000..bdc4c9701a7 --- /dev/null +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -0,0 +1,242 @@ +// Import the test harness FIRST — this installs the resource catalog so +// `chat.customAgent()` calls below register their task functions correctly. +import "../src/v3/test/index.js"; + +import { resourceCatalog, sessionStreams } from "@trigger.dev/core/v3"; +import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; +import { describe, expect, it } from "vitest"; +import { chat, type ChatMessageRecord, type ChatTaskWirePayload } from "../src/v3/ai.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; +} + +function userPayload(chatId: string, id: string): ChatTaskWirePayload { + return { + chatId, + trigger: "submit-message", + message: { + id, + role: "user", + parts: [{ type: "text", text: id }], + }, + }; +} + +describe("chat.messages mailbox", () => { + it("checks pending input without consuming and takes one buffered record at a time", async () => { + const chatId = "mailbox-buffered"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + initial?: boolean; + before?: boolean; + afterFirst?: boolean; + afterSecond?: boolean; + first?: ChatMessageRecord; + second?: ChatMessageRecord; + cursorAfterFirst?: number; + cursorAfterSecond?: number; + } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-buffered", + run: async () => { + observations.initial = await chat.messages.hasPending(); + ready.resolve(); + await inspect.promise; + + observations.before = await chat.messages.hasPending(); + observations.first = await chat.messages.next(); + observations.cursorAfterFirst = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.afterFirst = await chat.messages.hasPending(); + observations.second = await chat.messages.next(); + observations.cursorAfterSecond = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.afterSecond = await chat.messages.hasPending(); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u1") }, + "in", + { id: "part-1", seqNum: 10 } + ); + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u2") }, + "in", + { id: "part-2", seqNum: 11 } + ); + inspect.resolve(); + await runPromise; + }); + + expect(observations).toEqual({ + initial: false, + before: true, + first: { id: "part-1", seqNum: 10, payload: userPayload(chatId, "u1") }, + cursorAfterFirst: 10, + afterFirst: true, + second: { id: "part-2", seqNum: 11, payload: userPayload(chatId, "u2") }, + cursorAfterSecond: 11, + afterSecond: false, + }); + }); + + it("returns undefined when next times out", async () => { + let result: ChatMessageRecord | undefined; + const agent = chat.customAgent({ + id: "chat-messages-mailbox-timeout", + run: async () => { + result = await chat.messages.next({ timeoutInSeconds: 0.01 }); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext((drivers) => + run( + { chatId: "mailbox-timeout", trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ) + ); + + expect(result).toBeUndefined(); + }); + + it("leaves earlier non-message records for their own consumer", async () => { + const chatId = "mailbox-mixed-kinds"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + pending?: boolean; + blocked?: ChatMessageRecord; + cursorAfterBlocked?: number; + headAfterBlocked?: unknown; + control?: unknown; + message?: ChatMessageRecord; + cursorAfterMessage?: number; + } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-mixed-kinds", + run: async () => { + ready.resolve(); + await inspect.promise; + + observations.pending = await chat.messages.hasPending(); + observations.blocked = await chat.messages.next({ timeoutInSeconds: 0.01 }); + observations.cursorAfterBlocked = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.headAfterBlocked = sessionStreams.peekRecord(chatId, "in"); + + const control = await sessionStreams.onceRecord(chatId, "in"); + observations.control = control.ok ? control.output : undefined; + observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); + observations.cursorAfterMessage = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + + await drivers.sessions.in.send( + chatId, + { kind: "handover", partialAssistantMessage: [], isFinal: false }, + "in", + { id: "handover-1", seqNum: 30 } + ); + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u-after-handover") }, + "in", + { id: "message-1", seqNum: 31 } + ); + inspect.resolve(); + await runPromise; + }); + + expect(observations).toEqual({ + pending: true, + blocked: undefined, + cursorAfterBlocked: undefined, + headAfterBlocked: { + id: "handover-1", + seqNum: 30, + data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + }, + control: { + id: "handover-1", + seqNum: 30, + data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + }, + message: { + id: "message-1", + seqNum: 31, + payload: userPayload(chatId, "u-after-handover"), + }, + cursorAfterMessage: 31, + }); + }); + + it("keeps record id and sequence stable across redelivery", async () => { + const payload = userPayload("mailbox-redelivery", "u-redelivered"); + + async function consumeDelivery(agentId: string): Promise { + const ready = deferred(); + const consume = deferred(); + let result: ChatMessageRecord | undefined; + const agent = chat.customAgent({ + id: agentId, + run: async () => { + ready.resolve(); + await consume.promise; + result = await chat.messages.next(); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId: payload.chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + await drivers.sessions.in.send(payload.chatId, { kind: "message", payload }, "in", { + id: "part-redelivered", + seqNum: 27, + }); + consume.resolve(); + await runPromise; + }); + + return result; + } + + const first = await consumeDelivery("chat-messages-mailbox-first-delivery"); + const redelivered = await consumeDelivery("chat-messages-mailbox-redelivery"); + + expect(first).toEqual({ id: "part-redelivered", seqNum: 27, payload }); + expect(redelivered).toEqual(first); + }); +}); diff --git a/packages/trigger-sdk/test/mockChatAgent.test.ts b/packages/trigger-sdk/test/mockChatAgent.test.ts index 202c3923732..62437369a39 100644 --- a/packages/trigger-sdk/test/mockChatAgent.test.ts +++ b/packages/trigger-sdk/test/mockChatAgent.test.ts @@ -1878,11 +1878,10 @@ describe("mockChatAgent", () => { // The snapshot reflects the post-turn accumulator: 1 user + 1 assistant. const roles = snap!.messages.map((m) => m.role); expect(roles).toEqual(["user", "assistant"]); - // `lastInEventId` stays undefined here: TestSessionStreamManager - // deliberately has no seq numbers, so the committed `.in` cursor - // the production write site reads is undefined in harness runs. - // The cursor round-trip is covered by the live smoke instead. - expect(snap!.lastInEventId).toBeUndefined(); + // TestSessionStreamManager assigns the same zero-based sequence + // numbers as the durable channel, so the committed input cursor is + // represented in snapshots produced by the harness too. + expect(snap!.lastInEventId).toBe("0"); } finally { await harness.close(); } From 39b5c5706c06ecd188924033a9f5c41e20ad8b3d Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 21:50:29 -0700 Subject: [PATCH 02/20] fix(chat): fail loudly when record peeking is unsupported --- packages/core/src/v3/sessionStreams/index.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 073bb2ac514..49dc97586d3 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -79,7 +79,11 @@ export class SessionStreamsAPI implements SessionStreamManager { } public peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined { - return this.#getManager().peekRecord?.(sessionId, io); + const manager = this.#getManager(); + if (!manager.peekRecord) { + throw new Error("The configured Session stream manager does not support record metadata"); + } + return manager.peekRecord(sessionId, io); } public peekRecordWhere( From 6abc529a758762d3e29c1c12ad13e80155fd15b2 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 10:40:05 -0700 Subject: [PATCH 03/20] fix(chat): keep mailbox cursor behind pending input --- docs/ai-chat/custom-agents.mdx | 14 +- packages/core/src/v3/sessionStreams/index.ts | 12 - .../src/v3/sessionStreams/manager.test.ts | 216 ++++++++++++++++-- .../core/src/v3/sessionStreams/manager.ts | 124 ++++++---- .../core/src/v3/sessionStreams/noopManager.ts | 8 - packages/core/src/v3/sessionStreams/types.ts | 19 +- .../v3/test/test-session-stream-manager.ts | 64 +++++- packages/trigger-sdk/src/v3/ai.ts | 5 +- packages/trigger-sdk/src/v3/sessions.ts | 8 +- .../test/chat-messages-mailbox.test.ts | 138 ++++++++--- 10 files changed, 462 insertions(+), 146 deletions(-) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index e7ed3f1e5b2..59607aa64f3 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -228,14 +228,15 @@ For full control, skip `createSession` and compose the primitives directly: | Method | Behavior | | --- | --- | | `peek()` | Return the buffer head when it is a message, without consuming it; otherwise return `undefined` | -| `hasPending()` | Resolve `true` when any message is buffered; does not consume it | +| `hasPending()` | Resolve `true` when the buffer head is a message; does not consume it | | `next({ timeoutInSeconds? })` | Consume exactly one message record in channel order, or resolve `undefined` when the optional timeout elapses | | `on(handler)` | Consume messages as they arrive and invoke the handler | | `waitWithIdleTimeout(options)` | Wait warm, then suspend the run until the next message arrives | -`hasPending()` checks the local, already-delivered buffer. It does not query the -remote Session channel or start a subscription. Use `waitWithIdleTimeout()` when -the loop needs to idle until future input arrives. +`hasPending()` checks whether the local, already-delivered buffer head is a +message that `next()` can consume immediately. It does not query the remote +Session channel or start a subscription. Use `waitWithIdleTimeout()` when the +loop needs to idle until future input arrives. `next()` returns a readonly record envelope: @@ -258,8 +259,9 @@ contrast, `on()` commits a record as soon as it dispatches the handler; avoid mixing `on()` and `next()` when a single loop owns mailbox consumption. The Session `.in` channel also carries control records such as handovers. If one -comes before a message, `next()` leaves it for its own consumer and waits until -that record has been handled. +comes before a message, `hasPending()` stays `false` and `next()` leaves the +control record for its own consumer. After that record is handled, the message +becomes pending. A complete loop: diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 49dc97586d3..63082516889 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -86,18 +86,6 @@ export class SessionStreamsAPI implements SessionStreamManager { return manager.peekRecord(sessionId, io); } - public peekRecordWhere( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate - ): SessionStreamRecord | undefined { - const manager = this.#getManager(); - if (!manager.peekRecordWhere) { - throw new Error("The configured Session stream manager does not support selective records"); - } - return manager.peekRecordWhere(sessionId, io, predicate); - } - public lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.#getManager().lastSeqNum(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.test.ts b/packages/core/src/v3/sessionStreams/manager.test.ts index 29693e60fce..99b541fbb3b 100644 --- a/packages/core/src/v3/sessionStreams/manager.test.ts +++ b/packages/core/src/v3/sessionStreams/manager.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"; import { StandardSessionStreamManager } from "./manager.js"; import type { ApiClient } from "../apiClient/index.js"; import type { SSEStreamPart } from "../apiClient/runStream.js"; +import { InputStreamTimeoutError } from "../inputStreams/types.js"; // Single-shot mock that mimics S2's long-poll: delivers `records` once via // `onPart` on the first subscribe call, then keeps the returned async @@ -44,6 +45,31 @@ function singleShotApiClient( } as unknown as ApiClient; } +function repeatingApiClient(record: { + id: string; + recordId?: string; + chunk: unknown; + timestamp: number; +}): ApiClient { + return { + async subscribeToSessionStream( + _sessionIdOrExternalId: string, + _io: "out" | "in", + options?: { onPart?: (part: SSEStreamPart) => void; signal?: AbortSignal } + ) { + options?.onPart?.(record as SSEStreamPart); + const signal = options?.signal; + // eslint-disable-next-line require-yield + return (async function* () { + if (signal?.aborted) return; + await new Promise((resolve) => { + signal?.addEventListener("abort", () => resolve(), { once: true }); + }); + })() as unknown as Awaited>; + }, + } as unknown as ApiClient; +} + describe("StandardSessionStreamManager — minTimestamp filter", () => { const sessionId = "session-1"; const io = "in" as const; @@ -210,24 +236,39 @@ describe("StandardSessionStreamManager — record metadata", () => { }); it("returns the same envelope when a record is redelivered", async () => { - const firstDelivery = new StandardSessionStreamManager( - singleShotApiClient([records[0]!]), + const manager = new StandardSessionStreamManager( + repeatingApiClient(records[0]!), "http://localhost" ); - const redelivery = new StandardSessionStreamManager( - singleShotApiClient([records[0]!]), + + const first = await manager.onceRecord(sessionId, io); + manager.disconnectStream(sessionId, io); + const replayed = await manager.onceRecord(sessionId, io); + + expect(first).toEqual(replayed); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); + + it("returns immediately when the timeout is zero", async () => { + const manager = new StandardSessionStreamManager( + { + subscribeToSessionStream: () => { + throw new Error("zero-timeout reads must not subscribe"); + }, + } as unknown as ApiClient, "http://localhost" ); - const first = await firstDelivery.onceRecord(sessionId, io); - const replayed = await redelivery.onceRecord(sessionId, io); + const result = await manager.onceRecord(sessionId, io, { timeoutMs: 0 }); - expect(first).toEqual(replayed); + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.error).toBeInstanceOf(InputStreamTimeoutError); + } - firstDelivery.disconnectStream(sessionId, io); - firstDelivery.disconnect(); - redelivery.disconnectStream(sessionId, io); - redelivery.disconnect(); + manager.disconnect(); }); it("does not consume a matching record past an earlier unmatched record", async () => { @@ -256,10 +297,10 @@ describe("StandardSessionStreamManager — record metadata", () => { { timeoutMs: 200 } ); - expect(manager.peekRecordWhere(sessionId, io, (record) => record.id === "message-1")).toEqual({ - id: "message-1", - seqNum: 51, - data: { kind: "message", payload: { id: "u1" } }, + expect(manager.peekRecord(sessionId, io)).toEqual({ + id: "handover-1", + seqNum: 50, + data: { kind: "handover" }, }); expect(manager.lastDispatchedSeqNum(sessionId, io)).toBeUndefined(); @@ -281,4 +322,149 @@ describe("StandardSessionStreamManager — record metadata", () => { manager.disconnectStream(sessionId, io); manager.disconnect(); }); + + it("keeps the persisted cursor behind each earlier buffered record", async () => { + const manager = new StandardSessionStreamManager( + singleShotApiClient([ + { + id: "50", + recordId: "message-1", + chunk: { kind: "message", payload: { id: "u1" } }, + timestamp: 1000, + }, + { + id: "51", + recordId: "stop-1", + chunk: { kind: "stop" }, + timestamp: 2000, + }, + { + id: "52", + recordId: "message-2", + chunk: { kind: "message", payload: { id: "u2" } }, + timestamp: 3000, + }, + { + id: "53", + recordId: "stop-2", + chunk: { kind: "stop" }, + timestamp: 4000, + }, + ]), + "http://localhost" + ); + let resolveStop!: () => void; + let remainingStops = 2; + const stopConsumed = new Promise((resolve) => { + resolveStop = resolve; + }); + + manager.on(sessionId, io, (data) => { + if ((data as { kind?: string }).kind !== "stop") return; + remainingStops--; + if (remainingStops === 0) resolveStop(); + return true; + }); + await stopConsumed; + + expect(manager.peekRecord(sessionId, io)).toEqual({ + id: "message-1", + seqNum: 50, + data: { kind: "message", payload: { id: "u1" } }, + }); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49); + + const firstMessage = await manager.onceRecord(sessionId, io); + expect(firstMessage.ok && firstMessage.output.id).toBe("message-1"); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(51); + + const secondMessage = await manager.onceRecord(sessionId, io); + expect(secondMessage.ok && secondMessage.output.id).toBe("message-2"); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(53); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); + + it("retains cursor barriers when disconnect clears the buffer", async () => { + const manager = new StandardSessionStreamManager( + singleShotApiClient([ + { + id: "50", + recordId: "message-1", + chunk: { kind: "message", payload: { id: "u1" } }, + timestamp: 1000, + }, + { + id: "51", + recordId: "stop-1", + chunk: { kind: "stop" }, + timestamp: 2000, + }, + ]), + "http://localhost" + ); + let resolveStop!: () => void; + const stopConsumed = new Promise((resolve) => { + resolveStop = resolve; + }); + + manager.on(sessionId, io, (data) => { + if ((data as { kind?: string }).kind !== "stop") return; + resolveStop(); + return true; + }); + await stopConsumed; + + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49); + manager.disconnectStream(sessionId, io); + expect(manager.peekRecord(sessionId, io)).toBeUndefined(); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49); + + manager.setLastDispatchedSeqNum(sessionId, io, 51); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49); + + manager.reset(); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBeUndefined(); + }); + + it("does not expose a negative cursor when sequence zero is buffered", async () => { + const manager = new StandardSessionStreamManager( + singleShotApiClient([ + { + id: "0", + recordId: "message-0", + chunk: { kind: "message", payload: { id: "u0" } }, + timestamp: 1000, + }, + { + id: "1", + recordId: "stop-1", + chunk: { kind: "stop" }, + timestamp: 2000, + }, + ]), + "http://localhost" + ); + let resolveStop!: () => void; + const stopConsumed = new Promise((resolve) => { + resolveStop = resolve; + }); + + manager.on(sessionId, io, (data) => { + if ((data as { kind?: string }).kind !== "stop") return; + resolveStop(); + return true; + }); + await stopConsumed; + + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBeUndefined(); + + const message = await manager.onceRecord(sessionId, io); + expect(message.ok && message.output.id).toBe("message-0"); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(1); + + manager.disconnectStream(sessionId, io); + manager.disconnect(); + }); }); diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index 184d65d2443..dccb0baade2 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -66,14 +66,17 @@ export class StandardSessionStreamManager implements SessionStreamManager { // that's already being delivered out-of-band via the waitpoint. private explicitlyDisconnected = new Set(); private seqNums = new Map(); - // Highest seq_num that has been *consumed* (delivered to a once() - // waiter or shifted off the buffer into a once() caller) on a channel. + // Sequence numbers for records that were delivered but not consumed. + // Kept separately from `buffer` because `disconnectStream()` clears the + // local buffer before a waitpoint suspension, but those records must still + // hold the persisted consume cursor back. + private unconsumedSeqNums = new Map>(); + // High-water mark of seq_nums that have been *consumed* (delivered to a + // once() waiter or shifted off the buffer into a once() caller) on a channel. // Distinct from `seqNums`, which advances whenever any record is // received from SSE — even ones still sitting in the local buffer. - // The committed-consume cursor is what gets persisted on the - // turn-complete control record's `session-in-event-id` header so the - // next worker boot can resume `.in` from this point without - // re-delivering already-handled user messages. + // `lastDispatchedSeqNum()` clamps this behind any unconsumed barrier before + // it is persisted on a turn-complete control record. private lastDispatchedSeqNums = new Map(); // Reconnect attempt counter per key. Drives the exponential backoff // applied by `#ensureTailConnected`'s `.finally` so a persistent @@ -182,36 +185,30 @@ export class StandardSessionStreamManager implements SessionStreamManager { ): InputStreamOncePromise { const key = keyFor(sessionId, io); + if (options?.timeoutMs === 0) { + const record = this.#takeBufferedRecord(key, predicate); + return new InputStreamOncePromise((resolve) => { + resolve( + record + ? { ok: true, output: record } + : { ok: false, error: new InputStreamTimeoutError(key, 0) } + ); + }); + } + this.explicitlyDisconnected.delete(key); this.#ensureTailConnected(sessionId, io); - const buffered = this.buffer.get(key); - if (buffered && buffered.length > 0) { - const record = buffered[0]!; - if (!predicate || predicate(record)) { - buffered.shift(); - if (buffered.length === 0) { - this.buffer.delete(key); - } - this.#advanceLastDispatched(key, record.seqNum); - this.#drainOnceWaitersFromBuffer(key); - return new InputStreamOncePromise((resolve) => { - resolve({ ok: true, output: record }); - }); - } + const record = this.#takeBufferedRecord(key, predicate); + if (record) { + return new InputStreamOncePromise((resolve) => { + resolve({ ok: true, output: record }); + }); } return new InputStreamOncePromise((resolve, reject) => { const waiter: OnceWaiter = { resolve, reject, predicate }; - if (predicate && options?.timeoutMs === 0) { - resolve({ - ok: false, - error: new InputStreamTimeoutError(key, 0), - }); - return; - } - if (options?.signal) { if (options.signal.aborted) { reject(new Error("Aborted")); @@ -246,6 +243,25 @@ export class StandardSessionStreamManager implements SessionStreamManager { }); } + #takeBufferedRecord( + key: string, + predicate: SessionStreamRecordPredicate | undefined + ): SessionStreamRecord | undefined { + const buffered = this.buffer.get(key); + if (!buffered || buffered.length === 0) return undefined; + + const record = buffered[0]!; + if (predicate && !predicate(record)) return undefined; + + buffered.shift(); + if (buffered.length === 0) { + this.buffer.delete(key); + } + this.#advanceLastDispatched(key, record.seqNum); + this.#drainOnceWaitersFromBuffer(key); + return record; + } + peek(sessionId: string, io: SessionChannelIO): unknown | undefined { return this.peekRecord(sessionId, io)?.data; } @@ -254,14 +270,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { return this.buffer.get(keyFor(sessionId, io))?.[0]; } - peekRecordWhere( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate - ): SessionStreamRecord | undefined { - return this.buffer.get(keyFor(sessionId, io))?.find(predicate); - } - lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.seqNums.get(keyFor(sessionId, io)); } @@ -275,14 +283,30 @@ export class StandardSessionStreamManager implements SessionStreamManager { } lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - return this.lastDispatchedSeqNums.get(keyFor(sessionId, io)); + const key = keyFor(sessionId, io); + const highWatermark = this.lastDispatchedSeqNums.get(key); + if (highWatermark === undefined) return undefined; + + const unconsumedSeqNums = this.unconsumedSeqNums.get(key); + if (!unconsumedSeqNums || unconsumedSeqNums.size === 0) return highWatermark; + + let earliestUnconsumedSeqNum = Infinity; + for (const seqNum of unconsumedSeqNums) { + earliestUnconsumedSeqNum = Math.min(earliestUnconsumedSeqNum, seqNum); + } + + const safeCursor = Math.min(highWatermark, earliestUnconsumedSeqNum - 1); + return safeCursor >= 0 ? safeCursor : undefined; } setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void { + if (!Number.isFinite(seqNum)) return; + this.#advanceLastDispatched(keyFor(sessionId, io), seqNum); } #advanceLastDispatched(key: string, seqNum: number): void { + this.#removeUnconsumedRecord(key, seqNum); if (!Number.isFinite(seqNum)) return; const current = this.lastDispatchedSeqNums.get(key); if (current === undefined || seqNum > current) { @@ -290,6 +314,25 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } + #markUnconsumedRecord(key: string, seqNum: number): void { + if (!Number.isFinite(seqNum)) return; + + let unconsumedSeqNums = this.unconsumedSeqNums.get(key); + if (!unconsumedSeqNums) { + unconsumedSeqNums = new Set(); + this.unconsumedSeqNums.set(key, unconsumedSeqNums); + } + unconsumedSeqNums.add(seqNum); + } + + #removeUnconsumedRecord(key: string, seqNum: number): void { + const unconsumedSeqNums = this.unconsumedSeqNums.get(key); + unconsumedSeqNums?.delete(seqNum); + if (unconsumedSeqNums?.size === 0) { + this.unconsumedSeqNums.delete(key); + } + } + setMinTimestamp(sessionId: string, io: SessionChannelIO, minTimestamp: number | undefined): void { const key = keyFor(sessionId, io); if (minTimestamp === undefined) { @@ -366,6 +409,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.disconnect(); this.seqNums.clear(); this.lastDispatchedSeqNums.clear(); + this.unconsumedSeqNums.clear(); this.minTimestamps.clear(); this.handlers.clear(); this.reconnectAttempts.clear(); @@ -457,9 +501,8 @@ export class StandardSessionStreamManager implements SessionStreamManager { onPart: (part) => { if (signal.aborted) return; const seqNum = parseInt(part.id, 10); - if (Number.isFinite(seqNum)) { - this.seqNums.set(key, seqNum); - } + if (!Number.isFinite(seqNum)) return; + this.seqNums.set(key, seqNum); // Trigger control records (turn-complete, upgrade-required) // are dispatched out-of-band via `onControl` — they're not @@ -551,6 +594,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.buffer.set(key, buffered); } buffered.push(record); + this.#markUnconsumedRecord(key, record.seqNum); this.#drainOnceWaitersFromBuffer(key); } diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index 1a68c36e7e7..aeb2a9aeb44 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -55,14 +55,6 @@ export class NoopSessionStreamManager implements SessionStreamManager { return undefined; } - peekRecordWhere( - _sessionId: string, - _io: SessionChannelIO, - _predicate: SessionStreamRecordPredicate - ): SessionStreamRecord | undefined { - return undefined; - } - lastSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { return undefined; } diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index d80f8d6cec3..ce55cbee39a 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -82,13 +82,6 @@ export interface SessionStreamManager { /** Non-blocking peek at the head record, including its durable metadata. */ peekRecord?(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined; - /** Non-blocking peek at the first buffered record accepted by `predicate`. */ - peekRecordWhere?( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate - ): SessionStreamRecord | undefined; - /** Last S2 sequence number seen on the given channel. */ lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined; @@ -96,10 +89,11 @@ export interface SessionStreamManager { setLastSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void; /** - * Highest sequence number that has been *consumed* on the channel — - * delivered to a `once()` waiter or shifted off the buffer into one. - * Distinct from {@link lastSeqNum}, which advances on every received - * record regardless of whether anything consumed it. Used by + * Highest sequence number that is safe to persist as consumed. When a later + * record is handled while an earlier record remains unconsumed, this stays + * behind the earliest unconsumed record. Distinct from {@link lastSeqNum}, + * which advances on every received record regardless of whether anything + * consumed it. Used by * `chat.agent` to persist the `.in` resume cursor on each * `turn-complete` control record so the next worker boot can resume * the channel from this point without replaying processed messages. @@ -109,7 +103,8 @@ export interface SessionStreamManager { /** * Seed the committed-consume cursor at worker boot — e.g. from the * `session-in-event-id` header on the latest `turn-complete` on - * `.out`. Monotonic: only ever advances forward, never backwards. + * `.out`. Monotonic: only ever advances forward, never backwards. Existing + * unconsumed records still constrain {@link lastDispatchedSeqNum}. */ setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void; diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index c4865fcfe53..52de0b1571d 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -40,6 +40,7 @@ export class TestSessionStreamManager implements SessionStreamManager { private buffer = new Map(); private seqNums = new Map(); private dispatchedSeqNums = new Map(); + private unconsumedSeqNums = new Map>(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { const key = keyFor(sessionId, io); @@ -150,6 +151,14 @@ export class TestSessionStreamManager implements SessionStreamManager { } } + if (options?.timeoutMs === 0) { + resolve({ + ok: false, + error: new InputStreamTimeoutError(key, 0), + }); + return; + } + const waiter: OnceWaiter = { resolve, predicate, signal: options?.signal }; if (options?.timeoutMs !== undefined) { @@ -192,14 +201,6 @@ export class TestSessionStreamManager implements SessionStreamManager { return this.buffer.get(keyFor(sessionId, io))?.[0]; } - peekRecordWhere( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate - ): SessionStreamRecord | undefined { - return this.buffer.get(keyFor(sessionId, io))?.find(predicate); - } - lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.seqNums.get(keyFor(sessionId, io)); } @@ -209,20 +210,56 @@ export class TestSessionStreamManager implements SessionStreamManager { } lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - return this.dispatchedSeqNums.get(keyFor(sessionId, io)); + const key = keyFor(sessionId, io); + const highWatermark = this.dispatchedSeqNums.get(key); + if (highWatermark === undefined) return undefined; + + const unconsumedSeqNums = this.unconsumedSeqNums.get(key); + if (!unconsumedSeqNums || unconsumedSeqNums.size === 0) return highWatermark; + + let earliestUnconsumedSeqNum = Infinity; + for (const seqNum of unconsumedSeqNums) { + earliestUnconsumedSeqNum = Math.min(earliestUnconsumedSeqNum, seqNum); + } + + const safeCursor = Math.min(highWatermark, earliestUnconsumedSeqNum - 1); + return safeCursor >= 0 ? safeCursor : undefined; } setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void { + if (!Number.isFinite(seqNum)) return; + this.#advanceLastDispatched(keyFor(sessionId, io), seqNum); } #advanceLastDispatched(key: string, seqNum: number): void { + this.#removeUnconsumedRecord(key, seqNum); + if (!Number.isFinite(seqNum)) return; const current = this.dispatchedSeqNums.get(key); if (current === undefined || seqNum > current) { this.dispatchedSeqNums.set(key, seqNum); } } + #markUnconsumedRecord(key: string, seqNum: number): void { + if (!Number.isFinite(seqNum)) return; + + let unconsumedSeqNums = this.unconsumedSeqNums.get(key); + if (!unconsumedSeqNums) { + unconsumedSeqNums = new Set(); + this.unconsumedSeqNums.set(key, unconsumedSeqNums); + } + unconsumedSeqNums.add(seqNum); + } + + #removeUnconsumedRecord(key: string, seqNum: number): void { + const unconsumedSeqNums = this.unconsumedSeqNums.get(key); + unconsumedSeqNums?.delete(seqNum); + if (unconsumedSeqNums?.size === 0) { + this.unconsumedSeqNums.delete(key); + } + } + setMinTimestamp( _sessionId: string, _io: SessionChannelIO, @@ -245,8 +282,8 @@ export class TestSessionStreamManager implements SessionStreamManager { return false; } - disconnectStream(_sessionId: string, _io: SessionChannelIO): void { - // no-op — no real SSE tail in tests + disconnectStream(sessionId: string, io: SessionChannelIO): void { + this.buffer.delete(keyFor(sessionId, io)); } clearHandlers(): void { @@ -267,6 +304,7 @@ export class TestSessionStreamManager implements SessionStreamManager { this.buffer.clear(); this.seqNums.clear(); this.dispatchedSeqNums.clear(); + this.unconsumedSeqNums.clear(); } disconnect(): void { @@ -301,6 +339,9 @@ export class TestSessionStreamManager implements SessionStreamManager { ): Promise { const key = keyFor(sessionId, io); const seqNum = metadata?.seqNum ?? (this.seqNums.get(key) ?? -1) + 1; + if (!Number.isFinite(seqNum)) { + throw new TypeError("Test Session stream records require a finite sequence number"); + } const record: SessionStreamRecord = { id: metadata?.id ?? `test-record-${seqNum}`, seqNum, @@ -349,6 +390,7 @@ export class TestSessionStreamManager implements SessionStreamManager { this.buffer.set(key, buffered); } buffered.push(record); + this.#markUnconsumedRecord(key, record.seqNum); this.#drainOnceWaitersFromBuffer(key); } diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 52fcba19584..95b0f5267da 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1557,7 +1557,7 @@ export type ChatMessageRecord = Readonly<{ }>; export type ChatMessages = RealtimeDefinedInputStream & { - /** Whether a delivered message is waiting in the local buffer. Does not consume it. */ + /** Whether the local buffer head is a message that can be consumed immediately. */ hasPending(): Promise; /** Consume one message record, or return `undefined` when the optional timeout elapses. */ next(options?: { timeoutInSeconds?: number }): Promise; @@ -1632,8 +1632,7 @@ const messagesInput: ChatMessages = { return undefined; }, async hasPending() { - const session = getChatSession(); - return sessionStreams.peekRecordWhere(session.id, "in", isChatMessageRecord) !== undefined; + return messagesInput.peek() !== undefined; }, async next(options) { const timeoutInSeconds = options?.timeoutInSeconds; diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index 8a01f8293c4..8130d333700 100644 --- a/packages/trigger-sdk/src/v3/sessions.ts +++ b/packages/trigger-sdk/src/v3/sessions.ts @@ -683,10 +683,10 @@ export class SessionInputChannel { } /** - * The highest S2 sequence number of any record this channel has - * delivered to a `once()` / `wait()` consumer (or had shifted off its - * buffer into one). Distinct from "last received" — buffered-but-not- - * yet-consumed records don't count. + * The highest S2 sequence number that is safe to persist as consumed. + * This stays behind the earliest unconsumed record if a later record was + * handled first. Distinct from "last received", which advances for records + * that may still be pending. * * Used by `chat.agent` to persist the `.in` resume cursor on each * `turn-complete` control record, so the next worker boot can subscribe diff --git a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts index bdc4c9701a7..295b8f6c551 100644 --- a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -102,7 +102,7 @@ describe("chat.messages mailbox", () => { const agent = chat.customAgent({ id: "chat-messages-mailbox-timeout", run: async () => { - result = await chat.messages.next({ timeoutInSeconds: 0.01 }); + result = await chat.messages.next({ timeoutInSeconds: 0 }); }, }); const run = resourceCatalog.getTask(agent.id)?.fns.run; @@ -128,6 +128,7 @@ describe("chat.messages mailbox", () => { cursorAfterBlocked?: number; headAfterBlocked?: unknown; control?: unknown; + pendingAfterControl?: boolean; message?: ChatMessageRecord; cursorAfterMessage?: number; } = {}; @@ -139,12 +140,13 @@ describe("chat.messages mailbox", () => { await inspect.promise; observations.pending = await chat.messages.hasPending(); - observations.blocked = await chat.messages.next({ timeoutInSeconds: 0.01 }); + observations.blocked = await chat.messages.next({ timeoutInSeconds: 0 }); observations.cursorAfterBlocked = sessionStreams.lastDispatchedSeqNum(chatId, "in"); observations.headAfterBlocked = sessionStreams.peekRecord(chatId, "in"); const control = await sessionStreams.onceRecord(chatId, "in"); observations.control = control.ok ? control.output : undefined; + observations.pendingAfterControl = await chat.messages.hasPending(); observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); observations.cursorAfterMessage = sessionStreams.lastDispatchedSeqNum(chatId, "in"); }, @@ -176,7 +178,7 @@ describe("chat.messages mailbox", () => { }); expect(observations).toEqual({ - pending: true, + pending: false, blocked: undefined, cursorAfterBlocked: undefined, headAfterBlocked: { @@ -189,6 +191,7 @@ describe("chat.messages mailbox", () => { seqNum: 30, data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, }, + pendingAfterControl: true, message: { id: "message-1", seqNum: 31, @@ -198,43 +201,108 @@ describe("chat.messages mailbox", () => { }); }); + it("keeps the cursor behind a buffered message when a later stop is consumed", async () => { + const chatId = "mailbox-cursor-gap"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + cursorBefore?: number; + message?: ChatMessageRecord; + cursorAfter?: number; + } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-cursor-gap", + run: async () => { + const stop = chat.createStopSignal(); + ready.resolve(); + await inspect.promise; + + observations.cursorBefore = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); + observations.cursorAfter = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + stop.cleanup(); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u1") }, + "in", + { id: "message-1", seqNum: 50 } + ); + await drivers.sessions.in.send(chatId, { kind: "stop" }, "in", { + id: "stop-1", + seqNum: 51, + }); + inspect.resolve(); + await runPromise; + }); + + expect(observations).toEqual({ + cursorBefore: 49, + message: { + id: "message-1", + seqNum: 50, + payload: userPayload(chatId, "u1"), + }, + cursorAfter: 51, + }); + }); + it("keeps record id and sequence stable across redelivery", async () => { const payload = userPayload("mailbox-redelivery", "u-redelivered"); + const ready = deferred(); + const consumeFirst = deferred(); + const readyForRedelivery = deferred(); + const consumeRedelivery = deferred(); + let first: ChatMessageRecord | undefined; + let redelivered: ChatMessageRecord | undefined; + const agent = chat.customAgent({ + id: "chat-messages-mailbox-redelivery", + run: async () => { + ready.resolve(); + await consumeFirst.promise; + first = await chat.messages.next({ timeoutInSeconds: 0 }); - async function consumeDelivery(agentId: string): Promise { - const ready = deferred(); - const consume = deferred(); - let result: ChatMessageRecord | undefined; - const agent = chat.customAgent({ - id: agentId, - run: async () => { - ready.resolve(); - await consume.promise; - result = await chat.messages.next(); - }, - }); - const run = resourceCatalog.getTask(agent.id)?.fns.run; - if (!run) throw new Error("custom agent was not registered"); - - await runInMockTaskContext(async (drivers) => { - const runPromise = run( - { chatId: payload.chatId, trigger: "preload" }, - { ctx: drivers.ctx, signal: new AbortController().signal } - ); - await ready.promise; - await drivers.sessions.in.send(payload.chatId, { kind: "message", payload }, "in", { - id: "part-redelivered", - seqNum: 27, - }); - consume.resolve(); - await runPromise; - }); + sessionStreams.disconnectStream(payload.chatId, "in"); + readyForRedelivery.resolve(); + await consumeRedelivery.promise; + redelivered = await chat.messages.next({ timeoutInSeconds: 0 }); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); - return result; - } + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId: payload.chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + await drivers.sessions.in.send(payload.chatId, { kind: "message", payload }, "in", { + id: "part-redelivered", + seqNum: 27, + }); + consumeFirst.resolve(); - const first = await consumeDelivery("chat-messages-mailbox-first-delivery"); - const redelivered = await consumeDelivery("chat-messages-mailbox-redelivery"); + await readyForRedelivery.promise; + await drivers.sessions.in.send(payload.chatId, { kind: "message", payload }, "in", { + id: "part-redelivered", + seqNum: 27, + }); + consumeRedelivery.resolve(); + await runPromise; + }); expect(first).toEqual({ id: "part-redelivered", seqNum: 27, payload }); expect(redelivered).toEqual(first); From b334ab267e3d3ab87a3a1a4bc2d14ad781da1dea Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 14:45:57 -0700 Subject: [PATCH 04/20] fix(chat): preserve mailbox cursor across waitpoints --- ...uns.$runFriendlyId.session-streams.wait.ts | 15 +- ...ealtime.v1.sessions.$session.$io.append.ts | 17 +- ...ealtime.v1.sessions.$session.$io.append.ts | 21 +- .../sessionStreamWaitpointCache.server.ts | 79 ++++++- apps/webapp/app/v3/webhookEngine.server.ts | 15 +- packages/core/src/v3/schemas/api.ts | 2 + packages/core/src/v3/sessionStreams/index.ts | 8 + .../src/v3/sessionStreams/manager.test.ts | 19 +- .../core/src/v3/sessionStreams/manager.ts | 36 +-- .../core/src/v3/sessionStreams/noopManager.ts | 2 + packages/core/src/v3/sessionStreams/types.ts | 5 +- .../src/v3/sessionStreams/wireProtocol.ts | 47 ++++ .../src/v3/test/session-waitpoint-backend.ts | 36 ++- .../v3/test/test-session-stream-manager.ts | 21 +- packages/trigger-sdk/src/v3/sessions.ts | 74 ++++-- .../test/pending-message-drain.test.ts | 215 ++++++++++++++++-- 16 files changed, 517 insertions(+), 95 deletions(-) diff --git a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts index c00ff51b3be..3cb98f5847e 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts @@ -1,6 +1,8 @@ import { json } from "@remix-run/server-runtime"; import { CreateSessionStreamWaitpointRequestBody, + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, type CreateSessionStreamWaitpointResponseBody, } from "@trigger.dev/core/v3"; import { WaitpointId } from "@trigger.dev/core/v3/isomorphic"; @@ -125,7 +127,8 @@ const { action, loader } = createActionApiRoute( addressingKey, body.io, result.waitpoint.id, - ttlMs && ttlMs > 0 ? ttlMs : undefined + ttlMs && ttlMs > 0 ? ttlMs : undefined, + body.responseFormat ); // Race-check. If a record landed on the channel before this @@ -155,8 +158,14 @@ const { action, loader } = createActionApiRoute( await engine.completeWaitpoint({ id: result.waitpoint.id, output: { - value: record.data, - type: "application/json", + value: + body.responseFormat === "record-v1" + ? serializeSessionStreamWaitpointRecord(record.data, record.seqNum) + : record.data, + type: + body.responseFormat === "record-v1" + ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE + : "application/json", isError: false, }, }); diff --git a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts index d4dd1d9f19f..7ff85cde863 100644 --- a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts @@ -15,6 +15,7 @@ import { claimSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, + sessionStreamWaitpointOutput, } from "~/services/sessionStreamWaitpointCache.server"; import { anyResource, createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { engine } from "~/v3/runEngine.server"; @@ -201,7 +202,7 @@ const { action, loader } = createActionApiRoute( // keyed on the canonical addressing key the agent registered with via // `sessions.open(...).in.wait()`, so writers and readers converge // regardless of which URL form they used. - const [drainError, waitpointIds] = await tryCatch( + const [drainError, waitpoints] = await tryCatch( drainSessionStreamWaitpoints(authentication.environment.id, addressingKey, params.io) ); if (drainError) { @@ -210,24 +211,20 @@ const { action, loader } = createActionApiRoute( io: params.io, error: drainError, }); - } else if (waitpointIds && waitpointIds.length > 0) { + } else if (waitpoints && waitpoints.length > 0) { await Promise.all( - waitpointIds.map(async (waitpointId) => { + waitpoints.map(async (waitpoint) => { const [completeError] = await tryCatch( engine.completeWaitpoint({ - id: waitpointId, - output: { - value: part, - type: "application/json", - isError: false, - }, + id: waitpoint.id, + output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq), }) ); if (completeError) { logger.error("Failed to complete session stream waitpoint", { addressingKey, io: params.io, - waitpointId, + waitpointId: waitpoint.id, error: completeError, }); } diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts index ab318f31c7a..35bdf3a5dd9 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts @@ -13,7 +13,10 @@ import { resolveSessionByIdOrExternalId, } from "~/services/realtime/sessions.server"; import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server"; -import { drainSessionStreamWaitpoints } from "~/services/sessionStreamWaitpointCache.server"; +import { + drainSessionStreamWaitpoints, + sessionStreamWaitpointOutput, +} from "~/services/sessionStreamWaitpointCache.server"; import { requireUserId } from "~/services/session.server"; import { EnvironmentParamSchema } from "~/utils/pathBuilder"; import { engine } from "~/v3/runEngine.server"; @@ -114,7 +117,7 @@ export async function action({ request, params }: ActionFunctionArgs) { // Drain any waitpoints registered for this channel — same as the // public append. Best-effort; failure doesn't fail the append. - const [drainError, waitpointIds] = await tryCatch( + const [drainError, waitpoints] = await tryCatch( drainSessionStreamWaitpoints(environment.id, addressingKey, io) ); if (drainError) { @@ -123,24 +126,20 @@ export async function action({ request, params }: ActionFunctionArgs) { io, error: drainError, }); - } else if (waitpointIds && waitpointIds.length > 0) { + } else if (waitpoints && waitpoints.length > 0) { await Promise.all( - waitpointIds.map(async (waitpointId) => { + waitpoints.map(async (waitpoint) => { const [completeError] = await tryCatch( engine.completeWaitpoint({ - id: waitpointId, - output: { - value: part, - type: "application/json", - isError: false, - }, + id: waitpoint.id, + output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq ?? undefined), }) ); if (completeError) { logger.error("Failed to complete session stream waitpoint (playground)", { addressingKey, io, - waitpointId, + waitpointId: waitpoint.id, error: completeError, }); } diff --git a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts index 7b53042d8d3..0c21c10be1e 100644 --- a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts +++ b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts @@ -1,5 +1,9 @@ import { Redis } from "ioredis"; import { defaultReconnectOnError } from "@internal/redis"; +import { + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, +} from "@trigger.dev/core/v3"; import { env } from "~/env.server"; import { singleton } from "~/utils/singleton"; import { logger } from "./logger.server"; @@ -13,12 +17,35 @@ import { logger } from "./logger.server"; // is shared — without it, two environments using the same externalId // would drain each other's waitpoints. const KEY_PREFIX = "ssw:"; +const FORMAT_KEY_PREFIX = "sswf:"; const DEFAULT_TTL_MS = 7 * 24 * 60 * 60 * 1000; // 7 days +export type SessionStreamWaitpoint = { + id: string; + responseFormat?: "record-v1"; +}; + +export function sessionStreamWaitpointOutput( + waitpoint: SessionStreamWaitpoint, + data: string, + seqNum: number | undefined +): { value: string; type: string; isError: false } { + const hasRecordEnvelope = waitpoint.responseFormat === "record-v1" && seqNum !== undefined; + return { + value: hasRecordEnvelope ? serializeSessionStreamWaitpointRecord(data, seqNum) : data, + type: hasRecordEnvelope ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE : "application/json", + isError: false, + }; +} + function buildKey(environmentId: string, addressingKey: string, io: "out" | "in"): string { return `${KEY_PREFIX}${environmentId}:${addressingKey}:${io}`; } +function buildFormatKey(waitpointId: string): string { + return `${FORMAT_KEY_PREFIX}${waitpointId}`; +} + // Pre-env-scoping key format, drained for one release so waitpoints from the // previous deploy still wake. Removable once this has been live > turn timeout. function buildLegacyKey(addressingKey: string, io: "out" | "in"): string { @@ -81,13 +108,25 @@ export async function addSessionStreamWaitpoint( addressingKey: string, io: "out" | "in", waitpointId: string, - ttlMs?: number + ttlMs?: number, + responseFormat?: "record-v1" ): Promise { if (!redis) return; try { const key = buildKey(environmentId, addressingKey, io); - await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(ttlMs ?? DEFAULT_TTL_MS)); + const effectiveTtlMs = ttlMs ?? DEFAULT_TTL_MS; + + // Keep the set member as the plain waitpoint id so an older append + // instance can still drain it during a rolling deploy. New instances read + // the optional response format from this separate, TTL-bound key. + if (responseFormat) { + await redis.set(buildFormatKey(waitpointId), responseFormat, "PX", effectiveTtlMs); + } else { + await redis.del(buildFormatKey(waitpointId)); + } + + await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(effectiveTtlMs)); } catch (error) { logger.error("Failed to set session stream waitpoint cache", { environmentId, @@ -107,7 +146,7 @@ export async function drainSessionStreamWaitpoints( environmentId: string, addressingKey: string, io: "out" | "in" -): Promise { +): Promise { if (!redis) return []; try { @@ -129,7 +168,34 @@ export async function drainSessionStreamWaitpoints( if (err || !Array.isArray(members)) continue; for (const m of members as string[]) ids.add(m); } - return [...ids]; + const waitpointIds = [...ids]; + if (waitpointIds.length === 0) return []; + + let formatResults: Awaited> | null = null; + try { + const formatPipeline = redis.multi(); + for (const waitpointId of waitpointIds) { + formatPipeline.get(buildFormatKey(waitpointId)); + formatPipeline.del(buildFormatKey(waitpointId)); + } + formatResults = await formatPipeline.exec(); + } catch (error) { + // The waitpoint ids were already drained. Complete them with raw data + // rather than losing the wake-up because optional metadata was unavailable. + logger.error("Failed to read session stream waitpoint response formats", { + environmentId, + addressingKey, + io, + error, + }); + } + + return waitpointIds.map((id, index) => { + const formatEntry = formatResults?.[index * 2]; + const responseFormat = + formatEntry && !formatEntry[0] && formatEntry[1] === "record-v1" ? "record-v1" : undefined; + return { id, responseFormat }; + }); } catch (error) { logger.error("Failed to drain session stream waitpoint cache", { environmentId, @@ -240,7 +306,10 @@ export async function removeSessionStreamWaitpoint( try { const key = buildKey(environmentId, addressingKey, io); - await redis.srem(key, waitpointId); + const pipeline = redis.multi(); + pipeline.srem(key, waitpointId); + pipeline.del(buildFormatKey(waitpointId)); + await pipeline.exec(); } catch (error) { logger.error("Failed to remove session stream waitpoint cache entry", { environmentId, diff --git a/apps/webapp/app/v3/webhookEngine.server.ts b/apps/webapp/app/v3/webhookEngine.server.ts index d58a89919b9..6fbb2d48b91 100644 --- a/apps/webapp/app/v3/webhookEngine.server.ts +++ b/apps/webapp/app/v3/webhookEngine.server.ts @@ -17,6 +17,7 @@ import { claimSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, + sessionStreamWaitpointOutput, } from "~/services/sessionStreamWaitpointCache.server"; import { getSecretStore } from "~/services/secrets/secretStore.server"; import { singleton } from "~/utils/singleton"; @@ -229,10 +230,12 @@ function createWebhookEngine() { "in", deliveryId ); + let appendSeq: number | undefined; if (wonClaim) { - const [appendError] = await tryCatch( + const [appendError, seqNum] = await tryCatch( realtimeStream.appendPartToSessionStream(part, deliveryId, addressingKey, "in") ); + appendSeq = seqNum ?? undefined; if (appendError) { // Nothing landed — release the claim so a retry re-appends the same id. await releaseSessionStreamPart(environment.id, addressingKey, "in", deliveryId); @@ -245,7 +248,7 @@ function createWebhookEngine() { } // Wake any `.in` waitpoints the run registered (best-effort; the record is durable in S2). - const [drainError, waitpointIds] = await tryCatch( + const [drainError, waitpoints] = await tryCatch( drainSessionStreamWaitpoints(environment.id, addressingKey, "in") ); if (drainError) { @@ -253,13 +256,13 @@ function createWebhookEngine() { externalId, error: drainError, }); - } else if (waitpointIds && waitpointIds.length > 0) { + } else if (waitpoints && waitpoints.length > 0) { await Promise.all( - waitpointIds.map((waitpointId) => + waitpoints.map((waitpoint) => tryCatch( runEngine.completeWaitpoint({ - id: waitpointId, - output: { value: part, type: "application/json", isError: false }, + id: waitpoint.id, + output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq), }) ) ) diff --git a/packages/core/src/v3/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index 6cd100f7c3c..42690f79fca 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -1669,6 +1669,8 @@ export const CreateSessionStreamWaitpointRequestBody = z.object({ * Used to catch data that arrived before `.wait()` was called. */ lastSeqNum: z.number().optional(), + /** Internal capability flag: return the exact record sequence on resume. */ + responseFormat: z.literal("record-v1").optional(), }); export type CreateSessionStreamWaitpointRequestBody = z.infer< typeof CreateSessionStreamWaitpointRequestBody diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 63082516889..7d45c9248f2 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -94,6 +94,14 @@ export class SessionStreamsAPI implements SessionStreamManager { this.#getManager().setLastSeqNum(sessionId, io, seqNum); } + public consumeRecord(sessionId: string, io: SessionChannelIO, seqNum: number): void { + const manager = this.#getManager(); + if (!manager.consumeRecord) { + throw new Error("The configured Session stream manager does not support exact consumption"); + } + manager.consumeRecord(sessionId, io, seqNum); + } + public lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.#getManager().lastDispatchedSeqNum(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.test.ts b/packages/core/src/v3/sessionStreams/manager.test.ts index 99b541fbb3b..4262674bfe3 100644 --- a/packages/core/src/v3/sessionStreams/manager.test.ts +++ b/packages/core/src/v3/sessionStreams/manager.test.ts @@ -386,7 +386,7 @@ describe("StandardSessionStreamManager — record metadata", () => { manager.disconnect(); }); - it("retains cursor barriers when disconnect clears the buffer", async () => { + it("preserves buffered records across disconnect and consumes only the exact sequence", async () => { const manager = new StandardSessionStreamManager( singleShotApiClient([ { @@ -401,6 +401,12 @@ describe("StandardSessionStreamManager — record metadata", () => { chunk: { kind: "stop" }, timestamp: 2000, }, + { + id: "52", + recordId: "message-2", + chunk: { kind: "message", payload: { id: "u2" } }, + timestamp: 3000, + }, ]), "http://localhost" ); @@ -418,11 +424,16 @@ describe("StandardSessionStreamManager — record metadata", () => { expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49); manager.disconnectStream(sessionId, io); - expect(manager.peekRecord(sessionId, io)).toBeUndefined(); + expect(manager.peekRecord(sessionId, io)?.seqNum).toBe(50); expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49); - manager.setLastDispatchedSeqNum(sessionId, io, 51); - expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49); + manager.consumeRecord(sessionId, io, 50); + expect(manager.peekRecord(sessionId, io)?.seqNum).toBe(52); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(51); + + manager.consumeRecord(sessionId, io, 52); + expect(manager.peekRecord(sessionId, io)).toBeUndefined(); + expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(52); manager.reset(); expect(manager.lastDispatchedSeqNum(sessionId, io)).toBeUndefined(); diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index dccb0baade2..c4c1c0503a4 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -67,9 +67,8 @@ export class StandardSessionStreamManager implements SessionStreamManager { private explicitlyDisconnected = new Set(); private seqNums = new Map(); // Sequence numbers for records that were delivered but not consumed. - // Kept separately from `buffer` because `disconnectStream()` clears the - // local buffer before a waitpoint suspension, but those records must still - // hold the persisted consume cursor back. + // Kept separately from `buffer` so the committed cursor can be calculated + // without depending on buffer traversal. private unconsumedSeqNums = new Map>(); // High-water mark of seq_nums that have been *consumed* (delivered to a // once() waiter or shifted off the buffer into a once() caller) on a channel. @@ -282,6 +281,22 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } + consumeRecord(sessionId: string, io: SessionChannelIO, seqNum: number): void { + const key = keyFor(sessionId, io); + const buffered = this.buffer.get(key); + const index = buffered?.findIndex((record) => record.seqNum === seqNum) ?? -1; + + if (buffered && index !== -1) { + buffered.splice(index, 1); + if (buffered.length === 0) { + this.buffer.delete(key); + } + } + + this.#advanceLastDispatched(key, seqNum); + this.#drainOnceWaitersFromBuffer(key); + } + lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { const key = keyFor(sessionId, io); const highWatermark = this.lastDispatchedSeqNums.get(key); @@ -360,7 +375,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { disconnectStream(sessionId: string, io: SessionChannelIO): void { const key = keyFor(sessionId, io); const tail = this.tails.get(key); - const _bufferedSize = this.buffer.get(key)?.length ?? 0; // Mark as explicitly disconnected BEFORE we abort, so the tail's // `.finally` reconnect path sees the flag when it runs (which can be // synchronous in the AbortError catch). Cleared on the next explicit @@ -370,7 +384,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { tail.abortController.abort(); this.tails.delete(key); } - this.buffer.delete(key); // Reset the backoff counter so a future re-attach starts fresh — // an explicit disconnect is a deliberate teardown, not evidence of // a broken backend. @@ -442,15 +455,10 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.tails.delete(key); // If the tail was torn down explicitly via `disconnectStream`, - // honor that — the caller (typically `session.in.wait()`) is - // suspending the run and expects no records to be buffered or - // delivered until a fresh `on()` / `once()` re-attaches. Without - // this guard a run-level persistent handler (e.g. `chat.agent`'s - // `stopInput.on(...)`) would auto-reconnect during the suspend - // window, the resurrected tail would receive the same record the - // waitpoint just delivered, and that record would land in the - // buffer where the next turn's `messagesInput.on(...)` drains it - // and runs a duplicate turn. + // honor that until a fresh `on()` / `once()` re-attaches. Existing + // buffered records stay available across the suspension, but a + // run-level handler must not reconnect and receive another copy of + // the record being delivered through the waitpoint. if (this.explicitlyDisconnected.has(key)) { return; } diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index aeb2a9aeb44..1e5dbaebe9a 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -61,6 +61,8 @@ export class NoopSessionStreamManager implements SessionStreamManager { setLastSeqNum(_sessionId: string, _io: SessionChannelIO, _seqNum: number): void {} + consumeRecord(_sessionId: string, _io: SessionChannelIO, _seqNum: number): void {} + lastDispatchedSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { return undefined; } diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index ce55cbee39a..24b6f084512 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -88,6 +88,9 @@ export interface SessionStreamManager { /** Advance the last-seen sequence number (prevents SSE replay after `.wait` resume). */ setLastSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void; + /** Consume one exact record delivered through the waitpoint path. */ + consumeRecord?(sessionId: string, io: SessionChannelIO, seqNum: number): void; + /** * Highest sequence number that is safe to persist as consumed. When a later * record is handled while an earlier record remains unconsumed, this stays @@ -121,7 +124,7 @@ export interface SessionStreamManager { /** Remove and discard the first buffered record. Returns true if one was removed. */ shiftBuffer(sessionId: string, io: SessionChannelIO): boolean; - /** Abort the SSE tail and clear the buffer. Called before `.wait` suspends. */ + /** Abort the SSE tail while preserving buffered records. Called before `.wait` suspends. */ disconnectStream(sessionId: string, io: SessionChannelIO): void; /** Clear all `.on` handlers; abort tails without pending once-waiters. */ diff --git a/packages/core/src/v3/sessionStreams/wireProtocol.ts b/packages/core/src/v3/sessionStreams/wireProtocol.ts index 550e81a0af4..bb6aef3e1a7 100644 --- a/packages/core/src/v3/sessionStreams/wireProtocol.ts +++ b/packages/core/src/v3/sessionStreams/wireProtocol.ts @@ -40,6 +40,53 @@ export const SESSION_STATE_LAST_EVENT_ID_HEADER = "last-event-id" as const; */ export const SESSION_IN_EVENT_ID_HEADER = "session-in-event-id" as const; +/** + * Opt-in response format for Session stream waitpoints. Older SDKs omit this + * and continue receiving the raw record data. + */ +export const SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT = "record-v1" as const; + +/** Content type used only when a waitpoint actually returns a record-v1 envelope. */ +export const SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE = + "application/vnd.trigger.session-stream-record+json" as const; + +const SESSION_STREAM_WAITPOINT_RECORD_TYPE = "trigger-session-stream-record" as const; + +/** Internal envelope used to return an exact Session record from a waitpoint. */ +export type SessionStreamWaitpointRecord = Readonly<{ + type: typeof SESSION_STREAM_WAITPOINT_RECORD_TYPE; + version: 1; + seqNum: number; + data: unknown; +}>; + +export function serializeSessionStreamWaitpointRecord(data: unknown, seqNum: number): string { + return JSON.stringify({ + type: SESSION_STREAM_WAITPOINT_RECORD_TYPE, + version: 1, + seqNum, + data, + } satisfies SessionStreamWaitpointRecord); +} + +export function parseSessionStreamWaitpointRecord( + value: unknown +): SessionStreamWaitpointRecord | undefined { + if (!value || typeof value !== "object") return undefined; + + const record = value as Partial; + if ( + record.type !== SESSION_STREAM_WAITPOINT_RECORD_TYPE || + record.version !== 1 || + typeof record.seqNum !== "number" || + !Number.isFinite(record.seqNum) + ) { + return undefined; + } + + return record as SessionStreamWaitpointRecord; +} + export const TRIGGER_CONTROL_SUBTYPE = { TURN_COMPLETE: "turn-complete", UPGRADE_REQUIRED: "upgrade-required", diff --git a/packages/core/src/v3/test/session-waitpoint-backend.ts b/packages/core/src/v3/test/session-waitpoint-backend.ts index 8cae877f54e..fd7cd5d3609 100644 --- a/packages/core/src/v3/test/session-waitpoint-backend.ts +++ b/packages/core/src/v3/test/session-waitpoint-backend.ts @@ -1,6 +1,10 @@ import { ApiClient } from "../apiClient/index.js"; import { WaitpointId } from "../isomorphic/friendlyId.js"; import { NoopRuntimeManager } from "../runtime/noopRuntimeManager.js"; +import { + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, +} from "../sessionStreams/wireProtocol.js"; import type { CreateSessionStreamWaitpointRequestBody, CreateSessionStreamWaitpointResponseBody, @@ -13,6 +17,7 @@ type PendingWait = { io: "in" | "out"; lastSeqNum?: number; timeout?: string; + responseFormat?: "record-v1"; abort: AbortController; }; @@ -71,6 +76,7 @@ export class SessionWaitpointBackend { io: body.io, lastSeqNum: body.lastSeqNum, timeout: body.timeout, + responseFormat: body.responseFormat, abort: new AbortController(), }); return { waitpointId, isCached: false }; @@ -113,8 +119,20 @@ export class SessionWaitpointBackend { }; } - const output = typeof result === "string" ? result : JSON.stringify(result); - return { ok: true, output, outputType: "application/json" }; + const output = + pending.responseFormat === "record-v1" + ? serializeSessionStreamWaitpointRecord(result.data, result.seqNum) + : typeof result.data === "string" + ? result.data + : JSON.stringify(result.data); + return { + ok: true, + output, + outputType: + pending.responseFormat === "record-v1" + ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE + : "application/json", + }; } catch { return { ok: false, @@ -144,16 +162,23 @@ export class SessionWaitpointBackend { * which {@link wait} passes straight to the packet parser so it round-trips * to the same object `session.in.once()` returns. */ - private async readNextRecord(pending: PendingWait): Promise { + private async readNextRecord(pending: PendingWait): Promise<{ data: unknown; seqNum: number }> { const lastEventId = pending.lastSeqNum !== undefined && pending.lastSeqNum >= 0 ? String(pending.lastSeqNum) : undefined; + let deliveredSeqNum: number | undefined; const stream = await this.apiClient.subscribeToSessionStream(pending.session, pending.io, { lastEventId, signal: pending.abort.signal, timeoutInSeconds: 120, + onPart: (part) => { + const seqNum = Number.parseInt(part.id, 10); + if (Number.isFinite(seqNum)) { + deliveredSeqNum = seqNum; + } + }, }); const reader = stream.getReader(); @@ -162,7 +187,10 @@ export class SessionWaitpointBackend { if (done) { throw new Error("session stream closed"); } - return value; + if (deliveredSeqNum === undefined) { + throw new Error("session stream record is missing its sequence number"); + } + return { data: value, seqNum: deliveredSeqNum }; } finally { await reader.cancel().catch(() => {}); pending.abort.abort(); diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index 52de0b1571d..6d8ad0b5536 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -209,6 +209,22 @@ export class TestSessionStreamManager implements SessionStreamManager { this.seqNums.set(keyFor(sessionId, io), seqNum); } + consumeRecord(sessionId: string, io: SessionChannelIO, seqNum: number): void { + const key = keyFor(sessionId, io); + const buffered = this.buffer.get(key); + const index = buffered?.findIndex((record) => record.seqNum === seqNum) ?? -1; + + if (buffered && index !== -1) { + buffered.splice(index, 1); + if (buffered.length === 0) { + this.buffer.delete(key); + } + } + + this.#advanceLastDispatched(key, seqNum); + this.#drainOnceWaitersFromBuffer(key); + } + lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { const key = keyFor(sessionId, io); const highWatermark = this.dispatchedSeqNums.get(key); @@ -282,8 +298,9 @@ export class TestSessionStreamManager implements SessionStreamManager { return false; } - disconnectStream(sessionId: string, io: SessionChannelIO): void { - this.buffer.delete(keyFor(sessionId, io)); + disconnectStream(_sessionId: string, _io: SessionChannelIO): void { + // The production manager keeps buffered records reachable across a + // waitpoint suspension. The exact waitpoint record is removed on resume. } clearHandlers(): void { diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index 8130d333700..125991a0c44 100644 --- a/packages/trigger-sdk/src/v3/sessions.ts +++ b/packages/trigger-sdk/src/v3/sessions.ts @@ -25,6 +25,8 @@ import type { import { InputStreamOncePromise, ManualWaitpointPromise, + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT, SemanticInternalAttributes, SessionStreamInstance, WaitpointTimeoutError, @@ -32,6 +34,7 @@ import { apiClientManager, ensureReadableStream, mergeRequestOptions, + parseSessionStreamWaitpointRecord, runtime, sessionStreams, taskContext, @@ -713,6 +716,7 @@ export class SessionInputChannel { const apiClient = apiClientManager.clientOrThrow(); + const lastConsumedSeqNum = sessionStreams.lastDispatchedSeqNum(this.sessionId, "in"); const response = await apiClient.createSessionStreamWaitpoint(ctx.run.id, { session: this.sessionId, io: "in", @@ -720,7 +724,8 @@ export class SessionInputChannel { idempotencyKey: options?.idempotencyKey, idempotencyKeyTTL: options?.idempotencyKeyTTL, tags: options?.tags, - lastSeqNum: sessionStreams.lastSeqNum(this.sessionId, "in"), + lastSeqNum: lastConsumedSeqNum, + responseFormat: SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT, }); const result = await tracer.startActiveSpan( @@ -735,36 +740,77 @@ export class SessionInputChannel { throw new Error("Failed to block on session stream waitpoint"); } - // Drop the SSE tail + buffer before suspending so the record - // delivered via the waitpoint path isn't re-buffered on resume. + // Stop the SSE tail before suspending. Buffered records stay in + // place; the exact record returned by the waitpoint is removed on + // resume, while any later records remain available to consumers. sessionStreams.disconnectStream(this.sessionId, "in"); const waitResult = await runtime.waitUntil(response.waitpointId); + const hasRecordEnvelope = + waitResult.outputType === SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE; - const data = + const parsedOutput = waitResult.output !== undefined ? await conditionallyImportAndParsePacket( { data: waitResult.output, - dataType: waitResult.outputType ?? "application/json", + dataType: hasRecordEnvelope + ? "application/json" + : (waitResult.outputType ?? "application/json"), }, apiClient ) : undefined; if (waitResult.ok) { - // Advance both cursors past the record consumed via the - // waitpoint: the seq counter so the SSE tail doesn't replay - // it, and the consume cursor so turn-completes don't stamp a - // stale `session-in-event-id`. - const prevSeq = sessionStreams.lastSeqNum(this.sessionId, "in"); - const nextSeq = (prevSeq ?? -1) + 1; - sessionStreams.setLastSeqNum(this.sessionId, "in", nextSeq); - sessionStreams.setLastDispatchedSeqNum(this.sessionId, "in", nextSeq); + const record = hasRecordEnvelope + ? parseSessionStreamWaitpointRecord(parsedOutput) + : undefined; + let seqNum = record?.seqNum; + const data = record + ? await conditionallyImportAndParsePacket( + { + data: + typeof record.data === "string" ? record.data : JSON.stringify(record.data), + dataType: "application/json", + }, + apiClient + ) + : parsedOutput; + + // Older servers return only raw data. Recover its durable + // sequence from the channel instead of guessing and risking a + // cursor that skips or strands another record. + if (seqNum === undefined && waitResult.output !== undefined) { + try { + const response = await apiClient.readSessionStreamRecords(this.sessionId, "in", { + afterEventId: + lastConsumedSeqNum !== undefined ? String(lastConsumedSeqNum) : undefined, + }); + const matchingRecords = response.records.filter( + (candidate) => + candidate.data === waitResult.output || + (typeof candidate.data !== "string" && + JSON.stringify(candidate.data) === JSON.stringify(parsedOutput)) + ); + if (matchingRecords.length === 1) { + seqNum = matchingRecords[0]!.seqNum; + } + } catch { + // Leave the cursor behind when an older server cannot + // provide record metadata. At-least-once replay is safer + // than acknowledging an unknown sequence. + } + } + + if (seqNum !== undefined) { + sessionStreams.consumeRecord(this.sessionId, "in", seqNum); + sessionStreams.setLastSeqNum(this.sessionId, "in", seqNum); + } return { ok: true as const, output: data as T }; } else { - const error = new WaitpointTimeoutError(data?.message ?? "Timed out"); + const error = new WaitpointTimeoutError(parsedOutput?.message ?? "Timed out"); span.recordException(error); span.setStatus({ code: SpanStatusCode.ERROR }); return { ok: false as const, error }; diff --git a/packages/trigger-sdk/test/pending-message-drain.test.ts b/packages/trigger-sdk/test/pending-message-drain.test.ts index f5bd7057515..18ef2462bf2 100644 --- a/packages/trigger-sdk/test/pending-message-drain.test.ts +++ b/packages/trigger-sdk/test/pending-message-drain.test.ts @@ -5,7 +5,12 @@ import { mockChatAgent } from "../src/v3/test/index.js"; import { describe, expect, it, vi } from "vitest"; import { chat } from "../src/v3/ai.js"; import { __setSessionOpenImplForTests, sessions } from "../src/v3/sessions.js"; -import { apiClientManager, sessionStreams } from "@trigger.dev/core/v3"; +import { + apiClientManager, + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, + sessionStreams, +} from "@trigger.dev/core/v3"; import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; import { simulateReadableStream, streamText } from "ai"; import { MockLanguageModelV3 } from "ai/test"; @@ -70,6 +75,21 @@ async function waitFor(check: () => boolean, timeoutMs = 10_000) { throw new Error("waitFor timed out"); } +function runtimeWithWaitpointOutput(output: string, outputType = "application/json") { + return { + disable() {}, + waitForTask() { + throw new Error("Unexpected task wait"); + }, + waitForBatch() { + throw new Error("Unexpected batch wait"); + }, + waitForWaitpoint() { + return Promise.resolve({ ok: true, output, outputType }); + }, + }; +} + function streamedText(harness: { allChunks: unknown[] }): string { return (harness.allChunks as { type?: string; delta?: string }[]) .filter((c) => c.type === "text-delta") @@ -248,26 +268,179 @@ describe("chat.createSession stop + immediate send", () => { }); describe("session.in.wait() consume cursor", () => { - it("advances lastDispatchedSeqNum alongside lastSeqNum on waitpoint delivery", async () => { + it("keeps later input reachable across the suspend-and-resume race", async () => { __setSessionOpenImplForTests(undefined); - await runInMockTaskContext(async () => { - vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - createSessionStreamWaitpoint: async () => ({ waitpointId: "wp_test_1" }), - waitForWaitpointToken: async () => ({ success: true }), - } as never); - - const sessionId = "cursor-sess"; - // Simulate records 0..4 already received via SSE before the suspend. - sessionStreams.setLastSeqNum(sessionId, "in", 4); - - const result = await sessions.open(sessionId).in.wait(); - - expect(result.ok).toBe(true); - expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(5); - // The waitpoint-delivered record was consumed by this caller, so the - // committed-consume cursor (what turn-completes persist as - // `session-in-event-id`) must advance with it. - expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(5); - }); + const first = { kind: "message", payload: { id: "u1" } }; + const later = { kind: "message", payload: { id: "u2" } }; + const runtimeManager = runtimeWithWaitpointOutput( + serializeSessionStreamWaitpointRecord(JSON.stringify(first), 50), + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE + ); + let registeredLastSeqNum: number | undefined; + let registeredResponseFormat: string | undefined; + + await runInMockTaskContext( + async (drivers) => { + const sessionId = "cursor-sess"; + const channel = sessions.open(sessionId).in; + const stop = channel.on<{ kind: string }>((record) => record.kind === "stop"); + + sessionStreams.setLastSeqNum(sessionId, "in", 49); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 49); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async ( + _runId: string, + body: { lastSeqNum?: number; responseFormat?: string } + ) => { + registeredLastSeqNum = body.lastSeqNum; + registeredResponseFormat = body.responseFormat; + return { + waitpointId: "wp_test_1", + isCached: false, + }; + }, + waitForWaitpointToken: async () => { + // These records land after registration but before the tail is + // disconnected. The waitpoint resolves with seq 50, while the + // local tail has already consumed 51 and buffered 52. + await drivers.sessions.in.send(sessionId, first, "in", { seqNum: 50 }); + await drivers.sessions.in.send(sessionId, { kind: "stop" }, "in", { seqNum: 51 }); + await drivers.sessions.in.send(sessionId, later, "in", { seqNum: 52 }); + return { success: true }; + }, + } as never); + + const result = await channel.wait(); + + expect(result).toEqual({ ok: true, output: first }); + expect(registeredLastSeqNum).toBe(49); + expect(registeredResponseFormat).toBe("record-v1"); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(51); + expect(sessionStreams.peekRecord(sessionId, "in")?.seqNum).toBe(52); + + const next = await sessionStreams.onceRecord(sessionId, "in"); + expect(next).toEqual({ + ok: true, + output: { id: "test-record-52", seqNum: 52, data: later }, + }); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(52); + stop.off(); + }, + { runtimeManager } + ); + }); + + it("recovers the exact sequence from durable records for older servers", async () => { + __setSessionOpenImplForTests(undefined); + const payload = { kind: "message", payload: { id: "legacy" } }; + const rawPayload = JSON.stringify(payload); + const runtimeManager = runtimeWithWaitpointOutput(rawPayload); + let afterEventId: string | undefined; + + await runInMockTaskContext( + async () => { + const sessionId = "legacy-cursor-sess"; + sessionStreams.setLastSeqNum(sessionId, "in", 6); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ + waitpointId: "wp_legacy", + isCached: false, + }), + waitForWaitpointToken: async () => ({ success: true }), + readSessionStreamRecords: async ( + _sessionId: string, + _io: "in" | "out", + options?: { afterEventId?: string } + ) => { + afterEventId = options?.afterEventId; + return { + records: [{ id: "legacy-record", seqNum: 7, data: rawPayload }], + }; + }, + } as never); + + const result = await sessions.open(sessionId).in.wait(); + + expect(result).toEqual({ ok: true, output: payload }); + expect(afterEventId).toBe("6"); + expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(7); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); + }, + { runtimeManager } + ); + }); + + it("does not mistake an older server's user payload for the internal envelope", async () => { + __setSessionOpenImplForTests(undefined); + const payload = { + type: "trigger-session-stream-record", + version: 1, + seqNum: 999, + data: { user: "supplied" }, + }; + const rawPayload = JSON.stringify(payload); + + await runInMockTaskContext( + async () => { + const sessionId = "legacy-envelope-collision"; + sessionStreams.setLastSeqNum(sessionId, "in", 6); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ + waitpointId: "wp_legacy_collision", + isCached: false, + }), + waitForWaitpointToken: async () => ({ success: true }), + readSessionStreamRecords: async () => ({ + records: [{ id: "legacy-record", seqNum: 7, data: rawPayload }], + }), + } as never); + + const result = await sessions.open(sessionId).in.wait(); + + expect(result).toEqual({ ok: true, output: payload }); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); + }, + { runtimeManager: runtimeWithWaitpointOutput(rawPayload) } + ); + }); + + it("leaves the cursor behind when legacy payload matching is ambiguous", async () => { + __setSessionOpenImplForTests(undefined); + const payload = { kind: "message", payload: { id: "duplicate" } }; + const rawPayload = JSON.stringify(payload); + + await runInMockTaskContext( + async () => { + const sessionId = "legacy-duplicate-payload"; + sessionStreams.setLastSeqNum(sessionId, "in", 6); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ + waitpointId: "wp_legacy_duplicate", + isCached: false, + }), + waitForWaitpointToken: async () => ({ success: true }), + readSessionStreamRecords: async () => ({ + records: [ + { id: "duplicate-1", seqNum: 7, data: rawPayload }, + { id: "duplicate-2", seqNum: 8, data: rawPayload }, + ], + }), + } as never); + + const result = await sessions.open(sessionId).in.wait(); + + expect(result).toEqual({ ok: true, output: payload }); + expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(6); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(6); + }, + { runtimeManager: runtimeWithWaitpointOutput(rawPayload) } + ); }); }); From f4d18266306726fd4bcd6f22b6eea871f0176274 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 14:47:17 -0700 Subject: [PATCH 05/20] docs(chat): clarify non-blocking mailbox reads --- docs/ai-chat/custom-agents.mdx | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 59607aa64f3..cb1495fff44 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -238,6 +238,10 @@ message that `next()` can consume immediately. It does not query the remote Session channel or start a subscription. Use `waitWithIdleTimeout()` when the loop needs to idle until future input arrives. +`next({ timeoutInSeconds: 0 })` is also a local, non-blocking read. Call +`next()` without a timeout, or with a positive timeout, to subscribe for future +input. + `next()` returns a readonly record envelope: ```ts From 7655a5cbccc99f77ec8f1b9ab206b05b40b7ee1b Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 20 Aug 2026 15:13:34 +0100 Subject: [PATCH 06/20] fix(chat,sdk): stop an unconsumed control record wedging the mailbox `chat.messages.next()` and `hasPending()` only inspect the head of the `.in` buffer. A control record whose kind has no consumer on this boot therefore parked at the head forever, and every message queued behind it became undeliverable with no error: `hasPending()` stayed false and `next()` timed out on every call. Records of a kind nothing on the run consumes are now discarded at dispatch instead. Consuming at dispatch keeps the resume cursor exact, since the record never enters the buffer and so leaves no unconsumed barrier for `lastDispatchedSeqNum()` to clamp behind. `message` is always claimed, and handover kinds are claimed for the window in which a handover-prepare boot is actually waiting for them. The mixed-kinds test now builds its blocked-head state on a handover-prepare boot, where the handover kind is claimed and so stays buffered for the raw read it asserts. --- .changeset/tidy-mailboxes-wait.md | 2 + docs/ai-chat/custom-agents.mdx | 5 + packages/trigger-sdk/src/v3/ai.ts | 126 ++++++++++++++++-- .../test/chat-messages-mailbox.test.ts | 6 +- 4 files changed, 128 insertions(+), 11 deletions(-) diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md index 99013758d35..9535861607d 100644 --- a/.changeset/tidy-mailboxes-wait.md +++ b/.changeset/tidy-mailboxes-wait.md @@ -4,3 +4,5 @@ --- Custom agent loops can now inspect pending chat input without consuming it and consume one mailbox record at a time with `chat.messages.hasPending()` and `chat.messages.next()`. Mailbox records include stable identifiers for tracing and redelivery. + +A control record that nothing on the run consumes is now discarded rather than left at the head of the `.in` channel, where it would have made every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` now always means the mailbox is idle. diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index cb1495fff44..2f68014dc4f 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -267,6 +267,11 @@ comes before a message, `hasPending()` stays `false` and `next()` leaves the control record for its own consumer. After that record is handled, the message becomes pending. +A control record that nothing on the run consumes is discarded rather than left +at the head of the channel. `hasPending()` and `next()` only look at the head, so +a record parked there would make every message behind it undeliverable. This +means `next()` returning `undefined` always means the mailbox is idle. + A complete loop: ```ts trigger/my-chat-raw.ts diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 95b0f5267da..2b51b05a55c 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1885,17 +1885,118 @@ async function waitForHandover(options: { spanName?: string; }): Promise { if (options.payload.trigger !== "handover-prepare") return null; - const result = await handoverInput.waitWithIdleTimeout({ - idleTimeoutInSeconds: - options.idleTimeoutInSeconds ?? options.payload.idleTimeoutInSeconds ?? 60, - timeout: options.timeout, - spanName: options.spanName ?? "waiting for handover signal", + try { + const result = await handoverInput.waitWithIdleTimeout({ + idleTimeoutInSeconds: + options.idleTimeoutInSeconds ?? options.payload.idleTimeoutInSeconds ?? 60, + timeout: options.timeout, + spanName: options.spanName ?? "waiting for handover signal", + }); + // Non-ok = idle timeout or the warm handler crashed without signaling. + if (!result.ok) return null; + return result.output; + } finally { + // The handover window is over either way. A signal arriving after this + // point has no consumer, so hand it to the drain rather than letting it + // park at the head of the channel. + releaseChatInputKinds(CHAT_HANDOVER_KINDS); + } +} + +/** + * Record kinds on `session.in` that some consumer on THIS boot is responsible + * for. `"message"` is always claimed; handover kinds are claimed only for the + * window in which `waitForHandover` is actually waiting for them. + * + * Anything not in this set has no consumer on this boot, so leaving it buffered + * would park it at the head of the channel forever — `chat.messages.next()` and + * `hasPending()` only ever inspect the head, so every record queued behind it + * becomes undeliverable with no error. The drain below discards unclaimed kinds + * instead. + * @internal + */ +const chatClaimedKindsKey = locals.create>("chat.claimedKinds"); + +/** Kinds carried on `.in` that are not user messages. @internal */ +const CHAT_HANDOVER_KINDS = ["handover", "handover-skip"] as const; + +/** The run's attached drain subscription, so it can be re-offered the buffer. @internal */ +const chatInputDrainKey = locals.create<{ off: () => void }>("chat.inputDrain"); + +function chatClaimedKinds(): Set { + let claimed = locals.get(chatClaimedKindsKey); + if (!claimed) { + claimed = new Set(["message"]); + locals.set(chatClaimedKindsKey, claimed); + } + return claimed; +} + +/** + * Attach the unclaimed-control drain for this run. + * + * Consuming at dispatch (returning `true`) is what makes this safe for the + * resume cursor: the record is never buffered, so it leaves no unconsumed + * marker and `lastDispatchedSeqNum()` stays exact rather than being clamped + * behind a record nobody will ever take. + * + * `#dispatch` resolves a matching `once()` waiter BEFORE invoking handlers, so + * this can never take a record out from under a claimed consumer that is + * actively waiting for it. + * + * MUST be attached after `seedSessionInResumeCursorForCustomLoop`, like every + * other `.in` listener — attaching first would replay from seq 0. + * @internal + */ +function attachUnclaimedChatInputDrain(): { off: () => void } { + return getChatSession().in.on((chunk) => { + const kind = (chunk as { kind?: unknown } | undefined)?.kind; + // Malformed record: nothing can consume it, so don't let it wedge the head. + if (typeof kind !== "string") { + logger.warn("chat: discarded a malformed session.in record with no usable kind"); + return true; + } + if (chatClaimedKinds().has(kind)) return undefined; + logger.warn("chat: discarded a session.in record that no consumer handled on this boot", { + kind, + }); + return true; }); - // Non-ok = idle timeout or the warm handler crashed without signaling. - if (!result.ok) return null; - return result.output; } +/** + * Release claimed kinds and re-offer the buffer to the drain. + * + * Re-attaching is the sweep: `on()` re-offers every buffered record to the + * newly attached handler, so a record that was buffered while its kind was + * still claimed (the waiter-gap window between `once()` iterations) is + * discarded now and the cursor advances past it. + * @internal + */ +function releaseChatInputKinds(kinds: readonly string[]): void { + const claimed = chatClaimedKinds(); + let changed = false; + for (const kind of kinds) { + if (claimed.delete(kind)) changed = true; + } + if (!changed) return; + + const drain = locals.get(chatInputDrainKey); + if (!drain) return; + drain.off(); + locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); +} + +/** + * Declare that this run will consume the given `session.in` record kinds + * itself (via raw `session.in` reads). Claimed kinds are never discarded by + * the unclaimed-control drain: they stay buffered, block + * `chat.messages.next()` at the head of the channel, and hold the resume + * cursor behind them until consumed. Call `release()` when the loop stops + * consuming them — any still-buffered records of those kinds are then + * discarded and the cursor advances. + */ + /** * Per-turn deferred promises. Registered via `chat.defer()`, awaited * before `onTurnComplete` fires. Reset each turn. @@ -5428,6 +5529,14 @@ function chatCustomAgent< // listener — otherwise a continuation boot replays already-answered // messages into the loop's first wait. await seedSessionInResumeCursorForCustomLoop(payload); + // Claim the kinds this boot actually has a consumer for, then attach the + // drain for everything else. Handover kinds are only claimed on a + // handover-prepare boot, which is the only boot that waits for them. + const claimed = chatClaimedKinds(); + if (payload.trigger === "handover-prepare") { + for (const kind of CHAT_HANDOVER_KINDS) claimed.add(kind); + } + locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); return userRun(payload, runOptions); }, }); @@ -10765,6 +10874,7 @@ export const chat = { response: chatResponse, /** Pre-built input stream for receiving messages from the transport. */ messages: messagesInput, + /** Declare `session.in` record kinds this run consumes itself. See {@link chatClaimInputKinds}. */ /** Create a managed stop signal wired to the stop input stream. See {@link createStopSignal}. */ createStopSignal, /** Signal the frontend that the current turn is complete. See {@link chatWriteTurnComplete}. */ diff --git a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts index 295b8f6c551..a6d90744390 100644 --- a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -64,7 +64,7 @@ describe("chat.messages mailbox", () => { await runInMockTaskContext(async (drivers) => { const runPromise = run( - { chatId, trigger: "preload" }, + { chatId, trigger: "handover-prepare" }, { ctx: drivers.ctx, signal: new AbortController().signal } ); await ready.promise; @@ -156,7 +156,7 @@ describe("chat.messages mailbox", () => { await runInMockTaskContext(async (drivers) => { const runPromise = run( - { chatId, trigger: "preload" }, + { chatId, trigger: "handover-prepare" }, { ctx: drivers.ctx, signal: new AbortController().signal } ); await ready.promise; @@ -229,7 +229,7 @@ describe("chat.messages mailbox", () => { await runInMockTaskContext(async (drivers) => { const runPromise = run( - { chatId, trigger: "preload" }, + { chatId, trigger: "handover-prepare" }, { ctx: drivers.ctx, signal: new AbortController().signal } ); await ready.promise; From 0eb932f77ee8d4c9004685ec8c4b4e4b4c8ab85a Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 20 Aug 2026 15:31:04 +0100 Subject: [PATCH 07/20] fix(chat,sdk): drop dead claim-kind docs and quiet the drain on known kinds Removes a JSDoc block and a `{@link}` reference to a claim-kinds helper that is not part of this change, which also left `chat.createStopSignal` carrying two doc comments. The drain also warned on every record it discarded, including a stop that the stop facade had already handled: all handlers are invoked for a record regardless of whether an earlier one consumed it, so a stop with an active stop signal is both aborted and drained. A known kind with no active consumer is an expected state, so the warning is now limited to kinds this SDK version does not recognise, which is the case that indicates a newer server. Also corrects the docs and changeset, which claimed `next()` returning `undefined` always means the mailbox is idle. A control record that does have its own consumer can sit at the head while `next()` times out. --- .changeset/tidy-mailboxes-wait.md | 2 +- docs/ai-chat/custom-agents.mdx | 7 +++++-- packages/trigger-sdk/src/v3/ai.ts | 30 ++++++++++++++++-------------- 3 files changed, 22 insertions(+), 17 deletions(-) diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md index 9535861607d..24114e838c3 100644 --- a/.changeset/tidy-mailboxes-wait.md +++ b/.changeset/tidy-mailboxes-wait.md @@ -5,4 +5,4 @@ Custom agent loops can now inspect pending chat input without consuming it and consume one mailbox record at a time with `chat.messages.hasPending()` and `chat.messages.next()`. Mailbox records include stable identifiers for tracing and redelivery. -A control record that nothing on the run consumes is now discarded rather than left at the head of the `.in` channel, where it would have made every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` now always means the mailbox is idle. +A control record that nothing on the run consumes is now discarded rather than left at the head of the `.in` channel, where it would have made every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` means no message became consumable before the timeout. diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 2f68014dc4f..8c7a33b5dbd 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -269,8 +269,11 @@ becomes pending. A control record that nothing on the run consumes is discarded rather than left at the head of the channel. `hasPending()` and `next()` only look at the head, so -a record parked there would make every message behind it undeliverable. This -means `next()` returning `undefined` always means the mailbox is idle. +a record parked there would make every message behind it undeliverable. + +`next()` still returns `undefined` whenever no message became consumable before +the timeout, including while a control record that does have its own consumer +sits at the head. A complete loop: diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index a8124e5f9a0..99752c1b48b 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1920,6 +1920,19 @@ const chatClaimedKindsKey = locals.create>("chat.claimedKinds"); /** Kinds carried on `.in` that are not user messages. @internal */ const CHAT_HANDOVER_KINDS = ["handover", "handover-skip"] as const; +/** + * Every `ChatInputChunk` kind this SDK version knows about. A known kind with + * no active consumer on this boot is an expected, documented state; an unknown + * one means a newer server is sending something this worker cannot handle, + * which is worth surfacing. + * @internal + */ +const KNOWN_CHAT_INPUT_KINDS: ReadonlySet = new Set([ + "message", + "stop", + ...CHAT_HANDOVER_KINDS, +]); + /** The run's attached drain subscription, so it can be re-offered the buffer. @internal */ const chatInputDrainKey = locals.create<{ off: () => void }>("chat.inputDrain"); @@ -1957,9 +1970,9 @@ function attachUnclaimedChatInputDrain(): { off: () => void } { return true; } if (chatClaimedKinds().has(kind)) return undefined; - logger.warn("chat: discarded a session.in record that no consumer handled on this boot", { - kind, - }); + if (!KNOWN_CHAT_INPUT_KINDS.has(kind)) { + logger.warn("chat: discarded a session.in record of an unrecognised kind", { kind }); + } return true; }); } @@ -1987,16 +2000,6 @@ function releaseChatInputKinds(kinds: readonly string[]): void { locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); } -/** - * Declare that this run will consume the given `session.in` record kinds - * itself (via raw `session.in` reads). Claimed kinds are never discarded by - * the unclaimed-control drain: they stay buffered, block - * `chat.messages.next()` at the head of the channel, and hold the resume - * cursor behind them until consumed. Call `release()` when the loop stops - * consuming them — any still-buffered records of those kinds are then - * discarded and the cursor advances. - */ - /** * Per-turn deferred promises. Registered via `chat.defer()`, awaited * before `onTurnComplete` fires. Reset each turn. @@ -10861,7 +10864,6 @@ export const chat = { response: chatResponse, /** Pre-built input stream for receiving messages from the transport. */ messages: messagesInput, - /** Declare `session.in` record kinds this run consumes itself. See {@link chatClaimInputKinds}. */ /** Create a managed stop signal wired to the stop input stream. See {@link createStopSignal}. */ createStopSignal, /** Signal the frontend that the current turn is complete. See {@link chatWriteTurnComplete}. */ From e5c71237f35980c860982d6ac87224f72a504e73 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 20 Aug 2026 16:54:22 +0100 Subject: [PATCH 08/20] refactor(chat,sdk): resume from the channel instead of the waitpoint payload `session.in.wait()` needed the exact sequence of the record it returned, so that record could be acknowledged rather than guessed at. It got that by having the server attach the sequence to the waitpoint output, which meant a new versioned wire format, a capability flag on the waitpoint request, a second Redis key so a mid-deploy instance could still drain, and a payload-matching fallback for a server that does not send the envelope. That fallback gave up whenever two records on a channel shared a payload, and gave up by returning the record to the caller without acknowledging it, so the reconnecting tail delivered it a second time. None of that is necessary. The append route commits the record to the channel before it drains any waitpoint, so once the run wakes, the record is durably readable from the channel with its real sequence. The waitpoint is now treated as a wake signal only: its output is discarded, the tail re-attaches, and the record is read back through the normal buffer path, which acknowledges it and advances the cursor exactly. This removes the wire format, the capability flag, the format key, the fallback, and every webapp change, leaving no cross-service surface and no mixed-version behaviour to reason about. It also fixes the case the fallback could not: a duplicate append whose idempotency claim was lost wakes the run, the channel has nothing new, and the run waits instead of answering a stale record twice. Adds a regression test that the delivered record is acknowledged when identical payloads repeat on a channel, and one that a message queued behind a control record nothing consumes is still delivered. Drops the three tests that only described the removed envelope and its fallback. --- ...uns.$runFriendlyId.session-streams.wait.ts | 15 +-- ...ealtime.v1.sessions.$session.$io.append.ts | 17 ++- ...ealtime.v1.sessions.$session.$io.append.ts | 21 +-- .../sessionStreamWaitpointCache.server.ts | 79 +---------- apps/webapp/app/v3/webhookEngine.server.ts | 15 +-- packages/core/src/v3/schemas/api.ts | 2 - .../src/v3/sessionStreams/wireProtocol.ts | 47 ------- .../src/v3/test/session-waitpoint-backend.ts | 26 +--- packages/trigger-sdk/src/v3/sessions.ts | 113 ++++++---------- .../test/chat-messages-mailbox.test.ts | 49 +++++++ .../test/pending-message-drain.test.ts | 125 +++--------------- 11 files changed, 152 insertions(+), 357 deletions(-) diff --git a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts index 3cb98f5847e..c00ff51b3be 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts @@ -1,8 +1,6 @@ import { json } from "@remix-run/server-runtime"; import { CreateSessionStreamWaitpointRequestBody, - SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, - serializeSessionStreamWaitpointRecord, type CreateSessionStreamWaitpointResponseBody, } from "@trigger.dev/core/v3"; import { WaitpointId } from "@trigger.dev/core/v3/isomorphic"; @@ -127,8 +125,7 @@ const { action, loader } = createActionApiRoute( addressingKey, body.io, result.waitpoint.id, - ttlMs && ttlMs > 0 ? ttlMs : undefined, - body.responseFormat + ttlMs && ttlMs > 0 ? ttlMs : undefined ); // Race-check. If a record landed on the channel before this @@ -158,14 +155,8 @@ const { action, loader } = createActionApiRoute( await engine.completeWaitpoint({ id: result.waitpoint.id, output: { - value: - body.responseFormat === "record-v1" - ? serializeSessionStreamWaitpointRecord(record.data, record.seqNum) - : record.data, - type: - body.responseFormat === "record-v1" - ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE - : "application/json", + value: record.data, + type: "application/json", isError: false, }, }); diff --git a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts index 7ff85cde863..d4dd1d9f19f 100644 --- a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts @@ -15,7 +15,6 @@ import { claimSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, - sessionStreamWaitpointOutput, } from "~/services/sessionStreamWaitpointCache.server"; import { anyResource, createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { engine } from "~/v3/runEngine.server"; @@ -202,7 +201,7 @@ const { action, loader } = createActionApiRoute( // keyed on the canonical addressing key the agent registered with via // `sessions.open(...).in.wait()`, so writers and readers converge // regardless of which URL form they used. - const [drainError, waitpoints] = await tryCatch( + const [drainError, waitpointIds] = await tryCatch( drainSessionStreamWaitpoints(authentication.environment.id, addressingKey, params.io) ); if (drainError) { @@ -211,20 +210,24 @@ const { action, loader } = createActionApiRoute( io: params.io, error: drainError, }); - } else if (waitpoints && waitpoints.length > 0) { + } else if (waitpointIds && waitpointIds.length > 0) { await Promise.all( - waitpoints.map(async (waitpoint) => { + waitpointIds.map(async (waitpointId) => { const [completeError] = await tryCatch( engine.completeWaitpoint({ - id: waitpoint.id, - output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq), + id: waitpointId, + output: { + value: part, + type: "application/json", + isError: false, + }, }) ); if (completeError) { logger.error("Failed to complete session stream waitpoint", { addressingKey, io: params.io, - waitpointId: waitpoint.id, + waitpointId, error: completeError, }); } diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts index 35bdf3a5dd9..ab318f31c7a 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts @@ -13,10 +13,7 @@ import { resolveSessionByIdOrExternalId, } from "~/services/realtime/sessions.server"; import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server"; -import { - drainSessionStreamWaitpoints, - sessionStreamWaitpointOutput, -} from "~/services/sessionStreamWaitpointCache.server"; +import { drainSessionStreamWaitpoints } from "~/services/sessionStreamWaitpointCache.server"; import { requireUserId } from "~/services/session.server"; import { EnvironmentParamSchema } from "~/utils/pathBuilder"; import { engine } from "~/v3/runEngine.server"; @@ -117,7 +114,7 @@ export async function action({ request, params }: ActionFunctionArgs) { // Drain any waitpoints registered for this channel — same as the // public append. Best-effort; failure doesn't fail the append. - const [drainError, waitpoints] = await tryCatch( + const [drainError, waitpointIds] = await tryCatch( drainSessionStreamWaitpoints(environment.id, addressingKey, io) ); if (drainError) { @@ -126,20 +123,24 @@ export async function action({ request, params }: ActionFunctionArgs) { io, error: drainError, }); - } else if (waitpoints && waitpoints.length > 0) { + } else if (waitpointIds && waitpointIds.length > 0) { await Promise.all( - waitpoints.map(async (waitpoint) => { + waitpointIds.map(async (waitpointId) => { const [completeError] = await tryCatch( engine.completeWaitpoint({ - id: waitpoint.id, - output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq ?? undefined), + id: waitpointId, + output: { + value: part, + type: "application/json", + isError: false, + }, }) ); if (completeError) { logger.error("Failed to complete session stream waitpoint (playground)", { addressingKey, io, - waitpointId: waitpoint.id, + waitpointId, error: completeError, }); } diff --git a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts index 0c21c10be1e..7b53042d8d3 100644 --- a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts +++ b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts @@ -1,9 +1,5 @@ import { Redis } from "ioredis"; import { defaultReconnectOnError } from "@internal/redis"; -import { - SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, - serializeSessionStreamWaitpointRecord, -} from "@trigger.dev/core/v3"; import { env } from "~/env.server"; import { singleton } from "~/utils/singleton"; import { logger } from "./logger.server"; @@ -17,35 +13,12 @@ import { logger } from "./logger.server"; // is shared — without it, two environments using the same externalId // would drain each other's waitpoints. const KEY_PREFIX = "ssw:"; -const FORMAT_KEY_PREFIX = "sswf:"; const DEFAULT_TTL_MS = 7 * 24 * 60 * 60 * 1000; // 7 days -export type SessionStreamWaitpoint = { - id: string; - responseFormat?: "record-v1"; -}; - -export function sessionStreamWaitpointOutput( - waitpoint: SessionStreamWaitpoint, - data: string, - seqNum: number | undefined -): { value: string; type: string; isError: false } { - const hasRecordEnvelope = waitpoint.responseFormat === "record-v1" && seqNum !== undefined; - return { - value: hasRecordEnvelope ? serializeSessionStreamWaitpointRecord(data, seqNum) : data, - type: hasRecordEnvelope ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE : "application/json", - isError: false, - }; -} - function buildKey(environmentId: string, addressingKey: string, io: "out" | "in"): string { return `${KEY_PREFIX}${environmentId}:${addressingKey}:${io}`; } -function buildFormatKey(waitpointId: string): string { - return `${FORMAT_KEY_PREFIX}${waitpointId}`; -} - // Pre-env-scoping key format, drained for one release so waitpoints from the // previous deploy still wake. Removable once this has been live > turn timeout. function buildLegacyKey(addressingKey: string, io: "out" | "in"): string { @@ -108,25 +81,13 @@ export async function addSessionStreamWaitpoint( addressingKey: string, io: "out" | "in", waitpointId: string, - ttlMs?: number, - responseFormat?: "record-v1" + ttlMs?: number ): Promise { if (!redis) return; try { const key = buildKey(environmentId, addressingKey, io); - const effectiveTtlMs = ttlMs ?? DEFAULT_TTL_MS; - - // Keep the set member as the plain waitpoint id so an older append - // instance can still drain it during a rolling deploy. New instances read - // the optional response format from this separate, TTL-bound key. - if (responseFormat) { - await redis.set(buildFormatKey(waitpointId), responseFormat, "PX", effectiveTtlMs); - } else { - await redis.del(buildFormatKey(waitpointId)); - } - - await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(effectiveTtlMs)); + await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(ttlMs ?? DEFAULT_TTL_MS)); } catch (error) { logger.error("Failed to set session stream waitpoint cache", { environmentId, @@ -146,7 +107,7 @@ export async function drainSessionStreamWaitpoints( environmentId: string, addressingKey: string, io: "out" | "in" -): Promise { +): Promise { if (!redis) return []; try { @@ -168,34 +129,7 @@ export async function drainSessionStreamWaitpoints( if (err || !Array.isArray(members)) continue; for (const m of members as string[]) ids.add(m); } - const waitpointIds = [...ids]; - if (waitpointIds.length === 0) return []; - - let formatResults: Awaited> | null = null; - try { - const formatPipeline = redis.multi(); - for (const waitpointId of waitpointIds) { - formatPipeline.get(buildFormatKey(waitpointId)); - formatPipeline.del(buildFormatKey(waitpointId)); - } - formatResults = await formatPipeline.exec(); - } catch (error) { - // The waitpoint ids were already drained. Complete them with raw data - // rather than losing the wake-up because optional metadata was unavailable. - logger.error("Failed to read session stream waitpoint response formats", { - environmentId, - addressingKey, - io, - error, - }); - } - - return waitpointIds.map((id, index) => { - const formatEntry = formatResults?.[index * 2]; - const responseFormat = - formatEntry && !formatEntry[0] && formatEntry[1] === "record-v1" ? "record-v1" : undefined; - return { id, responseFormat }; - }); + return [...ids]; } catch (error) { logger.error("Failed to drain session stream waitpoint cache", { environmentId, @@ -306,10 +240,7 @@ export async function removeSessionStreamWaitpoint( try { const key = buildKey(environmentId, addressingKey, io); - const pipeline = redis.multi(); - pipeline.srem(key, waitpointId); - pipeline.del(buildFormatKey(waitpointId)); - await pipeline.exec(); + await redis.srem(key, waitpointId); } catch (error) { logger.error("Failed to remove session stream waitpoint cache entry", { environmentId, diff --git a/apps/webapp/app/v3/webhookEngine.server.ts b/apps/webapp/app/v3/webhookEngine.server.ts index 0b06f379dae..b06f2f0024b 100644 --- a/apps/webapp/app/v3/webhookEngine.server.ts +++ b/apps/webapp/app/v3/webhookEngine.server.ts @@ -17,7 +17,6 @@ import { claimSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, - sessionStreamWaitpointOutput, } from "~/services/sessionStreamWaitpointCache.server"; import { getSecretStore } from "~/services/secrets/secretStore.server"; import { singleton } from "~/utils/singleton"; @@ -228,12 +227,10 @@ function createWebhookEngine() { "in", deliveryId ); - let appendSeq: number | undefined; if (wonClaim) { - const [appendError, seqNum] = await tryCatch( + const [appendError] = await tryCatch( realtimeStream.appendPartToSessionStream(part, deliveryId, addressingKey, "in") ); - appendSeq = seqNum ?? undefined; if (appendError) { // Nothing landed — release the claim so a retry re-appends the same id. await releaseSessionStreamPart(environment.id, addressingKey, "in", deliveryId); @@ -246,7 +243,7 @@ function createWebhookEngine() { } // Wake any `.in` waitpoints the run registered (best-effort; the record is durable in S2). - const [drainError, waitpoints] = await tryCatch( + const [drainError, waitpointIds] = await tryCatch( drainSessionStreamWaitpoints(environment.id, addressingKey, "in") ); if (drainError) { @@ -254,13 +251,13 @@ function createWebhookEngine() { externalId, error: drainError, }); - } else if (waitpoints && waitpoints.length > 0) { + } else if (waitpointIds && waitpointIds.length > 0) { await Promise.all( - waitpoints.map((waitpoint) => + waitpointIds.map((waitpointId) => tryCatch( runEngine.completeWaitpoint({ - id: waitpoint.id, - output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq), + id: waitpointId, + output: { value: part, type: "application/json", isError: false }, }) ) ) diff --git a/packages/core/src/v3/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index 195e870a121..08dc4d9ca0a 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -1754,8 +1754,6 @@ export const CreateSessionStreamWaitpointRequestBody = z.object({ * Used to catch data that arrived before `.wait()` was called. */ lastSeqNum: z.number().optional(), - /** Internal capability flag: return the exact record sequence on resume. */ - responseFormat: z.literal("record-v1").optional(), }); export type CreateSessionStreamWaitpointRequestBody = z.infer< typeof CreateSessionStreamWaitpointRequestBody diff --git a/packages/core/src/v3/sessionStreams/wireProtocol.ts b/packages/core/src/v3/sessionStreams/wireProtocol.ts index bb6aef3e1a7..550e81a0af4 100644 --- a/packages/core/src/v3/sessionStreams/wireProtocol.ts +++ b/packages/core/src/v3/sessionStreams/wireProtocol.ts @@ -40,53 +40,6 @@ export const SESSION_STATE_LAST_EVENT_ID_HEADER = "last-event-id" as const; */ export const SESSION_IN_EVENT_ID_HEADER = "session-in-event-id" as const; -/** - * Opt-in response format for Session stream waitpoints. Older SDKs omit this - * and continue receiving the raw record data. - */ -export const SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT = "record-v1" as const; - -/** Content type used only when a waitpoint actually returns a record-v1 envelope. */ -export const SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE = - "application/vnd.trigger.session-stream-record+json" as const; - -const SESSION_STREAM_WAITPOINT_RECORD_TYPE = "trigger-session-stream-record" as const; - -/** Internal envelope used to return an exact Session record from a waitpoint. */ -export type SessionStreamWaitpointRecord = Readonly<{ - type: typeof SESSION_STREAM_WAITPOINT_RECORD_TYPE; - version: 1; - seqNum: number; - data: unknown; -}>; - -export function serializeSessionStreamWaitpointRecord(data: unknown, seqNum: number): string { - return JSON.stringify({ - type: SESSION_STREAM_WAITPOINT_RECORD_TYPE, - version: 1, - seqNum, - data, - } satisfies SessionStreamWaitpointRecord); -} - -export function parseSessionStreamWaitpointRecord( - value: unknown -): SessionStreamWaitpointRecord | undefined { - if (!value || typeof value !== "object") return undefined; - - const record = value as Partial; - if ( - record.type !== SESSION_STREAM_WAITPOINT_RECORD_TYPE || - record.version !== 1 || - typeof record.seqNum !== "number" || - !Number.isFinite(record.seqNum) - ) { - return undefined; - } - - return record as SessionStreamWaitpointRecord; -} - export const TRIGGER_CONTROL_SUBTYPE = { TURN_COMPLETE: "turn-complete", UPGRADE_REQUIRED: "upgrade-required", diff --git a/packages/core/src/v3/test/session-waitpoint-backend.ts b/packages/core/src/v3/test/session-waitpoint-backend.ts index fd7cd5d3609..71c0a3ddf93 100644 --- a/packages/core/src/v3/test/session-waitpoint-backend.ts +++ b/packages/core/src/v3/test/session-waitpoint-backend.ts @@ -1,10 +1,6 @@ import { ApiClient } from "../apiClient/index.js"; import { WaitpointId } from "../isomorphic/friendlyId.js"; import { NoopRuntimeManager } from "../runtime/noopRuntimeManager.js"; -import { - SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, - serializeSessionStreamWaitpointRecord, -} from "../sessionStreams/wireProtocol.js"; import type { CreateSessionStreamWaitpointRequestBody, CreateSessionStreamWaitpointResponseBody, @@ -17,7 +13,6 @@ type PendingWait = { io: "in" | "out"; lastSeqNum?: number; timeout?: string; - responseFormat?: "record-v1"; abort: AbortController; }; @@ -76,7 +71,6 @@ export class SessionWaitpointBackend { io: body.io, lastSeqNum: body.lastSeqNum, timeout: body.timeout, - responseFormat: body.responseFormat, abort: new AbortController(), }); return { waitpointId, isCached: false }; @@ -119,20 +113,12 @@ export class SessionWaitpointBackend { }; } - const output = - pending.responseFormat === "record-v1" - ? serializeSessionStreamWaitpointRecord(result.data, result.seqNum) - : typeof result.data === "string" - ? result.data - : JSON.stringify(result.data); - return { - ok: true, - output, - outputType: - pending.responseFormat === "record-v1" - ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE - : "application/json", - }; + // The waitpoint is a wake signal only. Production appends the record to + // the channel before draining any waitpoint, so the SDK re-attaches and + // reads it back from the channel with its real sequence. Returning the + // record here would let a test pass on output the SDK no longer reads. + void result; + return { ok: true }; } catch { return { ok: false, diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index 125991a0c44..86ac0389c81 100644 --- a/packages/trigger-sdk/src/v3/sessions.ts +++ b/packages/trigger-sdk/src/v3/sessions.ts @@ -25,8 +25,6 @@ import type { import { InputStreamOncePromise, ManualWaitpointPromise, - SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, - SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT, SemanticInternalAttributes, SessionStreamInstance, WaitpointTimeoutError, @@ -34,7 +32,6 @@ import { apiClientManager, ensureReadableStream, mergeRequestOptions, - parseSessionStreamWaitpointRecord, runtime, sessionStreams, taskContext, @@ -725,7 +722,6 @@ export class SessionInputChannel { idempotencyKeyTTL: options?.idempotencyKeyTTL, tags: options?.tags, lastSeqNum: lastConsumedSeqNum, - responseFormat: SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT, }); const result = await tracer.startActiveSpan( @@ -741,80 +737,59 @@ export class SessionInputChannel { } // Stop the SSE tail before suspending. Buffered records stay in - // place; the exact record returned by the waitpoint is removed on - // resume, while any later records remain available to consumers. + // place so nothing is lost across the suspend. sessionStreams.disconnectStream(this.sessionId, "in"); const waitResult = await runtime.waitUntil(response.waitpointId); - const hasRecordEnvelope = - waitResult.outputType === SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE; - - const parsedOutput = - waitResult.output !== undefined - ? await conditionallyImportAndParsePacket( - { - data: waitResult.output, - dataType: hasRecordEnvelope - ? "application/json" - : (waitResult.outputType ?? "application/json"), - }, - apiClient - ) - : undefined; - - if (waitResult.ok) { - const record = hasRecordEnvelope - ? parseSessionStreamWaitpointRecord(parsedOutput) - : undefined; - let seqNum = record?.seqNum; - const data = record - ? await conditionallyImportAndParsePacket( - { - data: - typeof record.data === "string" ? record.data : JSON.stringify(record.data), - dataType: "application/json", - }, - apiClient - ) - : parsedOutput; - - // Older servers return only raw data. Recover its durable - // sequence from the channel instead of guessing and risking a - // cursor that skips or strands another record. - if (seqNum === undefined && waitResult.output !== undefined) { - try { - const response = await apiClient.readSessionStreamRecords(this.sessionId, "in", { - afterEventId: - lastConsumedSeqNum !== undefined ? String(lastConsumedSeqNum) : undefined, - }); - const matchingRecords = response.records.filter( - (candidate) => - candidate.data === waitResult.output || - (typeof candidate.data !== "string" && - JSON.stringify(candidate.data) === JSON.stringify(parsedOutput)) - ); - if (matchingRecords.length === 1) { - seqNum = matchingRecords[0]!.seqNum; - } - } catch { - // Leave the cursor behind when an older server cannot - // provide record metadata. At-least-once replay is safer - // than acknowledging an unknown sequence. - } - } - if (seqNum !== undefined) { - sessionStreams.consumeRecord(this.sessionId, "in", seqNum); - sessionStreams.setLastSeqNum(this.sessionId, "in", seqNum); - } + if (!waitResult.ok) { + const parsed = + waitResult.output !== undefined + ? await conditionallyImportAndParsePacket( + { + data: waitResult.output, + dataType: waitResult.outputType ?? "application/json", + }, + apiClient + ) + : undefined; + const error = new WaitpointTimeoutError(parsed?.message ?? "Timed out"); + span.recordException(error); + span.setStatus({ code: SpanStatusCode.ERROR }); + return { ok: false as const, error }; + } - return { ok: true as const, output: data as T }; - } else { - const error = new WaitpointTimeoutError(parsedOutput?.message ?? "Timed out"); + // The waitpoint is only a wake signal. The append route commits the + // record to the channel before it drains any waitpoint, so by the + // time we are here the record is durably readable from the channel + // itself, carrying its real sequence. Reading it back that way is + // what keeps the cursor exact: the waitpoint payload cannot + // identify which record it corresponds to, and guessing or matching + // on payload equality both produce a cursor that strands or + // redelivers records. + const record = await sessionStreams.onceRecord(this.sessionId, "in"); + + if (!record.ok) { + const error = new WaitpointTimeoutError("Timed out"); span.recordException(error); span.setStatus({ code: SpanStatusCode.ERROR }); return { ok: false as const, error }; } + + sessionStreams.setLastSeqNum(this.sessionId, "in", record.output.seqNum); + + const data = await conditionallyImportAndParsePacket( + { + data: + typeof record.output.data === "string" + ? record.output.data + : JSON.stringify(record.output.data), + dataType: "application/json", + }, + apiClient + ); + + return { ok: true as const, output: data as T }; }, { attributes: { diff --git a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts index a6d90744390..10dafc47f0c 100644 --- a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -307,4 +307,53 @@ describe("chat.messages mailbox", () => { expect(first).toEqual({ id: "part-redelivered", seqNum: 27, payload }); expect(redelivered).toEqual(first); }); + + it("delivers a message queued behind a control record no consumer claimed", async () => { + const chatId = "mailbox-unclaimed-head"; + const ready = deferred(); + const inspect = deferred(); + const observed: { pending?: boolean; message?: ChatMessageRecord } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-unclaimed-head", + run: async () => { + ready.resolve(); + await inspect.promise; + observed.pending = await chat.messages.hasPending(); + observed.message = await chat.messages.next({ timeoutInSeconds: 0 }); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + + await drivers.sessions.in.send(chatId, { kind: "stop" }, "in", { + id: "unclaimed-stop", + seqNum: 60, + }); + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u-behind-stop") }, + "in", + { id: "behind-stop", seqNum: 61 } + ); + inspect.resolve(); + await runPromise; + }); + + expect(observed).toEqual({ + pending: true, + message: { + id: "behind-stop", + seqNum: 61, + payload: userPayload(chatId, "u-behind-stop"), + }, + }); + }); }); diff --git a/packages/trigger-sdk/test/pending-message-drain.test.ts b/packages/trigger-sdk/test/pending-message-drain.test.ts index 18ef2462bf2..7bb0d3d8249 100644 --- a/packages/trigger-sdk/test/pending-message-drain.test.ts +++ b/packages/trigger-sdk/test/pending-message-drain.test.ts @@ -5,12 +5,7 @@ import { mockChatAgent } from "../src/v3/test/index.js"; import { describe, expect, it, vi } from "vitest"; import { chat } from "../src/v3/ai.js"; import { __setSessionOpenImplForTests, sessions } from "../src/v3/sessions.js"; -import { - apiClientManager, - SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, - serializeSessionStreamWaitpointRecord, - sessionStreams, -} from "@trigger.dev/core/v3"; +import { apiClientManager, sessionStreams } from "@trigger.dev/core/v3"; import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; import { simulateReadableStream, streamText } from "ai"; import { MockLanguageModelV3 } from "ai/test"; @@ -272,12 +267,8 @@ describe("session.in.wait() consume cursor", () => { __setSessionOpenImplForTests(undefined); const first = { kind: "message", payload: { id: "u1" } }; const later = { kind: "message", payload: { id: "u2" } }; - const runtimeManager = runtimeWithWaitpointOutput( - serializeSessionStreamWaitpointRecord(JSON.stringify(first), 50), - SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE - ); + const runtimeManager = runtimeWithWaitpointOutput(JSON.stringify(first)); let registeredLastSeqNum: number | undefined; - let registeredResponseFormat: string | undefined; await runInMockTaskContext( async (drivers) => { @@ -289,12 +280,8 @@ describe("session.in.wait() consume cursor", () => { sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 49); vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - createSessionStreamWaitpoint: async ( - _runId: string, - body: { lastSeqNum?: number; responseFormat?: string } - ) => { + createSessionStreamWaitpoint: async (_runId: string, body: { lastSeqNum?: number }) => { registeredLastSeqNum = body.lastSeqNum; - registeredResponseFormat = body.responseFormat; return { waitpointId: "wp_test_1", isCached: false, @@ -315,7 +302,6 @@ describe("session.in.wait() consume cursor", () => { expect(result).toEqual({ ok: true, output: first }); expect(registeredLastSeqNum).toBe(49); - expect(registeredResponseFormat).toBe("record-v1"); expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(51); expect(sessionStreams.peekRecord(sessionId, "in")?.seqNum).toBe(52); @@ -331,116 +317,41 @@ describe("session.in.wait() consume cursor", () => { ); }); - it("recovers the exact sequence from durable records for older servers", async () => { + it("acknowledges the delivered record when identical payloads repeat on the channel", async () => { __setSessionOpenImplForTests(undefined); - const payload = { kind: "message", payload: { id: "legacy" } }; - const rawPayload = JSON.stringify(payload); - const runtimeManager = runtimeWithWaitpointOutput(rawPayload); - let afterEventId: string | undefined; + const chunk = { kind: "message", payload: { id: "repeated" } }; + const raw = JSON.stringify(chunk); + const sessionId = "ack-repeated-payload"; await runInMockTaskContext( - async () => { - const sessionId = "legacy-cursor-sess"; + async (drivers) => { sessionStreams.setLastSeqNum(sessionId, "in", 6); sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ createSessionStreamWaitpoint: async () => ({ - waitpointId: "wp_legacy", + waitpointId: "wp_ack_repeated", isCached: false, }), - waitForWaitpointToken: async () => ({ success: true }), - readSessionStreamRecords: async ( - _sessionId: string, - _io: "in" | "out", - options?: { afterEventId?: string } - ) => { - afterEventId = options?.afterEventId; - return { - records: [{ id: "legacy-record", seqNum: 7, data: rawPayload }], - }; + waitForWaitpointToken: async () => { + await drivers.sessions.in.send(sessionId, chunk, "in", { seqNum: 7 }); + await drivers.sessions.in.send(sessionId, chunk, "in", { seqNum: 8 }); + return { success: true }; }, - } as never); - - const result = await sessions.open(sessionId).in.wait(); - - expect(result).toEqual({ ok: true, output: payload }); - expect(afterEventId).toBe("6"); - expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(7); - expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); - }, - { runtimeManager } - ); - }); - - it("does not mistake an older server's user payload for the internal envelope", async () => { - __setSessionOpenImplForTests(undefined); - const payload = { - type: "trigger-session-stream-record", - version: 1, - seqNum: 999, - data: { user: "supplied" }, - }; - const rawPayload = JSON.stringify(payload); - - await runInMockTaskContext( - async () => { - const sessionId = "legacy-envelope-collision"; - sessionStreams.setLastSeqNum(sessionId, "in", 6); - sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); - - vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - createSessionStreamWaitpoint: async () => ({ - waitpointId: "wp_legacy_collision", - isCached: false, - }), - waitForWaitpointToken: async () => ({ success: true }), - readSessionStreamRecords: async () => ({ - records: [{ id: "legacy-record", seqNum: 7, data: rawPayload }], - }), - } as never); - - const result = await sessions.open(sessionId).in.wait(); - - expect(result).toEqual({ ok: true, output: payload }); - expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); - }, - { runtimeManager: runtimeWithWaitpointOutput(rawPayload) } - ); - }); - - it("leaves the cursor behind when legacy payload matching is ambiguous", async () => { - __setSessionOpenImplForTests(undefined); - const payload = { kind: "message", payload: { id: "duplicate" } }; - const rawPayload = JSON.stringify(payload); - - await runInMockTaskContext( - async () => { - const sessionId = "legacy-duplicate-payload"; - sessionStreams.setLastSeqNum(sessionId, "in", 6); - sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); - - vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - createSessionStreamWaitpoint: async () => ({ - waitpointId: "wp_legacy_duplicate", - isCached: false, - }), - waitForWaitpointToken: async () => ({ success: true }), readSessionStreamRecords: async () => ({ records: [ - { id: "duplicate-1", seqNum: 7, data: rawPayload }, - { id: "duplicate-2", seqNum: 8, data: rawPayload }, + { id: "repeated-1", seqNum: 7, data: raw }, + { id: "repeated-2", seqNum: 8, data: raw }, ], }), } as never); const result = await sessions.open(sessionId).in.wait(); + expect(result).toEqual({ ok: true, output: chunk }); - expect(result).toEqual({ ok: true, output: payload }); - expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(6); - expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(6); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); }, - { runtimeManager: runtimeWithWaitpointOutput(rawPayload) } + { runtimeManager: runtimeWithWaitpointOutput(raw) } ); }); }); From 90ec43e090b3e5435110cd9ba929bdef69ec924b Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 20 Aug 2026 19:39:53 +0100 Subject: [PATCH 09/20] fix(chat,sdk): hold the resume cursor only behind records that matter The `session-in-event-id` header has two consumers with opposite needs. A fresh boot reads it back as the `.in` resume cursor and needs it conservative, while a client compares it against the append sequence of its own send to recognise a turn boundary that predates that send, which needs it exact. Holding the cursor behind every unconsumed record served neither: an unconsumed control record pushed the header below the sequence of the message the turn had just answered, so a client discarded its own turn-complete and stayed streaming. The cursor is now held only behind records whose loss would matter, which for chat means messages. Replaying a stop or a handover on the next boot is benign, and a handover for a turn that never ran is discarded, so control records no longer need to hold the cursor. The manager takes the rule as a per-channel predicate and defaults to holding behind everything, so a missing or throwing predicate can only make the cursor more conservative. This also ends the case where one never-consumed record pinned the cursor for the rest of the run. Two further gaps in the same machinery: The unclaimed-kind drain and the cursor rule were installed only for `chat.customAgent`. `chat.agent` builds its task directly and got neither, so the managed agent, which is the common surface, kept accumulating barriers mid-turn. Both are now installed for both surfaces, with the drain attached after each one's resume cursor is seeded so it cannot open the subscribe at seq 0. A handover-prepare boot claims the handover kinds so a signal arriving before `waitForHandover` attaches is not drained, but the claim was released only inside `waitForHandover`. A loop that never called it held the claim for the life of the run, leaving a handover record parked at the head of the channel where it wedged `chat.messages.next()` permanently. The claim is now also released at the first turn boundary, by which point the handover window has closed either way. --- packages/core/src/v3/sessionStreams/index.ts | 8 +++ .../core/src/v3/sessionStreams/manager.ts | 44 ++++++++++++- .../core/src/v3/sessionStreams/noopManager.ts | 6 ++ packages/core/src/v3/sessionStreams/types.ts | 10 +++ .../v3/test/test-session-stream-manager.ts | 25 +++++++- packages/trigger-sdk/src/v3/ai.ts | 64 ++++++++++++++++--- 6 files changed, 147 insertions(+), 10 deletions(-) diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 7d45c9248f2..82a44c72b3c 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -86,6 +86,14 @@ export class SessionStreamsAPI implements SessionStreamManager { return manager.peekRecord(sessionId, io); } + public setCursorBarrier( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void { + this.#getManager().setCursorBarrier?.(sessionId, io, predicate); + } + public lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.#getManager().lastSeqNum(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index c4c1c0503a4..9360e224b11 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -70,6 +70,14 @@ export class StandardSessionStreamManager implements SessionStreamManager { // Kept separately from `buffer` so the committed cursor can be calculated // without depending on buffer traversal. private unconsumedSeqNums = new Map>(); + + /** + * Per-channel predicate deciding which buffered records hold the persisted + * cursor back. Absent means every record does, which is the conservative + * default. Consumers that know their record kinds narrow it so the cursor is + * only held behind records whose loss would matter. + */ + private cursorBarriers = new Map(); // High-water mark of seq_nums that have been *consumed* (delivered to a // once() waiter or shifted off the buffer into a once() caller) on a channel. // Distinct from `seqNums`, which advances whenever any record is @@ -329,6 +337,37 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } + setCursorBarrier( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void { + const key = keyFor(sessionId, io); + if (predicate) { + this.cursorBarriers.set(key, predicate); + } else { + this.cursorBarriers.delete(key); + } + } + + /** + * Fails safe: an absent or throwing predicate treats the record as a barrier, + * so a mistake here can only make the cursor more conservative, never skip a + * record. + */ + #isCursorBarrier(key: string, record: SessionStreamRecord): boolean { + const predicate = this.cursorBarriers.get(key); + if (!predicate) return true; + try { + return predicate(record); + } catch (error) { + if (this.debug) { + console.error("[SessionStreamManager] Cursor barrier predicate error:", error); + } + return true; + } + } + #markUnconsumedRecord(key: string, seqNum: number): void { if (!Number.isFinite(seqNum)) return; @@ -423,6 +462,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.seqNums.clear(); this.lastDispatchedSeqNums.clear(); this.unconsumedSeqNums.clear(); + this.cursorBarriers.clear(); this.minTimestamps.clear(); this.handlers.clear(); this.reconnectAttempts.clear(); @@ -602,7 +642,9 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.buffer.set(key, buffered); } buffered.push(record); - this.#markUnconsumedRecord(key, record.seqNum); + if (this.#isCursorBarrier(key, record)) { + this.#markUnconsumedRecord(key, record.seqNum); + } this.#drainOnceWaitersFromBuffer(key); } diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index 1e5dbaebe9a..38a8dc3f850 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -55,6 +55,12 @@ export class NoopSessionStreamManager implements SessionStreamManager { return undefined; } + setCursorBarrier( + _sessionId: string, + _io: SessionChannelIO, + _predicate: SessionStreamRecordPredicate | undefined + ): void {} + lastSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { return undefined; } diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index 24b6f084512..8e518d0d9c7 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -82,6 +82,16 @@ export interface SessionStreamManager { /** Non-blocking peek at the head record, including its durable metadata. */ peekRecord?(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined; + /** + * Narrow which buffered records hold the persisted cursor back. Absent means + * every record does. + */ + setCursorBarrier?( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void; + /** Last S2 sequence number seen on the given channel. */ lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined; diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index 6d8ad0b5536..9e4ece8f01c 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -41,6 +41,7 @@ export class TestSessionStreamManager implements SessionStreamManager { private seqNums = new Map(); private dispatchedSeqNums = new Map(); private unconsumedSeqNums = new Map>(); + private cursorBarriers = new Map(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { const key = keyFor(sessionId, io); @@ -197,6 +198,16 @@ export class TestSessionStreamManager implements SessionStreamManager { return this.peekRecord(sessionId, io)?.data; } + setCursorBarrier( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void { + const key = keyFor(sessionId, io); + if (predicate) this.cursorBarriers.set(key, predicate); + else this.cursorBarriers.delete(key); + } + peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined { return this.buffer.get(keyFor(sessionId, io))?.[0]; } @@ -257,6 +268,16 @@ export class TestSessionStreamManager implements SessionStreamManager { } } + #isCursorBarrier(key: string, record: SessionStreamRecord): boolean { + const predicate = this.cursorBarriers.get(key); + if (!predicate) return true; + try { + return predicate(record); + } catch { + return true; + } + } + #markUnconsumedRecord(key: string, seqNum: number): void { if (!Number.isFinite(seqNum)) return; @@ -407,7 +428,9 @@ export class TestSessionStreamManager implements SessionStreamManager { this.buffer.set(key, buffered); } buffered.push(record); - this.#markUnconsumedRecord(key, record.seqNum); + if (this.#isCursorBarrier(key, record)) { + this.#markUnconsumedRecord(key, record.seqNum); + } this.#drainOnceWaitersFromBuffer(key); } diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 99752c1b48b..4d842feaca2 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1563,6 +1563,23 @@ export type ChatMessages = RealtimeDefinedInputStream & { next(options?: { timeoutInSeconds?: number }): Promise; }; +/** + * Only message records hold the persisted `.in` cursor back. + * + * The `session-in-event-id` header serves two consumers with opposite needs: + * `findLatestSessionInCursor` reads it as a resume cursor and wants it + * conservative, while a client reads it to correlate its own send's + * turn-complete and wants it exact. Holding the cursor behind an unconsumed + * control record satisfies neither: resume safety does not need it (replaying a + * stop or a handover is benign, and a handover for a turn that never ran is + * discarded), while a client comparing the header against its own append + * sequence sees a value below its send and discards its own turn boundary. + * @internal + */ +function isChatCursorBarrier(record: { data: unknown }): boolean { + return (record.data as ChatInputChunk | undefined)?.kind === "message"; +} + function isChatMessageRecord(record: { data: unknown }): boolean { return (record.data as ChatInputChunk | undefined)?.kind === "message"; } @@ -2000,6 +2017,31 @@ function releaseChatInputKinds(kinds: readonly string[]): void { locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); } +/** + * Narrow what holds the persisted `.in` cursor back. Sets no listener, so it is + * safe to call before the resume cursor is seeded. + * @internal + */ +function setChatCursorBarrier(chatId: string): void { + sessionStreams.setCursorBarrier(chatId, "in", isChatCursorBarrier); +} + +/** + * Claim the kinds this boot has a consumer for and drain the rest. + * + * Attaches a `.in` listener, so it MUST run after the resume cursor is seeded; + * attaching first makes the subscribe open at seq 0 and replay every record the + * previous run already answered. + * @internal + */ +function attachChatInputDrain(payload: { trigger?: string }): void { + const claimed = chatClaimedKinds(); + if (payload.trigger === "handover-prepare") { + for (const kind of CHAT_HANDOVER_KINDS) claimed.add(kind); + } + locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); +} + /** * Per-turn deferred promises. Registered via `chat.defer()`, awaited * before `onTurnComplete` fires. Reset each turn. @@ -5527,19 +5569,13 @@ function chatCustomAgent< locals.set(lastTurnCompleteSeqNumKey, { value: undefined }); markChatAgentRunForStreamsWarning(); taskContext.setConversationId(payload.chatId); + setChatCursorBarrier(payload.chatId); stampConversationIdOnActiveSpan(payload.chatId); // Seed the `.in` resume cursor before user code attaches any `.in` // listener — otherwise a continuation boot replays already-answered // messages into the loop's first wait. await seedSessionInResumeCursorForCustomLoop(payload); - // Claim the kinds this boot actually has a consumer for, then attach the - // drain for everything else. Handover kinds are only claimed on a - // handover-prepare boot, which is the only boot that waits for them. - const claimed = chatClaimedKinds(); - if (payload.trigger === "handover-prepare") { - for (const kind of CHAT_HANDOVER_KINDS) claimed.add(kind); - } - locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); + attachChatInputDrain(payload); return userRun(payload, runOptions); }, }); @@ -5646,6 +5682,7 @@ function chatAgent< locals.set(lastTurnCompleteSeqNumKey, { value: undefined }); markChatAgentRunForStreamsWarning(); taskContext.setConversationId(payload.chatId); + setChatCursorBarrier(payload.chatId); // Stamp `gen_ai.conversation.id` on the run-level span. Every // nested span inherits the same attribute via @@ -5932,6 +5969,8 @@ function chatAgent< } } + attachChatInputDrain(payload); + // ── Recovery boot + chain reconstruction ──────────────────────── if (!hydrateMessages) { const settledMessages = mergeByIdReplaceWins( @@ -9109,6 +9148,15 @@ function createStopSignal(): { async function chatWriteTurnComplete(options?: { publicAccessToken?: string; }): Promise<{ lastEventId?: string; sessionInEventId?: string }> { + // A handover-prepare boot claims the handover kinds so a signal arriving + // before `waitForHandover` attaches is not drained. A loop that never calls + // `waitForHandover` would otherwise hold that claim for the life of the run, + // leaving any handover record parked at the head of the channel: it wedges + // `chat.messages.next()` and, before the cursor barrier narrowed, pinned the + // persisted cursor forever. By the time a turn completes the handover window + // is over either way. + releaseChatInputKinds(CHAT_HANDOVER_KINDS); + const result = await writeTurnCompleteChunk(undefined, options?.publicAccessToken); // Same cursor written to the `session-in-event-id` header inside // `writeTurnCompleteChunk`; surfaced here so the caller can persist it. From 839c046150a532fe9541a207ed2170461e94ee36 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 20 Aug 2026 23:08:42 +0100 Subject: [PATCH 10/20] fix(chat,sdk): release the handover claim on every surface's turn boundary The release sat in `chat.writeTurnComplete`, which only hand-rolled loops call. The managed agent reaches a turn boundary through the internal chunk writer, so a handover-prepare boot there kept the claim for the life of the run and a handover record stayed parked at the head of the channel, wedging `chat.messages.next()`. Moved to `writeTurnCompleteChunk`, which every surface goes through. --- packages/trigger-sdk/src/v3/ai.ts | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 4d842feaca2..ce14b607aec 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -9148,15 +9148,6 @@ function createStopSignal(): { async function chatWriteTurnComplete(options?: { publicAccessToken?: string; }): Promise<{ lastEventId?: string; sessionInEventId?: string }> { - // A handover-prepare boot claims the handover kinds so a signal arriving - // before `waitForHandover` attaches is not drained. A loop that never calls - // `waitForHandover` would otherwise hold that claim for the life of the run, - // leaving any handover record parked at the head of the channel: it wedges - // `chat.messages.next()` and, before the cursor barrier narrowed, pinned the - // persisted cursor forever. By the time a turn completes the handover window - // is over either way. - releaseChatInputKinds(CHAT_HANDOVER_KINDS); - const result = await writeTurnCompleteChunk(undefined, options?.publicAccessToken); // Same cursor written to the `session-in-event-id` header inside // `writeTurnCompleteChunk`; surfaced here so the caller can persist it. @@ -11010,6 +11001,17 @@ async function writeTurnCompleteChunk( ): Promise { const session = getChatSession(); + // A handover-prepare boot claims the handover kinds so a signal arriving + // before `waitForHandover` attaches is not drained. Released here rather than + // only in `waitForHandover`, because a loop that never calls it would + // otherwise hold the claim for the life of the run and leave a handover + // record parked at the head of the channel, where it wedges + // `chat.messages.next()`. Every surface reaches a turn boundary through this + // function, including the managed agent, which does not call the public + // `chat.writeTurnComplete`. By the time a turn completes the handover window + // is over either way. + releaseChatInputKinds(CHAT_HANDOVER_KINDS); + // 1. Write the turn-complete control record. The ack's `lastEventId` is // this record's seq_num — that's the trim target for the NEXT turn. // From 0a1abe7612aa7bd32390ff097fdb357f0c3e08ab Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Fri, 21 Aug 2026 08:19:14 +0100 Subject: [PATCH 11/20] docs(chat): describe the resume cursor accurately in the changeset and API docs The changeset only mentioned the new mailbox helpers, and led with them. The change a user is most likely to care about is that a chat could silently lose a message, which affected the managed agent too, so the release note now leads with that and with the retried-send duplicate. `chat.writeTurnComplete()` also promised that `sessionInEventId` identified the exact input record the turn acknowledged. It does not: it is the cursor that is safe to resume from, held back behind any message still waiting to be handled, so a value below the record just handled is expected rather than a fault. --- .changeset/tidy-mailboxes-wait.md | 17 +++++++++++++++-- packages/trigger-sdk/src/v3/ai.ts | 12 ++++++++---- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md index 24114e838c3..97710a72b29 100644 --- a/.changeset/tidy-mailboxes-wait.md +++ b/.changeset/tidy-mailboxes-wait.md @@ -3,6 +3,19 @@ "@trigger.dev/sdk": patch --- -Custom agent loops can now inspect pending chat input without consuming it and consume one mailbox record at a time with `chat.messages.hasPending()` and `chat.messages.next()`. Mailbox records include stable identifiers for tracing and redelivery. +Fixes a case where a chat could silently lose a message. If a message arrived while the agent was between turns and a stop arrived after it, the cursor the next boot resumed from could point past that message, so it was never answered and no error was raised. This affected `chat.agent`, not just custom agents. -A control record that nothing on the run consumes is now discarded rather than left at the head of the `.in` channel, where it would have made every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` means no message became consumable before the timeout. +Also fixes a retried send being answered twice. When a send was retried and its idempotency claim was lost, the agent could consume the same message a second time. + +Custom agent loops can now inspect pending chat input without consuming it, and consume one mailbox record at a time, with `chat.messages.hasPending()` and `chat.messages.next()`. Records carry stable identifiers so a redelivery is recognisable. + +```ts +if (await chat.messages.hasPending()) { + const record = await chat.messages.next({ timeoutInSeconds: 0 }); + if (record) handle(record.payload); +} +``` + +A control record that nothing on the run consumes is now discarded rather than left at the head of the input channel, where it would have made every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` means no message became consumable before the timeout. + +`chat.writeTurnComplete()`'s `sessionInEventId` is the cursor that is safe to resume from, not the sequence of the record the turn answered. It is held back behind any message still waiting to be handled, so a value below the record you just handled is expected. diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index ce14b607aec..5b807cc06c9 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -9131,10 +9131,14 @@ function createStopSignal(): { * task instead of round-tripping them back from the client: * - `lastEventId` — the turn-complete control record's seq_num on * `session.out`; where the next turn's output stream resumes. - * - `sessionInEventId` — the committed-consume cursor on `session.in` as of - * this turn-complete, letting a raw loop correlate the boundary with the - * exact input record it acknowledged. Trigger owns input-cursor recovery, - * so this is for correlation / out-of-sync detection, not required. + * - `sessionInEventId` — the safe-to-resume-from cursor on `session.in` as of + * this turn-complete. It is the highest sequence that can be resumed past + * without skipping an unhandled message, so it is held back behind any + * message still buffered unconsumed and is NOT necessarily the sequence of + * the record this turn answered. Trigger owns input-cursor recovery, so this + * is for correlation / out-of-sync detection, not required. Treat it as a + * lower bound: a value below the record you just handled is expected, not a + * sign of a lost turn. * * Either is `undefined` when the corresponding cursor isn't available. * From 36c37d56053505af22c0aaa2ef62f7b313ec576a Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Fri, 21 Aug 2026 23:00:52 +0100 Subject: [PATCH 12/20] fix(chat,sdk): stop an already-applied control record being applied again on resume Holding the resume cursor behind a message still waiting to be handled means resuming from it necessarily re-delivers every record after that point, including records the previous run had already handled. For a message that is the entire point. For a control record it is a fault: a stop carries no record of which turn it belonged to, so on redelivery it aborts whichever turn happens to be live, which is usually the turn answering the very message the cursor was held back to protect. The user got their answer cut off instead of never arriving, which is better but still wrong. Each turn boundary now also reports the highest sequence that run had actually consumed, unclamped, alongside the cursor that is safe to resume from. On boot a run reads it back and drops control records at or below it before any consumer sees them. Messages are never dropped, so recovery is unchanged. A chat whose turns predate the header reports nothing, drops nothing, and behaves as before. Driven on a stack: a message set aside, a stop consumed, a turn boundary, a SIGKILL, then a continuation. Before, the continuation consumed the message and was immediately aborted by the replayed stop. Now the stop is dropped and the turn proceeds, while a stop that arrives live still aborts its turn. --- packages/core/src/v3/sessionStreams/index.ts | 12 ++++ .../core/src/v3/sessionStreams/manager.ts | 47 ++++++++++++ .../core/src/v3/sessionStreams/noopManager.ts | 10 +++ packages/core/src/v3/sessionStreams/types.ts | 10 +++ .../src/v3/sessionStreams/wireProtocol.ts | 12 ++++ .../v3/test/test-session-stream-manager.ts | 15 ++++ packages/trigger-sdk/src/v3/ai.ts | 71 +++++++++++++++++++ 7 files changed, 177 insertions(+) diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 82a44c72b3c..6519628ce08 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -86,6 +86,18 @@ export class SessionStreamsAPI implements SessionStreamManager { return manager.peekRecord(sessionId, io); } + public highestConsumedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { + return this.#getManager().highestConsumedSeqNum?.(sessionId, io); + } + + public setDropPredicate( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void { + this.#getManager().setDropPredicate?.(sessionId, io, predicate); + } + public setCursorBarrier( sessionId: string, io: SessionChannelIO, diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index 9360e224b11..541e4182398 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -78,6 +78,16 @@ export class StandardSessionStreamManager implements SessionStreamManager { * only held behind records whose loss would matter. */ private cursorBarriers = new Map(); + + /** + * Per-channel predicate marking records that must not be delivered again. + * A resume cursor held back behind an unhandled record necessarily replays + * everything after it, including records that WERE handled before the + * previous run ended. Re-delivering those is not harmless: a control record + * applies a second time to whatever turn happens to be live. This lets the + * consumer name the ones to drop. + */ + private dropPredicates = new Map(); // High-water mark of seq_nums that have been *consumed* (delivered to a // once() waiter or shifted off the buffer into a once() caller) on a channel. // Distinct from `seqNums`, which advances whenever any record is @@ -337,6 +347,24 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } + /** The highest consumed sequence, unclamped. @see lastDispatchedSeqNum */ + highestConsumedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { + return this.lastDispatchedSeqNums.get(keyFor(sessionId, io)); + } + + setDropPredicate( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void { + const key = keyFor(sessionId, io); + if (predicate) { + this.dropPredicates.set(key, predicate); + } else { + this.dropPredicates.delete(key); + } + } + setCursorBarrier( sessionId: string, io: SessionChannelIO, @@ -463,6 +491,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.lastDispatchedSeqNums.clear(); this.unconsumedSeqNums.clear(); this.cursorBarriers.clear(); + this.dropPredicates.clear(); this.minTimestamps.clear(); this.handlers.clear(); this.reconnectAttempts.clear(); @@ -605,6 +634,24 @@ export class StandardSessionStreamManager implements SessionStreamManager { } #dispatch(key: string, record: SessionStreamRecord): void { + const drop = this.dropPredicates.get(key); + if (drop) { + let shouldDrop = false; + try { + shouldDrop = drop(record); + } catch (error) { + if (this.debug) { + console.error("[SessionStreamManager] Drop predicate error:", error); + } + } + if (shouldDrop) { + // Acknowledge it so the tail does not fetch it again, but never hand + // it to a waiter or a handler. + this.#advanceLastDispatched(key, record.seqNum); + return; + } + } + // Any record flowing through = healthy connection; reset the backoff // counter so the next disconnect starts fresh. this.reconnectAttempts.delete(key); diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index 38a8dc3f850..207434d92c0 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -55,6 +55,16 @@ export class NoopSessionStreamManager implements SessionStreamManager { return undefined; } + highestConsumedSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { + return undefined; + } + + setDropPredicate( + _sessionId: string, + _io: SessionChannelIO, + _predicate: SessionStreamRecordPredicate | undefined + ): void {} + setCursorBarrier( _sessionId: string, _io: SessionChannelIO, diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index 8e518d0d9c7..bd8901eb0a5 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -86,6 +86,16 @@ export interface SessionStreamManager { * Narrow which buffered records hold the persisted cursor back. Absent means * every record does. */ + /** The highest consumed sequence, unclamped. */ + highestConsumedSeqNum?(sessionId: string, io: SessionChannelIO): number | undefined; + + /** Mark records that must never be delivered again on this boot. */ + setDropPredicate?( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void; + setCursorBarrier?( sessionId: string, io: SessionChannelIO, diff --git a/packages/core/src/v3/sessionStreams/wireProtocol.ts b/packages/core/src/v3/sessionStreams/wireProtocol.ts index 550e81a0af4..9a5cfb0c710 100644 --- a/packages/core/src/v3/sessionStreams/wireProtocol.ts +++ b/packages/core/src/v3/sessionStreams/wireProtocol.ts @@ -40,6 +40,18 @@ export const SESSION_STATE_LAST_EVENT_ID_HEADER = "last-event-id" as const; */ export const SESSION_IN_EVENT_ID_HEADER = "session-in-event-id" as const; +/** + * Sibling of {@link SESSION_IN_EVENT_ID_HEADER}: the highest `.in` sequence this + * run had actually consumed at the turn boundary, unclamped. + * + * The resume cursor is held back behind records still waiting to be handled, so + * resuming from it necessarily re-delivers records that WERE handled. A message + * re-delivered that way is the point. A control record re-delivered that way is + * a bug: it applies a second time to whatever turn is live on the new run. On + * boot this bound tells the run which control records it has already seen. + */ +export const SESSION_IN_CONSUMED_ID_HEADER = "session-in-consumed-id" as const; + export const TRIGGER_CONTROL_SUBTYPE = { TURN_COMPLETE: "turn-complete", UPGRADE_REQUIRED: "upgrade-required", diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index 9e4ece8f01c..3626f117e0b 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -42,6 +42,7 @@ export class TestSessionStreamManager implements SessionStreamManager { private dispatchedSeqNums = new Map(); private unconsumedSeqNums = new Map>(); private cursorBarriers = new Map(); + private dropPredicates = new Map(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { const key = keyFor(sessionId, io); @@ -198,6 +199,20 @@ export class TestSessionStreamManager implements SessionStreamManager { return this.peekRecord(sessionId, io)?.data; } + highestConsumedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { + return this.dispatchedSeqNums.get(keyFor(sessionId, io)); + } + + setDropPredicate( + sessionId: string, + io: SessionChannelIO, + predicate: SessionStreamRecordPredicate | undefined + ): void { + const key = keyFor(sessionId, io); + if (predicate) this.dropPredicates.set(key, predicate); + else this.dropPredicates.delete(key); + } + setCursorBarrier( sessionId: string, io: SessionChannelIO, diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 5b807cc06c9..ebe5d1d7842 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -26,6 +26,7 @@ import { resourceCatalog, type SessionTriggerConfig, SemanticInternalAttributes, + SESSION_IN_CONSUMED_ID_HEADER, SESSION_IN_EVENT_ID_HEADER, sessionStreams, taskContext, @@ -222,6 +223,42 @@ async function findLatestSessionInCursor(chatId: string): Promise { + try { + return await findLatestSessionInConsumed(chatId); + } catch { + return undefined; + } +} + +/** + * The highest `.in` sequence a previous run had already consumed, read from the + * latest `turn-complete` on `.out`. + * + * Absent for a chat whose turns predate the header, in which case no record is + * dropped and behaviour matches the previous release. + * @internal + */ +async function findLatestSessionInConsumed(chatId: string): Promise { + const apiClient = apiClientManager.clientOrThrow(); + const response = await apiClient.readSessionStreamRecords(chatId, "out"); + let latest: number | undefined; + for (const record of response.records) { + if (controlSubtype(record.headers) !== TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE) continue; + const raw = headerValue(record.headers, SESSION_IN_CONSUMED_ID_HEADER); + if (!raw) continue; + const parsed = Number.parseInt(raw, 10); + if (Number.isFinite(parsed)) latest = parsed; + } + return latest; +} + /** Test-only entry point for the records-based cursor scan. @internal */ export async function __findLatestSessionInCursorForTests( chatId: string @@ -2026,6 +2063,34 @@ function setChatCursorBarrier(chatId: string): void { sessionStreams.setCursorBarrier(chatId, "in", isChatCursorBarrier); } +/** + * Stop already-handled control records from being applied a second time after a + * resume. + * + * The cursor is deliberately held back behind messages still waiting to be + * handled, so resuming from it re-delivers everything after that point. For a + * message that is the whole purpose. For a control record it is a fault: a stop + * carries no record of which turn it belonged to, so on redelivery it would + * abort whichever turn happens to be live, which is usually the turn answering + * the very message the cursor was held back to protect. + * + * `consumedThrough` is the highest sequence a previous run reported consuming. + * Control records at or below it have already been applied and are dropped + * before any consumer sees them. Messages are never dropped. Absent for a chat + * whose turns predate the header, in which case nothing is dropped. + * @internal + */ +function setChatReplayGuard(chatId: string, consumedThrough: number | undefined): void { + if (consumedThrough === undefined) { + sessionStreams.setDropPredicate(chatId, "in", undefined); + return; + } + sessionStreams.setDropPredicate(chatId, "in", (record) => { + if (record.seqNum > consumedThrough) return false; + return !isChatCursorBarrier(record); + }); +} + /** * Claim the kinds this boot has a consumer for and drain the rest. * @@ -5575,6 +5640,7 @@ function chatCustomAgent< // listener — otherwise a continuation boot replays already-answered // messages into the loop's first wait. await seedSessionInResumeCursorForCustomLoop(payload); + setChatReplayGuard(payload.chatId, await findLatestSessionInConsumedSafe(payload.chatId)); attachChatInputDrain(payload); return userRun(payload, runOptions); }, @@ -5969,6 +6035,7 @@ function chatAgent< } } + setChatReplayGuard(payload.chatId, await findLatestSessionInConsumedSafe(payload.chatId)); attachChatInputDrain(payload); // ── Recovery boot + chain reconstruction ──────────────────────── @@ -11034,6 +11101,10 @@ async function writeTurnCompleteChunk( if (inCursor !== undefined) { extraHeaders.push([SESSION_IN_EVENT_ID_HEADER, String(inCursor)]); } + const consumedCursor = sessionStreams.highestConsumedSeqNum(session.id, "in"); + if (consumedCursor !== undefined) { + extraHeaders.push([SESSION_IN_CONSUMED_ID_HEADER, String(consumedCursor)]); + } const result = await session.out.writeControl( TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE, extraHeaders From c1f5ba8ffe7132714e05b9558b3ad608b7f586d5 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sat, 22 Aug 2026 07:08:39 +0100 Subject: [PATCH 13/20] feat(core): route session channel records by declared delivery discipline Adds the session channel router: one reader that classifies every record once and gives it exactly one destination, instead of several kind-filtered facades each taking records off a shared buffer. A route declares two independent things: whether it queues a record when no consumer is ready, and whether a record it never handled has to survive into the next boot. Those two make the resume floor, the replay window and the discard-the-unowned behaviour derived properties rather than three predicates installed by hand. Not wired to anything yet. The chat layer keeps its current mechanisms until the reproduction catalog is green on the router. --- .../core/src/v3/sessionStreams/router.test.ts | 391 +++++++++++++++ packages/core/src/v3/sessionStreams/router.ts | 444 ++++++++++++++++++ 2 files changed, 835 insertions(+) create mode 100644 packages/core/src/v3/sessionStreams/router.test.ts create mode 100644 packages/core/src/v3/sessionStreams/router.ts diff --git a/packages/core/src/v3/sessionStreams/router.test.ts b/packages/core/src/v3/sessionStreams/router.test.ts new file mode 100644 index 00000000000..59006bc9bf4 --- /dev/null +++ b/packages/core/src/v3/sessionStreams/router.test.ts @@ -0,0 +1,391 @@ +import { describe, expect, it } from "vitest"; +import { SessionChannelRouter } from "./router.js"; +import type { SessionRouteTable } from "./router.js"; +import type { SessionStreamRecord } from "./types.js"; + +type Chunk = { kind: string; text?: string }; + +const CHAT_TABLE: SessionRouteTable = { + kindOf: (data) => (data as Chunk | undefined)?.kind, + routes: [ + { name: "messages", delivery: "queue", replayable: true, kinds: ["message"] }, + { name: "stop", delivery: "at-arrival", replayable: false, kinds: ["stop"] }, + { + name: "handover", + delivery: "queue", + replayable: false, + kinds: ["handover", "handover-skip"], + }, + ], +}; + +function router(onDrop?: Parameters[0]) { + return makeRouter(onDrop); +} + +function makeRouter( + onDrop?: (record: SessionStreamRecord, reason: string, route?: string) => void +) { + return new SessionChannelRouter(CHAT_TABLE, { onDrop }); +} + +function rec(seqNum: number, kind: string, text?: string): SessionStreamRecord { + return { id: `r${seqNum}`, seqNum, data: { kind, ...(text ? { text } : {}) } }; +} + +describe("SessionChannelRouter: table validation", () => { + it("rejects a route that is at-arrival and replayable", () => { + expect( + () => + new SessionChannelRouter({ + kindOf: () => "x", + routes: [{ name: "bad", delivery: "at-arrival", replayable: true, kinds: ["x"] }], + }) + ).toThrow(/at-arrival and replayable/); + }); + + it("rejects a kind claimed by two routes", () => { + expect( + () => + new SessionChannelRouter({ + kindOf: () => "x", + routes: [ + { name: "a", delivery: "queue", replayable: true, kinds: ["x"] }, + { name: "b", delivery: "queue", replayable: false, kinds: ["x"] }, + ], + }) + ).toThrow(/claimed by both/); + }); +}); + +describe("SessionChannelRouter: classification", () => { + it("queues a message when nobody is ready for it", () => { + const r = router(); + expect(r.ingest(rec(0, "message", "M0"))).toEqual({ action: "queue", route: "messages" }); + expect(r.hasPending("messages")).toBe(true); + }); + + it("discards a stop with no handler attached", () => { + const r = router(); + expect(r.ingest(rec(0, "stop"))).toEqual({ + action: "drop", + route: "stop", + reason: "no-handler", + }); + }); + + it("delivers a stop to a live handler", () => { + const r = router(); + const seen: number[] = []; + r.on("stop", (record) => seen.push(record.seqNum)); + expect(r.ingest(rec(3, "stop"))).toEqual({ action: "deliver", route: "stop" }); + expect(seen).toEqual([3]); + }); + + it("drops a kind no route claims, and reports it once", () => { + const drops: Array<[number, string]> = []; + const r = router((record, reason) => drops.push([record.seqNum, reason])); + expect(r.ingest(rec(1, "some-future-kind"))).toEqual({ + action: "drop", + reason: "unroutable", + }); + expect(drops).toEqual([[1, "unroutable"]]); + }); + + it("drops a record with no usable kind", () => { + const r = router(); + expect(r.ingest({ id: "x", seqNum: 0, data: { nope: true } })).toEqual({ + action: "drop", + reason: "malformed", + }); + }); + + it("does not let a throwing kindOf take the channel down", () => { + const r = new SessionChannelRouter({ + kindOf: () => { + throw new Error("boom"); + }, + routes: [{ name: "m", delivery: "queue", replayable: true, kinds: ["message"] }], + }); + expect(r.ingest(rec(0, "message"))).toEqual({ action: "drop", reason: "malformed" }); + }); +}); + +describe("SessionChannelRouter: the wedge cannot happen", () => { + it("delivers a message queued behind an unroutable record", async () => { + const r = router(); + r.ingest(rec(0, "mystery-kind")); + r.ingest(rec(1, "message", "M1")); + + expect(r.hasPending("messages")).toBe(true); + const taken = await r.next("messages", { timeoutMs: 0 }); + expect((taken!.data as Chunk).text).toBe("M1"); + }); + + it("reports pending for a message queued behind a stop", () => { + const r = router(); + r.ingest(rec(0, "stop")); + r.ingest(rec(1, "message", "M1")); + + expect(r.hasPending("messages")).toBe(true); + expect((r.peek("messages")!.data as Chunk).text).toBe("M1"); + }); +}); + +describe("SessionChannelRouter: delivery ordering", () => { + it("serves a parked waiter before a push handler", async () => { + const r = router(); + const handlerSaw: string[] = []; + const pending = r.next("messages"); + r.on("messages", (record) => handlerSaw.push((record.data as Chunk).text!)); + + r.ingest(rec(0, "message", "M0")); + + expect((await pending)?.seqNum).toBe(0); + expect(handlerSaw).toEqual([]); + }); + + it("re-offers a queued record to a handler attaching later, in order", () => { + const r = router(); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "message", "M1")); + + const seen: string[] = []; + r.on("messages", (record) => seen.push((record.data as Chunk).text!)); + + expect(seen).toEqual(["M0", "M1"]); + expect(r.hasPending("messages")).toBe(false); + }); + + it("resolves next() undefined on timeout without consuming anything", async () => { + const r = router(); + expect(await r.next("messages", { timeoutMs: 5 })).toBeUndefined(); + r.ingest(rec(0, "message", "M0")); + expect((await r.next("messages", { timeoutMs: 0 }))?.seqNum).toBe(0); + }); +}); + +describe("SessionChannelRouter: the resume floor", () => { + it("sits at the high water when nothing is owed", () => { + const r = router(); + r.on("stop", () => {}); + r.ingest(rec(0, "message")); + r.next("messages", { timeoutMs: 0 }); + r.ingest(rec(1, "stop")); + + expect(r.resumeFloor()).toBe(1); + expect(r.appliedThrough()).toBe(1); + }); + + it("is held below a message still queued, even as control records advance", () => { + const r = router(); + r.on("stop", () => {}); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "message", "M1")); + r.next("messages", { timeoutMs: 0 }); + r.ingest(rec(2, "stop")); + r.ingest(rec(3, "stop")); + + expect(r.resumeFloor()).toBe(0); + expect(r.appliedThrough()).toBe(3); + }); + + it("is undefined rather than negative when the very first record is owed", () => { + const r = router(); + r.ingest(rec(0, "message", "M0")); + expect(r.resumeFloor()).toBeUndefined(); + }); + + it("is not held back by a queued non-replayable record", () => { + const r = router(); + r.ingest(rec(0, "handover")); + r.ingest(rec(1, "message", "M1")); + r.next("messages", { timeoutMs: 0 }); + + expect(r.pendingCount("handover")).toBe(1); + expect(r.resumeFloor()).toBe(1); + }); + + it("advances once the owed message is taken", async () => { + const r = router(); + r.on("stop", () => {}); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "stop")); + expect(r.resumeFloor()).toBeUndefined(); + + await r.next("messages", { timeoutMs: 0 }); + expect(r.resumeFloor()).toBe(1); + }); + + it("tracks the earliest of several owed messages", () => { + const r = router(); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "message", "M1")); + r.ingest(rec(2, "message", "M2")); + r.next("messages", { timeoutMs: 0 }); + + expect(r.resumeFloor()).toBe(0); + expect(r.appliedThrough()).toBe(2); + }); +}); + +describe("SessionChannelRouter: resuming", () => { + it("does not apply a non-replayable record inside the replay window", () => { + const r = router(); + r.restore({ resumeFrom: 0, appliedThrough: 2 }); + const stops: number[] = []; + r.on("stop", (record) => stops.push(record.seqNum)); + + expect(r.ingest(rec(1, "message", "M1"))).toEqual({ action: "queue", route: "messages" }); + expect(r.ingest(rec(2, "stop"))).toEqual({ + action: "drop", + route: "stop", + reason: "replayed", + }); + expect(stops).toEqual([]); + }); + + it("applies the same kind arriving live, past the window", () => { + const r = router(); + r.restore({ resumeFrom: 0, appliedThrough: 2 }); + const stops: number[] = []; + r.on("stop", (record) => stops.push(record.seqNum)); + + r.ingest(rec(1, "message", "M1")); + r.ingest(rec(2, "stop")); + expect(r.ingest(rec(3, "stop"))).toEqual({ action: "deliver", route: "stop" }); + expect(stops).toEqual([3]); + }); + + it("treats an absent replay-window end as the floor", () => { + const r = router(); + r.restore({ resumeFrom: 4 }); + r.on("stop", () => {}); + + expect(r.ingest(rec(4, "stop")).action).toBe("drop"); + expect(r.ingest(rec(5, "stop"))).toEqual({ action: "deliver", route: "stop" }); + }); + + it("applies everything on a fresh session with no checkpoint", () => { + const r = router(); + r.on("stop", () => {}); + expect(r.ingest(rec(0, "stop"))).toEqual({ action: "deliver", route: "stop" }); + }); + + it("never drops a replayable record, however far inside the window", () => { + const r = router(); + r.restore({ resumeFrom: 0, appliedThrough: 9 }); + expect(r.ingest(rec(1, "message", "M1"))).toEqual({ action: "queue", route: "messages" }); + }); + + it("keeps the floor from moving backwards past the restored point", () => { + const r = router(); + r.restore({ resumeFrom: 7, appliedThrough: 7 }); + expect(r.resumeFloor()).toBe(7); + }); +}); + +describe("SessionChannelRouter: consumer windows", () => { + it("queues a handover that arrives before its consumer is ready", async () => { + const r = router(); + expect(r.ingest(rec(0, "handover")).action).toBe("queue"); + + const taken = await r.next("handover", { timeoutMs: 0 }); + expect(taken?.seqNum).toBe(0); + }); + + it("discards what is left on a route when its window closes", () => { + const r = router(); + r.ingest(rec(0, "handover")); + r.clearRoute("handover"); + + expect(r.pendingCount("handover")).toBe(0); + expect(r.resumeFloor()).toBe(0); + }); + + it("wakes a waiter empty when its window closes", async () => { + const r = router(); + const pending = r.next("handover"); + r.clearRoute("handover"); + expect(await pending).toBeUndefined(); + }); +}); + +/** + * The invariant the whole design exists to hold: across any interleaving and + * any crash point, a message is delivered exactly once and a stop is never + * applied twice. + * + * Runs every record sequence over a simulated two-boot lifecycle. The second + * boot resubscribes from the published floor, exactly as the tail does with + * `Last-Event-ID`, so a floor that is too high shows up as a lost message and + * one that replays a stop shows up as a duplicate application. + */ +describe("SessionChannelRouter: exactly-once across a crash", () => { + const KINDS = ["message", "stop", "handover", "junk"] as const; + + function interleavings(length: number): string[][] { + if (length === 0) return [[]]; + const shorter = interleavings(length - 1); + const out: string[][] = []; + for (const prefix of shorter) { + for (const kind of KINDS) out.push([...prefix, kind]); + } + return out; + } + + function runBoot( + records: SessionStreamRecord[], + checkpoint: { resumeFrom?: number; appliedThrough?: number }, + takeMessages: number, + attachStop: boolean + ) { + const r = router(); + r.restore(checkpoint); + const messages: number[] = []; + const stops: number[] = []; + if (attachStop) r.on("stop", (record) => stops.push(record.seqNum)); + + for (const record of records) { + if (checkpoint.resumeFrom !== undefined && record.seqNum <= checkpoint.resumeFrom) continue; + r.ingest(record); + } + + for (let i = 0; i < takeMessages; i++) { + const head = r.peek("messages"); + if (!head) break; + void r.next("messages", { timeoutMs: 0 }); + messages.push(head.seqNum); + } + + return { messages, stops, checkpoint: r.checkpoint() }; + } + + it("delivers every message exactly once and applies no stop twice", () => { + const cases = interleavings(4); + expect(cases.length).toBe(256); + + for (const kinds of cases) { + const records = kinds.map((kind, index) => rec(index, kind)); + const messageSeqs = records + .filter((record) => (record.data as Chunk).kind === "message") + .map((record) => record.seqNum); + + for (let crashAfter = 0; crashAfter <= kinds.length; crashAfter++) { + const first = runBoot(records, {}, crashAfter, true); + const second = runBoot(records, first.checkpoint, kinds.length, true); + + const delivered = [...first.messages, ...second.messages]; + expect(delivered, `messages for [${kinds.join(",")}] crashAfter=${crashAfter}`).toEqual( + messageSeqs + ); + + const appliedTwice = first.stops.filter((seq) => second.stops.includes(seq)); + expect( + appliedTwice, + `stops applied twice for [${kinds.join(",")}] crashAfter=${crashAfter}` + ).toEqual([]); + } + } + }); +}); diff --git a/packages/core/src/v3/sessionStreams/router.ts b/packages/core/src/v3/sessionStreams/router.ts new file mode 100644 index 00000000000..0563b3820fe --- /dev/null +++ b/packages/core/src/v3/sessionStreams/router.ts @@ -0,0 +1,444 @@ +import type { SessionStreamRecord } from "./types.js"; + +/** + * What happens to a record when no consumer is ready for it *right now*. + * + * - `queue`: it waits in the route's own queue until a consumer takes it. + * - `at-arrival`: it goes to a live handler or nowhere. A record that only + * means something to the turn that is live when it lands (a stop) is this. + */ +type RouteDelivery = "queue" | "at-arrival"; + +/** + * One route: which kinds it owns, whether it waits for a consumer, and whether + * a record it never handled has to survive into the next boot. + * + * The two properties are independent, and that is the point. Three of the four + * combinations are meaningful and cover everything `session.in` carries: + * + * | delivery | replayable | example | + * | --- | --- | --- | + * | `queue` | `true` | a user message: waits for a turn, and a crash must not lose it | + * | `at-arrival` | `false` | a stop: only the live turn cares, and a replayed one would abort the wrong turn | + * | `queue` | `false` | a handover signal: can arrive before its consumer is ready, but is meaningless to a later boot | + * + * The fourth is a contradiction (discard it when nobody is listening, yet + * recover it later) and the table rejects it. + */ +type SessionRoute = { + /** Unique within a table. */ + name: string; + delivery: RouteDelivery; + /** + * Whether a record this route never handled must be recovered by the next + * boot. Only `true` holds the resume floor back. + */ + replayable: boolean; + /** Record kinds this route owns. Every kind belongs to at most one route. */ + kinds: readonly string[]; +}; + +/** + * The complete statement of what a channel carries and who owns each kind. + * Intended to be a literal at the point of use, so "which kinds exist and what + * happens to each" is answerable by reading one object. + */ +export type SessionRouteTable = { + /** Extract a record's kind. Returning `undefined` marks it malformed. */ + kindOf: (data: unknown) => string | undefined; + routes: readonly SessionRoute[]; +}; + +export type RouterDropReason = + /** No route claims this kind: nothing on this boot can consume it. */ + | "unroutable" + /** No usable kind on the record at all. */ + | "malformed" + /** + * A non-replayable record inside the replay window. It was already observed + * by a previous run, and its route has declared that a later boot has no use + * for it. + */ + | "replayed" + /** An `at-arrival` record with no handler attached right now. */ + | "no-handler"; + +export type RouterDecision = + /** Handed to a consumer that was already waiting, or to a live handler. */ + | { action: "deliver"; route: string } + /** Parked in the route's queue for a future consumer. */ + | { action: "queue"; route: string } + | { action: "drop"; route?: string; reason: RouterDropReason }; + +/** + * The two numbers a turn boundary publishes, and that a boot reads back. + * + * `resumeFrom` is the floor: every record at or below it was terminally + * handled, so a boot subscribes from just past it. `appliedThrough` is the end + * of the replay window: the highest sequence a previous run observed. Records + * at or below it are being re-read rather than arriving live. + */ +export type RouterCheckpoint = { + resumeFrom?: number; + appliedThrough?: number; +}; + +type QueueWaiter = { + resolve: (record: SessionStreamRecord | undefined) => void; + timer?: ReturnType; +}; + +type RouteHandler = (record: SessionStreamRecord) => void; + +/** + * One route's live state: its queue of records nobody has taken yet, the + * consumers waiting for the next one, and any attached push handlers. + * + * An `at-arrival` route never fills `queue`; a route that is not `replayable` + * fills it but is skipped when the floor is computed. + */ +class RouteState { + readonly queue: SessionStreamRecord[] = []; + readonly waiters: QueueWaiter[] = []; + readonly handlers = new Set(); + + constructor(readonly route: SessionRoute) {} + + /** + * Lowest sequence a *later boot* would still have to recover. A route that + * is not replayable never holds anything back, however much it has queued. + */ + earliestUnrecovered(): number | undefined { + if (!this.route.replayable) return undefined; + return this.queue.length > 0 ? this.queue[0]!.seqNum : undefined; + } +} + +/** + * Reads one session channel and gives every record exactly one destination. + * + * The channel carries records for several independent consumers whose delivery + * needs differ: a message must be delivered eventually and so can lag + * arbitrarily far behind the newest record, while a stop only means anything to + * the turn that is live when it lands. Tracking that with a single scalar + * cursor is not possible — the true state is always "control applied through + * X, message Y still owed" — so the router tracks it per route and derives the + * published numbers from route state: + * + * - the resume floor is held back only by records a later boot would still have + * to recover, which is exactly the queued records on replayable routes; + * - a record whose route has declared it not replayable is never re-applied + * after a resume, because the route said a later boot has no use for it; + * - a record with no route is terminal by classification, so it never enters a + * queue and cannot park at the head of one. + * + * Those three properties are what previously needed a cursor-barrier + * predicate, a drop predicate with a second published header, and a + * discard-the-unclaimed drain respectively. + */ +export class SessionChannelRouter { + #routes = new Map(); + #kindToRoute = new Map(); + #highestSeq: number | undefined; + #resumeFrom: number | undefined; + #appliedThrough: number | undefined; + #onDrop?: (record: SessionStreamRecord, reason: RouterDropReason, route?: string) => void; + + constructor( + private table: SessionRouteTable, + options?: { + /** Called for every dropped record. Reporting only; never load-bearing. */ + onDrop?: (record: SessionStreamRecord, reason: RouterDropReason, route?: string) => void; + } + ) { + for (const route of table.routes) { + if (this.#routes.has(route.name)) { + throw new Error(`Duplicate route name "${route.name}" in session route table`); + } + if (route.delivery === "at-arrival" && route.replayable) { + throw new Error( + `Route "${route.name}" is at-arrival and replayable, which cannot both hold: a record discarded because nobody was listening cannot also be recovered later` + ); + } + this.#routes.set(route.name, new RouteState(route)); + for (const kind of route.kinds) { + const existing = this.#kindToRoute.get(kind); + if (existing) { + throw new Error( + `Kind "${kind}" is claimed by both "${existing}" and "${route.name}" in session route table` + ); + } + this.#kindToRoute.set(kind, route.name); + } + } + this.#onDrop = options?.onDrop; + } + + /** + * Seed the router from a previous run's turn boundary. + * + * An absent `appliedThrough` is treated as equal to the floor rather than as + * "nothing was applied". A boundary written before that value existed still + * tells us everything at or below the floor was terminal, and for anything + * above it the conservative choice for an `at-arrival` record is to not apply + * it: a missed stop is recoverable, while a stop applied to the wrong turn + * kills a live answer. + */ + restore(checkpoint: RouterCheckpoint): void { + this.#resumeFrom = checkpoint.resumeFrom; + this.#appliedThrough = checkpoint.appliedThrough ?? checkpoint.resumeFrom; + if (this.#resumeFrom !== undefined) { + this.#highestSeq = this.#resumeFrom; + } + } + + /** Where a boot should subscribe from: just past this sequence. */ + resumeFrom(): number | undefined { + return this.#resumeFrom; + } + + /** + * Classify one record and act on it. The record's destination is decided + * here, once, and never by whichever consumer happens to be waiting. + * + * Queued routes serve a waiting consumer before a push handler, so a handler + * can never take a record out from under a consumer that is actively awaiting + * one. A record handed straight to either never enters the queue, so it never + * holds the floor back. + */ + ingest(record: SessionStreamRecord): RouterDecision { + if (Number.isFinite(record.seqNum)) { + if (this.#highestSeq === undefined || record.seqNum > this.#highestSeq) { + this.#highestSeq = record.seqNum; + } + } + + const kind = this.#kindOf(record.data); + if (kind === undefined) { + return this.#drop(record, "malformed"); + } + + const routeName = this.#kindToRoute.get(kind); + const state = routeName ? this.#routes.get(routeName) : undefined; + if (!state || !routeName) { + return this.#drop(record, "unroutable"); + } + + if ( + !state.route.replayable && + this.#appliedThrough !== undefined && + record.seqNum <= this.#appliedThrough + ) { + return this.#drop(record, "replayed", routeName); + } + + if ( + state.route.delivery === "at-arrival" && + state.handlers.size === 0 && + state.waiters.length === 0 + ) { + return this.#drop(record, "no-handler", routeName); + } + + const waiter = state.waiters.shift(); + if (waiter) { + if (waiter.timer) clearTimeout(waiter.timer); + waiter.resolve(record); + return { action: "deliver", route: routeName }; + } + + if (state.handlers.size > 0) { + this.#invokeHandlers(state, record); + return { action: "deliver", route: routeName }; + } + + state.queue.push(record); + return { action: "queue", route: routeName }; + } + + #kindOf(data: unknown): string | undefined { + try { + const kind = this.table.kindOf(data); + return typeof kind === "string" && kind.length > 0 ? kind : undefined; + } catch { + return undefined; + } + } + + #drop(record: SessionStreamRecord, reason: RouterDropReason, route?: string): RouterDecision { + try { + this.#onDrop?.(record, reason, route); + } catch { + void 0; + } + return { action: "drop", route, reason }; + } + + #invokeHandlers(state: RouteState, record: SessionStreamRecord): void { + for (const handler of state.handlers) { + try { + handler(record); + } catch { + void 0; + } + } + } + + /** + * The highest sequence that can be resumed past without losing anything. + * + * Held back below the earliest record still queued anywhere, because those + * are exactly the records a replay has to recover. Everything else the + * router has seen was terminally decided, so the floor is free to sit at the + * high water when every queue is empty. + */ + resumeFloor(): number | undefined { + if (this.#highestSeq === undefined) return undefined; + + let earliestPending = Infinity; + for (const state of this.#routes.values()) { + const pending = state.earliestUnrecovered(); + if (pending !== undefined) earliestPending = Math.min(earliestPending, pending); + } + + if (earliestPending === Infinity) return this.#highestSeq; + + const floor = Math.min(this.#highestSeq, earliestPending - 1); + return floor >= 0 ? floor : undefined; + } + + /** + * The high water: highest sequence observed, never held back. Published as + * the end of the replay window so the next boot can tell a re-read + * `at-arrival` record from one arriving live. + */ + appliedThrough(): number | undefined { + return this.#highestSeq; + } + + /** Both published numbers for a turn boundary. */ + checkpoint(): RouterCheckpoint { + return { resumeFrom: this.resumeFloor(), appliedThrough: this.appliedThrough() }; + } + + #stateOrThrow(name: string): RouteState { + const state = this.#routes.get(name); + if (!state) throw new Error(`Unknown session route "${name}"`); + return state; + } + + /** + * Attach a push handler. A `queued` route with a handler attached delivers + * straight to it instead of queueing; an `at-arrival` route discards records + * whenever no handler is attached. + * + * Attaching re-offers anything already queued, so a consumer that attaches + * after records have piled up still sees them in order. + */ + on(name: string, handler: RouteHandler): { off: () => void } { + const state = this.#stateOrThrow(name); + state.handlers.add(handler); + + if (state.queue.length > 0) { + const pending = state.queue.splice(0, state.queue.length); + for (const record of pending) { + try { + handler(record); + } catch { + void 0; + } + } + } + + return { + off: () => { + state.handlers.delete(handler); + }, + }; + } + + /** Whether an `at-arrival` route currently has anywhere to deliver. */ + hasHandler(name: string): boolean { + return this.#stateOrThrow(name).handlers.size > 0; + } + + /** + * Take the next record on a route. + * + * `timeoutMs: 0` is a non-blocking take. Omitted means wait indefinitely. + * Resolves `undefined` on timeout. An `at-arrival` route delivers to a + * waiting caller when one is already parked here, which is what makes a pull + * consumer possible on such a route without weakening the discard rule. + */ + next(name: string, options?: { timeoutMs?: number }): Promise { + const state = this.#stateOrThrow(name); + const queued = state.queue.shift(); + if (queued) return Promise.resolve(queued); + + if (options?.timeoutMs === 0) return Promise.resolve(undefined); + + return new Promise((resolve) => { + const waiter: QueueWaiter = { resolve }; + if (options?.timeoutMs !== undefined) { + waiter.timer = setTimeout(() => { + const index = state.waiters.indexOf(waiter); + if (index !== -1) state.waiters.splice(index, 1); + resolve(undefined); + }, options.timeoutMs); + } + state.waiters.push(waiter); + }); + } + + /** Head of a route's queue without consuming it. */ + peek(name: string): SessionStreamRecord | undefined { + return this.#stateOrThrow(name).queue[0]; + } + + /** + * Whether a route has anything queued. Exact, because it reads that route's + * own queue rather than the head of a buffer shared with every other kind on + * the channel. + */ + hasPending(name: string): boolean { + return this.#stateOrThrow(name).queue.length > 0; + } + + /** Number of records queued on a route. Diagnostic use. */ + pendingCount(name: string): number { + return this.#stateOrThrow(name).queue.length; + } + + /** + * Discard whatever one route has queued and wake its waiters empty. + * + * Closes a consumer window: a route that is not replayable has nothing owed + * to a later boot, so once its window is over anything still queued on it is + * dead and must not sit at the head of the queue for the rest of the run. + */ + clearRoute(name: string): void { + const state = this.#stateOrThrow(name); + state.queue.length = 0; + for (const waiter of state.waiters) { + if (waiter.timer) clearTimeout(waiter.timer); + waiter.resolve(undefined); + } + state.waiters.length = 0; + } + + /** Drop every waiter and queue. Called between task executions. */ + reset(): void { + for (const state of this.#routes.values()) { + state.queue.length = 0; + state.handlers.clear(); + for (const waiter of state.waiters) { + if (waiter.timer) clearTimeout(waiter.timer); + waiter.resolve(undefined); + } + state.waiters.length = 0; + } + this.#highestSeq = undefined; + this.#resumeFrom = undefined; + this.#appliedThrough = undefined; + } +} From 86672f39e81de73e169c419e90cbb609909780f8 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sat, 22 Aug 2026 07:40:16 +0100 Subject: [PATCH 14/20] refactor(chat,core): route session.in instead of tracking one cursor for every consumer `session.in` carries records for several consumers whose delivery needs differ: a user message must be delivered eventually and so can lag arbitrarily, while a stop only means anything to the turn that is live when it lands. One scalar cursor cannot describe both, and each fix so far has been a workaround for that: a clamp so the cursor retreats below a queued message, a barrier predicate so control records do not make it retreat, a drop predicate and a second header so the retreat does not re-apply control records that were already applied, and a drain so a record with no consumer does not park at the head of a shared buffer. The router replaces all of it. Every record is classified once and handed to one route; each route declares whether it queues when no consumer is ready and whether an unhandled record must survive into the next boot. The resume floor, the replay window and the discard-the-unowned behaviour then fall out of those declarations, so `hasPending` and `next` read their own route's queue rather than the head of a buffer shared with three other kinds, and the checkpoint is read and the subscription opened in one call, closing the window where a listener could attach before the resume cursor was seeded. No wire change: both turn-boundary headers keep their meanings, so existing sessions resume exactly as before and no webapp change is involved. Also converges the test double on production semantics. It decided consumption after awaiting handlers, which left a window where a handler registered mid-dispatch saw a buffer the record had not been added to yet, and it carried its own copy of the cursor clamp. It now decides synchronously like production and stubs only the network boundary, so the wait path runs its real implementation in tests. --- packages/core/src/v3/session-streams-api.ts | 1 + packages/core/src/v3/sessionStreams/index.ts | 36 +- .../core/src/v3/sessionStreams/manager.ts | 150 ++-- .../core/src/v3/sessionStreams/noopManager.ts | 26 +- packages/core/src/v3/sessionStreams/router.ts | 4 +- packages/core/src/v3/sessionStreams/types.ts | 33 +- .../v3/test/test-session-stream-manager.ts | 189 ++--- packages/trigger-sdk/src/v3/ai.ts | 782 +++++++----------- packages/trigger-sdk/src/v3/sessions.ts | 129 +-- .../src/v3/test/mock-chat-agent.ts | 3 + .../src/v3/test/test-session-handle.ts | 49 +- .../test/chat-messages-mailbox.test.ts | 67 +- .../test/replay-session-in.test.ts | 36 +- 13 files changed, 641 insertions(+), 864 deletions(-) diff --git a/packages/core/src/v3/session-streams-api.ts b/packages/core/src/v3/session-streams-api.ts index 4f5c979aa3b..638a8674213 100644 --- a/packages/core/src/v3/session-streams-api.ts +++ b/packages/core/src/v3/session-streams-api.ts @@ -7,3 +7,4 @@ export const sessionStreams = SessionStreamsAPI.getInstance(); export * from "./sessionStreams/types.js"; export * from "./sessionStreams/wireProtocol.js"; export * from "./sessionStreams/chatSnapshot.js"; +export * from "./sessionStreams/router.js"; diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 6519628ce08..a1b6f840cf9 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -41,6 +41,18 @@ export class SessionStreamsAPI implements SessionStreamManager { return this.#getManager().on(sessionId, io, handler); } + public onRecord( + sessionId: string, + io: SessionChannelIO, + handler: (record: SessionStreamRecord) => void | boolean | Promise + ): { off: () => void } { + const manager = this.#getManager(); + if (!manager.onRecord) { + throw new Error("The configured Session stream manager does not support record handlers"); + } + return manager.onRecord(sessionId, io, handler); + } + public once( sessionId: string, io: SessionChannelIO, @@ -86,26 +98,6 @@ export class SessionStreamsAPI implements SessionStreamManager { return manager.peekRecord(sessionId, io); } - public highestConsumedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - return this.#getManager().highestConsumedSeqNum?.(sessionId, io); - } - - public setDropPredicate( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void { - this.#getManager().setDropPredicate?.(sessionId, io, predicate); - } - - public setCursorBarrier( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void { - this.#getManager().setCursorBarrier?.(sessionId, io, predicate); - } - public lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.#getManager().lastSeqNum(sessionId, io); } @@ -142,6 +134,10 @@ export class SessionStreamsAPI implements SessionStreamManager { return this.#getManager().shiftBuffer(sessionId, io); } + public reconnectStream(sessionId: string, io: SessionChannelIO): void { + this.#getManager().reconnectStream?.(sessionId, io); + } + public disconnectStream(sessionId: string, io: SessionChannelIO): void { this.#getManager().disconnectStream(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index 541e4182398..73c85f972e2 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -17,6 +17,17 @@ import { controlSubtype } from "./wireProtocol.js"; // available to other consumers. See `SessionStreamManager.on` in types.ts. type SessionStreamHandler = (data: unknown) => void | boolean | Promise; +/** + * A handler that sees the whole record rather than just its payload. Consumers + * that route by sequence number need the metadata, the same reason + * `onceRecord` exists alongside `once`. + */ +type SessionStreamRecordHandler = (record: SessionStreamRecord) => void | boolean | Promise; + +type RegisteredHandler = + | { kind: "data"; fn: SessionStreamHandler } + | { kind: "record"; fn: SessionStreamRecordHandler }; + type OnceWaiter = { resolve: (result: InputStreamOnceResult) => void; reject: (error: Error) => void; @@ -48,7 +59,7 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { * stream SSE. */ export class StandardSessionStreamManager implements SessionStreamManager { - private handlers = new Map>(); + private handlers = new Map>(); private onceWaiters = new Map(); private buffer = new Map(); private tails = new Map(); @@ -71,23 +82,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { // without depending on buffer traversal. private unconsumedSeqNums = new Map>(); - /** - * Per-channel predicate deciding which buffered records hold the persisted - * cursor back. Absent means every record does, which is the conservative - * default. Consumers that know their record kinds narrow it so the cursor is - * only held behind records whose loss would matter. - */ - private cursorBarriers = new Map(); - - /** - * Per-channel predicate marking records that must not be delivered again. - * A resume cursor held back behind an unhandled record necessarily replays - * everything after it, including records that WERE handled before the - * previous run ended. Re-delivering those is not harmless: a control record - * applies a second time to whatever turn happens to be live. This lets the - * consumer name the ones to drop. - */ - private dropPredicates = new Map(); // High-water mark of seq_nums that have been *consumed* (delivered to a // once() waiter or shifted off the buffer into a once() caller) on a channel. // Distinct from `seqNums`, which advances whenever any record is @@ -110,6 +104,26 @@ export class StandardSessionStreamManager implements SessionStreamManager { ) {} on(sessionId: string, io: SessionChannelIO, handler: SessionStreamHandler): { off: () => void } { + return this.#register(sessionId, io, { kind: "data", fn: handler }); + } + + /** + * Register a handler that receives the full record, including its sequence + * number. Same consume semantics as {@link on}: returning `true` consumes. + */ + onRecord( + sessionId: string, + io: SessionChannelIO, + handler: SessionStreamRecordHandler + ): { off: () => void } { + return this.#register(sessionId, io, { kind: "record", fn: handler }); + } + + #register( + sessionId: string, + io: SessionChannelIO, + handler: RegisteredHandler + ): { off: () => void } { const key = keyFor(sessionId, io); let handlerSet = this.handlers.get(key); @@ -139,7 +153,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { if (buffered && buffered.length > 0) { const keptRecords: SessionStreamRecord[] = []; for (const record of buffered) { - const consumed = this.#invokeHandler(handler, record.data); + const consumed = this.#invokeHandler(handler, record); if (consumed) { this.#advanceLastDispatched(key, record.seqNum); } else { @@ -347,55 +361,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } - /** The highest consumed sequence, unclamped. @see lastDispatchedSeqNum */ - highestConsumedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - return this.lastDispatchedSeqNums.get(keyFor(sessionId, io)); - } - - setDropPredicate( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void { - const key = keyFor(sessionId, io); - if (predicate) { - this.dropPredicates.set(key, predicate); - } else { - this.dropPredicates.delete(key); - } - } - - setCursorBarrier( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void { - const key = keyFor(sessionId, io); - if (predicate) { - this.cursorBarriers.set(key, predicate); - } else { - this.cursorBarriers.delete(key); - } - } - - /** - * Fails safe: an absent or throwing predicate treats the record as a barrier, - * so a mistake here can only make the cursor more conservative, never skip a - * record. - */ - #isCursorBarrier(key: string, record: SessionStreamRecord): boolean { - const predicate = this.cursorBarriers.get(key); - if (!predicate) return true; - try { - return predicate(record); - } catch (error) { - if (this.debug) { - console.error("[SessionStreamManager] Cursor barrier predicate error:", error); - } - return true; - } - } - #markUnconsumedRecord(key: string, seqNum: number): void { if (!Number.isFinite(seqNum)) return; @@ -457,6 +422,19 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.reconnectAttempts.delete(key); } + /** + * Re-open a channel that `disconnectStream` closed, without registering a + * new consumer. A single long-lived reader (the session channel router) has + * to be able to bring its own tail back after a suspend, and re-attaching + * its handler just to clear the suppression flag would replay the buffer at + * it. + */ + reconnectStream(sessionId: string, io: SessionChannelIO): void { + const key = keyFor(sessionId, io); + this.explicitlyDisconnected.delete(key); + this.#ensureTailConnected(sessionId, io); + } + clearHandlers(): void { this.handlers.clear(); @@ -490,8 +468,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.seqNums.clear(); this.lastDispatchedSeqNums.clear(); this.unconsumedSeqNums.clear(); - this.cursorBarriers.clear(); - this.dropPredicates.clear(); this.minTimestamps.clear(); this.handlers.clear(); this.reconnectAttempts.clear(); @@ -634,24 +610,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { } #dispatch(key: string, record: SessionStreamRecord): void { - const drop = this.dropPredicates.get(key); - if (drop) { - let shouldDrop = false; - try { - shouldDrop = drop(record); - } catch (error) { - if (this.debug) { - console.error("[SessionStreamManager] Drop predicate error:", error); - } - } - if (shouldDrop) { - // Acknowledge it so the tail does not fetch it again, but never hand - // it to a waiter or a handler. - this.#advanceLastDispatched(key, record.seqNum); - return; - } - } - // Any record flowing through = healthy connection; reset the backoff // counter so the next disconnect starts fresh. this.reconnectAttempts.delete(key); @@ -665,7 +623,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { // records advance the cursor in `once()` / `shiftBuffer()`. this.#advanceLastDispatched(key, record.seqNum); waiter.resolve({ ok: true, output: record }); - this.#invokeHandlers(key, record.data); + this.#invokeHandlers(key, record); return; } @@ -677,7 +635,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { // second turn. Records no handler consumed (e.g. a message arriving // while only the stop facade is attached during preload) are buffered // so a subsequent `once()` can still pick them up. - const consumed = this.#invokeHandlers(key, record.data); + const consumed = this.#invokeHandlers(key, record); if (consumed) { this.#advanceLastDispatched(key, record.seqNum); return; @@ -689,9 +647,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.buffer.set(key, buffered); } buffered.push(record); - if (this.#isCursorBarrier(key, record)) { - this.#markUnconsumedRecord(key, record.seqNum); - } + this.#markUnconsumedRecord(key, record.seqNum); this.#drainOnceWaitersFromBuffer(key); } @@ -736,12 +692,12 @@ export class StandardSessionStreamManager implements SessionStreamManager { } /** Returns true when any handler consumed the record. All handlers are invoked regardless. */ - #invokeHandlers(key: string, data: unknown): boolean { + #invokeHandlers(key: string, record: SessionStreamRecord): boolean { const handlers = this.handlers.get(key); if (!handlers) return false; let consumed = false; for (const handler of handlers) { - if (this.#invokeHandler(handler, data)) { + if (this.#invokeHandler(handler, record)) { consumed = true; } } @@ -749,9 +705,9 @@ export class StandardSessionStreamManager implements SessionStreamManager { } /** Returns true when the handler synchronously consumed the record (returned `true`). */ - #invokeHandler(handler: SessionStreamHandler, data: unknown): boolean { + #invokeHandler(handler: RegisteredHandler, record: SessionStreamRecord): boolean { try { - const result = handler(data); + const result = handler.kind === "record" ? handler.fn(record) : handler.fn(record.data); if (result === true) return true; if (result && typeof result === "object" && "catch" in result) { (result as Promise).catch((error) => { diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index 207434d92c0..8e894dfc6d3 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -16,6 +16,14 @@ export class NoopSessionStreamManager implements SessionStreamManager { return { off: () => {} }; } + onRecord( + _sessionId: string, + _io: SessionChannelIO, + _handler: (record: SessionStreamRecord) => void | boolean | Promise + ): { off: () => void } { + return { off: () => {} }; + } + once( _sessionId: string, _io: SessionChannelIO, @@ -55,22 +63,6 @@ export class NoopSessionStreamManager implements SessionStreamManager { return undefined; } - highestConsumedSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { - return undefined; - } - - setDropPredicate( - _sessionId: string, - _io: SessionChannelIO, - _predicate: SessionStreamRecordPredicate | undefined - ): void {} - - setCursorBarrier( - _sessionId: string, - _io: SessionChannelIO, - _predicate: SessionStreamRecordPredicate | undefined - ): void {} - lastSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { return undefined; } @@ -97,6 +89,8 @@ export class NoopSessionStreamManager implements SessionStreamManager { disconnectStream(_sessionId: string, _io: SessionChannelIO): void {} + reconnectStream(_sessionId: string, _io: SessionChannelIO): void {} + clearHandlers(): void {} reset(): void {} diff --git a/packages/core/src/v3/sessionStreams/router.ts b/packages/core/src/v3/sessionStreams/router.ts index 0563b3820fe..0adbe0475e4 100644 --- a/packages/core/src/v3/sessionStreams/router.ts +++ b/packages/core/src/v3/sessionStreams/router.ts @@ -7,7 +7,7 @@ import type { SessionStreamRecord } from "./types.js"; * - `at-arrival`: it goes to a live handler or nowhere. A record that only * means something to the turn that is live when it lands (a stop) is this. */ -type RouteDelivery = "queue" | "at-arrival"; +export type RouteDelivery = "queue" | "at-arrival"; /** * One route: which kinds it owns, whether it waits for a consumer, and whether @@ -25,7 +25,7 @@ type RouteDelivery = "queue" | "at-arrival"; * The fourth is a contradiction (discard it when nobody is listening, yet * recover it later) and the table rejects it. */ -type SessionRoute = { +export type SessionRoute = { /** Unique within a table. */ name: string; delivery: RouteDelivery; diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index bd8901eb0a5..cc6bde884cc 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -50,6 +50,16 @@ export interface SessionStreamManager { handler: (data: unknown) => void | boolean | Promise ): { off: () => void }; + /** + * Register a handler that receives the full record, including its sequence + * number. Same consume semantics as {@link on}. + */ + onRecord?( + sessionId: string, + io: SessionChannelIO, + handler: (record: SessionStreamRecord) => void | boolean | Promise + ): { off: () => void }; + /** Wait for the next record on the given channel (buffered or live). */ once( sessionId: string, @@ -82,26 +92,6 @@ export interface SessionStreamManager { /** Non-blocking peek at the head record, including its durable metadata. */ peekRecord?(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined; - /** - * Narrow which buffered records hold the persisted cursor back. Absent means - * every record does. - */ - /** The highest consumed sequence, unclamped. */ - highestConsumedSeqNum?(sessionId: string, io: SessionChannelIO): number | undefined; - - /** Mark records that must never be delivered again on this boot. */ - setDropPredicate?( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void; - - setCursorBarrier?( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void; - /** Last S2 sequence number seen on the given channel. */ lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined; @@ -147,6 +137,9 @@ export interface SessionStreamManager { /** Abort the SSE tail while preserving buffered records. Called before `.wait` suspends. */ disconnectStream(sessionId: string, io: SessionChannelIO): void; + /** Re-open a channel closed by {@link disconnectStream}, registering nothing. */ + reconnectStream?(sessionId: string, io: SessionChannelIO): void; + /** Clear all `.on` handlers; abort tails without pending once-waiters. */ clearHandlers(): void; diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index 3626f117e0b..5388c9dc40e 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -20,6 +20,9 @@ type OnceWaiter = { // returns `true` CONSUMES the record (not buffered, not re-delivered on a // future `on()` attach). See `SessionStreamManager.on` in types.ts. type Handler = (data: unknown) => void | boolean | Promise; +type RecordHandler = (record: SessionStreamRecord) => void | boolean | Promise; + +type RegisteredHandler = { kind: "data"; fn: Handler } | { kind: "record"; fn: RecordHandler }; function keyFor(sessionId: string, io: SessionChannelIO): string { return `${sessionId}:${io}`; @@ -35,16 +38,25 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { * registered are buffered so the first `once()` picks them up. */ export class TestSessionStreamManager implements SessionStreamManager { - private handlers = new Map>(); + private handlers = new Map>(); private onceWaiters = new Map(); private buffer = new Map(); private seqNums = new Map(); private dispatchedSeqNums = new Map(); - private unconsumedSeqNums = new Map>(); - private cursorBarriers = new Map(); - private dropPredicates = new Map(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { + return this.#register(sessionId, io, { kind: "data", fn: handler }); + } + + onRecord(sessionId: string, io: SessionChannelIO, handler: RecordHandler): { off: () => void } { + return this.#register(sessionId, io, { kind: "record", fn: handler }); + } + + #register( + sessionId: string, + io: SessionChannelIO, + handler: RegisteredHandler + ): { off: () => void } { const key = keyFor(sessionId, io); let set = this.handlers.get(key); @@ -68,7 +80,7 @@ export class TestSessionStreamManager implements SessionStreamManager { for (const record of buffered) { let consumed = false; try { - consumed = handler(record.data) === true; + consumed = this.#callHandler(handler, record) === true; } catch { // Never let a handler error break test state } @@ -199,30 +211,6 @@ export class TestSessionStreamManager implements SessionStreamManager { return this.peekRecord(sessionId, io)?.data; } - highestConsumedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - return this.dispatchedSeqNums.get(keyFor(sessionId, io)); - } - - setDropPredicate( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void { - const key = keyFor(sessionId, io); - if (predicate) this.dropPredicates.set(key, predicate); - else this.dropPredicates.delete(key); - } - - setCursorBarrier( - sessionId: string, - io: SessionChannelIO, - predicate: SessionStreamRecordPredicate | undefined - ): void { - const key = keyFor(sessionId, io); - if (predicate) this.cursorBarriers.set(key, predicate); - else this.cursorBarriers.delete(key); - } - peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined { return this.buffer.get(keyFor(sessionId, io))?.[0]; } @@ -252,20 +240,7 @@ export class TestSessionStreamManager implements SessionStreamManager { } lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - const key = keyFor(sessionId, io); - const highWatermark = this.dispatchedSeqNums.get(key); - if (highWatermark === undefined) return undefined; - - const unconsumedSeqNums = this.unconsumedSeqNums.get(key); - if (!unconsumedSeqNums || unconsumedSeqNums.size === 0) return highWatermark; - - let earliestUnconsumedSeqNum = Infinity; - for (const seqNum of unconsumedSeqNums) { - earliestUnconsumedSeqNum = Math.min(earliestUnconsumedSeqNum, seqNum); - } - - const safeCursor = Math.min(highWatermark, earliestUnconsumedSeqNum - 1); - return safeCursor >= 0 ? safeCursor : undefined; + return this.dispatchedSeqNums.get(keyFor(sessionId, io)); } setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void { @@ -275,7 +250,6 @@ export class TestSessionStreamManager implements SessionStreamManager { } #advanceLastDispatched(key: string, seqNum: number): void { - this.#removeUnconsumedRecord(key, seqNum); if (!Number.isFinite(seqNum)) return; const current = this.dispatchedSeqNums.get(key); if (current === undefined || seqNum > current) { @@ -283,35 +257,6 @@ export class TestSessionStreamManager implements SessionStreamManager { } } - #isCursorBarrier(key: string, record: SessionStreamRecord): boolean { - const predicate = this.cursorBarriers.get(key); - if (!predicate) return true; - try { - return predicate(record); - } catch { - return true; - } - } - - #markUnconsumedRecord(key: string, seqNum: number): void { - if (!Number.isFinite(seqNum)) return; - - let unconsumedSeqNums = this.unconsumedSeqNums.get(key); - if (!unconsumedSeqNums) { - unconsumedSeqNums = new Set(); - this.unconsumedSeqNums.set(key, unconsumedSeqNums); - } - unconsumedSeqNums.add(seqNum); - } - - #removeUnconsumedRecord(key: string, seqNum: number): void { - const unconsumedSeqNums = this.unconsumedSeqNums.get(key); - unconsumedSeqNums?.delete(seqNum); - if (unconsumedSeqNums?.size === 0) { - this.unconsumedSeqNums.delete(key); - } - } - setMinTimestamp( _sessionId: string, _io: SessionChannelIO, @@ -357,7 +302,6 @@ export class TestSessionStreamManager implements SessionStreamManager { this.buffer.clear(); this.seqNums.clear(); this.dispatchedSeqNums.clear(); - this.unconsumedSeqNums.clear(); } disconnect(): void { @@ -411,42 +355,24 @@ export class TestSessionStreamManager implements SessionStreamManager { if (waiter) { this.#advanceLastDispatched(key, record.seqNum); waiter.resolve({ ok: true, output: record }); - await this.#invokeHandlers(key, record.data); + await this.#invokeHandlers(key, record); return; } - const consumed = await this.#invokeHandlers(key, record.data); - if (consumed) { - this.#advanceLastDispatched(key, record.seqNum); - return; - } - - // Re-check waiters: handler invocation above is awaited (unlike the - // synchronous production dispatch), and the runtime commonly registers - // its next `once()` during that window — e.g. the turn loop reaching - // `waitWithIdleTimeout` while a handler settles. Without this second - // look the record would be buffered while the fresh waiter hangs. - const bufferedAfterHandlers = this.buffer.get(key); - const lateWaiter = - bufferedAfterHandlers && bufferedAfterHandlers.length > 0 - ? undefined - : this.#takeOnceWaiter(key, record); - if (lateWaiter) { + const { consumed, settled } = this.#invokeHandlersSync(key, record); + if (!consumed) { + let buffered = this.buffer.get(key); + if (!buffered) { + buffered = []; + this.buffer.set(key, buffered); + } + buffered.push(record); + this.#drainOnceWaitersFromBuffer(key); + } else { this.#advanceLastDispatched(key, record.seqNum); - lateWaiter.resolve({ ok: true, output: record }); - return; } - let buffered = this.buffer.get(key); - if (!buffered) { - buffered = []; - this.buffer.set(key, buffered); - } - buffered.push(record); - if (this.#isCursorBarrier(key, record)) { - this.#markUnconsumedRecord(key, record.seqNum); - } - this.#drainOnceWaitersFromBuffer(key); + await settled; } #takeOnceWaiter(key: string, record: SessionStreamRecord): OnceWaiter | undefined { @@ -492,26 +418,49 @@ export class TestSessionStreamManager implements SessionStreamManager { * Wrapped per-handler so a throwing/rejecting handler doesn't poison * Promise.all and break unrelated test state. */ - async #invokeHandlers(key: string, data: unknown): Promise { + async #invokeHandlers(key: string, record: SessionStreamRecord): Promise { + const { consumed, settled } = this.#invokeHandlersSync(key, record); + await settled; + return consumed; + } + + /** + * Decide consumption synchronously, exactly like the production dispatch, + * and hand back a promise for any async handler work so callers can still + * await it. Splitting the decision from the awaiting is what keeps a handler + * registered mid-dispatch from seeing an inconsistent buffer. + */ + #invokeHandlersSync( + key: string, + record: SessionStreamRecord + ): { consumed: boolean; settled: Promise } { const handlers = this.handlers.get(key); - if (!handlers || handlers.size === 0) return false; + if (!handlers || handlers.size === 0) { + return { consumed: false, settled: Promise.resolve() }; + } let consumed = false; - await Promise.all( - Array.from(handlers).map(async (h) => { - try { - const result = h(data); - if (result === true) { - consumed = true; - return; - } - await result; - } catch { - // Never let a handler error break test state + const pending: Array> = []; + for (const handler of Array.from(handlers)) { + try { + const result = this.#callHandler(handler, record); + if (result === true) { + consumed = true; + continue; } - }) - ); - return consumed; + if (result) pending.push(Promise.resolve(result).catch(() => {})); + } catch { + continue; + } + } + return { consumed, settled: Promise.all(pending) }; + } + + #callHandler( + handler: RegisteredHandler, + record: SessionStreamRecord + ): void | boolean | Promise { + return handler.kind === "record" ? handler.fn(record) : handler.fn(record.data); } /** diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index ebe5d1d7842..17cce3fce25 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -13,7 +13,6 @@ import { type inferSchemaIn, type inferSchemaOut, InputStreamOncePromise, - type InputStreamOnceResult, isAdditionalApiKey, isSchemaZodEsque, logger, @@ -29,6 +28,8 @@ import { SESSION_IN_CONSUMED_ID_HEADER, SESSION_IN_EVENT_ID_HEADER, sessionStreams, + SessionChannelRouter, + InputStreamTimeoutError, taskContext, type TaskIdentifier, type TaskOptions, @@ -37,6 +38,8 @@ import { type TaskWithSchema, TRIGGER_CONTROL_SUBTYPE, type StreamWriteResult, + type RouterCheckpoint, + type SessionRouteTable, } from "@trigger.dev/core/v3"; import type { FinishReason, @@ -223,89 +226,6 @@ async function findLatestSessionInCursor(chatId: string): Promise { - try { - return await findLatestSessionInConsumed(chatId); - } catch { - return undefined; - } -} - -/** - * The highest `.in` sequence a previous run had already consumed, read from the - * latest `turn-complete` on `.out`. - * - * Absent for a chat whose turns predate the header, in which case no record is - * dropped and behaviour matches the previous release. - * @internal - */ -async function findLatestSessionInConsumed(chatId: string): Promise { - const apiClient = apiClientManager.clientOrThrow(); - const response = await apiClient.readSessionStreamRecords(chatId, "out"); - let latest: number | undefined; - for (const record of response.records) { - if (controlSubtype(record.headers) !== TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE) continue; - const raw = headerValue(record.headers, SESSION_IN_CONSUMED_ID_HEADER); - if (!raw) continue; - const parsed = Number.parseInt(raw, 10); - if (Number.isFinite(parsed)) latest = parsed; - } - return latest; -} - -/** Test-only entry point for the records-based cursor scan. @internal */ -export async function __findLatestSessionInCursorForTests( - chatId: string -): Promise { - return findLatestSessionInCursor(chatId); -} - -/** - * Seed the `.in` resume cursor for custom-agent loops (`chat.customAgent` - * raw loops and `chat.createSession`) the way `chat.agent`'s boot does. - * - * MUST run before anything attaches a `.in` listener (`createStopSignal`, - * `chat.messages.on`, the first wait): attaching opens the SSE tail with - * `Last-Event-ID` from the seeded cursor, so attach-then-seed replays - * every record from seq 0 — already-answered user messages get delivered - * into the new run's first wait and the loop re-answers them. - * - * Seeds both cursors: `setLastSeqNum` controls the SSE `Last-Event-ID`, - * `setLastDispatchedSeqNum` gates waiter dispatch — seeding only the - * former still re-delivers records the manager buffered before the seed. - * - * No-ops on fresh boots and when a cursor is already seeded (e.g. the - * `chatCustomAgent` wrapper ran before a nested `createChatSession`). - * @internal - */ -async function seedSessionInResumeCursorForCustomLoop( - payload: Pick -): Promise { - if (sessionStreams.lastSeqNum(payload.chatId, "in") !== undefined) return; - // No continuation/attempt gate: the wire may omit `continuation` on a - // run that still has prior turns (chat.agent covers that case via its - // snapshot). The scan doubles as the prior-state probe — a fresh - // session has no turn-complete on `.out`, returns no cursor, and - // seeds nothing. Cost on fresh boots is one non-blocking records read. - try { - const cursor = await findLatestSessionInCursor(payload.chatId); - if (cursor !== undefined) { - sessionStreams.setLastSeqNum(payload.chatId, "in", cursor); - sessionStreams.setLastDispatchedSeqNum(payload.chatId, "in", cursor); - } - } catch (error) { - logger.warn("chat session: session.in resume cursor lookup failed; old messages may replay", { - error: error instanceof Error ? error.message : String(error), - }); - } -} - /** * Versioned blob written to S3 after every turn completes (when no * `hydrateMessages` hook is registered). Read at run boot to seed the @@ -1601,92 +1521,121 @@ export type ChatMessages = RealtimeDefinedInputStream & { }; /** - * Only message records hold the persisted `.in` cursor back. + * Read one record from a route, suspending the run if nothing is there yet. * - * The `session-in-event-id` header serves two consumers with opposite needs: - * `findLatestSessionInCursor` reads it as a resume cursor and wants it - * conservative, while a client reads it to correlate its own send's - * turn-complete and wants it exact. Holding the cursor behind an unconsumed - * control record satisfies neither: resume safety does not need it (replaying a - * stop or a handover is benign, and a handover for a turn that never ran is - * discarded), while a client comparing the header against its own append - * sequence sees a value below its send and discards its own turn boundary. + * The wake and the read are separate steps: the channel wakes the run, then the + * router hands over whatever it routed. Nothing else can take the record in + * between, which is what keeps the published cursors and the delivered record + * in agreement. * @internal */ -function isChatCursorBarrier(record: { data: unknown }): boolean { - return (record.data as ChatInputChunk | undefined)?.kind === "message"; -} +async function waitOnChatRoute( + route: string, + options: { + idleTimeoutInSeconds?: number; + timeout?: string; + spanName?: string; + skipSuspend?: boolean; + onSuspend?: () => Promise | void; + onResume?: () => Promise | void; + } +): Promise<{ ok: true; output: T } | { ok: false; error?: Error }> { + const router = chatInputRouter(); + const session = getChatSession(); -function isChatMessageRecord(record: { data: unknown }): boolean { - return (record.data as ChatInputChunk | undefined)?.kind === "message"; + return tracer.startActiveSpan( + options.spanName ?? `chat.${route}.wait()`, + async (span) => { + const idleMs = (options.idleTimeoutInSeconds ?? 0) * 1000; + if (idleMs > 0) { + const warm = await router.next(route, { timeoutMs: idleMs }); + if (warm) { + span.setAttribute("wait.resolved", "idle"); + return { ok: true as const, output: warm.data as T }; + } + } else { + const buffered = await router.next(route, { timeoutMs: 0 }); + if (buffered) { + span.setAttribute("wait.resolved", "buffered"); + return { ok: true as const, output: buffered.data as T }; + } + } + + if (options.skipSuspend) { + span.setAttribute("wait.resolved", "skipped"); + return { + ok: false as const, + error: new Error("Idle timeout elapsed and skipSuspend is set"), + }; + } + + if (options.onSuspend) await options.onSuspend(); + + span.setAttribute("wait.resolved", "suspended"); + while (true) { + const wake = await session.in.awaitWake({ + timeout: options.timeout, + lastSeqNum: router.resumeFloor(), + }); + if (!wake.ok) { + span.recordException(wake.error); + return { ok: false as const, error: wake.error }; + } + + const record = await router.next(route); + if (!record) continue; + + if (options.onResume) await options.onResume(); + return { ok: true as const, output: record.data as T }; + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "sessions", + session: session.id, + io: "in", + route, + ...accessoryAttributes({ + items: [{ text: `${session.id}.in:${route}`, variant: "normal" }], + style: "codepath", + }), + }, + } + ); } const messagesInput: ChatMessages = { id: "chat-messages", on(handler) { - return getChatSession().in.on((chunk) => { - if (chunk.kind === "message") { - // Returning `true` marks the record CONSUMED at the manager level: - // it is neither buffered for a later `once()` nor re-delivered by - // the buffer drain when the next turn re-attaches its handler. - // Without this, a message arriving mid-stream was delivered twice - // and ran a duplicate turn. - void Promise.resolve(handler(chunk.payload)).catch(() => {}); - return true; - } - return undefined; + return chatInputRouter().on(CHAT_ROUTE_MESSAGES, (record) => { + const chunk = record.data as Extract; + void Promise.resolve(handler(chunk.payload)).catch(() => {}); }); }, once(options) { - const ctx = taskContext.ctx; - const runId = ctx?.run.id; - return new InputStreamOncePromise((resolve, reject) => { - tracer - .startActiveSpan( - options?.spanName ?? `chat.messages.once()`, - async () => { - while (true) { - const result = await getChatSession().in.once(options); - if (!result.ok) { - resolve(result as InputStreamOnceResult); - return; - } - if (result.output.kind === "message") { - resolve({ ok: true, output: result.output.payload }); - return; - } - // Non-message chunks (stops) are handled by the stopInput - // facade's persistent listener; loop and wait for the next. - } - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "streams", - [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream", - ...(runId - ? { - [SemanticInternalAttributes.ENTITY_ID]: `${runId}:chat-messages`, - } - : {}), - streamId: "chat-messages", - ...accessoryAttributes({ - items: [{ text: "chat-messages", variant: "normal" }], - style: "codepath", - }), - }, + chatInputRouter() + .next(CHAT_ROUTE_MESSAGES, { timeoutMs: options?.timeoutMs }) + .then((record) => { + if (!record) { + resolve({ + ok: false, + error: new InputStreamTimeoutError("chat-messages", options?.timeoutMs ?? 0), + }); + return; } - ) - .catch(reject); + const chunk = record.data as Extract; + resolve({ ok: true, output: chunk.payload }); + }, reject); }); }, peek() { - const chunk = getChatSession().in.peek(); - if (chunk && chunk.kind === "message") return chunk.payload; - return undefined; + const record = chatInputRouter().peek(CHAT_ROUTE_MESSAGES); + if (!record) return undefined; + return (record.data as Extract).payload; }, async hasPending() { - return messagesInput.peek() !== undefined; + return chatInputRouter().hasPending(CHAT_ROUTE_MESSAGES); }, async next(options) { const timeoutInSeconds = options?.timeoutInSeconds; @@ -1699,59 +1648,41 @@ const messagesInput: ChatMessages = { ); } - const session = getChatSession(); - const result = await sessionStreams.onceRecordWhere( - session.id, - "in", - isChatMessageRecord, - timeoutInSeconds === undefined ? undefined : { timeoutMs: timeoutInSeconds * 1000 } - ); - if (!result.ok) return undefined; + const record = await chatInputRouter().next(CHAT_ROUTE_MESSAGES, { + timeoutMs: timeoutInSeconds === undefined ? undefined : timeoutInSeconds * 1000, + }); + if (!record) return undefined; - const chunk = result.output.data as Extract; - return { - id: result.output.id, - seqNum: result.output.seqNum, - payload: chunk.payload, - }; + const chunk = record.data as Extract; + return { id: record.id, seqNum: record.seqNum, payload: chunk.payload }; }, wait(options) { return new ManualWaitpointPromise(async (resolve, reject) => { try { - while (true) { - const result = await getChatSession().in.wait(options); - if (!result.ok) { - resolve(result); - return; - } - if (result.output.kind === "message") { - resolve({ ok: true, output: result.output.payload }); - return; - } - // Stop chunks are handled by the stopInput facade's persistent - // listener; loop back into the suspending wait. - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_MESSAGES, + { timeout: options?.timeout, spanName: options?.spanName } + ); + resolve( + result.ok + ? { ok: true, output: result.output.payload } + : { ok: false, error: result.error ?? new Error("Timed out") } + ); } catch (error) { reject(error); } }); }, async waitWithIdleTimeout(options) { - while (true) { - const result = await getChatSession().in.waitWithIdleTimeout(options); - if (!result.ok) return result; - if (result.output.kind === "message") { - return { ok: true, output: result.output.payload }; - } - // Swallow stop-kind chunks — persistent stop listener already handled - // the abort; we just loop for the next message. - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_MESSAGES, + options + ); + return result.ok + ? { ok: true, output: result.output.payload } + : { ok: false, error: result.error }; }, async send(_runId, data, options) { - // The `runId` argument is kept for signature parity with - // `RealtimeDefinedInputStream` but ignored — sessions are addressed - // by sessionId, not runId. Callers producing messages from outside - // the run should prefer the transport's `session.in.send(...)` path. await getChatSession().in.send( { kind: "message", payload: data } satisfies ChatInputChunk, options?.requestOptions @@ -1762,98 +1693,59 @@ const messagesInput: ChatMessages = { const stopInput: RealtimeDefinedInputStream<{ stop: true; message?: string }> = { id: "chat-stop", on(handler) { - return getChatSession().in.on((chunk) => { - if (chunk.kind === "stop") { - // Consume stop records (see the messages facade above). A stop is - // only meaningful to the turn it interrupts — buffering it would - // let a stale stop abort a future turn. - void Promise.resolve(handler({ stop: true, message: chunk.message })).catch(() => {}); - return true; - } - return undefined; + return chatInputRouter().on(CHAT_ROUTE_STOP, (record) => { + const chunk = record.data as Extract; + void Promise.resolve(handler({ stop: true, message: chunk.message })).catch(() => {}); }); }, once(options) { - const ctx = taskContext.ctx; - const runId = ctx?.run.id; - return new InputStreamOncePromise<{ stop: true; message?: string }>((resolve, reject) => { - tracer - .startActiveSpan( - options?.spanName ?? `chat.stop.once()`, - async () => { - while (true) { - const result = await getChatSession().in.once(options); - if (!result.ok) { - resolve(result as InputStreamOnceResult<{ stop: true; message?: string }>); - return; - } - if (result.output.kind === "stop") { - resolve({ - ok: true, - output: { stop: true, message: result.output.message }, - }); - return; - } - } - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "streams", - [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream", - ...(runId - ? { - [SemanticInternalAttributes.ENTITY_ID]: `${runId}:chat-stop`, - } - : {}), - streamId: "chat-stop", - ...accessoryAttributes({ - items: [{ text: "chat-stop", variant: "normal" }], - style: "codepath", - }), - }, + chatInputRouter() + .next(CHAT_ROUTE_STOP, { timeoutMs: options?.timeoutMs }) + .then((record) => { + if (!record) { + resolve({ + ok: false, + error: new InputStreamTimeoutError("chat-stop", options?.timeoutMs ?? 0), + }); + return; } - ) - .catch(reject); + const chunk = record.data as Extract; + resolve({ ok: true, output: { stop: true, message: chunk.message } }); + }, reject); }); }, peek() { - const chunk = getChatSession().in.peek(); - if (chunk && chunk.kind === "stop") { - return { stop: true, message: chunk.message }; - } - return undefined; + const record = chatInputRouter().peek(CHAT_ROUTE_STOP); + if (!record) return undefined; + const chunk = record.data as Extract; + return { stop: true, message: chunk.message }; }, wait(options) { return new ManualWaitpointPromise<{ stop: true; message?: string }>(async (resolve, reject) => { try { - while (true) { - const result = await getChatSession().in.wait(options); - if (!result.ok) { - resolve(result); - return; - } - if (result.output.kind === "stop") { - resolve({ - ok: true, - output: { stop: true, message: result.output.message }, - }); - return; - } - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_STOP, + { timeout: options?.timeout, spanName: options?.spanName } + ); + resolve( + result.ok + ? { ok: true, output: { stop: true, message: result.output.message } } + : { ok: false, error: result.error ?? new Error("Timed out") } + ); } catch (error) { reject(error); } }); }, async waitWithIdleTimeout(options) { - while (true) { - const result = await getChatSession().in.waitWithIdleTimeout(options); - if (!result.ok) return result; - if (result.output.kind === "stop") { - return { ok: true, output: { stop: true, message: result.output.message } }; - } - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_STOP, + options + ); + return result.ok + ? { ok: true as const, output: { stop: true, message: result.output.message } } + : { ok: false as const, error: result.error }; }, async send(_runId, data, options) { await getChatSession().in.send( @@ -1901,16 +1793,8 @@ const handoverInput = { spanName?: string; skipSuspend?: boolean; }) { - while (true) { - const result = await getChatSession().in.waitWithIdleTimeout(options); - if (!result.ok) return result; - if (result.output.kind === "handover" || result.output.kind === "handover-skip") { - return { ok: true as const, output: result.output as HandoverSignal }; - } - // Other kinds (message, stop) are not expected during handover-prepare. - // Loop back; the message and stop facades have their own listeners - // running so signals on those kinds aren't lost. - } + const result = await waitOnChatRoute(CHAT_ROUTE_HANDOVER, options); + return result.ok ? { ok: true as const, output: result.output } : result; }, }; @@ -1927,9 +1811,8 @@ const handoverInput = { * For the common case prefer `accumulator.consumeHandover()`, which also seeds * `payload.headStartMessages` and applies the partial for you. * - * Must be called at turn 0 before any `chat.messages.waitWithIdleTimeout` — - * that facade consumes and discards non-message chunks, which would swallow the - * handover signal. + * Safe to call at any point in turn 0: the handover signal has its own route, + * so a message facade waiting at the same time cannot take it. */ async function waitForHandover(options: { /** The run's wire payload (only `trigger` / `idleTimeoutInSeconds` are read). */ @@ -1950,161 +1833,183 @@ async function waitForHandover(options: { if (!result.ok) return null; return result.output; } finally { - // The handover window is over either way. A signal arriving after this - // point has no consumer, so hand it to the drain rather than letting it - // park at the head of the channel. - releaseChatInputKinds(CHAT_HANDOVER_KINDS); + chatInputRouter().clearRoute(CHAT_ROUTE_HANDOVER); } } /** - * Record kinds on `session.in` that some consumer on THIS boot is responsible - * for. `"message"` is always claimed; handover kinds are claimed only for the - * window in which `waitForHandover` is actually waiting for them. + * Everything `session.in` carries, and what happens to each kind. + * + * A route's two properties are what make the resume protocol derivable rather + * than hand-maintained. `messages` is replayable because losing a user message + * is data loss. `stop` is neither queued nor replayable: it only means anything + * to the turn that is live when it lands, and a replayed one would abort + * whichever turn happened to be running. `handover` is queued but not + * replayable, because it can arrive before its consumer is ready yet is + * meaningless to any later boot. * - * Anything not in this set has no consumer on this boot, so leaving it buffered - * would park it at the head of the channel forever — `chat.messages.next()` and - * `hasPending()` only ever inspect the head, so every record queued behind it - * becomes undeliverable with no error. The drain below discards unclaimed kinds - * instead. + * A kind absent from this table has no consumer, so the router discards it + * instead of letting it park at the head of a queue. * @internal */ -const chatClaimedKindsKey = locals.create>("chat.claimedKinds"); +const CHAT_INPUT_ROUTES: SessionRouteTable = { + kindOf: (data) => (data as ChatInputChunk | undefined)?.kind, + routes: [ + { name: "messages", delivery: "queue", replayable: true, kinds: ["message"] }, + { name: "stop", delivery: "at-arrival", replayable: false, kinds: ["stop"] }, + { + name: "handover", + delivery: "queue", + replayable: false, + kinds: ["handover", "handover-skip"], + }, + ], +}; -/** Kinds carried on `.in` that are not user messages. @internal */ -const CHAT_HANDOVER_KINDS = ["handover", "handover-skip"] as const; +const CHAT_ROUTE_MESSAGES = "messages"; +const CHAT_ROUTE_STOP = "stop"; +const CHAT_ROUTE_HANDOVER = "handover"; /** - * Every `ChatInputChunk` kind this SDK version knows about. A known kind with - * no active consumer on this boot is an expected, documented state; an unknown - * one means a newer server is sending something this worker cannot handle, - * which is worth surfacing. + * The `.in` router for the chat this worker is serving. + * + * One slot rather than a map: a worker process serves one chat, and the facades + * have to reach the same router the boot attached without depending on a locals + * scope being active. Tagged with its chat so a nested `chat.createSession` for + * the same chat reuses the attached router while a different chat gets a fresh + * one. * @internal */ -const KNOWN_CHAT_INPUT_KINDS: ReadonlySet = new Set([ - "message", - "stop", - ...CHAT_HANDOVER_KINDS, -]); +let currentChatInputRouter: + | { chatId: string; router: SessionChannelRouter; attached: boolean } + | undefined; -/** The run's attached drain subscription, so it can be re-offered the buffer. @internal */ -const chatInputDrainKey = locals.create<{ off: () => void }>("chat.inputDrain"); - -function chatClaimedKinds(): Set { - let claimed = locals.get(chatClaimedKindsKey); - if (!claimed) { - claimed = new Set(["message"]); - locals.set(chatClaimedKindsKey, claimed); +/** + * Both cursors from the latest `turn-complete` on `.out`, in one scan. + * + * Absent for a chat whose turns predate the headers, in which case the router + * starts from the beginning of the channel and nothing is treated as replayed. + * @internal + */ +async function findLatestSessionInCheckpoint(chatId: string): Promise { + const apiClient = apiClientManager.clientOrThrow(); + const response = await apiClient.readSessionStreamRecords(chatId, "out"); + const checkpoint: RouterCheckpoint = {}; + for (const record of response.records) { + if (controlSubtype(record.headers) !== TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE) continue; + const resumeFrom = Number.parseInt( + headerValue(record.headers, SESSION_IN_EVENT_ID_HEADER) ?? "", + 10 + ); + if (Number.isFinite(resumeFrom)) checkpoint.resumeFrom = resumeFrom; + const appliedThrough = Number.parseInt( + headerValue(record.headers, SESSION_IN_CONSUMED_ID_HEADER) ?? "", + 10 + ); + if (Number.isFinite(appliedThrough)) checkpoint.appliedThrough = appliedThrough; } - return claimed; + return checkpoint; } /** - * Attach the unclaimed-control drain for this run. - * - * Consuming at dispatch (returning `true`) is what makes this safe for the - * resume cursor: the record is never buffered, so it leaves no unconsumed - * marker and `lastDispatchedSeqNum()` stays exact rather than being clamped - * behind a record nobody will ever take. + * Attach the `.in` router for this run. * - * `#dispatch` resolves a matching `once()` waiter BEFORE invoking handlers, so - * this can never take a record out from under a claimed consumer that is - * actively waiting for it. + * Reads the checkpoint and subscribes in one call, so there is no window in + * which a listener is attached before the resume cursor is seeded. Attaching + * first would open the tail at sequence 0 and replay every record the previous + * run already answered, which is a mistake the previous shape of this code made + * possible and this shape does not. * - * MUST be attached after `seedSessionInResumeCursorForCustomLoop`, like every - * other `.in` listener — attaching first would replay from seq 0. + * The router consumes every record at dispatch, so the channel's own buffer + * stays empty for chat and its cursor bookkeeping never engages. Delivery, + * ordering and the published cursors are the router's, entirely. * @internal */ -function attachUnclaimedChatInputDrain(): { off: () => void } { - return getChatSession().in.on((chunk) => { - const kind = (chunk as { kind?: unknown } | undefined)?.kind; - // Malformed record: nothing can consume it, so don't let it wedge the head. - if (typeof kind !== "string") { - logger.warn("chat: discarded a malformed session.in record with no usable kind"); - return true; - } - if (chatClaimedKinds().has(kind)) return undefined; - if (!KNOWN_CHAT_INPUT_KINDS.has(kind)) { - logger.warn("chat: discarded a session.in record of an unrecognised kind", { kind }); - } +async function installChatInputRouter( + chatId: string, + options?: { fallbackResumeFrom?: number } +): Promise { + const entry = chatInputRouterEntry(chatId); + if (entry.attached) return entry.router; + + let checkpoint: RouterCheckpoint = {}; + try { + checkpoint = await findLatestSessionInCheckpoint(chatId); + } catch (error) { + logger.warn("chat session: session.in resume cursor lookup failed; old messages may replay", { + error: error instanceof Error ? error.message : String(error), + }); + } + if (checkpoint.resumeFrom === undefined && options?.fallbackResumeFrom !== undefined) { + checkpoint.resumeFrom = options.fallbackResumeFrom; + } + + const router = entry.router; + router.restore(checkpoint); + + const floor = router.resumeFrom(); + if (floor !== undefined) { + sessionStreams.setLastSeqNum(chatId, "in", floor); + sessionStreams.setLastDispatchedSeqNum(chatId, "in", floor); + } + + sessionStreams.onRecord(chatId, "in", (record) => { + router.ingest(record); return true; }); + + entry.attached = true; + return router; } -/** - * Release claimed kinds and re-offer the buffer to the drain. - * - * Re-attaching is the sweep: `on()` re-offers every buffered record to the - * newly attached handler, so a record that was buffered while its kind was - * still claimed (the waiter-gap window between `once()` iterations) is - * discarded now and the cursor advances past it. - * @internal - */ -function releaseChatInputKinds(kinds: readonly string[]): void { - const claimed = chatClaimedKinds(); - let changed = false; - for (const kind of kinds) { - if (claimed.delete(kind)) changed = true; - } - if (!changed) return; +function chatInputRouterEntry(chatId: string): { + chatId: string; + router: SessionChannelRouter; + attached: boolean; +} { + if (currentChatInputRouter?.chatId === chatId) return currentChatInputRouter; + + currentChatInputRouter = { + chatId, + router: new SessionChannelRouter(CHAT_INPUT_ROUTES, { + onDrop: (record, reason) => { + if (reason === "unroutable" || reason === "malformed") { + logger.warn("chat: discarded a session.in record nothing on this worker can consume", { + reason, + seqNum: record.seqNum, + }); + } + }, + }), + attached: false, + }; + return currentChatInputRouter; +} - const drain = locals.get(chatInputDrainKey); - if (!drain) return; - drain.off(); - locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); +/** Drop the router so the next boot attaches a fresh one. @internal */ +export function __resetChatInputRouterForTests(): void { + currentChatInputRouter = undefined; } -/** - * Narrow what holds the persisted `.in` cursor back. Sets no listener, so it is - * safe to call before the resume cursor is seeded. - * @internal - */ -function setChatCursorBarrier(chatId: string): void { - sessionStreams.setCursorBarrier(chatId, "in", isChatCursorBarrier); +/** The cursors this run would publish on its next turn boundary. @internal */ +export function __chatInputCheckpointForTests(): RouterCheckpoint { + return currentChatInputRouter?.router.checkpoint() ?? {}; } -/** - * Stop already-handled control records from being applied a second time after a - * resume. - * - * The cursor is deliberately held back behind messages still waiting to be - * handled, so resuming from it re-delivers everything after that point. For a - * message that is the whole purpose. For a control record it is a fault: a stop - * carries no record of which turn it belonged to, so on redelivery it would - * abort whichever turn happens to be live, which is usually the turn answering - * the very message the cursor was held back to protect. - * - * `consumedThrough` is the highest sequence a previous run reported consuming. - * Control records at or below it have already been applied and are dropped - * before any consumer sees them. Messages are never dropped. Absent for a chat - * whose turns predate the header, in which case nothing is dropped. - * @internal - */ -function setChatReplayGuard(chatId: string, consumedThrough: number | undefined): void { - if (consumedThrough === undefined) { - sessionStreams.setDropPredicate(chatId, "in", undefined); - return; - } - sessionStreams.setDropPredicate(chatId, "in", (record) => { - if (record.seqNum > consumedThrough) return false; - return !isChatCursorBarrier(record); - }); +/** Test-only entry point for the turn-boundary cursor scan. @internal */ +export async function __findLatestSessionInCheckpointForTests( + chatId: string +): Promise { + return findLatestSessionInCheckpoint(chatId); } /** - * Claim the kinds this boot has a consumer for and drain the rest. - * - * Attaches a `.in` listener, so it MUST run after the resume cursor is seeded; - * attaching first makes the subscribe open at seq 0 and replay every record the - * previous run already answered. + * This chat's router. Created on first use so a facade reached before the + * install still shares the one the install will attach. * @internal */ -function attachChatInputDrain(payload: { trigger?: string }): void { - const claimed = chatClaimedKinds(); - if (payload.trigger === "handover-prepare") { - for (const kind of CHAT_HANDOVER_KINDS) claimed.add(kind); - } - locals.set(chatInputDrainKey, attachUnclaimedChatInputDrain()); +function chatInputRouter(): SessionChannelRouter { + return chatInputRouterEntry(getChatSession().id).router; } /** @@ -5634,14 +5539,8 @@ function chatCustomAgent< locals.set(lastTurnCompleteSeqNumKey, { value: undefined }); markChatAgentRunForStreamsWarning(); taskContext.setConversationId(payload.chatId); - setChatCursorBarrier(payload.chatId); stampConversationIdOnActiveSpan(payload.chatId); - // Seed the `.in` resume cursor before user code attaches any `.in` - // listener — otherwise a continuation boot replays already-answered - // messages into the loop's first wait. - await seedSessionInResumeCursorForCustomLoop(payload); - setChatReplayGuard(payload.chatId, await findLatestSessionInConsumedSafe(payload.chatId)); - attachChatInputDrain(payload); + await installChatInputRouter(payload.chatId); return userRun(payload, runOptions); }, }); @@ -5748,7 +5647,6 @@ function chatAgent< locals.set(lastTurnCompleteSeqNumKey, { value: undefined }); markChatAgentRunForStreamsWarning(); taskContext.setConversationId(payload.chatId); - setChatCursorBarrier(payload.chatId); // Stamp `gen_ai.conversation.id` on the run-level span. Every // nested span inherits the same attribute via @@ -5987,56 +5885,15 @@ function chatAgent< ); } - // ── session.in resume cursor ─────────────────────────────────── + // ── session.in router ────────────────────────────────────────── // - // A fresh worker subscribes to `session.in` from seq 0 and would - // re-deliver every record ever appended — including user messages - // from turns already completed on a prior run. Without a cursor, - // the loop would re-process them as fresh turns and the slim-wire - // merge would replace-by-id against snapshot-restored copies, - // yielding no-op replaces while the customer's actual new message - // waits in the queue. - // - // The cursor is the seq_num of the last `.in` record the prior - // worker committed to processing, persisted on each `turn-complete` - // control record as a `session-in-event-id` sibling header. The - // boot scan reads the header off `.out`'s latest turn-complete and - // seeds the manager so the upcoming `.in` SSE subscribe opens with - // `Last-Event-ID: ` — S2 starts after that seq and old - // messages never reach this worker. - // - // Applies in three cases (any of which means `.in` has records - // belonging to completed turns the new run should skip): - // - OOM retry (`ctx.attempt.number > 1`) - // - Continuation run (`payload.continuation === true`) — prior run - // crashed / was canceled / requested upgrade - // - Snapshot exists at all (catches edge cases where the wire - // didn't set `continuation` but a snapshot indicates prior turns) - const needsResumeCursor = - ctx.attempt.number > 1 || payload.continuation === true || bootSnapshot !== undefined; - - if (needsResumeCursor) { - try { - // Reuse the cursor the boot block already resolved (snapshot - // field or records scan) — only scan here when the boot block - // was skipped (hydrateMessages, or snapshot-only signals). - const cursor = bootInCursorResolved - ? bootInCursor - : await findLatestSessionInCursor(payload.chatId); - if (cursor !== undefined) { - sessionStreams.setLastSeqNum(payload.chatId, "in", cursor); - sessionStreams.setLastDispatchedSeqNum(payload.chatId, "in", cursor); - } - } catch (error) { - logger.warn( - "chat.agent: session.in resume cursor lookup failed; old messages may replay", - { error: error instanceof Error ? error.message : String(error) } - ); - } - } - - setChatReplayGuard(payload.chatId, await findLatestSessionInConsumedSafe(payload.chatId)); - attachChatInputDrain(payload); + // Reads the turn boundary and subscribes in one call. `bootInCursor` is + // only a fallback: the boot block above may already have resolved a + // cursor from the snapshot, which is used when the boundary itself + // carries none. + await installChatInputRouter(payload.chatId, { + fallbackResumeFrom: bootInCursorResolved ? bootInCursor : undefined, + }); // ── Recovery boot + chain reconstruction ──────────────────────── if (!hydrateMessages) { @@ -7996,7 +7853,7 @@ function chatAgent< await tracer.startActiveSpan( "snapshot.write", async () => { - const snapshotInCursor = getChatSession().in.lastDispatchedSeqNum(); + const snapshotInCursor = chatInputRouter().resumeFloor(); await writeChatSnapshot(sessionIdForSnapshot, { version: 1, savedAt: Date.now(), @@ -8329,7 +8186,7 @@ function chatAgent< // neither the snapshot nor the replayable `.in` tail. if (!hydrateMessages) { try { - const errorSnapshotInCursor = getChatSession().in.lastDispatchedSeqNum(); + const errorSnapshotInCursor = chatInputRouter().resumeFloor(); await writeChatSnapshot(sessionIdForSnapshot, { version: 1, savedAt: Date.now(), @@ -9222,7 +9079,7 @@ async function chatWriteTurnComplete(options?: { const result = await writeTurnCompleteChunk(undefined, options?.publicAccessToken); // Same cursor written to the `session-in-event-id` header inside // `writeTurnCompleteChunk`; surfaced here so the caller can persist it. - const inCursor = getChatSession().in.lastDispatchedSeqNum(); + const inCursor = chatInputRouter().resumeFloor(); return { lastEventId: result?.lastEventId, ...(inCursor !== undefined ? { sessionInEventId: String(inCursor) } : {}), @@ -9866,7 +9723,7 @@ function createChatSession( activeMsgSub = undefined; if (!booted) { booted = true; - await seedSessionInResumeCursorForCustomLoop(currentPayload); + await installChatInputRouter(currentPayload.chatId); stop = createStopSignal(); } turn++; @@ -11081,7 +10938,7 @@ async function writeTurnCompleteChunk( // function, including the managed agent, which does not call the public // `chat.writeTurnComplete`. By the time a turn completes the handover window // is over either way. - releaseChatInputKinds(CHAT_HANDOVER_KINDS); + chatInputRouter().clearRoute(CHAT_ROUTE_HANDOVER); // 1. Write the turn-complete control record. The ack's `lastEventId` is // this record's seq_num — that's the trim target for the NEXT turn. @@ -11097,11 +10954,12 @@ async function writeTurnCompleteChunk( if (publicAccessToken) { extraHeaders.push(["public-access-token", publicAccessToken]); } - const inCursor = session.in.lastDispatchedSeqNum(); + const routerCheckpoint = chatInputRouter().checkpoint(); + const inCursor = routerCheckpoint.resumeFrom; if (inCursor !== undefined) { extraHeaders.push([SESSION_IN_EVENT_ID_HEADER, String(inCursor)]); } - const consumedCursor = sessionStreams.highestConsumedSeqNum(session.id, "in"); + const consumedCursor = routerCheckpoint.appliedThrough; if (consumedCursor !== undefined) { extraHeaders.push([SESSION_IN_CONSUMED_ID_HEADER, String(consumedCursor)]); } diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index 86ac0389c81..9758d534f21 100644 --- a/packages/trigger-sdk/src/v3/sessions.ts +++ b/packages/trigger-sdk/src/v3/sessions.ts @@ -702,71 +702,93 @@ export class SessionInputChannel { * run-engine waitpoint holds the run until the session append handler * fires it. Only callable from inside `task.run()`. */ + /** + * Suspend until the channel wakes this run, and read nothing. + * + * The waitpoint is only a wake signal: the append route commits the record to + * the channel before it drains any waitpoint, so once this resolves the + * record is durably readable from the channel itself, carrying its real + * sequence. Separating the wake from the read is what lets a consumer that + * owns its own delivery (the chat input router) reuse this without the + * channel also taking a record out from under it. + * + * @internal + */ + async awaitWake( + options?: InputStreamWaitOptions & { lastSeqNum?: number } + ): Promise<{ ok: true; waitpointId: string } | { ok: false; error: Error }> { + const ctx = taskContext.ctx; + + if (!ctx) { + throw new Error("session.in.wait() can only be used from inside a task.run()"); + } + + const apiClient = apiClientManager.clientOrThrow(); + + const lastConsumedSeqNum = + options?.lastSeqNum ?? sessionStreams.lastDispatchedSeqNum(this.sessionId, "in"); + const response = await apiClient.createSessionStreamWaitpoint(ctx.run.id, { + session: this.sessionId, + io: "in", + timeout: options?.timeout, + idempotencyKey: options?.idempotencyKey, + idempotencyKeyTTL: options?.idempotencyKeyTTL, + tags: options?.tags, + lastSeqNum: lastConsumedSeqNum, + }); + + const waitResponse = await apiClient.waitForWaitpointToken({ + runFriendlyId: ctx.run.id, + waitpointFriendlyId: response.waitpointId, + }); + + if (!waitResponse.success) { + throw new Error("Failed to block on session stream waitpoint"); + } + + sessionStreams.disconnectStream(this.sessionId, "in"); + + const waitResult = await runtime.waitUntil(response.waitpointId); + + if (!waitResult.ok) { + const parsed = + waitResult.output !== undefined + ? await conditionallyImportAndParsePacket( + { + data: waitResult.output, + dataType: waitResult.outputType ?? "application/json", + }, + apiClient + ) + : undefined; + return { + ok: false as const, + error: new WaitpointTimeoutError(parsed?.message ?? "Timed out"), + }; + } + + sessionStreams.reconnectStream(this.sessionId, "in"); + return { ok: true as const, waitpointId: response.waitpointId }; + } + wait(options?: InputStreamWaitOptions): ManualWaitpointPromise { return new ManualWaitpointPromise(async (resolve, reject) => { try { - const ctx = taskContext.ctx; - - if (!ctx) { - throw new Error("session.in.wait() can only be used from inside a task.run()"); - } - const apiClient = apiClientManager.clientOrThrow(); - const lastConsumedSeqNum = sessionStreams.lastDispatchedSeqNum(this.sessionId, "in"); - const response = await apiClient.createSessionStreamWaitpoint(ctx.run.id, { - session: this.sessionId, - io: "in", - timeout: options?.timeout, - idempotencyKey: options?.idempotencyKey, - idempotencyKeyTTL: options?.idempotencyKeyTTL, - tags: options?.tags, - lastSeqNum: lastConsumedSeqNum, - }); - const result = await tracer.startActiveSpan( options?.spanName ?? `sessions.open(${this.sessionId}).in.wait()`, async (span) => { - const waitResponse = await apiClient.waitForWaitpointToken({ - runFriendlyId: ctx.run.id, - waitpointFriendlyId: response.waitpointId, - }); + const wake = await this.awaitWake(options); - if (!waitResponse.success) { - throw new Error("Failed to block on session stream waitpoint"); - } - - // Stop the SSE tail before suspending. Buffered records stay in - // place so nothing is lost across the suspend. - sessionStreams.disconnectStream(this.sessionId, "in"); - - const waitResult = await runtime.waitUntil(response.waitpointId); - - if (!waitResult.ok) { - const parsed = - waitResult.output !== undefined - ? await conditionallyImportAndParsePacket( - { - data: waitResult.output, - dataType: waitResult.outputType ?? "application/json", - }, - apiClient - ) - : undefined; - const error = new WaitpointTimeoutError(parsed?.message ?? "Timed out"); - span.recordException(error); + if (!wake.ok) { + span.recordException(wake.error); span.setStatus({ code: SpanStatusCode.ERROR }); - return { ok: false as const, error }; + return { ok: false as const, error: wake.error }; } - // The waitpoint is only a wake signal. The append route commits the - // record to the channel before it drains any waitpoint, so by the - // time we are here the record is durably readable from the channel - // itself, carrying its real sequence. Reading it back that way is - // what keeps the cursor exact: the waitpoint payload cannot - // identify which record it corresponds to, and guessing or matching - // on payload equality both produce a cursor that strands or - // redelivers records. + span.setAttribute(SemanticInternalAttributes.ENTITY_ID, wake.waitpointId); + const record = await sessionStreams.onceRecord(this.sessionId, "in"); if (!record.ok) { @@ -795,7 +817,6 @@ export class SessionInputChannel { attributes: { [SemanticInternalAttributes.STYLE_ICON]: "wait", [SemanticInternalAttributes.ENTITY_TYPE]: "waitpoint", - [SemanticInternalAttributes.ENTITY_ID]: response.waitpointId, session: this.sessionId, io: "in", ...accessoryAttributes({ diff --git a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts index a669975f383..63768b9b3f2 100644 --- a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts +++ b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts @@ -4,6 +4,7 @@ import type { LocalsKey } from "@trigger.dev/core/v3"; import { runInMockTaskContext, type MockTaskContextOptions } from "@trigger.dev/core/v3/test"; import { __setSessionOpenImplForTests, __setSessionStartImplForTests } from "../sessions.js"; import { + __resetChatInputRouterForTests, __setReadChatSnapshotImplForTests, __setReplaySessionInTailImplForTests, __setReplaySessionOutTailImplForTests, @@ -390,6 +391,8 @@ export function mockChatAgent( let seededReplayPartial: UIMessage | undefined; let seededSessionInMessages: UIMessage[] = []; + __resetChatInputRouterForTests(); + __setReadChatSnapshotImplForTests((_id: string) => { return seededSnapshot as ChatSnapshotV1 | undefined; }); diff --git a/packages/trigger-sdk/src/v3/test/test-session-handle.ts b/packages/trigger-sdk/src/v3/test/test-session-handle.ts index 945cd231152..5150de87da3 100644 --- a/packages/trigger-sdk/src/v3/test/test-session-handle.ts +++ b/packages/trigger-sdk/src/v3/test/test-session-handle.ts @@ -4,7 +4,7 @@ import type { StreamWriteResult, WriterStreamOptions, } from "@trigger.dev/core/v3"; -import { ensureReadableStream, ManualWaitpointPromise } from "@trigger.dev/core/v3"; +import { ensureReadableStream } from "@trigger.dev/core/v3"; import type { SessionPipeStreamOptions, SessionSubscribeOptions } from "../sessions.js"; import { SessionHandle, SessionInputChannel, SessionOutputChannel } from "../sessions.js"; @@ -29,32 +29,27 @@ class TestSessionInputChannel extends SessionInputChannel { super(sessionId); } - // Override only the `wait` path. `on` / `once` / `peek` / `send` - // continue to flow through the real `sessionStreams` global, which - // the mock task context installs as a `TestSessionStreamManager`. - wait(): ManualWaitpointPromise { - return new ManualWaitpointPromise( - (resolve: (value: { ok: false; error: Error }) => void) => { - const signal = this.getAbortSignal(); - if (!signal) { - // Harness hasn't wired up its run signal yet — nothing to abort - // on. Stay pending; the run loop should never reach this state - // in practice but we don't want to throw here either. - return; - } - const onAbort = () => { - resolve({ - ok: false, - error: new Error("session.in.wait() aborted by test harness"), - }); - }; - if (signal.aborted) { - onAbort(); - return; - } - signal.addEventListener("abort", onAbort, { once: true }); - } - ); + /** + * Override the one step that talks to the network. Everything built on top + * of it (`wait`, and the chat facades' route waits) then runs its real + * implementation against the in-memory stream manager, so the harness stubs + * a boundary instead of reimplementing a composite. + */ + async awaitWake(): Promise<{ ok: true; waitpointId: string } | { ok: false; error: Error }> { + const signal = this.getAbortSignal(); + if (!signal) { + return new Promise(() => {}); + } + if (signal.aborted) { + return { ok: false, error: new Error("session.in.wait() aborted by test harness") }; + } + return new Promise((resolve) => { + signal.addEventListener( + "abort", + () => resolve({ ok: false, error: new Error("session.in.wait() aborted by test harness") }), + { once: true } + ); + }); } } diff --git a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts index 10dafc47f0c..cb94ae87f3c 100644 --- a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -5,7 +5,12 @@ import "../src/v3/test/index.js"; import { resourceCatalog, sessionStreams } from "@trigger.dev/core/v3"; import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; import { describe, expect, it } from "vitest"; -import { chat, type ChatMessageRecord, type ChatTaskWirePayload } from "../src/v3/ai.js"; +import { + __chatInputCheckpointForTests as chatInputCheckpoint, + chat, + type ChatMessageRecord, + type ChatTaskWirePayload, +} from "../src/v3/ai.js"; function deferred() { let resolve!: () => void; @@ -52,10 +57,10 @@ describe("chat.messages mailbox", () => { observations.before = await chat.messages.hasPending(); observations.first = await chat.messages.next(); - observations.cursorAfterFirst = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.cursorAfterFirst = chatInputCheckpoint().resumeFrom; observations.afterFirst = await chat.messages.hasPending(); observations.second = await chat.messages.next(); - observations.cursorAfterSecond = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.cursorAfterSecond = chatInputCheckpoint().resumeFrom; observations.afterSecond = await chat.messages.hasPending(); }, }); @@ -118,19 +123,15 @@ describe("chat.messages mailbox", () => { expect(result).toBeUndefined(); }); - it("leaves earlier non-message records for their own consumer", async () => { + it("delivers a message that arrived behind another kind, without losing that kind", async () => { const chatId = "mailbox-mixed-kinds"; const ready = deferred(); const inspect = deferred(); const observations: { pending?: boolean; - blocked?: ChatMessageRecord; - cursorAfterBlocked?: number; - headAfterBlocked?: unknown; - control?: unknown; - pendingAfterControl?: boolean; message?: ChatMessageRecord; - cursorAfterMessage?: number; + handover?: unknown; + cursorAfter?: number; } = {}; const agent = chat.customAgent({ @@ -140,15 +141,12 @@ describe("chat.messages mailbox", () => { await inspect.promise; observations.pending = await chat.messages.hasPending(); - observations.blocked = await chat.messages.next({ timeoutInSeconds: 0 }); - observations.cursorAfterBlocked = sessionStreams.lastDispatchedSeqNum(chatId, "in"); - observations.headAfterBlocked = sessionStreams.peekRecord(chatId, "in"); - - const control = await sessionStreams.onceRecord(chatId, "in"); - observations.control = control.ok ? control.output : undefined; - observations.pendingAfterControl = await chat.messages.hasPending(); observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); - observations.cursorAfterMessage = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.handover = await chat.waitForHandover({ + payload: { trigger: "handover-prepare" }, + idleTimeoutInSeconds: 0, + }); + observations.cursorAfter = chatInputCheckpoint().resumeFrom; }, }); const run = resourceCatalog.getTask(agent.id)?.fns.run; @@ -177,38 +175,30 @@ describe("chat.messages mailbox", () => { await runPromise; }); + // The handover has its own route, so it neither blocks the message behind + // it nor gets destroyed by the consumer that took that message. expect(observations).toEqual({ - pending: false, - blocked: undefined, - cursorAfterBlocked: undefined, - headAfterBlocked: { - id: "handover-1", - seqNum: 30, - data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, - }, - control: { - id: "handover-1", - seqNum: 30, - data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, - }, - pendingAfterControl: true, + pending: true, message: { id: "message-1", seqNum: 31, payload: userPayload(chatId, "u-after-handover"), }, - cursorAfterMessage: 31, + handover: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + cursorAfter: 31, }); }); - it("keeps the cursor behind a buffered message when a later stop is consumed", async () => { + it("holds the resume cursor behind a queued message while a later stop advances the replay window", async () => { const chatId = "mailbox-cursor-gap"; const ready = deferred(); const inspect = deferred(); const observations: { cursorBefore?: number; + appliedBefore?: number; message?: ChatMessageRecord; cursorAfter?: number; + appliedAfter?: number; } = {}; const agent = chat.customAgent({ @@ -218,9 +208,11 @@ describe("chat.messages mailbox", () => { ready.resolve(); await inspect.promise; - observations.cursorBefore = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.cursorBefore = chatInputCheckpoint().resumeFrom; + observations.appliedBefore = chatInputCheckpoint().appliedThrough; observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); - observations.cursorAfter = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.cursorAfter = chatInputCheckpoint().resumeFrom; + observations.appliedAfter = chatInputCheckpoint().appliedThrough; stop.cleanup(); }, }); @@ -249,13 +241,16 @@ describe("chat.messages mailbox", () => { }); expect(observations).toEqual({ + // Held below the queued message even though the stop after it was applied. cursorBefore: 49, + appliedBefore: 51, message: { id: "message-1", seqNum: 50, payload: userPayload(chatId, "u1"), }, cursorAfter: 51, + appliedAfter: 51, }); }); diff --git a/packages/trigger-sdk/test/replay-session-in.test.ts b/packages/trigger-sdk/test/replay-session-in.test.ts index 3d04974a9b6..35a06271374 100644 --- a/packages/trigger-sdk/test/replay-session-in.test.ts +++ b/packages/trigger-sdk/test/replay-session-in.test.ts @@ -4,7 +4,7 @@ import "../src/v3/test/index.js"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { apiClientManager } from "@trigger.dev/core/v3"; import { - __findLatestSessionInCursorForTests as findLatestSessionInCursor, + __findLatestSessionInCheckpointForTests as findLatestSessionInCheckpoint, __replaySessionInTailProductionPathForTests as replaySessionInTail, } from "../src/v3/ai.js"; @@ -167,8 +167,8 @@ function stubReadRecordsWithHeaders( return spy; } -describe("findLatestSessionInCursor", () => { - it("returns the LAST turn-complete's session-in-event-id", async () => { +describe("findLatestSessionInCheckpoint", () => { + it("returns the LAST turn-complete's cursors", async () => { const spy = stubReadRecordsWithHeaders([ { data: { type: "text-delta", delta: "hi" } }, { @@ -182,13 +182,15 @@ describe("findLatestSessionInCursor", () => { headers: [ ["trigger-control", "turn-complete"], ["session-in-event-id", "7"], + ["session-in-consumed-id", "9"], ], }, ]); - const cursor = await findLatestSessionInCursor("sess"); - expect(cursor).toBe(7); - // Non-blocking records read on `.out`, no SSE subscribe. + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint).toEqual({ resumeFrom: 7, appliedThrough: 9 }); + // One non-blocking records read on `.out` covers both cursors. + expect(spy).toHaveBeenCalledTimes(1); expect(spy).toHaveBeenCalledWith("sess", "out"); }); @@ -209,8 +211,22 @@ describe("findLatestSessionInCursor", () => { }, ]); - const cursor = await findLatestSessionInCursor("sess"); - expect(cursor).toBe(4); + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint.resumeFrom).toBe(4); + }); + + it("leaves the replay window absent when only the resume cursor was written", async () => { + stubReadRecordsWithHeaders([ + { + headers: [ + ["trigger-control", "turn-complete"], + ["session-in-event-id", "5"], + ], + }, + ]); + + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint).toEqual({ resumeFrom: 5 }); }); it("returns undefined when records carry no headers (older server)", async () => { @@ -219,7 +235,7 @@ describe("findLatestSessionInCursor", () => { { data: { type: "finish" } }, ]); - const cursor = await findLatestSessionInCursor("sess"); - expect(cursor).toBeUndefined(); + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint).toEqual({}); }); }); From 027eaf9e8dc9864dfebff54d37794f5c8defef3d Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sat, 22 Aug 2026 08:57:03 +0100 Subject: [PATCH 15/20] fix(chat,sdk): close the stale-stop gap for chats upgrading from an older SDK A turn boundary written before the replay window was published carries only the resume cursor. Resuming one, the run had no way to tell a control record it was re-reading from one arriving live, so a stop that a previous run had already applied was applied a second time, aborting the turn answering the very message the cursor was held back to protect. Queued messages were never at risk. Boundaries that predate the window now resolve it from the channel: everything already on `.in` at boot is by definition not arriving live on this run, so it is the end of that run's replay window. Bounded by `afterEventId`, so the read covers the replay window rather than the whole conversation, and only on the first turn after an upgrade. The trade is deliberate. A stop that landed legitimately in the moments before boot is now dropped along with the replayed ones. Missing a stop leaves the user able to press stop again; applying a stale one kills an answer they are waiting for. Verified on a deployed test-cloud project across two versions: a session driven to a crash on a build that writes the old boundary format, then continued on this build. Before, the continuation reported the stop aborting its turn; after, it declines the replayed stop and the turn answering the recovered message completes. --- .../core/src/v3/sessionStreams/router.test.ts | 20 +++++++++- packages/core/src/v3/sessionStreams/router.ts | 11 +++--- packages/trigger-sdk/src/v3/ai.ts | 38 +++++++++++++++++++ 3 files changed, 62 insertions(+), 7 deletions(-) diff --git a/packages/core/src/v3/sessionStreams/router.test.ts b/packages/core/src/v3/sessionStreams/router.test.ts index 59006bc9bf4..6f654269d1c 100644 --- a/packages/core/src/v3/sessionStreams/router.test.ts +++ b/packages/core/src/v3/sessionStreams/router.test.ts @@ -257,7 +257,7 @@ describe("SessionChannelRouter: resuming", () => { expect(stops).toEqual([3]); }); - it("treats an absent replay-window end as the floor", () => { + it("falls back to the floor when no replay-window end is supplied", () => { const r = router(); r.restore({ resumeFrom: 4 }); r.on("stop", () => {}); @@ -266,6 +266,24 @@ describe("SessionChannelRouter: resuming", () => { expect(r.ingest(rec(5, "stop"))).toEqual({ action: "deliver", route: "stop" }); }); + it("declines a control record inside a window resolved from the channel", () => { + const r = router(); + // What the chat layer supplies when the boundary predates the published + // window: everything already on the channel at boot counts as replayed. + r.restore({ resumeFrom: 4, appliedThrough: 6 }); + const stops: number[] = []; + r.on("stop", (record) => stops.push(record.seqNum)); + + expect(r.ingest(rec(5, "message", "M5"))).toEqual({ action: "queue", route: "messages" }); + expect(r.ingest(rec(6, "stop"))).toEqual({ + action: "drop", + route: "stop", + reason: "replayed", + }); + expect(r.ingest(rec(7, "stop"))).toEqual({ action: "deliver", route: "stop" }); + expect(stops).toEqual([7]); + }); + it("applies everything on a fresh session with no checkpoint", () => { const r = router(); r.on("stop", () => {}); diff --git a/packages/core/src/v3/sessionStreams/router.ts b/packages/core/src/v3/sessionStreams/router.ts index 0adbe0475e4..17b4cdaddb2 100644 --- a/packages/core/src/v3/sessionStreams/router.ts +++ b/packages/core/src/v3/sessionStreams/router.ts @@ -177,12 +177,11 @@ export class SessionChannelRouter { /** * Seed the router from a previous run's turn boundary. * - * An absent `appliedThrough` is treated as equal to the floor rather than as - * "nothing was applied". A boundary written before that value existed still - * tells us everything at or below the floor was terminal, and for anything - * above it the conservative choice for an `at-arrival` record is to not apply - * it: a missed stop is recoverable, while a stop applied to the wrong turn - * kills a live answer. + * An absent `appliedThrough` falls back to the floor, which leaves anything + * above the floor treated as live. A caller resuming a boundary that predates + * the published replay window should resolve the window from the channel + * instead, so it covers everything already there at boot: a missed stop is + * recoverable, while a stop applied to the wrong turn kills a live answer. */ restore(checkpoint: RouterCheckpoint): void { this.#resumeFrom = checkpoint.resumeFrom; diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 17cce3fce25..41f6da31408 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1911,6 +1911,41 @@ async function findLatestSessionInCheckpoint(chatId: string): Promise { + try { + const apiClient = apiClientManager.clientOrThrow(); + const response = await apiClient.readSessionStreamRecords(chatId, "in", { + afterEventId: String(afterSeqNum), + }); + let highest: number | undefined; + for (const record of response.records) { + const seqNum = typeof record.seqNum === "number" ? record.seqNum : Number.NaN; + if (Number.isFinite(seqNum) && (highest === undefined || seqNum > highest)) { + highest = seqNum; + } + } + return highest; + } catch { + return undefined; + } +} + /** * Attach the `.in` router for this run. * @@ -1943,6 +1978,9 @@ async function installChatInputRouter( if (checkpoint.resumeFrom === undefined && options?.fallbackResumeFrom !== undefined) { checkpoint.resumeFrom = options.fallbackResumeFrom; } + if (checkpoint.resumeFrom !== undefined && checkpoint.appliedThrough === undefined) { + checkpoint.appliedThrough = await findSessionInReplayWindowEnd(chatId, checkpoint.resumeFrom); + } const router = entry.router; router.restore(checkpoint); From e6bd62307cdb27c3ea8b271179b6c05509b5f8c9 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sat, 22 Aug 2026 09:30:58 +0100 Subject: [PATCH 16/20] docs(chat): describe the stop fix and the exact hasPending guarantee in the changeset --- .changeset/tidy-mailboxes-wait.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md index 97710a72b29..c26b63c97ce 100644 --- a/.changeset/tidy-mailboxes-wait.md +++ b/.changeset/tidy-mailboxes-wait.md @@ -5,9 +5,11 @@ Fixes a case where a chat could silently lose a message. If a message arrived while the agent was between turns and a stop arrived after it, the cursor the next boot resumed from could point past that message, so it was never answered and no error was raised. This affected `chat.agent`, not just custom agents. +Fixes a recovered answer being cut off. After a crash the agent replays the message it had not answered yet, but it was replaying the stop that arrived after that message too, so the turn answering it was aborted the moment it began. A stop is now only applied to the turn that was live when it arrived, including for chats whose most recent turn was completed by an older version of the SDK. + Also fixes a retried send being answered twice. When a send was retried and its idempotency claim was lost, the agent could consume the same message a second time. -Custom agent loops can now inspect pending chat input without consuming it, and consume one mailbox record at a time, with `chat.messages.hasPending()` and `chat.messages.next()`. Records carry stable identifiers so a redelivery is recognisable. +Custom agent loops can now inspect pending chat input without consuming it, and consume one record at a time, with `chat.messages.hasPending()` and `chat.messages.next()`. Records carry stable identifiers so a redelivery is recognisable. ```ts if (await chat.messages.hasPending()) { @@ -16,6 +18,6 @@ if (await chat.messages.hasPending()) { } ``` -A control record that nothing on the run consumes is now discarded rather than left at the head of the input channel, where it would have made every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` means no message became consumable before the timeout. +`hasPending()` answers for messages alone, so a message sitting behind a stop, or behind a record this version of the SDK does not recognise, still reports as pending and is still delivered. Anything the agent has no consumer for is discarded rather than left where it would make every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` means no message became consumable before the timeout. `chat.writeTurnComplete()`'s `sessionInEventId` is the cursor that is safe to resume from, not the sequence of the record the turn answered. It is held back behind any message still waiting to be handled, so a value below the record you just handled is expected. From dfed2bc6893d1df67b1ba2b9ade69dd788721529 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sun, 23 Aug 2026 07:35:22 +0100 Subject: [PATCH 17/20] docs(chat): correct the mailbox docs and describe the stop guarantee The mailbox section described head-of-line behaviour that no longer happens: it said a control record arriving before a message keeps `hasPending()` false and makes `next()` wait, when a message behind a stop or a handover is now reported and returned normally. Docs that state the opposite of what ships are worse than none, so those claims are replaced with what the delivery guarantee actually is. Also documents two things that had no coverage. A stop applies only to the turn that was live when it arrived, so a recovered message's turn is not aborted by a stop from before the crash. And `sessionInEventId` is a lower bound rather than the sequence of the record a turn answered, which is the mistake a client makes when it tries to match a turn boundary to its own send. --- docs/ai-chat/custom-agents.mdx | 37 ++++++++++++++++++---------------- docs/ai-chat/reference.mdx | 4 ++-- 2 files changed, 22 insertions(+), 19 deletions(-) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 8c7a33b5dbd..7aa52a5f358 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -179,6 +179,12 @@ for await (const turn of session) { The frontend stops a turn with [`transport.stopGeneration(chatId)`](/ai-chat/frontend#stop-generation), which writes a stop signal to the session's input stream. It aborts the current turn's generation but keeps the run alive, so the next message continues on the same session. +A stop only applies to the turn that was live when it arrived. If the run crashes +and a later run recovers a message that had not been answered yet, a stop that +was already applied before the crash is not applied again, so the turn answering +the recovered message runs to completion. A stop sent after the recovery is live +and aborts that turn as normal. + `turn.signal` is a combined stop-and-cancel `AbortSignal`, fresh each turn. Pass it to `streamText` so the stop reaches the model, then let `turn.complete()` finish the turn: ```ts trigger/my-chat.ts @@ -227,16 +233,16 @@ For full control, skip `createSession` and compose the primitives directly: | Method | Behavior | | --- | --- | -| `peek()` | Return the buffer head when it is a message, without consuming it; otherwise return `undefined` | -| `hasPending()` | Resolve `true` when the buffer head is a message; does not consume it | +| `peek()` | Return the next queued message without consuming it, or `undefined` when none is queued | +| `hasPending()` | Resolve `true` when a message is queued; does not consume it | | `next({ timeoutInSeconds? })` | Consume exactly one message record in channel order, or resolve `undefined` when the optional timeout elapses | | `on(handler)` | Consume messages as they arrive and invoke the handler | | `waitWithIdleTimeout(options)` | Wait warm, then suspend the run until the next message arrives | -`hasPending()` checks whether the local, already-delivered buffer head is a -message that `next()` can consume immediately. It does not query the remote -Session channel or start a subscription. Use `waitWithIdleTimeout()` when the -loop needs to idle until future input arrives. +`hasPending()` checks whether a message has already been delivered locally and is +waiting for `next()` to take it. It does not query the remote Session channel or +start a subscription. Use `waitWithIdleTimeout()` when the loop needs to idle +until future input arrives. `next({ timeoutInSeconds: 0 })` is also a local, non-blocking read. Call `next()` without a timeout, or with a positive timeout, to subscribe for future @@ -262,18 +268,15 @@ owns its own turn sequencing never advances past input it has not taken. By contrast, `on()` commits a record as soon as it dispatches the handler; avoid mixing `on()` and `next()` when a single loop owns mailbox consumption. -The Session `.in` channel also carries control records such as handovers. If one -comes before a message, `hasPending()` stays `false` and `next()` leaves the -control record for its own consumer. After that record is handled, the message -becomes pending. - -A control record that nothing on the run consumes is discarded rather than left -at the head of the channel. `hasPending()` and `next()` only look at the head, so -a record parked there would make every message behind it undeliverable. +The Session `.in` channel also carries control records such as stops and +handovers. Those are routed to their own consumers and never block messages: a +message that arrived behind one is still reported by `hasPending()` and still +returned by `next()`, in channel order. The same holds for a record kind this +version of the SDK does not recognise, which is discarded rather than left where +it would make every message behind it undeliverable. -`next()` still returns `undefined` whenever no message became consumable before -the timeout, including while a control record that does have its own consumer -sits at the head. +`next()` returns `undefined` when no message became consumable before the +timeout. A complete loop: diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index 65e98a05aa4..dd444961af9 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -504,7 +504,7 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`. | `chat.createSession(payload, options)` | Create an async iterator for chat turns | | `chat.pipe(source, options?)` | Pipe a stream to the frontend (from anywhere inside a task) | | `chat.pipeAndCapture(source, options?)` | Pipe and capture the response; returns `{ message, status, error }` | -| `chat.writeTurnComplete(options?)` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | +| `chat.writeTurnComplete(options?)` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors. `sessionInEventId` is a lower bound, not the sequence of the record the turn answered | | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | | `chat.messages` | Incoming message mailbox; supports non-consuming `.peek()` / `.hasPending()`, single-record `.next()`, `.on()`, and suspend-aware `.waitWithIdleTimeout()` | | `chat.local({ id })` | Create a per-run typed local (see [`chat.local`](/ai-chat/chat-local)) | @@ -645,7 +645,7 @@ The `onEvent` callback receives a `ChatTransportEvent` (exported from `@trigger. | `message-send-failed` | `messageId?`, `source`, `error`, `status?`, `durationMs`, `partId?`, `bodyBytes?` | A send definitively failed after internal retries. Fires in addition to `useChat`'s `onError`. | | `stream-connected` | `resumed`, `lastEventId?`, `messageId?` | The SSE subscription to the session's output stream started delivering. `resumed: true` when reconnecting from a stored cursor (page reload) rather than following a fresh send. `lastEventId` is the cursor it connected from. | | `first-chunk` | `chunkType?`, `lastEventId?`, `messageId?`, `sinceSendMs?` | The first response chunk of a turn arrived. `sinceSendMs` is the delta from the last turn-producing send — time to first token without any bookkeeping. | -| `turn-completed` | `lastEventId?`, `sessionInEventId?`, `messageId?`, `sinceSendMs?` | The agent's turn-complete control record arrived — the "finished answering" signal. `sinceSendMs` is the full turn latency; `sessionInEventId` is the agent's committed input-stream cursor. | +| `turn-completed` | `lastEventId?`, `sessionInEventId?`, `messageId?`, `sinceSendMs?` | The agent's turn-complete control record arrived — the "finished answering" signal. `sinceSendMs` is the full turn latency; `sessionInEventId` is the cursor the agent can safely resume its input stream from. Treat it as a lower bound: it is held back behind any message still waiting to be handled, so it can be below the sequence of the record this turn answered. Do not use it to decide whether a turn boundary belongs to your own send. | | `stream-error` | `error`, `status?` | The output stream failed unrecoverably. | `source` identifies the send path: `"submit-message"`, `"regenerate-message"`, `"steer"` (`sendPendingMessage`), `"action"` (`sendAction`), `"stop"` (`stopGeneration`), or `"head-start"`. From 89b5eb1a3ad04147853488fa742c60a2c5d71220 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sun, 23 Aug 2026 08:56:35 +0100 Subject: [PATCH 18/20] fix(chat,sdk): take the replay window from the channel, not just the turn boundary A stop applied by a run that then died was still being applied a second time on recovery whenever no turn boundary recorded it. A boundary is written when a turn ends, so a stop that arrives after the last one is not covered by it, and the previous fix only helped when the boundary carried no window at all. The result was a recovered turn aborted by a stop the user had pressed against a turn that no longer exists. Everything already on the channel when a run boots is, by definition, not arriving live on that run, so the window is now taken from the channel's own tail and the boundary's value is only a floor for it. Queued messages are unaffected, since they are replayable and never dropped. Found by extending the reproductions to the managed `chat.agent` surface, which reaches this case naturally: the managed loop writes its boundary at the end of each turn, so a stop arriving mid-turn lands after it. --- .changeset/tidy-mailboxes-wait.md | 2 +- packages/trigger-sdk/src/v3/ai.ts | 30 +++++++++++++++++++++--------- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md index c26b63c97ce..4205570fef5 100644 --- a/.changeset/tidy-mailboxes-wait.md +++ b/.changeset/tidy-mailboxes-wait.md @@ -5,7 +5,7 @@ Fixes a case where a chat could silently lose a message. If a message arrived while the agent was between turns and a stop arrived after it, the cursor the next boot resumed from could point past that message, so it was never answered and no error was raised. This affected `chat.agent`, not just custom agents. -Fixes a recovered answer being cut off. After a crash the agent replays the message it had not answered yet, but it was replaying the stop that arrived after that message too, so the turn answering it was aborted the moment it began. A stop is now only applied to the turn that was live when it arrived, including for chats whose most recent turn was completed by an older version of the SDK. +Fixes a recovered answer being cut off. After a crash the agent replays the message it had not answered yet, but it was replaying the stop that arrived after that message too, so the turn answering it was aborted the moment it began. A stop is now only applied to the turn that was live when it arrived. That holds however the stop got there: sent after the last completed turn, or sent to a chat whose most recent turn was completed by an older version of the SDK. Also fixes a retried send being answered twice. When a send was retried and its idempotency claim was lost, the agent could consume the same message a second time. diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 41f6da31408..64c0297a8c2 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1914,11 +1914,19 @@ async function findLatestSessionInCheckpoint(chatId: string): Promise { try { const apiClient = apiClientManager.clientOrThrow(); const response = await apiClient.readSessionStreamRecords(chatId, "in", { - afterEventId: String(afterSeqNum), + ...(afterSeqNum === undefined ? {} : { afterEventId: String(afterSeqNum) }), }); let highest: number | undefined; for (const record of response.records) { @@ -1978,8 +1986,12 @@ async function installChatInputRouter( if (checkpoint.resumeFrom === undefined && options?.fallbackResumeFrom !== undefined) { checkpoint.resumeFrom = options.fallbackResumeFrom; } - if (checkpoint.resumeFrom !== undefined && checkpoint.appliedThrough === undefined) { - checkpoint.appliedThrough = await findSessionInReplayWindowEnd(chatId, checkpoint.resumeFrom); + const replayWindowEnd = await findSessionInReplayWindowEnd(chatId, checkpoint.resumeFrom); + if (replayWindowEnd !== undefined) { + checkpoint.appliedThrough = Math.max( + checkpoint.appliedThrough ?? replayWindowEnd, + replayWindowEnd + ); } const router = entry.router; From 4693292ade7a9baa1e3d97e013943a514b989d41 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sun, 23 Aug 2026 10:13:31 +0100 Subject: [PATCH 19/20] fix(chat,sdk): only treat channel records as replayed when the run is resuming The previous commit took the replay window from the channel's tail so a stop that no turn boundary covered could not be applied twice. On a first boot that is wrong: nothing has been applied by anyone yet, so anything already on the channel was treated as replayed and dropped. That broke head starts. The client can signal a handover before the agent run has booted, which is the whole point of the flow, and the signal was discarded as already-applied. The agent then waited out its idle window and lost the warm partial. A replay window now only exists for a run that is resuming, which is either a run whose predecessor left a turn boundary or one the wire marks as a continuation or a retried attempt. A first boot has no window and applies what it finds. --- packages/trigger-sdk/src/v3/ai.ts | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 64c0297a8c2..713413914a1 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1970,7 +1970,7 @@ async function findSessionInReplayWindowEnd( */ async function installChatInputRouter( chatId: string, - options?: { fallbackResumeFrom?: number } + options?: { fallbackResumeFrom?: number; resuming?: boolean } ): Promise { const entry = chatInputRouterEntry(chatId); if (entry.attached) return entry.router; @@ -1986,12 +1986,19 @@ async function installChatInputRouter( if (checkpoint.resumeFrom === undefined && options?.fallbackResumeFrom !== undefined) { checkpoint.resumeFrom = options.fallbackResumeFrom; } - const replayWindowEnd = await findSessionInReplayWindowEnd(chatId, checkpoint.resumeFrom); - if (replayWindowEnd !== undefined) { - checkpoint.appliedThrough = Math.max( - checkpoint.appliedThrough ?? replayWindowEnd, - replayWindowEnd - ); + // Only a resuming run has a replay window. On a first boot nothing has been + // applied by anyone, so treating what is already on the channel as replayed + // would discard a signal that arrived before the agent got here, which is + // exactly how a head-start handover reaches a cold run. + const resuming = checkpoint.resumeFrom !== undefined || options?.resuming === true; + if (resuming) { + const replayWindowEnd = await findSessionInReplayWindowEnd(chatId, checkpoint.resumeFrom); + if (replayWindowEnd !== undefined) { + checkpoint.appliedThrough = Math.max( + checkpoint.appliedThrough ?? replayWindowEnd, + replayWindowEnd + ); + } } const router = entry.router; @@ -5590,7 +5597,9 @@ function chatCustomAgent< markChatAgentRunForStreamsWarning(); taskContext.setConversationId(payload.chatId); stampConversationIdOnActiveSpan(payload.chatId); - await installChatInputRouter(payload.chatId); + await installChatInputRouter(payload.chatId, { + resuming: Boolean(payload.continuation), + }); return userRun(payload, runOptions); }, }); @@ -5943,6 +5952,7 @@ function chatAgent< // carries none. await installChatInputRouter(payload.chatId, { fallbackResumeFrom: bootInCursorResolved ? bootInCursor : undefined, + resuming: Boolean(payload.continuation) || ctx.attempt.number > 1, }); // ── Recovery boot + chain reconstruction ──────────────────────── @@ -9773,7 +9783,9 @@ function createChatSession( activeMsgSub = undefined; if (!booted) { booted = true; - await installChatInputRouter(currentPayload.chatId); + await installChatInputRouter(currentPayload.chatId, { + resuming: Boolean(currentPayload.continuation), + }); stop = createStopSignal(); } turn++; From 471d17d812a8c795aaffa85a45833c7fad768c28 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sun, 23 Aug 2026 12:46:06 +0100 Subject: [PATCH 20/20] fix(chat,sdk): attach the input router per run, not per chat A worker process is reused across runs, and the executor tears the channel subscription down at the end of each one. The router was cached per chat, so a second run of the same chat in the same process skipped the attach and then had no subscription feeding it: no messages arrived, and the conversation hung with no error raised. Caching it per run instead means every run attaches its own subscription, while a nested `chat.createSession` inside one run still shares the router. Reported by review, reproduced first as a failing test that drives two runs through one process with the executor's teardown in between. --- packages/trigger-sdk/src/v3/ai.ts | 24 +++++--- .../test/chat-warm-process-reuse.test.ts | 60 +++++++++++++++++++ 2 files changed, 75 insertions(+), 9 deletions(-) create mode 100644 packages/trigger-sdk/test/chat-warm-process-reuse.test.ts diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 713413914a1..a44dd210cfe 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1871,17 +1871,18 @@ const CHAT_ROUTE_STOP = "stop"; const CHAT_ROUTE_HANDOVER = "handover"; /** - * The `.in` router for the chat this worker is serving. - * - * One slot rather than a map: a worker process serves one chat, and the facades - * have to reach the same router the boot attached without depending on a locals - * scope being active. Tagged with its chat so a nested `chat.createSession` for - * the same chat reuses the attached router while a different chat gets a fresh - * one. + * The `.in` router for the run this worker is currently serving. + * + * One slot rather than a map, because the facades have to reach the same router + * the boot attached without depending on a locals scope being active. Tagged + * with the run as well as the chat: a warm process is reused across runs and the + * executor tears the channel subscription down at the end of each one, so + * reusing a router across runs would leave the new run with no input at all. A + * nested `chat.createSession` within the same run still shares it. * @internal */ let currentChatInputRouter: - | { chatId: string; router: SessionChannelRouter; attached: boolean } + | { chatId: string; runId: string | undefined; router: SessionChannelRouter; attached: boolean } | undefined; /** @@ -2021,13 +2022,18 @@ async function installChatInputRouter( function chatInputRouterEntry(chatId: string): { chatId: string; + runId: string | undefined; router: SessionChannelRouter; attached: boolean; } { - if (currentChatInputRouter?.chatId === chatId) return currentChatInputRouter; + const runId = taskContext.ctx?.run.id; + if (currentChatInputRouter?.chatId === chatId && currentChatInputRouter.runId === runId) { + return currentChatInputRouter; + } currentChatInputRouter = { chatId, + runId, router: new SessionChannelRouter(CHAT_INPUT_ROUTES, { onDrop: (record, reason) => { if (reason === "unroutable" || reason === "malformed") { diff --git a/packages/trigger-sdk/test/chat-warm-process-reuse.test.ts b/packages/trigger-sdk/test/chat-warm-process-reuse.test.ts new file mode 100644 index 00000000000..266b10dfa5a --- /dev/null +++ b/packages/trigger-sdk/test/chat-warm-process-reuse.test.ts @@ -0,0 +1,60 @@ +import "../src/v3/test/index.js"; + +import { resourceCatalog, sessionStreams } from "@trigger.dev/core/v3"; +import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; +import { describe, expect, it } from "vitest"; +import { chat, type ChatTaskWirePayload } from "../src/v3/ai.js"; + +function userPayload(chatId: string, id: string): ChatTaskWirePayload { + return { + chatId, + trigger: "submit-message", + message: { id, role: "user", parts: [{ type: "text", text: id }] }, + }; +} + +/** + * A worker process is reused across runs. The end of every run tears down the + * channel subscription via `sessionStreams.clearHandlers()`, so a second run of + * the same chat in the same process has to attach a fresh one. Anything cached + * across runs that skips the attach leaves the new run with no input at all, and + * the conversation hangs with no error raised. + */ +describe("chat input across runs in one warm process", () => { + it("delivers messages to a second run of the same chat", async () => { + const chatId = "warm-reuse"; + const seen: string[] = []; + + const agent = chat.customAgent({ + id: "chat-warm-process-reuse", + run: async () => { + const record = await chat.messages.next({ timeoutInSeconds: 2 }); + const part = record?.payload.message?.parts?.[0]; + if (part && part.type === "text") seen.push(part.text); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + for (const attempt of ["first", "second"]) { + await runInMockTaskContext( + async (drivers) => { + const runPromise = run( + { chatId, trigger: "submit-message" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, attempt) }, + "in" + ); + await runPromise; + }, + { ctx: { run: { id: `run_${attempt}` } } } + ); + sessionStreams.clearHandlers(); + } + + expect(seen).toEqual(["first", "second"]); + }); +});