Wait for token is working

Author: matt-aitken

apps/webapp/app/routes/api.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts

@@ -0,0 +1,71 @@
+import { json } from "@remix-run/server-runtime";
+import {
+  CompleteWaitpointTokenRequestBody,
+  CompleteWaitpointTokenResponseBody,
+  conditionallyExportPacket,
+  CreateWaitpointTokenResponseBody,
+  stringifyIO,
+  WaitForWaitpointTokenRequestBody,
+  WaitForWaitpointTokenResponseBody,
+} from "@trigger.dev/core/v3";
+import { RunId, WaitpointId } from "@trigger.dev/core/v3/apps";
+import { z } from "zod";
+import { $replica } from "~/db.server";
+import { logger } from "~/services/logger.server";
+import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
+import { parseDelay } from "~/utils/delays";
+import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
+import { engine } from "~/v3/runEngine.server";
+
+const { action } = createActionApiRoute(
+  {
+    params: z.object({
+      runFriendlyId: z.string(),
+      waitpointFriendlyId: z.string(),
+    }),
+    body: WaitForWaitpointTokenRequestBody,
+    maxContentLength: 1024 * 10, // 10KB
+    method: "POST",
+  },
+  async ({ authentication, body, params }) => {
+    // Resume tokens are actually just waitpoints
+    const waitpointId = WaitpointId.toId(params.waitpointFriendlyId);
+    const runId = RunId.toId(params.runFriendlyId);
+
+    const timeout = await parseDelay(body.timeout);
+
+    try {
+      //check permissions
+      const waitpoint = await $replica.waitpoint.findFirst({
+        where: {
+          id: waitpointId,
+          environmentId: authentication.environment.id,
+        },
+      });
+
+      if (!waitpoint) {
+        throw json({ error: "Waitpoint not found" }, { status: 404 });
+      }
+
+      const result = await engine.blockRunWithWaitpoint({
+        runId,
+        waitpoints: [waitpointId],
+        environmentId: authentication.environment.id,
+        projectId: authentication.environment.project.id,
+        failAfter: timeout,
+      });
+
+      return json<WaitForWaitpointTokenResponseBody>(
+        {
+          success: true,
+        },
+        { status: 200 }
+      );
+    } catch (error) {
+      logger.error("Failed to wait for waitpoint", { runId, waitpointId, error });
+      throw json({ error: "Failed to wait for waitpoint token" }, { status: 500 });
+    }
+  }
+);
+
+export { action };

internal-packages/run-engine/src/engine/index.ts

@@ -2067,8 +2067,8 @@ export class RunEngine {
           id: `continueRunIfUnblocked:${runId}`,
           job: "continueRunIfUnblocked",
           payload: { runId: runId },
-          //100ms in the future
-          availableAt: new Date(Date.now() + 100),
+          //in the near future
+          availableAt: new Date(Date.now() + 50),
         });
       }
 
@@ -2099,13 +2099,15 @@ export class RunEngine {
         });
 
         if (affectedTaskRuns.length === 0) {
-          this.logger.warn(`No TaskRunWaitpoints found for waitpoint`, {
+          this.logger.warn(`completeWaitpoint: No TaskRunWaitpoints found for waitpoint`, {
             waitpointId: id,
           });
         }
 
         // 2. Update the waitpoint to completed (only if it's pending)
-        const waitpoint = await tx.waitpoint
+        let waitpoint: Waitpoint | null = null;
+        try {
+        waitpoint = await tx.waitpoint
           .update({
             where: { id, status: "PENDING" },
             data: {
@@ -2116,19 +2118,21 @@ export class RunEngine {
               outputIsError: output?.isError,
             },
           })
-          .catch(async (error) => {
-            if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2025") {
-              return tx.waitpoint.findUnique({
-                where: { id },
-              });
-            }
+        } catch(error) {
+          if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2025") {
+            waitpoint = await tx.waitpoint.findFirst({
+              where: { id },
+            });
+          } else {
+            this.logger.log('completeWaitpoint: error updating waitpoint:', {error});
             throw error;
-          });
+          }
+        };
 
         return { waitpoint, affectedTaskRuns };
       },
       (error) => {
-        this.logger.error(`Error completing waitpoint ${id}, retrying`, { error });
+        this.logger.error(`completeWaitpoint: Error completing waitpoint ${id}, retrying`, { error });
         throw error;
       }
     );

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

@@ -32,6 +32,8 @@ import {
   UpdateMetadataRequestBody,
   UpdateMetadataResponseBody,
   UpdateScheduleOptions,
+  WaitForWaitpointTokenRequestBody,
+  WaitForWaitpointTokenResponseBody,
 } from "../schemas/index.js";
 import { taskContext } from "../task-context-api.js";
 import { AnyRunTypes, TriggerJwtOptions } from "../types/tasks.js";
@@ -669,6 +671,24 @@ export class ApiClient {
     );
   }
 
+  waitForWaitpointToken(
+    runFriendlyId: string,
+    waitpointFriendlyId: string,
+    options?: WaitForWaitpointTokenRequestBody,
+    requestOptions?: ZodFetchOptions
+  ) {
+    return zodfetch(
+      WaitForWaitpointTokenResponseBody,
+      `${this.baseUrl}/api/v1/runs/${runFriendlyId}/waitpoints/tokens/${waitpointFriendlyId}/wait`,
+      {
+        method: "POST",
+        headers: this.#getHeaders(false),
+        body: JSON.stringify(options ?? {}),
+      },
+      mergeRequestOptions(this.defaultRequestOptions, requestOptions)
+    );
+  }
+
   subscribeToRun<TRunTypes extends AnyRunTypes>(
     runId: string,
     options?: {

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

@@ -4,6 +4,8 @@ import {
   BatchTaskRunExecutionResult,
   TaskRunContext,
   TaskRunExecutionResult,
+  WaitForWaitpointTokenRequestBody,
+  WaitpointTokenResult,
 } from "../schemas/index.js";
 import { getGlobal, registerGlobal, unregisterGlobal } from "../utils/globals.js";
 import { type RuntimeManager } from "./manager.js";
@@ -41,6 +43,15 @@ export class RuntimeAPI {
     return usage.pauseAsync(() => this.#getRuntimeManager().waitForTask(params));
   }
 
+  public waitForToken(
+    waitpointFriendlyId: string,
+    options?: WaitForWaitpointTokenRequestBody
+  ): Promise<WaitpointTokenResult> {
+    return usage.pauseAsync(() =>
+      this.#getRuntimeManager().waitForToken(waitpointFriendlyId, options)
+    );
+  }
+
   public waitForBatch(params: {
     id: string;
     runCount: number;

packages/core/src/v3/runtime/managedRuntimeManager.ts

@@ -6,6 +6,8 @@ import {
   TaskRunExecutionResult,
   TaskRunFailedExecutionResult,
   TaskRunSuccessfulExecutionResult,
+  WaitForWaitpointTokenRequestBody,
+  WaitpointTokenResult,
 } from "../schemas/index.js";
 import { ExecutorToWorkerProcessConnection } from "../zodIpc.js";
 import { RuntimeManager } from "./manager.js";
@@ -92,6 +94,23 @@ export class ManagedRuntimeManager implements RuntimeManager {
     };
   }
 
+  async waitForToken(
+    waitpointFriendlyId: string,
+    options?: WaitForWaitpointTokenRequestBody
+  ): Promise<WaitpointTokenResult> {
+    const promise = new Promise<CompletedWaitpoint>((resolve) => {
+      this.resolversByWaitId.set(waitpointFriendlyId, resolve);
+    });
+
+    const waitpoint = await promise;
+
+    return {
+      ok: !waitpoint.outputIsError,
+      output: waitpoint.output,
+      outputType: waitpoint.outputType,
+    };
+  }
+
   associateWaitWithWaitpoint(waitId: string, waitpointId: string) {
     this.resolversByWaitpoint.set(waitpointId, waitId);
   }
@@ -115,6 +134,8 @@ export class ManagedRuntimeManager implements RuntimeManager {
       //no waitpoint resolves associated with batch completions
       //a batch completion isn't when all the runs from a batch are completed
       return;
+    } else if (waitpoint.type === "MANUAL") {
+      waitId = waitpoint.friendlyId;
     } else {
       waitId = this.resolversByWaitpoint.get(waitpoint.id);
     }

packages/core/src/v3/runtime/manager.ts

@@ -2,6 +2,8 @@ import {
   BatchTaskRunExecutionResult,
   TaskRunContext,
   TaskRunExecutionResult,
+  WaitForWaitpointTokenRequestBody,
+  WaitpointTokenResult,
 } from "../schemas/index.js";
 
 export interface RuntimeManager {
@@ -14,4 +16,8 @@ export interface RuntimeManager {
     runCount: number;
     ctx: TaskRunContext;
   }): Promise<BatchTaskRunExecutionResult>;
+  waitForToken(
+    waitpointFriendlyId: string,
+    options?: WaitForWaitpointTokenRequestBody
+  ): Promise<WaitpointTokenResult>;
 }

packages/core/src/v3/runtime/noopRuntimeManager.ts

@@ -3,6 +3,8 @@ import {
   TaskRunContext,
   TaskRunErrorCodes,
   TaskRunExecutionResult,
+  WaitForWaitpointTokenRequestBody,
+  WaitpointTokenResult,
 } from "../schemas/index.js";
 import { RuntimeManager } from "./manager.js";
 
@@ -40,4 +42,14 @@ export class NoopRuntimeManager implements RuntimeManager {
       items: [],
     });
   }
+
+  waitForToken(
+    waitpointFriendlyId: string,
+    options?: WaitForWaitpointTokenRequestBody
+  ): Promise<WaitpointTokenResult> {
+    return Promise.resolve({
+      ok: true,
+      outputType: "application/json",
+    });
+  }
 }

packages/core/src/v3/schemas/api.ts

@@ -928,6 +928,22 @@ export const CompleteWaitpointTokenResponseBody = z.object({
 });
 export type CompleteWaitpointTokenResponseBody = z.infer<typeof CompleteWaitpointTokenResponseBody>;
 
+export const WaitForWaitpointTokenRequestBody = z.object({
+  /**
+   * The maximum amount of time to wait for the token to be completed.
+   * If this is exceeded, the waitpoint will timeout and will return `ok` of `false`.
+   *
+   * You can pass a `Date` object, or a string in this format: "30s", "1m", "2h", "3d", "4w".
+   */
+  timeout: TimePeriod.optional(),
+});
+export type WaitForWaitpointTokenRequestBody = z.infer<typeof WaitForWaitpointTokenRequestBody>;
+
+export const WaitForWaitpointTokenResponseBody = z.object({
+  success: z.boolean(),
+});
+export type WaitForWaitpointTokenResponseBody = z.infer<typeof WaitForWaitpointTokenResponseBody>;
+
 export const WAITPOINT_TIMEOUT_ERROR_CODE = "TRIGGER_WAITPOINT_TIMEOUT";
 
 export function isWaitpointOutputTimeout(output: string): boolean {

packages/core/src/v3/schemas/common.ts

@@ -385,6 +385,23 @@ export const BatchTaskRunExecutionResult = z.object({
 
 export type BatchTaskRunExecutionResult = z.infer<typeof BatchTaskRunExecutionResult>;
 
+export const WaitpointTokenResult = z.object({
+  ok: z.boolean(),
+  output: z.string().optional(),
+  outputType: z.string().optional(),
+});
+export type WaitpointTokenResult = z.infer<typeof WaitpointTokenResult>;
+
+export type WaitpointTokenTypedResult<T> =
+  | {
+      ok: true;
+      output: T;
+    }
+  | {
+      ok: false;
+      error: string;
+    };
+
 export const SerializedError = z.object({
   message: z.string(),
   name: z.string().optional(),