@vellumai/assistant 0.7.1 → 0.7.2

package/src/daemon/host-cu-proxy.ts

@@ -9,6 +9,9 @@
  * Unlike HostBashProxy/HostFileProxy/HostTransferProxy, this is NOT a
  * singleton — each conversation gets its own instance because CU state
  * (step count, AX tree history, loop detection) is per-conversation.
+ *
+ * RPC lifecycle (resolve/reject/timer/detachAbort) is stored in
+ * pendingInteractions alongside routing metadata.
  */
 
 import { v4 as uuid } from "uuid";
@@ -24,7 +27,6 @@ import * as pendingInteractions from "../runtime/pending-interactions.js";
 import type { ToolExecutionResult } from "../tools/types.js";
 import { AssistantError, ErrorCode } from "../util/errors.js";
 import { getLogger } from "../util/logger.js";
-import type { ServerMessage } from "./message-protocol.js";
 
 const log = getLogger("host-cu-proxy");
 
@@ -62,28 +64,19 @@ export interface ActionRecord {
   reasoning?: string;
 }
 
-interface PendingRequest {
-  resolve: (result: ToolExecutionResult) => void;
-  reject: (err: Error) => void;
-  timer: ReturnType<typeof setTimeout>;
-  conversationId: string;
-  /** Detach the abort listener from the caller's signal. No-op when no signal was passed. */
-  detachAbort: () => void;
-}
-
 // ---------------------------------------------------------------------------
 // HostCuProxy
 // ---------------------------------------------------------------------------
 
 export class HostCuProxy {
-  private pending = new Map<string, PendingRequest>();
-
   // CU state tracking (per-conversation)
   private _stepCount = 0;
   private _maxSteps: number;
   private _previousAXTree: string | undefined;
   private _consecutiveUnchangedSteps = 0;
   private _actionHistory: ActionRecord[] = [];
+  /** Request IDs owned by this instance — used to scope dispose(). */
+  private _ownedRequests = new Set<string>();
 
   constructor(maxSteps = loadConfig().maxStepsPerSession) {
     this._maxSteps = maxSteps;
@@ -121,17 +114,7 @@ export class HostCuProxy {
    * Whether a client with `host_cu` capability is connected.
    */
   isAvailable(): boolean {
-    return (
-      assistantEventHub.getMostRecentClientByCapability("host_cu") != null
-    );
-  }
-
-  // ---------------------------------------------------------------------------
-  // Send helper
-  // ---------------------------------------------------------------------------
-
-  private send(msg: ServerMessage): void {
-    broadcastMessage(msg, undefined, { targetCapability: "host_cu" });
+    return assistantEventHub.getMostRecentClientByCapability("host_cu") != null;
   }
 
   // ---------------------------------------------------------------------------
@@ -145,6 +128,7 @@ export class HostCuProxy {
     stepNumber: number,
     reasoning?: string,
     signal?: AbortSignal,
+    targetClientId?: string,
   ): Promise<ToolExecutionResult> {
     if (signal?.aborted) {
       return Promise.resolve({
@@ -153,7 +137,6 @@ export class HostCuProxy {
       });
     }
 
-    // Enforce step limit before sending to client
     if (this._stepCount > this._maxSteps) {
       return Promise.resolve({
         content: `Step limit (${this._maxSteps}) exceeded. Call computer_use_done to finish.`,
@@ -167,8 +150,7 @@ export class HostCuProxy {
       let detachAbort: () => void = () => {};
 
       const timer = setTimeout(() => {
-        this.pending.delete(requestId);
-        detachAbort();
+        this._ownedRequests.delete(requestId);
         pendingInteractions.resolve(requestId);
         log.warn({ requestId, toolName }, "Host CU proxy request timed out");
         resolve({
@@ -179,19 +161,24 @@ export class HostCuProxy {
 
       if (signal) {
         const onAbort = () => {
-          if (this.pending.has(requestId)) {
-            clearTimeout(timer);
-            this.pending.delete(requestId);
-            detachAbort();
+          if (pendingInteractions.get(requestId)) {
+            this._ownedRequests.delete(requestId);
             pendingInteractions.resolve(requestId);
             try {
-              this.send({
-                type: "host_cu_cancel",
-                requestId,
+              broadcastMessage(
+                {
+                  type: "host_cu_cancel",
+                  requestId,
+                  conversationId,
+                  ...(targetClientId != null
+                    ? { targetClientId }
+                    : {}),
+                },
                 conversationId,
-              });
+                { targetClientId },
+              );
             } catch {
-              // Best-effort cancel notification — connection may already be closed.
+              // Best-effort cancel notification
             }
             resolve({ content: "Aborted", isError: true });
           }
@@ -200,22 +187,36 @@ export class HostCuProxy {
         detachAbort = () => signal.removeEventListener("abort", onAbort);
       }
 
-      this.pending.set(requestId, { resolve, reject, timer, conversationId, detachAbort });
+      this._ownedRequests.add(requestId);
+
+      pendingInteractions.register(requestId, {
+        conversationId,
+        kind: "host_cu",
+        targetClientId,
+        rpcResolve: resolve,
+        rpcReject: reject,
+        timer,
+        detachAbort,
+      });
 
       try {
-        this.send({
-          type: "host_cu_request",
-          requestId,
+        broadcastMessage(
+          {
+            type: "host_cu_request",
+            requestId,
+            conversationId,
+            toolName,
+            input,
+            stepNumber,
+            reasoning,
+            // Include in body so receiving client can verify targeted endpoint.
+            ...(targetClientId != null ? { targetClientId } : {}),
+          },
           conversationId,
-          toolName,
-          input,
-          stepNumber,
-          reasoning,
-        });
+          { targetClientId },
+        );
       } catch (err) {
-        clearTimeout(timer);
-        this.pending.delete(requestId);
-        detachAbort();
+        this._ownedRequests.delete(requestId);
         pendingInteractions.resolve(requestId);
         log.warn({ requestId, toolName, err }, "Host CU proxy send failed");
         reject(err instanceof Error ? err : new Error(String(err)));
@@ -223,28 +224,27 @@ export class HostCuProxy {
     });
   }
 
-  resolve(requestId: string, observation: CuObservationResult): void {
-    const entry = this.pending.get(requestId);
-    if (!entry) {
+  /**
+   * Process a CU observation from the client and resolve the RPC.
+   * Updates CU state (step tracking, AX tree history) and formats
+   * the observation into a ToolExecutionResult.
+   */
+  processObservation(
+    requestId: string,
+    observation: CuObservationResult,
+  ): ToolExecutionResult | undefined {
+    this._ownedRequests.delete(requestId);
+    const interaction = pendingInteractions.resolve(requestId);
+    if (!interaction?.rpcResolve) {
       log.warn({ requestId }, "No pending host CU request for response");
-      return;
+      return undefined;
     }
-    clearTimeout(entry.timer);
-    entry.detachAbort();
-    this.pending.delete(requestId);
 
-    // Capture pre-update state so formatObservation sees the correct previous AX tree
     const prevAXTree = this._previousAXTree;
-
-    // Update CU state from observation
     this.updateStateFromObservation(observation);
-
     const result = this.formatObservation(observation, prevAXTree);
-    entry.resolve(result);
-  }
-
-  hasPendingRequest(requestId: string): boolean {
-    return this.pending.has(requestId);
+    interaction.rpcResolve(result);
+    return result;
   }
 
   // ---------------------------------------------------------------------------
@@ -267,7 +267,6 @@ export class HostCuProxy {
       input,
       reasoning,
     });
-    // Keep history bounded
     if (this._actionHistory.length > MAX_HISTORY_ENTRIES) {
       this._actionHistory = this._actionHistory.slice(-MAX_HISTORY_ENTRIES);
     }
@@ -297,7 +296,6 @@ export class HostCuProxy {
     const prevTree = previousAXTree;
     const parts: string[] = [];
 
-    // Surface user guidance prominently so the model sees it first
     if (obs.userGuidance) {
       parts.push(`USER GUIDANCE: ${obs.userGuidance}`);
       parts.push("");
@@ -308,12 +306,10 @@ export class HostCuProxy {
       parts.push("");
     }
 
-    // AX tree diff / unchanged warning
     if (obs.axDiff) {
       parts.push(obs.axDiff);
       parts.push("");
     } else if (prevTree != null && obs.axTree != null) {
-      // Skip unchanged warning after wait actions — they intentionally yield no immediate change
       const lastAction =
         this._actionHistory.length > 0
           ? this._actionHistory[this._actionHistory.length - 1]
@@ -321,7 +317,6 @@ export class HostCuProxy {
       const isWaitAction = lastAction?.toolName === "computer_use_wait";
 
       if (!isWaitAction) {
-        // No diff means the screen didn't change
         if (
           this._consecutiveUnchangedSteps >=
           CONSECUTIVE_UNCHANGED_WARNING_THRESHOLD
@@ -338,7 +333,6 @@ export class HostCuProxy {
       }
     }
 
-    // Loop detection: identical actions repeated
     if (this._actionHistory.length >= LOOP_DETECTION_WINDOW) {
       const recent = this._actionHistory.slice(-LOOP_DETECTION_WINDOW);
       const allIdentical = recent.every(
@@ -354,7 +348,6 @@ export class HostCuProxy {
       }
     }
 
-    // Current screen state wrapped in markers for history compaction
     if (obs.axTree) {
       parts.push("<ax-tree>");
       parts.push("CURRENT SCREEN STATE:");
@@ -362,7 +355,6 @@ export class HostCuProxy {
       parts.push("</ax-tree>");
     }
 
-    // Secondary windows for cross-app awareness
     if (obs.secondaryWindows) {
       parts.push("");
       parts.push(obs.secondaryWindows);
@@ -372,7 +364,6 @@ export class HostCuProxy {
       );
     }
 
-    // Screenshot metadata
     const screenshotMeta = this.formatScreenshotMetadata(obs);
     if (screenshotMeta.length > 0) {
       parts.push("");
@@ -381,7 +372,6 @@ export class HostCuProxy {
 
     const content = parts.join("\n").trim() || "Action executed";
 
-    // Build content blocks for screenshot
     const contentBlocks: ContentBlock[] = [];
     if (obs.screenshot) {
       contentBlocks.push({
@@ -410,31 +400,36 @@ export class HostCuProxy {
   // ---------------------------------------------------------------------------
 
   dispose(): void {
-    for (const [requestId, entry] of this.pending) {
-      clearTimeout(entry.timer);
-      entry.detachAbort();
-      pendingInteractions.resolve(requestId);
+    for (const requestId of this._ownedRequests) {
+      const entry = pendingInteractions.resolve(requestId);
+      if (!entry) continue;
       try {
-        this.send({
-          type: "host_cu_cancel",
-          requestId,
-          conversationId: entry.conversationId,
-        });
+        broadcastMessage(
+          {
+            type: "host_cu_cancel",
+            requestId,
+            conversationId: entry.conversationId,
+            ...(entry.targetClientId != null
+              ? { targetClientId: entry.targetClientId }
+              : {}),
+          },
+          entry.conversationId,
+          { targetClientId: entry.targetClientId as string | undefined },
+        );
       } catch {
-        // Best-effort cancel notification — connection may already be closed.
+        // Best-effort cancel notification
       }
-      entry.reject(
+      entry.rpcReject?.(
         new AssistantError("Host CU proxy disposed", ErrorCode.INTERNAL_ERROR),
       );
     }
-    this.pending.clear();
+    this._ownedRequests.clear();
   }
 
   // ---------------------------------------------------------------------------
   // Private helpers
   // ---------------------------------------------------------------------------
 
-  /** Update consecutive-unchanged tracking from an incoming observation. */
   private updateStateFromObservation(obs: CuObservationResult): void {
     if (this._stepCount > 0) {
       if (

package/src/daemon/host-file-proxy.ts

@@ -9,7 +9,6 @@ import { readImageBase64 } from "../tools/shared/filesystem/image-read.js";
 import type { ToolExecutionResult } from "../tools/types.js";
 import { AssistantError, ErrorCode } from "../util/errors.js";
 import { getLogger } from "../util/logger.js";
-import type { ServerMessage } from "./message-protocol.js";
 import type { HostFileRequest } from "./message-types/host-file.js";
 
 /** Distributive omit that preserves union variant fields. */
@@ -25,17 +24,6 @@ export type HostFileInput = DistributiveOmit<
 
 const log = getLogger("host-file-proxy");
 
-interface PendingRequest {
-  resolve: (result: ToolExecutionResult) => void;
-  reject: (err: Error) => void;
-  timer: ReturnType<typeof setTimeout>;
-  operation: HostFileInput["operation"];
-  path: string;
-  conversationId: string;
-  /** Detach the abort listener from the caller's signal. No-op when no signal was passed. */
-  detachAbort: () => void;
-}
-
 export class HostFileProxy {
   private static _instance: HostFileProxy | null = null;
 
@@ -65,8 +53,6 @@ export class HostFileProxy {
     HostFileProxy._instance = null;
   }
 
-  private pending = new Map<string, PendingRequest>();
-
   /**
    * Whether a client with `host_file` capability is connected.
    * Note: host_file covers both file operations and transfers.
@@ -77,56 +63,81 @@ export class HostFileProxy {
     );
   }
 
-  private send(msg: ServerMessage): void {
-    broadcastMessage(msg, undefined, { targetCapability: "host_file" });
-  }
-
   request(
     input: HostFileInput,
     conversationId: string,
     signal?: AbortSignal,
+    targetClientId?: string,
   ): Promise<ToolExecutionResult> {
     if (signal?.aborted) {
       return Promise.resolve({ content: "Aborted", isError: true });
     }
 
+    // Resolve targetClientId: explicit → validate; single capable client → auto-resolve.
+    // Callers may embed targetClientId in the input object (tool handlers) or pass it as
+    // the 4th parameter (legacy). Prefer the explicit param; fall back to input field.
+    let resolvedTargetClientId: string | undefined = targetClientId ?? input.targetClientId;
+    if (resolvedTargetClientId != null) {
+      const client = assistantEventHub.getClientById(resolvedTargetClientId);
+      if (!client) {
+        return Promise.resolve({
+          content: `No connected client with id '${resolvedTargetClientId}' supports host_file. Run \`assistant clients list --capability host_file\` to see available clients.`,
+          isError: true,
+        });
+      }
+      if (!client.capabilities.includes("host_file")) {
+        return Promise.resolve({
+          content: `Client '${resolvedTargetClientId}' does not support host_file. Run \`assistant clients list --capability host_file\` to see available clients.`,
+          isError: true,
+        });
+      }
+    } else {
+      const capable = assistantEventHub.listClientsByCapability("host_file");
+      if (capable.length === 1) {
+        resolvedTargetClientId = capable[0].clientId;
+      }
+    }
+
     const requestId = uuid();
 
     return new Promise<ToolExecutionResult>((resolve, reject) => {
-      // File operations should be fast — 30 second timeout.
       const timeoutSec = 30;
 
       let detachAbort: () => void = () => {};
 
       const timer = setTimeout(() => {
-        this.pending.delete(requestId);
-        detachAbort();
         pendingInteractions.resolve(requestId);
         log.warn(
           { requestId, operation: input.operation },
           "Host file proxy request timed out",
         );
         resolve({
-          content: "Host file proxy timed out waiting for client response",
+          content: resolvedTargetClientId
+            ? `Host file proxy timed out waiting for response from client '${resolvedTargetClientId}'`
+            : "Host file proxy timed out waiting for client response",
           isError: true,
         });
       }, timeoutSec * 1000);
 
       if (signal) {
         const onAbort = () => {
-          if (this.pending.has(requestId)) {
-            clearTimeout(timer);
-            this.pending.delete(requestId);
-            detachAbort();
+          if (pendingInteractions.get(requestId)) {
             pendingInteractions.resolve(requestId);
             try {
-              this.send({
-                type: "host_file_cancel",
-                requestId,
+              broadcastMessage(
+                {
+                  type: "host_file_cancel",
+                  requestId,
+                  conversationId,
+                  ...(resolvedTargetClientId != null
+                    ? { targetClientId: resolvedTargetClientId }
+                    : {}),
+                },
                 conversationId,
-              });
+                { targetClientId: resolvedTargetClientId },
+              );
             } catch {
-              // Best-effort cancel notification — connection may already be closed.
+              // Best-effort cancel notification
             }
             resolve({ content: "Aborted", isError: true });
           }
@@ -135,27 +146,34 @@ export class HostFileProxy {
         detachAbort = () => signal.removeEventListener("abort", onAbort);
       }
 
-      this.pending.set(requestId, {
-        resolve,
-        reject,
-        timer,
-        operation: input.operation,
-        path: input.path,
+      pendingInteractions.register(requestId, {
         conversationId,
+        kind: "host_file",
+        targetClientId: resolvedTargetClientId,
+        rpcResolve: resolve,
+        rpcReject: reject,
+        timer,
         detachAbort,
+        metadata: { operation: input.operation, path: input.path },
       });
 
       try {
-        this.send({
-          ...input,
-          type: "host_file_request",
-          requestId,
+        broadcastMessage(
+          {
+            ...input,
+            type: "host_file_request",
+            requestId,
+            conversationId,
+            // Always include in message body so the receiving client can verify
+            // which endpoint was targeted (even when auto-resolved).
+            ...(resolvedTargetClientId != null
+              ? { targetClientId: resolvedTargetClientId }
+              : {}),
+          },
           conversationId,
-        });
+          { targetClientId: resolvedTargetClientId },
+        );
       } catch (err) {
-        clearTimeout(timer);
-        this.pending.delete(requestId);
-        detachAbort();
         pendingInteractions.resolve(requestId);
         log.warn(
           { requestId, operation: input.operation, err },
@@ -166,55 +184,61 @@ export class HostFileProxy {
     });
   }
 
+  /**
+   * Process a client result and resolve the RPC. Called by route handlers.
+   */
   resolve(
     requestId: string,
     response: { content: string; isError: boolean; imageData?: string },
   ): void {
-    const entry = this.pending.get(requestId);
-    if (!entry) {
+    const interaction = pendingInteractions.resolve(requestId);
+    if (!interaction?.rpcResolve) {
       log.warn({ requestId }, "No pending host file request for response");
       return;
     }
-    clearTimeout(entry.timer);
-    entry.detachAbort();
-    this.pending.delete(requestId);
+    const meta = interaction.metadata ?? {};
     if (
-      entry.operation === "read" &&
+      meta.operation === "read" &&
       !response.isError &&
       typeof response.imageData === "string" &&
       response.imageData.length > 0
     ) {
-      entry.resolve(readImageBase64(response.imageData, entry.path));
+      interaction.rpcResolve(
+        readImageBase64(response.imageData, meta.path as string),
+      );
       return;
     }
-    entry.resolve({ content: response.content, isError: response.isError });
-  }
-
-  hasPendingRequest(requestId: string): boolean {
-    return this.pending.has(requestId);
+    interaction.rpcResolve({
+      content: response.content,
+      isError: response.isError,
+    });
   }
 
   dispose(): void {
-    for (const [requestId, entry] of this.pending) {
-      clearTimeout(entry.timer);
-      entry.detachAbort();
-      pendingInteractions.resolve(requestId);
+    for (const entry of pendingInteractions.getByKind("host_file")) {
+      pendingInteractions.resolve(entry.requestId);
       try {
-        this.send({
-          type: "host_file_cancel",
-          requestId,
-          conversationId: entry.conversationId,
-        });
+        broadcastMessage(
+          {
+            type: "host_file_cancel",
+            requestId: entry.requestId,
+            conversationId: entry.conversationId,
+            ...(entry.targetClientId != null
+              ? { targetClientId: entry.targetClientId }
+              : {}),
+          },
+          entry.conversationId,
+          { targetClientId: entry.targetClientId as string | undefined },
+        );
       } catch {
-        // Best-effort cancel notification — connection may already be closed.
+        // Best-effort cancel notification
       }
-      entry.reject(
+      entry.rpcReject?.(
         new AssistantError(
           "Host file proxy disposed",
           ErrorCode.INTERNAL_ERROR,
         ),
       );
     }
-    this.pending.clear();
   }
 }