🐛 fix: extract JSON from turn_end in PiAgent json mode NDJSON stream

src/agents/cli.ts

@@ -345,6 +345,45 @@ function extractTextFromJsonPayload(raw: string): string | undefined {
   return chunks.length ? chunks.join("") : undefined;
 }
 
+function extractTextFromPiNdjson(raw: string): string | undefined {
+  const trimmed = raw.trim();
+  if (!trimmed) return undefined;
+
+  const lines = trimmed.split(/\r?\n/).filter(Boolean);
+  let turnEndMessage: any = null;
+  let agentEndMessage: any = null;
+
+  for (let i = lines.length - 1; i >= 0; i--) {
+    try {
+      const parsed = JSON.parse(lines[i]!);
+      if (parsed.type === "turn_end" && parsed.message?.role === "assistant") {
+        turnEndMessage = parsed.message;
+        break;
+      }
+      if (parsed.type === "agent_end" && Array.isArray(parsed.messages)) {
+        for (let j = parsed.messages.length - 1; j >= 0; j--) {
+          const msg = parsed.messages[j];
+          if (msg?.role === "assistant") {
+            agentEndMessage = msg;
+            break;
+          }
+        }
+        if (agentEndMessage) break;
+      }
+    } catch {
+      continue;
+    }
+  }
+
+  const message = turnEndMessage ?? agentEndMessage;
+  if (message) {
+    const text = extractTextFromJsonValue(message);
+    if (text) return text;
+  }
+
+  return extractTextFromJsonPayload(raw);
+}
+
 function truncateToBytes(text: string, maxBytes?: number): string {
   if (!maxBytes || maxBytes <= 0) return text;
   const buf = Buffer.from(text, "utf8");
@@ -1318,8 +1357,13 @@ export class PiAgent extends BaseCliAgent {
       }
 
       const rawText = result.stdout.trim();
-      const output = mode === "json" ? tryParseJson(rawText) : rawText;
-      return buildGenerateResult(rawText, output, this.opts.model ?? "pi");
+      // In json mode, pi outputs NDJSON stream. Extract text from turn_end message
+      // rather than returning the first JSON object (session metadata).
+      const extractedText = mode === "json"
+        ? (extractTextFromPiNdjson(rawText) ?? rawText)
+        : rawText;
+      const output = tryParseJson(extractedText);
+      return buildGenerateResult(extractedText, output, this.opts.model ?? "pi");
     }
 
     // RPC mode

tests/pi-support.test.ts

@@ -281,4 +281,113 @@ import { afterEach, describe, expect, test } from "bun:test";
        await rm(fake.dir, { recursive: true, force: true });
      }
    });
+
+   test("PiAgent json mode extracts text from turn_end in NDJSON stream", async () => {
+     // Simulates real pi --mode json output: NDJSON stream with session metadata first,
+     // then message events, then turn_end containing the actual response.
+     const fake = await makeFakePi(`
+ const lines = [
+   JSON.stringify({ type: "session", version: 3, id: "test-session-id", timestamp: "2026-02-15T18:00:00.000Z", cwd: "/tmp" }),
+   JSON.stringify({ type: "agent_start" }),
+   JSON.stringify({ type: "turn_start" }),
+   JSON.stringify({ type: "message_start", message: { role: "user", content: [{ type: "text", text: "Hello" }] } }),
+   JSON.stringify({ type: "message_end", message: { role: "user", content: [{ type: "text", text: "Hello" }] } }),
+   JSON.stringify({ type: "message_start", message: { role: "assistant", content: [] } }),
+   JSON.stringify({ type: "message_update", assistantMessageEvent: { type: "text_delta", delta: "Here is" } }),
+   JSON.stringify({ type: "message_update", assistantMessageEvent: { type: "text_delta", delta: " your data" } }),
+   JSON.stringify({ type: "message_end", message: { role: "assistant", content: [{ type: "text", text: "Here is your data" }] } }),
+   JSON.stringify({ type: "turn_end", message: { role: "assistant", content: [{ type: "text", text: "Here is your data" }], stopReason: "stop" } }),
+   JSON.stringify({ type: "agent_end" })
+ ];
+ process.stdout.write(lines.join("\\n") + "\\n");
+ `);
+
+     try {
+       process.env.PATH = `${fake.dir}:${originalPath}`;
+
+       const agent = new PiAgent({
+         mode: "json",
+         model: "test-model",
+         env: { PATH: process.env.PATH! },
+       });
+
+       const result = await agent.generate({
+         messages: [{ role: "user", content: "Hello" }],
+       });
+
+       // Should extract text from turn_end, not from first JSON (session metadata)
+       expect(result.text).toBe("Here is your data");
+       // First JSON should NOT be parsed as output (would have "type: session")
+       expect(result.output).not.toHaveProperty("type", "session");
+     } finally {
+       await rm(fake.dir, { recursive: true, force: true });
+     }
+   });
+
+   test("PiAgent json mode extracts JSON from text content in turn_end", async () => {
+     // Simulates pi output where the agent returns JSON in the text content
+     const fake = await makeFakePi(`
+ const lines = [
+   JSON.stringify({ type: "session", version: 3, id: "test-session-id" }),
+   JSON.stringify({ type: "turn_end", message: { role: "assistant", content: [{ type: "text", text: '{"v":1,"tickets":[{"id":"task-1","title":"First task"}],"batchComplete":true}' }], stopReason: "stop" } }),
+   JSON.stringify({ type: "agent_end" })
+ ];
+ process.stdout.write(lines.join("\\n") + "\\n");
+ `);
+
+     try {
+       process.env.PATH = `${fake.dir}:${originalPath}`;
+
+       const agent = new PiAgent({
+         mode: "json",
+         model: "test-model",
+         env: { PATH: process.env.PATH! },
+       });
+
+       const result = await agent.generate({
+         messages: [{ role: "user", content: "Generate JSON" }],
+       });
+
+       expect(result.text).toContain('"v":1');
+       expect(result.output).toEqual({
+         v: 1,
+         tickets: [{ id: "task-1", title: "First task" }],
+         batchComplete: true,
+       });
+     } finally {
+       await rm(fake.dir, { recursive: true, force: true });
+     }
+   });
+
+   test("PiAgent json mode extracts text from agent_end when turn_end missing", async () => {
+     // Edge case: agent_end has messages array if turn_end is not present
+     const fake = await makeFakePi(`
+ const lines = [
+   JSON.stringify({ type: "session", version: 3, id: "test-session-id" }),
+   JSON.stringify({ type: "agent_end", messages: [
+     { role: "user", content: [{ type: "text", text: "Hello" }] },
+     { role: "assistant", content: [{ type: "text", text: "Response from agent_end" }] }
+   ]})
+ ];
+ process.stdout.write(lines.join("\\n") + "\\n");
+ `);
+
+     try {
+       process.env.PATH = `${fake.dir}:${originalPath}`;
+
+       const agent = new PiAgent({
+         mode: "json",
+         model: "test-model",
+         env: { PATH: process.env.PATH! },
+       });
+
+       const result = await agent.generate({
+         messages: [{ role: "user", content: "Hello" }],
+       });
+
+       expect(result.text).toBe("Response from agent_end");
+     } finally {
+       await rm(fake.dir, { recursive: true, force: true });
+     }
+   });
  });