From 8fc337ec0ff1c2ee549e99ade53b719b94386aca Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:42:03 +1000 Subject: [PATCH 1/4] fix(pi): use native Responses models safely --- src/ucode/agents/pi.py | 146 +++++++++++++++++++++++---- src/ucode/cli.py | 4 + src/ucode/databricks.py | 155 +++++++++++++++++++++++++++-- tests/test_agent_pi.py | 186 +++++++++++++++++++++++++++++++++-- tests/test_databricks.py | 92 ++++++++++++++++- tests/test_e2e.py | 4 +- tests/test_e2e_user_agent.py | 6 +- tests/test_state.py | 6 +- 8 files changed, 553 insertions(+), 46 deletions(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index a673a548..bb44efb1 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -1,11 +1,11 @@ -"""Pi coding agent: writes a ucode-private models.json with Databricks-backed providers. +"""Pi coding agent: writes the user's models.json with Databricks-backed providers. Pi (https://pi.dev) is a multi-provider coding agent. We register three providers in its `models.json`, each speaking the API dialect best suited to that family's gateway path: - `databricks-claude` (api: anthropic-messages) → /ai-gateway/anthropic -- `databricks-openai` (api: openai-responses) → /ai-gateway/codex/v1 +- `databricks-openai` (api: openai-responses) → /ai-gateway/openai/v1 - `databricks-gemini` (api: google-generative-ai) → /ai-gateway/gemini/v1beta Per-provider `compat` flags work around fields the gateway translators reject: @@ -31,6 +31,7 @@ import signal import subprocess import threading +from pathlib import Path from ucode.config_io import ( APP_DIR, @@ -45,17 +46,26 @@ TOKEN_REFRESH_INTERVAL_SECONDS, build_pi_base_urls, classify_model_family, + claude_model_capabilities, + discover_claude_models_unbucketed, get_databricks_token, + gpt_model_token_limits, + preferred_gpt_model, ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version -PI_UCODE_HOME = APP_DIR / "pi-home" -PI_CONFIG_DIR = PI_UCODE_HOME / ".pi" / "agent" +# Point Pi at its standard user configuration directory without replacing HOME. +# This lets `ucode pi` retain the user's installed extensions, packages and +# skills while ucode manages only its own provider keys and default selection. +PI_CONFIG_DIR = Path.home() / ".pi" / "agent" PI_CONFIG_PATH = PI_CONFIG_DIR / "models.json" PI_SETTINGS_PATH = PI_CONFIG_DIR / "settings.json" -PI_BACKUP_PATH = APP_DIR / "pi-models.backup.json" -PI_SETTINGS_BACKUP_PATH = APP_DIR / "pi-settings.backup.json" +# Do not reuse the legacy backup names from ucode's private Pi home. On upgrade, +# those files can contain an unrelated old private config and must never be +# restored over the user's standard ~/.pi/agent files. +PI_BACKUP_PATH = APP_DIR / "pi-agent-models.backup.json" +PI_SETTINGS_BACKUP_PATH = APP_DIR / "pi-agent-settings.backup.json" SPEC: ToolSpec = { "binary": "pi", @@ -83,12 +93,15 @@ def _resolve_model_selector( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + claude_model_ids: list[str] | None = None, ) -> str: """Return a Pi model selector in `/` form when possible.""" for name in PROVIDER_NAMES: if model.startswith(f"{name}/"): return model - if model in claude_models.values(): + all_claude_models = set(claude_models.values()) + all_claude_models.update(claude_model_ids or []) + if model in all_claude_models: return f"databricks-claude/{model}" if model in codex_models: return f"databricks-openai/{model}" @@ -97,6 +110,57 @@ def _resolve_model_selector( return model +def _pi_claude_model_entry(model_id: str) -> dict: + """Build a Claude entry with explicit context and thinking metadata.""" + capabilities = claude_model_capabilities(model_id) + entry: dict = { + "id": model_id, + "reasoning": True, + "input": ["text", "image"], + "contextWindow": capabilities.context, + "maxTokens": capabilities.output, + } + if capabilities.force_adaptive_thinking: + entry["compat"] = {"forceAdaptiveThinking": True} + entry["thinkingLevelMap"] = {"max": "max"} + if capabilities.supports_xhigh_thinking: + entry["thinkingLevelMap"]["xhigh"] = "xhigh" + return entry + + +def _pi_gpt_model_entry(model_id: str) -> dict: + """Build a Pi Responses model entry with explicit limits and reasoning.""" + limits = gpt_model_token_limits(model_id) + entry: dict = { + "id": model_id, + "contextWindow": limits["context"], + "maxTokens": limits["output"], + } + normalized_id = model_id.rsplit("/", 1)[-1].lower() + for prefix in ("system.ai.", "databricks-"): + if normalized_id.startswith(prefix): + normalized_id = normalized_id[len(prefix) :] + break + normalized_id = normalized_id.replace(".", "-") + if normalized_id == "grok-4-6": + # Grok 4.6 accepts exactly these reasoning levels. Hide Pi's unsupported + # off/minimal/max choices rather than translating them to invalid values. + entry["reasoning"] = True + entry["thinkingLevelMap"] = { + "off": None, + "minimal": None, + "xhigh": "xhigh", + "max": None, + } + elif "gpt-5" in normalized_id: + entry["reasoning"] = True + entry["input"] = ["text", "image"] + # Older GPT-5 routes reject `reasoning.effort: none`; None makes Pi omit + # the reasoning object entirely when thinking is off. + entry["thinkingLevelMap"] = {"off": None} + return entry + + def render_overlay( model: str, token: str, @@ -104,15 +168,16 @@ def render_overlay( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + claude_model_ids: list[str] | None = None, ) -> tuple[dict, list[list[str]]]: - """Return (overlay, managed_key_paths) for Pi's private agent config.""" + """Return (overlay, managed_key_paths) for Pi's user agent config.""" providers: dict = {} keys: list[list[str]] = [["model"]] # Pi expands header values that match an env var name. Our UA contains # `/` and a space so it can never collide — safe to pass as a literal. ua_headers = {"User-Agent": f"ucode/{ucode_version()} pi/{agent_version('pi')}"} - claude_ids = sorted(set(claude_models.values())) + claude_ids = sorted(set(claude_models.values()) | set(claude_model_ids or [])) if claude_ids: providers["databricks-claude"] = { "baseUrl": pi_base_urls["claude"], @@ -124,7 +189,7 @@ def render_overlay( # the legacy beta header instead when this is false. "compat": {"supportsEagerToolInputStreaming": False}, "headers": ua_headers, - "models": [{"id": m} for m in claude_ids], + "models": [_pi_claude_model_entry(m) for m in claude_ids], } keys.append(["providers", "databricks-claude"]) if codex_models: @@ -134,7 +199,7 @@ def render_overlay( "apiKey": token, "authHeader": True, "headers": ua_headers, - "models": [{"id": m} for m in codex_models], + "models": [_pi_gpt_model_entry(m) for m in codex_models], } keys.append(["providers", "databricks-openai"]) if gemini_models: @@ -148,7 +213,9 @@ def render_overlay( } keys.append(["providers", "databricks-gemini"]) overlay: dict = { - "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models), + "model": _resolve_model_selector( + model, claude_models, codex_models, gemini_models, claude_model_ids + ), } if providers: overlay["providers"] = providers @@ -169,11 +236,16 @@ def write_tool_config( ) pi_base_urls = state.get("base_urls", {}).get("pi") or build_pi_base_urls(state["workspace"]) managed_families = _managed_model_families(state) - claude_models, codex_models, gemini_models = managed_families or ( - state.get("claude_models") or {}, - state.get("codex_models") or [], - state.get("gemini_models") or [], - ) + if managed_families is None: + claude_models = state.get("claude_models") or {} + codex_models = state.get("codex_models") or [] + gemini_models = state.get("gemini_models") or [] + claude_model_ids = ( + _discover_pi_claude_models(state, token, claude_models) if claude_models else None + ) + else: + claude_models, codex_models, gemini_models = managed_families + claude_model_ids = _managed_pi_claude_models(state) overlay, managed_keys = render_overlay( model, token, @@ -181,6 +253,7 @@ def write_tool_config( claude_models, codex_models, gemini_models, + claude_model_ids, ) existing = read_json_safe(PI_CONFIG_PATH) providers = existing.get("providers") @@ -195,6 +268,39 @@ def write_tool_config( return state, token +def _managed_pi_claude_models(state: dict) -> list[str]: + """Return every Claude id explicitly allowed by a managed Pi config.""" + managed = state.get("pi_models") + if not isinstance(managed, list): + return [] + return [ + model + for model in managed + if isinstance(model, str) and classify_model_family(model) in ANTHROPIC_FAMILIES + ] + + +def _discover_pi_claude_models(state: dict, token: str, claude_models: dict[str, str]) -> list[str]: + """Supplement Pi's family pins with all enabled Claude model versions.""" + allowed_families = set(claude_models) + cached = state.get("pi_claude_models") + if isinstance(cached, list): + return [ + model + for model in cached + if isinstance(model, str) and classify_model_family(model) in allowed_families + ] + + try: + discovered, _ = discover_claude_models_unbucketed(state["workspace"], token) + except (RuntimeError, OSError): + discovered = [] + if discovered: + state["pi_claude_models"] = discovered + return [model for model in discovered if classify_model_family(model) in allowed_families] + return list(claude_models.values()) + + def _write_settings(model_selector: str) -> None: # Pin defaultProvider/defaultModel in settings.json so Pi doesn't fall # through to an env-key-backed provider (e.g. HF_TOKEN exposing @@ -251,9 +357,9 @@ def default_model(state: dict) -> str | None: for family in ("opus", "sonnet", "haiku"): if claude_models.get(family): return claude_models[family] - codex_models = state.get("codex_models") or [] - if codex_models: - return codex_models[0] + codex_model = preferred_gpt_model(state.get("codex_models") or []) + if codex_model: + return codex_model gemini_models = state.get("gemini_models") or [] return gemini_models[0] if gemini_models else None diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 9281e59b..02a32e68 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -585,6 +585,10 @@ def configure_shared_state( state.pop("fable_enabled", None) state["databricks_ai_tools_enabled"] = databricks_ai_tools_enabled state["base_urls"] = build_shared_base_urls(workspace) + # Refresh Pi's supplemental Claude inventory after discovery or a workspace + # change rather than carrying stale model ids into the next config write. + if not skip_preflight or previous_workspace != workspace: + state.pop("pi_claude_models", None) if skip_preflight: # A prior `ucode configure` created the profile; resolve it locally (no diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f875a5c7..d25d694c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -24,6 +24,7 @@ from concurrent.futures import ( TimeoutError as FutureTimeoutError, ) +from dataclasses import dataclass from decimal import Decimal, InvalidOperation from pathlib import Path from typing import Literal, NamedTuple, NoReturn, cast, overload @@ -1492,6 +1493,10 @@ def build_auth_shell_command( # support a new family. _OSS_MODEL_FAMILIES = ("kimi-", "glm-", "deepseek-") +# Models served through the OpenAI Responses route. Keep gpt-oss out: it is +# chat-completions-only and belongs to the MLflow provider. +_CODEX_MODEL_FAMILIES = ("gpt-", "grok-") + # Claude model families ucode buckets, newest tier first. Each maps to a # Claude Code family alias (ANTHROPIC_DEFAULT__MODEL). Add an entry to # support a new family in both discovery paths (`claude--*` via the @@ -1499,6 +1504,12 @@ def build_auth_shell_command( ANTHROPIC_FAMILIES = ("fable", "opus", "sonnet", "haiku") +def _is_codex_model(model_id: str) -> bool: + """Return whether a model id belongs on the OpenAI Responses route.""" + lowered = model_id.lower() + return any(family in lowered for family in _CODEX_MODEL_FAMILIES) and "gpt-oss" not in lowered + + def classify_model_family(model_id: str) -> str | None: """Bucket a model FQN into the family ucode keys its state by, or None if unrecognized. @@ -1507,14 +1518,15 @@ def classify_model_family(model_id: str) -> str | None: one of ``ANTHROPIC_FAMILIES``, ``"codex"``, ``"gemini"``, or ``"oss"``. Matching is by name substring because neither the listing nor the config records a model's API dialect. """ + lowered = model_id.lower() for family in ANTHROPIC_FAMILIES: - if f"claude-{family}-" in model_id: + if f"claude-{family}-" in lowered: return family - if "gpt-" in model_id: + if _is_codex_model(model_id): return "codex" - if "gemini-" in model_id: + if "gemini-" in lowered: return "gemini" - if any(oss in model_id for oss in _OSS_MODEL_FAMILIES): + if any(oss in lowered for oss in _OSS_MODEL_FAMILIES): return "oss" return None @@ -1544,6 +1556,135 @@ def model_token_limits(model_id: str) -> dict[str, int] | None: return None +# Gateway ids are custom models to Pi, so their limits cannot be inherited +# from Pi's built-in vendor catalogue. Entries are ordered most-specific first. +_GPT_TOKEN_LIMITS: tuple[tuple[str, dict[str, int]], ...] = ( + # Grok's output ceiling is not exposed structurally; retain the conservative + # Responses fallback while preserving its documented 500K context window. + ("grok-4-6", {"context": 500_000, "output": 16_384}), + ("gpt-5-6-sol", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-6-terra", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-6-luna", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-5-pro", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-4-pro", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-5", {"context": 272_000, "output": 128_000}), + ("gpt-5-4-mini", {"context": 400_000, "output": 128_000}), + ("gpt-5-4-nano", {"context": 400_000, "output": 128_000}), + ("gpt-5-4", {"context": 272_000, "output": 128_000}), + ("gpt-5", {"context": 400_000, "output": 128_000}), + ("gpt-4-1", {"context": 1_047_576, "output": 32_768}), + ("gpt-4o", {"context": 128_000, "output": 16_384}), + ("gpt-4-turbo", {"context": 128_000, "output": 4_096}), + ("gpt-4", {"context": 8_192, "output": 8_192}), +) +_GPT_FALLBACK_LIMITS = {"context": 128_000, "output": 16_384} + + +def _normalized_foundation_model_id(model_id: str) -> str: + """Strip route prefixes case-insensitively and normalize dotted versions.""" + tail = model_id.split("/")[-1].lower() + if tail.startswith("system.ai."): + tail = tail[len("system.ai.") :] + if tail.startswith("databricks-"): + tail = tail[len("databricks-") :] + return tail.replace(".", "-") + + +def gpt_model_token_limits(model_id: str) -> dict[str, int]: + """Return Pi metadata limits for a Responses gateway model.""" + tail = _normalized_foundation_model_id(model_id) + for family, limits in _GPT_TOKEN_LIMITS: + if tail == family or tail.startswith(f"{family}-"): + return dict(limits) + return dict(_GPT_FALLBACK_LIMITS) + + +def preferred_gpt_model(model_ids: list[str]) -> str | None: + """Prefer the newest numeric GPT id, then another Responses model.""" + eligible = [ + model_id + for model_id in model_ids + if not _normalized_foundation_model_id(model_id).startswith("gpt-oss") + ] + numeric_gpt = [ + model_id + for model_id in eligible + if re.match(r"^gpt-\d(?:-|$)", _normalized_foundation_model_id(model_id)) + ] + if numeric_gpt: + return min( + numeric_gpt, + key=lambda model_id: model_version_sort_key(_normalized_foundation_model_id(model_id)), + ) + return eligible[0] if eligible else None + + +@dataclass(frozen=True) +class ClaudeModelCapabilities: + context: int + output: int + supports_1m: bool = False + force_adaptive_thinking: bool = False + supports_xhigh_thinking: bool = False + + +_CLAUDE_FALLBACK_CAPABILITIES = ClaudeModelCapabilities(context=200_000, output=64_000) +_CLAUDE_MODEL_RE = re.compile(r"^claude-(fable|opus|sonnet|haiku)-(\d+)(?:-(\d+))?") + + +def claude_model_capabilities(model_id: str) -> ClaudeModelCapabilities: + """Return context, output, and thinking capabilities for a Claude model. + + Opus gained the opt-in 1M window in 4.6; Sonnet gained it in 4.5. + Fable 5 uses a 1M default window and therefore needs no ``[1m]`` suffix. + Extended thinking levels are explicit allowlists because later versions do + not necessarily retain a predecessor's accepted values. + """ + tail = _normalized_foundation_model_id(model_id) + match = _CLAUDE_MODEL_RE.match(tail) + if not match: + return _CLAUDE_FALLBACK_CAPABILITIES + family, major_raw, minor_raw = match.groups() + version = (int(major_raw), int(minor_raw or 0)) + if family == "opus" and version >= (4, 6): + return ClaudeModelCapabilities( + context=1_000_000, + output=128_000, + supports_1m=True, + force_adaptive_thinking=True, + supports_xhigh_thinking=version in {(4, 7), (4, 8)}, + ) + if family == "sonnet" and version >= (4, 6): + return ClaudeModelCapabilities( + context=1_000_000, + output=64_000, + supports_1m=True, + force_adaptive_thinking=True, + supports_xhigh_thinking=version == (5, 0), + ) + if family == "sonnet" and version >= (4, 5): + return ClaudeModelCapabilities(context=1_000_000, output=64_000, supports_1m=True) + if family == "fable" and version >= (5, 0): + return ClaudeModelCapabilities( + context=1_000_000, + output=128_000, + force_adaptive_thinking=True, + supports_xhigh_thinking=version == (5, 0), + ) + return _CLAUDE_FALLBACK_CAPABILITIES + + +def claude_model_supports_1m(model_id: str) -> bool: + """Whether Claude Code should request the model's opt-in ``[1m]`` tier.""" + return claude_model_capabilities(model_id).supports_1m + + +def claude_model_token_limits(model_id: str) -> dict[str, int]: + """Return Pi metadata limits from the shared Claude capability policy.""" + capabilities = claude_model_capabilities(model_id) + return {"context": capabilities.context, "output": capabilities.output} + + def _model_service_id(service: dict) -> str | None: """Extract the `system.ai.` id from one model-service entry. @@ -1796,7 +1937,7 @@ def discover_model_services( - ``claude_models`` maps ``fable``/``opus``/``sonnet``/``haiku`` to the newest matching ``system.ai.claude-*`` id (mirrors ``discover_claude_models``). - - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids, newest first. + - ``codex_models`` is the list of Responses-model ids, newest first. - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids. @@ -1823,7 +1964,7 @@ def discover_model_services( # newest-wins once the router accepts opus-5 (PR databricks-eng/universe#2365446). _prefer_opus_4_8(claude_models, ids) - codex_models = sorted([m for m in ids if "gpt-" in m], key=model_version_sort_key) + codex_models = sorted([m for m in ids if _is_codex_model(m)], key=model_version_sort_key) gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)] @@ -3362,7 +3503,7 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]: # only (MLflow rejects `store` and `tools[].function.strict`). return { "claude": build_tool_base_url("claude", workspace), - "openai": build_tool_base_url("codex", workspace), + "openai": f"{workspace}/ai-gateway/openai/v1", "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", } diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index ff7f172d..62f12b5c 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -4,8 +4,12 @@ import json from contextlib import nullcontext +from pathlib import Path from unittest.mock import patch +import pytest + +import ucode.config_io as config_io from ucode.agents import pi WS = "https://example.databricks.com" @@ -15,7 +19,7 @@ def _base_urls() -> dict[str, str]: # Native API per family — see agents/pi.py docstring for path conventions. return { "claude": f"{WS}/ai-gateway/anthropic", - "openai": f"{WS}/ai-gateway/codex/v1", + "openai": f"{WS}/ai-gateway/openai/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", } @@ -26,6 +30,7 @@ def _empty() -> dict: "claude_models": {}, "codex_models": [], "gemini_models": [], + "claude_model_ids": None, } @@ -39,6 +44,7 @@ def _overlay(model: str, token: str = "tok", **kwargs): bundle["claude_models"], bundle["codex_models"], bundle["gemini_models"], + bundle["claude_model_ids"], ) @@ -55,7 +61,31 @@ def test_display(self): def test_config_path_under_pi_agent_dir(self): assert pi.SPEC["config_path"].name == "models.json" assert pi.SPEC["config_path"].parent.name == "agent" - assert pi.PI_UCODE_HOME in pi.SPEC["config_path"].parents + assert pi.PI_CONFIG_DIR == Path.home() / ".pi" / "agent" + + @pytest.mark.parametrize( + ("new_name", "legacy_name"), + [ + (pi.PI_BACKUP_PATH.name, "pi-models.backup.json"), + (pi.PI_SETTINGS_BACKUP_PATH.name, "pi-settings.backup.json"), + ], + ) + def test_standard_config_backup_does_not_reuse_legacy_private_backup( + self, tmp_path, monkeypatch, new_name, legacy_name + ): + monkeypatch.setattr(config_io, "APP_DIR", tmp_path) + config = tmp_path / "standard.json" + current_backup = tmp_path / new_name + legacy_backup = tmp_path / legacy_name + config.write_text("user-standard-config") + legacy_backup.write_text("old-private-config-backup") + + assert config_io.backup_existing_file(config, current_backup) is True + config.write_text("ucode-overwrite") + assert config_io.restore_file(config, current_backup, managed=True) is True + + assert config.read_text() == "user-standard-config" + assert legacy_backup.read_text() == "old-private-config-backup" class TestRenderOverlayProviders: @@ -73,7 +103,28 @@ def test_openai_provider_uses_openai_responses(self): overlay, _ = _overlay("gpt-5", codex_models=["gpt-5"]) provider = overlay["providers"]["databricks-openai"] assert provider["api"] == "openai-responses" - assert provider["baseUrl"] == f"{WS}/ai-gateway/codex/v1" + assert provider["baseUrl"] == f"{WS}/ai-gateway/openai/v1" + + def test_claude_entries_pin_limits_and_extended_thinking_levels(self): + overlay, _ = _overlay( + "system.ai.claude-opus-4-8", + claude_models={ + "opus": "system.ai.claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-5", + "haiku": "system.ai.claude-haiku-4-5", + }, + ) + entries = {m["id"]: m for m in overlay["providers"]["databricks-claude"]["models"]} + opus = entries["system.ai.claude-opus-4-8"] + assert opus["contextWindow"] == 1_000_000 + assert opus["maxTokens"] == 128_000 + assert opus["compat"] == {"forceAdaptiveThinking": True} + assert opus["thinkingLevelMap"] == {"max": "max", "xhigh": "xhigh"} + assert entries["system.ai.claude-sonnet-5"]["thinkingLevelMap"] == { + "max": "max", + "xhigh": "xhigh", + } + assert "thinkingLevelMap" not in entries["system.ai.claude-haiku-4-5"] def test_gemini_provider_uses_google_generative_ai(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) @@ -153,11 +204,63 @@ def test_claude_models_listed(self): ids = {m["id"] for m in overlay["providers"]["databricks-claude"]["models"]} assert ids == {"claude-opus", "claude-sonnet"} + def test_pi_can_list_supplemental_claude_versions(self): + overlay, _ = _overlay( + "system.ai.claude-opus-5", + claude_models={"opus": "system.ai.claude-opus-4-8"}, + claude_model_ids=[ + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5", + ], + ) + provider = overlay["providers"]["databricks-claude"] + assert {model["id"] for model in provider["models"]} == { + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5", + } + assert overlay["model"] == "databricks-claude/system.ai.claude-opus-5" + def test_openai_models_listed(self): overlay, _ = _overlay("gpt-5", codex_models=["gpt-5", "gpt-5-mini"]) ids = {m["id"] for m in overlay["providers"]["databricks-openai"]["models"]} assert ids == {"gpt-5", "gpt-5-mini"} + def test_gpt_entries_pin_limits_and_omit_unsupported_off_effort(self): + overlay, _ = _overlay( + "system.ai.gpt-5-6-sol", + codex_models=["system.ai.gpt-5-6-sol", "system.ai.gpt-5"], + ) + entries = { + model["id"]: model for model in overlay["providers"]["databricks-openai"]["models"] + } + assert entries["system.ai.gpt-5-6-sol"]["contextWindow"] == 1_050_000 + assert entries["system.ai.gpt-5"]["contextWindow"] == 400_000 + assert entries["system.ai.gpt-5"]["thinkingLevelMap"] == {"off": None} + + def test_grok_appears_with_supported_thinking_levels(self): + grok = "system.ai.grok-4-6" + overlay, _ = _overlay(grok, codex_models=[grok]) + + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert entry["contextWindow"] == 500_000 + assert entry["maxTokens"] == 16_384 + assert entry["reasoning"] is True + assert entry["thinkingLevelMap"] == { + "off": None, + "minimal": None, + "xhigh": "xhigh", + "max": None, + } + assert overlay["model"] == f"databricks-openai/{grok}" + + def test_grok_preview_does_not_inherit_unverified_thinking_levels(self): + model = "system.ai.grok-4-6-preview" + overlay, _ = _overlay(model, codex_models=[model]) + + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert "reasoning" not in entry + assert "thinkingLevelMap" not in entry + def test_gemini_models_listed(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2", "gemini-2-pro"]) ids = {m["id"] for m in overlay["providers"]["databricks-gemini"]["models"]} @@ -220,9 +323,24 @@ def test_falls_back_to_haiku(self): state = {"claude_models": {"haiku": "h4"}} assert pi.default_model(state) == "h4" - def test_falls_back_to_codex(self): - state = {"claude_models": {}, "codex_models": ["gpt-5"]} - assert pi.default_model(state) == "gpt-5" + def test_falls_back_to_newest_gpt_model(self): + state = { + "claude_models": {}, + "codex_models": ["gpt-5", "system.ai.gpt-5-6-sol", "gpt-5-5"], + } + assert pi.default_model(state) == "system.ai.gpt-5-6-sol" + + def test_falls_back_to_grok_responses_endpoint(self): + grok = "system.ai.grok-4-6" + assert pi.default_model({"claude_models": {}, "codex_models": [grok]}) == grok + + def test_does_not_route_gpt_oss_to_responses(self): + state = { + "claude_models": {}, + "codex_models": ["system.ai.gpt-oss-120b"], + "gemini_models": ["gemini-2"], + } + assert pi.default_model(state) == "gemini-2" def test_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} @@ -240,7 +358,7 @@ def test_sets_oauth_token(self): env = pi.build_runtime_env("tok") assert env["OAUTH_TOKEN"] == "tok" - def test_sets_private_agent_dir_without_replacing_home(self, monkeypatch): + def test_sets_standard_agent_dir_without_replacing_home(self, monkeypatch): monkeypatch.setenv("HOME", "/real-user-home") env = pi.build_runtime_env("tok") @@ -360,6 +478,57 @@ def test_config_written_with_correct_model_and_token(self, tmp_path, monkeypatch assert written["model"] == "databricks-claude/claude-sonnet" assert written["providers"]["databricks-claude"]["apiKey"] == "tok" + def test_config_discovers_and_caches_supplemental_claude_versions(self, tmp_path, monkeypatch): + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) + state = self._state(claude_models={"opus": "system.ai.claude-opus-4-8"}) + discovered = ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"] + with ( + patch.object( + pi_mod, "discover_claude_models_unbucketed", return_value=(discovered, None) + ) as discover, + patch("ucode.agents.pi.save_state"), + ): + pi_mod.write_tool_config(state, "system.ai.claude-opus-4-8", token="tok") + + discover.assert_called_once_with(WS, "tok") + assert state["pi_claude_models"] == discovered + entries = json.loads(config_file.read_text())["providers"]["databricks-claude"]["models"] + assert {entry["id"] for entry in entries} == set(discovered) + + def test_failed_supplemental_discovery_keeps_shared_family_pins(self, tmp_path, monkeypatch): + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) + state = self._state(claude_models={"sonnet": "system.ai.claude-sonnet-4-6"}) + with ( + patch.object( + pi_mod, "discover_claude_models_unbucketed", side_effect=OSError("offline") + ), + patch("ucode.agents.pi.save_state"), + ): + pi_mod.write_tool_config(state, "system.ai.claude-sonnet-4-6", token="tok") + + entries = json.loads(config_file.read_text())["providers"]["databricks-claude"]["models"] + assert [entry["id"] for entry in entries] == ["system.ai.claude-sonnet-4-6"] + + def test_managed_pi_allowlist_keeps_same_family_claude_versions(self, tmp_path, monkeypatch): + pi_mod, config_file, settings_file, _ = self._setup(tmp_path, monkeypatch) + state = self._state( + pi_models=["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"], + pi_default_model="system.ai.claude-opus-5", + ) + + with patch("ucode.agents.pi.save_state"): + pi_mod.write_tool_config(state, pi.default_model(state), token="tok") + + written = json.loads(config_file.read_text()) + assert {model["id"] for model in written["providers"]["databricks-claude"]["models"]} == { + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5", + } + assert written["model"] == "databricks-claude/system.ai.claude-opus-5" + settings = json.loads(settings_file.read_text()) + assert settings["defaultProvider"] == "databricks-claude" + assert settings["defaultModel"] == "system.ai.claude-opus-5" + def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): # Without this, Pi's `findInitialModel` can fall through to a built-in # provider when an unrelated env var (e.g. HF_TOKEN) makes one look @@ -438,12 +607,13 @@ def test_managed_models_split_into_pis_per_provider_inputs(self): "pi_models": [ "system.ai.claude-opus-4-8", "system.ai.gpt-5", + "system.ai.grok-4-6", "system.ai.gemini-3-flash", ] } assert pi._managed_model_families(state) == ( {"opus": "system.ai.claude-opus-4-8"}, - ["system.ai.gpt-5"], + ["system.ai.gpt-5", "system.ai.grok-4-6"], ["system.ai.gemini-3-flash"], ) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 2d9f61a5..4b50d8d3 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -23,6 +23,7 @@ build_auth_token_argv, build_databricks_cli_env, build_opencode_base_urls, + build_pi_base_urls, build_shared_base_urls, build_skills_mcp_url, build_tool_base_url, @@ -130,6 +131,12 @@ def test_returns_anthropic_gemini_and_oss(self): assert urls["oss"] == f"{WS}/ai-gateway/mlflow/v1" +class TestBuildPiBaseUrls: + def test_returns_native_responses_gateway(self): + urls = build_pi_base_urls(WS) + assert urls["openai"] == f"{WS}/ai-gateway/openai/v1" + + class TestBuildSharedBaseUrls: def test_contains_all_tools(self): urls = build_shared_base_urls(WS) @@ -247,6 +254,83 @@ def test_uncapped_model_returns_none(self): assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") is None +class TestGptModelTokenLimits: + def test_gpt_and_grok_limits_across_id_forms(self): + assert db_mod.gpt_model_token_limits("SYSTEM.AI.GPT-5-6-SOL") == { + "context": 1_050_000, + "output": 128_000, + } + assert db_mod.gpt_model_token_limits("databricks-gpt-4-1") == { + "context": 1_047_576, + "output": 32_768, + } + assert db_mod.gpt_model_token_limits("system.ai.grok-4-6") == { + "context": 500_000, + "output": 16_384, + } + + def test_unknown_model_uses_conservative_fallback(self): + assert db_mod.gpt_model_token_limits("custom-responses") == { + "context": 128_000, + "output": 16_384, + } + + def test_preferred_gpt_model_uses_semantic_version_and_excludes_gpt_oss(self): + assert ( + db_mod.preferred_gpt_model(["gpt-5", "system.ai.gpt-5-6-sol", "databricks-gpt-5-5"]) + == "system.ai.gpt-5-6-sol" + ) + assert db_mod.preferred_gpt_model(["gpt-oss-120b", "system.ai.grok-4-6"]) == ( + "system.ai.grok-4-6" + ) + assert db_mod.preferred_gpt_model(["system.ai.gpt-oss-120b"]) is None + + +class TestClaudeModelCapabilities: + @pytest.mark.parametrize( + ("model_id", "context", "output", "supports_1m", "adaptive", "xhigh"), + [ + ("databricks-claude-opus-4-5", 200_000, 64_000, False, False, False), + ("databricks-claude-opus-4-6", 1_000_000, 128_000, True, True, False), + ("system.ai.claude-opus-4-8", 1_000_000, 128_000, True, True, True), + ("system.ai.claude-opus-5", 1_000_000, 128_000, True, True, False), + ("system.ai.claude-sonnet-4-5", 1_000_000, 64_000, True, False, False), + ("claude-sonnet-5", 1_000_000, 64_000, True, True, True), + ("claude-haiku-4-5", 200_000, 64_000, False, False, False), + ("system.ai.claude-fable-5", 1_000_000, 128_000, False, True, True), + ("claude-future", 200_000, 64_000, False, False, False), + ], + ) + def test_shared_capability_policy( + self, model_id, context, output, supports_1m, adaptive, xhigh + ): + capabilities = db_mod.claude_model_capabilities(model_id) + assert capabilities.context == context + assert capabilities.output == output + assert capabilities.supports_1m is supports_1m + assert capabilities.force_adaptive_thinking is adaptive + assert capabilities.supports_xhigh_thinking is xhigh + assert db_mod.claude_model_supports_1m(model_id) is supports_1m + assert db_mod.claude_model_token_limits(model_id) == { + "context": context, + "output": output, + } + + @pytest.mark.parametrize( + ("model_id", "expected"), + [ + ("claude-opus-4-7", True), + ("claude-opus-5", False), + ("claude-sonnet-5", True), + ("claude-sonnet-6", False), + ("claude-fable-5", True), + ("claude-fable-6", False), + ], + ) + def test_xhigh_thinking_is_explicitly_allowlisted(self, model_id, expected): + assert db_mod.claude_model_capabilities(model_id).supports_xhigh_thinking is expected + + class TestDiscoverModelServices: def test_buckets_families_by_name(self, monkeypatch): payload = { @@ -256,6 +340,8 @@ def test_buckets_families_by_name(self, monkeypatch): _model_service("system.ai.claude-opus-4-8"), _model_service("system.ai.claude-sonnet-4-6"), _model_service("system.ai.gpt-5"), + _model_service("system.ai.gpt-oss-120b"), + _model_service("system.ai.grok-4-6"), _model_service("system.ai.gemini-2-5-flash"), _model_service("system.ai.gemini-3-5-flash"), _model_service("system.ai.kimi-k2-7-code"), @@ -277,7 +363,8 @@ def test_buckets_families_by_name(self, monkeypatch): "opus": "system.ai.claude-opus-4-8", "sonnet": "system.ai.claude-sonnet-4-6", } - assert codex == ["system.ai.gpt-5"] + assert codex == ["system.ai.gpt-5", "system.ai.grok-4-6"] + assert "system.ai.gpt-oss-120b" not in codex # Gemini ordered newest-first via the shared sort key. assert gemini[0] == "system.ai.gemini-3-5-flash" # DeepSeek, GLM, and Kimi are allowlisted OSS families; Llama is not. @@ -2498,6 +2585,9 @@ class TestClassifyModelFamily: ("databricks-claude-haiku-4-5", "haiku"), ("system.ai.claude-fable-5", "fable"), ("system.ai.gpt-5-3-codex", "codex"), + ("system.ai.grok-4-6", "codex"), + ("SYSTEM.AI.GROK-4-6", "codex"), + ("system.ai.gpt-oss-120b", None), ("system.ai.gemini-3-flash", "gemini"), ("system.ai.kimi-k2-7-code", "oss"), ("system.ai.glm-4-6", "oss"), diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 3b9319eb..d304adca 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -1008,11 +1008,9 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) # Point PI_CODING_AGENT_DIR and ucode's config writer at the same # isolated directory without changing the process HOME. - pi_home = tmp_path / "pi-home" - pi_dir = pi_home / ".pi" / "agent" + pi_dir = tmp_path / ".pi" / "agent" config_path = pi_dir / "models.json" backup_path = tmp_path / "pi-models.backup.json" - monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json") diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index e6cec214..4ee92728 100644 --- a/tests/test_e2e_user_agent.py +++ b/tests/test_e2e_user_agent.py @@ -319,12 +319,10 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv from ucode.agents import pi _require_binary("pi") - pi_home = tmp_path / "pi-home" - pi_dir = pi_home / ".pi" / "agent" + pi_dir = tmp_path / ".pi" / "agent" config_path = pi_dir / "models.json" monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json") @@ -339,7 +337,7 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv "base_urls": { "pi": { "claude": f"{capture_server.base_url}/ai-gateway/anthropic", - "openai": f"{capture_server.base_url}/ai-gateway/codex/v1", + "openai": f"{capture_server.base_url}/ai-gateway/openai/v1", "gemini": f"{capture_server.base_url}/ai-gateway/gemini/v1beta", }, }, diff --git a/tests/test_state.py b/tests/test_state.py index 36c8ce4f..c031a0cb 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -33,7 +33,7 @@ "copilot": f"{FAKE_WS}/ai-gateway/mlflow/v1", "pi": { "claude": f"{FAKE_WS}/ai-gateway/anthropic", - "openai": f"{FAKE_WS}/ai-gateway/codex/v1", + "openai": f"{FAKE_WS}/ai-gateway/openai/v1", "gemini": f"{FAKE_WS}/ai-gateway/gemini/v1beta", }, } @@ -109,7 +109,7 @@ def test_round_trip(self): assert loaded["workspace"] == FAKE_WS assert loaded["claude_models"]["sonnet"] == "databricks-claude-sonnet-4" - def test_persists_codex_launcher_default_in_agent_state(self): + def test_persists_latest_gpt_pi_default_in_agent_state(self): save_state( { "workspace": FAKE_WS, @@ -124,7 +124,7 @@ def test_persists_codex_launcher_default_in_agent_state(self): persisted = load_full_state()["workspaces"][FAKE_WS] assert persisted["codex_models"][0] == "system.ai.gpt-5" assert "model" not in persisted["agents"]["codex"] - assert persisted["agents"]["pi"]["model"] == "system.ai.gpt-5" + assert persisted["agents"]["pi"]["model"] == "system.ai.gpt-5-6-luna" def test_save_respects_dry_run(self): import ucode.config_io as config_io_mod From 9d8ce94fc7ce0ab617d4bd5322812b360502f8a3 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:04:21 +1000 Subject: [PATCH 2/4] fix(discovery): harden Claude inventory fallbacks --- src/ucode/databricks.py | 139 ++++++++++++++++----- tests/test_databricks.py | 264 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 375 insertions(+), 28 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d25d694c..fb2564e5 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1777,8 +1777,8 @@ def list_model_services( ``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.`` 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 + is None on a complete success, otherwise it describes an HTTP/network error + or an empty or incomplete listing. 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``. @@ -1806,25 +1806,44 @@ def list_model_services( payload, reason = _get_model_services_page(url, token) if payload is None: # Surface the failure only if we have nothing yet; a mid-pagination - # blip still returns whatever we collected. + # blip still returns whatever we collected, but marks it incomplete + # so consumers can retry or use a fallback inventory. last_reason = reason break - data = cast(dict, payload) if isinstance(payload, dict) else {} - for service in data.get("model_services", []): + if not isinstance(payload, dict): + last_reason = "model-services listing returned invalid JSON" + break + data = cast(dict, payload) + raw_services = data.get("model_services", []) + if not isinstance(raw_services, list): + last_reason = "model-services listing returned invalid model_services" + break + for service in raw_services: if isinstance(service, dict): model_id = _model_service_id(service) if model_id: ids.append(model_id) - page_token = data.get("next_page_token") or None - if not page_token: + next_page_token = data.get("next_page_token") + if next_page_token is None or next_page_token == "": last_reason = None break - if page_token in seen_tokens: + if not isinstance(next_page_token, str): + last_reason = "model-services listing returned an invalid page token" break - seen_tokens.add(page_token) + if next_page_token in seen_tokens: + last_reason = "model-services listing repeated a page token" + break + seen_tokens.add(next_page_token) + page_token = next_page_token + else: + last_reason = "model-services listing exceeded the page limit" deduped = sorted(set(ids)) if deduped: + # Do not cache an incomplete walk: callers that need the full inventory + # can fall back to the legacy gateway listing or retry the UC walk. + if last_reason is not None: + return deduped, last_reason if use_cache: _MODEL_SERVICES_CACHE[workspace] = list(deduped) return deduped, None @@ -1899,18 +1918,76 @@ def model_service_exists( return False, None +_ANTHROPIC_MODELS_MAX_PAGES = 50 + + +def _discover_claude_gateway_ids(workspace: str, token: str) -> tuple[list[str], str | None]: + """Return all Claude model ids from the legacy AI Gateway listing. + + Uses the retrying Anthropic-models request helper so transient rate limits + and network blips don't empty the Claude inventory.""" + ids: list[str] = [] + after_id: str | None = None + seen_cursors: set[str] = set() + for _ in range(_ANTHROPIC_MODELS_MAX_PAGES): + payload, reason = _get_anthropic_models_json(workspace, token, after_id=after_id) + if payload is None: + return [], reason + if not isinstance(payload, dict): + return [], "AI Gateway returned invalid Claude model data" + data = cast(dict, payload) + # `data` is required on every page. Defaulting a missing member to an + # empty list would let a shape-regressed later page end the walk and + # report the ids gathered so far as a COMPLETE inventory. + raw_models = data.get("data") + if not isinstance(raw_models, list): + return [], "AI Gateway returned invalid Claude model data" + ids.extend( + model["id"] + for model in raw_models + if isinstance(model, dict) + and isinstance(model.get("id"), str) + and not model["id"].endswith("-anthropic") + ) + if not data.get("has_more"): + if ids: + return ids, None + return [], "AI Gateway returned no Claude model ids" + cursor = data.get("last_id") + if not isinstance(cursor, str) or not cursor or cursor in seen_cursors: + return [], "AI Gateway returned an invalid or repeated Claude model cursor" + seen_cursors.add(cursor) + after_id = cursor + # Page-budget exhaustion, unlike a malformed page, leaves every id we did + # read trustworthy — just not provably complete. Return the partial walk + # WITH a reason, matching `list_model_services` and the other paginated + # walkers in this module; a non-None reason already marks it incomplete so + # callers can union it with another view instead of losing the inventory. + return ids, f"AI Gateway Claude model listing exceeded {_ANTHROPIC_MODELS_MAX_PAGES} pages" + + def discover_claude_models_unbucketed(workspace: str, token: str) -> tuple[list[str], str | None]: - """Every `system.ai.claude-*` id on the workspace, unbucketed. + """Every Claude model id on the workspace, unbucketed. `discover_model_services` keeps only the newest id per family because the launch path pins one model per Claude family alias. An admin authoring a managed config needs the alternatives too (see `managed_setup.claude_family_candidates`), so this returns the full set without disturbing - that shape. + that shape. When UC model-services is unavailable, fall back to the legacy AI Gateway listing + so Pi and managed setup can still see all gateway models. """ ids, reason = list_model_services(workspace, token) - if not ids: - return [], reason - return [m for m in ids if "claude-" in m.lower()], None + uc_claude = [model for model in ids if "claude-" in model.lower()] + # A non-Claude UC result, or a partial UC walk, must not hide models from + # the legacy gateway inventory. Union both successful views when available. + if uc_claude and reason is None: + return uc_claude, None + gateway_ids, gateway_reason = _discover_claude_gateway_ids(workspace, token) + gateway_claude = [model for model in gateway_ids if "claude-" in model.lower()] + if gateway_claude: + return sorted(set(uc_claude) | set(gateway_claude)), None + if uc_claude: + return uc_claude, None + return [], gateway_reason or reason def _prefer_opus_4_8(models: dict[str, str], all_ids: list[str]) -> None: @@ -1963,6 +2040,14 @@ def discover_model_services( # routing works with the currently-deployed task_v1 router. Revert to # newest-wins once the router accepts opus-5 (PR databricks-eng/universe#2365446). _prefer_opus_4_8(claude_models, ids) + if reason is not None: + # A partial UC walk may omit an entire Claude family. Supplement the + # shared map too, not only Pi's unbucketed picker, so every agent gets + # the same routing-safe family inventory when the legacy listing works. + gateway_claude, _ = discover_claude_models(workspace, token) + for family, model in gateway_claude.items(): + claude_models.setdefault(family, model) + _prefer_opus_4_8(claude_models, [*ids, *gateway_claude.values()]) codex_models = sorted([m for m in ids if _is_codex_model(m)], key=model_version_sort_key) gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) @@ -2997,10 +3082,18 @@ def collect_services(result, _ref): return sorted(names), None -def _get_anthropic_models_json(workspace: str, token: str) -> tuple[dict | list | None, str | None]: +def _get_anthropic_models_json( + workspace: str, + token: str, + *, + after_id: str | None = None, +) -> tuple[dict | list | None, str | None]: hostname = workspace_hostname(workspace) + url = f"https://{hostname}{ANTHROPIC_MODELS_PATH}" + if after_id is not None: + url = f"{url}?{urlencode({'after_id': after_id})}" return _http_get_json( - f"https://{hostname}{ANTHROPIC_MODELS_PATH}", + url, token, max_retries=_ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES, ) @@ -3039,17 +3132,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], describes why the dict is empty (HTTP error, network error, or no models matching the expected naming convention). """ - payload, reason = _get_anthropic_models_json(workspace, token) - if payload is None: - return {}, reason - - data = cast(dict, payload) if isinstance(payload, dict) else {} - raw_ids = [ - m["id"] - for m in data.get("data", []) - if isinstance(m.get("id"), str) and not m["id"].endswith("-anthropic") - ] - + raw_ids, reason = _discover_claude_gateway_ids(workspace, token) result: dict[str, str] = {} for family in ANTHROPIC_FAMILIES: candidates = sorted( @@ -3063,7 +3146,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], if result: return result, None if not raw_ids: - return {}, "AI Gateway returned no Claude model ids" + return {}, reason or "AI Gateway returned no Claude model ids" sample = ", ".join(raw_ids[:5]) families = ",".join(ANTHROPIC_FAMILIES) return {}, ( diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 4b50d8d3..8ab459a8 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -231,6 +231,18 @@ def test_buckets_fable_family(self, monkeypatch): assert reason is None assert models["fable"] == "databricks-claude-fable-5" + def test_discovery_preserves_gateway_failure_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, **kwargs: (None, "HTTP 503 unavailable"), + ) + + models, reason = db_mod.discover_claude_models(WS, "token") + + assert models == {} + assert reason == "HTTP 503 unavailable" + def _model_service(model_id: str) -> dict: """A model-services entry whose `name` strips to `model_id`.""" @@ -425,6 +437,33 @@ def fake_get(url, token, timeout=10): assert codex == ["system.ai.gpt-5"] assert claude == {"opus": "system.ai.claude-opus-4-8"} + def test_partial_uc_listing_supplements_missing_claude_families(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "list_model_services", + lambda w, t: (["system.ai.claude-opus-4-8"], "UC page failed"), + ) + monkeypatch.setattr( + db_mod, + "discover_claude_models", + lambda w, t: ( + { + "opus": "databricks-claude-opus-4-8", + "sonnet": "databricks-claude-sonnet-5", + }, + None, + ), + ) + + claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert claude == { + "opus": "system.ai.claude-opus-4-8", + "sonnet": "databricks-claude-sonnet-5", + } + assert (codex, gemini, oss) == ([], [], []) + def test_http_failure_returns_reason(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token, timeout=10: (None, "HTTP 500 Server Error") @@ -2640,6 +2679,231 @@ def test_the_two_discovery_helpers_share_one_walk(self, monkeypatch): assert claude["opus"] == "system.ai.claude-opus-4-8" assert unbucketed == ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"] + def test_unbucketed_falls_back_to_legacy_gateway_inventory(self, monkeypatch): + monkeypatch.setattr(db_mod, "list_model_services", lambda w, t: ([], "UC unavailable")) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, **kwargs: ( + { + "data": [ + {"id": "databricks-claude-opus-4-8"}, + {"id": "databricks-claude-opus-5"}, + {"id": "databricks-claude-opus-5-anthropic"}, + ] + }, + None, + ), + ) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert reason is None + assert models == ["databricks-claude-opus-4-8", "databricks-claude-opus-5"] + + def test_unbucketed_paginates_legacy_gateway_inventory(self, monkeypatch): + monkeypatch.setattr(db_mod, "list_model_services", lambda w, t: ([], "UC unavailable")) + calls = [] + pages = [ + { + "data": [{"id": "databricks-claude-opus-4-8"}], + "has_more": True, + "last_id": "databricks-claude-opus-4-8", + }, + {"data": [{"id": "databricks-claude-sonnet-5"}], "has_more": False}, + ] + + def get_page(url, token, **kwargs): + calls.append(url) + return pages[len(calls) - 1], None + + monkeypatch.setattr(db_mod, "_http_get_json", get_page) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert reason is None + assert models == ["databricks-claude-opus-4-8", "databricks-claude-sonnet-5"] + assert len(calls) == 2 + assert calls[1].endswith("?after_id=databricks-claude-opus-4-8") + + def test_legacy_gateway_page_budget_keeps_what_it_read(self, monkeypatch): + # Exhausting the page budget leaves every id read trustworthy, just not + # provably complete: return the partial walk WITH a reason rather than + # discarding a large valid inventory (matches `list_model_services`). + monkeypatch.setattr(db_mod, "list_model_services", lambda w, t: ([], "UC unavailable")) + calls = 0 + + def get_page(url, token, **kwargs): + nonlocal calls + calls += 1 + return { + "data": [{"id": f"databricks-claude-opus-4-8-{calls}"}], + "has_more": True, + "last_id": f"cursor-{calls}", + }, None + + monkeypatch.setattr(db_mod, "_http_get_json", get_page) + + ids, reason = db_mod._discover_claude_gateway_ids(WS, "tok") + + assert len(ids) == db_mod._ANTHROPIC_MODELS_MAX_PAGES + assert reason == ( + f"AI Gateway Claude model listing exceeded {db_mod._ANTHROPIC_MODELS_MAX_PAGES} pages" + ) + + def test_legacy_gateway_cursor_cycle_is_rejected(self, monkeypatch): + monkeypatch.setattr(db_mod, "list_model_services", lambda w, t: ([], "UC unavailable")) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, **kwargs: ( + {"data": [], "has_more": True, "last_id": "same-cursor"}, + None, + ), + ) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert models == [] + assert reason == "AI Gateway returned an invalid or repeated Claude model cursor" + + def test_legacy_gateway_mid_pagination_error_discards_partial_inventory(self, monkeypatch): + monkeypatch.setattr(db_mod, "list_model_services", lambda w, t: ([], "UC unavailable")) + calls = 0 + + def get_page(url, token, **kwargs): + nonlocal calls + calls += 1 + if calls == 1: + return { + "data": [{"id": "databricks-claude-opus-4-8"}], + "has_more": True, + "last_id": "databricks-claude-opus-4-8", + }, None + return None, "HTTP 403: permission denied" + + monkeypatch.setattr(db_mod, "_http_get_json", get_page) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert models == [] + assert reason == "HTTP 403: permission denied" + + def test_legacy_gateway_malformed_later_page_is_not_a_complete_inventory(self, monkeypatch): + # A later page that omits the required `data` member must not end the + # walk and report the ids gathered so far as a complete inventory. + monkeypatch.setattr(db_mod, "list_model_services", lambda w, t: ([], "UC unavailable")) + calls = 0 + + def get_page(url, token, **kwargs): + nonlocal calls + calls += 1 + if calls == 1: + return { + "data": [{"id": "databricks-claude-opus-4-8"}], + "has_more": True, + "last_id": "databricks-claude-opus-4-8", + }, None + return {}, None + + monkeypatch.setattr(db_mod, "_http_get_json", get_page) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert models == [] + assert reason == "AI Gateway returned invalid Claude model data" + + def test_unbucketed_unions_legacy_inventory_after_partial_uc_walk(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "list_model_services", + lambda w, t: (["system.ai.claude-opus-4-8"], "UC page failed"), + ) + monkeypatch.setattr( + db_mod, + "_discover_claude_gateway_ids", + lambda w, t: (["databricks-claude-opus-5"], None), + ) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert reason is None + assert models == ["databricks-claude-opus-5", "system.ai.claude-opus-4-8"] + + def test_malformed_legacy_model_data_is_safe(self, monkeypatch): + monkeypatch.setattr(db_mod, "list_model_services", lambda w, t: ([], "UC unavailable")) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, **kwargs: ({"data": None}, None), + ) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert models == [] + assert reason == "AI Gateway returned invalid Claude model data" + + def test_repeated_uc_page_token_is_incomplete_and_uses_gateway_fallback(self, monkeypatch): + db_mod.clear_model_services_cache() + + def page(url, token): + return { + "model_services": [ + {"name": "model-services/system.ai.claude-opus-4-8"}, + ], + "next_page_token": "repeat", + }, None + + monkeypatch.setattr(db_mod, "_get_model_services_page", page) + monkeypatch.setattr( + db_mod, + "_discover_claude_gateway_ids", + lambda w, t: (["databricks-claude-opus-5"], None), + ) + + models, reason = db_mod.discover_claude_models_unbucketed(WS, "tok") + + assert reason is None + assert models == ["databricks-claude-opus-5", "system.ai.claude-opus-4-8"] + assert WS not in db_mod._MODEL_SERVICES_CACHE + + def test_malformed_uc_model_services_degrades_to_empty_result(self, monkeypatch): + db_mod.clear_model_services_cache() + monkeypatch.setattr( + db_mod, + "_get_model_services_page", + lambda url, token: ({"model_services": None}, None), + ) + + models, reason = db_mod.list_model_services(WS, "tok") + + assert models == [] + assert reason == "model-services listing returned invalid model_services" + assert WS not in db_mod._MODEL_SERVICES_CACHE + + @pytest.mark.parametrize("invalid_token", [[], {}, 0, False]) + def test_falsey_non_string_page_token_is_incomplete(self, monkeypatch, invalid_token): + db_mod.clear_model_services_cache() + monkeypatch.setattr( + db_mod, + "_get_model_services_page", + lambda url, token: ( + { + "model_services": [ + {"name": "model-services/system.ai.claude-opus-4-8"}, + ], + "next_page_token": invalid_token, + }, + None, + ), + ) + + models, reason = db_mod.list_model_services(WS, "tok") + + assert models == ["system.ai.claude-opus-4-8"] + assert reason == "model-services listing returned an invalid page token" + assert WS not in db_mod._MODEL_SERVICES_CACHE + def test_use_cache_false_forces_a_fresh_walk(self, monkeypatch): calls: dict = {} db_mod.clear_model_services_cache() From b728af041de2773229330d395d588d9377b227ce Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:17:29 +1000 Subject: [PATCH 3/4] feat(discovery): classify gateway models by capability --- src/ucode/agents/claude.py | 20 +- src/ucode/agents/opencode.py | 77 +++- src/ucode/cli.py | 35 ++ src/ucode/databricks.py | 707 ++++++++++++++++++++++++++++-- src/ucode/managed_resolve.py | 20 + tests/conftest.py | 37 +- tests/test_agent_claude.py | 33 ++ tests/test_agent_opencode.py | 138 +++++- tests/test_cli.py | 157 +++++++ tests/test_databricks.py | 786 +++++++++++++++++++++++++++++++++- tests/test_e2e_uc.py | 18 +- tests/test_managed_resolve.py | 39 ++ 12 files changed, 1972 insertions(+), 95 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 602c2538..c0f26233 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -31,6 +31,7 @@ from ucode.databricks import ( build_auth_shell_command, build_tool_base_url, + claude_model_supports_1m, get_databricks_token, ) from ucode.launcher import exec_or_spawn @@ -165,11 +166,6 @@ def _resolve_web_search_model(state: dict) -> str | None: WEB_SEARCH_MCP_NAME = "web_search" -# Matches both the AI Gateway form (`databricks-claude-opus-4-8`) and the UC -# model-services form (`system.ai.claude-opus-4-8`). -_CLAUDE_MODEL_RE = re.compile( - r"^(?:system\.ai\.)?(?:databricks-)?claude-(opus|sonnet)-(\d+)(?:-(\d+))?(.*)$" -) # Env keys the MLflow Stop hook reads to route traces. Written into the # settings `env` block alongside the hook itself. @@ -485,19 +481,9 @@ def render_overlay( def _maybe_add_1m_suffix(model: str) -> str: - if model.endswith("[1m]"): - return model - match = _CLAUDE_MODEL_RE.match(model) - if not match: + if model.endswith("[1m]") or not claude_model_supports_1m(model): return model - - family, major_raw, minor_raw, _ = match.groups() - major = int(major_raw) - minor = int(minor_raw or 0) - should_suffix = (family == "opus" and (major, minor) >= (4, 6)) or ( - family == "sonnet" and (major, minor) >= (4, 6) - ) - return f"{model}[1m]" if should_suffix else model + return f"{model}[1m]" def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool: diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index b7803d66..2bf10beb 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -6,6 +6,7 @@ import signal import subprocess import threading +from typing import cast from ucode.config_io import ( APP_DIR, @@ -64,17 +65,67 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) - return model -def _oss_model_overlay(model: str, ua_header: dict[str, str]) -> dict: - """Per-model overlay for an OSS model entry. +_OSS_SAFE_LIMITS = {"context": 128_000, "output": 8_192} - All OSS models carry the User-Agent header; models with known token limits - also pin `limit` (context + output) so OpenCode clamps `max_tokens` to a - value the gateway accepts. OpenCode's schema requires both fields together, - so the limits table always supplies both.""" + +def _positive_int(value: object) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None + + +def _oss_specs_by_id(raw_specs: object) -> dict[str, dict[str, object]]: + if not isinstance(raw_specs, list): + return {} + specs: dict[str, dict[str, object]] = {} + for raw_spec in raw_specs: + if not isinstance(raw_spec, dict): + continue + typed_spec = cast(dict[str, object], raw_spec) + model_id = typed_spec.get("id") + reasoning = typed_spec.get("reasoning") + context = typed_spec.get("context_window") + output = typed_spec.get("max_tokens") + valid_limits = all( + value is None or _positive_int(value) is not None for value in (context, output) + ) + if ( + isinstance(model_id, str) + and model_id + and isinstance(reasoning, bool) + and "context_window" in typed_spec + and "max_tokens" in typed_spec + and valid_limits + and model_id not in specs + ): + specs[model_id] = typed_spec + return specs + + +def _oss_model_overlay( + model: str, ua_header: dict[str, str], spec: dict[str, object] | None = None +) -> dict: + """Per-model OSS overlay from discovered or static capabilities. + + OpenCode requires context and output limits together. Every discovered spec + therefore receives a complete conservative pair. Missing specs retain + static GLM/Kimi/DeepSeek metadata, and unknown no-spec models remain uncapped. + """ overlay: dict = {"headers": ua_header} - limits = model_token_limits(model) - if limits is not None: - overlay["limit"] = limits + static_limits = model_token_limits(model) + context = _positive_int(spec.get("context_window")) if isinstance(spec, dict) else None + output = _positive_int(spec.get("max_tokens")) if isinstance(spec, dict) else None + if isinstance(spec, dict): + overlay["limit"] = { + "context": context + or (static_limits.get("context") if static_limits else _OSS_SAFE_LIMITS["context"]), + "output": output + or (static_limits.get("output") if static_limits else _OSS_SAFE_LIMITS["output"]), + } + elif static_limits is not None: + overlay["limit"] = static_limits + + reasoning = spec.get("reasoning") if isinstance(spec, dict) else None + if isinstance(reasoning, bool): + overlay["reasoning"] = reasoning return overlay @@ -83,6 +134,7 @@ def render_overlay( token: str, opencode_base_urls: dict[str, str], opencode_models: dict[str, list[str]], + oss_specs: list[dict] | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for opencode.json.""" auth_headers = {"Authorization": f"Bearer {token}"} @@ -132,14 +184,18 @@ def render_overlay( } keys.append(["provider", "databricks-google"]) if oss_models: + specs_by_id = _oss_specs_by_id(oss_specs) providers["databricks-oss"] = { "npm": "@ai-sdk/openai", "options": { "baseURL": opencode_base_urls["oss"], "apiKey": token, "headers": auth_headers, + # OpenCode otherwise adds `prompt_cache_key`, which the MLflow + # chat-completions gateway rejects as an unknown field. + "setCacheKey": False, }, - "models": {m: _oss_model_overlay(m, ua_header) for m in oss_models}, + "models": {m: _oss_model_overlay(m, ua_header, specs_by_id.get(m)) for m in oss_models}, } keys.append(["provider", "databricks-oss"]) @@ -169,6 +225,7 @@ def write_tool_config( token, opencode_base_urls, state.get("opencode_models") or {}, + state.get("oss_model_specs") or [], ) existing = read_json_safe(OPENCODE_CONFIG_PATH) providers = existing.get("provider") diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 02a32e68..bfac3565 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -45,6 +45,8 @@ discover_codex_models, discover_gemini_models, discover_model_services, + discover_oss_model_specs, + discover_responses_model_specs, ensure_ai_gateway, ensure_databricks_auth, ensure_pat_bearer, @@ -80,6 +82,7 @@ managed_provider_family_models, managed_provider_service, managed_supplies_models, + managed_unclassifiable_models, managed_unservable_models, recommended_agent, resolve_state, @@ -658,7 +661,9 @@ def configure_shared_state( claude_models = {} gemini_models = [] codex_models = [] + codex_specs: list[dict] = [] oss_models = [] + oss_specs: list[dict] = [] opencode_models: dict[str, list[str]] = {} web_search_model: str | None = None if skip_model_discovery: @@ -699,8 +704,26 @@ def configure_shared_state( codex_models, codex_reason = ms_codex, ms_reason if not codex_models: codex_models, codex_reason = discover_codex_models(workspace, token) + if codex_models: + codex_specs, _ = discover_responses_model_specs(workspace, token, codex_models) if want_oss: oss_models, oss_reason = ms_oss, ms_reason + if oss_models: + oss_specs, specs_reason = discover_oss_model_specs(workspace, token, oss_models) + # Keep IDs and specs aligned. Broad OSS families are admitted + # only by live capability validation; if that refresh fails, + # offering the stale IDs without safe metadata would regress + # them to uncapped client defaults. Static GLM/Kimi/DeepSeek + # fallback specs are still returned by discover_oss_model_specs. + oss_models = [spec["id"] for spec in oss_specs] + if not oss_specs and specs_reason: + oss_reason = specs_reason + else: + # The endpoint fallback returns ids and capabilities from + # the same validated listing, avoiding a second request + # whose transient failure could leave broad models uncapped. + oss_specs, oss_reason = discover_oss_model_specs(workspace, token) + oss_models = [spec["id"] for spec in oss_specs] if claude_models: opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: @@ -721,8 +744,10 @@ def configure_shared_state( state["gemini_models"] = gemini_models if want_codex: state["codex_models"] = codex_models + state["codex_model_specs"] = codex_specs if want_oss: state["oss_models"] = oss_models + state["oss_model_specs"] = oss_specs if fetch_all or "opencode" in tools: state["opencode_models"] = opencode_models save_state(state) @@ -1662,6 +1687,15 @@ def _migrate_legacy_smart_routing(state: dict) -> dict: return state +def _warn_unclassifiable_managed_models(managed: dict, tool: str) -> None: + """Explain when a managed model cannot be routed from its name alone.""" + for model in managed_unclassifiable_models(managed, tool): + print_warning( + f"Your workspace's managed config model {model} has an unrecognized model family " + "and will be ignored." + ) + + def _reject_disabled_agent(managed: dict | None, tool: str) -> None: """Refuse to launch ``tool`` when the managed config enables other agents but not this one. @@ -1986,6 +2020,7 @@ def _launch_tool( if managed is not None: state = resolve_state(managed, state, tool) print_success("Applied your workspace's managed coding agent config") + _warn_unclassifiable_managed_models(managed, tool) unservable = managed_unservable_models(managed, tool) if unservable: print_warning( diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index fb2564e5..a14a9ad0 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1489,14 +1489,37 @@ def build_auth_shell_command( # 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. +# OSS families with statically validated coding-agent behavior. They remain the +# safe fallback when the workspace's foundation-model capability listing is +# unavailable. Other model families are offered only after the listing confirms +# that they are served exclusively through MLflow chat completions. _OSS_MODEL_FAMILIES = ("kimi-", "glm-", "deepseek-") -# Models served through the OpenAI Responses route. Keep gpt-oss out: it is +# Models served through the OpenAI/Responses gateway route. UC model-service +# discovery cannot expose API dialects, so these known native families need a +# name-based classification alongside GPT. Keep gpt-oss out: it is # chat-completions-only and belongs to the MLflow provider. _CODEX_MODEL_FAMILIES = ("gpt-", "grok-") +# Native routes take precedence over the generic MLflow chat-completions route. +# A foundation model advertising any of these must not be duplicated as OSS. +_NATIVE_PROVIDER_API_TYPES = frozenset( + {"anthropic/v1/messages", "openai/v1/responses", "gemini/v1/generateContent"} +) + +# Non-chat services must never be offered to a chat agent, even if malformed +# metadata happens to advertise a chat API type. +_OSS_NON_CHAT_SUBSTRINGS = ("embedding", "embed", "rerank") + + +def _is_oss_chat_model(model_id: str) -> bool: + """True if the id matches an OSS chat family and isn't a non-chat service.""" + lowered = model_id.lower() + return not any(bad in lowered for bad in _OSS_NON_CHAT_SUBSTRINGS) and any( + family in lowered for family in _OSS_MODEL_FAMILIES + ) + + # Claude model families ucode buckets, newest tier first. Each maps to a # Claude Code family alias (ANTHROPIC_DEFAULT__MODEL). Add an entry to # support a new family in both discovery paths (`claude--*` via the @@ -1505,18 +1528,17 @@ def build_auth_shell_command( def _is_codex_model(model_id: str) -> bool: - """Return whether a model id belongs on the OpenAI Responses route.""" + """Return whether a model id belongs on the OpenAI/Responses route.""" lowered = model_id.lower() return any(family in lowered for family in _CODEX_MODEL_FAMILIES) and "gpt-oss" not in lowered def classify_model_family(model_id: str) -> str | None: - """Bucket a model FQN into the family ucode keys its state by, or None if unrecognized. + """Bucket a model FQN by recognized name family, or return None. - Mirrors how discovery buckets a model-services listing (see `discover_model_services`), so a - model named in a managed config lands in the same bucket it would have from discovery. Returns - one of ``ANTHROPIC_FAMILIES``, ``"codex"``, ``"gemini"``, or ``"oss"``. Matching is by name - substring because neither the listing nor the config records a model's API dialect. + Managed configs do not record API capabilities, so a capability-discovered model whose name + does not match a known family cannot be classified and is ignored when applying the config. + Admin-authored model lists must therefore use recognized family names. """ lowered = model_id.lower() for family in ANTHROPIC_FAMILIES: @@ -1526,7 +1548,7 @@ def classify_model_family(model_id: str) -> str | None: return "codex" if "gemini-" in lowered: return "gemini" - if any(oss in lowered for oss in _OSS_MODEL_FAMILIES): + if _is_oss_chat_model(lowered): return "oss" return None @@ -1538,22 +1560,490 @@ def classify_model_family(model_id: str) -> str | None: # config dialect. Both fields are provided because agents like OpenCode require # context and output together. Keyed by family substring; add an entry to bound # a new model. +# +# Output caps probed from the gateway 2026-07-16 (it 400s with "max_tokens (N) +# cannot exceed "); context windows from each model's docs/description +# (conservative when unstated). If the gateway raises a cap or ships a new +# model, update this table. _MODEL_TOKEN_LIMITS: dict[str, dict[str, int]] = { - # GLM-4.6: 200k context, but the gateway caps output well below the model's - # native 128k — pin 25k so requests aren't rejected. + # Keep the version-specific entry before the family fallback: GLM 5.2 has + # materially higher probed gateway limits than earlier/unknown variants. + "glm-5-2": {"context": 1_000_000, "output": 65_536}, "glm": {"context": 200_000, "output": 25_000}, + "kimi": {"context": 128_000, "output": 65_536}, } +# Conservative fallback for a future variant that matches a validated family +# but has no specific entry. Pinning a low output ceiling risks truncation, not +# a gateway 400, so it is the safe failure direction. +_OSS_FALLBACK_LIMITS = {"context": 128_000, "output": 8_192} + +# Validated families that emit reasoning. Pi renders their streamed +# reasoning_content as thinking when the model entry sets reasoning:true. +_OSS_REASONING_FAMILIES = ("glm", "kimi") + + +def model_is_reasoning(model_id: str) -> bool: + """True if the OSS model reports reasoning output (family-matched).""" + lowered = model_id.lower() + return _is_oss_chat_model(lowered) and any( + family in lowered for family in _OSS_REASONING_FAMILIES + ) + def model_token_limits(model_id: str) -> dict[str, int] | None: """Return ``{"context": ..., "output": ...}`` limits for ``model_id``, or None. - Matches by family substring (e.g. any ``*glm*`` id). None means the model - has no known limits and the agent should not pin any.""" + Prefers a specific `_MODEL_TOKEN_LIMITS` family entry (e.g. any ``*glm*`` + id). Any other OSS chat model falls back to a conservative floor so it is + never offered uncapped (which would 400). None only for non-OSS ids, where + the agent should not pin any limit.""" + lowered = model_id.lower() + if not _is_oss_chat_model(lowered): + return None for family, limits in _MODEL_TOKEN_LIMITS.items(): - if family in model_id: + if family in lowered: return dict(limits) - return None + return dict(_OSS_FALLBACK_LIMITS) + + +# The foundation-model API exposes context windows only in free-text +# descriptions (for example "context length of 1M tokens"). Keep parsing +# deliberately narrow: unrecognized or invalid text simply yields no override. +_CONTEXT_LENGTH_RES = ( + re.compile( + r"context (?:length|window) (?:of|is) ([\d.,]+)\s*(million|thousand|[MK])?" + r"(?:[-\s]*tokens?)?", + re.IGNORECASE, + ), + re.compile( + r"([\d.,]+)\s*(million|thousand|[MK])?[-\s]*tokens? context (?:length|window)", + re.IGNORECASE, + ), +) + + +def _parse_context_window(description: str) -> int | None: + if not isinstance(description, str): + return None + match = None + for pattern in _CONTEXT_LENGTH_RES: + match = pattern.search(description) + if match: + break + if not match: + return None + try: + value = float(match.group(1).replace(",", "")) + unit = (match.group(2) or "").lower() + multiplier = ( + 1_000_000 if unit in ("m", "million") else 1_000 if unit in ("k", "thousand") else 1 + ) + tokens = int(value * multiplier) + except (OverflowError, ValueError): + return None + return tokens if tokens > 0 else None + + +# Per-model output ceilings enforced by the MLflow gateway. There is no +# structured metadata field for these values; they were established by probing +# oversized requests. Keys omit route prefixes so the same entry applies to +# both `databricks-*` endpoint ids and `system.ai.*` model-service ids. +_OSS_MAX_OUTPUT_TOKENS: dict[str, int] = { + "glm-5-2": 65_536, + "inkling": 65_536, + "kimi-k2-7-code": 65_536, + "gpt-oss-120b": 25_000, + "gpt-oss-20b": 25_000, + "qwen35-122b-a10b": 25_000, + "qwen3-next-80b-a3b-instruct": 10_000, + "llama-4-maverick": 8_192, + "meta-llama-3-1-8b-instruct": 8_192, + "meta-llama-3-3-70b-instruct": 8_192, + "gemma-3-12b": 8_192, +} + + +def _canonical_oss_model_id(model_id: str) -> str: + """Normalize endpoint/model-service ids for capability matching.""" + tail = model_id.rsplit("/", 1)[-1].strip().lower() + if tail.startswith("system.ai."): + tail = tail[len("system.ai.") :] + if tail.startswith("databricks-"): + tail = tail[len("databricks-") :] + return tail + + +def _get_foundation_models_payload(workspace: str, token: str) -> tuple[dict | None, str | None]: + """Return one cached, structurally valid foundation-model catalog.""" + cached = _FOUNDATION_MODELS_CACHE.get(workspace) + if cached is not None: + return dict(cached), None + + hostname = workspace_hostname(workspace) + payload, reason = _http_get_json( + f"https://{hostname}/api/2.0/serving-endpoints:foundation-models", token + ) + if payload is None: + return None, reason + if not isinstance(payload, dict) or not isinstance(payload.get("endpoints"), list): + return None, "foundation-models listing returned malformed `endpoints`" + typed_payload = cast(dict, payload) + _FOUNDATION_MODELS_CACHE[workspace] = dict(typed_payload) + return dict(typed_payload), None + + +def _foundation_endpoint_is_ready(endpoint: dict[str, object]) -> bool: + """Treat only an explicit non-READY state as unavailable. + + Older foundation-model listings omit state, so missing or malformed state + remains compatible rather than hiding an otherwise valid endpoint. + """ + endpoint_state = endpoint.get("state") + if not isinstance(endpoint_state, dict): + return True + ready = cast(dict[str, object], endpoint_state).get("ready") + return not isinstance(ready, str) or ready.strip().upper() == "READY" + + +def _foundation_model_api_types(payload: object) -> dict[str, frozenset[str]]: + """Map canonical endpoint ids to their advertised AI Gateway V2 APIs. + + A valid endpoint is represented even when it has no V2 API types, allowing + explicit incompatible metadata to override a known-family name fallback. + """ + if not isinstance(payload, dict): + return {} + raw_endpoints = cast(dict[str, object], payload).get("endpoints") + if not isinstance(raw_endpoints, list): + return {} + + routes: dict[str, set[str]] = {} + for endpoint in raw_endpoints: + if not isinstance(endpoint, dict): + continue + endpoint_dict = cast(dict[str, object], endpoint) + name = endpoint_dict.get("name") + config = endpoint_dict.get("config") + if not isinstance(name, str) or not name.strip() or not isinstance(config, dict): + continue + api_types = routes.setdefault(_canonical_oss_model_id(name), set()) + if not _foundation_endpoint_is_ready(endpoint_dict): + # Preserve an empty entry so explicit unavailability suppresses + # known-family/static fallbacks for the same model. + continue + entities = cast(dict[str, object], config).get("served_entities") + if not isinstance(entities, list): + continue + for entity in entities: + if not isinstance(entity, dict): + continue + foundation_model = cast(dict[str, object], entity).get("foundation_model") + if not isinstance(foundation_model, dict): + continue + foundation_model_dict = cast(dict[str, object], foundation_model) + if foundation_model_dict.get("ai_gateway_v2_supported") is not True: + continue + raw_api_types = foundation_model_dict.get("api_types") + if isinstance(raw_api_types, list): + api_types.update(value for value in raw_api_types if isinstance(value, str)) + return {model_id: frozenset(api_types) for model_id, api_types in routes.items()} + + +def _foundation_model_v2_endpoint_ids(payload: object) -> list[str]: + """Return every READY, AI-Gateway-v2 endpoint id in the foundation-model catalog. + + The catalog is the gateway's own inventory and it leads UC model-services: a + foundation model can be live and routable as ``databricks-`` days before it + is registered as a ``system.ai.*`` model service (verified 2026-09: the gateway + served ``databricks-gemini-3-7-flash`` while ``system.ai.gemini-3-7-flash`` + 404'd). Names are returned verbatim because the endpoint id is exactly what the + gateway routes on. + + Endpoints are kept only when a served entity advertises + ``ai_gateway_v2_supported`` (ucode speaks only V2 routes) and the endpoint is not + reported as un-ready. Missing ``state`` metadata is treated as ready so a listing + that simply omits the field cannot hide a working model. + """ + if not isinstance(payload, dict): + return [] + raw_endpoints = cast(dict[str, object], payload).get("endpoints") + if not isinstance(raw_endpoints, list): + return [] + + names: list[str] = [] + for endpoint in raw_endpoints: + if not isinstance(endpoint, dict): + continue + endpoint_dict = cast(dict[str, object], endpoint) + name = endpoint_dict.get("name") + config = endpoint_dict.get("config") + if not isinstance(name, str) or not name.strip() or not isinstance(config, dict): + continue + if not _foundation_endpoint_is_ready(endpoint_dict): + continue + entities = cast(dict[str, object], config).get("served_entities") + if not isinstance(entities, list): + continue + supports_v2 = False + for entity in entities: + if not isinstance(entity, dict): + continue + foundation_model = cast(dict[str, object], entity).get("foundation_model") + if ( + isinstance(foundation_model, dict) + and cast(dict[str, object], foundation_model).get("ai_gateway_v2_supported") is True + ): + supports_v2 = True + break + if supports_v2: + names.append(name.strip()) + return sorted(set(names)) + + +def _gateway_only_model_ids(uc_ids: list[str], payload: object) -> list[str]: + """Catalog endpoint ids for models the UC ``system.ai`` listing doesn't have yet. + + UC-registered models keep their ``system.ai.*`` id (both spellings route, and the + UC id is what every existing config records), so a catalog entry is added only + when no UC id normalizes to the same model. Non-chat services stay excluded. + """ + known = {_canonical_oss_model_id(model_id) for model_id in uc_ids if isinstance(model_id, str)} + extra: list[str] = [] + for name in _foundation_model_v2_endpoint_ids(payload): + canonical_id = _canonical_oss_model_id(name) + if canonical_id in known or any(bad in canonical_id for bad in _OSS_NON_CHAT_SUBSTRINGS): + continue + known.add(canonical_id) + extra.append(name) + return extra + + +def _foundation_model_context_windows( + payload: object, *, api_type: str | None = None +) -> dict[str, int]: + """Return the largest advertised context window per V2 foundation model.""" + if not isinstance(payload, dict): + return {} + raw_endpoints = cast(dict[str, object], payload).get("endpoints") + if not isinstance(raw_endpoints, list): + return {} + + windows: dict[str, int] = {} + for endpoint in raw_endpoints: + if not isinstance(endpoint, dict): + continue + endpoint_dict = cast(dict[str, object], endpoint) + name = endpoint_dict.get("name") + config = endpoint_dict.get("config") + if not isinstance(name, str) or not name.strip() or not isinstance(config, dict): + continue + if not _foundation_endpoint_is_ready(endpoint_dict): + continue + entities = cast(dict[str, object], config).get("served_entities") + if not isinstance(entities, list): + continue + canonical_id = _canonical_oss_model_id(name) + for entity in entities: + if not isinstance(entity, dict): + continue + foundation_model = cast(dict[str, object], entity).get("foundation_model") + if not isinstance(foundation_model, dict): + continue + foundation_model_dict = cast(dict[str, object], foundation_model) + if foundation_model_dict.get("ai_gateway_v2_supported") is not True: + continue + raw_api_types = foundation_model_dict.get("api_types") + if api_type is not None and ( + not isinstance(raw_api_types, list) or api_type not in raw_api_types + ): + continue + description = foundation_model_dict.get("description") + context_window = ( + _parse_context_window(description) if isinstance(description, str) else None + ) + if context_window is not None: + windows[canonical_id] = max(windows.get(canonical_id, 0), context_window) + return windows + + +def discover_responses_model_specs( + workspace: str, token: str, model_ids: list[str] +) -> tuple[list[dict], str | None]: + """Project live context windows onto Responses-capable model IDs.""" + payload, reason = _get_foundation_models_payload(workspace, token) + if payload is None: + return [], reason + api_types_by_id = _foundation_model_api_types(payload) + context_by_id = _foundation_model_context_windows(payload, api_type="openai/v1/responses") + specs: list[dict] = [] + for model_id in model_ids: + canonical_id = _canonical_oss_model_id(model_id) + if "openai/v1/responses" not in api_types_by_id.get(canonical_id, frozenset()): + continue + context_window = context_by_id.get(canonical_id) + if context_window is not None: + specs.append({"id": model_id, "context_window": context_window}) + if specs or not model_ids: + return specs, None + return [], "Responses model metadata contained no context windows" + + +def _static_oss_spec(model_id: str) -> dict | None: + """Capability fallback for statically validated GLM/Kimi/DeepSeek families.""" + if not _is_oss_chat_model(model_id.lower()): + return None + limits = model_token_limits(model_id) or {} + return { + "id": model_id, + "reasoning": model_is_reasoning(model_id.lower()), + "context_window": limits.get("context"), + "max_tokens": limits.get("output"), + } + + +def _oss_specs_from_foundation_models(payload: object) -> list[dict]: + """Parse validated chat-completions-only model specs from a listing.""" + if not isinstance(payload, dict): + return [] + payload_dict = cast(dict[str, object], payload) + raw_endpoints = payload_dict.get("endpoints") + if not isinstance(raw_endpoints, list): + return [] + + specs: list[dict] = [] + for endpoint in raw_endpoints: + if not isinstance(endpoint, dict): + continue + endpoint_dict = cast(dict[str, object], endpoint) + name = endpoint_dict.get("name") + config = endpoint_dict.get("config") + if not isinstance(name, str) or not name.strip() or not isinstance(config, dict): + continue + if not _foundation_endpoint_is_ready(endpoint_dict): + continue + name = name.strip() + lowered_name = _canonical_oss_model_id(name) + is_native_family = ( + lowered_name.startswith("claude-") + or lowered_name.startswith("gemini-") + or _is_codex_model(name) + ) + if is_native_family or any(bad in lowered_name for bad in _OSS_NON_CHAT_SUBSTRINGS): + continue + config_dict = cast(dict[str, object], config) + entities = config_dict.get("served_entities") + if not isinstance(entities, list): + continue + + api_types: set[str] = set() + context_window: int | None = None + has_v2_entity = False + for entity in entities: + if not isinstance(entity, dict): + continue + entity_dict = cast(dict[str, object], entity) + foundation_model = entity_dict.get("foundation_model") + if not isinstance(foundation_model, dict): + continue + foundation_model_dict = cast(dict[str, object], foundation_model) + if foundation_model_dict.get("ai_gateway_v2_supported") is not True: + continue + has_v2_entity = True + raw_api_types = foundation_model_dict.get("api_types") + if isinstance(raw_api_types, list): + api_types.update(value for value in raw_api_types if isinstance(value, str)) + if "mlflow/v1/chat/completions" in raw_api_types: + raw_description = foundation_model_dict.get("description") + parsed_context = ( + _parse_context_window(raw_description) + if isinstance(raw_description, str) + else None + ) + if parsed_context is not None: + context_window = max(context_window or 0, parsed_context) + + if not has_v2_entity or "mlflow/v1/chat/completions" not in api_types: + continue + if api_types & _NATIVE_PROVIDER_API_TYPES: + continue + capabilities = endpoint_dict.get("capabilities") + capabilities_dict = ( + cast(dict[str, object], capabilities) if isinstance(capabilities, dict) else {} + ) + reasoning = capabilities_dict.get("openai_reasoning") is True + canonical_id = _canonical_oss_model_id(name) + max_tokens = _OSS_MAX_OUTPUT_TOKENS.get(canonical_id) + static_fallback = _static_oss_spec(name) + if static_fallback is not None: + # Missing/partial metadata must not regress the statically verified + # GLM/Kimi/DeepSeek capabilities used by existing installations. + reasoning = reasoning or static_fallback["reasoning"] + context_window = context_window or static_fallback["context_window"] + max_tokens = max_tokens or static_fallback["max_tokens"] + specs.append( + { + "id": name, + "reasoning": reasoning, + "context_window": context_window, + "max_tokens": max_tokens, + } + ) + # Foundation listings can repeat a served endpoint. Emit one stable spec + # per canonical model so downstream model lists/configs stay deduplicated. + deduped: dict[str, dict] = {} + for spec in sorted(specs, key=lambda item: item["id"]): + deduped.setdefault(_canonical_oss_model_id(spec["id"]), spec) + return list(deduped.values()) + + +def discover_oss_model_specs( + workspace: str, + token: str, + model_ids: list[str] | None = None, +) -> tuple[list[dict], str | None]: + """Discover validated MLflow chat-completions models and capabilities. + + With ``model_ids`` (the UC-first path), endpoint capabilities are projected + back onto those exact ids using their normalized model name. Statically + validated GLM/Kimi/DeepSeek ids remain available when capability discovery + fails or omits them. Without ``model_ids`` (the serving-endpoint fallback), only + models validated by the live API metadata are returned. + """ + payload, reason = _get_foundation_models_payload(workspace, token) + discovered = _oss_specs_from_foundation_models(payload) + + if model_ids is None: + if discovered: + return discovered, None + if payload is None: + return [], reason + return [], "no validated chat-completions-only OSS endpoints" + + discovered_by_id = {_canonical_oss_model_id(spec["id"]): spec for spec in discovered} + advertised_by_id = _foundation_model_api_types(payload) + specs: list[dict] = [] + for model_id in model_ids: + if not isinstance(model_id, str) or not model_id.strip(): + continue + canonical_id = _canonical_oss_model_id(model_id) + dynamic = discovered_by_id.get(canonical_id) + if dynamic is not None: + specs.append({**dynamic, "id": model_id}) + continue + # A matching catalog entry with no MLflow-chat-only spec is explicit + # evidence that this model belongs elsewhere or is incompatible. Static + # family fallback is only for missing/unavailable per-model metadata. + if canonical_id in advertised_by_id: + continue + fallback = _static_oss_spec(model_id) + if fallback is not None: + specs.append(fallback) + if specs: + return specs, None + if payload is None: + return [], reason + return [], "requested model ids matched no validated OSS endpoint" # Gateway ids are custom models to Pi, so their limits cannot be inherited @@ -1740,6 +2230,12 @@ def _get_model_services_page( # worth a second walk. Failures are never cached, so a transient error still retries. _MODEL_SERVICES_CACHE: dict[str, list[str]] = {} +# The foundation-model catalog carries the AI Gateway V2 API dialects needed to +# classify new system.ai models without vendor-name allowlists. Several +# discovery paths consume the same workspace-wide snapshot, so cache only +# successful, structurally valid responses for this short-lived process. +_FOUNDATION_MODELS_CACHE: dict[str, dict] = {} + # Same idea for the Model Provider Service listing (a different endpoint). It is workspace-wide and # filtered per agent afterwards, so `ucode setup` would otherwise re-list it once per MPS-capable # agent. Keyed by ``(workspace, parent)`` — a schema-scoped listing is a different result set than @@ -1750,6 +2246,7 @@ def _get_model_services_page( def clear_model_services_cache() -> None: """Forget cached model-service listings (used by tests, and after a workspace switch).""" _MODEL_SERVICES_CACHE.clear() + _FOUNDATION_MODELS_CACHE.clear() _MODEL_PROVIDER_SERVICES_CACHE.clear() @@ -2009,27 +2506,72 @@ def discover_model_services( ) -> tuple[dict[str, str], list[str], list[str], list[str], str | None]: """Discover models via UC model-services and bucket them by family name. + The inventory is the UC ``system.ai`` listing unioned with the AI Gateway's own + foundation-model catalog: models live on the gateway but not yet registered in UC + are included under their routable ``databricks-*`` endpoint id (see + :func:`_gateway_only_model_ids`), so a newly shipped gateway model is offered + immediately instead of waiting for UC registration. + Returns (claude_models, codex_models, gemini_models, oss_models, reason): - ``claude_models`` maps ``fable``/``opus``/``sonnet``/``haiku`` to the - newest matching ``system.ai.claude-*`` id (mirrors - ``discover_claude_models``). - - ``codex_models`` is the list of Responses-model ids, newest first. - - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. - - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids. + newest matching ``claude-*`` id (mirrors ``discover_claude_models``). + - ``codex_models`` contains every id whose live metadata advertises + ``openai/v1/responses``, newest first. GPT/Grok names are the safe fallback + when metadata is unavailable; ``gpt-oss-*`` stays excluded. + - ``gemini_models`` similarly contains every id advertising the native + Gemini API, with ``gemini-*`` as its metadata-unavailable fallback. + - ``oss_models`` is the list of OSS-model ids. ``reason`` is None on success, else explains why nothing was found. Family - bucketing is by name substring because the model-services API does not - expose per-model API dialects. + bucketing is by name substring because neither listing records a model's API + dialect. """ ids, reason = list_model_services(workspace, token) if not ids: return {}, [], [], [], reason + # The UC listing exposes names but not API dialects. Project the live + # foundation-model catalog onto those exact system.ai ids so future models + # are admitted by protocol capability rather than a vendor-name allowlist. + # A known-family name is only a fallback when that model is absent from the + # capability catalog (including a catalog outage); explicit incompatible + # metadata wins. Embedding/reranking services are never chat candidates. + foundation_payload, _ = _get_foundation_models_payload(workspace, token) + api_types_by_id = _foundation_model_api_types(foundation_payload) + + # The catalog also leads UC registration, so anything live on the gateway but + # not yet a `system.ai.*` model service joins the inventory under its routable + # `databricks-*` endpoint id. Without this, a newly shipped gateway model is + # invisible to every agent until UC catches up. + gateway_only_ids = _gateway_only_model_ids(ids, foundation_payload) + if gateway_only_ids: + ids = sorted({*ids, *gateway_only_ids}) + + def _claude_endpoint_unavailable(model_id: str) -> bool: + """True only when the catalog explicitly says this Claude model can't serve. + + `_foundation_model_api_types` keeps an EMPTY entry for an endpoint it saw + but which is not ready, and no entry at all for a model the catalog never + listed. Only the former is explicit unavailability; an absent entry stays + permissive so a workspace whose catalog omits Claude (or a UC-only or + legacy-gateway inventory) keeps working exactly as before. This mirrors + `supports_api`'s absent-vs-empty contract used for Codex/Gemini. + """ + advertised = api_types_by_id.get(_canonical_oss_model_id(model_id)) + return advertised is not None and not advertised + claude_models: dict[str, str] = {} for family in ANTHROPIC_FAMILIES: + # Sort on the canonical name so a mixed inventory still picks the newest + # version rather than the alphabetically-later id prefix. candidates = sorted( - [m for m in ids if f"claude-{family}-" in m], + [ + m + for m in ids + if f"claude-{family}-" in m.lower() and not _claude_endpoint_unavailable(m) + ], + key=lambda model_id: (_canonical_oss_model_id(model_id), model_id), reverse=True, ) if candidates: @@ -2044,15 +2586,48 @@ def discover_model_services( # A partial UC walk may omit an entire Claude family. Supplement the # shared map too, not only Pi's unbucketed picker, so every agent gets # the same routing-safe family inventory when the legacy listing works. - gateway_claude, _ = discover_claude_models(workspace, token) - for family, model in gateway_claude.items(): + legacy_claude, _ = discover_claude_models(workspace, token) + for family, model in legacy_claude.items(): claude_models.setdefault(family, model) - _prefer_opus_4_8(claude_models, [*ids, *gateway_claude.values()]) + _prefer_opus_4_8(claude_models, [*ids, *legacy_claude.values()]) - codex_models = sorted([m for m in ids if _is_codex_model(m)], key=model_version_sort_key) - gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) + def supports_api(model_id: str, api_type: str, *, known_family: bool) -> bool: + canonical_id = _canonical_oss_model_id(model_id) + if any(bad in canonical_id for bad in _OSS_NON_CHAT_SUBSTRINGS): + return False + advertised = api_types_by_id.get(canonical_id) + return api_type in advertised if advertised is not None else known_family + + codex_models = sorted( + [ + model_id + for model_id in ids + if supports_api( + model_id, + "openai/v1/responses", + known_family=_is_codex_model(model_id), + ) + ], + key=model_version_sort_key, + ) + gemini_models = sorted( + [ + model_id + for model_id in ids + if supports_api( + model_id, + "gemini/v1/generateContent", + known_family="gemini-" in model_id.lower(), + ) + ], + key=model_version_sort_key, + ) - oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)] + # Project the same cached capability snapshot onto the UC ids. This + # broadens discovery beyond the static GLM/Kimi/DeepSeek fallback only when the + # corresponding endpoint is validated as MLflow chat-completions-only. + oss_specs, _ = discover_oss_model_specs(workspace, token, ids) + oss_models = [spec["id"] for spec in oss_specs] if not (claude_models or codex_models or gemini_models or oss_models): sample = ", ".join(ids[:5]) @@ -3201,35 +3776,69 @@ def discover_endpoints_with_api_type( describes why the list is empty. `sort_key` overrides the default alphabetical ordering of the returned names. """ - hostname = workspace_hostname(workspace) - payload, reason = _http_get_json( - f"https://{hostname}/api/2.0/serving-endpoints:foundation-models", token - ) + payload, reason = _get_foundation_models_payload(workspace, token) if payload is None: return [], reason - data = cast(dict, payload) if isinstance(payload, dict) else {} - endpoints = data.get("endpoints", []) + data = cast(dict, payload) + raw_endpoints = data.get("endpoints", []) + endpoints = raw_endpoints if isinstance(raw_endpoints, list) else [] out: list[str] = [] saw_endpoint_without_v2 = False + saw_not_ready = False + saw_malformed = not isinstance(raw_endpoints, list) for ep in endpoints: - name = ep.get("name", "") - entities = ep.get("config", {}).get("served_entities", []) + if not isinstance(ep, dict): + saw_malformed = True + continue + name = ep.get("name") + config = ep.get("config") + if not isinstance(name, str) or not name or not isinstance(config, dict): + saw_malformed = True + continue + raw_entities = config.get("served_entities", []) + if not isinstance(raw_entities, list): + saw_malformed = True + continue api_types: set[str] = set() any_v2 = False - for se in entities: - fm = se.get("foundation_model", {}) + for se in raw_entities: + if not isinstance(se, dict): + saw_malformed = True + continue + fm = se.get("foundation_model") + if not isinstance(fm, dict): + saw_malformed = True + continue if fm.get("ai_gateway_v2_supported") is True: any_v2 = True - api_types.update(fm.get("api_types", [])) - if not any_v2 and entities: + raw_api_types = fm.get("api_types", []) + if isinstance(raw_api_types, list): + api_types.update(value for value in raw_api_types if isinstance(value, str)) + else: + saw_malformed = True + if not any_v2 and raw_entities: saw_endpoint_without_v2 = True if api_type in api_types: + # Readiness is judged only AFTER confirming this endpoint actually + # advertises the requested api_type. Flagging readiness earlier made + # an unready endpoint for a DIFFERENT api blame readiness for this + # one, reporting "no ready endpoint exposes X" when nothing exposed + # X at all. + if not _foundation_endpoint_is_ready(cast(dict[str, object], ep)): + saw_not_ready = True + continue out.append(name) if out: - return sorted(out, key=sort_key), None + return sorted(set(out), key=sort_key), None if not endpoints: + if saw_malformed: + return [], "foundation-models listing returned malformed `endpoints`" return [], "foundation-models listing returned no endpoints" + if saw_malformed: + return [], "foundation-models listing contained no valid matching endpoints" + if saw_not_ready: + return [], f"no ready endpoint exposes api_type `{api_type}`" if saw_endpoint_without_v2: return [], ( f"no endpoint exposes api_type `{api_type}` with " @@ -3260,6 +3869,18 @@ def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str | ) +def discover_oss_models(workspace: str, token: str) -> tuple[list[str], str | None]: + """Discover validated chat-completions-only serving endpoints. + + This is the fallback for workspaces without UC model-services. Unlike the + static family fallback used for UC ids, every endpoint returned here has + live metadata confirming AI Gateway v2 MLflow chat completions and no + competing native Anthropic, Responses, or Gemini route. + """ + specs, reason = discover_oss_model_specs(workspace, token) + return [spec["id"] for spec in specs], reason + + def fetch_gemini_models(workspace: str, token: str) -> list[str]: models, _ = discover_gemini_models(workspace, token) return models diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index b6658d4e..c7579ba7 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -87,6 +87,26 @@ def managed_state_overrides(managed: dict, tool: str) -> dict[str, object]: return overrides +def managed_unclassifiable_models(managed: dict, tool: str) -> list[str]: + """Models ignored because a name-based provider family cannot be identified. + + De-duplicated in first-seen order: a manifest may legitimately repeat an id, + and the caller warns once per returned entry. + """ + if tool not in ("opencode", "pi"): + return [] + models = _manifest_models(managed, tool) + if not isinstance(models, list): + return [] + seen: set[str] = set() + unclassifiable: list[str] = [] + for model in models: + if classify_model_family(model) is None and model not in seen: + seen.add(model) + unclassifiable.append(model) + return unclassifiable + + def managed_unservable_models(managed: dict, tool: str) -> list[str]: """The models the manifest names for ``tool`` when it has no provider to serve any of them. diff --git a/tests/conftest.py b/tests/conftest.py index 04d9638b..c955761f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,9 +8,12 @@ from ucode.databricks import ( build_shared_base_urls, - fetch_ai_gateway_claude_models, - fetch_codex_models, - fetch_gemini_models, + discover_claude_models, + discover_codex_models, + discover_gemini_models, + discover_model_services, + discover_oss_model_specs, + discover_responses_model_specs, get_databricks_token, ) from ucode.ui import normalize_workspace_url @@ -77,22 +80,42 @@ def e2e_token(e2e_workspace): @pytest.fixture(scope="session") def e2e_state(e2e_workspace, e2e_token): - """Full state dict mirroring what configure_shared_state produces.""" - claude_models = fetch_ai_gateway_claude_models(e2e_workspace, e2e_token) - gemini_models = fetch_gemini_models(e2e_workspace, e2e_token) - codex_models = fetch_codex_models(e2e_workspace, e2e_token) + """Full state dict mirroring configure's UC-first family discovery.""" + claude_models, codex_models, gemini_models, oss_models, _ = discover_model_services( + e2e_workspace, e2e_token + ) + if not claude_models: + claude_models, _ = discover_claude_models(e2e_workspace, e2e_token) + if not gemini_models: + gemini_models, _ = discover_gemini_models(e2e_workspace, e2e_token) + if not codex_models: + codex_models, _ = discover_codex_models(e2e_workspace, e2e_token) + codex_model_specs, _ = discover_responses_model_specs(e2e_workspace, e2e_token, codex_models) + if oss_models: + oss_model_specs, _ = discover_oss_model_specs(e2e_workspace, e2e_token, oss_models) + else: + oss_model_specs, _ = discover_oss_model_specs(e2e_workspace, e2e_token) + oss_models = [spec["id"] for spec in oss_model_specs] + + # E2E mirrors configure's default (Fable is premium and opt-in). + claude_models.pop("fable", None) opencode_models: dict = {} if claude_models: opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: opencode_models["gemini"] = gemini_models + if oss_models: + opencode_models["oss"] = oss_models return { "workspace": e2e_workspace, "claude_models": claude_models, "gemini_models": gemini_models, "codex_models": codex_models, + "codex_model_specs": codex_model_specs, + "oss_models": oss_models, + "oss_model_specs": oss_model_specs, "opencode_models": opencode_models, "base_urls": build_shared_base_urls(e2e_workspace), "managed_configs": {}, diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index c76b46e1..bca98fbb 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -117,6 +117,26 @@ def test_adds_1m_suffix_for_sonnet_4_6_and_later(self): overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-7[1m]" ) + def test_adds_1m_suffix_for_sonnet_4_5(self): + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-5"} + ) + assert ( + overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-5[1m]" + ) + + def test_does_not_add_1m_suffix_for_sonnet_4_4(self): + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-4"} + ) + assert overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-4" + + def test_does_not_add_1m_suffix_for_opus_4_5(self): + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"opus": "databricks-claude-opus-4-5"} + ) + assert overlay["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-5" + def test_does_not_add_1m_suffix_for_haiku(self): overlay, _ = claude.render_overlay( WS, "s4", claude_models={"haiku": "databricks-claude-haiku-4-6"} @@ -169,6 +189,19 @@ def test_custom_model_pins_fable_alias_only_when_fable_enabled(self): )[0]["env"] assert with_fable["ANTHROPIC_DEFAULT_FABLE_MODEL"] == "main.x.m" + @pytest.mark.parametrize( + ("model_id", "expected"), + [ + ("system.ai.claude-opus-5", "system.ai.claude-opus-5[1m]"), + ("databricks-claude-sonnet-5", "databricks-claude-sonnet-5[1m]"), + ("system.ai.claude-opus-4-5", "system.ai.claude-opus-4-5"), + ("system.ai.claude-fable-5", "system.ai.claude-fable-5"), + ("not-a-claude-model", "not-a-claude-model"), + ], + ) + def test_suffix_uses_shared_capability_policy(self, model_id, expected): + assert claude._maybe_add_1m_suffix(model_id) == expected + def test_sets_anthropic_base_url(self): overlay, _ = claude.render_overlay(WS, "s4") assert overlay["env"]["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic" diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index c83e8458..8c9a6c7f 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -63,6 +63,12 @@ def test_oss_provider_uses_ai_sdk_openai_package(self): ) assert overlay["provider"]["databricks-oss"]["npm"] == "@ai-sdk/openai" + def test_oss_provider_disables_unsupported_prompt_cache_key(self): + models = {"oss": ["system.ai.gpt-oss-120b"]} + overlay, _ = opencode.render_overlay("system.ai.gpt-oss-120b", "tok", _base_urls(), models) + options = overlay["provider"]["databricks-oss"]["options"] + assert options["setCacheKey"] is False + def test_deepseek_uses_oss_provider(self): model = "system.ai.deepseek-v4-pro" @@ -106,15 +112,111 @@ def test_glm_gets_token_limits(self): overlay, _ = opencode.render_overlay("system.ai.glm-5-2", "tok", _base_urls(), models) glm = overlay["provider"]["databricks-oss"]["models"]["system.ai.glm-5-2"] # OpenCode's schema requires both context and output on `limit`. - assert glm["limit"] == {"context": 200000, "output": 25000} + # Probed 2026-07-16: glm-5-2 is 1M context / 65536 output. + assert glm["limit"] == {"context": 1_000_000, "output": 65_536} - def test_non_glm_oss_model_has_no_output_cap(self): + def test_kimi_gets_token_limits(self): + # kimi is now a capped OSS family (128k context / 65536 output). models = {"oss": ["system.ai.kimi-k2-7-code"]} overlay, _ = opencode.render_overlay( "system.ai.kimi-k2-7-code", "tok", _base_urls(), models ) kimi = overlay["provider"]["databricks-oss"]["models"]["system.ai.kimi-k2-7-code"] - assert "limit" not in kimi + assert kimi["limit"] == {"context": 128_000, "output": 65_536} + + def test_uncapped_oss_model_has_no_limit(self): + # A model outside the limits table gets no `limit` (client default). + models = {"oss": ["system.ai.mystery-7b"]} + overlay, _ = opencode.render_overlay("system.ai.mystery-7b", "tok", _base_urls(), models) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.mystery-7b"] + assert "limit" not in entry + assert "reasoning" not in entry + + def test_dynamic_full_spec_sets_reasoning_and_limits(self): + models = {"oss": ["system.ai.qwen35-122b-a10b"]} + specs = [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 262_144, + "max_tokens": 25_000, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.qwen35-122b-a10b", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.qwen35-122b-a10b"] + assert entry["reasoning"] is True + assert entry["limit"] == {"context": 262_144, "output": 25_000} + + def test_dynamic_reasoning_false_is_respected(self): + models = {"oss": ["system.ai.glm-5-2"]} + specs = [ + { + "id": "system.ai.glm-5-2", + "reasoning": False, + "context_window": None, + "max_tokens": None, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.glm-5-2", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.glm-5-2"] + assert entry["reasoning"] is False + assert entry["limit"] == {"context": 1_000_000, "output": 65_536} + + def test_unknown_dynamic_spec_gets_safe_complete_limit_pair(self): + models = {"oss": ["system.ai.deepseek-v3"]} + specs = [ + { + "id": "system.ai.deepseek-v3", + "reasoning": False, + "context_window": None, + "max_tokens": None, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.deepseek-v3", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.deepseek-v3"] + assert entry["reasoning"] is False + assert entry["limit"] == {"context": 128_000, "output": 8_192} + + def test_partial_dynamic_limit_is_completed_as_valid_pair(self): + models = {"oss": ["system.ai.inkling"]} + specs = [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": None, + "max_tokens": 65_536, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.inkling", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.inkling"] + assert entry["limit"] == {"context": 128_000, "output": 65_536} + + def test_malformed_dynamic_spec_is_ignored_safely(self): + models = {"oss": ["system.ai.mystery-7b"]} + specs = [ + None, + {"id": 12, "reasoning": True}, + { + "id": "system.ai.mystery-7b", + "reasoning": "true", + "context_window": 0, + "max_tokens": True, + }, + ] + overlay, _ = opencode.render_overlay( + "system.ai.mystery-7b", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.mystery-7b"] + assert "reasoning" not in entry + assert "limit" not in entry def test_token_in_api_key(self): models = {"anthropic": ["claude-sonnet"]} @@ -423,3 +525,33 @@ def test_config_written_with_correct_model(self, tmp_path, monkeypatch): written = json.loads(config_file.read_text()) assert written["model"] == "databricks-anthropic/claude-sonnet" + + def test_state_oss_specs_reach_written_model_entry(self, tmp_path, monkeypatch): + import ucode.agents.opencode as oc_mod + + config_file = tmp_path / "opencode.json" + monkeypatch.setattr(oc_mod, "OPENCODE_CONFIG_PATH", config_file) + monkeypatch.setattr(oc_mod, "OPENCODE_BACKUP_PATH", tmp_path / "opencode-backup.json") + state = { + "workspace": WS, + "base_urls": {"opencode": _base_urls()}, + "opencode_models": {"oss": ["system.ai.inkling"]}, + "oss_model_specs": [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": 256_000, + "max_tokens": 65_536, + } + ], + "managed_configs": {}, + } + + with patch("ucode.agents.opencode.save_state"): + oc_mod.write_tool_config(state, "system.ai.inkling", token="tok") + + entry = json.loads(config_file.read_text())["provider"]["databricks-oss"]["models"][ + "system.ai.inkling" + ] + assert entry["reasoning"] is True + assert entry["limit"] == {"context": 256_000, "output": 65_536} diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bb1ea9b..b207c2ce 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2471,6 +2471,23 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "discover_responses_model_specs", lambda w, t, ids: ([], None)) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ( + [ + { + "id": model_id, + "reasoning": True, + "context_window": 128_000, + "max_tokens": 8_192, + } + for model_id in (model_ids or []) + ], + None, + ), + ) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) return cli_mod, logins, ensures, saved @@ -2561,6 +2578,82 @@ def test_uc_models_used_without_legacy_fallback(self, monkeypatch): assert legacy_called == [] assert "uc_enabled" not in state + def test_future_responses_model_persists_capability_spec(self, monkeypatch): + cli_mod, _, _, saved = self._stub_deps(monkeypatch, pat_token="dapi-pat") + model = "system.ai.future-coder-1" + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({}, [model], [], [], None), + ) + monkeypatch.setattr( + cli_mod, + "discover_responses_model_specs", + lambda w, t, ids: ([{"id": model, "context_window": 750_000}], None), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["codex_models"] == [model] + assert state["codex_model_specs"] == [{"id": model, "context_window": 750_000}] + assert saved[-1]["codex_models"] == [model] + + def test_uc_oss_ids_persist_matching_capability_specs(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({}, [], [], ["system.ai.qwen35-122b-a10b"], None), + ) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ( + [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 25_000, + } + ], + None, + ), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["oss_models"] == ["system.ai.qwen35-122b-a10b"] + assert state["oss_model_specs"] == [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 25_000, + } + ] + assert state["opencode_models"]["oss"] == ["system.ai.qwen35-122b-a10b"] + + def test_uc_dynamic_oss_ids_are_dropped_when_spec_refresh_fails(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({}, [], [], ["system.ai.inkling"], None), + ) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ([], "HTTP 503 unavailable"), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["oss_models"] == [] + assert state["oss_model_specs"] == [] + assert "oss" not in state["opencode_models"] + assert state["_discovery_reasons"]["oss"] == "HTTP 503 unavailable" + def test_codex_only_configure_persists_discovered_oss_models(self, monkeypatch): cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") monkeypatch.setattr( @@ -2700,6 +2793,7 @@ def test_ai_tools_disable_does_not_leak_across_workspaces(self, monkeypatch): def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): # No UC model-services: each family falls back to the legacy listing. cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + calls: list[str] = [] monkeypatch.setattr( cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], "no model services") ) @@ -2711,13 +2805,52 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): None, ), ) + monkeypatch.setattr( + cli_mod, + "discover_codex_models", + lambda w, t: (calls.append("codex") or ["databricks-gpt-5-6-sol"], None), + ) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ( + calls.append("oss") + or [ + { + "id": "databricks-glm-5-2", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 8_192, + } + ], + None, + ), + ) state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + assert calls == ["codex", "oss"] assert state["claude_models"] == { "opus": "databricks-claude-opus-4-8", "sonnet": "databricks-claude-sonnet-4-6", } + assert state["codex_models"] == ["databricks-gpt-5-6-sol"] + assert state["oss_models"] == ["databricks-glm-5-2"] + assert state["oss_model_specs"] == [ + { + "id": "databricks-glm-5-2", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 8_192, + } + ] + assert state["opencode_models"] == { + "anthropic": [ + "databricks-claude-opus-4-8", + "databricks-claude-sonnet-4-6", + ], + "oss": ["databricks-glm-5-2"], + } class TestConfigureSkipValidate: @@ -3011,6 +3144,30 @@ def test_a_config_naming_no_agents_blocks_nothing(self, managed): self._reject(managed, "gemini") +class TestManagedModelWarnings: + def test_warns_only_for_unclassifiable_models(self, monkeypatch): + from ucode import cli + + warnings = [] + monkeypatch.setattr(cli, "print_warning", warnings.append) + managed = { + "enabled_agents": { + "opencode": { + "model_config": { + "models": ["system.ai.future-chat-1", "system.ai.claude-opus-4-8"] + } + } + } + } + + cli._warn_unclassifiable_managed_models(managed, "opencode") + + assert len(warnings) == 1 + assert "system.ai.future-chat-1" in warnings[0] + assert "will be ignored" in warnings[0] + assert "claude-opus" not in warnings[0] + + class TestFetchManagedConfig: """The launch path's managed-config read, which gates both the allowlist and model discovery.""" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 8ab459a8..c6d81ccc 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -251,19 +251,57 @@ def _model_service(model_id: str) -> dict: class TestModelTokenLimits: def test_glm_is_capped(self): + # Probed 2026-07-16: glm-5-2 accepts 1M context / 65536 output. assert db_mod.model_token_limits("system.ai.glm-5-2") == { - "context": 200_000, - "output": 25_000, + "context": 1_000_000, + "output": 65_536, } - def test_glm_matches_any_version(self): - assert db_mod.model_token_limits("system.ai.glm-4-6-flash") == { + @pytest.mark.parametrize( + "model_id", + ["system.ai.glm-4-6-flash", "system.ai.glm-future"], + ) + def test_other_glm_versions_keep_conservative_limits(self, model_id): + assert db_mod.model_token_limits(model_id) == { "context": 200_000, "output": 25_000, } - def test_uncapped_model_returns_none(self): - assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") is None + def test_kimi_is_capped(self): + assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") == { + "context": 128_000, + "output": 65_536, + } + + def test_deepseek_uses_conservative_fallback(self): + assert db_mod.model_token_limits("system.ai.deepseek-v4-pro") == { + "context": 128_000, + "output": 8_192, + } + + def test_unvalidated_families_return_none(self): + for model_id in ( + "system.ai.inkling", + "system.ai.gpt-oss-120b", + "system.ai.llama-4-maverick", + "system.ai.qwen35-122b-a10b", + "system.ai.gemma-3-12b", + ): + assert db_mod.model_token_limits(model_id) is None + + def test_embedding_model_returns_none_not_fallback(self): + assert db_mod.model_token_limits("system.ai.qwen3-embedding-0-6b") is None + + +class TestModelIsReasoning: + def test_reasoning_families(self): + assert db_mod.model_is_reasoning("system.ai.glm-5-2") is True + assert db_mod.model_is_reasoning("system.ai.kimi-k2-7-code") is True + + def test_unvalidated_families_are_not_marked_reasoning(self): + assert db_mod.model_is_reasoning("system.ai.inkling") is False + assert db_mod.model_is_reasoning("system.ai.qwen35-122b-a10b") is False + assert db_mod.model_is_reasoning("system.ai.gpt-oss-120b") is False class TestGptModelTokenLimits: @@ -1328,6 +1366,65 @@ def test_unversioned_names_sort_last_alphabetically(self): assert ordered[1:] == ["another-endpoint", "custom-endpoint"] +class TestDiscoverEndpointsWithApiType: + @pytest.mark.parametrize( + "payload", + [ + {"endpoints": None}, + {"endpoints": "not-a-list"}, + {"endpoints": [None, "not-an-endpoint"]}, + {"endpoints": [{"name": 123, "config": {}}]}, + {"endpoints": [{"name": "model", "config": None}]}, + {"endpoints": [{"name": "model", "config": {"served_entities": None}}]}, + {"endpoints": [{"name": "model", "config": {"served_entities": [None, "bad"]}}]}, + { + "endpoints": [ + { + "name": "model", + "config": {"served_entities": [{"foundation_model": None}]}, + } + ] + }, + { + "endpoints": [ + { + "name": "model", + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": "openai/v1/responses", + } + } + ] + }, + } + ] + }, + ], + ) + def test_malformed_payload_records_are_skipped(self, monkeypatch, payload): + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_endpoints_with_api_type(WS, "token", "openai/v1/responses") + + assert models == [] + assert reason and ("malformed" in reason or "no valid" in reason) + + def test_malformed_records_do_not_hide_valid_endpoint(self, monkeypatch): + payload = _foundation_models_payload(["databricks-gemini-3-5-flash"]) + payload["endpoints"].insert(0, None) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_endpoints_with_api_type( + WS, "token", "gemini/v1/generateContent" + ) + + assert models == ["databricks-gemini-3-5-flash"] + assert reason is None + + class TestDiscoverGeminiModels: def test_returns_newest_flash_first(self, monkeypatch): payload = _foundation_models_payload( @@ -1372,6 +1469,632 @@ def test_codex_discovery_orders_newest_version_first(self, monkeypatch): assert reason is None assert models == ["databricks-gpt-5-2-codex", "databricks-gpt-4-1"] + @pytest.mark.parametrize( + ("name", "api_type", "discover"), + [ + ("databricks-gpt-5-9", "openai/v1/responses", db_mod.discover_codex_models), + ( + "databricks-gemini-3-5-flash", + "gemini/v1/generateContent", + db_mod.discover_gemini_models, + ), + ], + ) + def test_explicitly_not_ready_endpoints_are_excluded( + self, monkeypatch, name, api_type, discover + ): + payload = { + "endpoints": [ + { + "name": name, + "state": {"ready": "NOT_READY"}, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": [api_type], + } + } + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = discover(WS, "token") + + assert models == [] + assert reason and "no ready endpoint" in reason + + def test_unready_endpoint_for_another_api_does_not_blame_readiness(self, monkeypatch): + # An unready Gemini-only endpoint must not make a Responses lookup claim + # "no READY endpoint exposes Responses" — nothing exposed Responses at all. + payload = { + "endpoints": [ + { + "name": "databricks-gemini-3-5-flash", + "state": {"ready": "NOT_READY"}, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": ["gemini/v1/generateContent"], + } + } + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_codex_models(WS, "token") + + assert models == [] + assert reason == "no endpoint exposes api_type `openai/v1/responses`" + assert "no ready endpoint" not in reason + + def test_duplicate_endpoint_names_are_deduplicated(self, monkeypatch): + endpoint = _foundation_models_payload(["databricks-gpt-5"])["endpoints"][0] + endpoint["config"]["served_entities"][0]["foundation_model"]["api_types"] = [ + "openai/v1/responses" + ] + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: ({"endpoints": [endpoint, endpoint]}, None), + ) + + models, reason = db_mod.discover_codex_models(WS, "token") + + assert reason is None + assert models == ["databricks-gpt-5"] + + +def _foundation_endpoint(name, api_types, *, v2=True, description=None): + foundation_model = { + "ai_gateway_v2_supported": v2, + "api_types": api_types, + } + if description is not None: + foundation_model["description"] = description + return { + "name": name, + "config": {"served_entities": [{"foundation_model": foundation_model}]}, + } + + +def _mlflow_chat_payload(names, *, api_type="mlflow/v1/chat/completions", v2=True): + return {"endpoints": [_foundation_endpoint(name, [api_type], v2=v2) for name in names]} + + +class TestDiscoverOssModels: + def test_finds_oss_endpoints_via_foundation_models(self, monkeypatch): + # Mirrors a workspace with no system.ai UC model-services: OSS models are + # plain databricks-* serving endpoints under the mlflow chat dialect. + payload = _mlflow_chat_payload( + [ + "databricks-glm-5-2", + "databricks-kimi-k2-7-code", + "databricks-inkling", + "databricks-qwen35-122b-a10b", + "databricks-gemma-3-12b", + ] + ) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_oss_models(WS, "token") + + assert reason is None + assert models == [ + "databricks-gemma-3-12b", + "databricks-glm-5-2", + "databricks-inkling", + "databricks-kimi-k2-7-code", + "databricks-qwen35-122b-a10b", + ] + + def test_excludes_claude_and_gemini_sharing_the_mlflow_dialect(self, monkeypatch): + # On some workspaces every foundation model advertises the mlflow chat + # dialect, so the api_type filter alone is too broad — the OSS family + # filter must drop Claude/Gemini and keep only the OSS cohort. + payload = _mlflow_chat_payload( + [ + "databricks-claude-opus-4-8", + "databricks-gemini-2-5-pro", + "databricks-glm-5-2", + "databricks-qwen3-embedding-0-6b", + ] + ) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_oss_models(WS, "token") + + assert reason is None + assert models == ["databricks-glm-5-2"] + + def test_reports_reason_when_no_oss_family_matches(self, monkeypatch): + payload = _mlflow_chat_payload(["databricks-claude-opus-4-8"]) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_oss_models(WS, "token") + + assert models == [] + assert reason is not None + assert "OSS" in reason + + +class TestDiscoverOssModelSpecs: + def test_explicitly_unready_endpoint_is_excluded(self, monkeypatch): + not_ready = _foundation_endpoint("databricks-inkling", ["mlflow/v1/chat/completions"]) + not_ready["state"] = {"ready": "NOT_READY"} + payload = { + "endpoints": [ + not_ready, + # Missing state remains compatible with older listings. + _foundation_endpoint("databricks-future-chat-1", ["mlflow/v1/chat/completions"]), + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert reason is None + assert [spec["id"] for spec in specs] == ["databricks-future-chat-1"] + + def test_explicitly_unready_endpoint_suppresses_static_fallback(self, monkeypatch): + not_ready = _foundation_endpoint("databricks-glm-5-2", ["mlflow/v1/chat/completions"]) + not_ready["state"] = {"ready": "NOT_READY"} + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: ({"endpoints": [not_ready]}, None), + ) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token", ["system.ai.glm-5-2"]) + + assert specs == [] + assert reason is not None + + def test_parses_reasoning_context_and_known_output_cap(self, monkeypatch): + payload = { + "endpoints": [ + { + "name": "databricks-inkling", + "capabilities": {"openai_reasoning": True}, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": ["mlflow/v1/chat/completions"], + "description": "Supports a context window of 1.5M tokens.", + } + } + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert reason is None + assert specs == [ + { + "id": "databricks-inkling", + "reasoning": True, + "context_window": 1_500_000, + "max_tokens": 65_536, + } + ] + + def test_uses_largest_mlflow_entity_context(self, monkeypatch): + payload = { + "endpoints": [ + { + "name": "databricks-future-chat-1", + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": ["mlflow/v1/chat/completions"], + "description": "context window of 128K tokens", + } + }, + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": ["mlflow/v1/chat/completions"], + "description": "supports a 500,000-token context window", + } + }, + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert reason is None + assert specs[0]["context_window"] == 500_000 + + def test_excludes_endpoint_with_native_api_and_malformed_entries(self, monkeypatch): + payload = { + "endpoints": [ + None, + {"name": "broken", "config": {"served_entities": "bad"}}, + { + "name": "databricks-qwen35-122b-a10b", + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": [ + "mlflow/v1/chat/completions", + "openai/v1/responses", + ], + } + } + ] + }, + }, + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert specs == [] + assert reason is not None + + def test_v2_and_mlflow_type_must_belong_to_same_entity(self, monkeypatch): + payload = { + "endpoints": [ + { + "name": "databricks-inkling", + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": [], + } + }, + { + "foundation_model": { + "ai_gateway_v2_supported": False, + "api_types": ["mlflow/v1/chat/completions"], + } + }, + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert specs == [] + assert reason is not None + + def test_duplicate_endpoint_ids_are_deduplicated(self, monkeypatch): + payload = _mlflow_chat_payload(["databricks-inkling", "databricks-inkling"]) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert reason is None + assert [spec["id"] for spec in specs] == ["databricks-inkling"] + + def test_uc_ids_receive_matching_endpoint_capabilities(self, monkeypatch): + payload = { + "endpoints": [ + { + "name": "databricks-qwen35-122b-a10b", + "capabilities": {"openai_reasoning": True}, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": ["mlflow/v1/chat/completions"], + "description": "context length of 128K tokens", + } + } + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token", ["system.ai.qwen35-122b-a10b"]) + + assert reason is None + assert specs == [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 25_000, + } + ] + + def test_explicit_native_metadata_suppresses_static_oss_fallback(self, monkeypatch): + payload = { + "endpoints": [_foundation_endpoint("databricks-glm-5-2", ["openai/v1/responses"])] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token", ["system.ai.glm-5-2"]) + + assert specs == [] + assert reason is not None + + def test_unavailable_metadata_keeps_static_glm_kimi_fallback(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token: (None, "HTTP 503 unavailable") + ) + + specs, reason = db_mod.discover_oss_model_specs( + WS, + "token", + ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code", "system.ai.inkling"], + ) + + assert reason is None + assert [spec["id"] for spec in specs] == [ + "system.ai.glm-5-2", + "system.ai.kimi-k2-7-code", + ] + + @pytest.mark.parametrize( + ("description", "expected"), + [ + ("supports a context window of 1.5M tokens", 1_500_000), + ("supports a 500,000-token context window", 500_000), + ("context length is 1 million tokens", 1_000_000), + ], + ) + def test_context_description_formats(self, description, expected): + assert db_mod._parse_context_window(description) == expected + + @pytest.mark.parametrize("description", ["", "context length of nope", "context window of 0K"]) + def test_malformed_context_description_degrades_to_none(self, monkeypatch, description): + payload = _mlflow_chat_payload(["databricks-inkling"]) + payload["endpoints"][0]["config"]["served_entities"][0]["foundation_model"][ + "description" + ] = description + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, _ = db_mod.discover_oss_model_specs(WS, "token") + + assert specs[0]["context_window"] is None + + +class TestDiscoverModelServicesDynamicOss: + def test_uc_first_broad_model_requires_matching_endpoint_validation(self, monkeypatch): + model_services = { + "model_services": [ + _model_service("system.ai.qwen35-122b-a10b"), + _model_service("system.ai.inkling"), + ] + } + foundation_models = _mlflow_chat_payload(["databricks-qwen35-122b-a10b"]) + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return foundation_models, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + _, _, _, oss, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert oss == ["system.ai.qwen35-122b-a10b"] + + def test_unknown_system_models_are_bucketed_by_live_api(self, monkeypatch): + model_services = { + "model_services": [ + _model_service("system.ai.future-coder-1"), + _model_service("system.ai.orion-1"), + _model_service("system.ai.future-chat-1"), + _model_service("system.ai.future-embed-1"), + ] + } + foundation_models = { + "endpoints": [ + _foundation_endpoint( + "databricks-future-coder-1", + ["openai/v1/responses", "mlflow/v1/chat/completions"], + description="supports a context window of 750K tokens", + ), + _foundation_endpoint( + "databricks-orion-1", + ["gemini/v1/generateContent", "mlflow/v1/chat/completions"], + ), + _foundation_endpoint("databricks-future-chat-1", ["mlflow/v1/chat/completions"]), + # Even misleading chat metadata cannot admit a non-chat service. + _foundation_endpoint("databricks-future-embed-1", ["openai/v1/responses"]), + ] + } + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return foundation_models, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + _, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert codex == ["system.ai.future-coder-1"] + assert gemini == ["system.ai.orion-1"] + assert oss == ["system.ai.future-chat-1"] + specs, spec_reason = db_mod.discover_responses_model_specs(WS, "token", codex) + assert spec_reason is None + assert specs == [{"id": "system.ai.future-coder-1", "context_window": 750_000}] + + def test_explicit_capabilities_override_known_name_fallback(self, monkeypatch): + model_services = { + "model_services": [ + _model_service("system.ai.grok-4-6"), + _model_service("system.ai.gemini-future"), + ] + } + foundation_models = { + "endpoints": [ + _foundation_endpoint("databricks-grok-4-6", []), + _foundation_endpoint("databricks-gemini-future", []), + ] + } + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return foundation_models, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + + assert (claude, codex, gemini, oss) == ({}, [], [], []) + assert reason is not None + + def test_claude_honours_explicit_endpoint_unavailability(self, monkeypatch): + # Claude must not be offered when the catalog explicitly says its endpoint + # is not ready — the same policy Codex/Gemini/OSS already apply. + model_services = {"model_services": [_model_service("system.ai.claude-sonnet-5")]} + endpoint = _foundation_endpoint("databricks-claude-sonnet-5", ["anthropic/v1/messages"]) + endpoint["state"] = {"ready": "NOT_READY"} + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return {"endpoints": [endpoint]}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + claude, _, _, _, _ = db_mod.discover_model_services(WS, "token") + + assert claude == {} + + def test_claude_survives_a_catalog_that_never_lists_it(self, monkeypatch): + # An ABSENT catalog entry is not explicit unavailability: a workspace whose + # foundation catalog omits Claude (or is unreadable) must keep working. + model_services = {"model_services": [_model_service("system.ai.claude-sonnet-5")]} + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return {"endpoints": []}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + claude, _, _, _, _ = db_mod.discover_model_services(WS, "token") + + assert claude == {"sonnet": "system.ai.claude-sonnet-5"} + + def test_gateway_models_missing_from_uc_are_offered(self, monkeypatch): + # The gateway catalog leads UC registration: these models are routable as + # `databricks-*` today, so they must not wait for a `system.ai.*` entry. + model_services = { + "model_services": [ + _model_service("system.ai.gemini-3-6-flash"), + _model_service("system.ai.kimi-k3"), + ] + } + foundation_models = { + "endpoints": [ + _foundation_endpoint("databricks-gemini-3-6-flash", ["gemini/v1/generateContent"]), + _foundation_endpoint("databricks-gemini-3-7-flash", ["gemini/v1/generateContent"]), + _foundation_endpoint("databricks-kimi-k3", ["mlflow/v1/chat/completions"]), + _foundation_endpoint("databricks-kimi-k3-neo", ["mlflow/v1/chat/completions"]), + _foundation_endpoint("databricks-grok-4-7", ["openai/v1/responses"]), + ] + } + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return foundation_models, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + _, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + # UC-registered models keep their system.ai id (no databricks-* twin); + # gemini/codex stay ordered newest-version-first. + assert gemini == ["databricks-gemini-3-7-flash", "system.ai.gemini-3-6-flash"] + assert oss == ["databricks-kimi-k3-neo", "system.ai.kimi-k3"] + assert codex == ["databricks-grok-4-7"] + + def test_unready_or_v1_only_gateway_endpoints_are_ignored(self, monkeypatch): + model_services = {"model_services": [_model_service("system.ai.gpt-5")]} + not_ready = _foundation_endpoint("databricks-gpt-5-9", ["openai/v1/responses"]) + not_ready["state"] = {"ready": "NOT_READY"} + ready = _foundation_endpoint("databricks-gpt-5-8", ["openai/v1/responses"]) + ready["state"] = {"ready": "READY"} + foundation_models = { + "endpoints": [ + _foundation_endpoint("databricks-gpt-5", ["openai/v1/responses"]), + not_ready, + ready, + # v1-only endpoints can't serve ucode's V2 routes. + _foundation_endpoint("databricks-gpt-5-7", ["openai/v1/responses"], v2=False), + # Non-chat services are never candidates. + _foundation_endpoint("databricks-bge-large-embed", ["openai/v1/responses"]), + ] + } + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return foundation_models, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + _, codex, _, _, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert codex == ["databricks-gpt-5-8", "system.ai.gpt-5"] + + def test_newest_claude_wins_across_mixed_id_spellings(self, monkeypatch): + # A gateway-only newer opus must beat the alphabetically-later system.ai id. + model_services = {"model_services": [_model_service("system.ai.claude-sonnet-4-6")]} + foundation_models = { + "endpoints": [ + _foundation_endpoint("databricks-claude-sonnet-4-6", ["anthropic/v1/messages"]), + _foundation_endpoint("databricks-claude-sonnet-5", ["anthropic/v1/messages"]), + ] + } + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return foundation_models, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + claude, _, _, oss, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert claude == {"sonnet": "databricks-claude-sonnet-5"} + assert oss == [] + class TestResolvePatToken: def test_reads_pat_profile_token_from_cfg(self, monkeypatch, tmp_path): @@ -2638,6 +3361,57 @@ def test_buckets_by_family(self, model_id, expected): assert classify_model_family(model_id) == expected +class TestFoundationModelsCache: + def test_discovery_consumers_share_one_successful_snapshot(self, monkeypatch): + calls = {"foundation": 0} + db_mod.clear_model_services_cache() + monkeypatch.setattr( + db_mod, + "_get_model_services_page", + lambda url, token: ( + {"model_services": [_model_service("system.ai.future-chat-1")]}, + None, + ), + ) + payload = _mlflow_chat_payload(["databricks-future-chat-1"]) + + def fake_get(url, token, timeout=10): + calls["foundation"] += 1 + return payload, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + _, _, _, oss, _ = db_mod.discover_model_services(WS, "tok") + specs, _ = db_mod.discover_oss_model_specs(WS, "tok", oss) + endpoints, _ = db_mod.discover_endpoints_with_api_type( + WS, "tok", "mlflow/v1/chat/completions" + ) + + assert oss == ["system.ai.future-chat-1"] + assert [spec["id"] for spec in specs] == oss + assert endpoints == ["databricks-future-chat-1"] + assert calls["foundation"] == 1 + + def test_failed_catalog_fetch_is_retried(self, monkeypatch): + calls = {"foundation": 0} + payload = _mlflow_chat_payload(["databricks-future-chat-1"]) + + def fake_get(url, token): + calls["foundation"] += 1 + if calls["foundation"] == 1: + return None, "HTTP 503 unavailable" + return payload, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + first, first_reason = db_mod._get_foundation_models_payload(WS, "tok") + second, second_reason = db_mod._get_foundation_models_payload(WS, "tok") + + assert first is None and first_reason == "HTTP 503 unavailable" + assert second == payload and second_reason is None + assert calls["foundation"] == 2 + + class TestModelServicesCache: """A successful listing is memoized per workspace: several callers want different views of the same paginated walk (bucketed families vs the raw Claude ids), so one `ucode setup` run would diff --git a/tests/test_e2e_uc.py b/tests/test_e2e_uc.py index c716dcb8..d7e20bdf 100644 --- a/tests/test_e2e_uc.py +++ b/tests/test_e2e_uc.py @@ -2,8 +2,8 @@ Verifies that `configure_shared_state` discovers models via UC model-services (`system.ai.*`) by default, falls back to the legacy per-family AI Gateway -listings when UC model-services are absent, and surfaces only `system.ai.*` -entries from the UC primitives. +listings when UC model-services are absent, and surfaces only `system.ai.*` UC +names or `databricks-*` gateway endpoint names from the discovery primitives. Run with: UCODE_TEST_WORKSPACE=https://your-workspace.databricks.com \ @@ -36,17 +36,17 @@ def _all_resolved_model_ids(state: dict) -> list[str]: # --------------------------------------------------------------------------- -# UC discovery primitives — verify the endpoints return only `system.ai.*` -# entries (the per-family/connection filters drop everything else). +# Model discovery returns UC `system.ai.*` names plus gateway-only +# `databricks-*` endpoint names; per-family filters must drop everything else. # --------------------------------------------------------------------------- class TestDiscoverModelServicesE2E: - def test_returns_only_system_ai_models(self, e2e_workspace, e2e_token): + def test_returns_only_uc_or_gateway_models(self, e2e_workspace, e2e_token): claude, codex, gemini, oss, reason = discover_model_services(e2e_workspace, e2e_token) if not (claude or codex or gemini or oss): - pytest.skip(f"No system.ai.* model services on workspace: {reason}") - non_system = sorted( + pytest.skip(f"No model services on workspace: {reason}") + unrelated = sorted( { m for m in _all_resolved_model_ids( @@ -57,10 +57,10 @@ def test_returns_only_system_ai_models(self, e2e_workspace, e2e_token): "oss_models": oss, } ) - if not m.startswith("system.ai.") + if not (m.startswith("system.ai.") or m.startswith("databricks-")) } ) - assert not non_system, f"Non-system.ai entries leaked through: {non_system[:5]}" + assert not unrelated, f"Unrelated model entries leaked through: {unrelated[:5]}" class TestListMcpServicesE2E: diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index cf5d5cdb..2d11832b 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -17,6 +17,7 @@ managed_provider_service, managed_state_overrides, managed_supplies_models, + managed_unclassifiable_models, managed_unservable_models, recommended_agent, resolve_state, @@ -461,6 +462,44 @@ def test_unclassifiable_models_are_dropped_from_buckets(self): "opencode_models": {"anthropic": ["system.ai.claude-opus-4-8"]} } + def test_unclassifiable_models_are_reported_without_affecting_known_families(self): + managed = { + "enabled_agents": { + "opencode": { + "model_config": { + "models": ["system.ai.future-chat-1", "system.ai.claude-opus-4-8"] + } + } + } + } + + assert managed_unclassifiable_models(managed, "opencode") == ["system.ai.future-chat-1"] + assert managed_state_overrides(managed, "opencode") == { + "opencode_models": {"anthropic": ["system.ai.claude-opus-4-8"]} + } + + def test_repeated_unclassifiable_model_is_reported_once(self): + # A manifest may legitimately repeat an id; the caller warns per entry, so + # duplicates here would mean duplicate identical warnings. + managed = { + "enabled_agents": { + "opencode": { + "model_config": { + "models": [ + "system.ai.future-chat-1", + "system.ai.future-chat-1", + "system.ai.other-unknown", + ] + } + } + } + } + + assert managed_unclassifiable_models(managed, "opencode") == [ + "system.ai.future-chat-1", + "system.ai.other-unknown", + ] + def test_no_override_when_nothing_is_servable(self): # An all-unservable list must not replace the developer's buckets with an empty dict — # that would leave OpenCode with no models at all. From 598eb1fe16166778d7a00f297aba400142dc0dc0 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:42:21 +1000 Subject: [PATCH 4/4] feat(opencode): route Responses models natively --- src/ucode/agents/__init__.py | 2 +- src/ucode/agents/opencode.py | 75 ++++++++++++++++++++++++++- src/ucode/cli.py | 8 ++- src/ucode/databricks.py | 3 ++ src/ucode/managed_resolve.py | 12 +++-- tests/conftest.py | 2 + tests/test_agent_opencode.py | 98 +++++++++++++++++++++++++++++++++++ tests/test_agents_init.py | 7 +++ tests/test_cli.py | 8 ++- tests/test_databricks.py | 3 +- tests/test_e2e.py | 12 +++++ tests/test_managed_resolve.py | 15 ++++-- 12 files changed, 231 insertions(+), 14 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 578aa208..49d17d1c 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -414,7 +414,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: _TOOL_DISCOVERY_SOURCES: dict[str, tuple[str, ...]] = { "claude": ("claude",), - "opencode": ("claude", "gemini", "oss"), + "opencode": ("claude", "codex", "gemini", "oss"), "codex": ("codex",), "gemini": ("gemini",), "copilot": ("claude", "codex"), diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 2bf10beb..2785ec9f 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -20,7 +20,9 @@ TOKEN_REFRESH_INTERVAL_SECONDS, build_opencode_base_urls, get_databricks_token, + gpt_model_token_limits, model_token_limits, + preferred_gpt_model, ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version @@ -41,13 +43,21 @@ PROVIDER_KEYS: list[list[str]] = [ ["provider", "databricks-anthropic"], ["provider", "databricks-google"], + ["provider", "databricks-openai"], ["provider", "databricks-oss"], ] def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -> str: """Return an OpenCode model selector in provider/model form when possible.""" - if model.startswith(("databricks-anthropic/", "databricks-google/", "databricks-oss/")): + if model.startswith( + ( + "databricks-anthropic/", + "databricks-google/", + "databricks-openai/", + "databricks-oss/", + ) + ): return model anthropic_models = opencode_models.get("anthropic") or [] @@ -58,6 +68,10 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) - if model in gemini_models: return f"databricks-google/{model}" + openai_models = opencode_models.get("openai") or [] + if model in openai_models: + return f"databricks-openai/{model}" + oss_models = opencode_models.get("oss") or [] if model in oss_models: return f"databricks-oss/{model}" @@ -129,12 +143,45 @@ def _oss_model_overlay( return overlay +def _responses_specs_by_id(raw_specs: object) -> dict[str, dict[str, object]]: + if not isinstance(raw_specs, list): + return {} + specs: dict[str, dict[str, object]] = {} + for raw_spec in raw_specs: + if not isinstance(raw_spec, dict): + continue + typed_spec = cast(dict[str, object], raw_spec) + model_id = typed_spec.get("id") + context = _positive_int(typed_spec.get("context_window")) + if isinstance(model_id, str) and model_id and context is not None: + specs.setdefault(model_id, typed_spec) + return specs + + +def _openai_model_overlay( + model: str, ua_header: dict[str, str], spec: dict[str, object] | None = None +) -> dict: + """Per-model Responses API options and explicit token limits.""" + limits = gpt_model_token_limits(model) + discovered_context = ( + _positive_int(spec.get("context_window")) if isinstance(spec, dict) else None + ) + if discovered_context is not None: + limits["context"] = discovered_context + return { + "headers": ua_header, + "limit": limits, + "options": {"useResponsesApi": True}, + } + + def render_overlay( model: str, token: str, opencode_base_urls: dict[str, str], opencode_models: dict[str, list[str]], oss_specs: list[dict] | None = None, + codex_specs: list[dict] | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for opencode.json.""" auth_headers = {"Authorization": f"Bearer {token}"} @@ -148,6 +195,7 @@ def render_overlay( anthropic_models = opencode_models.get("anthropic") or [] gemini_models = opencode_models.get("gemini") or [] + openai_models = opencode_models.get("openai") or [] oss_models = opencode_models.get("oss") or [] providers: dict = {} @@ -183,6 +231,27 @@ def render_overlay( "models": {m: {"headers": ua_header} for m in gemini_models}, } keys.append(["provider", "databricks-google"]) + if openai_models: + codex_specs_by_id = _responses_specs_by_id(codex_specs) + # @ai-sdk/openai supports both Responses and legacy chat completions. + # These models use the native `/ai-gateway/openai/v1/responses` route, + # so `useResponsesApi: true` lives in models..options where OpenCode + # reads it (provider-level options is read by the SDK only). + providers["databricks-openai"] = { + "npm": "@ai-sdk/openai", + "options": { + "baseURL": opencode_base_urls["openai"], + "apiKey": token, + "headers": auth_headers, + }, + "models": { + model_id: _openai_model_overlay( + model_id, ua_header, codex_specs_by_id.get(model_id) + ) + for model_id in openai_models + }, + } + keys.append(["provider", "databricks-openai"]) if oss_models: specs_by_id = _oss_specs_by_id(oss_specs) providers["databricks-oss"] = { @@ -226,6 +295,7 @@ def write_tool_config( opencode_base_urls, state.get("opencode_models") or {}, state.get("oss_model_specs") or [], + state.get("codex_model_specs") or [], ) existing = read_json_safe(OPENCODE_CONFIG_PATH) providers = existing.get("provider") @@ -286,6 +356,9 @@ def default_model(state: dict) -> str | None: anthropic = opencode_models.get("anthropic") or [] if anthropic: return anthropic[0] + openai = preferred_gpt_model(opencode_models.get("openai") or []) + if openai: + return openai gemini = opencode_models.get("gemini") or [] if gemini: return gemini[0] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index bfac3565..fd9abc19 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -147,7 +147,7 @@ _DISCOVERY_CONSUMERS: dict[str, tuple[str, ...]] = { "claude": ("claude", "opencode", "copilot", "pi"), - "codex": ("codex", "copilot", "pi"), + "codex": ("codex", "copilot", "opencode", "pi"), "gemini": ("gemini", "opencode", "pi"), "oss": ("opencode",), } @@ -649,7 +649,9 @@ def configure_shared_state( fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools ) want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools - want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools + want_codex = ( + fetch_all or "codex" in tools or "copilot" in tools or "opencode" in tools or "pi" in tools + ) # Codex smart routing can select OSS models such as GLM, so a Codex-only # configure must persist that discovered family too. want_oss = fetch_all or "opencode" in tools or "codex" in tools @@ -728,6 +730,8 @@ def configure_shared_state( opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: opencode_models["gemini"] = gemini_models + if codex_models: + opencode_models["openai"] = codex_models if oss_models: opencode_models["oss"] = oss_models diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a14a9ad0..f7d72ca2 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -4189,6 +4189,9 @@ def build_opencode_base_urls(workspace: str) -> dict[str, str]: return { "anthropic": build_tool_base_url("claude", workspace) + "/v1", "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", + # @ai-sdk/openai appends `/responses`. OpenCode speaks the native + # OpenAI-compatible API, not Codex CLI's coding-agent route. + "openai": f"{workspace}/ai-gateway/openai/v1", "oss": f"{workspace}/ai-gateway/mlflow/v1", } diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index c7579ba7..fecf21b7 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -112,8 +112,8 @@ def managed_unservable_models(managed: dict, tool: str) -> list[str]: Only non-empty when *every* named model is unservable, which is when the translation yields nothing and the developer's own models stand — so the caller can say why the admin's list had no - effect. opencode has no OpenAI provider and pi has no OSS provider, so each can be handed a - valid model FQN it cannot route. + effect. An agent can be handed a valid model FQN that none of its own + providers route, so the manifest names models it cannot serve. """ if tool not in ("opencode", "pi"): return [] @@ -151,15 +151,17 @@ def _manifest_models(managed: dict, tool: str) -> dict | list | None: def _bucket_by_provider(models: list[str]) -> dict[str, list[str]]: """Group model FQNs into OpenCode's provider buckets, mirroring how discovery builds them. - Discovery derives these from the per-family lists (claude -> anthropic, and gemini/oss as-is), so - the same family classification recovers them from a flat manifest list. Models whose family - can't be identified are dropped. + Discovery derives these from the per-family lists (claude -> anthropic, codex -> openai, and + gemini/oss as-is), so the same family classification recovers them from a flat manifest list. + Models whose family can't be identified are dropped. """ buckets: dict[str, list[str]] = {} for model in models: family = classify_model_family(model) if family in ANTHROPIC_FAMILIES: buckets.setdefault("anthropic", []).append(model) + elif family == "codex": + buckets.setdefault("openai", []).append(model) elif family in ("gemini", "oss"): buckets.setdefault(family, []).append(model) return buckets diff --git a/tests/conftest.py b/tests/conftest.py index c955761f..47113fa0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -105,6 +105,8 @@ def e2e_state(e2e_workspace, e2e_token): opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: opencode_models["gemini"] = gemini_models + if codex_models: + opencode_models["openai"] = codex_models if oss_models: opencode_models["oss"] = oss_models diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 8c9a6c7f..a3fce3e8 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -14,6 +14,7 @@ def _base_urls() -> dict[str, str]: return { "anthropic": f"{WS}/ai-gateway/anthropic/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", + "openai": f"{WS}/ai-gateway/openai/v1", "oss": f"{WS}/ai-gateway/mlflow/v1", } @@ -312,6 +313,81 @@ def test_prefixes_oss_model_with_provider_id(self): assert overlay["model"] == "databricks-oss/system.ai.kimi-k2-7-code" +class TestOpenAIProvider: + """OpenCode uses the native Responses route for GPT and Grok models.""" + + def test_openai_provider_added_when_codex_models_present(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + assert "databricks-openai" in overlay["provider"] + + def test_openai_provider_uses_native_sdk_and_gateway(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + provider = overlay["provider"]["databricks-openai"] + assert provider["npm"] == "@ai-sdk/openai" + assert provider["options"]["baseURL"] == f"{WS}/ai-gateway/openai/v1" + assert provider["options"]["headers"]["Authorization"] == "Bearer tok" + + def test_use_responses_api_set_on_every_codex_model(self): + model_ids = ["databricks-gpt-5-6-sol", "system.ai.grok-4-6"] + overlay, _ = opencode.render_overlay( + model_ids[0], "tok", _base_urls(), {"openai": model_ids} + ) + entries = overlay["provider"]["databricks-openai"]["models"] + assert all(entry["options"]["useResponsesApi"] is True for entry in entries.values()) + assert entries["databricks-gpt-5-6-sol"]["limit"] == { + "context": 1_050_000, + "output": 128_000, + } + assert entries["system.ai.grok-4-6"]["limit"]["context"] == 500_000 + + def test_live_responses_context_overrides_static_fallback(self): + model = "system.ai.future-coder-1" + specs = [{"id": model, "context_window": 750_000}] + overlay, _ = opencode.render_overlay( + model, "tok", _base_urls(), {"openai": [model]}, [], specs + ) + entry = overlay["provider"]["databricks-openai"]["models"][model] + assert entry["limit"] == {"context": 750_000, "output": 16_384} + + def test_openai_user_agent_header(self, monkeypatch): + monkeypatch.setattr(opencode, "ucode_version", lambda: "0.1.0") + monkeypatch.setattr(opencode, "agent_version", lambda binary: "0.74.0") + model = "databricks-gpt-5-6-sol" + overlay, _ = opencode.render_overlay(model, "tok", _base_urls(), {"openai": [model]}) + entry = overlay["provider"]["databricks-openai"]["models"][model] + assert entry["headers"]["User-Agent"] == "ucode/0.1.0 opencode/0.74.0" + + def test_selector_and_managed_keys_use_openai_provider(self): + model = "databricks-gpt-5-6-sol" + overlay, keys = opencode.render_overlay(model, "tok", _base_urls(), {"openai": [model]}) + assert overlay["model"] == f"databricks-openai/{model}" + assert ["provider", "databricks-openai"] in keys + assert ["provider", "databricks-openai"] in opencode.PROVIDER_KEYS + + def test_already_prefixed_openai_model_is_preserved(self): + model = "databricks-gpt-5-6-sol" + selector = f"databricks-openai/{model}" + overlay, _ = opencode.render_overlay(selector, "tok", _base_urls(), {"openai": [model]}) + assert overlay["model"] == selector + + def test_all_four_providers_can_coexist(self): + models = { + "anthropic": ["claude-sonnet"], + "gemini": ["gemini-2"], + "openai": ["databricks-gpt-5-6-sol"], + "oss": ["system.ai.kimi-k2-7-code"], + } + overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + assert set(overlay["provider"]) == { + "databricks-anthropic", + "databricks-google", + "databricks-openai", + "databricks-oss", + } + + class TestMcpServerConfig: # ucode registers the `ucode mcp-proxy ...` bridge as a `local` (stdio) MCP # server; the proxy handles token refresh, so no URL/bearer header here. @@ -420,6 +496,28 @@ def test_prefers_anthropic(self): state = {"opencode_models": {"anthropic": ["claude-sonnet"], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "claude-sonnet" + def test_falls_back_to_openai_before_gemini(self): + state = { + "opencode_models": { + "anthropic": [], + "openai": ["databricks-gpt-5-6-sol"], + "gemini": ["gemini-2"], + } + } + assert opencode.default_model(state) == "databricks-gpt-5-6-sol" + + def test_openai_fallback_chooses_newest_gpt(self): + state = { + "opencode_models": { + "openai": ["databricks-gpt-4-1", "databricks-gpt-5-5", "databricks-gpt-5-4"] + } + } + assert opencode.default_model(state) == "databricks-gpt-5-5" + + def test_gpt_oss_is_not_selected_for_responses(self): + state = {"opencode_models": {"openai": ["gpt-oss-120b"], "gemini": ["gemini-2"]}} + assert opencode.default_model(state) == "gemini-2" + def test_falls_back_to_gemini(self): state = {"opencode_models": {"anthropic": [], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "gemini-2" diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 5977cc0c..e039897c 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -193,6 +193,13 @@ def test_opencode_available(self): state = {"opencode_models": {"anthropic": ["claude-sonnet"]}} assert check_gateway_endpoint(state, "opencode") is True + def test_opencode_failure_detail_includes_codex_discovery(self): + detail = agents_mod._availability_failure_detail( + "opencode", {"_discovery_reasons": {"codex": "HTTP 503 unavailable"}} + ) + + assert "codex discovery: HTTP 503 unavailable" in detail + def test_copilot_available_with_claude(self): assert check_gateway_endpoint({"claude_models": {"sonnet": "s4"}}, "copilot") is True diff --git a/tests/test_cli.py b/tests/test_cli.py index b207c2ce..ea3a40ae 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2592,10 +2592,15 @@ def test_future_responses_model_persists_capability_spec(self, monkeypatch): lambda w, t, ids: ([{"id": model, "context_window": 750_000}], None), ) - state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + state = cli_mod.configure_shared_state( + self.WS, + profile="DEFAULT", + tools=["opencode"], + ) assert state["codex_models"] == [model] assert state["codex_model_specs"] == [{"id": model, "context_window": 750_000}] + assert state["opencode_models"]["openai"] == [model] assert saved[-1]["codex_models"] == [model] def test_uc_oss_ids_persist_matching_capability_specs(self, monkeypatch): @@ -2849,6 +2854,7 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): "databricks-claude-opus-4-8", "databricks-claude-sonnet-4-6", ], + "openai": ["databricks-gpt-5-6-sol"], "oss": ["databricks-glm-5-2"], } diff --git a/tests/test_databricks.py b/tests/test_databricks.py index c6d81ccc..e954e2f4 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -124,10 +124,11 @@ def test_unsupported_tool_raises(self): class TestBuildOpencodeBaseUrls: - def test_returns_anthropic_gemini_and_oss(self): + def test_returns_anthropic_gemini_openai_and_oss(self): urls = build_opencode_base_urls(WS) assert urls["anthropic"] == f"{WS}/ai-gateway/anthropic/v1" assert urls["gemini"] == f"{WS}/ai-gateway/gemini/v1beta" + assert urls["openai"] == f"{WS}/ai-gateway/openai/v1" assert urls["oss"] == f"{WS}/ai-gateway/mlflow/v1" diff --git a/tests/test_e2e.py b/tests/test_e2e.py index d304adca..85b4fed5 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -768,6 +768,18 @@ def _all_models(self, e2e_state: dict) -> list[tuple[str, str]]: out.append((provider, model)) return out + def test_all_models_includes_native_openai_provider(self): + models = self._all_models( + { + "opencode_models": { + "anthropic": ["claude-sonnet"], + "openai": ["system.ai.gpt-5-6-sol"], + } + } + ) + + assert ("openai", "system.ai.gpt-5-6-sol") in models + def test_launch_opencode_per_model( self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token ): diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index 2d11832b..6faf1286 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -407,6 +407,7 @@ def test_opencode_gets_provider_buckets_not_a_flat_list(self): "model_config": { "models": [ "system.ai.claude-opus-4-8", + "system.ai.gpt-5", "system.ai.gemini-3-flash", "system.ai.kimi-k2-7-code", ] @@ -417,6 +418,7 @@ def test_opencode_gets_provider_buckets_not_a_flat_list(self): assert managed_state_overrides(managed, "opencode") == { "opencode_models": { "anthropic": ["system.ai.claude-opus-4-8"], + "openai": ["system.ai.gpt-5"], "gemini": ["system.ai.gemini-3-flash"], "oss": ["system.ai.kimi-k2-7-code"], } @@ -433,6 +435,14 @@ def test_opencode_buckets_are_usable_by_its_own_writer(self): "databricks-anthropic/system.ai.claude-opus-4-8" ) + gpt_managed = { + "enabled_agents": {"opencode": {"model_config": {"models": ["system.ai.gpt-5"]}}} + } + gpt_buckets = managed_state_overrides(gpt_managed, "opencode")["opencode_models"] + assert opencode._resolve_model_selector("system.ai.gpt-5", gpt_buckets) == ( + "databricks-openai/system.ai.gpt-5" + ) + @pytest.mark.parametrize("tool", ["pi", "copilot"]) def test_pi_and_copilot_get_their_own_key(self, tool): # They compose from claude_models/codex_models/gemini_models, which claude, codex, and gemini @@ -568,10 +578,9 @@ def test_pi_oss_only_is_unservable(self): self._managed("pi", ["system.ai.kimi-k2-7-code"]), "pi" ) == ["system.ai.kimi-k2-7-code"] - def test_opencode_gpt_only_is_unservable(self): - # OpenCode has no OpenAI provider block. + def test_opencode_gpt_only_is_servable(self): managed = self._managed("opencode", ["system.ai.gpt-5"]) - assert managed_unservable_models(managed, "opencode") == ["system.ai.gpt-5"] + assert managed_unservable_models(managed, "opencode") == [] @pytest.mark.parametrize( ("tool", "models"),