duron 0.3.0-beta.9 → 0.3.0

package/src/step-manager.ts

@@ -1,3 +1,13 @@
+import {
+  type Context,
+  context,
+  type Span,
+  SpanKind,
+  type SpanOptions,
+  SpanStatusCode,
+  type Tracer,
+  trace,
+} from '@opentelemetry/api'
 import fastq from 'fastq'
 import type { Logger } from 'pino'
 import type { z } from 'zod'
@@ -24,45 +34,118 @@ import {
   serializeError,
   UnhandledChildStepsError,
 } from './errors.js'
-import type { ObserveContext, Span, TelemetryAdapter, Tracer, TracerSpan } from './telemetry/adapter.js'
 
-// ============================================================================
-// Noop Tracer (for fallback observe context)
-// ============================================================================
+/**
+ * Telemetry context provided to action and step handlers.
+ * Provides access to OpenTelemetry APIs for recording traces and metrics.
+ */
+export interface TelemetryContext {
+  /**
+   * Get the active OpenTelemetry span for the current job/step.
+   */
+  getActiveSpan(): Span | undefined
 
-const noopTracerSpan: TracerSpan = {
-  setAttribute() {
-    // No-op
-  },
-  setAttributes() {
-    // No-op
-  },
-  addEvent() {
-    // No-op
-  },
-  recordException() {
-    // No-op
-  },
-  setStatusOk() {
-    // No-op
-  },
-  setStatusError() {
-    // No-op
-  },
-  end() {
-    // No-op
-  },
-  isRecording() {
-    return false
-  },
+  /**
+   * Get an OpenTelemetry tracer for creating custom spans.
+   */
+  getTracer(name: string): Tracer
+
+  /**
+   * Start a new span as a child of the current job/step span.
+   * This is a convenience method that properly links the span to the trace hierarchy.
+   */
+  startSpan(name: string, options?: { attributes?: Record<string, any> }): Span
+
+  /**
+   * Record a custom metric as a span event.
+   */
+  recordMetric(name: string, value: number, attributes?: Record<string, any>): void
+}
+
+/**
+ * Inject parent span into a context if we have one.
+ */
+function injectParentSpan(ctx: Context, parentSpan: Span | null): Context {
+  return parentSpan ? trace.setSpan(ctx, parentSpan) : ctx
+}
+
+/**
+ * Create a context-aware tracer wrapper that automatically injects the parent span.
+ * This ensures spans created by external libraries (like AI SDK) are properly linked
+ * to the current job/step trace hierarchy.
+ */
+function createContextAwareTracer(tracer: Tracer, parentSpan: Span | null): Tracer {
+  return {
+    startSpan(name: string, options?: SpanOptions, ctx?: Context): Span {
+      // Always inject our parent span into the context, regardless of what context is passed.
+      // This is necessary because without global registration, context.active() returns
+      // ROOT_CONTEXT, so external libraries (like AI SDK) that pass context.active()
+      // would otherwise create orphan spans.
+      const baseContext = ctx ?? context.active()
+      const effectiveContext = injectParentSpan(baseContext, parentSpan)
+      return tracer.startSpan(name, options, effectiveContext)
+    },
+    // startActiveSpan has multiple overloads, we need to handle them all
+    startActiveSpan<F extends (span: Span) => unknown>(
+      name: string,
+      optionsOrFn: SpanOptions | F,
+      ctxOrFn?: Context | F,
+      fn?: F,
+    ): ReturnType<F> {
+      // Parse the overloaded arguments
+      let options: SpanOptions | undefined
+      let ctx: Context | undefined
+      let callback: F
+
+      if (typeof optionsOrFn === 'function') {
+        // startActiveSpan(name, fn)
+        callback = optionsOrFn
+      } else if (typeof ctxOrFn === 'function') {
+        // startActiveSpan(name, options, fn)
+        options = optionsOrFn
+        callback = ctxOrFn
+      } else {
+        // startActiveSpan(name, options, context, fn)
+        options = optionsOrFn
+        ctx = ctxOrFn
+        callback = fn!
+      }
+
+      const baseContext = ctx ?? context.active()
+      const effectiveContext = injectParentSpan(baseContext, parentSpan)
+
+      return tracer.startActiveSpan(name, options ?? {}, effectiveContext, callback)
+    },
+  } as Tracer
 }
 
-const createNoopTracer = (name: string): Tracer => ({
-  name,
-  startSpan(): TracerSpan {
-    return noopTracerSpan
-  },
-})
+/**
+ * Create a TelemetryContext that wraps an OTel span.
+ */
+function createTelemetryContext(span: Span | null, tracer: Tracer): TelemetryContext {
+  return {
+    getActiveSpan(): Span | undefined {
+      return span ?? undefined
+    },
+    getTracer(_name: string): Tracer {
+      // Return a context-aware tracer that automatically links spans to the current trace
+      return createContextAwareTracer(tracer, span)
+    },
+    startSpan(name: string, options?: { attributes?: Record<string, any> }): Span {
+      // Create a child span linked to the current span (job or step)
+      const parentContext = span ? trace.setSpan(context.active(), span) : context.active()
+      return tracer.startSpan(name, { attributes: options?.attributes }, parentContext)
+    },
+    recordMetric(name: string, value: number, attributes?: Record<string, any>): void {
+      if (span) {
+        span.addEvent(`metric:${name}`, {
+          'metric.value': value,
+          ...attributes,
+        })
+      }
+    },
+  }
+}
 
 import pRetry from './utils/p-retry.js'
 import waitForAbort from './utils/wait-for-abort.js'
@@ -172,7 +255,7 @@ export interface StepManagerOptions {
   jobId: string
   actionName: string
   adapter: Adapter
-  telemetry: TelemetryAdapter
+  tracer: Tracer
   logger: Logger
   concurrencyLimit: number
 }
@@ -185,7 +268,7 @@ export class StepManager {
   #jobId: string
   #actionName: string
   #stepStore: StepStore
-  #telemetry: TelemetryAdapter
+  #tracer: Tracer
   #queue: fastq.queueAsPromised<TaskStep, any>
   #logger: Logger
   // each step name should be executed only once per parent (name + parentStepId)
@@ -193,7 +276,7 @@ export class StepManager {
   // Store step spans for nested step tracking
   #stepSpans = new Map<string, Span>()
   // Store the job span for creating step spans
-  #jobSpan: Span | null = null
+  #jobSpan!: Span
   // Factory function to create run functions with the correct parent step ID and abort signal
   #runFnFactory: ((parentStepId: string | null, abortSignal: AbortSignal) => StepHandlerContext['run']) | null = null
 
@@ -210,7 +293,7 @@ export class StepManager {
     this.#jobId = options.jobId
     this.#actionName = options.actionName
     this.#logger = options.logger
-    this.#telemetry = options.telemetry
+    this.#tracer = options.tracer
     this.#stepStore = new StepStore(options.adapter)
     this.#queue = fastq.promise(async (task: TaskStep) => {
       // Create composite key: name + parentStepId (allows same name under different parents)
@@ -254,7 +337,6 @@ 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>>(
@@ -263,38 +345,21 @@ export class StepManager {
     variables: TVariables,
     abortSignal: AbortSignal,
     logger: Logger,
-    observeContext: ObserveContext,
   ): ActionHandlerContext<TInput, TVariables> {
-    return new ActionContext(this, job, action, variables, abortSignal, logger, observeContext)
+    const telemetryContext = createTelemetryContext(this.#jobSpan, this.#tracer)
+    return new ActionContext(this, job, action, variables, abortSignal, logger, telemetryContext)
   }
 
   /**
-   * Create an observe context for a step.
+   * Create a telemetry context for a step.
    */
-  createStepObserveContext(stepId: string): ObserveContext {
+  createStepTelemetryContext(stepId: string): TelemetryContext {
     const stepSpan = this.#stepSpans.get(stepId)
     if (stepSpan) {
-      return this.#telemetry.createObserveContext(this.#jobId, stepId, stepSpan)
+      return createTelemetryContext(stepSpan, this.#tracer)
     }
     // 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
-      },
-      getTracer: (name: string) => {
-        return createNoopTracer(name)
-      },
-    }
+    return createTelemetryContext(this.#jobSpan, this.#tracer)
   }
 
   /**
@@ -381,15 +446,22 @@ export class StepManager {
 
         step = newStep
 
-        // Start step telemetry span
+        // Start step span - uses no-op tracer if no SDK is configured
         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,
-        })
+        const parentContext = parentSpan ? trace.setSpan(context.active(), parentSpan) : context.active()
+        const stepSpan = this.#tracer.startSpan(
+          `step:${name}`,
+          {
+            kind: SpanKind.INTERNAL,
+            attributes: {
+              'duron.job.id': this.#jobId,
+              'duron.step.id': step.id,
+              'duron.step.name': name,
+              'duron.step.parent_id': parentStepId ?? undefined,
+            },
+          },
+          parentContext,
+        )
         this.#stepSpans.set(step.id, stepSpan)
 
         if (abortSignal.aborted) {
@@ -427,7 +499,11 @@ export class StepManager {
       // Create abort controller for this step's timeout
       const stepAbortController = new AbortController()
       const timeoutId = setTimeout(() => {
-        const timeoutError = new StepTimeoutError(name, this.#jobId, expire)
+        const timeoutError = new StepTimeoutError(name, this.#jobId, expire, {
+          stepId: step?.id,
+          parentStepId,
+          actionName: this.#actionName,
+        })
         stepAbortController.abort(timeoutError)
       }, expire)
 
@@ -447,15 +523,15 @@ export class StepManager {
       const childAbortController = new AbortController()
       const childSignal = AbortSignal.any([stepSignal, childAbortController.signal])
 
-      // Create observe context for this step
-      const stepObserveContext = this.createStepObserveContext(step.id)
+      // Create telemetry context for this step
+      const stepTelemetryContext = this.createStepTelemetryContext(step.id)
 
       // Create StepHandlerContext with nested step support
       const stepContext: StepHandlerContext = {
         signal: stepSignal,
         stepId: step.id,
         parentStepId,
-        observe: stepObserveContext,
+        telemetry: stepTelemetryContext,
         step: <TChildResult>(
           childName: string,
           childCb: (ctx: StepHandlerContext) => Promise<TChildResult>,
@@ -505,10 +581,15 @@ export class StepManager {
       try {
         // Race between abort signal and callback execution
         const abortPromise = waitForAbort(stepSignal)
-        const callbackPromise = cb(stepContext)
+
+        // Execute callback within the span context so that child spans inherit the trace
+        const currentStepSpan = step?.id ? this.#stepSpans.get(step.id) : undefined
+        const spanContext = currentStepSpan ? trace.setSpan(context.active(), currentStepSpan) : context.active()
+        const callbackPromise = context.with(spanContext, () => cb(stepContext))
 
         let result: any = null
         let aborted = false
+        let callbackError: Error | null = null
 
         await Promise.race([
           abortPromise.promise.then(() => {
@@ -520,11 +601,25 @@ export class StepManager {
                 result = res
               }
             })
+            .catch((err) => {
+              callbackError = err
+            })
             .finally(() => {
               abortPromise.release()
             }),
         ])
 
+        // If callback threw an error, abort children and wait for them before re-throwing
+        if (callbackError) {
+          if (childSteps.length > 0) {
+            // Abort all children with the callback error as reason
+            childAbortController.abort(callbackError)
+            // Wait for all children to settle
+            await Promise.allSettled(childSteps.map((c) => c.promise))
+          }
+          throw callbackError
+        }
+
         // 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)
@@ -550,7 +645,12 @@ export class StepManager {
           )
 
           // Abort all pending children
-          const unhandledError = new UnhandledChildStepsError(name, unsettledChildren.length)
+          const unhandledError = new UnhandledChildStepsError(name, unsettledChildren.length, {
+            stepId: step.id,
+            parentStepId,
+            jobId: this.#jobId,
+            actionName: this.#actionName,
+          })
           childAbortController.abort(unhandledError)
 
           // Wait for all children to settle (they'll reject with cancellation)
@@ -566,10 +666,11 @@ export class StepManager {
           throw new Error(`Failed to complete step "${name}" for job "${this.#jobId}" action "${this.#actionName}"`)
         }
 
-        // End step telemetry span successfully
+        // End step span successfully
         const stepSpan = this.#stepSpans.get(step.id)
         if (stepSpan) {
-          await this.#telemetry.endStepSpan(stepSpan, { status: 'ok' })
+          stepSpan.setStatus({ code: SpanStatusCode.OK })
+          stepSpan.end()
           this.#stepSpans.delete(step.id)
         }
 
@@ -599,12 +700,27 @@ export class StepManager {
         if (
           isNonRetriableError(error) ||
           (error.cause && isNonRetriableError(error.cause)) ||
-          (error instanceof Error && error.name === 'AbortError' && isNonRetriableError(error.cause))
+          (error instanceof Error && error.name === 'AbortError')
         ) {
-          throw error
+          const err = isNonRetriableError(error)
+            ? error
+            : error instanceof Error && error.name === 'AbortError'
+              ? new NonRetriableError(error.message, { cause: error.cause })
+              : (error.cause as NonRetriableError)
+
+          if (Object.keys(err.metadata).length === 0) {
+            err.setMetadata({
+              stepId: step?.id,
+              parentStepId,
+              jobId: this.#jobId,
+              actionName: this.#actionName,
+            })
+          }
+          throw err
         }
 
         if (ctx.retriesLeft > 0 && step) {
+          this.#clearHistoryForStep(step.id)
           const delayed = await this.#stepStore.delay(step.id, ctx.finalDelay, serializeError(error))
           if (!delayed) {
             throw new Error(`Failed to delay step "${name}" for job "${this.#jobId}" action "${this.#actionName}"`)
@@ -614,22 +730,35 @@ export class StepManager {
             ;(error as any).nonRetriable = true
             throw error
           }
-          throw new NonRetriableError(
-            `Failed to execute step="${name}", jobId="${this.#jobId}", action="${this.#actionName}", parentStepId="${parentStepId}"`,
-            { cause: error },
-          )
+
+          const errorMessage = error instanceof Error ? error.message : String(error)
+          const err = new NonRetriableError(errorMessage, { cause: error })
+          err.setMetadata({
+            stepId: step?.id,
+            parentStepId,
+            jobId: this.#jobId,
+            actionName: this.#actionName,
+          })
+          throw err
         }
       },
     }).catch(async (error) => {
       if (step) {
-        // End step telemetry span with error/cancelled status
+        // End step span with error/cancelled status
         const stepSpan = this.#stepSpans.get(step.id)
         if (stepSpan) {
           if (isCancelError(error)) {
-            await this.#telemetry.endStepSpan(stepSpan, { status: 'cancelled' })
+            stepSpan.setStatus({ code: SpanStatusCode.ERROR, message: 'Step cancelled' })
           } else {
-            await this.#telemetry.endStepSpan(stepSpan, { status: 'error', error })
+            stepSpan.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: error instanceof Error ? error.message : String(error),
+            })
+            if (error instanceof Error) {
+              stepSpan.recordException(error)
+            }
           }
+          stepSpan.end()
           this.#stepSpans.delete(step.id)
         }
 
@@ -642,6 +771,19 @@ export class StepManager {
       throw error
     })
   }
+
+  /**
+   * Clear the history of nested steps for a given step.
+   * We do't need to clear the history for the root step because it's not a parent step, it's the action itself.
+   * @param stepId - The ID of the step to clear the history for
+   */
+  #clearHistoryForStep(stepId: string): void {
+    this.#historySteps.forEach((stepKey) => {
+      if (stepKey.startsWith(stepId)) {
+        this.#historySteps.delete(stepKey)
+      }
+    })
+  }
 }
 
 // ============================================================================
@@ -663,7 +805,7 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
   #jobId: string
   #groupKey: string = '@default'
   #action: Action<TInput, TOutput, TVariables>
-  #observeContext: ObserveContext
+  #telemetryContext: TelemetryContext
 
   // ============================================================================
   // Constructor
@@ -676,7 +818,7 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
     variables: TVariables,
     abortSignal: AbortSignal,
     logger: Logger,
-    observeContext: ObserveContext,
+    telemetryContext: TelemetryContext,
   ) {
     this.#stepManager = stepManager
     this.#variables = variables
@@ -685,7 +827,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
+    this.#telemetryContext = telemetryContext
     if (action.input) {
       this.#input = action.input.parse(job.input, {
         error: () => 'Error parsing action input',
@@ -745,10 +887,10 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
   }
 
   /**
-   * Get the observability context for recording metrics and span data.
+   * Get the telemetry context for recording metrics and span data.
    */
-  get observe(): ObserveContext {
-    return this.#observeContext
+  get telemetry(): TelemetryContext {
+    return this.#telemetryContext
   }
 
   /**
@@ -818,7 +960,16 @@ class ActionContext<TInput extends z.ZodObject, TOutput extends z.ZodObject, TVa
       : (input as z.infer<TStepInput>)
 
     // Resolve step name (static or dynamic)
-    const stepName = typeof stepDef.name === 'function' ? stepDef.name({ input: validatedInput }) : stepDef.name
+    // If it's a function, pass the full context including input, variables, jobId, and parentStepId
+    const stepName =
+      typeof stepDef.name === 'function'
+        ? stepDef.name({
+            input: validatedInput,
+            var: this.#variables,
+            jobId: this.#jobId,
+            parentStepId,
+          })
+        : stepDef.name
 
     // Merge options: action defaults -> step definition -> call-time overrides
     const mergedOptions: z.input<typeof StepOptionsSchema> = {

package/src/telemetry/index.ts

@@ -1,20 +1,2 @@
-// Re-export telemetry adapters and types
-
-export {
-  type AddSpanAttributeOptions,
-  type AddSpanEventOptions,
-  type EndSpanOptions,
-  type ObserveContext,
-  type RecordMetricOptions,
-  type Span,
-  type StartDatabaseSpanOptions,
-  type StartJobSpanOptions,
-  type StartSpanOptions,
-  type StartStepSpanOptions,
-  TelemetryAdapter,
-  type Tracer,
-  type TracerSpan,
-} from './adapter.js'
-export { LocalTelemetryAdapter, type LocalTelemetryAdapterOptions, localTelemetryAdapter } from './local.js'
-export { NoopTelemetryAdapter, noopTelemetryAdapter } from './noop.js'
-export { OpenTelemetryAdapter, type OpenTelemetryAdapterOptions, openTelemetryAdapter } from './opentelemetry.js'
+// OpenTelemetry-based span exporter for local database persistence
+export { LocalSpanExporter, type LocalSpanExporterOptions } from './local-span-exporter.js'