From bf8e1533cd82b8e0e32571e0050e05993eb02ed0 Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Sat, 12 Sep 2026 23:38:12 +0000 Subject: [PATCH] [skills] Record the workspace org id in downloaded-skill attribution Capture the numeric workspace (org) id and store it as `workspace_id` on each downloaded-skill record, giving attribution a rename-proof workspace identifier alongside the workspace URL. Databricks stamps every authenticated response with an `X-Databricks-Org-Id` header, so the download's existing API calls already carry it; the HTTP GET helpers now capture it into a session cache (keyed by hostname, mirroring the model-service listing caches) with no extra request. A record omits the field when no response has revealed the id yet. Co-authored-by: Isaac --- src/ucode/databricks.py | 27 +++++++++++++++++++++++++++ src/ucode/skills_download.py | 3 +++ src/ucode/skills_state.py | 2 ++ tests/conftest.py | 1 + tests/test_databricks.py | 32 ++++++++++++++++++++++++++++++++ tests/test_skills_download.py | 13 +++++++++++++ 6 files changed, 78 insertions(+) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index aa998c53..f0566d71 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -26,6 +26,7 @@ ) from dataclasses import dataclass from decimal import Decimal, InvalidOperation +from email.message import Message from enum import Enum from pathlib import Path from typing import Literal, NamedTuple, NoReturn, cast, overload @@ -260,6 +261,30 @@ def _http_get_retry_delay(retry_after: str | None, retry_index: int) -> float: return backoff + random.uniform(0, min(backoff * 0.25, 0.5)) +# Databricks stamps every authenticated API response with the caller's numeric workspace (org) id in +# this header, so any call ucode already makes reveals it with no dedicated lookup. Captured by +# hostname as responses go by; session-only, like the listing caches below. +_ORG_ID_HEADER = "X-Databricks-Org-Id" +_WORKSPACE_ORG_IDS: dict[str, str] = {} + + +def _capture_org_id(url: str, headers: Message | None) -> None: + org_id = headers.get(_ORG_ID_HEADER) if headers is not None else None + hostname = urlparse(url).hostname + if org_id and hostname: + _WORKSPACE_ORG_IDS[hostname] = org_id + + +def workspace_org_id(workspace: str) -> str | None: + """The numeric workspace (org) id for ``workspace``, or None if no response has revealed it yet.""" + return _WORKSPACE_ORG_IDS.get(workspace_hostname(workspace)) + + +def clear_workspace_org_id_cache() -> None: + """Forget captured workspace org ids (used by tests, and after a workspace switch).""" + _WORKSPACE_ORG_IDS.clear() + + def _http_get_json( url: str, token: str, @@ -286,6 +311,7 @@ def _http_get_json( try: with urllib_request.urlopen(request, timeout=timeout) as response: body = response.read().decode("utf-8") + _capture_org_id(url, getattr(response, "headers", None)) _debug(f"GET {url}", f"HTTP 200, {len(body)} bytes") if _debug_enabled(): _debug("body", body[:4000]) @@ -434,6 +460,7 @@ def _http_get_bytes(url: str, token: str, *, timeout: int = 10) -> tuple[bytes | try: with urllib_request.urlopen(request, timeout=timeout) as response: body = response.read() + _capture_org_id(url, getattr(response, "headers", None)) _debug(f"GET {url}", f"HTTP 200, {len(body)} bytes") return body, None except urllib_error.HTTPError as exc: diff --git a/src/ucode/skills_download.py b/src/ucode/skills_download.py index bfb26f42..299c1bd6 100644 --- a/src/ucode/skills_download.py +++ b/src/ucode/skills_download.py @@ -17,6 +17,7 @@ get_databricks_token, walk_catalog_schemas, workspace_hostname, + workspace_org_id, ) from ucode.mcp import register_schemaless_skills_connection, setup_mcp_clients from ucode.skills_state import ( @@ -323,6 +324,7 @@ def _skill_installs( """Attribution records for ``refs`` written into ``roots`` (see ``skills_state``).""" base = path or str(Path.home()) scope = "project" if path else "user" + org_id = workspace_org_id(workspace) return [ SkillInstall( fqn=ref.fqn, @@ -332,6 +334,7 @@ def _skill_installs( base=base, dirs=tuple(str(root / ref.bundle_name) for root in roots), metastore_id=ref.metastore_id, + workspace_id=org_id, skill_id=ref.skill_id, uc_update_time=ref.uc_update_time, ) diff --git a/src/ucode/skills_state.py b/src/ucode/skills_state.py index fba8f07b..a2faadfd 100644 --- a/src/ucode/skills_state.py +++ b/src/ucode/skills_state.py @@ -38,6 +38,7 @@ class SkillInstall: base: str dirs: tuple[str, ...] metastore_id: str | None = None + workspace_id: str | None = None skill_id: str | None = None uc_update_time: str | None = None @@ -111,6 +112,7 @@ def _to_record(install: SkillInstall) -> dict: "bundle_name": install.bundle_name, "metastore_id": install.metastore_id, "workspace": install.workspace, + "workspace_id": install.workspace_id, "scope": install.scope, "base": install.base, "dirs": list(install.dirs), diff --git a/tests/conftest.py b/tests/conftest.py index 64b94eaf..a1a7419f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -65,6 +65,7 @@ def reject_privileged_write(path, _desired_text): # The model-services listing is memoized for the life of the process, so without this a cached # result would leak into the next test and make a stubbed listing look like it was never called. databricks_mod.clear_model_services_cache() + databricks_mod.clear_workspace_org_id_cache() def _workspace() -> str: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 9a727989..f3e7ba2b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -183,6 +183,38 @@ def test_invalid_url_raises(self): workspace_hostname("") +class _FakeResponseWithHeaders(_FakeResponse): + def __init__(self, payload: dict, headers: dict): + super().__init__(payload) + self.headers = headers + + +class TestWorkspaceOrgId: + def _stub_response(self, monkeypatch, headers: dict) -> None: + monkeypatch.setattr( + db_mod.urllib_request, + "urlopen", + lambda request, timeout=None: _FakeResponseWithHeaders({"ok": True}, headers), + ) + + def test_captures_org_id_header_from_get(self, monkeypatch): + self._stub_response(monkeypatch, {"X-Databricks-Org-Id": "1234567890"}) + + db_mod._http_get_json(f"{WS}/api/2.1/unity-catalog/skills", "token") + + assert db_mod.workspace_org_id(WS) == "1234567890" + + def test_absent_until_a_response_reveals_it(self): + assert db_mod.workspace_org_id(WS) is None + + def test_missing_header_leaves_it_absent(self, monkeypatch): + self._stub_response(monkeypatch, {}) + + db_mod._http_get_json(f"{WS}/api/x", "token") + + assert db_mod.workspace_org_id(WS) is None + + class TestBuildDatabricksCliEnv: def test_sets_databricks_host(self): env = build_databricks_cli_env(WS) diff --git a/tests/test_skills_download.py b/tests/test_skills_download.py index dd20ae0b..404af58d 100644 --- a/tests/test_skills_download.py +++ b/tests/test_skills_download.py @@ -740,6 +740,19 @@ def test_records_downloaded_skills(self, tmp_path, monkeypatch): assert record["fqn"] == "main.default.triage" assert record["scope"] == "project" assert record["base"] == str(tmp_path) + assert "workspace_id" not in record + + def test_records_workspace_id_when_known(self, tmp_path, monkeypatch): + monkeypatch.setattr(sd, "get_skill", lambda ws, tok, fqn: ref(fqn.rsplit(".", 1)[-1])) + monkeypatch.setattr( + sd, "fetch_skill_bundle", lambda ws, tok, c, s, leaf: ({"SKILL.md": b"x"}, None) + ) + monkeypatch.setattr(sd, "workspace_org_id", lambda ws: "org-42") + + sd.download_selected_skills(WS, "token", ["main.default.triage"], str(tmp_path)) + + record = skills_state.attribution_for_dir(tmp_path / ".claude/skills/triage") + assert record["workspace_id"] == "org-42" class TestSkillRefMetadata: