feat: stream reply blocks immediately

src/agents/pi-embedded-runner.ts

@@ -327,6 +327,10 @@ export async function runEmbeddedPiAgent(params: {
     text?: string;
     mediaUrls?: string[];
   }) => void | Promise<void>;
+  onBlockReply?: (payload: {
+    text?: string;
+    mediaUrls?: string[];
+  }) => void | Promise<void>;
   onToolResult?: (payload: {
     text?: string;
     mediaUrls?: string[];
@@ -489,6 +493,7 @@ export async function runEmbeddedPiAgent(params: {
           verboseLevel: params.verboseLevel,
           shouldEmitToolResult: params.shouldEmitToolResult,
           onToolResult: params.onToolResult,
+          onBlockReply: params.onBlockReply,
           onPartialReply: params.onPartialReply,
           onAgentEvent: params.onAgentEvent,
           enforceFinalTag: params.enforceFinalTag,

src/agents/pi-embedded-subscribe.test.ts

@@ -97,6 +97,37 @@ describe("subscribeEmbeddedPiSession", () => {
     expect(payload.text).toBe("Hello world");
   });
 
+  it("emits block replies on message_end", () => {
+    let handler: ((evt: unknown) => void) | undefined;
+    const session: StubSession = {
+      subscribe: (fn) => {
+        handler = fn;
+        return () => {};
+      },
+    };
+
+    const onBlockReply = vi.fn();
+
+    subscribeEmbeddedPiSession({
+      session: session as unknown as Parameters<
+        typeof subscribeEmbeddedPiSession
+      >[0]["session"],
+      runId: "run",
+      onBlockReply,
+    });
+
+    const assistantMessage = {
+      role: "assistant",
+      content: [{ type: "text", text: "Hello block" }],
+    } as AssistantMessage;
+
+    handler?.({ type: "message_end", message: assistantMessage });
+
+    expect(onBlockReply).toHaveBeenCalled();
+    const payload = onBlockReply.mock.calls[0][0];
+    expect(payload.text).toBe("Hello block");
+  });
+
   it("waits for auto-compaction retry and clears buffered text", async () => {
     const listeners: SessionEventHandler[] = [];
     const session = {

src/agents/pi-embedded-subscribe.ts

@@ -57,6 +57,10 @@ export function subscribeEmbeddedPiSession(params: {
     text?: string;
     mediaUrls?: string[];
   }) => void | Promise<void>;
+  onBlockReply?: (payload: {
+    text?: string;
+    mediaUrls?: string[];
+  }) => void | Promise<void>;
   onPartialReply?: (payload: {
     text?: string;
     mediaUrls?: string[];
@@ -314,6 +318,15 @@ export function subscribeEmbeddedPiSession(params: {
               ? (extractFinalText(cleaned)?.trim() ?? cleaned)
               : cleaned;
           if (text) assistantTexts.push(text);
+          if (text && params.onBlockReply) {
+            const { text: cleanedText, mediaUrls } = splitMediaFromOutput(text);
+            if (cleanedText || (mediaUrls && mediaUrls.length > 0)) {
+              void params.onBlockReply({
+                text: cleanedText,
+                mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
+              });
+            }
+          }
           deltaBuffer = "";
         }
       }