@vellumai/assistant 0.7.1 → 0.7.2

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

@@ -12,32 +12,29 @@ mock.module("../config/loader.js", () => ({
   getConfig: () => mockConfig,
   loadConfig: () => mockConfig,
   invalidateConfigCache: () => {},
-  saveConfig: () => {},
 }));
 
 const sentMessages: unknown[] = [];
-const resolvedInteractionIds: string[] = [];
+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) => sentMessages.push(msg),
+  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),
   },
 }));
 
-mock.module("../runtime/pending-interactions.js", () => ({
-  resolve: (requestId: string) => {
-    resolvedInteractionIds.push(requestId);
-    return undefined;
-  },
-  get: () => undefined,
-  getByKind: () => [],
-  getByConversation: () => [],
-  removeByConversation: () => {},
-}));
-
+// 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", () => {
@@ -45,14 +42,34 @@ describe("HostBashProxy", () => {
 
   function setup() {
     sentMessages.length = 0;
-    resolvedInteractionIds.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)", () => {
@@ -74,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,
@@ -87,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 () => {
@@ -107,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,
@@ -129,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,
@@ -147,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,
@@ -170,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,
@@ -196,10 +213,10 @@ 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();
 
       // Resolve to avoid test hanging
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "",
         stderr: "",
         exitCode: 0,
@@ -226,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 () => {
@@ -300,11 +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);
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
       expect(resultPromise).rejects.toThrow("Host bash proxy disposed");
     });
 
@@ -336,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();
@@ -353,23 +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("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,
@@ -401,7 +418,7 @@ describe("HostBashProxy", () => {
       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);
@@ -417,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,
@@ -465,8 +482,8 @@ describe("HostBashProxy", () => {
     });
   });
 
-  describe("pendingInteractions.resolve callback", () => {
-    test("fires on abort", async () => {
+  describe("pendingInteractions cleanup", () => {
+    test("cleans up on abort", async () => {
       setup();
 
       const controller = new AbortController();
@@ -478,14 +495,15 @@ describe("HostBashProxy", () => {
 
       const sent = sentMessages[0] as Record<string, unknown>;
       const requestId = sent.requestId as string;
+      expect(pendingInteractions.get(requestId)).toBeDefined();
 
       controller.abort();
       await resultPromise;
 
-      expect(resolvedInteractionIds).toEqual([requestId]);
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
     });
 
-    test("fires for each pending request on dispose", () => {
+    test("cleans up for each pending request on dispose", () => {
       setup();
 
       const p1 = proxy.request({ command: "echo a" }, "session-1");
@@ -497,15 +515,16 @@ describe("HostBashProxy", () => {
         (m) => m.requestId as string,
       );
       expect(ids).toHaveLength(2);
+      expect(pendingInteractions.get(ids[0])).toBeDefined();
+      expect(pendingInteractions.get(ids[1])).toBeDefined();
 
       proxy.dispose();
 
-      expect(resolvedInteractionIds).toHaveLength(2);
-      expect(resolvedInteractionIds).toContain(ids[0]);
-      expect(resolvedInteractionIds).toContain(ids[1]);
+      expect(pendingInteractions.get(ids[0])).toBeUndefined();
+      expect(pendingInteractions.get(ids[1])).toBeUndefined();
     });
 
-    test("does not fire on normal client-initiated resolve", async () => {
+    test("cleans up on normal client-initiated resolveResult", async () => {
       setup();
 
       const resultPromise = proxy.request(
@@ -515,9 +534,9 @@ describe("HostBashProxy", () => {
 
       const sent = sentMessages[0] as Record<string, unknown>;
       const requestId = sent.requestId as string;
+      expect(pendingInteractions.get(requestId)).toBeDefined();
 
-      // Normal resolve from client — should NOT trigger pendingInteractions.resolve
-      proxy.resolve(requestId, {
+      proxy.resolveResult(requestId, {
         stdout: "hello",
         stderr: "",
         exitCode: 0,
@@ -525,7 +544,187 @@ describe("HostBashProxy", () => {
       });
 
       await resultPromise;
-      expect(resolvedInteractionIds).toEqual([]);
+      expect(pendingInteractions.get(requestId)).toBeUndefined();
+    });
+  });
+
+  describe("target client routing", () => {
+    test("auto-resolves when exactly one capable client is connected", async () => {
+      setup();
+      setupSingleClient("client-abc");
+
+      const resultPromise = proxy.request(
+        { command: "echo hello" },
+        "session-1",
+      );
+
+      expect(sentMessages).toHaveLength(1);
+      const sent = sentMessages[0] as Record<string, unknown>;
+      expect(sent.targetClientId).toBe("client-abc");
+
+      // Options passed to broadcastMessage should also have targetClientId
+      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: "hello\n",
+        stderr: "",
+        exitCode: 0,
+        timedOut: false,
+      });
+
+      const result = await resultPromise;
+      expect(result.isError).toBe(false);
+    });
+
+    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);
+
+      const resultPromise = proxy.request(
+        { command: "echo hello", targetClientId: "client-abc" },
+        "session-1",
+      );
+
+      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",
+      );
+
+      // 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: [],
+      });
+
+      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("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,
+      });
+
+      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;
+    });
+
+    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();
+      }
     });
   });
 });