diff --git a/.changeset/eighty-donkeys-shake.md b/.changeset/eighty-donkeys-shake.md new file mode 100644 index 00000000000..bb6a1035653 --- /dev/null +++ b/.changeset/eighty-donkeys-shake.md @@ -0,0 +1,6 @@ +--- +"@trigger.dev/core": patch +"@trigger.dev/react-hooks": patch +--- + +Realtime stream subscriptions can now refresh an expired access token and reconnect, via a new optional `refreshAccessToken` option on the client configuration and the React hooks. diff --git a/.changeset/realtime-streams-from-latest.md b/.changeset/realtime-streams-from-latest.md new file mode 100644 index 00000000000..103403705e7 --- /dev/null +++ b/.changeset/realtime-streams-from-latest.md @@ -0,0 +1,19 @@ +--- +"@trigger.dev/react-hooks": patch +"@trigger.dev/core": patch +"@trigger.dev/sdk": patch +--- + +Subscribe to a realtime stream from its latest record instead of replaying the whole history. Pass `from: "latest"` to `useRealtimeStream`, `streams.read()`, or `fetchStream` to start at the current tail (the latest record, then live updates) instead of replaying (a live "last value" view), and `maxParts` to keep the accumulated `parts` array bounded. A reconnect or remount resumes from the last record it saw, so no records are missed and none are replayed. `from: "latest"` needs a server that supports it; older servers safely fall back to a full replay. + +`useRealtimeStream` also gains a `lastEventId` option and returns the `lastEventId` of the last part seen, so you can persist the cursor (for example across a page reload) and resume exactly where you left off. An `onParts` callback delivers each throttled batch of parts with their event ids. + +```tsx +const { parts, lastEventId } = useRealtimeStream(runId, "frames", { + from: "latest", // skip history, start at the current tail + maxParts: 1, // keep only the most recent frame + lastEventId: savedCursor, // resume from a persisted cursor + onParts: (batch) => save(batch.at(-1)?.id), // track the cursor + accessToken, +}); +``` diff --git a/.changeset/spotty-pillows-visit.md b/.changeset/spotty-pillows-visit.md new file mode 100644 index 00000000000..6e3b27bbe52 --- /dev/null +++ b/.changeset/spotty-pillows-visit.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/react-hooks": patch +--- + +Added a `useSessionStream` React hook for reading a session's output or input channel in realtime. It accumulates records with automatic resume from the last record you received, and supports `from: "latest"` (start at the current tail, only new records after you connect), `maxRecords` (keep a bounded number of records in memory), a `lastEventId` resume cursor, and an `onRecords` callback that delivers each throttled batch of records with their event ids. diff --git a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.ts b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.ts index b4ff0ba9500..437000faa8b 100644 --- a/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.ts +++ b/apps/webapp/app/routes/realtime.v1.sessions.$session.$io.ts @@ -1,4 +1,5 @@ import { json } from "@remix-run/server-runtime"; +import { STREAM_START_HEADER } from "@trigger.dev/core/v3"; import { z } from "zod"; import { getRequestAbortSignal } from "~/services/httpAsyncStorage.server"; import { S2RealtimeStreams } from "~/services/realtime/s2realtimeStreams.server"; @@ -185,12 +186,15 @@ const loader = createLoaderApiRoute( // turn's first chunk and the SSE closes before records land. const peekSettled = request.headers.get("X-Peek-Settled") === "1"; + const startFrom = + request.headers.get(STREAM_START_HEADER)?.toLowerCase() === "latest" ? "latest" : undefined; + return realtimeStream.streamResponseFromSessionStream( request, resource.addressingKey, params.io, getRequestAbortSignal(), - { lastEventId, timeoutInSeconds, peekSettled } + { lastEventId, timeoutInSeconds, peekSettled, startFrom } ); } ); diff --git a/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts b/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts index ab941fd05ad..3eb061c83bc 100644 --- a/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts +++ b/apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts @@ -1,3 +1,4 @@ +import { STREAM_START_HEADER } from "@trigger.dev/core/v3"; import { z } from "zod"; import { $replica } from "~/db.server"; import { getRequestAbortSignal } from "~/services/httpAsyncStorage.server"; @@ -59,6 +60,9 @@ export const loader = createLoaderApiRoute( // Get Last-Event-ID header for resuming from a specific position const lastEventId = request.headers.get("Last-Event-ID") || undefined; + const startFrom = + request.headers.get(STREAM_START_HEADER)?.toLowerCase() === "latest" ? "latest" : undefined; + const timeoutInSecondsRaw = request.headers.get("Timeout-Seconds") ?? undefined; const timeoutInSeconds = timeoutInSecondsRaw ? parseInt(timeoutInSecondsRaw) : undefined; @@ -88,6 +92,7 @@ export const loader = createLoaderApiRoute( { lastEventId, timeoutInSeconds, + startFrom, } ); } diff --git a/apps/webapp/app/services/realtime/redisRealtimeStreams.server.ts b/apps/webapp/app/services/realtime/redisRealtimeStreams.server.ts index 6fbb26e5c9c..f3ddc65b542 100644 --- a/apps/webapp/app/services/realtime/redisRealtimeStreams.server.ts +++ b/apps/webapp/app/services/realtime/redisRealtimeStreams.server.ts @@ -70,8 +70,7 @@ export class RedisRealtimeStreams implements StreamIngestor, StreamResponder { const stream = new ReadableStream({ start: async (controller) => { - // Start from lastEventId if provided, otherwise from beginning - let lastId = options?.lastEventId ?? "0"; + let lastId = options?.lastEventId ?? (options?.startFrom === "latest" ? "$" : "0"); let retryCount = 0; const maxRetries = 3; let lastDataTime = Date.now(); diff --git a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts index e3485a63ebd..174040b2053 100644 --- a/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts +++ b/apps/webapp/app/services/realtime/s2realtimeStreams.server.ts @@ -527,12 +527,19 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { ): Promise { const startSeq = this.parseLastEventId(options?.lastEventId); - this.logger.info(`S2 streaming records from stream`, { stream: s2Stream, startSeq }); + const tailFromLatest = startSeq == null && options?.startFrom === "latest"; + + this.logger.info(`S2 streaming records from stream`, { + stream: s2Stream, + startSeq, + tailFromLatest, + }); // Request SSE stream from S2 and return it directly const s2Response = await this.s2StreamRecords(s2Stream, { - seq_num: startSeq ?? 0, - clamp: true, + ...(tailFromLatest + ? { tail_offset: 1, clamp: true } + : { seq_num: startSeq ?? 0, clamp: true }), wait: options?.timeoutInSeconds ?? this.s2WaitSeconds, // S2 will keep the connection open and stream new records signal, // Pass abort signal so S2 connection is cleaned up when client disconnects }); @@ -672,6 +679,7 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { stream: string, opts: { seq_num?: number; + tail_offset?: number; clamp?: boolean; wait?: number; signal?: AbortSignal; @@ -680,6 +688,7 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor { // GET /v1/streams/{stream}/records with Accept: text/event-stream for SSE streaming const qs = new URLSearchParams(); if (opts.seq_num != null) qs.set("seq_num", String(opts.seq_num)); + if (opts.tail_offset != null) qs.set("tail_offset", String(opts.tail_offset)); if (opts.clamp != null) qs.set("clamp", String(opts.clamp)); if (opts.wait != null) qs.set("wait", String(opts.wait)); diff --git a/apps/webapp/app/services/realtime/types.ts b/apps/webapp/app/services/realtime/types.ts index 3121f0e18df..04eeb74446e 100644 --- a/apps/webapp/app/services/realtime/types.ts +++ b/apps/webapp/app/services/realtime/types.ts @@ -36,6 +36,13 @@ export interface StreamIngestor { export type StreamResponseOptions = { timeoutInSeconds?: number; lastEventId?: string; + /** + * Where a fresh subscription (no `lastEventId`) starts reading. `"latest"` + * starts at the current tail so the subscriber sees only records appended + * after it connects; `"beginning"` (the default when unset) replays history. + * Ignored when `lastEventId` is set. + */ + startFrom?: "beginning" | "latest"; /** * Session-stream-only. When `true`, the responder MAY peek the tail * of `.out` and short-circuit to `wait=0` + `X-Session-Settled: true` diff --git a/apps/webapp/test/redisRealtimeStreams.test.ts b/apps/webapp/test/redisRealtimeStreams.test.ts index 306e6303bb8..e767aa0dd6f 100644 --- a/apps/webapp/test/redisRealtimeStreams.test.ts +++ b/apps/webapp/test/redisRealtimeStreams.test.ts @@ -1521,4 +1521,146 @@ describe("RedisRealtimeStreams", () => { await redis.quit(); } ); + + redisTest( + "startFrom 'latest' skips the backlog and delivers only new records", + { timeout: 30_000 }, + async ({ redisOptions }) => { + const redis = new Redis(redisOptions); + const redisRealtimeStreams = new RedisRealtimeStreams({ redis: redisOptions }); + + const runId = "run_latest_test"; + const streamId = "latest-stream"; + const encoder = new TextEncoder(); + + const ingest = async (line: string) => { + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(line + "\n")); + controller.close(); + }, + }); + await redisRealtimeStreams.ingestData(stream, runId, streamId, "default"); + }; + + await ingest("old-0"); + await ingest("old-1"); + + const abortController = new AbortController(); + const response = await redisRealtimeStreams.streamResponse( + new Request("http://localhost/test"), + runId, + streamId, + abortController.signal, + { startFrom: "latest", timeoutInSeconds: 10 } + ); + + const reader = response.body!.getReader(); + const decoder = new TextDecoder(); + const receivedData: string[] = []; + + const readLoop = (async () => { + let done = false; + while (!done && receivedData.length < 1) { + const { value, done: streamDone } = await reader.read().catch(() => ({ + value: undefined, + done: true, + })); + done = streamDone; + if (value) { + const events = decoder + .decode(value) + .split("\n\n") + .filter((event) => event.trim()); + for (const event of events) { + for (const l of event.split("\n")) { + if (l.startsWith("data: ")) { + const data = l.substring(6).trim(); + if (data) receivedData.push(data); + } + } + } + } + } + })(); + + await new Promise((resolve) => setTimeout(resolve, 500)); + await ingest("new-0"); + + await readLoop; + abortController.abort(); + reader.releaseLock(); + + expect(receivedData).toContain("new-0"); + expect(receivedData).not.toContain("old-0"); + expect(receivedData).not.toContain("old-1"); + + await redis.del(`stream:${runId}:${streamId}`); + await redis.quit(); + } + ); + + redisTest( + "default start replays the backlog from the beginning", + { timeout: 30_000 }, + async ({ redisOptions }) => { + const redis = new Redis(redisOptions); + const redisRealtimeStreams = new RedisRealtimeStreams({ redis: redisOptions }); + + const runId = "run_beginning_test"; + const streamId = "beginning-stream"; + const encoder = new TextEncoder(); + + const ingestStream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode("old-0\n")); + controller.enqueue(encoder.encode("old-1\n")); + controller.close(); + }, + }); + await redisRealtimeStreams.ingestData(ingestStream, runId, streamId, "default"); + + const abortController = new AbortController(); + const response = await redisRealtimeStreams.streamResponse( + new Request("http://localhost/test"), + runId, + streamId, + abortController.signal, + { timeoutInSeconds: 10 } + ); + + const reader = response.body!.getReader(); + const decoder = new TextDecoder(); + const receivedData: string[] = []; + + let done = false; + while (!done && receivedData.length < 2) { + const { value, done: streamDone } = await reader.read(); + done = streamDone; + if (value) { + const events = decoder + .decode(value) + .split("\n\n") + .filter((event) => event.trim()); + for (const event of events) { + for (const l of event.split("\n")) { + if (l.startsWith("data: ")) { + const data = l.substring(6).trim(); + if (data) receivedData.push(data); + } + } + } + } + } + + abortController.abort(); + reader.releaseLock(); + + expect(receivedData).toContain("old-0"); + expect(receivedData).toContain("old-1"); + + await redis.del(`stream:${runId}:${streamId}`); + await redis.quit(); + } + ); }); diff --git a/docs/docs.json b/docs/docs.json index 98c3dc46322..669bc1ebc9d 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -217,6 +217,7 @@ "realtime/react-hooks/triggering", "realtime/react-hooks/subscribe", "realtime/react-hooks/streams", + "realtime/react-hooks/session-stream", "realtime/react-hooks/swr", "realtime/react-hooks/use-wait-token" ] diff --git a/docs/realtime/auth.mdx b/docs/realtime/auth.mdx index 783e8b1c38a..b7cd067a8b4 100644 --- a/docs/realtime/auth.mdx +++ b/docs/realtime/auth.mdx @@ -135,6 +135,24 @@ When using non-root API keys (recommended), the expiration cannot be more than 3 The format used for a time span is the same as the [jose package](https://github.com/panva/jose), which is a number followed by a unit. Valid units are: "sec", "secs", "second", "seconds", "s", "minute", "minutes", "min", "mins", "m", "hour", "hours", "hr", "hrs", "h", "day", "days", "d", "week", "weeks", "w", "year", "years", "yr", "yrs", and "y". It is not possible to specify months. 365.25 days is used as an alias for a year. If the string is suffixed with "ago", or prefixed with a "-", the resulting time span gets subtracted from the current unix timestamp. A "from now" suffix can also be used for readability when adding to the current unix timestamp. +### Refreshing an expired token + +A realtime stream subscription can outlive its token. Pass a `refreshAccessToken` callback and a subscription rejected with a 401/403 re-mints once and reconnects, instead of failing. With no refresher, auth errors stay terminal. Mint the fresh token from your backend, where your secret key lives: + +```tsx +import { useRealtimeStream } from "@trigger.dev/react-hooks"; + +const { parts } = useRealtimeStream<{ url: string }>(runId, "frames", { + accessToken, + refreshAccessToken: async () => { + const res = await fetch("/api/realtime-token"); // your backend calls auth.createPublicToken + return (await res.json()).token; + }, +}); +``` + +`refreshAccessToken` is available on every realtime hook, and on `useApiClient` and `TriggerAuthContext` so hooks under a provider share one refresher. + ### Auto-generated tokens When you [trigger tasks](/triggering) from your backend, the `handle` received includes a `publicAccessToken` field. This token can be used to authenticate real-time requests in your frontend application. diff --git a/docs/realtime/react-hooks/session-stream.mdx b/docs/realtime/react-hooks/session-stream.mdx new file mode 100644 index 00000000000..1d08f5f1300 --- /dev/null +++ b/docs/realtime/react-hooks/session-stream.mdx @@ -0,0 +1,109 @@ +--- +title: "Read a session channel in React" +sidebarTitle: "Session streams" +description: "Subscribe to a session's output or input channel from React with useSessionStream: accumulate records, resume from a cursor, and read only the latest." +--- + +**`useSessionStream` subscribes to one channel of a [session](/ai-chat/sessions) and updates a `records` array as new records arrive.** It reads the `out` channel by default (the agent's output) or `in` (the input channel). It is read-only; `useSession` is reserved for two-way (read and write) communication. + + + Requires a Public Access Token with the `read:sessions:{id}` scope. See [Realtime + auth](/realtime/auth) for generating one. + + +## Basic usage + +Pass the session id (or external id) and an `accessToken`. The hook returns the `records` received so far, the last control record, the cursor of the last record seen, and any error: + +```tsx +"use client"; + +import { useSessionStream } from "@trigger.dev/react-hooks"; + +export function SessionViewer({ + sessionId, + accessToken, +}: { + sessionId: string; + accessToken: string; +}) { + const { records, error } = useSessionStream(sessionId, { accessToken }); + + if (error) return
Error: {error.message}
; + + return
{records.join("")}
; +} +``` + +## Options + +```tsx +const { records, lastEventId, lastControl, error, stop } = useSessionStream(sessionId, { + accessToken: "pk_...", // Required: public access token with read:sessions:{id} + io: "out", // Optional: "out" (default) or "in" + from: "beginning", // Optional: "beginning" (default) or "latest" + maxRecords: 100, // Optional: keep only the most recent N records (default: unbounded) + lastEventId: undefined, // Optional: resume cursor + timeoutInSeconds: 60, // Optional: close after this long with no new data (default: 60) + throttleInMs: 16, // Optional: throttle record updates (default: 16ms) + onRecords: (batch) => {}, // Optional: callback per throttled batch, each with its event id + onControl: (event) => {}, // Optional: callback for control records (e.g. turn-complete) +}); +``` + +The return value: + +- **`records`**: every data record received so far, in arrival order. Control records are delivered to `onControl` instead. +- **`lastEventId`**: the cursor of the last record seen. Persist it and pass it back as the `lastEventId` option to resume. +- **`lastControl`**: the last control record (for example `turn-complete`). +- **`stop`**: abort the subscription, keeping the records received so far. + +## Start from the latest record + +By default the hook replays the channel history, then live-tails. Pass `from: "latest"` to start at the current tail (the latest record, then live updates) instead of replaying, and `maxRecords` to bound memory: + +```tsx +const { records } = useSessionStream<{ url: string }>(sessionId, { + accessToken, + io: "out", + from: "latest", // start at the latest record, then live updates + maxRecords: 1, // keep just the most recent record +}); +``` + + + `from: "latest"` requires a server that supports it. Against an older server a client that passes + it degrades safely to a full replay. + + +## Resume from a cursor + +The hook resumes automatically across a component remount. A full page reload clears in-memory state, so to resume there, persist the returned `lastEventId` and pass it back on the next load. The channel then continues after that record with no replay and no gap: + +```tsx +const cursorKey = `session-cursor:${sessionId}:out`; // scope the key to this session and channel +const saved = localStorage.getItem(cursorKey) ?? undefined; + +const { records, lastEventId } = useSessionStream(sessionId, { + accessToken, + lastEventId: saved, + onRecords: (batch) => localStorage.setItem(cursorKey, batch.at(-1)!.id), +}); +``` + +## React to control records + +Control records (such as `turn-complete`) never enter `records`. Handle them with `onControl`, or read the latest from `lastControl`: + +```tsx +const { records, lastControl } = useSessionStream(sessionId, { + accessToken, + onControl: (event) => { + if (event.subtype === "turn-complete") { + console.log("The turn is complete"); + } + }, +}); +``` + +For an expiring token on a long-lived subscription, pass `refreshAccessToken` (see [Realtime auth](/realtime/auth)). To read a session channel outside React, use [`session.out.read()`](/ai-chat/sessions). diff --git a/docs/realtime/react-hooks/streams.mdx b/docs/realtime/react-hooks/streams.mdx index c3c7031d44a..c55f610280d 100644 --- a/docs/realtime/react-hooks/streams.mdx +++ b/docs/realtime/react-hooks/streams.mdx @@ -130,16 +130,84 @@ export function AIStreamViewer({ The `useRealtimeStream` hook accepts the following options: ```tsx -const { parts, error } = useRealtimeStream(streamOrRunId, streamKeyOrOptions, { +const { parts, lastEventId, error } = useRealtimeStream(streamOrRunId, streamKeyOrOptions, { accessToken: "pk_...", // Required: Public access token baseURL: "https://api.trigger.dev", // Optional: Custom API URL timeoutInSeconds: 60, // Optional: Timeout (default: 60) - startIndex: 0, // Optional: Start from specific chunk + from: "beginning", // Optional: "beginning" (default) or "latest" + maxParts: 100, // Optional: keep only the most recent N parts (default: unbounded) + lastEventId: undefined, // Optional: resume cursor (takes precedence over startIndex) + startIndex: 0, // Optional: start from a specific chunk index throttleInMs: 16, // Optional: Throttle updates (default: 16ms) - onData: (chunk) => {}, // Optional: Callback for each chunk + onData: (chunk) => {}, // Optional: callback for each chunk + onParts: (batch) => {}, // Optional: callback per throttled batch, each with its event id + refreshAccessToken: async () => "pk_...", // Optional: mint a fresh token on expiry }); ``` +The hook returns `lastEventId`, the cursor of the last part it received. Persist it and pass it back as the `lastEventId` option to resume later. + +### Live view: start from the latest record + +By default a subscriber replays the whole stream history, then live-tails. Pass `from: "latest"` to start at the current tail (the latest record, then live updates) instead of replaying, and `maxParts` to keep memory bounded. Together they give a last-value view: + +```tsx +"use client"; + +import { useRealtimeStream } from "@trigger.dev/react-hooks"; + +export function LatestFrame({ runId, accessToken }: { runId: string; accessToken: string }) { + const { parts } = useRealtimeStream<{ url: string }>(runId, "frames", { + accessToken, + from: "latest", // start at the latest frame, then live updates + maxParts: 1, // keep just the most recent frame + }); + + const frame = parts.at(-1); + return frame ? latest frame : null; +} +``` + + + `from: "latest"` requires a server that supports it. Against an older server a client that passes + it degrades safely to a full replay. + + +### Resume across a page reload + +The hook resumes automatically across a component remount. A full page reload clears in-memory +state, so to resume there, persist the returned `lastEventId` and pass it back on the next load. The +subscription then continues after that record with no replay and no gap: + +```tsx +const cursorKey = `frames-cursor:${runId}`; // scope the key to this stream +const saved = localStorage.getItem(cursorKey) ?? undefined; + +const { parts, lastEventId } = useRealtimeStream<{ url: string }>(runId, "frames", { + accessToken, + lastEventId: saved, // resume where the previous session left off + onParts: (batch) => localStorage.setItem(cursorKey, batch.at(-1)!.id), +}); +``` + +### Refresh an expired access token + +Public access tokens are short-lived. For a long-running subscription, pass `refreshAccessToken` to +mint a fresh token when the server rejects the connection with a 401/403. The subscription re-mints +once and reconnects; with no refresher, auth errors stay terminal: + +```tsx +const { parts } = useRealtimeStream<{ url: string }>(runId, "frames", { + accessToken, + refreshAccessToken: async () => { + const res = await fetch("/api/realtime-token"); // your backend mints a fresh public token + return (await res.json()).token; + }, +}); +``` + +`refreshAccessToken` is also available on [`useApiClient` and `TriggerAuthContext`](/realtime/auth), so every hook under a provider shares one refresher. + ### Using Default Stream You can omit the stream key to use the default stream: diff --git a/docs/tasks/streams.mdx b/docs/tasks/streams.mdx index 9c43ca2d4d1..611041092f2 100644 --- a/docs/tasks/streams.mdx +++ b/docs/tasks/streams.mdx @@ -144,9 +144,12 @@ With options: const stream = await aiStream.read(runId, { timeoutInSeconds: 60, // Stop if no data for 60 seconds startIndex: 10, // Start from the 10th chunk + from: "latest", // Or skip history and read only new records from now }); ``` +Pass `from: "latest"` to start at the current tail and receive only records appended after the read connects, instead of replaying from the beginning. + #### Appending to a Stream Use the defined stream's `append()` method to add a single chunk: diff --git a/packages/core/src/v3/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts index e6b4125222f..c270f86ea62 100644 --- a/packages/core/src/v3/apiClient/index.ts +++ b/packages/core/src/v3/apiClient/index.ts @@ -110,6 +110,7 @@ import { zodfetchOffsetLimitPage, } from "./core.js"; import { ApiConnectionError, ApiError, BatchNotSealedError } from "./errors.js"; +import { refreshAccessTokenOnce, type RefreshAccessTokenFn } from "./refreshAccessToken.js"; import { type AnyRealtimeRun, type AnyRunShape, @@ -123,6 +124,7 @@ import { SSEStreamSubscriptionFactory, runShapeStream, type SSEStreamPart, + STREAM_START_HEADER, } from "./runStream.js"; import type { CreateBulkActionOptions, @@ -190,11 +192,12 @@ export type ApiClientFutureFlags = { v2RealtimeStreams?: boolean; }; -export { SSEStreamSubscription, isRequestOptions }; +export { SSEStreamSubscription, STREAM_START_HEADER, isRequestOptions }; export type { AnyRealtimeRun, AnyRunShape, ApiRequestOptions, + ControlEvent, RealtimeRun, RunShape, RunStreamCallback, @@ -224,6 +227,7 @@ export class ApiClient { public readonly futureFlags: ApiClientFutureFlags; private readonly additionalHeaders?: Record; private readonly defaultRequestOptions: ZodFetchOptions; + private readonly refreshAccessToken?: RefreshAccessTokenFn; constructor( baseUrl: string, @@ -232,9 +236,11 @@ export class ApiClient { // x-trigger-branch header, and the server disambiguates by the token's env. previewBranch?: string, requestOptions: ApiRequestOptions = {}, - futureFlags: ApiClientFutureFlags = {} + futureFlags: ApiClientFutureFlags = {}, + refreshAccessToken?: RefreshAccessTokenFn ) { this.accessToken = accessToken; + this.refreshAccessToken = refreshAccessToken; this.baseUrl = baseUrl.replace(/\/$/, ""); this.previewBranch = previewBranch; const { additionalHeaders, ...restRequestOptions } = requestOptions; @@ -280,6 +286,32 @@ export class ApiClient { return this.#getHeaders(false); } + /** + * Header resolver handed to stream subscriptions so a connection rejected with + * a 401/403 can reconnect with a freshly minted token. `undefined` when no + * `refreshAccessToken` was configured, which keeps auth errors terminal. + */ + #resolveStreamHeaders(): (() => Promise>) | undefined { + const refreshAccessToken = this.refreshAccessToken; + if (!refreshAccessToken) return undefined; + + return async () => { + const accessToken = await refreshAccessTokenOnce(refreshAccessToken); + return this.#getHeaders(false, { Authorization: `Bearer ${accessToken}` }); + }; + } + + /** As {@link ApiClient.#resolveStreamHeaders}, for the leaner realtime header set. */ + #resolveRealtimeHeaders(): (() => Promise>) | undefined { + const refreshAccessToken = this.refreshAccessToken; + if (!refreshAccessToken) return undefined; + + return async () => { + const accessToken = await refreshAccessTokenOnce(refreshAccessToken); + return { ...this.#getRealtimeHeaders(), Authorization: `Bearer ${accessToken}` }; + }; + } + async getRunResult( runId: string, requestOptions?: ZodFetchOptions @@ -1449,6 +1481,12 @@ export class ApiClient { onComplete?: () => void; onError?: (error: Error) => void; lastEventId?: string; + /** + * Where a fresh subscription (no `lastEventId`) starts reading. `"latest"` + * starts at the current tail (only records after connect); `"beginning"` + * (default) replays history. + */ + from?: "beginning" | "latest"; onPart?: (part: SSEStreamPart) => void; /** * Fires when a `trigger-control` record arrives on the stream (e.g. @@ -1462,11 +1500,13 @@ export class ApiClient { const subscription = new SSEStreamSubscription(url, { headers: this.getHeaders(), + resolveHeaders: this.#resolveStreamHeaders(), signal: options?.signal, onComplete: options?.onComplete, onError: options?.onError, timeoutInSeconds: options?.timeoutInSeconds, lastEventId: options?.lastEventId, + from: options?.from, }); const stream = await subscription.subscribe(); @@ -1665,6 +1705,7 @@ export class ApiClient { closeOnComplete: typeof options?.closeOnComplete === "boolean" ? options.closeOnComplete : true, headers: this.#getRealtimeHeaders(), + resolveHeaders: this.#resolveRealtimeHeaders(), client: this, signal: options?.signal, onFetchError: options?.onFetchError, @@ -1688,6 +1729,7 @@ export class ApiClient { { closeOnComplete: false, headers: this.#getRealtimeHeaders(), + resolveHeaders: this.#resolveRealtimeHeaders(), client: this, signal: options?.signal, onFetchError: options?.onFetchError, @@ -1714,6 +1756,7 @@ export class ApiClient { { closeOnComplete: false, headers: this.#getRealtimeHeaders(), + resolveHeaders: this.#resolveRealtimeHeaders(), client: this, signal: options?.signal, onFetchError: options?.onFetchError, @@ -1778,6 +1821,12 @@ export class ApiClient { onComplete?: () => void; onError?: (error: Error) => void; lastEventId?: string; + /** + * Where a fresh subscription (no `lastEventId`) starts reading. `"latest"` + * starts at the current tail (only records after connect); `"beginning"` + * (default) replays history. + */ + from?: "beginning" | "latest"; /** Called for each SSE event with the full event metadata (id, timestamp). */ onPart?: (part: SSEStreamPart) => void; } @@ -1785,6 +1834,7 @@ export class ApiClient { const streamFactory = new SSEStreamSubscriptionFactory(options?.baseUrl ?? this.baseUrl, { headers: this.getHeaders(), signal: options?.signal, + resolveHeaders: this.#resolveStreamHeaders(), }); const subscription = streamFactory.createSubscription(runId, streamKey, { @@ -1792,6 +1842,7 @@ export class ApiClient { onError: options?.onError, timeoutInSeconds: options?.timeoutInSeconds, lastEventId: options?.lastEventId, + from: options?.from, }); const stream = await subscription.subscribe(); diff --git a/packages/core/src/v3/apiClient/refreshAccessToken.test.ts b/packages/core/src/v3/apiClient/refreshAccessToken.test.ts new file mode 100644 index 00000000000..7c28c25ff56 --- /dev/null +++ b/packages/core/src/v3/apiClient/refreshAccessToken.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it } from "vitest"; +import { refreshAccessTokenOnce } from "./refreshAccessToken.js"; + +describe("refreshAccessTokenOnce", () => { + it("shares one in-flight mint between concurrent callers", async () => { + let calls = 0; + let release: (token: string) => void = () => {}; + const refresh = () => { + calls++; + return new Promise((resolve) => { + release = resolve; + }); + }; + + const results = Promise.all([ + refreshAccessTokenOnce(refresh), + refreshAccessTokenOnce(refresh), + refreshAccessTokenOnce(refresh), + ]); + release("fresh"); + + expect(await results).toEqual(["fresh", "fresh", "fresh"]); + expect(calls).toBe(1); + }); + + it("does not share a mint between different refreshers", async () => { + const a = async () => "a"; + const b = async () => "b"; + + expect(await Promise.all([refreshAccessTokenOnce(a), refreshAccessTokenOnce(b)])).toEqual([ + "a", + "b", + ]); + }); + + it("mints again once the previous call has settled", async () => { + let calls = 0; + const refresh = async () => `token-${++calls}`; + + expect(await refreshAccessTokenOnce(refresh)).toBe("token-1"); + expect(await refreshAccessTokenOnce(refresh)).toBe("token-2"); + }); + + it("rejects every concurrent caller and does not poison later calls", async () => { + let calls = 0; + const refresh = async () => { + calls++; + if (calls === 1) throw new Error("mint failed"); + return "recovered"; + }; + + const first = refreshAccessTokenOnce(refresh); + const second = refreshAccessTokenOnce(refresh); + + await expect(first).rejects.toThrow("mint failed"); + await expect(second).rejects.toThrow("mint failed"); + expect(calls).toBe(1); + + expect(await refreshAccessTokenOnce(refresh)).toBe("recovered"); + expect(calls).toBe(2); + }); +}); diff --git a/packages/core/src/v3/apiClient/refreshAccessToken.ts b/packages/core/src/v3/apiClient/refreshAccessToken.ts new file mode 100644 index 00000000000..f48b1a6ea84 --- /dev/null +++ b/packages/core/src/v3/apiClient/refreshAccessToken.ts @@ -0,0 +1,22 @@ +export type RefreshAccessTokenFn = () => Promise; + +const pendingRefreshes = new WeakMap>(); + +/** + * Call `refreshAccessToken`, deduping concurrent calls to the same function. + * Several subscriptions (or several React hooks sharing one refresher) can hit + * an expired token at once; they reuse the in-flight mint instead of firing one + * per caller. Keyed on the refresher itself so callers only share a mint when + * they share a token owner. + */ +export function refreshAccessTokenOnce(refreshAccessToken: RefreshAccessTokenFn): Promise { + const pending = pendingRefreshes.get(refreshAccessToken); + if (pending) return pending; + + const promise = refreshAccessToken().finally(() => { + pendingRefreshes.delete(refreshAccessToken); + }); + pendingRefreshes.set(refreshAccessToken, promise); + + return promise; +} diff --git a/packages/core/src/v3/apiClient/runStream.test.ts b/packages/core/src/v3/apiClient/runStream.test.ts index ee3f3df22a6..f91d377a501 100644 --- a/packages/core/src/v3/apiClient/runStream.test.ts +++ b/packages/core/src/v3/apiClient/runStream.test.ts @@ -1,5 +1,5 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { SSEStreamSubscription } from "./runStream.js"; +import { SSEStreamSubscription, STREAM_START_HEADER } from "./runStream.js"; vi.setConfig({ testTimeout: 10_000 }); @@ -43,6 +43,33 @@ describe("SSEStreamSubscription retry behavior", () => { }); } + /** An accepted connection that dies before delivering a single record. */ + function makeDroppedResponse() { + const body = new ReadableStream({ + start(controller) { + controller.error(new Error("connection dropped")); + }, + }); + return new Response(body, { + status: 200, + headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v1" }, + }); + } + + /** One delivered record, then the connection dies. */ + function makeChunkThenDropResponse() { + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(`id: 1\ndata: {"hello":1}\n\n`)); + setTimeout(() => controller.error(new Error("connection dropped")), 20); + }, + }); + 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 }>) { @@ -427,6 +454,195 @@ describe("SSEStreamSubscription retry behavior", () => { expect(result.error).toBeDefined(); }); + it("fails the stream on a 401 when no resolveHeaders is supplied", async () => { + let attempts = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + attempts++; + return new Response("unauthorized", { status: 401 }); + }); + + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + }); + + const result = await sub.subscribe().then(drain); + expect(attempts).toBe(1); + expect(result.error).toBeDefined(); + }); + + it("retries a 401 once with the headers from resolveHeaders", async () => { + const seenTokens: Array = []; + globalThis.fetch = vi.fn().mockImplementation(async (_url: string, init: RequestInit) => { + const token = new Headers(init.headers).get("Authorization"); + seenTokens.push(token); + if (token !== "Bearer fresh") return new Response("unauthorized", { status: 401 }); + return makeSSEResponse(); + }); + + let refreshes = 0; + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + resolveHeaders: async () => { + refreshes++; + return { Authorization: "Bearer fresh" }; + }, + }); + + const result = await sub.subscribe().then(drain); + expect(seenTokens).toEqual(["Bearer expired", "Bearer fresh"]); + expect(refreshes).toBe(1); + expect(result.error).toBeUndefined(); + expect(result.chunks).toHaveLength(1); + }); + + it("fails the stream when the refreshed headers are rejected too", async () => { + let attempts = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + attempts++; + return new Response("unauthorized", { status: 401 }); + }); + + let refreshes = 0; + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + resolveHeaders: async () => { + refreshes++; + return { Authorization: "Bearer also-expired" }; + }, + }); + + const result = await sub.subscribe().then(drain); + expect(attempts).toBe(2); + expect(refreshes).toBe(1); + expect(result.error).toBeDefined(); + }); + + it("retries a 403 once with the headers from resolveHeaders", async () => { + const seenTokens: Array = []; + globalThis.fetch = vi.fn().mockImplementation(async (_url: string, init: RequestInit) => { + const token = new Headers(init.headers).get("Authorization"); + seenTokens.push(token); + if (token !== "Bearer fresh") return new Response("forbidden", { status: 403 }); + return makeSSEResponse(); + }); + + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + resolveHeaders: async () => ({ Authorization: "Bearer fresh" }), + }); + + const result = await sub.subscribe().then(drain); + expect(seenTokens).toEqual(["Bearer expired", "Bearer fresh"]); + expect(result.error).toBeUndefined(); + expect(result.chunks).toHaveLength(1); + }); + + it("does not report a 401 that the refresh recovered from", async () => { + let attempts = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + attempts++; + if (attempts === 1) return new Response("unauthorized", { status: 401 }); + return makeSSEResponse(); + }); + + const errors: Error[] = []; + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + onError: (e) => errors.push(e), + resolveHeaders: async () => ({ Authorization: "Bearer fresh" }), + }); + + const result = await sub.subscribe().then(drain); + expect(errors).toHaveLength(0); + expect(result.error).toBeUndefined(); + }); + + it("terminates on a 401 when the refresher itself throws", async () => { + let attempts = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + attempts++; + return new Response("unauthorized", { status: 401 }); + }); + + const errors: Error[] = []; + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + onError: (e) => errors.push(e), + resolveHeaders: async () => { + throw new Error("mint failed"); + }, + }); + + const result = await sub.subscribe().then(drain); + expect(attempts).toBe(1); + expect(errors).toHaveLength(1); + expect(result.error).toBeDefined(); + }); + + it("does not re-mint for a connection that is accepted but delivers nothing", async () => { + let attempts = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + attempts++; + if (attempts === 2) return makeDroppedResponse(); + return new Response("unauthorized", { status: 401 }); + }); + + let refreshes = 0; + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + resolveHeaders: async () => { + refreshes++; + return { Authorization: `Bearer fresh-${refreshes}` }; + }, + }); + + const result = await sub.subscribe().then(drain); + expect(refreshes).toBe(1); + expect(attempts).toBe(3); + expect(result.error).toBeDefined(); + }); + + it("allows another refresh once a connection has delivered a record", async () => { + let attempts = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + attempts++; + if (attempts === 2) return makeChunkThenDropResponse(); + if (attempts === 4) return makeSSEResponse(); + return new Response("unauthorized", { status: 401 }); + }); + + let refreshes = 0; + const sub = new SSEStreamSubscription("http://example.test/sse", { + headers: { Authorization: "Bearer expired" }, + retryDelayMs: 1, + maxRetryDelayMs: 5, + resolveHeaders: async () => { + refreshes++; + return { Authorization: `Bearer fresh-${refreshes}` }; + }, + }); + + const result = await sub.subscribe().then(drain); + expect(refreshes).toBe(2); + expect(attempts).toBe(4); + expect(result.error).toBeUndefined(); + expect(result.chunks).toHaveLength(2); + }); + it("retries on 503 (caller-tunable nonRetryableStatuses)", async () => { let attempts = 0; globalThis.fetch = vi.fn().mockImplementation(async () => { @@ -642,3 +858,101 @@ describe("SSEStreamSubscription v2 batch parsing — record kinds", () => { expect((parts[1]!.chunk as any).delta).toBe("x"); }); }); + +describe("SSEStreamSubscription start position (from)", () => { + const originalFetch = globalThis.fetch; + + afterEach(() => { + globalThis.fetch = originalFetch; + vi.restoreAllMocks(); + }); + + async function drainReader(reader: ReadableStreamDefaultReader) { + let next = await reader.read(); + while (!next.done) { + next = await reader.read(); + } + } + + function makeClosedSSEResponse(id: string) { + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(`id: ${id}\ndata: {"hello":1}\n\n`)); + controller.close(); + }, + }); + return new Response(body, { + status: 200, + headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v1" }, + }); + } + + it('from: "latest" sends the start header and no Last-Event-ID on first connect', async () => { + const seenHeaders: Array> = []; + globalThis.fetch = vi.fn().mockImplementation(async (_url: string, init?: RequestInit) => { + seenHeaders.push((init?.headers as Record) ?? {}); + return makeClosedSSEResponse("42"); + }); + + const sub = new SSEStreamSubscription("http://example.test/sse", { from: "latest" }); + await drainReader((await sub.subscribe()).getReader()); + + expect(seenHeaders[0]![STREAM_START_HEADER]).toBe("latest"); + expect(seenHeaders[0]!["Last-Event-ID"]).toBeUndefined(); + }); + + it('from: "latest" drops the start header and resumes with Last-Event-ID after a record', async () => { + let attempts = 0; + const seenHeaders: Array> = []; + globalThis.fetch = vi.fn().mockImplementation(async (_url: string, init?: RequestInit) => { + attempts++; + seenHeaders.push((init?.headers as Record) ?? {}); + if (attempts === 1) { + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(`id: 7\ndata: {"first":true}\n\n`)); + init?.signal?.addEventListener("abort", () => controller.error(new Error("aborted"))); + }, + }); + return new Response(body, { + status: 200, + headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v1" }, + }); + } + return makeClosedSSEResponse("8"); + }); + + const sub = new SSEStreamSubscription("http://example.test/sse", { + from: "latest", + retryDelayMs: 1, + maxRetryDelayMs: 5, + fetchTimeoutMs: 60_000, + }); + + const reader = (await sub.subscribe()).getReader(); + const first = await reader.read(); + expect(first.done).toBe(false); + + sub.forceReconnect(); + await drainReader(reader); + + expect(attempts).toBe(2); + expect(seenHeaders[0]![STREAM_START_HEADER]).toBe("latest"); + expect(seenHeaders[0]!["Last-Event-ID"]).toBeUndefined(); + expect(seenHeaders[1]![STREAM_START_HEADER]).toBeUndefined(); + expect(seenHeaders[1]!["Last-Event-ID"]).toBe("7"); + }); + + it("default (no from) never sends the start header", async () => { + const seenHeaders: Array> = []; + globalThis.fetch = vi.fn().mockImplementation(async (_url: string, init?: RequestInit) => { + seenHeaders.push((init?.headers as Record) ?? {}); + return makeClosedSSEResponse("1"); + }); + + const sub = new SSEStreamSubscription("http://example.test/sse", {}); + await drainReader((await sub.subscribe()).getReader()); + + expect(seenHeaders[0]![STREAM_START_HEADER]).toBeUndefined(); + }); +}); diff --git a/packages/core/src/v3/apiClient/runStream.ts b/packages/core/src/v3/apiClient/runStream.ts index d39ac7fcaa6..f0ec1267f97 100644 --- a/packages/core/src/v3/apiClient/runStream.ts +++ b/packages/core/src/v3/apiClient/runStream.ts @@ -15,6 +15,14 @@ import { ApiError, isTriggerRealtimeAuthError } from "./errors.js"; import type { ApiClient } from "./index.js"; import { zodShapeStream } from "./stream.js"; +/** + * Request header carrying the start position for a fresh realtime-stream + * subscription. Value `"latest"` asks the server to start at the current tail + * (only records appended after connect). Only sent when there is no + * `Last-Event-ID`. Read by the realtime streams route on the server. + */ +export const STREAM_START_HEADER = "X-Trigger-Stream-Start"; + export type RunShape = TRunTypes extends AnyRunTypes ? { id: string; @@ -82,6 +90,7 @@ export type RunStreamCallback = ( export type RunShapeStreamOptions = { headers?: Record; + resolveHeaders?: () => Promise>; fetchClient?: typeof fetch; closeOnComplete?: boolean; signal?: AbortSignal; @@ -114,6 +123,7 @@ export function runShapeStream( getEnvVar("TRIGGER_STREAM_URL", getEnvVar("TRIGGER_API_URL")) ?? "https://api.trigger.dev", { headers: options?.headers, + resolveHeaders: options?.resolveHeaders, signal: abortController.signal, } ); @@ -159,6 +169,17 @@ export type CreateStreamSubscriptionOptions = { onError?: (error: Error) => void; timeoutInSeconds?: number; lastEventId?: string; + /** + * Where a fresh subscription (no `lastEventId`) starts reading from. + * + * - `"beginning"` (default): replay the full stream history, then live-tail. + * - `"latest"`: skip history and start at the current tail — the subscriber + * sees only records appended after it connects (a last-value / live view). + * + * Ignored once `lastEventId` is set: a reconnect always resumes from the last + * seen record, so `"latest"` only governs the very first connect. + */ + from?: "beginning" | "latest"; }; export interface StreamSubscriptionFactory { @@ -196,6 +217,7 @@ type PumpItem = { type: "part"; part: SSEStreamPart }; // Real implementation for production export class SSEStreamSubscription implements StreamSubscription { private lastEventId: string | undefined; + private from: "beginning" | "latest"; private retryCount = 0; private maxRetries: number; private retryDelayMs: number; @@ -208,6 +230,12 @@ export class SSEStreamSubscription implements StreamSubscription { private internalAbort: AbortController | null = null; private cancelledByConsumer = false; private completeNotified = false; + /** Headers for the next attempt. Replaced by `resolveHeaders` after an auth failure. */ + private currentHeaders: Record | undefined; + /** A refresh has already been tried on this connection; a second auth failure is terminal. */ + private authRefreshed = false; + /** Headers were just refreshed for the auth error now unwinding; retry once instead of failing. */ + private retryAfterAuthRefresh = false; /** * True when the most recent response carried `X-Session-Settled: true` — @@ -225,6 +253,7 @@ export class SSEStreamSubscription implements StreamSubscription { onError?: (error: Error) => void; timeoutInSeconds?: number; lastEventId?: string; + from?: "beginning" | "latest"; // Retry knobs. Defaults: retry forever, 100ms initial backoff, // capped at 5s with 50% jitter. Keeps mobile clients reconnecting // through transient drops without giving up after a fixed window @@ -257,9 +286,12 @@ export class SSEStreamSubscription implements StreamSubscription { // the SSE connect through a custom path (proxy, custom headers, // tracing). Defaults to global `fetch`. fetchClient?: typeof fetch; + resolveHeaders?: () => Promise>; } ) { + this.currentHeaders = options.headers; this.lastEventId = options.lastEventId; + this.from = options.from ?? "beginning"; this.maxRetries = options.maxRetries ?? Infinity; this.retryDelayMs = options.retryDelayMs ?? 100; this.maxRetryDelayMs = options.maxRetryDelayMs ?? 5000; @@ -388,9 +420,10 @@ export class SSEStreamSubscription implements StreamSubscription { try { const headers: Record = { Accept: "text/event-stream", - ...this.options.headers, + ...this.currentHeaders, }; if (this.lastEventId) headers["Last-Event-ID"] = this.lastEventId; + else if (this.from === "latest") headers[STREAM_START_HEADER] = "latest"; if (this.options.timeoutInSeconds) { headers["Timeout-Seconds"] = this.options.timeoutInSeconds.toString(); } @@ -409,11 +442,14 @@ export class SSEStreamSubscription implements StreamSubscription { "Could not subscribe to stream", Object.fromEntries(response.headers) ); - this.options.onError?.(error); if (this.nonRetryableStatuses.has(response.status)) { + this.options.onError?.(error); controller.error(error); return; } + if (!(await this.refreshHeadersForAuthError(response.status))) { + this.options.onError?.(error); + } throw error; } @@ -541,6 +577,7 @@ export class SSEStreamSubscription implements StreamSubscription { } armStall(); // any chunk (including server keepalives) resets the silence timer + this.authRefreshed = false; controller.enqueue(value); } } catch (error) { @@ -556,11 +593,15 @@ export class SSEStreamSubscription implements StreamSubscription { } if (isTriggerRealtimeAuthError(error)) { - // `onError` was already invoked in the `!response.ok` branch above - // (where the auth ApiError was originally constructed and thrown). - // Auth errors are non-retryable: terminate the stream cleanly. - controller.error(error as Error); - return; + if (this.retryAfterAuthRefresh) { + this.retryAfterAuthRefresh = false; + } else { + // `onError` was already invoked in the `!response.ok` branch above + // (where the auth ApiError was originally constructed and thrown). + // Auth errors are non-retryable: terminate the stream cleanly. + controller.error(error as Error); + return; + } } cleanupAttempt(); @@ -570,6 +611,29 @@ export class SSEStreamSubscription implements StreamSubscription { } } + /** + * Re-resolve the headers after a 401/403 so the retry carries a fresh token. + * At most once per live connection: if the refreshed token is rejected too, + * the auth error stays terminal. A refresher that can't mint leaves the + * rejected token in place so the auth error stays terminal too. Returns true + * when a retry should follow. + */ + private async refreshHeadersForAuthError(status: number): Promise { + if (status !== 401 && status !== 403) return false; + if (!this.options.resolveHeaders || this.authRefreshed) return false; + + this.authRefreshed = true; + + try { + this.currentHeaders = await this.options.resolveHeaders(); + } catch { + return false; + } + + this.retryAfterAuthRefresh = true; + return true; + } + private async retryConnection( controller: ReadableStreamDefaultController, error?: Error @@ -648,6 +712,7 @@ export class SSEStreamSubscriptionFactory implements StreamSubscriptionFactory { private options: { headers?: Record; signal?: AbortSignal; + resolveHeaders?: () => Promise>; } ) {} diff --git a/packages/core/src/v3/apiClientManager/index.ts b/packages/core/src/v3/apiClientManager/index.ts index e9ec9fd1e4e..9cb15997bec 100644 --- a/packages/core/src/v3/apiClientManager/index.ts +++ b/packages/core/src/v3/apiClientManager/index.ts @@ -102,6 +102,7 @@ export class APIClientManagerAPI { getEnvVar("TRIGGER_SECRET_KEY") ?? getEnvVar("TRIGGER_ACCESS_TOKEN"), secretKey: partial.secretKey, + refreshAccessToken: partial.refreshAccessToken, previewBranch: partial.previewBranch ?? getEnvVar("TRIGGER_PREVIEW_BRANCH") ?? @@ -128,7 +129,8 @@ export class APIClientManagerAPI { this.accessToken, this.branchName, requestOptions, - futureFlags + futureFlags, + source?.refreshAccessToken ); } @@ -146,7 +148,14 @@ export class APIClientManagerAPI { const requestOptions = config?.requestOptions ?? source?.requestOptions; const futureFlags = config?.future ?? source?.future; - return new ApiClient(baseURL, accessToken, branchName, requestOptions, futureFlags); + return new ApiClient( + baseURL, + accessToken, + branchName, + requestOptions, + futureFlags, + config?.refreshAccessToken ?? source?.refreshAccessToken + ); } runWithConfig Promise>( diff --git a/packages/core/src/v3/apiClientManager/types.ts b/packages/core/src/v3/apiClientManager/types.ts index fe738182279..cf562d05ca9 100644 --- a/packages/core/src/v3/apiClientManager/types.ts +++ b/packages/core/src/v3/apiClientManager/types.ts @@ -10,6 +10,12 @@ export type ApiClientConfiguration = { * The access token to authenticate with the Trigger API. */ accessToken?: string; + /** + * Mints a fresh access token. Called when a realtime stream subscription is + * rejected with a 401/403, so a long-lived subscription can survive the + * expiry of the token it started with. + */ + refreshAccessToken?: () => Promise; /** * The preview branch name (for preview environments) */ diff --git a/packages/core/src/v3/index.ts b/packages/core/src/v3/index.ts index f2916f88576..47bdeeaed7d 100644 --- a/packages/core/src/v3/index.ts +++ b/packages/core/src/v3/index.ts @@ -3,6 +3,7 @@ export * from "./apiClient/types.js"; export * from "./apiClient/pagination.js"; export type { ApiPromise, OffsetLimitPagePromise, CursorPagePromise } from "./apiClient/core.js"; export * from "./apiClient/errors.js"; +export * from "./apiClient/refreshAccessToken.js"; export * from "./clock-api.js"; export * from "./errors.js"; export * from "./externalDeploymentId.js"; diff --git a/packages/core/src/v3/realtimeStreams/types.ts b/packages/core/src/v3/realtimeStreams/types.ts index 06c397e45c2..5f3562c4b7b 100644 --- a/packages/core/src/v3/realtimeStreams/types.ts +++ b/packages/core/src/v3/realtimeStreams/types.ts @@ -125,6 +125,17 @@ export type ReadStreamOptions = { * @default 0 (start from beginning) */ startIndex?: number; + + /** + * Where a fresh read starts. + * + * - `"beginning"` (default): replay the full stream history, then live-tail. + * - `"latest"`: skip history and start at the current tail — only records + * appended after this read connects are delivered (a last-value / live view). + * + * Ignored when `startIndex` is set (which pins an absolute start position). + */ + from?: "beginning" | "latest"; }; /** diff --git a/packages/react-hooks/src/hooks/useApiClient.ts b/packages/react-hooks/src/hooks/useApiClient.ts index 02206f4a2f8..d9c538742b0 100644 --- a/packages/react-hooks/src/hooks/useApiClient.ts +++ b/packages/react-hooks/src/hooks/useApiClient.ts @@ -1,7 +1,8 @@ "use client"; import type { ApiRequestOptions } from "@trigger.dev/core/v3"; -import { ApiClient } from "@trigger.dev/core/v3"; +import { ApiClient, refreshAccessTokenOnce } from "@trigger.dev/core/v3"; +import { useCallback, useEffect, useRef } from "react"; import { useTriggerAuthContextOptional } from "../contexts.js"; /** @@ -16,6 +17,11 @@ export type UseApiClientOptions = { previewBranch?: string; /** Optional additional request configuration */ requestOptions?: ApiRequestOptions; + /** + * Optional callback that mints a fresh access token. Used to reconnect a + * realtime stream that the server rejected because its token expired. + */ + refreshAccessToken?: () => Promise; /** * Enable or disable the API client instance. @@ -51,6 +57,22 @@ export function useApiClient(options?: UseApiClientOptions): ApiClient | undefin const baseUrl = options?.baseURL ?? auth?.baseURL ?? "https://api.trigger.dev"; const accessToken = options?.accessToken ?? auth?.accessToken; const previewBranch = options?.previewBranch ?? auth?.previewBranch; + const refreshAccessToken = options?.refreshAccessToken ?? auth?.refreshAccessToken; + + const refreshAccessTokenRef = useRef(refreshAccessToken); + useEffect(() => { + refreshAccessTokenRef.current = refreshAccessToken; + }, [refreshAccessToken]); + const stableRefreshAccessToken = useCallback(async () => { + const refresh = refreshAccessTokenRef.current; + + if (!refresh) { + throw new Error("Missing refreshAccessToken in TriggerAuthContext or useApiClient options"); + } + + return refreshAccessTokenOnce(refresh); + }, []); + if (!accessToken) { if (options?.enabled === false) { return undefined; @@ -64,5 +86,12 @@ export function useApiClient(options?: UseApiClientOptions): ApiClient | undefin ...options?.requestOptions, }; - return new ApiClient(baseUrl, accessToken, previewBranch, requestOptions); + return new ApiClient( + baseUrl, + accessToken, + previewBranch, + requestOptions, + undefined, + refreshAccessToken ? stableRefreshAccessToken : undefined + ); } diff --git a/packages/react-hooks/src/hooks/useRealtime.ts b/packages/react-hooks/src/hooks/useRealtime.ts index b07b7359648..f76869c719e 100644 --- a/packages/react-hooks/src/hooks/useRealtime.ts +++ b/packages/react-hooks/src/hooks/useRealtime.ts @@ -8,6 +8,7 @@ import type { RealtimeDefinedStream, RealtimeRun, RealtimeRunSkipColumns, + SSEStreamPart, } from "@trigger.dev/core/v3"; import { useCallback, useEffect, useId, useRef, useState } from "react"; import type { KeyedMutator } from "../utils/trigger-swr.js"; @@ -15,17 +16,7 @@ import { useSWR } from "../utils/trigger-swr.js"; import type { UseApiClientOptions } from "./useApiClient.js"; import { useApiClient } from "./useApiClient.js"; import { createThrottledQueue } from "../utils/throttle.js"; - -// Keep subscription lifecycles controlled by their effects while using the latest request inputs. -function useStableRequestCallback(callback: () => Promise) { - const callbackRef = useRef(callback); - - useEffect(() => { - callbackRef.current = callback; - }, [callback]); - - return useCallback(() => callbackRef.current(), []); -} +import { useStableRequestCallback } from "../utils/useStableRequestCallback.js"; export type UseRealtimeRunOptions = UseApiClientOptions & { id?: string; @@ -625,6 +616,13 @@ export function useRealtimeBatch( export type UseRealtimeStreamInstance = { parts: Array; + /** + * The event id of the last part seen. Persist this (e.g. to localStorage) and + * pass it back as the `lastEventId` option to resume the stream where you left + * off after a page reload. Updated on each throttled flush. + */ + lastEventId: string | undefined; + error: Error | undefined; /** @@ -657,10 +655,50 @@ export type UseRealtimeStreamOptions = UseApiClientOptions & { */ startIndex?: number; + /** + * The event id to resume from, as returned in `lastEventId`. Persist it across + * a page reload and pass it back to continue where the previous session left + * off, with no replay and no gap. Takes precedence over `startIndex` and + * `from`. + */ + lastEventId?: string | number; + + /** + * Where a fresh subscription starts reading. + * + * - `"beginning"` (default): replay the full stream history, then live-tail. + * - `"latest"`: start at the current tail (the latest record, then live + * updates) instead of replaying history, for a last-value / live view. On + * reconnect or remount the subscription resumes from the last record it + * saw, so no frames are missed and none are replayed. + * + * Ignored when `startIndex` is set (which pins an absolute start position). + */ + from?: "beginning" | "latest"; + + /** + * Cap the number of parts kept in the accumulated `parts` array. When more + * than `maxParts` parts have been received, only the most recent `maxParts` + * are retained (older parts are dropped). Use `maxParts: 1` together with + * `from: "latest"` for a pure last-value view with bounded memory. + * + * When unset, `parts` accumulates every record for the lifetime of the + * subscription (the default). + */ + maxParts?: number; + /** * Callback this is called when new data is received. */ onData?: (data: TPart) => void; + + /** + * Callback invoked once per throttled flush with the batch of parts in that + * flush, each carrying its event `id`, `chunk` and `timestamp`. Use it to + * track the resume cursor without re-rendering on every record. Fires at the + * `throttleInMs` cadence, not per record. + */ + onParts?: (parts: Array>) => void; }; export function useRealtimeStream>( @@ -834,6 +872,20 @@ function useRealtimeStreamImplementation( partsRef.current = parts || ([] as Array); }, [parts]); + const { data: persistedLastEventId, mutate: mutateLastEventId } = useSWR( + [idKey, runId, streamKey, "lastEventId"], + null + ); + const lastEventIdRef = useRef(persistedLastEventId); + const streamIdentityRef = useRef(`${idKey}:${runId}:${streamKey}`); + useEffect(() => { + const identity = `${idKey}:${runId}:${streamKey}`; + if (streamIdentityRef.current !== identity) { + streamIdentityRef.current = identity; + lastEventIdRef.current = persistedLastEventId; + } + }, [idKey, runId, streamKey, persistedLastEventId]); + // Add state to track when the subscription is complete const { data: _isComplete = false, mutate: setIsComplete } = useSWR( [idKey, runId, streamKey, "complete"], @@ -865,10 +917,32 @@ function useRealtimeStreamImplementation( [onDataCallback] ); + const onPartsCallback = options?.onParts; + const onParts = useCallback( + (partsBatch: Array>) => { + if (onPartsCallback) { + onPartsCallback(partsBatch); + } + }, + [onPartsCallback] + ); + const apiClient = useApiClient(options); const timeoutInSeconds = options?.timeoutInSeconds; const startIndex = options?.startIndex; + const startEventId = options?.lastEventId; const throttleInMs = options?.throttleInMs; + const from = options?.from; + const maxParts = options?.maxParts; + + useEffect(() => { + if (maxParts != null && maxParts >= 0) { + const current = partsRef.current; + if (current.length > maxParts) { + mutateParts(current.slice(current.length - maxParts)); + } + } + }, [maxParts, mutateParts]); const triggerRequest = useCallback(async () => { try { @@ -890,7 +964,13 @@ function useRealtimeStreamImplementation( abortControllerRef, timeoutInSeconds, startIndex, - throttleInMs ?? 16 + throttleInMs ?? 16, + from, + maxParts, + lastEventIdRef, + (id) => mutateLastEventId(id, false), + startEventId !== undefined ? String(startEventId) : undefined, + onParts ); } catch (err) { // Ignore abort errors as they are expected. @@ -916,9 +996,14 @@ function useRealtimeStreamImplementation( setError, setIsComplete, onData, + onParts, timeoutInSeconds, startIndex, + startEventId, throttleInMs, + from, + maxParts, + mutateLastEventId, ]); const requestSubscription = useStableRequestCallback(triggerRequest); @@ -938,7 +1023,7 @@ function useRealtimeStreamImplementation( }; }, [runId, stop, options?.enabled, requestSubscription]); - return { parts: parts ?? initialPartsFallback, error, stop }; + return { parts: parts ?? initialPartsFallback, lastEventId: persistedLastEventId, error, stop }; } async function processRealtimeBatch( @@ -1114,24 +1199,52 @@ async function processRealtimeStream( abortControllerRef: React.MutableRefObject, timeoutInSeconds?: number, startIndex?: number, - throttleInMs?: number + throttleInMs?: number, + from?: "beginning" | "latest", + maxParts?: number, + lastEventIdRef?: React.MutableRefObject, + persistLastEventId?: (id: string) => void, + userLastEventId?: string, + onParts?: (parts: Array>) => void ) { try { + const resumeFromEventId = + lastEventIdRef?.current ?? + userLastEventId ?? + (startIndex ? (startIndex - 1).toString() : undefined); + + const partsQueue = createThrottledQueue>(async (batch) => { + const combined = [...existingPartsRef.current, ...batch.map((part) => part.chunk)]; + const bounded = + maxParts != null && maxParts >= 0 && combined.length > maxParts + ? combined.slice(combined.length - maxParts) + : combined; + existingPartsRef.current = bounded; + mutatePartsData(bounded); + if (persistLastEventId && lastEventIdRef?.current) { + persistLastEventId(lastEventIdRef.current); + } + onParts?.(batch); + }, throttleInMs); + const stream = await apiClient.fetchStream(runId, streamKey, { signal: abortControllerRef.current?.signal, timeoutInSeconds, - lastEventId: startIndex ? (startIndex - 1).toString() : undefined, + lastEventId: resumeFromEventId, + from: startIndex !== undefined ? undefined : from, + onPart: (part) => { + if (part.id && lastEventIdRef) { + lastEventIdRef.current = part.id; + } + partsQueue.add(part); + }, }); - // Throttle the stream - const streamQueue = createThrottledQueue(async (parts) => { - mutatePartsData([...existingPartsRef.current, ...parts]); - }, throttleInMs); - for await (const part of stream) { onData(part); - streamQueue.add(part); } + + await partsQueue.flush(); } catch (err) { if ((err as any).name === "AbortError") { return; diff --git a/packages/react-hooks/src/hooks/useSessionStream.ts b/packages/react-hooks/src/hooks/useSessionStream.ts new file mode 100644 index 00000000000..686c2c6fcc7 --- /dev/null +++ b/packages/react-hooks/src/hooks/useSessionStream.ts @@ -0,0 +1,405 @@ +"use client"; + +import type { ApiClient, ControlEvent, SSEStreamPart } from "@trigger.dev/core/v3"; +import { useCallback, useEffect, useId, useRef, useState } from "react"; +import { createThrottledQueue } from "../utils/throttle.js"; +import type { KeyedMutator } from "../utils/trigger-swr.js"; +import { useSWR } from "../utils/trigger-swr.js"; +import { useStableRequestCallback } from "../utils/useStableRequestCallback.js"; +import type { UseApiClientOptions } from "./useApiClient.js"; +import { useApiClient } from "./useApiClient.js"; + +export type UseSessionStreamInstance = { + /** + * The records received so far on the channel, in arrival order. Control records are + * never included here, they are delivered to `onControl` instead. + */ + records: Array; + + /** + * The cursor of the last record seen. Persist this and pass it back as the `lastEventId` + * option to resume the channel where you left off. + */ + lastEventId: string | undefined; + + /** + * The last control record seen on the channel (e.g. `turn-complete`). + */ + lastControl: ControlEvent | undefined; + + error: Error | undefined; + + /** + * Abort the current request immediately, keep the records received so far. + */ + stop: () => void; +}; + +export type UseSessionStreamOptions = UseApiClientOptions & { + id?: string; + enabled?: boolean; + /** + * Which channel of the session to read. + * + * @default "out" + */ + io?: "out" | "in"; + /** + * The number of milliseconds to throttle the record updates. + * + * @default 16 + */ + throttleInMs?: number; + /** + * The number of seconds to wait for new data to be available, + * If no data arrives within the timeout, the stream will be closed. + * + * @default 60 seconds + */ + timeoutInSeconds?: number; + /** + * The cursor to resume from. If not provided, the channel is read from the beginning. + */ + lastEventId?: string | number; + /** + * Where a fresh subscription (no `lastEventId`) starts reading. + * + * - `"beginning"` (default): replay the full channel history, then live-tail. + * - `"latest"`: start at the current tail (the latest record, then live + * updates) instead of replaying history, for a last-value / live view. + * + * Ignored when `lastEventId` is set. + */ + from?: "beginning" | "latest"; + /** + * Cap the number of records kept in the accumulated `records` array. When more + * than `maxRecords` have been received, only the most recent `maxRecords` are + * retained. Use `maxRecords: 1` with `from: "latest"` for a last-value view + * with bounded memory. When unset, `records` accumulates without bound. + */ + maxRecords?: number; + /** + * Callback invoked once per throttled flush with the batch of records in that + * flush, each carrying its event `id`, `chunk` and `timestamp`. Fires at the + * `throttleInMs` cadence (not per record) and includes control records, so it + * can track the resume cursor for everything on the channel. + */ + onRecords?: (records: Array>) => void; + /** + * Callback this is called when a control record is received (e.g. `turn-complete`). + */ + onControl?: (event: ControlEvent) => void; +}; + +/** + * Hook to read one channel of a Session's realtime stream. + * + * This hook subscribes to one of the session's channels (`out` by default, or `in`) and + * updates the `records` array as new records arrive. It is read-only: use `useSession` for + * two-way (read and write) communication. The subscription is automatically managed: it + * starts when the component mounts (or when `enabled` becomes `true`) and stops when the + * component unmounts or when `stop()` is called. + * + * Requires a Public Access Token with the `read:sessions:{id}` scope. + * + * @template TRecord - The type of each record on the channel + * @param sessionIdOrExternalId - The id or external id of the session to subscribe to + * @param options - Optional configuration for the subscription + * @returns An object containing: + * - `records`: An array of all the records received so far (accumulates over time) + * - `lastEventId`: The cursor of the last record seen, for resuming later + * - `lastControl`: The last control record seen + * - `error`: Any error that occurred during subscription + * - `stop`: A function to manually stop the subscription + * + * @example + * ```tsx + * "use client"; + * import { useSessionStream } from "@trigger.dev/react-hooks"; + * + * function SessionViewer({ sessionId }: { sessionId: string }) { + * const { records, error } = useSessionStream(sessionId, { + * accessToken: publicAccessToken, + * }); + * + * if (error) return
Error: {error.message}
; + * + * return
{records.join("")}
; + * } + * ``` + * + * @example + * ```tsx + * // Read the input channel, resuming from a persisted cursor + * const { records, lastEventId, stop } = useSessionStream(sessionId, { + * accessToken: publicAccessToken, + * io: "in", + * lastEventId: persistedCursor, + * onControl: (event) => { + * if (event.subtype === "turn-complete") { + * console.log("The turn is complete"); + * } + * }, + * }); + * ``` + */ +export function useSessionStream( + sessionIdOrExternalId?: string, + options?: UseSessionStreamOptions +): UseSessionStreamInstance { + const hookId = useId(); + const idKey = options?.id ?? hookId; + const io = options?.io ?? "out"; + + const [initialRecordsFallback] = useState([] as Array); + + const { data: records, mutate: mutateRecords } = useSWR>( + [idKey, sessionIdOrExternalId, io, "records"], + null, + { + fallbackData: initialRecordsFallback, + } + ); + + const recordsRef = useRef>(records ?? ([] as Array)); + useEffect(() => { + recordsRef.current = records || ([] as Array); + }, [records]); + + const { data: lastEventId = undefined, mutate: setLastEventId } = useSWR( + [idKey, sessionIdOrExternalId, io, "lastEventId"], + null + ); + const lastEventIdRef = useRef(lastEventId); + const channelIdentityRef = useRef(`${idKey}:${sessionIdOrExternalId}:${io}`); + useEffect(() => { + const identity = `${idKey}:${sessionIdOrExternalId}:${io}`; + if (channelIdentityRef.current !== identity) { + channelIdentityRef.current = identity; + lastEventIdRef.current = lastEventId; + } + }, [idKey, sessionIdOrExternalId, io, lastEventId]); + + const { data: lastControl = undefined, mutate: setLastControl } = useSWR< + undefined | ControlEvent + >([idKey, sessionIdOrExternalId, io, "lastControl"], null); + + const { data: _isComplete = false, mutate: setIsComplete } = useSWR( + [idKey, sessionIdOrExternalId, io, "complete"], + null + ); + + const { data: error = undefined, mutate: setError } = useSWR( + [idKey, sessionIdOrExternalId, io, "error"], + null + ); + + const abortControllerRef = useRef(null); + + const stop = useCallback(() => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + abortControllerRef.current = null; + } + }, []); + + const onRecordsCallback = options?.onRecords; + const onRecords = useCallback( + (recordsBatch: Array>) => { + if (onRecordsCallback) { + onRecordsCallback(recordsBatch); + } + }, + [onRecordsCallback] + ); + + const onControlCallback = options?.onControl; + const onControl = useCallback( + (event: ControlEvent) => { + if (onControlCallback) { + onControlCallback(event); + } + }, + [onControlCallback] + ); + + const apiClient = useApiClient(options); + const timeoutInSeconds = options?.timeoutInSeconds; + const startEventId = options?.lastEventId; + const throttleInMs = options?.throttleInMs; + const from = options?.from; + const maxRecords = options?.maxRecords; + + useEffect(() => { + if (maxRecords != null && maxRecords >= 0) { + const current = recordsRef.current; + if (current.length > maxRecords) { + mutateRecords(current.slice(current.length - maxRecords)); + } + } + }, [maxRecords, mutateRecords]); + + const triggerRequest = useCallback(async () => { + try { + if (!sessionIdOrExternalId || !apiClient) { + return; + } + + const abortController = new AbortController(); + abortControllerRef.current = abortController; + + await processSessionStream( + sessionIdOrExternalId, + io, + apiClient, + mutateRecords, + recordsRef, + setLastEventId, + setLastControl, + setError, + onRecords, + onControl, + abortControllerRef, + timeoutInSeconds, + startEventId !== undefined ? String(startEventId) : lastEventIdRef.current, + throttleInMs ?? 16, + from, + maxRecords + ); + } catch (err) { + if ((err as any).name === "AbortError") { + abortControllerRef.current = null; + return; + } + + setError(err as Error); + } finally { + if (abortControllerRef.current) { + abortControllerRef.current = null; + } + + setIsComplete(true); + } + }, [ + sessionIdOrExternalId, + io, + apiClient, + mutateRecords, + setLastEventId, + setLastControl, + setError, + setIsComplete, + onRecords, + onControl, + timeoutInSeconds, + startEventId, + throttleInMs, + from, + maxRecords, + ]); + const requestSubscription = useStableRequestCallback(triggerRequest); + + useEffect(() => { + if (typeof options?.enabled === "boolean" && !options.enabled) { + return; + } + + if (!sessionIdOrExternalId) { + return; + } + + requestSubscription().finally(() => {}); + + return () => { + stop(); + }; + }, [sessionIdOrExternalId, io, stop, options?.enabled, requestSubscription]); + + return { records: records ?? initialRecordsFallback, lastEventId, lastControl, error, stop }; +} + +async function processSessionStream( + sessionIdOrExternalId: string, + io: "out" | "in", + apiClient: ApiClient, + mutateRecordsData: KeyedMutator>, + existingRecordsRef: React.MutableRefObject>, + setLastEventId: KeyedMutator, + setLastControl: KeyedMutator, + onError: (e: Error) => void, + onRecords: (records: Array>) => void, + onControl: (event: ControlEvent) => void, + abortControllerRef: React.MutableRefObject, + timeoutInSeconds?: number, + lastEventId?: string, + throttleInMs?: number, + from?: "beginning" | "latest", + maxRecords?: number +) { + let lastSeenEventId: string | undefined; + let publishedEventId: string | undefined; + let partsBatch: Array> = []; + + const publishLastEventId = () => { + if (lastSeenEventId !== publishedEventId) { + publishedEventId = lastSeenEventId; + setLastEventId(lastSeenEventId); + } + }; + + const flushParts = () => { + if (partsBatch.length === 0) return; + const batch = partsBatch; + partsBatch = []; + onRecords(batch); + }; + + try { + const stream = await apiClient.subscribeToSessionStream(sessionIdOrExternalId, io, { + signal: abortControllerRef.current?.signal, + timeoutInSeconds, + lastEventId, + from, + onPart: (part) => { + lastSeenEventId = part.id; + partsBatch.push(part); + }, + onControl: (event) => { + setLastControl(event); + onControl(event); + }, + }); + + const recordsQueue = createThrottledQueue(async (newRecords) => { + const combined = [...existingRecordsRef.current, ...newRecords]; + const bounded = + maxRecords != null && maxRecords >= 0 && combined.length > maxRecords + ? combined.slice(combined.length - maxRecords) + : combined; + existingRecordsRef.current = bounded; + mutateRecordsData(bounded); + publishLastEventId(); + flushParts(); + }, throttleInMs); + + for await (const record of stream) { + recordsQueue.add(record); + } + + await recordsQueue.flush(); + publishLastEventId(); + flushParts(); + } catch (err) { + if ((err as any).name === "AbortError") { + return; + } + + if (err instanceof Error) { + onError(err); + } else { + onError(new Error(String(err))); + } + + throw err; + } +} diff --git a/packages/react-hooks/src/index.ts b/packages/react-hooks/src/index.ts index 23c8ca947d5..57aa3b16877 100644 --- a/packages/react-hooks/src/index.ts +++ b/packages/react-hooks/src/index.ts @@ -5,3 +5,4 @@ export * from "./hooks/useRealtime.js"; export * from "./hooks/useTaskTrigger.js"; export * from "./hooks/useWaitToken.js"; export * from "./hooks/useInputStreamSend.js"; +export * from "./hooks/useSessionStream.js"; diff --git a/packages/react-hooks/src/utils/useStableRequestCallback.ts b/packages/react-hooks/src/utils/useStableRequestCallback.ts new file mode 100644 index 00000000000..61af8812255 --- /dev/null +++ b/packages/react-hooks/src/utils/useStableRequestCallback.ts @@ -0,0 +1,17 @@ +"use client"; + +import { useCallback, useEffect, useRef } from "react"; + +/** + * Keep subscription lifecycles controlled by their effects while using the + * latest request inputs. + */ +export function useStableRequestCallback(callback: () => Promise) { + const callbackRef = useRef(callback); + + useEffect(() => { + callbackRef.current = callback; + }, [callback]); + + return useCallback(() => callbackRef.current(), []); +} diff --git a/packages/trigger-sdk/src/v3/streams.ts b/packages/trigger-sdk/src/v3/streams.ts index 81d19e128e7..6c2307c7788 100644 --- a/packages/trigger-sdk/src/v3/streams.ts +++ b/packages/trigger-sdk/src/v3/streams.ts @@ -382,6 +382,7 @@ async function readStreamImpl( signal: options?.signal, timeoutInSeconds: options?.timeoutInSeconds ?? 60, lastEventId: options?.startIndex ? (options.startIndex - 1).toString() : undefined, + from: options?.startIndex !== undefined ? undefined : options?.from, onComplete: () => { span.end(); },