diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md new file mode 100644 index 00000000000..24114e838c3 --- /dev/null +++ b/.changeset/tidy-mailboxes-wait.md @@ -0,0 +1,8 @@ +--- +"@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. + +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 197bff6b5e1..8c7a33b5dbd 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,60 @@ 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 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 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({ 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 +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, `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. + +`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: ```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 ffd9bb18084..d39ac7fcaa6 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..82a44c72b3c 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 { + 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 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); } @@ -55,6 +102,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 9b489616f74..4262674bfe3 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 @@ -11,7 +12,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 { @@ -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; @@ -160,3 +186,296 @@ 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 manager = new StandardSessionStreamManager( + repeatingApiClient(records[0]!), + "http://localhost" + ); + + 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 result = await manager.onceRecord(sessionId, io, { timeoutMs: 0 }); + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.error).toBeInstanceOf(InputStreamTimeoutError); + } + + manager.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.peekRecord(sessionId, io)).toEqual({ + id: "handover-1", + seqNum: 50, + data: { kind: "handover" }, + }); + 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(); + }); + + 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("preserves buffered records across disconnect and consumes only the exact sequence", 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, + }, + ]), + "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)?.seqNum).toBe(50); + 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(); + }); + + 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 fb87b211643..9360e224b11 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 @@ -72,14 +66,24 @@ 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` 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 // 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 @@ -123,28 +127,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,30 +159,62 @@ 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); + 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 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 = this.#takeBufferedRecord(key, predicate); + if (record) { return new InputStreamOncePromise((resolve) => { - resolve({ ok: true, output: data }); + resolve({ ok: true, output: record }); }); } - return new InputStreamOncePromise((resolve, reject) => { - const waiter: OnceWaiter = { resolve, reject }; + return new InputStreamOncePromise((resolve, reject) => { + const waiter: OnceWaiter = { resolve, reject, predicate }; if (options?.signal) { if (options.signal.aborted) { @@ -221,10 +250,31 @@ 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 { - 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]; } lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { @@ -239,21 +289,104 @@ 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 { - 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) { this.lastDispatchedSeqNums.set(key, seqNum); } } + 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; + + 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) { @@ -267,16 +400,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; @@ -285,7 +414,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 @@ -295,8 +423,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { tail.abortController.abort(); 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. @@ -335,6 +461,8 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.disconnect(); this.seqNums.clear(); this.lastDispatchedSeqNums.clear(); + this.unconsumedSeqNums.clear(); + this.cursorBarriers.clear(); this.minTimestamps.clear(); this.handlers.clear(); this.reconnectAttempts.clear(); @@ -350,7 +478,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { } this.onceWaiters.clear(); this.buffer.clear(); - this.bufferSeqNums.clear(); } #ensureTailConnected(sessionId: string, io: SessionChannelIO): void { @@ -368,15 +495,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; } @@ -427,9 +549,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 @@ -454,7 +575,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 +604,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 +630,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 +641,51 @@ 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); + if (this.#isCursorBarrier(key, record)) { + this.#markUnconsumedRecord(key, record.seqNum); + } + 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..38a8dc3f850 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,16 +26,49 @@ 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; + } + + setCursorBarrier( + _sessionId: string, + _io: SessionChannelIO, + _predicate: SessionStreamRecordPredicate | undefined + ): void {} + lastSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { return undefined; } 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 ae24259b3fc..8e518d0d9c7 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,20 +57,56 @@ 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; + + /** + * 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; /** 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 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. @@ -65,7 +116,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; @@ -82,7 +134,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/test/mock-task-context.ts b/packages/core/src/v3/test/mock-task-context.ts index 085acf999f1..21e90755cfa 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/session-waitpoint-backend.ts b/packages/core/src/v3/test/session-waitpoint-backend.ts index 8cae877f54e..71c0a3ddf93 100644 --- a/packages/core/src/v3/test/session-waitpoint-backend.ts +++ b/packages/core/src/v3/test/session-waitpoint-backend.ts @@ -113,8 +113,12 @@ export class SessionWaitpointBackend { }; } - const output = typeof result === "string" ? result : JSON.stringify(result); - return { ok: true, output, outputType: "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, @@ -144,16 +148,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 +173,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 0e08441d4c3..9e4ece8f01c 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,9 +37,11 @@ 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(); + private unconsumedSeqNums = new Map>(); + private cursorBarriers = new Map(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { const key = keyFor(sessionId, io); @@ -55,21 +63,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 +97,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 +141,26 @@ 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 }); + 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; + } + } + + if (options?.timeoutMs === 0) { + resolve({ + ok: false, + error: new InputStreamTimeoutError(key, 0), + }); 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 +195,21 @@ 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; + } + + 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]; } lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { @@ -151,22 +220,83 @@ 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 { - // `__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)); + 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 { - const key = keyFor(sessionId, io); + 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); } } + #isCursorBarrier(key: string, record: SessionStreamRecord): boolean { + const predicate = this.cursorBarriers.get(key); + if (!predicate) return true; + try { + return predicate(record); + } catch { + return true; + } + } + + #markUnconsumedRecord(key: string, seqNum: number): void { + if (!Number.isFinite(seqNum)) return; + + let unconsumedSeqNums = this.unconsumedSeqNums.get(key); + if (!unconsumedSeqNums) { + unconsumedSeqNums = new Set(); + this.unconsumedSeqNums.set(key, unconsumedSeqNums); + } + unconsumedSeqNums.add(seqNum); + } + + #removeUnconsumedRecord(key: string, seqNum: number): void { + const unconsumedSeqNums = this.unconsumedSeqNums.get(key); + unconsumedSeqNums?.delete(seqNum); + if (unconsumedSeqNums?.size === 0) { + this.unconsumedSeqNums.delete(key); + } + } + setMinTimestamp( _sessionId: string, _io: SessionChannelIO, @@ -180,15 +310,18 @@ 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; } disconnectStream(_sessionId: string, _io: SessionChannelIO): void { - // no-op — no real SSE tail in tests + // The production manager keeps buffered records reachable across a + // waitpoint suspension. The exact waitpoint record is removed on resume. } clearHandlers(): void { @@ -209,6 +342,7 @@ export class TestSessionStreamManager implements SessionStreamManager { this.buffer.clear(); this.seqNums.clear(); this.dispatchedSeqNums.clear(); + this.unconsumedSeqNums.clear(); } disconnect(): void { @@ -235,39 +369,56 @@ 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; + 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, + 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 +427,48 @@ export class TestSessionStreamManager implements SessionStreamManager { buffered = []; this.buffer.set(key, buffered); } - buffered.push(data); + buffered.push(record); + if (this.#isCursorBarrier(key, record)) { + this.#markUnconsumedRecord(key, record.seqNum); + } + 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 c241930323b..ce14b607aec 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1543,7 +1543,48 @@ 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 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; +}; + +/** + * 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"; +} + +const messagesInput: ChatMessages = { id: "chat-messages", on(handler) { return getChatSession().in.on((chunk) => { @@ -1607,6 +1648,36 @@ const messagesInput: RealtimeDefinedInputStream = { if (chunk && chunk.kind === "message") return chunk.payload; return undefined; }, + async hasPending() { + return messagesInput.peek() !== 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 { @@ -1831,15 +1902,144 @@ 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; + +/** + * 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"); + +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; + if (!KNOWN_CHAT_INPUT_KINDS.has(kind)) { + logger.warn("chat: discarded a session.in record of an unrecognised kind", { 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()); +} + +/** + * 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()); } /** @@ -5369,11 +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); + attachChatInputDrain(payload); return userRun(payload, runOptions); }, }); @@ -5480,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 @@ -5766,6 +5969,8 @@ function chatAgent< } } + attachChatInputDrain(payload); + // ── Recovery boot + chain reconstruction ──────────────────────── if (!hydrateMessages) { const settledMessages = mergeByIdReplaceWins( @@ -10796,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. // diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index 8a01f8293c4..86ac0389c81 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 @@ -713,6 +713,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 +721,7 @@ export class SessionInputChannel { idempotencyKey: options?.idempotencyKey, idempotencyKeyTTL: options?.idempotencyKeyTTL, tags: options?.tags, - lastSeqNum: sessionStreams.lastSeqNum(this.sessionId, "in"), + lastSeqNum: lastConsumedSeqNum, }); const result = await tracer.startActiveSpan( @@ -735,40 +736,60 @@ 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 so nothing is lost across the suspend. sessionStreams.disconnectStream(this.sessionId, "in"); const waitResult = await runtime.waitUntil(response.waitpointId); - const data = - waitResult.output !== undefined - ? await conditionallyImportAndParsePacket( - { - data: waitResult.output, - dataType: 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); - - return { ok: true as const, output: data as T }; - } else { - const error = new WaitpointTimeoutError(data?.message ?? "Timed out"); + 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 }; } + + // 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 new file mode 100644 index 00000000000..10dafc47f0c --- /dev/null +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -0,0 +1,359 @@ +// 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: "handover-prepare" }, + { 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 }); + }, + }); + 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; + pendingAfterControl?: boolean; + 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 }); + 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"); + }, + }); + 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: "handover-prepare" }, + { 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: false, + blocked: undefined, + cursorAfterBlocked: undefined, + headAfterBlocked: { + id: "handover-1", + seqNum: 30, + data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + }, + control: { + id: "handover-1", + seqNum: 30, + data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + }, + pendingAfterControl: true, + message: { + id: "message-1", + seqNum: 31, + payload: userPayload(chatId, "u-after-handover"), + }, + cursorAfterMessage: 31, + }); + }); + + 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: "handover-prepare" }, + { 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 }); + + 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"); + + 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(); + + 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); + }); + + 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/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(); } diff --git a/packages/trigger-sdk/test/pending-message-drain.test.ts b/packages/trigger-sdk/test/pending-message-drain.test.ts index f5bd7057515..7bb0d3d8249 100644 --- a/packages/trigger-sdk/test/pending-message-drain.test.ts +++ b/packages/trigger-sdk/test/pending-message-drain.test.ts @@ -70,6 +70,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 +263,95 @@ 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(JSON.stringify(first)); + let registeredLastSeqNum: number | 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 }) => { + registeredLastSeqNum = body.lastSeqNum; + 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(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("acknowledges the delivered record when identical payloads repeat on the channel", async () => { + __setSessionOpenImplForTests(undefined); + const chunk = { kind: "message", payload: { id: "repeated" } }; + const raw = JSON.stringify(chunk); + const sessionId = "ack-repeated-payload"; + + await runInMockTaskContext( + async (drivers) => { + sessionStreams.setLastSeqNum(sessionId, "in", 6); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ + waitpointId: "wp_ack_repeated", + isCached: false, + }), + 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 }; + }, + readSessionStreamRecords: async () => ({ + records: [ + { 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(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); + }, + { runtimeManager: runtimeWithWaitpointOutput(raw) } + ); }); });