fix: start typing on message start

src/agents/pi-embedded-runner.ts

@@ -1383,6 +1383,7 @@ export async function runEmbeddedPiAgent(params: {
     text?: string;
     mediaUrls?: string[];
   }) => void | Promise<void>;
+  onAssistantMessageStart?: () => void | Promise<void>;
   onBlockReply?: (payload: {
     text?: string;
     mediaUrls?: string[];
@@ -1774,6 +1775,7 @@ export async function runEmbeddedPiAgent(params: {
               blockReplyBreak: params.blockReplyBreak,
               blockReplyChunking: params.blockReplyChunking,
               onPartialReply: params.onPartialReply,
+              onAssistantMessageStart: params.onAssistantMessageStart,
               onAgentEvent: params.onAgentEvent,
               enforceFinalTag: params.enforceFinalTag,
             });

src/agents/pi-embedded-subscribe.ts

@@ -146,7 +146,7 @@ function extractMessagingToolSend(
     : undefined;
 }
 
-export function subscribeEmbeddedPiSession(params: {
+export type SubscribeEmbeddedPiSessionParams = {
   session: AgentSession;
   runId: string;
   verboseLevel?: "off" | "on";
@@ -173,12 +173,17 @@ export function subscribeEmbeddedPiSession(params: {
     text?: string;
     mediaUrls?: string[];
   }) => void | Promise<void>;
+  onAssistantMessageStart?: () => void | Promise<void>;
   onAgentEvent?: (evt: {
     stream: string;
     data: Record<string, unknown>;
   }) => void;
   enforceFinalTag?: boolean;
-}) {
+};
+
+export function subscribeEmbeddedPiSession(
+  params: SubscribeEmbeddedPiSessionParams,
+) {
   const assistantTexts: string[] = [];
   const toolMetas: Array<{ toolName?: string; meta?: string }> = [];
   const toolMetaById = new Map<string, string | undefined>();
@@ -492,6 +497,8 @@ export function subscribeEmbeddedPiSession(params: {
           // may deliver late text_end updates after message_end, which would
           // otherwise re-trigger block replies.
           resetAssistantMessageState(assistantTexts.length);
+          // Use assistant message_start as the earliest "writing" signal for typing.
+          void params.onAssistantMessageStart?.();
         }
       }
 

src/auto-reply/reply/agent-runner.heartbeat-typing.test.ts

@@ -51,6 +51,7 @@ type EmbeddedPiAgentParams = {
     text?: string;
     mediaUrls?: string[];
   }) => Promise<void> | void;
+  onAssistantMessageStart?: () => Promise<void> | void;
   onBlockReply?: (payload: {
     text?: string;
     mediaUrls?: string[];
@@ -212,19 +213,21 @@ describe("runReplyAgent typing (heartbeat)", () => {
     expect(typing.startTypingLoop).not.toHaveBeenCalled();
   });
 
-  it("starts typing only on deltas in message mode", async () => {
-    runEmbeddedPiAgentMock.mockImplementationOnce(async () => ({
-      payloads: [{ text: "final" }],
-      meta: {},
-    }));
+  it("starts typing on assistant message start in message mode", async () => {
+    runEmbeddedPiAgentMock.mockImplementationOnce(
+      async (params: EmbeddedPiAgentParams) => {
+        await params.onAssistantMessageStart?.();
+        return { payloads: [{ text: "final" }], meta: {} };
+      },
+    );
 
     const { run, typing } = createMinimalRun({
       typingMode: "message",
     });
     await run();
 
+    expect(typing.startTypingLoop).toHaveBeenCalled();
     expect(typing.startTypingOnText).not.toHaveBeenCalled();
-    expect(typing.startTypingLoop).not.toHaveBeenCalled();
   });
 
   it("starts typing from reasoning stream in thinking mode", async () => {

src/auto-reply/reply/agent-runner.ts

@@ -689,6 +689,9 @@ export async function runReplyAgent(params: {
                     });
                   }
                 : undefined,
+              onAssistantMessageStart: async () => {
+                await typingSignals.signalMessageStart();
+              },
               onReasoningStream:
                 typingSignals.shouldStartOnReasoning || opts?.onReasoningStream
                   ? async (payload) => {

src/auto-reply/reply/typing-mode.test.ts

@@ -96,6 +96,20 @@ describe("createTypingSignaler", () => {
     expect(typing.startTypingLoop).not.toHaveBeenCalled();
   });
 
+  it("signals on message start for message mode", async () => {
+    const typing = createMockTypingController();
+    const signaler = createTypingSignaler({
+      typing,
+      mode: "message",
+      isHeartbeat: false,
+    });
+
+    await signaler.signalMessageStart();
+
+    expect(typing.startTypingLoop).toHaveBeenCalled();
+    expect(typing.startTypingOnText).not.toHaveBeenCalled();
+  });
+
   it("signals on reasoning for thinking mode", async () => {
     const typing = createMockTypingController();
     const signaler = createTypingSignaler({

src/auto-reply/reply/typing-mode.ts

@@ -25,9 +25,11 @@ export function resolveTypingMode({
 export type TypingSignaler = {
   mode: TypingMode;
   shouldStartImmediately: boolean;
+  shouldStartOnMessageStart: boolean;
   shouldStartOnText: boolean;
   shouldStartOnReasoning: boolean;
   signalRunStart: () => Promise<void>;
+  signalMessageStart: () => Promise<void>;
   signalTextDelta: (text?: string) => Promise<void>;
   signalReasoningDelta: () => Promise<void>;
   signalToolStart: () => Promise<void>;
@@ -40,6 +42,7 @@ export function createTypingSignaler(params: {
 }): TypingSignaler {
   const { typing, mode, isHeartbeat } = params;
   const shouldStartImmediately = mode === "instant";
+  const shouldStartOnMessageStart = mode === "message";
   const shouldStartOnText = mode === "message" || mode === "instant";
   const shouldStartOnReasoning = mode === "thinking";
   const disabled = isHeartbeat || mode === "never";
@@ -49,6 +52,11 @@ export function createTypingSignaler(params: {
     await typing.startTypingLoop();
   };
 
+  const signalMessageStart = async () => {
+    if (disabled || !shouldStartOnMessageStart) return;
+    await typing.startTypingLoop();
+  };
+
   const signalTextDelta = async (text?: string) => {
     if (disabled) return;
     if (shouldStartOnText) {
@@ -80,9 +88,11 @@ export function createTypingSignaler(params: {
   return {
     mode,
     shouldStartImmediately,
+    shouldStartOnMessageStart,
     shouldStartOnText,
     shouldStartOnReasoning,
     signalRunStart,
+    signalMessageStart,
     signalTextDelta,
     signalReasoningDelta,
     signalToolStart,