From dc94e0067cec416ce98616f467283164c1c594a2 Mon Sep 17 00:00:00 2001 From: XinweiHe Date: Tue, 30 Jun 2026 22:37:04 -0700 Subject: [PATCH] fix(codex-plugin): live streaming, final-output capture, root-span duration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Defer the root span to the Stop hook so a running turn leaves no completed root in ClickHouse — the backend keeps the live SSE stream open instead of treating the trace as done on the first hook (no more manual refresh). The "Codex Turn" title is preserved during the live window via traceroot.span.path. - Capture the turn's final output from the agent_message event (Codex writes it before the terminal task_complete, which lands just after the Stop hook fires), so live runs no longer show empty output; task_complete still refines it. - Compute the root span end from the latest step/tool activity so aborted or live-snapshot turns span their children instead of rendering a 0ms bar. - Fix the README install steps (codex plugin add + features.hooks). Verified against prod and with unit tests (72 pass). Rebuilt dist bundle. Refs traceroot-ai/traceroot#1401 Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 13 +++++- plugins/tracing/dist/index.mjs | 36 ++++++++++++---- plugins/tracing/src/emit.ts | 17 ++++++-- plugins/tracing/src/spans.ts | 57 ++++++++++++++++++++----- plugins/tracing/src/transcript.ts | 12 +++++- plugins/tracing/test/emit.test.ts | 24 +++++------ plugins/tracing/test/spans.test.ts | 47 +++++++++++++++++++- plugins/tracing/test/transcript.test.ts | 15 +++++++ 8 files changed, 184 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index fbfd0b6..4b1a976 100644 --- a/README.md +++ b/README.md @@ -4,14 +4,23 @@ Trace OpenAI Codex sessions to [TraceRoot](https://traceroot.ai). Every query be ## Install +Add the marketplace, then install the plugin from it: + ```bash codex plugin marketplace add traceroot-ai/traceroot-codex-plugin +codex plugin add tracing@traceroot-codex-plugin +``` + +Enable Codex hooks (the plugin runs as a hook): + +```bash +codex features enable hooks ``` -Enable it in `~/.codex/config.toml`: +This writes `features.hooks = true` to `~/.codex/config.toml`. The `codex plugin add` step above also enables the plugin there: ```toml -plugin_hooks = true +features.hooks = true [plugins."tracing@traceroot-codex-plugin"] enabled = true diff --git a/plugins/tracing/dist/index.mjs b/plugins/tracing/dist/index.mjs index be14d93..9a48e1f 100644 --- a/plugins/tracing/dist/index.mjs +++ b/plugins/tracing/dist/index.mjs @@ -32004,6 +32004,7 @@ function mapUsage(u) { //#endregion //#region src/spans.ts init_esm$2(); +const ROOT_NAME = "Codex Turn"; const str = (v, max) => truncate(typeof v === "string" ? v : JSON.stringify(v ?? ""), max); function commonTrace(attrs, sessionMeta, ctx) { attrs["traceroot.sdk.name"] = SDK_NAME; @@ -32019,7 +32020,10 @@ function planTurnSpans(sessionMeta, turn, ctx, opts) { const rootId = makeSpanId(seed + turn.turnId + ":root"); const rootParentSpanId = opts?.rootParentSpanId !== void 0 ? opts.rootParentSpanId : null; const out = []; - const rootAttrs = { "traceroot.span.type": "AGENT" }; + const rootAttrs = { + "traceroot.span.type": "AGENT", + "traceroot.span.path": [ROOT_NAME] + }; commonTrace(rootAttrs, sessionMeta, ctx); if (turn.userInput) rootAttrs["traceroot.span.input"] = str(turn.userInput, ctx.maxChars); if (turn.finalOutput) rootAttrs["traceroot.span.output"] = str(turn.finalOutput, ctx.maxChars); @@ -32033,16 +32037,21 @@ function planTurnSpans(sessionMeta, turn, ctx, opts) { })); if (ctx.git?.repo) rootAttrs["traceroot.git.repo"] = ctx.git.repo; if (ctx.git?.ref) rootAttrs["traceroot.git.ref"] = ctx.git.ref; + let rootEnd = turn.endTime; + for (const step of turn.steps) { + if (step.endTime > rootEnd) rootEnd = step.endTime; + for (const tc of step.toolCalls) if (tc.endTime !== void 0 && tc.endTime > rootEnd) rootEnd = tc.endTime; + } out.push({ spanId: rootId, parentSpanId: rootParentSpanId, traceId, kind: "AGENT", - name: "Codex Turn", + name: ROOT_NAME, startTime: turn.startTime, - endTime: turn.endTime, + endTime: rootEnd, attributes: rootAttrs, - complete: rootParentSpanId === null ? true : turn.completed + complete: rootParentSpanId === null ? turn.completed || !!ctx.turnEnding : turn.completed }); let prevToolResults; for (const step of turn.steps) { @@ -32075,7 +32084,10 @@ function buildStepOutput(step, max) { return truncate(JSON.stringify(o), max); } function planStepSpan(sessionMeta, turn, step, spanId, parentSpanId, traceId, ctx, input) { - const attrs = { "traceroot.span.type": "LLM" }; + const attrs = { + "traceroot.span.type": "LLM", + "traceroot.span.path": [ROOT_NAME, turn.model ?? "model"] + }; commonTrace(attrs, sessionMeta, ctx); if (turn.model) attrs["traceroot.llm.model"] = turn.model; if (input) attrs["traceroot.span.input"] = truncate(input, ctx.maxChars); @@ -32096,7 +32108,10 @@ function planStepSpan(sessionMeta, turn, step, spanId, parentSpanId, traceId, ct }; } function planToolSpan(sessionMeta, tc, seed, parentSpanId, traceId, ctx) { - const attrs = { "traceroot.span.type": "TOOL" }; + const attrs = { + "traceroot.span.type": "TOOL", + "traceroot.span.path": [ROOT_NAME, tc.name] + }; commonTrace(attrs, sessionMeta, ctx); attrs["traceroot.span.input"] = str(tc.args, ctx.maxChars); if (tc.output !== void 0) attrs["traceroot.span.output"] = str(tc.output, ctx.maxChars); @@ -32312,6 +32327,9 @@ function parseRollout(lines) { case "user_message": if (turn && typeof p.message === "string") turn.userInput = p.message; break; + case "agent_message": + if (turn && typeof p.message === "string") turn.finalOutput = p.message; + break; case "token_count": if (turn && p.info?.last_token_usage) { const s = ensureStep(turn, at); @@ -32325,7 +32343,7 @@ function parseRollout(lines) { turn.completed = true; turn.endTime = at; const lastText = [...turn.steps].reverse().find((s) => s.text)?.text; - turn.finalOutput = p.last_agent_message ?? lastText ?? void 0; + turn.finalOutput = p.last_agent_message ?? turn.finalOutput ?? lastText ?? void 0; } break; default: @@ -32459,6 +32477,7 @@ async function dispatch(hook, config, deps) { const buildTracingFn = deps?.buildTracing ?? buildTracing; const getGit = deps?.getGit ?? getGitContext; const findSubagentFn = deps?.findSubagent ?? locateSubagentRollout; + const turnEnding = hook.hook_event_name === "Stop"; const { sessionMeta, turns } = parseRollout(await readRollout(transcript)); if (sessionMeta.threadSource === "subagent" || sessionMeta.parentThreadId) { debugLog(`subagent session ${sessionMeta.sessionId}; skipping standalone emit (nested under parent)`); @@ -32469,7 +32488,8 @@ async function dispatch(hook, config, deps) { environment: config.environment, userId: config.userId, git, - maxChars: config.maxChars + maxChars: config.maxChars, + turnEnding }; const already = await loadEmittedSpanIds(transcript); const tracing = buildTracingFn(config); diff --git a/plugins/tracing/src/emit.ts b/plugins/tracing/src/emit.ts index e0c94cd..2a196ce 100644 --- a/plugins/tracing/src/emit.ts +++ b/plugins/tracing/src/emit.ts @@ -110,8 +110,17 @@ export async function dispatch( const getGit = deps?.getGit ?? getGitContext; const findSubagentFn = deps?.findSubagent ?? locateSubagentRollout; - const lines = await readRollout(transcript); - const { sessionMeta, turns } = parseRollout(lines); + // The Stop hook is the "turn is ending" signal — only then is the root span + // emitted (see planTurnSpans). Mid-turn PostToolUse hooks leave the root absent + // so the trace stays live. SubagentStop carries the PARENT transcript while the + // parent turn is still running, so it must NOT count as the parent turn ending. + // The Stop hook is the "turn is ending" signal — only then is the root span + // emitted (see planTurnSpans). Mid-turn PostToolUse hooks leave the root absent + // so the trace stays live. SubagentStop carries the PARENT transcript while the + // parent turn is still running, so it must NOT count as the parent turn ending. + const turnEnding = hook.hook_event_name === "Stop"; + + const { sessionMeta, turns } = parseRollout(await readRollout(transcript)); // A subagent session is emitted NESTED under its parent's trace (via the // parent's spawn_agent tool). Codex fires PostToolUse/Stop on the subagent @@ -126,7 +135,9 @@ export async function dispatch( const cwd = sessionMeta.cwd ?? hook.cwd ?? process.cwd(); const git = await getGit(cwd).catch(() => ({})); - const ctx: EmitCtx = { environment: config.environment, userId: config.userId, git, maxChars: config.maxChars }; + const ctx: EmitCtx = { + environment: config.environment, userId: config.userId, git, maxChars: config.maxChars, turnEnding, + }; const already = await loadEmittedSpanIds(transcript); const tracing = buildTracingFn(config); diff --git a/plugins/tracing/src/spans.ts b/plugins/tracing/src/spans.ts index 453e645..e6752b8 100644 --- a/plugins/tracing/src/spans.ts +++ b/plugins/tracing/src/spans.ts @@ -6,13 +6,24 @@ import { mapUsage } from "./tokens.js"; import type { ModelStep, SessionMeta, ToolCall, Turn } from "./types.js"; import { truncate } from "./util.js"; +// The trace title. Used as the root span name AND as path[0] on every span so +// the trace reads as "Codex Turn" even before the root span exists (see below). +export const ROOT_NAME = "Codex Turn"; + export type EmitCtx = { environment?: string; userId?: string; git?: { repo?: string; ref?: string }; maxChars: number; + // True only on the terminal Stop hook (the turn is ending). The root span is + // emitted only when the turn ends, so a still-running turn has NO completed + // root in ClickHouse — which is how the backend's live-stream detector knows + // the trace is in progress and keeps streaming (matching the SDK's model). + turnEnding?: boolean; }; +type SpanAttrs = Record; + export type EmittableSpan = { spanId: string; parentSpanId: string | null; @@ -21,14 +32,14 @@ export type EmittableSpan = { name: string; startTime: number; endTime: number; - attributes: Record; + attributes: SpanAttrs; complete: boolean; }; const str = (v: unknown, max: number): string => truncate(typeof v === "string" ? v : JSON.stringify(v ?? ""), max); -function commonTrace(attrs: Record, sessionMeta: SessionMeta, ctx: EmitCtx): void { +function commonTrace(attrs: SpanAttrs, sessionMeta: SessionMeta, ctx: EmitCtx): void { attrs["traceroot.sdk.name"] = SDK_NAME; attrs["traceroot.sdk.version"] = SDK_VERSION; // Child subagent spans intentionally carry the CHILD session id here (sessionMeta is @@ -57,7 +68,7 @@ export function planTurnSpans(sessionMeta: SessionMeta, turn: Turn, ctx: EmitCtx const out: EmittableSpan[] = []; // Root AGENT span (trace-level attrs live here). - const rootAttrs: Record = { "traceroot.span.type": "AGENT" }; + const rootAttrs: SpanAttrs = { "traceroot.span.type": "AGENT", "traceroot.span.path": [ROOT_NAME] }; commonTrace(rootAttrs, sessionMeta, ctx); if (turn.userInput) rootAttrs["traceroot.span.input"] = str(turn.userInput, ctx.maxChars); if (turn.finalOutput) rootAttrs["traceroot.span.output"] = str(turn.finalOutput, ctx.maxChars); @@ -71,14 +82,31 @@ export function planTurnSpans(sessionMeta: SessionMeta, turn: Turn, ctx: EmitCtx })); if (ctx.git?.repo) rootAttrs["traceroot.git.repo"] = ctx.git.repo; if (ctx.git?.ref) rootAttrs["traceroot.git.ref"] = ctx.git.ref; + // Root end = the latest activity in the turn, not just the task_complete time. + // turn.endTime only advances on task_complete, so a turn that was aborted or + // snapshotted live (no task_complete) would have endTime === startTime and + // render as a 0ms bar. Cover the turn's own steps and tool calls (including + // wait_agent, which ends when a spawned subagent finishes) so the root always + // spans its children — and grows correctly across live hooks. + let rootEnd = turn.endTime; + for (const step of turn.steps) { + if (step.endTime > rootEnd) rootEnd = step.endTime; + for (const tc of step.toolCalls) { + if (tc.endTime !== undefined && tc.endTime > rootEnd) rootEnd = tc.endTime; + } + } out.push({ spanId: rootId, parentSpanId: rootParentSpanId, traceId, kind: "AGENT", - name: "Codex Turn", startTime: turn.startTime, endTime: turn.endTime, + name: ROOT_NAME, startTime: turn.startTime, endTime: rootEnd, attributes: rootAttrs, - // The trace root (no parent) is emittable immediately so the trace is named - // "Codex Turn" from the first live hook; emit.ts re-emits it to refine - // end/output. Subagent roots stay gated on completion. - complete: rootParentSpanId === null ? true : turn.completed, + // The root is emitted only when the turn ENDS (task_complete, or the terminal + // Stop hook signalled via ctx.turnEnding) — never mid-turn. A completed root in + // ClickHouse is the backend's "trace is done" signal, so emitting it early would + // close the live SSE stream on the first hook (the trace would look finished and + // stop updating until refresh). Deferring it keeps the trace live while the turn + // runs; meanwhile traceroot.span.path keeps the title "Codex Turn". Subagent roots + // (rootParentSpanId !== null) stay gated on their own completion. + complete: rootParentSpanId === null ? (turn.completed || !!ctx.turnEnding) : turn.completed, }); // LLM-step input is the user prompt (first step) or the prior step's tool @@ -117,7 +145,13 @@ function planStepSpan( sessionMeta: SessionMeta, turn: Turn, step: ModelStep, spanId: string, parentSpanId: string, traceId: string, ctx: EmitCtx, input: string | undefined, ): EmittableSpan { - const attrs: Record = { "traceroot.span.type": "LLM" }; + const attrs: SpanAttrs = { + "traceroot.span.type": "LLM", + // path[0] = ROOT_NAME so the trace title stays "Codex Turn" while the root + // span is still deferred (the backend names the trace from the shallowest + // span's path[0] until the real root arrives). + "traceroot.span.path": [ROOT_NAME, turn.model ?? "model"], + }; commonTrace(attrs, sessionMeta, ctx); if (turn.model) attrs["traceroot.llm.model"] = turn.model; if (input) attrs["traceroot.span.input"] = truncate(input, ctx.maxChars); @@ -135,7 +169,10 @@ function planStepSpan( function planToolSpan( sessionMeta: SessionMeta, tc: ToolCall, seed: string, parentSpanId: string, traceId: string, ctx: EmitCtx, ): EmittableSpan { - const attrs: Record = { "traceroot.span.type": "TOOL" }; + const attrs: SpanAttrs = { + "traceroot.span.type": "TOOL", + "traceroot.span.path": [ROOT_NAME, tc.name], + }; commonTrace(attrs, sessionMeta, ctx); attrs["traceroot.span.input"] = str(tc.args, ctx.maxChars); if (tc.output !== undefined) attrs["traceroot.span.output"] = str(tc.output, ctx.maxChars); diff --git a/plugins/tracing/src/transcript.ts b/plugins/tracing/src/transcript.ts index d22287d..1adbef4 100644 --- a/plugins/tracing/src/transcript.ts +++ b/plugins/tracing/src/transcript.ts @@ -120,6 +120,14 @@ export function parseRollout(lines: RolloutLine[]): { sessionMeta: SessionMeta; case "user_message": if (turn && typeof p.message === "string") turn.userInput = p.message; break; + case "agent_message": + // The agent's final message. Codex writes this BEFORE the terminal + // task_complete line (which lands just AFTER the Stop hook fires), so + // capturing the output here is what lets the root span carry the output + // on a live run — relying only on task_complete, our hook reads too early + // to see it. Last one wins. + if (turn && typeof p.message === "string") turn.finalOutput = p.message; + break; case "token_count": if (turn && p.info?.last_token_usage) { const s = ensureStep(turn, at); @@ -133,7 +141,9 @@ export function parseRollout(lines: RolloutLine[]): { sessionMeta: SessionMeta; turn.completed = true; turn.endTime = at; const lastText = [...turn.steps].reverse().find((s) => s.text)?.text; - turn.finalOutput = (p.last_agent_message ?? lastText) ?? undefined; + // Prefer an explicit last_agent_message; otherwise keep what the + // agent_message event already captured, then fall back to step text. + turn.finalOutput = (p.last_agent_message ?? turn.finalOutput ?? lastText) ?? undefined; } break; default: diff --git a/plugins/tracing/test/emit.test.ts b/plugins/tracing/test/emit.test.ts index 13799e3..dc5dd97 100644 --- a/plugins/tracing/test/emit.test.ts +++ b/plugins/tracing/test/emit.test.ts @@ -50,7 +50,7 @@ describe("dispatch", () => { expect(r2.emitted).toBe(1); }); - it("live: trace root + completed tool/llm emit before the turn closes; root refreshes on completion", async () => { + it("live: completed LLM/TOOL stream mid-turn with NO root (trace stays live); Stop emits the root", async () => { dir = await fs.mkdtemp(path.join(os.tmpdir(), "tr-live-")); const transcript = path.join(dir, "rollout.jsonl"); await fs.copyFile(path.join(here, "fixtures", "rollout-inprogress.jsonl"), transcript); @@ -66,19 +66,19 @@ describe("dispatch", () => { getGit: async () => ({}), }; - const r1 = await dispatch({ transcript_path: transcript }, config, deps); - // Trace root (named "Codex Turn") emits immediately for live naming, plus the - // already-complete LLM step and its TOOL — before the turn closes. - expect(r1.emitted).toBe(3); + // Mid-turn PostToolUse hook: the already-complete LLM step + its TOOL emit, + // but the root does NOT — so ClickHouse holds no completed root and the + // backend keeps the live SSE stream open (the trace updates without refresh). + const r1 = await dispatch({ transcript_path: transcript, hook_event_name: "PostToolUse" }, config, deps); + expect(r1.emitted).toBe(2); expect(mem.getFinishedSpans().map((s) => s.attributes["traceroot.span.type"]).sort()) - .toEqual(["AGENT", "LLM", "TOOL"]); + .toEqual(["LLM", "TOOL"]); - // Now the turn completes: append task_complete and re-run. The root re-emits - // (refreshed end + finalOutput); LLM/TOOL already emitted (sidecar dedup). - await fs.appendFile(transcript, - '{"timestamp":"2026-06-23T23:21:36.100Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-1","last_agent_message":"done","completed_at":1782256896}}\n'); - const r2 = await dispatch({ transcript_path: transcript }, config, deps); - expect(r2.emitted).toBe(1); // the AGENT root re-emitted + // The terminal Stop hook ends the turn: the root emits (even though this + // in-progress rollout never got a task_complete), which signals trace_complete. + const r2 = await dispatch({ transcript_path: transcript, hook_event_name: "Stop" }, config, deps); + expect(r2.emitted).toBe(1); + expect(mem.getFinishedSpans().filter((s) => s.attributes["traceroot.span.type"] === "AGENT")).toHaveLength(1); }); it("skips standalone emission for a subagent session (it nests under its parent)", async () => { diff --git a/plugins/tracing/test/spans.test.ts b/plugins/tracing/test/spans.test.ts index f875e33..fd70df7 100644 --- a/plugins/tracing/test/spans.test.ts +++ b/plugins/tracing/test/spans.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { planTurnSpans } from "../src/spans.js"; +import { planTurnSpans, ROOT_NAME } from "../src/spans.js"; import { makeSpanId, makeTraceId } from "../src/ids.js"; import type { SessionMeta, ToolCall, Turn } from "../src/types.js"; @@ -72,6 +72,51 @@ describe("planTurnSpans", () => { expect(String(tool.attributes["traceroot.span.output"])).toContain("a.txt"); expect(tool.complete).toBe(true); }); + + it("every span carries traceroot.span.path with ROOT_NAME first (live trace title)", () => { + // While the root span is deferred, the backend titles the trace from the + // shallowest span's path[0] — so path[0] must be the root name on all spans. + for (const s of spans) { + const path = s.attributes["traceroot.span.path"]; + expect(Array.isArray(path)).toBe(true); + expect((path as string[])[0]).toBe(ROOT_NAME); + } + }); + + // Part 1 — the root is emitted (complete) only when the turn ENDS, so a + // running turn leaves no completed root in ClickHouse and the live stream stays + // open. ctx.turnEnding (the Stop hook) forces it even without task_complete. + const rootComplete = (t: Turn, turnEnding: boolean) => + planTurnSpans(sessionMeta, t, { ...ctx, turnEnding }).find((s) => s.kind === "AGENT")!.complete; + + it("root is NOT complete mid-turn (not completed, not the Stop hook)", () => { + expect(rootComplete({ ...turn, completed: false }, false)).toBe(false); + }); + it("root IS complete on the Stop hook even if task_complete was missed", () => { + expect(rootComplete({ ...turn, completed: false }, true)).toBe(true); + }); + it("root IS complete when the turn completed cleanly", () => { + expect(rootComplete({ ...turn, completed: true }, false)).toBe(true); + }); + + // Regression: a turn with no task_complete (aborted / live snapshot) has + // endTime === startTime. The root must still span its latest child activity, + // not render as a 0ms bar. + it("root end covers latest child activity when the turn never completed", () => { + const incomplete: Turn = { + turnId: "turn-x", startTime: 1000, endTime: 1000, // never advanced (no task_complete) + completed: false, aborted: false, subagents: [], + steps: [{ + index: 0, startTime: 1200, endTime: 1800, toolCalls: [ + // e.g. wait_agent finishing when a subagent completes, long after start + { callId: "c", name: "wait_agent", args: {}, startTime: 1300, endTime: 9000 }, + ], + }], + }; + const root = planTurnSpans(sessionMeta, incomplete, ctx).find((s) => s.kind === "AGENT")!; + expect(root.startTime).toBe(1000); + expect(root.endTime).toBe(9000); // covers the latest tool end, not 1000 + }); }); describe("planToolSpan enriched metadata", () => { diff --git a/plugins/tracing/test/transcript.test.ts b/plugins/tracing/test/transcript.test.ts index 698119b..4aa576e 100644 --- a/plugins/tracing/test/transcript.test.ts +++ b/plugins/tracing/test/transcript.test.ts @@ -55,6 +55,21 @@ describe("parseRollout", () => { expect(t.steps[0]!.text).toBeUndefined(); }); + it("captures final output from agent_message even without task_complete (live-read race)", () => { + // Codex writes agent_message BEFORE the terminal task_complete (which lands + // just after the Stop hook fires). finalOutput must come from agent_message + // so a live run's root span isn't empty. No task_complete line on purpose. + const lines: RolloutLine[] = [ + { timestamp: ts(100), type: "session_meta", payload: { id: "sess-1" } }, + { timestamp: ts(101), type: "event_msg", payload: { type: "task_started", turn_id: "turn-1" } }, + { timestamp: ts(102), type: "event_msg", payload: { type: "user_message", message: "hi" } }, + { timestamp: ts(103), type: "event_msg", payload: { type: "agent_message", message: "Hey there" } }, + ]; + const { turns } = parseRollout(lines); + expect(turns[0]!.finalOutput).toBe("Hey there"); + expect(turns[0]!.completed).toBe(false); // still not marked complete (no task_complete) + }); + it("enriches tool calls from *_end event_msg events", async () => { const lines = await readRollout(toolDetailFixture); const { turns } = parseRollout(lines);