WIP

ericallam · ericallam · commit 6ca8cb711851 · 2025-03-27T06:37:44.000Z
diff --git a/apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.complete.ts b/apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.complete.ts
@@ -13,14 +13,20 @@ import { logger } from "~/services/logger.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { engine } from "~/v3/runEngine.server";
 
-const { action } = createActionApiRoute(
+const { action, loader } = createActionApiRoute(
   {
     params: z.object({
       waitpointFriendlyId: z.string(),
     }),
     body: CompleteWaitpointTokenRequestBody,
     maxContentLength: env.TASK_PAYLOAD_MAXIMUM_SIZE,
-    method: "POST",
+    allowJWT: true,
+    authorization: {
+      action: "write",
+      resource: (params) => ({ waitpoints: params.waitpointFriendlyId }),
+      superScopes: ["write:waitpoints", "admin"],
+    },
+    corsStrategy: "all",
   },
   async ({ authentication, body, params }) => {
     // Resume tokens are actually just waitpoints
@@ -65,4 +71,4 @@ const { action } = createActionApiRoute(
   }
 );
 
-export { action };
+export { action, loader };
diff --git a/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts b/apps/webapp/app/routes/api.v1.waitpoints.tokens.ts
@@ -13,6 +13,7 @@ import {
   createActionApiRoute,
   createLoaderApiRoute,
 } from "~/services/routeBuilders/apiBuilder.server";
+import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { parseDelay } from "~/utils/delays";
 import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
 import { engine } from "~/v3/runEngine.server";
@@ -77,12 +78,14 @@ const { action } = createActionApiRoute(
         tags: bodyTags,
       });
 
+      const $responseHeaders = await responseHeaders(authentication.environment);
+
       return json<CreateWaitpointTokenResponseBody>(
         {
           id: WaitpointId.toFriendlyId(result.waitpoint.id),
           isCached: result.isCached,
         },
-        { status: 200 }
+        { status: 200, headers: $responseHeaders }
       );
     } catch (error) {
       if (error instanceof ServiceValidationError) {
@@ -96,4 +99,17 @@ const { action } = createActionApiRoute(
   }
 );
 
+async function responseHeaders(
+  environment: AuthenticatedEnvironment
+): Promise<Record<string, string>> {
+  const claimsHeader = JSON.stringify({
+    sub: environment.id,
+    pub: true,
+  });
+
+  return {
+    "x-trigger-jwt-claims": claimsHeader,
+  };
+}
+
 export { action };
diff --git a/apps/webapp/app/services/authorization.server.ts b/apps/webapp/app/services/authorization.server.ts
@@ -1,6 +1,6 @@
 export type AuthorizationAction = "read" | "write" | string; // Add more actions as needed
 
-const ResourceTypes = ["tasks", "tags", "runs", "batch"] as const;
+const ResourceTypes = ["tasks", "tags", "runs", "batch", "waitpoints"] as const;
 
 export type AuthorizationResources = {
   [key in (typeof ResourceTypes)[number]]?: string | string[];
diff --git a/packages/core/src/v3/apiClient/core.ts b/packages/core/src/v3/apiClient/core.ts
@@ -5,7 +5,7 @@ import { calculateNextRetryDelay } from "../utils/retries.js";
 import { ApiConnectionError, ApiError, ApiSchemaValidationError } from "./errors.js";
 
 import { Attributes, context, propagation, Span } from "@opentelemetry/api";
-import {suppressTracing} from "@opentelemetry/core"
+import { suppressTracing } from "@opentelemetry/core";
 import { SemanticInternalAttributes } from "../semanticInternalAttributes.js";
 import type { TriggerTracer } from "../tracer.js";
 import { accessoryAttributes } from "../utils/styleAttributes.js";
@@ -27,14 +27,14 @@ export const defaultRetryOptions = {
   randomize: false,
 } satisfies RetryOptions;
 
-export type ZodFetchOptions<T = unknown> = {
+export type ZodFetchOptions<TInput = unknown, TOutput = TInput> = {
   retry?: RetryOptions;
   tracer?: TriggerTracer;
   name?: string;
   attributes?: Attributes;
   icon?: string;
-  onResponseBody?: (body: T, span: Span) => void;
-  prepareData?: (data: T) => Promise<T> | T;
+  onResponseBody?: (body: TInput, span: Span) => void;
+  prepareData?: (data: TInput, response: Response) => Promise<TOutput> | TOutput;
 };
 
 export type AnyZodFetchOptions = ZodFetchOptions<any>;
@@ -67,12 +67,15 @@ interface FetchOffsetLimitPageParams extends OffsetLimitPageParams {
   query?: URLSearchParams;
 }
 
-export function zodfetch<TResponseBodySchema extends z.ZodTypeAny>(
+export function zodfetch<
+  TResponseBodySchema extends z.ZodTypeAny,
+  TOutput = z.output<TResponseBodySchema>,
+>(
   schema: TResponseBodySchema,
   url: string,
   requestInit?: RequestInit,
-  options?: ZodFetchOptions<z.output<TResponseBodySchema>>
-): ApiPromise<z.output<TResponseBodySchema>> {
+  options?: ZodFetchOptions<z.output<TResponseBodySchema>, TOutput>
+): ApiPromise<TOutput> {
   return new ApiPromise(_doZodFetch(schema, url, requestInit, options));
 }
 
@@ -110,7 +113,14 @@ export function zodfetchCursorPage<TItemSchema extends z.ZodTypeAny>(
 
   const fetchResult = _doZodFetch(cursorPageSchema, $url.href, requestInit, options);
 
-  return new CursorPagePromise(fetchResult, schema, url, params, requestInit, options);
+  return new CursorPagePromise(
+    fetchResult as Promise<ZodFetchResult<CursorPageResponse<z.output<TItemSchema>>>>,
+    schema,
+    url,
+    params,
+    requestInit,
+    options
+  );
 }
 
 export function zodfetchOffsetLimitPage<TItemSchema extends z.ZodTypeAny>(
@@ -144,7 +154,14 @@ export function zodfetchOffsetLimitPage<TItemSchema extends z.ZodTypeAny>(
 
   const fetchResult = _doZodFetch(offsetLimitPageSchema, $url.href, requestInit, options);
 
-  return new OffsetLimitPagePromise(fetchResult, schema, url, params, requestInit, options);
+  return new OffsetLimitPagePromise(
+    fetchResult as Promise<ZodFetchResult<OffsetLimitPageResponse<z.output<TItemSchema>>>>,
+    schema,
+    url,
+    params,
+    requestInit,
+    options
+  );
 }
 
 type ZodFetchResult<T> = {
@@ -184,12 +201,15 @@ async function traceZodFetch<T>(
   );
 }
 
-async function _doZodFetch<TResponseBodySchema extends z.ZodTypeAny>(
+async function _doZodFetch<
+  TResponseBodySchema extends z.ZodTypeAny,
+  TOutput = z.output<TResponseBodySchema>,
+>(
   schema: TResponseBodySchema,
   url: string,
   requestInit?: PromiseOrValue<RequestInit>,
-  options?: ZodFetchOptions
-): Promise<ZodFetchResult<z.output<TResponseBodySchema>>> {
+  options?: ZodFetchOptions<z.output<TResponseBodySchema>, TOutput>
+): Promise<ZodFetchResult<TOutput>> {
   let $requestInit = await requestInit;
 
   return traceZodFetch({ url, requestInit: $requestInit, options }, async (span) => {
@@ -202,7 +222,7 @@ async function _doZodFetch<TResponseBodySchema extends z.ZodTypeAny>(
     }
 
     if (options?.prepareData) {
-      result.data = await options.prepareData(result.data);
+      result.data = await options.prepareData(result.data, result.response);
     }
 
     return result;
@@ -707,7 +727,7 @@ export async function wrapZodFetch<T extends z.ZodTypeAny>(
   schema: T,
   url: string,
   requestInit?: RequestInit,
-  options?: ZodFetchOptions<z.output<T>>
+  options?: ZodFetchOptions<z.output<T>, z.infer<T>>
 ): Promise<ApiResult<z.infer<T>>> {
   try {
     const response = await zodfetch(schema, url, requestInit, {
diff --git a/packages/core/src/v3/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts
@@ -214,36 +214,37 @@ export class ApiClient {
         headers: this.#getHeaders(clientOptions?.spanParentAsLink ?? false),
         body: JSON.stringify(body),
       },
-      mergeRequestOptions(this.defaultRequestOptions, requestOptions)
-    )
-      .withResponse()
-      .then(async ({ response, data }) => {
-        const jwtHeader = response.headers.get("x-trigger-jwt");
+      {
+        ...mergeRequestOptions(this.defaultRequestOptions, requestOptions),
+        prepareData: async (data, response) => {
+          const jwtHeader = response.headers.get("x-trigger-jwt");
+
+          if (typeof jwtHeader === "string") {
+            return {
+              ...data,
+              publicAccessToken: jwtHeader,
+            };
+          }
+
+          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,
+              scopes: [`read:runs:${data.id}`],
+            },
+            expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
+          });
 
-        if (typeof jwtHeader === "string") {
           return {
             ...data,
-            publicAccessToken: jwtHeader,
+            publicAccessToken: jwt,
           };
-        }
-
-        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,
-            scopes: [`read:runs:${data.id}`],
-          },
-          expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
-        });
-
-        return {
-          ...data,
-          publicAccessToken: jwt,
-        };
-      });
+        },
+      }
+    );
   }
 
   batchTriggerV3(
@@ -261,27 +262,28 @@ 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,
-            scopes: [`read:batch:${data.id}`],
-          },
-          expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
-        });
-
-        return {
-          ...data,
-          publicAccessToken: jwt,
-        };
-      });
+      {
+        ...mergeRequestOptions(this.defaultRequestOptions, requestOptions),
+        prepareData: async (data, response) => {
+          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,
+              scopes: [`read:batch:${data.id}`],
+            },
+            expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
+          });
+
+          return {
+            ...data,
+            publicAccessToken: jwt,
+          };
+        },
+      }
+    );
   }
 
   createUploadPayloadUrl(filename: string, requestOptions?: ZodFetchOptions) {
@@ -669,7 +671,36 @@ export class ApiClient {
         headers: this.#getHeaders(false),
         body: JSON.stringify(options),
       },
-      mergeRequestOptions(this.defaultRequestOptions, requestOptions)
+      {
+        ...mergeRequestOptions(this.defaultRequestOptions, requestOptions),
+        prepareData: async (data, response) => {
+          const jwtHeader = response.headers.get("x-trigger-jwt");
+
+          if (typeof jwtHeader === "string") {
+            return {
+              ...data,
+              publicAccessToken: jwtHeader,
+            };
+          }
+
+          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,
+              scopes: [`write:waitpoints:${data.id}`],
+            },
+            expirationTime: "1h",
+          });
+
+          return {
+            ...data,
+            publicAccessToken: jwt,
+          };
+        },
+      }
     );
   }
 
diff --git a/packages/trigger-sdk/src/v3/wait.ts b/packages/trigger-sdk/src/v3/wait.ts
@@ -48,10 +48,14 @@ import { SpanStatusCode } from "@opentelemetry/api";
  * @param requestOptions - The request options for the waitpoint token.
  * @returns The waitpoint token.
  */
+export type CreateWaitpointTokenResponse = CreateWaitpointTokenResponseBody & {
+  publicAccessToken: string;
+};
+
 function createToken(
   options?: CreateWaitpointTokenRequestBody,
   requestOptions?: ApiRequestOptions
-): ApiPromise<CreateWaitpointTokenResponseBody> {
+): ApiPromise<CreateWaitpointTokenResponse> {
   const apiClient = apiClientManager.clientOrThrow();
 
   const $requestOptions = mergeRequestOptions(
