Paginate the model-provider-service listing, and resolve by name as a fallback

src/ucode/databricks.py

@@ -1495,7 +1495,15 @@ def _provider_type_tag(provider_type: str | None) -> str:
     return tag.lower()
 
 
-def list_model_provider_services(workspace: str, token: str) -> tuple[list[dict], str | None]:
+# The listing is paginated; a metastore with more services than one page silently truncated before
+# this was honored, making services on later pages look nonexistent.
+_PROVIDER_SERVICES_PAGE_SIZE = 100
+_PROVIDER_SERVICES_MAX_PAGES = 50
+
+
+def list_model_provider_services(
+    workspace: str, token: str, *, parent: str | None = None
+) -> tuple[list[dict], str | None]:
     """List Unity Catalog Model Provider Services on the workspace.
 
     Returns ``(services, reason)`` where each service is
@@ -1505,46 +1513,107 @@ def list_model_provider_services(workspace: str, token: str) -> tuple[list[dict]
     Bedrock model names). ``relayed`` is True for a credential-less Anthropic
     service (Claude Max/Team/Enterprise subscription relay). A non-None
     ``reason`` means the listing call itself failed.
+
+    Pages through the endpoint: a metastore with more services than fit on one page used to have the
+    remainder silently dropped, so a service that plainly existed looked absent. ``parent`` scopes
+    the listing to one ``catalog.schema`` — the metastore-wide default is documented as an internal,
+    likely-to-be-deprecated scope, so prefer passing it when the schema is known.
     """
     hostname = workspace_hostname(workspace)
-    url = f"https://{hostname}/api/2.1/unity-catalog/model-provider-services"
-    payload, reason = _http_get_json(url, token, timeout=30)
-    if payload is None:
-        return [], reason
-    data = cast(dict, payload) if isinstance(payload, dict) else {}
     services: list[dict] = []
-    for service in data.get("model_provider_services") or []:
-        if not isinstance(service, dict):
-            continue
-        raw_name = service.get("name")
-        if not isinstance(raw_name, str) or not raw_name:
-            continue
-        # The API returns `model-provider-services/<catalog>.<schema>.<name>`.
-        full_name = raw_name.split("/", 1)[1] if "/" in raw_name else raw_name
-        config = service.get("config") if isinstance(service.get("config"), dict) else {}
-        targets = []
-        for target in config.get("targets") or []:
-            model_id = target.get("model") if isinstance(target, dict) else None
-            if isinstance(model_id, str) and model_id:
-                targets.append(model_id)
-        # Relayed = credential-less Anthropic (subscription relay). Only whether
-        # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed
-        # server-side, so both launch identically.
-        anthropic_cfg = config.get("anthropic")
-        relayed = isinstance(anthropic_cfg, dict) and "relayed" in anthropic_cfg
-        services.append(
-            {
-                "name": full_name,
-                "provider_type": _provider_type_tag(config.get("provider_type")),
-                "targets": targets,
-                "allow_all_targets": bool(config.get("allow_all_targets")),
-                "relayed": relayed,
-            }
+    page_token: str | None = None
+    seen_tokens: set[str] = set()
+    last_reason: str | None = None
+    for _ in range(_PROVIDER_SERVICES_MAX_PAGES):
+        params: dict[str, str] = {"page_size": str(_PROVIDER_SERVICES_PAGE_SIZE)}
+        if parent:
+            params["parent"] = f"schemas/{parent}"
+        if page_token:
+            params["page_token"] = page_token
+        url = (
+            f"https://{hostname}/api/2.1/unity-catalog/model-provider-services?{urlencode(params)}"
         )
+        payload, reason = _http_get_json(url, token, timeout=30)
+        if payload is None:
+            # Surface the failure only if we have nothing yet; a mid-pagination blip still
+            # returns whatever was collected.
+            last_reason = reason
+            break
+        data = cast(dict, payload) if isinstance(payload, dict) else {}
+        for service in data.get("model_provider_services") or []:
+            entry = _provider_service_entry(service)
+            if entry is not None:
+                services.append(entry)
+        page_token = data.get("next_page_token") or None
+        if not page_token:
+            last_reason = None
+            break
+        if page_token in seen_tokens:
+            break
+        seen_tokens.add(page_token)
+
+    if not services and last_reason is not None:
+        return [], last_reason
     services.sort(key=lambda s: s["name"])
     return services, None
 
 
+def _provider_service_entry(raw_service: object) -> dict | None:
+    """Normalize one listing entry, or None when it isn't usable."""
+    if not isinstance(raw_service, dict):
+        return None
+    # A bare isinstance narrows to dict[Never, Never], which rejects string keys.
+    service = cast("dict[str, object]", raw_service)
+    raw_name = service.get("name")
+    if not isinstance(raw_name, str) or not raw_name:
+        return None
+    # The API returns `model-provider-services/<catalog>.<schema>.<name>`.
+    full_name = raw_name.split("/", 1)[1] if "/" in raw_name else raw_name
+    raw_config = service.get("config")
+    config = cast("dict[str, object]", raw_config) if isinstance(raw_config, dict) else {}
+    targets: list[str] = []
+    raw_targets = config.get("targets")
+    for target in raw_targets if isinstance(raw_targets, list) else []:
+        if not isinstance(target, dict):
+            continue
+        model_id = cast("dict[str, object]", target).get("model")
+        if isinstance(model_id, str) and model_id:
+            targets.append(model_id)
+    # Relayed = credential-less Anthropic (subscription relay). Only whether
+    # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed
+    # server-side, so both launch identically.
+    anthropic_cfg = config.get("anthropic")
+    relayed = isinstance(anthropic_cfg, dict) and "relayed" in anthropic_cfg
+    raw_type = config.get("provider_type")
+    return {
+        "name": full_name,
+        "provider_type": _provider_type_tag(raw_type if isinstance(raw_type, str) else None),
+        "targets": targets,
+        "allow_all_targets": bool(config.get("allow_all_targets")),
+        "relayed": relayed,
+    }
+
+
+def get_model_provider_service(
+    service_name: str, workspace: str, token: str
+) -> tuple[dict | None, str | None]:
+    """Fetch one provider service by its full `catalog.schema.name`, bypassing the listing.
+
+    The listing is paginated and metastore-wide, so any gap in it (a page we failed to fetch, a
+    server-side filter) makes a service that plainly exists look absent. Addressing it directly
+    removes that whole class of false negative.
+    """
+    hostname = workspace_hostname(workspace)
+    url = f"https://{hostname}/api/2.1/unity-catalog/model-provider-services/{service_name}"
+    payload, reason = _http_get_json(url, token, timeout=30)
+    if payload is None:
+        return None, reason
+    entry = _provider_service_entry(payload)
+    if entry is None:
+        return None, "model-provider-service response had an unexpected shape"
+    return entry, None
+
+
 def is_model_provider_feature_unavailable(reason: str | None) -> bool:
     """True when a model-provider-services API failure means the workspace
     simply hasn't enabled the feature (HTTP 400 "feature is not available"),
@@ -1603,11 +1672,16 @@ def resolve_provider_service(
         return None, f"Could not list model provider services: {reason}"
     match = next((s for s in services if s["name"] == service_name), None)
     if match is None:
-        usable = [
-            s["name"] for s in services if tool_supports_provider_type(tool, s["provider_type"])
-        ]
-        suffix = f" Available for {tool}: {', '.join(usable)}." if usable else ""
-        return None, f"Model provider service '{service_name}' was not found.{suffix}"
+        # Don't conclude "not found" from a listing that may be incomplete — a named service can be
+        # fetched directly. Only when that 404s is it really absent.
+        match, get_reason = get_model_provider_service(service_name, workspace, token)
+        if match is None:
+            usable = [
+                s["name"] for s in services if tool_supports_provider_type(tool, s["provider_type"])
+            ]
+            suffix = f" Available for {tool}: {', '.join(usable)}." if usable else ""
+            detail = f" ({get_reason})" if get_reason and "404" not in get_reason else ""
+            return None, f"Model provider service '{service_name}' was not found.{detail}{suffix}"
     provider_type = match["provider_type"]
     if not tool_supports_provider_type(tool, provider_type):
         supported = ", ".join(_TOOL_PROVIDER_TYPES.get(tool, ())) or "none"

tests/test_databricks.py

@@ -499,6 +499,136 @@ def test_empty_when_no_claude(self):
         assert db_mod.map_bedrock_claude_models(["amazon.titan-text-express-v1"]) == {}
 
 
+class TestProviderServicePagination:
+    """The listing is paginated; ignoring next_page_token hid services on later pages entirely."""
+
+    @staticmethod
+    def _page(names, next_token=None):
+        payload = {
+            "model_provider_services": [
+                {
+                    "name": f"model-provider-services/{n}",
+                    "config": {"provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI"},
+                }
+                for n in names
+            ]
+        }
+        if next_token:
+            payload["next_page_token"] = next_token
+        return payload
+
+    def test_follows_next_page_token(self, monkeypatch):
+        pages = [
+            self._page(["main.s.one"], next_token="tok2"),
+            self._page(["main.s.two"]),
+        ]
+        seen: list[str] = []
+
+        def fake_get(url, token, **kwargs):
+            seen.append(url)
+            return pages[len(seen) - 1], None
+
+        monkeypatch.setattr(db_mod, "_http_get_json", fake_get)
+        services, reason = db_mod.list_model_provider_services("https://ws", "tok")
+
+        assert reason is None
+        assert [s["name"] for s in services] == ["main.s.one", "main.s.two"]
+        assert "page_token=tok2" in seen[1]
+
+    def test_stops_on_a_repeated_token(self, monkeypatch):
+        # A server that echoes the same token would otherwise spin forever.
+        monkeypatch.setattr(
+            db_mod,
+            "_http_get_json",
+            lambda url, token, **kw: (self._page(["main.s.one"], next_token="same"), None),
+        )
+        services, reason = db_mod.list_model_provider_services("https://ws", "tok")
+        assert reason is None
+        assert len(services) >= 1
+
+    def test_keeps_earlier_pages_when_a_later_one_fails(self, monkeypatch):
+        calls = {"n": 0}
+
+        def fake_get(url, token, **kwargs):
+            calls["n"] += 1
+            if calls["n"] == 1:
+                return self._page(["main.s.one"], next_token="tok2"), None
+            return None, "HTTP 500 Server Error"
+
+        monkeypatch.setattr(db_mod, "_http_get_json", fake_get)
+        services, reason = db_mod.list_model_provider_services("https://ws", "tok")
+
+        # A mid-pagination blip should degrade to partial results, not to an error.
+        assert reason is None
+        assert [s["name"] for s in services] == ["main.s.one"]
+
+    def test_reports_the_failure_when_nothing_was_collected(self, monkeypatch):
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, **kw: (None, "HTTP 403 Forbidden")
+        )
+        services, reason = db_mod.list_model_provider_services("https://ws", "tok")
+        assert services == []
+        assert reason == "HTTP 403 Forbidden"
+
+    def test_parent_scopes_the_listing(self, monkeypatch):
+        seen: dict = {}
+
+        def fake_get(url, token, **kwargs):
+            seen["url"] = url
+            return self._page(["main.tien_le.openai"]), None
+
+        monkeypatch.setattr(db_mod, "_http_get_json", fake_get)
+        db_mod.list_model_provider_services("https://ws", "tok", parent="main.tien_le")
+
+        assert "parent=schemas%2Fmain.tien_le" in seen["url"]
+
+    def test_page_size_is_always_sent(self, monkeypatch):
+        seen: dict = {}
+
+        def fake_get(url, token, **kwargs):
+            seen["url"] = url
+            return self._page(["main.s.one"]), None
+
+        monkeypatch.setattr(db_mod, "_http_get_json", fake_get)
+        db_mod.list_model_provider_services("https://ws", "tok")
+
+        assert "page_size=" in seen["url"]
+
+
+class TestGetModelProviderService:
+    def test_addresses_the_service_directly(self, monkeypatch):
+        seen: dict = {}
+
+        def fake_get(url, token, **kwargs):
+            seen["url"] = url
+            return {
+                "name": "model-provider-services/main.tien_le.openai_all",
+                "config": {
+                    "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI",
+                    "allow_all_targets": True,
+                },
+            }, None
+
+        monkeypatch.setattr(db_mod, "_http_get_json", fake_get)
+        service, reason = db_mod.get_model_provider_service(
+            "main.tien_le.openai_all", "https://ws", "tok"
+        )
+
+        assert reason is None
+        assert service is not None
+        assert service["name"] == "main.tien_le.openai_all"
+        assert service["allow_all_targets"] is True
+        assert seen["url"].endswith("/model-provider-services/main.tien_le.openai_all")
+
+    def test_missing_service_returns_the_reason(self, monkeypatch):
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, **kw: (None, "HTTP 404 Not Found")
+        )
+        service, reason = db_mod.get_model_provider_service("main.a.b", "https://ws", "tok")
+        assert service is None
+        assert "404" in (reason or "")
+
+
 class TestResolveProviderService:
     _PAYLOAD = TestListModelProviderServices._PAYLOAD