From eafdc493902f78cf26708457322068cf638dcb19 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 11:49:12 -0400 Subject: [PATCH 1/4] fix(pi): continue legacy traced sessions --- bt-daemon/src/translate/pi.rs | 46 ++++++++++++ bt-daemon/tests/pi_translator.rs | 70 +++++++++++++++++++ src/plugins/pi/content/README.md | 6 +- .../pi/content/src/daemon-adapter.test.ts | 37 ++++++++++ src/plugins/pi/content/src/index.ts | 7 ++ .../pi/content/src/legacy-session.test.ts | 54 ++++++++++++++ src/plugins/pi/content/src/legacy-session.ts | 64 +++++++++++++++++ 7 files changed, 283 insertions(+), 1 deletion(-) create mode 100644 src/plugins/pi/content/src/legacy-session.test.ts create mode 100644 src/plugins/pi/content/src/legacy-session.ts diff --git a/bt-daemon/src/translate/pi.rs b/bt-daemon/src/translate/pi.rs index f41d241..95151c0 100644 --- a/bt-daemon/src/translate/pi.rs +++ b/bt-daemon/src/translate/pi.rs @@ -493,6 +493,17 @@ impl PiTranslator { return Vec::new(); } self.opened = true; + 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; + // 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 +877,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_continuation")?; + let root_span_id = value.get("root_span_id")?.as_str()?.to_owned(); + if root_span_id.is_empty() { + return None; + } + let trace_root_span_id = value + .get("trace_root_span_id") + .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("total_turns")?.as_u64()?).ok()?; + let total_tool_calls = u32::try_from(value.get("total_tool_calls")?.as_u64()?).ok()?; + Some(LegacyContinuation { + root_span_id, + trace_root_span_id, + parent_span_id: value + .get("parent_span_id") + .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/pi_translator.rs b/bt-daemon/tests/pi_translator.rs index fb1dd73..7800ae4 100644 --- a/bt-daemon/tests/pi_translator.rs +++ b/bt-daemon/tests/pi_translator.rs @@ -294,6 +294,76 @@ 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_continuation"] = json!({ + "root_span_id":"legacy-root", + "trace_root_span_id":"legacy-trace-root", + "parent_span_id":"upstream-parent", + "total_turns":3, + "total_tool_calls":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_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..e536cb3 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -5,6 +5,7 @@ const mockState = vi.hoisted(() => ({ flushes: [] as string[], closed: 0, claim: true, + legacyContinuation: undefined as Record | undefined, logGate: undefined as Promise | undefined, statusGate: undefined as Promise | undefined, })); @@ -61,12 +62,17 @@ vi.mock("./config.ts", () => ({ }), })); +vi.mock("./legacy-session.ts", () => ({ + legacyContinuationFor: () => 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.logGate = undefined; mockState.statusGate = undefined; }); @@ -295,4 +301,35 @@ 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 = { + root_span_id: "legacy-root", + trace_root_span_id: "legacy-trace", + total_turns: 3, + total_tool_calls: 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_continuation: mockState.legacyContinuation, + }); + }); }); diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 7eefaa2..0fdb709 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"; @@ -50,6 +51,7 @@ function sessionDescriptor(ctx: ExtensionContext): { sessionId: string; sessionFile?: string; nativeSessionId?: string; + legacyContinuation?: LegacyContinuation; } { const sessionFile = ctx.sessionManager.getSessionFile(); const nativeSessionId = ctx.sessionManager.getSessionId(); @@ -57,6 +59,7 @@ function sessionDescriptor(ctx: ExtensionContext): { sessionId: sessionKeyFor(sessionFile, nativeSessionId, ctx.cwd), sessionFile, nativeSessionId, + legacyContinuation: legacyContinuationFor(sessionFile), }; } @@ -66,6 +69,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { if (!claimManagedTracingInstance("pi")) return; let sessionId: string | undefined; + let legacyContinuation: LegacyContinuation | undefined; let lastContext: ExtensionContext | undefined; let awaitingFirstToken = false; let uiGeneration = 0; @@ -91,6 +95,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { const descriptor = sessionDescriptor(ctx); if (sessionId !== descriptor.sessionId) uiGeneration += 1; sessionId = descriptor.sessionId; + legacyContinuation = descriptor.legacyContinuation; return descriptor; }; @@ -136,6 +141,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { extension_version: EXTENSION_VERSION, session_file: descriptor?.sessionFile, native_session_id: descriptor?.nativeSessionId, + legacy_continuation: descriptor?.legacyContinuation ?? legacyContinuation, cwd: ctx?.cwd, model: nativePayload(ctx?.model), }, @@ -184,6 +190,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { // request that finishes during shutdown cannot restore stale state. uiGeneration += 1; sessionId = undefined; + legacyContinuation = 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..dab2bdd --- /dev/null +++ b/src/plugins/pi/content/src/legacy-session.test.ts @@ -0,0 +1,54 @@ +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } 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("returns only valid continuation state for the matching session file", 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, + }, + }, + }), + ); + + expect(legacyContinuationFor("/tmp/live.jsonl")).toEqual({ + root_span_id: "legacy-root", + trace_root_span_id: "legacy-trace", + parent_span_id: "upstream", + total_turns: 3, + total_tool_calls: 7, + }); + expect(legacyContinuationFor("/tmp/other.jsonl")).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"); + expect(legacyContinuationFor("/tmp/live.jsonl")).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..bd73d34 --- /dev/null +++ b/src/plugins/pi/content/src/legacy-session.ts @@ -0,0 +1,64 @@ +import { existsSync, readFileSync } from "node:fs"; +import { homedir } from "node:os"; +import { join } from "node:path"; + +export interface LegacyContinuation { + root_span_id: string; + trace_root_span_id: string; + parent_span_id?: string; + total_turns: number; + total_tool_calls: 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 function legacyContinuationFor( + sessionFile: string | undefined, +): LegacyContinuation | undefined { + 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"); + if (!existsSync(stateFile)) return undefined; + + try { + const parsed: unknown = JSON.parse(readFileSync(stateFile, "utf8")); + const sessions = objectValue(parsed)?.sessions; + const session = objectValue(objectValue(sessions)?.[`file:${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 { + root_span_id: rootSpanId, + trace_root_span_id: traceRootSpanId, + ...(stringValue(session?.parentSpanId) + ? { parent_span_id: stringValue(session?.parentSpanId) } + : {}), + total_turns: totalTurns, + total_tool_calls: 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; +} From 63505eb167d6cf3aa7f608abc8542c645d77b10a Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 14:12:44 -0400 Subject: [PATCH 2/4] fix(pi): validate legacy session migration --- bt-daemon/src/translate/pi.rs | 12 ++++----- bt-daemon/tests/codex_translator.rs | 2 ++ bt-daemon/tests/pi_translator.rs | 12 ++++----- .../pi/content/src/daemon-adapter.test.ts | 10 +++---- src/plugins/pi/content/src/index.ts | 2 +- .../pi/content/src/legacy-session.test.ts | 19 ++++++++----- src/plugins/pi/content/src/legacy-session.ts | 27 +++++++++---------- 7 files changed, 45 insertions(+), 39 deletions(-) diff --git a/bt-daemon/src/translate/pi.rs b/bt-daemon/src/translate/pi.rs index 95151c0..f9334d6 100644 --- a/bt-daemon/src/translate/pi.rs +++ b/bt-daemon/src/translate/pi.rs @@ -886,24 +886,24 @@ struct LegacyContinuation { } fn legacy_continuation(payload: &Value) -> Option { - let value = payload.get("legacy_continuation")?; - let root_span_id = value.get("root_span_id")?.as_str()?.to_owned(); + 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_root_span_id") + .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("total_turns")?.as_u64()?).ok()?; - let total_tool_calls = u32::try_from(value.get("total_tool_calls")?.as_u64()?).ok()?; + 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_span_id") + .get("parent") .and_then(Value::as_str) .filter(|id| !id.is_empty()) .map(str::to_owned), 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 7800ae4..844c38b 100644 --- a/bt-daemon/tests/pi_translator.rs +++ b/bt-daemon/tests/pi_translator.rs @@ -303,12 +303,12 @@ fn pi_adopts_a_legacy_root_and_continues_its_turn_sequence() { config: None, }; let mut start = event("session_start", 1, json!({"reason":"resume"})); - start.payload["legacy_continuation"] = json!({ - "root_span_id":"legacy-root", - "trace_root_span_id":"legacy-trace-root", - "parent_span_id":"upstream-parent", - "total_turns":3, - "total_tool_calls":7, + 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(); diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index e536cb3..36a0029 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -304,10 +304,10 @@ describe("Pi daemon adapter", () => { it("forwards legacy continuation state with the first daemon event", async () => { mockState.legacyContinuation = { - root_span_id: "legacy-root", - trace_root_span_id: "legacy-trace", - total_turns: 3, - total_tool_calls: 7, + span: "legacy-root", + trace: "legacy-trace", + turns: 3, + tools: 7, }; const handlers = new Map Promise>(); const pi = { @@ -329,7 +329,7 @@ describe("Pi daemon adapter", () => { await handlers.get("session_start")?.({ reason: "resume" }, ctx); expect(mockState.logs[0]?.payload).toMatchObject({ - legacy_continuation: mockState.legacyContinuation, + legacy_resume: mockState.legacyContinuation, }); }); }); diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 0fdb709..fcc979c 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -141,7 +141,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { extension_version: EXTENSION_VERSION, session_file: descriptor?.sessionFile, native_session_id: descriptor?.nativeSessionId, - legacy_continuation: descriptor?.legacyContinuation ?? legacyContinuation, + legacy_resume: descriptor?.legacyContinuation ?? legacyContinuation, cwd: ctx?.cwd, model: nativePayload(ctx?.model), }, diff --git a/src/plugins/pi/content/src/legacy-session.test.ts b/src/plugins/pi/content/src/legacy-session.test.ts index dab2bdd..5a477d9 100644 --- a/src/plugins/pi/content/src/legacy-session.test.ts +++ b/src/plugins/pi/content/src/legacy-session.test.ts @@ -1,6 +1,6 @@ import { mkdtemp, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; -import { join } from "node:path"; +import { join, relative } from "node:path"; import { afterEach, describe, expect, it } from "vitest"; import { legacyContinuationFor } from "./legacy-session.ts"; @@ -14,7 +14,7 @@ afterEach(async () => { }); describe("legacy Pi session state", () => { - it("returns only valid continuation state for the matching session file", async () => { + 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; @@ -34,12 +34,17 @@ describe("legacy Pi session state", () => { }), ); + // @braintrust/pi-extension 0.9.0 persisted this exact sessions.json + // shape, including camelCase fields and an absolute file session key. expect(legacyContinuationFor("/tmp/live.jsonl")).toEqual({ - root_span_id: "legacy-root", - trace_root_span_id: "legacy-trace", - parent_span_id: "upstream", - total_turns: 3, - total_tool_calls: 7, + span: "legacy-root", + trace: "legacy-trace", + parent: "upstream", + turns: 3, + tools: 7, + }); + expect(legacyContinuationFor(relative(process.cwd(), "/tmp/live.jsonl"))).toMatchObject({ + span: "legacy-root", }); expect(legacyContinuationFor("/tmp/other.jsonl")).toBeUndefined(); }); diff --git a/src/plugins/pi/content/src/legacy-session.ts b/src/plugins/pi/content/src/legacy-session.ts index bd73d34..c000bd9 100644 --- a/src/plugins/pi/content/src/legacy-session.ts +++ b/src/plugins/pi/content/src/legacy-session.ts @@ -1,13 +1,13 @@ import { existsSync, readFileSync } from "node:fs"; import { homedir } from "node:os"; -import { join } from "node:path"; +import { join, resolve } from "node:path"; export interface LegacyContinuation { - root_span_id: string; - trace_root_span_id: string; - parent_span_id?: string; - total_turns: number; - total_tool_calls: number; + span: string; + trace: string; + parent?: string; + turns: number; + tools: number; } /** @@ -28,7 +28,8 @@ export function legacyContinuationFor( try { const parsed: unknown = JSON.parse(readFileSync(stateFile, "utf8")); const sessions = objectValue(parsed)?.sessions; - const session = objectValue(objectValue(sessions)?.[`file:${sessionFile}`]); + // 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; @@ -36,13 +37,11 @@ export function legacyContinuationFor( const totalToolCalls = countValue(session?.totalToolCalls); if (totalTurns === undefined || totalToolCalls === undefined) return undefined; return { - root_span_id: rootSpanId, - trace_root_span_id: traceRootSpanId, - ...(stringValue(session?.parentSpanId) - ? { parent_span_id: stringValue(session?.parentSpanId) } - : {}), - total_turns: totalTurns, - total_tool_calls: totalToolCalls, + span: rootSpanId, + trace: traceRootSpanId, + ...(stringValue(session?.parentSpanId) ? { parent: stringValue(session?.parentSpanId) } : {}), + turns: totalTurns, + tools: totalToolCalls, }; } catch { return undefined; From 7feb844643861ef9c68f44981250f75ebdaa95d0 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 15:33:06 -0400 Subject: [PATCH 3/4] fix(pi): cache legacy migration state --- .../pi/content/src/daemon-adapter.test.ts | 33 ++++++++++++++++++- src/plugins/pi/content/src/index.ts | 17 ++++++---- .../pi/content/src/legacy-session.test.ts | 10 +++--- src/plugins/pi/content/src/legacy-session.ts | 10 +++--- 4 files changed, 53 insertions(+), 17 deletions(-) diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index 36a0029..90d8928 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -6,6 +6,7 @@ const mockState = vi.hoisted(() => ({ closed: 0, claim: true, legacyContinuation: undefined as Record | undefined, + legacyContinuationFor: vi.fn(), logGate: undefined as Promise | undefined, statusGate: undefined as Promise | undefined, })); @@ -63,7 +64,10 @@ vi.mock("./config.ts", () => ({ })); vi.mock("./legacy-session.ts", () => ({ - legacyContinuationFor: () => mockState.legacyContinuation, + legacyContinuationFor: (...args: unknown[]) => { + mockState.legacyContinuationFor(...args); + return mockState.legacyContinuation; + }, })); describe("Pi daemon adapter", () => { @@ -73,6 +77,7 @@ describe("Pi daemon adapter", () => { mockState.closed = 0; mockState.claim = true; mockState.legacyContinuation = undefined; + mockState.legacyContinuationFor.mockClear(); mockState.logGate = undefined; mockState.statusGate = undefined; }); @@ -332,4 +337,30 @@ describe("Pi daemon adapter", () => { 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 fcc979c..98b4f6d 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -51,7 +51,6 @@ function sessionDescriptor(ctx: ExtensionContext): { sessionId: string; sessionFile?: string; nativeSessionId?: string; - legacyContinuation?: LegacyContinuation; } { const sessionFile = ctx.sessionManager.getSessionFile(); const nativeSessionId = ctx.sessionManager.getSessionId(); @@ -59,7 +58,6 @@ function sessionDescriptor(ctx: ExtensionContext): { sessionId: sessionKeyFor(sessionFile, nativeSessionId, ctx.cwd), sessionFile, nativeSessionId, - legacyContinuation: legacyContinuationFor(sessionFile), }; } @@ -70,6 +68,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { let sessionId: string | undefined; let legacyContinuation: LegacyContinuation | undefined; + let legacyContinuationSessionId: string | undefined; let lastContext: ExtensionContext | undefined; let awaitingFirstToken = false; let uiGeneration = 0; @@ -90,12 +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; - legacyContinuation = descriptor.legacyContinuation; + 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; }; @@ -128,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", @@ -141,7 +145,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { extension_version: EXTENSION_VERSION, session_file: descriptor?.sessionFile, native_session_id: descriptor?.nativeSessionId, - legacy_resume: descriptor?.legacyContinuation ?? legacyContinuation, + legacy_resume: legacyContinuation, cwd: ctx?.cwd, model: nativePayload(ctx?.model), }, @@ -191,6 +195,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { 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 index 5a477d9..60be30b 100644 --- a/src/plugins/pi/content/src/legacy-session.test.ts +++ b/src/plugins/pi/content/src/legacy-session.test.ts @@ -36,17 +36,19 @@ describe("legacy Pi session state", () => { // @braintrust/pi-extension 0.9.0 persisted this exact sessions.json // shape, including camelCase fields and an absolute file session key. - expect(legacyContinuationFor("/tmp/live.jsonl")).toEqual({ + await expect(legacyContinuationFor("/tmp/live.jsonl")).resolves.toEqual({ span: "legacy-root", trace: "legacy-trace", parent: "upstream", turns: 3, tools: 7, }); - expect(legacyContinuationFor(relative(process.cwd(), "/tmp/live.jsonl"))).toMatchObject({ + await expect( + legacyContinuationFor(relative(process.cwd(), "/tmp/live.jsonl")), + ).resolves.toMatchObject({ span: "legacy-root", }); - expect(legacyContinuationFor("/tmp/other.jsonl")).toBeUndefined(); + await expect(legacyContinuationFor("/tmp/other.jsonl")).resolves.toBeUndefined(); }); it("fails open when legacy state is malformed or incomplete", async () => { @@ -54,6 +56,6 @@ describe("legacy Pi session state", () => { stateDirs.push(stateDir); process.env.BRAINTRUST_STATE_DIR = stateDir; await writeFile(join(stateDir, "sessions.json"), "not json"); - expect(legacyContinuationFor("/tmp/live.jsonl")).toBeUndefined(); + 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 index c000bd9..48747a3 100644 --- a/src/plugins/pi/content/src/legacy-session.ts +++ b/src/plugins/pi/content/src/legacy-session.ts @@ -1,4 +1,4 @@ -import { existsSync, readFileSync } from "node:fs"; +import { readFile } from "node:fs/promises"; import { homedir } from "node:os"; import { join, resolve } from "node:path"; @@ -15,18 +15,16 @@ export interface LegacyContinuation { * This is deliberately a compatibility boundary: invalid or unavailable state * must leave a new daemon session untouched. */ -export function legacyContinuationFor( +export async function legacyContinuationFor( sessionFile: string | undefined, -): LegacyContinuation | 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"); - if (!existsSync(stateFile)) return undefined; - try { - const parsed: unknown = JSON.parse(readFileSync(stateFile, "utf8")); + 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)}`]); From dea5ae015b5e3370c2f1288263bde02d2df7fc45 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 17:00:07 -0400 Subject: [PATCH 4/4] fix(pi): preserve migrated session counters --- bt-daemon/src/translate/pi.rs | 9 ++++ bt-daemon/tests/pi_translator.rs | 75 ++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/bt-daemon/src/translate/pi.rs b/bt-daemon/src/translate/pi.rs index f9334d6..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,12 +498,16 @@ 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. diff --git a/bt-daemon/tests/pi_translator.rs b/bt-daemon/tests/pi_translator.rs index 844c38b..bc393d5 100644 --- a/bt-daemon/tests/pi_translator.rs +++ b/bt-daemon/tests/pi_translator.rs @@ -364,6 +364,81 @@ fn pi_adopts_a_legacy_root_and_continues_its_turn_sequence() { 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();