Fix multi-workspace host auth: append o=<workspace_id> to discovery calls

src/ucode/databricks.py

@@ -204,13 +204,60 @@ def _log_auth_diagnostics() -> None:
         _debug(f"databrickscfg ({cfg_path})", f"read error: {exc}")
 
 
+def _workspace_id_for_hostname(hostname: str) -> str | None:
+    """Return the ``workspace_id`` for the ~/.databrickscfg profile matching a host.
+
+    Databricks CLI writes a numeric ``workspace_id`` for workspace-scoped
+    profiles. It is absent (or the literal ``none``) for account-scoped
+    profiles, in which case there is nothing to disambiguate with."""
+    cfg_path = Path(os.environ.get("DATABRICKS_CONFIG_FILE") or "~/.databrickscfg").expanduser()
+    parser = configparser.ConfigParser(default_section="@ucode-no-defaults@", interpolation=None)
+    try:
+        if not parser.read(cfg_path, encoding="utf-8"):
+            return None
+    except (configparser.Error, OSError):
+        return None
+    for section in parser.sections():
+        host = (parser.get(section, "host", fallback="") or "").strip()
+        if not host:
+            continue
+        try:
+            if urlparse(normalize_workspace_url(host)).hostname != hostname:
+                continue
+        except RuntimeError:
+            continue
+        wid = (parser.get(section, "workspace_id", fallback="") or "").strip()
+        if wid and wid.lower() != "none":
+            return wid
+    return None
+
+
+def _with_workspace_disambiguator(url: str) -> str:
+    """Append ``o=<workspace_id>`` to a workspace API URL when needed.
+
+    A host that serves multiple workspaces rejects an account-audience OAuth
+    token unless the request carries the workspace disambiguator: without it
+    the gateway/UC APIs 303-redirect to ``/login`` and discovery reads the
+    login page as "response was not valid JSON". A single-workspace host
+    ignores the extra param, so this is always safe."""
+    parsed = urlparse(url)
+    if not parsed.hostname or "o=" in (parsed.query or ""):
+        return url
+    workspace_id = _workspace_id_for_hostname(parsed.hostname)
+    if not workspace_id:
+        return url
+    separator = "&" if parsed.query else "?"
+    return f"{url}{separator}o={workspace_id}"
+
+
 def _http_get_json(
     url: str, token: str, *, timeout: int = 10
 ) -> tuple[dict | list | None, str | None]:
     """GET a JSON endpoint. Returns (payload, None) on success, (None, reason) on failure.
 
     Honors UCODE_DEBUG=1 to append status + truncated body to ~/.ucode/debug.log.
     """
+    url = _with_workspace_disambiguator(url)
     request = urllib_request.Request(
         url,
         headers={"Authorization": f"Bearer {token}", "Accept": "application/json"},
@@ -257,6 +304,7 @@ def _http_post_json(
 ) -> tuple[dict | list | None, str | None]:
     """POST a JSON body to an endpoint. Returns (payload, None) on success,
     (None, reason) on failure. Mirrors `_http_get_json`."""
+    url = _with_workspace_disambiguator(url)
     body_bytes = json.dumps(payload).encode("utf-8")
     request = urllib_request.Request(
         url,

tests/test_databricks.py

@@ -16,6 +16,8 @@
     _run_databricks_cli_installer,
     _scrub_databrickscfg,
     _scrub_json,
+    _with_workspace_disambiguator,
+    _workspace_id_for_hostname,
     build_auth_shell_command,
     build_auth_token_argv,
     build_databricks_cli_env,
@@ -1940,3 +1942,71 @@ def test_failure_surfaces_cli_stderr(self, monkeypatch):
         install_ai_tools(["copilot"])
         assert len(warnings) == 1
         assert "copilot: cli-not-on-path: could not resolve copilot" in warnings[0]
+
+
+DBCFG = """\
+[account-only]
+host = https://example.databricks.com
+workspace_id = none
+
+[workspace-scoped]
+host = https://example.databricks.com
+workspace_id = 1234567890123456
+
+[other]
+host = https://other.databricks.com
+workspace_id = 42
+"""
+
+
+class TestWorkspaceIdForHostname:
+    def _cfg(self, tmp_path, monkeypatch, contents=DBCFG):
+        cfg = tmp_path / "databrickscfg"
+        cfg.write_text(contents)
+        monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
+        return cfg
+
+    def test_resolves_workspace_id_ignoring_account_only_profile(self, tmp_path, monkeypatch):
+        self._cfg(tmp_path, monkeypatch)
+        # The account-only profile (workspace_id = none) for the same host must
+        # be skipped in favor of the workspace-scoped one.
+        assert _workspace_id_for_hostname("example.databricks.com") == "1234567890123456"
+
+    def test_matches_other_host(self, tmp_path, monkeypatch):
+        self._cfg(tmp_path, monkeypatch)
+        assert _workspace_id_for_hostname("other.databricks.com") == "42"
+
+    def test_unknown_host_returns_none(self, tmp_path, monkeypatch):
+        self._cfg(tmp_path, monkeypatch)
+        assert _workspace_id_for_hostname("unknown.databricks.com") is None
+
+    def test_missing_config_returns_none(self, tmp_path, monkeypatch):
+        monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(tmp_path / "does-not-exist"))
+        assert _workspace_id_for_hostname("example.databricks.com") is None
+
+
+class TestWithWorkspaceDisambiguator:
+    def _cfg(self, tmp_path, monkeypatch, contents=DBCFG):
+        cfg = tmp_path / "databrickscfg"
+        cfg.write_text(contents)
+        monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
+
+    def test_appends_o_param_when_no_query(self, tmp_path, monkeypatch):
+        self._cfg(tmp_path, monkeypatch)
+        url = "https://example.databricks.com/api/ai-gateway/v2/endpoints"
+        assert _with_workspace_disambiguator(url) == url + "?o=1234567890123456"
+
+    def test_appends_o_param_preserving_existing_query(self, tmp_path, monkeypatch):
+        self._cfg(tmp_path, monkeypatch)
+        url = "https://example.databricks.com/api/ai-gateway/v2/endpoints?page_size=1"
+        assert _with_workspace_disambiguator(url) == url + "&o=1234567890123456"
+
+    def test_leaves_url_untouched_when_o_already_present(self, tmp_path, monkeypatch):
+        self._cfg(tmp_path, monkeypatch)
+        url = "https://example.databricks.com/api/ai-gateway/v2/endpoints?o=999"
+        assert _with_workspace_disambiguator(url) == url
+
+    def test_leaves_url_untouched_when_host_unknown(self, tmp_path, monkeypatch):
+        self._cfg(tmp_path, monkeypatch)
+        url = "https://unknown.databricks.com/api/ai-gateway/v2/endpoints"
+        assert _with_workspace_disambiguator(url) == url