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

[jcampabadal-db]

Problem

On a Databricks host that serves multiple workspaces, the CLI issues an account-audience OAuth token. The workspace APIs reject that token at the authz layer unless the request carries the workspace disambiguator o=<workspace_id>. Without it they respond 303 to /login, and ucode parses the returned sign-in HTML as JSON.

ucode configure therefore aborts before writing any agent config:

✔ Databricks authentication complete
ERROR Databricks Unity AI Gateway probe failed on this workspace
      (response was not valid JSON (Expecting value)).
      See https://docs.databricks.com/aws/en/ai-gateway/overview-beta

Patching only the probe moves the failure one step later, to model discovery:

✔ Unity AI Gateway detected
ERROR No coding agents are available on this workspace.
• Claude models (needed for: claude, opencode, copilot, pi): response was not valid JSON (Expecting value)
• Gemini models (needed for: gemini, opencode, pi): no models returned

The affected requests all flow through _http_get_json / _http_post_json:

• GET /api/ai-gateway/v2/endpoints — the ensure_ai_gateway_v2 probe
• GET /api/2.1/unity-catalog/model-services — model discovery
• GET /ai-gateway/anthropic/v1/models
• GET /api/2.0/sql/warehouses, the UC catalog/schema/function listings, etc.

The workspace is fully gateway-enabled and the token is valid — the request just can't be attributed to a workspace.

Observed against a multi-workspace staging host:

$ curl -si -H "Authorization: Bearer $TOKEN" \
    "https://$HOST/api/ai-gateway/v2/endpoints?page_size=1"
HTTP/2 303
x-databricks-apiproxy-response-code-details: ext_authz_denied
location: /login?next_url=%2Fapi%2Fai-gateway%2Fv2%2Fendpoints%3Fpage_size%3D1

$ curl -so /dev/null -w '%{http_code}\n' -H "Authorization: Bearer $TOKEN" \
    "https://$HOST/api/ai-gateway/v2/endpoints?page_size=1&o=$WORKSPACE_ID"
200

An X-Databricks-Org-Id: <workspace_id> header works equivalently.

Fix

Centralize it in the two HTTP helpers rather than at the ~30 individual call sites that each build f"https://{hostname}/...".

• _workspace_id_for_hostname(hostname) reads ~/.databrickscfg (honoring DATABRICKS_CONFIG_FILE), normalizes each profile's host to a hostname, and returns the workspace_id of the first match. Blank values and the literal none are skipped, so a workspace-scoped profile wins over an account-only profile on the same host — a common layout, since an account-level databricks auth login writes workspace_id = none.
• _with_workspace_disambiguator(url) appends ?o=<id> or &o=<id> as appropriate.
• _http_get_json and _http_post_json each call it on entry.

It is a no-op wherever it isn't needed: a single-workspace host ignores the extra param, an account-only profile resolves to None, a URL already carrying o= is untouched, and a malformed or missing config degrades to current behavior instead of raising. No authentication logic changes — same token, same headers, same audience.

Verification

After the change, against the same host:

✔ Databricks authentication complete
✔ Unity AI Gateway detected
✔ Databricks AI Tools installed
╭───────── Configuration Complete ─────────╮
│ Claude Code: configured (Provider: Databricks) │
╰──────────────────────────────────────────╯
✔ Claude Code is working

ucode claude -p "..." then runs headless against system.ai.claude-opus-5 through the gateway.

Tests

8 focused tests in tests/test_databricks.py (TestWorkspaceIdForHostname, TestWithWorkspaceDisambiguator): resolving past an account-only profile on the same host, a second distinct host, unknown host, missing config file, append with and without an existing query string, and the no-op when o= is already present.

uv run pytest -q                          # 1055 passed, 40 skipped
uv run ruff check .                       # All checks passed!
uv run ruff format --check src/ tests/    # 55 files already formatted

Note for reviewers

ci.yml triggers on pull_request, and GitHub withholds repo secrets from fork-based PRs, so the steps needing UCODE_TEST_WORKSPACE / DATABRICKS_BEARER won't have credentials here. Unit tests and lint run normally; the e2e job may need a workflow_dispatch re-run from a branch in this repo.

[jcampabadal-db]

The following monkey patch also fixes this error:

ERROR Databricks Unity AI Gateway probe failed on this workspace (response was not valid JSON (Expecting value)

python3 - <<'EOF'
import glob, pathlib, subprocess
tools = subprocess.check_output(["uv","tool","dir"], text=True).strip()
hits = glob.glob(f"{tools}/ucode/lib/python3.*/site-packages/ucode/databricks.py")
p = pathlib.Path(hits[0]); src = p.read_text()

HELPERS = '''
def _workspace_id_for_hostname(hostname):
    import configparser
    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 Exception:
        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 Exception:
            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):
    parsed = urlparse(url)
    if not parsed.hostname or "o=" in (parsed.query or ""):
        return url
    wid = _workspace_id_for_hostname(parsed.hostname)
    if not wid:
        return url
    return f"{url}{'&' if parsed.query else '?'}o={wid}"

def _http_get_json('''

anchor = "\ndef _http_get_json("
assert anchor in src, "anchor not found - already patched or version changed"
src = src.replace(anchor, HELPERS, 1)

for fn in ("_http_get_json", "_http_post_json"):
    marker = {"_http_get_json": "    request = urllib_request.Request(\n        url,\n        headers={\"Authorization\": f\"Bearer {token}\", \"Accept\": \"application/json\"},",
              "_http_post_json": "    body_bytes = json.dumps(payload).encode(\"utf-8\")"}[fn]
    assert marker in src, f"{fn} marker missing"
    src = src.replace(marker, "    url = _with_workspace_disambiguator(url)\n" + marker, 1)

p.with_suffix(".py.orig-bak").write_text(pathlib.Path(hits[0]).read_text())
p.write_text(src)
print("patched", p)
EOF