Fix/acp agent proxy remote

frontend/src/components/chat/acp/__tests__/state.test.ts

@@ -2,6 +2,11 @@
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import * as shortcuts from "@/core/hotkeys/shortcuts";
+import {
+  DEFAULT_RUNTIME_CONFIG,
+  runtimeConfigAtom,
+} from "@/core/runtime/config";
+import { store } from "@/core/state/jotai";
 import {
   type AgentSession,
   type AgentSessionState,
@@ -700,69 +705,51 @@ describe("state utility functions", () => {
   });
 
   describe("getAgentWebSocketUrl", () => {
-    const originalLocation = window.location;
-
     afterEach(() => {
-      // Restore original window.location
-      Object.defineProperty(window, "location", {
-        value: originalLocation,
-        writable: true,
-      });
-    });
-
-    it("should return ws:// URL with localhost for http protocol", () => {
-      Object.defineProperty(window, "location", {
-        value: {
-          hostname: "localhost",
-          protocol: "http:",
-        },
-        writable: true,
-      });
+      store.set(runtimeConfigAtom, DEFAULT_RUNTIME_CONFIG);
+    });
+
+    const setRuntimeUrl = (url: string) => {
+      store.set(runtimeConfigAtom, { url, lazy: true });
+    };
+
+    it("should return a ws:// URL for http runtimes", () => {
+      setRuntimeUrl("http://localhost:2718/");
 
       expect(getAgentWebSocketUrl("claude")).toMatchInlineSnapshot(
-        `"ws://localhost:3017/message"`,
+        `"ws://localhost:2718/acp/claude"`,
       );
     });
 
-    it("should return wss:// URL for https protocol", () => {
-      Object.defineProperty(window, "location", {
-        value: {
-          hostname: "example.com",
-          protocol: "https:",
-        },
-        writable: true,
-      });
+    it("should return a wss:// URL for https runtimes", () => {
+      setRuntimeUrl("https://example.com/");
 
       expect(getAgentWebSocketUrl("claude")).toMatchInlineSnapshot(
-        `"wss://example.com:3017/message"`,
+        `"wss://example.com/acp/claude"`,
       );
     });
 
-    it("should work with IP addresses", () => {
-      Object.defineProperty(window, "location", {
-        value: {
-          hostname: "192.168.1.100",
-          protocol: "http:",
-        },
-        writable: true,
-      });
+    it("should work with IP addresses and ports", () => {
+      setRuntimeUrl("http://192.168.1.100:8080/");
 
       expect(getAgentWebSocketUrl("claude")).toMatchInlineSnapshot(
-        `"ws://192.168.1.100:3017/message"`,
+        `"ws://192.168.1.100:8080/acp/claude"`,
       );
     });
 
     it("should work with remote hostnames", () => {
-      Object.defineProperty(window, "location", {
-        value: {
-          hostname: "marimo.example.com",
-          protocol: "https:",
-        },
-        writable: true,
-      });
+      setRuntimeUrl("https://marimo.example.com/");
 
       expect(getAgentWebSocketUrl("gemini")).toMatchInlineSnapshot(
-        `"wss://marimo.example.com:3019/message"`,
+        `"wss://marimo.example.com/acp/gemini"`,
+      );
+    });
+
+    it("should preserve a base path prefix", () => {
+      setRuntimeUrl("https://example.com/notebooks/abc/");
+
+      expect(getAgentWebSocketUrl("codex")).toMatchInlineSnapshot(
+        `"wss://example.com/notebooks/abc/acp/codex"`,
       );
     });
   });

frontend/src/components/chat/acp/state.ts

@@ -3,6 +3,7 @@
 import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import { isPlatformWindows } from "@/core/hotkeys/shortcuts";
+import { getRuntimeManager } from "@/core/runtime/config";
 import { jotaiJsonStorage } from "@/utils/storage/jotai";
 import { capitalize } from "@/utils/strings";
 import type { TypedString } from "@/utils/typed";
@@ -238,15 +239,19 @@ export function getAgentDisplayName(agentId: ExternalAgentId): string {
 }
 
 export function getAgentWebSocketUrl(agentId: ExternalAgentId): string {
-  const port = AGENT_CONFIG[agentId].port;
-  // Use the current page's hostname so the agent is reachable when
-  // marimo is accessed remotely (e.g. via direct IP or reverse proxy).
-  const hostname = window.location.hostname;
-  const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
-  return `${protocol}//${hostname}:${port}/message` as const;
+  // Connect through the marimo server (`<base>/acp/<agentId>`) rather than
+  // directly to the agent's port, which isn't reachable from the browser when
+  // marimo is served behind a reverse proxy.
+  return getRuntimeManager().getAgentWsURL(agentId).toString();
 }
 
 interface AgentConfig {
+  /**
+   * Port the agent's websocket bridge listens on. Only used to render the
+   * connection command shown to the user -- the browser reaches the agent
+   * through the marimo server, not this port. Must stay in sync with
+   * `ACP_AGENT_PORTS` in `marimo/_server/main.py`.
+   */
   port: number;
   command: string;
   sessionSupport: SessionSupportType;

frontend/src/core/runtime/__tests__/runtime.test.ts

@@ -343,6 +343,70 @@ describe("RuntimeManager", () => {
     });
   });
 
+  describe("getAgentWsURL", () => {
+    it("should return the agent URL proxied through the server", () => {
+      const runtime = new RuntimeManager(mockConfig);
+      const url = runtime.getAgentWsURL("claude");
+
+      expect(url.protocol).toBe("wss:");
+      expect(url.pathname).toBe("/acp/claude");
+      expect(url.host).toBe("example.com");
+    });
+
+    it("should respect a base path prefix", () => {
+      const runtime = new RuntimeManager({
+        url: "http://example.com/prefix/",
+        lazy: true,
+      });
+      const url = runtime.getAgentWsURL("gemini");
+
+      expect(url.toString()).toBe("ws://example.com/prefix/acp/gemini");
+    });
+
+    it("should strip non-auth query params", () => {
+      const runtime = new RuntimeManager({
+        url: "https://example.com?foo=bar&baz=qux",
+        lazy: true,
+      });
+      const url = runtime.getAgentWsURL("codex");
+
+      expect(url.pathname).toBe("/acp/codex");
+      expect(url.search).toBe("");
+    });
+
+    it("should preserve access_token when cross-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com?foo=bar",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getAgentWsURL("claude");
+
+      expect(url.pathname).toBe("/acp/claude");
+      expect(url.searchParams.get("access_token")).toBe("my-secret-token");
+      expect(url.searchParams.get("foo")).toBeNull();
+    });
+
+    it("should not have access_token when same-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: window.location.origin,
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getAgentWsURL("claude");
+
+      expect(url.protocol).toBe("ws:");
+      expect(url.pathname).toBe("/acp/claude");
+      expect(url.search).toBe("");
+    });
+  });
+
   describe("getAiURL", () => {
     it("should return AI completion URL", () => {
       const runtime = new RuntimeManager(mockConfig);

frontend/src/core/runtime/runtime.ts

@@ -206,6 +206,19 @@ export class RuntimeManager {
     return this.formatWsURL(`/lsp/${lsp}`);
   }
 
+  /**
+   * The WebSocket URL of an external ACP agent.
+   */
+  getAgentWsURL(agentId: string): URL {
+    const url = this.formatWsURL(`/acp/${encodeURIComponent(agentId)}`);
+    const accessToken = url.searchParams.get(KnownQueryParams.accessToken);
+    url.search = "";
+    if (accessToken) {
+      url.searchParams.set(KnownQueryParams.accessToken, accessToken);
+    }
+    return url;
+  }
+
   getAiURL(path: "completion" | "chat"): URL {
     return this.formatHttpURL({ path: `/api/ai/${path}` });
   }

marimo/_server/main.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Final
 
 from starlette.applications import Starlette
 from starlette.exceptions import HTTPException
@@ -23,6 +23,7 @@
     ProxyMiddleware,
     SkewProtectionMiddleware,
     TimeoutMiddleware,
+    create_proxy_error_handler,
 )
 from marimo._server.api.router import build_routes
 from marimo._server.errors import handle_error
@@ -46,6 +47,17 @@ class LspPorts:
     copilot: int | None
 
 
+# This must stay in sync with AGENT_CONFIG in
+# frontend/src/components/chat/acp/state.ts
+ACP_AGENT_PORTS: Final[dict[str, int]] = {
+    "claude": 3017,
+    "gemini": 3019,
+    "codex": 3021,
+    "opencode": 3023,
+    "cursor": 3025,
+}
+
+
 # Create app
 def create_starlette_app(
     *,
@@ -56,6 +68,7 @@ def create_starlette_app(
     enable_auth: bool = True,
     allow_origins: tuple[str, ...] | None = None,
     lsp_servers: list[LspServer] | None = None,
+    enable_acp_proxy: bool = False,
     skew_protection: bool = True,
     timeout: float | None = None,
 ) -> Starlette:
@@ -109,6 +122,11 @@ def create_starlette_app(
             )
         )
 
+    if enable_acp_proxy:
+        final_middlewares.extend(
+            _create_acp_proxy_middleware(base_url=base_url)
+        )
+
     if middleware:
         final_middlewares.extend(middleware)
 
@@ -145,3 +163,31 @@ def _create_lsps_proxy_middleware(
         )
         for server in servers
     )
+
+
+def _create_acp_proxy_middleware(base_url: str) -> Iterator[Middleware]:
+    """Proxy the external ACP agents' websockets through the marimo server.
+
+    The agents listen on fixed localhost ports, which aren't reachable from
+    the browser when marimo is served behind a reverse proxy. Proxying them
+    under `<base_url>/acp/<agent_id>` keeps the connection same-origin, so it
+    only needs the port marimo is already served on.
+    """
+    return (
+        Middleware(
+            ProxyMiddleware,
+            proxy_path=f"{base_url}/acp/{agent_id}",
+            target_url=f"http://127.0.0.1:{port}",
+            path_rewrite=_rewrite_acp_path,
+            connection_error_handler=create_proxy_error_handler(
+                f"The {agent_id} agent is not running. "
+                "Start it with the command shown in the agent panel."
+            ),
+        )
+        for agent_id, port in ACP_AGENT_PORTS.items()
+    )
+
+
+def _rewrite_acp_path(_path: str) -> str:
+    """ACP agents serve a single endpoint, regardless of the proxied path."""
+    return "/message"

marimo/_server/start.py

@@ -354,6 +354,7 @@ def start(
         lsp_servers=list(lsp_composite_server.servers.values())
         if lsp_composite_server is not None
         else None,
+        enable_acp_proxy=mode == SessionMode.EDIT,
         skew_protection=skew_protection,
         timeout=timeout,
     )