@pensar/apex 2.5.0 → 2.6.1-canary.025b1ecf

package/src/core/observability/active-spans.ts

@@ -0,0 +1,80 @@
+import { diag, SpanStatusCode } from "@opentelemetry/api";
+import type {
+  ReadableSpan,
+  Span as SdkSpan,
+  SpanProcessor,
+} from "@opentelemetry/sdk-trace-base";
+
+export const INTERRUPTED_SPAN_ATTRIBUTE = "pensar.telemetry.interrupted";
+export const INTERRUPTED_ERROR_TYPE = "ApexProcessInterrupted";
+export const INTERRUPTED_SPAN_REASON =
+  "Apex shut down before this span completed";
+const MAX_TRACKED_ACTIVE_SPANS = 10_000;
+
+function spanKey(span: Pick<ReadableSpan, "spanContext">): string {
+  const context = span.spanContext();
+  return `${context.traceId}:${context.spanId}`;
+}
+
+/**
+ * Tracks every recording span owned by Apex's standalone SDK so shutdown can
+ * finish spans that their instrumentation still has open. OpenTelemetry only
+ * queues spans for export from `onEnd`; forceFlush alone cannot recover them.
+ *
+ * This processor must remain first in the provider's processor list: its
+ * shutdown hook is the final backstop for spans created during forceFlush,
+ * before downstream processors stop accepting ended spans.
+ */
+export class ActiveSpanProcessor implements SpanProcessor {
+  private readonly active = new Map<string, SdkSpan>();
+  private overflowWarned = false;
+
+  onStart(span: SdkSpan): void {
+    if (this.active.size >= MAX_TRACKED_ACTIVE_SPANS) {
+      if (!this.overflowWarned) {
+        this.overflowWarned = true;
+        diag.warn(
+          `Apex active-span shutdown tracking reached ${MAX_TRACKED_ACTIVE_SPANS} spans; newer spans will not be retained`,
+        );
+      }
+      return;
+    }
+    this.active.set(spanKey(span), span);
+  }
+
+  onEnd(span: ReadableSpan): void {
+    this.active.delete(spanKey(span));
+  }
+
+  endAll(reason = INTERRUPTED_SPAN_REASON): void {
+    // Spans start parent-first, so reverse insertion order closes the deepest
+    // children before their parents. Copy first because span.end() calls onEnd.
+    const spans = [...this.active.entries()].reverse();
+    for (const [key, span] of spans) {
+      try {
+        if (span.isRecording()) {
+          span.setAttribute(INTERRUPTED_SPAN_ATTRIBUTE, true);
+          if (span.attributes["error.type"] === undefined) {
+            span.setAttribute("error.type", INTERRUPTED_ERROR_TYPE);
+          }
+          if (span.status.code !== SpanStatusCode.ERROR) {
+            span.setStatus({ code: SpanStatusCode.ERROR, message: reason });
+          }
+          span.end();
+        }
+      } catch {
+        // One faulty instrumentation span must not block the remaining flush.
+      } finally {
+        // Delete only the copied entry. A downstream onEnd hook may start a
+        // new span re-entrantly, which must survive for the shutdown backstop.
+        this.active.delete(key);
+      }
+    }
+  }
+
+  async forceFlush(): Promise<void> {}
+
+  async shutdown(): Promise<void> {
+    this.endAll();
+  }
+}

package/src/core/observability/index.ts

@@ -0,0 +1,91 @@
+/**
+ * OpenTelemetry helpers. The tracer is a no-op unless the host process
+ * registers an OTel SDK; payload capture is off unless
+ * `AI_TRACE_RECORD_PAYLOADS=true`.
+ */
+
+import { AsyncLocalStorage } from "node:async_hooks";
+import { context, propagation, type Tracer, trace } from "@opentelemetry/api";
+import { isSessionId } from "../id/id";
+
+export function getApexTracer(): Tracer {
+  return trace.getTracer("apex");
+}
+
+/**
+ * Run-scoped override for AI payload capture. The durable runtime is
+ * multi-tenant (one process serves many workspaces concurrently), so payload
+ * capture cannot be a process-wide env toggle without cross-tenant races. Each
+ * run binds its own value here and {@link shouldRecordAiPayloads} reads it
+ * before falling back to the env.
+ */
+const aiPayloadCaptureStore = new AsyncLocalStorage<boolean>();
+
+/**
+ * Bind AI payload capture for the duration of `fn`, overriding the
+ * `AI_TRACE_RECORD_PAYLOADS` env for that async context only.
+ * @public Consumed by Console's durable agent runtime for per-workspace gating.
+ */
+export function runWithAiPayloadCapture<T>(enabled: boolean, fn: () => T): T {
+  return aiPayloadCaptureStore.run(enabled, fn);
+}
+
+export function shouldRecordAiPayloads(): boolean {
+  const override = aiPayloadCaptureStore.getStore();
+  if (override !== undefined) return override;
+  return process.env.AI_TRACE_RECORD_PAYLOADS === "true";
+}
+
+export {
+  registerActiveRootSpan,
+  unregisterActiveRootSpan,
+} from "./active-root-spans";
+export {
+  type AiTelemetryOperation,
+  type AiTelemetrySettings,
+  type CreateAiTelemetryInput,
+  createAiTelemetrySettings,
+  createGenerationSpanTracker,
+  type GenerationSpanTracker,
+  withModelCallDiagnostics,
+} from "./telemetry";
+
+/**
+ * OTel baggage key for the execution-session id. Console's SpanProcessor copies
+ * `pensar.*` baggage onto every span — must stay in sync with it.
+ */
+export const SESSION_BAGGAGE_KEY = "pensar.session.id";
+
+/**
+ * Run `fn` inside an OTel context that overrides `pensar.session.id` with
+ * `sessionId` (only that key — `pensar.root_session.id` is preserved), so the
+ * subagent's span subtree carries its own execution-session id. No-ops unless
+ * `sessionId` is a real `ses_` id.
+ */
+export function withSubagentSessionBaggage<T>(
+  sessionId: string | undefined,
+  fn: () => T,
+): T {
+  // Composite routing ids (`${ses_}-plan`) pass isSessionId but aren't real
+  // agent_sessions ids; the hyphen marks them derived — reject to protect the join.
+  if (!sessionId || !isSessionId(sessionId) || sessionId.includes("-")) {
+    return fn();
+  }
+  const baggage = (
+    propagation.getActiveBaggage() ?? propagation.createBaggage()
+  ).setEntry(SESSION_BAGGAGE_KEY, { value: sessionId });
+  return context.with(propagation.setBaggage(context.active(), baggage), fn);
+}
+
+/**
+ * The active span's correlation ids, when a valid recording span is active.
+ * Null outside spans (or under a no-op SDK) — callers skip the fields.
+ */
+export function getActiveTraceCorrelation(): {
+  traceId: string;
+  spanId: string;
+} | null {
+  const spanContext = trace.getSpan(context.active())?.spanContext();
+  if (!spanContext || !trace.isSpanContextValid(spanContext)) return null;
+  return { traceId: spanContext.traceId, spanId: spanContext.spanId };
+}

package/src/core/observability/log-correlation.test.ts

@@ -0,0 +1,276 @@
+// O5 — correlation between OTel spans, structured logs, and trace.jsonl
+// records. The active span's ids land in JSON logs (trace_id/span_id) and
+// on every StepTraceWriter record (correlation), flowing to W&B verbatim
+// through the existing trace-record event. Legacy records without
+// correlation stay readable.
+
+import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { type StepRecord, StepTraceWriter } from "../agents/offSecAgent/trace";
+import { AgentEventBus } from "../eventBus";
+import { createLogger, setLogSink } from "../logger/structured";
+import { getApexTracer } from "../observability";
+import {
+  type OtelTestHarness,
+  requireSpan,
+  startOtelTestHarness,
+} from "./testkit";
+
+let otel: OtelTestHarness;
+let tmpDir: string;
+const savedEnv: Record<string, string | undefined> = {};
+const ENV_KEYS = ["PENSAR_LOG_FORMAT", "PENSAR_LOG_LEVEL"];
+
+beforeEach(() => {
+  otel = startOtelTestHarness();
+  tmpDir = mkdtempSync(join(tmpdir(), "correlation-test-"));
+  for (const key of ENV_KEYS) {
+    savedEnv[key] = process.env[key];
+  }
+  process.env.PENSAR_LOG_FORMAT = "json";
+  process.env.PENSAR_LOG_LEVEL = "INFO";
+});
+
+afterEach(async () => {
+  await otel.shutdown();
+  rmSync(tmpDir, { recursive: true, force: true });
+  for (const key of ENV_KEYS) {
+    if (savedEnv[key] === undefined) delete process.env[key];
+    else process.env[key] = savedEnv[key];
+  }
+  setLogSink(null);
+});
+
+function captureLogLines(): string[] {
+  const lines: string[] = [];
+  setLogSink((line) => lines.push(line));
+  return lines;
+}
+
+// ---------------------------------------------------------------------------
+// Structured logger
+// ---------------------------------------------------------------------------
+
+describe("logger correlation", () => {
+  it("active spans add valid trace and span IDs to JSON logs", async () => {
+    const lines = captureLogLines();
+    const logger = createLogger("correlated");
+
+    await getApexTracer().startActiveSpan("log-span", {}, async (span) => {
+      logger.info("inside span", { extra: 1 });
+      span.end();
+    });
+
+    expect(lines).toHaveLength(1);
+    const record = JSON.parse(lines[0] ?? "") as Record<string, string>;
+    // 32-hex trace id, 16-hex span id — valid OTel formats.
+    expect(record.trace_id).toMatch(/^[0-9a-f]{32}$/);
+    expect(record.span_id).toMatch(/^[0-9a-f]{16}$/);
+    expect(record.msg).toBe("inside span");
+    expect(record.extra).toBe(1);
+  });
+
+  it("logs outside spans remain unchanged", () => {
+    const lines = captureLogLines();
+    createLogger("plain").info("outside span");
+
+    expect(lines).toHaveLength(1);
+    const record = JSON.parse(lines[0] ?? "") as Record<string, unknown>;
+    expect(record.trace_id).toBeUndefined();
+    expect(record.span_id).toBeUndefined();
+  });
+
+  it("child logs carry the child span id and the shared trace id", async () => {
+    const lines = captureLogLines();
+    const logger = createLogger("nested");
+
+    await getApexTracer().startActiveSpan("root", {}, async (root) => {
+      await getApexTracer().startActiveSpan("child", {}, async (child) => {
+        logger.info("in child");
+        child.end();
+      });
+      root.end();
+    });
+
+    const record = JSON.parse(lines[0] ?? "") as Record<string, string>;
+    const childSpan = requireSpan(otel.getFinishedSpans(), "child");
+    const rootSpan = requireSpan(otel.getFinishedSpans(), "root");
+    expect(record.span_id).toBe(childSpan.spanContext().spanId);
+    expect(record.trace_id).toBe(rootSpan.spanContext().traceId);
+  });
+
+  it("JSON logs remain valid one-line records", async () => {
+    const lines = captureLogLines();
+    const logger = createLogger("oneline");
+
+    await getApexTracer().startActiveSpan("s", {}, async (span) => {
+      logger.info("multi", { nested: { deep: "value" } });
+      span.end();
+    });
+
+    expect(lines).toHaveLength(1);
+    expect(lines[0]?.includes("\n")).toBe(false);
+    expect(() => JSON.parse(lines[0] ?? "")).not.toThrow();
+  });
+
+  it("pretty terminal logs remain readable (no trace ids)", () => {
+    process.env.PENSAR_LOG_FORMAT = "pretty";
+    const lines = captureLogLines();
+    const logger = createLogger("pretty");
+
+    return getApexTracer()
+      .startActiveSpan("pretty-span", {}, async (span) => {
+        logger.info("pretty message");
+        span.end();
+      })
+      .then(() => {
+        expect(lines).toHaveLength(1);
+        expect(lines[0]).toContain("pretty message");
+        expect(lines[0]).not.toContain("trace_id");
+        expect(lines[0]).not.toMatch(/[0-9a-f]{32}/);
+      });
+  });
+});
+
+// ---------------------------------------------------------------------------
+// StepTraceWriter / trace.jsonl records
+// ---------------------------------------------------------------------------
+
+function makeWriter(eventBus?: AgentEventBus): {
+  writer: StepTraceWriter;
+  tracePath: string;
+} {
+  const tracePath = join(tmpDir, "trace.jsonl");
+  return {
+    writer: new StepTraceWriter({
+      tracePath,
+      agentId: "ses_agent_1",
+      ...(eventBus ? { eventBus } : {}),
+    }),
+    tracePath,
+  };
+}
+
+describe("trace record correlation", () => {
+  it("root and child records share the same trace ID; child carries its own span ID", async () => {
+    const { writer, tracePath } = makeWriter();
+
+    await getApexTracer().startActiveSpan("root", {}, async (root) => {
+      writer.recordStep(
+        [{ role: "assistant", content: [{ type: "text", text: "root step" }] }],
+        { inputTokens: 10, outputTokens: 2 },
+      );
+      await getApexTracer().startActiveSpan("child", {}, async (child) => {
+        writer.recordStep(
+          [
+            {
+              role: "assistant",
+              content: [{ type: "text", text: "child step" }],
+            },
+          ],
+          { inputTokens: 5, outputTokens: 1 },
+        );
+        child.end();
+      });
+      root.end();
+    });
+
+    const lines = readFileSync(tracePath, "utf-8").trim().split("\n");
+    expect(lines).toHaveLength(2);
+    const [rootRecord, childRecord] = lines.map(
+      (line) => JSON.parse(line) as StepRecord,
+    );
+
+    const rootSpan = requireSpan(otel.getFinishedSpans(), "root");
+    const childSpan = requireSpan(otel.getFinishedSpans(), "child");
+    // Same trace…
+    expect(rootRecord.correlation?.traceId).toBe(
+      rootSpan.spanContext().traceId,
+    );
+    expect(childRecord.correlation?.traceId).toBe(
+      rootSpan.spanContext().traceId,
+    );
+    // …but each record names the span it was written under.
+    expect(rootRecord.correlation?.spanId).toBe(rootSpan.spanContext().spanId);
+    expect(childRecord.correlation?.spanId).toBe(
+      childSpan.spanContext().spanId,
+    );
+  });
+
+  it("records written outside spans carry no correlation", () => {
+    const { writer, tracePath } = makeWriter();
+    writer.recordStep(
+      [{ role: "assistant", content: [{ type: "text", text: "bare" }] }],
+      { inputTokens: 1, outputTokens: 1 },
+    );
+
+    const record = JSON.parse(
+      readFileSync(tracePath, "utf-8").trim(),
+    ) as Record<string, unknown>;
+    expect(record.correlation).toBeUndefined();
+  });
+
+  it("JSONL readers accept old and new records", () => {
+    const { writer, tracePath } = makeWriter();
+    // A legacy record line, exactly as older Apex wrote it (no correlation).
+    const legacy = JSON.stringify({
+      type: "step",
+      timestamp: new Date().toISOString(),
+      agentId: null,
+      stepIndex: 0,
+      text: "legacy",
+    });
+    writeFileSync(tracePath, `${legacy}\n`);
+
+    return getApexTracer()
+      .startActiveSpan("modern", {}, async (span) => {
+        writer.recordStep(
+          [{ role: "assistant", content: [{ type: "text", text: "modern" }] }],
+          { inputTokens: 1, outputTokens: 1 },
+        );
+        span.end();
+      })
+      .then(() => {
+        const records = readFileSync(tracePath, "utf-8")
+          .trim()
+          .split("\n")
+          .map((l) => JSON.parse(l) as Record<string, unknown>);
+        expect(records).toHaveLength(2);
+        expect(records[0]?.correlation).toBeUndefined(); // legacy untouched
+        expect(records[1]?.correlation).toBeDefined(); // new correlated
+      });
+  });
+
+  it("W&B receives correlation through its existing trace-record input", async () => {
+    const bus = new AgentEventBus();
+    const received: Array<{
+      record: { correlation?: { traceId?: string; spanId?: string } };
+      subagentId?: string;
+    }> = [];
+    bus.on("trace-record", (e) => {
+      received.push(
+        e as {
+          record: { correlation?: { traceId?: string; spanId?: string } };
+          subagentId?: string;
+        },
+      );
+    });
+
+    const { writer } = makeWriter(bus);
+    await getApexTracer().startActiveSpan("wandb-span", {}, async (span) => {
+      writer.recordStep(
+        [{ role: "assistant", content: [{ type: "text", text: "uploaded" }] }],
+        { inputTokens: 3, outputTokens: 1 },
+      );
+      span.end();
+    });
+
+    // The event payload IS what attachWandbToEventBus hands the uploader —
+    // no additional upload logic needed for correlation to reach W&B.
+    expect(received).toHaveLength(1);
+    expect(received[0]?.record.correlation?.traceId).toMatch(/^[0-9a-f]{32}$/);
+    expect(received[0]?.subagentId).toBe("ses_agent_1");
+  });
+});

package/src/core/observability/payload-capture.test.ts

@@ -0,0 +1,59 @@
+import { afterEach, describe, expect, it } from "vitest";
+import { runWithAiPayloadCapture, shouldRecordAiPayloads } from "./index";
+import { createAiTelemetrySettings } from "./telemetry";
+
+afterEach(() => {
+  process.env.AI_TRACE_RECORD_PAYLOADS = undefined;
+});
+
+describe("runWithAiPayloadCapture", () => {
+  it("falls back to the env when no run-scoped override is set", () => {
+    process.env.AI_TRACE_RECORD_PAYLOADS = undefined;
+    expect(shouldRecordAiPayloads()).toBe(false);
+    process.env.AI_TRACE_RECORD_PAYLOADS = "true";
+    expect(shouldRecordAiPayloads()).toBe(true);
+  });
+
+  it("enables capture inside the run even when the env is unset", () => {
+    process.env.AI_TRACE_RECORD_PAYLOADS = undefined;
+    runWithAiPayloadCapture(true, () => {
+      expect(shouldRecordAiPayloads()).toBe(true);
+      const settings = createAiTelemetrySettings({
+        operation: "apex.agent.stream",
+      });
+      expect(settings.recordInputs).toBe(true);
+      expect(settings.recordOutputs).toBe(true);
+    });
+    // Override does not leak past the run.
+    expect(shouldRecordAiPayloads()).toBe(false);
+  });
+
+  it("disables capture inside the run even when the env forces it on", () => {
+    process.env.AI_TRACE_RECORD_PAYLOADS = "true";
+    runWithAiPayloadCapture(false, () => {
+      expect(shouldRecordAiPayloads()).toBe(false);
+      const settings = createAiTelemetrySettings({
+        operation: "apex.structured.generate",
+      });
+      expect(settings.recordInputs).toBe(false);
+      expect(settings.recordOutputs).toBe(false);
+    });
+    expect(shouldRecordAiPayloads()).toBe(true);
+  });
+
+  it("isolates concurrent runs across awaits (multi-tenant safety)", async () => {
+    process.env.AI_TRACE_RECORD_PAYLOADS = undefined;
+    const readAfterTick = (enabled: boolean) =>
+      runWithAiPayloadCapture(enabled, async () => {
+        await Promise.resolve();
+        return shouldRecordAiPayloads();
+      });
+
+    const [on, off] = await Promise.all([
+      readAfterTick(true),
+      readAfterTick(false),
+    ]);
+    expect(on).toBe(true);
+    expect(off).toBe(false);
+  });
+});