From 171a45707979fdfe073b145d9a15243ee17035eb Mon Sep 17 00:00:00 2001 From: Zaid Dreakh Date: Tue, 6 Oct 2026 13:04:08 +0200 Subject: [PATCH 1/2] fix: stream large Codex rollouts with durable checkpoints --- plugins/tracing/src/index.ts | 12 +- plugins/tracing/src/instrumentation.ts | 69 +++- plugins/tracing/src/parse.ts | 64 ++-- plugins/tracing/src/rollout.ts | 112 +++++++ plugins/tracing/src/sidecar.ts | 19 +- plugins/tracing/src/trace.ts | 300 +++++++++++------- plugins/tracing/test/instrumentation.test.ts | 11 +- plugins/tracing/test/rollout.test.ts | 123 +++++++ plugins/tracing/test/streaming-hook.test.ts | 240 ++++++++++++++ .../tracing/test/turn-emit-failure.test.ts | 2 +- scripts/benchmark-rollout.mjs | 126 ++++++++ 11 files changed, 920 insertions(+), 158 deletions(-) create mode 100644 plugins/tracing/src/rollout.ts create mode 100644 plugins/tracing/test/rollout.test.ts create mode 100644 plugins/tracing/test/streaming-hook.test.ts create mode 100644 scripts/benchmark-rollout.mjs diff --git a/plugins/tracing/src/index.ts b/plugins/tracing/src/index.ts index 2d535d8..17fa44c 100644 --- a/plugins/tracing/src/index.ts +++ b/plugins/tracing/src/index.ts @@ -51,14 +51,15 @@ export async function runHook(): Promise { } 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 { 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) => { diff --git a/plugins/tracing/src/instrumentation.ts b/plugins/tracing/src/instrumentation.ts index f7a3f05..14c69f6 100644 --- a/plugins/tracing/src/instrumentation.ts +++ b/plugins/tracing/src/instrumentation.ts @@ -2,6 +2,9 @@ import { createHash, randomBytes } from "node:crypto"; import { LangfuseSpanProcessor } from "@langfuse/otel"; import { setLangfuseTracerProvider } from "@langfuse/tracing"; +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,6 +40,8 @@ function seededId(kind: "trace" | "span"): string { } export type Instrumentation = { + /** Wait for delivery; reject even if an automatic batch failed earlier. */ + flush: () => Promise; /** Flush buffered spans and tear down the tracer provider. */ shutdown: () => Promise; }; @@ -57,13 +62,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 within 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 { + let ended = 0; + let delivered = 0; + 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 ?? "", + "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) delivered += spans.length; + 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 +100,27 @@ export function setupInstrumentation(config: Config): Instrumentation { shouldExportSpan: () => true, }); + const trackedProcessor: SpanProcessor = { + onStart: (span, context) => spanProcessor.onStart(span, context), + onEnd: (span) => { + ended++; + spanProcessor.onEnd(span); + }, + forceFlush: () => spanProcessor.forceFlush(), + shutdown: () => spanProcessor.shutdown(), + }; + + const flush = async () => { + await spanProcessor.forceFlush(); + if (exportFailure) throw exportFailure; + if (delivered !== ended) { + throw new Error(`Incomplete span export: confirmed ${delivered} of ${ended} spans`); + } + }; + const provider = new NodeTracerProvider({ resource: buildResource(), - spanProcessors: [spanProcessor], + spanProcessors: [trackedProcessor], idGenerator: { generateTraceId: () => seededId("trace"), generateSpanId: () => seededId("span"), @@ -86,10 +132,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); + } + } }, }; } diff --git a/plugins/tracing/src/parse.ts b/plugins/tracing/src/parse.ts index 5c70645..6340bf0 100644 --- a/plugins/tracing/src/parse.ts +++ b/plugins/tracing/src/parse.ts @@ -112,7 +112,7 @@ function newTurn(startTime: number): MutableTurn { }; } -function sessionMetaFrom(line: RolloutLine): SessionMeta { +export function sessionMetaFrom(line: RolloutLine): SessionMeta { const p = line.payload as RolloutLine["payload"] & { id?: string; cli_version?: string; @@ -131,15 +131,15 @@ function sessionMetaFrom(line: RolloutLine): SessionMeta { }; } -export function parseSession(lines: RolloutLine[]): { - sessionMeta: SessionMeta; - turns: Turn[]; -} { - const ownHeader = lines.find((line) => line.type === "session_meta"); - const sessionMeta: SessionMeta = ownHeader - ? sessionMetaFrom(ownHeader) - : { sessionId: "unknown" }; - const turns: Turn[] = []; +/** Incremental state machine; callers drain completed turns after each record. */ +export function createSessionParser(sessionMeta: SessionMeta, maxContextBytes = Infinity) { + let turns: Turn[] = []; + let contextBytes = Buffer.byteLength(sessionMeta.baseInstructions ?? ""); + const addContextBytes = (bytes: number) => { + contextBytes += bytes; + if (contextBytes > maxContextBytes) throw new Error("Rollout context exceeds memory limit"); + }; + addContextBytes(0); let turn: MutableTurn | null = null; let step: ModelStep | null = null; @@ -152,7 +152,10 @@ export function parseSession(lines: RolloutLine[]): { let exportedSegmentCount = 0; const collect = (into: string[], text: string) => { - if (!into.includes(text)) into.push(text); + if (!into.includes(text)) { + addContextBytes(Buffer.byteLength(text)); + into.push(text); + } }; const systemPromptFor = (): SystemPrompt | undefined => { @@ -216,7 +219,7 @@ export function parseSession(lines: RolloutLine[]): { toolCallsById = new Map(); }; - for (const line of lines) { + const processLine = (line: RolloutLine) => { const ts = Number.isFinite(Date.parse(line.timestamp)) ? Date.parse(line.timestamp) : lastTimestamp; @@ -229,7 +232,7 @@ export function parseSession(lines: RolloutLine[]): { const effort = typeof p.effort === "string" ? p.effort : p.reasoning_effort; if (typeof effort === "string") t.reasoningEffort = effort; t.invocationParams = line.payload as Record; - continue; + return; } if (line.type === "response_item") { @@ -339,6 +342,7 @@ export function parseSession(lines: RolloutLine[]): { const tool = raw as { name?: unknown; description?: unknown; parameters?: unknown }; if (typeof tool.name !== "string" || !tool.name) continue; if (toolDefinitions.some((known) => known.name === tool.name)) continue; + addContextBytes(Buffer.byteLength(JSON.stringify(tool))); toolDefinitions.push({ name: tool.name, ...(typeof tool.description === "string" ? { description: tool.description } : {}), @@ -355,7 +359,7 @@ export function parseSession(lines: RolloutLine[]): { s.reasoning = s.reasoning ? `${s.reasoning}\n${reasoning}` : reasoning; } } - continue; + return; } if (line.type === "event_msg") { @@ -366,11 +370,11 @@ export function parseSession(lines: RolloutLine[]): { if (turn) finishTurn(ts, { completed: false, aborted: false }); turn = newTurn(ts); turn.turnId = typeof p.turn_id === "string" ? p.turn_id : undefined; - continue; + return; } if (TURN_OPENING_EVENTS.has(et)) ensureTurn(ts); - else if (!turn) continue; + else if (!turn) return; if (et === "user_message" && typeof p.message === "string") { if (!turn!.userInput) turn!.userInput = p.message; @@ -434,11 +438,33 @@ export function parseSession(lines: RolloutLine[]): { } } } - continue; + return; } - } + }; - if (turn) finishTurn(lastTimestamp, { completed: false, aborted: false }); + const drain = (): Turn[] => { + const ready = turns; + turns = []; + return ready; + }; + return { + push(line: RolloutLine): Turn[] { + processLine(line); + return drain(); + }, + finish(): Turn[] { + if (turn) finishTurn(lastTimestamp, { completed: false, aborted: false }); + return drain(); + }, + }; +} +export function parseSession(lines: RolloutLine[]): { sessionMeta: SessionMeta; turns: Turn[] } { + const header = lines.find((line) => line.type === "session_meta"); + const sessionMeta = header ? sessionMetaFrom(header) : { sessionId: "unknown" }; + const parser = createSessionParser(sessionMeta); + const turns: Turn[] = []; + for (const line of lines) turns.push(...parser.push(line)); + turns.push(...parser.finish()); return { sessionMeta, turns }; } diff --git a/plugins/tracing/src/rollout.ts b/plugins/tracing/src/rollout.ts new file mode 100644 index 0000000..33b9e2d --- /dev/null +++ b/plugins/tracing/src/rollout.ts @@ -0,0 +1,112 @@ +import { createReadStream } from "node:fs"; + +import { createSessionParser, sessionMetaFrom } from "./parse.js"; +import type { RolloutLine, SessionMeta, Turn } from "./types.js"; + +// Limits apply before JSON decoding / retaining a turn. Fail with a retryable +// error rather than letting an individual record or turn exhaust the heap. +export const ROLLOUT_LIMITS = { + recordBytes: 8 * 1024 * 1024, + turnBytes: 32 * 1024 * 1024, + contextBytes: 4 * 1024 * 1024, + turns: 100_000, +}; + +/** A snapshot prevents a concurrently appended turn changing ordinals between passes. */ +export async function* readRollout( + file: string, + size: number, + maxRecordBytes = ROLLOUT_LIMITS.recordBytes, +): AsyncGenerator<{ line: RolloutLine; bytes: number }> { + if (size === 0) return; + const input = createReadStream(file, { end: size - 1, highWaterMark: 64 * 1024 }); + let parts: Buffer[] = []; + let bytes = 0; + let lineNumber = 1; + const append = (part: Buffer) => { + bytes += part.length; + if (bytes > maxRecordBytes) { + throw new Error(`Rollout record ${lineNumber} exceeds ${maxRecordBytes} bytes: ${file}`); + } + if (part.length) parts.push(part); + }; + const decode = () => { + const raw = Buffer.concat(parts, bytes).toString("utf8"); + const recordBytes = bytes; + parts = []; + bytes = 0; + lineNumber++; + let line: RolloutLine; + try { + line = JSON.parse(raw) as RolloutLine; + } catch { + return undefined; // Preserve the existing malformed-line tolerance. + } + if (!line || typeof line !== "object" || !line.payload || typeof line.payload !== "object") { + return undefined; + } + // Compaction carries potentially large context snapshots the parser never uses. + if (!["session_meta", "turn_context", "response_item", "event_msg"].includes(line.type)) { + return undefined; + } + return { line, bytes: recordBytes }; + }; + try { + for await (const chunk of input) { + const buffer = chunk as Buffer; + let start = 0; + let newline: number; + while ((newline = buffer.indexOf(10, start)) !== -1) { + append(buffer.subarray(start, newline)); + const record = decode(); + if (record) yield record; + start = newline + 1; + } + // Copy the trailing fragment so it does not pin an entire input chunk. + append(Buffer.from(buffer.subarray(start))); + } + if (bytes) { + const record = decode(); + if (record) yield record; + } + } finally { + input.destroy(); + } +} + +export async function readRolloutMeta(file: string, size: number): Promise { + for await (const { line } of readRollout(file, size)) { + if (line.type === "session_meta") return sessionMetaFrom(line); + } + return { sessionId: "unknown" }; +} + +export async function* readTurns( + file: string, + size: number, + sessionMeta: SessionMeta, +): AsyncGenerator { + const parser = createSessionParser(sessionMeta, ROLLOUT_LIMITS.contextBytes); + let turnBytes = 0; + let turnCount = 0; + for await (const { line, bytes } of readRollout(file, size)) { + if (line.type === "session_meta") continue; + if (line.type === "event_msg" && line.payload.type === "task_started") turnBytes = 0; + turnBytes += bytes; + if (turnBytes > ROLLOUT_LIMITS.turnBytes) { + throw new Error(`Rollout turn exceeds ${ROLLOUT_LIMITS.turnBytes} bytes: ${file}`); + } + const ready = parser.push(line); + for (const turn of ready) { + if (++turnCount > ROLLOUT_LIMITS.turns) throw new Error("Rollout has too many turns"); + yield turn; + } + if (ready.length && !(line.type === "event_msg" && line.payload.type === "task_started")) { + turnBytes = 0; + } + } + for (const turn of parser.finish()) { + if (++turnCount > ROLLOUT_LIMITS.turns) throw new Error("Rollout has too many turns"); + yield turn; + } +} diff --git a/plugins/tracing/src/sidecar.ts b/plugins/tracing/src/sidecar.ts index c525379..ff82ac1 100644 --- a/plugins/tracing/src/sidecar.ts +++ b/plugins/tracing/src/sidecar.ts @@ -1,4 +1,5 @@ import * as fs from "node:fs/promises"; +import { debugLog } from "./utils.js"; /** * The `Stop` hook reads the rollout file again after each Codex turn. @@ -7,8 +8,14 @@ import * as fs from "node:fs/promises"; */ export async function loadUploadedTurnIds(rolloutFile: string): Promise> { try { + if ((await fs.stat(`${rolloutFile}.langfuse`)).size > 16 * 1024 * 1024) { + throw new Error("Rollout upload sidecar exceeds memory limit"); + } const data = await fs.readFile(`${rolloutFile}.langfuse`, "utf-8"); - return new Set(data.split("\n").filter(Boolean)); + const lines = data.split("\n"); + // An interrupted append must not acknowledge a partial id. + lines.pop(); + return new Set(lines.filter(Boolean)); } catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") return new Set(); throw error; @@ -16,9 +23,15 @@ export async function loadUploadedTurnIds(rolloutFile: string): Promise { + let handle; try { - await fs.appendFile(`${rolloutFile}.langfuse`, `${turnId}\n`, "utf-8"); - } catch { + handle = await fs.open(`${rolloutFile}.langfuse`, "a"); + await handle.writeFile(`${turnId}\n`, "utf-8"); + await handle.sync(); + } catch (error) { // Best-effort: a failed write only risks a duplicate upload next time. + debugLog("failed to checkpoint uploaded turn; it may be retried:", error); + } finally { + await handle?.close(); } } diff --git a/plugins/tracing/src/trace.ts b/plugins/tracing/src/trace.ts index b67396a..11a8112 100644 --- a/plugins/tracing/src/trace.ts +++ b/plugins/tracing/src/trace.ts @@ -1,7 +1,6 @@ -import { createReadStream, type Dirent } from "node:fs"; +import { type Dirent } from "node:fs"; import * as fs from "node:fs/promises"; import * as path from "node:path"; -import { createInterface } from "node:readline"; import { createTraceId, @@ -15,7 +14,8 @@ import { TraceFlags, type SpanContext } from "@opentelemetry/api"; import type { Config } from "./config.js"; import { currentIdSeed, seedIds } from "./instrumentation.js"; -import { parseArgs, parseSession } from "./parse.js"; +import { parseArgs } from "./parse.js"; +import { readRollout, readRolloutMeta, readTurns } from "./rollout.js"; import { loadUploadedTurnIds } from "./sidecar.js"; import { skillsForToolCall, traceTags } from "./skills.js"; import type { @@ -31,20 +31,7 @@ import type { } from "./types.js"; import { debugLog, toText } from "./utils.js"; -async function loadSession(file: string): Promise { - const lines: RolloutLine[] = []; - const input = createReadStream(file, { encoding: "utf-8" }); - for await (const raw of createInterface({ input, crlfDelay: Infinity })) { - const trimmed = raw.trim(); - if (!trimmed) continue; - try { - lines.push(JSON.parse(trimmed) as RolloutLine); - } catch { - // skip malformed lines rather than aborting the whole upload - } - } - return lines; -} +const MAX_ROUTING_BYTES = 8 * 1024 * 1024; type SubagentRollout = { threadId: string; @@ -104,6 +91,7 @@ export async function buildSubagentIndex( const fromDay = path.relative(root, path.dirname(rolloutFile)); const bounded = !options.includeEarlierDays && /^\d{4}\/\d{2}\/\d{2}$/.test(fromDay); const index: SubagentIndex = { byParent: new Map(), byThread: new Map() }; + let indexBytes = 0; async function walk(dir: string, rel: string): Promise { if (bounded && rel && rel < fromDay.slice(0, rel.length)) return; @@ -130,6 +118,8 @@ export async function buildSubagentIndex( startTime: meta.startTime, nickname: meta.nickname, }; + indexBytes += Buffer.byteLength(JSON.stringify(rollout)); + if (indexBytes > MAX_ROUTING_BYTES) throw new Error("Subagent index exceeds memory limit"); index.byThread.set(meta.threadId, rollout); if (!meta.parentThreadId) continue; index.byParent.set(meta.parentThreadId, [ @@ -143,18 +133,31 @@ export async function buildSubagentIndex( return index; } -function turnIndexByNickname(turns: Turn[]): Map { - const byNickname = new Map(); - const ambiguous = new Set(); - turns.forEach((turn, index) => { - for (const tc of turn.steps.flatMap((s) => s.toolCalls)) { +type TurnSummary = Pick & { + nicknames: string[]; +}; + +function spawnNicknames(turn: Turn): string[] { + const names: string[] = []; + for (const step of turn.steps) { + for (const tc of step.toolCalls) { if (tc.name !== "spawn_agent" || tc.output == null) continue; const out = parseArgs(toText(tc.output)); const nickname = out !== null && typeof out === "object" ? (out as { nickname?: unknown }).nickname : undefined; - if (typeof nickname !== "string" || !nickname) continue; + if (typeof nickname === "string" && nickname) names.push(nickname); + } + } + return names; +} + +function turnIndexByNickname(turns: TurnSummary[]): Map { + const byNickname = new Map(); + const ambiguous = new Set(); + turns.forEach((turn, index) => { + for (const nickname of turn.nicknames) { if (byNickname.has(nickname) && byNickname.get(nickname) !== index) { ambiguous.add(nickname); } @@ -165,7 +168,7 @@ function turnIndexByNickname(turns: Turn[]): Map { return byNickname; } -function turnIndexAt(turns: Turn[], startTime: number): number { +function turnIndexAt(turns: TurnSummary[], startTime: number): number { const running = turns.findIndex((t) => startTime >= t.startTime && startTime <= t.endTime); if (running !== -1) return running; let last = 0; @@ -367,22 +370,33 @@ function turnHistoryMessages(turn: Turn): ChatMlMessage[] { return messages; } -function generationInput( - systemMessage: string | undefined, - historyPrefix: ChatMlMessage[], - turn: Turn, - stepIndex: number, -): ChatMlMessage[] | undefined { - const messages: ChatMlMessage[] = []; - if (systemMessage) messages.push({ role: "system", content: systemMessage }); - messages.push(...historyPrefix); - const user = userMessage(turn); - if (user) messages.push(user); - for (let j = 0; j < stepIndex; j++) { - messages.push(assistantMessage(turn.steps[j])); - messages.push(...toolMessages(turn.steps[j])); +export const MAX_GENERATION_INPUT_BYTES = 256 * 1024; +const OMITTED_INPUT = [ + { + role: "system" as const, + content: + "[Generation input omitted from trace: context exceeds 256 KiB. See turn and tool spans for evidence.]", + }, +]; + +/** Keep ordinary transcripts exact; stop retaining repeated history after the budget. */ +class GenerationHistory { + messages: ChatMlMessage[] = []; + bytes = 2; + omitted = false; + + append(messages: ChatMlMessage[]): void { + if (this.omitted) return; + for (const message of messages) { + this.bytes += Buffer.byteLength(JSON.stringify(message)) + 1; + if (this.bytes > MAX_GENERATION_INPUT_BYTES) { + this.messages = []; + this.omitted = true; + return; + } + this.messages.push(message); + } } - return messages.length > 0 ? messages : undefined; } type ContentPart = @@ -434,9 +448,11 @@ async function emitTurn( seenThreadIds: Set; unannouncedSubagents?: SubagentRollout[]; inheritableTurnIds?: ReadonlySet; - historyPrefix?: ChatMlMessage[]; + historyPrefix?: GenerationHistory; + flush?: () => Promise; + depth?: number; }, -): Promise { +): Promise { const isSubagent = sessionMeta.isSubagentThread === true || ctx.parentObservation != null; const outerSeed = currentIdSeed(); @@ -485,18 +501,29 @@ async function emitTurn( let failure: unknown; try { const systemMessage = systemPromptText(turn.systemPrompt); - const historyPrefix = ctx.historyPrefix ?? []; + const history = new GenerationHistory(); + if (systemMessage) history.append([{ role: "system", content: systemMessage }]); + if (ctx.historyPrefix?.omitted) history.omitted = true; + else history.append(ctx.historyPrefix?.messages ?? []); + const user = userMessage(turn); + if (user) history.append([user]); + // Tool schemas are attached to every generation too. + history.bytes += Buffer.byteLength(JSON.stringify(turn.toolDefinitions)); + if (history.bytes > MAX_GENERATION_INPUT_BYTES) history.omitted = true; const spawnObservations = new Map(); + let pendingSpans = 0; for (let i = 0; i < turn.steps.length; i++) { const step = turn.steps[i]; const generation = startObservation( isSubagent ? "LLM Subagent" : "LLM", { - input: attachToolDefinitions( - generationInput(systemMessage, historyPrefix, turn, i), - turn.toolDefinitions, - ), + input: history.omitted + ? OMITTED_INPUT + : attachToolDefinitions( + history.messages.length ? history.messages : undefined, + turn.toolDefinitions, + ), output: buildGenerationOutput(step), model: turn.model, ...(turn.reasoningEffort @@ -505,6 +532,7 @@ async function emitTurn( usageDetails: toUsageDetails(step.usage), metadata: { "codex.step_index": i, + ...(history.omitted ? { "codex.generation_input.omitted": true } : {}), "codex.reasoning_effort": turn.reasoningEffort, }, }, @@ -518,9 +546,18 @@ async function emitTurn( for (const tc of step.toolCalls) { const observation = emitToolCall(tc, root, step.endTime); if (tc.name === "spawn_agent") spawnObservations.set(tc.callId, observation); + if (++pendingSpans >= 32 && ctx.flush) { + await ctx.flush(); + pendingSpans = 0; + } } generation.end(new Date(generationEnd(step))); + if (++pendingSpans >= 32 && ctx.flush) { + await ctx.flush(); + pendingSpans = 0; + } + if (!history.omitted) history.append([assistantMessage(step), ...toolMessages(step)]); } const announced: SubagentRollout[] = []; @@ -542,6 +579,8 @@ async function emitTurn( subagentIndex: ctx.subagentIndex, seenThreadIds: ctx.seenThreadIds, ancestorTurnIds: ctx.inheritableTurnIds, + flush: ctx.flush, + depth: (ctx.depth ?? 0) + 1, }); seedIds(seedBeforeChild); } @@ -559,6 +598,7 @@ async function emitTurn( root.end(new Date(turn.endTime)); seedIds(outerSeed); if (failure && ctx.config.fail_on_error) throw failure; + return failure === undefined; } function emitToolCall( @@ -596,7 +636,7 @@ function isFinal( async function readTurnIds(file: string): Promise> { const ids = new Set(); - for (const line of await loadSession(file)) { + for await (const { line } of readRollout(file, (await fs.stat(file)).size)) { if (line.type !== "event_msg") continue; const p = line.payload as EventMsgPayload; if (p.type === "task_started" && typeof p.turn_id === "string") ids.add(p.turn_id); @@ -639,19 +679,32 @@ export async function convertRollout( seenThreadIds?: Set; ancestorTurnIds?: ReadonlySet; stoppedTurnId?: string; + flush?: () => Promise; + onTurnExported?: (turnId: string) => Promise; + depth?: number; }, ): Promise { - const { sessionMeta, turns } = parseSession(await loadSession(rolloutFile)); - - const historyPrefixes: ChatMlMessage[][] = []; - { - const seen: ChatMlMessage[] = []; - for (const turn of turns) { - historyPrefixes.push([...seen]); - seen.push(...turnHistoryMessages(turn)); - } + if ((options.depth ?? 0) > 16) throw new Error("Subagent nesting exceeds memory limit"); + const size = (await fs.stat(rolloutFile)).size; + const sessionMeta = await readRolloutMeta(rolloutFile, size); + // The first pass keeps only routing metadata. A later spawn announcement can + // disambiguate an earlier child, so routing must be resolved before emission. + const turns: TurnSummary[] = []; + let routingBytes = 0; + for await (const turn of readTurns(rolloutFile, size, sessionMeta)) { + const summary = { + turnId: turn.turnId, + startTime: turn.startTime, + endTime: turn.endTime, + subagentThreadIds: turn.subagentThreadIds, + nicknames: spawnNicknames(turn), + }; + routingBytes += Buffer.byteLength(JSON.stringify(summary)); + if (routingBytes > MAX_ROUTING_BYTES) throw new Error("Turn index exceeds memory limit"); + turns.push(summary); } - debugLog(`parsed ${turns.length} turn(s) from ${path.basename(rolloutFile)}`); + const history = new GenerationHistory(); + debugLog(`indexed ${turns.length} turn(s) from ${path.basename(rolloutFile)}`); const subagentIndex = options.subagentIndex ?? @@ -704,76 +757,87 @@ export async function convertRollout( if (turn.turnId) inheritableTurnIds.add(turn.turnId); } - if (options.parentObservation) { - for (let turnIndex = 0; turnIndex < turns.length; turnIndex++) { - const turn = turns[turnIndex]; - if (skipInherited(turn, turnIndex)) continue; - await emitTurn(turn, sessionMeta, { - config: options.config, - rolloutFile, - parentObservation: options.parentObservation, - subagentIndex, - seenThreadIds, - unannouncedSubagents: subagentsFor(turnIndex), - inheritableTurnIds, - historyPrefix: historyPrefixes[turnIndex], - }); - } - return []; - } - const uploaded = await loadUploadedTurnIds(rolloutFile); const exportedTurnIds: string[] = []; const attached = options.parentSpanContext != null; - for (let turnIndex = 0; turnIndex < turns.length; turnIndex++) { - const turn = turns[turnIndex]; + let turnIndex = -1; + for await (const turn of readTurns(rolloutFile, size, sessionMeta)) { + turnIndex++; + try { + if (skipInherited(turn, turnIndex)) continue; - if (skipInherited(turn, turnIndex)) continue; + if (options.parentObservation) { + const success = await emitTurn(turn, sessionMeta, { + config: options.config, + rolloutFile, + parentObservation: options.parentObservation, + subagentIndex, + seenThreadIds, + unannouncedSubagents: subagentsFor(turnIndex), + inheritableTurnIds, + historyPrefix: history, + flush: options.flush, + depth: options.depth, + }); + await options.flush?.(); + if (!success) throw new Error(`Subagent turn ${turn.turnId ?? "(no turn id)"} failed`); + continue; + } - if (!isFinal(turn, options.stoppedTurnId, turnIndex < turns.length - 1)) { - debugLog(`skipping turn ${turn.turnId ?? "(no turn id)"}: not final`); - continue; - } - if (uploaded.has(turn.turnId)) { - continue; - } + if (!isFinal(turn, options.stoppedTurnId, turnIndex < turns.length - 1)) { + debugLog(`skipping turn ${turn.turnId ?? "(no turn id)"}: not final`); + continue; + } + if (uploaded.has(turn.turnId)) { + continue; + } - const parentSpanContext = - options.parentSpanContext ?? - (await seededTraceParent(options.config, sessionMeta, turnIndex + 1)); - - const emit = () => - emitTurn(turn, sessionMeta, { - config: options.config, - rolloutFile, - parentSpanContext, - attached, - subagentIndex, - seenThreadIds, - unannouncedSubagents: subagentsFor(turnIndex), - inheritableTurnIds, - historyPrefix: historyPrefixes[turnIndex], - }); + const parentSpanContext = + options.parentSpanContext ?? + (await seededTraceParent(options.config, sessionMeta, turnIndex + 1)); + + const emit = () => + emitTurn(turn, sessionMeta, { + config: options.config, + rolloutFile, + parentSpanContext, + attached, + subagentIndex, + seenThreadIds, + unannouncedSubagents: subagentsFor(turnIndex), + inheritableTurnIds, + historyPrefix: history, + flush: options.flush, + depth: options.depth, + }); + + let success: boolean; + if (attached) { + success = await emit(); + } else { + const tags = traceTags(options.config, turn); + success = await propagateAttributes( + { + sessionId: sessionMeta.sessionId, + traceName: sessionMeta.isSubagentThread ? "Codex Subagent Turn" : "Codex Turn", + ...(options.config.user_id ? { userId: options.config.user_id } : {}), + ...(tags.length > 0 ? { tags } : {}), + ...(options.config.metadata ? { metadata: options.config.metadata } : {}), + }, + emit, + ); + } - if (attached) { - await emit(); - } else { - const tags = traceTags(options.config, turn); - await propagateAttributes( - { - sessionId: sessionMeta.sessionId, - traceName: sessionMeta.isSubagentThread ? "Codex Subagent Turn" : "Codex Turn", - ...(options.config.user_id ? { userId: options.config.user_id } : {}), - ...(tags.length > 0 ? { tags } : {}), - ...(options.config.metadata ? { metadata: options.config.metadata } : {}), - }, - emit, - ); + await options.flush?.(); + if (success) { + await options.onTurnExported?.(turn.turnId); + uploaded.add(turn.turnId); + exportedTurnIds.push(turn.turnId); + } + } finally { + if (!history.omitted) history.append(turnHistoryMessages(turn)); } - - uploaded.add(turn.turnId); - exportedTurnIds.push(turn.turnId); } return exportedTurnIds; diff --git a/plugins/tracing/test/instrumentation.test.ts b/plugins/tracing/test/instrumentation.test.ts index 444335d..4cd3428 100644 --- a/plugins/tracing/test/instrumentation.test.ts +++ b/plugins/tracing/test/instrumentation.test.ts @@ -4,7 +4,7 @@ import * as path from "node:path"; import { fileURLToPath } from "node:url"; import { getLangfuseTracerProvider, setLangfuseTracerProvider } from "@langfuse/tracing"; -import type { ReadableSpan, SpanProcessor } from "@opentelemetry/sdk-trace-base"; +import type { ReadableSpan, SpanProcessor, SpanExporter } from "@opentelemetry/sdk-trace-base"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { Config } from "../src/config.js"; @@ -17,15 +17,24 @@ const captured: ReadableSpan[] = []; vi.mock("@langfuse/otel", () => ({ LangfuseSpanProcessor: class implements SpanProcessor { + constructor(private params: { exporter: SpanExporter }) {} onStart(): void {} onEnd(span: ReadableSpan): void { captured.push(span); + this.params.exporter.export([span], () => {}); } async forceFlush(): Promise {} async shutdown(): Promise {} }, })); +vi.mock("@opentelemetry/exporter-trace-otlp-http", () => ({ + OTLPTraceExporter: class { + export: SpanExporter["export"] = (_spans, callback) => callback({ code: 0 }); + async shutdown(): Promise {} + }, +})); + const OTEL_API_KEY = Symbol.for("opentelemetry.js.api.1"); const baseConfig: Config = { diff --git a/plugins/tracing/test/rollout.test.ts b/plugins/tracing/test/rollout.test.ts new file mode 100644 index 0000000..c5bfc78 --- /dev/null +++ b/plugins/tracing/test/rollout.test.ts @@ -0,0 +1,123 @@ +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { fileURLToPath } from "node:url"; +import { afterEach, describe, expect, it } from "vitest"; + +import { parseSession } from "../src/parse.js"; +import { readRollout, readRolloutMeta, readTurns, ROLLOUT_LIMITS } from "../src/rollout.js"; +import type { RolloutLine } from "../src/types.js"; + +const dirs: string[] = []; +const limits = { ...ROLLOUT_LIMITS }; +afterEach(async () => { + Object.assign(ROLLOUT_LIMITS, limits); + await Promise.all(dirs.splice(0).map((dir) => fs.rm(dir, { recursive: true, force: true }))); +}); +const line = (type: string, payload: Record) => + JSON.stringify({ timestamp: "2026-10-05T00:00:00.000Z", type, payload }); +async function stage(contents: string) { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "lf-stream-")); + dirs.push(dir); + const file = path.join(dir, "rollout.jsonl"); + await fs.writeFile(file, contents); + return { file, size: Buffer.byteLength(contents) }; +} +async function records(file: string, size: number, limit?: number) { + const result = []; + for await (const record of readRollout(file, size, limit)) result.push(record.line); + return result; +} +async function turns(file: string, size: number) { + const meta = await readRolloutMeta(file, size); + const result = []; + for await (const turn of readTurns(file, size, meta)) result.push(turn); + return { sessionMeta: meta, turns: result }; +} + +describe("bounded rollout reader", () => { + it("matches the existing parser for every fixture including inherited and unfinished turns", async () => { + const root = path.join(path.dirname(fileURLToPath(import.meta.url)), "fixtures/sessions"); + const files = (await fs.readdir(root, { recursive: true })).filter((p) => p.endsWith(".jsonl")); + for (const relative of files) { + const file = path.join(root, relative); + const contents = await fs.readFile(file, "utf8"); + const lines = contents + .trim() + .split("\n") + .map((raw) => JSON.parse(raw) as RolloutLine); + expect(await turns(file, Buffer.byteLength(contents)), relative).toEqual(parseSession(lines)); + } + }); + + it("handles UTF-8 across chunk boundaries, CRLF, blank/malformed lines, and an unterminated final record", async () => { + const contents = + line("session_meta", { id: "test" }) + + "\r\n\nnot-json\nnull\n" + + line("event_msg", { type: "user_message", message: "界".repeat(70_000) }); + const { file, size } = await stage(contents); + const result = await records(file, size); + expect(result).toHaveLength(2); + expect(result[1].payload).toMatchObject({ message: "界".repeat(70_000) }); + }); + + it("does not read records appended after the snapshot", async () => { + const { file, size } = await stage(line("session_meta", { id: "first" }) + "\n"); + await fs.appendFile(file, line("session_meta", { id: "second" }) + "\n"); + expect(await records(file, size)).toHaveLength(1); + }); + + it("rejects an oversized record before JSON decoding, including a record without a newline", async () => { + for (const suffix of ["", "\n"]) { + const { file, size } = await stage("x".repeat(200_000) + suffix); + await expect(records(file, size, 100_000)).rejects.toThrow("record 1 exceeds 100000 bytes"); + } + }); + + it("does not retain compaction snapshots", async () => { + const { file, size } = await stage(line("compacted", { text: "x".repeat(100_000) })); + expect(await records(file, size)).toEqual([]); + }); + + it("bounds a turn composed of many individually small records", async () => { + ROLLOUT_LIMITS.turnBytes = 400; + const { file, size } = await stage( + [ + line("event_msg", { type: "task_started", turn_id: "t1" }), + ...Array.from({ length: 20 }, () => + line("event_msg", { type: "user_message", message: "hi" }), + ), + ].join("\n"), + ); + await expect(turns(file, size)).rejects.toThrow("Rollout turn exceeds 400 bytes"); + }); + + it("bounds context retained across turns", async () => { + ROLLOUT_LIMITS.contextBytes = 20; + const { file, size } = await stage( + [ + line("session_meta", { id: "s", base_instructions: { text: "1234567890" } }), + line("response_item", { + type: "message", + role: "developer", + content: [{ type: "text", text: "abcdefghijk" }], + }), + ].join("\n"), + ); + await expect(turns(file, size)).rejects.toThrow("Rollout context exceeds memory limit"); + }); + + it("bounds the number of indexed turns and propagates file errors", async () => { + ROLLOUT_LIMITS.turns = 1; + const { file, size } = await stage( + [ + line("event_msg", { type: "task_started", turn_id: "t1" }), + line("event_msg", { type: "task_complete" }), + line("event_msg", { type: "task_started", turn_id: "t2" }), + ].join("\n"), + ); + await expect(turns(file, size)).rejects.toThrow("too many turns"); + await fs.unlink(file); + await expect(records(file, size)).rejects.toThrow("ENOENT"); + }); +}); diff --git a/plugins/tracing/test/streaming-hook.test.ts b/plugins/tracing/test/streaming-hook.test.ts new file mode 100644 index 0000000..e9abaef --- /dev/null +++ b/plugins/tracing/test/streaming-hook.test.ts @@ -0,0 +1,240 @@ +import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process"; +import * as fs from "node:fs/promises"; +import * as http from "node:http"; +import * as os from "node:os"; +import * as path from "node:path"; +import { fileURLToPath } from "node:url"; +import { afterEach, describe, expect, it } from "vitest"; + +const bundle = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../dist/index.mjs"); +const dirs: string[] = []; +afterEach(async () => { + await Promise.all(dirs.splice(0).map((dir) => fs.rm(dir, { recursive: true, force: true }))); +}); + +async function stage(count = 3, outputBytes = 32) { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "lf-streaming-hook-")); + dirs.push(dir); + const sessions = path.join(dir, "sessions/2026/10/05"); + await fs.mkdir(sessions, { recursive: true }); + const file = path.join(sessions, "rollout.jsonl"); + const handle = await fs.open(file, "w"); + let n = 0; + const line = async (type: string, payload: Record) => { + await handle.writeFile( + JSON.stringify({ + timestamp: new Date(Date.UTC(2026, 9, 5) + n++ * 100).toISOString(), + type, + payload, + }) + "\n", + ); + }; + await line("session_meta", { id: "streaming-session" }); + for (let i = 1; i <= count; i++) { + await line("event_msg", { type: "task_started", turn_id: `turn-${i}` }); + await line("event_msg", { type: "user_message", message: `question ${i}` }); + await line("response_item", { + type: "function_call", + name: "exec_command", + call_id: `call-${i}`, + arguments: "{}", + }); + await line("response_item", { + type: "function_call_output", + call_id: `call-${i}`, + output: "x".repeat(outputBytes), + }); + await line("event_msg", { type: "token_count" }); + await line("response_item", { + type: "message", + role: "assistant", + content: [{ type: "text", text: `answer ${i}` }], + }); + await line("event_msg", { type: "task_complete" }); + } + await handle.close(); + return { dir, file }; +} + +type WireSpan = { + traceId: string; + spanId: string; + name: string; + attributes: Array<{ key: string; value: { stringValue?: string; boolValue?: boolean } }>; +}; +function parseSpans(body: string): WireSpan[] { + const payload = JSON.parse(body) as { + resourceSpans: Array<{ scopeSpans: Array<{ spans: WireSpan[] }> }>; + }; + return payload.resourceSpans.flatMap((r) => r.scopeSpans.flatMap((s) => s.spans)); +} +async function sidecar(file: string) { + return (await fs.readFile(file + ".langfuse", "utf8").catch(() => "")) + .split("\n") + .filter(Boolean); +} +function run(dir: string, file: string, port: number) { + const child = spawn(process.execPath, ["--max-old-space-size=128", bundle], { + cwd: dir, + env: { + ...process.env, + CODEX_HOME: dir, + OTEL_TRACES_SAMPLER: "always_on", + TRACE_TO_LANGFUSE: "true", + LANGFUSE_PUBLIC_KEY: "pk-test", + LANGFUSE_SECRET_KEY: "sk-test", + LANGFUSE_BASE_URL: `http://127.0.0.1:${port}`, + LANGFUSE_CODEX_FAIL_ON_ERROR: "true", + LANGFUSE_CODEX_DEBUG: "true", + LANGFUSE_CODEX_TRACE_SEED: "streaming-test", + }, + stdio: ["pipe", "pipe", "pipe"], + }); + let stderr = ""; + child.stderr.on("data", (chunk) => { + stderr += chunk; + }); + const result = new Promise<{ + code: number | null; + signal: NodeJS.Signals | null; + stderr: string; + }>((resolve, reject) => { + const timeout = setTimeout(() => { + child.kill("SIGKILL"); + reject(new Error("hook timed out")); + }, 15_000); + child.once("error", (error) => { + clearTimeout(timeout); + reject(error); + }); + child.once("exit", (code, signal) => { + clearTimeout(timeout); + resolve({ code, signal, stderr }); + }); + }); + child.stdin.end(JSON.stringify({ transcript_path: file })); + return { child, result }; +} +async function receiver( + onRequest: (body: string, response: http.ServerResponse) => void | Promise, +) { + const server = http.createServer((req, res) => { + const chunks: Buffer[] = []; + req.on("data", (chunk: Buffer) => chunks.push(chunk)); + req.on("end", () => { + void onRequest(Buffer.concat(chunks).toString(), res); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + return { + port: (server.address() as { port: number }).port, + close: async () => { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); + }, + }; +} +const ok = (res: http.ServerResponse) => { + res.writeHead(200, { "content-type": "application/json" }); + res.end("{}"); +}; + +describe("streaming delivery checkpoints", () => { + for (const mode of ["rejection", "interruption"] as const) { + it(`retries unfinished turns after ${mode}, preserving ids and confirmed checkpoints`, async () => { + const { dir, file } = await stage(); + let child: ChildProcessWithoutNullStreams; + let attempt = 0; + const requests: WireSpan[][][] = [[], []]; + let marksAtFailure: string[] = []; + const server = await receiver(async (body, res) => { + requests[attempt].push(parseSpans(body)); + if (attempt === 0 && requests[0].length === 2) { + marksAtFailure = await sidecar(file); + if (mode === "interruption") { + child.kill("SIGTERM"); + res.destroy(); + } else { + res.writeHead(400); + res.end("rejected"); + } + } else ok(res); + }); + try { + const first = run(dir, file, server.port); + child = first.child; + const failure = await first.result; + expect(failure.code).not.toBe(0); + expect(marksAtFailure).toEqual(["turn-1"]); + expect(await sidecar(file)).toEqual(["turn-1"]); + attempt = 1; + const retry = await run(dir, file, server.port).result; + expect(retry.code, retry.stderr).toBe(0); + expect(await sidecar(file)).toEqual(["turn-1", "turn-2", "turn-3"]); + const identities = (spans: WireSpan[]) => + spans.map(({ traceId, spanId }) => ({ traceId, spanId })); + expect(identities(requests[1][0])).toEqual(identities(requests[0][1])); + expect(requests[1]).toHaveLength(2); + } finally { + await server.close(); + } + }); + } + + it("processes growing history under a 128 MiB heap and explicitly labels omitted repeated input", async () => { + const { dir, file } = await stage(120, 100_000); + let spanCount = 0; + let omitted = 0; + let intactToolOutputs = 0; + const server = await receiver((body, res) => { + const spans = parseSpans(body); + spanCount += spans.length; + omitted += spans.filter((s) => + s.attributes.some( + (a) => + a.key.endsWith("codex.generation_input.omitted") && + (a.value.boolValue === true || a.value.stringValue === "true"), + ), + ).length; + intactToolOutputs += spans.filter( + (s) => + s.name === "exec_command" && + s.attributes.some( + (a) => + a.key === "langfuse.observation.output" && + a.value.stringValue?.includes("x".repeat(100_000)), + ), + ).length; + ok(res); + }); + try { + const result = await run(dir, file, server.port).result; + expect(result.code, result.stderr).toBe(0); + expect(spanCount).toBe(120 * 4); + expect(intactToolOutputs).toBe(120); + expect(omitted).toBeGreaterThan(0); + expect(await sidecar(file)).toHaveLength(120); + } finally { + await server.close(); + } + }, 20_000); + + it("reports an oversized record without acknowledging the turn", async () => { + const { dir, file } = await stage(1); + await fs.appendFile(file, "x".repeat(8 * 1024 * 1024 + 1)); + let requests = 0; + const server = await receiver((_body, res) => { + requests++; + ok(res); + }); + try { + const result = await run(dir, file, server.port).result; + expect(result.code).toBe(1); + expect(result.stderr).toContain("exceeds 8388608 bytes"); + expect(requests).toBe(0); + expect(await sidecar(file)).toEqual([]); + } finally { + await server.close(); + } + }); +}); diff --git a/plugins/tracing/test/turn-emit-failure.test.ts b/plugins/tracing/test/turn-emit-failure.test.ts index 6c81fa9..bcc8d67 100644 --- a/plugins/tracing/test/turn-emit-failure.test.ts +++ b/plugins/tracing/test/turn-emit-failure.test.ts @@ -80,7 +80,7 @@ describe("a turn whose conversion fails", () => { it("is exported as an ERROR trace and does not take the other turns with it", async () => { const exported = await convertRollout(stageRollout(), { config: baseConfig }); - expect(exported).toEqual(["turn-1", "turn-2", "turn-3"]); + expect(exported).toEqual(["turn-1", "turn-3"]); const roots = exporter.getFinishedSpans().filter((s) => s.name === "Codex Turn"); const byTurn = Object.fromEntries( diff --git a/scripts/benchmark-rollout.mjs b/scripts/benchmark-rollout.mjs new file mode 100644 index 0000000..656491a --- /dev/null +++ b/scripts/benchmark-rollout.mjs @@ -0,0 +1,126 @@ +// Usage: node scripts/benchmark-rollout.mjs [bundle] [turns] [tool-output-bytes] +// Runs the real bundled hook under a 128 MiB heap against a local OTLP receiver. +import { spawn } from "node:child_process"; +import { createServer } from "node:http"; +import { mkdtemp, mkdir, open, rm, stat, readFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { pathToFileURL } from "node:url"; + +const bundle = resolve(process.argv[2] ?? "plugins/tracing/dist/index.mjs"); +const turns = Number(process.argv[3] ?? 200); +const outputBytes = Number(process.argv[4] ?? 200_000); +const dir = await mkdtemp(join(tmpdir(), "codex-rollout-benchmark-")); +const sessions = join(dir, "sessions/2026/10/05"); +await mkdir(sessions, { recursive: true }); +const rollout = join(sessions, "rollout.jsonl"); +const file = await open(rollout, "w"); +let recordIndex = 0; +const line = async (type, payload) => + file.writeFile( + JSON.stringify({ + timestamp: new Date(Date.UTC(2026, 9, 5) + recordIndex++ * 100).toISOString(), + type, + payload, + }) + "\n", + ); +await line("session_meta", { id: "benchmark-session" }); +for (let i = 0; i < turns; i++) { + await line("event_msg", { type: "task_started", turn_id: `turn-${i}` }); + await line("event_msg", { type: "user_message", message: `Inspect file ${i}` }); + await line("response_item", { + type: "function_call", + name: "exec_command", + call_id: `call-${i}`, + arguments: '{"cmd":"cat example.txt"}', + }); + await line("response_item", { + type: "function_call_output", + call_id: `call-${i}`, + output: "x".repeat(outputBytes), + }); + await line("event_msg", { type: "token_count" }); + await line("response_item", { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: `File ${i} inspected.` }], + }); + await line("event_msg", { type: "task_complete" }); +} +await file.close(); +let requests = 0; +let spans = 0; +const server = createServer((req, res) => { + const chunks = []; + req.on("data", (chunk) => chunks.push(chunk)); + req.on("end", () => { + requests++; + const payload = JSON.parse(Buffer.concat(chunks).toString()); + for (const resource of payload.resourceSpans ?? []) { + for (const scope of resource.scopeSpans ?? []) spans += scope.spans.length; + } + res.writeHead(200, { "Content-Type": "application/json" }); + res.end("{}"); + }); +}); +await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); +try { + const started = performance.now(); + const child = spawn( + process.execPath, + [ + "--max-old-space-size=128", + "--input-type=module", + "-e", + `process.on('exit', () => console.error('MAX_RSS_KIB=' + process.resourceUsage().maxRSS)); await import(${JSON.stringify(pathToFileURL(bundle).href)});`, + ], + { + cwd: dir, + env: { + ...process.env, + CODEX_HOME: dir, + OTEL_TRACES_SAMPLER: "always_on", + TRACE_TO_LANGFUSE: "true", + LANGFUSE_PUBLIC_KEY: "pk-test", + LANGFUSE_SECRET_KEY: "sk-test", + LANGFUSE_BASE_URL: `http://127.0.0.1:${server.address().port}`, + LANGFUSE_CODEX_FAIL_ON_ERROR: "true", + LANGFUSE_CODEX_DEBUG: "true", + }, + stdio: ["pipe", "ignore", "pipe"], + }, + ); + let stderr = ""; + child.stderr.on("data", (chunk) => { + stderr = (stderr + chunk).slice(-16000); + }); + child.stdin.end(JSON.stringify({ transcript_path: rollout, turn_id: `turn-${turns - 1}` })); + const timeout = setTimeout(() => child.kill("SIGKILL"), 120_000); + const result = await new Promise((resolve, reject) => { + child.once("error", reject); + child.once("exit", (code, signal) => resolve({ code, signal })); + }); + clearTimeout(timeout); + const marks = await readFile(rollout + ".langfuse", "utf8").catch(() => ""); + console.log( + JSON.stringify( + { + bundle, + turns, + rolloutBytes: (await stat(rollout)).size, + elapsedMs: Math.round(performance.now() - started), + ...result, + peakRssKiB: Number(/MAX_RSS_KIB=(\d+)/.exec(stderr)?.[1]) || null, + requests, + spans, + checkpointedTurns: marks.split("\n").filter(Boolean).length, + ...(result.code !== 0 ? { error: stderr.slice(-3000) } : {}), + }, + null, + 2, + ), + ); +} finally { + await new Promise((resolve) => server.close(resolve)); + await rm(dir, { recursive: true, force: true }); +} From 2fc48db773a6017369ce04a9f1f83ed0e8231787 Mon Sep 17 00:00:00 2001 From: Zaid Dreakh Date: Tue, 6 Oct 2026 13:57:18 +0200 Subject: [PATCH 2/2] fix: preserve tracing through rollout limits and batch delivery --- README.md | 8 + plugins/tracing/src/export-batch.ts | 64 ++++++ plugins/tracing/src/instrumentation.ts | 41 +++- plugins/tracing/src/parse.ts | 25 ++- plugins/tracing/src/rollout.ts | 41 +++- plugins/tracing/src/trace.ts | 232 ++++++++++++++------ plugins/tracing/src/types.ts | 3 + plugins/tracing/test/delivery.test.ts | 110 ++++++++++ plugins/tracing/test/export-batch.test.ts | 84 +++++++ plugins/tracing/test/rollout.test.ts | 43 +++- plugins/tracing/test/streaming-hook.test.ts | 105 +++++++-- plugins/tracing/test/trace.test.ts | 53 ++++- scripts/benchmark-rollout.mjs | 7 +- 13 files changed, 693 insertions(+), 123 deletions(-) create mode 100644 plugins/tracing/src/export-batch.ts create mode 100644 plugins/tracing/test/delivery.test.ts create mode 100644 plugins/tracing/test/export-batch.test.ts diff --git a/README.md b/README.md index 774f3e7..2001790 100644 --- a/README.md +++ b/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). diff --git a/plugins/tracing/src/export-batch.ts b/plugins/tracing/src/export-batch.ts new file mode 100644 index 0000000..8a180d6 --- /dev/null +++ b/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; + success: boolean; + checkpoint: (id: string) => Promise; +}; + +/** 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(); + private transportFailure: Error | undefined; + incomplete = false; + + constructor(private deliver?: () => Promise) {} + + async ended(observation: LangfuseObservation, spanIds: Set): Promise { + spanIds.add(spanIdentity(observation.otelSpan.spanContext())); + this.spans++; + const { attributes } = observation.otelSpan as unknown as Pick; + 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 { + // 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. + } +} diff --git a/plugins/tracing/src/instrumentation.ts b/plugins/tracing/src/instrumentation.ts index 14c69f6..e9d4c7f 100644 --- a/plugins/tracing/src/instrumentation.ts +++ b/plugins/tracing/src/instrumentation.ts @@ -2,6 +2,7 @@ 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"; @@ -41,11 +42,17 @@ function seededId(kind: "trace" | "span"): string { export type Instrumentation = { /** Wait for delivery; reject even if an automatic batch failed earlier. */ - flush: () => Promise; + flush: () => Promise; /** Flush buffered spans and tear down the tracer provider. */ shutdown: () => Promise; }; +export type ExportReceipt = { failedSpanIds: ReadonlySet }; + +export function spanIdentity(context: SpanContext): string { + return `${context.traceId}:${context.spanId}`; +} + const DEFAULT_SERVICE_NAME = "codex"; function buildResource() { @@ -62,14 +69,13 @@ function buildResource() { * instrumentation loading. Registering the provider also installs the * AsyncLocalStorage context manager that `propagateAttributes` relies on. * - * Batch within turns, with explicit checkpoints to bound the export backlog. + * 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 { - let ended = 0; - let delivered = 0; + const pending = new Set(); let exportFailure: Error | undefined; const exporter = new OTLPTraceExporter({ url: `${config.base_url.replace(/\/$/, "")}/api/public/otel/v1/traces`, @@ -83,8 +89,9 @@ export function setupInstrumentation(config: Config): Instrumentation { const trackedExporter: SpanExporter = { export: (spans, callback) => { exporter.export(spans, (result) => { - if (result.code === ExportResultCode.SUCCESS) delivered += spans.length; - else exportFailure ??= result.error ?? new Error("Span export failed"); + 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); }); }, @@ -103,7 +110,9 @@ export function setupInstrumentation(config: Config): Instrumentation { const trackedProcessor: SpanProcessor = { onStart: (span, context) => spanProcessor.onStart(span, context), onEnd: (span) => { - ended++; + if (span.spanContext().traceFlags & TraceFlags.SAMPLED) { + pending.add(spanIdentity(span.spanContext())); + } spanProcessor.onEnd(span); }, forceFlush: () => spanProcessor.forceFlush(), @@ -111,11 +120,21 @@ export function setupInstrumentation(config: Config): Instrumentation { }; const flush = async () => { - await spanProcessor.forceFlush(); - if (exportFailure) throw exportFailure; - if (delivered !== ended) { - throw new Error(`Incomplete span export: confirmed ${delivered} of ${ended} spans`); + 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({ diff --git a/plugins/tracing/src/parse.ts b/plugins/tracing/src/parse.ts index 6340bf0..15ae8b9 100644 --- a/plugins/tracing/src/parse.ts +++ b/plugins/tracing/src/parse.ts @@ -135,11 +135,17 @@ export function sessionMetaFrom(line: RolloutLine): SessionMeta { export function createSessionParser(sessionMeta: SessionMeta, maxContextBytes = Infinity) { let turns: Turn[] = []; let contextBytes = Buffer.byteLength(sessionMeta.baseInstructions ?? ""); + let contextTruncated = contextBytes > maxContextBytes; + const baseInstructions = contextTruncated ? undefined : sessionMeta.baseInstructions; + if (contextTruncated) contextBytes = 0; const addContextBytes = (bytes: number) => { + if (contextBytes + bytes > maxContextBytes) { + contextTruncated = true; + return false; + } contextBytes += bytes; - if (contextBytes > maxContextBytes) throw new Error("Rollout context exceeds memory limit"); + return true; }; - addContextBytes(0); let turn: MutableTurn | null = null; let step: ModelStep | null = null; @@ -152,21 +158,21 @@ export function createSessionParser(sessionMeta: SessionMeta, maxContextBytes = let exportedSegmentCount = 0; const collect = (into: string[], text: string) => { - if (!into.includes(text)) { - addContextBytes(Buffer.byteLength(text)); + if (!into.includes(text) && addContextBytes(Buffer.byteLength(text))) { into.push(text); } }; const systemPromptFor = (): SystemPrompt | undefined => { - const base = sessionMeta.baseInstructions; + const base = baseInstructions; const segmentCount = (base ? 1 : 0) + developerMessages.length + injectedContext.length; - if (segmentCount === 0) return undefined; + if (segmentCount === 0 && !contextTruncated) return undefined; const snapshot: SystemPrompt = { baseInstructions: base, developerMessages: [...developerMessages], injectedContext: [...injectedContext], changed: segmentCount !== exportedSegmentCount, + ...(contextTruncated ? { truncated: true } : {}), }; exportedSegmentCount = segmentCount; return snapshot; @@ -342,7 +348,7 @@ export function createSessionParser(sessionMeta: SessionMeta, maxContextBytes = const tool = raw as { name?: unknown; description?: unknown; parameters?: unknown }; if (typeof tool.name !== "string" || !tool.name) continue; if (toolDefinitions.some((known) => known.name === tool.name)) continue; - addContextBytes(Buffer.byteLength(JSON.stringify(tool))); + if (!addContextBytes(Buffer.byteLength(JSON.stringify(tool)))) continue; toolDefinitions.push({ name: tool.name, ...(typeof tool.description === "string" ? { description: tool.description } : {}), @@ -366,6 +372,11 @@ export function createSessionParser(sessionMeta: SessionMeta, maxContextBytes = const p = line.payload as EventMsgPayload; const et = p.type; + if (et === "langfuse_record_omitted") { + ensureTurn(ts).truncated = true; + return; + } + if (et === "task_started") { if (turn) finishTurn(ts, { completed: false, aborted: false }); turn = newTurn(ts); diff --git a/plugins/tracing/src/rollout.ts b/plugins/tracing/src/rollout.ts index 33b9e2d..8668499 100644 --- a/plugins/tracing/src/rollout.ts +++ b/plugins/tracing/src/rollout.ts @@ -2,9 +2,10 @@ import { createReadStream } from "node:fs"; import { createSessionParser, sessionMetaFrom } from "./parse.js"; import type { RolloutLine, SessionMeta, Turn } from "./types.js"; +import { debugLog } from "./utils.js"; -// Limits apply before JSON decoding / retaining a turn. Fail with a retryable -// error rather than letting an individual record or turn exhaust the heap. +// Limits apply before JSON decoding / retaining a turn. Omit oversized data, +// not the whole session, and explicitly label affected turns. export const ROLLOUT_LIMITS = { recordBytes: 8 * 1024 * 1024, turnBytes: 32 * 1024 * 1024, @@ -22,15 +23,34 @@ export async function* readRollout( const input = createReadStream(file, { end: size - 1, highWaterMark: 64 * 1024 }); let parts: Buffer[] = []; let bytes = 0; + let oversized = false; let lineNumber = 1; const append = (part: Buffer) => { bytes += part.length; + if (oversized) return; if (bytes > maxRecordBytes) { - throw new Error(`Rollout record ${lineNumber} exceeds ${maxRecordBytes} bytes: ${file}`); + debugLog(`Omitting rollout record ${lineNumber}: exceeds ${maxRecordBytes} bytes: ${file}`); + oversized = true; + parts = []; + return; } if (part.length) parts.push(part); }; const decode = () => { + if (oversized) { + oversized = false; + parts = []; + bytes = 0; + lineNumber++; + return { + line: { + timestamp: "", + type: "event_msg", + payload: { type: "langfuse_record_omitted" }, + } as RolloutLine, + bytes: 0, + }; + } const raw = Buffer.concat(parts, bytes).toString("utf8"); const recordBytes = bytes; parts = []; @@ -88,17 +108,23 @@ export async function* readTurns( ): AsyncGenerator { const parser = createSessionParser(sessionMeta, ROLLOUT_LIMITS.contextBytes); let turnBytes = 0; - let turnCount = 0; for await (const { line, bytes } of readRollout(file, size)) { if (line.type === "session_meta") continue; if (line.type === "event_msg" && line.payload.type === "task_started") turnBytes = 0; turnBytes += bytes; - if (turnBytes > ROLLOUT_LIMITS.turnBytes) { - throw new Error(`Rollout turn exceeds ${ROLLOUT_LIMITS.turnBytes} bytes: ${file}`); + const boundary = + line.type === "event_msg" && + ["task_started", "task_complete", "turn_aborted"].includes(String(line.payload.type)); + if (turnBytes > ROLLOUT_LIMITS.turnBytes && !boundary) { + parser.push({ + timestamp: line.timestamp, + type: "event_msg", + payload: { type: "langfuse_record_omitted" }, + } as RolloutLine); + continue; } const ready = parser.push(line); for (const turn of ready) { - if (++turnCount > ROLLOUT_LIMITS.turns) throw new Error("Rollout has too many turns"); yield turn; } if (ready.length && !(line.type === "event_msg" && line.payload.type === "task_started")) { @@ -106,7 +132,6 @@ export async function* readTurns( } } for (const turn of parser.finish()) { - if (++turnCount > ROLLOUT_LIMITS.turns) throw new Error("Rollout has too many turns"); yield turn; } } diff --git a/plugins/tracing/src/trace.ts b/plugins/tracing/src/trace.ts index 11a8112..1eea6c6 100644 --- a/plugins/tracing/src/trace.ts +++ b/plugins/tracing/src/trace.ts @@ -13,9 +13,10 @@ import { import { TraceFlags, type SpanContext } from "@opentelemetry/api"; import type { Config } from "./config.js"; -import { currentIdSeed, seedIds } from "./instrumentation.js"; +import { ExportBatch } from "./export-batch.js"; +import { currentIdSeed, seedIds, type ExportReceipt } from "./instrumentation.js"; import { parseArgs } from "./parse.js"; -import { readRollout, readRolloutMeta, readTurns } from "./rollout.js"; +import { readRollout, readRolloutMeta, readTurns, ROLLOUT_LIMITS } from "./rollout.js"; import { loadUploadedTurnIds } from "./sidecar.js"; import { skillsForToolCall, traceTags } from "./skills.js"; import type { @@ -44,6 +45,7 @@ type SubagentRollout = { export type SubagentIndex = { byParent: Map; byThread: Map; + truncated?: boolean; }; async function readSessionMeta( @@ -118,8 +120,12 @@ export async function buildSubagentIndex( startTime: meta.startTime, nickname: meta.nickname, }; - indexBytes += Buffer.byteLength(JSON.stringify(rollout)); - if (indexBytes > MAX_ROUTING_BYTES) throw new Error("Subagent index exceeds memory limit"); + const bytes = Buffer.byteLength(JSON.stringify(rollout)); + if (indexBytes + bytes > MAX_ROUTING_BYTES) { + index.truncated = true; + continue; + } + indexBytes += bytes; index.byThread.set(meta.threadId, rollout); if (!meta.parentThreadId) continue; index.byParent.set(meta.parentThreadId, [ @@ -309,6 +315,7 @@ function systemPromptMetadata(systemPrompt: SystemPrompt): Record + messages.reduce((bytes, message) => bytes + Buffer.byteLength(JSON.stringify(message)) + 1, 0); + +/** Evict whole turns/steps so a retained tool result keeps its assistant call. */ class GenerationHistory { - messages: ChatMlMessage[] = []; - bytes = 2; - omitted = false; + private groups: HistoryGroup[] = []; + private bytes = 0; + truncated = false; append(messages: ChatMlMessage[]): void { - if (this.omitted) return; - for (const message of messages) { - this.bytes += Buffer.byteLength(JSON.stringify(message)) + 1; - if (this.bytes > MAX_GENERATION_INPUT_BYTES) { - this.messages = []; - this.omitted = true; - return; - } - this.messages.push(message); + const bytes = messageBytes(messages); + if (!messages.length) return; + if (bytes > MAX_GENERATION_INPUT_BYTES - 2) { + // One oversized group cannot fit, but later small groups still can. + this.truncated = true; + return; + } + this.groups.push({ messages, bytes }); + this.bytes += bytes; + while (this.bytes > MAX_GENERATION_INPUT_BYTES - 2) { + this.bytes -= this.groups.shift()!.bytes; + this.truncated = true; } } + + recent(budget: number): { messages: ChatMlMessage[]; bytes: number; truncated: boolean } { + let bytes = 0; + let first = this.groups.length; + while (first > 0 && bytes + this.groups[first - 1].bytes <= budget) { + bytes += this.groups[--first].bytes; + } + return { + messages: this.groups.slice(first).flatMap((group) => group.messages), + bytes, + truncated: this.truncated || first > 0, + }; + } +} + +function generationInput( + system: string | undefined, + user: ChatMlMessage | undefined, + prefix: GenerationHistory | undefined, + steps: GenerationHistory, + definitions: ToolDefinition[], +): { input: unknown; truncated: boolean } { + let pinned: ChatMlMessage[] = [ + ...(system ? [{ role: "system" as const, content: system }] : []), + ...(user ? [user] : []), + ]; + let truncated = false; + // Reserve room for recent history even when a single prompt is enormous. + if (messageBytes(pinned) > MAX_GENERATION_INPUT_BYTES / 2) { + pinned = pinned.map((message) => { + if (messageBytes([message]) <= MAX_GENERATION_INPUT_BYTES / 4) return message; + let text = "content" in message ? toText(message.content) : toText(message); + let bounded: ChatMlMessage; + do { + text = text.slice(0, Math.floor(text.length / 2)); + bounded = { + role: message.role === "system" ? "system" : "user", + content: `${text}\n[Input truncated; see turn span.]`, + }; + } while (messageBytes([bounded]) > MAX_GENERATION_INPUT_BYTES / 4); + return bounded; + }); + truncated = true; + } + let tools = definitions; + let toolsBytes = tools.length ? Buffer.byteLength(JSON.stringify(tools)) + 10 : 0; + if (toolsBytes > MAX_GENERATION_INPUT_BYTES / 4) { + tools = []; + toolsBytes = 0; + truncated = true; + } + const budget = MAX_GENERATION_INPUT_BYTES - messageBytes(pinned) - toolsBytes - 2; + const current = steps.recent(budget); + const previous = prefix?.recent(budget - current.bytes); + const messages = [ + ...pinned.filter((message) => message.role === "system"), + ...(previous?.messages ?? []), + ...pinned.filter((message) => message.role === "user"), + ...current.messages, + ]; + return { + input: attachToolDefinitions(messages.length ? messages : undefined, tools), + truncated: truncated || current.truncated || previous?.truncated === true, + }; } type ContentPart = @@ -449,7 +520,9 @@ async function emitTurn( unannouncedSubagents?: SubagentRollout[]; inheritableTurnIds?: ReadonlySet; historyPrefix?: GenerationHistory; - flush?: () => Promise; + batch: ExportBatch; + spanIds: Set; + routingTruncated?: boolean; depth?: number; }, ): Promise { @@ -474,6 +547,10 @@ async function emitTurn( "codex.cli_version": sessionMeta.cliVersion, "codex.aborted": turn.aborted, "codex.tool_call_count": turn.steps.reduce((n, s) => n + s.toolCalls.length, 0), + ...(turn.truncated ? { "codex.rollout.truncated": true } : {}), + ...(ctx.routingTruncated || ctx.subagentIndex.truncated + ? { "codex.routing.truncated": true } + : {}), ...(ctx.attached && ctx.parentSpanContext ? { "codex.parent_trace_id": ctx.parentSpanContext.traceId, @@ -502,28 +579,22 @@ async function emitTurn( try { const systemMessage = systemPromptText(turn.systemPrompt); const history = new GenerationHistory(); - if (systemMessage) history.append([{ role: "system", content: systemMessage }]); - if (ctx.historyPrefix?.omitted) history.omitted = true; - else history.append(ctx.historyPrefix?.messages ?? []); const user = userMessage(turn); - if (user) history.append([user]); - // Tool schemas are attached to every generation too. - history.bytes += Buffer.byteLength(JSON.stringify(turn.toolDefinitions)); - if (history.bytes > MAX_GENERATION_INPUT_BYTES) history.omitted = true; const spawnObservations = new Map(); - let pendingSpans = 0; for (let i = 0; i < turn.steps.length; i++) { const step = turn.steps[i]; + const input = generationInput( + systemMessage, + user, + ctx.historyPrefix, + history, + turn.toolDefinitions, + ); const generation = startObservation( isSubagent ? "LLM Subagent" : "LLM", { - input: history.omitted - ? OMITTED_INPUT - : attachToolDefinitions( - history.messages.length ? history.messages : undefined, - turn.toolDefinitions, - ), + input: input.input, output: buildGenerationOutput(step), model: turn.model, ...(turn.reasoningEffort @@ -532,7 +603,7 @@ async function emitTurn( usageDetails: toUsageDetails(step.usage), metadata: { "codex.step_index": i, - ...(history.omitted ? { "codex.generation_input.omitted": true } : {}), + ...(input.truncated ? { "codex.generation_input.truncated": true } : {}), "codex.reasoning_effort": turn.reasoningEffort, }, }, @@ -546,18 +617,12 @@ async function emitTurn( for (const tc of step.toolCalls) { const observation = emitToolCall(tc, root, step.endTime); if (tc.name === "spawn_agent") spawnObservations.set(tc.callId, observation); - if (++pendingSpans >= 32 && ctx.flush) { - await ctx.flush(); - pendingSpans = 0; - } + await ctx.batch.ended(observation, ctx.spanIds); } generation.end(new Date(generationEnd(step))); - if (++pendingSpans >= 32 && ctx.flush) { - await ctx.flush(); - pendingSpans = 0; - } - if (!history.omitted) history.append([assistantMessage(step), ...toolMessages(step)]); + await ctx.batch.ended(generation, ctx.spanIds); + history.append([assistantMessage(step), ...toolMessages(step)]); } const announced: SubagentRollout[] = []; @@ -579,7 +644,8 @@ async function emitTurn( subagentIndex: ctx.subagentIndex, seenThreadIds: ctx.seenThreadIds, ancestorTurnIds: ctx.inheritableTurnIds, - flush: ctx.flush, + batch: ctx.batch, + spanIds: ctx.spanIds, depth: (ctx.depth ?? 0) + 1, }); seedIds(seedBeforeChild); @@ -596,7 +662,11 @@ async function emitTurn( } root.end(new Date(turn.endTime)); - seedIds(outerSeed); + try { + await ctx.batch.ended(root, ctx.spanIds); + } finally { + seedIds(outerSeed); + } if (failure && ctx.config.fail_on_error) throw failure; return failure === undefined; } @@ -679,9 +749,11 @@ export async function convertRollout( seenThreadIds?: Set; ancestorTurnIds?: ReadonlySet; stoppedTurnId?: string; - flush?: () => Promise; + flush?: () => Promise; onTurnExported?: (turnId: string) => Promise; depth?: number; + batch?: ExportBatch; + spanIds?: Set; }, ): Promise { if ((options.depth ?? 0) > 16) throw new Error("Subagent nesting exceeds memory limit"); @@ -691,7 +763,11 @@ export async function convertRollout( // disambiguate an earlier child, so routing must be resolved before emission. const turns: TurnSummary[] = []; let routingBytes = 0; + let totalTurns = 0; + let routingTruncated = false; for await (const turn of readTurns(rolloutFile, size, sessionMeta)) { + totalTurns++; + if (routingTruncated) continue; const summary = { turnId: turn.turnId, startTime: turn.startTime, @@ -699,12 +775,16 @@ export async function convertRollout( subagentThreadIds: turn.subagentThreadIds, nicknames: spawnNicknames(turn), }; - routingBytes += Buffer.byteLength(JSON.stringify(summary)); - if (routingBytes > MAX_ROUTING_BYTES) throw new Error("Turn index exceeds memory limit"); + const bytes = Buffer.byteLength(JSON.stringify(summary)); + if (turns.length >= ROLLOUT_LIMITS.turns || routingBytes + bytes > MAX_ROUTING_BYTES) { + routingTruncated = true; + continue; + } + routingBytes += bytes; turns.push(summary); } const history = new GenerationHistory(); - debugLog(`indexed ${turns.length} turn(s) from ${path.basename(rolloutFile)}`); + debugLog(`indexed ${turns.length} of ${totalTurns} turn(s) from ${path.basename(rolloutFile)}`); const subagentIndex = options.subagentIndex ?? @@ -716,7 +796,10 @@ export async function convertRollout( const announced = new Set(turns.flatMap((t) => t.subagentThreadIds)); const unannounced = (subagentIndex.byParent.get(sessionMeta.sessionId) ?? []).filter( - (s) => !announced.has(s.threadId) && !seenThreadIds.has(s.threadId), + (s) => + !announced.has(s.threadId) && + !seenThreadIds.has(s.threadId) && + (!routingTruncated || s.startTime <= (turns.at(-1)?.endTime ?? 0)), ); const spawnTurnOf = unannounced.length > 0 ? turnIndexByNickname(turns) : new Map(); @@ -760,6 +843,7 @@ export async function convertRollout( const uploaded = await loadUploadedTurnIds(rolloutFile); const exportedTurnIds: string[] = []; const attached = options.parentSpanContext != null; + const batch = options.batch ?? new ExportBatch(options.flush); let turnIndex = -1; for await (const turn of readTurns(rolloutFile, size, sessionMeta)) { @@ -777,15 +861,16 @@ export async function convertRollout( unannouncedSubagents: subagentsFor(turnIndex), inheritableTurnIds, historyPrefix: history, - flush: options.flush, + batch, + spanIds: options.spanIds ?? new Set(), + routingTruncated, depth: options.depth, }); - await options.flush?.(); if (!success) throw new Error(`Subagent turn ${turn.turnId ?? "(no turn id)"} failed`); continue; } - if (!isFinal(turn, options.stoppedTurnId, turnIndex < turns.length - 1)) { + if (!isFinal(turn, options.stoppedTurnId, turnIndex < totalTurns - 1)) { debugLog(`skipping turn ${turn.turnId ?? "(no turn id)"}: not final`); continue; } @@ -796,6 +881,7 @@ export async function convertRollout( const parentSpanContext = options.parentSpanContext ?? (await seededTraceParent(options.config, sessionMeta, turnIndex + 1)); + const spanIds = new Set(); const emit = () => emitTurn(turn, sessionMeta, { @@ -808,7 +894,9 @@ export async function convertRollout( unannouncedSubagents: subagentsFor(turnIndex), inheritableTurnIds, historyPrefix: history, - flush: options.flush, + batch, + spanIds, + routingTruncated, depth: options.depth, }); @@ -829,16 +917,26 @@ export async function convertRollout( ); } - await options.flush?.(); - if (success) { - await options.onTurnExported?.(turn.turnId); - uploaded.add(turn.turnId); - exportedTurnIds.push(turn.turnId); - } + batch.complete({ + id: turn.turnId, + spanIds, + success, + checkpoint: async (id) => { + await options.onTurnExported?.(id); + uploaded.add(id); + exportedTurnIds.push(id); + }, + }); } finally { - if (!history.omitted) history.append(turnHistoryMessages(turn)); + history.append(turnHistoryMessages(turn)); } } + if (!options.batch) { + await batch.flush(); + if (batch.incomplete && options.config.fail_on_error) { + throw new Error("Some turns were not completely delivered; they will be retried"); + } + } return exportedTurnIds; } diff --git a/plugins/tracing/src/types.ts b/plugins/tracing/src/types.ts index 298bcc8..42ff67f 100644 --- a/plugins/tracing/src/types.ts +++ b/plugins/tracing/src/types.ts @@ -185,6 +185,7 @@ export type SystemPrompt = { developerMessages: string[]; injectedContext: string[]; changed: boolean; + truncated?: boolean; }; export type ToolCall = { @@ -226,4 +227,6 @@ export type Turn = { systemPrompt?: SystemPrompt; userImages: string[]; toolDefinitions: ToolDefinition[]; + /** Some transcript data was omitted to stay within the memory budget. */ + truncated?: boolean; }; diff --git a/plugins/tracing/test/delivery.test.ts b/plugins/tracing/test/delivery.test.ts new file mode 100644 index 0000000..05daac6 --- /dev/null +++ b/plugins/tracing/test/delivery.test.ts @@ -0,0 +1,110 @@ +import * as path from "node:path"; + +import { LangfuseSpanProcessor } from "@langfuse/otel"; +import { startObservation, setLangfuseTracerProvider } from "@langfuse/tracing"; +import type { SpanExporter, ReadableSpan } from "@opentelemetry/sdk-trace-base"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { setupInstrumentation, spanIdentity } from "../src/instrumentation.js"; +import { convertRollout } from "../src/trace.js"; + +const state = vi.hoisted(() => ({ fail: false })); +vi.mock("@opentelemetry/exporter-trace-otlp-http", () => ({ + OTLPTraceExporter: class { + export: SpanExporter["export"] = (_spans, callback) => { + callback(state.fail ? { code: 1, error: new Error("export rejected") } : { code: 0 }); + state.fail = false; + }; + async shutdown(): Promise {} + }, +})); + +const config = { + enabled: true, + public_key: "pk-test", + secret_key: "sk-test", + base_url: "http://localhost", + skill_tags: true, + debug: false, + fail_on_error: false, +}; +beforeEach(() => { + state.fail = false; + vi.stubEnv("OTEL_TRACES_SAMPLER", "always_on"); +}); +afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllEnvs(); + setLangfuseTracerProvider(null); +}); + +describe("per-window span delivery receipts", () => { + it("checkpoints healthy turns in the same batch as a preprocessing failure", async () => { + const prototype = LangfuseSpanProcessor.prototype as unknown as { + processEndedSpan(span: ReadableSpan): Promise; + }; + vi.spyOn(prototype, "processEndedSpan").mockRejectedValueOnce( + new Error("preprocessing failed"), + ); + const instrumentation = setupInstrumentation(config); + const checkpoint = vi.fn(async (_id: string) => {}); + try { + const file = path.resolve( + "plugins/tracing/test/fixtures/sessions/2026/06/03/rollout-two-turns-main.jsonl", + ); + const exported = await convertRollout(file, { + config, + flush: instrumentation.flush, + onTurnExported: checkpoint, + }); + expect(exported).toEqual(["turn-b"]); + expect(checkpoint).toHaveBeenCalledExactlyOnceWith("turn-b"); + } finally { + await instrumentation.shutdown(); + } + }); + + it("reports a preprocessing drop once and recovers for a later span in the same process", async () => { + const prototype = LangfuseSpanProcessor.prototype as unknown as { + processEndedSpan(span: ReadableSpan): Promise; + }; + vi.spyOn(prototype, "processEndedSpan").mockRejectedValueOnce( + new Error("preprocessing failed"), + ); + const instrumentation = setupInstrumentation(config); + try { + const dropped = startObservation("dropped"); + dropped.end(); + const receipt = await instrumentation.flush(); + expect([...receipt.failedSpanIds]).toEqual([spanIdentity(dropped.otelSpan.spanContext())]); + startObservation("healthy").end(); + expect((await instrumentation.flush()).failedSpanIds.size).toBe(0); + } finally { + await instrumentation.shutdown(); + } + }); + + it("does not treat intentionally unsampled spans as failed deliveries", async () => { + vi.stubEnv("OTEL_TRACES_SAMPLER", "always_off"); + const instrumentation = setupInstrumentation(config); + try { + startObservation("not sampled").end(); + expect((await instrumentation.flush()).failedSpanIds.size).toBe(0); + } finally { + await instrumentation.shutdown(); + } + }); + + it("rejects a transport failure but resets instrumentation for a subsequent independent batch", async () => { + const instrumentation = setupInstrumentation(config); + try { + state.fail = true; + startObservation("rejected").end(); + await expect(instrumentation.flush()).rejects.toThrow("export rejected"); + startObservation("healthy").end(); + expect((await instrumentation.flush()).failedSpanIds.size).toBe(0); + } finally { + await instrumentation.shutdown(); + } + }); +}); diff --git a/plugins/tracing/test/export-batch.test.ts b/plugins/tracing/test/export-batch.test.ts new file mode 100644 index 0000000..948a0eb --- /dev/null +++ b/plugins/tracing/test/export-batch.test.ts @@ -0,0 +1,84 @@ +import type { LangfuseObservation } from "@langfuse/tracing"; +import { describe, expect, it, vi } from "vitest"; + +import { ExportBatch } from "../src/export-batch.js"; +import { spanIdentity } from "../src/instrumentation.js"; + +function observation(id: number, bytes = 0): LangfuseObservation { + return { + otelSpan: { + spanContext: () => ({ + traceId: "a".repeat(32), + spanId: id.toString(16).padStart(16, "0"), + traceFlags: 1, + }), + attributes: { output: "x".repeat(bytes) }, + }, + } as unknown as LangfuseObservation; +} + +describe("bounded delivery batches", () => { + it("isolates a dropped span to its own turn and does not poison later batches", async () => { + const missing = spanIdentity(observation(2).otelSpan.spanContext()); + let first = true; + const batch = new ExportBatch(async () => { + const failedSpanIds = new Set(first ? [missing] : []); + first = false; + return { failedSpanIds }; + }); + const checkpoint = vi.fn(async (_id: string) => {}); + for (const n of [1, 2, 3]) { + const spanIds = new Set(); + await batch.ended(observation(n), spanIds); + batch.complete({ id: `turn-${n}`, spanIds, success: true, checkpoint }); + } + await batch.flush(); + expect(checkpoint.mock.calls.map((call) => call[0])).toEqual(["turn-1", "turn-3"]); + const spanIds = new Set(); + await batch.ended(observation(4), spanIds); + batch.complete({ id: "turn-4", spanIds, success: true, checkpoint }); + await batch.flush(); + expect(checkpoint).toHaveBeenLastCalledWith("turn-4"); + }); + + it("remembers failures across intermediate flushes until an open turn completes", async () => { + const missing = spanIdentity(observation(1).otelSpan.spanContext()); + let first = true; + const batch = new ExportBatch(async () => { + const failedSpanIds = new Set(first ? [missing] : []); + first = false; + return { failedSpanIds }; + }); + const spanIds = new Set(); + for (let n = 1; n <= 64; n++) await batch.ended(observation(n), spanIds); + const checkpoint = vi.fn(async (_id: string) => {}); + batch.complete({ id: "large-turn", spanIds, success: true, checkpoint }); + await batch.flush(); + expect(checkpoint).not.toHaveBeenCalled(); + expect(batch.incomplete).toBe(true); + }); + + it("flushes by bytes as well as by count", async () => { + const deliver = vi.fn(async () => ({ failedSpanIds: new Set() })); + const batch = new ExportBatch(deliver); + await batch.ended(observation(1, 3 * 1024 * 1024), new Set()); + expect(deliver).not.toHaveBeenCalled(); + await batch.ended(observation(2, 3 * 1024 * 1024), new Set()); + expect(deliver).toHaveBeenCalledOnce(); + }); + + it("cannot checkpoint a rejected batch after a later empty flush", async () => { + const deliver = vi.fn(async () => { + throw new Error("network failed"); + }); + const batch = new ExportBatch(deliver); + const checkpoint = vi.fn(async (_id: string) => {}); + const spanIds = new Set(); + await batch.ended(observation(1), spanIds); + batch.complete({ id: "rejected", spanIds, success: true, checkpoint }); + await expect(batch.flush()).rejects.toThrow("network failed"); + await expect(batch.flush()).rejects.toThrow("network failed"); + expect(deliver).toHaveBeenCalledOnce(); + expect(checkpoint).not.toHaveBeenCalled(); + }); +}); diff --git a/plugins/tracing/test/rollout.test.ts b/plugins/tracing/test/rollout.test.ts index c5bfc78..106340e 100644 --- a/plugins/tracing/test/rollout.test.ts +++ b/plugins/tracing/test/rollout.test.ts @@ -67,11 +67,17 @@ describe("bounded rollout reader", () => { expect(await records(file, size)).toHaveLength(1); }); - it("rejects an oversized record before JSON decoding, including a record without a newline", async () => { + it("discards an oversized record before JSON decoding and resumes at the next newline", async () => { for (const suffix of ["", "\n"]) { const { file, size } = await stage("x".repeat(200_000) + suffix); - await expect(records(file, size, 100_000)).rejects.toThrow("record 1 exceeds 100000 bytes"); + expect(await records(file, size, 100_000)).toEqual([ + expect.objectContaining({ payload: { type: "langfuse_record_omitted" } }), + ]); } + const { file, size } = await stage( + "x".repeat(200_000) + "\n" + line("session_meta", { id: "after" }), + ); + expect(await records(file, size, 100_000)).toHaveLength(2); }); it("does not retain compaction snapshots", async () => { @@ -89,7 +95,9 @@ describe("bounded rollout reader", () => { ), ].join("\n"), ); - await expect(turns(file, size)).rejects.toThrow("Rollout turn exceeds 400 bytes"); + const result = await turns(file, size); + expect(result.turns).toHaveLength(1); + expect(result.turns[0]).toMatchObject({ turnId: "t1", truncated: true }); }); it("bounds context retained across turns", async () => { @@ -97,6 +105,7 @@ describe("bounded rollout reader", () => { const { file, size } = await stage( [ line("session_meta", { id: "s", base_instructions: { text: "1234567890" } }), + line("event_msg", { type: "task_started", turn_id: "context-turn" }), line("response_item", { type: "message", role: "developer", @@ -104,10 +113,11 @@ describe("bounded rollout reader", () => { }), ].join("\n"), ); - await expect(turns(file, size)).rejects.toThrow("Rollout context exceeds memory limit"); + const result = await turns(file, size); + expect(result.turns[0].systemPrompt).toMatchObject({ truncated: true }); }); - it("bounds the number of indexed turns and propagates file errors", async () => { + it("keeps reading turns beyond the routing-index limit and propagates file errors", async () => { ROLLOUT_LIMITS.turns = 1; const { file, size } = await stage( [ @@ -116,8 +126,29 @@ describe("bounded rollout reader", () => { line("event_msg", { type: "task_started", turn_id: "t2" }), ].join("\n"), ); - await expect(turns(file, size)).rejects.toThrow("too many turns"); + expect((await turns(file, size)).turns).toHaveLength(2); await fs.unlink(file); await expect(records(file, size)).rejects.toThrow("ENOENT"); }); + + it("keeps turn boundaries and later turns after the per-turn budget is exceeded", async () => { + ROLLOUT_LIMITS.turnBytes = 400; + const { file, size } = await stage( + [ + line("event_msg", { type: "task_started", turn_id: "large" }), + ...Array.from({ length: 20 }, () => + line("event_msg", { type: "user_message", message: "hi" }), + ), + line("event_msg", { type: "task_complete" }), + line("event_msg", { type: "task_started", turn_id: "small" }), + line("event_msg", { type: "user_message", message: "next" }), + line("event_msg", { type: "task_complete" }), + ].join("\n"), + ); + const result = await turns(file, size); + expect(result.turns.map((t) => t.turnId)).toEqual(["large", "small"]); + expect(result.turns[0]).toMatchObject({ truncated: true, completed: true }); + expect(result.turns[1]).toMatchObject({ userInput: "next", completed: true }); + expect(result.turns[1].truncated).toBeUndefined(); + }); }); diff --git a/plugins/tracing/test/streaming-hook.test.ts b/plugins/tracing/test/streaming-hook.test.ts index e9abaef..bae7c46 100644 --- a/plugins/tracing/test/streaming-hook.test.ts +++ b/plugins/tracing/test/streaming-hook.test.ts @@ -29,7 +29,10 @@ async function stage(count = 3, outputBytes = 32) { }) + "\n", ); }; - await line("session_meta", { id: "streaming-session" }); + await line("session_meta", { + id: "streaming-session", + base_instructions: { text: "You are a helpful coding assistant." }, + }); for (let i = 1; i <= count; i++) { await line("event_msg", { type: "task_started", turn_id: `turn-${i}` }); await line("event_msg", { type: "user_message", message: `question ${i}` }); @@ -142,7 +145,7 @@ const ok = (res: http.ServerResponse) => { describe("streaming delivery checkpoints", () => { for (const mode of ["rejection", "interruption"] as const) { it(`retries unfinished turns after ${mode}, preserving ids and confirmed checkpoints`, async () => { - const { dir, file } = await stage(); + const { dir, file } = await stage(20); let child: ChildProcessWithoutNullStreams; let attempt = 0; const requests: WireSpan[][][] = [[], []]; @@ -165,37 +168,43 @@ describe("streaming delivery checkpoints", () => { child = first.child; const failure = await first.result; expect(failure.code).not.toBe(0); - expect(marksAtFailure).toEqual(["turn-1"]); - expect(await sidecar(file)).toEqual(["turn-1"]); + const confirmed = Array.from({ length: 7 }, (_, i) => `turn-${i + 1}`); + expect(marksAtFailure).toEqual(confirmed); + expect(await sidecar(file)).toEqual(confirmed); attempt = 1; const retry = await run(dir, file, server.port).result; expect(retry.code, retry.stderr).toBe(0); - expect(await sidecar(file)).toEqual(["turn-1", "turn-2", "turn-3"]); - const identities = (spans: WireSpan[]) => - spans.map(({ traceId, spanId }) => ({ traceId, spanId })); - expect(identities(requests[1][0])).toEqual(identities(requests[0][1])); - expect(requests[1]).toHaveLength(2); + expect(await sidecar(file)).toEqual(Array.from({ length: 20 }, (_, i) => `turn-${i + 1}`)); + const retried = new Set(requests[1].flat().map((s) => `${s.traceId}:${s.spanId}`)); + expect(requests[0][1].every((s) => retried.has(`${s.traceId}:${s.spanId}`))).toBe(true); + expect(requests[1].length).toBeLessThan(20); } finally { await server.close(); } }); } - it("processes growing history under a 128 MiB heap and explicitly labels omitted repeated input", async () => { + it("keeps current prompts and recent history under a 128 MiB heap instead of permanently omitting input", async () => { const { dir, file } = await stage(120, 100_000); let spanCount = 0; - let omitted = 0; + let truncated = 0; + const inputs: string[] = []; let intactToolOutputs = 0; const server = await receiver((body, res) => { const spans = parseSpans(body); spanCount += spans.length; - omitted += spans.filter((s) => + truncated += spans.filter((s) => s.attributes.some( (a) => - a.key.endsWith("codex.generation_input.omitted") && + a.key.endsWith("codex.generation_input.truncated") && (a.value.boolValue === true || a.value.stringValue === "true"), ), ).length; + for (const span of spans.filter((s) => s.name === "LLM")) { + const input = span.attributes.find((a) => a.key === "langfuse.observation.input")?.value + .stringValue; + if (input) inputs.push(input); + } intactToolOutputs += spans.filter( (s) => s.name === "exec_command" && @@ -212,29 +221,83 @@ describe("streaming delivery checkpoints", () => { expect(result.code, result.stderr).toBe(0); expect(spanCount).toBe(120 * 4); expect(intactToolOutputs).toBe(120); - expect(omitted).toBeGreaterThan(0); + expect(truncated).toBeGreaterThan(0); + expect(inputs).toHaveLength(240); + for (const input of inputs) { + expect(Buffer.byteLength(input)).toBeLessThanOrEqual(256 * 1024); + expect(input).toContain("You are a helpful coding assistant."); + expect(input).toMatch(/question \d+/); + expect(input).not.toContain("Generation input omitted"); + } + expect(inputs.at(-1)).toContain("question 120"); + expect(inputs.at(-1)).toContain("question 119"); + expect(inputs.at(-1)).not.toContain('"question 1"'); expect(await sidecar(file)).toHaveLength(120); } finally { await server.close(); } }, 20_000); - it("reports an oversized record without acknowledging the turn", async () => { - const { dir, file } = await stage(1); - await fs.appendFile(file, "x".repeat(8 * 1024 * 1024 + 1)); + it("labels an oversized tool output and delivers earlier and later turns on every invocation", async () => { + const { dir, file } = await stage(3); + const contents = + (await fs.readFile(file, "utf8")) + .split("\n") + .filter(Boolean) + .map((raw) => { + const record = JSON.parse(raw); + if ( + record.payload.type === "function_call_output" && + record.payload.call_id === "call-2" + ) { + record.payload.output = "x".repeat(8 * 1024 * 1024 + 1); + } + return JSON.stringify(record); + }) + .join("\n") + "\n"; + await fs.writeFile(file, contents); let requests = 0; - const server = await receiver((_body, res) => { + const received: WireSpan[] = []; + const server = await receiver((body, res) => { requests++; + received.push(...parseSpans(body)); ok(res); }); try { const result = await run(dir, file, server.port).result; - expect(result.code).toBe(1); + expect(result.code, result.stderr).toBe(0); expect(result.stderr).toContain("exceeds 8388608 bytes"); - expect(requests).toBe(0); - expect(await sidecar(file)).toEqual([]); + expect(requests).toBeGreaterThan(0); + expect(await sidecar(file)).toEqual(["turn-1", "turn-2", "turn-3"]); + expect(received.filter((s) => s.name === "Codex Turn")).toHaveLength(3); + expect( + received.filter((s) => s.attributes.some((a) => a.key.endsWith("codex.rollout.truncated"))), + ).toHaveLength(1); + const before = requests; + expect((await run(dir, file, server.port).result).code).toBe(0); + expect(requests).toBe(before); } finally { await server.close(); } }); + + it("batches a 200-turn catch-up within the 30-second hook budget at 200 ms request latency", async () => { + const { dir, file } = await stage(200); + let requests = 0; + const server = await receiver(async (_body, res) => { + requests++; + await new Promise((resolve) => setTimeout(resolve, 200)); + ok(res); + }); + try { + const started = performance.now(); + const result = await run(dir, file, server.port).result; + expect(result.code, result.stderr).toBe(0); + expect(requests).toBeLessThanOrEqual(26); + expect(performance.now() - started).toBeLessThan(15_000); + expect(await sidecar(file)).toHaveLength(200); + } finally { + await server.close(); + } + }, 20_000); }); diff --git a/plugins/tracing/test/trace.test.ts b/plugins/tracing/test/trace.test.ts index 36e015e..a6deeea 100644 --- a/plugins/tracing/test/trace.test.ts +++ b/plugins/tracing/test/trace.test.ts @@ -15,7 +15,8 @@ import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest"; import type { Config } from "../src/config.js"; import { markTurnUploaded } from "../src/sidecar.js"; -import { buildSubagentIndex, convertRollout } from "../src/trace.js"; +import { buildSubagentIndex, convertRollout, MAX_GENERATION_INPUT_BYTES } from "../src/trace.js"; +import { ROLLOUT_LIMITS } from "../src/rollout.js"; const exporter = new InMemorySpanExporter(); let provider: NodeTracerProvider; @@ -86,6 +87,56 @@ beforeEach(() => { }); describe("convertRollout", () => { + it("continues exporting and labels routing degradation when the bounded index fills", async () => { + const limit = ROLLOUT_LIMITS.turns; + ROLLOUT_LIMITS.turns = 1; + try { + const dir = stageFixtures(); + expect( + await convertRollout(path.join(dir, "rollout-two-turns-main.jsonl"), { + config: baseConfig, + }), + ).toEqual(["turn-a", "turn-b"]); + expect(turnRoots()).toHaveLength(2); + for (const root of turnRoots()) + expect(attr(root, "langfuse.observation.metadata.codex.routing.truncated")).toBe("true"); + } finally { + ROLLOUT_LIMITS.turns = limit; + } + }); + + it("bounds huge UTF-8/escaped prompts without losing the current user or system messages", async () => { + const dir = stageFixtures(); + const file = path.join(dir, "rollout-two-turns-main.jsonl"); + const records = fs + .readFileSync(file, "utf8") + .trim() + .split("\n") + .map((raw) => JSON.parse(raw)); + records[0].payload.base_instructions = { text: "system-prefix " + "界\\\n".repeat(100_000) }; + records.find((record) => record.payload.type === "user_message").payload.message = + "current-user " + "界\\\n".repeat(100_000); + fs.writeFileSync(file, records.map((record) => JSON.stringify(record)).join("\n") + "\n"); + await convertRollout(file, { config: baseConfig }); + const generations = exporter + .getFinishedSpans() + .filter((span) => obsType(span) === "generation"); + expect(generations).toHaveLength(2); + const inputs = generations.map((span) => attr(span, "langfuse.observation.input")); + for (const input of inputs) { + expect(Buffer.byteLength(input)).toBeLessThanOrEqual(MAX_GENERATION_INPUT_BYTES); + expect(JSON.parse(input).some((message: { role: string }) => message.role === "system")).toBe( + true, + ); + expect(JSON.parse(input).some((message: { role: string }) => message.role === "user")).toBe( + true, + ); + expect(input).toContain("system-prefix"); + } + expect(inputs[0]).toContain("current-user"); + expect(inputs[1]).toContain("And 2 + 2?"); + }); + it("emits an agent → generation → tool tree with backdated timestamps", async () => { const dir = stageFixtures(); await convertRollout(path.join(dir, "rollout-basic-main.jsonl"), { config: baseConfig }); diff --git a/scripts/benchmark-rollout.mjs b/scripts/benchmark-rollout.mjs index 656491a..f96cd2a 100644 --- a/scripts/benchmark-rollout.mjs +++ b/scripts/benchmark-rollout.mjs @@ -1,4 +1,4 @@ -// Usage: node scripts/benchmark-rollout.mjs [bundle] [turns] [tool-output-bytes] +// Usage: node scripts/benchmark-rollout.mjs [bundle] [turns] [tool-output-bytes] [request-delay-ms] // Runs the real bundled hook under a 128 MiB heap against a local OTLP receiver. import { spawn } from "node:child_process"; import { createServer } from "node:http"; @@ -10,6 +10,7 @@ import { pathToFileURL } from "node:url"; const bundle = resolve(process.argv[2] ?? "plugins/tracing/dist/index.mjs"); const turns = Number(process.argv[3] ?? 200); const outputBytes = Number(process.argv[4] ?? 200_000); +const requestDelayMs = Number(process.argv[5] ?? 0); const dir = await mkdtemp(join(tmpdir(), "codex-rollout-benchmark-")); const sessions = join(dir, "sessions/2026/10/05"); await mkdir(sessions, { recursive: true }); @@ -53,12 +54,13 @@ let spans = 0; const server = createServer((req, res) => { const chunks = []; req.on("data", (chunk) => chunks.push(chunk)); - req.on("end", () => { + req.on("end", async () => { requests++; const payload = JSON.parse(Buffer.concat(chunks).toString()); for (const resource of payload.resourceSpans ?? []) { for (const scope of resource.scopeSpans ?? []) spans += scope.spans.length; } + if (requestDelayMs) await new Promise((resolve) => setTimeout(resolve, requestDelayMs)); res.writeHead(200, { "Content-Type": "application/json" }); res.end("{}"); }); @@ -107,6 +109,7 @@ try { { bundle, turns, + requestDelayMs, rolloutBytes: (await stat(rollout)).size, elapsedMs: Math.round(performance.now() - started), ...result,