[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 <no-reply@databricks.com>

Author: xsh310

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:

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 (
@@ -333,6 +334,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,
@@ -342,6 +344,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,
         )

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
 
@@ -121,6 +122,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),

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:

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)

tests/test_skills_download.py

@@ -766,6 +766,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: