pi-zulip-extension/src/index.ts

@@ -213,10 +213,9 @@ async function createZulipQueue(config: ZulipConfig) {
       to: number | string;
       subject?: string;
       content: string;
-    }): Promise<void> {
+    }): Promise<number> {
       const formBody = new URLSearchParams();
       formBody.append("type", params.type);
-      // Zulip API: private messages need JSON array of user IDs, stream messages accept string/number
       formBody.append("to", params.type === "private" ? JSON.stringify([params.to]) : String(params.to));
       if (params.subject) formBody.append("subject", params.subject);
       formBody.append("content", params.content);
@@ -235,6 +234,32 @@ async function createZulipQueue(config: ZulipConfig) {
       if (!res.ok) {
         throw new Error(`Send message API returned ${res.status}: ${await res.text()}`);
       }
+
+      const body = (await res.json()) as { id: number };
+      return body.id;
+    },
+
+    /**
+     * Edit a previously sent Zulip message (streaming update).
+     */
+    async editMessage(messageId: number, content: string): Promise<void> {
+      const formBody = new URLSearchParams();
+      formBody.append("content", content);
+
+      const res = await fetch(`${config.zulip.site}/api/v1/messages/${messageId}`, {
+        method: "PATCH",
+        headers: {
+          Authorization:
+            "Basic " +
+            Buffer.from(`${config.zulip.email}:${config.zulip.api_key}`).toString("base64"),
+          "Content-Type": "application/x-www-form-urlencoded",
+        },
+        body: formBody.toString(),
+      });
+
+      if (!res.ok && res.status !== 400) {
+        console.warn(`[zulip-ext] Edit message ${messageId} returned ${res.status}`);
+      }
     },
 
     /**
@@ -331,6 +356,7 @@ export default function (pi: ExtensionAPI) {
   }
 
   // Pending Zulip replies awaiting the LLM response from agent_end.
+  // zulipMessageId is set after the placeholder is sent.
   const pendingZulipReplies: Array<{
     id: number;
     type: "private" | "stream";
@@ -339,12 +365,18 @@ export default function (pi: ExtensionAPI) {
     streamName?: string;
     topic?: string;
     streamId?: number;
+    zulipMessageId?: number;
   }> = [];
   let zulipReplyIdCounter = 0;
+  let isAgentBusy = false;
+  let streamingTimer: ReturnType<typeof setTimeout> | null = null;
+  let streamingAccumulator = "";
+  let streamingReplyId: number | null = null;
 
   /**
    * Queue a Zulip message into the current pi session via sendUserMessage
    * and schedule the reply to be sent back on the next agent_end event.
+   * Sends a "Thinking..." placeholder to Zulip immediately for streaming UX.
    */
   function injectIntoSession(
     msg: ZulipMessage,
@@ -361,15 +393,41 @@ export default function (pi: ExtensionAPI) {
   ): void {
     // Send typing indicator (fire-and-forget)
     queue!.sendTypingNotification([senderId], "start").catch(() => {});
-    // Track the pending reply so agent_end knows where to send the response
+
     const replyId = ++zulipReplyIdCounter;
-    pendingZulipReplies.push({ id: replyId, ...replyInfo });
-    // Inject directly into this pi session — no subprocess overhead
     const label = replyInfo.type === "stream"
       ? `@all-bots in #${replyInfo.streamName!}`
       : `DM from @${senderName}`;
-    pi.sendUserMessage(`[Zulip ${label}]: ${msg.content}`);
-    console.log(`[zulip-ext] Injected into session [${replyId}]: from ${senderName}`);
+
+    // Send a placeholder "Thinking..." message to Zulip immediately
+    const placeholders = [
+      ":robot: _Processing your message..._",
+      ":hourglass_flowing_sand: _Thinking..._",
+      ":brain: _Generating response..._",
+    ];
+    const placeholder = placeholders[replyId % placeholders.length];
+    const placeholderTarget = replyInfo.type === "private"
+      ? { type: "private" as const, to: senderId }
+      : { type: "stream" as const, to: replyInfo.streamId ?? replyInfo.streamName!, subject: replyInfo.topic! };
+
+    queue!.sendMessage({ ...placeholderTarget, content: placeholder })
+      .then((msgId) => {
+        pendingZulipReplies.push({ id: replyId, ...replyInfo, zulipMessageId: msgId });
+        console.log(`[zulip-ext] Placed [${replyId}] (msg ${msgId}) for ${senderName}`);
+      })
+      .catch(() => {
+        pendingZulipReplies.push({ id: replyId, ...replyInfo });
+        console.log(`[zulip-ext] Queued [${replyId}] for ${senderName} (no placeholder)`);
+      });
+
+    // Inject into the pi session — use steer mode if agent is busy
+    if (isAgentBusy) {
+      pi.sendUserMessage(`[Zulip ${label}]: ${msg.content}`, { deliverAs: "steer" });
+      console.log(`[zulip-ext] Steer [${replyId}]: queued while agent busy`);
+    } else {
+      pi.sendUserMessage(`[Zulip ${label}]: ${msg.content}`);
+      console.log(`[zulip-ext] Injected [${replyId}]: from ${senderName}`);
+    }
   }
 
   /**
@@ -481,25 +539,86 @@ export default function (pi: ExtensionAPI) {
   pi.on("session_start", async (_event) => {
     console.log(`[zulip-ext] Starting Zulip extension for ${config.agent.name} (${config.zulip.email})`);
 
-    // Start health endpoint
     if (!healthServer) {
       healthServer = startHealthServer(config.health_port ?? 9200, getState);
     }
 
-    // Begin Zulip event polling
     startPolling();
   });
 
+  // -----------------------------------------------------------------------
+  // Streaming + steering lifecycle
+  // -----------------------------------------------------------------------
+
   /**
-   * agent_end captures the LLM response from this session and relays it
-   * back to the Zulip user/stream that sent the original message.
+   * turn_start fires when a new LLM turn begins. Mark the agent as busy
+   * so subsequent Zulip DMs are steered instead of starting new turns.
+   */
+  pi.on("turn_start", async () => {
+    isAgentBusy = true;
+    streamingAccumulator = "";
+    streamingReplyId = pendingZulipReplies.length > 0 ? pendingZulipReplies[0].id : null;
+  });
+
+  /**
+   * message_update fires on each streaming token. Accumulate partial text
+   * and periodically edit the Zulip placeholder message so the user sees
+   * the response come in live.
+   */
+  pi.on("message_update", async (event) => {
+    const msg = event.message as any;
+    const content = msg.content;
+    let partialText = "";
+
+    if (typeof content === "string") {
+      partialText = content;
+    } else if (Array.isArray(content)) {
+      partialText = (content as Array<{ type: string; text: string }>)
+        .filter((c) => c.type === "text")
+        .map((c) => c.text)
+        .join("\n");
+    }
+
+    streamingAccumulator = partialText;
+
+    if (!streamingTimer && streamingAccumulator.length > 10) {
+      streamingTimer = setTimeout(async () => {
+        streamingTimer = null;
+        if (!streamingAccumulator || streamingReplyId === null) return;
+
+        const pending = pendingZulipReplies.find((r) => r.id === streamingReplyId);
+        if (!pending || !pending.zulipMessageId) return;
+
+        try {
+          const preview = streamingAccumulator.length > 1800
+            ? streamingAccumulator.slice(0, 1800) + "\n\n_… still generating…_"
+            : streamingAccumulator;
+          await queue!.editMessage(pending.zulipMessageId, preview);
+        } catch {
+          // Non-critical during streaming
+        }
+      }, 2000);
+    }
+  });
+
+  /**
+   * agent_end captures the final LLM response and relays it back to Zulip.
+   * If we sent a placeholder message earlier, edit it with the final content
+   * (instead of sending a new message) for a seamless streaming experience.
    */
   pi.on("agent_end", async (event) => {
+    isAgentBusy = false;
+
+    if (streamingTimer) {
+      clearTimeout(streamingTimer);
+      streamingTimer = null;
+    }
+
     if (pendingZulipReplies.length === 0) return;
 
     const reply = pendingZulipReplies.shift()!;
 
-    // Extract assistant response text from the completed turn
+    // Extract final assistant response text
     const allMsgs = event.messages as unknown as Array<Record<string, unknown>>;
     const assistantMsgs = allMsgs.filter((m) => m.role === "assistant");
     const lastAssistant = assistantMsgs[assistantMsgs.length - 1];
@@ -522,31 +641,43 @@ export default function (pi: ExtensionAPI) {
       queue!.sendTypingNotification([reply.senderId], "stop").catch(() => {});
     }
 
-    if (responseText.trim()) {
-      const truncated = responseText.length > 1500
-        ? responseText.slice(0, 1500) + "\n\n[...truncated, see session for full output]"
-        : responseText;
+    if (!responseText.trim()) {
+      console.log(`[zulip-ext] No assistant response for [${reply.id}] from ${reply.senderName}`);
+      return;
+    }
 
-      const target = reply.type === "private"
-        ? { type: "private" as const, to: reply.senderId! }
-        : { type: "stream" as const, to: reply.streamId ?? reply.streamName!, subject: reply.topic! };
+    const truncated = responseText.length > 1500
+      ? responseText.slice(0, 1500) + "\n\n[...truncated, see session for full output]"
+      : responseText;
 
-      try {
+    try {
+      if (reply.zulipMessageId) {
+        // Edit the placeholder with the final response (streaming UX)
+        await queue!.editMessage(reply.zulipMessageId, truncated);
+        console.log(`[zulip-ext] Finalized [${reply.id}] for ${reply.senderName} (${truncated.length} chars)`);
+      } else {
+        // No placeholder — send as new message
+        const target = reply.type === "private"
+          ? { type: "private" as const, to: reply.senderId! }
+          : { type: "stream" as const, to: reply.streamId ?? reply.streamName!, subject: reply.topic! };
         await queue!.sendMessage({ ...target, content: truncated });
         console.log(`[zulip-ext] Replied to [${reply.id}] ${reply.senderName} (${truncated.length} chars)`);
-      } catch (err) {
-        const errMsg = err instanceof Error ? err.message : String(err);
-        console.error(`[zulip-ext] Failed to send reply [${reply.id}]: ${errMsg}`);
-        lastError = errMsg;
       }
-    } else {
-      console.log(`[zulip-ext] No assistant response for [${reply.id}] from ${reply.senderName}`);
+    } catch (err) {
+      const errMsg = err instanceof Error ? err.message : String(err);
+      console.error(`[zulip-ext] Failed to finalize reply [${reply.id}]: ${errMsg}`);
+      lastError = errMsg;
     }
   });
 
   pi.on("session_shutdown", async () => {
     console.log(`[zulip-ext] Shutting down...`);
 
+    if (streamingTimer) {
+      clearTimeout(streamingTimer);
+      streamingTimer = null;
+    }
+
     if (pollTimer) {
       clearInterval(pollTimer);
       pollTimer = null;