suspension-handler.ts

import type { Span } from '@opentelemetry/api';
import { waitUntil } from '@vercel/functions';
import { WorkflowAPIError } from '@workflow/errors';
import {
  type CreateEventRequest,
  type SerializedData,
  SPEC_VERSION_CURRENT,
  type WorkflowRun,
  type World,
} from '@workflow/world';
import type {
  HookInvocationQueueItem,
  StepInvocationQueueItem,
  WaitInvocationQueueItem,
  WorkflowSuspension,
} from '../global.js';
import { importKey } from '../encryption.js';
import { runtimeLogger } from '../logger.js';
import { dehydrateStepArguments } from '../serialization.js';
import * as Attribute from '../telemetry/semantic-conventions.js';
import { serializeTraceCarrier } from '../telemetry.js';
import { queueMessage } from './helpers.js';

/**
 * Extracts W3C trace context headers from a trace carrier for HTTP propagation.
 * Returns an object with `traceparent` and optionally `tracestate` headers.
 */
function extractTraceHeaders(
  traceCarrier: Record<string, string>
): Record<string, string> {
  const headers: Record<string, string> = {};
  if (traceCarrier.traceparent) {
    headers.traceparent = traceCarrier.traceparent;
  }
  if (traceCarrier.tracestate) {
    headers.tracestate = traceCarrier.tracestate;
  }
  return headers;
}

export interface SuspensionHandlerParams {
  suspension: WorkflowSuspension;
  world: World;
  run: WorkflowRun;
  span?: Span;
}

export interface SuspensionHandlerResult {
  timeoutSeconds?: number;
}

/**
 * Handles a workflow suspension by processing all pending operations (hooks, steps, waits).
 * Uses an event-sourced architecture where entities (steps, hooks) are created atomically
 * with their corresponding events via events.create().
 *
 * Processing order:
 * 1. Hooks are processed first to prevent race conditions with webhook receivers
 * 2. Steps and waits are processed in parallel after hooks complete
 */
export async function handleSuspension({
  suspension,
  world,
  run,
  span,
}: SuspensionHandlerParams): Promise<SuspensionHandlerResult> {
  const runId = run.runId;
  const workflowName = run.workflowName;
  const workflowStartedAt = run.startedAt ? +run.startedAt : Date.now();
  // Separate queue items by type
  const stepItems = suspension.steps.filter(
    (item): item is StepInvocationQueueItem => item.type === 'step'
  );
  const allHookItems = suspension.steps.filter(
    (item): item is HookInvocationQueueItem => item.type === 'hook'
  );
  const waitItems = suspension.steps.filter(
    (item): item is WaitInvocationQueueItem => item.type === 'wait'
  );

  // Split hooks by what actions they need
  const hooksNeedingCreation = allHookItems.filter(
    (item) => !item.hasCreatedEvent
  );
  // Hooks needing disposal: any disposed hook (including those needing creation first)
  // Hooks are created before disposal in the processing order below
  const hooksNeedingDisposal = allHookItems.filter((item) => item.disposed);

  // Resolve encryption key for this run
  const rawKey = await world.getEncryptionKeyForRun?.(run);
  const encryptionKey = rawKey ? await importKey(rawKey) : undefined;

  // Build hook_created events (World will atomically create hook entities)
  const hookEvents: CreateEventRequest[] = await Promise.all(
    hooksNeedingCreation.map(async (queueItem) => {
      const hookMetadata: SerializedData | undefined =
        typeof queueItem.metadata === 'undefined'
          ? undefined
          : ((await dehydrateStepArguments(
              queueItem.metadata,
              runId,
              encryptionKey,
              suspension.globalThis
            )) as SerializedData);
      return {
        eventType: 'hook_created' as const,
        specVersion: SPEC_VERSION_CURRENT,
        correlationId: queueItem.correlationId,
        eventData: {
          token: queueItem.token,
          metadata: hookMetadata,
        },
      };
    })
  );

  // Process hooks first to prevent race conditions with webhook receivers
  // All hook creations run in parallel
  // Track any hook conflicts that occur - these will be handled by re-enqueueing the workflow
  let hasHookConflict = false;

  if (hookEvents.length > 0) {
    await Promise.all(
      hookEvents.map(async (hookEvent) => {
        try {
          const result = await world.events.create(runId, hookEvent);
          // Check if the world returned a hook_conflict event instead of hook_created
          // The hook_conflict event is stored in the event log and will be replayed
          // on the next workflow invocation, causing the hook's promise to reject
          // Note: hook events always create an event (legacy runs throw, not return undefined)
          if (result.event!.eventType === 'hook_conflict') {
            hasHookConflict = true;
          }
        } catch (err) {
          if (WorkflowAPIError.is(err)) {
            if (err.status === 410) {
              runtimeLogger.info(
                'Workflow run already completed, skipping hook',
                {
                  workflowRunId: runId,
                  message: err.message,
                }
              );
            } else {
              throw err;
            }
          } else {
            throw err;
          }
        }
      })
    );
  }

  // Process hook disposals - these release hook tokens for reuse by other workflows
  if (hooksNeedingDisposal.length > 0) {
    await Promise.all(
      hooksNeedingDisposal.map(async (queueItem) => {
        const hookDisposedEvent: CreateEventRequest = {
          eventType: 'hook_disposed' as const,
          specVersion: SPEC_VERSION_CURRENT,
          correlationId: queueItem.correlationId,
        };
        try {
          await world.events.create(runId, hookDisposedEvent);
        } catch (err) {
          if (WorkflowAPIError.is(err)) {
            if (err.status === 410) {
              runtimeLogger.info(
                'Workflow run already completed, skipping hook disposal',
                {
                  workflowRunId: runId,
                  correlationId: queueItem.correlationId,
                  message: err.message,
                }
              );
            } else if (err.status === 404) {
              // Hook may have already been disposed or never created
              runtimeLogger.info('Hook not found for disposal, continuing', {
                workflowRunId: runId,
                correlationId: queueItem.correlationId,
                message: err.message,
              });
            } else {
              throw err;
            }
          } else {
            throw err;
          }
        }
      })
    );
  }

  // Build a map of stepId -> step event for steps that need creation
  const stepsNeedingCreation = new Set(
    stepItems
      .filter((queueItem) => !queueItem.hasCreatedEvent)
      .map((queueItem) => queueItem.correlationId)
  );

  // Process steps and waits in parallel
  // Each step: create event (if needed) -> queue message
  // Each wait: create event (if needed)
  const ops: Promise<void>[] = [];

  // Steps: create event then queue message, all in parallel
  for (const queueItem of stepItems) {
    ops.push(
      (async () => {
        // Create step event if not already created
        if (stepsNeedingCreation.has(queueItem.correlationId)) {
          const dehydratedInput = await dehydrateStepArguments(
            {
              args: queueItem.args,
              closureVars: queueItem.closureVars,
              thisVal: queueItem.thisVal,
            },
            runId,
            encryptionKey,
            suspension.globalThis
          );
          const stepEvent: CreateEventRequest = {
            eventType: 'step_created' as const,
            specVersion: SPEC_VERSION_CURRENT,
            correlationId: queueItem.correlationId,
            eventData: {
              stepName: queueItem.stepName,
              input: dehydratedInput as SerializedData,
            },
          };
          try {
            await world.events.create(runId, stepEvent);
          } catch (err) {
            if (WorkflowAPIError.is(err) && err.status === 409) {
              runtimeLogger.info('Step already exists, continuing', {
                workflowRunId: runId,
                correlationId: queueItem.correlationId,
                message: err.message,
              });
            } else {
              throw err;
            }
          }
        }

        // Queue step execution message
        // Serialize trace context once and include in both payload and headers
        // Payload: for manual context restoration in step handler
        // Headers: for automatic trace propagation by Vercel's infrastructure
        const traceCarrier = await serializeTraceCarrier();
        await queueMessage(
          world,
          `__wkf_step_${queueItem.stepName}`,
          {
            workflowName,
            workflowRunId: runId,
            workflowStartedAt,
            stepId: queueItem.correlationId,
            traceCarrier,
            requestedAt: new Date(),
          },
          {
            idempotencyKey: queueItem.correlationId,
            headers: {
              ...extractTraceHeaders(traceCarrier),
            },
          }
        );
      })()
    );
  }

  // Waits: create events in parallel (no queueing needed for waits)
  for (const queueItem of waitItems) {
    if (!queueItem.hasCreatedEvent) {
      ops.push(
        (async () => {
          const waitEvent: CreateEventRequest = {
            eventType: 'wait_created' as const,
            specVersion: SPEC_VERSION_CURRENT,
            correlationId: queueItem.correlationId,
            eventData: {
              resumeAt: queueItem.resumeAt,
            },
          };
          try {
            await world.events.create(runId, waitEvent);
          } catch (err) {
            if (WorkflowAPIError.is(err) && err.status === 409) {
              runtimeLogger.info('Wait already exists, continuing', {
                workflowRunId: runId,
                correlationId: queueItem.correlationId,
                message: err.message,
              });
            } else {
              throw err;
            }
          }
        })()
      );
    }
  }

  // Wait for all step and wait operations to complete
  waitUntil(
    Promise.all(ops).catch((opErr) => {
      const isAbortError =
        opErr?.name === 'AbortError' || opErr?.name === 'ResponseAborted';
      if (!isAbortError) throw opErr;
    })
  );
  await Promise.all(ops);

  // Calculate minimum timeout from waits
  const now = Date.now();
  const minTimeoutSeconds = waitItems.reduce<number | null>(
    (min, queueItem) => {
      const resumeAtMs = queueItem.resumeAt.getTime();
      const delayMs = Math.max(1000, resumeAtMs - now);
      const timeoutSeconds = Math.ceil(delayMs / 1000);
      if (min === null) return timeoutSeconds;
      return Math.min(min, timeoutSeconds);
    },
    null
  );

  span?.setAttributes({
    ...Attribute.WorkflowRunStatus('workflow_suspended'),
    ...Attribute.WorkflowStepsCreated(stepItems.length),
    ...Attribute.WorkflowHooksCreated(hooksNeedingCreation.length),
    ...Attribute.WorkflowWaitsCreated(waitItems.length),
  });

  // If any hook conflicts occurred, re-enqueue the workflow immediately
  // On the next iteration, the hook consumer will see the hook_conflict event
  // and reject the promise with a WorkflowRuntimeError
  // We do this after processing all other operations (steps, waits) to ensure
  // they are recorded in the event log before the re-execution
  if (hasHookConflict) {
    return { timeoutSeconds: 1 };
  }

  if (minTimeoutSeconds !== null) {
    return { timeoutSeconds: minTimeoutSeconds };
  }

  return {};
}