diff --git a/bt-daemon/src/translate/pi.rs b/bt-daemon/src/translate/pi.rs index f41d241..dc7605d 100644 --- a/bt-daemon/src/translate/pi.rs +++ b/bt-daemon/src/translate/pi.rs @@ -34,6 +34,7 @@ impl TranslatorFactory for PiTranslatorFactory { effective_root_span_id: String::new(), external_parent: None, opened: false, + legacy_root_adopted: false, turn: None, turn_seq: 0, llm_seq: 0, @@ -308,6 +309,10 @@ struct PiTranslator { effective_root_span_id: String, external_parent: Option, opened: bool, + // A legacy state file remains available after migration, so its counters + // are initialization input only. Later reopen/replay events must retain + // the daemon's accumulated counters and deterministic turn sequence. + legacy_root_adopted: bool, turn: Option<(String, Value)>, turn_seq: u32, llm_seq: u32, @@ -493,6 +498,21 @@ impl PiTranslator { return Vec::new(); } self.opened = true; + if self.legacy_root_adopted { + return Vec::new(); + } + if let Some(legacy) = legacy_continuation(&envelope.payload) { + self.root_span_id = legacy.root_span_id; + self.effective_root_span_id = legacy.trace_root_span_id; + self.external_parent = legacy.parent_span_id; + self.turn_seq = legacy.total_turns; + self.total_tools = legacy.total_tool_calls; + self.legacy_root_adopted = true; + // The legacy extension already created this root. Re-emitting an + // insert could replace its metadata and attachment, so only emit + // descendants and terminal aggregate merges from this point on. + return Vec::new(); + } let attached = ctx .config .as_ref() @@ -866,6 +886,41 @@ impl PiTranslator { } } +struct LegacyContinuation { + root_span_id: String, + trace_root_span_id: String, + parent_span_id: Option, + total_turns: u32, + total_tool_calls: u32, +} + +fn legacy_continuation(payload: &Value) -> Option { + let value = payload.get("legacy_resume")?; + let root_span_id = value.get("span")?.as_str()?.to_owned(); + if root_span_id.is_empty() { + return None; + } + let trace_root_span_id = value + .get("trace") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .unwrap_or(&root_span_id) + .to_owned(); + let total_turns = u32::try_from(value.get("turns")?.as_u64()?).ok()?; + let total_tool_calls = u32::try_from(value.get("tools")?.as_u64()?).ok()?; + Some(LegacyContinuation { + root_span_id, + trace_root_span_id, + parent_span_id: value + .get("parent") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned), + total_turns, + total_tool_calls, + }) +} + fn compaction_message(event: &SessionCompact) -> Option { let entry = event.compaction_entry.as_ref()?; let summary = entry.summary.clone()?; diff --git a/bt-daemon/tests/codex_translator.rs b/bt-daemon/tests/codex_translator.rs index 2908241..a7d1208 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -453,6 +453,7 @@ fn attached_codex_root_merge_preserves_external_parent() { flush_mode: FlushMode::FireAndForget, additional_metadata: None, tags: Vec::new(), + span_plugins: Vec::new(), }), }; let registry = Registry::default_agents(); @@ -1664,6 +1665,7 @@ fn codex_root_source_merge_after_stop_keeps_external_parent() { flush_mode: FlushMode::FireAndForget, additional_metadata: None, tags: Vec::new(), + span_plugins: Vec::new(), }), }; let registry = Registry::default_agents(); diff --git a/bt-daemon/tests/pi_translator.rs b/bt-daemon/tests/pi_translator.rs index fb1dd73..bc393d5 100644 --- a/bt-daemon/tests/pi_translator.rs +++ b/bt-daemon/tests/pi_translator.rs @@ -294,6 +294,151 @@ fn pi_additional_metadata_reaches_roots_without_overriding_session_fields() { assert_eq!(root.tags, Some(vec!["ci".into(), "docs".into()])); } +#[test] +fn pi_adopts_a_legacy_root_and_continues_its_turn_sequence() { + let registry = Registry::default_agents(); + let mut translator = registry.create("pi", "daemon-session"); + let ctx = SessionCtx { + session_id: "daemon-session".into(), + config: None, + }; + let mut start = event("session_start", 1, json!({"reason":"resume"})); + start.payload["legacy_resume"] = json!({ + "span":"legacy-root", + "trace":"legacy-trace-root", + "parent":"upstream-parent", + "turns":3, + "tools":7, + }); + + let mut ops = translator.handle(&start, &ctx).unwrap(); + assert!( + ops.is_empty(), + "the existing legacy root must not be reinserted" + ); + ops.extend( + translator + .handle( + &event("before_agent_start", 2, json!({"prompt":"continue"})), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle(&event("agent_end", 3, json!({"messages":[]})), &ctx) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event("session_shutdown", 4, json!({"reason":"quit"})), + &ctx, + ) + .unwrap(), + ); + + let turn = ops + .iter() + .find_map(|op| match op { + SpanOp::Insert(row) if row.name == "Turn 4" => Some(row), + _ => None, + }) + .expect("first daemon turn continues the legacy sequence"); + assert_eq!(turn.root_span_id, "legacy-trace-root"); + assert_eq!(turn.parent_span_ids, ["legacy-root"]); + assert!(ops + .iter() + .all(|op| !matches!(op, SpanOp::Insert(row) if row.name == "Pi"))); + + let root = ops + .iter() + .find_map(|op| match op { + SpanOp::Merge(row) if row.span_id == "legacy-root" => Some(row), + _ => None, + }) + .expect("shutdown updates the adopted root"); + assert_eq!(root.root_span_id, "legacy-trace-root"); + assert_eq!(root.parent_span_ids, ["upstream-parent"]); + assert_eq!(root.metadata.as_ref().unwrap()["total_turns"], 4); + assert_eq!(root.metadata.as_ref().unwrap()["total_tool_calls"], 7); +} + +#[test] +fn pi_does_not_reapply_legacy_counters_when_a_migrated_session_reopens() { + let registry = Registry::default_agents(); + let mut translator = registry.create("pi", "daemon-session"); + let ctx = SessionCtx { + session_id: "daemon-session".into(), + config: None, + }; + let mut legacy_start = event("session_start", 1, json!({"reason":"resume"})); + legacy_start.payload["legacy_resume"] = json!({ + "span":"legacy-root", + "trace":"legacy-trace-root", + "turns":3, + "tools":7, + }); + + translator.handle(&legacy_start, &ctx).unwrap(); + translator + .handle( + &event( + "before_agent_start", + 2, + json!({"prompt":"first daemon turn"}), + ), + &ctx, + ) + .unwrap(); + translator + .handle( + &event("session_shutdown", 3, json!({"reason":"quit"})), + &ctx, + ) + .unwrap(); + + let mut reopened = event("session_start", 4, json!({"reason":"resume"})); + reopened.payload["legacy_resume"] = legacy_start.payload["legacy_resume"].clone(); + let mut ops = translator.handle(&reopened, &ctx).unwrap(); + ops.extend( + translator + .handle( + &event( + "before_agent_start", + 5, + json!({"prompt":"second daemon turn"}), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event("session_shutdown", 6, json!({"reason":"quit"})), + &ctx, + ) + .unwrap(), + ); + + assert!(ops + .iter() + .any(|op| matches!(op, SpanOp::Insert(row) if row.name == "Turn 5"))); + assert!(ops + .iter() + .all(|op| !matches!(op, SpanOp::Insert(row) if row.name == "Turn 4"))); + let root = ops + .iter() + .find_map(|op| match op { + SpanOp::Merge(row) if row.span_id == "legacy-root" => Some(row), + _ => None, + }) + .expect("reopened session updates the adopted root"); + assert_eq!(root.metadata.as_ref().unwrap()["total_turns"], 5); + assert_eq!(root.metadata.as_ref().unwrap()["total_tool_calls"], 7); +} + #[test] fn pi_checkpoint_preserves_the_open_session_and_turn() { let registry = Registry::default_agents(); diff --git a/src/plugins/pi/content/README.md b/src/plugins/pi/content/README.md index 37d88ec..ea0d540 100644 --- a/src/plugins/pi/content/README.md +++ b/src/plugins/pi/content/README.md @@ -56,7 +56,11 @@ bt trace run --project my-coding-agent pi -- -p "summarize this repository" The `bt trace run` routing and metadata flags also accept their matching `BRAINTRUST_*` environment variables; a plain `pi` session's extension does not. -Historical import and live attach are not supported for Pi. +Historical import and live attach are not supported for Pi. Upgrading an active +session from the pre-daemon extension is supported: when its local legacy state +file is still present, the daemon continues the existing Braintrust root and +turn sequence. If that state was removed or is invalid, tracing safely starts a +new daemon session instead. ## Compatibility diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index 9608218..90d8928 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -5,6 +5,8 @@ const mockState = vi.hoisted(() => ({ flushes: [] as string[], closed: 0, claim: true, + legacyContinuation: undefined as Record | undefined, + legacyContinuationFor: vi.fn(), logGate: undefined as Promise | undefined, statusGate: undefined as Promise | undefined, })); @@ -61,12 +63,21 @@ vi.mock("./config.ts", () => ({ }), })); +vi.mock("./legacy-session.ts", () => ({ + legacyContinuationFor: (...args: unknown[]) => { + mockState.legacyContinuationFor(...args); + return mockState.legacyContinuation; + }, +})); + describe("Pi daemon adapter", () => { beforeEach(() => { mockState.logs.length = 0; mockState.flushes.length = 0; mockState.closed = 0; mockState.claim = true; + mockState.legacyContinuation = undefined; + mockState.legacyContinuationFor.mockClear(); mockState.logGate = undefined; mockState.statusGate = undefined; }); @@ -295,4 +306,61 @@ describe("Pi daemon adapter", () => { expect(statuses.length).toBeGreaterThan(0); expect(mockState.closed).toBe(2); }); + + it("forwards legacy continuation state with the first daemon event", async () => { + mockState.legacyContinuation = { + span: "legacy-root", + trace: "legacy-trace", + turns: 3, + tools: 7, + }; + const handlers = new Map Promise>(); + const pi = { + on: (name: string, handler: (...args: unknown[]) => Promise) => + handlers.set(name, handler), + }; + const ctx = { + cwd: "/tmp/project", + hasUI: false, + ui: { setStatus: vi.fn(), setWidget: vi.fn() }, + sessionManager: { + getSessionFile: () => "/tmp/session.jsonl", + getSessionId: () => "native-session", + }, + }; + const { default: extension } = await import("./index.ts"); + extension(pi as never); + + await handlers.get("session_start")?.({ reason: "resume" }, ctx); + + expect(mockState.logs[0]?.payload).toMatchObject({ + legacy_resume: mockState.legacyContinuation, + }); + }); + + it("reads legacy continuation state once per Pi session", async () => { + const handlers = new Map Promise>(); + const pi = { + on: (name: string, handler: (...args: unknown[]) => Promise) => + handlers.set(name, handler), + }; + const ctx = { + cwd: "/tmp/project", + hasUI: false, + ui: { setStatus: vi.fn(), setWidget: vi.fn() }, + sessionManager: { + getSessionFile: () => "/tmp/session.jsonl", + getSessionId: () => "native-session", + }, + }; + const { default: extension } = await import("./index.ts"); + extension(pi as never); + + await handlers.get("session_start")?.({}, ctx); + await handlers.get("context")?.({}, ctx); + await handlers.get("tool_execution_end")?.({}, ctx); + + expect(mockState.legacyContinuationFor).toHaveBeenCalledTimes(1); + expect(mockState.legacyContinuationFor).toHaveBeenCalledWith("/tmp/session.jsonl"); + }); }); diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 7eefaa2..98b4f6d 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -2,6 +2,7 @@ import { createHash, randomUUID } from "node:crypto"; import { resolve } from "node:path"; import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; import { loadConfig } from "./config.ts"; +import { legacyContinuationFor, type LegacyContinuation } from "./legacy-session.ts"; import { loadPiPackageMetadata } from "./pi-package.ts"; import { claimManagedTracingInstance, DaemonClient } from "./runtime/daemon-client.ts"; import { EXTENSION_VERSION } from "./version.ts"; @@ -66,6 +67,8 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { if (!claimManagedTracingInstance("pi")) return; let sessionId: string | undefined; + let legacyContinuation: LegacyContinuation | undefined; + let legacyContinuationSessionId: string | undefined; let lastContext: ExtensionContext | undefined; let awaitingFirstToken = false; let uiGeneration = 0; @@ -86,11 +89,17 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { requestTimeoutMs: UI_STATUS_TIMEOUT_MS, }); - const remember = (ctx: ExtensionContext): ReturnType => { + const remember = async (ctx: ExtensionContext): Promise> => { lastContext = ctx; const descriptor = sessionDescriptor(ctx); if (sessionId !== descriptor.sessionId) uiGeneration += 1; sessionId = descriptor.sessionId; + if (legacyContinuationSessionId !== descriptor.sessionId) { + // The legacy file is immutable migration input. Read it asynchronously + // once per Pi session before forwarding its first event. + legacyContinuation = await legacyContinuationFor(descriptor.sessionFile); + legacyContinuationSessionId = descriptor.sessionId; + } return descriptor; }; @@ -123,7 +132,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { ctx?: ExtensionContext, updateUi = false, ): Promise => { - const descriptor = ctx ? remember(ctx) : undefined; + const descriptor = ctx ? await remember(ctx) : undefined; if (!sessionId) return; await client.log({ source: "pi", @@ -136,6 +145,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { extension_version: EXTENSION_VERSION, session_file: descriptor?.sessionFile, native_session_id: descriptor?.nativeSessionId, + legacy_resume: legacyContinuation, cwd: ctx?.cwd, model: nativePayload(ctx?.model), }, @@ -184,6 +194,8 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { // request that finishes during shutdown cannot restore stale state. uiGeneration += 1; sessionId = undefined; + legacyContinuation = undefined; + legacyContinuationSessionId = undefined; lastContext = undefined; if (ctx.hasUI) { ctx.ui.setStatus(STATUS_KEY, undefined); diff --git a/src/plugins/pi/content/src/legacy-session.test.ts b/src/plugins/pi/content/src/legacy-session.test.ts new file mode 100644 index 0000000..60be30b --- /dev/null +++ b/src/plugins/pi/content/src/legacy-session.test.ts @@ -0,0 +1,61 @@ +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join, relative } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { legacyContinuationFor } from "./legacy-session.ts"; + +const previousStateDir = process.env.BRAINTRUST_STATE_DIR; +const stateDirs: string[] = []; + +afterEach(async () => { + if (previousStateDir === undefined) delete process.env.BRAINTRUST_STATE_DIR; + else process.env.BRAINTRUST_STATE_DIR = previousStateDir; + await Promise.all(stateDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +describe("legacy Pi session state", () => { + it("reads the published 0.9.0 legacy state shape for the matching session", async () => { + const stateDir = await mkdtemp(join(tmpdir(), "braintrust-pi-state-")); + stateDirs.push(stateDir); + process.env.BRAINTRUST_STATE_DIR = stateDir; + await writeFile( + join(stateDir, "sessions.json"), + JSON.stringify({ + version: 1, + sessions: { + "file:/tmp/live.jsonl": { + rootSpanId: "legacy-root", + traceRootSpanId: "legacy-trace", + parentSpanId: "upstream", + totalTurns: 3, + totalToolCalls: 7, + }, + }, + }), + ); + + // @braintrust/pi-extension 0.9.0 persisted this exact sessions.json + // shape, including camelCase fields and an absolute file session key. + await expect(legacyContinuationFor("/tmp/live.jsonl")).resolves.toEqual({ + span: "legacy-root", + trace: "legacy-trace", + parent: "upstream", + turns: 3, + tools: 7, + }); + await expect( + legacyContinuationFor(relative(process.cwd(), "/tmp/live.jsonl")), + ).resolves.toMatchObject({ + span: "legacy-root", + }); + await expect(legacyContinuationFor("/tmp/other.jsonl")).resolves.toBeUndefined(); + }); + + it("fails open when legacy state is malformed or incomplete", async () => { + const stateDir = await mkdtemp(join(tmpdir(), "braintrust-pi-state-")); + stateDirs.push(stateDir); + process.env.BRAINTRUST_STATE_DIR = stateDir; + await writeFile(join(stateDir, "sessions.json"), "not json"); + await expect(legacyContinuationFor("/tmp/live.jsonl")).resolves.toBeUndefined(); + }); +}); diff --git a/src/plugins/pi/content/src/legacy-session.ts b/src/plugins/pi/content/src/legacy-session.ts new file mode 100644 index 0000000..48747a3 --- /dev/null +++ b/src/plugins/pi/content/src/legacy-session.ts @@ -0,0 +1,61 @@ +import { readFile } from "node:fs/promises"; +import { homedir } from "node:os"; +import { join, resolve } from "node:path"; + +export interface LegacyContinuation { + span: string; + trace: string; + parent?: string; + turns: number; + tools: number; +} + +/** + * Read the small, local state file written by the pre-daemon Pi extension. + * This is deliberately a compatibility boundary: invalid or unavailable state + * must leave a new daemon session untouched. + */ +export async function legacyContinuationFor( + sessionFile: string | undefined, +): Promise { + if (!sessionFile) return undefined; + const stateDir = + process.env.BRAINTRUST_STATE_DIR ?? + join(homedir(), ".pi", "agent", "state", "braintrust-pi-extension"); + const stateFile = join(stateDir, "sessions.json"); + try { + const parsed: unknown = JSON.parse(await readFile(stateFile, "utf8")); + const sessions = objectValue(parsed)?.sessions; + // Version 0.9.0 stored the key with an absolute session-file path. + const session = objectValue(objectValue(sessions)?.[`file:${resolve(sessionFile)}`]); + const rootSpanId = stringValue(session?.rootSpanId); + if (!rootSpanId) return undefined; + const traceRootSpanId = stringValue(session?.traceRootSpanId) ?? rootSpanId; + const totalTurns = countValue(session?.totalTurns); + const totalToolCalls = countValue(session?.totalToolCalls); + if (totalTurns === undefined || totalToolCalls === undefined) return undefined; + return { + span: rootSpanId, + trace: traceRootSpanId, + ...(stringValue(session?.parentSpanId) ? { parent: stringValue(session?.parentSpanId) } : {}), + turns: totalTurns, + tools: totalToolCalls, + }; + } catch { + return undefined; + } +} + +function objectValue(value: unknown): Record | undefined { + return value !== null && typeof value === "object" && !Array.isArray(value) + ? (value as Record) + : undefined; +} + +function stringValue(value: unknown): string | undefined { + return typeof value === "string" && value.length > 0 ? value : undefined; +} + +function countValue(value: unknown): number | undefined { + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined; +}