@vellumai/assistant 0.7.1 → 0.7.2

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

@@ -13,25 +13,24 @@ 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-transfer-proxy");
 
-interface PendingTransfer {
-  resolve: (result: ToolExecutionResult) => void;
-  reject: (err: Error) => void;
-  timer: ReturnType<typeof setTimeout>;
+/**
+ * Lightweight entry for the transfers map (keyed by transferId).
+ * Points back to the requestId so route handlers can correlate
+ * content endpoints with the pending interaction.
+ */
+interface TransferEntry {
   requestId: string;
   transferId: string;
-  conversationId: string;
   direction: "to_host" | "to_sandbox";
   filePath: string;
   overwrite?: boolean;
   sizeBytes?: number;
   sha256?: string;
   fileBuffer?: Buffer;
-  /** Detach the abort listener from the caller's signal. No-op when no signal was passed. */
-  detachAbort: () => void;
+  targetClientId?: string;
 }
 
 /**
@@ -82,10 +81,22 @@ export class HostTransferProxy {
     HostTransferProxy._instance = null;
   }
 
-  /** Pending transfers keyed by requestId (for resolution from client results). */
-  private pending = new Map<string, PendingTransfer>();
   /** Pending transfers keyed by transferId (for content endpoint lookups). */
-  private transfers = new Map<string, PendingTransfer>();
+  private transfers = new Map<string, TransferEntry>();
+  /**
+   * Briefly retains size/sha256 of a just-consumed transfer so the GET-content
+   * route's `resolveResponseHeaders` callback (which the HTTP adapter invokes
+   * AFTER the request handler) can still set `Content-Length` and
+   * `X-Transfer-SHA256` headers. Without this, the handler's `getTransferContent`
+   * call deletes the entry before the header resolver runs, and the resolver
+   * silently falls back to default headers — meaning the documented response
+   * headers were never actually sent. Entries here self-clear on read; a 30s
+   * fallback timer prevents long-term retention if the resolver never runs.
+   */
+  private justConsumedMetadata = new Map<
+    string,
+    { sizeBytes: number; sha256: string }
+  >();
 
   /**
    * Whether a client with `host_file` capability is connected.
@@ -97,10 +108,6 @@ export class HostTransferProxy {
     );
   }
 
-  private send(msg: ServerMessage): void {
-    broadcastMessage(msg, undefined, { targetCapability: "host_file" });
-  }
-
   /**
    * Request a file transfer from the sandbox to the host machine.
    *
@@ -114,6 +121,7 @@ export class HostTransferProxy {
       destPath: string;
       overwrite: boolean;
       conversationId: string;
+      targetClientId?: string;
     },
     signal?: AbortSignal,
   ): Promise<ToolExecutionResult> {
@@ -121,13 +129,32 @@ export class HostTransferProxy {
       return Promise.resolve({ content: "Aborted", isError: true });
     }
 
+    let resolvedTargetClientId: string | undefined = 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();
     const transferId = uuid();
 
     return new Promise<ToolExecutionResult>((resolve, reject) => {
       readFile(input.sourcePath)
         .then((fileBuffer) => {
-          // Check again after async read in case signal fired during I/O.
           if (signal?.aborted) {
             resolve({ content: "Aborted", isError: true });
             return;
@@ -142,37 +169,40 @@ export class HostTransferProxy {
           let detachAbort: () => void = () => {};
 
           const timer = setTimeout(() => {
-            this.pending.delete(requestId);
             this.transfers.delete(transferId);
-            detachAbort();
             pendingInteractions.resolve(requestId);
             log.warn(
               { requestId, transferId, direction: "to_host" },
               "Host transfer proxy request timed out",
             );
             resolve({
-              content:
-                "Host transfer proxy timed out waiting for client response",
+              content: resolvedTargetClientId
+                ? `Host transfer proxy timed out waiting for response from client '${resolvedTargetClientId}'`
+                : "Host transfer proxy timed out waiting for client response",
               isError: true,
             });
           }, timeoutMs);
 
           if (signal) {
             const onAbort = () => {
-              if (this.pending.has(requestId)) {
-                clearTimeout(timer);
-                this.pending.delete(requestId);
+              if (pendingInteractions.get(requestId)) {
                 this.transfers.delete(transferId);
-                detachAbort();
                 pendingInteractions.resolve(requestId);
                 try {
-                  this.send({
-                    type: "host_transfer_cancel",
-                    requestId,
-                    conversationId: input.conversationId,
-                  });
+                  broadcastMessage(
+                    {
+                      type: "host_transfer_cancel",
+                      requestId,
+                      conversationId: input.conversationId,
+                      ...(resolvedTargetClientId != null
+                        ? { targetClientId: resolvedTargetClientId }
+                        : {}),
+                    },
+                    input.conversationId,
+                    { targetClientId: resolvedTargetClientId },
+                  );
                 } catch {
-                  // Best-effort cancel notification — connection may already be closed.
+                  // Best-effort cancel notification
                 }
                 resolve({ content: "Aborted", isError: true });
               }
@@ -181,40 +211,49 @@ export class HostTransferProxy {
             detachAbort = () => signal.removeEventListener("abort", onAbort);
           }
 
-          const entry: PendingTransfer = {
-            resolve,
-            reject,
-            timer,
+          this.transfers.set(transferId, {
             requestId,
             transferId,
-            conversationId: input.conversationId,
             direction: "to_host",
             filePath: input.destPath,
             sizeBytes,
             sha256,
             fileBuffer,
+            targetClientId: resolvedTargetClientId,
+          });
+
+          pendingInteractions.register(requestId, {
+            conversationId: input.conversationId,
+            kind: "host_transfer",
+            targetClientId: resolvedTargetClientId,
+            rpcResolve: resolve,
+            rpcReject: reject,
+            timer,
             detachAbort,
-          };
-          this.pending.set(requestId, entry);
-          this.transfers.set(transferId, entry);
+            metadata: { transferId },
+          });
 
           try {
-            this.send({
-              type: "host_transfer_request",
-              requestId,
-              conversationId: input.conversationId,
-              direction: "to_host",
-              transferId,
-              destPath: input.destPath,
-              sizeBytes,
-              sha256,
-              overwrite: input.overwrite,
-            });
+            broadcastMessage(
+              {
+                type: "host_transfer_request",
+                requestId,
+                conversationId: input.conversationId,
+                direction: "to_host",
+                transferId,
+                destPath: input.destPath,
+                sizeBytes,
+                sha256,
+                overwrite: input.overwrite,
+                ...(resolvedTargetClientId != null
+                  ? { targetClientId: resolvedTargetClientId }
+                  : {}),
+              },
+              input.conversationId,
+              { targetClientId: resolvedTargetClientId },
+            );
           } catch (err) {
-            clearTimeout(timer);
-            this.pending.delete(requestId);
             this.transfers.delete(transferId);
-            detachAbort();
             pendingInteractions.resolve(requestId);
             log.warn(
               { requestId, transferId, err },
@@ -249,6 +288,7 @@ export class HostTransferProxy {
       destPath: string;
       overwrite?: boolean;
       conversationId: string;
+      targetClientId?: string;
     },
     signal?: AbortSignal,
   ): Promise<ToolExecutionResult> {
@@ -256,6 +296,26 @@ export class HostTransferProxy {
       return Promise.resolve({ content: "Aborted", isError: true });
     }
 
+    let resolvedTargetClientId: string | undefined = 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();
     const transferId = uuid();
 
@@ -265,36 +325,40 @@ export class HostTransferProxy {
       let detachAbort: () => void = () => {};
 
       const timer = setTimeout(() => {
-        this.pending.delete(requestId);
         this.transfers.delete(transferId);
-        detachAbort();
         pendingInteractions.resolve(requestId);
         log.warn(
           { requestId, transferId, direction: "to_sandbox" },
           "Host transfer proxy request timed out",
         );
         resolve({
-          content: "Host transfer proxy timed out waiting for client response",
+          content: resolvedTargetClientId
+            ? `Host transfer proxy timed out waiting for response from client '${resolvedTargetClientId}'`
+            : "Host transfer proxy timed out waiting for client response",
           isError: true,
         });
       }, timeoutMs);
 
       if (signal) {
         const onAbort = () => {
-          if (this.pending.has(requestId)) {
-            clearTimeout(timer);
-            this.pending.delete(requestId);
+          if (pendingInteractions.get(requestId)) {
             this.transfers.delete(transferId);
-            detachAbort();
             pendingInteractions.resolve(requestId);
             try {
-              this.send({
-                type: "host_transfer_cancel",
-                requestId,
-                conversationId: input.conversationId,
-              });
+              broadcastMessage(
+                {
+                  type: "host_transfer_cancel",
+                  requestId,
+                  conversationId: input.conversationId,
+                  ...(resolvedTargetClientId != null
+                    ? { targetClientId: resolvedTargetClientId }
+                    : {}),
+                },
+                input.conversationId,
+                { targetClientId: resolvedTargetClientId },
+              );
             } catch {
-              // Best-effort cancel notification — connection may already be closed.
+              // Best-effort cancel notification
             }
             resolve({ content: "Aborted", isError: true });
           }
@@ -303,35 +367,44 @@ export class HostTransferProxy {
         detachAbort = () => signal.removeEventListener("abort", onAbort);
       }
 
-      const entry: PendingTransfer = {
-        resolve,
-        reject,
-        timer,
+      this.transfers.set(transferId, {
         requestId,
         transferId,
-        conversationId: input.conversationId,
         direction: "to_sandbox",
         filePath: input.destPath,
         overwrite: input.overwrite,
+        targetClientId: resolvedTargetClientId,
+      });
+
+      pendingInteractions.register(requestId, {
+        conversationId: input.conversationId,
+        kind: "host_transfer",
+        targetClientId: resolvedTargetClientId,
+        rpcResolve: resolve,
+        rpcReject: reject,
+        timer,
         detachAbort,
-      };
-      this.pending.set(requestId, entry);
-      this.transfers.set(transferId, entry);
+        metadata: { transferId },
+      });
 
       try {
-        this.send({
-          type: "host_transfer_request",
-          requestId,
-          conversationId: input.conversationId,
-          direction: "to_sandbox",
-          transferId,
-          sourcePath: input.sourcePath,
-        });
+        broadcastMessage(
+          {
+            type: "host_transfer_request",
+            requestId,
+            conversationId: input.conversationId,
+            direction: "to_sandbox",
+            transferId,
+            sourcePath: input.sourcePath,
+            ...(resolvedTargetClientId != null
+              ? { targetClientId: resolvedTargetClientId }
+              : {}),
+          },
+          input.conversationId,
+          { targetClientId: resolvedTargetClientId },
+        );
       } catch (err) {
-        clearTimeout(timer);
-        this.pending.delete(requestId);
         this.transfers.delete(transferId);
-        detachAbort();
         pendingInteractions.resolve(requestId);
         log.warn(
           { requestId, transferId, err },
@@ -353,23 +426,21 @@ export class HostTransferProxy {
       errorMessage?: string;
     },
   ): void {
-    const entry = this.pending.get(requestId);
-    if (!entry) {
+    const interaction = pendingInteractions.resolve(requestId);
+    if (!interaction?.rpcResolve) {
       log.warn({ requestId }, "No pending host transfer request for response");
       return;
     }
-    clearTimeout(entry.timer);
-    entry.detachAbort();
-    this.pending.delete(requestId);
-    this.transfers.delete(entry.transferId);
+    const transferId = interaction.metadata?.transferId as string | undefined;
+    if (transferId) this.transfers.delete(transferId);
 
     if (result.isError) {
-      entry.resolve({
+      interaction.rpcResolve({
         content: result.errorMessage ?? "Host transfer failed",
         isError: true,
       });
     } else {
-      entry.resolve({
+      interaction.rpcResolve({
         content: `File transferred successfully${result.bytesWritten != null ? ` (${result.bytesWritten} bytes)` : ""}`,
         isError: false,
       });
@@ -394,8 +465,21 @@ export class HostTransferProxy {
     ) {
       return null;
     }
-    // Single-use: consume the transfer from the transfers map.
-    // The pending map entry stays alive for the result resolution.
+    // Stash size/sha256 so the GET-content route's `resolveResponseHeaders`
+    // callback can still set `Content-Length` and `X-Transfer-SHA256` on the
+    // response. The HTTP adapter invokes the handler (this method) BEFORE the
+    // response-header resolver, so without this stash the resolver sees a
+    // deleted entry and silently falls back to default headers.
+    this.justConsumedMetadata.set(transferId, {
+      sizeBytes: entry.sizeBytes,
+      sha256: entry.sha256,
+    });
+    // Fallback cleanup: if the resolver never reads (e.g., handler error after
+    // consume, request abort), drop the metadata after a short grace window.
+    // `unref()` so the timer never holds the process open.
+    setTimeout(() => {
+      this.justConsumedMetadata.delete(transferId);
+    }, 30_000).unref?.();
     this.transfers.delete(transferId);
     return {
       buffer: entry.fileBuffer,
@@ -404,6 +488,23 @@ export class HostTransferProxy {
     };
   }
 
+  /**
+   * Returns and clears the size/sha256 metadata for a transfer that was just
+   * consumed by `getTransferContent`. Intended for use by the GET-content
+   * route's `resolveResponseHeaders` callback to populate `Content-Length` and
+   * `X-Transfer-SHA256` response headers. Returns null if no metadata is
+   * cached (e.g., transfer was never consumed, or already read by a previous
+   * resolver call).
+   */
+  takeJustConsumedTransferMetadata(
+    transferId: string,
+  ): { sizeBytes: number; sha256: string } | null {
+    const meta = this.justConsumedMetadata.get(transferId);
+    if (!meta) return null;
+    this.justConsumedMetadata.delete(transferId);
+    return meta;
+  }
+
   /**
    * Receive file content from the client for a to_sandbox transfer (the PUT content endpoint).
    *
@@ -437,29 +538,25 @@ export class HostTransferProxy {
 
     const { requestId } = entry;
 
-    // Enforce overwrite policy before writing.
     if (entry.overwrite !== true && existsSync(entry.filePath)) {
       const errorMsg = `Destination file already exists: ${entry.filePath}. Set overwrite to true to replace it.`;
-      clearTimeout(entry.timer);
-      entry.detachAbort();
-      this.pending.delete(requestId);
+      const interaction = pendingInteractions.resolve(requestId);
       this.transfers.delete(transferId);
-      entry.resolve({ content: errorMsg, isError: true });
+      interaction?.rpcResolve?.({ content: errorMsg, isError: true });
       return { accepted: false, error: errorMsg };
     }
 
     const cleanup = () => {
-      clearTimeout(entry.timer);
-      entry.detachAbort();
-      this.pending.delete(requestId);
+      pendingInteractions.resolve(requestId);
       this.transfers.delete(transferId);
     };
 
     try {
       await mkdir(dirname(entry.filePath), { recursive: true });
       await writeFile(entry.filePath, data);
+      const interaction = pendingInteractions.get(requestId);
       cleanup();
-      entry.resolve({
+      interaction?.rpcResolve?.({
         content: `File received and written to ${entry.filePath} (${data.length} bytes)`,
         isError: false,
       });
@@ -470,31 +567,37 @@ export class HostTransferProxy {
         { transferId, filePath: entry.filePath, err },
         "Failed to write received transfer content",
       );
+      const interaction = pendingInteractions.get(requestId);
       cleanup();
-      entry.resolve({ content: errorMsg, isError: true });
+      interaction?.rpcResolve?.({ content: errorMsg, isError: true });
       return { accepted: false, error: errorMsg };
     }
   }
 
   /** Cancel a pending transfer by requestId. */
   cancel(requestId: string): void {
-    const entry = this.pending.get(requestId);
-    if (!entry) return;
-    clearTimeout(entry.timer);
-    entry.detachAbort();
-    this.pending.delete(requestId);
-    this.transfers.delete(entry.transferId);
+    const interaction = pendingInteractions.get(requestId);
+    if (!interaction) return;
+    const transferId = interaction.metadata?.transferId as string | undefined;
+    if (transferId) this.transfers.delete(transferId);
     pendingInteractions.resolve(requestId);
     try {
-      this.send({
-        type: "host_transfer_cancel",
-        requestId,
-        conversationId: entry.conversationId,
-      });
+      broadcastMessage(
+        {
+          type: "host_transfer_cancel",
+          requestId,
+          conversationId: interaction.conversationId,
+          ...(interaction.targetClientId != null
+            ? { targetClientId: interaction.targetClientId }
+            : {}),
+        },
+        interaction.conversationId,
+        { targetClientId: interaction.targetClientId },
+      );
     } catch {
-      // Best-effort cancel notification — connection may already be closed.
+      // Best-effort cancel notification
     }
-    entry.resolve({ content: "Transfer cancelled", isError: true });
+    interaction.rpcResolve?.({ content: "Transfer cancelled", isError: true });
   }
 
   hasPendingTransfer(transferId: string): boolean {
@@ -511,28 +614,43 @@ export class HostTransferProxy {
     return entry?.requestId ?? null;
   }
 
+  /**
+   * Look up the targetClientId for a given transferId without consuming the entry.
+   * Routes call this to verify ownership without affecting the transfer state.
+   * Returns null when untargeted (no validation needed).
+   */
+  getTargetClientIdForTransfer(transferId: string): string | null {
+    return this.transfers.get(transferId)?.targetClientId ?? null;
+  }
+
   dispose(): void {
-    for (const [requestId, entry] of this.pending) {
-      clearTimeout(entry.timer);
-      entry.detachAbort();
-      pendingInteractions.resolve(requestId);
+    for (const entry of pendingInteractions.getByKind("host_transfer")) {
+      const transferId = entry.metadata?.transferId as string | undefined;
+      if (transferId) this.transfers.delete(transferId);
+      pendingInteractions.resolve(entry.requestId);
       try {
-          this.send({
+        broadcastMessage(
+          {
             type: "host_transfer_cancel",
-            requestId,
+            requestId: entry.requestId,
             conversationId: entry.conversationId,
-          });
-        } catch {
-          // Best-effort cancel notification — connection may already be closed.
-        }
-      entry.reject(
+            ...(entry.targetClientId != null
+              ? { targetClientId: entry.targetClientId }
+              : {}),
+          },
+          entry.conversationId,
+          { targetClientId: entry.targetClientId },
+        );
+      } catch {
+        // Best-effort cancel notification
+      }
+      entry.rpcReject?.(
         new AssistantError(
           "Host transfer proxy disposed",
           ErrorCode.INTERNAL_ERROR,
         ),
       );
     }
-    this.pending.clear();
     this.transfers.clear();
   }
 }