From d4aa1ad2b5793eb8cb85cd9606d4f5cbcc0ddb39 Mon Sep 17 00:00:00 2001 From: Julian Campabadal Date: Fri, 31 Jul 2026 16:42:28 +0000 Subject: [PATCH] Fix multi-workspace host auth: append o= to discovery calls On a Databricks host that serves multiple workspaces, an account-audience OAuth token is rejected at the authz layer unless the request carries the workspace disambiguator. Without it, the AI Gateway v2 probe (`GET /api/ai-gateway/v2/endpoints`) and the model-discovery calls (`/api/2.1/unity-catalog/model-services`, `/ai-gateway/anthropic/v1/models`) 303-redirect to `/login`, and ucode reads the returned login HTML/empty body as "response was not valid JSON". `ucode configure` then aborts with "Unity AI Gateway probe failed" / "No coding agents are available", even though the workspace is fully gateway-enabled. Centralize the fix in the HTTP helpers: `_with_workspace_disambiguator` appends `o=` (resolved from the matching ~/.databrickscfg profile) to any workspace URL that lacks it. Single-workspace hosts ignore the extra param, so the change is a no-op there. Verified end-to-end against a multi-workspace staging host: the probe and model discovery return HTTP 200 with the param and 303 to /login without it, and `ucode configure --agents claude` completes and validates. Co-authored-by: Isaac --- src/ucode/databricks.py | 48 +++++++++++++++++++++++++++ tests/test_databricks.py | 70 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 1d32f31..f4423ac 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -204,6 +204,52 @@ 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=`` 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]: @@ -211,6 +257,7 @@ def _http_get_json( 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, diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7e1a73a..d23326c 100644 --- a/tests/test_databricks.py +++ b/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