duron 0.2.1 → 0.3.0-beta.0

package/src/step-manager.ts

@@ -19,7 +19,9 @@ import {
   StepAlreadyExecutedError,
   StepTimeoutError,
   serializeError,
+  UnhandledChildStepsError,
 } from './errors.js'
+import type { ObserveContext, Span, TelemetryAdapter } from './telemetry/adapter.js'
 import pRetry from './utils/p-retry.js'
 import waitForAbort from './utils/wait-for-abort.js'
 
@@ -28,6 +30,8 @@ export interface TaskStep {
   cb: (ctx: StepHandlerContext) => Promise<any>
   options: StepOptions
   abortSignal: AbortSignal
+  parentStepId: string | null
+  parallel: boolean
 }
 
 /**
@@ -61,16 +65,27 @@ export class StepStore {
    * @param name - The name of the step
    * @param timeoutMs - Timeout in milliseconds for the step
    * @param retriesLimit - Maximum number of retries for the step
+   * @param parentStepId - The ID of the parent step (null for root steps)
+   * @param parallel - Whether this step runs in parallel (independent from siblings during time travel)
    * @returns Promise resolving to the created step ID
    * @throws Error if step creation fails
    */
-  async getOrCreate(jobId: string, name: string, timeoutMs: number, retriesLimit: number) {
+  async getOrCreate(
+    jobId: string,
+    name: string,
+    timeoutMs: number,
+    retriesLimit: number,
+    parentStepId: string | null = null,
+    parallel: boolean = false,
+  ) {
     try {
       return await this.#adapter.createOrRecoverJobStep({
         jobId,
         name,
         timeoutMs,
         retriesLimit,
+        parentStepId,
+        parallel,
       })
     } catch (error) {
       throw new NonRetriableError(`Failed to get or create step "${name}" for job "${jobId}"`, { cause: error })
@@ -115,6 +130,7 @@ export interface StepManagerOptions {
   jobId: string
   actionName: string
   adapter: Adapter
+  telemetry: TelemetryAdapter
   logger: Logger
   concurrencyLimit: number
 }
@@ -127,10 +143,15 @@ export class StepManager {
   #jobId: string
   #actionName: string
   #stepStore: StepStore
+  #telemetry: TelemetryAdapter
   #queue: fastq.queueAsPromised<TaskStep, any>
   #logger: Logger
   // each step name should be executed only once per action job
   #historySteps = new Set<string>()
+  // Store step spans for nested step tracking
+  #stepSpans = new Map<string, Span>()
+  // Store the job span for creating step spans
+  #jobSpan: Span | null = null
 
   // ============================================================================
   // Constructor
@@ -145,16 +166,25 @@ export class StepManager {
     this.#jobId = options.jobId
     this.#actionName = options.actionName
     this.#logger = options.logger
+    this.#telemetry = options.telemetry
     this.#stepStore = new StepStore(options.adapter)
     this.#queue = fastq.promise(async (task: TaskStep) => {
       if (this.#historySteps.has(task.name)) {
         throw new StepAlreadyExecutedError(task.name, this.#jobId, this.#actionName)
       }
       this.#historySteps.add(task.name)
-      return this.#executeStep(task.name, task.cb, task.options, task.abortSignal)
+      return this.#executeStep(task.name, task.cb, task.options, task.abortSignal, task.parentStepId, task.parallel)
     }, options.concurrencyLimit)
   }
 
+  /**
+   * Set the job span for this step manager.
+   * Called from ActionJob after the job span is created.
+   */
+  setJobSpan(span: Span): void {
+    this.#jobSpan = span
+  }
+
   // ============================================================================
   // Public API Methods
   // ============================================================================
@@ -168,6 +198,7 @@ export class StepManager {
    * @param variables - Variables available to the action
    * @param abortSignal - Abort signal for cancelling the action
    * @param logger - Pino child logger for this job
+   * @param observeContext - Observability context for telemetry
    * @returns ActionHandlerContext instance
    */
   createActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVariables = Record<string, unknown>>(
@@ -176,8 +207,35 @@ export class StepManager {
     variables: TVariables,
     abortSignal: AbortSignal,
     logger: Logger,
+    observeContext: ObserveContext,
   ): ActionHandlerContext<TInput, TVariables> {
-    return new ActionContext(this, job, action, variables, abortSignal, logger)
+    return new ActionContext(this, job, action, variables, abortSignal, logger, observeContext)
+  }
+
+  /**
+   * Create an observe context for a step.
+   */
+  createStepObserveContext(stepId: string): ObserveContext {
+    const stepSpan = this.#stepSpans.get(stepId)
+    if (stepSpan) {
+      return this.#telemetry.createObserveContext(this.#jobId, stepId, stepSpan)
+    }
+    // Fallback to job span if step span not found
+    if (this.#jobSpan) {
+      return this.#telemetry.createObserveContext(this.#jobId, stepId, this.#jobSpan)
+    }
+    // No-op observe context
+    return {
+      recordMetric: () => {
+        // No-op
+      },
+      addSpanAttribute: () => {
+        // No-op
+      },
+      addSpanEvent: () => {
+        // No-op
+      },
+    }
   }
 
   /**
@@ -206,9 +264,11 @@ export class StepManager {
    * @param cb - The step handler function
    * @param options - Step options including concurrency, retry, and expire settings
    * @param abortSignal - Abort signal for cancelling the step
+   * @param parentStepId - The ID of the parent step (null for root steps)
    * @returns Promise resolving to the step result
    * @throws StepTimeoutError if the step times out
    * @throws StepCancelError if the step is cancelled
+   * @throws UnhandledChildStepsError if child steps are not awaited
    * @throws Error if the step fails
    */
   async #executeStep<TResult>(
@@ -216,6 +276,8 @@ export class StepManager {
     cb: (ctx: StepHandlerContext) => Promise<TResult>,
     options: StepOptions,
     abortSignal: AbortSignal,
+    parentStepId: string | null,
+    parallel: boolean,
   ): Promise<TResult> {
     const expire = options.expire
     const retryOptions = options.retry
@@ -227,8 +289,15 @@ export class StepManager {
           throw new ActionCancelError(this.#actionName, this.#jobId, { cause: 'step cancelled before create step' })
         }
 
-        // Create step record
-        const newStep = await this.#stepStore.getOrCreate(this.#jobId, name, expire, retryOptions.limit)
+        // Create step record with parentStepId and parallel
+        const newStep = await this.#stepStore.getOrCreate(
+          this.#jobId,
+          name,
+          expire,
+          retryOptions.limit,
+          parentStepId,
+          parallel,
+        )
         if (!newStep) {
           throw new NonRetriableError(
             `Failed to create step "${name}" for job "${this.#jobId}" action "${this.#actionName}"`,
@@ -238,6 +307,17 @@ export class StepManager {
 
         step = newStep
 
+        // Start step telemetry span
+        const parentSpan = parentStepId ? this.#stepSpans.get(parentStepId) : this.#jobSpan
+        const stepSpan = await this.#telemetry.startStepSpan({
+          jobId: this.#jobId,
+          stepId: step.id,
+          stepName: name,
+          parentSpan: parentSpan ?? undefined,
+          parentStepId,
+        })
+        this.#stepSpans.set(step.id, stepSpan)
+
         if (abortSignal.aborted) {
           throw new ActionCancelError(this.#actionName, this.#jobId, { cause: 'step cancelled after create step' })
         }
@@ -245,7 +325,7 @@ export class StepManager {
         if (step.status === STEP_STATUS_COMPLETED) {
           // this is how we recover a completed step
           this.#logger.debug(
-            { jobId: this.#jobId, actionName: this.#actionName, stepName: name, stepId: step.id },
+            { jobId: this.#jobId, actionName: this.#actionName, stepName: name, stepId: step.id, parentStepId },
             '[StepManager] Step recovered (already completed)',
           )
           return step.output as TResult
@@ -265,11 +345,12 @@ export class StepManager {
 
         // Log step start
         this.#logger.debug(
-          { jobId: this.#jobId, actionName: this.#actionName, stepName: name, stepId: step.id },
+          { jobId: this.#jobId, actionName: this.#actionName, stepName: name, stepId: step.id, parentStepId },
           '[StepManager] Step started executing',
         )
       }
 
+      // Create abort controller for this step's timeout
       const stepAbortController = new AbortController()
       const timeoutId = setTimeout(() => {
         const timeoutError = new StepTimeoutError(name, this.#jobId, expire)
@@ -278,18 +359,85 @@ export class StepManager {
 
       timeoutId?.unref?.()
 
-      // Combine abort signals
-      const signal = AbortSignal.any([abortSignal, stepAbortController.signal])
+      // Combine abort signals: parent chain + this step's timeout
+      const stepSignal = AbortSignal.any([abortSignal, stepAbortController.signal])
+
+      // Track child steps for enforcement
+      interface TrackedChildStep {
+        promise: Promise<any>
+        settled: boolean
+      }
+      const childSteps: TrackedChildStep[] = []
+
+      // Create abort controller for child steps (used when parent returns with pending children)
+      const childAbortController = new AbortController()
+      const childSignal = AbortSignal.any([stepSignal, childAbortController.signal])
+
+      // Create observe context for this step
+      const stepObserveContext = this.createStepObserveContext(step.id)
+
+      // Create StepHandlerContext with nested step support
+      const stepContext: StepHandlerContext = {
+        signal: stepSignal,
+        stepId: step.id,
+        parentStepId,
+        observe: stepObserveContext,
+        step: <TChildResult>(
+          childName: string,
+          childCb: (ctx: StepHandlerContext) => Promise<TChildResult>,
+          childOptions: z.input<typeof StepOptionsSchema> = {},
+        ): Promise<TChildResult> => {
+          // Inherit parent step options EXCEPT parallel (each step's parallel status is independent)
+          const { parallel: _parentParallel, ...inheritableOptions } = options
+          const parsedChildOptions = StepOptionsSchema.parse({
+            ...inheritableOptions,
+            ...childOptions,
+          })
+
+          // Push child step with this step as parent
+          const childPromise = this.push({
+            name: childName,
+            cb: childCb,
+            options: parsedChildOptions,
+            abortSignal: childSignal, // Child uses composed signal
+            parentStepId: step!.id, // This step is the parent
+            parallel: parsedChildOptions.parallel, // Pass parallel option
+          })
+
+          // Track the child promise
+          const trackedChild: TrackedChildStep = {
+            promise: childPromise,
+            settled: false,
+          }
+          childSteps.push(trackedChild)
+
+          // Mark as settled when done (success or failure)
+          // Use .then/.catch instead of .finally to properly handle rejections
+          childPromise
+            .then(() => {
+              trackedChild.settled = true
+            })
+            .catch(() => {
+              trackedChild.settled = true
+              // Swallow the error here - it will be re-thrown to the caller via the returned promise
+            })
+
+          return childPromise
+        },
+      }
 
       try {
         // Race between abort signal and callback execution
-        const abortPromise = waitForAbort(signal)
-        const callbackPromise = cb({ signal })
+        const abortPromise = waitForAbort(stepSignal)
+        const callbackPromise = cb(stepContext)
 
         let result: any = null
+        let aborted = false
 
         await Promise.race([
-          abortPromise.promise,
+          abortPromise.promise.then(() => {
+            aborted = true
+          }),
           callbackPromise
             .then((res) => {
               if (res !== undefined && res !== null) {
@@ -301,12 +449,54 @@ export class StepManager {
             }),
         ])
 
+        // If aborted, wait for child steps to settle before propagating
+        if (aborted) {
+          // Wait for all child steps to settle (they'll be aborted via signal propagation)
+          if (childSteps.length > 0) {
+            await Promise.allSettled(childSteps.map((c) => c.promise))
+          }
+          // Re-throw the abort reason
+          throw stepSignal.reason
+        }
+
+        // After parent callback returns, check for pending children
+        const unsettledChildren = childSteps.filter((c) => !c.settled)
+        if (unsettledChildren.length > 0) {
+          this.#logger.warn(
+            {
+              jobId: this.#jobId,
+              actionName: this.#actionName,
+              stepName: name,
+              stepId: step.id,
+              pendingCount: unsettledChildren.length,
+            },
+            '[StepManager] Parent step completed with unhandled child steps - aborting children',
+          )
+
+          // Abort all pending children
+          const unhandledError = new UnhandledChildStepsError(name, unsettledChildren.length)
+          childAbortController.abort(unhandledError)
+
+          // Wait for all children to settle (they'll reject with cancellation)
+          await Promise.allSettled(unsettledChildren.map((c) => c.promise))
+
+          // Now throw the error
+          throw unhandledError
+        }
+
         // Update step as completed
         const completed = await this.#stepStore.updateStatus(step.id, 'completed', result)
         if (!completed) {
           throw new Error(`Failed to complete step "${name}" for job "${this.#jobId}" action "${this.#actionName}"`)
         }
 
+        // End step telemetry span successfully
+        const stepSpan = this.#stepSpans.get(step.id)
+        if (stepSpan) {
+          await this.#telemetry.endStepSpan(stepSpan, { status: 'ok' })
+          this.#stepSpans.delete(step.id)
+        }
+
         // Log step completion
         this.#logger.debug(
           { jobId: this.#jobId, actionName: this.#actionName, stepName: name, stepId: step.id },
@@ -347,6 +537,17 @@ export class StepManager {
       },
     }).catch(async (error) => {
       if (step) {
+        // End step telemetry span with error/cancelled status
+        const stepSpan = this.#stepSpans.get(step.id)
+        if (stepSpan) {
+          if (isCancelError(error)) {
+            await this.#telemetry.endStepSpan(stepSpan, { status: 'cancelled' })
+          } else {
+            await this.#telemetry.endStepSpan(stepSpan, { status: 'error', error })
+          }
+          this.#stepSpans.delete(step.id)
+        }
+
         if (isCancelError(error)) {
           await this.#stepStore.updateStatus(step.id, 'cancelled')
         } else {
@@ -377,6 +578,7 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
   #jobId: string
   #groupKey: string = '@default'
   #action: Action<TInput, TOutput, TVariables>
+  #observeContext: ObserveContext
 
   // ============================================================================
   // Constructor
@@ -389,6 +591,7 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
     variables: TVariables,
     abortSignal: AbortSignal,
     logger: Logger,
+    observeContext: ObserveContext,
   ) {
     this.#stepManager = stepManager
     this.#variables = variables
@@ -397,6 +600,7 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
     this.#action = action
     this.#jobId = job.id
     this.#groupKey = job.groupKey ?? '@default'
+    this.#observeContext = observeContext
     if (action.input) {
       this.#input = action.input.parse(job.input, {
         error: () => 'Error parsing action input',
@@ -450,8 +654,16 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
     return this.#logger
   }
 
+  /**
+   * Get the observability context for recording metrics and span data.
+   */
+  get observe(): ObserveContext {
+    return this.#observeContext
+  }
+
   /**
    * Execute a step within the action.
+   * This creates a root step (no parent).
    *
    * @param name - The name of the step
    * @param cb - The step handler function
@@ -467,6 +679,13 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
       ...this.#action.steps,
       ...options,
     })
-    return this.#stepManager.push({ name, cb, options: parsedOptions, abortSignal: this.#abortSignal })
+    return this.#stepManager.push({
+      name,
+      cb,
+      options: parsedOptions,
+      abortSignal: this.#abortSignal,
+      parentStepId: null, // Root steps have no parent
+      parallel: parsedOptions.parallel, // Pass parallel option
+    })
   }
 }