From 843777fbd472114b3812733c64be9df1fc9d4568 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 20:26:17 -0700 Subject: [PATCH 01/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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. //