diff --git a/.changeset/tidy-mailboxes-wait.md b/.changeset/tidy-mailboxes-wait.md new file mode 100644 index 0000000000..99013758d3 --- /dev/null +++ b/.changeset/tidy-mailboxes-wait.md @@ -0,0 +1,6 @@ +--- +"@trigger.dev/core": patch +"@trigger.dev/sdk": patch +--- + +Custom agent loops can now inspect pending chat input without consuming it and consume one mailbox record at a time with `chat.messages.hasPending()` and `chat.messages.next()`. Mailbox records include stable identifiers for tracing and redelivery. diff --git a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts index c00ff51b3b..3cb98f5847 100644 --- a/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts +++ b/apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts @@ -1,6 +1,8 @@ import { json } from "@remix-run/server-runtime"; import { CreateSessionStreamWaitpointRequestBody, + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, type CreateSessionStreamWaitpointResponseBody, } from "@trigger.dev/core/v3"; import { WaitpointId } from "@trigger.dev/core/v3/isomorphic"; @@ -125,7 +127,8 @@ const { action, loader } = createActionApiRoute( addressingKey, body.io, result.waitpoint.id, - ttlMs && ttlMs > 0 ? ttlMs : undefined + ttlMs && ttlMs > 0 ? ttlMs : undefined, + body.responseFormat ); // Race-check. If a record landed on the channel before this @@ -155,8 +158,14 @@ const { action, loader } = createActionApiRoute( await engine.completeWaitpoint({ id: result.waitpoint.id, output: { - value: record.data, - type: "application/json", + value: + body.responseFormat === "record-v1" + ? serializeSessionStreamWaitpointRecord(record.data, record.seqNum) + : record.data, + type: + body.responseFormat === "record-v1" + ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE + : "application/json", isError: false, }, }); diff --git a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts index d4dd1d9f19..7ff85cde86 100644 --- a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts @@ -15,6 +15,7 @@ import { claimSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, + sessionStreamWaitpointOutput, } from "~/services/sessionStreamWaitpointCache.server"; import { anyResource, createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { engine } from "~/v3/runEngine.server"; @@ -201,7 +202,7 @@ const { action, loader } = createActionApiRoute( // keyed on the canonical addressing key the agent registered with via // `sessions.open(...).in.wait()`, so writers and readers converge // regardless of which URL form they used. - const [drainError, waitpointIds] = await tryCatch( + const [drainError, waitpoints] = await tryCatch( drainSessionStreamWaitpoints(authentication.environment.id, addressingKey, params.io) ); if (drainError) { @@ -210,24 +211,20 @@ const { action, loader } = createActionApiRoute( io: params.io, error: drainError, }); - } else if (waitpointIds && waitpointIds.length > 0) { + } else if (waitpoints && waitpoints.length > 0) { await Promise.all( - waitpointIds.map(async (waitpointId) => { + waitpoints.map(async (waitpoint) => { const [completeError] = await tryCatch( engine.completeWaitpoint({ - id: waitpointId, - output: { - value: part, - type: "application/json", - isError: false, - }, + id: waitpoint.id, + output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq), }) ); if (completeError) { logger.error("Failed to complete session stream waitpoint", { addressingKey, io: params.io, - waitpointId, + waitpointId: waitpoint.id, error: completeError, }); } diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts index ab318f31c7..35bdf3a5dd 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts @@ -13,7 +13,10 @@ import { resolveSessionByIdOrExternalId, } from "~/services/realtime/sessions.server"; import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server"; -import { drainSessionStreamWaitpoints } from "~/services/sessionStreamWaitpointCache.server"; +import { + drainSessionStreamWaitpoints, + sessionStreamWaitpointOutput, +} from "~/services/sessionStreamWaitpointCache.server"; import { requireUserId } from "~/services/session.server"; import { EnvironmentParamSchema } from "~/utils/pathBuilder"; import { engine } from "~/v3/runEngine.server"; @@ -114,7 +117,7 @@ export async function action({ request, params }: ActionFunctionArgs) { // Drain any waitpoints registered for this channel — same as the // public append. Best-effort; failure doesn't fail the append. - const [drainError, waitpointIds] = await tryCatch( + const [drainError, waitpoints] = await tryCatch( drainSessionStreamWaitpoints(environment.id, addressingKey, io) ); if (drainError) { @@ -123,24 +126,20 @@ export async function action({ request, params }: ActionFunctionArgs) { io, error: drainError, }); - } else if (waitpointIds && waitpointIds.length > 0) { + } else if (waitpoints && waitpoints.length > 0) { await Promise.all( - waitpointIds.map(async (waitpointId) => { + waitpoints.map(async (waitpoint) => { const [completeError] = await tryCatch( engine.completeWaitpoint({ - id: waitpointId, - output: { - value: part, - type: "application/json", - isError: false, - }, + id: waitpoint.id, + output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq ?? undefined), }) ); if (completeError) { logger.error("Failed to complete session stream waitpoint (playground)", { addressingKey, io, - waitpointId, + waitpointId: waitpoint.id, error: completeError, }); } diff --git a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts index 7b53042d8d..0c21c10be1 100644 --- a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts +++ b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts @@ -1,5 +1,9 @@ import { Redis } from "ioredis"; import { defaultReconnectOnError } from "@internal/redis"; +import { + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, +} from "@trigger.dev/core/v3"; import { env } from "~/env.server"; import { singleton } from "~/utils/singleton"; import { logger } from "./logger.server"; @@ -13,12 +17,35 @@ import { logger } from "./logger.server"; // is shared — without it, two environments using the same externalId // would drain each other's waitpoints. const KEY_PREFIX = "ssw:"; +const FORMAT_KEY_PREFIX = "sswf:"; const DEFAULT_TTL_MS = 7 * 24 * 60 * 60 * 1000; // 7 days +export type SessionStreamWaitpoint = { + id: string; + responseFormat?: "record-v1"; +}; + +export function sessionStreamWaitpointOutput( + waitpoint: SessionStreamWaitpoint, + data: string, + seqNum: number | undefined +): { value: string; type: string; isError: false } { + const hasRecordEnvelope = waitpoint.responseFormat === "record-v1" && seqNum !== undefined; + return { + value: hasRecordEnvelope ? serializeSessionStreamWaitpointRecord(data, seqNum) : data, + type: hasRecordEnvelope ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE : "application/json", + isError: false, + }; +} + function buildKey(environmentId: string, addressingKey: string, io: "out" | "in"): string { return `${KEY_PREFIX}${environmentId}:${addressingKey}:${io}`; } +function buildFormatKey(waitpointId: string): string { + return `${FORMAT_KEY_PREFIX}${waitpointId}`; +} + // Pre-env-scoping key format, drained for one release so waitpoints from the // previous deploy still wake. Removable once this has been live > turn timeout. function buildLegacyKey(addressingKey: string, io: "out" | "in"): string { @@ -81,13 +108,25 @@ export async function addSessionStreamWaitpoint( addressingKey: string, io: "out" | "in", waitpointId: string, - ttlMs?: number + ttlMs?: number, + responseFormat?: "record-v1" ): Promise { if (!redis) return; try { const key = buildKey(environmentId, addressingKey, io); - await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(ttlMs ?? DEFAULT_TTL_MS)); + const effectiveTtlMs = ttlMs ?? DEFAULT_TTL_MS; + + // Keep the set member as the plain waitpoint id so an older append + // instance can still drain it during a rolling deploy. New instances read + // the optional response format from this separate, TTL-bound key. + if (responseFormat) { + await redis.set(buildFormatKey(waitpointId), responseFormat, "PX", effectiveTtlMs); + } else { + await redis.del(buildFormatKey(waitpointId)); + } + + await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(effectiveTtlMs)); } catch (error) { logger.error("Failed to set session stream waitpoint cache", { environmentId, @@ -107,7 +146,7 @@ export async function drainSessionStreamWaitpoints( environmentId: string, addressingKey: string, io: "out" | "in" -): Promise { +): Promise { if (!redis) return []; try { @@ -129,7 +168,34 @@ export async function drainSessionStreamWaitpoints( if (err || !Array.isArray(members)) continue; for (const m of members as string[]) ids.add(m); } - return [...ids]; + const waitpointIds = [...ids]; + if (waitpointIds.length === 0) return []; + + let formatResults: Awaited> | null = null; + try { + const formatPipeline = redis.multi(); + for (const waitpointId of waitpointIds) { + formatPipeline.get(buildFormatKey(waitpointId)); + formatPipeline.del(buildFormatKey(waitpointId)); + } + formatResults = await formatPipeline.exec(); + } catch (error) { + // The waitpoint ids were already drained. Complete them with raw data + // rather than losing the wake-up because optional metadata was unavailable. + logger.error("Failed to read session stream waitpoint response formats", { + environmentId, + addressingKey, + io, + error, + }); + } + + return waitpointIds.map((id, index) => { + const formatEntry = formatResults?.[index * 2]; + const responseFormat = + formatEntry && !formatEntry[0] && formatEntry[1] === "record-v1" ? "record-v1" : undefined; + return { id, responseFormat }; + }); } catch (error) { logger.error("Failed to drain session stream waitpoint cache", { environmentId, @@ -240,7 +306,10 @@ export async function removeSessionStreamWaitpoint( try { const key = buildKey(environmentId, addressingKey, io); - await redis.srem(key, waitpointId); + const pipeline = redis.multi(); + pipeline.srem(key, waitpointId); + pipeline.del(buildFormatKey(waitpointId)); + await pipeline.exec(); } catch (error) { logger.error("Failed to remove session stream waitpoint cache entry", { environmentId, diff --git a/apps/webapp/app/v3/webhookEngine.server.ts b/apps/webapp/app/v3/webhookEngine.server.ts index d58a89919b..6fbb2d48b9 100644 --- a/apps/webapp/app/v3/webhookEngine.server.ts +++ b/apps/webapp/app/v3/webhookEngine.server.ts @@ -17,6 +17,7 @@ import { claimSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, + sessionStreamWaitpointOutput, } from "~/services/sessionStreamWaitpointCache.server"; import { getSecretStore } from "~/services/secrets/secretStore.server"; import { singleton } from "~/utils/singleton"; @@ -229,10 +230,12 @@ function createWebhookEngine() { "in", deliveryId ); + let appendSeq: number | undefined; if (wonClaim) { - const [appendError] = await tryCatch( + const [appendError, seqNum] = await tryCatch( realtimeStream.appendPartToSessionStream(part, deliveryId, addressingKey, "in") ); + appendSeq = seqNum ?? undefined; if (appendError) { // Nothing landed — release the claim so a retry re-appends the same id. await releaseSessionStreamPart(environment.id, addressingKey, "in", deliveryId); @@ -245,7 +248,7 @@ function createWebhookEngine() { } // Wake any `.in` waitpoints the run registered (best-effort; the record is durable in S2). - const [drainError, waitpointIds] = await tryCatch( + const [drainError, waitpoints] = await tryCatch( drainSessionStreamWaitpoints(environment.id, addressingKey, "in") ); if (drainError) { @@ -253,13 +256,13 @@ function createWebhookEngine() { externalId, error: drainError, }); - } else if (waitpointIds && waitpointIds.length > 0) { + } else if (waitpoints && waitpoints.length > 0) { await Promise.all( - waitpointIds.map((waitpointId) => + waitpoints.map((waitpoint) => tryCatch( runEngine.completeWaitpoint({ - id: waitpointId, - output: { value: part, type: "application/json", isError: false }, + id: waitpoint.id, + output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq), }) ) ) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 197bff6b5e..cb1495fff4 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -213,7 +213,7 @@ For full control, skip `createSession` and compose the primitives directly: | Primitive | Description | | ------------------------------- | -------------------------------------------------------------------------------------------- | -| `chat.messages` | Input stream for incoming messages — use `.waitWithIdleTimeout()` to wait for the next turn | +| `chat.messages` | Mailbox for incoming messages — inspect buffered input, consume one record, or suspend until the next turn | | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | | `chat.pipeAndCapture(result)` | Pipe a stream and capture the response; returns `{ message, status, error }` | | `chat.writeTurnComplete()` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | @@ -221,6 +221,52 @@ For full control, skip `createSession` and compose the primitives directly: | `chat.pipe(stream)` | Pipe a stream to the frontend (no response capture) | | `chat.cleanupAbortedParts(msg)` | Clean up incomplete parts from a stopped response | +### `chat.messages` mailbox + +`chat.messages` exposes the incoming message mailbox for hand-rolled loops: + +| Method | Behavior | +| --- | --- | +| `peek()` | Return the buffer head when it is a message, without consuming it; otherwise return `undefined` | +| `hasPending()` | Resolve `true` when the buffer head is a message; does not consume it | +| `next({ timeoutInSeconds? })` | Consume exactly one message record in channel order, or resolve `undefined` when the optional timeout elapses | +| `on(handler)` | Consume messages as they arrive and invoke the handler | +| `waitWithIdleTimeout(options)` | Wait warm, then suspend the run until the next message arrives | + +`hasPending()` checks whether the local, already-delivered buffer head is a +message that `next()` can consume immediately. It does not query the remote +Session channel or start a subscription. Use `waitWithIdleTimeout()` when the +loop needs to idle until future input arrives. + +`next({ timeoutInSeconds: 0 })` is also a local, non-blocking read. Call +`next()` without a timeout, or with a positive timeout, to subscribe for future +input. + +`next()` returns a readonly record envelope: + +```ts +const record = await chat.messages.next({ timeoutInSeconds: 5 }); +if (record) { + console.log(record.id, record.seqNum); + currentPayload = record.payload; +} +``` + +- `id` is the append's stable idempotency key. +- `seqNum` is the monotonic sequence on this Session's `.in` channel. +- `payload` is the existing `ChatTaskWirePayload` delivered by the other mailbox methods. + +Both identifiers remain the same if the record is delivered again after a +reconnect. Each `next()` call commits only the record it returns, so a loop that +owns its own turn sequencing never advances past input it has not taken. By +contrast, `on()` commits a record as soon as it dispatches the handler; avoid +mixing `on()` and `next()` when a single loop owns mailbox consumption. + +The Session `.in` channel also carries control records such as handovers. If one +comes before a message, `hasPending()` stays `false` and `next()` leaves the +control record for its own consumer. After that record is handled, the message +becomes pending. + A complete loop: ```ts trigger/my-chat-raw.ts diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index da08f9a047..65e98a05aa 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -506,7 +506,7 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`. | `chat.pipeAndCapture(source, options?)` | Pipe and capture the response; returns `{ message, status, error }` | | `chat.writeTurnComplete(options?)` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | -| `chat.messages` | Input stream for incoming messages — use `.waitWithIdleTimeout()` | +| `chat.messages` | Incoming message mailbox; supports non-consuming `.peek()` / `.hasPending()`, single-record `.next()`, `.on()`, and suspend-aware `.waitWithIdleTimeout()` | | `chat.local({ id })` | Create a per-run typed local (see [`chat.local`](/ai-chat/chat-local)) | | `chat.createStartSessionAction(taskId, options?)` | Returns a server action that creates a chat Session + triggers the first run + returns a session-scoped PAT. Idempotent on `(env, externalId)`. | | `chat.waitForHandover(options)` | Wait for a [`chat.headStart`](/ai-chat/fast-starts#handover-with-custom-agents) handover signal in a custom loop. Returns the signal or `null`. `chat.MessageAccumulator` wraps this as `consumeHandover()` / `applyHandover()` | diff --git a/packages/core/src/v3/apiClient/runStream.test.ts b/packages/core/src/v3/apiClient/runStream.test.ts index 3a266f2a91..ee3f3df22a 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 b0d43ef3f9..b01fc6e964 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/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index 6cd100f7c3..42690f79fc 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -1669,6 +1669,8 @@ export const CreateSessionStreamWaitpointRequestBody = z.object({ * Used to catch data that arrived before `.wait()` was called. */ lastSeqNum: z.number().optional(), + /** Internal capability flag: return the exact record sequence on resume. */ + responseFormat: z.literal("record-v1").optional(), }); export type CreateSessionStreamWaitpointRequestBody = z.infer< typeof CreateSessionStreamWaitpointRequestBody diff --git a/packages/core/src/v3/sessionStreams/index.ts b/packages/core/src/v3/sessionStreams/index.ts index 21e2e8d245..7d45c9248f 100644 --- a/packages/core/src/v3/sessionStreams/index.ts +++ b/packages/core/src/v3/sessionStreams/index.ts @@ -1,6 +1,12 @@ import { getGlobal, registerGlobal } from "../utils/globals.js"; import { NoopSessionStreamManager } from "./noopManager.js"; -import type { InputStreamOncePromise, SessionChannelIO, SessionStreamManager } from "./types.js"; +import type { + InputStreamOncePromise, + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "./types.js"; import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; const API_NAME = "session-streams"; @@ -43,10 +49,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 +94,14 @@ export class SessionStreamsAPI implements SessionStreamManager { this.#getManager().setLastSeqNum(sessionId, io, seqNum); } + public consumeRecord(sessionId: string, io: SessionChannelIO, seqNum: number): void { + const manager = this.#getManager(); + if (!manager.consumeRecord) { + throw new Error("The configured Session stream manager does not support exact consumption"); + } + manager.consumeRecord(sessionId, io, seqNum); + } + public lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { return this.#getManager().lastDispatchedSeqNum(sessionId, io); } diff --git a/packages/core/src/v3/sessionStreams/manager.test.ts b/packages/core/src/v3/sessionStreams/manager.test.ts index 9b489616f7..4262674bfe 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 fb87b21164..c4c1c0503a 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -3,7 +3,12 @@ import type { InputStreamOnceResult } from "../inputStreams/types.js"; import { InputStreamOncePromise, InputStreamTimeoutError } from "../inputStreams/types.js"; import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; import { computeReconnectDelayMs } from "../utils/reconnectBackoff.js"; -import type { SessionChannelIO, SessionStreamManager } from "./types.js"; +import type { + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "./types.js"; import { controlSubtype } from "./wireProtocol.js"; // A handler that synchronously returns `true` CONSUMES the record: it is @@ -13,8 +18,9 @@ import { controlSubtype } from "./wireProtocol.js"; type SessionStreamHandler = (data: unknown) => void | boolean | Promise; type OnceWaiter = { - resolve: (result: InputStreamOnceResult) => void; + resolve: (result: InputStreamOnceResult) => void; reject: (error: Error) => void; + predicate?: SessionStreamRecordPredicate; timeoutHandle?: ReturnType; // The abort signal and its handler are tracked on the waiter so any // resolution path (dispatch / timeout / explicit removal) can detach @@ -44,19 +50,7 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { export class StandardSessionStreamManager implements SessionStreamManager { private handlers = new Map>(); private onceWaiters = new Map(); - private buffer = new Map(); - // Parallel to `buffer`: the SSE seq_num of each buffered record. Same - // length and order as `buffer[key]`. Used so that when `once()` shifts - // a buffered record into a waiter, the cursor (`lastDispatchedSeqNums`) - // can advance to that record's seq. Kept as a separate map so the - // existing `peek()` shape (returns `unknown`) stays unchanged. - // - // Entries are `number | undefined` so the array stays length-locked - // with `buffer` even if a record arrives without a parseable seq — - // shifting `undefined` is just a no-op for the cursor advance, but - // the slot still gets consumed. Drifting lengths would map seq_nums - // to the wrong records on subsequent shifts. - private bufferSeqNums = new Map>(); + private buffer = new Map(); private tails = new Map(); // Per-stream lower-bound timestamp filter. When set, records whose // SSE timestamp is <= the bound are dropped before dispatch — used by @@ -72,14 +66,16 @@ 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 @@ -123,28 +119,21 @@ export class StandardSessionStreamManager implements SessionStreamManager { // duplicating turns. const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - const seqList = this.bufferSeqNums.get(key) ?? []; - const keptRecords: unknown[] = []; - // Kept in lock-step with `keptRecords` — drifting lengths would map - // seq_nums to the wrong records on subsequent shifts. - const keptSeqNums: Array = []; - for (let i = 0; i < buffered.length; i++) { - const consumed = this.#invokeHandler(handler, buffered[i]); + const keptRecords: SessionStreamRecord[] = []; + for (const record of buffered) { + const consumed = this.#invokeHandler(handler, record.data); if (consumed) { - const s = seqList[i]; - if (s !== undefined) this.#advanceLastDispatched(key, s); + this.#advanceLastDispatched(key, record.seqNum); } else { - keptRecords.push(buffered[i]); - keptSeqNums.push(seqList[i]); + keptRecords.push(record); } } if (keptRecords.length > 0) { this.buffer.set(key, keptRecords); - this.bufferSeqNums.set(key, keptSeqNums); } else { this.buffer.delete(key); - this.bufferSeqNums.delete(key); } + this.#drainOnceWaitersFromBuffer(key); } return { @@ -162,30 +151,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 +242,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 +281,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 +361,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 +375,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { disconnectStream(sessionId: string, io: SessionChannelIO): void { const key = keyFor(sessionId, io); const tail = this.tails.get(key); - const _bufferedSize = this.buffer.get(key)?.length ?? 0; // Mark as explicitly disconnected BEFORE we abort, so the tail's // `.finally` reconnect path sees the flag when it runs (which can be // synchronous in the AbortError catch). Cleared on the next explicit @@ -295,8 +384,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { tail.abortController.abort(); this.tails.delete(key); } - this.buffer.delete(key); - this.bufferSeqNums.delete(key); // Reset the backoff counter so a future re-attach starts fresh — // an explicit disconnect is a deliberate teardown, not evidence of // a broken backend. @@ -335,6 +422,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 +438,6 @@ export class StandardSessionStreamManager implements SessionStreamManager { } this.onceWaiters.clear(); this.buffer.clear(); - this.bufferSeqNums.clear(); } #ensureTailConnected(sessionId: string, io: SessionChannelIO): void { @@ -368,15 +455,10 @@ export class StandardSessionStreamManager implements SessionStreamManager { this.tails.delete(key); // If the tail was torn down explicitly via `disconnectStream`, - // honor that — the caller (typically `session.in.wait()`) is - // suspending the run and expects no records to be buffered or - // delivered until a fresh `on()` / `once()` re-attaches. Without - // this guard a run-level persistent handler (e.g. `chat.agent`'s - // `stopInput.on(...)`) would auto-reconnect during the suspend - // window, the resurrected tail would receive the same record the - // waitpoint just delivered, and that record would land in the - // buffer where the next turn's `messagesInput.on(...)` drains it - // and runs a duplicate turn. + // honor that until a fresh `on()` / `once()` re-attaches. Existing + // buffered records stay available across the suspension, but a + // run-level handler must not reconnect and receive another copy of + // the record being delivered through the waitpoint. if (this.explicitlyDisconnected.has(key)) { return; } @@ -427,9 +509,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 +535,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 +564,21 @@ export class StandardSessionStreamManager implements SessionStreamManager { } } - #dispatch(key: string, data: unknown, seqNum: number | undefined): void { + #dispatch(key: string, record: SessionStreamRecord): void { // Any record flowing through = healthy connection; reset the backoff // counter so the next disconnect starts fresh. this.reconnectAttempts.delete(key); - const waiters = this.onceWaiters.get(key); - if (waiters && waiters.length > 0) { - const waiter = waiters.shift()!; - if (waiters.length === 0) this.onceWaiters.delete(key); - if (waiter.timeoutHandle) clearTimeout(waiter.timeoutHandle); - if (waiter.signal && waiter.abortHandler) { - waiter.signal.removeEventListener("abort", waiter.abortHandler); - } + const existingBuffer = this.buffer.get(key); + const waiter = + existingBuffer && existingBuffer.length > 0 ? undefined : this.#takeOnceWaiter(key, record); + if (waiter) { // Record was consumed directly by a waiter — advance the // committed-consume cursor immediately. Buffered-then-shifted // records advance the cursor in `once()` / `shiftBuffer()`. - if (seqNum !== undefined) { - this.#advanceLastDispatched(key, seqNum); - } - waiter.resolve({ ok: true, output: data }); - this.#invokeHandlers(key, data); + this.#advanceLastDispatched(key, record.seqNum); + waiter.resolve({ ok: true, output: record }); + this.#invokeHandlers(key, record.data); return; } @@ -511,11 +590,9 @@ export class StandardSessionStreamManager implements SessionStreamManager { // second turn. Records no handler consumed (e.g. a message arriving // while only the stop facade is attached during preload) are buffered // so a subsequent `once()` can still pick them up. - const consumed = this.#invokeHandlers(key, data); + const consumed = this.#invokeHandlers(key, record.data); if (consumed) { - if (seqNum !== undefined) { - this.#advanceLastDispatched(key, seqNum); - } + this.#advanceLastDispatched(key, record.seqNum); return; } @@ -524,17 +601,49 @@ 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. */ diff --git a/packages/core/src/v3/sessionStreams/noopManager.ts b/packages/core/src/v3/sessionStreams/noopManager.ts index f2d355d24e..1e5dbaebe9 100644 --- a/packages/core/src/v3/sessionStreams/noopManager.ts +++ b/packages/core/src/v3/sessionStreams/noopManager.ts @@ -1,6 +1,11 @@ import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; import { InputStreamOncePromise } from "../inputStreams/types.js"; -import type { SessionChannelIO, SessionStreamManager } from "./types.js"; +import type { + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "./types.js"; export class NoopSessionStreamManager implements SessionStreamManager { on( @@ -21,16 +26,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; } diff --git a/packages/core/src/v3/sessionStreams/types.ts b/packages/core/src/v3/sessionStreams/types.ts index ae24259b3f..24b6f08451 100644 --- a/packages/core/src/v3/sessionStreams/types.ts +++ b/packages/core/src/v3/sessionStreams/types.ts @@ -12,6 +12,21 @@ export type { InputStreamOnceResult }; export type SessionChannelIO = "out" | "in"; +/** + * One durable Session channel record. + * + * `id` is the append's stable idempotency key. `seqNum` is the record's + * monotonic S2 sequence within the Session channel. Both stay stable when + * the same record is delivered again after a reconnect. + */ +export type SessionStreamRecord = Readonly<{ + id: string; + seqNum: number; + data: T; +}>; + +export type SessionStreamRecordPredicate = (record: SessionStreamRecord) => boolean; + /** * Manager for Session channel reads: a session-scoped parallel to * {@link InputStreamManager} keyed on `(sessionId, io)` instead of @@ -42,20 +57,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 +106,8 @@ export interface SessionStreamManager { /** * Seed the committed-consume cursor at worker boot — e.g. from the * `session-in-event-id` header on the latest `turn-complete` on - * `.out`. Monotonic: only ever advances forward, never backwards. + * `.out`. Monotonic: only ever advances forward, never backwards. Existing + * unconsumed records still constrain {@link lastDispatchedSeqNum}. */ setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void; @@ -82,7 +124,7 @@ export interface SessionStreamManager { /** Remove and discard the first buffered record. Returns true if one was removed. */ shiftBuffer(sessionId: string, io: SessionChannelIO): boolean; - /** Abort the SSE tail and clear the buffer. Called before `.wait` suspends. */ + /** Abort the SSE tail while preserving buffered records. Called before `.wait` suspends. */ disconnectStream(sessionId: string, io: SessionChannelIO): void; /** Clear all `.on` handlers; abort tails without pending once-waiters. */ diff --git a/packages/core/src/v3/sessionStreams/wireProtocol.ts b/packages/core/src/v3/sessionStreams/wireProtocol.ts index 550e81a0af..bb6aef3e1a 100644 --- a/packages/core/src/v3/sessionStreams/wireProtocol.ts +++ b/packages/core/src/v3/sessionStreams/wireProtocol.ts @@ -40,6 +40,53 @@ export const SESSION_STATE_LAST_EVENT_ID_HEADER = "last-event-id" as const; */ export const SESSION_IN_EVENT_ID_HEADER = "session-in-event-id" as const; +/** + * Opt-in response format for Session stream waitpoints. Older SDKs omit this + * and continue receiving the raw record data. + */ +export const SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT = "record-v1" as const; + +/** Content type used only when a waitpoint actually returns a record-v1 envelope. */ +export const SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE = + "application/vnd.trigger.session-stream-record+json" as const; + +const SESSION_STREAM_WAITPOINT_RECORD_TYPE = "trigger-session-stream-record" as const; + +/** Internal envelope used to return an exact Session record from a waitpoint. */ +export type SessionStreamWaitpointRecord = Readonly<{ + type: typeof SESSION_STREAM_WAITPOINT_RECORD_TYPE; + version: 1; + seqNum: number; + data: unknown; +}>; + +export function serializeSessionStreamWaitpointRecord(data: unknown, seqNum: number): string { + return JSON.stringify({ + type: SESSION_STREAM_WAITPOINT_RECORD_TYPE, + version: 1, + seqNum, + data, + } satisfies SessionStreamWaitpointRecord); +} + +export function parseSessionStreamWaitpointRecord( + value: unknown +): SessionStreamWaitpointRecord | undefined { + if (!value || typeof value !== "object") return undefined; + + const record = value as Partial; + if ( + record.type !== SESSION_STREAM_WAITPOINT_RECORD_TYPE || + record.version !== 1 || + typeof record.seqNum !== "number" || + !Number.isFinite(record.seqNum) + ) { + return undefined; + } + + return record as SessionStreamWaitpointRecord; +} + export const TRIGGER_CONTROL_SUBTYPE = { TURN_COMPLETE: "turn-complete", UPGRADE_REQUIRED: "upgrade-required", diff --git a/packages/core/src/v3/test/mock-task-context.ts b/packages/core/src/v3/test/mock-task-context.ts index 5fbe195761..cd65ac24d4 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 8cae877f54..fd7cd5d360 100644 --- a/packages/core/src/v3/test/session-waitpoint-backend.ts +++ b/packages/core/src/v3/test/session-waitpoint-backend.ts @@ -1,6 +1,10 @@ import { ApiClient } from "../apiClient/index.js"; import { WaitpointId } from "../isomorphic/friendlyId.js"; import { NoopRuntimeManager } from "../runtime/noopRuntimeManager.js"; +import { + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, +} from "../sessionStreams/wireProtocol.js"; import type { CreateSessionStreamWaitpointRequestBody, CreateSessionStreamWaitpointResponseBody, @@ -13,6 +17,7 @@ type PendingWait = { io: "in" | "out"; lastSeqNum?: number; timeout?: string; + responseFormat?: "record-v1"; abort: AbortController; }; @@ -71,6 +76,7 @@ export class SessionWaitpointBackend { io: body.io, lastSeqNum: body.lastSeqNum, timeout: body.timeout, + responseFormat: body.responseFormat, abort: new AbortController(), }); return { waitpointId, isCached: false }; @@ -113,8 +119,20 @@ export class SessionWaitpointBackend { }; } - const output = typeof result === "string" ? result : JSON.stringify(result); - return { ok: true, output, outputType: "application/json" }; + const output = + pending.responseFormat === "record-v1" + ? serializeSessionStreamWaitpointRecord(result.data, result.seqNum) + : typeof result.data === "string" + ? result.data + : JSON.stringify(result.data); + return { + ok: true, + output, + outputType: + pending.responseFormat === "record-v1" + ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE + : "application/json", + }; } catch { return { ok: false, @@ -144,16 +162,23 @@ export class SessionWaitpointBackend { * which {@link wait} passes straight to the packet parser so it round-trips * to the same object `session.in.once()` returns. */ - private async readNextRecord(pending: PendingWait): Promise { + private async readNextRecord(pending: PendingWait): Promise<{ data: unknown; seqNum: number }> { const lastEventId = pending.lastSeqNum !== undefined && pending.lastSeqNum >= 0 ? String(pending.lastSeqNum) : undefined; + let deliveredSeqNum: number | undefined; const stream = await this.apiClient.subscribeToSessionStream(pending.session, pending.io, { lastEventId, signal: pending.abort.signal, timeoutInSeconds: 120, + onPart: (part) => { + const seqNum = Number.parseInt(part.id, 10); + if (Number.isFinite(seqNum)) { + deliveredSeqNum = seqNum; + } + }, }); const reader = stream.getReader(); @@ -162,7 +187,10 @@ export class SessionWaitpointBackend { if (done) { throw new Error("session stream closed"); } - return value; + if (deliveredSeqNum === undefined) { + throw new Error("session stream record is missing its sequence number"); + } + return { data: value, seqNum: deliveredSeqNum }; } finally { await reader.cancel().catch(() => {}); pending.abort.abort(); diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index 0e08441d4c..6d8ad0b553 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -1,10 +1,16 @@ import type { InputStreamOnceResult } from "../inputStreams/types.js"; import { InputStreamOncePromise, InputStreamTimeoutError } from "../inputStreams/types.js"; import type { InputStreamOnceOptions } from "../realtimeStreams/types.js"; -import type { SessionChannelIO, SessionStreamManager } from "../sessionStreams/types.js"; +import type { + SessionChannelIO, + SessionStreamManager, + SessionStreamRecord, + SessionStreamRecordPredicate, +} from "../sessionStreams/types.js"; type OnceWaiter = { - resolve: (value: InputStreamOnceResult) => void; + resolve: (value: InputStreamOnceResult) => void; + predicate?: SessionStreamRecordPredicate; timer?: ReturnType; signal?: AbortSignal; abortHandler?: () => void; @@ -31,9 +37,10 @@ function keyFor(sessionId: string, io: SessionChannelIO): string { export class TestSessionStreamManager implements SessionStreamManager { private handlers = new Map>(); private onceWaiters = new Map(); - private buffer = new Map(); + private buffer = new Map(); private seqNums = new Map(); private dispatchedSeqNums = new Map(); + private unconsumedSeqNums = new Map>(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { const key = keyFor(sessionId, io); @@ -55,21 +62,26 @@ export class TestSessionStreamManager implements SessionStreamManager { // messages into every newly attached per-turn handler. const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - const kept: unknown[] = []; - for (const data of buffered) { + const kept: SessionStreamRecord[] = []; + for (const record of buffered) { let consumed = false; try { - consumed = handler(data) === true; + consumed = handler(record.data) === true; } catch { // Never let a handler error break test state } - if (!consumed) kept.push(data); + if (consumed) { + this.#advanceLastDispatched(key, record.seqNum); + } else { + kept.push(record); + } } if (kept.length > 0) { this.buffer.set(key, kept); } else { this.buffer.delete(key); } + this.#drainOnceWaitersFromBuffer(key); } return { @@ -84,9 +96,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 +140,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 +194,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,22 +209,73 @@ export class TestSessionStreamManager implements SessionStreamManager { this.seqNums.set(keyFor(sessionId, io), seqNum); } + consumeRecord(sessionId: string, io: SessionChannelIO, seqNum: number): void { + const key = keyFor(sessionId, io); + const buffered = this.buffer.get(key); + const index = buffered?.findIndex((record) => record.seqNum === seqNum) ?? -1; + + if (buffered && index !== -1) { + buffered.splice(index, 1); + if (buffered.length === 0) { + this.buffer.delete(key); + } + } + + this.#advanceLastDispatched(key, seqNum); + this.#drainOnceWaitersFromBuffer(key); + } + lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { - // `__sendFromTest` carries no seq numbers, so this only reflects - // explicit `setLastDispatchedSeqNum` calls (e.g. the waitpoint - // delivery path). Full cursor behaviour is exercised via the real - // manager. - return this.dispatchedSeqNums.get(keyFor(sessionId, io)); + const key = keyFor(sessionId, io); + const highWatermark = this.dispatchedSeqNums.get(key); + if (highWatermark === undefined) return undefined; + + const unconsumedSeqNums = this.unconsumedSeqNums.get(key); + if (!unconsumedSeqNums || unconsumedSeqNums.size === 0) return highWatermark; + + let earliestUnconsumedSeqNum = Infinity; + for (const seqNum of unconsumedSeqNums) { + earliestUnconsumedSeqNum = Math.min(earliestUnconsumedSeqNum, seqNum); + } + + const safeCursor = Math.min(highWatermark, earliestUnconsumedSeqNum - 1); + return safeCursor >= 0 ? safeCursor : undefined; } setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void { - const key = keyFor(sessionId, io); + if (!Number.isFinite(seqNum)) return; + + this.#advanceLastDispatched(keyFor(sessionId, io), seqNum); + } + + #advanceLastDispatched(key: string, seqNum: number): void { + this.#removeUnconsumedRecord(key, seqNum); + if (!Number.isFinite(seqNum)) return; const current = this.dispatchedSeqNums.get(key); if (current === undefined || seqNum > current) { this.dispatchedSeqNums.set(key, seqNum); } } + #markUnconsumedRecord(key: string, seqNum: number): void { + if (!Number.isFinite(seqNum)) return; + + let unconsumedSeqNums = this.unconsumedSeqNums.get(key); + if (!unconsumedSeqNums) { + unconsumedSeqNums = new Set(); + this.unconsumedSeqNums.set(key, unconsumedSeqNums); + } + unconsumedSeqNums.add(seqNum); + } + + #removeUnconsumedRecord(key: string, seqNum: number): void { + const unconsumedSeqNums = this.unconsumedSeqNums.get(key); + unconsumedSeqNums?.delete(seqNum); + if (unconsumedSeqNums?.size === 0) { + this.unconsumedSeqNums.delete(key); + } + } + setMinTimestamp( _sessionId: string, _io: SessionChannelIO, @@ -180,15 +289,18 @@ export class TestSessionStreamManager implements SessionStreamManager { const key = keyFor(sessionId, io); const buffered = this.buffer.get(key); if (buffered && buffered.length > 0) { - buffered.shift(); + const record = buffered.shift()!; if (buffered.length === 0) this.buffer.delete(key); + this.#advanceLastDispatched(key, record.seqNum); + this.#drainOnceWaitersFromBuffer(key); return true; } return false; } disconnectStream(_sessionId: string, _io: SessionChannelIO): void { - // no-op — no real SSE tail in tests + // The production manager keeps buffered records reachable across a + // waitpoint suspension. The exact waitpoint record is removed on resume. } clearHandlers(): void { @@ -209,6 +321,7 @@ export class TestSessionStreamManager implements SessionStreamManager { this.buffer.clear(); this.seqNums.clear(); this.dispatchedSeqNums.clear(); + this.unconsumedSeqNums.clear(); } disconnect(): void { @@ -235,39 +348,56 @@ export class TestSessionStreamManager implements SessionStreamManager { * resolves. Consumption is decided on the synchronous return value, * exactly like production. */ - async __sendFromTest(sessionId: string, io: SessionChannelIO, data: unknown): Promise { + async __sendFromTest( + sessionId: string, + io: SessionChannelIO, + data: unknown, + metadata?: { id?: string; seqNum?: number } + ): Promise { const key = keyFor(sessionId, io); + const seqNum = metadata?.seqNum ?? (this.seqNums.get(key) ?? -1) + 1; + if (!Number.isFinite(seqNum)) { + throw new TypeError("Test Session stream records require a finite sequence number"); + } + const record: SessionStreamRecord = { + id: metadata?.id ?? `test-record-${seqNum}`, + seqNum, + data, + }; + const lastSeqNum = this.seqNums.get(key); + if (lastSeqNum === undefined || seqNum > lastSeqNum) { + this.seqNums.set(key, seqNum); + } - const waiters = this.onceWaiters.get(key); - if (waiters && waiters.length > 0) { - const w = waiters.shift()!; - if (waiters.length === 0) this.onceWaiters.delete(key); - if (w.timer) clearTimeout(w.timer); - if (w.signal && w.abortHandler) { - w.signal.removeEventListener("abort", w.abortHandler); - } - w.resolve({ ok: true, output: data }); - await this.#invokeHandlers(key, data); + const existingBuffer = this.buffer.get(key); + const waiter = + existingBuffer && existingBuffer.length > 0 ? undefined : this.#takeOnceWaiter(key, record); + if (waiter) { + this.#advanceLastDispatched(key, record.seqNum); + waiter.resolve({ ok: true, output: record }); + await this.#invokeHandlers(key, record.data); return; } - const consumed = await this.#invokeHandlers(key, data); - if (consumed) return; + const consumed = await this.#invokeHandlers(key, record.data); + if (consumed) { + this.#advanceLastDispatched(key, record.seqNum); + return; + } // Re-check waiters: handler invocation above is awaited (unlike the // synchronous production dispatch), and the runtime commonly registers // its next `once()` during that window — e.g. the turn loop reaching // `waitWithIdleTimeout` while a handler settles. Without this second // look the record would be buffered while the fresh waiter hangs. - const lateWaiters = this.onceWaiters.get(key); - if (lateWaiters && lateWaiters.length > 0) { - const w = lateWaiters.shift()!; - if (lateWaiters.length === 0) this.onceWaiters.delete(key); - if (w.timer) clearTimeout(w.timer); - if (w.signal && w.abortHandler) { - w.signal.removeEventListener("abort", w.abortHandler); - } - w.resolve({ ok: true, output: data }); + const bufferedAfterHandlers = this.buffer.get(key); + const lateWaiter = + bufferedAfterHandlers && bufferedAfterHandlers.length > 0 + ? undefined + : this.#takeOnceWaiter(key, record); + if (lateWaiter) { + this.#advanceLastDispatched(key, record.seqNum); + lateWaiter.resolve({ ok: true, output: record }); return; } @@ -276,7 +406,46 @@ export class TestSessionStreamManager implements SessionStreamManager { buffered = []; this.buffer.set(key, buffered); } - buffered.push(data); + 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 { + return false; + } + }); + if (index === -1) return undefined; + + const [waiter] = waiters.splice(index, 1); + if (waiters.length === 0) this.onceWaiters.delete(key); + if (waiter!.timer) clearTimeout(waiter!.timer); + if (waiter!.signal && waiter!.abortHandler) { + waiter!.signal.removeEventListener("abort", waiter!.abortHandler); + } + return waiter; + } + + #drainOnceWaitersFromBuffer(key: string): void { + const buffered = this.buffer.get(key); + while (buffered && buffered.length > 0) { + const record = buffered[0]!; + const waiter = this.#takeOnceWaiter(key, record); + if (!waiter) return; + + buffered.shift(); + if (buffered.length === 0) this.buffer.delete(key); + this.#advanceLastDispatched(key, record.seqNum); + waiter.resolve({ ok: true, output: record }); + } } /** diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 844d506079..95b0f5267d 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1543,7 +1543,31 @@ export type ChatTaskRunPayload< // keep their original shape. Each accessor resolves the session handle // lazily via `getChatSession()` so the module-level references stay // compatible with the pre-migration wiring. -const messagesInput: RealtimeDefinedInputStream = { +/** + * One message record delivered through {@link chat.messages}. + * + * `id` is the append's stable idempotency key and `seqNum` is its monotonic + * sequence on the Session `.in` channel. Both remain stable if the record is + * delivered again after a reconnect. + */ +export type ChatMessageRecord = Readonly<{ + id: string; + seqNum: number; + payload: ChatTaskWirePayload; +}>; + +export type ChatMessages = RealtimeDefinedInputStream & { + /** Whether 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; +}; + +function isChatMessageRecord(record: { data: unknown }): boolean { + return (record.data as ChatInputChunk | undefined)?.kind === "message"; +} + +const messagesInput: ChatMessages = { id: "chat-messages", on(handler) { return getChatSession().in.on((chunk) => { @@ -1607,6 +1631,36 @@ const messagesInput: RealtimeDefinedInputStream = { if (chunk && chunk.kind === "message") return chunk.payload; return undefined; }, + async hasPending() { + return messagesInput.peek() !== undefined; + }, + async next(options) { + const timeoutInSeconds = options?.timeoutInSeconds; + if ( + timeoutInSeconds !== undefined && + (!Number.isFinite(timeoutInSeconds) || timeoutInSeconds < 0) + ) { + throw new TypeError( + "chat.messages.next() timeoutInSeconds must be a finite non-negative number" + ); + } + + const session = getChatSession(); + const result = await sessionStreams.onceRecordWhere( + session.id, + "in", + isChatMessageRecord, + timeoutInSeconds === undefined ? undefined : { timeoutMs: timeoutInSeconds * 1000 } + ); + if (!result.ok) return undefined; + + const chunk = result.output.data as Extract; + return { + id: result.output.id, + seqNum: result.output.seqNum, + payload: chunk.payload, + }; + }, wait(options) { return new ManualWaitpointPromise(async (resolve, reject) => { try { diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index 8a01f8293c..125991a0c4 100644 --- a/packages/trigger-sdk/src/v3/sessions.ts +++ b/packages/trigger-sdk/src/v3/sessions.ts @@ -25,6 +25,8 @@ import type { import { InputStreamOncePromise, ManualWaitpointPromise, + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT, SemanticInternalAttributes, SessionStreamInstance, WaitpointTimeoutError, @@ -32,6 +34,7 @@ import { apiClientManager, ensureReadableStream, mergeRequestOptions, + parseSessionStreamWaitpointRecord, runtime, sessionStreams, taskContext, @@ -683,10 +686,10 @@ export class SessionInputChannel { } /** - * The highest S2 sequence number of any record this channel has - * delivered to a `once()` / `wait()` consumer (or had shifted off its - * buffer into one). Distinct from "last received" — buffered-but-not- - * yet-consumed records don't count. + * The highest S2 sequence number that is safe to persist as consumed. + * This stays behind the earliest unconsumed record if a later record was + * handled first. Distinct from "last received", which advances for records + * that may still be pending. * * Used by `chat.agent` to persist the `.in` resume cursor on each * `turn-complete` control record, so the next worker boot can subscribe @@ -713,6 +716,7 @@ export class SessionInputChannel { const apiClient = apiClientManager.clientOrThrow(); + const lastConsumedSeqNum = sessionStreams.lastDispatchedSeqNum(this.sessionId, "in"); const response = await apiClient.createSessionStreamWaitpoint(ctx.run.id, { session: this.sessionId, io: "in", @@ -720,7 +724,8 @@ export class SessionInputChannel { idempotencyKey: options?.idempotencyKey, idempotencyKeyTTL: options?.idempotencyKeyTTL, tags: options?.tags, - lastSeqNum: sessionStreams.lastSeqNum(this.sessionId, "in"), + lastSeqNum: lastConsumedSeqNum, + responseFormat: SESSION_STREAM_WAITPOINT_RESPONSE_FORMAT, }); const result = await tracer.startActiveSpan( @@ -735,36 +740,77 @@ export class SessionInputChannel { throw new Error("Failed to block on session stream waitpoint"); } - // Drop the SSE tail + buffer before suspending so the record - // delivered via the waitpoint path isn't re-buffered on resume. + // Stop the SSE tail before suspending. Buffered records stay in + // place; the exact record returned by the waitpoint is removed on + // resume, while any later records remain available to consumers. sessionStreams.disconnectStream(this.sessionId, "in"); const waitResult = await runtime.waitUntil(response.waitpointId); + const hasRecordEnvelope = + waitResult.outputType === SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE; - const data = + const parsedOutput = waitResult.output !== undefined ? await conditionallyImportAndParsePacket( { data: waitResult.output, - dataType: waitResult.outputType ?? "application/json", + dataType: hasRecordEnvelope + ? "application/json" + : (waitResult.outputType ?? "application/json"), }, apiClient ) : undefined; if (waitResult.ok) { - // Advance both cursors past the record consumed via the - // waitpoint: the seq counter so the SSE tail doesn't replay - // it, and the consume cursor so turn-completes don't stamp a - // stale `session-in-event-id`. - const prevSeq = sessionStreams.lastSeqNum(this.sessionId, "in"); - const nextSeq = (prevSeq ?? -1) + 1; - sessionStreams.setLastSeqNum(this.sessionId, "in", nextSeq); - sessionStreams.setLastDispatchedSeqNum(this.sessionId, "in", nextSeq); + const record = hasRecordEnvelope + ? parseSessionStreamWaitpointRecord(parsedOutput) + : undefined; + let seqNum = record?.seqNum; + const data = record + ? await conditionallyImportAndParsePacket( + { + data: + typeof record.data === "string" ? record.data : JSON.stringify(record.data), + dataType: "application/json", + }, + apiClient + ) + : parsedOutput; + + // Older servers return only raw data. Recover its durable + // sequence from the channel instead of guessing and risking a + // cursor that skips or strands another record. + if (seqNum === undefined && waitResult.output !== undefined) { + try { + const response = await apiClient.readSessionStreamRecords(this.sessionId, "in", { + afterEventId: + lastConsumedSeqNum !== undefined ? String(lastConsumedSeqNum) : undefined, + }); + const matchingRecords = response.records.filter( + (candidate) => + candidate.data === waitResult.output || + (typeof candidate.data !== "string" && + JSON.stringify(candidate.data) === JSON.stringify(parsedOutput)) + ); + if (matchingRecords.length === 1) { + seqNum = matchingRecords[0]!.seqNum; + } + } catch { + // Leave the cursor behind when an older server cannot + // provide record metadata. At-least-once replay is safer + // than acknowledging an unknown sequence. + } + } + + if (seqNum !== undefined) { + sessionStreams.consumeRecord(this.sessionId, "in", seqNum); + sessionStreams.setLastSeqNum(this.sessionId, "in", seqNum); + } return { ok: true as const, output: data as T }; } else { - const error = new WaitpointTimeoutError(data?.message ?? "Timed out"); + const error = new WaitpointTimeoutError(parsedOutput?.message ?? "Timed out"); span.recordException(error); span.setStatus({ code: SpanStatusCode.ERROR }); return { ok: false as const, error }; diff --git a/packages/trigger-sdk/test/chat-messages-mailbox.test.ts b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts new file mode 100644 index 0000000000..295b8f6c55 --- /dev/null +++ b/packages/trigger-sdk/test/chat-messages-mailbox.test.ts @@ -0,0 +1,310 @@ +// Import the test harness FIRST — this installs the resource catalog so +// `chat.customAgent()` calls below register their task functions correctly. +import "../src/v3/test/index.js"; + +import { resourceCatalog, sessionStreams } from "@trigger.dev/core/v3"; +import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; +import { describe, expect, it } from "vitest"; +import { chat, type ChatMessageRecord, type ChatTaskWirePayload } from "../src/v3/ai.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; +} + +function userPayload(chatId: string, id: string): ChatTaskWirePayload { + return { + chatId, + trigger: "submit-message", + message: { + id, + role: "user", + parts: [{ type: "text", text: id }], + }, + }; +} + +describe("chat.messages mailbox", () => { + it("checks pending input without consuming and takes one buffered record at a time", async () => { + const chatId = "mailbox-buffered"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + initial?: boolean; + before?: boolean; + afterFirst?: boolean; + afterSecond?: boolean; + first?: ChatMessageRecord; + second?: ChatMessageRecord; + cursorAfterFirst?: number; + cursorAfterSecond?: number; + } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-buffered", + run: async () => { + observations.initial = await chat.messages.hasPending(); + ready.resolve(); + await inspect.promise; + + observations.before = await chat.messages.hasPending(); + observations.first = await chat.messages.next(); + observations.cursorAfterFirst = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.afterFirst = await chat.messages.hasPending(); + observations.second = await chat.messages.next(); + observations.cursorAfterSecond = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.afterSecond = await chat.messages.hasPending(); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u1") }, + "in", + { id: "part-1", seqNum: 10 } + ); + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u2") }, + "in", + { id: "part-2", seqNum: 11 } + ); + inspect.resolve(); + await runPromise; + }); + + expect(observations).toEqual({ + initial: false, + before: true, + first: { id: "part-1", seqNum: 10, payload: userPayload(chatId, "u1") }, + cursorAfterFirst: 10, + afterFirst: true, + second: { id: "part-2", seqNum: 11, payload: userPayload(chatId, "u2") }, + cursorAfterSecond: 11, + afterSecond: false, + }); + }); + + it("returns undefined when next times out", async () => { + let result: ChatMessageRecord | undefined; + const agent = chat.customAgent({ + id: "chat-messages-mailbox-timeout", + run: async () => { + result = await chat.messages.next({ timeoutInSeconds: 0 }); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext((drivers) => + run( + { chatId: "mailbox-timeout", trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ) + ); + + expect(result).toBeUndefined(); + }); + + it("leaves earlier non-message records for their own consumer", async () => { + const chatId = "mailbox-mixed-kinds"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + pending?: boolean; + blocked?: ChatMessageRecord; + cursorAfterBlocked?: number; + headAfterBlocked?: unknown; + control?: unknown; + pendingAfterControl?: boolean; + message?: ChatMessageRecord; + cursorAfterMessage?: number; + } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-mixed-kinds", + run: async () => { + ready.resolve(); + await inspect.promise; + + observations.pending = await chat.messages.hasPending(); + observations.blocked = await chat.messages.next({ timeoutInSeconds: 0 }); + observations.cursorAfterBlocked = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.headAfterBlocked = sessionStreams.peekRecord(chatId, "in"); + + const control = await sessionStreams.onceRecord(chatId, "in"); + observations.control = control.ok ? control.output : undefined; + observations.pendingAfterControl = await chat.messages.hasPending(); + observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); + observations.cursorAfterMessage = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + + await drivers.sessions.in.send( + chatId, + { kind: "handover", partialAssistantMessage: [], isFinal: false }, + "in", + { id: "handover-1", seqNum: 30 } + ); + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u-after-handover") }, + "in", + { id: "message-1", seqNum: 31 } + ); + inspect.resolve(); + await runPromise; + }); + + expect(observations).toEqual({ + pending: false, + blocked: undefined, + cursorAfterBlocked: undefined, + headAfterBlocked: { + id: "handover-1", + seqNum: 30, + data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + }, + control: { + id: "handover-1", + seqNum: 30, + data: { kind: "handover", partialAssistantMessage: [], isFinal: false }, + }, + pendingAfterControl: true, + message: { + id: "message-1", + seqNum: 31, + payload: userPayload(chatId, "u-after-handover"), + }, + cursorAfterMessage: 31, + }); + }); + + it("keeps the cursor behind a buffered message when a later stop is consumed", async () => { + const chatId = "mailbox-cursor-gap"; + const ready = deferred(); + const inspect = deferred(); + const observations: { + cursorBefore?: number; + message?: ChatMessageRecord; + cursorAfter?: number; + } = {}; + + const agent = chat.customAgent({ + id: "chat-messages-mailbox-cursor-gap", + run: async () => { + const stop = chat.createStopSignal(); + ready.resolve(); + await inspect.promise; + + observations.cursorBefore = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + observations.message = await chat.messages.next({ timeoutInSeconds: 0 }); + observations.cursorAfter = sessionStreams.lastDispatchedSeqNum(chatId, "in"); + stop.cleanup(); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + await runInMockTaskContext(async (drivers) => { + const runPromise = run( + { chatId, trigger: "preload" }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ); + await ready.promise; + + await drivers.sessions.in.send( + chatId, + { kind: "message", payload: userPayload(chatId, "u1") }, + "in", + { id: "message-1", seqNum: 50 } + ); + await drivers.sessions.in.send(chatId, { kind: "stop" }, "in", { + id: "stop-1", + seqNum: 51, + }); + inspect.resolve(); + await runPromise; + }); + + expect(observations).toEqual({ + cursorBefore: 49, + message: { + id: "message-1", + seqNum: 50, + payload: userPayload(chatId, "u1"), + }, + cursorAfter: 51, + }); + }); + + it("keeps record id and sequence stable across redelivery", async () => { + const payload = userPayload("mailbox-redelivery", "u-redelivered"); + const ready = deferred(); + const consumeFirst = deferred(); + const readyForRedelivery = deferred(); + const consumeRedelivery = deferred(); + let first: ChatMessageRecord | undefined; + let redelivered: ChatMessageRecord | undefined; + const agent = chat.customAgent({ + id: "chat-messages-mailbox-redelivery", + run: async () => { + ready.resolve(); + await consumeFirst.promise; + first = await chat.messages.next({ timeoutInSeconds: 0 }); + + 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); + }); +}); diff --git a/packages/trigger-sdk/test/mockChatAgent.test.ts b/packages/trigger-sdk/test/mockChatAgent.test.ts index 202c392373..62437369a3 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 f5bd705751..18ef2462bf 100644 --- a/packages/trigger-sdk/test/pending-message-drain.test.ts +++ b/packages/trigger-sdk/test/pending-message-drain.test.ts @@ -5,7 +5,12 @@ import { mockChatAgent } from "../src/v3/test/index.js"; import { describe, expect, it, vi } from "vitest"; import { chat } from "../src/v3/ai.js"; import { __setSessionOpenImplForTests, sessions } from "../src/v3/sessions.js"; -import { apiClientManager, sessionStreams } from "@trigger.dev/core/v3"; +import { + apiClientManager, + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE, + serializeSessionStreamWaitpointRecord, + sessionStreams, +} from "@trigger.dev/core/v3"; import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; import { simulateReadableStream, streamText } from "ai"; import { MockLanguageModelV3 } from "ai/test"; @@ -70,6 +75,21 @@ async function waitFor(check: () => boolean, timeoutMs = 10_000) { throw new Error("waitFor timed out"); } +function runtimeWithWaitpointOutput(output: string, outputType = "application/json") { + return { + disable() {}, + waitForTask() { + throw new Error("Unexpected task wait"); + }, + waitForBatch() { + throw new Error("Unexpected batch wait"); + }, + waitForWaitpoint() { + return Promise.resolve({ ok: true, output, outputType }); + }, + }; +} + function streamedText(harness: { allChunks: unknown[] }): string { return (harness.allChunks as { type?: string; delta?: string }[]) .filter((c) => c.type === "text-delta") @@ -248,26 +268,179 @@ describe("chat.createSession stop + immediate send", () => { }); describe("session.in.wait() consume cursor", () => { - it("advances lastDispatchedSeqNum alongside lastSeqNum on waitpoint delivery", async () => { + it("keeps later input reachable across the suspend-and-resume race", async () => { __setSessionOpenImplForTests(undefined); - await runInMockTaskContext(async () => { - vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - createSessionStreamWaitpoint: async () => ({ waitpointId: "wp_test_1" }), - waitForWaitpointToken: async () => ({ success: true }), - } as never); - - const sessionId = "cursor-sess"; - // Simulate records 0..4 already received via SSE before the suspend. - sessionStreams.setLastSeqNum(sessionId, "in", 4); - - const result = await sessions.open(sessionId).in.wait(); - - expect(result.ok).toBe(true); - expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(5); - // The waitpoint-delivered record was consumed by this caller, so the - // committed-consume cursor (what turn-completes persist as - // `session-in-event-id`) must advance with it. - expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(5); - }); + const first = { kind: "message", payload: { id: "u1" } }; + const later = { kind: "message", payload: { id: "u2" } }; + const runtimeManager = runtimeWithWaitpointOutput( + serializeSessionStreamWaitpointRecord(JSON.stringify(first), 50), + SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE + ); + let registeredLastSeqNum: number | undefined; + let registeredResponseFormat: string | undefined; + + await runInMockTaskContext( + async (drivers) => { + const sessionId = "cursor-sess"; + const channel = sessions.open(sessionId).in; + const stop = channel.on<{ kind: string }>((record) => record.kind === "stop"); + + sessionStreams.setLastSeqNum(sessionId, "in", 49); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 49); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async ( + _runId: string, + body: { lastSeqNum?: number; responseFormat?: string } + ) => { + registeredLastSeqNum = body.lastSeqNum; + registeredResponseFormat = body.responseFormat; + return { + waitpointId: "wp_test_1", + isCached: false, + }; + }, + waitForWaitpointToken: async () => { + // These records land after registration but before the tail is + // disconnected. The waitpoint resolves with seq 50, while the + // local tail has already consumed 51 and buffered 52. + await drivers.sessions.in.send(sessionId, first, "in", { seqNum: 50 }); + await drivers.sessions.in.send(sessionId, { kind: "stop" }, "in", { seqNum: 51 }); + await drivers.sessions.in.send(sessionId, later, "in", { seqNum: 52 }); + return { success: true }; + }, + } as never); + + const result = await channel.wait(); + + expect(result).toEqual({ ok: true, output: first }); + expect(registeredLastSeqNum).toBe(49); + expect(registeredResponseFormat).toBe("record-v1"); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(51); + expect(sessionStreams.peekRecord(sessionId, "in")?.seqNum).toBe(52); + + const next = await sessionStreams.onceRecord(sessionId, "in"); + expect(next).toEqual({ + ok: true, + output: { id: "test-record-52", seqNum: 52, data: later }, + }); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(52); + stop.off(); + }, + { runtimeManager } + ); + }); + + it("recovers the exact sequence from durable records for older servers", async () => { + __setSessionOpenImplForTests(undefined); + const payload = { kind: "message", payload: { id: "legacy" } }; + const rawPayload = JSON.stringify(payload); + const runtimeManager = runtimeWithWaitpointOutput(rawPayload); + let afterEventId: string | undefined; + + await runInMockTaskContext( + async () => { + const sessionId = "legacy-cursor-sess"; + sessionStreams.setLastSeqNum(sessionId, "in", 6); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ + waitpointId: "wp_legacy", + isCached: false, + }), + waitForWaitpointToken: async () => ({ success: true }), + readSessionStreamRecords: async ( + _sessionId: string, + _io: "in" | "out", + options?: { afterEventId?: string } + ) => { + afterEventId = options?.afterEventId; + return { + records: [{ id: "legacy-record", seqNum: 7, data: rawPayload }], + }; + }, + } as never); + + const result = await sessions.open(sessionId).in.wait(); + + expect(result).toEqual({ ok: true, output: payload }); + expect(afterEventId).toBe("6"); + expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(7); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); + }, + { runtimeManager } + ); + }); + + it("does not mistake an older server's user payload for the internal envelope", async () => { + __setSessionOpenImplForTests(undefined); + const payload = { + type: "trigger-session-stream-record", + version: 1, + seqNum: 999, + data: { user: "supplied" }, + }; + const rawPayload = JSON.stringify(payload); + + await runInMockTaskContext( + async () => { + const sessionId = "legacy-envelope-collision"; + sessionStreams.setLastSeqNum(sessionId, "in", 6); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ + waitpointId: "wp_legacy_collision", + isCached: false, + }), + waitForWaitpointToken: async () => ({ success: true }), + readSessionStreamRecords: async () => ({ + records: [{ id: "legacy-record", seqNum: 7, data: rawPayload }], + }), + } as never); + + const result = await sessions.open(sessionId).in.wait(); + + expect(result).toEqual({ ok: true, output: payload }); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(7); + }, + { runtimeManager: runtimeWithWaitpointOutput(rawPayload) } + ); + }); + + it("leaves the cursor behind when legacy payload matching is ambiguous", async () => { + __setSessionOpenImplForTests(undefined); + const payload = { kind: "message", payload: { id: "duplicate" } }; + const rawPayload = JSON.stringify(payload); + + await runInMockTaskContext( + async () => { + const sessionId = "legacy-duplicate-payload"; + sessionStreams.setLastSeqNum(sessionId, "in", 6); + sessionStreams.setLastDispatchedSeqNum(sessionId, "in", 6); + + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ + waitpointId: "wp_legacy_duplicate", + isCached: false, + }), + waitForWaitpointToken: async () => ({ success: true }), + readSessionStreamRecords: async () => ({ + records: [ + { id: "duplicate-1", seqNum: 7, data: rawPayload }, + { id: "duplicate-2", seqNum: 8, data: rawPayload }, + ], + }), + } as never); + + const result = await sessions.open(sessionId).in.wait(); + + expect(result).toEqual({ ok: true, output: payload }); + expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(6); + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(6); + }, + { runtimeManager: runtimeWithWaitpointOutput(rawPayload) } + ); }); });