@marimo-team/frontend 0.22.5-dev21 → 0.22.5-dev24

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

@@ -86,6 +86,74 @@ describe("RuntimeManager", () => {
     });
   });
 
+  describe("cross-origin auth token in WS URLs", () => {
+    it("should add access_token to WS URL when cross-origin with authToken", () => {
+      // example.com is cross-origin relative to the test environment (localhost)
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getWsURL("s_123" as SessionId);
+
+      expect(url.searchParams.get("access_token")).toBe("my-secret-token");
+      expect(url.searchParams.get("session_id")).toBe("s_123");
+    });
+
+    it("should not add access_token to WS URL when same-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: window.location.origin,
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getWsURL("s_123" as SessionId);
+
+      expect(url.searchParams.get("access_token")).toBeNull();
+    });
+
+    it("should not add access_token when no authToken is configured", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com",
+          lazy: true,
+        },
+        true,
+      );
+      const url = runtime.getWsURL("s_123" as SessionId);
+
+      expect(url.searchParams.get("access_token")).toBeNull();
+    });
+
+    it("should add access_token to all WS URL types when cross-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+
+      const wsUrl = runtime.getWsURL("s_123" as SessionId);
+      const wsSyncUrl = runtime.getWsSyncURL("s_123" as SessionId);
+      const terminalUrl = runtime.getTerminalWsURL();
+
+      expect(wsUrl.searchParams.get("access_token")).toBe("my-secret-token");
+      expect(wsSyncUrl.searchParams.get("access_token")).toBe(
+        "my-secret-token",
+      );
+      expect(terminalUrl.searchParams.get("access_token")).toBe(
+        "my-secret-token",
+      );
+    });
+  });
+
   describe("getWsSyncURL", () => {
     it("should return WebSocket Sync URL", () => {
       const runtime = new RuntimeManager(mockConfig);
@@ -117,12 +185,52 @@ describe("RuntimeManager", () => {
       expect(url.pathname).toBe("/lsp/pylsp");
     });
 
-    it("should return copilot URL", () => {
-      const runtime = new RuntimeManager(mockConfig);
+    it("should return copilot URL without non-auth query params", () => {
+      const runtime = new RuntimeManager({
+        url: "https://example.com?foo=bar&baz=qux",
+        lazy: true,
+      });
       const url = runtime.getLSPURL("copilot");
 
       expect(url.protocol).toBe("wss:");
       expect(url.pathname).toBe("/lsp/copilot");
+      expect(url.searchParams.get("foo")).toBeNull();
+      expect(url.searchParams.get("baz")).toBeNull();
+    });
+
+    it("should preserve access_token on copilot URL when cross-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com?foo=bar",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getLSPURL("copilot");
+
+      expect(url.protocol).toBe("wss:");
+      expect(url.pathname).toBe("/lsp/copilot");
+      expect(url.searchParams.get("access_token")).toBe("my-secret-token");
+      // Other params should be stripped
+      expect(url.searchParams.get("foo")).toBeNull();
+    });
+
+    it("should not have access_token on copilot URL when same-origin", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: window.location.origin,
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getLSPURL("copilot");
+
+      expect(url.protocol).toBe("ws:");
+      expect(url.pathname).toBe("/lsp/copilot");
+      expect(url.searchParams.get("access_token")).toBeNull();
+      expect(url.search).toBe("");
     });
   });
 
@@ -200,6 +308,34 @@ describe("RuntimeManager", () => {
 
       expect(result).toBe(false);
     });
+
+    it("should update config.url on redirect, stripping /health from pathname", async () => {
+      global.fetch = vi.fn().mockResolvedValue({
+        ok: true,
+        redirected: true,
+        url: "https://sandbox.example.com/health?some_value=abc123",
+      });
+
+      const runtime = new RuntimeManager(
+        {
+          ...mockConfig,
+          url: "https://backend.example.com/lazy?some_value=abc123",
+        },
+        true, // lazy — don't call init() in constructor
+      );
+      const result = await runtime.isHealthy();
+
+      expect(result).toBe(true);
+      // Should strip /health from pathname but preserve query params
+      const wsUrl = runtime.getWsURL("s_test" as SessionId);
+      expect(wsUrl.pathname).toBe("/ws");
+      expect(wsUrl.hostname).toBe("sandbox.example.com");
+      expect(wsUrl.searchParams.get("some_value")).toBe("abc123");
+
+      // Clean up side effects
+      document.querySelectorAll("base").forEach((el) => el.remove());
+      global.fetch = vi.fn().mockResolvedValue({ ok: false });
+    });
   });
 
   describe("waitForHealthy", () => {

package/src/core/runtime/runtime.ts

@@ -88,6 +88,15 @@ export class RuntimeManager {
       searchParams,
       /* restrictToKnownQueryParams =*/ false,
     );
+
+    // For cross-origin runtimes, pass the auth token as a query parameter.
+    // WebSocket connections cannot send custom headers (no Authorization
+    // header), and cross-origin cookies are blocked by browsers, so the
+    // access_token query param is the only way to authenticate.
+    if (!this.isSameOrigin && this.config.authToken) {
+      url.searchParams.set(KnownQueryParams.accessToken, this.config.authToken);
+    }
+
     return asWsUrl(url.toString());
   }
 
@@ -143,9 +152,15 @@ export class RuntimeManager {
    */
   getLSPURL(lsp: "pylsp" | "basedpyright" | "copilot" | "ty" | "pyrefly"): URL {
     if (lsp === "copilot") {
-      // For copilot, don't include any query parameters
+      // For copilot, strip all query parameters except the auth token.
+      // Copilot doesn't understand arbitrary query params, but we still
+      // need access_token for cross-origin authentication.
       const url = this.formatWsURL(`/lsp/${lsp}`);
+      const accessToken = url.searchParams.get(KnownQueryParams.accessToken);
       url.search = "";
+      if (accessToken) {
+        url.searchParams.set(KnownQueryParams.accessToken, accessToken);
+      }
       return url;
     }
     return this.formatWsURL(`/lsp/${lsp}`);
@@ -173,9 +188,10 @@ export class RuntimeManager {
       // If there is a redirect, update the URL in the config
       if (response.redirected) {
         Logger.debug(`Runtime redirected to ${response.url}`);
-        // strip /health from the URL
-        const baseUrl = response.url.replace(/\/health$/, "");
-        this.config.url = baseUrl;
+        // strip /health from the URL, using URL parsing to handle query params
+        const redirected = new URL(response.url);
+        redirected.pathname = redirected.pathname.replace(/\/health$/, "");
+        this.config.url = redirected.toString();
       }
 
       const success = response.ok;
@@ -183,7 +199,11 @@ export class RuntimeManager {
         this.setDOMBaseUri(this.config.url);
       }
       return success;
-    } catch {
+    } catch (error) {
+      Logger.error(
+        `Failed to check health: ${error instanceof Error ? error.message : "Unknown error"}`,
+        { cause: error },
+      );
       return false;
     }
   }

package/dist/assets/config-Cgj0Ahvb.js

@@ -1 +0,0 @@
-var w=Object.defineProperty;var S=(t,e,r)=>e in t?w(t,e,{enumerable:!0,configurable:!0,writable:!0,value:r}):t[e]=r;var R=(t,e,r)=>S(t,typeof e!="symbol"?e+"":e,r);import{i as T,l as d,n as g,o as C,p as i,u as y}from"./useEvent-D91BmmQi.js";import{t as D}from"./compiler-runtime-B3qBwwSJ.js";import{_ as c}from"./useEventListener-DGjKht0c.js";import{t as I}from"./utils-8btzWeZg.js";import{r as l}from"./constants-tOPFFcLZ.js";import{t as A}from"./Deferred-DxQeE5uh.js";import{t as f}from"./session-DdnWW30b.js";function m(){return typeof document<"u"&&document.querySelector("marimo-wasm")!==null}const n={NOT_STARTED:"NOT_STARTED",CONNECTING:"CONNECTING",OPEN:"OPEN",CLOSING:"CLOSING",CLOSED:"CLOSED"},_={KERNEL_DISCONNECTED:"KERNEL_DISCONNECTED",ALREADY_RUNNING:"ALREADY_RUNNING",MALFORMED_QUERY:"MALFORMED_QUERY",KERNEL_STARTUP_ERROR:"KERNEL_STARTUP_ERROR"},o=i({state:n.NOT_STARTED});function P(){return C(o,t=>t.state===n.OPEN)}const b=i(t=>t(o).state===n.CONNECTING);i(t=>t(o).state===n.OPEN);const G=i(t=>{let e=t(o);return e.state===n.OPEN||e.state===n.NOT_STARTED}),W=i(t=>t(o).state===n.CLOSED),k=i(t=>t(o).state===n.NOT_STARTED);function H(t){return t===n.CLOSED}function $(t){return t===n.CONNECTING}function M(t){return t===n.OPEN}function v(t){return t===n.CLOSING}function L(t){return t===n.NOT_STARTED}function z(t){return t===n.CLOSED||t===n.CLOSING||t===n.CONNECTING}function F(t){switch(t){case n.CLOSED:return"App disconnected";case n.CONNECTING:return"Connecting to a runtime ...";case n.CLOSING:return"App disconnecting...";case n.OPEN:return"";case n.NOT_STARTED:return"Click to connect to a runtime";default:return""}}var x=class{constructor(t,e=!1){R(this,"initialHealthyCheck",new A);this.config=t,this.lazy=e;try{new URL(this.config.url)}catch(r){throw Error(`Invalid runtime URL: ${this.config.url}. ${r instanceof Error?r.message:"Unknown error"}`,{cause:r})}this.lazy||this.init()}get isLazy(){return this.lazy}get httpURL(){return new URL(this.config.url)}get isSameOrigin(){return this.httpURL.origin===window.location.origin}formatHttpURL(t,e,r=!0){t||(t="");let a=this.httpURL,s=new URLSearchParams(window.location.search);if(e)for(let[h,u]of e.entries())a.searchParams.set(h,u);for(let[h,u]of s.entries())r&&!Object.values(l).includes(h)||a.searchParams.set(h,u);return a.pathname=`${a.pathname.replace(/\/$/,"")}/${t.replace(/^\//,"")}`,a.hash="",a}formatWsURL(t,e){return K(this.formatHttpURL(t,e,!1).toString())}getWsURL(t){let e=new URL(this.config.url),r=new URLSearchParams(e.search);return new URLSearchParams(window.location.search).forEach((a,s)=>{r.has(s)||r.set(s,a)}),r.set(l.sessionId,t),this.formatWsURL("/ws",r)}getWsSyncURL(t){let e=new URL(this.config.url),r=new URLSearchParams(e.search);return new URLSearchParams(window.location.search).forEach((a,s)=>{r.has(s)||r.set(s,a)}),r.set(l.sessionId,t),this.formatWsURL("/ws_sync",r)}getTerminalWsURL(){return this.formatWsURL("/terminal/ws")}getLSPURL(t){if(t==="copilot"){let e=this.formatWsURL(`/lsp/${t}`);return e.search="",e}return this.formatWsURL(`/lsp/${t}`)}getAiURL(t){return this.formatHttpURL(`/api/ai/${t}`)}healthURL(){return this.formatHttpURL("/health")}async isHealthy(){if(m()||I())return!0;try{let t=await fetch(this.healthURL().toString());if(t.redirected){c.debug(`Runtime redirected to ${t.url}`);let r=t.url.replace(/\/health$/,"");this.config.url=r}let e=t.ok;return e&&this.setDOMBaseUri(this.config.url),e}catch{return!1}}setDOMBaseUri(t){t=t.split("?",1)[0],t.endsWith("/")||(t+="/");let e=document.querySelector("base");e?e.setAttribute("href",t):(e=document.createElement("base"),e.setAttribute("href",t),document.head.append(e))}async init(t){c.debug("Initializing runtime...");let e=0;for(;!await this.isHealthy();){if(e>=25){c.error("Failed to connect after 25 retries"),this.initialHealthyCheck.reject(Error("Failed to connect after 25 retries"));return}if(!(t!=null&&t.disableRetryDelay)){let r=Math.min(100*1.2**e,2e3);await new Promise(a=>setTimeout(a,r))}e++}c.debug("Runtime is healthy"),this.initialHealthyCheck.resolve()}async waitForHealthy(){return this.initialHealthyCheck.promise}headers(){let t={"Marimo-Session-Id":f(),"Marimo-Server-Token":this.config.serverToken??"","x-runtime-url":this.httpURL.toString()};return this.config.authToken&&(t.Authorization=`Bearer ${this.config.authToken}`),t}sessionHeaders(){return{"Marimo-Session-Id":f()}}};function K(t){if(!t.startsWith("http")){c.warn(`URL must start with http: ${t}`);let e=new URL(t);return e.protocol="ws",e}return new URL(t.replace(/^http/,"ws"))}var Y=D();function j(){let t=new URL(document.baseURI);return t.search="",t.hash="",t.toString()}const N={lazy:!0,url:j()},E=i(N);var U=i(t=>{let e=t(E);return new x(e,e.lazy)});function p(){return y(U)}function B(){let t=(0,Y.c)(4),e=p(),[r,a]=d(o),s;return t[0]!==r.state||t[1]!==e||t[2]!==a?(s=async()=>{L(r.state)?(a({state:n.CONNECTING}),await e.init()):c.log("Runtime already started or starting...")},t[0]=r.state,t[1]=e,t[2]=a,t[3]=s):s=t[3],g(s)}function O(){return T.get(U)}function q(t){if(t.startsWith("http"))return new URL(t);let e=O().httpURL.toString();return e.startsWith("blob:")&&(e=e.replace("blob:","")),new URL(t,e)}export{m as S,b as _,B as a,_ as b,H as c,$ as d,z as f,W as g,o as h,E as i,v as l,G as m,q as n,p as o,L as p,O as r,F as s,N as t,M as u,k as v,n as x,P as y};