@vellumai/assistant 0.12.1 → 0.12.2-staging.1

package/src/runtime/http-server.ts

@@ -32,6 +32,12 @@ import {
 } from "../daemon/daemon-readiness.js";
 import { processMessage } from "../daemon/process-message.js";
 import { makeAddrInUseError } from "../daemon/startup-error.js";
+import { isAssistantDesktopEnabled } from "../desktop/desktop-feature.js";
+import {
+  DESKTOP_CLOSE,
+  destroyDesktopSessionManager,
+} from "../desktop/desktop-session-manager.js";
+import { DesktopStreamBridge } from "../desktop/desktop-stream-bridge.js";
 import {
   createLiveVoiceConnection,
   type LiveVoiceConnection,
@@ -229,6 +235,32 @@ interface WatchStreamWebSocketData {
   session?: WatchStreamSession;
 }
 
+/**
+ * WebSocket data for `/v1/desktop/stream`: raw RFB bytes to the assistant
+ * desktop.
+ */
+interface DesktopStreamWebSocketData {
+  wsType: "desktop-stream";
+  /** Bound at open time so message/close handlers reach this socket's pump. */
+  bridge?: DesktopStreamBridge;
+}
+
+type AllWebSocketData =
+  | MediaStreamWebSocketData
+  | SttStreamWebSocketData
+  | LiveVoiceWebSocketData
+  | WatchStreamWebSocketData
+  | DesktopStreamWebSocketData;
+
+function assistantDesktopEnabled(): boolean {
+  try {
+    return isAssistantDesktopEnabled(getConfig());
+  } catch (err) {
+    log.warn({ err }, "Failed to read config for desktop stream gate");
+    return false;
+  }
+}
+
 export class RuntimeHttpServer {
   private server: ReturnType<typeof Bun.serve> | null = null;
   private port: number;
@@ -264,11 +296,6 @@ export class RuntimeHttpServer {
   }
 
   async start(): Promise<void> {
-    type AllWebSocketData =
-      | MediaStreamWebSocketData
-      | SttStreamWebSocketData
-      | LiveVoiceWebSocketData
-      | WatchStreamWebSocketData;
     this.server = Bun.serve<AllWebSocketData>({
       port: this.port,
       hostname: this.hostname,
@@ -424,6 +451,22 @@ export class RuntimeHttpServer {
             void session.start();
             return;
           }
+          if (data.wsType === "desktop-stream") {
+            log.info("Desktop stream WebSocket opened");
+            if (!assistantDesktopEnabled()) {
+              ws.close(
+                DESKTOP_CLOSE.unavailable,
+                "Desktop is not available on this assistant",
+              );
+              return;
+            }
+            const bridge = new DesktopStreamBridge(ws, {
+              isEnabled: assistantDesktopEnabled,
+            });
+            data.bridge = bridge;
+            void bridge.start();
+            return;
+          }
           log.warn("WebSocket opened with unknown data type — closing");
           ws.close(1008, "Unknown WebSocket type");
         },
@@ -474,6 +517,10 @@ export class RuntimeHttpServer {
             }
             return;
           }
+          if (data.wsType === "desktop-stream") {
+            data.bridge?.handleClientFrame(message);
+            return;
+          }
           log.warn("WebSocket message on unknown data type — closing");
           ws.close(1008, "Unknown WebSocket type");
         },
@@ -567,6 +614,14 @@ export class RuntimeHttpServer {
             }
             return;
           }
+          if (data.wsType === "desktop-stream") {
+            log.info(
+              { code, reason: reason?.toString() },
+              "Desktop stream WebSocket closed",
+            );
+            data.bridge?.handleClose();
+            return;
+          }
           log.warn(
             { code, reason: reason?.toString() },
             "WebSocket with unknown data type closed",
@@ -707,6 +762,9 @@ export class RuntimeHttpServer {
     // that a socket closing just before shutdown had already under way.
     await drainWatchRetros();
 
+    // Bounded by the manager's kill grace, so shutdown never hangs on X.
+    await destroyDesktopSessionManager();
+
     const liveVoiceManager = getLiveVoiceSessionManager();
     const liveVoiceSessionId = liveVoiceManager.activeSessionId;
     if (liveVoiceSessionId) {
@@ -810,6 +868,15 @@ export class RuntimeHttpServer {
       return this.handleWatchStreamUpgrade(req, server);
     }
 
+    // WebSocket upgrade for the assistant desktop RFB stream, under the same
+    // private-network restrictions and gateway-service token verification.
+    if (
+      path === "/v1/desktop/stream" &&
+      req.headers.get("upgrade")?.toLowerCase() === "websocket"
+    ) {
+      return this.handleDesktopStreamUpgrade(req, server);
+    }
+
     // Twilio webhook endpoints — before auth check because Twilio
     // webhook POSTs don't include bearer tokens.
     const twilioResponse = await this.handleTwilioWebhook(req, path);
@@ -949,199 +1016,141 @@ export class RuntimeHttpServer {
     return null;
   }
 
-  private handleMediaStreamUpgrade(
+  /**
+   * Shared path for gateway-proxied WebSocket upgrades: private network peers
+   * and origins only, then a gateway service token, then `parse` builds the
+   * socket data (or a 4xx Response) and the upgrade happens. The gateway owns
+   * downstream client auth and dials these upstreams on the client's behalf.
+   * Returns a Response on failure. On success Bun has consumed the request and
+   * nothing is returned; the type is `Response` only to fit the router.
+   */
+  private upgradeRuntimeStream(
     req: Request,
     server: ReturnType<typeof Bun.serve>,
+    label: string,
+    parse: (query: URLSearchParams) => AllWebSocketData | Response,
   ): Response {
     if (!isPrivateNetworkPeer(server, req) || !isPrivateNetworkOrigin(req)) {
       return httpError(
         "FORBIDDEN",
-        "Direct media-stream access disabled — only private network peers allowed",
+        `Direct ${label} access disabled: only private network peers allowed`,
         403,
       );
     }
-
-    // Verify the gateway service token before accepting the upgrade.
     const tokenError = this.verifyGatewayServiceToken(req);
     if (tokenError) {
       return tokenError;
     }
-
-    const wsUrl = new URL(req.url);
-    const callSessionId = wsUrl.searchParams.get("callSessionId");
-    if (!callSessionId) {
-      return new Response("Missing callSessionId", { status: 400 });
+    const data = parse(new URL(req.url).searchParams);
+    if (data instanceof Response) {
+      return data;
     }
-    // Media-stream connections use a distinct wsType so the open/message/close
-    // handlers route them to MediaStreamCallSession.
-    const upgraded = server.upgrade(req, {
-      data: {
-        wsType: "media-stream",
-        callSessionId,
-      } satisfies MediaStreamWebSocketData,
-    });
-    if (!upgraded) {
+    if (!server.upgrade(req, { data })) {
       return new Response("WebSocket upgrade failed", { status: 500 });
     }
-    // Bun's WebSocket upgrade consumes the request — no Response is sent.
     return undefined!;
   }
 
-  /**
-   * Handle WebSocket upgrade for `/v1/stt/stream`.
-   *
-   * Private-network restrictions apply (same as media-stream) so the
-   * runtime remains unreachable from the public internet. The gateway
-   * authenticates the downstream client and proxies the upgrade with a
-   * short-lived gateway service token.
-   */
-  private handleSttStreamUpgrade(
+  private handleMediaStreamUpgrade(
     req: Request,
     server: ReturnType<typeof Bun.serve>,
   ): Response {
-    if (!isPrivateNetworkPeer(server, req) || !isPrivateNetworkOrigin(req)) {
-      return httpError(
-        "FORBIDDEN",
-        "Direct STT stream access disabled — only private network peers allowed",
-        403,
-      );
-    }
-
-    // Verify the gateway service token before accepting the upgrade.
-    const tokenError = this.verifyGatewayServiceToken(req);
-    if (tokenError) {
-      return tokenError;
-    }
-
-    const wsUrl = new URL(req.url);
-    // provider is optional compatibility metadata — the runtime resolves
-    // the streaming transcriber from config (`services.stt.provider`).
-    const provider = wsUrl.searchParams.get("provider") ?? undefined;
-    const mimeType = wsUrl.searchParams.get("mimeType");
-    if (!mimeType) {
-      return new Response("Missing required query parameter: mimeType", {
-        status: 400,
-      });
-    }
-
-    const sampleRateRaw = wsUrl.searchParams.get("sampleRate");
-    const sampleRate = sampleRateRaw ? parseInt(sampleRateRaw, 10) : undefined;
+    return this.upgradeRuntimeStream(req, server, "media-stream", (query) => {
+      const callSessionId = query.get("callSessionId");
+      if (!callSessionId) {
+        return new Response("Missing callSessionId", { status: 400 });
+      }
+      // A distinct wsType routes these to MediaStreamCallSession.
+      return {
+        wsType: "media-stream",
+        callSessionId,
+      } satisfies MediaStreamWebSocketData;
+    });
+  }
 
-    const sessionId = crypto.randomUUID();
-    const upgraded = server.upgrade(req, {
-      data: {
+  /** Handle WebSocket upgrade for `/v1/stt/stream`. */
+  private handleSttStreamUpgrade(
+    req: Request,
+    server: ReturnType<typeof Bun.serve>,
+  ): Response {
+    return this.upgradeRuntimeStream(req, server, "STT stream", (query) => {
+      // provider is optional compatibility metadata; the runtime resolves
+      // the streaming transcriber from config (`services.stt.provider`).
+      const provider = query.get("provider") ?? undefined;
+      const mimeType = query.get("mimeType");
+      if (!mimeType) {
+        return new Response("Missing required query parameter: mimeType", {
+          status: 400,
+        });
+      }
+      const sampleRateRaw = query.get("sampleRate");
+      return {
         wsType: "stt-stream",
         provider,
         mimeType,
-        sampleRate,
-        sessionId,
-      } satisfies SttStreamWebSocketData,
+        sampleRate: sampleRateRaw ? parseInt(sampleRateRaw, 10) : undefined,
+        sessionId: crypto.randomUUID(),
+      } satisfies SttStreamWebSocketData;
     });
-    if (!upgraded) {
-      return new Response("WebSocket upgrade failed", { status: 500 });
-    }
-    // Bun's WebSocket upgrade consumes the request — no Response is sent.
-    return undefined!;
   }
 
-  /**
-   * Handle WebSocket upgrade for `/v1/live-voice`.
-   *
-   * The gateway owns downstream client auth and forwards this upstream with
-   * a short-lived gateway service token. The runtime accepts only private
-   * network peers/origins so the shell is not publicly reachable.
-   */
+  /** Handle WebSocket upgrade for `/v1/live-voice`. */
   private handleLiveVoiceUpgrade(
     req: Request,
     server: ReturnType<typeof Bun.serve>,
   ): Response {
-    if (!isPrivateNetworkPeer(server, req) || !isPrivateNetworkOrigin(req)) {
-      return httpError(
-        "FORBIDDEN",
-        "Direct live voice access disabled — only private network peers allowed",
-        403,
-      );
-    }
-
-    const tokenError = this.verifyGatewayServiceToken(req);
-    if (tokenError) {
-      return tokenError;
-    }
-
-    const guardianPrincipalId =
-      new URL(req.url).searchParams.get("guardianPrincipalId")?.trim() ||
-      undefined;
-
-    const upgraded = server.upgrade(req, {
-      data: {
+    return this.upgradeRuntimeStream(req, server, "live voice", (query) => {
+      const guardianPrincipalId = query.get("guardianPrincipalId")?.trim();
+      return {
         wsType: "live-voice",
         ...(guardianPrincipalId ? { guardianPrincipalId } : {}),
-      } satisfies LiveVoiceWebSocketData,
+      } satisfies LiveVoiceWebSocketData;
     });
-    if (!upgraded) {
-      return new Response("WebSocket upgrade failed", { status: 500 });
-    }
-    return undefined!;
   }
 
-  /**
-   * Handle WebSocket upgrade for `/v1/watch/stream`.
-   *
-   * Gated exactly as `/v1/stt/stream` is: private network peers and origins
-   * only, then a gateway service token. The gateway owns downstream client
-   * auth and dials this upstream on the client's behalf.
-   */
+  /** Handle WebSocket upgrade for `/v1/watch/stream`. */
   private handleWatchStreamUpgrade(
     req: Request,
     server: ReturnType<typeof Bun.serve>,
   ): Response {
-    if (!isPrivateNetworkPeer(server, req) || !isPrivateNetworkOrigin(req)) {
-      return httpError(
-        "FORBIDDEN",
-        "Direct watch stream access disabled: only private network peers allowed",
-        403,
-      );
-    }
-
-    const tokenError = this.verifyGatewayServiceToken(req);
-    if (tokenError) {
-      return tokenError;
-    }
-
-    const wsUrl = new URL(req.url);
-    const mimeType = wsUrl.searchParams.get("mimeType");
-    if (!mimeType) {
-      return new Response("Missing required query parameter: mimeType", {
-        status: 400,
-      });
-    }
-
-    const sampleRateRaw = wsUrl.searchParams.get("sampleRate");
-    const sampleRate = sampleRateRaw ? parseInt(sampleRateRaw, 10) : undefined;
-    const conversationId =
-      wsUrl.searchParams.get("conversationId")?.trim() || undefined;
-    const clientId = wsUrl.searchParams.get("clientId")?.trim() || undefined;
-    const parsedTarget = parseWatchCaptureTarget(wsUrl.searchParams);
-    if ("error" in parsedTarget) {
-      return new Response(parsedTarget.error, { status: 400 });
-    }
-
-    const upgraded = server.upgrade(req, {
-      data: {
+    return this.upgradeRuntimeStream(req, server, "watch stream", (query) => {
+      const mimeType = query.get("mimeType");
+      if (!mimeType) {
+        return new Response("Missing required query parameter: mimeType", {
+          status: 400,
+        });
+      }
+      const sampleRateRaw = query.get("sampleRate");
+      const parsedTarget = parseWatchCaptureTarget(query);
+      if ("error" in parsedTarget) {
+        return new Response(parsedTarget.error, { status: 400 });
+      }
+      return {
         wsType: "watch-stream",
         mimeType,
-        sampleRate,
-        conversationId,
-        clientId,
+        sampleRate: sampleRateRaw ? parseInt(sampleRateRaw, 10) : undefined,
+        conversationId: query.get("conversationId")?.trim() || undefined,
+        clientId: query.get("clientId")?.trim() || undefined,
         captureTarget: parsedTarget.captureTarget,
         sessionId: crypto.randomUUID(),
-      } satisfies WatchStreamWebSocketData,
+      } satisfies WatchStreamWebSocketData;
     });
-    if (!upgraded) {
-      return new Response("WebSocket upgrade failed", { status: 500 });
-    }
-    // Bun's WebSocket upgrade consumes the request, so no Response is sent.
-    return undefined!;
+  }
+
+  /**
+   * Handle WebSocket upgrade for `/v1/desktop/stream`. The feature gate runs
+   * in the open handler, as close code 4008.
+   */
+  private handleDesktopStreamUpgrade(
+    req: Request,
+    server: ReturnType<typeof Bun.serve>,
+  ): Response {
+    return this.upgradeRuntimeStream(
+      req,
+      server,
+      "desktop stream",
+      () => ({ wsType: "desktop-stream" }) satisfies DesktopStreamWebSocketData,
+    );
   }
 
   private async handleTwilioWebhook(

package/src/runtime/routes/__tests__/acp-routes.test.ts

@@ -11,7 +11,8 @@
  * silently auto-installs allowlisted adapter packages via a sandboxed `bun`
  * global install before failing with the install hint. `execFile` is stubbed
  * via the shared `installExecFileStub` helper so tests can script
- * `bun add --global` outcomes.
+ * `bun add --global` outcomes. It also threads the optional `model` through to
+ * the session manager and relays the warning a refused model comes back with.
  */
 
 import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
@@ -43,6 +44,10 @@ interface FakeSessionState {
   completedAt?: number;
   error?: string;
   stopReason?: string;
+  model?: string;
+  availableModels?: Array<{ value: string; label: string }>;
+  modelRevisionEpoch?: string;
+  modelRevision?: number;
   latestUsage?: {
     usedTokens: number;
     contextSize: number;
@@ -55,10 +60,27 @@ interface FakeSessionState {
 
 let fakeInMemorySessions: FakeSessionState[] = [];
 
-const spawnMock = mock(async () => ({
+interface SpawnResult {
+  acpSessionId: string;
+  protocolSessionId: string;
+  requestedModel?: string;
+  effectiveModel?: string;
+  modelWarning?: string;
+}
+
+const DEFAULT_SPAWN_RESULT: SpawnResult = {
   acpSessionId: "acp-route-session",
   protocolSessionId: "proto-route-session",
-}));
+  effectiveModel: "opus",
+};
+let spawnResult: SpawnResult = DEFAULT_SPAWN_RESULT;
+const spawnMock = mock(async (...args: unknown[]) => {
+  const options = args[6] as { model?: string } | undefined;
+  return {
+    ...spawnResult,
+    requestedModel: options?.model?.trim() || undefined,
+  };
+});
 
 const defaultSteerOrResumeImpl = async (
   _id: string,
@@ -182,6 +204,10 @@ interface ResponseShape {
     outputTokens?: number;
     eventLog?: unknown[];
     authErrorCode?: string;
+    model?: string;
+    availableModels?: Array<{ value: string; label: string }>;
+    modelRevisionEpoch?: string;
+    modelRevision?: number;
   }>;
 }
 
@@ -190,6 +216,7 @@ beforeEach(() => {
   clearHistory();
   resetExecFileStub();
   spawnMock.mockClear();
+  spawnResult = DEFAULT_SPAWN_RESULT;
   steerOrResumeMock.mockClear();
   steerOrResumeImpl = defaultSteerOrResumeImpl;
   _resetAdapterInstallCacheForTests();
@@ -328,6 +355,60 @@ describe("GET /v1/acp/sessions — merged in-memory + history", () => {
     expect(s!.outputTokens).toBeUndefined();
   });
 
+  test("carries the model for both layers and the picker for live ones only", async () => {
+    fakeInMemorySessions = [
+      {
+        id: "live-model",
+        agentId: "agent-live",
+        acpSessionId: "proto-live",
+        parentConversationId: "conv-model",
+        status: "running",
+        startedAt: 9000,
+        model: "opus",
+        modelRevisionEpoch: "01900000-0000-7000-8000-000000000001",
+        modelRevision: 17,
+        availableModels: [
+          { value: "sonnet", label: "Sonnet" },
+          { value: "opus", label: "Opus" },
+        ],
+      },
+    ];
+    insertHistoryRow({
+      id: "hist-model",
+      agentId: "agent-hist",
+      acpSessionId: "proto-hist",
+      parentConversationId: "conv-model",
+      startedAt: 1000,
+      status: "completed",
+    });
+
+    const handler = getSessionsHandler();
+    const body = (await handler({})) as ResponseShape;
+    const live = body.sessions.find((s) => s.id === "live-model");
+    const hist = body.sessions.find((s) => s.id === "hist-model");
+    expect(live).toMatchObject({
+      model: "opus",
+      modelRevisionEpoch: "01900000-0000-7000-8000-000000000001",
+      modelRevision: 17,
+      availableModels: [
+        { value: "sonnet", label: "Sonnet" },
+        { value: "opus", label: "Opus" },
+      ],
+    });
+    // A finished run has no live process to ask, so it reports neither.
+    expect(hist?.model).toBeUndefined();
+    expect(hist?.availableModels).toBeUndefined();
+  });
+
+  test("omits the model for history rows written without one", async () => {
+    insertHistoryRow({ id: "hist-no-model", status: "completed" });
+
+    const handler = getSessionsHandler();
+    const body = (await handler({})) as ResponseShape;
+    const s = body.sessions.find((row) => row.id === "hist-no-model");
+    expect(s?.model).toBeUndefined();
+  });
+
   test("dedupes by id with in-memory winning on collision", async () => {
     // Same id in both layers — in-memory entry should win and its eventLog
     // comes from the live ring buffer (empty here), not the stale history row.
@@ -579,6 +660,8 @@ describe("POST /v1/acp/spawn: sandboxed bun auto-install on missing binary", ()
       acpSessionId: "acp-route-session",
       protocolSessionId: "proto-route-session",
       agent: "claude",
+      requestedModel: null,
+      effectiveModel: "opus",
     });
     expect(spawnMock).toHaveBeenCalledTimes(1);
     // The real adapter binary is spawned, not a `bun x` wrapper.
@@ -593,7 +676,7 @@ describe("POST /v1/acp/spawn: sandboxed bun auto-install on missing binary", ()
     expect(args).toEqual([
       "add",
       "--global",
-      "@agentclientprotocol/claude-agent-acp",
+      "@agentclientprotocol/claude-agent-acp@0.75.1",
     ]);
   });
 
@@ -680,6 +763,78 @@ describe("POST /v1/acp/spawn: sandboxed bun auto-install on missing binary", ()
   });
 });
 
+describe("POST /v1/acp/spawn: model selection", () => {
+  test("threads the requested model to the session manager", async () => {
+    const handler = getSpawnHandler();
+    const body = (await handler({
+      body: { ...SPAWN_BODY, model: "opus" },
+    })) as Record<string, unknown>;
+
+    expect(confirmationRequests).toHaveLength(1);
+    expect(confirmationRequests[0]?.input).toEqual({
+      agent: "claude",
+      task: "do something",
+      cwd: process.cwd(),
+      model: "opus",
+    });
+    expect(spawnMock).toHaveBeenCalledTimes(1);
+    expect((spawnMock.mock.calls[0] as unknown[])[6]).toEqual({
+      model: "opus",
+    });
+
+    expect(body.requestedModel).toBe("opus");
+    expect(body.effectiveModel).toBe("opus");
+  });
+
+  test("reports the manager-normalized requested model", async () => {
+    const handler = getSpawnHandler();
+    const body = (await handler({
+      body: { ...SPAWN_BODY, model: "  opus  " },
+    })) as Record<string, unknown>;
+
+    expect((spawnMock.mock.calls[0] as unknown[])[6]).toEqual({
+      model: "  opus  ",
+    });
+    expect(body.requestedModel).toBe("opus");
+  });
+
+  test("a spawn that names no model asks the manager for none", async () => {
+    const handler = getSpawnHandler();
+    await handler({ body: SPAWN_BODY });
+
+    const options = (spawnMock.mock.calls[0] as unknown[])[6] as {
+      model?: string;
+    };
+    const approvalInput = confirmationRequests[0]?.input as Record<
+      string,
+      unknown
+    >;
+    expect(approvalInput).not.toHaveProperty("model");
+    expect(options.model).toBeUndefined();
+  });
+
+  test("a refused model is a warning on an otherwise successful spawn", async () => {
+    spawnResult = {
+      ...DEFAULT_SPAWN_RESULT,
+      modelWarning: "Invalid value for config option model: nope",
+    };
+
+    const handler = getSpawnHandler();
+    const body = (await handler({
+      body: { ...SPAWN_BODY, model: "nope" },
+    })) as Record<string, unknown>;
+
+    expect(body).toEqual({
+      acpSessionId: "acp-route-session",
+      protocolSessionId: "proto-route-session",
+      agent: "claude",
+      requestedModel: "nope",
+      effectiveModel: "opus",
+      modelWarning: "Invalid value for config option model: nope",
+    });
+  });
+});
+
 // ---------------------------------------------------------------------------
 // POST /v1/acp/:id/steer: transparent resume of sessions not in memory
 // ---------------------------------------------------------------------------