fix(observability): retain provider scope after admission handoff

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/services/heartbeat/runtime-execution.test.ts

@@ -7,7 +7,7 @@ import { eq, sql } from "drizzle-orm";
 import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import { trace } from "@opentelemetry/api";
 import { InMemorySpanExporter, SimpleSpanProcessor, TracerProvider } from "@opentelemetry/sdk-trace";
-import { measureRunPhase, measureWorkspaceSyncPhase, withRunExecutionTiming } from "../execution-preparation-timing.js";
+import { createRunPreparationTiming, measureRunPhase, measureWorkspaceSyncPhase, withRunExecutionTiming } from "../execution-preparation-timing.js";
 import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js";
 import * as legacyRecovery from "../legacy-execution-recovery.js";
 import { createAgentIdentityRedactor } from "../agent-identity-redaction.js";
@@ -267,8 +267,8 @@ describe.skipIf(!support.supported)("heartbeat runtime execution boundary", () =
     }
   });
 
-  it("ends admission timing while the admitted provider is still running", async () => {
-    const { input, run } = await fixture(true);
+  it.each([false, true])("ends admission timing without losing later provider spans (native: %s)", async native => {
+    const { input, adapter, run } = await fixture(native);
     const endpoint = process.env.OTEL_EXPORTER_OTLP_ENDPOINT;
     vi.stubEnv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.test");
     const exporter = new InMemorySpanExporter();
@@ -276,15 +276,29 @@ describe.skipIf(!support.supported)("heartbeat runtime execution boundary", () =
     expect(trace.setGlobalTracerProvider(provider)).toBe(true);
     let release!: (result: Awaited<ReturnType<typeof executePaperclipNativeSession>>) => void;
     const providerResult = new Promise<Awaited<ReturnType<typeof executePaperclipNativeSession>>>(resolve => { release = resolve; });
-    vi.mocked(executePaperclipNativeSession).mockReturnValueOnce(providerResult);
+    const providerWork = async () => {
+      const result = await providerResult;
+      await createRunPreparationTiming(run.id, "native.controller")("coordinator_release", async () => {});
+      return result;
+    };
+    if (native) vi.mocked(executePaperclipNativeSession).mockImplementationOnce(providerWork);
+    else adapter.execute.mockImplementationOnce(providerWork);
     const running = withRunExecutionTiming(run.id, () => executeHeartbeatRuntime(db, input));
     try {
       await vi.waitFor(() => {
         expect(exporter.getFinishedSpans().filter(span => span.name === "execution.dispatch_admission")).toHaveLength(1);
       }, { timeout: 5_000 });
-      expect(executePaperclipNativeSession).toHaveBeenCalledOnce();
+      expect(native ? executePaperclipNativeSession : adapter.execute).toHaveBeenCalledOnce();
       expect(exporter.getFinishedSpans().find(span => span.name === "task.execution_attempt")).toBeUndefined();
       expect(await finalizations()).toHaveLength(0);
+      release({ exitCode: 0, signal: null, timedOut: false });
+      await running;
+      const spans = exporter.getFinishedSpans();
+      const attempt = spans.find(span => span.name === "task.execution_attempt")!;
+      const later = spans.filter(span => span.name === "native.controller.coordinator_release");
+      expect(later).toHaveLength(1);
+      expect(later[0].spanContext().traceId).toBe(attempt.spanContext().traceId);
+      expect(later[0].parentSpanContext?.spanId).toBe(attempt.spanContext().spanId);
     } finally {
       release({ exitCode: 0, signal: null, timedOut: false });
       await running;

server/src/services/heartbeat/runtime-execution.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from "node:async_hooks";
 import { measureRunPhase } from "../execution-preparation-timing.js";
 import { parseNativeSessionGoalControl } from "./scheduling.js";
 import { isHeartbeatRunTerminalStatus } from "./run-lifecycle.js";
@@ -231,8 +232,10 @@ export async function executeHeartbeatRuntime(db: Db, input: HeartbeatRuntimeExe
   const { nativeInstructionWorkingCopy, collectStoppedInstructions, nativeInstructionReservation, releaseInstructionCopy } = input.instructions;
   const { providerTraceCapture, nativeRunnerPreparationSpans, attestedQuestionResponseAtMs, attemptStartedAtMs, environmentAcquireStartedAtMs, environmentRealizeEndedAtMs } = input.trace;
   const { executionControl, executionPhaseContext, dispatchResolvedInteractionContinuationWithAtomicGate } = input.control;
-  const dispatchWithTiming: HeartbeatRuntimeExecutionInput["control"]["dispatchResolvedInteractionContinuationWithAtomicGate"] = dispatch =>
-    measureRunPhase("execution.dispatch_admission", () => dispatchResolvedInteractionContinuationWithAtomicGate(dispatch));
+  const dispatchWithTiming: HeartbeatRuntimeExecutionInput["control"]["dispatchResolvedInteractionContinuationWithAtomicGate"] = dispatch => {
+    const providerDispatch = AsyncLocalStorage.bind(dispatch);
+    return measureRunPhase("execution.dispatch_admission", () => dispatchResolvedInteractionContinuationWithAtomicGate(providerDispatch));
+  };
   const { identityRedactor, appendIdentityRedactedLog, onLog, onAdapterEvent, onAdapterMeta } = input.output;
   const { executionWorkspacesSvc, instanceSettings, envOrchestrator, upsertTaskSession, getRun, cancelRunInternal, enqueueWakeup, persistRunProcessMetadata, recordCurrentHeartbeatRunRuntimeProgress, dispatchPendingNativeStatusWakeups } = input.services;
   const localAgentJwtScope =