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);