Add optional retryAfter property to Step interface (#142)

.changeset/chatty-bees-sing.md

@@ -0,0 +1,5 @@
+---
+"@workflow/world-vercel": patch
+---
+
+Enforce the Vercel Queue max visibility limit

.changeset/slick-rabbits-travel.md

@@ -0,0 +1,5 @@
+---
+"@workflow/world-local": patch
+---
+
+Allow `WORKFLOW_LOCAL_QUEUE_MAX_VISIBILITY` env var to set max queue visibility timeout

.changeset/social-paths-swim.md

@@ -0,0 +1,7 @@
+---
+"@workflow/web": patch
+"@workflow/world": patch
+"@workflow/world-postgres": patch
+---
+
+Add optional `retryAfter` property to `Step` interface

.changeset/two-cooks-unite.md

@@ -0,0 +1,5 @@
+---
+"@workflow/core": patch
+---
+
+Respect the `retryAfter` property in the step function callback handler

packages/ai/src/workflow-chat-transport.ts

@@ -47,7 +47,7 @@ export interface WorkflowChatTransportOptions<UI_MESSAGE extends UIMessage> {
    * Defaults to /api/chat if not provided
    */
   api?: string;
-  
+
   /**
    * Custom fetch implementation to use for HTTP requests.
    * Defaults to the global fetch function if not provided.

packages/core/src/runtime.ts

@@ -532,6 +532,24 @@ export const stepEntrypoint =
             ...Attribute.StepStatus(step.status),
           });
 
+          // Check if the step has a `retryAfter` timestamp that hasn't been reached yet
+          const now = Date.now();
+          if (step.retryAfter && step.retryAfter.getTime() > now) {
+            const timeoutSeconds = Math.ceil(
+              (step.retryAfter.getTime() - now) / 1000
+            );
+            span?.setAttributes({
+              ...Attribute.StepRetryTimeoutSeconds(timeoutSeconds),
+            });
+            runtimeLogger.debug('Step retryAfter timestamp not yet reached', {
+              stepName,
+              stepId: step.stepId,
+              retryAfter: step.retryAfter,
+              timeoutSeconds,
+            });
+            return { timeoutSeconds };
+          }
+
           let result: unknown;
           const attempt = step.attempt + 1;
           try {
@@ -710,6 +728,13 @@ export const stepEntrypoint =
                   },
                 });
 
+                await world.steps.update(workflowRunId, stepId, {
+                  status: 'pending', // TODO: Should be "retrying" once we have that status
+                  ...(RetryableError.is(err) && {
+                    retryAfter: err.retryAfter,
+                  }),
+                });
+
                 const timeoutSeconds = Math.max(
                   1,
                   RetryableError.is(err)
@@ -717,10 +742,6 @@ export const stepEntrypoint =
                     : 1
                 );
 
-                await world.steps.update(workflowRunId, stepId, {
-                  status: 'pending', // TODO: Should be "retrying" once we have that status
-                });
-
                 span?.setAttributes({
                   ...Attribute.StepRetryTimeoutSeconds(timeoutSeconds),
                   ...Attribute.StepRetryWillRetry(true),

packages/web-shared/src/sidebar/attribute-panel.tsx

@@ -50,6 +50,7 @@ const attributeOrder: AttributeKey[] = [
   'startedAt',
   'updatedAt',
   'completedAt',
+  'retryAfter',
   'error',
   'errorCode',
   'metadata',
@@ -98,6 +99,7 @@ const attributeToDisplayFn: Record<
   startedAt: (value: unknown) => new Date(String(value)).toLocaleString(),
   updatedAt: (value: unknown) => new Date(String(value)).toLocaleString(),
   completedAt: (value: unknown) => new Date(String(value)).toLocaleString(),
+  retryAfter: (value: unknown) => new Date(String(value)).toLocaleString(),
   // Resolved attributes, won't actually use this function
   metadata: JsonBlock,
   input: (value: unknown) => {

packages/world-local/src/queue.ts

@@ -4,6 +4,12 @@ import { MessageId, type Queue, ValidQueueName } from '@workflow/world';
 import { monotonicFactory } from 'ulid';
 import z from 'zod';
 
+// For local queue, there is no technical limit on the message visibility lifespan,
+// but the environment variable can be used for testing purposes to set a max visibility limit.
+const LOCAL_QUEUE_MAX_VISIBILITY =
+  parseInt(process.env.WORKFLOW_LOCAL_QUEUE_MAX_VISIBILITY ?? '0', 10) ||
+  Infinity;
+
 export function createQueue(port?: number): Queue {
   const transport = new JsonTransport();
   const generateId = monotonicFactory();
@@ -69,8 +75,8 @@ export function createQueue(port?: number): Queue {
 
         if (response.status === 503) {
           try {
-            const retryIn = Number(JSON.parse(text).retryIn);
-            await setTimeout(retryIn * 1000);
+            const timeoutSeconds = Number(JSON.parse(text).timeoutSeconds);
+            await setTimeout(timeoutSeconds * 1000);
             defaultRetriesLeft++;
             continue;
           } catch {}
@@ -124,12 +130,18 @@ export function createQueue(port?: number): Queue {
 
       const body = await new JsonTransport().deserialize(req.body);
       try {
-        const response = await handler(body, { attempt, queueName, messageId });
-        const retryIn =
-          typeof response === 'undefined' ? null : response.timeoutSeconds;
+        const result = await handler(body, { attempt, queueName, messageId });
 
-        if (retryIn) {
-          return Response.json({ retryIn }, { status: 503 });
+        let timeoutSeconds: number | null = null;
+        if (typeof result?.timeoutSeconds === 'number') {
+          timeoutSeconds = Math.min(
+            result.timeoutSeconds,
+            LOCAL_QUEUE_MAX_VISIBILITY
+          );
+        }
+
+        if (timeoutSeconds) {
+          return Response.json({ timeoutSeconds }, { status: 503 });
         }
 
         return Response.json({ ok: true });

packages/world-postgres/src/drizzle/schema.ts

@@ -110,6 +110,7 @@ export const steps = pgTable(
       .defaultNow()
       .$onUpdateFn(() => new Date())
       .notNull(),
+    retryAfter: timestamp('retry_after'),
   } satisfies DrizzlishOfType<Step>,
   (tb) => ({
     runFk: index().on(tb.runId),

packages/world-vercel/src/queue.ts

@@ -13,6 +13,8 @@ const MessageWrapper = z.object({
   queueName: ValidQueueName,
 });
 
+const VERCEL_QUEUE_MAX_VISIBILITY = 82800; // 23 hours in seconds
+
 export function createQueue(config?: APIConfig): Queue {
   const { baseUrl, usingProxy } = getHttpUrl(config);
   const headers = getHeaders(config);
@@ -45,13 +47,26 @@ export function createQueue(config?: APIConfig): Queue {
   const createQueueHandler: Queue['createQueueHandler'] = (prefix, handler) => {
     return handleCallback({
       [`${prefix}*`]: {
-        default: (body, meta) => {
+        default: async (body, meta) => {
           const { payload, queueName } = MessageWrapper.parse(body);
-          return handler(payload, {
+          const result = await handler(payload, {
             queueName,
             messageId: MessageId.parse(meta.messageId),
             attempt: meta.deliveryCount,
           });
+          if (typeof result?.timeoutSeconds === 'number') {
+            // For Vercel Queue, enforce the max visibility limit:
+            //   - When a step function throws a `RetryableError`, the retryAfter timestamp is updated and stored on the Step document
+            const adjustedTimeoutSeconds = Math.min(
+              result.timeoutSeconds,
+              VERCEL_QUEUE_MAX_VISIBILITY
+            );
+
+            if (adjustedTimeoutSeconds !== result.timeoutSeconds) {
+              result.timeoutSeconds = adjustedTimeoutSeconds;
+            }
+          }
+          return result;
         },
       },
     });

packages/world/src/steps.ts

@@ -25,6 +25,7 @@ export const StepSchema = z.object({
   completedAt: z.coerce.date().optional(),
   createdAt: z.coerce.date(),
   updatedAt: z.coerce.date(),
+  retryAfter: z.coerce.date().optional(),
 });
 
 // Inferred types
@@ -44,6 +45,7 @@ export interface UpdateStepRequest {
   output?: SerializedData;
   error?: string;
   errorCode?: string;
+  retryAfter?: Date;
 }
 
 export interface GetStepParams {