@vellumai/assistant 0.7.0 → 0.7.2

package/src/__tests__/host-bash-proxy.test.ts

@@ -12,24 +12,64 @@ mock.module("../config/loader.js", () => ({
   getConfig: () => mockConfig,
   loadConfig: () => mockConfig,
   invalidateConfigCache: () => {},
-  saveConfig: () => {},
 }));
 
+const sentMessages: unknown[] = [];
+const sentMessageOptions: unknown[] = [];
+let mockHasClient = false;
+let mockCapableClients: Array<{ clientId: string; capabilities: string[] }> = [];
+let mockClientRegistry: Map<string, { clientId: string; capabilities: string[] }> = new Map();
+
+mock.module("../runtime/assistant-event-hub.js", () => ({
+  broadcastMessage: (msg: unknown, _conversationId?: string, options?: unknown) => {
+    sentMessages.push(msg);
+    sentMessageOptions.push(options);
+  },
+  assistantEventHub: {
+    getMostRecentClientByCapability: (cap: string) =>
+      cap === "host_bash" && mockHasClient ? { id: "mock-client" } : null,
+    listClientsByCapability: (_cap: string) => mockCapableClients,
+    getClientById: (clientId: string) => mockClientRegistry.get(clientId),
+  },
+}));
+
+// Use the REAL pending-interactions module — the proxy self-registers here.
+const pendingInteractions = await import("../runtime/pending-interactions.js");
 const { HostBashProxy } = await import("../daemon/host-bash-proxy.js");
 
 describe("HostBashProxy", () => {
   let proxy: InstanceType<typeof HostBashProxy>;
-  let sentMessages: unknown[];
-  let sendToClient: (msg: unknown) => void;
 
-  function setup(onInternalResolve?: (requestId: string) => void) {
-    sentMessages = [];
-    sendToClient = (msg: unknown) => sentMessages.push(msg);
-    proxy = new HostBashProxy(sendToClient, onInternalResolve);
+  function setup() {
+    sentMessages.length = 0;
+    sentMessageOptions.length = 0;
+    mockHasClient = false;
+    mockCapableClients = [];
+    mockClientRegistry = new Map();
+    pendingInteractions.clear();
+    proxy = new (HostBashProxy as any)();
+  }
+
+  function setupSingleClient(clientId = "client-1") {
+    const entry = { clientId, capabilities: ["host_bash"] };
+    mockCapableClients = [entry];
+    mockClientRegistry.set(clientId, entry);
+  }
+
+  function setupMultipleClients(clientIds: string[]) {
+    mockCapableClients = clientIds.map((id) => ({
+      clientId: id,
+      capabilities: ["host_bash"],
+    }));
+    for (const entry of mockCapableClients) {
+      mockClientRegistry.set(entry.clientId, entry);
+    }
   }
 
   afterEach(() => {
     proxy?.dispose();
+    HostBashProxy.reset();
+    pendingInteractions.clear();
   });
 
   describe("request/resolve lifecycle (happy path)", () => {
@@ -41,7 +81,7 @@ describe("HostBashProxy", () => {
         "session-1",
       );
 
-      // Verify the request was sent to the client
+      // Verify the request was sent via broadcastMessage
       expect(sentMessages).toHaveLength(1);
       const sent = sentMessages[0] as Record<string, unknown>;
       expect(sent.type).toBe("host_bash_request");
@@ -51,10 +91,10 @@ describe("HostBashProxy", () => {
       expect(typeof sent.requestId).toBe("string");
 
       const requestId = sent.requestId as string;
-      expect(proxy.hasPendingRequest(requestId)).toBe(true);
+      expect(pendingInteractions.get(requestId)).toBeDefined();
 
       // Simulate client response
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "hello\n",
         stderr: "",
         exitCode: 0,
@@ -64,7 +104,7 @@ describe("HostBashProxy", () => {
       const result = await resultPromise;
       expect(result.content).toContain("hello");
       expect(result.isError).toBe(false);
-      expect(proxy.hasPendingRequest(requestId)).toBe(false);
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
     });
 
     test("forwards env field in host_bash_request message", async () => {
@@ -84,7 +124,7 @@ describe("HostBashProxy", () => {
       expect(sent.env).toEqual({ VELLUM_UNTRUSTED_SHELL: "1" });
 
       const requestId = sent.requestId as string;
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "locked\n",
         stderr: "",
         exitCode: 0,
@@ -106,7 +146,7 @@ describe("HostBashProxy", () => {
       expect(sent.env).toBeUndefined();
 
       const requestId = sent.requestId as string;
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "normal\n",
         stderr: "",
         exitCode: 0,
@@ -124,7 +164,7 @@ describe("HostBashProxy", () => {
       const sent = sentMessages[0] as Record<string, unknown>;
       const requestId = sent.requestId as string;
 
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "",
         stderr: "command not found",
         exitCode: 127,
@@ -147,7 +187,7 @@ describe("HostBashProxy", () => {
       const sent = sentMessages[0] as Record<string, unknown>;
       const requestId = sent.requestId as string;
 
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "partial",
         stderr: "",
         exitCode: null,
@@ -171,15 +211,12 @@ describe("HostBashProxy", () => {
         "session-1",
       );
 
-      // The proxy timeout is shellMaxTimeoutSec + 30 seconds.
-      // With shellMaxTimeoutSec=0 that's 30 seconds which is too long for a test.
-      // Instead, just verify the pending state and resolve it.
       const sent = sentMessages[0] as Record<string, unknown>;
       const requestId = sent.requestId as string;
-      expect(proxy.hasPendingRequest(requestId)).toBe(true);
+      expect(pendingInteractions.get(requestId)).toBeDefined();
 
       // Resolve to avoid test hanging
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "",
         stderr: "",
         exitCode: 0,
@@ -206,13 +243,13 @@ describe("HostBashProxy", () => {
 
       const sent = sentMessages[0] as Record<string, unknown>;
       const requestId = sent.requestId as string;
-      expect(proxy.hasPendingRequest(requestId)).toBe(true);
+      expect(pendingInteractions.get(requestId)).toBeDefined();
 
       controller.abort();
 
       const result = await resultPromise;
       expect(result.content).toContain("Aborted");
-      expect(proxy.hasPendingRequest(requestId)).toBe(false);
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
     });
 
     test("sends host_bash_cancel to client on abort", async () => {
@@ -256,24 +293,17 @@ describe("HostBashProxy", () => {
   });
 
   describe("isAvailable", () => {
-    test("returns false by default (no client connected)", () => {
+    test("returns false when no client with host_bash capability is connected", () => {
       setup();
+      mockHasClient = false;
       expect(proxy.isAvailable()).toBe(false);
     });
 
-    test("returns true after updateSender with clientConnected=true", () => {
+    test("returns true when a client with host_bash capability is connected", () => {
       setup();
-      proxy.updateSender(sendToClient, true);
+      mockHasClient = true;
       expect(proxy.isAvailable()).toBe(true);
     });
-
-    test("returns false after updateSender with clientConnected=false", () => {
-      setup();
-      proxy.updateSender(sendToClient, true);
-      expect(proxy.isAvailable()).toBe(true);
-      proxy.updateSender(sendToClient, false);
-      expect(proxy.isAvailable()).toBe(false);
-    });
   });
 
   describe("dispose", () => {
@@ -287,12 +317,11 @@ describe("HostBashProxy", () => {
 
       const sent = sentMessages[0] as Record<string, unknown>;
       const requestId = sent.requestId as string;
-      expect(proxy.hasPendingRequest(requestId)).toBe(true);
+      expect(pendingInteractions.get(requestId)).toBeDefined();
 
       proxy.dispose();
 
-      expect(proxy.hasPendingRequest(requestId)).toBe(false);
-      // The promise should reject since dispose rejects pending
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
       expect(resultPromise).rejects.toThrow("Host bash proxy disposed");
     });
 
@@ -324,7 +353,7 @@ describe("HostBashProxy", () => {
   });
 
   describe("late resolve after abort", () => {
-    test("resolve is a no-op after abort (entry already deleted)", async () => {
+    test("resolveResult is a no-op after abort (entry already deleted)", async () => {
       setup();
 
       const controller = new AbortController();
@@ -341,50 +370,23 @@ describe("HostBashProxy", () => {
       const result = await resultPromise;
       expect(result.content).toContain("Aborted");
 
-      // Late resolve should be silently ignored (no throw, no double-resolve)
-      proxy.resolve(requestId, {
+      // Late resolveResult should be silently ignored (no throw, no double-resolve)
+      proxy.resolveResult(requestId, {
         stdout: "late",
         stderr: "",
         exitCode: 0,
         timedOut: false,
       });
 
-      expect(proxy.hasPendingRequest(requestId)).toBe(false);
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
     });
   });
 
-  describe("updateSender", () => {
-    test("uses updated sender for new requests", async () => {
-      setup();
-
-      const newMessages: unknown[] = [];
-      proxy.updateSender((msg) => newMessages.push(msg), true);
-
-      const resultPromise = proxy.request(
-        { command: "echo updated" },
-        "session-1",
-      );
-
-      expect(sentMessages).toHaveLength(0); // Old sender not used
-      expect(newMessages).toHaveLength(1); // New sender used
-
-      const sent = newMessages[0] as Record<string, unknown>;
-      proxy.resolve(sent.requestId as string, {
-        stdout: "updated",
-        stderr: "",
-        exitCode: 0,
-        timedOut: false,
-      });
-
-      await resultPromise;
-    });
-  });
-
-  describe("resolve with unknown requestId", () => {
+  describe("resolveResult with unknown requestId", () => {
     test("silently ignores unknown requestId", () => {
       setup();
       // Should not throw
-      proxy.resolve("unknown-id", {
+      proxy.resolveResult("unknown-id", {
         stdout: "",
         stderr: "",
         exitCode: 0,
@@ -394,9 +396,6 @@ describe("HostBashProxy", () => {
   });
 
   describe("abort listener lifecycle", () => {
-    // Helper that wraps an AbortSignal to observe add/removeEventListener
-    // invocations without tripping over tsc's strict overload matching on
-    // AbortSignal itself.
     type Spied = {
       signal: AbortSignal;
       addCalls: string[];
@@ -408,24 +407,18 @@ describe("HostBashProxy", () => {
       const s = source as any;
       const origAdd = source.addEventListener.bind(source);
       const origRemove = source.removeEventListener.bind(source);
-      s.addEventListener = (
-        type: string,
-        ...rest: any[]
-      ) => {
+      s.addEventListener = (type: string, ...rest: any[]) => {
         addCalls.push(type);
         return (origAdd as any)(type, ...rest);
       };
-      s.removeEventListener = (
-        type: string,
-        ...rest: any[]
-      ) => {
+      s.removeEventListener = (type: string, ...rest: any[]) => {
         removeCalls.push(type);
         return (origRemove as any)(type, ...rest);
       };
       return { signal: source, addCalls, removeCalls };
     }
 
-    test("removes abort listener from signal after resolve completes", async () => {
+    test("removes abort listener from signal after resolveResult completes", async () => {
       setup();
       const controller = new AbortController();
       const spy = spySignal(controller.signal);
@@ -441,7 +434,7 @@ describe("HostBashProxy", () => {
 
       const requestId = (sentMessages[0] as Record<string, unknown>)
         .requestId as string;
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "hello\n",
         stderr: "",
         exitCode: 0,
@@ -449,12 +442,9 @@ describe("HostBashProxy", () => {
       });
       await resultPromise;
 
-      // Listener is detached after normal completion.
       expect(spy.removeCalls).toEqual(["abort"]);
 
-      // Subsequent aborts are harmless no-ops (no side effects on the proxy).
       controller.abort();
-      // No additional emitted envelopes from the late abort.
       expect(sentMessages).toHaveLength(1);
     });
 
@@ -482,10 +472,8 @@ describe("HostBashProxy", () => {
         expect(result.isError).toBe(true);
         expect(result.content).toContain("Host bash proxy timed out");
 
-        // Listener is detached after the timer fires.
         expect(spy.removeCalls).toEqual(["abort"]);
 
-        // Subsequent aborts should be harmless — no cancel emitted.
         controller.abort();
         expect(sentMessages).toHaveLength(1);
       } finally {
@@ -494,114 +482,249 @@ describe("HostBashProxy", () => {
     });
   });
 
-  describe("sender throws synchronously", () => {
-    test("rejects the promise, clears pending state and timer, invokes onInternalResolve", async () => {
-      const resolvedIds: string[] = [];
-      sentMessages = [];
-      sendToClient = () => {
-        throw new Error("transport down");
-      };
-      proxy = new HostBashProxy(sendToClient, (id) => resolvedIds.push(id));
+  describe("pendingInteractions cleanup", () => {
+    test("cleans up on abort", async () => {
+      setup();
 
-      // request() synchronously calls sendToClient inside the Promise
-      // executor. A throw there surfaces as a rejected promise.
+      const controller = new AbortController();
       const resultPromise = proxy.request(
         { command: "echo hello" },
         "session-1",
+        controller.signal,
       );
 
-      await expect(resultPromise).rejects.toThrow("transport down");
+      const sent = sentMessages[0] as Record<string, unknown>;
+      const requestId = sent.requestId as string;
+      expect(pendingInteractions.get(requestId)).toBeDefined();
 
-      // The internal resolve should fire exactly once as part of cleanup.
-      expect(resolvedIds).toHaveLength(1);
+      controller.abort();
+      await resultPromise;
 
-      // Issue a new request on a fresh (non-throwing) sender and verify
-      // the proxy is still functional — no stale timers or bookkeeping
-      // from the failed request.
-      sentMessages = [];
-      proxy.updateSender((msg) => sentMessages.push(msg), true);
-      const okPromise = proxy.request({ command: "echo ok" }, "session-1");
-      expect(sentMessages).toHaveLength(1);
-      const okRequestId = (sentMessages[0] as Record<string, unknown>)
-        .requestId as string;
-      expect(proxy.hasPendingRequest(okRequestId)).toBe(true);
-      proxy.resolve(okRequestId, {
-        stdout: "ok\n",
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
+    });
+
+    test("cleans up for each pending request on dispose", () => {
+      setup();
+
+      const p1 = proxy.request({ command: "echo a" }, "session-1");
+      const p2 = proxy.request({ command: "echo b" }, "session-1");
+      p1.catch(() => {}); // Expected rejection on dispose
+      p2.catch(() => {}); // Expected rejection on dispose
+
+      const ids = (sentMessages as Array<Record<string, unknown>>).map(
+        (m) => m.requestId as string,
+      );
+      expect(ids).toHaveLength(2);
+      expect(pendingInteractions.get(ids[0])).toBeDefined();
+      expect(pendingInteractions.get(ids[1])).toBeDefined();
+
+      proxy.dispose();
+
+      expect(pendingInteractions.get(ids[0])).toBeUndefined();
+      expect(pendingInteractions.get(ids[1])).toBeUndefined();
+    });
+
+    test("cleans up on normal client-initiated resolveResult", async () => {
+      setup();
+
+      const resultPromise = proxy.request(
+        { command: "echo hello" },
+        "session-1",
+      );
+
+      const sent = sentMessages[0] as Record<string, unknown>;
+      const requestId = sent.requestId as string;
+      expect(pendingInteractions.get(requestId)).toBeDefined();
+
+      proxy.resolveResult(requestId, {
+        stdout: "hello",
         stderr: "",
         exitCode: 0,
         timedOut: false,
       });
-      const okResult = await okPromise;
-      expect(okResult.content).toContain("ok");
-      expect(okResult.isError).toBe(false);
+
+      await resultPromise;
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
     });
   });
 
-  describe("onInternalResolve callback", () => {
-    test("fires on abort", async () => {
-      const resolvedIds: string[] = [];
-      setup((id) => resolvedIds.push(id));
+  describe("target client routing", () => {
+    test("auto-resolves when exactly one capable client is connected", async () => {
+      setup();
+      setupSingleClient("client-abc");
 
-      const controller = new AbortController();
       const resultPromise = proxy.request(
         { command: "echo hello" },
         "session-1",
-        controller.signal,
       );
 
+      expect(sentMessages).toHaveLength(1);
       const sent = sentMessages[0] as Record<string, unknown>;
-      const requestId = sent.requestId as string;
+      expect(sent.targetClientId).toBe("client-abc");
 
-      controller.abort();
-      await resultPromise;
+      // Options passed to broadcastMessage should also have targetClientId
+      const opts = sentMessageOptions[0] as Record<string, unknown> | undefined;
+      expect(opts?.targetClientId).toBe("client-abc");
 
-      expect(resolvedIds).toEqual([requestId]);
+      const requestId = sent.requestId as string;
+      proxy.resolveResult(requestId, {
+        stdout: "hello\n",
+        stderr: "",
+        exitCode: 0,
+        timedOut: false,
+      });
+
+      const result = await resultPromise;
+      expect(result.isError).toBe(false);
     });
 
-    test("fires for each pending request on dispose", () => {
-      const resolvedIds: string[] = [];
-      setup((id) => resolvedIds.push(id));
+    test("uses explicit targetClientId when it is valid", async () => {
+      setup();
+      setupSingleClient("client-abc");
+      // Also register a second client so we're sure explicit targeting works
+      const entry2 = { clientId: "client-xyz", capabilities: ["host_bash"] };
+      mockCapableClients.push(entry2);
+      mockClientRegistry.set("client-xyz", entry2);
 
-      // Create two pending requests and catch rejections from dispose
-      const p1 = proxy.request({ command: "echo a" }, "session-1");
-      const p2 = proxy.request({ command: "echo b" }, "session-1");
-      p1.catch(() => {}); // Expected rejection on dispose
-      p2.catch(() => {}); // Expected rejection on dispose
+      const resultPromise = proxy.request(
+        { command: "echo hello", targetClientId: "client-abc" },
+        "session-1",
+      );
 
-      const ids = (sentMessages as Array<Record<string, unknown>>).map(
-        (m) => m.requestId as string,
+      expect(sentMessages).toHaveLength(1);
+      const sent = sentMessages[0] as Record<string, unknown>;
+      expect(sent.targetClientId).toBe("client-abc");
+
+      const opts = sentMessageOptions[0] as Record<string, unknown> | undefined;
+      expect(opts?.targetClientId).toBe("client-abc");
+
+      const requestId = sent.requestId as string;
+      proxy.resolveResult(requestId, {
+        stdout: "ok\n",
+        stderr: "",
+        exitCode: 0,
+        timedOut: false,
+      });
+
+      const result = await resultPromise;
+      expect(result.isError).toBe(false);
+    });
+
+    test("returns error for explicit targetClientId that is not connected", async () => {
+      setup();
+      setupSingleClient("client-abc");
+
+      const result = await proxy.request(
+        { command: "echo hello", targetClientId: "client-unknown" },
+        "session-1",
       );
-      expect(ids).toHaveLength(2);
 
-      proxy.dispose();
+      // Should return error without broadcasting
+      expect(result.isError).toBe(true);
+      expect(result.content).toContain("client-unknown");
+      expect(result.content).toContain("assistant clients list --capability host_bash");
+      expect(sentMessages).toHaveLength(0);
+    });
+
+    test("returns error for explicit targetClientId that is connected but lacks host_bash", async () => {
+      setup();
+      // Register a client without host_bash capability
+      mockClientRegistry.set("client-no-bash", {
+        clientId: "client-no-bash",
+        capabilities: [],
+      });
 
-      expect(resolvedIds).toHaveLength(2);
-      expect(resolvedIds).toContain(ids[0]);
-      expect(resolvedIds).toContain(ids[1]);
+      const result = await proxy.request(
+        { command: "echo hello", targetClientId: "client-no-bash" },
+        "session-1",
+      );
+
+      expect(result.isError).toBe(true);
+      expect(result.content).toContain("client-no-bash");
+      expect(result.content).toContain("does not support host_bash");
+      expect(sentMessages).toHaveLength(0);
     });
 
-    test("does not fire on normal client-initiated resolve", async () => {
-      const resolvedIds: string[] = [];
-      setup((id) => resolvedIds.push(id));
+    test("falls through to untargeted broadcast when multiple capable clients are connected and no targetClientId", async () => {
+      setup();
+      setupMultipleClients(["client-1", "client-2", "client-3"]);
 
       const resultPromise = proxy.request(
         { command: "echo hello" },
         "session-1",
       );
 
+      // Should broadcast without an early error return
+      expect(sentMessages).toHaveLength(1);
       const sent = sentMessages[0] as Record<string, unknown>;
+      expect(sent.type).toBe("host_bash_request");
+      // No target client resolved — untargeted broadcast
+      expect(sent.targetClientId).toBeUndefined();
+
+      const opts = sentMessageOptions[0] as Record<string, unknown> | undefined;
+      expect(opts?.targetClientId).toBeUndefined();
+
+      // Manually resolve to clean up
       const requestId = sent.requestId as string;
+      proxy.resolveResult(requestId, {
+        stdout: "hello\n",
+        stderr: "",
+        exitCode: 0,
+        timedOut: false,
+      });
 
-      // Normal resolve from client — should NOT trigger onInternalResolve
-      proxy.resolve(requestId, {
-        stdout: "hello",
+      const result = await resultPromise;
+      expect(result.isError).toBe(false);
+    });
+
+    test("falls through to broadcast when zero capable clients (existing timeout path)", async () => {
+      setup();
+      // mockCapableClients is empty (default), so capableClients.length === 0
+
+      const resultPromise = proxy.request(
+        { command: "echo hello" },
+        "session-1",
+      );
+
+      // Should still broadcast (no early return)
+      expect(sentMessages).toHaveLength(1);
+      const sent = sentMessages[0] as Record<string, unknown>;
+      expect(sent.type).toBe("host_bash_request");
+      // targetClientId is undefined when no clients present
+      expect(sent.targetClientId).toBeUndefined();
+
+      // Manually resolve to clean up
+      const requestId = sent.requestId as string;
+      proxy.resolveResult(requestId, {
+        stdout: "",
         stderr: "",
         exitCode: 0,
         timedOut: false,
       });
 
       await resultPromise;
-      expect(resolvedIds).toEqual([]);
+    });
+
+    test("includes targetClientId in timeout error message when client was resolved", async () => {
+      setup();
+      setupSingleClient("client-mac");
+
+      jest.useFakeTimers();
+      try {
+        const resultPromise = proxy.request(
+          { command: "echo slow", timeout_seconds: 30 },
+          "session-1",
+        );
+
+        // Proxy timeout = 33s; advance past it
+        jest.advanceTimersByTime(34 * 1000);
+
+        const result = await resultPromise;
+        expect(result.isError).toBe(true);
+        expect(result.content).toContain("client-mac");
+      } finally {
+        jest.useRealTimers();
+      }
     });
   });
 });