Fix Pi custom turn inputs and long responses

bt-daemon/src/translate/pi.rs

@@ -55,6 +55,8 @@ impl TranslatorFactory for crate::agents::Pi {
 struct BeforeAgentStart {
     #[serde(default)]
     prompt: Option<Value>,
+    #[serde(default)]
+    message: Option<Value>,
 }
 
 /// `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,

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

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<Record<string, unknown>>,
@@ -106,6 +106,8 @@ function fakeContext(ui: Record<string, unknown> = { setStatus: vi.fn(), setWidg
 }
 
 describe("Pi daemon adapter", () => {
+  afterEach(() => vi.restoreAllMocks());
+
   beforeEach(() => {
     mockState.logs.length = 0;
     mockState.flushes.length = 0;
@@ -367,6 +369,27 @@ 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",

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,26 @@ 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;
-    awaitingFirstToken = false;
+    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 (!isFirstToken && now - lastStreamingKeepaliveAt < 10_000) return;
+    if (isFirstToken) 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));