diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 88935dd..06bd831 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2842,6 +2842,10 @@ def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str | return discover_endpoints_with_api_type(workspace, token, "openai/v1/responses") +def discover_anthropic_messages_models(workspace: str, token: str) -> tuple[list[str], str | None]: + return discover_endpoints_with_api_type(workspace, token, "anthropic/v1/messages") + + def fetch_gemini_models(workspace: str, token: str) -> list[str]: models, _ = discover_gemini_models(workspace, token) return models diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index a26e185..be5f902 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -373,7 +373,13 @@ def _known_models(state: dict) -> set[str]: claude_models = state.get("claude_models") if isinstance(claude_models, dict): known.update(m for m in claude_models.values() if isinstance(m, str) and m) - for key in ("codex_models", "gemini_models", "oss_models", "all_claude_models"): + for key in ( + "codex_models", + "gemini_models", + "oss_models", + "all_claude_models", + "anthropic_messages_models", + ): models = state.get(key) if isinstance(models, list): known.update(m for m in models if isinstance(m, str) and m) diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 2940b8b..4289d37 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -25,6 +25,7 @@ all_users_can_use_schema, create_coding_agent_config, delete_coding_agent_config, + discover_anthropic_messages_models, discover_claude_models_unbucketed, ensure_databricks_auth, get_databricks_token, @@ -1366,21 +1367,37 @@ def _with_claude_inventory(state: dict, workspace: str, profile: str | None) -> what the wizard happened to leave behind, which also covers a hand-edited or ``--from-file`` manifest authored on another machine. + Also fetches workspace endpoints exposing ``anthropic/v1/messages`` (custom Model Serving the + UC listing misses) onto ``state["anthropic_messages_models"]`` for the same reason. + Best-effort: a failed listing returns ``state`` untouched, leaving validation on the narrower inventory rather than blocking a publish on a transient API error. """ - if isinstance(state.get("all_claude_models"), list) and state["all_claude_models"]: - return state try: token = get_databricks_token(workspace, profile) - all_claude, _ = discover_claude_models_unbucketed(workspace, token) except (RuntimeError, OSError): # OSError covers a missing `databricks` binary: `get_databricks_token` shells out, so a # machine without the CLI on PATH raises FileNotFoundError rather than RuntimeError. return state - if not all_claude: - return state - return {**state, "all_claude_models": all_claude} + updated = state + if not (isinstance(state.get("all_claude_models"), list) and state["all_claude_models"]): + try: + all_claude, _ = discover_claude_models_unbucketed(workspace, token) + except (RuntimeError, OSError): + all_claude = [] + if all_claude: + updated = {**updated, "all_claude_models": all_claude} + if not ( + isinstance(state.get("anthropic_messages_models"), list) + and state["anthropic_messages_models"] + ): + try: + messages_models, _ = discover_anthropic_messages_models(workspace, token) + except (RuntimeError, OSError): + messages_models = [] + if messages_models: + updated = {**updated, "anthropic_messages_models": messages_models} + return updated def apply_command(*, yes: bool = False) -> int: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 64b4dcc..795b0fe 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1162,6 +1162,37 @@ def test_codex_discovery_keeps_alphabetical_order(self, monkeypatch): assert reason is None assert models == ["databricks-gpt-4-1", "databricks-gpt-5-2-codex"] + def test_anthropic_messages_returns_messages_api_endpoints(self, monkeypatch): + # Mirrors discover_codex_models but for anthropic/v1/messages — surfaces + # custom Model Serving the name-keyed Claude discoveries miss. + payload = { + "endpoints": [ + { + "name": name, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": [api_type], + } + } + ] + }, + } + for name, api_type in [ + ("main.default.my_claude_ms", "anthropic/v1/messages"), + ("databricks-gpt-5-2-codex", "openai/v1/responses"), + ] + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_anthropic_messages_models(WS, "token") + + assert reason is None + assert models == ["main.default.my_claude_ms"] + class TestResolvePatToken: def test_reads_pat_profile_token_from_cfg(self, monkeypatch, tmp_path): diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index d4425a6..dee13f0 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -582,6 +582,19 @@ def test_unknown_claude_version_is_still_rejected(self): errors = validate_manifest(manifest, state) assert any("not available on this workspace" in e for e in errors) + def test_custom_messages_api_ms_is_accepted(self): + # A custom Model Serving endpoint exposing anthropic/v1/messages isn't in the + # UC `all_claude_models` listing, but is routable — `anthropic_messages_models` + # vouches for it the way `custom_models` does for a hand-typed id. + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "main.default.my_claude_ms"}} + }, + } + state = {**STATE, "anthropic_messages_models": ["main.default.my_claude_ms"]} + assert validate_manifest(manifest, state) == [] + def test_custom_model_is_accepted_via_the_marker(self): # A hand-typed model service outside the discovered inventory is listed in `custom_models` # (it was verified to exist when entered), so the inventory check must not reject it. diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 4872b79..4e3ecaf 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -2142,6 +2142,39 @@ def fake_create(workspace, token, payload): ) assert published + def test_a_custom_messages_api_ms_publishes(self): + # A hand-edited manifest pinning a custom Model Serving (exposing + # anthropic/v1/messages) used to be rejected as unknown — the UC listing + # doesn't contain it. `apply` now re-fetches the Messages-API endpoints so + # validation accepts it. + managed_config_mod.save_managed_state( + WORKSPACE, + { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "main.default.my_claude_ms"}} + }, + }, + ) + published: dict = {} + + def fake_create(workspace, token, payload): + published["payload"] = payload + return {"name": "coding-agent-configs/new"}, None + + assert ( + self._run( + discover_claude_models_unbucketed=lambda *a, **k: ([], None), + discover_anthropic_messages_models=lambda *a, **k: ( + ["main.default.my_claude_ms"], + None, + ), + create_coding_agent_config=fake_create, + ) + == 0 + ) + assert published, "the manifest should have been published" + def test_declining_the_prompt_publishes_nothing(self): managed_config_mod.save_managed_state(WORKSPACE, self.MANIFEST) created = {"called": False}