@kronos-ts/messaging 0.5.1 → 0.7.0

package/src/tracing-command-bus.ts

@@ -1,18 +1,26 @@
 import type { CommandBus } from "./command-bus.js"
 import type { CommandMessage } from "./message.js"
-import type { SpanFactory } from "./span-factory.js"
+import type { SpanFactory, Span } from "./span-factory.js"
+
+/** Run `fn` with `span` active, falling back to a plain call when the span lacks runActive. */
+function runActive<R>(span: Span, fn: () => R): R {
+  return span.runActive ? span.runActive(fn) : fn()
+}
 
 /**
- * A {@link CommandBus} decorator that wraps dispatch and handler invocations
- * with tracing spans.
+ * A {@link CommandBus} decorator that traces command dispatch (the producer
+ * side).
  *
  * Dispatch creates a "dispatch" span and propagates trace context into the
- * message metadata. Subscribe wraps each handler with a "handle" span.
- *
+ * message metadata, so the handler links back to it across the bus boundary.
+ * The handler ("handle") span is created by {@link tracingHandlerEnhancerDefinition},
+ * the single authority for handler-side spans across command/query/event
+ * handlers — so this decorator does not wrap subscribe, avoiding a duplicate
+ * command handle span.
  *
  * @param delegate  The underlying command bus to decorate.
  * @param spanFactory  The span factory for creating tracing spans.
- * @returns A decorated command bus with tracing instrumentation.
+ * @returns A decorated command bus with dispatch tracing.
  */
 export function createTracingCommandBus(
   delegate: CommandBus,
@@ -22,8 +30,12 @@ export function createTracingCommandBus(
     async dispatch(message: CommandMessage): Promise<unknown> {
       const span = spanFactory.createDispatchSpan(`dispatch(${String(message.name)})`, message).start()
       try {
-        const propagated = spanFactory.propagateContext(message)
-        const result = await delegate.dispatch(propagated)
+        // Propagate inside the active dispatch span so the outgoing message
+        // carries this span's trace context and the handler links to it.
+        const result = await runActive(span, () => {
+          const propagated = spanFactory.propagateContext(message)
+          return delegate.dispatch(propagated)
+        })
         span.end()
         return result
       } catch (err) {
@@ -36,17 +48,9 @@ export function createTracingCommandBus(
       commandName: string,
       handler: (message: CommandMessage) => Promise<unknown>,
     ): void {
-      delegate.subscribe(commandName, async (msg: CommandMessage) => {
-        const span = spanFactory.createHandlerSpan(`handle(${commandName})`, msg).start()
-        try {
-          const result = await handler(msg)
-          span.end()
-          return result
-        } catch (err) {
-          span.recordException(err instanceof Error ? err : new Error(String(err)))
-          throw err
-        }
-      })
+      // Handler-side spans are owned by tracingHandlerEnhancerDefinition; pass
+      // the handler through untouched so commands get exactly one handle span.
+      delegate.subscribe(commandName, handler)
     },
   }
 }

package/src/tracing-handler-enhancer.ts

@@ -1,12 +1,20 @@
 import type { HandlerEnhancerDefinition, HandlerMetadata } from "./handler-enhancer.js"
-import type { SpanFactory } from "./span-factory.js"
+import type { SpanFactory, Span } from "./span-factory.js"
+import type { Message } from "./message.js"
+import { contributeCorrelationData } from "./correlation-data.js"
 
 /**
  * Handler enhancer that wraps message handler invocations with tracing spans.
  *
- * Creates an internal span per handler invocation, recording the handler
- * name and message type as context. Errors are recorded on the span.
+ * The span is created from the message being handled (extracting any trace
+ * context from its metadata) so the handler re-parents onto the dispatcher's
+ * trace across the message boundary. Event handlers start a new trace linked to
+ * the triggering event; command/query handlers continue the current trace.
  *
+ * The handler runs inside the span's active context, and the active trace
+ * context is captured onto the UnitOfWork (via contributeCorrelationData) so
+ * appended and dispatched messages carry it — including events published at
+ * commit time, after the span has ended.
  */
 export function tracingHandlerEnhancerDefinition(
   spanFactory: SpanFactory,
@@ -19,9 +27,27 @@ export function tracingHandlerEnhancerDefinition(
       const spanName = `${metadata.handlerGroup}.${metadata.messageName}`
 
       return (async (...args: any[]) => {
-        const span = spanFactory.createInternalSpan(spanName).start()
+        const message = args[0]
+        const span = createSpan(spanFactory, spanName, message, metadata).start()
+        const runActive: <R>(fn: () => R) => R = span.runActive
+          ? span.runActive.bind(span)
+          : (fn) => fn()
+
         try {
-          const result = await handler(...args)
+          const result = await runActive(() => {
+            // Store the active trace context on the UnitOfWork so outgoing and
+            // appended messages carry it. Best-effort: tracing must never break
+            // handling, and there may be no active UnitOfWork.
+            try {
+              const traceContext = spanFactory.currentTraceContext?.()
+              if (traceContext && Object.keys(traceContext).length > 0) {
+                contributeCorrelationData(traceContext)
+              }
+            } catch {
+              // no active UnitOfWork or no tracing context — skip
+            }
+            return handler(...args)
+          })
           span.end()
           return result
         } catch (error) {
@@ -32,3 +58,28 @@ export function tracingHandlerEnhancerDefinition(
     },
   }
 }
+
+function isMessage(value: unknown): value is Message {
+  return (
+    typeof value === "object" &&
+    value !== null &&
+    "metadata" in value &&
+    "identifier" in value
+  )
+}
+
+function createSpan(
+  spanFactory: SpanFactory,
+  spanName: string,
+  message: unknown,
+  metadata: HandlerMetadata,
+): Span {
+  if (!isMessage(message)) {
+    // No message to re-parent from (defensive — wired handlers always receive one).
+    return spanFactory.createInternalSpan(spanName)
+  }
+  if (metadata.messageType === "event" && spanFactory.createLinkedHandlerSpan) {
+    return spanFactory.createLinkedHandlerSpan(spanName, message)
+  }
+  return spanFactory.createHandlerSpan(spanName, message)
+}

package/src/tracking-event-processor.ts

@@ -22,6 +22,7 @@ import {
   advanceToken,
 } from "./tracking-token.js"
 import { REPLAY_STATE_KEY } from "./replay-token.js"
+import { applyCorrelationData, type CorrelationDataProvider } from "./correlation-data.js"
 import { setResource, onPrepareCommit } from "./processing-state.js"
 import type { HandlerEnhancerDefinition } from "./handler-enhancer.js"
 import type { CommandBus } from "./command-bus.js"
@@ -76,6 +77,12 @@ export interface TrackingEventProcessorOptions {
   queryBus?: QueryBus
   /** Event scheduler injected into ALS at handler-invocation entry (read by schedule()). */
   eventScheduler?: EventScheduler
+  /**
+   * Correlation data providers run against each event before its handlers are
+   * invoked, so commands/events dispatched from an event handler inherit the
+   * triggering event's correlationId/causationId.
+   */
+  correlationDataProviders?: ReadonlyArray<CorrelationDataProvider>
   /** Optional per-event callback fired inside the UoW before handler invocation (e.g. monitoring). */
   onEventDelivery?: () => void
   unitOfWorkRunner?: UoWRunner
@@ -151,6 +158,7 @@ export function createTrackingEventProcessor(
     commandBus,
     queryBus,
     eventScheduler,
+    correlationDataProviders,
     onEventDelivery,
     unitOfWorkRunner = runInNewUoW,
     tokenStore,
@@ -341,6 +349,11 @@ export function createTrackingEventProcessor(
     if (commandBus !== undefined) setResource(COMMAND_BUS_KEY, commandBus)
     if (queryBus !== undefined) setResource(QUERY_BUS_KEY, queryBus)
     if (eventScheduler !== undefined) setResource(EVENT_SCHEDULER_KEY, eventScheduler)
+    // Seed correlation data from the triggering event so an automation's
+    // outgoing commands/events inherit its lineage.
+    if (correlationDataProviders && correlationDataProviders.length > 0) {
+      applyCorrelationData(event, correlationDataProviders)
+    }
     // Optional per-event callback (e.g. monitoring hooks registered inside the UoW).
     if (onEventDelivery) onEventDelivery()
 
@@ -375,6 +388,9 @@ export function createTrackingEventProcessor(
     if (commandBus !== undefined) setResource(COMMAND_BUS_KEY, commandBus)
     if (queryBus !== undefined) setResource(QUERY_BUS_KEY, queryBus)
     if (eventScheduler !== undefined) setResource(EVENT_SCHEDULER_KEY, eventScheduler)
+    if (correlationDataProviders && correlationDataProviders.length > 0) {
+      applyCorrelationData(event, correlationDataProviders)
+    }
 
     const position =
       typeof letter.diagnostics.position === "number" ? BigInt(letter.diagnostics.position) : 0n