diff --git a/assistant/src/__tests__/host-bash-proxy.test.ts b/assistant/src/__tests__/host-bash-proxy.test.ts index fe9c4dc5246..834d001a0e5 100644 --- a/assistant/src/__tests__/host-bash-proxy.test.ts +++ b/assistant/src/__tests__/host-bash-proxy.test.ts @@ -16,14 +16,22 @@ mock.module("../config/loader.js", () => ({ })); const sentMessages: unknown[] = []; +const sentMessageOptions: unknown[] = []; const resolvedInteractionIds: string[] = []; let mockHasClient = false; +let mockCapableClients: Array<{ clientId: string; capabilities: string[] }> = []; +let mockClientRegistry: Map = 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), }, })); @@ -45,11 +53,30 @@ describe("HostBashProxy", () => { function setup() { sentMessages.length = 0; + sentMessageOptions.length = 0; resolvedInteractionIds.length = 0; mockHasClient = false; + mockCapableClients = []; + mockClientRegistry = new Map(); 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(); @@ -528,4 +555,184 @@ describe("HostBashProxy", () => { expect(resolvedInteractionIds).toEqual([]); }); }); + + 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; + expect(sent.targetClientId).toBe("client-abc"); + + // Options passed to broadcastMessage should also have targetClientId + const opts = sentMessageOptions[0] as Record | undefined; + expect(opts?.targetClientId).toBe("client-abc"); + + const requestId = sent.requestId as string; + proxy.resolve(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; + expect(sent.targetClientId).toBe("client-abc"); + + const opts = sentMessageOptions[0] as Record | undefined; + expect(opts?.targetClientId).toBe("client-abc"); + + const requestId = sent.requestId as string; + proxy.resolve(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; + expect(sent.type).toBe("host_bash_request"); + // No target client resolved — untargeted broadcast + expect(sent.targetClientId).toBeUndefined(); + + const opts = sentMessageOptions[0] as Record | undefined; + expect(opts?.targetClientId).toBeUndefined(); + + // Manually resolve to clean up + const requestId = sent.requestId as string; + proxy.resolve(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; + 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.resolve(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(); + } + }); + }); }); diff --git a/assistant/src/daemon/host-bash-proxy.ts b/assistant/src/daemon/host-bash-proxy.ts index 1efa1ceb753..d08ca459c72 100644 --- a/assistant/src/daemon/host-bash-proxy.ts +++ b/assistant/src/daemon/host-bash-proxy.ts @@ -21,6 +21,7 @@ interface PendingRequest { conversationId: string; /** Detach the abort listener from the caller's signal. No-op when no signal was passed. */ detachAbort: () => void; + targetClientId?: string; } export class HostBashProxy { @@ -69,6 +70,7 @@ export class HostBashProxy { working_dir?: string; timeout_seconds?: number; env?: Record; + targetClientId?: string; }, conversationId: string, signal?: AbortSignal, @@ -78,6 +80,28 @@ export class HostBashProxy { return Promise.resolve(result); } + const capableClients = assistantEventHub.listClientsByCapability("host_bash"); + + let resolvedTargetClientId: string | undefined; + + if (input.targetClientId) { + const target = assistantEventHub.getClientById(input.targetClientId); + if (!target || !target.capabilities.includes("host_bash")) { + return Promise.resolve({ + content: `Error: client "${input.targetClientId}" is not connected or does not support host_bash. Run \`assistant clients list --capability host_bash\` to see available clients.`, + isError: true, + }); + } + resolvedTargetClientId = input.targetClientId; + } else if (capableClients.length === 1) { + // Auto-resolve when exactly one capable client is connected. + resolvedTargetClientId = capableClients[0].clientId; + } + // capableClients.length === 0 or > 1 without explicit target: resolvedTargetClientId + // stays undefined and falls through to untargeted broadcast — the existing timeout/error + // path handles the zero-client case, and multi-client ambiguity is enforced at the tool + // executor layer (not here) once target_client_id is exposed in the tool schema. + const requestId = uuid(); return new Promise((resolve, reject) => { @@ -98,10 +122,13 @@ export class HostBashProxy { { requestId, command: input.command }, "Host bash proxy request timed out", ); + const timeoutMessage = resolvedTargetClientId + ? `Host bash proxy timed out waiting for response from client ${resolvedTargetClientId}` + : "Host bash proxy timed out waiting for client response"; resolve( formatShellOutput( "", - "Host bash proxy timed out waiting for client response", + timeoutMessage, null, true, timeoutSec, @@ -139,20 +166,26 @@ export class HostBashProxy { timeoutSec, conversationId, detachAbort, + targetClientId: resolvedTargetClientId, }); try { - broadcastMessage({ - type: "host_bash_request", - requestId, + broadcastMessage( + { + type: "host_bash_request", + requestId, + conversationId, + command: input.command, + working_dir: input.working_dir, + timeout_seconds: input.timeout_seconds, + targetClientId: resolvedTargetClientId, + ...(input.env && Object.keys(input.env).length > 0 + ? { env: input.env } + : {}), + }, conversationId, - command: input.command, - working_dir: input.working_dir, - timeout_seconds: input.timeout_seconds, - ...(input.env && Object.keys(input.env).length > 0 - ? { env: input.env } - : {}), - }); + { targetClientId: resolvedTargetClientId }, + ); } catch (err) { clearTimeout(timer); this.pending.delete(requestId);