@vellumai/assistant 0.8.0 → 0.8.2

package/src/daemon/host-app-control-proxy.ts

@@ -15,9 +15,17 @@
  * session at a time, and that session is bound to a specific target app.
  * The lock is module-level (`activeAppControlSession`) because the session
  * targets the user's actual desktop application, which is a host-wide
- * resource. It is acquired on a successful `app_control_start` (storing
- * `(conversationId, app)`) and released when the owning proxy's
- * `dispose()` fires.
+ * resource. It is acquired optimistically when `app_control_start` is
+ * dispatched (storing `(conversationId, app)`) so that the synchronous
+ * guard and the asynchronous host round-trip cannot race. A separate
+ * `confirmedAppControlSession` tracks the last session the host reported
+ * `running`; the rollback path on a failed `start` restores from the
+ * current confirmed pointer (not from a per-call snapshot of a sibling
+ * optimistic write), so two overlapping starts that both fail cannot
+ * leave a phantom lock — and a late-arriving `running` for an older
+ * overlapping start still updates the confirmed pointer so the lock
+ * survives a subsequent rollback of the newer start. The lock is
+ * released outright when the owning proxy's `dispose()` fires.
  *
  * `app_control_start` is the only tool that can acquire the lock — the
  * user's medium-risk approval at start time is the consent boundary. All
@@ -79,34 +87,47 @@ export interface ActiveAppControlSession {
 }
 
 /**
- * Currently active session, or `undefined` when no session is held.
+ * Currently active session, or `undefined` when no session is held. This
+ * is the optimistic value: it is set the moment a `start` is dispatched
+ * and only promoted to {@link confirmedAppControlSession} when the host
+ * returns `running`.
  *
  * Exported for test inspection only. Production code paths must not read
  * or mutate this directly — use the proxy methods.
  */
 let activeAppControlSession: ActiveAppControlSession | undefined;
 
-/** Test-only helper: read current session. */
+/**
+ * Last session whose `start` was confirmed by the host (`payload.state ===
+ * "running"`). Used as the rollback baseline so a failed start never
+ * restores a sibling in-flight optimistic write — only a session that was
+ * actually running can re-emerge from a rollback.
+ */
+let confirmedAppControlSession: ActiveAppControlSession | undefined;
+
+/** Test-only helper: read current (optimistic) session. */
 export function _getActiveAppControlSession():
   | ActiveAppControlSession
   | undefined {
   return activeAppControlSession;
 }
 
-/** Test-only helper: clear session between test cases. */
+/** Test-only helper: clear both session pointers between test cases. */
 export function _resetActiveAppControlSession(): void {
   activeAppControlSession = undefined;
+  confirmedAppControlSession = undefined;
 }
 
 /**
- * Test-only helper: prime the active session without a full `start` round-trip.
- * Useful for tests that exercise non-start tool paths and don't need to
- * verify the start flow itself.
+ * Test-only helper: prime both session pointers without a full `start`
+ * round-trip. Useful for tests that exercise non-start tool paths and
+ * don't need to verify the start flow itself.
  */
 export function _setActiveAppControlSession(
   session: ActiveAppControlSession,
 ): void {
   activeAppControlSession = session;
+  confirmedAppControlSession = session;
 }
 
 /**
@@ -144,12 +165,12 @@ function checkNonStartAuthorization(
   // except `stop`, and `stop` short-circuits in conversation-surfaces and
   // does not reach this method in production. A stop reaching here would
   // be a defensive bug — surface it explicitly rather than dispatch.
-  const requestedApp = (input as { app?: string }).app;
-  if (requestedApp == null) {
+  const requestedApp = (input as { app?: unknown }).app;
+  if (typeof requestedApp !== "string") {
     return {
       content:
-        "Tool input missing required 'app' field; cannot validate against " +
-        "the active app-control session.",
+        "Tool input missing required string 'app' field; cannot validate " +
+        "against the active app-control session.",
       isError: true,
     };
   }
@@ -222,6 +243,8 @@ export class HostAppControlProxy extends HostProxyBase<
     input: HostAppControlInput,
     conversationId: string,
     signal: AbortSignal,
+    sourceActorPrincipalId?: string,
+    targetClientId?: string,
   ): Promise<ToolExecutionResult> {
     if (signal.aborted) {
       return { content: "Aborted", isError: true };
@@ -232,10 +255,11 @@ export class HostAppControlProxy extends HostProxyBase<
     // belong to the active session and target the same `app`. Without this
     // gate, prompt-injected calls would bypass the start-time approval and
     // send raw input to arbitrary apps.
+    let attemptedSession: ActiveAppControlSession | undefined;
     if (input.tool === "start") {
       if (
         activeAppControlSession != null &&
-        activeAppControlSession.conversationId !== conversationId
+        activeAppControlSession.conversationId !== this.conversationId
       ) {
         return {
           content:
@@ -245,8 +269,24 @@ export class HostAppControlProxy extends HostProxyBase<
           isError: true,
         };
       }
+      // Acquire optimistically to close the TOCTOU window between this
+      // synchronous guard and the asynchronous `dispatchRequest` below. Two
+      // concurrent starts from different conversations would otherwise both
+      // see `activeAppControlSession == null` and both pass the guard. The
+      // lock is rolled back below if dispatch fails or the host returns a
+      // non-running state — keyed on object identity so that a later
+      // overlapping start that has already replaced our write is not
+      // clobbered by a stale rollback.
+      attemptedSession = {
+        conversationId: this.conversationId,
+        app: input.app,
+      };
+      activeAppControlSession = attemptedSession;
     } else {
-      const sessionError = checkNonStartAuthorization(input, conversationId);
+      const sessionError = checkNonStartAuthorization(
+        input,
+        this.conversationId,
+      );
       if (sessionError != null) {
         return sessionError;
       }
@@ -258,9 +298,21 @@ export class HostAppControlProxy extends HostProxyBase<
         input,
         conversationId,
         signal,
+        undefined,
+        targetClientId,
       );
-      return this.handleSuccess(input, payload);
+      if (input.tool === "start") {
+        if (payload.state === "running") {
+          this.promoteStartIfCurrent(attemptedSession);
+        } else {
+          this.rollbackStartIfCurrent(attemptedSession);
+        }
+      }
+      return this.handleSuccess(payload);
     } catch (err) {
+      if (input.tool === "start") {
+        this.rollbackStartIfCurrent(attemptedSession);
+      }
       if (err instanceof HostProxyRequestError) {
         if (err.reason === "timeout") {
           log.warn({ toolName }, "Host app-control proxy request timed out");
@@ -279,12 +331,71 @@ export class HostAppControlProxy extends HostProxyBase<
     }
   }
 
+  /**
+   * Roll back the optimistic overwrite performed by a `start` when the
+   * dispatch fails or the host returns non-running. Keyed on the
+   * `attempted` reference, not just `conversationId`, so that an
+   * out-of-order failure does not clobber a newer overlapping start that
+   * already replaced our write — e.g. start A → start B (pending) →
+   * start C (success); when B later fails, the live session is C and the
+   * identity check makes our rollback a no-op rather than restoring A.
+   *
+   * Restores from the *current* `confirmedAppControlSession`, not a
+   * per-call snapshot of it. This matters when a late-arriving `running`
+   * for an older overlapping start has updated `confirmedAppControlSession`
+   * in the meantime: if A is dispatched, then C is dispatched (overwriting
+   * active), then A returns `running` (confirming A), then C returns
+   * non-running, the rollback must restore active to A — not undefined.
+   */
+  private rollbackStartIfCurrent(
+    attempted: ActiveAppControlSession | undefined,
+  ): void {
+    if (attempted != null && activeAppControlSession === attempted) {
+      activeAppControlSession = confirmedAppControlSession;
+    }
+  }
+
+  /**
+   * Promote this start's optimistic write to the confirmed pointer when
+   * the host returns `running`. Gated on conversation ownership rather
+   * than object identity: a newer overlapping start in the same
+   * conversation may have superseded our optimistic write while we were
+   * waiting on the host, but the host's `running` response for our
+   * `attempted` is still ground-truth that the lock should be held.
+   * The conversation-ownership check ensures we don't resurrect a session
+   * after `dispose()` cleared the lock or after another conversation
+   * acquired it.
+   */
+  private promoteStartIfCurrent(
+    attempted: ActiveAppControlSession | undefined,
+  ): void {
+    if (attempted == null) return;
+    if (activeAppControlSession?.conversationId !== attempted.conversationId) {
+      return;
+    }
+    confirmedAppControlSession = attempted;
+  }
+
+  /**
+   * Release both the optimistic and confirmed module-level session
+   * pointers if this proxy is the current holder. Used by `dispose()` —
+   * distinct from `rollbackStartIfCurrent` because dispose is keyed on
+   * ownership (conversationId) rather than on a specific in-flight start.
+   */
+  private releaseSessionIfHeld(): void {
+    if (activeAppControlSession?.conversationId === this.conversationId) {
+      activeAppControlSession = undefined;
+    }
+    if (confirmedAppControlSession?.conversationId === this.conversationId) {
+      confirmedAppControlSession = undefined;
+    }
+  }
+
   // ---------------------------------------------------------------------------
   // Result handling
   // ---------------------------------------------------------------------------
 
   private handleSuccess(
-    input: HostAppControlInput,
     payload: HostAppControlResultPayload,
   ): ToolExecutionResult {
     // Update PNG-hash loop tracking only for the "running" state — other
@@ -304,15 +415,6 @@ export class HostAppControlProxy extends HostProxyBase<
       }
     }
 
-    // Store the exact `app` form for validation against subsequent
-    // non-start tool calls.
-    if (input.tool === "start" && payload.state === "running") {
-      activeAppControlSession = {
-        conversationId: this.conversationId,
-        app: input.app,
-      };
-    }
-
     return this.formatResult(payload, stuck);
   }
 
@@ -382,8 +484,6 @@ export class HostAppControlProxy extends HostProxyBase<
    */
   override dispose(): void {
     super.dispose();
-    if (activeAppControlSession?.conversationId === this.conversationId) {
-      activeAppControlSession = undefined;
-    }
+    this.releaseSessionIfHeld();
   }
 }

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

@@ -1,10 +1,5 @@
-import { v4 as uuid } from "uuid";
-
 import { getConfig } from "../config/loader.js";
-import {
-  assistantEventHub,
-  broadcastMessage,
-} from "../runtime/assistant-event-hub.js";
+import { assistantEventHub } from "../runtime/assistant-event-hub.js";
 import {
   ambiguousSameUserError,
   enforceSameActorOrErrorResult,
@@ -13,14 +8,42 @@ import {
 import * as pendingInteractions from "../runtime/pending-interactions.js";
 import { formatShellOutput } from "../tools/shared/shell-output.js";
 import type { ToolExecutionResult } from "../tools/types.js";
-import { AssistantError, ErrorCode } from "../util/errors.js";
 import { getLogger } from "../util/logger.js";
+import { HostProxyBase, HostProxyRequestError } from "./host-proxy-base.js";
 
 const log = getLogger("host-bash-proxy");
 
-export class HostBashProxy {
+export type HostBashInput = {
+  command: string;
+  working_dir?: string;
+  timeout_seconds?: number;
+  env?: Record<string, string>;
+  targetClientId?: string;
+};
+
+type HostBashResultPayload = {
+  stdout: string;
+  stderr: string;
+  exitCode: number | null;
+  timedOut: boolean;
+};
+
+export class HostBashProxy extends HostProxyBase<
+  Record<string, unknown>,
+  HostBashResultPayload
+> {
   private static _instance: HostBashProxy | null = null;
 
+  constructor() {
+    super({
+      capabilityName: "host_bash",
+      requestEventName: "host_bash_request",
+      cancelEventName: "host_bash_cancel",
+      resultPendingKind: "host_bash",
+      disposedMessage: "Host bash proxy disposed",
+    });
+  }
+
   /**
    * Lazily-initialized singleton. Availability of an actual desktop
    * connection is checked at send time via the assistant event hub,
@@ -47,31 +70,15 @@ export class HostBashProxy {
     HostBashProxy._instance = null;
   }
 
-  /**
-   * Whether a client with `host_bash` capability is connected.
-   */
-  isAvailable(): boolean {
-    return (
-      assistantEventHub.getMostRecentClientByCapability("host_bash") != null
-    );
-  }
-
-  request(
-    input: {
-      command: string;
-      working_dir?: string;
-      timeout_seconds?: number;
-      env?: Record<string, string>;
-      targetClientId?: string;
-    },
+  async request(
+    input: HostBashInput,
     conversationId: string,
     signal?: AbortSignal,
     // Principal ID of the actor on whose behalf this request is initiated.
     sourceActorPrincipalId?: string,
   ): Promise<ToolExecutionResult> {
     if (signal?.aborted) {
-      const result = formatShellOutput("", "Aborted", null, false, 0);
-      return Promise.resolve(result);
+      return formatShellOutput("", "Aborted", null, false, 0);
     }
 
     let resolvedTargetClientId: string | undefined;
@@ -79,10 +86,10 @@ export class HostBashProxy {
     if (input.targetClientId) {
       const target = assistantEventHub.getClientById(input.targetClientId);
       if (!target || !target.capabilities.includes("host_bash")) {
-        return Promise.resolve({
+        return {
           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 {
@@ -97,7 +104,7 @@ export class HostBashProxy {
         sourceActorPrincipalId,
       });
       if (resolved.kind === "ambiguous") {
-        return Promise.resolve(ambiguousSameUserError("host_bash"));
+        return ambiguousSameUserError("host_bash");
       }
       resolvedTargetClientId =
         resolved.kind === "match" ? resolved.clientId : undefined;
@@ -114,98 +121,56 @@ export class HostBashProxy {
         targetClientId: resolvedTargetClientId,
         op: "host_bash",
       });
-      if (rejection) return Promise.resolve(rejection);
+      if (rejection) return rejection;
     }
 
-    const requestId = uuid();
-
-    return new Promise<ToolExecutionResult>((resolve, reject) => {
-      const shellMaxTimeoutSec = getConfig().timeouts.shellMaxTimeoutSec;
-      const timeoutSec = input.timeout_seconds ?? shellMaxTimeoutSec;
-      const proxyTimeoutSec = timeoutSec + 3;
-
-      let detachAbort: () => void = () => {};
-
-      const timer = setTimeout(() => {
-        pendingInteractions.resolve(requestId);
-        log.warn(
-          { 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("", timeoutMessage, null, true, timeoutSec));
-      }, proxyTimeoutSec * 1000);
-
-      if (signal) {
-        const onAbort = () => {
-          if (pendingInteractions.get(requestId)) {
-            pendingInteractions.resolve(requestId);
-            try {
-              broadcastMessage(
-                {
-                  type: "host_bash_cancel",
-                  requestId,
-                  conversationId,
-                  targetClientId: resolvedTargetClientId,
-                },
-                conversationId,
-                { targetClientId: resolvedTargetClientId },
-              );
-            } catch {
-              // Best-effort cancel notification
-            }
-            resolve(formatShellOutput("", "Aborted", null, false, 0));
-          }
-        };
-        signal.addEventListener("abort", onAbort, { once: true });
-        detachAbort = () => signal.removeEventListener("abort", onAbort);
-      }
-
-      pendingInteractions.register(requestId, {
+    const shellMaxTimeoutSec = getConfig().timeouts.shellMaxTimeoutSec;
+    const timeoutSec = input.timeout_seconds ?? shellMaxTimeoutSec;
+    const proxyTimeoutMs = (timeoutSec + 3) * 1000;
+
+    // Spread command fields at the top level of the envelope so the desktop
+    // client receives the same flat message shape it has always expected.
+    const extraFields: Record<string, unknown> = { command: input.command };
+    if (input.working_dir !== undefined) extraFields.working_dir = input.working_dir;
+    if (input.timeout_seconds !== undefined)
+      extraFields.timeout_seconds = input.timeout_seconds;
+    if (input.env && Object.keys(input.env).length > 0) extraFields.env = input.env;
+
+    try {
+      const payload = await this.dispatchRequest(
+        "host_bash",
+        {},
         conversationId,
-        kind: "host_bash",
-        rpcResolve: resolve as (v: unknown) => void,
-        rpcReject: reject,
-        timer,
-        detachAbort,
-        targetClientId: resolvedTargetClientId,
-        targetActorPrincipalId:
-          resolvedTargetClientId != null
-            ? assistantEventHub.getActorPrincipalIdForClient(
-                resolvedTargetClientId,
-              )
-            : undefined,
-        metadata: { timeoutSec },
-      });
-
-      try {
-        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,
-          { targetClientId: resolvedTargetClientId },
-        );
-      } catch (err) {
-        pendingInteractions.resolve(requestId);
-        log.warn(
-          { requestId, command: input.command, err },
-          "Host bash proxy send failed",
-        );
-        reject(err instanceof Error ? err : new Error(String(err)));
+        signal,
+        extraFields,
+        resolvedTargetClientId,
+        proxyTimeoutMs,
+      );
+      return formatShellOutput(
+        payload.stdout,
+        payload.stderr,
+        payload.exitCode,
+        payload.timedOut,
+        timeoutSec,
+      );
+    } catch (err) {
+      if (err instanceof HostProxyRequestError) {
+        if (err.reason === "timeout") {
+          log.warn(
+            { command: input.command },
+            "Host bash proxy request timed out",
+          );
+          const msg = resolvedTargetClientId
+            ? `Host bash proxy timed out waiting for response from client ${resolvedTargetClientId}`
+            : "Host bash proxy timed out waiting for client response";
+          return formatShellOutput("", msg, null, true, timeoutSec);
+        }
+        if (err.reason === "aborted") {
+          return formatShellOutput("", "Aborted", null, false, 0);
+        }
       }
-    });
+      throw err;
+    }
   }
 
   /**
@@ -220,45 +185,7 @@ export class HostBashProxy {
       timedOut: boolean;
     },
   ): void {
-    const interaction = pendingInteractions.resolve(requestId);
-    if (!interaction?.rpcResolve) {
-      log.warn({ requestId }, "No pending host bash request for response");
-      return;
-    }
-    const timeoutSec = (interaction.metadata?.timeoutSec as number) ?? 0;
-    const result = formatShellOutput(
-      response.stdout,
-      response.stderr,
-      response.exitCode,
-      response.timedOut,
-      timeoutSec,
-    );
-    interaction.rpcResolve(result);
-  }
-
-  dispose(): void {
-    for (const entry of pendingInteractions.getByKind("host_bash")) {
-      pendingInteractions.resolve(entry.requestId);
-      try {
-        broadcastMessage(
-          {
-            type: "host_bash_cancel",
-            requestId: entry.requestId,
-            conversationId: entry.conversationId,
-            targetClientId: entry.targetClientId,
-          },
-          entry.conversationId,
-          { targetClientId: entry.targetClientId },
-        );
-      } catch {
-        // Best-effort cancel notification — connection may already be closed.
-      }
-      entry.rpcReject?.(
-        new AssistantError(
-          "Host bash proxy disposed",
-          ErrorCode.INTERNAL_ERROR,
-        ),
-      );
-    }
+    pendingInteractions.resolve(requestId);
+    this.resolve(requestId, response);
   }
 }