From 97cecde25672ac493fb75a77f6b4cfccb5bdc60a Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:07:25 +0000 Subject: [PATCH 1/2] claude: accept custom Messages-API MS in managed-config validation Add discover_anthropic_messages_models (mirroring discover_codex_models) for anthropic/v1/messages endpoints, wire it into _known_models so validate_manifest accepts them, and extend _with_claude_inventory so `ucode apply` re-fetches them for a hand-edited manifest. Custom Model Serving exposing the Messages API is now validated without the custom_models escape hatch. Co-authored-by: Isaac --- src/ucode/databricks.py | 12 ++++++++++++ src/ucode/managed_setup.py | 8 +++++++- src/ucode/managed_wizard.py | 29 +++++++++++++++++++++++------ tests/test_databricks.py | 31 +++++++++++++++++++++++++++++++ tests/test_managed_setup.py | 13 +++++++++++++ tests/test_managed_wizard.py | 35 +++++++++++++++++++++++++++++++++++ 6 files changed, 121 insertions(+), 7 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 88935dd..ac4c414 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2842,6 +2842,18 @@ 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]: + """Workspace endpoints exposing the Anthropic Messages API (``anthropic/v1/messages``). + + Mirrors :func:`discover_codex_models` for the Messages API, surfacing custom + Model Serving the name-keyed Claude discoveries miss (neither + ``databricks-claude-`` nor ``system.ai.claude-``). + """ + 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..f712502 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -2142,6 +2142,41 @@ 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} From 5b2424f3611437abe4bd02bbaa2fa8506df5621a Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:19:33 +0000 Subject: [PATCH 2/2] claude: drop verbose docstring, fix ruff format Remove the discover_anthropic_messages_models docstring and collapse multi-line signatures to satisfy ruff format (CI failure on #348). Co-authored-by: Isaac --- src/ucode/databricks.py | 10 +--------- tests/test_managed_wizard.py | 4 +--- 2 files changed, 2 insertions(+), 12 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index ac4c414..06bd831 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2842,15 +2842,7 @@ 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]: - """Workspace endpoints exposing the Anthropic Messages API (``anthropic/v1/messages``). - - Mirrors :func:`discover_codex_models` for the Messages API, surfacing custom - Model Serving the name-keyed Claude discoveries miss (neither - ``databricks-claude-`` nor ``system.ai.claude-``). - """ +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") diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index f712502..4e3ecaf 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -2152,9 +2152,7 @@ def test_a_custom_messages_api_ms_publishes(self): { "default_agent": "claude", "enabled_agents": { - "claude": { - "model_config": {"default_model": "main.default.my_claude_ms"} - } + "claude": {"model_config": {"default_model": "main.default.my_claude_ms"}} }, }, )