Speed up ucode setup and ucode apply

src/ucode/databricks.py

@@ -1328,10 +1328,18 @@ def build_auth_shell_command(
 # part after the prefix is exactly the model string agents send (no
 # `databricks-` infix — that only appears on the inner destination name).
 _MODEL_SERVICE_NAME_PREFIX = "model-services/"
-# The metastore-scope listing returns services from EVERY schema (e.g.
-# `main.user.foo`, `temp.*`, internal DLT schemas). We only want the
-# Databricks-managed foundation models under `system.ai`.
+# The listing can return services from EVERY schema (e.g. `main.user.foo`,
+# `temp.*`, internal DLT schemas). We only want the Databricks-managed
+# foundation models under `system.ai`.
 _MODEL_SERVICE_REQUIRED_PREFIX = "system.ai."
+# Scope the listing to the `system.ai` schema via the `parent` query param
+# (`schemas/{catalog}.{schema}`). Without it the endpoint walks the ENTIRE
+# metastore — hundreds of unrelated services across dozens of ~2s pages, then
+# discards all but `system.ai.*` client-side (a ~50s walk on a busy workspace).
+# Parent-scoped, the same set comes back in a single page (~1s). The endpoint
+# ignores the other filters (`catalog_name`/`schema_name`/`filter`), so `parent`
+# is the only server-side narrowing that works.
+_MODEL_SERVICE_PARENT_SCHEMA = "schemas/system.ai"
 
 # Supported OSS chat families, matched by name substring. Add an entry to
 # support a new family.
@@ -1473,11 +1481,14 @@ def list_model_services(
 ) -> tuple[list[str], str | None]:
     """List all `system.ai.*` model ids via the UC model-services API.
 
-    Pages through ``/api/2.1/unity-catalog/model-services`` (metastore scope)
-    with a bounded ``page_size`` (the endpoint 499s without one) and returns the
-    de-duplicated, sorted list of ``system.ai.<model-name>`` ids. Returns
-    (ids, reason); reason is None on success, otherwise it describes why the
-    list is empty (HTTP/network error or no services).
+    Pages through ``/api/2.1/unity-catalog/model-services`` scoped to the
+    ``system.ai`` schema (``parent=schemas/system.ai``) with a bounded
+    ``page_size`` (the endpoint 499s without one) and returns the de-duplicated,
+    sorted list of ``system.ai.<model-name>`` ids. Returns (ids, reason); reason
+    is None on success, otherwise it describes why the list is empty (HTTP/network
+    error or no services). Scoping matters: the unscoped metastore listing walks
+    every schema across dozens of ~2s pages (~50s on a busy workspace) only to
+    keep the same ``system.ai.*`` subset — see ``_MODEL_SERVICE_PARENT_SCHEMA``.
 
     A successful result is memoized per workspace for the life of the process; pass
     ``use_cache=False`` to force a fresh walk.
@@ -1493,7 +1504,10 @@ def list_model_services(
     seen_tokens: set[str] = set()
     last_reason: str | None = None
     for _ in range(max_pages):
-        params: dict[str, str] = {"page_size": str(page_size)}
+        params: dict[str, str] = {
+            "parent": _MODEL_SERVICE_PARENT_SCHEMA,
+            "page_size": str(page_size),
+        }
         if page_token:
             params["page_token"] = page_token
         url = f"https://{hostname}/api/2.1/unity-catalog/model-services?{urlencode(params)}"

tests/test_databricks.py

@@ -355,6 +355,9 @@ def fake_get(url, token, timeout=10):
         assert ids == ["system.ai.gpt-5"]
         assert reason is None
         assert all("page_size=" in u for u in urls)
+        # Scope to the `system.ai` schema so the endpoint returns just the
+        # foundation models rather than walking the whole metastore.
+        assert all("parent=schemas%2Fsystem.ai" in u for u in urls)
 
     def test_retries_page_before_giving_up(self, monkeypatch):
         payload = {"model_services": [_model_service("system.ai.gpt-5")]}