fix: activate activity trace context

CHANGELOG.md

@@ -40,6 +40,8 @@
 
 ### Fixes
 
+- Make the Durable Task activity execution span active while user activity code runs, so
+  spans created by activities inherit the correct parent across awaits and concurrent executions.
 - Preserve received structured failure `properties` in task errors, retry callbacks, client
   state/history, entity failures, and testing. Forward them with rethrown task failure chains.
   Empty wire maps remain `undefined`; arbitrary JavaScript `Error` fields are not collected.

packages/durabletask-js/src/worker/task-hub-grpc-worker.ts

@@ -38,6 +38,7 @@ import {
   setSpanError,
   setSpanOk,
   endSpan,
+  getOtelApi,
 } from "../tracing";
 
 /** Default timeout in milliseconds for graceful shutdown. */
@@ -1351,13 +1352,19 @@ export class TaskHubGrpcWorker {
         );
       } else {
         const executor = new ActivityExecutor(this._registry, this._logger);
-        const result = await executor.execute(
-          instanceId,
-          req.getName(),
-          req.getTaskid(),
-          req.getInput()?.getValue() ?? "",
-          req.getVersion()?.getValue(),
-        );
+        const executeActivity = () =>
+          executor.execute(
+            instanceId,
+            req.getName(),
+            req.getTaskid(),
+            req.getInput()?.getValue() ?? "",
+            req.getVersion()?.getValue(),
+          );
+        const otel = getOtelApi();
+        const result =
+          otel && activitySpan
+            ? await otel.context.with(otel.trace.setSpan(otel.context.active(), activitySpan), executeActivity)
+            : await executeActivity();
         res.setResult(new StringValue().setValue(result ?? ""));
         setSpanOk(activitySpan);
       }

packages/durabletask-js/test/worker-tracing.spec.ts

@@ -12,16 +12,20 @@
  */
 
 import * as pb from "../src/proto/orchestrator_service_pb";
+import * as stubs from "../src/proto/orchestrator_service_grpc_pb";
 import {
   newOrchestratorStartedEvent,
   newExecutionStartedEvent,
 } from "../src/utils/pb-helper.util";
 import { OrchestrationExecutor } from "../src/worker/orchestration-executor";
 import { Registry } from "../src/worker/registry";
+import { TaskHubGrpcWorker } from "../src/worker/task-hub-grpc-worker";
 import { NoOpLogger } from "../src/types/logger.type";
 import { TOrchestrator } from "../src/types/orchestrator.type";
 import { OrchestrationContext } from "../src/task/context/orchestration-context";
 import { ActivityContext } from "../src/task/context/activity-context";
+import * as traceContextUtils from "../src/tracing/trace-context-utils";
+import { AsyncLocalStorage } from "async_hooks";
 
 import {
   startSpanForOrchestrationExecution,
@@ -50,17 +54,48 @@ let exporter: InMemorySpanExporter;
 let provider: BasicTracerProvider;
 let previousTracerProvider: otel.TracerProvider | undefined;
 
+class TestAsyncContextManager implements otel.ContextManager {
+  private readonly storage = new AsyncLocalStorage<otel.Context>();
+
+  active(): otel.Context {
+    return this.storage.getStore() ?? otel.ROOT_CONTEXT;
+  }
+
+  with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
+    context: otel.Context,
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    return this.storage.run(context, () => fn.call(thisArg, ...args));
+  }
+
+  bind<T>(_context: otel.Context, target: T): T {
+    return target;
+  }
+
+  enable(): this {
+    return this;
+  }
+
+  disable(): this {
+    this.storage.disable();
+    return this;
+  }
+}
+
 beforeAll(() => {
   exporter = new InMemorySpanExporter();
   provider = new BasicTracerProvider();
   provider.addSpanProcessor(new SimpleSpanProcessor(exporter));
 
   previousTracerProvider = otel.trace.getTracerProvider();
-  provider.register();
+  provider.register({ contextManager: new TestAsyncContextManager() });
 });
 
 afterAll(async () => {
   await provider.shutdown();
+  otel.context.disable();
 
   if (previousTracerProvider) {
     otel.trace.setGlobalTracerProvider(previousTracerProvider);
@@ -71,6 +106,38 @@ beforeEach(() => {
   exporter.reset();
 });
 
+function createActivityRequest(
+  name: string,
+  taskId: number,
+  traceId: string,
+  parentSpanId: string,
+): pb.ActivityRequest {
+  return new pb.ActivityRequest()
+    .setName(name)
+    .setTaskid(taskId)
+    .setOrchestrationinstance(new pb.OrchestrationInstance().setInstanceid(`${name}-instance`))
+    .setParenttracecontext(createPbTraceContext(`00-${traceId}-${parentSpanId}-01`));
+}
+
+function createActivityStub(): {
+  stub: stubs.TaskHubSidecarServiceClient;
+  response: () => pb.ActivityResponse | undefined;
+} {
+  let capturedResponse: pb.ActivityResponse | undefined;
+  const stub = {
+    completeActivityTask: (
+      response: pb.ActivityResponse,
+      _metadata: unknown,
+      callback: (error: Error | null, response: pb.CompleteTaskResponse) => void,
+    ) => {
+      capturedResponse = response;
+      callback(null, new pb.CompleteTaskResponse());
+    },
+  } as unknown as stubs.TaskHubSidecarServiceClient;
+
+  return { stub, response: () => capturedResponse };
+}
+
 /**
  * Helper that mirrors the worker's _executeOrchestratorInternal tracing flow.
  * This replicates the exact sequence of tracing calls made by the worker so
@@ -432,3 +499,159 @@ describe("Worker Tracing - Activity Span Lifecycle", () => {
     expect(actSpan!.endTime).toBeDefined();
   });
 });
+
+describe("Worker Tracing - Activity Active Context", () => {
+  const tracer = otel.trace.getTracer("worker-tracing-user-code");
+
+  it("parents user spans to the activity execution span before and after await", async () => {
+    const worker = new TaskHubGrpcWorker({ logger: testLogger });
+    worker.addNamedActivity("parentedActivity", async () => {
+      const beforeAwait = tracer.startSpan("user-before-await");
+      beforeAwait.end();
+
+      await Promise.resolve();
+
+      const afterAwait = tracer.startSpan("user-after-await");
+      afterAwait.end();
+      return "done";
+    });
+    const request = createActivityRequest(
+      "parentedActivity",
+      1,
+      "11111111111111111111111111111111",
+      "1111111111111111",
+    );
+    const { stub } = createActivityStub();
+    const callerSpan = tracer.startSpan("success-caller");
+    const callerContext = otel.trace.setSpan(otel.ROOT_CONTEXT, callerSpan);
+
+    await otel.context.with(callerContext, async () => {
+      expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan);
+      await (worker as any)._executeActivityInternal(request, "completion-token", stub);
+      expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan);
+    });
+    callerSpan.end();
+
+    const spans = exporter.getFinishedSpans();
+    const activitySpan = spans.find((span) => span.name === "activity:parentedActivity")!;
+    const beforeAwait = spans.find((span) => span.name === "user-before-await")!;
+    const afterAwait = spans.find((span) => span.name === "user-after-await")!;
+
+    expect(beforeAwait.parentSpanId).toBe(activitySpan.spanContext().spanId);
+    expect(afterAwait.parentSpanId).toBe(activitySpan.spanContext().spanId);
+    expect(beforeAwait.spanContext().traceId).toBe(activitySpan.spanContext().traceId);
+    expect(afterAwait.spanContext().traceId).toBe(activitySpan.spanContext().traceId);
+  });
+
+  it("keeps concurrent activity contexts isolated", async () => {
+    const worker = new TaskHubGrpcWorker({ logger: testLogger });
+    let started = 0;
+    let releaseActivities!: () => void;
+    let reportBothStarted!: () => void;
+    const activitiesStarted = new Promise<void>((resolve) => {
+      reportBothStarted = resolve;
+    });
+    const release = new Promise<void>((resolve) => {
+      releaseActivities = resolve;
+    });
+
+    for (const name of ["activityA", "activityB"]) {
+      worker.addNamedActivity(name, async () => {
+        started++;
+        if (started === 2) {
+          reportBothStarted();
+        }
+        await release;
+        const child = tracer.startSpan(`user-${name}`);
+        child.end();
+        return name;
+      });
+    }
+
+    const first = createActivityRequest("activityA", 1, "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "aaaaaaaaaaaaaaaa");
+    const second = createActivityRequest("activityB", 2, "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", "bbbbbbbbbbbbbbbb");
+    const firstStub = createActivityStub();
+    const secondStub = createActivityStub();
+
+    const firstExecution = (worker as any)._executeActivityInternal(first, "first-completion-token", firstStub.stub);
+    const secondExecution = (worker as any)._executeActivityInternal(
+      second,
+      "second-completion-token",
+      secondStub.stub,
+    );
+    await activitiesStarted;
+    releaseActivities();
+    await Promise.all([firstExecution, secondExecution]);
+
+    const spans = exporter.getFinishedSpans();
+    for (const name of ["activityA", "activityB"]) {
+      const activitySpan = spans.find((span) => span.name === `activity:${name}`)!;
+      const childSpan = spans.find((span) => span.name === `user-${name}`)!;
+      expect(childSpan.parentSpanId).toBe(activitySpan.spanContext().spanId);
+      expect(childSpan.spanContext().traceId).toBe(activitySpan.spanContext().traceId);
+    }
+  });
+
+  it("restores the caller context and preserves failure tracing when activity code throws", async () => {
+    const worker = new TaskHubGrpcWorker({ logger: testLogger });
+    worker.addNamedActivity("throwingActivity", async () => {
+      await Promise.resolve();
+      const child = tracer.startSpan("user-throwing-activity");
+      child.end();
+      throw new Error("activity exploded");
+    });
+    const request = createActivityRequest(
+      "throwingActivity",
+      3,
+      "33333333333333333333333333333333",
+      "3333333333333333",
+    );
+    const activityStub = createActivityStub();
+    const callerSpan = tracer.startSpan("caller");
+    const callerContext = otel.trace.setSpan(otel.ROOT_CONTEXT, callerSpan);
+
+    await otel.context.with(callerContext, async () => {
+      expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan);
+      await (worker as any)._executeActivityInternal(request, "failure-completion-token", activityStub.stub);
+      expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan);
+    });
+    callerSpan.end();
+
+    const spans = exporter.getFinishedSpans();
+    const activitySpan = spans.find((span) => span.name === "activity:throwingActivity")!;
+    const childSpan = spans.find((span) => span.name === "user-throwing-activity")!;
+
+    expect(childSpan.parentSpanId).toBe(activitySpan.spanContext().spanId);
+    expect(activitySpan.status).toEqual({
+      code: otel.SpanStatusCode.ERROR,
+      message: "activity exploded",
+    });
+    expect(activityStub.response()?.getFailuredetails()?.getErrormessage()).toBe("activity exploded");
+  });
+
+  it("executes and completes activities when tracing is unavailable", async () => {
+    const worker = new TaskHubGrpcWorker({ logger: testLogger });
+    worker.addNamedActivity("untracedActivity", async () => {
+      await Promise.resolve();
+      return "untraced result";
+    });
+    const request = createActivityRequest(
+      "untracedActivity",
+      4,
+      "44444444444444444444444444444444",
+      "4444444444444444",
+    );
+    const activityStub = createActivityStub();
+    const getOtelApi = jest.spyOn(traceContextUtils, "getOtelApi").mockReturnValue(undefined);
+
+    try {
+      await (worker as any)._executeActivityInternal(request, "untraced-completion-token", activityStub.stub);
+    } finally {
+      getOtelApi.mockRestore();
+    }
+
+    expect(activityStub.response()?.getResult()?.getValue()).toBe(JSON.stringify("untraced result"));
+    expect(activityStub.response()?.getFailuredetails()).toBeUndefined();
+    expect(exporter.getFinishedSpans()).toHaveLength(0);
+  });
+});