From a42ac91d7f8fb9a836444b4bbacdd07c38c01511 Mon Sep 17 00:00:00 2001 From: Dotta Date: Thu, 15 Oct 2026 19:58:24 -0500 Subject: [PATCH] fix(observability): retain provider scope after admission handoff Co-Authored-By: Paperclip --- .../heartbeat/runtime-execution.test.ts | 24 +++++++++++++++---- .../services/heartbeat/runtime-execution.ts | 7 ++++-- 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/server/src/services/heartbeat/runtime-execution.test.ts b/server/src/services/heartbeat/runtime-execution.test.ts index f8d9ee53c02..361fe422b69 100644 --- a/server/src/services/heartbeat/runtime-execution.test.ts +++ b/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>) => void; const providerResult = new Promise>>(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; diff --git a/server/src/services/heartbeat/runtime-execution.ts b/server/src/services/heartbeat/runtime-execution.ts index bdd459d45cc..761c31b2cb2 100644 --- a/server/src/services/heartbeat/runtime-execution.ts +++ b/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 =