diff --git a/.changeset/chat-agent-preload-oom-recovery.md b/.changeset/chat-agent-preload-oom-recovery.md new file mode 100644 index 00000000000..7885a3530ab --- /dev/null +++ b/.changeset/chat-agent-preload-oom-recovery.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Fix a preloaded `chat.agent` run dropping an in-flight message when it retries after an out-of-memory error. The message being processed when the run hit the OOM is now recovered and re-run on the retry, instead of being skipped while the run waited for a new message. diff --git a/.changeset/chat-session-caught-up-resume.md b/.changeset/chat-session-caught-up-resume.md new file mode 100644 index 00000000000..600ae7c9aed --- /dev/null +++ b/.changeset/chat-session-caught-up-resume.md @@ -0,0 +1,6 @@ +--- +"@trigger.dev/core": patch +"@trigger.dev/sdk": patch +--- + +Chat sessions now close a resumed stream as soon as it has caught up to the latest output, instead of holding the connection open for the full long-poll window. Reloading or reconnecting to an idle chat settles faster. diff --git a/.github/workflows/e2e-webapp.yml b/.github/workflows/e2e-webapp.yml index 1931cfdd336..4b9b984edc0 100644 --- a/.github/workflows/e2e-webapp.yml +++ b/.github/workflows/e2e-webapp.yml @@ -80,6 +80,8 @@ jobs: docker pull postgres:14 docker pull redis:7.2 docker pull testcontainers/ryuk:0.11.0 + docker pull ghcr.io/s2-streamstore/s2:0.40.0@sha256:b26249e2ede0949755f5af8028185dc2bcfc3aa2db21eb9610543d144eb6ee9d + docker pull minio/minio:latest echo "Image pre-pull complete" - name: ๐Ÿ“ฅ Download deps @@ -91,6 +93,9 @@ jobs: - name: ๐Ÿ—๏ธ Build Webapp run: pnpm run build --filter webapp + - name: ๐ŸŽญ Install Playwright Chromium + run: cd apps/webapp && pnpm exec playwright install chromium + - name: ๐Ÿงช Run Webapp E2E Tests run: cd apps/webapp && pnpm exec vitest run --config vitest.e2e.config.ts --reporter=default env: diff --git a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts index 787b3f1ab44..428163d3f8d 100644 --- a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts +++ b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts @@ -107,8 +107,8 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { constructor(opts: S2RealtimeStreamsOptions) { this.basin = opts.basin; - this.baseUrl = opts.endpoint ?? `https://${this.basin}.b.aws.s2.dev/v1`; - this.accountUrl = opts.endpoint ?? `https://aws.s2.dev/v1`; + this.baseUrl = opts.endpoint ?? `https://${this.basin}.b.s2.dev/v1`; + this.accountUrl = opts.endpoint ?? `https://a.s2.dev/v1`; this.endpoint = opts.endpoint; this.token = opts.accessToken; this.streamPrefix = opts.streamPrefix ?? ""; diff --git a/apps/webapp/app/services/realtime/streamBasinProvisioner.server.ts b/apps/webapp/app/services/realtime/streamBasinProvisioner.server.ts index 98316b84fe8..d4a6d1738ee 100644 --- a/apps/webapp/app/services/realtime/streamBasinProvisioner.server.ts +++ b/apps/webapp/app/services/realtime/streamBasinProvisioner.server.ts @@ -185,7 +185,7 @@ type CreateBasinOptions = { }; async function s2CreateBasin(name: string, opts: CreateBasinOptions): Promise { - const url = `https://aws.s2.dev/v1/basins`; + const url = `https://a.s2.dev/v1/basins`; const body = { basin: name, config: { @@ -222,7 +222,7 @@ type ReconfigureBasinOptions = { }; async function s2ReconfigureBasin(name: string, opts: ReconfigureBasinOptions): Promise { - const url = `https://aws.s2.dev/v1/basins/${encodeURIComponent(name)}`; + const url = `https://a.s2.dev/v1/basins/${encodeURIComponent(name)}`; const body = { default_stream_config: { retention_policy: { age: parseDuration(opts.retentionPolicy) }, diff --git a/apps/webapp/package.json b/apps/webapp/package.json index 678eeb76033..5e2003d955d 100644 --- a/apps/webapp/package.json +++ b/apps/webapp/package.json @@ -103,7 +103,7 @@ "@remix-run/react": "2.17.5", "@remix-run/router": "^1.23.3", "@remix-run/server-runtime": "2.17.5", - "@s2-dev/streamstore": "^0.22.10", + "@s2-dev/streamstore": "^0.25.0", "@sentry/remix": "9.46.0", "@slack/web-api": "7.16.0", "@socket.io/redis-adapter": "^8.3.0", @@ -219,6 +219,7 @@ "@internal/clickhouse": "workspace:*", "@internal/replication": "workspace:*", "@internal/testcontainers": "workspace:*", + "@playwright/test": "^1.36.2", "@remix-run/dev": "2.17.5", "@remix-run/testing": "^2.17.5", "@sentry/cli": "2.50.2", diff --git a/apps/webapp/test/helpers/agentHarness.ts b/apps/webapp/test/helpers/agentHarness.ts new file mode 100644 index 00000000000..1d616958231 --- /dev/null +++ b/apps/webapp/test/helpers/agentHarness.ts @@ -0,0 +1,195 @@ +import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3"; +import type { LocalsKey } from "@trigger.dev/core/v3"; +import type { LanguageModel } from "ai"; +import { + installSessionWaitpointBackend, + runInMockTaskContext, + StandardSessionStreamManager, +} from "@trigger.dev/core/v3/test"; + +export type RunRealChatAgentOptions = { + agentId: string; + baseUrl: string; + addressingKey: string; + /** + * The environment secret key. The agent writes `.out` and reads `.in` as the + * backend (PRIVATE auth) โ€” the `.out` channel rejects client session tokens. + */ + secretKey: string; + model: LanguageModel; + modelLocal: LocalsKey; + runId?: string; + /** + * Boot as a continuation of a previous run for the same session. Gates the + * snapshot + `.out`/`.in` replay boot path, so the agent restores prior + * history instead of treating the chat as brand new. + */ + continuation?: boolean; + previousRunId?: string; + /** + * Idle window (seconds) before the turn loop falls through from the SSE + * once() to the suspending `session.in.wait()`. Set this low to force the + * suspend/resume path in a test. + */ + idleTimeoutInSeconds?: number; + /** + * `ctx.attempt.number`. A value greater than 1 makes the boot treat the run + * as a retry (`couldHavePriorState`), restoring from the snapshot + `.in` + * replay. Used to model an OOM retry re-dispatch. + */ + attemptNumber?: number; +}; + +export type RunningAgent = { + done: Promise; + close: () => Promise; +}; + +/** + * Run the real `chat.agent` turn loop in-process, wired to a running webapp: + * `apiClientManager` + a real `StandardSessionStreamManager` point the agent's + * `.in`/`.out` at the webapp's Session streams (real S2 + SSE), the model is + * injected via locals (so it survives without serialization), and turns are + * driven by appending to `.in` over HTTP. Callers keep each message inside the + * idle window and `close()` promptly so the run-engine suspend path is never + * reached. + */ +export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent { + apiClientManager.setGlobalAPIClientConfiguration({ + baseURL: opts.baseUrl, + accessToken: opts.secretKey, + }); + const apiClient = apiClientManager.clientOrThrow(); + const manager = new StandardSessionStreamManager(apiClient, opts.baseUrl); + const { runtimeManager, restore } = installSessionWaitpointBackend(apiClient); + + const taskEntry = resourceCatalog.getTask(opts.agentId); + if (!taskEntry) { + restore(); + throw new Error(`runRealChatAgent: agent "${opts.agentId}" is not registered`); + } + const runFn = taskEntry.fns.run as ( + payload: unknown, + params: { ctx: unknown; signal: AbortSignal } + ) => Promise; + + const runSignal = new AbortController(); + const runId = opts.runId ?? `run_${opts.addressingKey}`; + + const idle = + opts.idleTimeoutInSeconds !== undefined + ? { idleTimeoutInSeconds: opts.idleTimeoutInSeconds } + : {}; + + const done = ( + runInMockTaskContext( + async (drivers) => { + drivers.locals.set(opts.modelLocal, opts.model); + const payload = opts.continuation + ? { + chatId: opts.addressingKey, + continuation: true, + metadata: {}, + ...idle, + ...(opts.previousRunId ? { previousRunId: opts.previousRunId } : {}), + } + : { chatId: opts.addressingKey, trigger: "preload", metadata: {}, ...idle }; + await runFn(payload, { ctx: drivers.ctx, signal: runSignal.signal }); + }, + { + ctx: { + run: { id: runId }, + ...(opts.attemptNumber !== undefined ? { attempt: { number: opts.attemptNumber } } : {}), + }, + sessionStreamManager: manager, + runtimeManager, + } + ) as Promise + ).finally(restore); + + return { + done, + close: async () => { + try { + await fetch( + `${opts.baseUrl}/realtime/v1/sessions/${encodeURIComponent(opts.addressingKey)}/in/append`, + { + method: "POST", + headers: { + Authorization: `Bearer ${opts.secretKey}`, + "Content-Type": "application/json", + "X-Part-Id": "close", + }, + body: JSON.stringify({ + kind: "message", + payload: { chatId: opts.addressingKey, trigger: "close" }, + }), + } + ); + } catch {} + runSignal.abort(); + await Promise.race([ + done.catch(() => {}), + new Promise((resolve) => setTimeout(resolve, 10_000)), + ]); + }, + }; +} + +export type ChatAgentSessionOptions = Omit< + RunRealChatAgentOptions, + "runId" | "continuation" | "previousRunId" +>; + +export type ChatAgentSession = { + /** How many runs the session has spawned so far (1 fresh + N continuations). */ + runCount: () => number; + /** Close the currently-active run and stop spawning continuations. */ + close: () => Promise; +}; + +/** + * A session-scoped orchestrator that stands in for the run-engine's run + * lifecycle: it starts a run, and whenever that run exits on its own + * (`chat.endRun()` / `chat.requestUpgrade()`), spawns the next run as a + * continuation (new run id, `continuation: true`, `previousRunId` threaded) + * for the same session. That mirrors the server triggering a fresh run on the + * next append after the previous run went terminal, and lets each continuation + * restore prior history from the persisted snapshot. Runs never overlap: the + * next spawn is chained on the previous run's `done` (after its manager + * teardown), so the process-global managers are never installed twice at once. + */ +export function runChatAgentSession(opts: ChatAgentSessionOptions): ChatAgentSession { + let closed = false; + let index = 0; + let current: RunningAgent | undefined; + let previousRunId: string | undefined; + + const spawn = () => { + index += 1; + const runId = `run_${opts.addressingKey}_${index}`; + current = runRealChatAgent({ + ...opts, + runId, + continuation: index > 1, + previousRunId, + }); + previousRunId = runId; + const settle = () => { + if (!closed) { + spawn(); + } + }; + current.done.then(settle, settle); + }; + + spawn(); + + return { + runCount: () => index, + close: async () => { + closed = true; + await current?.close(); + }, + }; +} diff --git a/apps/webapp/test/helpers/sessionStream.ts b/apps/webapp/test/helpers/sessionStream.ts new file mode 100644 index 00000000000..beac405c750 --- /dev/null +++ b/apps/webapp/test/helpers/sessionStream.ts @@ -0,0 +1,280 @@ +import { AppendInput, AppendRecord, S2 } from "@s2-dev/streamstore"; +import { generateJWT } from "@trigger.dev/core/v3/jwt"; +import { SSEStreamSubscription } from "@trigger.dev/core/v3"; + +export type SessionAddressing = { + orgId: string; + envSlug: string; + envId: string; + addressingKey: string; + io?: "out" | "in"; +}; + +/** + * The full, prefixed S2 stream name for a session channel on the shared basin + * (per-org basins disabled), matching `toSessionStreamName` + + * `streamPrefixFor` in the webapp. A test appending with the root S2 token uses + * this literal name. + */ +export function sessionStreamName(p: SessionAddressing): string { + return `org/${p.orgId}/env/${p.envSlug}/${p.envId}/sessions/${p.addressingKey}/${p.io ?? "out"}`; +} + +/** + * Mint a session-scoped public access token the way `mintSessionToken.server.ts` + * does: a JWT signed with the environment secret, `sub` = env id, `pub` true, + * scoped to read/write the given addressing key. + */ +export function mintSessionToken(p: { + apiKey: string; + envId: string; + addressingKey: string; +}): Promise { + return generateJWT({ + secretKey: p.apiKey, + payload: { + pub: true, + sub: p.envId, + scopes: [`read:sessions:${p.addressingKey}`, `write:sessions:${p.addressingKey}`], + }, + expirationTime: "1h", + }); +} + +/** + * Writes `.out` records straight to S2 (the "agent simulator"). Uses the same + * `@s2-dev/streamstore` primitives the real agent runtime uses, so data, + * `trigger-control` and `trim` command records land in exactly the shapes the + * client + proxy expect. + */ +export class SessionStreamProducer { + private stream; + + constructor(p: { endpoint: string; basin: string; streamName: string; accessToken?: string }) { + const s2 = new S2({ + accessToken: p.accessToken ?? "ignored", + endpoints: { account: p.endpoint, basin: p.endpoint }, + }); + this.stream = s2.basin(p.basin).stream(p.streamName); + } + + /** Append one data record (`{data, id}` envelope). Returns its seq_num. */ + async appendData(data: unknown, id: string): Promise { + const ack = await this.stream.append( + AppendInput.create([AppendRecord.string({ body: JSON.stringify({ data, id }) })]) + ); + return Number(ack.start.seqNum); + } + + /** Append a `trigger-control: turn-complete` record (empty body). */ + async appendTurnComplete(publicAccessToken?: string): Promise { + const headers: Array<[string, string]> = [["trigger-control", "turn-complete"]]; + if (publicAccessToken) headers.push(["public-access-token", publicAccessToken]); + const ack = await this.stream.append( + AppendInput.create([AppendRecord.string({ body: "", headers })]) + ); + return Number(ack.start.seqNum); + } + + /** Append an S2 `trim` command record, trimming below `earliestSeqNum`. */ + async trim(earliestSeqNum: number): Promise { + await this.stream.append(AppendInput.create([AppendRecord.trim(earliestSeqNum)])); + } +} + +export type CollectedPart = { + id: string; + chunk: unknown; + headers?: ReadonlyArray; +}; + +export function isTurnComplete(part: CollectedPart): boolean { + return (part.headers ?? []).some(([k, v]) => k === "trigger-control" && v === "turn-complete"); +} + +export function isUpgradeRequired(part: CollectedPart): boolean { + return (part.headers ?? []).some(([k, v]) => k === "trigger-control" && v === "upgrade-required"); +} + +export type SubscribeOptions = { + baseUrl: string; + addressingKey: string; + token: string; + lastEventId?: string; + timeoutInSeconds?: number; + peekSettled?: boolean; + io?: "out" | "in"; +}; + +function sessionChannelUrl(baseUrl: string, addressingKey: string, io: "out" | "in"): string { + return `${baseUrl}/realtime/v1/sessions/${encodeURIComponent(addressingKey)}/${io}`; +} + +/** + * Append a record to a session's `.in` channel through the webapp route the + * browser uses (client -> agent). Returns the status, the reflected + * Access-Control-Allow-Origin (when an Origin was sent), and the parsed body. + */ +export async function appendInput(opts: { + baseUrl: string; + addressingKey: string; + token: string; + body: string; + partId?: string; + origin?: string; +}): Promise<{ status: number; acao: string | null; json: unknown }> { + const url = `${sessionChannelUrl(opts.baseUrl, opts.addressingKey, "in")}/append`; + const res = await fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${opts.token}`, + "Content-Type": "application/json", + ...(opts.partId ? { "X-Part-Id": opts.partId } : {}), + ...(opts.origin ? { Origin: opts.origin } : {}), + }, + body: opts.body, + }); + let json: unknown; + try { + json = await res.json(); + } catch {} + return { status: res.status, acao: res.headers.get("access-control-allow-origin"), json }; +} + +/** + * Raw SSE GET against a session channel, exposing the response status + + * headers (which `SSEStreamSubscription` hides). Reads the body to close and + * reports how long that took, for asserting the server-side peek fast-close. + */ +export async function openChannelRaw( + opts: SubscribeOptions & { maxMs?: number } +): Promise<{ status: number; sessionSettled: string | null; closedMs: number; body: string }> { + const url = sessionChannelUrl(opts.baseUrl, opts.addressingKey, opts.io ?? "out"); + const started = performance.now(); + const res = await fetch(url, { + headers: { + Authorization: `Bearer ${opts.token}`, + Accept: "text/event-stream", + ...(opts.lastEventId ? { "Last-Event-ID": opts.lastEventId } : {}), + ...(opts.timeoutInSeconds ? { "Timeout-Seconds": String(opts.timeoutInSeconds) } : {}), + ...(opts.peekSettled ? { "X-Peek-Settled": "1" } : {}), + }, + }); + const sessionSettled = res.headers.get("x-session-settled"); + let body = ""; + if (res.body) { + const reader = res.body.getReader(); + const decoder = new TextDecoder(); + const deadline = started + (opts.maxMs ?? 15_000); + try { + while (performance.now() < deadline) { + const { done, value } = await reader.read(); + if (done) break; + body += decoder.decode(value, { stream: true }); + } + } finally { + await reader.cancel().catch(() => {}); + } + } + return { status: res.status, sessionSettled, closedMs: performance.now() - started, body }; +} + +export function subscribeSessionOut(opts: SubscribeOptions): SSEStreamSubscription { + const url = sessionChannelUrl(opts.baseUrl, opts.addressingKey, opts.io ?? "out"); + return new SSEStreamSubscription(url, { + headers: { + Authorization: `Bearer ${opts.token}`, + ...(opts.peekSettled ? { "X-Peek-Settled": "1" } : {}), + }, + timeoutInSeconds: opts.timeoutInSeconds ?? 30, + lastEventId: opts.lastEventId, + maxRetries: 0, + }); +} + +/** + * Subscribe + drain parts into an array, stopping when `until(parts)` is true, + * the stream closes, or `maxMs` elapses. Cancels the reader on exit. Returns + * the parts plus how many distinct SSE connections/opens the subscription made + * (for asserting the round-trip count on a reconnect). + */ +export async function collectSessionOut( + opts: SubscribeOptions & { until?: (parts: CollectedPart[]) => boolean; maxMs?: number } +): Promise<{ parts: CollectedPart[]; durationMs: number; subscription: SSEStreamSubscription }> { + const subscription = subscribeSessionOut(opts); + const stream = await subscription.subscribe(); + const reader = stream.getReader(); + const parts: CollectedPart[] = []; + const started = performance.now(); + const deadline = started + (opts.maxMs ?? 30_000); + + try { + while (true) { + if (opts.until && opts.until(parts)) break; + const remaining = deadline - performance.now(); + if (remaining <= 0) break; + const next = await Promise.race([ + reader.read(), + new Promise<"timeout">((r) => setTimeout(() => r("timeout"), remaining)), + ]); + if (next === "timeout") break; + if (next.done) break; + parts.push(next.value as CollectedPart); + } + } finally { + await reader.cancel().catch(() => {}); + } + + return { parts, durationMs: performance.now() - started, subscription }; +} + +/** + * Subscribe + drain while awaiting the client's `caughtUp()`. Resolves once the + * client reports it has drained to the live tail (or `maxMs` elapses). Used by + * the GREEN legs: the drain is what lets the wrapper mark the tail boundary, so + * `parts` holds everything delivered up to the moment caught-up fires โ€” the + * data-loss guard. + */ +export async function collectUntilCaughtUp(opts: SubscribeOptions & { maxMs?: number }): Promise<{ + parts: CollectedPart[]; + caughtUp: boolean; + tailSeqNum?: number; + settleMs: number; +}> { + const subscription = subscribeSessionOut(opts); + const stream = await subscription.subscribe(); + const reader = stream.getReader(); + const parts: CollectedPart[] = []; + const started = performance.now(); + + let caughtUp = false; + let tailSeqNum: number | undefined; + let settleMs = 0; + subscription + .caughtUp() + .then((tail) => { + caughtUp = true; + tailSeqNum = tail.seqNum; + settleMs = performance.now() - started; + }) + .catch(() => {}); + + const deadline = started + (opts.maxMs ?? 30_000); + try { + while (!caughtUp) { + const remaining = deadline - performance.now(); + if (remaining <= 0) break; + const next = await Promise.race([ + reader.read(), + new Promise<"tick">((r) => setTimeout(() => r("tick"), Math.min(remaining, 250))), + ]); + if (next === "tick") continue; + if (next.done) break; + parts.push(next.value as CollectedPart); + } + } finally { + await reader.cancel().catch(() => {}); + } + + return { parts, caughtUp, tailSeqNum, settleMs: settleMs || performance.now() - started }; +} diff --git a/apps/webapp/test/helpers/testChatAgent.ts b/apps/webapp/test/helpers/testChatAgent.ts new file mode 100644 index 00000000000..8aebb713a96 --- /dev/null +++ b/apps/webapp/test/helpers/testChatAgent.ts @@ -0,0 +1,341 @@ +import { chat } from "@trigger.dev/sdk/ai"; +import { locals, OutOfMemoryError } from "@trigger.dev/core/v3"; +import { stepCountIs, streamText, tool, type LanguageModel, type UIMessage } from "ai"; +import { z } from "zod"; + +/** + * The model is injected through locals (an in-process, by-reference value) + * rather than clientData, because a `MockLanguageModelV3` can't survive the + * JSON round-trip through the real `.in/append` route. The harness sets this + * before the run starts. + */ +export const testChatModelLocal = locals.create("e2e-test-chat.model"); + +export type TestChatClientData = { hydrated?: UIMessage[] }; + +function firstText(m: UIMessage): string { + const p = m.parts?.[0]; + return p?.type === "text" ? p.text : ""; +} + +export const testChatAgent = chat + .withClientData({ + schema: z.custom((v) => v == null || typeof v === "object"), + }) + .agent({ + id: "e2e-test-chat", + + onValidateMessages: async ({ messages }) => { + for (const m of messages) { + if (m.role === "user" && firstText(m).toLowerCase().includes("blocked-word")) { + throw new Error("Message blocked by content filter"); + } + } + return messages; + }, + + hydrateMessages: async ({ clientData, incomingMessages }) => { + if (!clientData?.hydrated) return incomingMessages; + const merged = [...clientData.hydrated]; + for (const m of incomingMessages) { + const idx = merged.findIndex((x) => x.id === m.id); + if (idx === -1) merged.push(m); + else merged[idx] = m; + } + return merged; + }, + + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("undo") })]), + + onAction: async ({ action }) => { + if (action.type === "undo") { + chat.history.slice(0, -2); + } + }, + + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, + }); + +/** + * Records suspend/resume lifecycle-hook fires (keyed by chatId so tests can + * filter to their own session). Appended to by {@link testSuspendHooksChatAgent}. + */ +export const suspendResumeEvents: Array<{ + chatId: string; + kind: "suspend" | "resume"; + phase: string; +}> = []; + +/** + * Short-idle agent that records `onChatSuspend` / `onChatResume` fires, so a + * test can assert the lifecycle hooks run around a real suspend/resume. + */ +export const testSuspendHooksChatAgent = chat.agent({ + id: "e2e-test-chat-suspend-hooks", + idleTimeoutInSeconds: 1, + preloadIdleTimeoutInSeconds: 1, + onChatSuspend: async ({ chatId, phase }) => { + suspendResumeEvents.push({ chatId, kind: "suspend", phase }); + }, + onChatResume: async ({ chatId, phase }) => { + suspendResumeEvents.push({ chatId, kind: "resume", phase }); + }, + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, +}); + +/** + * Throws `OutOfMemoryError` on the first attempt so the run fails the way the + * runtime detects for a machine swap, then succeeds on attempt 2 (the retry + * boots through the restore path because `ctx.attempt.number > 1`). + */ +export const testOomChatAgent = chat.agent({ + id: "e2e-test-chat-oom", + idleTimeoutInSeconds: 1, + preloadIdleTimeoutInSeconds: 1, + run: async ({ messages, signal, ctx }) => { + if (ctx.attempt.number === 1) { + throw new OutOfMemoryError(); + } + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, +}); + +/** + * Short idle window and a short `turnTimeout`, so a run with no incoming + * message suspends and then times out on the waitpoint, ending the run. + */ +export const testTimeoutChatAgent = chat.agent({ + id: "e2e-test-chat-timeout", + idleTimeoutInSeconds: 1, + preloadIdleTimeoutInSeconds: 1, + turnTimeout: "3s", + preloadTimeout: "3s", + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, +}); + +/** + * A minimal agent with no lifecycle hooks. With `hydrateMessages` absent, the + * default snapshot + `.out`/`.in` replay boot path is what restores prior + * history on a continuation run. + */ +export const testPlainChatAgent = chat.agent({ + id: "e2e-test-chat-plain", + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, +}); + +/** + * A plain agent with a 1-second idle window, so the turn loop falls through to + * the suspending `session.in.wait()` almost immediately instead of catching the + * next message in the warm once() window. Used to exercise the suspend/resume + * waitpoint path. + */ +export const testIdleChatAgent = chat.agent({ + id: "e2e-test-chat-idle", + idleTimeoutInSeconds: 1, + preloadIdleTimeoutInSeconds: 1, + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, +}); + +/** + * Ends the run after every turn via `chat.endRun()`. The next message on the + * same chat starts a fresh continuation run (the orchestrator drives that). + */ +export const testEndRunChatAgent = chat.agent({ + id: "e2e-test-chat-endrun", + idleTimeoutInSeconds: 2, + preloadIdleTimeoutInSeconds: 2, + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + const result = streamText({ model, messages, abortSignal: signal }); + chat.endRun(); + return result; + }, +}); + +/** + * Requests an upgrade from `onTurnStart` (the pre-turn path): `run()` is + * skipped, an `upgrade-required` control record lands on `.out`, and the run + * exits so a fresh run on the new version handles the message. + */ +export const testUpgradeChatAgent = chat.agent({ + id: "e2e-test-chat-upgrade", + idleTimeoutInSeconds: 2, + preloadIdleTimeoutInSeconds: 2, + onTurnStart: async () => { + chat.requestUpgrade(); + }, + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, +}); + +/** + * Requests an upgrade only on the fresh (non-continuation) run. The first run + * defers the message via `upgrade-required`; the continuation run treats it as + * the "new version" and processes the deferred message instead of upgrading + * again (which would loop). + */ +export const testUpgradeOnceChatAgent = chat.agent({ + id: "e2e-test-chat-upgrade-once", + idleTimeoutInSeconds: 2, + preloadIdleTimeoutInSeconds: 2, + onTurnStart: async ({ continuation }) => { + if (!continuation) { + chat.requestUpgrade(); + } + }, + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ model, messages, abortSignal: signal }); + }, +}); + +/** + * A tool with a server-side `execute`: the agent runs it automatically and + * feeds the result back to the model, so a single turn covers the whole + * tool loop. + */ +const weatherTool = tool({ + description: "Get the current weather for a city.", + inputSchema: z.object({ city: z.string() }), + execute: async ({ city }) => ({ city, tempC: 21, summary: "clear" }), +}); + +export const testToolChatAgent = chat.agent({ + id: "e2e-test-chat-tool", + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ + model, + messages, + tools: { getWeather: weatherTool }, + stopWhen: stepCountIs(5), + abortSignal: signal, + }); + }, +}); + +/** + * A tool with no `execute`: the model's call parks the turn on a + * human-in-the-loop round-trip. The client supplies the tool output on the + * next message and the agent continues from there. + */ +const askUserTool = tool({ + description: "Ask the user a question and wait for their answer.", + inputSchema: z.object({ question: z.string() }), +}); + +export const testHitlChatAgent = chat.agent({ + id: "e2e-test-chat-hitl", + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ + model, + messages, + tools: { askUser: askUserTool }, + abortSignal: signal, + }); + }, +}); + +/** + * Same HITL tool but with a 1-second idle window, so the run suspends on the + * waitpoint while waiting for the human's tool answer. Exercises a HITL + * round-trip that crosses a suspend/resume boundary. + */ +export const testHitlIdleChatAgent = chat.agent({ + id: "e2e-test-chat-hitl-idle", + idleTimeoutInSeconds: 1, + preloadIdleTimeoutInSeconds: 1, + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ + model, + messages, + tools: { askUser: askUserTool }, + abortSignal: signal, + }); + }, +}); + +/** + * A tool that both executes and requires approval. The model's call parks on + * an approval request; the client approves (or denies) before the `execute` + * runs. + */ +const deleteResourceTool = tool({ + description: "Delete a resource. Requires human approval before running.", + inputSchema: z.object({ resource: z.string() }), + needsApproval: true, + execute: async ({ resource }) => ({ deleted: resource }), +}); + +export const testApprovalChatAgent = chat.agent({ + id: "e2e-test-chat-approval", + run: async ({ messages, signal }) => { + const model = locals.get(testChatModelLocal); + if (!model) { + throw new Error("test model not injected via locals"); + } + return streamText({ + model, + messages, + tools: { deleteResource: deleteResourceTool }, + stopWhen: stepCountIs(5), + abortSignal: signal, + }); + }, +}); diff --git a/apps/webapp/test/session-agent.e2e.test.ts b/apps/webapp/test/session-agent.e2e.test.ts new file mode 100644 index 00000000000..1ed5c701afe --- /dev/null +++ b/apps/webapp/test/session-agent.e2e.test.ts @@ -0,0 +1,1518 @@ +/** + * Tier-1 agent e2e: the real chat.agent turn loop against the full stack. + * + * Boots the testcontainer stack (webapp + Postgres + Redis + s2-lite + MinIO), + * runs the genuine `chat.agent` run loop in-process wired to the real webapp + * Session streams (real S2 `.out`/`.in` + object-store snapshots), with the + * language model injected as a deterministic MockLanguageModelV3. Turns are + * driven by appending to `.in` over HTTP; output is read back through the real + * SSE proxy. No real LLM, no `trigger dev`. + * + * Requires a pre-built webapp: pnpm run build --filter webapp + */ +import { randomBytes } from "crypto"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import "@trigger.dev/sdk/ai/test"; +import { MockLanguageModelV3 } from "ai/test"; +import { simulateReadableStream } from "ai"; +import type { SessionStreamTestServer } from "@internal/testcontainers/webapp"; +import { startSessionStreamTestServer } from "@internal/testcontainers/webapp"; +import { seedTestEnvironment } from "./helpers/seedTestEnvironment"; +import { + appendInput, + collectSessionOut, + isTurnComplete, + isUpgradeRequired, + mintSessionToken, + subscribeSessionOut, + type CollectedPart, +} from "./helpers/sessionStream"; +import { runChatAgentSession, runRealChatAgent } from "./helpers/agentHarness"; +import { + suspendResumeEvents, + testApprovalChatAgent, + testChatAgent, + testChatModelLocal, + testEndRunChatAgent, + testHitlChatAgent, + testHitlIdleChatAgent, + testIdleChatAgent, + testOomChatAgent, + testPlainChatAgent, + testSuspendHooksChatAgent, + testTimeoutChatAgent, + testToolChatAgent, + testUpgradeChatAgent, + testUpgradeOnceChatAgent, +} from "./helpers/testChatAgent"; + +async function waitFor(predicate: () => boolean, maxMs: number): Promise { + const deadline = performance.now() + maxMs; + while (performance.now() < deadline) { + if (predicate()) return; + await new Promise((r) => setTimeout(r, 100)); + } +} + +vi.setConfig({ testTimeout: 120_000, hookTimeout: 240_000 }); + +let server: SessionStreamTestServer; + +beforeAll(async () => { + server = await startSessionStreamTestServer(); +}, 240_000); + +afterAll(async () => { + await server?.stop(); +}, 120_000); + +function textModel(text: string) { + const chunks = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 5, reasoning: undefined }, + }, + }, + ]; + return new MockLanguageModelV3({ + doStream: async () => ({ stream: simulateReadableStream({ chunks: chunks as never }) }), + }); +} + +async function setupSession(agentId: string = testChatAgent.id) { + const { organization, project, environment, apiKey } = await seedTestEnvironment(server.prisma); + const addressingKey = `chat-${randomBytes(6).toString("hex")}`; + await server.prisma.session.create({ + data: { + friendlyId: `session_${randomBytes(8).toString("hex")}`, + externalId: addressingKey, + type: "chat.agent", + projectId: project.id, + runtimeEnvironmentId: environment.id, + environmentType: environment.type, + organizationId: organization.id, + taskIdentifier: agentId, + triggerConfig: { basePayload: {} }, + }, + }); + const token = await mintSessionToken({ apiKey, envId: environment.id, addressingKey }); + return { addressingKey, token, apiKey, baseUrl: server.webapp.baseUrl }; +} + +function promptText(prompt: unknown): string { + if (!Array.isArray(prompt)) return ""; + let out = ""; + for (const m of prompt) { + const c = (m as { content?: unknown }).content; + if (typeof c === "string") out += `${c} `; + else if (Array.isArray(c)) { + for (const part of c) { + if (part && typeof part === "object" && (part as { type?: string }).type === "text") { + out += `${(part as { text?: string }).text ?? ""} `; + } + } + } + } + return out; +} + +function echoModel() { + return new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + const text = promptText(prompt).trim(); + const chunks = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: `echo:${text}` }, + { type: "text-end", id: "t1" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, + }, + }, + ]; + return { stream: simulateReadableStream({ chunks: chunks as never }) }; + }, + }); +} + +function userMessage(text: string, id = "u0") { + return { id, role: "user", parts: [{ type: "text", text }] }; +} + +function submitBody(addressingKey: string, message: unknown, metadata: unknown = {}) { + return JSON.stringify({ + kind: "message", + payload: { message, chatId: addressingKey, trigger: "submit-message", metadata }, + }); +} + +function actionBody(addressingKey: string, action: unknown, metadata: unknown = {}) { + return JSON.stringify({ + kind: "message", + payload: { chatId: addressingKey, trigger: "action", action, metadata }, + }); +} + +const stopBody = JSON.stringify({ kind: "stop" }); + +/** + * Streams `words` as separate text-delta chunks with a delay between each, so + * a stop signal sent after the first chunk reliably truncates the turn before + * the model finishes. + */ +function slowModel(words: string[], delayMs: number) { + const chunks = [ + { type: "text-start", id: "t1" }, + ...words.map((delta) => ({ type: "text-delta", id: "t1", delta })), + { type: "text-end", id: "t1" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: words.length, text: words.length, reasoning: undefined }, + }, + }, + ]; + return new MockLanguageModelV3({ + doStream: async () => ({ + stream: simulateReadableStream({ chunks: chunks as never, chunkDelayInMs: delayMs }), + }), + }); +} + +function chunkType(part: CollectedPart): string | undefined { + const c = part.chunk as { type?: unknown } | null; + return c && typeof c === "object" && typeof c.type === "string" ? c.type : undefined; +} + +const FINISH_STOP = { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 5, reasoning: undefined }, + }, +}; + +/** + * A stateful model that emits a tool-call on its first `doStream` and a plain + * text response on every call after that. Drives both the automatic + * tool-execute loop (one turn) and the HITL round-trip (two turns). + */ +function toolCallThenText(opts: { + toolName: string; + toolCallId: string; + input: Record; + finalText: string; +}) { + const inputJson = JSON.stringify(opts.input); + const call1 = [ + { type: "tool-input-start", id: opts.toolCallId, toolName: opts.toolName }, + { type: "tool-input-delta", id: opts.toolCallId, delta: inputJson }, + { type: "tool-input-end", id: opts.toolCallId }, + { type: "tool-call", toolCallId: opts.toolCallId, toolName: opts.toolName, input: inputJson }, + { + type: "finish", + finishReason: { unified: "tool-calls", raw: "tool_calls" }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 0, reasoning: undefined }, + }, + }, + ]; + const call2 = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: opts.finalText }, + { type: "text-end", id: "t1" }, + FINISH_STOP, + ]; + let idx = 0; + return new MockLanguageModelV3({ + doStream: async () => ({ + stream: simulateReadableStream({ chunks: (idx++ === 0 ? call1 : call2) as never }), + }), + }); +} + +function joinChunks(parts: CollectedPart[]): string { + return parts + .filter((p) => p.chunk != null) + .map((p) => JSON.stringify(p.chunk)) + .join(""); +} + +/** A model that returns `texts[i]` on its i-th `doStream` call. */ +function sequenceModel(texts: string[]) { + let idx = 0; + return new MockLanguageModelV3({ + doStream: async () => { + const text = texts[Math.min(idx, texts.length - 1)]!; + idx++; + const chunks = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + FINISH_STOP, + ]; + return { stream: simulateReadableStream({ chunks: chunks as never }) }; + }, + }); +} + +function regenerateBody(addressingKey: string) { + return JSON.stringify({ + kind: "message", + payload: { chatId: addressingKey, trigger: "regenerate-message", metadata: {} }, + }); +} + +function findApprovalRequest( + parts: CollectedPart[] +): { approvalId: string; toolCallId: string } | undefined { + for (const p of parts) { + const c = p.chunk as { type?: string; approvalId?: string; toolCallId?: string } | null; + if (c && c.type === "tool-approval-request" && c.approvalId && c.toolCallId) { + return { approvalId: c.approvalId, toolCallId: c.toolCallId }; + } + } + return undefined; +} + +describe("session agent e2e (real chat.agent loop)", () => { + it("EA1: a real agent turn streams assistant text to .out", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(); + const agent = runRealChatAgent({ + agentId: testChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("hello from the agent"), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u0", + body: submitBody(addressingKey, userMessage("hi")), + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + const text = parts + .filter((p) => p.chunk != null) + .map((p) => JSON.stringify(p.chunk)) + .join(""); + expect(text).toContain("hello from the agent"); + expect(parts.some(isTurnComplete)).toBe(true); + } finally { + await agent.close(); + } + }); + + it("EA2: two turns continue on the same run", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(); + const agent = runRealChatAgent({ + agentId: testChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("reply"), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("one", "u1")), + }); + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, userMessage("two", "u2")), + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 40_000, + }); + + expect(parts.filter(isTurnComplete)).toHaveLength(2); + const replyText = parts + .filter((p) => p.chunk != null) + .map((p) => JSON.stringify(p.chunk)) + .join(""); + expect((replyText.match(/reply/g) ?? []).length).toBeGreaterThanOrEqual(2); + const seqs = parts.map((p) => Number(p.id)); + expect(seqs).toEqual([...seqs].sort((a, b) => a - b)); + } finally { + await agent.close(); + } + }); + + it("EA3: hydrateMessages injects prior history into the turn", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(); + const agent = runRealChatAgent({ + agentId: testChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + }); + + try { + const hydrated = [ + { id: "h1", role: "user", parts: [{ type: "text", text: "HISTORY-MARKER" }] }, + { id: "h2", role: "assistant", parts: [{ type: "text", text: "prior reply" }] }, + ]; + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("current question", "u1"), { hydrated }), + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + const text = parts + .filter((p) => p.chunk != null) + .map((p) => JSON.stringify(p.chunk)) + .join(""); + expect(text).toContain("HISTORY-MARKER"); + expect(text).toContain("current question"); + } finally { + await agent.close(); + } + }); + + it("EA4: onValidateMessages rejection writes an error chunk, keeps the run alive", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(); + const agent = runRealChatAgent({ + agentId: testChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("should not appear"), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("this has a blocked-word in it", "u1")), + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + const errorChunk = parts.find((p) => chunkType(p) === "error"); + expect(errorChunk, "an error chunk should be written to .out").toBeTruthy(); + expect(JSON.stringify((errorChunk as CollectedPart).chunk)).toContain("content filter"); + expect(parts.some(isTurnComplete), "the turn still completes after the error").toBe(true); + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, userMessage("hello now", "u2")), + }); + const follow = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 30_000, + }); + expect(follow.parts.filter(isTurnComplete).length).toBeGreaterThanOrEqual(2); + } finally { + await agent.close(); + } + }); + + it("EA5: a custom action fires onAction and completes without a turn", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(); + const agent = runRealChatAgent({ + agentId: testChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("turn reply"), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("first", "u1")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "act1", + body: actionBody(addressingKey, { type: "undo" }), + }); + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 30_000, + }); + + expect(parts.filter(isTurnComplete).length).toBeGreaterThanOrEqual(2); + expect( + parts.find((p) => chunkType(p) === "error"), + "no error on the action" + ).toBeFalsy(); + } finally { + await agent.close(); + } + }); + + it("EA6: a stop signal aborts the turn mid-stream but still completes it", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(); + const words = Array.from({ length: 12 }, (_, i) => ` w${i}`); + const agent = runRealChatAgent({ + agentId: testChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: slowModel(words, 250), + modelLocal: testChatModelLocal, + }); + + try { + const subscription = subscribeSessionOut({ baseUrl, addressingKey, token }); + const stream = await subscription.subscribe(); + const reader = stream.getReader(); + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("stream something long", "u1")), + }); + + const parts: CollectedPart[] = []; + let sentStop = false; + const deadline = performance.now() + 40_000; + try { + while (performance.now() < deadline) { + const remaining = deadline - performance.now(); + const next = await Promise.race([ + reader.read(), + new Promise<"timeout">((r) => setTimeout(() => r("timeout"), remaining)), + ]); + if (next === "timeout" || next.done) break; + const part = next.value as CollectedPart; + parts.push(part); + if (!sentStop && chunkType(part) === "text-delta") { + sentStop = true; + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "stop1", + body: stopBody, + }); + } + if (parts.some(isTurnComplete)) break; + } + } finally { + await reader.cancel().catch(() => {}); + } + + expect(sentStop, "a text-delta arrived so a stop could be sent").toBe(true); + expect(parts.some(isTurnComplete), "the aborted turn still writes turn-complete").toBe(true); + const deltas = parts.filter((p) => chunkType(p) === "text-delta").length; + expect(deltas, "generation was cut short before all deltas streamed").toBeLessThan( + words.length + ); + } finally { + await agent.close(); + } + }); + + it("EA7: the agent runs a server-side tool and streams the final answer", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testToolChatAgent.id); + const agent = runRealChatAgent({ + agentId: testToolChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: toolCallThenText({ + toolName: "getWeather", + toolCallId: "tc_weather_e2e", + input: { city: "Paris" }, + finalText: "It is 21C and clear in Paris", + }), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("what is the weather in Paris?", "u1")), + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 40_000, + }); + + const blob = joinChunks(parts); + expect(blob, "the tool call is streamed to .out").toContain("getWeather"); + expect(blob, "the tool result feeds a final answer").toContain( + "It is 21C and clear in Paris" + ); + expect(parts.some(isTurnComplete)).toBe(true); + } finally { + await agent.close(); + } + }); + + it("EA8: a HITL tool round-trip parks on the call and resumes from the client answer", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testHitlChatAgent.id); + const toolCallId = "tc_ask_e2e"; + const agent = runRealChatAgent({ + agentId: testHitlChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: toolCallThenText({ + toolName: "askUser", + toolCallId, + input: { question: "what color?" }, + finalText: "blue it is", + }), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("pick a color", "u1")), + }); + const turn1 = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + const blob1 = joinChunks(turn1.parts); + expect(blob1, "turn 1 streams the tool call").toContain("askUser"); + expect(blob1).toContain(toolCallId); + + const answer = { + id: "a-answer", + role: "assistant", + parts: [ + { + type: "tool-askUser", + toolCallId, + state: "output-available", + input: { question: "what color?" }, + output: { color: "blue" }, + }, + ], + }; + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, answer), + }); + const turn2 = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 30_000, + }); + expect(joinChunks(turn2.parts), "turn 2 resumes with the final answer").toContain( + "blue it is" + ); + } finally { + await agent.close(); + } + }); + + it("EA9: a continuation run restores prior history from the persisted snapshot", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testPlainChatAgent.id); + const runId1 = `run_persist_1_${addressingKey}`; + const runId2 = `run_persist_2_${addressingKey}`; + + const firstRun = runRealChatAgent({ + agentId: testPlainChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("assistant acknowledges"), + modelLocal: testChatModelLocal, + runId: runId1, + }); + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("REMEMBER-THIS-42", "u1")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + } finally { + await firstRun.close(); + } + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, userMessage("second question", "u2")), + }); + + const secondRun = runRealChatAgent({ + agentId: testPlainChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + runId: runId2, + continuation: true, + previousRunId: runId1, + }); + try { + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 40_000, + }); + const blob = joinChunks(parts); + expect(blob, "the restored prior user message is back in the model prompt").toContain( + "REMEMBER-THIS-42" + ); + expect(blob, "the new turn also ran").toContain("second question"); + } finally { + await secondRun.close(); + } + }); + + it("EA10: regenerate re-runs the last user turn with a fresh model call", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testPlainChatAgent.id); + const agent = runRealChatAgent({ + agentId: testPlainChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: sequenceModel(["first-answer", "regenerated-answer"]), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("explain it", "u1")), + }); + const first = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + expect(joinChunks(first.parts)).toContain("first-answer"); + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "regen1", + body: regenerateBody(addressingKey), + }); + const second = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 30_000, + }); + expect(joinChunks(second.parts), "the regenerated turn produced a fresh answer").toContain( + "regenerated-answer" + ); + } finally { + await agent.close(); + } + }); + + it("EA11: a needsApproval tool parks on an approval request and runs once approved", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testApprovalChatAgent.id); + const toolCallId = "tc_delete_e2e"; + const agent = runRealChatAgent({ + agentId: testApprovalChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: toolCallThenText({ + toolName: "deleteResource", + toolCallId, + input: { resource: "widget-1" }, + finalText: "deleted widget-1", + }), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("delete widget-1", "u1")), + }); + const turn1 = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + const approval = findApprovalRequest(turn1.parts); + expect(approval, "turn 1 emits an approval request instead of executing").toBeTruthy(); + expect(joinChunks(turn1.parts), "the tool did not execute before approval").not.toContain( + "deleted widget-1" + ); + + const answer = { + id: "a-approval", + role: "assistant", + parts: [ + { + type: "tool-deleteResource", + toolCallId: approval!.toolCallId, + state: "approval-responded", + approval: { id: approval!.approvalId, approved: true }, + }, + ], + }; + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, answer), + }); + const turn2 = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 30_000, + }); + expect( + joinChunks(turn2.parts), + "after approval the tool runs and the agent answers" + ).toContain("deleted widget-1"); + } finally { + await agent.close(); + } + }); + + it("EA12: two chats stay isolated - neither run bleeds into the other's .out", async () => { + const chatA = await setupSession(testPlainChatAgent.id); + const chatB = await setupSession(testPlainChatAgent.id); + + const agentA = runRealChatAgent({ + agentId: testPlainChatAgent.id, + baseUrl: chatA.baseUrl, + addressingKey: chatA.addressingKey, + secretKey: chatA.apiKey, + model: textModel("answer-for-A"), + modelLocal: testChatModelLocal, + }); + try { + await appendInput({ + baseUrl: chatA.baseUrl, + addressingKey: chatA.addressingKey, + token: chatA.token, + partId: "a1", + body: submitBody(chatA.addressingKey, userMessage("hi from A", "a1")), + }); + await collectSessionOut({ + baseUrl: chatA.baseUrl, + addressingKey: chatA.addressingKey, + token: chatA.token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + } finally { + await agentA.close(); + } + + const agentB = runRealChatAgent({ + agentId: testPlainChatAgent.id, + baseUrl: chatB.baseUrl, + addressingKey: chatB.addressingKey, + secretKey: chatB.apiKey, + model: textModel("answer-for-B"), + modelLocal: testChatModelLocal, + }); + try { + await appendInput({ + baseUrl: chatB.baseUrl, + addressingKey: chatB.addressingKey, + token: chatB.token, + partId: "b1", + body: submitBody(chatB.addressingKey, userMessage("hi from B", "b1")), + }); + await collectSessionOut({ + baseUrl: chatB.baseUrl, + addressingKey: chatB.addressingKey, + token: chatB.token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + } finally { + await agentB.close(); + } + + const outA = joinChunks( + ( + await collectSessionOut({ + baseUrl: chatA.baseUrl, + addressingKey: chatA.addressingKey, + token: chatA.token, + until: (p) => p.some(isTurnComplete), + maxMs: 15_000, + }) + ).parts + ); + const outB = joinChunks( + ( + await collectSessionOut({ + baseUrl: chatB.baseUrl, + addressingKey: chatB.addressingKey, + token: chatB.token, + until: (p) => p.some(isTurnComplete), + maxMs: 15_000, + }) + ).parts + ); + + expect(outA).toContain("answer-for-A"); + expect(outA, "A's stream never sees B's output").not.toContain("answer-for-B"); + expect(outB).toContain("answer-for-B"); + expect(outB, "B's stream never sees A's output").not.toContain("answer-for-A"); + }); + + it("EA13: the run suspends on the idle waitpoint, then the next message resumes it in place", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testIdleChatAgent.id); + const agent = runRealChatAgent({ + agentId: testIdleChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("resumed after suspend"), + modelLocal: testChatModelLocal, + }); + + try { + await new Promise((r) => setTimeout(r, 2500)); + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("hi after idle", "u1")), + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + expect( + joinChunks(parts), + "a message sent well after the idle window still produces a turn, so the waitpoint resume delivered it" + ).toContain("resumed after suspend"); + expect(parts.some(isTurnComplete)).toBe(true); + } finally { + await agent.close(); + } + }); + + it("EA14: chat.endRun() ends the run; the next message continues on a fresh run", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testEndRunChatAgent.id); + const session = runChatAgentSession({ + agentId: testEndRunChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("MARKER-ONE", "u1")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + await waitFor(() => session.runCount() >= 2, 10_000); + expect( + session.runCount(), + "endRun exited run 1, so the orchestrator spawned a continuation" + ).toBeGreaterThanOrEqual(2); + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, userMessage("second msg", "u2")), + }); + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 30_000, + }); + + const blob = joinChunks(parts); + expect(blob, "the continuation run restored run 1's history").toContain("MARKER-ONE"); + expect(blob, "the continuation run ran the new turn").toContain("second msg"); + } finally { + await session.close(); + } + }); + + it("EA15: chat.requestUpgrade() emits upgrade-required on .out and exits the run", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testUpgradeChatAgent.id); + const agent = runRealChatAgent({ + agentId: testUpgradeChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + }); + + let exitedOnItsOwn = false; + agent.done.then(() => { + exitedOnItsOwn = true; + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("upgrade me", "u1")), + }); + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isUpgradeRequired), + maxMs: 30_000, + }); + + expect( + parts.some(isUpgradeRequired), + "an upgrade-required control record is written to .out" + ).toBe(true); + + await waitFor(() => exitedOnItsOwn, 10_000); + expect(exitedOnItsOwn, "the run exits itself after requesting the upgrade").toBe(true); + } finally { + await agent.close(); + } + }); + + it("EA16: a HITL tool approval survives a suspend/resume boundary", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testHitlIdleChatAgent.id); + const toolCallId = "tc_ask_suspend"; + const agent = runRealChatAgent({ + agentId: testHitlIdleChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: toolCallThenText({ + toolName: "askUser", + toolCallId, + input: { question: "what color?" }, + finalText: "blue it is", + }), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("pick a color", "u1")), + }); + const turn1 = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + expect(joinChunks(turn1.parts), "turn 1 parks on the tool call").toContain("askUser"); + + await new Promise((r) => setTimeout(r, 2500)); + + const answer = { + id: "a-answer", + role: "assistant", + parts: [ + { + type: "tool-askUser", + toolCallId, + state: "output-available", + input: { question: "what color?" }, + output: { color: "blue" }, + }, + ], + }; + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, answer), + }); + const turn2 = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => joinChunks(p).includes("blue it is"), + maxMs: 30_000, + }); + expect( + joinChunks(turn2.parts), + "the answer sent after the idle window resumed the suspended run" + ).toContain("blue it is"); + } finally { + await agent.close(); + } + }); + + it("EA17: the run suspends and resumes across multiple turns", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testIdleChatAgent.id); + const replies = ["reply-one", "reply-two", "reply-three"]; + const agent = runRealChatAgent({ + agentId: testIdleChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: sequenceModel(replies), + modelLocal: testChatModelLocal, + }); + + try { + for (let i = 0; i < replies.length; i++) { + await new Promise((r) => setTimeout(r, i === 0 ? 1500 : 2000)); + await appendInput({ + baseUrl, + addressingKey, + token, + partId: `u${i + 1}`, + body: submitBody(addressingKey, userMessage(`turn ${i + 1}`, `u${i + 1}`)), + }); + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => joinChunks(p).includes(replies[i]!), + maxMs: 30_000, + }); + expect( + joinChunks(parts), + `turn ${i + 1} resumed from a suspend and produced its reply` + ).toContain(replies[i]!); + } + } finally { + await agent.close(); + } + }); + + it("EA18: requestUpgrade defers the message; the upgraded run processes it", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession( + testUpgradeOnceChatAgent.id + ); + const session = runChatAgentSession({ + agentId: testUpgradeOnceChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("DEFER-ME", "u1")), + }); + const turn1 = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isUpgradeRequired), + maxMs: 30_000, + }); + expect( + turn1.parts.some(isUpgradeRequired), + "the fresh run defers the message with upgrade-required" + ).toBe(true); + + await waitFor(() => session.runCount() >= 2, 10_000); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => joinChunks(p).includes("DEFER-ME"), + maxMs: 30_000, + }); + expect( + joinChunks(parts), + "the continuation run processed the deferred message instead of upgrading again" + ).toContain("DEFER-ME"); + } finally { + await session.close(); + } + }); + + it("EA19: endRun continuation restores history across multiple hops", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testEndRunChatAgent.id); + const session = runChatAgentSession({ + agentId: testEndRunChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("HOP-MARKER", "u1")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + await waitFor(() => session.runCount() >= 2, 10_000); + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, userMessage("second hop", "u2")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => joinChunks(p).includes("second hop"), + maxMs: 30_000, + }); + + await waitFor(() => session.runCount() >= 3, 10_000); + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u3", + body: submitBody(addressingKey, userMessage("third hop", "u3")), + }); + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => joinChunks(p).includes("third hop"), + maxMs: 30_000, + }); + + const blob = joinChunks(parts); + expect(blob, "history from the first hop is restored two continuations later").toContain( + "HOP-MARKER" + ); + expect(blob).toContain("third hop"); + expect(session.runCount()).toBeGreaterThanOrEqual(3); + } finally { + await session.close(); + } + }); + + it("EA20: onChatSuspend and onChatResume fire around a suspend/resume", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession( + testSuspendHooksChatAgent.id + ); + const agent = runRealChatAgent({ + agentId: testSuspendHooksChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: sequenceModel(["turn-a", "turn-b"]), + modelLocal: testChatModelLocal, + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("first", "u1")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + await new Promise((r) => setTimeout(r, 2500)); + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u2", + body: submitBody(addressingKey, userMessage("second", "u2")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => joinChunks(p).includes("turn-b"), + maxMs: 30_000, + }); + + const mine = suspendResumeEvents.filter((e) => e.chatId === addressingKey); + expect( + mine.some((e) => e.kind === "suspend"), + "onChatSuspend fired when the run suspended" + ).toBe(true); + expect( + mine.some((e) => e.kind === "resume"), + "onChatResume fired when the next message resumed it" + ).toBe(true); + } finally { + await agent.close(); + } + }); + + it("EA21: a preloaded run's OOM retry recovers the in-flight message", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testOomChatAgent.id); + const runId = `run_oom_${addressingKey}`; + + const attempt1 = runRealChatAgent({ + agentId: testOomChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + runId, + attemptNumber: 1, + }); + let attempt1Error: unknown; + attempt1.done.catch((e) => { + attempt1Error = e; + }); + + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("OOM-THEN-OK", "u1")), + }); + await waitFor(() => attempt1Error !== undefined, 15_000); + expect( + String(attempt1Error), + "attempt 1 fails with an OOM so the runtime can swap machines" + ).toMatch(/OutOfMemory/i); + + const attempt2 = runRealChatAgent({ + agentId: testOomChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: echoModel(), + modelLocal: testChatModelLocal, + runId, + attemptNumber: 2, + }); + try { + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => joinChunks(p).includes("OOM-THEN-OK"), + maxMs: 30_000, + }); + expect( + joinChunks(parts), + "the attempt-2 retry restored the unprocessed message and ran it" + ).toContain("OOM-THEN-OK"); + } finally { + await attempt2.close(); + } + }); + + it("EA22: the idle wait times out and ends the run when no message arrives", async () => { + const { addressingKey, token, apiKey, baseUrl } = await setupSession(testTimeoutChatAgent.id); + const agent = runRealChatAgent({ + agentId: testTimeoutChatAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("turn one"), + modelLocal: testChatModelLocal, + }); + + let ended = false; + agent.done.then(() => { + ended = true; + }); + + try { + await appendInput({ + baseUrl, + addressingKey, + token, + partId: "u1", + body: submitBody(addressingKey, userMessage("hello", "u1")), + }); + await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + + await waitFor(() => ended, 20_000); + expect( + ended, + "with no next message the between-turns wait times out and the run exits itself" + ).toBe(true); + } finally { + await agent.close(); + } + }); +}); diff --git a/apps/webapp/test/session-stream.browser.e2e.test.ts b/apps/webapp/test/session-stream.browser.e2e.test.ts new file mode 100644 index 00000000000..4eeebf826ae --- /dev/null +++ b/apps/webapp/test/session-stream.browser.e2e.test.ts @@ -0,0 +1,133 @@ +/** + * Browser-level session-stream e2e. + * + * Same full stack as session-stream.e2e.test.ts, but the client is a real + * Chromium page on a DIFFERENT origin than the webapp, driven via + * page.evaluate. It exercises what node-fetch can't: a real browser doing a + * cross-origin fetch + streaming read of the session `.out` SSE through the + * webapp proxy, so the browser enforces the CORS preflight + response headers + * the customer scenario depends on (a frontend on its own origin subscribing + * to the API). The webapp runs production-mode here, so this checks the + * production CORS path, not the dev-server one. + * + * Requires a pre-built webapp (pnpm run build --filter webapp) and Chromium + * (pnpm exec playwright install chromium). + */ +import { randomBytes } from "crypto"; +import { createServer, type Server } from "http"; +import { chromium, type Browser } from "@playwright/test"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import type { SessionStreamTestServer } from "@internal/testcontainers/webapp"; +import { startSessionStreamTestServer } from "@internal/testcontainers/webapp"; +import { seedTestEnvironment } from "./helpers/seedTestEnvironment"; +import { + mintSessionToken, + SessionStreamProducer, + sessionStreamName, +} from "./helpers/sessionStream"; + +vi.setConfig({ testTimeout: 120_000, hookTimeout: 240_000 }); + +let server: SessionStreamTestServer; +let browser: Browser | undefined; +let origin: Server; +let originUrl: string; + +beforeAll(async () => { + server = await startSessionStreamTestServer(); + try { + browser = await chromium.launch(); + } catch (error) { + console.warn("[browser-e2e] Chromium unavailable, leg will skip:", String(error)); + } + origin = createServer((_req, res) => { + res.writeHead(200, { "content-type": "text/html" }); + res.end("origin"); + }); + await new Promise((resolve) => origin.listen(0, "127.0.0.1", () => resolve())); + const addr = origin.address(); + const port = typeof addr === "object" && addr ? addr.port : 0; + originUrl = `http://127.0.0.1:${port}`; +}, 240_000); + +afterAll(async () => { + await browser?.close().catch(() => {}); + await new Promise((resolve) => (origin ? origin.close(() => resolve()) : resolve())); + await server?.stop(); +}, 120_000); + +describe("session stream browser e2e", () => { + it("EB1: a cross-origin browser subscribes to .out and streams records", async (ctx) => { + if (!browser) { + ctx.skip(); + return; + } + const { organization, environment, apiKey } = await seedTestEnvironment(server.prisma); + const addressingKey = `sess-${randomBytes(6).toString("hex")}`; + const token = await mintSessionToken({ apiKey, envId: environment.id, addressingKey }); + const streamName = sessionStreamName({ + orgId: organization.id, + envSlug: environment.slug, + envId: environment.id, + addressingKey, + }); + const producer = new SessionStreamProducer({ + endpoint: server.s2.endpoint, + basin: server.s2.basin, + streamName, + }); + + await producer.appendData({ n: 0 }, "p0"); + await producer.appendData({ n: 1 }, "p1"); + await producer.appendTurnComplete(); + + const sseUrl = `${server.webapp.baseUrl}/realtime/v1/sessions/${encodeURIComponent( + addressingKey + )}/out`; + + const page = await browser.newPage(); + try { + await page.goto(originUrl); + const result = await page.evaluate( + async ({ url, token }) => { + const ac = new AbortController(); + const deadlineTimer = setTimeout(() => ac.abort(), 8000); + try { + const res = await fetch(url, { + headers: { Authorization: `Bearer ${token}`, Accept: "text/event-stream" }, + signal: ac.signal, + }); + const reader = (res.body as ReadableStream).getReader(); + const decoder = new TextDecoder(); + let text = ""; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + text += decoder.decode(value, { stream: true }); + if (text.includes('"records"')) break; + } + ac.abort(); + return { + ok: res.ok, + status: res.status, + sawBatch: text.includes('"records"'), + pageOrigin: location.origin, + }; + } catch (e) { + return { error: String(e) }; + } finally { + clearTimeout(deadlineTimer); + } + }, + { url: sseUrl, token } + ); + + expect("error" in result ? result.error : undefined).toBeUndefined(); + expect(result).toMatchObject({ ok: true, status: 200, sawBatch: true }); + expect((result as { pageOrigin: string }).pageOrigin).toBe(originUrl); + expect(originUrl).not.toBe(server.webapp.baseUrl); + } finally { + await page.close(); + } + }); +}); diff --git a/apps/webapp/test/session-stream.e2e.test.ts b/apps/webapp/test/session-stream.e2e.test.ts new file mode 100644 index 00000000000..63e68a8ce48 --- /dev/null +++ b/apps/webapp/test/session-stream.e2e.test.ts @@ -0,0 +1,308 @@ +/** + * Full-stack session-stream e2e. + * + * Boots the real webapp + Postgres + Redis + s2-lite 0.40.0 (via + * startSessionStreamTestServer), then drives the session `.out` wire protocol + * directly: a producer appends records straight to S2 (the agent simulator), + * and the client subscribes through the webapp SSE proxy using the real + * `SSEStreamSubscription` from `@trigger.dev/core` โ€” the same code the browser + * runs, minus the DOM. + * + * Requires a pre-built webapp: pnpm run build --filter webapp + */ +import { randomBytes } from "crypto"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import type { SessionStreamTestServer } from "@internal/testcontainers/webapp"; +import { startSessionStreamTestServer } from "@internal/testcontainers/webapp"; +import { seedTestEnvironment } from "./helpers/seedTestEnvironment"; +import { + appendInput, + collectSessionOut, + collectUntilCaughtUp, + isTurnComplete, + mintSessionToken, + openChannelRaw, + SessionStreamProducer, + sessionStreamName, +} from "./helpers/sessionStream"; + +vi.setConfig({ testTimeout: 120_000, hookTimeout: 180_000 }); + +let server: SessionStreamTestServer; + +beforeAll(async () => { + server = await startSessionStreamTestServer(); +}, 180_000); + +afterAll(async () => { + await server?.stop(); +}, 120_000); + +async function setupSession() { + const { organization, project, environment, apiKey } = await seedTestEnvironment(server.prisma); + const addressingKey = `sess-${randomBytes(6).toString("hex")}`; + await server.prisma.session.create({ + data: { + friendlyId: `session_${randomBytes(8).toString("hex")}`, + externalId: addressingKey, + type: "chat.agent", + projectId: project.id, + runtimeEnvironmentId: environment.id, + environmentType: environment.type, + organizationId: organization.id, + taskIdentifier: "chat-agent", + triggerConfig: {}, + }, + }); + const token = await mintSessionToken({ apiKey, envId: environment.id, addressingKey }); + const streamName = sessionStreamName({ + orgId: organization.id, + envSlug: environment.slug, + envId: environment.id, + addressingKey, + }); + const producer = new SessionStreamProducer({ + endpoint: server.s2.endpoint, + basin: server.s2.basin, + streamName, + }); + return { addressingKey, token, producer, baseUrl: server.webapp.baseUrl }; +} + +describe("session stream e2e", () => { + it("E1 basic: data records + turn-complete are delivered in order", async () => { + const { addressingKey, token, producer, baseUrl } = await setupSession(); + + await producer.appendData({ n: 0 }, "p0"); + await producer.appendData({ n: 1 }, "p1"); + await producer.appendData({ n: 2 }, "p2"); + await producer.appendTurnComplete(); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.some(isTurnComplete), + maxMs: 20_000, + }); + + const dataChunks = parts + .filter((p) => !isTurnComplete(p) && p.chunk != null) + .map((p) => (p.chunk as { n: number }).n); + + expect(dataChunks).toEqual([0, 1, 2]); + expect(parts.some(isTurnComplete)).toBe(true); + + const seqs = parts.map((p) => Number(p.id)); + expect(seqs).toEqual([...seqs].sort((a, b) => a - b)); + }); + + it("E2 continuation: seq is monotonic across two turns on the same .out", async () => { + const { addressingKey, token, producer, baseUrl } = await setupSession(); + + await producer.appendData({ n: 0 }, "p0"); + await producer.appendTurnComplete(); + await producer.appendData({ n: 1 }, "p1"); + await producer.appendData({ n: 2 }, "p2"); + const tc2 = await producer.appendTurnComplete(); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter(isTurnComplete).length >= 2, + maxMs: 20_000, + }); + + expect(parts.filter(isTurnComplete)).toHaveLength(2); + const seqs = parts.map((p) => Number(p.id)); + expect(seqs).toEqual([...seqs].sort((a, b) => a - b)); + expect(Math.max(...seqs)).toBe(tc2); + const dataChunks = parts + .filter((p) => !isTurnComplete(p) && p.chunk != null) + .map((p) => (p.chunk as { n: number }).n); + expect(dataChunks).toEqual([0, 1, 2]); + }); + + it("E5 no-dup reconnect: resume via Last-Event-ID delivers each record once", async () => { + const { addressingKey, token, producer, baseUrl } = await setupSession(); + + await producer.appendData({ n: 0 }, "p0"); + await producer.appendData({ n: 1 }, "p1"); + const s2 = await producer.appendData({ n: 2 }, "p2"); + + const first = await collectSessionOut({ + baseUrl, + addressingKey, + token, + until: (p) => p.filter((x) => x.chunk != null).length >= 3, + maxMs: 20_000, + }); + const firstSeqs = first.parts.map((p) => Number(p.id)); + expect(firstSeqs).toContain(s2); + + await producer.appendData({ n: 3 }, "p3"); + await producer.appendData({ n: 4 }, "p4"); + + const second = await collectSessionOut({ + baseUrl, + addressingKey, + token, + lastEventId: String(s2), + until: (p) => p.filter((x) => x.chunk != null).length >= 2, + maxMs: 20_000, + }); + const secondSeqs = second.parts.map((p) => Number(p.id)); + + const all = [...firstSeqs, ...secondSeqs]; + expect(new Set(all).size).toBe(all.length); + const allData = [...first.parts, ...second.parts] + .filter((p) => p.chunk != null) + .map((p) => (p.chunk as { n: number }).n) + .sort((a, b) => a - b); + expect(allData).toEqual([0, 1, 2, 3, 4]); + }); + + it("E7 trim-at-tail: trim command record is filtered; caught-up still fires", async () => { + const { addressingKey, token, producer, baseUrl } = await setupSession(); + + await producer.appendData({ n: 0 }, "p0"); + const keep = await producer.appendData({ n: 1 }, "p1"); + await producer.appendTurnComplete(); + await producer.trim(keep); + + const { parts, caughtUp } = await collectUntilCaughtUp({ + baseUrl, + addressingKey, + token, + maxMs: 15_000, + }); + + expect(caughtUp).toBe(true); + const commandRecords = parts.filter((p) => (p.headers ?? []).some(([k]) => k === "")); + expect(commandRecords).toHaveLength(0); + expect(parts.filter((p) => p.chunk != null).length).toBeGreaterThanOrEqual(1); + }); + + it("E3 quiescent reconnect (GREEN): client caught-up close settles at the tail", async () => { + const { addressingKey, token, producer, baseUrl } = await setupSession(); + + await producer.appendData({ n: 0 }, "p0"); + await producer.appendData({ n: 1 }, "p1"); + const tc = await producer.appendTurnComplete(); + + const { parts, caughtUp, settleMs } = await collectUntilCaughtUp({ + baseUrl, + addressingKey, + token, + lastEventId: String(tc), + timeoutInSeconds: 30, + maxMs: 15_000, + }); + + expect(caughtUp).toBe(true); + expect(settleMs).toBeLessThan(5_000); + expect(parts.filter((p) => p.chunk != null)).toHaveLength(0); + }); + + it("E4 backlog reaches tail (GREEN): every record is delivered before caught-up", async () => { + const { addressingKey, token, producer, baseUrl } = await setupSession(); + + await producer.appendData({ n: 0 }, "p0"); + await producer.appendData({ n: 1 }, "p1"); + await producer.appendData({ n: 2 }, "p2"); + const tc = await producer.appendTurnComplete(); + + const { parts, caughtUp, tailSeqNum } = await collectUntilCaughtUp({ + baseUrl, + addressingKey, + token, + maxMs: 15_000, + }); + + expect(caughtUp).toBe(true); + expect(tailSeqNum).toBe(tc + 1); + const dataChunks = parts + .filter((p) => !isTurnComplete(p) && p.chunk != null) + .map((p) => (p.chunk as { n: number }).n); + expect(dataChunks).toEqual([0, 1, 2]); + expect(parts.some(isTurnComplete)).toBe(true); + }); + + it("E8 in/append 413 for an oversized body still carries CORS headers", async () => { + const { addressingKey, token, baseUrl } = await setupSession(); + + const oversized = "x".repeat(2 * 1024 * 1024); + const { status, acao } = await appendInput({ + baseUrl, + addressingKey, + token, + origin: "http://example.com", + body: JSON.stringify({ kind: "message", payload: { big: oversized } }), + }); + + expect(status).toBe(413); + expect(acao).not.toBeNull(); + }); + + it("E9 server peek fast-closes at a turn-complete tail with X-Session-Settled", async () => { + const { addressingKey, token, producer, baseUrl } = await setupSession(); + + await producer.appendData({ n: 0 }, "p0"); + const tc = await producer.appendTurnComplete(); + + const { status, sessionSettled, closedMs } = await openChannelRaw({ + baseUrl, + addressingKey, + token, + lastEventId: String(tc), + peekSettled: true, + timeoutInSeconds: 30, + maxMs: 10_000, + }); + + expect(status).toBe(200); + expect(sessionSettled).toBe("true"); + expect(closedMs).toBeLessThan(5_000); + }); + + it("E11 in/append delivers the record on the .in channel", async () => { + const { addressingKey, token, baseUrl } = await setupSession(); + + const payload = JSON.stringify({ kind: "message", text: "hello from client" }); + const appended = await appendInput({ + baseUrl, + addressingKey, + token, + partId: "in-0", + body: payload, + }); + expect(appended.status).toBe(200); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token, + io: "in", + until: (p) => p.some((x) => x.chunk != null), + maxMs: 15_000, + }); + + const got = parts.find((p) => p.chunk != null); + expect(got).toBeTruthy(); + expect(String(got?.chunk)).toContain("hello from client"); + }); + + it("E12 subscribe with an invalid token is rejected", async () => { + const { addressingKey, baseUrl } = await setupSession(); + + const { status } = await openChannelRaw({ + baseUrl, + addressingKey, + token: "tr_pub_invalid_not_a_real_token", + maxMs: 5_000, + }); + + expect([401, 403]).toContain(status); + }); +}); diff --git a/apps/webapp/vite.config.ts b/apps/webapp/vite.config.ts index b008ae04ed0..0b33a4edb6f 100644 --- a/apps/webapp/vite.config.ts +++ b/apps/webapp/vite.config.ts @@ -30,6 +30,7 @@ export default defineConfig({ }, }, server: { + cors: false, warmup: { clientFiles: ["./app/entry.client.tsx", "./app/root.tsx", "./app/components/**/*.tsx"], ssrFiles: ["./app/entry.server.tsx", "./app/root.tsx"], diff --git a/docs/ai-chat/client-protocol.mdx b/docs/ai-chat/client-protocol.mdx index fd8326e4ba9..f0adf3a0812 100644 --- a/docs/ai-chat/client-protocol.mdx +++ b/docs/ai-chat/client-protocol.mdx @@ -320,7 +320,7 @@ The output stream uses [S2](https://s2.dev) under the hood and follows the stand | Event | Meaning | | --- | --- | | `batch` | One or more records. The records you actually care about. | -| `ping` | Keepalive (~every 5s on idle). Body is `{"timestamp": }`. Ignore it. | +| `ping` | Keepalive (~every 5s on idle). Body is `{"timestamp": }`, and on backends that report it a `tail` (same `{seq_num, timestamp}` shape as a batch tail). The tail is what lets you tell you have caught up to the live edge; if you don't need caught-up detection, ignore the ping. | | _(no `event:`, just `data: [DONE]`)_ | Stream is closing โ€” server sends this once before EOF. | A `batch` event in raw SSE format looks like this โ€” note the `data` is a single line of JSON, no embedded newlines (per the SSE spec): @@ -359,7 +359,7 @@ Decoded `data` payload: | `records[].timestamp` | Unix ms when the record was written to S2. | | `records[].body` | For data records: a JSON-encoded **string** wrapping `{ data: UIMessageChunk, id: string }`. For control records: an empty string (semantics live in `headers`). For S2 command records: opaque bytes. See [Records on session.out](#records-on-session-out). | | `records[].headers` | Optional `[name, value]` pairs. Empty for data records; a `trigger-control` entry for control records; a single empty-name `["", ""]` entry for S2 command records. | -| `tail.seq_num` | Latest known tail of the S2 stream โ€” useful for detecting how far behind the live edge you are. Skip if you don't need it. | +| `tail.seq_num` | Latest known tail of the S2 stream, useful for detecting how far behind the live edge you are. Track the **highest `seq_num` you have received**, counting every record in the batch (including the command records you skip, see below), not just the last application-visible one. When `highest received seq_num + 1 === tail.seq_num` you have drained the backlog and are caught up to the live edge. The same `tail` also rides on `ping` events. Skip if you don't need it. | | `tail.timestamp` | Timestamp of `tail.seq_num`. | ### Records on `session.out` @@ -652,6 +652,8 @@ On **reconnect-on-reload** paths (resuming a chat where nothing may be streaming **Do not send `X-Peek-Settled` on the active-send response-stream path.** The peek would race the newly-triggered turn's first chunk โ€” if the agent hasn't written the new turn's first record yet, the peek sees the prior turn's `turn-complete` and closes the SSE before the response lands on S2. The built-in `TriggerChatTransport.reconnectToStream` sets the header; `sendMessages โ†’ subscribeToStream` does not. +If you use `TriggerChatTransport` (or `useChat`, which builds on it), the transport settles the stream for you: on the reconnect path it watches the `tail` on batch and ping events via `SSEStreamSubscription.caughtUp()` and closes the resumed stream as soon as it reaches the live edge, so a settled idle reconnect closes promptly without waiting out the long poll. `SSEStreamSubscription` on its own only exposes the caught-up signal, it does not close itself: consuming it directly, call `caughtUp()` (or compare your last received `seq_num` to the ping/batch `tail`) and close your own stream. On older self-hosted backends whose `ping` carries no `tail`, this falls back to the `X-Peek-Settled` behavior above. + ```ts // Reconnect path (page reload) const response = await fetch(sseUrl, { diff --git a/internal-packages/testcontainers/src/s2.ts b/internal-packages/testcontainers/src/s2.ts new file mode 100644 index 00000000000..fe2a0494774 --- /dev/null +++ b/internal-packages/testcontainers/src/s2.ts @@ -0,0 +1,67 @@ +import type { StartedNetwork, StartedTestContainer } from "testcontainers"; +import { GenericContainer, Wait } from "testcontainers"; +import { withCiResourceLimits } from "./utils"; + +/** + * s2-lite 0.40.0 is the first image that emits the tail on the heartbeat ping + * (matching cloud S2 0.25.0), which the caught-up client keys off. Older images + * fire a bare keepalive with no tail, so the digest is pinned to keep the e2e + * deterministic. + */ +const S2_LITE_IMAGE = + "ghcr.io/s2-streamstore/s2:0.40.0@sha256:b26249e2ede0949755f5af8028185dc2bcfc3aa2db21eb9610543d144eb6ee9d"; + +const DEFAULT_BASIN = "trigger-local"; + +function specJson(basin: string): string { + return JSON.stringify({ + basins: [ + { + name: basin, + config: { create_stream_on_append: true, create_stream_on_read: true }, + }, + ], + }); +} + +export interface StartedS2Container { + container: StartedTestContainer; + /** Base URL of the s2-lite HTTP API, e.g. `http://localhost:49xxx`. */ + endpoint: string; + /** The single basin s2-lite is initialised with. */ + basin: string; +} + +/** + * Boots s2-lite (the open-source S2 server) in `lite` mode with a single basin, + * for full-stack realtime session-stream e2e. The spec is copied in rather than + * bind-mounted so the helper carries no external file dependency. + */ +export async function createS2Container( + network?: StartedNetwork, + opts: { basin?: string } = {} +): Promise { + const basin = opts.basin ?? DEFAULT_BASIN; + + let builder = withCiResourceLimits(new GenericContainer(S2_LITE_IMAGE)) + .withExposedPorts(80) + .withCopyContentToContainer([ + { content: specJson(basin), target: "/s2-spec.json", mode: 0o444 }, + ]) + .withCommand(["lite", "--init-file", "/s2-spec.json"]) + .withWaitStrategy(Wait.forLogMessage(/starting plain http server/)) + .withStartupTimeout(60_000); + + if (network) { + builder = builder.withNetwork(network).withNetworkAliases("s2"); + } + + const container = await builder.start(); + const mappedPort = container.getMappedPort(80); + + return { + container, + endpoint: `http://${container.getHost()}:${mappedPort}`, + basin, + }; +} diff --git a/internal-packages/testcontainers/src/webapp.ts b/internal-packages/testcontainers/src/webapp.ts index fc587abd696..2685b6fa7d7 100644 --- a/internal-packages/testcontainers/src/webapp.ts +++ b/internal-packages/testcontainers/src/webapp.ts @@ -3,7 +3,9 @@ import { createServer } from "net"; import { delimiter, resolve } from "path"; import { Network } from "testcontainers"; import { PrismaClient } from "@trigger.dev/database"; -import { createPostgresContainer, createRedisContainer } from "./utils"; +import { createPostgresContainer, createRedisContainer, withCiResourceLimits } from "./utils"; +import { createS2Container, type StartedS2Container } from "./s2"; +import { MinIOContainer, type StartedMinIOContainer } from "./minio"; const WEBAPP_ROOT = resolve(__dirname, "../../../apps/webapp"); // pnpm hoists transitive deps to node_modules/.pnpm/node_modules but does NOT symlink them @@ -65,6 +67,13 @@ export interface StartWebappOptions { * behaviour when the loader is short-circuited). */ requirePlugins?: boolean; + + /** + * Extra environment variables merged into the spawned webapp, after the + * defaults so they override them (e.g. the `REALTIME_STREAMS_S2_*` vars for + * session-stream e2e). `NODE_PATH` and the worker-disable vars still win. + */ + extraEnv?: Record; } export async function startWebapp( @@ -115,6 +124,7 @@ export async function startWebapp( REDIS_HOST: redis.host, REDIS_PORT: String(redis.port), REDIS_TLS_DISABLED: "true", // all *_REDIS_TLS_DISABLED vars default to this; test Redis has no TLS + ...(options.extraEnv ?? {}), // Disable all background workers. Each worker has its own env var and its own // check idiom ("0" vs "false" vs boolean), so we set all of them explicitly. WORKER_ENABLED: "false", // disables workerQueue.initialize() (checked === "true") @@ -255,3 +265,87 @@ export async function startTestServer(options: StartWebappOptions = {}): Promise return { webapp, prisma: prisma!, databaseUrl: pgUrl!, stop }; } + +export { createS2Container } from "./s2"; +export type { StartedS2Container } from "./s2"; + +export interface SessionStreamTestServer extends TestServer { + s2: StartedS2Container; + minio: StartedMinIOContainer; +} + +/** + * Full-stack session-stream harness: postgres + redis + s2-lite 0.40.0 + the + * real webapp, wired for S2 v2 realtime session streams. The webapp is a host + * process reaching every container over its mapped port, so the S2 endpoint is + * the mapped localhost URL (the docker-network alias is unusable from the host). + */ +export async function startSessionStreamTestServer(): Promise { + const network = await new Network().start(); + + let pgContainer: Awaited>["container"] | undefined; + let pgUrl: string | undefined; + let redisContainer: Awaited>["container"] | undefined; + let s2: StartedS2Container | undefined; + let minio: StartedMinIOContainer | undefined; + let prisma: PrismaClient | undefined; + let stopWebapp: (() => Promise) | undefined; + let webapp: WebappInstance; + + try { + const pg = await createPostgresContainer(network); + pgContainer = pg.container; + pgUrl = pg.url; + + const { container: rc } = await createRedisContainer({ network }); + redisContainer = rc; + + s2 = await createS2Container(network); + minio = await withCiResourceLimits(new MinIOContainer()).withNetwork(network).start(); + const minioConfig = minio.getConnectionConfig(); + + prisma = new PrismaClient({ datasources: { db: { url: pg.url } } }); + await prisma.$connect(); + + const started = await startWebapp( + pg.url, + { host: rc.getHost(), port: rc.getPort() }, + { + extraEnv: { + REALTIME_STREAMS_S2_BASIN: s2.basin, + REALTIME_STREAMS_S2_ENDPOINT: `${s2.endpoint}/v1`, + REALTIME_STREAMS_S2_SKIP_ACCESS_TOKENS: "true", + REALTIME_STREAMS_DEFAULT_VERSION: "v2", + OBJECT_STORE_BASE_URL: minioConfig.baseUrl, + OBJECT_STORE_BUCKET: "packets", + OBJECT_STORE_ACCESS_KEY_ID: minioConfig.accessKeyId, + OBJECT_STORE_SECRET_ACCESS_KEY: minioConfig.secretAccessKey, + OBJECT_STORE_REGION: minioConfig.region, + }, + } + ); + webapp = started.instance; + stopWebapp = started.stop; + } catch (err) { + await stopWebapp?.().catch(() => {}); + await prisma?.$disconnect().catch(() => {}); + await minio?.stop().catch(() => {}); + await s2?.container.stop().catch(() => {}); + await pgContainer?.stop().catch(() => {}); + await redisContainer?.stop().catch(() => {}); + await network.stop().catch(() => {}); + throw err; + } + + const stop = async () => { + await stopWebapp!().catch((err) => console.error("stopWebapp failed:", err)); + await prisma!.$disconnect().catch((err) => console.error("prisma.$disconnect failed:", err)); + await minio!.stop().catch((err) => console.error("minio.stop failed:", err)); + await s2!.container.stop().catch((err) => console.error("s2.stop failed:", err)); + await pgContainer!.stop().catch((err) => console.error("pgContainer.stop failed:", err)); + await redisContainer!.stop().catch((err) => console.error("redisContainer.stop failed:", err)); + await network.stop().catch((err) => console.error("network.stop failed:", err)); + }; + + return { webapp, prisma: prisma!, databaseUrl: pgUrl!, s2: s2!, minio: minio!, stop }; +} diff --git a/packages/cli-v3/package.json b/packages/cli-v3/package.json index c6d8d9719dc..318ff9268d5 100644 --- a/packages/cli-v3/package.json +++ b/packages/cli-v3/package.json @@ -95,7 +95,7 @@ "@opentelemetry/resources": "2.7.1", "@opentelemetry/sdk-trace-node": "2.7.1", "@opentelemetry/semantic-conventions": "1.41.1", - "@s2-dev/streamstore": "^0.22.10", + "@s2-dev/streamstore": "^0.25.0", "@trigger.dev/build": "workspace:4.5.7", "@trigger.dev/core": "workspace:4.5.7", "@trigger.dev/schema-to-json": "workspace:4.5.7", diff --git a/packages/core/package.json b/packages/core/package.json index f0c3adfcf2f..7eb81b899ed 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -207,7 +207,7 @@ "@opentelemetry/sdk-trace-base": "2.7.1", "@opentelemetry/sdk-trace-node": "2.7.1", "@opentelemetry/semantic-conventions": "1.41.1", - "@s2-dev/streamstore": "0.22.10", + "@s2-dev/streamstore": "0.25.0", "dequal": "^2.0.3", "eventsource": "^3.0.5", "eventsource-parser": "^3.0.0", diff --git a/packages/core/src/v3/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts index deda5d01f9a..58e8e60b7a7 100644 --- a/packages/core/src/v3/apiClient/index.ts +++ b/packages/core/src/v3/apiClient/index.ts @@ -1410,6 +1410,13 @@ export class ApiClient { * enqueued into the consumer stream โ€” handle the event here. */ onControl?: (event: ControlEvent) => void; + /** + * Fires once when the session reaches the live tail (backlog drained), + * with the observed tail position. No-op when the backend omits the + * tail-carrying heartbeat. Mirrors the browser transport's caught-up + * signal on the worker / apiClient read path. + */ + onCaughtUp?: (tail: { seqNum: number; timestamp: Date }) => void; } ): Promise> { const url = `${options?.baseUrl ?? this.baseUrl}/realtime/v1/sessions/${encodeURIComponent(sessionIdOrExternalId)}/${io}`; @@ -1424,6 +1431,12 @@ export class ApiClient { }); const stream = await subscription.subscribe(); + if (options?.onCaughtUp) { + subscription + .caughtUp() + .then(options.onCaughtUp) + .catch(() => {}); + } const onPart = options?.onPart; const onControl = options?.onControl; diff --git a/packages/core/src/v3/apiClient/runStream.test.ts b/packages/core/src/v3/apiClient/runStream.test.ts index 0dca73779af..3226f6d9c65 100644 --- a/packages/core/src/v3/apiClient/runStream.test.ts +++ b/packages/core/src/v3/apiClient/runStream.test.ts @@ -27,6 +27,22 @@ describe("SSEStreamSubscription retry behavior", () => { }); } + /** + * A response.body that emits one SSE event then stays open (never closes), + * so the connection sits mid-read until the fetch signal aborts. + */ + function makeOpenSSEResponse() { + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(`id: 1\ndata: {"hello":1}\n\n`)); + }, + }); + return new Response(body, { + status: 200, + headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v1" }, + }); + } + // Drain a ReadableStream until it closes or errors. // Returns received chunks plus terminal state. async function drain(stream: ReadableStream<{ id: string; chunk: unknown }>) { @@ -74,6 +90,30 @@ describe("SSEStreamSubscription retry behavior", () => { expect(result.chunks).toHaveLength(1); }); + it("does not reconnect after the consumer cancels the returned stream", async () => { + let attempts = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + attempts++; + return makeOpenSSEResponse(); + }); + + const sub = new SSEStreamSubscription("http://example.test/sse", { + retryDelayMs: 1, + maxRetryDelayMs: 5, + maxRetries: 10, + }); + + const stream = await sub.subscribe(); + const reader = stream.getReader(); + const first = await reader.read(); + expect(first.done).toBe(false); + + await reader.cancel(); + await new Promise((r) => setTimeout(r, 50)); + + expect(attempts).toBe(1); + }); + it("caps the exponential backoff at maxRetryDelayMs", async () => { let attempts = 0; const callTimes: number[] = []; @@ -600,3 +640,206 @@ describe("SSEStreamSubscription v2 batch parsing โ€” record kinds", () => { expect((parts[1]!.chunk as any).delta).toBe("x"); }); }); + +describe("SSEStreamSubscription caught-up tracking", () => { + const originalFetch = globalThis.fetch; + + afterEach(() => { + globalThis.fetch = originalFetch; + vi.restoreAllMocks(); + }); + + type Rec = { + body: string; + seq_num: number; + timestamp: number; + headers?: Array<[string, string]>; + }; + type Tail = { seq_num: number; timestamp: number }; + + function dataRec(seq: number): Rec { + return { + body: JSON.stringify({ data: { type: "text-delta", delta: "x" }, id: `p${seq}` }), + seq_num: seq, + timestamp: 1, + headers: [], + }; + } + + function batchEvent(records: Rec[], tail?: Tail): string { + const data = tail ? { records, tail } : { records }; + return `event: batch\ndata: ${JSON.stringify(data)}\n\n`; + } + + function pingEvent(tail?: Tail): string { + const data = tail ? { timestamp: 1, tail } : { timestamp: 1 }; + return `event: ping\ndata: ${JSON.stringify(data)}\n\n`; + } + + function makeEventsResponse(events: string[]) { + const body = new ReadableStream({ + start(controller) { + for (const e of events) controller.enqueue(new TextEncoder().encode(e)); + controller.close(); + }, + }); + return new Response(body, { + status: 200, + headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v2" }, + }); + } + + async function drain(stream: ReadableStream<{ id: string; chunk: unknown }>) { + const reader = stream.getReader(); + const parts: Array<{ id: string; chunk: unknown }> = []; + while (true) { + const { done, value } = await reader.read(); + if (done) { + reader.releaseLock(); + return parts; + } + parts.push(value); + } + } + + it("resolves caughtUp() when a batch reaches the reported tail", async () => { + globalThis.fetch = vi + .fn() + .mockResolvedValue( + makeEventsResponse([ + batchEvent([dataRec(0), dataRec(1), dataRec(2)], { seq_num: 3, timestamp: 1 }), + ]) + ); + const sub = new SSEStreamSubscription("http://x", { maxRetries: 0 }); + const stream = await sub.subscribe(); + const cu = sub.caughtUp(); + await drain(stream); + const tail = await cu; + expect(tail.seqNum).toBe(3); + expect(sub.isCaughtUp()).toBe(true); + }); + + it("stays behind when the batch does not reach the tail", async () => { + globalThis.fetch = vi + .fn() + .mockResolvedValue( + makeEventsResponse([batchEvent([dataRec(0)], { seq_num: 3, timestamp: 1 })]) + ); + const sub = new SSEStreamSubscription("http://x", { maxRetries: 0 }); + const stream = await sub.subscribe(); + await drain(stream); + expect(sub.isCaughtUp()).toBe(false); + }); + + it("resolves caughtUp() from a ping tail with an empty backlog (open-at-tail)", async () => { + globalThis.fetch = vi + .fn() + .mockResolvedValue(makeEventsResponse([pingEvent({ seq_num: 3, timestamp: 1 })])); + const sub = new SSEStreamSubscription("http://x", { maxRetries: 0 }); + const stream = await sub.subscribe(); + const cu = sub.caughtUp(); + await drain(stream); + const tail = await cu; + expect(tail.seqNum).toBe(3); + expect(sub.isCaughtUp()).toBe(true); + }); + + it("reaches caught-up when the tail is a trim command record (raw counts include it)", async () => { + globalThis.fetch = vi.fn().mockResolvedValue( + makeEventsResponse([ + batchEvent( + [ + dataRec(0), + { + body: "", + seq_num: 1, + timestamp: 1, + headers: [["trigger-control", "turn-complete"]], + }, + { body: "AAAAAAAAAAQ=", seq_num: 2, timestamp: 1, headers: [["", "trim"]] }, + ], + { seq_num: 3, timestamp: 1 } + ), + ]) + ); + const sub = new SSEStreamSubscription("http://x", { maxRetries: 0 }); + const stream = await sub.subscribe(); + const cu = sub.caughtUp(); + const parts = await drain(stream); + const tail = await cu; + expect(tail.seqNum).toBe(3); + expect(sub.isCaughtUp()).toBe(true); + expect(parts).toHaveLength(2); + }); + + it("never reaches caught-up when the wire carries no tail (feature-detect fallback)", async () => { + globalThis.fetch = vi + .fn() + .mockResolvedValue( + makeEventsResponse([batchEvent([dataRec(0), dataRec(1), dataRec(2)]), pingEvent()]) + ); + const sub = new SSEStreamSubscription("http://x", { maxRetries: 0 }); + const stream = await sub.subscribe(); + await drain(stream); + expect(sub.isCaughtUp()).toBe(false); + }); + + it("does not resolve caughtUp() until the consumer drains the tail", async () => { + globalThis.fetch = vi + .fn() + .mockResolvedValue( + makeEventsResponse([ + batchEvent([dataRec(0), dataRec(1), dataRec(2)], { seq_num: 3, timestamp: 1 }), + ]) + ); + const sub = new SSEStreamSubscription("http://x", { maxRetries: 0 }); + const stream = await sub.subscribe(); + const cu = sub.caughtUp(); + let resolved = false; + cu.then(() => { + resolved = true; + }).catch(() => {}); + + await new Promise((r) => setTimeout(r, 20)); + expect(resolved).toBe(false); + expect(sub.isCaughtUp()).toBe(false); + + const parts = await drain(stream); + const tail = await cu; + expect(tail.seqNum).toBe(3); + expect(parts).toHaveLength(3); + expect(sub.isCaughtUp()).toBe(true); + }); + + it("holds caughtUp() until the final record before the tail is consumed", async () => { + globalThis.fetch = vi + .fn() + .mockResolvedValue( + makeEventsResponse([ + batchEvent([dataRec(0), dataRec(1), dataRec(2)], { seq_num: 3, timestamp: 1 }), + ]) + ); + const sub = new SSEStreamSubscription("http://x", { maxRetries: 0 }); + const stream = await sub.subscribe(); + const cu = sub.caughtUp(); + let resolved = false; + cu.then(() => { + resolved = true; + }).catch(() => {}); + const reader = stream.getReader(); + + const first = await reader.read(); + expect(first.value?.id).toBe("0"); + await reader.read(); + await reader.read(); + await new Promise((r) => setTimeout(r, 20)); + expect(resolved).toBe(false); + + const done = await reader.read(); + expect(done.done).toBe(true); + await new Promise((r) => setTimeout(r, 20)); + expect(resolved).toBe(true); + expect(sub.isCaughtUp()).toBe(true); + reader.releaseLock(); + }); +}); diff --git a/packages/core/src/v3/apiClient/runStream.ts b/packages/core/src/v3/apiClient/runStream.ts index 5bd7986c86a..17dbfd83306 100644 --- a/packages/core/src/v3/apiClient/runStream.ts +++ b/packages/core/src/v3/apiClient/runStream.ts @@ -14,6 +14,8 @@ import { conditionallyImportAndParsePacket, parsePacket } from "../utils/ioSeria import { ApiError, isTriggerRealtimeAuthError } from "./errors.js"; import type { ApiClient } from "./index.js"; import { zodShapeStream } from "./stream.js"; +import type { CaughtUpBoundary } from "@s2-dev/streamstore"; +import { CaughtUpTracker } from "@s2-dev/streamstore"; export type RunShape = TRunTypes extends AnyRunTypes ? { @@ -184,6 +186,16 @@ export type SSEStreamPart = { headers?: Array<[string, string]>; }; +/** + * Internal item flowing from the decode transform to the consumer-facing + * stream. A `boundary` entry is emitted after the visible records of the batch + * (or ping) it belongs to, so the wrapper marks it delivered only once the + * consumer has drained past those records. + */ +type PumpItem = + | { type: "part"; part: SSEStreamPart } + | { type: "boundary"; boundary: CaughtUpBoundary }; + // Real implementation for production export class SSEStreamSubscription implements StreamSubscription { private lastEventId: string | undefined; @@ -197,6 +209,8 @@ export class SSEStreamSubscription implements StreamSubscription { private nonRetryableStatuses: ReadonlySet; private retryNowController: AbortController | null = null; private internalAbort: AbortController | null = null; + private cancelledByConsumer = false; + private caughtUpTracker = new CaughtUpTracker(); constructor( private url: string, @@ -278,22 +292,91 @@ export class SSEStreamSubscription implements StreamSubscription { this.retryNowController?.abort(); } + /** + * True once this session has consumed everything up to the tail the server + * last reported (via a batch `tail` or a heartbeat `ping`). Resets to false + * on reconnect and while records remain before the reported tail. Backed by + * S2's `CaughtUpTracker`, fed from the v2 batch/ping wire signals. + */ + isCaughtUp(): boolean { + return this.caughtUpTracker.isCaughtUp(); + } + + /** + * Resolves when this session reaches the live tail (backlog drained), + * carrying the last observed tail position. Resolves immediately if already + * caught up; call again after falling behind. Rejects if the stream ends + * before catching up. Stays pending across internal reconnects. + */ + caughtUp(): ReturnType { + return this.caughtUpTracker.caughtUp(); + } + + /** + * The transport pumps decoded items (records and caught-up boundaries) into + * an internal stream; the returned stream drains it on demand. It uses + * `highWaterMark: 0` so each consumer read pulls exactly one item, which lets + * a caught-up boundary be marked delivered only once every visible record + * preceding it has actually been consumed. The terminal `end()` fires here, + * after that drain, so a caught-up `caughtUp()` resolves rather than being + * rejected by an early end from the transport. + */ async subscribe(): Promise> { // eslint-disable-next-line no-this-alias const self = this; - return new ReadableStream({ + const internal = new ReadableStream({ async start(controller) { await self.connectStream(controller); }, cancel() { - self.options.onComplete?.(); + self.cancelledByConsumer = true; + self.internalAbort?.abort(); + self.retryNowController?.abort(); }, }); + const internalReader = internal.getReader(); + + return new ReadableStream( + { + async pull(controller) { + while (true) { + let result: ReadableStreamReadResult; + try { + result = await internalReader.read(); + } catch (err) { + self.caughtUpTracker.end(); + controller.error(err); + return; + } + if (result.done) { + self.caughtUpTracker.end(); + self.options.onComplete?.(); + controller.close(); + return; + } + const item = result.value; + if (item.type === "boundary") { + item.boundary.markDelivered(); + continue; + } + controller.enqueue(item.part); + return; + } + }, + cancel(reason) { + self.cancelledByConsumer = true; + self.caughtUpTracker.end(); + self.options.onComplete?.(); + internalReader.cancel(reason).catch(() => {}); + }, + }, + { highWaterMark: 0 } + ); } private async connectStream( - controller: ReadableStreamDefaultController + controller: ReadableStreamDefaultController ): Promise { // Two abort sources flow through `internalAbort.signal`: // - this.options.signal: caller cancel โ€” bypass retry, exit cleanly. @@ -386,7 +469,7 @@ export class SSEStreamSubscription implements StreamSubscription { .pipeThrough(new TextDecoderStream()) .pipeThrough(new EventSourceParserStream()) .pipeThrough( - new TransformStream({ + new TransformStream({ transform: (chunk, chunkController) => { if (streamVersion === "v1") { if (chunk.id) { @@ -394,9 +477,12 @@ export class SSEStreamSubscription implements StreamSubscription { } const timestamp = parseRedisStreamIdTimestamp(chunk.id); chunkController.enqueue({ - id: chunk.id ?? "unknown", - chunk: safeParseJSON(chunk.data), - timestamp, + type: "part", + part: { + id: chunk.id ?? "unknown", + chunk: safeParseJSON(chunk.data), + timestamp, + }, }); } else { if (chunk.event === "batch") { @@ -407,9 +493,18 @@ export class SSEStreamSubscription implements StreamSubscription { timestamp: number; headers?: Array<[string, string]>; }>; + tail?: { seq_num: number; timestamp: number }; }; if (!data || !Array.isArray(data.records)) return; + const boundary = this.caughtUpTracker.observeBatch({ + recordCount: data.records.length, + lastSeqNum: data.records.at(-1)?.seq_num, + tail: data.tail + ? { seqNum: data.tail.seq_num, timestamp: new Date(data.tail.timestamp) } + : undefined, + }); + for (const record of data.records) { // Always advance the resume cursor โ€” even for records we // skip โ€” so a future Last-Event-ID reconnect lands past @@ -438,12 +533,35 @@ export class SSEStreamSubscription implements StreamSubscription { rememberSeen(parsedBody.id); } chunkController.enqueue({ - id: record.seq_num.toString(), - chunk: parsedBody?.data, - timestamp: record.timestamp, - headers: record.headers ?? [], + type: "part", + part: { + id: record.seq_num.toString(), + chunk: parsedBody?.data, + timestamp: record.timestamp, + headers: record.headers ?? [], + }, }); } + + if (boundary) { + chunkController.enqueue({ type: "boundary", boundary }); + } + } else if (chunk.event === "ping") { + const ping = safeParseJSON(chunk.data) as + | { tail?: { seq_num: number; timestamp: number } } + | undefined; + if (ping?.tail) { + const pingBoundary = this.caughtUpTracker.observeBatch({ + recordCount: 0, + tail: { + seqNum: ping.tail.seq_num, + timestamp: new Date(ping.tail.timestamp), + }, + }); + if (pingBoundary) { + chunkController.enqueue({ type: "boundary", boundary: pingBoundary }); + } + } } } }, @@ -459,7 +577,6 @@ export class SSEStreamSubscription implements StreamSubscription { if (done) { reader.releaseLock(); controller.close(); - this.options.onComplete?.(); return; } @@ -467,7 +584,6 @@ export class SSEStreamSubscription implements StreamSubscription { reader.cancel(); reader.releaseLock(); controller.close(); - this.options.onComplete?.(); return; } @@ -479,10 +595,9 @@ export class SSEStreamSubscription implements StreamSubscription { throw error; } } catch (error) { - if (this.options.signal?.aborted) { + if (this.options.signal?.aborted || this.cancelledByConsumer) { // User cancel โ€” exit cleanly, don't retry. controller.close(); - this.options.onComplete?.(); return; } @@ -505,9 +620,8 @@ export class SSEStreamSubscription implements StreamSubscription { controller: ReadableStreamDefaultController, error?: Error ): Promise { - if (this.options.signal?.aborted) { + if (this.options.signal?.aborted || this.cancelledByConsumer) { controller.close(); - this.options.onComplete?.(); return; } @@ -547,13 +661,13 @@ export class SSEStreamSubscription implements StreamSubscription { }); this.retryNowController = null; - if (this.options.signal?.aborted) { + if (this.options.signal?.aborted || this.cancelledByConsumer) { controller.close(); - this.options.onComplete?.(); return; } // Reconnect + this.caughtUpTracker.reconnect(); await this.connectStream(controller); } } diff --git a/packages/core/src/v3/sessionStreams/manager.ts b/packages/core/src/v3/sessionStreams/manager.ts index fb87b211643..6f22757e589 100644 --- a/packages/core/src/v3/sessionStreams/manager.ts +++ b/packages/core/src/v3/sessionStreams/manager.ts @@ -466,6 +466,11 @@ export class StandardSessionStreamManager implements SessionStreamManager { console.error(`[SessionStreamManager] Tail error for "${key}":`, error); } }, + onCaughtUp: (tail) => { + if (this.debug) { + console.log(`[SessionStreamManager] Caught up on "${key}" at seq ${tail.seqNum}`); + } + }, }); // Drain to keep the pipeThrough flowing. Records were already diff --git a/packages/core/src/v3/test/index.ts b/packages/core/src/v3/test/index.ts index 402f618c01b..ec7ba85d913 100644 --- a/packages/core/src/v3/test/index.ts +++ b/packages/core/src/v3/test/index.ts @@ -7,3 +7,10 @@ export { TestInputStreamManager } from "./test-input-stream-manager.js"; export { TestRealtimeStreamsManager } from "./test-realtime-streams-manager.js"; export { TestRunMetadataManager } from "./test-run-metadata-manager.js"; export { TestSessionStreamManager } from "./test-session-stream-manager.js"; +export { StandardSessionStreamManager } from "../sessionStreams/manager.js"; +export type { SessionStreamManager } from "../sessionStreams/types.js"; +export { + SessionWaitpointBackend, + TestRuntimeManager, + installSessionWaitpointBackend, +} from "./session-waitpoint-backend.js"; diff --git a/packages/core/src/v3/test/mock-task-context.ts b/packages/core/src/v3/test/mock-task-context.ts index 0da633c6a7e..5fbe1957613 100644 --- a/packages/core/src/v3/test/mock-task-context.ts +++ b/packages/core/src/v3/test/mock-task-context.ts @@ -9,10 +9,11 @@ import { runtime } from "../runtime-api.js"; import { StandardLocalsManager } from "../locals/manager.js"; import { StandardLifecycleHooksManager } from "../lifecycleHooks/manager.js"; import { NoopRuntimeManager } from "../runtime/noopRuntimeManager.js"; +import type { RuntimeManager } from "../runtime/manager.js"; import { unregisterGlobal } from "../utils/globals.js"; import type { ServerBackgroundWorker, TaskRunContext } from "../schemas/index.js"; import type { LocalsKey } from "../locals/types.js"; -import type { SessionChannelIO } from "../sessionStreams/types.js"; +import type { SessionChannelIO, SessionStreamManager } from "../sessionStreams/types.js"; import { TestInputStreamManager } from "./test-input-stream-manager.js"; import { TestRealtimeStreamsManager } from "./test-realtime-streams-manager.js"; import { TestRunMetadataManager } from "./test-run-metadata-manager.js"; @@ -45,6 +46,20 @@ export type MockTaskContextOptions = { worker?: Partial; /** Whether this is a warm start. */ isWarmStart?: boolean; + /** + * Session-streams manager installed as the `session-streams` global. Defaults + * to an in-memory {@link TestSessionStreamManager}. Pass a real + * `StandardSessionStreamManager` (with an ApiClient pointed at a running + * webapp) to drive the task's `.in`/`.out` against real streams. + */ + sessionStreamManager?: SessionStreamManager; + /** + * Runtime manager installed as the `runtime` global. Defaults to a + * {@link NoopRuntimeManager}. Pass a `TestRuntimeManager` (wired to a + * {@link SessionWaitpointBackend}) to make `session.in.wait()` suspend and + * resume in place against real streams, without the run-engine. + */ + runtimeManager?: RuntimeManager; }; /** @@ -215,11 +230,11 @@ export async function runInMockTaskContext( const localsManager = new StandardLocalsManager(); const lifecycleManager = new StandardLifecycleHooksManager(); - const runtimeManager = new NoopRuntimeManager(); + const runtimeManager = options?.runtimeManager ?? new NoopRuntimeManager(); const metadataManager = new TestRunMetadataManager(); const inputManager = new TestInputStreamManager(); const outputManager = new TestRealtimeStreamsManager(); - const sessionStreamManager = new TestSessionStreamManager(); + const sessionStreamManager = options?.sessionStreamManager ?? new TestSessionStreamManager(); // Unregister any previously-installed managers so `setGlobal*` wins โ€” // `registerGlobal` returns false silently if an entry already exists. @@ -263,8 +278,16 @@ export async function runInMockTaskContext( sessions: { in: { send: (sessionId, data, io = "in") => - sessionStreamManager.__sendFromTest(sessionId, io, data), - close: (sessionId, io = "in") => sessionStreamManager.__closeFromTest(sessionId, io), + sessionStreamManager instanceof TestSessionStreamManager + ? sessionStreamManager.__sendFromTest(sessionId, io, data) + : Promise.reject( + new Error("drivers.sessions.in.send requires the default TestSessionStreamManager") + ), + close: (sessionId, io = "in") => { + if (sessionStreamManager instanceof TestSessionStreamManager) { + sessionStreamManager.__closeFromTest(sessionId, io); + } + }, }, }, ctx, @@ -286,7 +309,9 @@ export async function runInMockTaskContext( localsManager.reset(); inputManager.reset(); outputManager.reset(); - sessionStreamManager.reset(); + if (sessionStreamManager instanceof TestSessionStreamManager) { + sessionStreamManager.reset(); + } metadataManager.reset(); } } diff --git a/packages/core/src/v3/test/session-waitpoint-backend.ts b/packages/core/src/v3/test/session-waitpoint-backend.ts new file mode 100644 index 00000000000..092c90ba6c2 --- /dev/null +++ b/packages/core/src/v3/test/session-waitpoint-backend.ts @@ -0,0 +1,257 @@ +import { ApiClient } from "../apiClient/index.js"; +import { WaitpointId } from "../isomorphic/friendlyId.js"; +import { NoopRuntimeManager } from "../runtime/noopRuntimeManager.js"; +import type { + CreateSessionStreamWaitpointRequestBody, + CreateSessionStreamWaitpointResponseBody, + WaitForWaitpointTokenResponseBody, +} from "../schemas/api.js"; +import type { WaitpointTokenResult } from "../schemas/common.js"; + +type PendingWait = { + session: string; + io: "in" | "out"; + lastSeqNum?: number; + timeout?: string; + abort: AbortController; +}; + +const TIMED_OUT = Symbol("session-waitpoint-timeout"); + +function parseTimeoutMs(timeout: string | undefined): number | undefined { + if (!timeout) { + return undefined; + } + const match = /^(\d+)(ms|s|m|h)$/.exec(timeout.trim()); + if (!match) { + return undefined; + } + const value = Number(match[1]); + switch (match[2]) { + case "ms": + return value; + case "s": + return value * 1_000; + case "m": + return value * 60_000; + default: + return value * 3_600_000; + } +} + +/** + * In-process stand-in for the run-engine's session-stream waitpoint machinery. + * + * `session.in.wait()` suspends a run on a waitpoint; in production the run-engine + * completes that waitpoint (and the supervisor resumes the run) when the next + * `.in` record arrives. This backend reproduces the task-observable half in one + * process: `register()` mints a waitpoint id when the SDK creates the waitpoint, + * and `wait()` opens its own short-lived tail on the session channel and resolves + * with the next record. That record is handed back through `TestRuntimeManager` + * exactly as the real runtime hands back a completed waitpoint's output, so the + * same `run()` invocation continues in place (the faithful dev / deployed + * task-observable semantics: same promise, same record, no re-invocation). + * + * It deliberately does not model server-side waitpoint bookkeeping, process + * checkpoint/restore, or client-vs-server timeout, none of which are visible to + * task code on the resume path. + */ +export class SessionWaitpointBackend { + private readonly pending = new Map(); + + constructor(private readonly apiClient: ApiClient) {} + + register( + body: CreateSessionStreamWaitpointRequestBody + ): CreateSessionStreamWaitpointResponseBody { + const waitpointId = WaitpointId.generate().friendlyId; + this.pending.set(waitpointId, { + session: body.session, + io: body.io, + lastSeqNum: body.lastSeqNum, + timeout: body.timeout, + abort: new AbortController(), + }); + return { waitpointId, isCached: false }; + } + + async wait(waitpointFriendlyId: string): Promise { + const pending = this.pending.get(waitpointFriendlyId); + if (!pending) { + return { ok: true }; + } + this.pending.delete(waitpointFriendlyId); + + const timeoutMs = parseTimeoutMs(pending.timeout); + let timer: ReturnType | undefined; + + try { + const recordPromise = this.readNextRecord(pending); + const result = + timeoutMs === undefined + ? await recordPromise + : await Promise.race([ + recordPromise, + new Promise((resolve) => { + timer = setTimeout(() => resolve(TIMED_OUT), timeoutMs); + }), + ]); + + if (result === TIMED_OUT) { + pending.abort.abort(); + return { + ok: false, + output: JSON.stringify({ message: "Timed out" }), + outputType: "application/json", + }; + } + + const output = typeof result === "string" ? result : JSON.stringify(result); + return { ok: true, output, outputType: "application/json" }; + } catch { + return { + ok: false, + output: JSON.stringify({ message: "Session stream wait ended before a record arrived" }), + outputType: "application/json", + }; + } finally { + if (timer) { + clearTimeout(timer); + } + } + } + + disable(): void { + for (const pending of this.pending.values()) { + pending.abort.abort(); + } + this.pending.clear(); + } + + /** + * Opens an independent tail on the session channel starting after the + * caller's last-seen seq and resolves with the next record. For `.in` the + * tail yields the raw appended JSON string (already `JSON.stringify(chunk)`), + * 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 { + const lastEventId = + pending.lastSeqNum !== undefined && pending.lastSeqNum >= 0 + ? String(pending.lastSeqNum) + : undefined; + + const stream = await this.apiClient.subscribeToSessionStream(pending.session, pending.io, { + lastEventId, + signal: pending.abort.signal, + timeoutInSeconds: 120, + }); + + const reader = stream.getReader(); + try { + const { value, done } = await reader.read(); + if (done) { + throw new Error("session stream closed"); + } + return value; + } finally { + await reader.cancel().catch(() => {}); + pending.abort.abort(); + } + } +} + +/** + * A {@link NoopRuntimeManager} whose `waitForWaitpoint` resolves in place from a + * {@link SessionWaitpointBackend} instead of hanging. Install it as the `runtime` + * global via `runInMockTaskContext({ runtimeManager })`. + */ +export class TestRuntimeManager extends NoopRuntimeManager { + constructor(private readonly backend: SessionWaitpointBackend) { + super(); + } + + override waitForWaitpoint(params: { + waitpointFriendlyId: string; + finishDate?: Date; + }): Promise { + return this.backend.wait(params.waitpointFriendlyId); + } +} + +let activeBackend: SessionWaitpointBackend | undefined; +let patchDepth = 0; +let originalCreate: ApiClient["createSessionStreamWaitpoint"] | undefined; +let originalWait: ApiClient["waitForWaitpointToken"] | undefined; + +/** + * Route the two waitpoint apiClient calls through an in-process backend for the + * duration of a harness run, and hand back a {@link TestRuntimeManager} to + * install as the `runtime` global. + * + * `apiClientManager.clientOrThrow()` builds a fresh `ApiClient` per call, so + * there is no instance to swap; the two methods are patched on the prototype + * (ref-counted, restored by `restore()`). `createSessionStreamWaitpoint` + * registers a pending wait and returns a real-shaped `{ waitpointId, isCached }`; + * `waitForWaitpointToken` returns `{ success: true }` exactly like the real route, + * which only creates the block edge and returns synchronously. + */ +export function installSessionWaitpointBackend(apiClient: ApiClient): { + backend: SessionWaitpointBackend; + runtimeManager: TestRuntimeManager; + restore: () => void; +} { + const backend = new SessionWaitpointBackend(apiClient); + activeBackend = backend; + + if (patchDepth === 0) { + originalCreate = ApiClient.prototype.createSessionStreamWaitpoint; + originalWait = ApiClient.prototype.waitForWaitpointToken; + + ApiClient.prototype.createSessionStreamWaitpoint = function ( + this: ApiClient, + runFriendlyId: string, + body: CreateSessionStreamWaitpointRequestBody + ): Promise { + if (activeBackend) { + return Promise.resolve(activeBackend.register(body)); + } + return originalCreate!.call(this, runFriendlyId, body); + } as unknown as ApiClient["createSessionStreamWaitpoint"]; + + ApiClient.prototype.waitForWaitpointToken = function ( + this: ApiClient + ): Promise { + if (activeBackend) { + return Promise.resolve({ success: true }); + } + // eslint-disable-next-line prefer-rest-params + return (originalWait as ApiClient["waitForWaitpointToken"]).apply(this, arguments as never); + } as unknown as ApiClient["waitForWaitpointToken"]; + } + + patchDepth += 1; + let restored = false; + + return { + backend, + runtimeManager: new TestRuntimeManager(backend), + restore() { + if (restored) { + return; + } + restored = true; + backend.disable(); + if (activeBackend === backend) { + activeBackend = undefined; + } + patchDepth -= 1; + if (patchDepth === 0 && originalCreate && originalWait) { + ApiClient.prototype.createSessionStreamWaitpoint = originalCreate; + ApiClient.prototype.waitForWaitpointToken = originalWait; + originalCreate = undefined; + originalWait = undefined; + } + }, + }; +} diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 4612302c4ae..73268ab4c39 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -6062,8 +6062,21 @@ function chatAgent< ); } + /** + * A retry (e.g. OOM) or continuation whose boot recovery restored an + * in-flight `.in` message dispatches it as the first turn. The recovery + * block already advanced the `.in` cursor past the recovered message, + * so the preload wait below would otherwise strand it (the run would + * sit waiting for a "first message" that already arrived). + */ + let dispatchedRecoveredFirstTurn = false; + if (preloaded && bootInjectedQueue.length > 0) { + currentWirePayload = bootInjectedQueue.shift()!; + dispatchedRecoveredFirstTurn = true; + } + // Handle preloaded runs โ€” fire onPreload, then wait for the first real message - if (preloaded) { + if (preloaded && !dispatchedRecoveredFirstTurn) { if (activeSpan) { activeSpan.setAttribute("chat.preloaded", true); } diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 9dddac1c0eb..e4e965c5973 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -1725,6 +1725,8 @@ export class TriggerChatTransport implements ChatTransport { ); }) as typeof fetch) : undefined; + let sawFirstChunk = false; + const connectSseOnce = async (token: string) => { const subscription = new SSEStreamSubscription(streamUrl, { headers: { @@ -1744,13 +1746,29 @@ export class TriggerChatTransport implements ChatTransport { currentSubscription = subscription; const sseStream = await subscription.subscribe(); const reader = sseStream.getReader(); + /** + * Register the caught-up close before the priming read below. A + * quiescent reconnect delivers only a tail-bearing ping, which + * yields no visible record, so waiting for the first visible read + * here would mean caught-up is never observed. + */ + if (options?.peekSettled) { + subscription + .caughtUp() + .then(() => { + if (!sawFirstChunk && !combinedSignal.aborted) { + internalAbort.abort(); + } + }) + .catch(() => {}); + } try { const first = await reader.read(); if (first.done) { reader.releaseLock(); return null; } - return { reader, primed: first.value }; + return { reader, primed: first.value, subscription }; } catch (readErr) { reader.releaseLock(); throw readErr; @@ -1798,7 +1816,6 @@ export class TriggerChatTransport implements ChatTransport { lastEventId: state.lastEventId, messageId: this.lastTurnSends.get(chatId)?.messageId, }); - let sawFirstChunk = false; while (true) { let value: { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index cd2d2d93fba..4236f83a2c4 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -436,8 +436,8 @@ importers: specifier: 2.17.5 version: 2.17.5(typescript@6.0.3) '@s2-dev/streamstore': - specifier: ^0.22.10 - version: 0.22.10(supports-color@10.0.0) + specifier: ^0.25.0 + version: 0.25.0(supports-color@10.0.0) '@sentry/remix': specifier: 9.46.0 version: 9.46.0(patch_hash=146126b032581925294aaed63ab53ce3f5e0356a755f1763d7a9a76b9846943b)(@remix-run/node@2.17.5(typescript@6.0.3))(@remix-run/react@2.17.5(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(typescript@6.0.3))(@remix-run/server-runtime@2.17.5(typescript@6.0.3))(encoding@0.1.13)(react@18.3.1) @@ -778,6 +778,9 @@ importers: '@internal/testcontainers': specifier: workspace:* version: link:../../internal-packages/testcontainers + '@playwright/test': + specifier: ^1.36.2 + version: 1.37.0 '@remix-run/dev': specifier: 2.17.5 version: 2.17.5(@remix-run/react@2.17.5(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(typescript@6.0.3))(@remix-run/serve@2.17.5(typescript@6.0.3))(@types/node@24.13.3)(bufferutil@4.0.9)(jiti@2.7.0)(lightningcss@1.32.0)(terser@5.46.1)(tsx@4.20.6)(typescript@6.0.3)(vite@6.4.2(@types/node@24.13.3)(jiti@2.7.0)(lightningcss@1.32.0)(terser@5.46.1)(tsx@4.20.6)(yaml@2.9.0))(yaml@2.9.0) @@ -1490,8 +1493,8 @@ importers: specifier: 1.41.1 version: 1.41.1 '@s2-dev/streamstore': - specifier: ^0.22.10 - version: 0.22.10(supports-color@10.0.0) + specifier: ^0.25.0 + version: 0.25.0(supports-color@10.0.0) '@trigger.dev/build': specifier: workspace:4.5.7 version: link:../build @@ -1761,8 +1764,8 @@ importers: specifier: 1.41.1 version: 1.41.1 '@s2-dev/streamstore': - specifier: 0.22.10 - version: 0.22.10(supports-color@10.0.0) + specifier: 0.25.0 + version: 0.25.0(supports-color@10.0.0) dequal: specifier: ^2.0.3 version: 2.0.3 @@ -6755,8 +6758,8 @@ packages: cpu: [x64] os: [win32] - '@s2-dev/streamstore@0.22.10': - resolution: {integrity: sha512-dtm+oFHVE8szINwOUoNQdx9xpGSJOrcAEvsxspPFvomjYKGnmhIRmU4OX8o6kxcPoiK76S1tPeU0smjZdmOngA==} + '@s2-dev/streamstore@0.25.0': + resolution: {integrity: sha512-oB8OJObT/s2Q68Tr01NLfDnQPFJp/Kn2pP3seWybgP7VrjQ+j1efXReOsyVCDsUBVMxkyPlzt75o/+GuL5Q3vg==} '@sec-ant/readable-stream@0.4.1': resolution: {integrity: sha512-831qok9r2t8AlxLko40y2ebgSDhenenCatLVeW/uBtnHPyhHOvG0C7TvfgecV+wHzIm5KUICgzmVpWS+IMEAeg==} @@ -21177,7 +21180,7 @@ snapshots: '@rollup/rollup-win32-x64-msvc@4.60.1': optional: true - '@s2-dev/streamstore@0.22.10(supports-color@10.0.0)': + '@s2-dev/streamstore@0.25.0(supports-color@10.0.0)': dependencies: '@protobuf-ts/runtime': 2.11.1 debug: 4.4.3(supports-color@10.0.0)