diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md new file mode 100644 index 00000000000..4205570fef5 --- /dev/null +++ b/.changeset/tidy-mailboxes-wait.md @@ -0,0 +1,23 @@ +--- +"@trigger.dev/core": patch +"@trigger.dev/sdk": patch +--- + +Fixes a case where a chat could silently lose a message. If a message arrived while the agent was between turns and a stop arrived after it, the cursor the next boot resumed from could point past that message, so it was never answered and no error was raised. This affected `chat.agent`, not just custom agents. + +Fixes a recovered answer being cut off. After a crash the agent replays the message it had not answered yet, but it was replaying the stop that arrived after that message too, so the turn answering it was aborted the moment it began. A stop is now only applied to the turn that was live when it arrived. That holds however the stop got there: sent after the last completed turn, or sent to a chat whose most recent turn was completed by an older version of the SDK. + +Also fixes a retried send being answered twice. When a send was retried and its idempotency claim was lost, the agent could consume the same message a second time. + +Custom agent loops can now inspect pending chat input without consuming it, and consume one record at a time, with `chat.messages.hasPending()` and `chat.messages.next()`. Records carry stable identifiers so a redelivery is recognisable. + +```ts +if (await chat.messages.hasPending()) { + const record = await chat.messages.next({ timeoutInSeconds: 0 }); + if (record) handle(record.payload); +} +``` + +`hasPending()` answers for messages alone, so a message sitting behind a stop, or behind a record this version of the SDK does not recognise, still reports as pending and is still delivered. Anything the agent has no consumer for is discarded rather than left where it would make every message queued behind it undeliverable. `chat.messages.next()` returning `undefined` means no message became consumable before the timeout. + +`chat.writeTurnComplete()`'s `sessionInEventId` is the cursor that is safe to resume from, not the sequence of the record the turn answered. It is held back behind any message still waiting to be handled, so a value below the record you just handled is expected. diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 197bff6b5e1..7aa52a5f358 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -179,6 +179,12 @@ for await (const turn of session) { The frontend stops a turn with [`transport.stopGeneration(chatId)`](/ai-chat/frontend#stop-generation), which writes a stop signal to the session's input stream. It aborts the current turn's generation but keeps the run alive, so the next message continues on the same session. +A stop only applies to the turn that was live when it arrived. If the run crashes +and a later run recovers a message that had not been answered yet, a stop that +was already applied before the crash is not applied again, so the turn answering +the recovered message runs to completion. A stop sent after the recovery is live +and aborts that turn as normal. + `turn.signal` is a combined stop-and-cancel `AbortSignal`, fresh each turn. Pass it to `streamText` so the stop reaches the model, then let `turn.complete()` finish the turn: ```ts trigger/my-chat.ts @@ -213,7 +219,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 +227,57 @@ 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 next queued message without consuming it, or `undefined` when none is queued | +| `hasPending()` | Resolve `true` when a message is queued; does not consume it | +| `next({ timeoutInSeconds? })` | Consume exactly one message record in channel order, or resolve `undefined` when the optional timeout elapses | +| `on(handler)` | Consume messages as they arrive and invoke the handler | +| `waitWithIdleTimeout(options)` | Wait warm, then suspend the run until the next message arrives | + +`hasPending()` checks whether a message has already been delivered locally and is +waiting for `next()` to take it. It does not query the remote Session channel or +start a subscription. Use `waitWithIdleTimeout()` when the loop needs to idle +until future input arrives. + +`next({ timeoutInSeconds: 0 })` is also a local, non-blocking read. Call +`next()` without a timeout, or with a positive timeout, to subscribe for future +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 stops and +handovers. Those are routed to their own consumers and never block messages: a +message that arrived behind one is still reported by `hasPending()` and still +returned by `next()`, in channel order. The same holds for a record kind this +version of the SDK does not recognise, which is discarded rather than left where +it would make every message behind it undeliverable. + +`next()` returns `undefined` when no message became consumable before the +timeout. + 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..dd444961af9 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -504,9 +504,9 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`. | `chat.createSession(payload, options)` | Create an async iterator for chat turns | | `chat.pipe(source, options?)` | Pipe a stream to the frontend (from anywhere inside a task) | | `chat.pipeAndCapture(source, options?)` | Pipe and capture the response; returns `{ message, status, error }` | -| `chat.writeTurnComplete(options?)` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | +| `chat.writeTurnComplete(options?)` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors. `sessionInEventId` is a lower bound, not the sequence of the record the turn answered | | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | -| `chat.messages` | 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()` | @@ -645,7 +645,7 @@ The `onEvent` callback receives a `ChatTransportEvent` (exported from `@trigger. | `message-send-failed` | `messageId?`, `source`, `error`, `status?`, `durationMs`, `partId?`, `bodyBytes?` | A send definitively failed after internal retries. Fires in addition to `useChat`'s `onError`. | | `stream-connected` | `resumed`, `lastEventId?`, `messageId?` | The SSE subscription to the session's output stream started delivering. `resumed: true` when reconnecting from a stored cursor (page reload) rather than following a fresh send. `lastEventId` is the cursor it connected from. | | `first-chunk` | `chunkType?`, `lastEventId?`, `messageId?`, `sinceSendMs?` | The first response chunk of a turn arrived. `sinceSendMs` is the delta from the last turn-producing send — time to first token without any bookkeeping. | -| `turn-completed` | `lastEventId?`, `sessionInEventId?`, `messageId?`, `sinceSendMs?` | The agent's turn-complete control record arrived — the "finished answering" signal. `sinceSendMs` is the full turn latency; `sessionInEventId` is the agent's committed input-stream cursor. | +| `turn-completed` | `lastEventId?`, `sessionInEventId?`, `messageId?`, `sinceSendMs?` | The agent's turn-complete control record arrived — the "finished answering" signal. `sinceSendMs` is the full turn latency; `sessionInEventId` is the cursor the agent can safely resume its input stream from. Treat it as a lower bound: it is held back behind any message still waiting to be handled, so it can be below the sequence of the record this turn answered. Do not use it to decide whether a turn boundary belongs to your own send. | | `stream-error` | `error`, `status?` | The output stream failed unrecoverably. | `source` identifies the send path: `"submit-message"`, `"regenerate-message"`, `"steer"` (`sendPendingMessage`), `"action"` (`sendAction`), `"stop"` (`stopGeneration`), or `"head-start"`. 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/session-streams-api.ts b/packages/core/src/v3/session-streams-api.ts index 4f5c979aa3b..638a8674213 100644 --- a/packages/core/src/v3/session-streams-api.ts +++ b/packages/core/src/v3/session-streams-api.ts @@ -7,3 +7,4 @@ export const sessionStreams = SessionStreamsAPI.getInstance(); export * from "./sessionStreams/types.js"; export * from "./sessionStreams/wireProtocol.js"; export * from "./sessionStreams/chatSnapshot.js"; +export * from "./sessionStreams/router.js"; diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 21e2e8d2450..a1b6f840cf9 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"; @@ -35,6 +41,18 @@ export class SessionStreamsAPI implements SessionStreamManager { return this.#getManager().on(sessionId, io, handler); } + public onRecord( + sessionId: string, + io: SessionChannelIO, + handler: (record: SessionStreamRecord) => void | boolean | Promise + ): { off: () => void } { + const manager = this.#getManager(); + if (!manager.onRecord) { + throw new Error("The configured Session stream manager does not support record handlers"); + } + return manager.onRecord(sessionId, io, handler); + } + public once( sessionId: string, io: SessionChannelIO, @@ -43,10 +61,43 @@ 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 lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.#getManager().lastSeqNum(sessionId, io); } @@ -55,6 +106,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); } @@ -75,6 +134,10 @@ export class SessionStreamsAPI implements SessionStreamManager { return this.#getManager().shiftBuffer(sessionId, io); } + public reconnectStream(sessionId: string, io: SessionChannelIO): void { + this.#getManager().reconnectStream?.(sessionId, io); + } + public disconnectStream(sessionId: string, io: SessionChannelIO): void { this.#getManager().disconnectStream(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.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..73c85f972e2 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 @@ -12,9 +17,21 @@ import { controlSubtype } from "./wireProtocol.js"; // available to other consumers. See `SessionStreamManager.on` in types.ts. type SessionStreamHandler = (data: unknown) => void | boolean | Promise; +/** + * A handler that sees the whole record rather than just its payload. Consumers + * that route by sequence number need the metadata, the same reason + * `onceRecord` exists alongside `once`. + */ +type SessionStreamRecordHandler = (record: SessionStreamRecord) => void | boolean | Promise; + +type RegisteredHandler = + | { kind: "data"; fn: SessionStreamHandler } + | { kind: "record"; fn: SessionStreamRecordHandler }; + type OnceWaiter = { - resolve: (result: InputStreamOnceResult) => void; + 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 @@ -42,21 +59,9 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { * stream SSE. */ export class StandardSessionStreamManager implements SessionStreamManager { - private handlers = new Map>(); + private handlers = new Map>(); private onceWaiters = new Map(); - private buffer = new Map(); - // 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 +77,17 @@ export class StandardSessionStreamManager implements SessionStreamManager { // that's already being delivered out-of-band via the waitpoint. private explicitlyDisconnected = new Set(); private seqNums = new Map(); - // Highest seq_num that has been *consumed* (delivered to a once() - // waiter or shifted off the buffer into a once() caller) on a channel. + // Sequence numbers for records that were delivered but not consumed. + // Kept separately from `buffer` so the committed cursor can be calculated + // without depending on buffer traversal. + private unconsumedSeqNums = new Map>(); + + // High-water mark of seq_nums that have been *consumed* (delivered to a + // once() waiter or shifted off the buffer into a once() caller) on a channel. // 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 @@ -96,6 +104,26 @@ export class StandardSessionStreamManager implements SessionStreamManager { ) {} on(sessionId: string, io: SessionChannelIO, handler: SessionStreamHandler): { off: () => void } { + return this.#register(sessionId, io, { kind: "data", fn: handler }); + } + + /** + * Register a handler that receives the full record, including its sequence + * number. Same consume semantics as {@link on}: returning `true` consumes. + */ + onRecord( + sessionId: string, + io: SessionChannelIO, + handler: SessionStreamRecordHandler + ): { off: () => void } { + return this.#register(sessionId, io, { kind: "record", fn: handler }); + } + + #register( + sessionId: string, + io: SessionChannelIO, + handler: RegisteredHandler + ): { off: () => void } { const key = keyFor(sessionId, io); let handlerSet = this.handlers.get(key); @@ -123,28 +151,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); 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 +183,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 +274,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 +313,73 @@ 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); } } + #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 +393,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 +407,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,14 +416,25 @@ 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. this.reconnectAttempts.delete(key); } + /** + * Re-open a channel that `disconnectStream` closed, without registering a + * new consumer. A single long-lived reader (the session channel router) has + * to be able to bring its own tail back after a suspend, and re-attaching + * its handler just to clear the suppression flag would replay the buffer at + * it. + */ + reconnectStream(sessionId: string, io: SessionChannelIO): void { + const key = keyFor(sessionId, io); + this.explicitlyDisconnected.delete(key); + this.#ensureTailConnected(sessionId, io); + } + clearHandlers(): void { this.handlers.clear(); @@ -335,6 +467,7 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.disconnect(); this.seqNums.clear(); this.lastDispatchedSeqNums.clear(); + this.unconsumedSeqNums.clear(); this.minTimestamps.clear(); this.handlers.clear(); this.reconnectAttempts.clear(); @@ -350,7 +483,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { } this.onceWaiters.clear(); this.buffer.clear(); - this.bufferSeqNums.clear(); } #ensureTailConnected(sessionId: string, io: SessionChannelIO): void { @@ -368,15 +500,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 +554,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 +580,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 +609,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); return; } @@ -511,11 +635,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); if (consumed) { - if (seqNum !== undefined) { - this.#advanceLastDispatched(key, seqNum); - } + this.#advanceLastDispatched(key, record.seqNum); return; } @@ -524,26 +646,58 @@ export class StandardSessionStreamManager implements SessionStreamManager { buffered = []; this.buffer.set(key, buffered); } - buffered.push(data); - let bufferedSeqs = this.bufferSeqNums.get(key); - if (!bufferedSeqs) { - bufferedSeqs = []; - this.bufferSeqNums.set(key, bufferedSeqs); + buffered.push(record); + this.#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. */ - #invokeHandlers(key: string, data: unknown): boolean { + #invokeHandlers(key: string, record: SessionStreamRecord): boolean { const handlers = this.handlers.get(key); if (!handlers) return false; let consumed = false; for (const handler of handlers) { - if (this.#invokeHandler(handler, data)) { + if (this.#invokeHandler(handler, record)) { consumed = true; } } @@ -551,9 +705,9 @@ export class StandardSessionStreamManager implements SessionStreamManager { } /** Returns true when the handler synchronously consumed the record (returned `true`). */ - #invokeHandler(handler: SessionStreamHandler, data: unknown): boolean { + #invokeHandler(handler: RegisteredHandler, record: SessionStreamRecord): boolean { try { - const result = handler(data); + const result = handler.kind === "record" ? handler.fn(record) : handler.fn(record.data); if (result === true) return true; if (result && typeof result === "object" && "catch" in result) { (result as Promise).catch((error) => { diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index f2d355d24ef..8e894dfc6d3 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( @@ -11,6 +16,14 @@ export class NoopSessionStreamManager implements SessionStreamManager { return { off: () => {} }; } + onRecord( + _sessionId: string, + _io: SessionChannelIO, + _handler: (record: SessionStreamRecord) => void | boolean | Promise + ): { off: () => void } { + return { off: () => {} }; + } + once( _sessionId: string, _io: SessionChannelIO, @@ -21,16 +34,43 @@ export class NoopSessionStreamManager implements SessionStreamManager { }); } + onceRecord( + _sessionId: string, + _io: SessionChannelIO, + _options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return new InputStreamOncePromise(() => { + // Never resolves in noop mode. + }); + } + + onceRecordWhere( + _sessionId: string, + _io: SessionChannelIO, + _predicate: SessionStreamRecordPredicate, + _options?: InputStreamOnceOptions + ): InputStreamOncePromise { + return new InputStreamOncePromise(() => { + // Never resolves in noop mode. + }); + } + peek(_sessionId: string, _io: SessionChannelIO): unknown | undefined { return undefined; } + peekRecord(_sessionId: string, _io: SessionChannelIO): SessionStreamRecord | undefined { + return undefined; + } + 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; } @@ -49,6 +89,8 @@ export class NoopSessionStreamManager implements SessionStreamManager { disconnectStream(_sessionId: string, _io: SessionChannelIO): void {} + reconnectStream(_sessionId: string, _io: SessionChannelIO): void {} + clearHandlers(): void {} reset(): void {} diff --git a/packages/core/src/v3/sessionStreams/router.test.ts b/packages/core/src/v3/sessionStreams/router.test.ts new file mode 100644 index 00000000000..6f654269d1c --- /dev/null +++ b/packages/core/src/v3/sessionStreams/router.test.ts @@ -0,0 +1,409 @@ +import { describe, expect, it } from "vitest"; +import { SessionChannelRouter } from "./router.js"; +import type { SessionRouteTable } from "./router.js"; +import type { SessionStreamRecord } from "./types.js"; + +type Chunk = { kind: string; text?: string }; + +const CHAT_TABLE: SessionRouteTable = { + kindOf: (data) => (data as Chunk | undefined)?.kind, + routes: [ + { name: "messages", delivery: "queue", replayable: true, kinds: ["message"] }, + { name: "stop", delivery: "at-arrival", replayable: false, kinds: ["stop"] }, + { + name: "handover", + delivery: "queue", + replayable: false, + kinds: ["handover", "handover-skip"], + }, + ], +}; + +function router(onDrop?: Parameters[0]) { + return makeRouter(onDrop); +} + +function makeRouter( + onDrop?: (record: SessionStreamRecord, reason: string, route?: string) => void +) { + return new SessionChannelRouter(CHAT_TABLE, { onDrop }); +} + +function rec(seqNum: number, kind: string, text?: string): SessionStreamRecord { + return { id: `r${seqNum}`, seqNum, data: { kind, ...(text ? { text } : {}) } }; +} + +describe("SessionChannelRouter: table validation", () => { + it("rejects a route that is at-arrival and replayable", () => { + expect( + () => + new SessionChannelRouter({ + kindOf: () => "x", + routes: [{ name: "bad", delivery: "at-arrival", replayable: true, kinds: ["x"] }], + }) + ).toThrow(/at-arrival and replayable/); + }); + + it("rejects a kind claimed by two routes", () => { + expect( + () => + new SessionChannelRouter({ + kindOf: () => "x", + routes: [ + { name: "a", delivery: "queue", replayable: true, kinds: ["x"] }, + { name: "b", delivery: "queue", replayable: false, kinds: ["x"] }, + ], + }) + ).toThrow(/claimed by both/); + }); +}); + +describe("SessionChannelRouter: classification", () => { + it("queues a message when nobody is ready for it", () => { + const r = router(); + expect(r.ingest(rec(0, "message", "M0"))).toEqual({ action: "queue", route: "messages" }); + expect(r.hasPending("messages")).toBe(true); + }); + + it("discards a stop with no handler attached", () => { + const r = router(); + expect(r.ingest(rec(0, "stop"))).toEqual({ + action: "drop", + route: "stop", + reason: "no-handler", + }); + }); + + it("delivers a stop to a live handler", () => { + const r = router(); + const seen: number[] = []; + r.on("stop", (record) => seen.push(record.seqNum)); + expect(r.ingest(rec(3, "stop"))).toEqual({ action: "deliver", route: "stop" }); + expect(seen).toEqual([3]); + }); + + it("drops a kind no route claims, and reports it once", () => { + const drops: Array<[number, string]> = []; + const r = router((record, reason) => drops.push([record.seqNum, reason])); + expect(r.ingest(rec(1, "some-future-kind"))).toEqual({ + action: "drop", + reason: "unroutable", + }); + expect(drops).toEqual([[1, "unroutable"]]); + }); + + it("drops a record with no usable kind", () => { + const r = router(); + expect(r.ingest({ id: "x", seqNum: 0, data: { nope: true } })).toEqual({ + action: "drop", + reason: "malformed", + }); + }); + + it("does not let a throwing kindOf take the channel down", () => { + const r = new SessionChannelRouter({ + kindOf: () => { + throw new Error("boom"); + }, + routes: [{ name: "m", delivery: "queue", replayable: true, kinds: ["message"] }], + }); + expect(r.ingest(rec(0, "message"))).toEqual({ action: "drop", reason: "malformed" }); + }); +}); + +describe("SessionChannelRouter: the wedge cannot happen", () => { + it("delivers a message queued behind an unroutable record", async () => { + const r = router(); + r.ingest(rec(0, "mystery-kind")); + r.ingest(rec(1, "message", "M1")); + + expect(r.hasPending("messages")).toBe(true); + const taken = await r.next("messages", { timeoutMs: 0 }); + expect((taken!.data as Chunk).text).toBe("M1"); + }); + + it("reports pending for a message queued behind a stop", () => { + const r = router(); + r.ingest(rec(0, "stop")); + r.ingest(rec(1, "message", "M1")); + + expect(r.hasPending("messages")).toBe(true); + expect((r.peek("messages")!.data as Chunk).text).toBe("M1"); + }); +}); + +describe("SessionChannelRouter: delivery ordering", () => { + it("serves a parked waiter before a push handler", async () => { + const r = router(); + const handlerSaw: string[] = []; + const pending = r.next("messages"); + r.on("messages", (record) => handlerSaw.push((record.data as Chunk).text!)); + + r.ingest(rec(0, "message", "M0")); + + expect((await pending)?.seqNum).toBe(0); + expect(handlerSaw).toEqual([]); + }); + + it("re-offers a queued record to a handler attaching later, in order", () => { + const r = router(); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "message", "M1")); + + const seen: string[] = []; + r.on("messages", (record) => seen.push((record.data as Chunk).text!)); + + expect(seen).toEqual(["M0", "M1"]); + expect(r.hasPending("messages")).toBe(false); + }); + + it("resolves next() undefined on timeout without consuming anything", async () => { + const r = router(); + expect(await r.next("messages", { timeoutMs: 5 })).toBeUndefined(); + r.ingest(rec(0, "message", "M0")); + expect((await r.next("messages", { timeoutMs: 0 }))?.seqNum).toBe(0); + }); +}); + +describe("SessionChannelRouter: the resume floor", () => { + it("sits at the high water when nothing is owed", () => { + const r = router(); + r.on("stop", () => {}); + r.ingest(rec(0, "message")); + r.next("messages", { timeoutMs: 0 }); + r.ingest(rec(1, "stop")); + + expect(r.resumeFloor()).toBe(1); + expect(r.appliedThrough()).toBe(1); + }); + + it("is held below a message still queued, even as control records advance", () => { + const r = router(); + r.on("stop", () => {}); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "message", "M1")); + r.next("messages", { timeoutMs: 0 }); + r.ingest(rec(2, "stop")); + r.ingest(rec(3, "stop")); + + expect(r.resumeFloor()).toBe(0); + expect(r.appliedThrough()).toBe(3); + }); + + it("is undefined rather than negative when the very first record is owed", () => { + const r = router(); + r.ingest(rec(0, "message", "M0")); + expect(r.resumeFloor()).toBeUndefined(); + }); + + it("is not held back by a queued non-replayable record", () => { + const r = router(); + r.ingest(rec(0, "handover")); + r.ingest(rec(1, "message", "M1")); + r.next("messages", { timeoutMs: 0 }); + + expect(r.pendingCount("handover")).toBe(1); + expect(r.resumeFloor()).toBe(1); + }); + + it("advances once the owed message is taken", async () => { + const r = router(); + r.on("stop", () => {}); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "stop")); + expect(r.resumeFloor()).toBeUndefined(); + + await r.next("messages", { timeoutMs: 0 }); + expect(r.resumeFloor()).toBe(1); + }); + + it("tracks the earliest of several owed messages", () => { + const r = router(); + r.ingest(rec(0, "message", "M0")); + r.ingest(rec(1, "message", "M1")); + r.ingest(rec(2, "message", "M2")); + r.next("messages", { timeoutMs: 0 }); + + expect(r.resumeFloor()).toBe(0); + expect(r.appliedThrough()).toBe(2); + }); +}); + +describe("SessionChannelRouter: resuming", () => { + it("does not apply a non-replayable record inside the replay window", () => { + const r = router(); + r.restore({ resumeFrom: 0, appliedThrough: 2 }); + const stops: number[] = []; + r.on("stop", (record) => stops.push(record.seqNum)); + + expect(r.ingest(rec(1, "message", "M1"))).toEqual({ action: "queue", route: "messages" }); + expect(r.ingest(rec(2, "stop"))).toEqual({ + action: "drop", + route: "stop", + reason: "replayed", + }); + expect(stops).toEqual([]); + }); + + it("applies the same kind arriving live, past the window", () => { + const r = router(); + r.restore({ resumeFrom: 0, appliedThrough: 2 }); + const stops: number[] = []; + r.on("stop", (record) => stops.push(record.seqNum)); + + r.ingest(rec(1, "message", "M1")); + r.ingest(rec(2, "stop")); + expect(r.ingest(rec(3, "stop"))).toEqual({ action: "deliver", route: "stop" }); + expect(stops).toEqual([3]); + }); + + it("falls back to the floor when no replay-window end is supplied", () => { + const r = router(); + r.restore({ resumeFrom: 4 }); + r.on("stop", () => {}); + + expect(r.ingest(rec(4, "stop")).action).toBe("drop"); + expect(r.ingest(rec(5, "stop"))).toEqual({ action: "deliver", route: "stop" }); + }); + + it("declines a control record inside a window resolved from the channel", () => { + const r = router(); + // What the chat layer supplies when the boundary predates the published + // window: everything already on the channel at boot counts as replayed. + r.restore({ resumeFrom: 4, appliedThrough: 6 }); + const stops: number[] = []; + r.on("stop", (record) => stops.push(record.seqNum)); + + expect(r.ingest(rec(5, "message", "M5"))).toEqual({ action: "queue", route: "messages" }); + expect(r.ingest(rec(6, "stop"))).toEqual({ + action: "drop", + route: "stop", + reason: "replayed", + }); + expect(r.ingest(rec(7, "stop"))).toEqual({ action: "deliver", route: "stop" }); + expect(stops).toEqual([7]); + }); + + it("applies everything on a fresh session with no checkpoint", () => { + const r = router(); + r.on("stop", () => {}); + expect(r.ingest(rec(0, "stop"))).toEqual({ action: "deliver", route: "stop" }); + }); + + it("never drops a replayable record, however far inside the window", () => { + const r = router(); + r.restore({ resumeFrom: 0, appliedThrough: 9 }); + expect(r.ingest(rec(1, "message", "M1"))).toEqual({ action: "queue", route: "messages" }); + }); + + it("keeps the floor from moving backwards past the restored point", () => { + const r = router(); + r.restore({ resumeFrom: 7, appliedThrough: 7 }); + expect(r.resumeFloor()).toBe(7); + }); +}); + +describe("SessionChannelRouter: consumer windows", () => { + it("queues a handover that arrives before its consumer is ready", async () => { + const r = router(); + expect(r.ingest(rec(0, "handover")).action).toBe("queue"); + + const taken = await r.next("handover", { timeoutMs: 0 }); + expect(taken?.seqNum).toBe(0); + }); + + it("discards what is left on a route when its window closes", () => { + const r = router(); + r.ingest(rec(0, "handover")); + r.clearRoute("handover"); + + expect(r.pendingCount("handover")).toBe(0); + expect(r.resumeFloor()).toBe(0); + }); + + it("wakes a waiter empty when its window closes", async () => { + const r = router(); + const pending = r.next("handover"); + r.clearRoute("handover"); + expect(await pending).toBeUndefined(); + }); +}); + +/** + * The invariant the whole design exists to hold: across any interleaving and + * any crash point, a message is delivered exactly once and a stop is never + * applied twice. + * + * Runs every record sequence over a simulated two-boot lifecycle. The second + * boot resubscribes from the published floor, exactly as the tail does with + * `Last-Event-ID`, so a floor that is too high shows up as a lost message and + * one that replays a stop shows up as a duplicate application. + */ +describe("SessionChannelRouter: exactly-once across a crash", () => { + const KINDS = ["message", "stop", "handover", "junk"] as const; + + function interleavings(length: number): string[][] { + if (length === 0) return [[]]; + const shorter = interleavings(length - 1); + const out: string[][] = []; + for (const prefix of shorter) { + for (const kind of KINDS) out.push([...prefix, kind]); + } + return out; + } + + function runBoot( + records: SessionStreamRecord[], + checkpoint: { resumeFrom?: number; appliedThrough?: number }, + takeMessages: number, + attachStop: boolean + ) { + const r = router(); + r.restore(checkpoint); + const messages: number[] = []; + const stops: number[] = []; + if (attachStop) r.on("stop", (record) => stops.push(record.seqNum)); + + for (const record of records) { + if (checkpoint.resumeFrom !== undefined && record.seqNum <= checkpoint.resumeFrom) continue; + r.ingest(record); + } + + for (let i = 0; i < takeMessages; i++) { + const head = r.peek("messages"); + if (!head) break; + void r.next("messages", { timeoutMs: 0 }); + messages.push(head.seqNum); + } + + return { messages, stops, checkpoint: r.checkpoint() }; + } + + it("delivers every message exactly once and applies no stop twice", () => { + const cases = interleavings(4); + expect(cases.length).toBe(256); + + for (const kinds of cases) { + const records = kinds.map((kind, index) => rec(index, kind)); + const messageSeqs = records + .filter((record) => (record.data as Chunk).kind === "message") + .map((record) => record.seqNum); + + for (let crashAfter = 0; crashAfter <= kinds.length; crashAfter++) { + const first = runBoot(records, {}, crashAfter, true); + const second = runBoot(records, first.checkpoint, kinds.length, true); + + const delivered = [...first.messages, ...second.messages]; + expect(delivered, `messages for [${kinds.join(",")}] crashAfter=${crashAfter}`).toEqual( + messageSeqs + ); + + const appliedTwice = first.stops.filter((seq) => second.stops.includes(seq)); + expect( + appliedTwice, + `stops applied twice for [${kinds.join(",")}] crashAfter=${crashAfter}` + ).toEqual([]); + } + } + }); +}); diff --git a/packages/core/src/v3/sessionStreams/router.ts b/packages/core/src/v3/sessionStreams/router.ts new file mode 100644 index 00000000000..17b4cdaddb2 --- /dev/null +++ b/packages/core/src/v3/sessionStreams/router.ts @@ -0,0 +1,443 @@ +import type { SessionStreamRecord } from "./types.js"; + +/** + * What happens to a record when no consumer is ready for it *right now*. + * + * - `queue`: it waits in the route's own queue until a consumer takes it. + * - `at-arrival`: it goes to a live handler or nowhere. A record that only + * means something to the turn that is live when it lands (a stop) is this. + */ +export type RouteDelivery = "queue" | "at-arrival"; + +/** + * One route: which kinds it owns, whether it waits for a consumer, and whether + * a record it never handled has to survive into the next boot. + * + * The two properties are independent, and that is the point. Three of the four + * combinations are meaningful and cover everything `session.in` carries: + * + * | delivery | replayable | example | + * | --- | --- | --- | + * | `queue` | `true` | a user message: waits for a turn, and a crash must not lose it | + * | `at-arrival` | `false` | a stop: only the live turn cares, and a replayed one would abort the wrong turn | + * | `queue` | `false` | a handover signal: can arrive before its consumer is ready, but is meaningless to a later boot | + * + * The fourth is a contradiction (discard it when nobody is listening, yet + * recover it later) and the table rejects it. + */ +export type SessionRoute = { + /** Unique within a table. */ + name: string; + delivery: RouteDelivery; + /** + * Whether a record this route never handled must be recovered by the next + * boot. Only `true` holds the resume floor back. + */ + replayable: boolean; + /** Record kinds this route owns. Every kind belongs to at most one route. */ + kinds: readonly string[]; +}; + +/** + * The complete statement of what a channel carries and who owns each kind. + * Intended to be a literal at the point of use, so "which kinds exist and what + * happens to each" is answerable by reading one object. + */ +export type SessionRouteTable = { + /** Extract a record's kind. Returning `undefined` marks it malformed. */ + kindOf: (data: unknown) => string | undefined; + routes: readonly SessionRoute[]; +}; + +export type RouterDropReason = + /** No route claims this kind: nothing on this boot can consume it. */ + | "unroutable" + /** No usable kind on the record at all. */ + | "malformed" + /** + * A non-replayable record inside the replay window. It was already observed + * by a previous run, and its route has declared that a later boot has no use + * for it. + */ + | "replayed" + /** An `at-arrival` record with no handler attached right now. */ + | "no-handler"; + +export type RouterDecision = + /** Handed to a consumer that was already waiting, or to a live handler. */ + | { action: "deliver"; route: string } + /** Parked in the route's queue for a future consumer. */ + | { action: "queue"; route: string } + | { action: "drop"; route?: string; reason: RouterDropReason }; + +/** + * The two numbers a turn boundary publishes, and that a boot reads back. + * + * `resumeFrom` is the floor: every record at or below it was terminally + * handled, so a boot subscribes from just past it. `appliedThrough` is the end + * of the replay window: the highest sequence a previous run observed. Records + * at or below it are being re-read rather than arriving live. + */ +export type RouterCheckpoint = { + resumeFrom?: number; + appliedThrough?: number; +}; + +type QueueWaiter = { + resolve: (record: SessionStreamRecord | undefined) => void; + timer?: ReturnType; +}; + +type RouteHandler = (record: SessionStreamRecord) => void; + +/** + * One route's live state: its queue of records nobody has taken yet, the + * consumers waiting for the next one, and any attached push handlers. + * + * An `at-arrival` route never fills `queue`; a route that is not `replayable` + * fills it but is skipped when the floor is computed. + */ +class RouteState { + readonly queue: SessionStreamRecord[] = []; + readonly waiters: QueueWaiter[] = []; + readonly handlers = new Set(); + + constructor(readonly route: SessionRoute) {} + + /** + * Lowest sequence a *later boot* would still have to recover. A route that + * is not replayable never holds anything back, however much it has queued. + */ + earliestUnrecovered(): number | undefined { + if (!this.route.replayable) return undefined; + return this.queue.length > 0 ? this.queue[0]!.seqNum : undefined; + } +} + +/** + * Reads one session channel and gives every record exactly one destination. + * + * The channel carries records for several independent consumers whose delivery + * needs differ: a message must be delivered eventually and so can lag + * arbitrarily far behind the newest record, while a stop only means anything to + * the turn that is live when it lands. Tracking that with a single scalar + * cursor is not possible — the true state is always "control applied through + * X, message Y still owed" — so the router tracks it per route and derives the + * published numbers from route state: + * + * - the resume floor is held back only by records a later boot would still have + * to recover, which is exactly the queued records on replayable routes; + * - a record whose route has declared it not replayable is never re-applied + * after a resume, because the route said a later boot has no use for it; + * - a record with no route is terminal by classification, so it never enters a + * queue and cannot park at the head of one. + * + * Those three properties are what previously needed a cursor-barrier + * predicate, a drop predicate with a second published header, and a + * discard-the-unclaimed drain respectively. + */ +export class SessionChannelRouter { + #routes = new Map(); + #kindToRoute = new Map(); + #highestSeq: number | undefined; + #resumeFrom: number | undefined; + #appliedThrough: number | undefined; + #onDrop?: (record: SessionStreamRecord, reason: RouterDropReason, route?: string) => void; + + constructor( + private table: SessionRouteTable, + options?: { + /** Called for every dropped record. Reporting only; never load-bearing. */ + onDrop?: (record: SessionStreamRecord, reason: RouterDropReason, route?: string) => void; + } + ) { + for (const route of table.routes) { + if (this.#routes.has(route.name)) { + throw new Error(`Duplicate route name "${route.name}" in session route table`); + } + if (route.delivery === "at-arrival" && route.replayable) { + throw new Error( + `Route "${route.name}" is at-arrival and replayable, which cannot both hold: a record discarded because nobody was listening cannot also be recovered later` + ); + } + this.#routes.set(route.name, new RouteState(route)); + for (const kind of route.kinds) { + const existing = this.#kindToRoute.get(kind); + if (existing) { + throw new Error( + `Kind "${kind}" is claimed by both "${existing}" and "${route.name}" in session route table` + ); + } + this.#kindToRoute.set(kind, route.name); + } + } + this.#onDrop = options?.onDrop; + } + + /** + * Seed the router from a previous run's turn boundary. + * + * An absent `appliedThrough` falls back to the floor, which leaves anything + * above the floor treated as live. A caller resuming a boundary that predates + * the published replay window should resolve the window from the channel + * instead, so it covers everything already there at boot: a missed stop is + * recoverable, while a stop applied to the wrong turn kills a live answer. + */ + restore(checkpoint: RouterCheckpoint): void { + this.#resumeFrom = checkpoint.resumeFrom; + this.#appliedThrough = checkpoint.appliedThrough ?? checkpoint.resumeFrom; + if (this.#resumeFrom !== undefined) { + this.#highestSeq = this.#resumeFrom; + } + } + + /** Where a boot should subscribe from: just past this sequence. */ + resumeFrom(): number | undefined { + return this.#resumeFrom; + } + + /** + * Classify one record and act on it. The record's destination is decided + * here, once, and never by whichever consumer happens to be waiting. + * + * Queued routes serve a waiting consumer before a push handler, so a handler + * can never take a record out from under a consumer that is actively awaiting + * one. A record handed straight to either never enters the queue, so it never + * holds the floor back. + */ + ingest(record: SessionStreamRecord): RouterDecision { + if (Number.isFinite(record.seqNum)) { + if (this.#highestSeq === undefined || record.seqNum > this.#highestSeq) { + this.#highestSeq = record.seqNum; + } + } + + const kind = this.#kindOf(record.data); + if (kind === undefined) { + return this.#drop(record, "malformed"); + } + + const routeName = this.#kindToRoute.get(kind); + const state = routeName ? this.#routes.get(routeName) : undefined; + if (!state || !routeName) { + return this.#drop(record, "unroutable"); + } + + if ( + !state.route.replayable && + this.#appliedThrough !== undefined && + record.seqNum <= this.#appliedThrough + ) { + return this.#drop(record, "replayed", routeName); + } + + if ( + state.route.delivery === "at-arrival" && + state.handlers.size === 0 && + state.waiters.length === 0 + ) { + return this.#drop(record, "no-handler", routeName); + } + + const waiter = state.waiters.shift(); + if (waiter) { + if (waiter.timer) clearTimeout(waiter.timer); + waiter.resolve(record); + return { action: "deliver", route: routeName }; + } + + if (state.handlers.size > 0) { + this.#invokeHandlers(state, record); + return { action: "deliver", route: routeName }; + } + + state.queue.push(record); + return { action: "queue", route: routeName }; + } + + #kindOf(data: unknown): string | undefined { + try { + const kind = this.table.kindOf(data); + return typeof kind === "string" && kind.length > 0 ? kind : undefined; + } catch { + return undefined; + } + } + + #drop(record: SessionStreamRecord, reason: RouterDropReason, route?: string): RouterDecision { + try { + this.#onDrop?.(record, reason, route); + } catch { + void 0; + } + return { action: "drop", route, reason }; + } + + #invokeHandlers(state: RouteState, record: SessionStreamRecord): void { + for (const handler of state.handlers) { + try { + handler(record); + } catch { + void 0; + } + } + } + + /** + * The highest sequence that can be resumed past without losing anything. + * + * Held back below the earliest record still queued anywhere, because those + * are exactly the records a replay has to recover. Everything else the + * router has seen was terminally decided, so the floor is free to sit at the + * high water when every queue is empty. + */ + resumeFloor(): number | undefined { + if (this.#highestSeq === undefined) return undefined; + + let earliestPending = Infinity; + for (const state of this.#routes.values()) { + const pending = state.earliestUnrecovered(); + if (pending !== undefined) earliestPending = Math.min(earliestPending, pending); + } + + if (earliestPending === Infinity) return this.#highestSeq; + + const floor = Math.min(this.#highestSeq, earliestPending - 1); + return floor >= 0 ? floor : undefined; + } + + /** + * The high water: highest sequence observed, never held back. Published as + * the end of the replay window so the next boot can tell a re-read + * `at-arrival` record from one arriving live. + */ + appliedThrough(): number | undefined { + return this.#highestSeq; + } + + /** Both published numbers for a turn boundary. */ + checkpoint(): RouterCheckpoint { + return { resumeFrom: this.resumeFloor(), appliedThrough: this.appliedThrough() }; + } + + #stateOrThrow(name: string): RouteState { + const state = this.#routes.get(name); + if (!state) throw new Error(`Unknown session route "${name}"`); + return state; + } + + /** + * Attach a push handler. A `queued` route with a handler attached delivers + * straight to it instead of queueing; an `at-arrival` route discards records + * whenever no handler is attached. + * + * Attaching re-offers anything already queued, so a consumer that attaches + * after records have piled up still sees them in order. + */ + on(name: string, handler: RouteHandler): { off: () => void } { + const state = this.#stateOrThrow(name); + state.handlers.add(handler); + + if (state.queue.length > 0) { + const pending = state.queue.splice(0, state.queue.length); + for (const record of pending) { + try { + handler(record); + } catch { + void 0; + } + } + } + + return { + off: () => { + state.handlers.delete(handler); + }, + }; + } + + /** Whether an `at-arrival` route currently has anywhere to deliver. */ + hasHandler(name: string): boolean { + return this.#stateOrThrow(name).handlers.size > 0; + } + + /** + * Take the next record on a route. + * + * `timeoutMs: 0` is a non-blocking take. Omitted means wait indefinitely. + * Resolves `undefined` on timeout. An `at-arrival` route delivers to a + * waiting caller when one is already parked here, which is what makes a pull + * consumer possible on such a route without weakening the discard rule. + */ + next(name: string, options?: { timeoutMs?: number }): Promise { + const state = this.#stateOrThrow(name); + const queued = state.queue.shift(); + if (queued) return Promise.resolve(queued); + + if (options?.timeoutMs === 0) return Promise.resolve(undefined); + + return new Promise((resolve) => { + const waiter: QueueWaiter = { resolve }; + if (options?.timeoutMs !== undefined) { + waiter.timer = setTimeout(() => { + const index = state.waiters.indexOf(waiter); + if (index !== -1) state.waiters.splice(index, 1); + resolve(undefined); + }, options.timeoutMs); + } + state.waiters.push(waiter); + }); + } + + /** Head of a route's queue without consuming it. */ + peek(name: string): SessionStreamRecord | undefined { + return this.#stateOrThrow(name).queue[0]; + } + + /** + * Whether a route has anything queued. Exact, because it reads that route's + * own queue rather than the head of a buffer shared with every other kind on + * the channel. + */ + hasPending(name: string): boolean { + return this.#stateOrThrow(name).queue.length > 0; + } + + /** Number of records queued on a route. Diagnostic use. */ + pendingCount(name: string): number { + return this.#stateOrThrow(name).queue.length; + } + + /** + * Discard whatever one route has queued and wake its waiters empty. + * + * Closes a consumer window: a route that is not replayable has nothing owed + * to a later boot, so once its window is over anything still queued on it is + * dead and must not sit at the head of the queue for the rest of the run. + */ + clearRoute(name: string): void { + const state = this.#stateOrThrow(name); + state.queue.length = 0; + for (const waiter of state.waiters) { + if (waiter.timer) clearTimeout(waiter.timer); + waiter.resolve(undefined); + } + state.waiters.length = 0; + } + + /** Drop every waiter and queue. Called between task executions. */ + reset(): void { + for (const state of this.#routes.values()) { + state.queue.length = 0; + state.handlers.clear(); + for (const waiter of state.waiters) { + if (waiter.timer) clearTimeout(waiter.timer); + waiter.resolve(undefined); + } + state.waiters.length = 0; + } + this.#highestSeq = undefined; + this.#resumeFrom = undefined; + this.#appliedThrough = undefined; + } +} diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index ae24259b3fc..cc6bde884cc 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 @@ -35,6 +50,16 @@ export interface SessionStreamManager { handler: (data: unknown) => void | boolean | Promise ): { off: () => void }; + /** + * Register a handler that receives the full record, including its sequence + * number. Same consume semantics as {@link on}. + */ + onRecord?( + sessionId: string, + io: SessionChannelIO, + handler: (record: SessionStreamRecord) => void | boolean | Promise + ): { off: () => void }; + /** Wait for the next record on the given channel (buffered or live). */ once( sessionId: string, @@ -42,20 +67,46 @@ 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; + /** 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,9 +134,12 @@ 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; + /** Re-open a channel closed by {@link disconnectStream}, registering nothing. */ + reconnectStream?(sessionId: string, io: SessionChannelIO): void; + /** Clear all `.on` handlers; abort tails without pending once-waiters. */ clearHandlers(): void; diff --git a/packages/core/src/v3/sessionStreams/wireProtocol.ts b/packages/core/src/v3/sessionStreams/wireProtocol.ts index 550e81a0af4..9a5cfb0c710 100644 --- a/packages/core/src/v3/sessionStreams/wireProtocol.ts +++ b/packages/core/src/v3/sessionStreams/wireProtocol.ts @@ -40,6 +40,18 @@ export const SESSION_STATE_LAST_EVENT_ID_HEADER = "last-event-id" as const; */ export const SESSION_IN_EVENT_ID_HEADER = "session-in-event-id" as const; +/** + * Sibling of {@link SESSION_IN_EVENT_ID_HEADER}: the highest `.in` sequence this + * run had actually consumed at the turn boundary, unclamped. + * + * The resume cursor is held back behind records still waiting to be handled, so + * resuming from it necessarily re-delivers records that WERE handled. A message + * re-delivered that way is the point. A control record re-delivered that way is + * a bug: it applies a second time to whatever turn is live on the new run. On + * boot this bound tells the run which control records it has already seen. + */ +export const SESSION_IN_CONSUMED_ID_HEADER = "session-in-consumed-id" as const; + export const TRIGGER_CONTROL_SUBTYPE = { TURN_COMPLETE: "turn-complete", UPGRADE_REQUIRED: "upgrade-required", diff --git a/packages/core/src/v3/test/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..5388c9dc40e 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; @@ -14,6 +20,9 @@ type OnceWaiter = { // returns `true` CONSUMES the record (not buffered, not re-delivered on a // future `on()` attach). See `SessionStreamManager.on` in types.ts. type Handler = (data: unknown) => void | boolean | Promise; +type RecordHandler = (record: SessionStreamRecord) => void | boolean | Promise; + +type RegisteredHandler = { kind: "data"; fn: Handler } | { kind: "record"; fn: RecordHandler }; function keyFor(sessionId: string, io: SessionChannelIO): string { return `${sessionId}:${io}`; @@ -29,13 +38,25 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { * registered are buffered so the first `once()` picks them up. */ export class TestSessionStreamManager implements SessionStreamManager { - private handlers = new Map>(); + private handlers = new Map>(); private onceWaiters = new Map(); - private buffer = new Map(); + private buffer = new Map(); private seqNums = new Map(); private dispatchedSeqNums = new Map(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { + return this.#register(sessionId, io, { kind: "data", fn: handler }); + } + + onRecord(sessionId: string, io: SessionChannelIO, handler: RecordHandler): { off: () => void } { + return this.#register(sessionId, io, { kind: "record", fn: handler }); + } + + #register( + sessionId: string, + io: SessionChannelIO, + handler: RegisteredHandler + ): { off: () => void } { const key = keyFor(sessionId, io); let set = this.handlers.get(key); @@ -55,21 +76,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 = this.#callHandler(handler, record) === 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 +110,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 +154,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 +208,11 @@ export class TestSessionStreamManager implements SessionStreamManager { } peek(sessionId: string, io: SessionChannelIO): unknown | undefined { - const buffered = this.buffer.get(keyFor(sessionId, io)); - if (buffered && buffered.length > 0) return buffered[0]; - return undefined; + return this.peekRecord(sessionId, io)?.data; + } + + peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined { + return this.buffer.get(keyFor(sessionId, io))?.[0]; } lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { @@ -151,16 +223,34 @@ 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)); } 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 { + if (!Number.isFinite(seqNum)) return; const current = this.dispatchedSeqNums.get(key); if (current === undefined || seqNum > current) { this.dispatchedSeqNums.set(key, seqNum); @@ -180,15 +270,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 { @@ -235,48 +328,88 @@ 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); return; } - const consumed = await this.#invokeHandlers(key, data); - if (consumed) 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); + const { consumed, settled } = this.#invokeHandlersSync(key, record); + if (!consumed) { + let buffered = this.buffer.get(key); + if (!buffered) { + buffered = []; + this.buffer.set(key, buffered); } - w.resolve({ ok: true, output: data }); - return; + buffered.push(record); + this.#drainOnceWaitersFromBuffer(key); + } else { + this.#advanceLastDispatched(key, record.seqNum); } - let buffered = this.buffer.get(key); - if (!buffered) { - buffered = []; - this.buffer.set(key, buffered); + await settled; + } + + #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 }); } - buffered.push(data); } /** @@ -285,26 +418,49 @@ export class TestSessionStreamManager implements SessionStreamManager { * Wrapped per-handler so a throwing/rejecting handler doesn't poison * Promise.all and break unrelated test state. */ - async #invokeHandlers(key: string, data: unknown): Promise { + async #invokeHandlers(key: string, record: SessionStreamRecord): Promise { + const { consumed, settled } = this.#invokeHandlersSync(key, record); + await settled; + return consumed; + } + + /** + * Decide consumption synchronously, exactly like the production dispatch, + * and hand back a promise for any async handler work so callers can still + * await it. Splitting the decision from the awaiting is what keeps a handler + * registered mid-dispatch from seeing an inconsistent buffer. + */ + #invokeHandlersSync( + key: string, + record: SessionStreamRecord + ): { consumed: boolean; settled: Promise } { const handlers = this.handlers.get(key); - if (!handlers || handlers.size === 0) return false; + if (!handlers || handlers.size === 0) { + return { consumed: false, settled: Promise.resolve() }; + } let consumed = false; - await Promise.all( - Array.from(handlers).map(async (h) => { - try { - const result = h(data); - if (result === true) { - consumed = true; - return; - } - await result; - } catch { - // Never let a handler error break test state + const pending: Array> = []; + for (const handler of Array.from(handlers)) { + try { + const result = this.#callHandler(handler, record); + if (result === true) { + consumed = true; + continue; } - }) - ); - return consumed; + if (result) pending.push(Promise.resolve(result).catch(() => {})); + } catch { + continue; + } + } + return { consumed, settled: Promise.all(pending) }; + } + + #callHandler( + handler: RegisteredHandler, + record: SessionStreamRecord + ): void | boolean | Promise { + return handler.kind === "record" ? handler.fn(record) : handler.fn(record.data); } /** diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index c241930323b..a44dd210cfe 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -13,7 +13,6 @@ import { type inferSchemaIn, type inferSchemaOut, InputStreamOncePromise, - type InputStreamOnceResult, isAdditionalApiKey, isSchemaZodEsque, logger, @@ -26,8 +25,11 @@ import { resourceCatalog, type SessionTriggerConfig, SemanticInternalAttributes, + SESSION_IN_CONSUMED_ID_HEADER, SESSION_IN_EVENT_ID_HEADER, sessionStreams, + SessionChannelRouter, + InputStreamTimeoutError, taskContext, type TaskIdentifier, type TaskOptions, @@ -36,6 +38,8 @@ import { type TaskWithSchema, TRIGGER_CONTROL_SUBTYPE, type StreamWriteResult, + type RouterCheckpoint, + type SessionRouteTable, } from "@trigger.dev/core/v3"; import type { FinishReason, @@ -222,53 +226,6 @@ async function findLatestSessionInCursor(chatId: string): Promise { - return findLatestSessionInCursor(chatId); -} - -/** - * Seed the `.in` resume cursor for custom-agent loops (`chat.customAgent` - * raw loops and `chat.createSession`) the way `chat.agent`'s boot does. - * - * MUST run before anything attaches a `.in` listener (`createStopSignal`, - * `chat.messages.on`, the first wait): attaching opens the SSE tail with - * `Last-Event-ID` from the seeded cursor, so attach-then-seed replays - * every record from seq 0 — already-answered user messages get delivered - * into the new run's first wait and the loop re-answers them. - * - * Seeds both cursors: `setLastSeqNum` controls the SSE `Last-Event-ID`, - * `setLastDispatchedSeqNum` gates waiter dispatch — seeding only the - * former still re-delivers records the manager buffered before the seed. - * - * No-ops on fresh boots and when a cursor is already seeded (e.g. the - * `chatCustomAgent` wrapper ran before a nested `createChatSession`). - * @internal - */ -async function seedSessionInResumeCursorForCustomLoop( - payload: Pick -): Promise { - if (sessionStreams.lastSeqNum(payload.chatId, "in") !== undefined) return; - // No continuation/attempt gate: the wire may omit `continuation` on a - // run that still has prior turns (chat.agent covers that case via its - // snapshot). The scan doubles as the prior-state probe — a fresh - // session has no turn-complete on `.out`, returns no cursor, and - // seeds nothing. Cost on fresh boots is one non-blocking records read. - try { - const cursor = await findLatestSessionInCursor(payload.chatId); - if (cursor !== undefined) { - sessionStreams.setLastSeqNum(payload.chatId, "in", cursor); - sessionStreams.setLastDispatchedSeqNum(payload.chatId, "in", cursor); - } - } catch (error) { - logger.warn("chat session: session.in resume cursor lookup failed; old messages may replay", { - error: error instanceof Error ? error.message : String(error), - }); - } -} - /** * Versioned blob written to S3 after every turn completes (when no * `hydrateMessages` hook is registered). Read at run boot to seed the @@ -1543,107 +1500,189 @@ 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; +}; + +/** + * Read one record from a route, suspending the run if nothing is there yet. + * + * The wake and the read are separate steps: the channel wakes the run, then the + * router hands over whatever it routed. Nothing else can take the record in + * between, which is what keeps the published cursors and the delivered record + * in agreement. + * @internal + */ +async function waitOnChatRoute( + route: string, + options: { + idleTimeoutInSeconds?: number; + timeout?: string; + spanName?: string; + skipSuspend?: boolean; + onSuspend?: () => Promise | void; + onResume?: () => Promise | void; + } +): Promise<{ ok: true; output: T } | { ok: false; error?: Error }> { + const router = chatInputRouter(); + const session = getChatSession(); + + return tracer.startActiveSpan( + options.spanName ?? `chat.${route}.wait()`, + async (span) => { + const idleMs = (options.idleTimeoutInSeconds ?? 0) * 1000; + if (idleMs > 0) { + const warm = await router.next(route, { timeoutMs: idleMs }); + if (warm) { + span.setAttribute("wait.resolved", "idle"); + return { ok: true as const, output: warm.data as T }; + } + } else { + const buffered = await router.next(route, { timeoutMs: 0 }); + if (buffered) { + span.setAttribute("wait.resolved", "buffered"); + return { ok: true as const, output: buffered.data as T }; + } + } + + if (options.skipSuspend) { + span.setAttribute("wait.resolved", "skipped"); + return { + ok: false as const, + error: new Error("Idle timeout elapsed and skipSuspend is set"), + }; + } + + if (options.onSuspend) await options.onSuspend(); + + span.setAttribute("wait.resolved", "suspended"); + while (true) { + const wake = await session.in.awaitWake({ + timeout: options.timeout, + lastSeqNum: router.resumeFloor(), + }); + if (!wake.ok) { + span.recordException(wake.error); + return { ok: false as const, error: wake.error }; + } + + const record = await router.next(route); + if (!record) continue; + + if (options.onResume) await options.onResume(); + return { ok: true as const, output: record.data as T }; + } + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "sessions", + session: session.id, + io: "in", + route, + ...accessoryAttributes({ + items: [{ text: `${session.id}.in:${route}`, variant: "normal" }], + style: "codepath", + }), + }, + } + ); +} + +const messagesInput: ChatMessages = { id: "chat-messages", on(handler) { - return getChatSession().in.on((chunk) => { - if (chunk.kind === "message") { - // Returning `true` marks the record CONSUMED at the manager level: - // it is neither buffered for a later `once()` nor re-delivered by - // the buffer drain when the next turn re-attaches its handler. - // Without this, a message arriving mid-stream was delivered twice - // and ran a duplicate turn. - void Promise.resolve(handler(chunk.payload)).catch(() => {}); - return true; - } - return undefined; + return chatInputRouter().on(CHAT_ROUTE_MESSAGES, (record) => { + const chunk = record.data as Extract; + void Promise.resolve(handler(chunk.payload)).catch(() => {}); }); }, once(options) { - const ctx = taskContext.ctx; - const runId = ctx?.run.id; - return new InputStreamOncePromise((resolve, reject) => { - tracer - .startActiveSpan( - options?.spanName ?? `chat.messages.once()`, - async () => { - while (true) { - const result = await getChatSession().in.once(options); - if (!result.ok) { - resolve(result as InputStreamOnceResult); - return; - } - if (result.output.kind === "message") { - resolve({ ok: true, output: result.output.payload }); - return; - } - // Non-message chunks (stops) are handled by the stopInput - // facade's persistent listener; loop and wait for the next. - } - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "streams", - [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream", - ...(runId - ? { - [SemanticInternalAttributes.ENTITY_ID]: `${runId}:chat-messages`, - } - : {}), - streamId: "chat-messages", - ...accessoryAttributes({ - items: [{ text: "chat-messages", variant: "normal" }], - style: "codepath", - }), - }, + chatInputRouter() + .next(CHAT_ROUTE_MESSAGES, { timeoutMs: options?.timeoutMs }) + .then((record) => { + if (!record) { + resolve({ + ok: false, + error: new InputStreamTimeoutError("chat-messages", options?.timeoutMs ?? 0), + }); + return; } - ) - .catch(reject); + const chunk = record.data as Extract; + resolve({ ok: true, output: chunk.payload }); + }, reject); }); }, peek() { - const chunk = getChatSession().in.peek(); - if (chunk && chunk.kind === "message") return chunk.payload; - return undefined; + const record = chatInputRouter().peek(CHAT_ROUTE_MESSAGES); + if (!record) return undefined; + return (record.data as Extract).payload; + }, + async hasPending() { + return chatInputRouter().hasPending(CHAT_ROUTE_MESSAGES); + }, + 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 record = await chatInputRouter().next(CHAT_ROUTE_MESSAGES, { + timeoutMs: timeoutInSeconds === undefined ? undefined : timeoutInSeconds * 1000, + }); + if (!record) return undefined; + + const chunk = record.data as Extract; + return { id: record.id, seqNum: record.seqNum, payload: chunk.payload }; }, wait(options) { return new ManualWaitpointPromise(async (resolve, reject) => { try { - while (true) { - const result = await getChatSession().in.wait(options); - if (!result.ok) { - resolve(result); - return; - } - if (result.output.kind === "message") { - resolve({ ok: true, output: result.output.payload }); - return; - } - // Stop chunks are handled by the stopInput facade's persistent - // listener; loop back into the suspending wait. - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_MESSAGES, + { timeout: options?.timeout, spanName: options?.spanName } + ); + resolve( + result.ok + ? { ok: true, output: result.output.payload } + : { ok: false, error: result.error ?? new Error("Timed out") } + ); } catch (error) { reject(error); } }); }, async waitWithIdleTimeout(options) { - while (true) { - const result = await getChatSession().in.waitWithIdleTimeout(options); - if (!result.ok) return result; - if (result.output.kind === "message") { - return { ok: true, output: result.output.payload }; - } - // Swallow stop-kind chunks — persistent stop listener already handled - // the abort; we just loop for the next message. - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_MESSAGES, + options + ); + return result.ok + ? { ok: true, output: result.output.payload } + : { ok: false, error: result.error }; }, async send(_runId, data, options) { - // The `runId` argument is kept for signature parity with - // `RealtimeDefinedInputStream` but ignored — sessions are addressed - // by sessionId, not runId. Callers producing messages from outside - // the run should prefer the transport's `session.in.send(...)` path. await getChatSession().in.send( { kind: "message", payload: data } satisfies ChatInputChunk, options?.requestOptions @@ -1654,98 +1693,59 @@ const messagesInput: RealtimeDefinedInputStream = { const stopInput: RealtimeDefinedInputStream<{ stop: true; message?: string }> = { id: "chat-stop", on(handler) { - return getChatSession().in.on((chunk) => { - if (chunk.kind === "stop") { - // Consume stop records (see the messages facade above). A stop is - // only meaningful to the turn it interrupts — buffering it would - // let a stale stop abort a future turn. - void Promise.resolve(handler({ stop: true, message: chunk.message })).catch(() => {}); - return true; - } - return undefined; + return chatInputRouter().on(CHAT_ROUTE_STOP, (record) => { + const chunk = record.data as Extract; + void Promise.resolve(handler({ stop: true, message: chunk.message })).catch(() => {}); }); }, once(options) { - const ctx = taskContext.ctx; - const runId = ctx?.run.id; - return new InputStreamOncePromise<{ stop: true; message?: string }>((resolve, reject) => { - tracer - .startActiveSpan( - options?.spanName ?? `chat.stop.once()`, - async () => { - while (true) { - const result = await getChatSession().in.once(options); - if (!result.ok) { - resolve(result as InputStreamOnceResult<{ stop: true; message?: string }>); - return; - } - if (result.output.kind === "stop") { - resolve({ - ok: true, - output: { stop: true, message: result.output.message }, - }); - return; - } - } - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "streams", - [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream", - ...(runId - ? { - [SemanticInternalAttributes.ENTITY_ID]: `${runId}:chat-stop`, - } - : {}), - streamId: "chat-stop", - ...accessoryAttributes({ - items: [{ text: "chat-stop", variant: "normal" }], - style: "codepath", - }), - }, + chatInputRouter() + .next(CHAT_ROUTE_STOP, { timeoutMs: options?.timeoutMs }) + .then((record) => { + if (!record) { + resolve({ + ok: false, + error: new InputStreamTimeoutError("chat-stop", options?.timeoutMs ?? 0), + }); + return; } - ) - .catch(reject); + const chunk = record.data as Extract; + resolve({ ok: true, output: { stop: true, message: chunk.message } }); + }, reject); }); }, peek() { - const chunk = getChatSession().in.peek(); - if (chunk && chunk.kind === "stop") { - return { stop: true, message: chunk.message }; - } - return undefined; + const record = chatInputRouter().peek(CHAT_ROUTE_STOP); + if (!record) return undefined; + const chunk = record.data as Extract; + return { stop: true, message: chunk.message }; }, wait(options) { return new ManualWaitpointPromise<{ stop: true; message?: string }>(async (resolve, reject) => { try { - while (true) { - const result = await getChatSession().in.wait(options); - if (!result.ok) { - resolve(result); - return; - } - if (result.output.kind === "stop") { - resolve({ - ok: true, - output: { stop: true, message: result.output.message }, - }); - return; - } - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_STOP, + { timeout: options?.timeout, spanName: options?.spanName } + ); + resolve( + result.ok + ? { ok: true, output: { stop: true, message: result.output.message } } + : { ok: false, error: result.error ?? new Error("Timed out") } + ); } catch (error) { reject(error); } }); }, async waitWithIdleTimeout(options) { - while (true) { - const result = await getChatSession().in.waitWithIdleTimeout(options); - if (!result.ok) return result; - if (result.output.kind === "stop") { - return { ok: true, output: { stop: true, message: result.output.message } }; - } - } + const result = await waitOnChatRoute>( + CHAT_ROUTE_STOP, + options + ); + return result.ok + ? { ok: true as const, output: { stop: true, message: result.output.message } } + : { ok: false as const, error: result.error }; }, async send(_runId, data, options) { await getChatSession().in.send( @@ -1793,16 +1793,8 @@ const handoverInput = { spanName?: string; skipSuspend?: boolean; }) { - while (true) { - const result = await getChatSession().in.waitWithIdleTimeout(options); - if (!result.ok) return result; - if (result.output.kind === "handover" || result.output.kind === "handover-skip") { - return { ok: true as const, output: result.output as HandoverSignal }; - } - // Other kinds (message, stop) are not expected during handover-prepare. - // Loop back; the message and stop facades have their own listeners - // running so signals on those kinds aren't lost. - } + const result = await waitOnChatRoute(CHAT_ROUTE_HANDOVER, options); + return result.ok ? { ok: true as const, output: result.output } : result; }, }; @@ -1819,9 +1811,8 @@ const handoverInput = { * For the common case prefer `accumulator.consumeHandover()`, which also seeds * `payload.headStartMessages` and applies the partial for you. * - * Must be called at turn 0 before any `chat.messages.waitWithIdleTimeout` — - * that facade consumes and discards non-message chunks, which would swallow the - * handover signal. + * Safe to call at any point in turn 0: the handover signal has its own route, + * so a message facade waiting at the same time cannot take it. */ async function waitForHandover(options: { /** The run's wire payload (only `trigger` / `idleTimeoutInSeconds` are read). */ @@ -1831,15 +1822,257 @@ 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 { + chatInputRouter().clearRoute(CHAT_ROUTE_HANDOVER); + } +} + +/** + * Everything `session.in` carries, and what happens to each kind. + * + * A route's two properties are what make the resume protocol derivable rather + * than hand-maintained. `messages` is replayable because losing a user message + * is data loss. `stop` is neither queued nor replayable: it only means anything + * to the turn that is live when it lands, and a replayed one would abort + * whichever turn happened to be running. `handover` is queued but not + * replayable, because it can arrive before its consumer is ready yet is + * meaningless to any later boot. + * + * A kind absent from this table has no consumer, so the router discards it + * instead of letting it park at the head of a queue. + * @internal + */ +const CHAT_INPUT_ROUTES: SessionRouteTable = { + kindOf: (data) => (data as ChatInputChunk | undefined)?.kind, + routes: [ + { name: "messages", delivery: "queue", replayable: true, kinds: ["message"] }, + { name: "stop", delivery: "at-arrival", replayable: false, kinds: ["stop"] }, + { + name: "handover", + delivery: "queue", + replayable: false, + kinds: ["handover", "handover-skip"], + }, + ], +}; + +const CHAT_ROUTE_MESSAGES = "messages"; +const CHAT_ROUTE_STOP = "stop"; +const CHAT_ROUTE_HANDOVER = "handover"; + +/** + * The `.in` router for the run this worker is currently serving. + * + * One slot rather than a map, because the facades have to reach the same router + * the boot attached without depending on a locals scope being active. Tagged + * with the run as well as the chat: a warm process is reused across runs and the + * executor tears the channel subscription down at the end of each one, so + * reusing a router across runs would leave the new run with no input at all. A + * nested `chat.createSession` within the same run still shares it. + * @internal + */ +let currentChatInputRouter: + | { chatId: string; runId: string | undefined; router: SessionChannelRouter; attached: boolean } + | undefined; + +/** + * Both cursors from the latest `turn-complete` on `.out`, in one scan. + * + * Absent for a chat whose turns predate the headers, in which case the router + * starts from the beginning of the channel and nothing is treated as replayed. + * @internal + */ +async function findLatestSessionInCheckpoint(chatId: string): Promise { + const apiClient = apiClientManager.clientOrThrow(); + const response = await apiClient.readSessionStreamRecords(chatId, "out"); + const checkpoint: RouterCheckpoint = {}; + for (const record of response.records) { + if (controlSubtype(record.headers) !== TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE) continue; + const resumeFrom = Number.parseInt( + headerValue(record.headers, SESSION_IN_EVENT_ID_HEADER) ?? "", + 10 + ); + if (Number.isFinite(resumeFrom)) checkpoint.resumeFrom = resumeFrom; + const appliedThrough = Number.parseInt( + headerValue(record.headers, SESSION_IN_CONSUMED_ID_HEADER) ?? "", + 10 + ); + if (Number.isFinite(appliedThrough)) checkpoint.appliedThrough = appliedThrough; + } + return checkpoint; +} + +/** + * The highest `.in` sequence already on the channel at boot, above the floor. + * + * Everything already on the channel when a run boots is, by definition, not + * arriving live on this run, so the channel's own tail is the end of this run's + * replay window. + * + * The turn boundary's value is not enough on its own. A boundary is written when + * a turn ends, so a control record that arrived after the last boundary is not + * covered by it, and a boundary written by an older SDK does not carry one at + * all. Both cases leave an already-applied record looking live. + * + * Bounded by `afterEventId` when a floor is known, which is the common case, so + * the read covers the replay window rather than the conversation. An absent + * floor means no boundary has committed a cursor yet, and the channel is still + * short. + * + * Returns `undefined` if the read fails; the caller then falls back to the + * floor, which is the previous release's behaviour. + * @internal + */ +async function findSessionInReplayWindowEnd( + chatId: string, + afterSeqNum: number | undefined +): Promise { + try { + const apiClient = apiClientManager.clientOrThrow(); + const response = await apiClient.readSessionStreamRecords(chatId, "in", { + ...(afterSeqNum === undefined ? {} : { afterEventId: String(afterSeqNum) }), + }); + let highest: number | undefined; + for (const record of response.records) { + const seqNum = typeof record.seqNum === "number" ? record.seqNum : Number.NaN; + if (Number.isFinite(seqNum) && (highest === undefined || seqNum > highest)) { + highest = seqNum; + } + } + return highest; + } catch { + return undefined; + } +} + +/** + * Attach the `.in` router for this run. + * + * Reads the checkpoint and subscribes in one call, so there is no window in + * which a listener is attached before the resume cursor is seeded. Attaching + * first would open the tail at sequence 0 and replay every record the previous + * run already answered, which is a mistake the previous shape of this code made + * possible and this shape does not. + * + * The router consumes every record at dispatch, so the channel's own buffer + * stays empty for chat and its cursor bookkeeping never engages. Delivery, + * ordering and the published cursors are the router's, entirely. + * @internal + */ +async function installChatInputRouter( + chatId: string, + options?: { fallbackResumeFrom?: number; resuming?: boolean } +): Promise { + const entry = chatInputRouterEntry(chatId); + if (entry.attached) return entry.router; + + let checkpoint: RouterCheckpoint = {}; + try { + checkpoint = await findLatestSessionInCheckpoint(chatId); + } catch (error) { + logger.warn("chat session: session.in resume cursor lookup failed; old messages may replay", { + error: error instanceof Error ? error.message : String(error), + }); + } + if (checkpoint.resumeFrom === undefined && options?.fallbackResumeFrom !== undefined) { + checkpoint.resumeFrom = options.fallbackResumeFrom; + } + // Only a resuming run has a replay window. On a first boot nothing has been + // applied by anyone, so treating what is already on the channel as replayed + // would discard a signal that arrived before the agent got here, which is + // exactly how a head-start handover reaches a cold run. + const resuming = checkpoint.resumeFrom !== undefined || options?.resuming === true; + if (resuming) { + const replayWindowEnd = await findSessionInReplayWindowEnd(chatId, checkpoint.resumeFrom); + if (replayWindowEnd !== undefined) { + checkpoint.appliedThrough = Math.max( + checkpoint.appliedThrough ?? replayWindowEnd, + replayWindowEnd + ); + } + } + + const router = entry.router; + router.restore(checkpoint); + + const floor = router.resumeFrom(); + if (floor !== undefined) { + sessionStreams.setLastSeqNum(chatId, "in", floor); + sessionStreams.setLastDispatchedSeqNum(chatId, "in", floor); + } + + sessionStreams.onRecord(chatId, "in", (record) => { + router.ingest(record); + return true; }); - // Non-ok = idle timeout or the warm handler crashed without signaling. - if (!result.ok) return null; - return result.output; + + entry.attached = true; + return router; +} + +function chatInputRouterEntry(chatId: string): { + chatId: string; + runId: string | undefined; + router: SessionChannelRouter; + attached: boolean; +} { + const runId = taskContext.ctx?.run.id; + if (currentChatInputRouter?.chatId === chatId && currentChatInputRouter.runId === runId) { + return currentChatInputRouter; + } + + currentChatInputRouter = { + chatId, + runId, + router: new SessionChannelRouter(CHAT_INPUT_ROUTES, { + onDrop: (record, reason) => { + if (reason === "unroutable" || reason === "malformed") { + logger.warn("chat: discarded a session.in record nothing on this worker can consume", { + reason, + seqNum: record.seqNum, + }); + } + }, + }), + attached: false, + }; + return currentChatInputRouter; +} + +/** Drop the router so the next boot attaches a fresh one. @internal */ +export function __resetChatInputRouterForTests(): void { + currentChatInputRouter = undefined; +} + +/** The cursors this run would publish on its next turn boundary. @internal */ +export function __chatInputCheckpointForTests(): RouterCheckpoint { + return currentChatInputRouter?.router.checkpoint() ?? {}; +} + +/** Test-only entry point for the turn-boundary cursor scan. @internal */ +export async function __findLatestSessionInCheckpointForTests( + chatId: string +): Promise { + return findLatestSessionInCheckpoint(chatId); +} + +/** + * This chat's router. Created on first use so a facade reached before the + * install still shares the one the install will attach. + * @internal + */ +function chatInputRouter(): SessionChannelRouter { + return chatInputRouterEntry(getChatSession().id).router; } /** @@ -5370,10 +5603,9 @@ function chatCustomAgent< markChatAgentRunForStreamsWarning(); taskContext.setConversationId(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); + await installChatInputRouter(payload.chatId, { + resuming: Boolean(payload.continuation), + }); return userRun(payload, runOptions); }, }); @@ -5718,53 +5950,16 @@ function chatAgent< ); } - // ── session.in resume cursor ─────────────────────────────────── - // - // A fresh worker subscribes to `session.in` from seq 0 and would - // re-deliver every record ever appended — including user messages - // from turns already completed on a prior run. Without a cursor, - // the loop would re-process them as fresh turns and the slim-wire - // merge would replace-by-id against snapshot-restored copies, - // yielding no-op replaces while the customer's actual new message - // waits in the queue. + // ── session.in router ────────────────────────────────────────── // - // The cursor is the seq_num of the last `.in` record the prior - // worker committed to processing, persisted on each `turn-complete` - // control record as a `session-in-event-id` sibling header. The - // boot scan reads the header off `.out`'s latest turn-complete and - // seeds the manager so the upcoming `.in` SSE subscribe opens with - // `Last-Event-ID: ` — S2 starts after that seq and old - // messages never reach this worker. - // - // Applies in three cases (any of which means `.in` has records - // belonging to completed turns the new run should skip): - // - OOM retry (`ctx.attempt.number > 1`) - // - Continuation run (`payload.continuation === true`) — prior run - // crashed / was canceled / requested upgrade - // - Snapshot exists at all (catches edge cases where the wire - // didn't set `continuation` but a snapshot indicates prior turns) - const needsResumeCursor = - ctx.attempt.number > 1 || payload.continuation === true || bootSnapshot !== undefined; - - if (needsResumeCursor) { - try { - // Reuse the cursor the boot block already resolved (snapshot - // field or records scan) — only scan here when the boot block - // was skipped (hydrateMessages, or snapshot-only signals). - const cursor = bootInCursorResolved - ? bootInCursor - : await findLatestSessionInCursor(payload.chatId); - if (cursor !== undefined) { - sessionStreams.setLastSeqNum(payload.chatId, "in", cursor); - sessionStreams.setLastDispatchedSeqNum(payload.chatId, "in", cursor); - } - } catch (error) { - logger.warn( - "chat.agent: session.in resume cursor lookup failed; old messages may replay", - { error: error instanceof Error ? error.message : String(error) } - ); - } - } + // Reads the turn boundary and subscribes in one call. `bootInCursor` is + // only a fallback: the boot block above may already have resolved a + // cursor from the snapshot, which is used when the boundary itself + // carries none. + await installChatInputRouter(payload.chatId, { + fallbackResumeFrom: bootInCursorResolved ? bootInCursor : undefined, + resuming: Boolean(payload.continuation) || ctx.attempt.number > 1, + }); // ── Recovery boot + chain reconstruction ──────────────────────── if (!hydrateMessages) { @@ -7724,7 +7919,7 @@ function chatAgent< await tracer.startActiveSpan( "snapshot.write", async () => { - const snapshotInCursor = getChatSession().in.lastDispatchedSeqNum(); + const snapshotInCursor = chatInputRouter().resumeFloor(); await writeChatSnapshot(sessionIdForSnapshot, { version: 1, savedAt: Date.now(), @@ -8057,7 +8252,7 @@ function chatAgent< // neither the snapshot nor the replayable `.in` tail. if (!hydrateMessages) { try { - const errorSnapshotInCursor = getChatSession().in.lastDispatchedSeqNum(); + const errorSnapshotInCursor = chatInputRouter().resumeFloor(); await writeChatSnapshot(sessionIdForSnapshot, { version: 1, savedAt: Date.now(), @@ -8926,10 +9121,14 @@ function createStopSignal(): { * task instead of round-tripping them back from the client: * - `lastEventId` — the turn-complete control record's seq_num on * `session.out`; where the next turn's output stream resumes. - * - `sessionInEventId` — the committed-consume cursor on `session.in` as of - * this turn-complete, letting a raw loop correlate the boundary with the - * exact input record it acknowledged. Trigger owns input-cursor recovery, - * so this is for correlation / out-of-sync detection, not required. + * - `sessionInEventId` — the safe-to-resume-from cursor on `session.in` as of + * this turn-complete. It is the highest sequence that can be resumed past + * without skipping an unhandled message, so it is held back behind any + * message still buffered unconsumed and is NOT necessarily the sequence of + * the record this turn answered. Trigger owns input-cursor recovery, so this + * is for correlation / out-of-sync detection, not required. Treat it as a + * lower bound: a value below the record you just handled is expected, not a + * sign of a lost turn. * * Either is `undefined` when the corresponding cursor isn't available. * @@ -8946,7 +9145,7 @@ async function chatWriteTurnComplete(options?: { const result = await writeTurnCompleteChunk(undefined, options?.publicAccessToken); // Same cursor written to the `session-in-event-id` header inside // `writeTurnCompleteChunk`; surfaced here so the caller can persist it. - const inCursor = getChatSession().in.lastDispatchedSeqNum(); + const inCursor = chatInputRouter().resumeFloor(); return { lastEventId: result?.lastEventId, ...(inCursor !== undefined ? { sessionInEventId: String(inCursor) } : {}), @@ -9590,7 +9789,9 @@ function createChatSession( activeMsgSub = undefined; if (!booted) { booted = true; - await seedSessionInResumeCursorForCustomLoop(currentPayload); + await installChatInputRouter(currentPayload.chatId, { + resuming: Boolean(currentPayload.continuation), + }); stop = createStopSignal(); } turn++; @@ -10796,6 +10997,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. + chatInputRouter().clearRoute(CHAT_ROUTE_HANDOVER); + // 1. Write the turn-complete control record. The ack's `lastEventId` is // this record's seq_num — that's the trim target for the NEXT turn. // @@ -10810,10 +11022,15 @@ async function writeTurnCompleteChunk( if (publicAccessToken) { extraHeaders.push(["public-access-token", publicAccessToken]); } - const inCursor = session.in.lastDispatchedSeqNum(); + const routerCheckpoint = chatInputRouter().checkpoint(); + const inCursor = routerCheckpoint.resumeFrom; if (inCursor !== undefined) { extraHeaders.push([SESSION_IN_EVENT_ID_HEADER, String(inCursor)]); } + const consumedCursor = routerCheckpoint.appliedThrough; + if (consumedCursor !== undefined) { + extraHeaders.push([SESSION_IN_CONSUMED_ID_HEADER, String(consumedCursor)]); + } const result = await session.out.writeControl( TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE, extraHeaders diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index 8a01f8293c4..9758d534f21 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 @@ -702,79 +702,121 @@ export class SessionInputChannel { * run-engine waitpoint holds the run until the session append handler * fires it. Only callable from inside `task.run()`. */ + /** + * Suspend until the channel wakes this run, and read nothing. + * + * The waitpoint is only a wake signal: the append route commits the record to + * the channel before it drains any waitpoint, so once this resolves the + * record is durably readable from the channel itself, carrying its real + * sequence. Separating the wake from the read is what lets a consumer that + * owns its own delivery (the chat input router) reuse this without the + * channel also taking a record out from under it. + * + * @internal + */ + async awaitWake( + options?: InputStreamWaitOptions & { lastSeqNum?: number } + ): Promise<{ ok: true; waitpointId: string } | { ok: false; error: Error }> { + const ctx = taskContext.ctx; + + if (!ctx) { + throw new Error("session.in.wait() can only be used from inside a task.run()"); + } + + const apiClient = apiClientManager.clientOrThrow(); + + const lastConsumedSeqNum = + options?.lastSeqNum ?? sessionStreams.lastDispatchedSeqNum(this.sessionId, "in"); + const response = await apiClient.createSessionStreamWaitpoint(ctx.run.id, { + session: this.sessionId, + io: "in", + timeout: options?.timeout, + idempotencyKey: options?.idempotencyKey, + idempotencyKeyTTL: options?.idempotencyKeyTTL, + tags: options?.tags, + lastSeqNum: lastConsumedSeqNum, + }); + + const waitResponse = await apiClient.waitForWaitpointToken({ + runFriendlyId: ctx.run.id, + waitpointFriendlyId: response.waitpointId, + }); + + if (!waitResponse.success) { + throw new Error("Failed to block on session stream waitpoint"); + } + + sessionStreams.disconnectStream(this.sessionId, "in"); + + const waitResult = await runtime.waitUntil(response.waitpointId); + + if (!waitResult.ok) { + const parsed = + waitResult.output !== undefined + ? await conditionallyImportAndParsePacket( + { + data: waitResult.output, + dataType: waitResult.outputType ?? "application/json", + }, + apiClient + ) + : undefined; + return { + ok: false as const, + error: new WaitpointTimeoutError(parsed?.message ?? "Timed out"), + }; + } + + sessionStreams.reconnectStream(this.sessionId, "in"); + return { ok: true as const, waitpointId: response.waitpointId }; + } + wait(options?: InputStreamWaitOptions): ManualWaitpointPromise { return new ManualWaitpointPromise(async (resolve, reject) => { try { - const ctx = taskContext.ctx; - - if (!ctx) { - throw new Error("session.in.wait() can only be used from inside a task.run()"); - } - const apiClient = apiClientManager.clientOrThrow(); - const response = await apiClient.createSessionStreamWaitpoint(ctx.run.id, { - session: this.sessionId, - io: "in", - timeout: options?.timeout, - idempotencyKey: options?.idempotencyKey, - idempotencyKeyTTL: options?.idempotencyKeyTTL, - tags: options?.tags, - lastSeqNum: sessionStreams.lastSeqNum(this.sessionId, "in"), - }); - const result = await tracer.startActiveSpan( options?.spanName ?? `sessions.open(${this.sessionId}).in.wait()`, async (span) => { - const waitResponse = await apiClient.waitForWaitpointToken({ - runFriendlyId: ctx.run.id, - waitpointFriendlyId: response.waitpointId, - }); + const wake = await this.awaitWake(options); - if (!waitResponse.success) { - throw new Error("Failed to block on session stream waitpoint"); + if (!wake.ok) { + span.recordException(wake.error); + span.setStatus({ code: SpanStatusCode.ERROR }); + return { ok: false as const, error: wake.error }; } - // Drop the SSE tail + buffer before suspending so the record - // delivered via the waitpoint path isn't re-buffered on resume. - 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"); + span.setAttribute(SemanticInternalAttributes.ENTITY_ID, wake.waitpointId); + + 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: { [SemanticInternalAttributes.STYLE_ICON]: "wait", [SemanticInternalAttributes.ENTITY_TYPE]: "waitpoint", - [SemanticInternalAttributes.ENTITY_ID]: response.waitpointId, session: this.sessionId, io: "in", ...accessoryAttributes({ diff --git a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts index a669975f383..63768b9b3f2 100644 --- a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts +++ b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts @@ -4,6 +4,7 @@ import type { LocalsKey } from "@trigger.dev/core/v3"; import { runInMockTaskContext, type MockTaskContextOptions } from "@trigger.dev/core/v3/test"; import { __setSessionOpenImplForTests, __setSessionStartImplForTests } from "../sessions.js"; import { + __resetChatInputRouterForTests, __setReadChatSnapshotImplForTests, __setReplaySessionInTailImplForTests, __setReplaySessionOutTailImplForTests, @@ -390,6 +391,8 @@ export function mockChatAgent( let seededReplayPartial: UIMessage | undefined; let seededSessionInMessages: UIMessage[] = []; + __resetChatInputRouterForTests(); + __setReadChatSnapshotImplForTests((_id: string) => { return seededSnapshot as ChatSnapshotV1 | undefined; }); diff --git a/packages/trigger-sdk/src/v3/test/test-session-handle.ts b/packages/trigger-sdk/src/v3/test/test-session-handle.ts index 945cd231152..5150de87da3 100644 --- a/packages/trigger-sdk/src/v3/test/test-session-handle.ts +++ b/packages/trigger-sdk/src/v3/test/test-session-handle.ts @@ -4,7 +4,7 @@ import type { StreamWriteResult, WriterStreamOptions, } from "@trigger.dev/core/v3"; -import { ensureReadableStream, ManualWaitpointPromise } from "@trigger.dev/core/v3"; +import { ensureReadableStream } from "@trigger.dev/core/v3"; import type { SessionPipeStreamOptions, SessionSubscribeOptions } from "../sessions.js"; import { SessionHandle, SessionInputChannel, SessionOutputChannel } from "../sessions.js"; @@ -29,32 +29,27 @@ class TestSessionInputChannel extends SessionInputChannel { super(sessionId); } - // Override only the `wait` path. `on` / `once` / `peek` / `send` - // continue to flow through the real `sessionStreams` global, which - // the mock task context installs as a `TestSessionStreamManager`. - wait(): ManualWaitpointPromise { - return new ManualWaitpointPromise( - (resolve: (value: { ok: false; error: Error }) => void) => { - const signal = this.getAbortSignal(); - if (!signal) { - // Harness hasn't wired up its run signal yet — nothing to abort - // on. Stay pending; the run loop should never reach this state - // in practice but we don't want to throw here either. - return; - } - const onAbort = () => { - resolve({ - ok: false, - error: new Error("session.in.wait() aborted by test harness"), - }); - }; - if (signal.aborted) { - onAbort(); - return; - } - signal.addEventListener("abort", onAbort, { once: true }); - } - ); + /** + * Override the one step that talks to the network. Everything built on top + * of it (`wait`, and the chat facades' route waits) then runs its real + * implementation against the in-memory stream manager, so the harness stubs + * a boundary instead of reimplementing a composite. + */ + async awaitWake(): Promise<{ ok: true; waitpointId: string } | { ok: false; error: Error }> { + const signal = this.getAbortSignal(); + if (!signal) { + return new Promise(() => {}); + } + if (signal.aborted) { + return { ok: false, error: new Error("session.in.wait() aborted by test harness") }; + } + return new Promise((resolve) => { + signal.addEventListener( + "abort", + () => resolve({ ok: false, error: new Error("session.in.wait() aborted by test harness") }), + { once: true } + ); + }); } } diff --git a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts new file mode 100644 index 00000000000..cb94ae87f3c --- /dev/null +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -0,0 +1,354 @@ +// 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 { + __chatInputCheckpointForTests as chatInputCheckpoint, + 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 = chatInputCheckpoint().resumeFrom; + observations.afterFirst = await chat.messages.hasPending(); + observations.second = await chat.messages.next(); + observations.cursorAfterSecond = chatInputCheckpoint().resumeFrom; + 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("delivers a message that arrived behind another kind, without losing that kind", async () => { + const chatId = "mailbox-mixed-kinds"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + pending?: boolean; + message?: ChatMessageRecord; + handover?: unknown; + cursorAfter?: 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.message = await chat.messages.next({ timeoutInSeconds: 0 }); + observations.handover = await chat.waitForHandover({ + payload: { trigger: "handover-prepare" }, + idleTimeoutInSeconds: 0, + }); + observations.cursorAfter = chatInputCheckpoint().resumeFrom; + }, + }); + 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; + }); + + // The handover has its own route, so it neither blocks the message behind + // it nor gets destroyed by the consumer that took that message. + expect(observations).toEqual({ + pending: true, + message: { + id: "message-1", + seqNum: 31, + payload: userPayload(chatId, "u-after-handover"), + }, + handover: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + cursorAfter: 31, + }); + }); + + it("holds the resume cursor behind a queued message while a later stop advances the replay window", async () => { + const chatId = "mailbox-cursor-gap"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + cursorBefore?: number; + appliedBefore?: number; + message?: ChatMessageRecord; + cursorAfter?: number; + appliedAfter?: number; + } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-cursor-gap", + run: async () => { + const stop = chat.createStopSignal(); + ready.resolve(); + await inspect.promise; + + observations.cursorBefore = chatInputCheckpoint().resumeFrom; + observations.appliedBefore = chatInputCheckpoint().appliedThrough; + observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); + observations.cursorAfter = chatInputCheckpoint().resumeFrom; + observations.appliedAfter = chatInputCheckpoint().appliedThrough; + 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({ + // Held below the queued message even though the stop after it was applied. + cursorBefore: 49, + appliedBefore: 51, + message: { + id: "message-1", + seqNum: 50, + payload: userPayload(chatId, "u1"), + }, + cursorAfter: 51, + appliedAfter: 51, + }); + }); + + 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/chat-warm-process-reuse.test.ts b/packages/trigger-sdk/test/chat-warm-process-reuse.test.ts new file mode 100644 index 00000000000..266b10dfa5a --- /dev/null +++ b/packages/trigger-sdk/test/chat-warm-process-reuse.test.ts @@ -0,0 +1,60 @@ +import "../src/v3/test/index.js"; + +import { resourceCatalog, sessionStreams } from "@trigger.dev/core/v3"; +import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; +import { describe, expect, it } from "vitest"; +import { chat, type ChatTaskWirePayload } from "../src/v3/ai.js"; + +function userPayload(chatId: string, id: string): ChatTaskWirePayload { + return { + chatId, + trigger: "submit-message", + message: { id, role: "user", parts: [{ type: "text", text: id }] }, + }; +} + +/** + * A worker process is reused across runs. The end of every run tears down the + * channel subscription via `sessionStreams.clearHandlers()`, so a second run of + * the same chat in the same process has to attach a fresh one. Anything cached + * across runs that skips the attach leaves the new run with no input at all, and + * the conversation hangs with no error raised. + */ +describe("chat input across runs in one warm process", () => { + it("delivers messages to a second run of the same chat", async () => { + const chatId = "warm-reuse"; + const seen: string[] = []; + + const agent = chat.customAgent({ + id: "chat-warm-process-reuse", + run: async () => { + const record = await chat.messages.next({ timeoutInSeconds: 2 }); + const part = record?.payload.message?.parts?.[0]; + if (part && part.type === "text") seen.push(part.text); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + for (const attempt of ["first", "second"]) { + await runInMockTaskContext( + async (drivers) => { + const runPromise = run( + { chatId, trigger: "submit-message" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, attempt) }, + "in" + ); + await runPromise; + }, + { ctx: { run: { id: `run_${attempt}` } } } + ); + sessionStreams.clearHandlers(); + } + + expect(seen).toEqual(["first", "second"]); + }); +}); 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) } + ); }); }); diff --git a/packages/trigger-sdk/test/replay-session-in.test.ts b/packages/trigger-sdk/test/replay-session-in.test.ts index 3d04974a9b6..35a06271374 100644 --- a/packages/trigger-sdk/test/replay-session-in.test.ts +++ b/packages/trigger-sdk/test/replay-session-in.test.ts @@ -4,7 +4,7 @@ import "../src/v3/test/index.js"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { apiClientManager } from "@trigger.dev/core/v3"; import { - __findLatestSessionInCursorForTests as findLatestSessionInCursor, + __findLatestSessionInCheckpointForTests as findLatestSessionInCheckpoint, __replaySessionInTailProductionPathForTests as replaySessionInTail, } from "../src/v3/ai.js"; @@ -167,8 +167,8 @@ function stubReadRecordsWithHeaders( return spy; } -describe("findLatestSessionInCursor", () => { - it("returns the LAST turn-complete's session-in-event-id", async () => { +describe("findLatestSessionInCheckpoint", () => { + it("returns the LAST turn-complete's cursors", async () => { const spy = stubReadRecordsWithHeaders([ { data: { type: "text-delta", delta: "hi" } }, { @@ -182,13 +182,15 @@ describe("findLatestSessionInCursor", () => { headers: [ ["trigger-control", "turn-complete"], ["session-in-event-id", "7"], + ["session-in-consumed-id", "9"], ], }, ]); - const cursor = await findLatestSessionInCursor("sess"); - expect(cursor).toBe(7); - // Non-blocking records read on `.out`, no SSE subscribe. + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint).toEqual({ resumeFrom: 7, appliedThrough: 9 }); + // One non-blocking records read on `.out` covers both cursors. + expect(spy).toHaveBeenCalledTimes(1); expect(spy).toHaveBeenCalledWith("sess", "out"); }); @@ -209,8 +211,22 @@ describe("findLatestSessionInCursor", () => { }, ]); - const cursor = await findLatestSessionInCursor("sess"); - expect(cursor).toBe(4); + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint.resumeFrom).toBe(4); + }); + + it("leaves the replay window absent when only the resume cursor was written", async () => { + stubReadRecordsWithHeaders([ + { + headers: [ + ["trigger-control", "turn-complete"], + ["session-in-event-id", "5"], + ], + }, + ]); + + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint).toEqual({ resumeFrom: 5 }); }); it("returns undefined when records carry no headers (older server)", async () => { @@ -219,7 +235,7 @@ describe("findLatestSessionInCursor", () => { { data: { type: "finish" } }, ]); - const cursor = await findLatestSessionInCursor("sess"); - expect(cursor).toBeUndefined(); + const checkpoint = await findLatestSessionInCheckpoint("sess"); + expect(checkpoint).toEqual({}); }); });