fix: prevent OOM when processing large Codex rollouts

README.md

@@ -70,6 +70,14 @@ The remaining options are settable both ways, as a `langfuse.json` key or as the
 
 Everything the plugin traces is uploaded to Langfuse, including prompts and tool inputs and outputs, so do not enable it for sessions containing data you do not want stored there.
 
+### Large rollouts
+
+Rollouts are read as a stream using a fixed file-size snapshot. Records over 8 MiB are discarded without decoding, and further non-boundary records in a turn are omitted after its 32 MiB budget. Turn boundaries are still processed, so an oversized tool output does not prevent earlier or later turns from being traced. Affected turn spans carry `codex.rollout.truncated=true`. Retained system/tool context is capped at 4 MiB and labelled `codex.system_prompt.truncated` when incomplete. Routing metadata is bounded too; `codex.routing.truncated` indicates degraded subagent attribution rather than aborting the session.
+
+Generation inputs use a 256 KiB recent-history window. Whole old turns/steps are evicted to preserve assistant/tool-call relationships. The system message and current user message are retained separately, with labelled excerpts if either is too large; oversized history groups or tool schemas may be omitted. `codex.generation_input.truncated=true` makes this degradation explicit. Later small inputs remain useful rather than being permanently replaced with an omission placeholder.
+
+Exports are flushed after at most 32 ended spans or approximately 4 MiB of span attributes, and at the end of the hook. A turn is checkpointed only after it has completed and all its sampled spans have been confirmed exported. Preprocessing drops leave only the affected turns retryable; transport failures leave the unconfirmed batch retryable. This batches catch-up work within the hook's 30-second timeout budget instead of requiring one HTTP request per turn. It is not a guarantee against slow networks or unusually large sessions; confirmed progress survives interruption.
+
 ## Contributing
 
 See the [contributing guide](./CONTRIBUTING.md).

plugins/tracing/src/export-batch.ts

@@ -0,0 +1,64 @@
+import type { LangfuseObservation } from "@langfuse/tracing";
+import type { ReadableSpan } from "@opentelemetry/sdk-trace-base";
+
+import { spanIdentity, type ExportReceipt } from "./instrumentation.js";
+import { debugLog } from "./utils.js";
+
+type PendingTurn = {
+  id: string;
+  spanIds: Set<string>;
+  success: boolean;
+  checkpoint: (id: string) => Promise<void>;
+};
+
+/** Bound the queue across turns; acknowledge only complete, delivered turns. */
+export class ExportBatch {
+  private spans = 0;
+  private bytes = 0;
+  private turns: PendingTurn[] = [];
+  private failedSpanIds = new Set<string>();
+  private transportFailure: Error | undefined;
+  incomplete = false;
+
+  constructor(private deliver?: () => Promise<ExportReceipt | void>) {}
+
+  async ended(observation: LangfuseObservation, spanIds: Set<string>): Promise<void> {
+    spanIds.add(spanIdentity(observation.otelSpan.spanContext()));
+    this.spans++;
+    const { attributes } = observation.otelSpan as unknown as Pick<ReadableSpan, "attributes">;
+    this.bytes += Buffer.byteLength(JSON.stringify(attributes ?? {}));
+    if (this.spans >= 32 || this.bytes >= 4 * 1024 * 1024) await this.flush();
+  }
+
+  complete(turn: PendingTurn): void {
+    this.turns.push(turn);
+  }
+
+  async flush(): Promise<void> {
+    // A transport failure invalidates this entire unacknowledged batch. Never
+    // let a later empty flush acknowledge spans that its exporter discarded.
+    if (this.transportFailure) throw this.transportFailure;
+    let receipt: ExportReceipt | void;
+    try {
+      receipt = await this.deliver?.();
+    } catch (error) {
+      this.transportFailure = error instanceof Error ? error : new Error(String(error));
+      throw this.transportFailure;
+    }
+    for (const id of receipt?.failedSpanIds ?? []) this.failedSpanIds.add(id);
+    this.spans = 0;
+    this.bytes = 0;
+    const turns = this.turns;
+    this.turns = [];
+    for (const turn of turns) {
+      const missing = [...turn.spanIds].some((id) => this.failedSpanIds.has(id));
+      for (const id of turn.spanIds) this.failedSpanIds.delete(id);
+      if (turn.success && !missing) await turn.checkpoint(turn.id);
+      else {
+        this.incomplete = true;
+        debugLog(`turn ${turn.id} was not completely delivered; leaving it retryable`);
+      }
+    }
+    // Failures for a still-open parent/current turn remain until it completes.
+  }
+}

plugins/tracing/src/index.ts

@@ -51,14 +51,15 @@ export async function runHook(): Promise<void> {
   }
 
   const instrumentation = setupInstrumentation(config);
-  let exportedTurnIds: string[] = [];
   let failure: unknown;
 
   try {
-    exportedTurnIds = await convertRollout(hookInput.transcript_path, {
+    await convertRollout(hookInput.transcript_path, {
       config,
       parentSpanContext,
       stoppedTurnId: hookInput.turn_id ?? undefined,
+      flush: instrumentation.flush,
+      onTurnExported: (turnId) => markTurnUploaded(hookInput.transcript_path!, turnId),
     });
   } catch (error) {
     failure = error;
@@ -72,17 +73,12 @@ export async function runHook(): Promise<void> {
     debugLog("error during flush/shutdown:", error);
   }
 
-  // Record delivery only after the flush: a turn marked despite a failed export
-  // is lost silently, while an unmarked one is retried.
+  // Earlier checkpoints remain delivered; failed and unfinished turns retry.
   if (failure) {
     if (config.fail_on_error) throw failure;
     debugLog("export incomplete; leaving turns unmarked so they stay retryable");
     return;
   }
-
-  for (const turnId of exportedTurnIds) {
-    await markTurnUploaded(hookInput.transcript_path, turnId);
-  }
 }
 
 runHook().catch((error) => {

plugins/tracing/src/instrumentation.ts

@@ -2,6 +2,10 @@ import { createHash, randomBytes } from "node:crypto";
 
 import { LangfuseSpanProcessor } from "@langfuse/otel";
 import { setLangfuseTracerProvider } from "@langfuse/tracing";
+import { TraceFlags, type SpanContext } from "@opentelemetry/api";
+import { ExportResultCode } from "@opentelemetry/core";
+import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
+import type { SpanExporter, SpanProcessor } from "@opentelemetry/sdk-trace-base";
 import {
   defaultResource,
   detectResources,
@@ -37,10 +41,18 @@ function seededId(kind: "trace" | "span"): string {
 }
 
 export type Instrumentation = {
+  /** Wait for delivery; reject even if an automatic batch failed earlier. */
+  flush: () => Promise<ExportReceipt>;
   /** Flush buffered spans and tear down the tracer provider. */
   shutdown: () => Promise<void>;
 };
 
+export type ExportReceipt = { failedSpanIds: ReadonlySet<string> };
+
+export function spanIdentity(context: SpanContext): string {
+  return `${context.traceId}:${context.spanId}`;
+}
+
 const DEFAULT_SERVICE_NAME = "codex";
 
 function buildResource() {
@@ -57,13 +69,36 @@ function buildResource() {
  * instrumentation loading. Registering the provider also installs the
  * AsyncLocalStorage context manager that `propagateAttributes` relies on.
  *
- * We use `exportMode: "batched"` and flush once at the end: the whole rollout
- * is converted in-process, so batching every span into one (or a few) requests
- * is far faster than one request per span — important for the hook's timeout
- * budget. `shutdown()` below calls `forceFlush()` before the process exits.
+ * Batch across turns, with explicit checkpoints to bound the export backlog.
+ * Track delivery independently: a background batch can fail before forceFlush,
+ * and the SDK can swallow preprocessing errors or discard an overflowing queue.
+ * Neither case may cause the hook to mark an undelivered turn uploaded.
  */
 export function setupInstrumentation(config: Config): Instrumentation {
+  const pending = new Set<string>();
+  let exportFailure: Error | undefined;
+  const exporter = new OTLPTraceExporter({
+    url: `${config.base_url.replace(/\/$/, "")}/api/public/otel/v1/traces`,
+    headers: {
+      Authorization: `Basic ${Buffer.from(`${config.public_key}:${config.secret_key}`).toString("base64")}`,
+      "x-langfuse-public-key": config.public_key ?? "<missing>",
+      "x-langfuse-sdk-name": "javascript",
+    },
+    timeoutMillis: Number(process.env.LANGFUSE_TIMEOUT ?? 5) * 1000,
+  });
+  const trackedExporter: SpanExporter = {
+    export: (spans, callback) => {
+      exporter.export(spans, (result) => {
+        if (result.code === ExportResultCode.SUCCESS) {
+          for (const span of spans) pending.delete(spanIdentity(span.spanContext()));
+        } else exportFailure ??= result.error ?? new Error("Span export failed");
+        callback(result);
+      });
+    },
+    shutdown: () => exporter.shutdown(),
+  };
   const spanProcessor = new LangfuseSpanProcessor({
+    exporter: trackedExporter,
     publicKey: config.public_key,
     secretKey: config.secret_key,
     baseUrl: config.base_url,
@@ -72,9 +107,39 @@ export function setupInstrumentation(config: Config): Instrumentation {
     shouldExportSpan: () => true,
   });
 
+  const trackedProcessor: SpanProcessor = {
+    onStart: (span, context) => spanProcessor.onStart(span, context),
+    onEnd: (span) => {
+      if (span.spanContext().traceFlags & TraceFlags.SAMPLED) {
+        pending.add(spanIdentity(span.spanContext()));
+      }
+      spanProcessor.onEnd(span);
+    },
+    forceFlush: () => spanProcessor.forceFlush(),
+    shutdown: () => spanProcessor.shutdown(),
+  };
+
+  const flush = async () => {
+    let flushFailure: unknown;
+    try {
+      await spanProcessor.forceFlush();
+    } catch (error) {
+      flushFailure = error;
+    }
+    // forceFlush waits for preprocessing and the export queue. Anything still
+    // pending was dropped; report it only for this window, not every future turn.
+    const failedSpanIds = new Set(pending);
+    pending.clear();
+    const failure = exportFailure;
+    exportFailure = undefined;
+    if (flushFailure) throw flushFailure;
+    if (failure) throw failure;
+    return { failedSpanIds };
+  };
+
   const provider = new NodeTracerProvider({
     resource: buildResource(),
-    spanProcessors: [spanProcessor],
+    spanProcessors: [trackedProcessor],
     idGenerator: {
       generateTraceId: () => seededId("trace"),
       generateSpanId: () => seededId("span"),
@@ -86,10 +151,17 @@ export function setupInstrumentation(config: Config): Instrumentation {
   setLangfuseTracerProvider(provider);
 
   return {
+    flush,
     shutdown: async () => {
-      await spanProcessor.forceFlush();
-      await provider.shutdown();
-      setLangfuseTracerProvider(null);
+      try {
+        await flush();
+      } finally {
+        try {
+          await provider.shutdown();
+        } finally {
+          setLangfuseTracerProvider(null);
+        }
+      }
     },
   };
 }