fix: retain background preparation traces from unsampled requests (#16205)

Start preparation phase root traces independently of request sampling and link each phase to its trigger. Preserve operator root sampling, concurrent child context, and silent unconfigured behavior.

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

doc/observability.md

@@ -170,6 +170,12 @@ and launcher preparation, run output initialization, workspace validation,
 runtime services, instruction files, and runtime selection. Database and HTTP
 instrumentation can attach child spans during each awaited operation.
 
+Each phase starts a new root trace and links to the active trigger span when
+one exists. An unsampled HTTP request must not suppress diagnostics for the
+background run it starts. Root traces still follow the operator's configured
+sampler. The trigger's trace remains unchanged, and overlapping runs keep
+their own phase and child contexts.
+
 These spans exist before provider dispatch, including setup failures. Each has
 only the existing `paperclip.task.run.run_id` attribute: the first 12 hexadecimal
 characters of the run UUID's SHA-256 hash. Failed operations set error status

server/src/services/heartbeat/preparation-timing-context.test.ts

@@ -1,11 +1,102 @@
 import { createHash } from "node:crypto";
 import { setImmediate } from "node:timers/promises";
-import { context, trace } from "@opentelemetry/api";
+import { context, trace, TraceFlags } from "@opentelemetry/api";
 import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks";
-import { TracerProvider, InMemorySpanExporter, SimpleSpanProcessor } from "@opentelemetry/sdk-trace";
-import { expect, it } from "vitest";
+import { AlwaysOffSampler, ParentBasedSampler, TracerProvider, InMemorySpanExporter, SimpleSpanProcessor } from "@opentelemetry/sdk-trace";
+import { expect, it, vi } from "vitest";
 import { createHeartbeatPreparationTiming } from "./preparation-timing.js";
 
+it.each([TraceFlags.NONE, TraceFlags.SAMPLED])("records background preparation independently of trigger flags %i and links back to it", async traceFlags => {
+  const manager = new AsyncLocalStorageContextManager().enable();
+  const exporter = new InMemorySpanExporter();
+  const provider = new TracerProvider({ spanProcessors: [new SimpleSpanProcessor({ exporter })] });
+  expect(context.setGlobalContextManager(manager)).toBe(true);
+  expect(trace.setGlobalTracerProvider(provider)).toBe(true);
+  const trigger = {
+    traceId: "11111111111111111111111111111111",
+    spanId: "2222222222222222",
+    traceFlags,
+    isRemote: true,
+  };
+  try {
+    const result = {};
+    await context.with(trace.setSpanContext(context.active(), trigger), async () => {
+      await expect(createHeartbeatPreparationTiming("background-run")("workspace", async () => {
+        await setImmediate();
+        const child = trace.getTracer("preparation-context-test").startSpan("background.child");
+        expect(child.isRecording()).toBe(true);
+        child.end();
+        return result;
+      })).resolves.toBe(result);
+      expect(trace.getActiveSpan()?.spanContext()).toEqual(trigger);
+    });
+    await provider.forceFlush();
+    const spans = exporter.getFinishedSpans();
+    expect(spans).toHaveLength(2);
+    const phase = spans.find(span => span.name === "heartbeat.prepare.workspace")!;
+    const child = spans.find(span => span.name === "background.child")!;
+    expect(phase.parentSpanContext).toBeUndefined();
+    expect(phase.links).toHaveLength(1);
+    expect(phase.links[0].context).toEqual(trigger);
+    expect(phase.spanContext().traceId).not.toBe(trigger.traceId);
+    expect(child.parentSpanContext?.spanId).toBe(phase.spanContext().spanId);
+    expect(child.spanContext().traceId).toBe(phase.spanContext().traceId);
+  } finally {
+    trace.disable();
+    context.disable();
+    manager.disable();
+    await provider.shutdown();
+  }
+});
+
+it("honors an operator sampler that disables root traces even when the trigger is sampled", async () => {
+  const manager = new AsyncLocalStorageContextManager().enable();
+  const exporter = new InMemorySpanExporter();
+  const provider = new TracerProvider({
+    sampler: new ParentBasedSampler({ root: new AlwaysOffSampler() }),
+    spanProcessors: [new SimpleSpanProcessor({ exporter })],
+  });
+  expect(context.setGlobalContextManager(manager)).toBe(true);
+  expect(trace.setGlobalTracerProvider(provider)).toBe(true);
+  try {
+    await context.with(trace.setSpanContext(context.active(), {
+      traceId: "11111111111111111111111111111111", spanId: "2222222222222222",
+      traceFlags: TraceFlags.SAMPLED, isRemote: true,
+    }), async () => {
+      await createHeartbeatPreparationTiming("disabled-run")("instructions", async () => {
+        expect(trace.getActiveSpan()?.isRecording()).toBe(false);
+      });
+    });
+    await provider.forceFlush();
+    expect(exporter.getFinishedSpans()).toEqual([]);
+  } finally {
+    trace.disable();
+    context.disable();
+    manager.disable();
+    await provider.shutdown();
+  }
+});
+
+it("stays silent and preserves values and errors without a tracer provider", async () => {
+  trace.disable();
+  context.disable();
+  const messages = ["log", "info", "warn", "error", "debug"] as const;
+  const spies = messages.map(method => vi.spyOn(console, method));
+  try {
+    const time = createHeartbeatPreparationTiming("unconfigured-run");
+    const result = {};
+    await expect(time("workspace", async () => {
+      expect(trace.getActiveSpan()?.isRecording()).not.toBe(true);
+      return result;
+    })).resolves.toBe(result);
+    const error = new Error("original failure");
+    await expect(time("instructions", async () => { throw error; })).rejects.toBe(error);
+    for (const spy of spies) expect(spy).not.toHaveBeenCalled();
+  } finally {
+    for (const spy of spies) spy.mockRestore();
+  }
+});
+
 it("parents child spans to their own preparation phase across overlapping awaited runs", async () => {
   const manager = new AsyncLocalStorageContextManager().enable();
   const exporter = new InMemorySpanExporter();

server/src/services/heartbeat/preparation-timing.test.ts

@@ -3,14 +3,15 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { createHeartbeatPreparationTiming } from "./preparation-timing.js";
 
 const handles = vi.hoisted(() => ({ spans: [] as Array<{ setStatus: ReturnType<typeof vi.fn>; end: ReturnType<typeof vi.fn> }> }));
-const startActiveSpan = vi.hoisted(() => vi.fn((_name, _options, callback) => {
+const startActiveSpan = vi.hoisted(() => vi.fn((_name, _options, _context, callback) => {
   const span = { setStatus: vi.fn(), end: vi.fn() };
   handles.spans.push(span);
   return callback(span);
 }));
 vi.mock("@opentelemetry/api", () => ({
   SpanStatusCode: { ERROR: 2 },
-  trace: { getTracer: () => ({ startActiveSpan }) },
+  ROOT_CONTEXT: "root-context",
+  trace: { getTracer: () => ({ startActiveSpan }), getActiveSpan: () => undefined },
 }));
 
 describe("heartbeat preparation timing", () => {
@@ -23,7 +24,7 @@ describe("heartbeat preparation timing", () => {
     await expect(time("instructions", async () => result)).resolves.toBe(result);
     expect(startActiveSpan).toHaveBeenCalledExactlyOnceWith("heartbeat.prepare.instructions", {
       attributes: { "paperclip.task.run.run_id": createHash("sha256").update(runId).digest("hex").slice(0, 12) },
-    }, expect.any(Function));
+    }, "root-context", expect.any(Function));
     expect(handles.spans[0].end).toHaveBeenCalledOnce();
     expect(handles.spans[0].setStatus).not.toHaveBeenCalled();
   });

server/src/services/heartbeat/preparation-timing.ts

@@ -1,5 +1,5 @@
 import { createHash } from "node:crypto";
-import { SpanStatusCode, trace } from "@opentelemetry/api";
+import { isSpanContextValid, ROOT_CONTEXT, SpanStatusCode, trace } from "@opentelemetry/api";
 
 type PreparationPhase = "environment_selection" | "execution_configuration" | "session"
   | "workspace" | "github_environment" | "github_launchers" | "run_output"
@@ -7,10 +7,15 @@ type PreparationPhase = "environment_selection" | "execution_configuration" | "s
 
 export function createHeartbeatPreparationTiming(runId: string) {
   const runIdHash = createHash("sha256").update(runId).digest("hex").slice(0, 12);
-  return <T>(phase: PreparationPhase, work: () => Promise<T>): Promise<T> =>
-    trace.getTracer("paperclip.heartbeat-preparation").startActiveSpan(
+  return <T>(phase: PreparationPhase, work: () => Promise<T>): Promise<T> => {
+    const trigger = trace.getActiveSpan()?.spanContext();
+    return trace.getTracer("paperclip.heartbeat-preparation").startActiveSpan(
       `heartbeat.prepare.${phase}`,
-      { attributes: { "paperclip.task.run.run_id": runIdHash } },
+      {
+        attributes: { "paperclip.task.run.run_id": runIdHash },
+        ...(trigger && isSpanContextValid(trigger) ? { links: [{ context: trigger }] } : {}),
+      },
+      ROOT_CONTEXT,
       async (span) => {
         try {
           return await work();
@@ -22,4 +27,5 @@ export function createHeartbeatPreparationTiming(runId: string) {
         }
       },
     );
+  };
 }