From 04c93a35099be9bd1e98db03900fcf931d2159a9 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 27 Oct 2026 21:05:53 +0800 Subject: [PATCH 1/3] fix(pi): preserve custom inputs and long responses --- bt-daemon/src/translate/pi.rs | 4 +++- src/plugins/pi/content/src/index.ts | 7 ++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/bt-daemon/src/translate/pi.rs b/bt-daemon/src/translate/pi.rs index 584b4bd2..477df5a0 100644 --- a/bt-daemon/src/translate/pi.rs +++ b/bt-daemon/src/translate/pi.rs @@ -55,6 +55,8 @@ impl TranslatorFactory for crate::agents::Pi { struct BeforeAgentStart { #[serde(default)] prompt: Option, + #[serde(default)] + message: Option, } /// `ctx.model`: the selected model, which may be a Pi 1.0 virtual model that @@ -649,7 +651,7 @@ impl PiTranslator { self.pending_llms.clear(); self.tools.clear(); let id = ids::span_id(&self.session_id, &format!("turn:{}", self.turn_seq)); - let input = event.prompt.unwrap_or(Value::Null); + let input = event.prompt.or(event.message).unwrap_or(Value::Null); let skills = input.as_str().map(explicit_skills).unwrap_or_default(); let mut metadata = without_nulls(json!({ "turn_number": self.turn_seq, diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 756cff8c..1b0d0f05 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -77,6 +77,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { let legacyContinuationSessionId: string | undefined; let lastContext: ExtensionContext | undefined; let awaitingFirstToken = false; + let lastStreamingKeepaliveAt = 0; let uiGeneration = 0; // The session_start refresh runs before the daemon has delivered a trace, // so agent runs poll status until the current session shows its link. @@ -196,20 +197,24 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { pi.on("before_agent_start", async (event, ctx) => forward("before_agent_start", event, ctx)); pi.on("context", async (event, ctx) => { awaitingFirstToken = true; + lastStreamingKeepaliveAt = 0; await forward("context", event, ctx); }); pi.on("before_provider_request", async (event) => forward("before_provider_request", event)); pi.on("after_provider_response", async (event) => forward("after_provider_response", event)); pi.on("message_update", async (event) => { - if (!awaitingFirstToken) return; const kind = streamingUpdateKind(event); if (!kind || !["text_delta", "thinking_delta", "text", "thinking"].includes(kind)) return; + const now = Date.now(); + if (!awaitingFirstToken && now - lastStreamingKeepaliveAt < 10_000) return; awaitingFirstToken = false; + lastStreamingKeepaliveAt = now; await forward("message_update", event); }); pi.on("thinking_level_select", async (event) => forward("thinking_level_select", event)); pi.on("message_end", async (event) => { awaitingFirstToken = false; + lastStreamingKeepaliveAt = 0; await forward("message_end", event); }); pi.on("tool_execution_start", async (event) => forward("tool_execution_start", event)); From 18435183c0cf4695a1ade81ed942ce104583e760 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 27 Oct 2026 21:14:08 +0800 Subject: [PATCH 2/3] test(pi): cover custom turns and long tool calls --- bt-daemon/tests/pi_translator.rs | 50 +++++++++++++++++++ .../pi/content/src/daemon-adapter.test.ts | 24 ++++++++- src/plugins/pi/content/src/index.ts | 8 +-- 3 files changed, 78 insertions(+), 4 deletions(-) diff --git a/bt-daemon/tests/pi_translator.rs b/bt-daemon/tests/pi_translator.rs index 953b1a5d..0055a952 100644 --- a/bt-daemon/tests/pi_translator.rs +++ b/bt-daemon/tests/pi_translator.rs @@ -99,6 +99,56 @@ fn pi_records_the_running_host_identity() { assert_eq!(root.metadata.as_ref().unwrap()["pi_version"], "18.12.6"); } +#[test] +fn pi_captures_custom_trigger_input_and_keeps_following_spans_in_its_turn() { + let mut translator = Registry::default_agents().create("pi", "pi-session"); + let ctx = SessionCtx { + session_id: "pi-session".into(), + config: None, + }; + let events = [ + event( + "before_agent_start", + 1, + json!({"message":"run the custom task"}), + ), + event( + "context", + 2, + json!({"messages":[{"role":"user","content":"run the custom task"}]}), + ), + event( + "message_end", + 3, + json!({"message":{"role":"assistant","provider":"openai","model":"gpt-5","content":[{"type":"text","text":"working"}]}}), + ), + event( + "tool_execution_start", + 4, + json!({"toolCallId":"call-1","toolName":"read","args":{"path":"README.md"}}), + ), + event( + "tool_execution_end", + 5, + json!({"toolCallId":"call-1","toolName":"read","result":"contents","isError":false}), + ), + ]; + let rows = reduce( + events + .iter() + .flat_map(|event| translator.handle(event, &ctx).unwrap()) + .collect(), + ); + let turn = rows.values().find(|row| row.is_turn).unwrap(); + assert_eq!(turn.input, Some(json!("run the custom task"))); + for row in rows + .values() + .filter(|row| matches!(row.span_type, SpanType::Llm | SpanType::Tool)) + { + assert_eq!(row.parent_span_ids, vec![turn.span_id.clone()]); + } +} + #[test] fn pi_builds_turn_llm_tool_compaction_and_shutdown_spans() { let registry = Registry::default_agents(); diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index 3467fc49..d245b30b 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -1,4 +1,4 @@ -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; const mockState = vi.hoisted(() => ({ logs: [] as Array>, @@ -106,6 +106,8 @@ function fakeContext(ui: Record = { setStatus: vi.fn(), setWidg } describe("Pi daemon adapter", () => { + afterEach(() => vi.restoreAllMocks()); + beforeEach(() => { mockState.logs.length = 0; mockState.flushes.length = 0; @@ -367,6 +369,26 @@ describe("Pi daemon adapter", () => { } }); + it("uses throttled tool-call deltas to keep long generations alive", async () => { + let now = 100_000; + vi.spyOn(Date, "now").mockImplementation(() => now); + const handlers = await loadExtension(); + const ctx = fakeContext(); + await handlers.get("context")?.({ messages: [] }, ctx); + const update = () => handlers.get("message_update")?.({ + type: "message_update", + assistantMessageEvent: { type: "toolcall_delta", delta: "{" }, + }); + + await update(); + now += 9_999; + await update(); + now += 1; + await update(); + + expect(mockState.logs.filter((log) => log.event === "message_update")).toHaveLength(2); + }); + it("forwards legacy continuation state, read once per Pi session", async () => { mockState.legacyContinuation = { span: "legacy-root", diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 1b0d0f05..95bc7adf 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -204,10 +204,12 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { pi.on("after_provider_response", async (event) => forward("after_provider_response", event)); pi.on("message_update", async (event) => { const kind = streamingUpdateKind(event); - if (!kind || !["text_delta", "thinking_delta", "text", "thinking"].includes(kind)) return; + const tokenKinds = ["text_delta", "thinking_delta", "text", "thinking"]; + const isFirstToken = awaitingFirstToken && !!kind && tokenKinds.includes(kind); + if (!kind || (!tokenKinds.includes(kind) && kind !== "toolcall_delta")) return; const now = Date.now(); - if (!awaitingFirstToken && now - lastStreamingKeepaliveAt < 10_000) return; - awaitingFirstToken = false; + if (!isFirstToken && now - lastStreamingKeepaliveAt < 10_000) return; + if (isFirstToken) awaitingFirstToken = false; lastStreamingKeepaliveAt = now; await forward("message_update", event); }); From 5dff895083338c060ff620e1ad463449d073f8c1 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 27 Oct 2026 21:28:05 +0800 Subject: [PATCH 3/3] style(pi): format keepalive regression test --- src/plugins/pi/content/src/daemon-adapter.test.ts | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index d245b30b..09ba906a 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -375,10 +375,11 @@ describe("Pi daemon adapter", () => { const handlers = await loadExtension(); const ctx = fakeContext(); await handlers.get("context")?.({ messages: [] }, ctx); - const update = () => handlers.get("message_update")?.({ - type: "message_update", - assistantMessageEvent: { type: "toolcall_delta", delta: "{" }, - }); + const update = () => + handlers.get("message_update")?.({ + type: "message_update", + assistantMessageEvent: { type: "toolcall_delta", delta: "{" }, + }); await update(); now += 9_999;