@agent-marshal/marshal-prime-agent 0.54.0 → 0.55.0

package/package.json

@@ -50,5 +50,5 @@
     "test": "bun test"
   },
   "type": "module",
-  "version": "0.54.0"
+  "version": "0.55.0"
 }

package/src/daemon.ts

@@ -41,9 +41,11 @@ import {
   type SessionNicknameItem,
 } from "./entities.ts";
 import type { Identity } from "./identity.ts";
+import { InboxPuller, type InboxPull } from "./inbox-pull.ts";
 
 const LIVENESS_INTERVAL_MS = 5_000;
 const INBOX_PULL_LIMIT = 20;
+const INBOX_COMMAND_TIMEOUT_MS = 1_500;
 const SOURCE_ID = "marshal-prime-agent";
 
 /** myko's CellServer serves its WebSocket at the `/myko` path. */
@@ -78,7 +80,11 @@ export class MarshalDaemon {
   private roster: SessionItem[] = [];
   private nicknames: SessionNicknameItem[] = [];
 
-  private readonly draining = new Map<string, Promise<string | null>>();
+  private readonly inbox = new InboxPuller<MessageView>({
+    read: (sessionId) => this.readUnread(sessionId),
+    ack: async (sessionId, messageIds) => { await this.sendForInbox(ackMessages(sessionId, messageIds)); },
+    render: (messages) => this.renderInbox(messages),
+  });
   private notifyHandler: NotifyHandler | null = null;
   private livenessTimer: ReturnType<typeof setInterval> | null = null;
   private started = false;
@@ -250,12 +256,14 @@ export class MarshalDaemon {
     return `- [${m.messageId}] ${this.senderLabel(m.fromSessionId)} [response_required=${m.responseRequired ?? false}]: ${m.body}`;
   }
 
-  drainInbox(sessionId: string): Promise<string | null> {
-    const prev = this.draining.get(sessionId) ?? Promise.resolve<string | null>(null);
-    const next = prev.catch(() => null).then(() => this.drainInboxInner(sessionId));
-    this.draining.set(sessionId, next);
-    void next.finally(() => { if (this.draining.get(sessionId) === next) this.draining.delete(sessionId); });
-    return next;
+  /**
+   * Read the unread inbox without acking it. Commit the pull once its text is
+   * in the conversation, or release it if it never gets there, so a message is
+   * never marked read without having been shown.
+   */
+  pullInbox(sessionId: string): Promise<InboxPull | null> {
+    if (!this.connected) return Promise.resolve(null);
+    return this.inbox.pull(sessionId).catch(() => null);
   }
 
   nicknameFor(sessionId: string): string {
@@ -270,11 +278,11 @@ export class MarshalDaemon {
 
   // ── Internals ──────────────────────────────────────────────────────────
 
-  private async drainInboxInner(sessionId: string): Promise<string | null> {
+  private async readUnread(sessionId: string): Promise<MessageView[] | null> {
     let result: ReadMessagesResult | undefined;
     for (let attempt = 0; attempt < 3; attempt++) {
       try {
-        result = await this.send(readMessages({
+        result = await this.sendForInbox(readMessages({
           asSession: sessionId,
           toSession: sessionId,
           inbox: false,
@@ -288,13 +296,7 @@ export class MarshalDaemon {
       if (result.messages && result.messages.length > 0) break;
       if (attempt < 2) await new Promise((r) => setTimeout(r, 100 * (attempt + 1)));
     }
-    if (!result?.messages?.length) return null;
-
-    const block = this.renderInbox(result.messages);
-    try {
-      await this.send(ackMessages(sessionId, result.messages.map((m) => m.messageId)));
-    } catch { this.log("inbox ack failed"); }
-    return block;
+    return result?.messages?.length ? result.messages : null;
   }
 
   private renderInbox(messages: MessageView[]): string {
@@ -352,6 +354,38 @@ export class MarshalDaemon {
     return this.client.sendCommand(command as never) as Promise<R>;
   }
 
+  /** Bounded send for the inbox path: reject at once while the WS is down, and
+   *  otherwise time out, so a read that never answers can't hold up every later
+   *  pull for the session. */
+  private sendForInbox<R>(command: MarshalCommand<R>): Promise<R> {
+    if (!this.connected) return Promise.reject(new Error("marshal is disconnected"));
+
+    return new Promise<R>((resolve, reject) => {
+      let settled = false;
+      const timer = setTimeout(() => {
+        if (settled) return;
+        settled = true;
+        reject(new Error(`marshal inbox command timed out after ${INBOX_COMMAND_TIMEOUT_MS}ms`));
+      }, INBOX_COMMAND_TIMEOUT_MS);
+      if (typeof (timer as any).unref === "function") (timer as any).unref();
+
+      void this.send(command).then(
+        (value) => {
+          if (settled) return;
+          settled = true;
+          clearTimeout(timer);
+          resolve(value);
+        },
+        (error: unknown) => {
+          if (settled) return;
+          settled = true;
+          clearTimeout(timer);
+          reject(error);
+        },
+      );
+    });
+  }
+
   private watch<I>(query: MarshalQuery<I>) {
     return this.client.watchQuery(query as never) as unknown as import("rxjs").Observable<I[]>;
   }

package/src/inbox-pull.ts

@@ -0,0 +1,86 @@
+// Inbox pulls that are acked only once their text has been shown.
+//
+// Acking in the same step as the read loses messages whenever the caller ends
+// up not showing what it read: a per-turn pull that loses its timeout race, or
+// an injection that fails. A pull here hands the messages back unacked; the
+// caller commits it once the text is in the conversation, or releases it so
+// the next pull hands the same messages out again.
+
+/** Unread messages read but not yet acked. */
+export interface InboxPull {
+  /** The rendered messages, ready to inject. */
+  readonly text: string;
+  /** Ack the messages. Call once `text` has been shown. */
+  commit(): Promise<void>;
+  /** Give the messages back unacked, so the next pull hands them out again. */
+  release(): void;
+}
+
+/** Where a pull reads, acks and renders messages. */
+export interface InboxSource<M extends { messageId: string }> {
+  read(sessionId: string): Promise<M[] | null>;
+  ack(sessionId: string, messageIds: string[]): Promise<void>;
+  render(messages: M[]): string;
+}
+
+export class InboxPuller<M extends { messageId: string }> {
+  /** Per session, the ids an outstanding pull holds: still unread, so a read returns them. */
+  private readonly held = new Map<string, Set<string>>();
+  /** Per session, the pull in flight. Pulls run one at a time, so each sees what the last one holds. */
+  private readonly pulling = new Map<string, Promise<unknown>>();
+
+  private readonly source: InboxSource<M>;
+
+  constructor(source: InboxSource<M>) {
+    this.source = source;
+  }
+
+  pull(sessionId: string): Promise<InboxPull | null> {
+    const previous = this.pulling.get(sessionId) ?? Promise.resolve();
+    const next = previous.catch(() => {}).then(() => this.pullNow(sessionId));
+    this.pulling.set(sessionId, next);
+    void next.catch(() => {}).finally(() => {
+      if (this.pulling.get(sessionId) === next) this.pulling.delete(sessionId);
+    });
+    return next;
+  }
+
+  private async pullNow(sessionId: string): Promise<InboxPull | null> {
+    const held = this.held.get(sessionId) ?? new Set<string>();
+    const messages = (await this.source.read(sessionId))?.filter((m) => !held.has(m.messageId));
+    if (!messages?.length) return null;
+    // Render before holding anything, so a render that throws holds nothing.
+    const text = this.source.render(messages);
+    const ids = messages.map((m) => m.messageId);
+    for (const id of ids) held.add(id);
+    this.held.set(sessionId, held);
+    const settle = () => { for (const id of ids) held.delete(id); };
+    return {
+      text,
+      commit: async () => {
+        try {
+          await this.source.ack(sessionId, ids);
+        } finally {
+          settle();
+        }
+      },
+      release: settle,
+    };
+  }
+}
+
+/**
+ * Pull for a turn, giving up after `timeoutMs` so the turn is never held up.
+ * A pull that lands after that is released, not acked: its messages were not
+ * shown, so they stay unread for the next turn.
+ */
+export async function pullForTurn(pull: () => Promise<InboxPull | null>, timeoutMs: number): Promise<InboxPull | null> {
+  const pending = pull().catch(() => null);
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const deadline = new Promise<"late">((resolve) => { timer = setTimeout(() => resolve("late"), timeoutMs); });
+  const winner = await Promise.race([pending, deadline]);
+  clearTimeout(timer);
+  if (winner !== "late") return winner;
+  void pending.then((landed) => landed?.release());
+  return null;
+}

package/src/index.ts

@@ -23,6 +23,7 @@ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { Type } from "typebox";
 
 import { MarshalDaemon } from "./daemon.ts";
+import { pullForTurn } from "./inbox-pull.ts";
 import { resolveIdentity, type Identity } from "./identity.ts";
 import type { NotifyChannelMeta, SessionItem } from "./entities.ts";
 import { shouldRegisterPrimeAgentSession } from "./runtime.ts";
@@ -111,12 +112,19 @@ async function init(pi: ExtensionAPI) {
           );
         } else {
           // Compatibility fallback for older daemons that omit body metadata.
-          const inbox = await daemon!.drainInbox(sid).catch(() => null);
-          if (!inbox) return;
-          await pi.sendMessage(
-            { customType: "marshal-channel", content: inbox, display: true, details: meta },
-            { deliverAs: "steer", triggerTurn: true },
-          );
+          // The pull is acked only once it has been injected.
+          const pull = await daemon!.pullInbox(sid);
+          if (!pull) return;
+          try {
+            await pi.sendMessage(
+              { customType: "marshal-channel", content: pull.text, display: true, details: meta },
+              { deliverAs: "steer", triggerTurn: true },
+            );
+          } catch (e) {
+            pull.release();
+            throw e;
+          }
+          await pull.commit().catch(() => log("inbox ack failed"));
         }
       } catch (e) {
         log(`live push injection failed: ${String(e)}`);
@@ -328,19 +336,21 @@ async function init(pi: ExtensionAPI) {
     // Add marshal system prompt context (once per turn, before the LLM call).
     result.systemPrompt = event.systemPrompt + "\n\n" + SYSTEM_PROMPT_BLOCK;
 
-    // Drain any missed inbox messages — only if the daemon is connected.
-    if (daemon && sessionId && daemon.isConnected()) {
-      daemon.registerSession(sessionId); // self-heal a dropped roster row
-      const inbox = await Promise.race([
-        daemon.drainInbox(sessionId).catch(() => null),
-        new Promise<null>((r) => setTimeout(() => r(null), 2000)),
-      ]);
-      if (inbox) {
+    // Pull any missed inbox messages — only if the daemon is connected. The
+    // pull is acked only once its text goes into this turn; one that misses
+    // the deadline is released, so its messages stay unread for the next turn.
+    const d = daemon;
+    const sid = sessionId;
+    if (d && sid && d.isConnected()) {
+      d.registerSession(sid); // self-heal a dropped roster row
+      const pull = await pullForTurn(() => d.pullInbox(sid), 2000);
+      if (pull) {
         result.message = {
           customType: "marshal-inbox",
-          content: inbox,
+          content: pull.text,
           display: true,
         };
+        void pull.commit().catch(() => log("inbox ack failed"));
       }
     }