diff --git a/.changeset/spotty-pillows-visit.md b/.changeset/spotty-pillows-visit.md new file mode 100644 index 00000000000..4e62b419b07 --- /dev/null +++ b/.changeset/spotty-pillows-visit.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/react-hooks": patch +--- + +Added a `useSession` React hook for reading a session's output or input channel in realtime, with automatic resume from the last record you received. diff --git a/packages/core/src/v3/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts index e6b4125222f..446b3042863 100644 --- a/packages/core/src/v3/apiClient/index.ts +++ b/packages/core/src/v3/apiClient/index.ts @@ -195,6 +195,7 @@ export type { AnyRealtimeRun, AnyRunShape, ApiRequestOptions, + ControlEvent, RealtimeRun, RunShape, RunStreamCallback, diff --git a/packages/react-hooks/src/hooks/useRealtime.ts b/packages/react-hooks/src/hooks/useRealtime.ts index b07b7359648..afb355be200 100644 --- a/packages/react-hooks/src/hooks/useRealtime.ts +++ b/packages/react-hooks/src/hooks/useRealtime.ts @@ -15,17 +15,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; diff --git a/packages/react-hooks/src/hooks/useSession.ts b/packages/react-hooks/src/hooks/useSession.ts new file mode 100644 index 00000000000..2cdf06db086 --- /dev/null +++ b/packages/react-hooks/src/hooks/useSession.ts @@ -0,0 +1,353 @@ +"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 UseSessionInstance = { + /** + * 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 UseSessionOptions = 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; + /** + * Callback this is called when a record is received, before throttling. This fires for + * control records too, so you can track the cursor for every record on the channel. + */ + onRecord?: (record: SSEStreamPart) => void; + /** + * Callback this is called when a control record is received (e.g. `turn-complete`). + */ + onControl?: (event: ControlEvent) => void; +}; + +/** + * Hook to subscribe to a Session channel. + * + * This hook automatically subscribes to one of the session's channels and updates the `records` + * array as new records arrive. 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 { useSession } from "@trigger.dev/react-hooks"; + * + * function SessionViewer({ sessionId }: { sessionId: string }) { + * const { records, error } = useSession(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 } = useSession(sessionId, { + * accessToken: publicAccessToken, + * io: "in", + * lastEventId: persistedCursor, + * onControl: (event) => { + * if (event.subtype === "turn-complete") { + * console.log("The turn is complete"); + * } + * }, + * }); + * ``` + */ +export function useSession( + sessionIdOrExternalId?: string, + options?: UseSessionOptions +): UseSessionInstance { + const hookId = useId(); + const idKey = options?.id ?? hookId; + const io = options?.io ?? "out"; + + const [initialRecordsFallback] = useState([] as Array); + + // Store the records state in SWR, using the idKey as the key to share states. + const { data: records, mutate: mutateRecords } = useSWR>( + [idKey, sessionIdOrExternalId, io, "records"], + null, + { + fallbackData: initialRecordsFallback, + } + ); + + // Keep the latest records in a ref. + 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 { data: lastControl = undefined, mutate: setLastControl } = useSWR< + undefined | ControlEvent + >([idKey, sessionIdOrExternalId, io, "lastControl"], null); + + // Add state to track when the subscription is complete + const { data: _isComplete = false, mutate: setIsComplete } = useSWR( + [idKey, sessionIdOrExternalId, io, "complete"], + null + ); + + const { data: error = undefined, mutate: setError } = useSWR( + [idKey, sessionIdOrExternalId, io, "error"], + null + ); + + // Abort controller to cancel the current API call. + const abortControllerRef = useRef(null); + + const stop = useCallback(() => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + abortControllerRef.current = null; + } + }, []); + + const onRecordCallback = options?.onRecord; + const onRecord = useCallback( + (record: SSEStreamPart) => { + if (onRecordCallback) { + onRecordCallback(record); + } + }, + [onRecordCallback] + ); + + 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 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, + onRecord, + onControl, + abortControllerRef, + timeoutInSeconds, + startEventId !== undefined ? String(startEventId) : undefined, + throttleInMs ?? 16 + ); + } catch (err) { + // Ignore abort errors as they are expected. + if ((err as any).name === "AbortError") { + abortControllerRef.current = null; + return; + } + + setError(err as Error); + } finally { + if (abortControllerRef.current) { + abortControllerRef.current = null; + } + + // Mark the subscription as complete + setIsComplete(true); + } + }, [ + sessionIdOrExternalId, + io, + apiClient, + mutateRecords, + setLastEventId, + setLastControl, + setError, + setIsComplete, + onRecord, + onControl, + timeoutInSeconds, + startEventId, + throttleInMs, + ]); + 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, + onRecord: (record: SSEStreamPart) => void, + onControl: (event: ControlEvent) => void, + abortControllerRef: React.MutableRefObject, + timeoutInSeconds?: number, + lastEventId?: string, + throttleInMs?: number +) { + // Published with the throttled record flush, so consumers re-render once per + // batch instead of once per record. + let lastSeenEventId: string | undefined; + let publishedEventId: string | undefined; + + const publishLastEventId = () => { + if (lastSeenEventId !== publishedEventId) { + publishedEventId = lastSeenEventId; + setLastEventId(lastSeenEventId); + } + }; + + try { + const stream = await apiClient.subscribeToSessionStream(sessionIdOrExternalId, io, { + signal: abortControllerRef.current?.signal, + timeoutInSeconds, + lastEventId, + onPart: (part) => { + lastSeenEventId = part.id; + onRecord(part); + }, + onControl: (event) => { + setLastControl(event); + onControl(event); + }, + }); + + // Throttle the records + const recordsQueue = createThrottledQueue(async (newRecords) => { + mutateRecordsData([...existingRecordsRef.current, ...newRecords]); + publishLastEventId(); + }, throttleInMs); + + for await (const record of stream) { + recordsQueue.add(record); + } + + // The last batch can be smaller than the throttle window, so flush it. The + // cursor is published even when that batch is empty (control records only). + await recordsQueue.flush(); + publishLastEventId(); + } 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..af9902ec9e0 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/useSession.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..538e20d8bfa --- /dev/null +++ b/packages/react-hooks/src/utils/useStableRequestCallback.ts @@ -0,0 +1,14 @@ +"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(), []); +}