Add streaming support for the realtime react hooks

Author: ericallam

apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts

@@ -111,13 +111,13 @@ async function responseHeaders(
   triggerClient?: string | null
 ): Promise<Record<string, string>> {
   const claimsHeader = JSON.stringify({
-    sub: run.runtimeEnvironmentId,
+    sub: environment.id,
     pub: true,
   });
 
   if (triggerClient === "browser") {
     const claims = {
-      sub: run.runtimeEnvironmentId,
+      sub: environment.id,
       pub: true,
       scopes: [`read:runs:${run.friendlyId}`],
     };

apps/webapp/app/services/apiAuth.server.ts

@@ -41,6 +41,7 @@ export async function authenticateApiRequest(
   options: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
 ): Promise<ApiAuthenticationResult | undefined> {
   const apiKey = getApiKeyFromRequest(request);
+
   if (!apiKey) {
     return;
   }

apps/webapp/app/services/routeBuiilders/apiBuilder.server.ts

@@ -506,5 +506,7 @@ export function createActionApiRoute<
 }
 
 function wrapResponse(request: Request, response: Response, useCors: boolean) {
-  return useCors ? apiCors(request, response) : response;
+  return useCors
+    ? apiCors(request, response, { exposedHeaders: ["x-trigger-jwt", "x-trigger-jwt-claims"] })
+    : response;
 }

apps/webapp/app/utils/apiCors.ts

@@ -8,6 +8,7 @@ type CorsOptions = {
   maxAge?: number;
   origin?: boolean | string;
   credentials?: boolean;
+  exposedHeaders?: string[];
 };
 
 export async function apiCors(

packages/core/src/v3/apiClient/index.ts

@@ -45,6 +45,7 @@ import {
   RunStreamCallback,
   RunSubscription,
   TaskRunShape,
+  RealtimeRun,
 } from "./runStream.js";
 import {
   CreateEnvironmentVariableParams,
@@ -88,7 +89,14 @@ const DEFAULT_ZOD_FETCH_OPTIONS: ZodFetchOptions = {
 
 export { isRequestOptions };
 export type { ApiRequestOptions };
-export type { RunShape, AnyRunShape, TaskRunShape, RunStreamCallback, RunSubscription };
+export type {
+  RunShape,
+  AnyRunShape,
+  TaskRunShape,
+  RealtimeRun,
+  RunStreamCallback,
+  RunSubscription,
+};
 
 /**
  * Trigger.dev v3 API client
@@ -603,15 +611,19 @@ export class ApiClient {
     );
   }
 
-  subscribeToRun<TRunTypes extends AnyRunTypes>(runId: string) {
+  subscribeToRun<TRunTypes extends AnyRunTypes>(runId: string, options?: { signal?: AbortSignal }) {
     return runShapeStream<TRunTypes>(`${this.baseUrl}/realtime/v1/runs/${runId}`, {
       closeOnComplete: true,
       headers: this.#getRealtimeHeaders(),
       client: this,
+      signal: options?.signal,
     });
   }
 
-  subscribeToRunsWithTag<TRunTypes extends AnyRunTypes>(tag: string | string[]) {
+  subscribeToRunsWithTag<TRunTypes extends AnyRunTypes>(
+    tag: string | string[],
+    options?: { signal?: AbortSignal }
+  ) {
     const searchParams = createSearchQueryForSubscribeToRuns({
       tags: tag,
     });
@@ -622,15 +634,20 @@ export class ApiClient {
         closeOnComplete: false,
         headers: this.#getRealtimeHeaders(),
         client: this,
+        signal: options?.signal,
       }
     );
   }
 
-  subscribeToBatch<TRunTypes extends AnyRunTypes>(batchId: string) {
+  subscribeToBatch<TRunTypes extends AnyRunTypes>(
+    batchId: string,
+    options?: { signal?: AbortSignal }
+  ) {
     return runShapeStream<TRunTypes>(`${this.baseUrl}/realtime/v1/batches/${batchId}`, {
       closeOnComplete: false,
       headers: this.#getRealtimeHeaders(),
       client: this,
+      signal: options?.signal,
     });
   }
 

packages/core/src/v3/apiClient/runStream.ts

@@ -42,6 +42,7 @@ export type RunShape<TRunTypes extends AnyRunTypes> = TRunTypes extends AnyRunTy
 export type AnyRunShape = RunShape<AnyRunTypes>;
 
 export type TaskRunShape<TTask extends AnyTask> = RunShape<InferRunTypes<TTask>>;
+export type RealtimeRun<TTask extends AnyTask> = TaskRunShape<TTask>;
 
 export type RunStreamCallback<TRunTypes extends AnyRunTypes> = (
   run: RunShape<TRunTypes>
@@ -99,7 +100,7 @@ export interface StreamSubscription {
 }
 
 export interface StreamSubscriptionFactory {
-  createSubscription(runId: string, streamKey: string): StreamSubscription;
+  createSubscription(runId: string, streamKey: string, baseUrl?: string): StreamSubscription;
 }
 
 // Real implementation for production
@@ -145,8 +146,8 @@ export class SSEStreamSubscriptionFactory implements StreamSubscriptionFactory {
     private options: { headers?: Record<string, string>; signal?: AbortSignal }
   ) {}
 
-  createSubscription(runId: string, streamKey: string): StreamSubscription {
-    const url = `${this.baseUrl}/realtime/v1/streams/${runId}/${streamKey}`;
+  createSubscription(runId: string, streamKey: string, baseUrl?: string): StreamSubscription {
+    const url = `${baseUrl ?? this.baseUrl}/realtime/v1/streams/${runId}/${streamKey}`;
     return new SSEStreamSubscription(url, this.options);
   }
 }
@@ -243,7 +244,8 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
 
                 const subscription = this.options.streamFactory.createSubscription(
                   run.id,
-                  streamKey.toString()
+                  streamKey.toString(),
+                  this.options.client?.baseUrl
                 );
 
                 await subscription.subscribe(async (chunk) => {

packages/react-hooks/package.json

@@ -38,7 +38,8 @@
   },
   "dependencies": {
     "@trigger.dev/core": "workspace:^3.1.2",
-    "swr": "^2.2.5"
+    "swr": "^2.2.5",
+    "throttleit": "^2.1.0"
   },
   "devDependencies": {
     "@arethetypeswrong/cli": "^0.15.4",

packages/react-hooks/src/contexts.tsx

@@ -4,7 +4,7 @@ import React from "react";
 import { createContextAndHook } from "./utils/createContextAndHook.js";
 import type { ApiClientConfiguration } from "@trigger.dev/core/v3";
 
-const [TriggerAuthContext, useTriggerAuthContext] =
+const [TriggerAuthContext, useTriggerAuthContext, useTriggerAuthContextOptional] =
   createContextAndHook<ApiClientConfiguration>("TriggerAuthContext");
 
-export { TriggerAuthContext, useTriggerAuthContext };
+export { TriggerAuthContext, useTriggerAuthContext, useTriggerAuthContextOptional };

packages/react-hooks/src/hooks/useApiClient.ts

@@ -1,16 +1,28 @@
 "use client";
 
-import { ApiClient } from "@trigger.dev/core/v3";
-import { useTriggerAuthContext } from "../contexts.js";
+import { ApiClient, ApiRequestOptions } from "@trigger.dev/core/v3";
+import { useTriggerAuthContextOptional } from "../contexts.js";
 
-export function useApiClient() {
-  const auth = useTriggerAuthContext();
+export type UseApiClientOptions = {
+  accessToken?: string;
+  baseURL?: string;
+  requestOptions?: ApiRequestOptions;
+};
 
-  const baseUrl = auth.baseURL ?? "https://api.trigger.dev";
+export function useApiClient(options?: UseApiClientOptions): ApiClient {
+  const auth = useTriggerAuthContextOptional();
 
-  if (!auth.accessToken) {
-    throw new Error("Missing accessToken in TriggerAuthContext");
+  const baseUrl = auth?.baseURL ?? options?.baseURL ?? "https://api.trigger.dev";
+  const accessToken = auth?.accessToken ?? options?.accessToken;
+
+  if (!accessToken) {
+    throw new Error("Missing accessToken in TriggerAuthContext or useApiClient options");
   }
 
-  return new ApiClient(baseUrl, auth.accessToken, auth.requestOptions);
+  const requestOptions: ApiRequestOptions = {
+    ...auth?.requestOptions,
+    ...options?.requestOptions,
+  };
+
+  return new ApiClient(baseUrl, accessToken, requestOptions);
 }