fix: query parameters in ws (marimo-team#5381)

This fixes a regression that caused `mo.query_params()` to stop working.
We need to send all query parameters to the backend, not just known
ones, as part of the websocket URL.

Tested by playing with the query parameters smoke test
(`_marimo/smoke_tests`).

Fixes marimo-team#5376.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 2 people

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

@@ -271,6 +271,50 @@ describe("RuntimeManager", () => {
       expect(aiUrl.pathname).toBe("/nested/path/api/ai/completion");
     });
 
+    it("should preserve all window-level query parameters in URLs", () => {
+      // Mock window.location.search with custom query parameters
+      const originalLocation = window.location;
+      Object.defineProperty(window, "location", {
+        value: {
+          ...originalLocation,
+          search: "?custom_param=value123&user_token=abc&theme=dark",
+        },
+        writable: true,
+      });
+
+      const runtime = new RuntimeManager({
+        url: "https://example.com/path?base_param=existing",
+      });
+
+      const wsUrl = runtime.getWsURL("test" as SessionId);
+      const httpUrl = runtime.formatHttpURL(
+        "api/test",
+        new URLSearchParams(),
+        false,
+      );
+
+      // Should preserve base URL query params
+      expect(wsUrl.searchParams.get("base_param")).toBe("existing");
+      expect(httpUrl.searchParams.get("base_param")).toBe("existing");
+
+      // Should preserve all window-level query params (including custom ones)
+      expect(wsUrl.searchParams.get("custom_param")).toBe("value123");
+      expect(wsUrl.searchParams.get("user_token")).toBe("abc");
+      expect(wsUrl.searchParams.get("theme")).toBe("dark");
+      expect(httpUrl.searchParams.get("custom_param")).toBe("value123");
+      expect(httpUrl.searchParams.get("user_token")).toBe("abc");
+      expect(httpUrl.searchParams.get("theme")).toBe("dark");
+
+      // Should also include session_id for WebSocket URLs
+      expect(wsUrl.searchParams.get("session_id")).toBe("test");
+
+      // Restore original location
+      Object.defineProperty(window, "location", {
+        value: originalLocation,
+        writable: true,
+      });
+    });
+
     it("should handle URLs with query parameters and fragments", () => {
       const runtime = new RuntimeManager({
         url: "https://example.com/path?existing=param#fragment",

frontend/src/core/runtime/runtime.ts

@@ -38,7 +38,11 @@ export class RuntimeManager {
   /**
    * The base URL of the runtime.
    */
-  formatHttpURL(path?: string, searchParams?: URLSearchParams): URL {
+  formatHttpURL(
+    path?: string,
+    searchParams?: URLSearchParams,
+    restrictToKnownQueryParams: boolean = true,
+  ): URL {
     if (!path) {
       path = "";
     }
@@ -52,14 +56,14 @@ export class RuntimeManager {
       }
     }
 
-    // Move over window level parameters to the WebSocket URL
-    // if they are "known" query params.
-    for (const lookup in KnownQueryParams) {
-      const key = KnownQueryParams[lookup as keyof typeof KnownQueryParams];
-      const value = currentParams.get(key);
-      if (value !== null) {
-        baseUrl.searchParams.set(key, value);
+    for (const [key, value] of currentParams.entries()) {
+      if (
+        restrictToKnownQueryParams &&
+        !Object.values(KnownQueryParams).includes(key)
+      ) {
+        continue;
       }
+      baseUrl.searchParams.set(key, value);
     }
 
     const cleanPath = baseUrl.pathname.replace(/\/$/, "");
@@ -69,7 +73,13 @@ export class RuntimeManager {
   }
 
   formatWsURL(path: string, searchParams?: URLSearchParams): URL {
-    const url = this.formatHttpURL(path, searchParams);
+    // We don't restrict to known query parameters, since mo.query_params()
+    // can accept arbitrary parameters.
+    const url = this.formatHttpURL(
+      path,
+      searchParams,
+      /* restrictToKnownQueryParams =*/ false,
+    );
     return asWsUrl(url.toString());
   }