Added useBatch hook

Author: ericallam

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

@@ -104,10 +104,20 @@ export async function action({ request, params }: ActionFunctionArgs) {
       return json({ error: "Task not found" }, { status: 404 });
     }
 
-    return json({
-      batchId: result.batch.friendlyId,
-      runs: result.runs,
-    });
+    return json(
+      {
+        batchId: result.batch.friendlyId,
+        runs: result.runs,
+      },
+      {
+        headers: {
+          "x-trigger-jwt-claims": JSON.stringify({
+            sub: authenticationResult.environment.id,
+            pub: true,
+          }),
+        },
+      }
+    );
   } catch (error) {
     if (error instanceof Error) {
       return json({ error: error.message }, { status: 400 });

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

@@ -107,9 +107,19 @@ export async function action({ request, params }: ActionFunctionArgs) {
       return json({ error: "Task not found" }, { status: 404 });
     }
 
-    return json({
-      id: run.friendlyId,
-    });
+    return json(
+      {
+        id: run.friendlyId,
+      },
+      {
+        headers: {
+          "x-trigger-jwt-claims": JSON.stringify({
+            sub: authenticationResult.environment.id,
+            pub: true,
+          }),
+        },
+      }
+    );
   } catch (error) {
     if (error instanceof ServiceValidationError) {
       return json({ error: error.message }, { status: 422 });

apps/webapp/app/routes/realtime.v1.batches.$batchId.ts

@@ -0,0 +1,81 @@
+import { type ActionFunctionArgs, json } from "@remix-run/server-runtime";
+import { z } from "zod";
+import { $replica } from "~/db.server";
+import { env } from "~/env.server";
+import { permittedToReadBatch } from "~/services/accessControl.server";
+import { authenticateApiRequest } from "~/services/apiAuth.server";
+import { logger } from "~/services/logger.server";
+import { makeApiCors } from "~/utils/apiCors";
+import { longPollingFetch } from "~/utils/longPollingFetch";
+
+const ParamsSchema = z.object({
+  batchId: z.string(),
+});
+
+export async function loader({ request, params }: ActionFunctionArgs) {
+  const apiCors = makeApiCors(request);
+
+  if (request.method.toUpperCase() === "OPTIONS") {
+    return apiCors(json({}));
+  }
+
+  // Authenticate the request
+  const authenticationResult = await authenticateApiRequest(request, { allowJWT: true });
+
+  if (!authenticationResult) {
+    return apiCors(json({ error: "Invalid or Missing API Key" }, { status: 401 }));
+  }
+
+  const parsedParams = ParamsSchema.safeParse(params);
+
+  if (!parsedParams.success) {
+    return apiCors(
+      json(
+        { error: "Invalid request parameters", issues: parsedParams.error.issues },
+        { status: 400 }
+      )
+    );
+  }
+
+  if (!permittedToReadBatch(authenticationResult, parsedParams.data.batchId)) {
+    return apiCors(json({ error: "Unauthorized" }, { status: 403 }));
+  }
+
+  try {
+    const batchRun = await $replica.batchTaskRun.findFirst({
+      where: {
+        friendlyId: parsedParams.data.batchId,
+        runtimeEnvironmentId: authenticationResult.environment.id,
+      },
+    });
+
+    if (!batchRun) {
+      return apiCors(json({ error: "Batch Run not found" }, { status: 404 }));
+    }
+
+    const url = new URL(request.url);
+    const originUrl = new URL(`${env.ELECTRIC_ORIGIN}/v1/shape/public."TaskRun"`);
+    url.searchParams.forEach((value, key) => {
+      originUrl.searchParams.set(key, value);
+    });
+
+    originUrl.searchParams.set("where", `"batchId"='${batchRun.id}'`);
+
+    const finalUrl = originUrl.toString();
+
+    return longPollingFetch(finalUrl);
+  } catch (error) {
+    if (error instanceof Response) {
+      // Error responses from longPollingFetch
+      return apiCors(error);
+    } else if (error instanceof TypeError) {
+      // Unexpected errors
+      logger.error("Unexpected error in loader:", { error: error.message });
+      return apiCors(new Response("An unexpected error occurred", { status: 500 }));
+    } else {
+      // Unknown errors
+      logger.error("Unknown error occurred in loader, not Error", { error: JSON.stringify(error) });
+      return apiCors(new Response("An unknown error occurred", { status: 500 }));
+    }
+  }
+}

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

@@ -41,3 +41,40 @@ export function permittedToReadRun(
 
   return false;
 }
+
+export function permittedToReadBatch(
+  authenticationResult: ApiAuthenticationResult,
+  batchId: string
+): boolean {
+  if (authenticationResult.type === "PRIVATE") {
+    return true;
+  }
+
+  if (authenticationResult.type === "PUBLIC") {
+    return true;
+  }
+
+  if (!authenticationResult.claims) {
+    return false;
+  }
+
+  const parsedClaims = ClaimsSchema.safeParse(authenticationResult.claims);
+
+  if (!parsedClaims.success) {
+    return false;
+  }
+
+  if (parsedClaims.data.permissions?.includes("read:runs")) {
+    return true;
+  }
+
+  if (parsedClaims.data.permissions?.includes(`read:batches:${batchId}`)) {
+    return true;
+  }
+
+  if (parsedClaims.data.permissions?.includes(batchId)) {
+    return true;
+  }
+
+  return false;
+}

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

@@ -53,6 +53,7 @@ import {
   ListRunsQueryParams,
   UpdateEnvironmentVariableParams,
 } from "./types.js";
+import { generateJWT } from "../jwt.js";
 
 export type {
   CreateEnvironmentVariableParams,
@@ -167,7 +168,26 @@ export class ApiClient {
         body: JSON.stringify(body),
       },
       mergeRequestOptions(this.defaultRequestOptions, requestOptions)
-    );
+    )
+      .withResponse()
+      .then(async ({ response, data }) => {
+        const claimsHeader = response.headers.get("x-trigger-jwt-claims");
+        const claims = claimsHeader ? JSON.parse(claimsHeader) : undefined;
+
+        const jwt = await generateJWT({
+          secretKey: this.accessToken,
+          payload: {
+            ...claims,
+            permissions: [data.id],
+          },
+          expirationTime: "1h",
+        });
+
+        return {
+          ...data,
+          jwt,
+        };
+      });
   }
 
   batchTriggerTask(
@@ -187,7 +207,26 @@ export class ApiClient {
         body: JSON.stringify(body),
       },
       mergeRequestOptions(this.defaultRequestOptions, requestOptions)
-    );
+    )
+      .withResponse()
+      .then(async ({ response, data }) => {
+        const claimsHeader = response.headers.get("x-trigger-jwt-claims");
+        const claims = claimsHeader ? JSON.parse(claimsHeader) : undefined;
+
+        const jwt = await generateJWT({
+          secretKey: this.accessToken,
+          payload: {
+            ...claims,
+            permissions: [data.batchId],
+          },
+          expirationTime: "1h",
+        });
+
+        return {
+          ...data,
+          jwt,
+        };
+      });
   }
 
   createUploadPayloadUrl(filename: string, requestOptions?: ZodFetchOptions) {
@@ -545,7 +584,20 @@ export class ApiClient {
   subscribeToRunChanges<TPayload = any, TOutput = any>(runId: string) {
     return runShapeStream<TPayload, TOutput>(
       `${this.baseUrl}/realtime/v1/runs/${runId}`,
-      this.fetchClient
+      this.fetchClient,
+      {
+        closeOnComplete: true,
+      }
+    );
+  }
+
+  subscribeToBatchChanges<TPayload = any, TOutput = any>(batchId: string) {
+    return runShapeStream<TPayload, TOutput>(
+      `${this.baseUrl}/realtime/v1/batches/${batchId}`,
+      this.fetchClient,
+      {
+        closeOnComplete: false,
+      }
     );
   }
 

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

@@ -42,11 +42,16 @@ export type RunStreamCallback<TPayload = any, TOutput = any> = (
   run: RunShape<TPayload, TOutput>
 ) => void | Promise<void>;
 
+export type RunShapeStreamOptions = {
+  closeOnComplete?: boolean;
+};
+
 export function runShapeStream<TPayload = any, TOutput = any>(
   url: string,
-  fetchClient: typeof fetch
+  fetchClient: typeof fetch,
+  options?: RunShapeStreamOptions
 ): RunSubscription<TPayload, TOutput> {
-  const subscription = new RunSubscription<TPayload, TOutput>(url, fetchClient);
+  const subscription = new RunSubscription<TPayload, TOutput>(url, fetchClient, options);
 
   return subscription.init();
 }
@@ -59,7 +64,8 @@ export class RunSubscription<TPayload = any, TOutput = any> {
 
   constructor(
     private url: string,
-    private fetchClient: typeof fetch
+    private fetchClient: typeof fetch,
+    private options?: RunShapeStreamOptions
   ) {
     this.abortController = new AbortController();
   }
@@ -72,7 +78,11 @@ export class RunSubscription<TPayload = any, TOutput = any> {
           this.url,
           async (shape) => {
             controller.enqueue(shape);
-            if (shape.completedAt && !this.abortController.signal.aborted) {
+            if (
+              this.options?.closeOnComplete &&
+              shape.completedAt &&
+              !this.abortController.signal.aborted
+            ) {
               controller.close();
               this.abortController.abort();
             }

packages/core/src/v3/types/tasks.ts

@@ -341,6 +341,7 @@ export type BrandedRun<T, P, O> = T & BrandRun<O, P>;
 export type RunHandle<TPayload, TOutput> = BrandedRun<
   {
     id: string;
+    jwt: string;
   },
   TPayload,
   TOutput
@@ -355,6 +356,7 @@ export type BatchRunHandle<TPayload, TOutput> = BrandedRun<
   {
     batchId: string;
     runs: Array<RunHandle<TPayload, TOutput>>;
+    jwt: string;
   },
   TOutput,
   TPayload

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

@@ -0,0 +1,52 @@
+"use client";
+
+import { AnyTask, TaskRunShape } from "@trigger.dev/core/v3";
+import { useEffect, useState } from "react";
+import { useApiClient } from "./useApiClient.js";
+
+export function useBatch<TTask extends AnyTask>(batchId: string) {
+  const [runShapes, setRunShapes] = useState<TaskRunShape<TTask>[]>([]);
+  const [error, setError] = useState<Error | null>(null);
+  const apiClient = useApiClient();
+
+  useEffect(() => {
+    const subscription = apiClient.subscribeToBatchChanges(batchId);
+
+    async function iterateUpdates() {
+      for await (const run of subscription) {
+        setRunShapes((prevRuns) => {
+          return insertRunShapeInOrder(prevRuns, run);
+        });
+      }
+    }
+
+    iterateUpdates().catch((err) => {
+      setError(err);
+    });
+
+    return () => {
+      subscription.unsubscribe();
+    };
+  }, [batchId]);
+
+  return { runs: runShapes, error };
+}
+
+// Inserts and then orders by the run number, and ensures that the run is not duplicated
+function insertRunShapeInOrder<TTask extends AnyTask>(
+  previousRuns: TaskRunShape<TTask>[],
+  run: TaskRunShape<TTask>
+) {
+  const existingRun = previousRuns.find((r) => r.id === run.id);
+  if (existingRun) {
+    return previousRuns.map((r) => (r.id === run.id ? run : r));
+  }
+
+  const runNumber = run.number;
+  const index = previousRuns.findIndex((r) => r.number > runNumber);
+  if (index === -1) {
+    return [...previousRuns, run];
+  }
+
+  return [...previousRuns.slice(0, index), run, ...previousRuns.slice(index)];
+}

packages/react-hooks/src/index.ts

@@ -1,3 +1,4 @@
 export * from "./contexts.js";
 export * from "./hooks/useApiClient.js";
 export * from "./hooks/useRun.js";
+export * from "./hooks/useBatch.js";

packages/trigger-sdk/src/v3/index.ts

@@ -32,7 +32,7 @@ export {
   type LogLevel,
 } from "@trigger.dev/core/v3";
 
-export { runs, type RunShape, type AnyRunShape } from "./runs.js";
+export { runs, type RunShape, type AnyRunShape, type TaskRunShape } from "./runs.js";
 export * as schedules from "./schedules/index.js";
 export * as envvars from "./envvars.js";
 export type { ImportEnvironmentVariablesParams } from "./envvars.js";

packages/trigger-sdk/src/v3/runs.ts

@@ -5,6 +5,7 @@ import type {
   RescheduleRunRequestBody,
   RunShape,
   AnyRunShape,
+  TaskRunShape,
   RunSubscription,
   Prettify,
 } from "@trigger.dev/core/v3";
@@ -32,7 +33,7 @@ export type RetrieveRunResult<TPayload = any, TOutput = any> = Prettify<
   }
 >;
 
-export type { RunShape, AnyRunShape };
+export type { RunShape, AnyRunShape, TaskRunShape };
 
 export const runs = {
   replay: replayRun,