SDK-391 Continue legacy Pi sessions in bt-daemon

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<String>,
     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<String>,
+    total_turns: u32,
+    total_tool_calls: u32,
+}
+
+fn legacy_continuation(payload: &Value) -> Option<LegacyContinuation> {
+    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<Value> {
     let entry = event.compaction_entry.as_ref()?;
     let summary = entry.summary.clone()?;

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

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

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
 

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<string, unknown> | undefined,
+  legacyContinuationFor: vi.fn(),
   logGate: undefined as Promise<void> | undefined,
   statusGate: undefined as Promise<void> | 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<string, (...args: unknown[]) => Promise<unknown>>();
+    const pi = {
+      on: (name: string, handler: (...args: unknown[]) => Promise<unknown>) =>
+        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<string, (...args: unknown[]) => Promise<unknown>>();
+    const pi = {
+      on: (name: string, handler: (...args: unknown[]) => Promise<unknown>) =>
+        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");
+  });
 });