diff --git a/.changeset/calm-chat-reconnects.md b/.changeset/calm-chat-reconnects.md new file mode 100644 index 0000000000..a1af4aa43d --- /dev/null +++ b/.changeset/calm-chat-reconnects.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Browser chats now keep the active turn open across page reloads when older completion records are replayed. 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..9b686caf08 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 @@ -13,6 +13,7 @@ import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.s import { stripClientWebhookActionSource } from "~/services/realtime/sanitizeSessionInput.server"; import { claimSessionStreamPart, + commitSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, } from "~/services/sessionStreamWaitpointCache.server"; @@ -147,22 +148,65 @@ const { action, loader } = createActionApiRoute( const partId = clientPartId ?? nanoid(7); // Idempotency on client-supplied part ids: atomically claim the id before - // appending. A concurrent or retried POST that loses the claim skips the - // append (no duplicate record) but still falls through to the drain below, - // so a retry whose first attempt died before waking the waitpoint can still - // recover it. The claim is released on append failure so a genuine retry - // can re-claim and proceed. - const wonClaim = clientPartId + // appending, then store the resulting sequence with the claim. A retry + // returns that original sequence and still falls through to the drain, so + // a first attempt that lost its HTTP response cannot lose turn correlation + // or strand a waitpoint. The claim is released on append failure. + let claim = clientPartId ? await claimSessionStreamPart( authentication.environment.id, addressingKey, params.io, clientPartId ) - : true; + : ({ status: "claimed", claimValue: undefined } as const); - let appendSeq: number | undefined; - if (wonClaim) { + let appendSeq = claim.status === "committed" ? claim.seq : undefined; + if (claim.status === "pending") { + const pendingClaim = claim; + const [recoveryError, recoveredSeq] = await tryCatch( + realtimeStream.findSessionStreamPartSequence( + addressingKey, + params.io, + partId, + pendingClaim.claimedAt + ) + ); + + if (recoveryError) { + logger.error("Failed to recover pending session stream append", { + sessionId: session.id, + io: params.io, + partId, + error: recoveryError, + }); + } else if (recoveredSeq !== undefined) { + await commitSessionStreamPart( + authentication.environment.id, + addressingKey, + params.io, + partId, + pendingClaim.claimValue, + recoveredSeq + ); + // The S2 record is durable even if Redis expired or evicted the + // pending claim before the best-effort sequence write-back. + appendSeq = recoveredSeq; + claim = { status: "committed", seq: recoveredSeq }; + } + + // `.out` data may have been trimmed after it was accepted, so a + // successful search with no record preserves the prior no-op response. + // A failed search proves nothing and remains retryable for both channels. + if (claim.status === "pending" && (params.io === "in" || recoveryError)) { + return json( + { ok: false, error: "This append is still in progress, please retry." }, + { status: 409, headers: { "Retry-After": "1" } } + ); + } + } + + if (claim.status === "claimed") { const [appendError, seq] = await tryCatch( realtimeStream.appendPartToSessionStream(part, partId, addressingKey, params.io) ); @@ -174,7 +218,8 @@ const { action, loader } = createActionApiRoute( authentication.environment.id, addressingKey, params.io, - clientPartId + clientPartId, + claim.claimValue ); } if (appendError instanceof ServiceValidationError) { @@ -193,6 +238,17 @@ const { action, loader } = createActionApiRoute( { status: 500 } ); } + + if (clientPartId && appendSeq !== undefined) { + await commitSessionStreamPart( + authentication.environment.id, + addressingKey, + params.io, + clientPartId, + claim.claimValue, + appendSeq + ); + } } // Fire any run-scoped waitpoints registered against this channel. Best diff --git a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts index aff543d305..56d206da84 100644 --- a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts +++ b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts @@ -94,6 +94,18 @@ type S2AppendAck = { end: { seq_num: number; timestamp: number }; tail: { seq_num: number; timestamp: number }; }; +type S2ReadRecord = { + body: string; + seq_num: number; + timestamp: number; + headers?: Array<[string, string]>; +}; +type S2TailResponse = { + tail: { seq_num: number; timestamp: number }; +}; + +const PART_RECOVERY_CLOCK_SKEW_MS = 10_000; +const PART_RECOVERY_PAGE_SIZE = 1_000; export class S2RealtimeStreams implements StreamResponder, StreamIngestor { private readonly basin: string; @@ -264,6 +276,102 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { return this.#readRecordsByName(this.toSessionStreamName(friendlyId, io), afterSeqNum); } + /** + * Find an append whose Redis claim was left pending after S2 accepted it. + * Snapshots the stream tail, then reads forward from the claim timestamp. + * The fixed tail keeps an active stream from extending the recovery search. + */ + async findSessionStreamPartSequence( + friendlyId: string, + io: "out" | "in", + partId: string, + claimedAt?: number + ): Promise { + const s2Stream = this.toSessionStreamName(friendlyId, io); + const tailResponse = await fetch( + `${this.baseUrl}/streams/${encodeURIComponent(s2Stream)}/records/tail`, + { + method: "GET", + headers: { + Authorization: `Bearer ${this.token}`, + Accept: "application/json", + "S2-Basin": this.basin, + }, + } + ); + + if (!tailResponse.ok) { + if (tailResponse.status === 404) return undefined; + const text = await tailResponse.text().catch(() => ""); + throw new Error( + `S2 findSessionStreamPartSequence tail check failed: ${tailResponse.status} ${tailResponse.statusText} ${text}` + ); + } + + const { tail } = (await tailResponse.json()) as S2TailResponse; + const snapshotTail = tail.seq_num; + if (snapshotTail === 0) return undefined; + + let nextSeq: number | undefined; + + while (true) { + const qs = new URLSearchParams(); + if (nextSeq !== undefined) { + qs.set("seq_num", String(nextSeq)); + } else if (claimedAt !== undefined) { + qs.set("timestamp", String(Math.max(0, claimedAt - PART_RECOVERY_CLOCK_SKEW_MS))); + } else { + qs.set("seq_num", "0"); + } + qs.set("clamp", "true"); + qs.set("count", String(PART_RECOVERY_PAGE_SIZE)); + qs.set("bytes", String(S2_MAX_METERED_BYTES)); + qs.set("wait", "0"); + + const res = await fetch( + `${this.baseUrl}/streams/${encodeURIComponent(s2Stream)}/records?${qs}`, + { + method: "GET", + headers: { + Authorization: `Bearer ${this.token}`, + Accept: "application/json", + "S2-Format": "raw", + "S2-Basin": this.basin, + }, + } + ); + + if (!res.ok) { + if (res.status === 404 || res.status === 416) return undefined; + const text = await res.text().catch(() => ""); + throw new Error( + `S2 findSessionStreamPartSequence failed: ${res.status} ${res.statusText} ${text}` + ); + } + + const json = (await res.json()) as { + records?: S2ReadRecord[]; + }; + const records = json.records ?? []; + + for (const record of records) { + if (record.seq_num >= snapshotTail) break; + if (record.headers?.[0]?.[0] === "") continue; + try { + const envelope = JSON.parse(record.body) as { id?: unknown }; + if (envelope.id === partId) return record.seq_num; + } catch { + // A malformed or control record cannot be the requested data part. + } + } + + const lastRecord = records.at(-1); + if (!lastRecord) return undefined; + nextSeq = lastRecord.seq_num + 1; + if (nextSeq >= snapshotTail) return undefined; + } + } + async #readRecordsByName(s2Stream: string, afterSeqNum?: number): Promise { const startSeq = afterSeqNum != null ? afterSeqNum + 1 : 0; @@ -320,14 +428,7 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { if (eventType === "batch" && data) { try { - const parsed = JSON.parse(data) as { - records: Array<{ - body: string; - seq_num: number; - timestamp: number; - headers?: Array<[string, string]>; - }>; - }; + const parsed = JSON.parse(data) as { records: S2ReadRecord[] }; for (const record of parsed.records) { // S2 command records (trim/fence) have a single header with diff --git a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts index 7b53042d8d..8f59c5bfea 100644 --- a/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts +++ b/apps/webapp/app/services/sessionStreamWaitpointCache.server.ts @@ -1,3 +1,4 @@ +import { randomUUID } from "node:crypto"; import { Redis } from "ioredis"; import { defaultReconnectOnError } from "@internal/redis"; import { env } from "~/env.server"; @@ -146,15 +147,44 @@ export async function drainSessionStreamWaitpoints( * where `.wait()` completed the waitpoint from pre-arrived data. */ // "ssa" — session-stream-append. Idempotency claim for the append route: -// when a caller supplies an `X-Part-Id`, the first request atomically claims -// the key (SET NX) before appending; a concurrent or retried POST with the -// same id fails the claim and skips the append, so it never produces a -// duplicate record (or double-fires the waitpoint drain). The claim is -// released if the append fails, so a retry of a genuinely failed append -// still goes through. 5-minute window — covers retry storms, not a -// permanent idempotency store. +// when a caller supplies an `X-Part-Id`, the first request atomically stores a +// tokened pending claim before appending, then replaces it with the committed +// S2 sequence. A retry can therefore return the original sequence without +// producing a duplicate record. The pending value includes its creation time +// so recovery can search the relevant section of the S2 stream. This remains +// a 5-minute retry window, not a permanent idempotency store. const APPEND_DEDUPE_PREFIX = "ssa:"; const APPEND_DEDUPE_TTL_SECONDS = 5 * 60; +const APPEND_DEDUPE_PENDING_PREFIX = "pending:"; +const APPEND_DEDUPE_SEQUENCE_PREFIX = "seq:"; + +const COMMIT_APPEND_DEDUPE_SCRIPT = ` + local current = redis.call("GET", KEYS[1]) + if current == ARGV[2] then + return 1 + end + if current and current ~= ARGV[1] then + return 0 + end + local restoredMissingClaim = not current + redis.call("SET", KEYS[1], ARGV[2], "EX", ARGV[3]) + if restoredMissingClaim then + return 2 + end + return 1 +`; + +const RELEASE_APPEND_DEDUPE_SCRIPT = ` + if redis.call("GET", KEYS[1]) ~= ARGV[1] then + return 0 + end + return redis.call("DEL", KEYS[1]) +`; + +type SessionStreamPartClaim = + | { status: "claimed"; claimValue: string | undefined } + | { status: "committed"; seq: number } + | { status: "pending"; claimValue: string; claimedAt: number | undefined }; function buildAppendDedupeKey( environmentId: string, @@ -171,40 +201,116 @@ function buildAppendDedupeKey( } /** - * Atomically claim a part id before appending. Returns true if this caller - * won the claim (first to see this id) and should perform the append, false - * if the id was already claimed (a concurrent or retried POST) and the append - * should be skipped. Fails open (returns true) when Redis is unavailable — - * appends degrade to at-least-once, never to dropped. + * Atomically claim a part id before appending. Returns the caller's claim + * value, the sequence from an already committed append, or pending when the + * winning request has not committed yet. Fails open as an untracked claim + * when Redis is unavailable — appends degrade to at-least-once, never to + * dropped. */ export async function claimSessionStreamPart( environmentId: string, addressingKey: string, io: "out" | "in", partId: string +): Promise { + if (!redis) return { status: "claimed", claimValue: undefined }; + + try { + const key = buildAppendDedupeKey(environmentId, addressingKey, io, partId); + for (let attempt = 0; attempt < 3; attempt++) { + const claimedAt = Date.now(); + const pendingValue = `${APPEND_DEDUPE_PENDING_PREFIX}${claimedAt}:${randomUUID()}`; + // SET NX is the atomic claim: "OK" when set (we won), null when the key + // already exists (someone else owns this id). + const result = await redis.set(key, pendingValue, "EX", APPEND_DEDUPE_TTL_SECONDS, "NX"); + if (result === "OK") { + return { status: "claimed", claimValue: pendingValue }; + } + + const existing = await redis.get(key); + if (existing?.startsWith(APPEND_DEDUPE_SEQUENCE_PREFIX)) { + const seq = Number(existing.slice(APPEND_DEDUPE_SEQUENCE_PREFIX.length)); + if (Number.isSafeInteger(seq) && seq >= 0) { + return { status: "committed", seq }; + } + } + + if (existing) { + let pendingClaimedAt = parsePendingClaimedAt(existing); + if (pendingClaimedAt === undefined) { + // Claims from the previous release stored only "1". Estimate their + // creation time from the remaining TTL so recovery starts near the + // original append during a rolling deploy. + const ttlMs = await redis.pttl(key); + if (ttlMs >= 0) { + pendingClaimedAt = Date.now() - Math.max(0, APPEND_DEDUPE_TTL_SECONDS * 1000 - ttlMs); + } + } + return { + status: "pending", + claimValue: existing, + claimedAt: pendingClaimedAt, + }; + } + + // The key expired or was evicted between SET NX and GET. Retry the + // reservation a small number of times without growing the call stack. + } + + // Redis kept losing the reservation. Fail open like other cache errors so + // this append degrades to at-least-once instead of failing the request. + return { status: "claimed", claimValue: undefined }; + } catch (error) { + logger.error("Failed to claim session stream append part", { + environmentId, + addressingKey, + io, + partId, + error, + }); + return { status: "claimed", claimValue: undefined }; + } +} + +function parsePendingClaimedAt(value: string): number | undefined { + if (!value.startsWith(APPEND_DEDUPE_PENDING_PREFIX)) return undefined; + const separator = value.indexOf(":", APPEND_DEDUPE_PENDING_PREFIX.length); + if (separator === -1) return undefined; + const claimedAt = Number(value.slice(APPEND_DEDUPE_PENDING_PREFIX.length, separator)); + return Number.isSafeInteger(claimedAt) && claimedAt >= 0 ? claimedAt : undefined; +} + +/** Persist a committed S2 sequence if the expected claim still owns the key or the key vanished. */ +export async function commitSessionStreamPart( + environmentId: string, + addressingKey: string, + io: "out" | "in", + partId: string, + claimValue: string | undefined, + seq: number ): Promise { - if (!redis) return true; + if (!redis || !claimValue) return false; try { - // SET NX is the atomic claim: "OK" when set (we won), null when the key - // already exists (someone else owns this id). - const result = await redis.set( + const result = await redis.eval( + COMMIT_APPEND_DEDUPE_SCRIPT, + 1, buildAppendDedupeKey(environmentId, addressingKey, io, partId), - "1", - "EX", - APPEND_DEDUPE_TTL_SECONDS, - "NX" + claimValue, + `${APPEND_DEDUPE_SEQUENCE_PREFIX}${seq}`, + String(APPEND_DEDUPE_TTL_SECONDS) ); - return result === "OK"; + return result === 1 || result === 2; } catch (error) { - logger.error("Failed to claim session stream append part", { + logger.error("Failed to commit session stream append part", { environmentId, addressingKey, io, partId, + seq, error, }); - return true; + return false; } } @@ -213,12 +319,18 @@ export async function releaseSessionStreamPart( environmentId: string, addressingKey: string, io: "out" | "in", - partId: string + partId: string, + claimValue: string | undefined ): Promise { - if (!redis) return; + if (!redis || !claimValue) return; try { - await redis.del(buildAppendDedupeKey(environmentId, addressingKey, io, partId)); + await redis.eval( + RELEASE_APPEND_DEDUPE_SCRIPT, + 1, + buildAppendDedupeKey(environmentId, addressingKey, io, partId), + claimValue + ); } catch (error) { logger.error("Failed to release session stream append part", { environmentId, diff --git a/apps/webapp/app/v3/webhookEngine.server.ts b/apps/webapp/app/v3/webhookEngine.server.ts index d58a89919b..019e0f72ec 100644 --- a/apps/webapp/app/v3/webhookEngine.server.ts +++ b/apps/webapp/app/v3/webhookEngine.server.ts @@ -15,6 +15,7 @@ import { findOrCreateSession, findSessionByExternalId } from "~/services/realtim import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server"; import { claimSessionStreamPart, + commitSessionStreamPart, drainSessionStreamWaitpoints, releaseSessionStreamPart, } from "~/services/sessionStreamWaitpointCache.server"; @@ -221,27 +222,65 @@ function createWebhookEngine() { }; const part = JSON.stringify({ kind: "message", payload }); - // deliveryId as the part id → a deliver-job retry re-claims the same id and skips a duplicate - // append. The S2 record is durable, so a run that boots later still reads it. - const wonClaim = await claimSessionStreamPart( - environment.id, - addressingKey, - "in", - deliveryId - ); - if (wonClaim) { - const [appendError] = await tryCatch( + // deliveryId as the part id → a deliver-job retry reuses the committed claim and skips a + // duplicate append. The S2 record is durable, so a run that boots later still reads it. + let claim = await claimSessionStreamPart(environment.id, addressingKey, "in", deliveryId); + if (claim.status === "pending") { + const pendingClaim = claim; + const recoveredSeq = await realtimeStream.findSessionStreamPartSequence( + addressingKey, + "in", + deliveryId, + pendingClaim.claimedAt + ); + if (recoveredSeq !== undefined) { + await commitSessionStreamPart( + environment.id, + addressingKey, + "in", + deliveryId, + pendingClaim.claimValue, + recoveredSeq + ); + // The S2 record is durable even if Redis expired or evicted the + // pending claim before the best-effort sequence write-back. + claim = { status: "committed", seq: recoveredSeq }; + } + + if (claim.status === "pending") { + throw new Error("Session stream append is still in progress"); + } + } + + if (claim.status === "claimed") { + const [appendError, appendSeq] = await tryCatch( realtimeStream.appendPartToSessionStream(part, deliveryId, addressingKey, "in") ); if (appendError) { // Nothing landed — release the claim so a retry re-appends the same id. - await releaseSessionStreamPart(environment.id, addressingKey, "in", deliveryId); + await releaseSessionStreamPart( + environment.id, + addressingKey, + "in", + deliveryId, + claim.claimValue + ); // A ServiceValidationError (e.g. record too large) is terminal; anything else is transient. if (appendError instanceof ServiceValidationError) { return { success: false, error: appendError.message }; } throw appendError; } + if (typeof appendSeq === "number") { + await commitSessionStreamPart( + environment.id, + addressingKey, + "in", + deliveryId, + claim.claimValue, + appendSeq + ); + } } // Wake any `.in` waitpoints the run registered (best-effort; the record is durable in S2). diff --git a/apps/webapp/test/helpers/sessionStream.ts b/apps/webapp/test/helpers/sessionStream.ts index 9e0f4e7c86..ce56d478da 100644 --- a/apps/webapp/test/helpers/sessionStream.ts +++ b/apps/webapp/test/helpers/sessionStream.ts @@ -122,8 +122,9 @@ export async function appendInput(opts: { body: string; partId?: string; origin?: string; + io?: "out" | "in"; }): Promise<{ status: number; acao: string | null; json: unknown }> { - const url = `${sessionChannelUrl(opts.baseUrl, opts.addressingKey, "in")}/append`; + const url = `${sessionChannelUrl(opts.baseUrl, opts.addressingKey, opts.io ?? "in")}/append`; const res = await fetch(url, { method: "POST", headers: { diff --git a/apps/webapp/test/session-stream.e2e.test.ts b/apps/webapp/test/session-stream.e2e.test.ts index af1b631ee1..8f55cdaa7b 100644 --- a/apps/webapp/test/session-stream.e2e.test.ts +++ b/apps/webapp/test/session-stream.e2e.test.ts @@ -11,6 +11,7 @@ * Requires a pre-built webapp: pnpm run build --filter webapp */ import { randomBytes } from "crypto"; +import { Redis } from "ioredis"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { SessionStreamInstance } from "@trigger.dev/core/v3"; import type { SessionStreamTestServer } from "@internal/testcontainers/webapp"; @@ -61,12 +62,33 @@ async function setupSession() { envId: environment.id, addressingKey, }); + const inStreamName = sessionStreamName({ + orgId: organization.id, + envSlug: environment.slug, + envId: environment.id, + addressingKey, + io: "in", + }); const producer = new SessionStreamProducer({ endpoint: server.s2.endpoint, basin: server.s2.basin, streamName, }); - return { addressingKey, token, producer, streamName, baseUrl: server.webapp.baseUrl }; + const inProducer = new SessionStreamProducer({ + endpoint: server.s2.endpoint, + basin: server.s2.basin, + streamName: inStreamName, + }); + return { + addressingKey, + environmentId: environment.id, + apiKey, + token, + producer, + inProducer, + streamName, + baseUrl: server.webapp.baseUrl, + }; } function readableFrom(chunks: T[]): ReadableStream { @@ -311,7 +333,261 @@ describe("session stream e2e", () => { expect(String(got?.chunk)).toContain("hello from client"); }); - it("E12 subscribe with an invalid token is rejected", async () => { + it("E12 in/append retry returns the original seq without duplicating the record", async () => { + const { addressingKey, token, baseUrl } = await setupSession(); + + const payload = JSON.stringify({ kind: "message", text: "lost response" }); + const partId = `retry-${randomBytes(6).toString("hex")}`; + + // The first append commits. Model its response being lost by issuing the + // same idempotent request again; keep the first seq only as a test oracle. + const first = await appendInput({ baseUrl, addressingKey, token, partId, body: payload }); + expect(first.status).toBe(200); + expect(first.json).toEqual({ ok: true, seq: expect.any(Number) }); + const firstSeq = (first.json as { seq: number }).seq; + + const retry = await appendInput({ baseUrl, addressingKey, token, partId, body: payload }); + expect(retry.status).toBe(200); + expect(retry.json).toEqual({ ok: true, seq: firstSeq }); + const retrySeq = (retry.json as { seq: number }).seq; + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + io: "in", + timeoutInSeconds: 1, + maxMs: 5_000, + }); + const dataParts = parts.filter((part) => part.chunk != null); + expect(dataParts).toHaveLength(1); + expect(Number(dataParts[0]!.id)).toBe(retrySeq); + }); + + it("E13 in/append recovers a crash after S2 accepted the record", async () => { + const { addressingKey, environmentId, token, inProducer, baseUrl } = await setupSession(); + const redis = new Redis({ ...server.redis, keyPrefix: "tr:" }); + + try { + const payload = JSON.stringify({ kind: "message", text: "accepted before crash" }); + const partId = `crash-${randomBytes(6).toString("hex")}`; + const claimKey = `ssa:${encodeURIComponent(environmentId)}:${encodeURIComponent( + addressingKey + )}:in:${encodeURIComponent(partId)}`; + const waitpointKey = `ssw:${environmentId}:${addressingKey}:in`; + + // Seed the exact durable state left by a process dying after S2 accepts + // the append but before Redis records its sequence. + await redis.set(claimKey, `pending:${Date.now()}:crashed-owner`, "EX", 5 * 60); + await redis.sadd(waitpointKey, "waitpoint_crash_test"); + const originalSeq = await inProducer.appendData(payload, partId); + + const retry = await appendInput({ baseUrl, addressingKey, token, partId, body: payload }); + expect(retry.status).toBe(200); + expect(retry.json).toEqual({ ok: true, seq: originalSeq }); + expect(await redis.get(claimKey)).toBe(`seq:${originalSeq}`); + expect(await redis.exists(waitpointKey)).toBe(0); + + const committedRetry = await appendInput({ + baseUrl, + addressingKey, + token, + partId, + body: payload, + }); + expect(committedRetry.status).toBe(200); + expect(committedRetry.json).toEqual({ ok: true, seq: originalSeq }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + io: "in", + timeoutInSeconds: 1, + maxMs: 5_000, + }); + const dataParts = parts.filter((part) => part.chunk != null); + expect(dataParts).toHaveLength(1); + expect(Number(dataParts[0]!.id)).toBe(originalSeq); + } finally { + await redis.quit(); + } + }); + + it("E14 in/append recovers a sequence from a legacy claim", async () => { + const { addressingKey, environmentId, token, inProducer, baseUrl } = await setupSession(); + const redis = new Redis({ ...server.redis, keyPrefix: "tr:" }); + + try { + const payload = JSON.stringify({ kind: "message", text: "rolling deploy retry" }); + const partId = `legacy-${randomBytes(6).toString("hex")}`; + const claimKey = `ssa:${encodeURIComponent(environmentId)}:${encodeURIComponent( + addressingKey + )}:in:${encodeURIComponent(partId)}`; + + // The previous release stored literal "1" for both pending and accepted + // appends. Its remaining TTL supplies the recovery window after deploy. + await redis.set(claimKey, "1", "EX", 5 * 60); + const originalSeq = await inProducer.appendData(payload, partId); + + const retry = await appendInput({ baseUrl, addressingKey, token, partId, body: payload }); + expect(retry.status).toBe(200); + expect(retry.json).toEqual({ ok: true, seq: originalSeq }); + expect(await redis.get(claimKey)).toBe(`seq:${originalSeq}`); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + io: "in", + timeoutInSeconds: 1, + maxMs: 5_000, + }); + expect(parts.filter((part) => part.chunk != null)).toHaveLength(1); + } finally { + await redis.quit(); + } + }); + + it("E15 in/append recovery searches beyond 32 full read pages", async () => { + const { addressingKey, environmentId, token, inProducer, baseUrl } = await setupSession(); + const redis = new Redis({ ...server.redis, keyPrefix: "tr:" }); + + try { + const payload = JSON.stringify({ kind: "message", text: "after large records" }); + const partId = `deep-${randomBytes(6).toString("hex")}`; + const claimKey = `ssa:${encodeURIComponent(environmentId)}:${encodeURIComponent( + addressingKey + )}:in:${encodeURIComponent(partId)}`; + + await redis.set(claimKey, `pending:${Date.now()}:crashed-owner`, "EX", 5 * 60); + + // Each record is over half the 1 MiB read budget, forcing one record per + // page. The target therefore sits on page 33 of the recovery scan. + const largeDecoy = "x".repeat(530_000); + for (let index = 0; index < 32; index++) { + await inProducer.appendData(largeDecoy, `decoy-${index}`); + } + const originalSeq = await inProducer.appendData(payload, partId); + + const retry = await appendInput({ baseUrl, addressingKey, token, partId, body: payload }); + expect(retry.status).toBe(200); + expect(retry.json).toEqual({ ok: true, seq: originalSeq }); + expect(await redis.get(claimKey)).toBe(`seq:${originalSeq}`); + } finally { + await redis.quit(); + } + }); + + it("E16 in/append recovery survives an expired Redis claim", async () => { + const { addressingKey, environmentId, token, inProducer, baseUrl } = await setupSession(); + const redis = new Redis({ ...server.redis, keyPrefix: "tr:" }); + const monitorSource = new Redis(server.redis); + let monitor: Redis | undefined; + + try { + const payload = JSON.stringify({ kind: "message", text: "expired during recovery" }); + const partId = `expired-${randomBytes(6).toString("hex")}`; + const claimKey = `ssa:${encodeURIComponent(environmentId)}:${encodeURIComponent( + addressingKey + )}:in:${encodeURIComponent(partId)}`; + const waitpointKey = `ssw:${environmentId}:${addressingKey}:in`; + + await redis.set(claimKey, `pending:${Date.now()}:crashed-owner`, "EX", 5 * 60); + await redis.sadd(waitpointKey, "waitpoint_expired_claim_test"); + const originalSeq = await inProducer.appendData(payload, partId); + + monitor = await monitorSource.monitor(); + const claimRead = new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("Timed out waiting for claim read")), + 5_000 + ); + const onMonitor = (_time: string, args: string[]) => { + if (args[0]?.toLowerCase() !== "get" || args[1] !== `tr:${claimKey}`) return; + monitor!.off("monitor", onMonitor); + void redis.del(claimKey).then( + () => { + clearTimeout(timeout); + resolve(); + }, + (error) => { + clearTimeout(timeout); + reject(error); + } + ); + }; + monitor.on("monitor", onMonitor); + }); + + const retryPromise = appendInput({ baseUrl, addressingKey, token, partId, body: payload }); + await claimRead; + const retry = await retryPromise; + + expect(retry.status).toBe(200); + expect(retry.json).toEqual({ ok: true, seq: originalSeq }); + expect(await redis.get(claimKey)).toBe(`seq:${originalSeq}`); + expect(await redis.exists(waitpointKey)).toBe(0); + + const committedRetry = await appendInput({ + baseUrl, + addressingKey, + token, + partId, + body: payload, + }); + expect(committedRetry.status).toBe(200); + expect(committedRetry.json).toEqual({ ok: true, seq: originalSeq }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + io: "in", + timeoutInSeconds: 1, + maxMs: 5_000, + }); + expect(parts.filter((part) => part.chunk != null)).toHaveLength(1); + } finally { + monitor?.disconnect(); + await monitorSource.quit(); + await redis.quit(); + } + }); + + it("E17 out/append keeps a missing recovered record as a successful no-op", async () => { + const { addressingKey, environmentId, apiKey, baseUrl } = await setupSession(); + const redis = new Redis({ ...server.redis, keyPrefix: "tr:" }); + + try { + const payload = JSON.stringify({ kind: "message", text: "already trimmed" }); + const partId = `trimmed-${randomBytes(6).toString("hex")}`; + const claimKey = `ssa:${encodeURIComponent(environmentId)}:${encodeURIComponent( + addressingKey + )}:out:${encodeURIComponent(partId)}`; + const waitpointKey = `ssw:${environmentId}:${addressingKey}:out`; + + await redis.set(claimKey, `pending:${Date.now()}:crashed-owner`, "EX", 5 * 60); + await redis.sadd(waitpointKey, "waitpoint_trimmed_claim_test"); + + const retry = await appendInput({ + baseUrl, + addressingKey, + token: apiKey, + partId, + body: payload, + io: "out", + }); + + expect(retry.status).toBe(200); + expect(retry.json).toEqual({ ok: true }); + expect(await redis.exists(waitpointKey)).toBe(0); + } finally { + await redis.quit(); + } + }); + + it("E18 subscribe with an invalid token is rejected", async () => { const { addressingKey, baseUrl } = await setupSession(); const { status } = await openChannelRaw({ diff --git a/packages/core/src/v3/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index 6cd100f7c3..73cc8aefec 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -2181,6 +2181,7 @@ export type CreateStreamResponseBody = z.infer; export const AppendToStreamResponseBody = z.object({ ok: z.boolean(), message: z.string().optional(), + seq: z.number().optional(), }); export type AppendToStreamResponseBody = z.infer; diff --git a/packages/trigger-sdk/src/v3/chat-server.test.ts b/packages/trigger-sdk/src/v3/chat-server.test.ts index 539fe0247a..75382da2b5 100644 --- a/packages/trigger-sdk/src/v3/chat-server.test.ts +++ b/packages/trigger-sdk/src/v3/chat-server.test.ts @@ -119,8 +119,8 @@ function createSessionResponse(externalId: string): Response { ); } -function appendOkResponse(): Response { - return new Response(JSON.stringify({ ok: true }), { +function appendOkResponse(seq?: number): Response { + return new Response(JSON.stringify(seq === undefined ? { ok: true } : { ok: true, seq }), { status: 200, headers: { "content-type": "application/json" }, }); @@ -289,7 +289,7 @@ describe("chat.headStart (route handler)", () => { return createSessionResponse("chat-final"); } if (urlStr.includes("/realtime/v1/sessions/") && urlStr.endsWith("/in/append")) { - return appendOkResponse(); + return appendOkResponse(17); } // Stitched response subscribes to `.out` after handover. if (/\/realtime\/v1\/sessions\/[^/]+\/out$/.test(urlStr)) { @@ -336,6 +336,9 @@ describe("chat.headStart (route handler)", () => { // Drain the SSE body so handoverWhenDone observes finishReason. const chunks = await readSSEBodyToChunks(res); expect(chunks.some((c) => c.type === "text-delta")).toBe(true); + expect(chunks).toContainEqual( + expect.objectContaining({ type: "trigger:session-state", activeInputSeq: 17 }) + ); // Give the deferred handoverWhenDone a tick to dispatch. await new Promise((r) => setTimeout(r, 30)); diff --git a/packages/trigger-sdk/src/v3/chat-server.ts b/packages/trigger-sdk/src/v3/chat-server.ts index 5e48e3b24b..c16818c120 100644 --- a/packages/trigger-sdk/src/v3/chat-server.ts +++ b/packages/trigger-sdk/src/v3/chat-server.ts @@ -662,7 +662,7 @@ async function openHandoverSession(opts: { } }; - const handover = async (args: { + const dispatchHandover = async (args: { partialAssistantMessage: ModelMessage[]; messageId?: string; isFinal: boolean; @@ -673,7 +673,16 @@ async function openHandoverSession(opts: { messageId: args.messageId, isFinal: args.isFinal, }; - await apiClient.appendToSessionStream(chatId, "in", JSON.stringify(chunk)); + const result = await apiClient.appendToSessionStream(chatId, "in", JSON.stringify(chunk)); + return result.seq; + }; + + const handover = async (args: { + partialAssistantMessage: ModelMessage[]; + messageId?: string; + isFinal: boolean; + }) => { + await dispatchHandover(args); }; /** @@ -707,7 +716,7 @@ async function openHandoverSession(opts: { // and dispatches the handover decision. The stitched response stream // awaits this to know whether to close (skip) or pull more chunks // from session.out (handover). - type HandoverDecision = { kind: "handover" | "handover-skip" }; + type HandoverDecision = { kind: "handover"; activeInputSeq?: number } | { kind: "handover-skip" }; let resolveDecision!: (decision: HandoverDecision) => void; const decisionPromise = new Promise((resolve) => { resolveDecision = resolve; @@ -737,11 +746,12 @@ async function openHandoverSession(opts: { // so the agent's `streamText` resumes by executing them // before the step-2 LLM call. const reshaped = reshapeForHandoverResume(responseMessages); - await handover({ + const activeInputSeq = await dispatchHandover({ partialAssistantMessage: reshaped, messageId: turnMessageId, isFinal: false, }); + resolveDecision({ kind: "handover", activeInputSeq }); } else { // Pure-text (or any non-tool-calls) finish — customer's step 1 // IS the final response. The agent runs the turn-loop hooks @@ -749,13 +759,13 @@ async function openHandoverSession(opts: { // this partial as the response, but skips the LLM call. That // way persistence (`onTurnComplete` writing to DB), self- // review, and any post-turn work all fire normally. - await handover({ + const activeInputSeq = await dispatchHandover({ partialAssistantMessage: responseMessages, messageId: turnMessageId, isFinal: true, }); + resolveDecision({ kind: "handover", activeInputSeq }); } - resolveDecision({ kind: "handover" }); } catch (err) { // Dispatch failed before we could send the handover signal. // Tell the agent to exit clean (no hooks fire) and close the @@ -812,6 +822,18 @@ async function openHandoverSession(opts: { return; } + // The handover append happens after the HTTP response has started, + // so its `.in` sequence cannot be returned in a response header. + // Send it as transport-only state before any agent output so a + // reload during the remainder of this turn can reject stale + // turn-complete records. + if (decision.activeInputSeq !== undefined) { + controller.enqueue({ + type: "trigger:session-state", + activeInputSeq: decision.activeInputSeq, + } as unknown as UIMessageChunk); + } + // Phase 2b: agent is taking over. Resume from session.out // starting AFTER the customer tee's last write, so we don't // re-emit chunks the browser already saw. diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index 452c8090b7..3dfc98a410 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -35,6 +35,10 @@ function sseEncode(chunks: (UIMessageChunk | Record)[]): string const headers: Array<[string, string]> = [["trigger-control", "turn-complete"]]; const token = (chunk as { publicAccessToken?: string }).publicAccessToken; if (token) headers.push(["public-access-token", token]); + const sessionInEventId = (chunk as { sessionInEventId?: string | number }).sessionInEventId; + if (sessionInEventId !== undefined) { + headers.push(["session-in-event-id", String(sessionInEventId)]); + } return { body: "", seq_num: nextSeq++, @@ -228,6 +232,7 @@ describe("TriggerChatTransport", () => { "chat-1": { publicAccessToken: "hydrated-pat", lastEventId: "42", + activeInputSeq: 41, isStreaming: false, }, }, @@ -237,6 +242,7 @@ describe("TriggerChatTransport", () => { expect(session).toEqual({ publicAccessToken: "hydrated-pat", lastEventId: "42", + activeInputSeq: 41, isStreaming: false, }); }); @@ -262,15 +268,21 @@ describe("TriggerChatTransport", () => { transport.setSession("chat-x", { publicAccessToken: "tok", lastEventId: "10", + activeInputSeq: 9, }); expect(transport.getSession("chat-x")).toMatchObject({ publicAccessToken: "tok", lastEventId: "10", + activeInputSeq: 9, }); expect(onSessionChange).toHaveBeenCalledWith( "chat-x", - expect.objectContaining({ publicAccessToken: "tok", lastEventId: "10" }) + expect.objectContaining({ + publicAccessToken: "tok", + lastEventId: "10", + activeInputSeq: 9, + }) ); }); @@ -977,7 +989,9 @@ describe("TriggerChatTransport", () => { it("marks the session streaming and notifies before subscribing", async () => { global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { const urlStr = typeof url === "string" ? url : url.toString(); - if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse(); + if (isSessionStreamAppendUrl(urlStr)) { + return new Response(JSON.stringify({ ok: true, seq: 7 }), { status: 200 }); + } if (isSessionOutSubscribeUrl(urlStr)) return defaultSseResponse(); throw new Error(`Unexpected URL: ${urlStr}`); }); @@ -994,7 +1008,9 @@ describe("TriggerChatTransport", () => { // isStreaming:true must be observed during the action — otherwise a reload // mid-action sees a persisted isStreaming:false and never resumes. expect( - onSessionChange.mock.calls.some(([, session]) => session && session.isStreaming === true) + onSessionChange.mock.calls.some( + ([, session]) => session && session.isStreaming === true && session.activeInputSeq === 7 + ) ).toBe(true); await drainChunks(stream); }); @@ -1750,14 +1766,17 @@ describe("TriggerChatTransport", () => { * does: `data: \n\n` per chunk. The transport's * `parseUIMessageSseTransform` parses this back into chunk objects. */ - function handoverSseBody(chunks: UIMessageChunk[]): ReadableStream { + function handoverSseBody( + chunks: (UIMessageChunk | Record)[], + close = true + ): ReadableStream { const encoder = new TextEncoder(); return new ReadableStream({ start(controller) { for (const chunk of chunks) { controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`)); } - controller.close(); + if (close) controller.close(); }, }); } @@ -1765,9 +1784,10 @@ describe("TriggerChatTransport", () => { function handoverResponse(args: { chatId: string; accessToken: string; - chunks: UIMessageChunk[]; + chunks: (UIMessageChunk | Record)[]; + close?: boolean; }): Response { - return new Response(handoverSseBody(args.chunks), { + return new Response(handoverSseBody(args.chunks, args.close), { status: 200, headers: { "content-type": "text/event-stream", @@ -1902,6 +1922,70 @@ describe("TriggerChatTransport", () => { expect(subscribe).toBeDefined(); }); + it("persists the handover input sequence for a reload while the first turn is active", async () => { + global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (urlStr === "https://my-app.example/api/chat") { + return handoverResponse({ + chatId: "chat-handover-reload", + accessToken: "handover-pat-reload", + chunks: [ + { type: "trigger:session-state", activeInputSeq: 7 }, + { type: "text-delta", id: "part-1", delta: "working" }, + ], + close: false, + }); + } + if (isSessionOutSubscribeUrl(urlStr)) { + return defaultSseResponse([ + { type: "trigger:turn-complete", sessionInEventId: 6 }, + { type: "text-delta", id: "part-2", delta: "current" }, + { type: "trigger:turn-complete", sessionInEventId: 7 }, + ]); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + headStart: "https://my-app.example/api/chat", + }); + + const firstTurn = await transport.sendMessages({ + trigger: "submit-message", + chatId: "chat-handover-reload", + messageId: "m1", + messages: [createUserMessage("first")], + abortSignal: undefined, + }); + const firstTurnReader = firstTurn.getReader(); + await expect(firstTurnReader.read()).resolves.toMatchObject({ + value: { type: "text-delta", delta: "working" }, + }); + + const persisted = transport.getSession("chat-handover-reload"); + expect(persisted).toMatchObject({ + publicAccessToken: "handover-pat-reload", + activeInputSeq: 7, + isStreaming: true, + }); + await firstTurnReader.cancel(); + + const rehydrated = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + sessions: { "chat-handover-reload": persisted! }, + }); + const resumed = await rehydrated.reconnectToStream({ chatId: "chat-handover-reload" }); + + expect(resumed).not.toBeNull(); + await expect(drainChunks(resumed!)).resolves.toEqual([ + { type: "text-delta", id: "part-2", delta: "current" }, + ]); + expect(rehydrated.getSession("chat-handover-reload")?.activeInputSeq).toBeUndefined(); + }); + it("bypasses endpoint when a session is already hydrated (page reload after first turn)", async () => { const requests: Array<{ url: string; init?: RequestInit }> = []; global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => { diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index a7c7125575..15ab70e96d 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -422,11 +422,13 @@ export type StartSessionResult = { * Public surface of {@link TriggerChatTransport}'s session state. Everything * the customer should persist for resumption across page reloads. The * transport addresses by `chatId` everywhere, so this is light: just a PAT, - * the last SSE event id, and a couple of UX-state flags. + * resume cursors, and a couple of UX-state flags. */ export type ChatSessionPersistedState = { publicAccessToken: string; lastEventId?: string; + /** The `.in` append sequence of the last send this client owned; reused as `sinceInSeq` on reconnect. */ + activeInputSeq?: number; isStreaming?: boolean; }; @@ -631,6 +633,8 @@ type ChatSessionState = { publicAccessToken: string; /** Last SSE event ID — used to resume the stream without replaying old events. */ lastEventId?: string; + /** `.in` append sequence used to filter stale turn boundaries after reconnecting. */ + activeInputSeq?: number; /** Set when the stream was aborted mid-turn (stop). On reconnect, skip chunks until trigger:turn-complete. */ skipToTurnComplete?: boolean; /** Whether the agent is currently streaming a response. Set on first chunk, cleared on turn-complete. */ @@ -718,6 +722,7 @@ export class TriggerChatTransport implements ChatTransport { this.sessions.set(chatId, { publicAccessToken: session.publicAccessToken, lastEventId: session.lastEventId, + activeInputSeq: session.activeInputSeq, isStreaming: session.isStreaming, }); } @@ -870,6 +875,7 @@ export class TriggerChatTransport implements ChatTransport { this.activeStreams.delete(chatId); } + state.activeInputSeq = inSeq; state.isStreaming = true; this.notifySessionChange(chatId, state); @@ -975,10 +981,10 @@ export class TriggerChatTransport implements ChatTransport { // useChat resume / reconnectToStream path doesn't open a // second `session.out` subscription on top of our stitched // response. - // - On `trigger:session-state`, hydrate `state.lastEventId` - // with the agent's final S2 event id. Without this, turn 2's - // `session.out` subscribe reads from the start and replays - // turn 1's chunks back into the UI. + // - On `trigger:session-state`, hydrate `state.activeInputSeq` + // as soon as the handover append commits and `state.lastEventId` + // when the agent stream finishes. These keep reloads correlated + // to the active turn and keep turn 2 from replaying turn 1. // - On stream end (handover-skip case — no // `trigger:turn-complete` arrives, customer's stream just // ends), also clear `isStreaming` for the same reason. @@ -987,9 +993,17 @@ export class TriggerChatTransport implements ChatTransport { this.notifySessionChange(id, state); const TRIGGER_TURN_COMPLETE = "trigger:turn-complete"; const TRIGGER_SESSION_STATE = "trigger:session-state"; + const clearActiveTurn = () => { + const state = sessions.get(chatId); + if (state && (state.isStreaming || state.activeInputSeq !== undefined)) { + state.activeInputSeq = undefined; + state.isStreaming = false; + notifyChange(chatId, state); + } + }; const clearStreaming = () => { const state = sessions.get(chatId); - if (state && state.isStreaming) { + if (state?.isStreaming) { state.isStreaming = false; notifyChange(chatId, state); } @@ -1001,6 +1015,13 @@ export class TriggerChatTransport implements ChatTransport { notifyChange(chatId, state); } }; + const setActiveInputSeq = (activeInputSeq: number) => { + const state = sessions.get(chatId); + if (state) { + state.activeInputSeq = activeInputSeq; + notifyChange(chatId, state); + } + }; const emit = (event: ChatTransportEvent) => this.emitEvent(event); const attribution = () => this.turnAttribution(chatId); let sawFirstChunk = false; @@ -1022,7 +1043,7 @@ export class TriggerChatTransport implements ChatTransport { if (chunk && typeof chunk === "object") { const type = (chunk as { type?: unknown }).type; if (type === TRIGGER_TURN_COMPLETE) { - clearStreaming(); + clearActiveTurn(); emit({ type: "turn-completed", chatId, @@ -1033,10 +1054,17 @@ export class TriggerChatTransport implements ChatTransport { return; // drop — not a real UIMessageChunk } if (type === TRIGGER_SESSION_STATE) { - const lastEventId = (chunk as { lastEventId?: unknown }).lastEventId; + const sessionState = chunk as { + lastEventId?: unknown; + activeInputSeq?: unknown; + }; + const lastEventId = sessionState.lastEventId; if (typeof lastEventId === "string") { setLastEventId(lastEventId); } + if (typeof sessionState.activeInputSeq === "number") { + setActiveInputSeq(sessionState.activeInputSeq); + } return; // drop } } @@ -1177,13 +1205,14 @@ export class TriggerChatTransport implements ChatTransport { return this.subscribeToSessionStream(state, abortSignal, options.chatId, { resumed: true, sendStopOnAbort: options.stopOnAbort ?? false, + sinceInSeq: state.activeInputSeq, // Reconnect-on-reload opts into the server's settled-peek shortcut - // so the SSE doesn't hang for 60s when no turn is in flight. Active - // send-a-message paths must keep wait=60 to avoid racing the - // freshly-triggered turn's first chunk. Watch mode must NOT peek: a - // settled peek between turns sets sessionSettled and closes the + // so the SSE doesn't hang for 60s when no turn is in flight. A known + // active input must not peek because the previous turn's completion + // can remain at the tail until the current turn writes its first chunk. + // Watch mode must NOT peek: a settled peek between turns closes the // standing subscription, so the viewer never sees the next turn. - peekSettled: !this.watchMode, + peekSettled: !this.watchMode && state.activeInputSeq === undefined, }); }; @@ -1283,6 +1312,7 @@ export class TriggerChatTransport implements ChatTransport { // Mark streaming + persist so a reload mid-action resumes (reconnectToStream // no-ops when the persisted session says isStreaming: false). + state.activeInputSeq = inSeq; state.isStreaming = true; this.notifySessionChange(chatId, state); @@ -1307,6 +1337,7 @@ export class TriggerChatTransport implements ChatTransport { this.sessions.set(chatId, { publicAccessToken: session.publicAccessToken, lastEventId: session.lastEventId, + activeInputSeq: session.activeInputSeq, isStreaming: session.isStreaming, }); this.notifySessionChange(chatId, this.toPersisted(this.sessions.get(chatId)!)); @@ -1441,6 +1472,7 @@ export class TriggerChatTransport implements ChatTransport { private toPersisted = (state: ChatSessionState): ChatSessionPersistedState => ({ publicAccessToken: state.publicAccessToken, lastEventId: state.lastEventId, + activeInputSeq: state.activeInputSeq, isStreaming: state.isStreaming, }); @@ -1750,6 +1782,7 @@ export class TriggerChatTransport implements ChatTransport { }) as typeof fetch) : undefined; let sawFirstChunk = false; + let sinceInSeq = options?.sinceInSeq; const connectSseOnce = async (token: string) => { const subscription = new SSEStreamSubscription(streamUrl, { @@ -1983,10 +2016,10 @@ export class TriggerChatTransport implements ChatTransport { if (controlValue === TRIGGER_CONTROL_SUBTYPE.TURN_COMPLETE) { // Skip a turn-complete from an earlier turn (committed `.in` cursor // below this send's seq), e.g. an undo action that raced this send. - if (options?.sinceInSeq !== undefined) { + if (sinceInSeq !== undefined) { const cursorRaw = headerValue(value.headers, SESSION_IN_EVENT_ID_HEADER); const cursor = cursorRaw !== undefined ? Number.parseInt(cursorRaw, 10) : NaN; - if (!Number.isNaN(cursor) && cursor < options.sinceInSeq) { + if (!Number.isNaN(cursor) && cursor < sinceInSeq) { continue; } } @@ -2004,6 +2037,8 @@ export class TriggerChatTransport implements ChatTransport { sessionInEventId: headerValue(value.headers, SESSION_IN_EVENT_ID_HEADER), ...this.turnAttribution(chatId), }); + state.activeInputSeq = undefined; + sinceInSeq = undefined; state.isStreaming = false; this.notifySessionChange(chatId, state); this.coordinator?.release(chatId); diff --git a/packages/trigger-sdk/test/chat-turn-correlation.test.ts b/packages/trigger-sdk/test/chat-turn-correlation.test.ts index 7295ecba79..fc827a9b0b 100644 --- a/packages/trigger-sdk/test/chat-turn-correlation.test.ts +++ b/packages/trigger-sdk/test/chat-turn-correlation.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import type { UIMessage } from "ai"; import { TriggerChatTransport, type TriggerChatTransportOptions } from "../src/v3/chat.js"; @@ -17,13 +17,18 @@ type BatchRecord = { headers?: Array<[string, string]>; }; -function batchResponse(records: BatchRecord[]): Response { +function batchResponse(records: BatchRecord[], settled = false): Response { const frames = records .map((r) => `event: batch\ndata: ${JSON.stringify({ records: [r] })}\n\n`) .join(""); + const headers: Record = { + "Content-Type": "text/event-stream", + "X-Stream-Version": "v2", + }; + if (settled) headers["X-Session-Settled"] = "true"; return new Response(frames, { status: 200, - headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v2" }, + headers, }); } @@ -42,7 +47,10 @@ function turnComplete(seqNum: number, inCursor: number): BatchRecord { function textDelta(seqNum: number, text: string): BatchRecord { return { - body: JSON.stringify({ data: { type: "text-delta", id: "t1", delta: text }, id: "m1" }), + body: JSON.stringify({ + data: { type: "text-delta", id: "t1", delta: text }, + id: `m${seqNum}`, + }), seq_num: seqNum, timestamp: seqNum, headers: [], @@ -88,6 +96,36 @@ async function submit(transport: TriggerChatTransport): Promise { } describe("transport turn correlation", () => { + it("persists the owned send's input sequence before subscribing", async () => { + const onSessionChange = vi.fn(); + const transport = new TriggerChatTransport({ + task: "test-task", + accessToken: async () => "tok_test", + sessions: { c1: { publicAccessToken: "tok_test", isStreaming: false } }, + onSessionChange, + fetch: async (_url, _init, ctx) => + ctx.endpoint === "in" ? inResponse(5) : batchResponse([turnComplete(10, 5)]), + }); + + const stream = await transport.sendMessages({ + trigger: "submit-message", + chatId: "c1", + messageId: undefined, + messages: [user("hi", "u-1")], + abortSignal: undefined, + }); + + expect(onSessionChange).toHaveBeenCalledWith("c1", { + publicAccessToken: "tok_test", + lastEventId: undefined, + activeInputSeq: 5, + isStreaming: true, + }); + expect(transport.getSession("c1")?.activeInputSeq).toBe(5); + await readDeltas(stream); + expect(transport.getSession("c1")?.activeInputSeq).toBeUndefined(); + }); + it("skips an earlier turn's turn-complete and closes on its own", async () => { // Append seq 5; the undo turn's complete (cursor 4) must be skipped. const out = batchResponse([turnComplete(10, 4), textDelta(11, "56"), turnComplete(12, 5)]); @@ -107,4 +145,191 @@ describe("transport turn correlation", () => { const deltas = await submit(makeTransport(out, undefined)); expect(deltas).toEqual([]); }); + + it("reuses a hydrated input sequence to skip stale turn-completes after reconnecting", async () => { + const transport = new TriggerChatTransport({ + task: "test-task", + accessToken: async () => "tok_test", + sessions: { + c1: { publicAccessToken: "tok_test", isStreaming: true, activeInputSeq: 5 }, + }, + fetch: async () => + batchResponse([turnComplete(10, 4), textDelta(11, "current"), turnComplete(12, 5)]), + }); + + const stream = await transport.reconnectToStream({ chatId: "c1" }); + + expect(stream).not.toBeNull(); + await expect(readDeltas(stream!)).resolves.toEqual(["current"]); + expect(transport.getSession("c1")?.isStreaming).toBe(false); + expect(transport.getSession("c1")?.activeInputSeq).toBeUndefined(); + }); + + it("does not request a settled peek while reconnecting a known active input", async () => { + vi.useFakeTimers(); + try { + const subscribeHeaders: Headers[] = []; + const transport = new TriggerChatTransport({ + task: "test-task", + accessToken: async () => "tok_test", + sessions: { + c1: { publicAccessToken: "tok_test", isStreaming: true, activeInputSeq: 5 }, + }, + fetch: async (_url, init) => { + const headers = new Headers(init?.headers); + subscribeHeaders.push(headers); + + if (subscribeHeaders.length === 1) { + // Match the server shortcut: a peek sees the previous turn's + // boundary at the tail and marks this otherwise-normal EOF settled. + return batchResponse([turnComplete(10, 4)], headers.has("X-Peek-Settled")); + } + + return batchResponse([textDelta(11, "current"), turnComplete(12, 5)]); + }, + }); + + const stream = await transport.reconnectToStream({ chatId: "c1" }); + + expect(stream).not.toBeNull(); + const deltas = readDeltas(stream!); + await vi.advanceTimersByTimeAsync(1_000); + await expect(deltas).resolves.toEqual(["current"]); + expect(subscribeHeaders).toHaveLength(2); + expect(subscribeHeaders[0]?.get("X-Peek-Settled")).toBeNull(); + expect(transport.getSession("c1")?.isStreaming).toBe(false); + expect(transport.getSession("c1")?.activeInputSeq).toBeUndefined(); + } finally { + vi.useRealTimers(); + } + }); + + it("keeps turn correlation through a lost append response, retry, and reload", async () => { + const committed = new Map(); + const appendPartIds: string[] = []; + const appendAttempt = (init?: RequestInit) => { + const partId = new Headers(init?.headers).get("X-Part-Id"); + expect(partId).not.toBeNull(); + appendPartIds.push(partId!); + const existing = committed.get(partId!); + if (existing !== undefined) return existing; + committed.set(partId!, 5); + return 5; + }; + + const transport = new TriggerChatTransport({ + task: "test-task", + accessToken: async () => "tok_test", + sessions: { c1: { publicAccessToken: "tok_test", isStreaming: false } }, + fetch: async (_url, init, ctx) => { + if (ctx.endpoint === "in") { + appendAttempt(init); // committed response is lost + return inResponse(appendAttempt(init)); // transparent retry + } + + return new Response( + new ReadableStream({ + start(controller) { + const signal = init?.signal; + if (signal?.aborted) { + controller.close(); + return; + } + signal?.addEventListener("abort", () => controller.close(), { once: true }); + }, + }), + { status: 200, headers: { "X-Stream-Version": "v2" } } + ); + }, + }); + + const liveStream = await transport.sendMessages({ + trigger: "submit-message", + chatId: "c1", + messageId: undefined, + messages: [user("hi", "u-1")], + abortSignal: undefined, + }); + const persisted = transport.getSession("c1"); + + expect(appendPartIds).toHaveLength(2); + expect(appendPartIds[0]).toBe(appendPartIds[1]); + expect(persisted).toMatchObject({ activeInputSeq: 5, isStreaming: true }); + await liveStream.cancel(); + + const completed: number[] = []; + const rehydrated = new TriggerChatTransport({ + task: "test-task", + accessToken: async () => "tok_test", + sessions: { c1: persisted! }, + onEvent: (event) => { + if (event.type === "turn-completed") completed.push(Number(event.sessionInEventId)); + }, + fetch: async () => + batchResponse([turnComplete(10, 4), textDelta(11, "current"), turnComplete(12, 5)]), + }); + + const resumed = await rehydrated.reconnectToStream({ chatId: "c1" }); + + expect(resumed).not.toBeNull(); + await expect(readDeltas(resumed!)).resolves.toEqual(["current"]); + expect(completed).toEqual([5]); + expect(rehydrated.getSession("c1")?.activeInputSeq).toBeUndefined(); + }); + + it.each([5, 6])( + "accepts a reconnected turn-complete at or after the active input sequence (%i)", + async (inCursor) => { + const transport = new TriggerChatTransport({ + task: "test-task", + accessToken: async () => "tok_test", + sessions: { + c1: { publicAccessToken: "tok_test", isStreaming: true, activeInputSeq: 5 }, + }, + fetch: async () => batchResponse([turnComplete(10, inCursor), textDelta(11, "late")]), + }); + + const stream = await transport.reconnectToStream({ chatId: "c1" }); + + expect(stream).not.toBeNull(); + await expect(readDeltas(stream!)).resolves.toEqual([]); + expect(transport.getSession("c1")?.isStreaming).toBe(false); + expect(transport.getSession("c1")?.activeInputSeq).toBeUndefined(); + } + ); + + it("uses the input sequence for one accepted watch turn only", async () => { + let outCalls = 0; + const turnCompleted: number[] = []; + const transport = new TriggerChatTransport({ + task: "test-task", + accessToken: async () => "tok_test", + watch: true, + sessions: { + c1: { publicAccessToken: "tok_test", isStreaming: true, activeInputSeq: 5 }, + }, + onEvent: (event) => { + if (event.type === "turn-completed") turnCompleted.push(Number(event.sessionInEventId)); + }, + fetch: async () => { + outCalls++; + return outCalls === 1 + ? batchResponse([ + turnComplete(10, 4), + textDelta(11, "first"), + turnComplete(12, 5), + textDelta(13, "second"), + turnComplete(14, 4), + ]) + : batchResponse([], true); + }, + }); + + const stream = await transport.reconnectToStream({ chatId: "c1" }); + + expect(stream).not.toBeNull(); + await expect(readDeltas(stream!)).resolves.toEqual(["first", "second"]); + expect(turnCompleted).toEqual([5, 4]); + expect(transport.getSession("c1")?.activeInputSeq).toBeUndefined(); + }); });