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 98cde76a6be519043997b220b85993e89debfff8 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:59:21 +1000 Subject: [PATCH 4/4] feat(pi): add validated MLflow OSS streaming --- src/ucode/agents/__init__.py | 3 +- src/ucode/agents/_mlflow_proxy.py | 277 ++++++++++ src/ucode/agents/pi.py | 449 ++++++++++++++-- src/ucode/cli.py | 21 +- src/ucode/databricks.py | 1 + src/ucode/managed_resolve.py | 6 +- tests/test_agent_pi.py | 845 ++++++++++++++++++++++++++++-- tests/test_agents_init.py | 17 + tests/test_cli.py | 28 + tests/test_databricks.py | 11 +- tests/test_e2e.py | 15 + tests/test_managed_resolve.py | 12 +- tests/test_mlflow_proxy.py | 336 ++++++++++++ 13 files changed, 1895 insertions(+), 126 deletions(-) create mode 100644 src/ucode/agents/_mlflow_proxy.py create mode 100644 tests/test_mlflow_proxy.py diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 578aa208..c33626d2 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -408,6 +408,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: bool(state.get("claude_models")) or bool(state.get("codex_models")) or bool(state.get("gemini_models")) + or bool(state.get("oss_models")) ) return False @@ -418,7 +419,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: "codex": ("codex",), "gemini": ("gemini",), "copilot": ("claude", "codex"), - "pi": ("claude", "codex", "gemini"), + "pi": ("claude", "codex", "gemini", "oss"), } diff --git a/src/ucode/agents/_mlflow_proxy.py b/src/ucode/agents/_mlflow_proxy.py new file mode 100644 index 00000000..1112a6db --- /dev/null +++ b/src/ucode/agents/_mlflow_proxy.py @@ -0,0 +1,277 @@ +"""Loopback SSE-repair proxy for Pi's MLflow chat-completions provider. + +Some MLflow-served models omit the terminal OpenAI ``finish_reason``. Pi's +strict ``openai-completions`` parser rejects those streams. This loopback-only +proxy forwards requests without logging credentials or bodies and repairs only +successful SSE responses that have already produced data. Healthy SSE and all +non-streaming/error responses pass through unchanged. +""" + +from __future__ import annotations + +import json +from email.message import Message +from http.client import IncompleteRead +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import IO +from urllib import error as urllib_error +from urllib import request as urllib_request +from urllib.parse import urlsplit + +from ucode.gateway_proxy import HOP_BY_HOP_HEADERS +from ucode.ui import print_warning + +_STREAM_CHUNK = 8192 +_CHAT_COMPLETIONS_PATH = "/ai-gateway/mlflow/v1/chat/completions" +_SKIP_REQUEST_HEADERS = HOP_BY_HOP_HEADERS | {"accept-encoding"} +_ERROR_BODY = b'{"error":"MLflow proxy upstream unavailable"}\n' + + +class _NoRedirect(urllib_request.HTTPRedirectHandler): + """Keep authenticated requests pinned to the configured workspace origin.""" + + def redirect_request(self, req, fp, code, msg, headers, newurl): + return None + + +def _finish_chunk(chunk_id: str | None) -> bytes: + payload: dict = { + "object": "chat.completion.chunk", + "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}], + } + if chunk_id is not None: + payload["id"] = chunk_id + return json.dumps(payload, separators=(",", ":")).encode("utf-8") + + +def _data_payload(raw_line: bytes) -> bytes | None: + """Return an SSE data field's payload, accepting the optional one space.""" + stripped = raw_line.rstrip(b"\r\n") + if not stripped.startswith(b"data:"): + return None + payload = stripped[5:] + return payload[1:] if payload.startswith(b" ") else payload + + +def _forwarded_request_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]: + return { + key: value + for key, value in handler.headers.items() + if key.lower() not in _SKIP_REQUEST_HEADERS + } + + +def _safe_response_headers(headers: Message, *, streaming: bool) -> list[tuple[str, str]]: + safe: list[tuple[str, str]] = [] + for key, value in headers.items(): + lowered = key.lower() + if lowered in HOP_BY_HOP_HEADERS: + # A non-streaming body is unchanged, so preserving Content-Length + # avoids relying on EOF framing. Repaired streams can change size. + if lowered == "content-length" and not streaming: + safe.append((key, value)) + continue + safe.append((key, value)) + return safe + + +class _ProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + upstream_origin: str + + def log_message(self, format: str, *args: object) -> None: + return + + def do_POST(self) -> None: # noqa: N802 (stdlib handler API) + raw_length = self.headers.get("Content-Length", "0") + try: + length = int(raw_length) + if length < 0: + raise ValueError + except (TypeError, ValueError): + self._fixed_response(400, b'{"error":"invalid Content-Length"}\n') + return + + try: + body = self.rfile.read(length) if length else b"" + except OSError: + return + parsed_target = urlsplit(self.path) + if ( + parsed_target.scheme + or parsed_target.netloc + or parsed_target.fragment + or parsed_target.path != _CHAT_COMPLETIONS_PATH + ): + self._fixed_response(400, b'{"error":"invalid MLflow proxy request target"}\n') + return + target = self.upstream_origin.rstrip("/") + parsed_target.path + if parsed_target.query: + target += f"?{parsed_target.query}" + request = urllib_request.Request( + target, + data=body, + method="POST", + headers=_forwarded_request_headers(self), + ) + try: + opener = urllib_request.build_opener(_NoRedirect) + with opener.open(request, timeout=600) as response: # noqa: S310 + content_type = response.headers.get_content_type().lower() + if content_type == "text/event-stream": + self._relay_sse(response.status, response.headers, response) + else: + self._relay_verbatim(response.status, response.headers, response) + except urllib_error.HTTPError as exc: + # Relay upstream status, headers, and bytes verbatim. Never turn an + # upstream rejection into a successful repaired stream. + self._relay_verbatim(exc.code, exc.headers, exc) + except (urllib_error.URLError, OSError): + self._fixed_response(502, _ERROR_BODY) + + def _send_headers(self, status: int, headers: Message, *, streaming: bool) -> bool: + try: + self.send_response(status) + for key, value in _safe_response_headers(headers, streaming=streaming): + self.send_header(key, value) + # The proxy never reuses downstream connections. EOF framing is + # therefore safe for responses without Content-Length (including + # 204s and repaired SSE), and shutdown cannot leave a keep-alive + # client waiting on an otherwise complete response. + self.send_header("Connection", "close") + self.close_connection = True + self.end_headers() + return True + except (BrokenPipeError, ConnectionResetError, OSError): + return False + + def _relay_verbatim(self, status: int, headers: Message, stream: IO[bytes]) -> None: + if not self._send_headers(status, headers, streaming=False): + return + try: + while chunk := stream.read(_STREAM_CHUNK): + self.wfile.write(chunk) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError, OSError, IncompleteRead): + return + + def _relay_sse(self, status: int, headers: Message, stream: IO[bytes]) -> None: + if not self._send_headers(status, headers, streaming=True): + return + saw_data = False + saw_finish = False + saw_done = False + saw_error = False + last_id: str | None = None + event_data: list[bytes] = [] + + def inspect_event() -> None: + nonlocal saw_error, saw_finish, last_id + if not event_data: + return + payload = b"\n".join(event_data) + event_data.clear() + try: + event = json.loads(payload) + except (UnicodeDecodeError, json.JSONDecodeError): + return + if not isinstance(event, dict): + return + if "error" in event: + saw_error = True + event_id = event.get("id") + if isinstance(event_id, str): + last_id = event_id + choices = event.get("choices") + if isinstance(choices, list) and any( + isinstance(choice, dict) and choice.get("finish_reason") is not None + for choice in choices + ): + saw_finish = True + + try: + for raw_line in stream: + payload = _data_payload(raw_line) + event_line = raw_line.rstrip(b"\r\n") + if not event_line: + inspect_event() + elif event_line.lower().startswith(b"event:"): + event_name = event_line[6:] + if event_name.startswith(b" "): + event_name = event_name[1:] + if event_name.lower() == b"error": + saw_error = True + if payload == b"[DONE]": + inspect_event() + if saw_data and not saw_finish and not saw_error: + self._write(b"data: " + _finish_chunk(last_id) + b"\n\n") + saw_finish = True + self._write(raw_line) + saw_done = True + continue + if payload is not None and payload: + saw_data = True + event_data.append(payload) + self._write(raw_line) + except (BrokenPipeError, ConnectionResetError, OSError, IncompleteRead): + # Never turn a transport-failed partial stream into a successful + # synthetic completion. EOF without a transport exception remains + # repairable below because affected gateways can end cleanly after + # their final data event. + return + + # ``HTTPResponse`` line iteration can end without raising even when a + # declared Content-Length was not satisfied. A positive remainder is + # still a transport truncation, not a clean finish-reason omission. + remaining = getattr(stream, "length", None) + if isinstance(remaining, int) and remaining > 0: + return + + inspect_event() + if saw_data and not saw_error: + try: + if not saw_finish: + self._write(b"data: " + _finish_chunk(last_id) + b"\n\n") + if not saw_done: + self._write(b"data: [DONE]\n\n") + except (BrokenPipeError, ConnectionResetError, OSError): + return + + def _write(self, data: bytes) -> None: + self.wfile.write(data) + self.wfile.flush() + + def _fixed_response(self, status: int, body: bytes) -> None: + try: + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError, OSError): + return + + +class _Server(ThreadingHTTPServer): + daemon_threads = True + allow_reuse_address = True + + +def start(upstream_origin: str) -> tuple[ThreadingHTTPServer, str] | None: + """Bind a repair proxy to a fresh loopback port; the caller owns its lifecycle.""" + if not isinstance(upstream_origin, str) or not upstream_origin: + print_warning("MLflow stream repair proxy was not started: invalid upstream URL.") + return None + parsed_origin = urlsplit(upstream_origin) + if parsed_origin.scheme not in {"http", "https"} or not parsed_origin.netloc: + print_warning("MLflow stream repair proxy was not started: invalid upstream URL.") + return None + handler = type("_BoundProxyHandler", (_ProxyHandler,), {"upstream_origin": upstream_origin}) + try: + server = _Server(("127.0.0.1", 0), handler) + except OSError as exc: + print_warning(f"MLflow stream repair proxy was not started ({exc}).") + return None + port = int(server.server_address[1]) + return server, f"http://127.0.0.1:{port}" diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index bb44efb1..3691a16d 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -1,12 +1,13 @@ -"""Pi coding agent: writes the user's models.json with Databricks-backed providers. +"""Pi coding agent: writes the user's ~/.pi/agent/models.json with Databricks-backed providers. -Pi (https://pi.dev) is a multi-provider coding agent. We register three +Pi (https://pi.dev) is a multi-provider coding agent. We register four 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/openai/v1 - `databricks-gemini` (api: google-generative-ai) → /ai-gateway/gemini/v1beta +- `databricks-mlflow` (api: openai-completions) → /ai-gateway/mlflow/v1 Per-provider `compat` flags work around fields the gateway translators reject: @@ -15,11 +16,19 @@ pi uses for every request. With this flag pi omits the per-tool field and sends the legacy `anthropic-beta: fine-grained-tool-streaming-...` header instead, which the gateway accepts. - -OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via -pi today — they live behind /ai-gateway/mlflow/v1 with per-model -`max_tokens` caps that pi has no global way to honor without per-model -config we don't currently maintain. +- mlflow: `supportsStore: false` and `supportsStrictMode: false` — the MLflow + chat-completions gateway rejects OpenAI's `store` field and + `tools[].function.strict`. +- openai: no `compat` flags needed, but the per-model `thinkingLevelMap` + matters — see `_pi_gpt_model_entry`. Declaring `reasoning: true` without an + off-state makes Pi send `reasoning: {effort: "none"}`, which `gpt-5`, + `gpt-5-mini`, `gpt-5-nano` and `gpt-5-5-pro` reject with a 400. + +The `databricks-mlflow` provider carries validated MLflow chat-completions +models discovered upstream. Per-model reasoning and token metadata comes from +the persisted gateway capability specs, with conservative static fallback. +At launch this provider alone is routed through a loopback repair proxy because +some models (notably Inkling) omit the terminal `finish_reason` Pi requires. The bearer token is baked into the file and refreshed by a background thread while the session runs (same pattern as OpenCode/Copilot). @@ -27,12 +36,18 @@ from __future__ import annotations +import ipaddress import os import signal import subprocess import threading +from collections.abc import Callable from pathlib import Path +from types import FrameType +from typing import Any, cast +from urllib.parse import urlparse +from ucode.agents import _mlflow_proxy from ucode.config_io import ( APP_DIR, ToolSpec, @@ -50,14 +65,16 @@ discover_claude_models_unbucketed, get_databricks_token, gpt_model_token_limits, + model_is_reasoning, + model_token_limits, preferred_gpt_model, ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version +from ucode.ui import print_warning -# 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. +# Use Pi's standard user configuration directory. Keep ucode's backups and +# state under APP_DIR, but let Pi share the user's normal ~/.pi/agent config. PI_CONFIG_DIR = Path.home() / ".pi" / "agent" PI_CONFIG_PATH = PI_CONFIG_DIR / "models.json" PI_SETTINGS_PATH = PI_CONFIG_DIR / "settings.json" @@ -66,6 +83,7 @@ # 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" +_CONFIG_WRITE_LOCK = threading.RLock() SPEC: ToolSpec = { "binary": "pi", @@ -79,6 +97,7 @@ "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-mlflow", ) PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES] @@ -93,6 +112,7 @@ def _resolve_model_selector( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + oss_models: list[str], claude_model_ids: list[str] | None = None, ) -> str: """Return a Pi model selector in `/` form when possible.""" @@ -107,11 +127,17 @@ def _resolve_model_selector( return f"databricks-openai/{model}" if model in gemini_models: return f"databricks-gemini/{model}" + if model in oss_models: + return f"databricks-mlflow/{model}" return model def _pi_claude_model_entry(model_id: str) -> dict: - """Build a Claude entry with explicit context and thinking metadata.""" + """Build a Claude entry with explicit limits. + + Databricks model ids do not match Pi's built-in Anthropic ids, so a bare + custom entry silently gets Pi's 128k context / 4k output defaults. + """ capabilities = claude_model_capabilities(model_id) entry: dict = { "id": model_id, @@ -122,18 +148,121 @@ def _pi_claude_model_entry(model_id: str) -> dict: } if capabilities.force_adaptive_thinking: entry["compat"] = {"forceAdaptiveThinking": True} + # Pi hides its extended levels unless custom models declare them. All + # adaptive Claude models support `max`; native `xhigh` is limited to + # Opus 4.7/4.8, Sonnet 5, and Fable 5. 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.""" +_OSS_SAFE_LIMITS = {"context": 128_000, "output": 8_192} + + +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 _pi_oss_model_entry(model_id: str, spec: dict[str, object] | None = None) -> dict: + """Build a Pi MLflow model entry from discovered or static capabilities. + + A valid discovered boolean overrides static reasoning. Any discovered spec + receives a complete conservative limit pair, so missing capability fields + cannot leave a validated model effectively uncapped. Missing specs retain + the existing static GLM/Kimi/DeepSeek behavior, while unknown models remain bare. + """ + entry: dict = {"id": model_id} + static_limits = model_token_limits(model_id) + static_reasoning = model_is_reasoning(model_id) + + reasoning = spec.get("reasoning") if isinstance(spec, dict) else None + if not isinstance(reasoning, bool): + reasoning = static_reasoning + if reasoning: + entry["reasoning"] = True + + 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): + entry["contextWindow"] = context or ( + static_limits.get("context") if static_limits else _OSS_SAFE_LIMITS["context"] + ) + entry["maxTokens"] = output or ( + static_limits.get("output") if static_limits else _OSS_SAFE_LIMITS["output"] + ) + elif static_limits: + entry["contextWindow"] = static_limits["context"] + entry["maxTokens"] = static_limits["output"] + return entry + + +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 _pi_gpt_model_entry(model_id: str, spec: dict[str, object] | None = None) -> dict: + """Build a Pi Responses model entry with explicit limits and reasoning. + + Gateway ids aren't in Pi's built-in catalog, so capabilities must be + declared here rather than inherited from Pi's vendor model registry. + + `thinkingLevelMap: {"off": None}` is required alongside `reasoning: True`. + When a model declares `reasoning` but no off-state, Pi's Responses builder + falls back to `reasoning: {effort: "none"}` for the thinking-off case + (`pi-ai/dist/api/openai-responses.js`, the `thinkingLevelMap?.off !== null` + branch). `"none"` is only valid on gpt-5.1+, so `gpt-5`, `gpt-5-mini`, + `gpt-5-nano` and `gpt-5-5-pro` reject every request with + `BAD_REQUEST: Unsupported value: 'none' is not supported with the 'gpt-5' + model`. An explicit `None` makes Pi omit `reasoning` entirely, which the + gateway accepts for all ids. + """ limits = gpt_model_token_limits(model_id) + discovered_context = ( + _positive_int(spec.get("context_window")) if isinstance(spec, dict) else None + ) entry: dict = { "id": model_id, - "contextWindow": limits["context"], + "contextWindow": discovered_context or limits["context"], "maxTokens": limits["output"], } normalized_id = model_id.rsplit("/", 1)[-1].lower() @@ -143,8 +272,9 @@ def _pi_gpt_model_entry(model_id: str) -> dict: 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. + # Grok 4.6 is reasoning-only. Databricks supports low/medium/high/xhigh; + # hide Pi's unsupported off/minimal/max choices instead of mapping them + # to values the gateway would reject. entry["reasoning"] = True entry["thinkingLevelMap"] = { "off": None, @@ -155,8 +285,6 @@ def _pi_gpt_model_entry(model_id: str) -> dict: 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 @@ -168,7 +296,10 @@ def render_overlay( claude_models: dict[str, str], codex_models: list[str], gemini_models: list[str], + oss_models: list[str], + oss_specs: list[dict] | None = None, claude_model_ids: list[str] | None = None, + codex_specs: list[dict] | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Pi's user agent config.""" providers: dict = {} @@ -186,20 +317,30 @@ def render_overlay( "authHeader": True, # Gateway's Anthropic translator rejects per-tool # `eager_input_streaming` on the streaming + tools path. Pi sends - # the legacy beta header instead when this is false. - "compat": {"supportsEagerToolInputStreaming": False}, + # the legacy beta header instead when this is false. Session + # affinity keeps a Pi conversation on one AI Gateway destination, + # which is important when traffic splitting is configured because + # prompt caches are destination-local. + "compat": { + "supportsEagerToolInputStreaming": False, + "sendSessionAffinityHeaders": True, + }, "headers": ua_headers, "models": [_pi_claude_model_entry(m) for m in claude_ids], } keys.append(["providers", "databricks-claude"]) if codex_models: + codex_specs_by_id = _responses_specs_by_id(codex_specs) providers["databricks-openai"] = { "baseUrl": pi_base_urls["openai"], "api": "openai-responses", "apiKey": token, "authHeader": True, "headers": ua_headers, - "models": [_pi_gpt_model_entry(m) for m in codex_models], + "models": [ + _pi_gpt_model_entry(model_id, codex_specs_by_id.get(model_id)) + for model_id in codex_models + ], } keys.append(["providers", "databricks-openai"]) if gemini_models: @@ -212,9 +353,28 @@ def render_overlay( "models": [{"id": m} for m in gemini_models], } keys.append(["providers", "databricks-gemini"]) + if oss_models: + specs_by_id = _oss_specs_by_id(oss_specs) + providers["databricks-mlflow"] = { + "baseUrl": pi_base_urls["oss"], + "api": "openai-completions", + "apiKey": token, + "authHeader": True, + # MLflow chat-completions gateway rejects OpenAI's `store` field + # and per-tool `strict`. Pi omits both when these are false. + "compat": {"supportsStore": False, "supportsStrictMode": False}, + "headers": ua_headers, + "models": [_pi_oss_model_entry(m, specs_by_id.get(m)) for m in oss_models], + } + keys.append(["providers", "databricks-mlflow"]) overlay: dict = { "model": _resolve_model_selector( - model, claude_models, codex_models, gemini_models, claude_model_ids + model, + claude_models, + codex_models, + gemini_models, + oss_models, + claude_model_ids, ), } if providers: @@ -228,6 +388,17 @@ def write_tool_config( token: str | None = None, *, force_refresh: bool = False, +) -> tuple[dict, str]: + with _CONFIG_WRITE_LOCK: + return _write_tool_config_unlocked(state, model, token, force_refresh=force_refresh) + + +def _write_tool_config_unlocked( + state: dict, + model: str, + token: str | None = None, + *, + force_refresh: bool = False, ) -> tuple[dict, str]: backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH) if token is None: @@ -236,15 +407,19 @@ 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_model_ids: list[str] | None = None 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 [] + oss_models = state.get("oss_models") or [] + # The shared map keeps one model per family; Pi supplements it with + # the full inventory for every family that discovery enabled. 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_models, codex_models, gemini_models, oss_models = managed_families claude_model_ids = _managed_pi_claude_models(state) overlay, managed_keys = render_overlay( model, @@ -253,7 +428,10 @@ def write_tool_config( claude_models, codex_models, gemini_models, + oss_models, + state.get("oss_model_specs") or [], claude_model_ids, + state.get("codex_model_specs") or [], ) existing = read_json_safe(PI_CONFIG_PATH) providers = existing.get("providers") @@ -262,7 +440,7 @@ def write_tool_config( providers.pop(stale, None) merged = deep_merge_dict(existing, overlay) write_json_file(PI_CONFIG_PATH, merged) - _write_settings(overlay["model"]) + _write_settings(overlay["model"], clear_unresolved=managed_families is not None) state = mark_tool_managed(state, "pi", managed_keys) save_state(state) return state, token @@ -281,15 +459,24 @@ def _managed_pi_claude_models(state: dict) -> list[str]: 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.""" + """Return all Claude ids Pi may offer without changing shared routing state. + + Shared discovery intentionally keeps one model per Claude family and pins + Opus to 4.8 while the smart router requires that arm. Pi is a model picker, + so it can safely expose newer versions as long as its default remains the + shared pinned id. + Cache the supplemental inventory in state after the first Pi config write; + a failed supplemental request degrades to the shared family picks. + """ allowed_families = set(claude_models) cached = state.get("pi_claude_models") if isinstance(cached, list): - return [ + cached_models = [ model for model in cached if isinstance(model, str) and classify_model_family(model) in allowed_families ] + return list(dict.fromkeys([*claude_models.values(), *cached_models])) try: discovered, _ = discover_claude_models_unbucketed(state["workspace"], token) @@ -297,29 +484,42 @@ def _discover_pi_claude_models(state: dict, token: str, claude_models: dict[str, discovered = [] if discovered: state["pi_claude_models"] = discovered - return [model for model in discovered if classify_model_family(model) in allowed_families] + discovered_models = [ + model for model in discovered if classify_model_family(model) in allowed_families + ] + return list(dict.fromkeys([*claude_models.values(), *discovered_models])) return list(claude_models.values()) -def _write_settings(model_selector: str) -> None: +def _write_settings(model_selector: str, *, clear_unresolved: bool = False) -> None: # Pin defaultProvider/defaultModel in settings.json so Pi doesn't fall # through to an env-key-backed provider (e.g. HF_TOKEN exposing # huggingface) in `findInitialModel` when no --model is passed. provider, _, model_id = model_selector.partition("/") - if not model_id: + if not model_id and not clear_unresolved: return backup_existing_file(PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH) existing = read_json_safe(PI_SETTINGS_PATH) + if not model_id: + # A non-empty managed allowlist with no servable model must not retain a + # stale default that bypasses the current administrator policy. + existing.pop("defaultProvider", None) + existing.pop("defaultModel", None) + write_json_file(PI_SETTINGS_PATH, existing) + return merged = deep_merge_dict(existing, {"defaultProvider": provider, "defaultModel": model_id}) write_json_file(PI_SETTINGS_PATH, merged) -def _managed_model_families(state: dict) -> tuple[dict[str, str], list[str], list[str]] | None: - """Split a managed config's ``pi_models`` into the per-family inputs Pi's providers need. +def _managed_model_families( + state: dict, +) -> tuple[dict[str, str], list[str], list[str], list[str]] | None: + """Split a managed config's ``pi_models`` into Pi provider inputs. - Pi builds one provider block per family, so a flat list has to be classified back out. Returns - None when the managed models yield no family Pi can serve, leaving the workspace-wide discovery - lists in play rather than writing a config with no usable provider. + ``None`` means there is no managed allowlist and workspace discovery may be + used. A present non-empty allowlist always returns four collections—even if + every entry is malformed or unsupported—so unlisted discovered models can + never leak back into managed Pi configuration. """ managed = state.get("pi_models") if not isinstance(managed, list) or not managed: @@ -327,6 +527,7 @@ def _managed_model_families(state: dict) -> tuple[dict[str, str], list[str], lis claude: dict[str, str] = {} codex: list[str] = [] gemini: list[str] = [] + oss: list[str] = [] for model in managed: if not isinstance(model, str) or not model.strip(): continue @@ -337,13 +538,14 @@ def _managed_model_families(state: dict) -> tuple[dict[str, str], list[str], lis codex.append(model) elif family == "gemini": gemini.append(model) - if not (claude or codex or gemini): - return None - return claude, codex, gemini + elif family == "oss": + oss.append(model) + return claude, codex, gemini, oss def default_model(state: dict) -> str | None: - """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini. + """Prefer a managed Pi default/allowlist, then Claude opus → sonnet → haiku; + fall back to codex, Gemini, then OSS. A managed config's ``pi_default_model`` and ``pi_models`` both win outright: the former is the admin's chosen session start, the latter their allowlist. Workspace-wide discovery falls back. @@ -361,7 +563,10 @@ def default_model(state: dict) -> str | None: if codex_model: return codex_model gemini_models = state.get("gemini_models") or [] - return gemini_models[0] if gemini_models else None + if gemini_models: + return gemini_models[0] + oss_models = state.get("oss_models") or [] + return oss_models[0] if oss_models else None def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: @@ -387,27 +592,161 @@ def build_runtime_env(token: str) -> dict[str, str]: return env -def launch(state: dict, tool_args: list[str]) -> None: - token = _refresh_token_once(state) - env = build_runtime_env(token) +def _is_loopback_origin(origin: str) -> bool: + hostname = urlparse(origin).hostname + if not hostname: + return True + if hostname.lower() == "localhost": + return True + try: + return ipaddress.ip_address(hostname).is_loopback + except ValueError: + return False - stop_event = threading.Event() - refresher = threading.Thread( - target=_refresh_forever, - args=(state, stop_event), - daemon=True, + +def _start_oss_proxy( + state: dict, +) -> tuple[threading.Thread, _mlflow_proxy.ThreadingHTTPServer] | None: + """Start Pi's MLflow repair proxy and rewrite only the in-memory base URL. + + The upstream is always derived from ``workspace`` rather than an existing + base URL, so a stale loopback port from an old config can never become the + next proxy's upstream. + """ + managed_families = _managed_model_families(state) + oss_models = state.get("oss_models") or [] if managed_families is None else managed_families[3] + if not oss_models: + return None + direct_oss_url = build_pi_base_urls(state["workspace"])["oss"] + origin = direct_oss_url.split("/ai-gateway/", 1)[0] + if _is_loopback_origin(origin): + print_warning("MLflow stream repair proxy skipped for a loopback workspace URL.") + return None + started = _mlflow_proxy.start(origin) + if started is None: + return None + server, proxy_origin = started + thread = threading.Thread(target=server.serve_forever, daemon=True) + try: + thread.start() + except RuntimeError: + server.server_close() + print_warning("MLflow stream repair proxy could not start; using the direct gateway URL.") + return None + pi_urls = state.setdefault("base_urls", {}).setdefault( + "pi", build_pi_base_urls(state["workspace"]) ) - refresher.start() + pi_urls["oss"] = f"{proxy_origin}/ai-gateway/mlflow/v1" + return thread, server + + +def _restore_direct_oss_config(state: dict, token: str | None) -> None: + """Replace the session-only proxy URL after any in-flight config write.""" + with _CONFIG_WRITE_LOCK: + _restore_direct_oss_config_unlocked(state, token) + + +def _restore_direct_oss_config_unlocked(state: dict, token: str | None) -> None: + pi_urls = state.setdefault("base_urls", {}).setdefault( + "pi", build_pi_base_urls(state["workspace"]) + ) + pi_urls["oss"] = build_pi_base_urls(state["workspace"])["oss"] + model = default_model(state) + if model and token is not None: + write_tool_config(state, model, token=token) + return + + # Token acquisition can fail before the normal config rewrite returns a + # token. Repair an existing generated provider in place without changing + # its credential, so a stale loopback port is never left behind. + existing = read_json_safe(PI_CONFIG_PATH) + providers = existing.get("providers") + mlflow = providers.get("databricks-mlflow") if isinstance(providers, dict) else None + if isinstance(mlflow, dict): + mlflow["baseUrl"] = pi_urls["oss"] + write_json_file(PI_CONFIG_PATH, existing) + + +def launch(state: dict, tool_args: list[str]) -> None: + proxy: tuple[threading.Thread, _mlflow_proxy.ThreadingHTTPServer] | None = None + stop_event = threading.Event() + refresher: threading.Thread | None = None + proc: subprocess.Popen | None = None + token: str | None = None + primary_error: BaseException | None = None + previous_signal_handlers: dict[ + signal.Signals, int | Callable[[int, FrameType | None], Any] | None + ] = {} + termination_requested = False + + def handle_termination(signum: int, _frame: FrameType | None) -> None: + """Forward catchable terminal shutdowns and unwind through proxy cleanup.""" + nonlocal termination_requested + if termination_requested: + return + termination_requested = True + if proc is not None and proc.poll() is None: + try: + proc.send_signal(signum) + except OSError: + pass + raise SystemExit(128 + signum) - proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env) try: - returncode = proc.wait() - except KeyboardInterrupt: - proc.send_signal(signal.SIGINT) - returncode = proc.wait() + for signal_name in ("SIGTERM", "SIGHUP"): + signum = getattr(signal, signal_name, None) + if signum is None: + continue + # Record the previous handler only AFTER a successful install, so the + # restore loop below can never put back a handler we never replaced + # (and so a failed install isn't masked by a failing restore). + previous_handler = signal.getsignal(signum) + signal.signal(signum, handle_termination) + previous_signal_handlers[signum] = previous_handler + try: + # The proxy must be live and its URL in state before the first config + # write; refreshes then keep writing the same live loopback endpoint. + proxy = _start_oss_proxy(state) + token = _refresh_token_once(state) + env = build_runtime_env(token) + + refresher = threading.Thread( + target=_refresh_forever, + args=(state, stop_event), + daemon=True, + ) + refresher.start() + + proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() + except BaseException as exc: + primary_error = exc + raise + finally: + stop_event.set() + if refresher is not None: + refresher.join(timeout=1) + if proxy is not None: + proxy_thread, server = proxy + restore_error: Exception | None = None + try: + _restore_direct_oss_config(state, token) + except Exception as exc: + restore_error = exc + print_warning(f"Pi MLflow direct configuration could not be restored ({exc}).") + finally: + server.shutdown() + server.server_close() + proxy_thread.join(timeout=1) + if restore_error is not None and primary_error is None: + raise restore_error finally: - stop_event.set() - refresher.join(timeout=1) + for signum, previous_handler in previous_signal_handlers.items(): + signal.signal(signum, previous_handler) raise SystemExit(returncode) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index bfac3565..ffb6f417 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -81,6 +81,7 @@ managed_launch_model, managed_provider_family_models, managed_provider_service, + managed_state_overrides, managed_supplies_models, managed_unclassifiable_models, managed_unservable_models, @@ -149,7 +150,7 @@ "claude": ("claude", "opencode", "copilot", "pi"), "codex": ("codex", "copilot", "pi"), "gemini": ("gemini", "opencode", "pi"), - "oss": ("opencode",), + "oss": ("codex", "opencode", "pi"), } @@ -652,7 +653,7 @@ def configure_shared_state( want_codex = fetch_all or "codex" in tools or "copilot" 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 + want_oss = fetch_all or "opencode" in tools or "codex" in tools or "pi" in tools claude_reason: str | None = None gemini_reason: str | None = None @@ -2023,9 +2024,23 @@ def _launch_tool( _warn_unclassifiable_managed_models(managed, tool) unservable = managed_unservable_models(managed, tool) if unservable: + # Only agents whose translation yields no state override actually + # fall back to discovery. Pi keeps its managed `pi_models` list + # even when nothing in it is servable (so unlisted models can + # never leak into a managed config), so claiming a fallback there + # would be false. + falls_back = f"{tool}_models" not in managed_state_overrides(managed, tool) + detail = ( + "using your discovered models instead." + if falls_back + else ( + "and its list is authoritative, so no model can be configured. " + "Ask your workspace admin to publish a supported model." + ) + ) print_warning( f"Your workspace's managed config lists no {TOOL_SPECS[tool]['display']}-servable " - f"models ({', '.join(unservable)}); using your discovered models instead." + f"models ({', '.join(unservable)}); {detail}" ) elif managed_agent_config_enabled(): print_note("No managed coding agent config found; using your own settings") diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a14a9ad0..13150acb 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -4209,6 +4209,7 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]: "claude": build_tool_base_url("claude", workspace), "openai": f"{workspace}/ai-gateway/openai/v1", "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", + "oss": f"{workspace}/ai-gateway/mlflow/v1", } diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index c7579ba7..9d9c5656 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 [] @@ -126,7 +126,7 @@ def managed_unservable_models(managed: dict, tool: str) -> list[str]: else [ m for m in models - if classify_model_family(m) in (*ANTHROPIC_FAMILIES, "codex", "gemini") + if classify_model_family(m) in (*ANTHROPIC_FAMILIES, "codex", "gemini", "oss") ] ) return [] if servable else models diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 62f12b5c..653bacb4 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -3,9 +3,10 @@ from __future__ import annotations import json +import threading from contextlib import nullcontext from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -21,6 +22,7 @@ def _base_urls() -> dict[str, str]: "claude": f"{WS}/ai-gateway/anthropic", "openai": f"{WS}/ai-gateway/openai/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", + "oss": f"{WS}/ai-gateway/mlflow/v1", } @@ -30,7 +32,9 @@ def _empty() -> dict: "claude_models": {}, "codex_models": [], "gemini_models": [], - "claude_model_ids": None, + "oss_models": [], + "oss_specs": [], + "codex_specs": [], } @@ -44,7 +48,10 @@ def _overlay(model: str, token: str = "tok", **kwargs): bundle["claude_models"], bundle["codex_models"], bundle["gemini_models"], - bundle["claude_model_ids"], + bundle["oss_models"], + bundle["oss_specs"], + None, + bundle["codex_specs"], ) @@ -58,10 +65,10 @@ def test_package(self): def test_display(self): assert pi.SPEC["display"] == "Pi" - def test_config_path_under_pi_agent_dir(self): + def test_config_path_uses_standard_pi_agent_dir(self): assert pi.SPEC["config_path"].name == "models.json" - assert pi.SPEC["config_path"].parent.name == "agent" assert pi.PI_CONFIG_DIR == Path.home() / ".pi" / "agent" + assert pi.SPEC["config_path"].parent == pi.PI_CONFIG_DIR @pytest.mark.parametrize( ("new_name", "legacy_name"), @@ -105,26 +112,117 @@ def test_openai_provider_uses_openai_responses(self): assert provider["api"] == "openai-responses" assert provider["baseUrl"] == f"{WS}/ai-gateway/openai/v1" - def test_claude_entries_pin_limits_and_extended_thinking_levels(self): + def test_gpt56_sol_model_entry_pins_1m_context(self): + # Gateway ids are custom to Pi, so explicit metadata is required to + # avoid its 128k custom-model default. + overlay, _ = _overlay("gpt-5-6-sol", codex_models=["gpt-5-6-sol"]) + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert entry["id"] == "gpt-5-6-sol" + assert entry["contextWindow"] == 1_050_000 + assert entry["maxTokens"] == 128_000 + assert entry["reasoning"] is True + assert entry["input"] == ["text", "image"] + + def test_gpt_entries_pin_off_thinking_level_to_none(self): + # `reasoning: True` without an off-state makes Pi send + # `reasoning: {effort: "none"}`, which gpt-5 / -mini / -nano / -5-5-pro + # reject with a 400. `{"off": None}` makes Pi omit `reasoning`. overlay, _ = _overlay( - "system.ai.claude-opus-4-8", + "system.ai.gpt-5", + codex_models=[ + "system.ai.gpt-5", + "system.ai.gpt-5-mini", + "system.ai.gpt-5-nano", + "system.ai.gpt-5-5-pro", + "system.ai.gpt-5-6-luna", + ], + ) + entries = overlay["providers"]["databricks-openai"]["models"] + assert entries, "expected gpt entries" + for entry in entries: + assert entry["reasoning"] is True + assert entry["thinkingLevelMap"] == {"off": None}, entry["id"] + + def test_non_gpt_codex_entry_has_no_thinking_level_map(self): + # Only the gpt-5 family declares `reasoning`, so only it needs the + # off-state override. + overlay, _ = _overlay("gpt-oss-120b", codex_models=["gpt-oss-120b"]) + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert "reasoning" not in entry + assert "thinkingLevelMap" not in entry + + def test_live_responses_context_overrides_static_fallback(self): + model = "system.ai.future-coder-1" + overlay, _ = _overlay( + model, + codex_models=[model], + codex_specs=[{"id": model, "context_window": 750_000}], + ) + + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert entry["contextWindow"] == 750_000 + assert entry["maxTokens"] == 16_384 + + def test_gpt_model_entries_use_model_specific_windows(self): + overlay, _ = _overlay( + "system.ai.gpt-5-2", + codex_models=[ + "system.ai.gpt-5-2", + "databricks-gpt-5-4-nano", + "databricks-gpt-5-6-sol", + ], + ) + windows = { + m["id"]: m["contextWindow"] for m in overlay["providers"]["databricks-openai"]["models"] + } + assert windows == { + "system.ai.gpt-5-2": 400_000, + "databricks-gpt-5-4-nano": 400_000, + "databricks-gpt-5-6-sol": 1_050_000, + } + + def test_claude_entries_pin_limits_and_capabilities(self): + overlay, _ = _overlay( + "databricks-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", + "opus": "databricks-claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-4-5", + "haiku": "databricks-claude-haiku-4-5", + "fable": "system.ai.claude-fable-5", }, ) entries = {m["id"]: m for m in overlay["providers"]["databricks-claude"]["models"]} - opus = entries["system.ai.claude-opus-4-8"] + opus = entries["databricks-claude-opus-4-8"] assert opus["contextWindow"] == 1_000_000 assert opus["maxTokens"] == 128_000 + assert opus["reasoning"] is True + assert opus["input"] == ["text", "image"] assert opus["compat"] == {"forceAdaptiveThinking": True} assert opus["thinkingLevelMap"] == {"max": "max", "xhigh": "xhigh"} + assert entries["system.ai.claude-sonnet-4-5"]["contextWindow"] == 1_000_000 + assert entries["system.ai.claude-sonnet-4-5"]["maxTokens"] == 64_000 + assert "thinkingLevelMap" not in entries["system.ai.claude-sonnet-4-5"] + assert entries["databricks-claude-haiku-4-5"]["contextWindow"] == 200_000 + fable = entries["system.ai.claude-fable-5"] + assert fable["contextWindow"] == 1_000_000 + assert fable["maxTokens"] == 128_000 + assert fable["compat"] == {"forceAdaptiveThinking": True} + assert fable["thinkingLevelMap"] == {"max": "max", "xhigh": "xhigh"} + + def test_claude_extended_levels_follow_model_capabilities(self): + overlay, _ = _overlay( + "claude-sonnet-5", + claude_models={ + "opus": "system.ai.claude-opus-4-6", + "sonnet": "system.ai.claude-sonnet-5", + }, + ) + entries = {m["id"]: m for m in overlay["providers"]["databricks-claude"]["models"]} + assert entries["system.ai.claude-opus-4-6"]["thinkingLevelMap"] == {"max": "max"} 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"]) @@ -132,17 +230,131 @@ def test_gemini_provider_uses_google_generative_ai(self): assert provider["api"] == "google-generative-ai" assert provider["baseUrl"] == f"{WS}/ai-gateway/gemini/v1beta" - def test_all_three_providers_when_all_present(self): + def test_mlflow_provider_uses_openai_completions(self): + overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"]) + provider = overlay["providers"]["databricks-mlflow"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{WS}/ai-gateway/mlflow/v1" + assert provider["compat"] == {"supportsStore": False, "supportsStrictMode": False} + + def test_no_mlflow_provider_when_no_oss_models(self): + overlay, _ = _overlay("gpt-5", codex_models=["gpt-5"]) + assert "databricks-mlflow" not in overlay.get("providers", {}) + + def test_all_four_providers_when_all_present(self): overlay, _ = _overlay( "claude-sonnet", claude_models={"sonnet": "claude-sonnet"}, codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.glm-5-2"], ) assert set(overlay["providers"].keys()) == { "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-mlflow", + } + + +class TestRenderOverlayOssEnrichment: + """OSS mlflow model entries carry reasoning + contextWindow + maxTokens + from the shared databricks.model_token_limits / model_is_reasoning tables.""" + + def test_reasoning_model_enriched(self): + overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"]) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry["id"] == "system.ai.glm-5-2" + assert entry["reasoning"] is True + assert entry["contextWindow"] == 1_000_000 + assert entry["maxTokens"] == 65_536 + + def test_unvalidated_model_has_no_inferred_metadata(self): + # Discovery does not offer this model; even if supplied directly, Pi + # must not infer capabilities for an unvalidated coding model. + overlay, _ = _overlay("system.ai.inkling", oss_models=["system.ai.inkling"]) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry == {"id": "system.ai.inkling"} + + def test_unknown_oss_model_bare(self): + # No limits/reasoning table entry -> only id, client keeps defaults. + overlay, _ = _overlay("system.ai.mystery-7b", oss_models=["system.ai.mystery-7b"]) + assert overlay["providers"]["databricks-mlflow"]["models"][0] == { + "id": "system.ai.mystery-7b" + } + + def test_dynamic_full_spec_overrides_static_metadata(self): + specs = [ + { + "id": "system.ai.glm-5-2", + "reasoning": False, + "context_window": 256_000, + "max_tokens": 12_345, + } + ] + overlay, _ = _overlay( + "system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"], oss_specs=specs + ) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry == { + "id": "system.ai.glm-5-2", + "contextWindow": 256_000, + "maxTokens": 12_345, + } + + def test_dynamic_reasoning_true_is_applied_with_safe_unknown_limits(self): + specs = [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": None, + "max_tokens": None, + } + ] + overlay, _ = _overlay( + "system.ai.inkling", oss_models=["system.ai.inkling"], oss_specs=specs + ) + assert overlay["providers"]["databricks-mlflow"]["models"][0] == { + "id": "system.ai.inkling", + "reasoning": True, + "contextWindow": 128_000, + "maxTokens": 8_192, + } + + def test_partial_dynamic_limits_are_completed_conservatively(self): + specs = [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": None, + "max_tokens": 65_536, + } + ] + overlay, _ = _overlay( + "system.ai.inkling", oss_models=["system.ai.inkling"], oss_specs=specs + ) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry["contextWindow"] == 128_000 + assert entry["maxTokens"] == 65_536 + + def test_malformed_spec_is_ignored_safely(self): + specs = [ + None, + {"id": 12, "reasoning": True}, + { + "id": "system.ai.mystery-7b", + "reasoning": "yes", + "context_window": -1, + "max_tokens": True, + }, + ] + overlay, _ = _overlay( + "system.ai.mystery-7b", + oss_models=["system.ai.mystery-7b"], + oss_specs=specs, + ) + assert overlay["providers"]["databricks-mlflow"]["models"][0] == { + "id": "system.ai.mystery-7b" } @@ -169,6 +381,7 @@ def test_claude_disables_eager_tool_input_streaming(self): overlay, _ = _overlay("claude-sonnet", claude_models={"sonnet": "claude-sonnet"}) compat = overlay["providers"]["databricks-claude"]["compat"] assert compat["supportsEagerToolInputStreaming"] is False + assert compat["sendSessionAffinityHeaders"] is True def test_openai_and_gemini_have_no_compat_flags(self): # Their gateway routes accept pi's request shape as-is. @@ -205,13 +418,18 @@ def test_claude_models_listed(self): assert ids == {"claude-opus", "claude-sonnet"} def test_pi_can_list_supplemental_claude_versions(self): - overlay, _ = _overlay( + # Shared discovery pins the opus family to 4.8 for smart routing. Pi's + # model picker still needs to expose newer versions such as Opus 5. + overlay, _ = pi.render_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", - ], + "tok", + _base_urls(), + {"opus": "system.ai.claude-opus-4-8"}, + [], + [], + [], + [], + ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"], ) provider = overlay["providers"]["databricks-claude"] assert {model["id"] for model in provider["models"]} == { @@ -225,25 +443,15 @@ def test_openai_models_listed(self): 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): + def test_grok_appears_in_openai_model_picker(self): grok = "system.ai.grok-4-6" overlay, _ = _overlay(grok, codex_models=[grok]) - entry = overlay["providers"]["databricks-openai"]["models"][0] + provider = overlay["providers"]["databricks-openai"] + assert provider["api"] == "openai-responses" + assert [model["id"] for model in provider["models"]] == [grok] + entry = provider["models"][0] assert entry["contextWindow"] == 500_000 - assert entry["maxTokens"] == 16_384 assert entry["reasoning"] is True assert entry["thinkingLevelMap"] == { "off": None, @@ -253,6 +461,12 @@ def test_grok_appears_with_supported_thinking_levels(self): } assert overlay["model"] == f"databricks-openai/{grok}" + def test_databricks_grok_id_gets_same_thinking_levels(self): + entry = pi._pi_gpt_model_entry("databricks-grok-4-6") + assert entry["reasoning"] is True + assert entry["thinkingLevelMap"]["xhigh"] == "xhigh" + assert entry["thinkingLevelMap"]["off"] is None + def test_grok_preview_does_not_inherit_unverified_thinking_levels(self): model = "system.ai.grok-4-6-preview" overlay, _ = _overlay(model, codex_models=[model]) @@ -296,6 +510,10 @@ def test_prefixes_gemini_model(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) assert overlay["model"] == "databricks-gemini/gemini-2" + def test_prefixes_oss_model(self): + overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"]) + assert overlay["model"] == "databricks-mlflow/system.ai.glm-5-2" + def test_preserves_already_prefixed_model(self): overlay, _ = _overlay( "databricks-claude/claude-sonnet", @@ -323,21 +541,22 @@ def test_falls_back_to_haiku(self): state = {"claude_models": {"haiku": "h4"}} assert pi.default_model(state) == "h4" - def test_falls_back_to_newest_gpt_model(self): + def test_falls_back_to_newest_codex_model(self): state = { "claude_models": {}, - "codex_models": ["gpt-5", "system.ai.gpt-5-6-sol", "gpt-5-5"], + "codex_models": ["databricks-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 + state = {"claude_models": {}, "codex_models": [grok]} + assert pi.default_model(state) == grok def test_does_not_route_gpt_oss_to_responses(self): state = { "claude_models": {}, - "codex_models": ["system.ai.gpt-oss-120b"], + "codex_models": ["gpt-oss-120b"], "gemini_models": ["gemini-2"], } assert pi.default_model(state) == "gemini-2" @@ -346,6 +565,15 @@ def test_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} assert pi.default_model(state) == "gemini-2" + def test_falls_back_to_oss_last(self): + state = { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + "oss_models": ["system.ai.glm-5-2"], + } + assert pi.default_model(state) == "system.ai.glm-5-2" + def test_returns_none_when_empty(self): assert pi.default_model({}) is None assert ( @@ -358,7 +586,7 @@ def test_sets_oauth_token(self): env = pi.build_runtime_env("tok") assert env["OAUTH_TOKEN"] == "tok" - def test_sets_standard_agent_dir_without_replacing_home(self, monkeypatch): + def test_sets_private_agent_dir_without_replacing_home(self, monkeypatch): monkeypatch.setenv("HOME", "/real-user-home") env = pi.build_runtime_env("tok") @@ -418,6 +646,7 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc "databricks-claude": {"old": True}, "databricks-openai": {"old": True}, "databricks-gemini": {"old": True}, + "databricks-mlflow": {"old": True}, "user-provider": {"keep": True}, } } @@ -433,6 +662,7 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc providers = written.get("providers", {}) assert providers.get("databricks-claude") != {"old": True} assert "old" not in providers.get("databricks-claude", {}) + assert "databricks-mlflow" not in providers assert providers.get("user-provider") == {"keep": True} def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): @@ -478,41 +708,129 @@ 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): + def test_cached_pi_inventory_always_keeps_shared_default(self): + state = { + "workspace": WS, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + "pi_claude_models": ["system.ai.claude-opus-5"], + } + + models = pi._discover_pi_claude_models(state, "tok", state["claude_models"]) + + assert models == ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"] + + def test_config_discovers_supplemental_claude_versions_for_pi(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"] + state = self._state( + claude_models={"opus": "system.ai.claude-opus-4-8"}, + ) with ( patch.object( - pi_mod, "discover_claude_models_unbucketed", return_value=(discovered, None) + pi_mod, + "discover_claude_models_unbucketed", + return_value=( + ["system.ai.claude-opus-4-8", "system.ai.claude-opus-5"], + 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) + assert state["pi_claude_models"] == [ + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5", + ] + ids = { + model["id"] + for model in json.loads(config_file.read_text())["providers"]["databricks-claude"][ + "models" + ] + } + assert ids == {"system.ai.claude-opus-4-8", "system.ai.claude-opus-5"} - def test_failed_supplemental_discovery_keeps_shared_family_pins(self, tmp_path, monkeypatch): + def test_config_discovers_supplemental_claude_versions_without_opus( + 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"}) + 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") - ), + pi_mod, + "discover_claude_models_unbucketed", + return_value=( + [ + "system.ai.claude-sonnet-4-5", + "system.ai.claude-sonnet-4-6", + ], + None, + ), + ) as discover, 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"] + discover.assert_called_once_with(WS, "tok") + assert { + model["id"] + for model in json.loads(config_file.read_text())["providers"]["databricks-claude"][ + "models" + ] + } == { + "system.ai.claude-sonnet-4-5", + "system.ai.claude-sonnet-4-6", + } + + def test_state_oss_specs_reach_written_model_entry(self, tmp_path, monkeypatch): + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) + state = self._state( + claude_models={}, + oss_models=["system.ai.inkling"], + oss_model_specs=[ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": 256_000, + "max_tokens": 65_536, + } + ], + ) + + with patch("ucode.agents.pi.save_state"): + pi_mod.write_tool_config(state, "system.ai.inkling", token="tok") + + entry = json.loads(config_file.read_text())["providers"]["databricks-mlflow"]["models"][0] + assert entry["reasoning"] is True + assert entry["contextWindow"] == 256_000 + assert entry["maxTokens"] == 65_536 + + def test_managed_oss_allowlist_excludes_unlisted_discovered_models(self, tmp_path, monkeypatch): + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) + state = self._state( + pi_models=["system.ai.deepseek-v4-pro"], + claude_models={"sonnet": "unlisted-claude"}, + oss_models=["system.ai.deepseek-v4-pro", "system.ai.glm-5-2"], + ) + + with patch("ucode.agents.pi.save_state"): + pi_mod.write_tool_config(state, "system.ai.deepseek-v4-pro", token="tok") + + providers = json.loads(config_file.read_text())["providers"] + assert set(providers) == {"databricks-mlflow"} + assert [model["id"] for model in providers["databricks-mlflow"]["models"]] == [ + "system.ai.deepseek-v4-pro" + ] 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_models=[ + "system.ai.claude-opus-4-8", + "system.ai.claude-opus-5", + ], pi_default_model="system.ai.claude-opus-5", ) @@ -545,6 +863,30 @@ def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): assert settings["defaultProvider"] == "databricks-claude" assert settings["defaultModel"] == "claude-sonnet" + def test_unservable_managed_allowlist_clears_stale_default(self, tmp_path, monkeypatch): + pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) + settings_file.write_text( + json.dumps( + { + "defaultProvider": "databricks-claude", + "defaultModel": "stale-unlisted-model", + "theme": "Default Dark", + } + ) + ) + + with patch("ucode.agents.pi.save_state"): + pi_mod.write_tool_config( + self._state(pi_models=["system.ai.unsupported-model"]), + "system.ai.unsupported-model", + token="tok", + ) + + settings = json.loads(settings_file.read_text()) + assert "defaultProvider" not in settings + assert "defaultModel" not in settings + assert settings["theme"] == "Default Dark" + def test_pre_existing_settings_are_backed_up_before_first_write(self, tmp_path, monkeypatch): pi_mod, _, settings_file, settings_backup_file = self._setup(tmp_path, monkeypatch) @@ -609,27 +951,414 @@ def test_managed_models_split_into_pis_per_provider_inputs(self): "system.ai.gpt-5", "system.ai.grok-4-6", "system.ai.gemini-3-flash", + "system.ai.deepseek-v4-pro", ] } assert pi._managed_model_families(state) == ( {"opus": "system.ai.claude-opus-4-8"}, ["system.ai.gpt-5", "system.ai.grok-4-6"], ["system.ai.gemini-3-flash"], + ["system.ai.deepseek-v4-pro"], ) def test_no_split_without_managed_models(self): assert pi._managed_model_families({"claude_models": {"opus": "x"}}) is None - def test_none_when_no_managed_model_is_servable(self): - # Pi has no OSS provider, so an oss-only list yields no families. Returning an all-empty - # tuple would be truthy and suppress the fallback, writing a config with zero providers. - assert pi._managed_model_families({"pi_models": ["system.ai.kimi-k2-7-code"]}) is None + def test_oss_only_allowlist_does_not_fall_back_to_discovery(self): + assert pi._managed_model_families({"pi_models": ["system.ai.kimi-k2-7-code"]}) == ( + {}, + [], + [], + ["system.ai.kimi-k2-7-code"], + ) + + def test_unsupported_nonempty_allowlist_stays_empty(self): + assert pi._managed_model_families({"pi_models": ["system.ai.unsupported-model"]}) == ( + {}, + [], + [], + [], + ) def test_partially_servable_list_still_splits(self): families = pi._managed_model_families( {"pi_models": ["system.ai.kimi-k2-7-code", "system.ai.claude-opus-4-8"]} ) - assert families == ({"opus": "system.ai.claude-opus-4-8"}, [], []) + assert families == ( + {"opus": "system.ai.claude-opus-4-8"}, + [], + [], + ["system.ai.kimi-k2-7-code"], + ) + + +class TestMlflowProxyLifecycle: + def test_not_started_without_oss_models(self): + state = {"workspace": WS, "oss_models": [], "base_urls": {"pi": _base_urls()}} + with patch.object(pi._mlflow_proxy, "start") as start: + assert pi._start_oss_proxy(state) is None + start.assert_not_called() + + def test_managed_oss_allowlist_starts_proxy_without_discovery_state(self): + server = MagicMock() + state = { + "workspace": WS, + "pi_models": ["system.ai.kimi-k2-7-code"], + "base_urls": {"pi": _base_urls()}, + } + with patch.object( + pi._mlflow_proxy, + "start", + return_value=(server, "http://127.0.0.1:60000"), + ) as start: + running = pi._start_oss_proxy(state) + assert running is not None + start.assert_called_once_with(WS) + assert state["base_urls"]["pi"]["oss"] == ("http://127.0.0.1:60000/ai-gateway/mlflow/v1") + + def test_stale_loopback_url_is_replaced_and_real_workspace_is_upstream(self): + server = MagicMock() + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with patch.object( + pi._mlflow_proxy, + "start", + return_value=(server, "http://127.0.0.1:60000"), + ) as start: + running = pi._start_oss_proxy(state) + assert running is not None + start.assert_called_once_with(WS) + assert state["base_urls"]["pi"]["oss"] == ("http://127.0.0.1:60000/ai-gateway/mlflow/v1") + + def test_loopback_workspace_is_not_recursively_proxied(self): + state = { + "workspace": "http://127.0.0.1:9999", + "oss_models": ["system.ai.inkling"], + } + with patch.object(pi._mlflow_proxy, "start") as start: + assert pi._start_oss_proxy(state) is None + start.assert_not_called() + + @staticmethod + def _proxy_pair(): + proxy_thread = MagicMock() + server = MagicMock() + return (proxy_thread, server), proxy_thread, server + + def test_restore_rewrites_persistent_config_to_direct_gateway(self): + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with patch.object(pi, "write_tool_config") as write: + pi._restore_direct_oss_config(state, "tok") + assert state["base_urls"]["pi"]["oss"] == f"{WS}/ai-gateway/mlflow/v1" + write.assert_called_once_with(state, "system.ai.inkling", token="tok") + + def test_restore_without_token_clears_state_and_existing_config_url( + self, tmp_path, monkeypatch + ): + config_path = tmp_path / "models.json" + config_path.write_text( + json.dumps( + { + "providers": { + "databricks-mlflow": { + "baseUrl": "http://127.0.0.1:54321/ai-gateway/mlflow/v1", + "apiKey": "existing-token", + } + } + } + ) + ) + monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with patch.object(pi, "write_tool_config") as write: + pi._restore_direct_oss_config(state, None) + assert state["base_urls"]["pi"]["oss"] == f"{WS}/ai-gateway/mlflow/v1" + write.assert_not_called() + restored = json.loads(config_path.read_text()) + assert restored["providers"]["databricks-mlflow"]["baseUrl"] == ( + f"{WS}/ai-gateway/mlflow/v1" + ) + assert restored["providers"]["databricks-mlflow"]["apiKey"] == "existing-token" + + def test_refresh_keeps_live_proxy_url(self, tmp_path, monkeypatch): + config_path = tmp_path / "models.json" + settings_path = tmp_path / "settings.json" + monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", settings_path) + monkeypatch.setattr(pi, "PI_BACKUP_PATH", tmp_path / "models.backup.json") + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "settings.backup.json") + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with ( + patch.object(pi, "get_databricks_token", return_value="refreshed-token"), + patch.object(pi, "save_state"), + ): + pi._refresh_token_once(state, force_refresh=True) + provider = json.loads(config_path.read_text())["providers"]["databricks-mlflow"] + assert provider["baseUrl"] == "http://127.0.0.1:54321/ai-gateway/mlflow/v1" + assert provider["apiKey"] == "refreshed-token" + + def test_restore_waits_for_inflight_refresh_and_writes_direct_last(self, monkeypatch): + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + refresh_started = threading.Event() + release_refresh = threading.Event() + restore_done = threading.Event() + write_order: list[str] = [] + + def fake_write(current, model, token=None, *, force_refresh=False): + write_order.append(str(token)) + if token == "refresh-token": + refresh_started.set() + assert release_refresh.wait(timeout=2) + return current, str(token) + + monkeypatch.setattr(pi, "_write_tool_config_unlocked", fake_write) + refresher = threading.Thread( + target=pi.write_tool_config, + args=(state, "system.ai.inkling", "refresh-token"), + ) + refresher.start() + assert refresh_started.wait(timeout=2) + restorer = threading.Thread( + target=lambda: ( + pi._restore_direct_oss_config(state, "restore-token"), + restore_done.set(), + ) + ) + restorer.start() + assert not restore_done.wait(timeout=0.05) + release_refresh.set() + refresher.join(timeout=2) + restorer.join(timeout=2) + assert not refresher.is_alive() + assert not restorer.is_alive() + assert write_order == ["refresh-token", "restore-token"] + assert state["base_urls"]["pi"]["oss"] == f"{WS}/ai-gateway/mlflow/v1" + + def test_proxy_precedes_first_config_write_and_is_cleaned_on_normal_exit(self): + proxy, proxy_thread, server = self._proxy_pair() + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + order = [] + + def start(current): + current.setdefault("base_urls", {}).setdefault("pi", {})["oss"] = "http://live" + order.append("proxy") + return proxy + + def refresh(current, *, force_refresh=False): + assert current["base_urls"]["pi"]["oss"] == "http://live" + order.append("config") + return "tok" + + proc = MagicMock() + proc.wait.return_value = 0 + with ( + patch.object(pi, "_start_oss_proxy", side_effect=start), + patch.object(pi, "_refresh_token_once", side_effect=refresh), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config") as restore, + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(SystemExit) as exit_info, + ): + pi.launch(state, []) + assert exit_info.value.code == 0 + assert order[:2] == ["proxy", "config"] + restore.assert_called_once_with(state, "tok") + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + proxy_thread.join.assert_called_once_with(timeout=1) + + def test_interrupt_forwards_sigint_and_cleans_proxy(self): + proxy, _, server = self._proxy_pair() + proc = MagicMock() + proc.wait.side_effect = [KeyboardInterrupt, 130] + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config") as restore, + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(SystemExit) as exit_info, + ): + pi.launch(state, []) + assert exit_info.value.code == 130 + proc.send_signal.assert_called_once_with(pi.signal.SIGINT) + restore.assert_called_once_with(state, "tok") + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + def test_missing_sighup_still_installs_and_restores_sigterm(self, monkeypatch): + proxy, _, server = self._proxy_pair() + proc = MagicMock() + proc.poll.return_value = None + installed = {} + previous_sigterm = object() + signal_calls = [] + + def install(signum, handler): + signal_calls.append((signum, handler)) + installed[signum] = handler + + def wait(): + handler = installed[pi.signal.SIGTERM] + handler(pi.signal.SIGTERM, None) + + proc.wait.side_effect = wait + monkeypatch.delattr(pi.signal, "SIGHUP") + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config"), + patch.object(pi.signal, "getsignal", return_value=previous_sigterm), + patch.object(pi.signal, "signal", side_effect=install), + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(SystemExit) as exit_info, + ): + pi.launch(state, []) + + assert exit_info.value.code == 128 + pi.signal.SIGTERM + proc.send_signal.assert_called_once_with(pi.signal.SIGTERM) + assert signal_calls[-1] == (pi.signal.SIGTERM, previous_sigterm) + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + @pytest.mark.parametrize( + "signum", + [getattr(pi.signal, name) for name in ("SIGTERM", "SIGHUP") if hasattr(pi.signal, name)], + ) + def test_termination_signal_forwards_and_cleans_proxy(self, signum): + proxy, _, server = self._proxy_pair() + proc = MagicMock() + proc.poll.return_value = None + installed: dict[int, object] = {} + previous = { + getattr(pi.signal, name): object() + for name in ("SIGTERM", "SIGHUP") + if hasattr(pi.signal, name) + } + signal_calls: list[tuple[int, object]] = [] + + def install(current_signum, handler): + signal_calls.append((current_signum, handler)) + installed[current_signum] = handler + + def wait(): + handler = installed[signum] + assert callable(handler) + handler(signum, None) + + proc.wait.side_effect = wait + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config") as restore, + patch.object(pi.signal, "getsignal", side_effect=previous.__getitem__), + patch.object(pi.signal, "signal", side_effect=install), + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(SystemExit) as exit_info, + ): + pi.launch(state, []) + assert exit_info.value.code == 128 + signum + proc.send_signal.assert_called_once_with(signum) + restore.assert_called_once_with(state, "tok") + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + assert signal_calls[-2:] == list(previous.items()) + + @pytest.mark.parametrize("failure_stage", ["config", "popen"]) + def test_setup_failure_still_cleans_proxy(self, failure_stage): + proxy, _, server = self._proxy_pair() + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + refresh = MagicMock(return_value="tok") + popen = MagicMock(return_value=MagicMock()) + if failure_stage == "config": + refresh.side_effect = RuntimeError("token failed") + else: + popen.side_effect = OSError("binary missing") + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", refresh), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config") as restore, + patch.object(pi.subprocess, "Popen", popen), + pytest.raises((RuntimeError, OSError)) as exc_info, + ): + pi.launch(state, []) + expected = "token failed" if failure_stage == "config" else "binary missing" + assert str(exc_info.value) == expected + expected_token = None if failure_stage == "config" else "tok" + restore.assert_called_once_with(state, expected_token) + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + def test_setup_failure_remains_primary_when_restore_also_fails(self): + proxy, _, server = self._proxy_pair() + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object( + pi, "_restore_direct_oss_config", side_effect=RuntimeError("restore failed") + ), + patch.object(pi, "print_warning") as warning, + patch.object(pi.subprocess, "Popen", side_effect=OSError("binary missing")), + pytest.raises(OSError, match="binary missing"), + ): + pi.launch(state, []) + warning.assert_called_once() + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + def test_restore_failure_does_not_skip_proxy_shutdown(self): + proxy, _, server = self._proxy_pair() + proc = MagicMock() + proc.wait.return_value = 0 + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config", side_effect=OSError("restore failed")), + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(OSError, match="restore failed"), + ): + pi.launch(state, []) + server.shutdown.assert_called_once() + server.server_close.assert_called_once() class TestManagedDefaultModel: diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 5977cc0c..274efe2b 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -215,6 +215,14 @@ def test_pi_available_with_codex(self): def test_pi_available_with_gemini(self): assert check_gateway_endpoint({"gemini_models": ["gemini-2"]}, "pi") is True + def test_pi_available_with_oss(self): + assert check_gateway_endpoint({"oss_models": ["system.ai.glm-5-2"]}, "pi") is True + + def test_pi_oss_discovery_reason_is_reported(self): + state = {"_discovery_reasons": {"oss": "no validated OSS models"}} + detail = agents_mod._availability_failure_detail("pi", state) + assert detail == " (oss discovery: no validated OSS models)" + def test_pi_unavailable_when_no_models(self): assert check_gateway_endpoint({}, "pi") is False @@ -269,6 +277,15 @@ def test_pi_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} assert default_model_for_tool("pi", state) == "gemini-2" + def test_pi_falls_back_to_oss(self): + state = { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + "oss_models": ["system.ai.glm-5-2"], + } + assert default_model_for_tool("pi", state) == "system.ai.glm-5-2" + def test_pi_returns_none_when_no_models(self): assert default_model_for_tool("pi", {}) is None diff --git a/tests/test_cli.py b/tests/test_cli.py index b207c2ce..4919747c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -30,6 +30,17 @@ def _strip_ansi(text: str) -> str: TOOLS = ["codex", "claude", "gemini", "opencode"] +def test_oss_discovery_diagnostic_names_all_consumers(monkeypatch): + import ucode.cli as cli_mod + + notes = [] + monkeypatch.setattr(cli_mod, "print_note", notes.append) + + cli_mod._print_discovery_diagnostics({"_discovery_reasons": {"oss": "not found"}}) + + assert notes[0] == "OSS models (needed for: codex, opencode, pi): not found" + + def _jwt(expires_at: float) -> str: payload = base64.urlsafe_b64encode(json.dumps({"exp": expires_at}).encode()).decode() return f"header.{payload.rstrip('=')}.signature" @@ -3167,6 +3178,23 @@ def test_warns_only_for_unclassifiable_models(self, monkeypatch): assert "will be ignored" in warnings[0] assert "claude-opus" not in warnings[0] + def test_unservable_message_only_claims_fallback_when_one_happens(self): + # Pi keeps its managed `pi_models` even when nothing in it is servable, so + # discovery does NOT stand in and the message must not claim it does. + # OpenCode withholds the override entirely, so for it the claim is true. + from ucode.managed_resolve import managed_state_overrides, managed_unservable_models + + def managed_for(tool): + return {"enabled_agents": {tool: {"model_config": {"models": ["system.ai.mystery"]}}}} + + pi_managed = managed_for("pi") + assert managed_unservable_models(pi_managed, "pi") == ["system.ai.mystery"] + assert "pi_models" in managed_state_overrides(pi_managed, "pi") + + oc_managed = managed_for("opencode") + assert managed_unservable_models(oc_managed, "opencode") == ["system.ai.mystery"] + assert "opencode_models" not in managed_state_overrides(oc_managed, "opencode") + 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 c6d81ccc..f075a5b7 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -289,8 +289,12 @@ def test_unvalidated_families_return_none(self): ): 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 + @pytest.mark.parametrize( + "model_id", + ["system.ai.glm-embedding-0-6b", "system.ai.qwen3-embedding-0-6b"], + ) + def test_embedding_model_returns_none_not_fallback(self, model_id): + assert db_mod.model_token_limits(model_id) is None class TestModelIsReasoning: @@ -298,6 +302,9 @@ 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_matching_is_case_insensitive(self): + assert db_mod.model_is_reasoning("SYSTEM.AI.GLM-5-2") 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 diff --git a/tests/test_e2e.py b/tests/test_e2e.py index d304adca..d87fa587 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -994,8 +994,23 @@ def _all_models(self, e2e_state: dict) -> list[tuple[str, str]]: out.append(("codex", model)) for model in e2e_state.get("gemini_models") or []: out.append(("gemini", model)) + for model in e2e_state.get("oss_models") or []: + out.append(("oss", model)) return out + def test_all_models_includes_oss_provider(self): + models = self._all_models( + { + "claude_models": {"sonnet": "claude-sonnet"}, + "codex_models": ["gpt-5"], + "gemini_models": ["gemini-3"], + "oss_models": ["system.ai.glm-5-2"], + } + ) + + assert ("oss", "system.ai.glm-5-2") in models + assert len(models) == 4 + def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token): import ucode.config_io as config_io_mod from ucode.agents import pi diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index 2d11832b..f6a30041 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -562,11 +562,15 @@ class TestManagedUnservableModels: def _managed(tool, models): return {"enabled_agents": {tool: {"model_config": {"models": models}}}} - def test_pi_oss_only_is_unservable(self): - # Pi has no OSS provider block. + def test_pi_oss_only_is_servable(self): + assert ( + managed_unservable_models(self._managed("pi", ["system.ai.kimi-k2-7-code"]), "pi") == [] + ) + + def test_pi_unknown_only_is_unservable(self): assert managed_unservable_models( - self._managed("pi", ["system.ai.kimi-k2-7-code"]), "pi" - ) == ["system.ai.kimi-k2-7-code"] + self._managed("pi", ["system.ai.unknown-model"]), "pi" + ) == ["system.ai.unknown-model"] def test_opencode_gpt_only_is_unservable(self): # OpenCode has no OpenAI provider block. diff --git a/tests/test_mlflow_proxy.py b/tests/test_mlflow_proxy.py new file mode 100644 index 00000000..6338606b --- /dev/null +++ b/tests/test_mlflow_proxy.py @@ -0,0 +1,336 @@ +"""Behavioral tests for Pi's MLflow SSE-repair proxy.""" + +from __future__ import annotations + +import json +import socket +import threading +import urllib.error +import urllib.request +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest + +from ucode.agents import _mlflow_proxy + +_STREAM_NO_FINISH = ( + b'data: {"id":"c1","choices":[{"delta":{"content":"ok"},"index":0}]}\n\ndata: [DONE]\n\n' +) +_STREAM_WITH_FINISH = ( + b'data: {"id":"c2","choices":[{"delta":{"content":"ok"},"index":0}]}\n\n' + b'data: {"id":"c2","choices":[{"delta":{},"finish_reason":"stop","index":0}]}\n\n' + b"data: [DONE]\n\n" +) + + +class _Gateway(HTTPServer): + response_status = 200 + response_type = "text/event-stream" + response_body = b"" + truncate = False + received_headers: dict[str, str] + + +def _gateway( + body: bytes, + *, + status: int = 200, + content_type: str = "text/event-stream", + truncate: bool = False, +) -> tuple[str, _Gateway, threading.Thread]: + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self): # noqa: N802 + self.server.received_headers = dict(self.headers.items()) # type: ignore[attr-defined] + length = int(self.headers.get("Content-Length", 0)) + self.rfile.read(length) + self.send_response(self.server.response_status) # type: ignore[attr-defined] + self.send_header("Content-Type", self.server.response_type) # type: ignore[attr-defined] + advertised = len(self.server.response_body) + (20 if self.server.truncate else 0) # type: ignore[attr-defined] + self.send_header("Content-Length", str(advertised)) + self.send_header("X-Upstream", "yes") + self.end_headers() + self.wfile.write(self.server.response_body) # type: ignore[attr-defined] + self.wfile.flush() + if self.server.truncate: # type: ignore[attr-defined] + self.close_connection = True + + def log_message(self, format: str, *args: object) -> None: + return + + server = _Gateway(("127.0.0.1", 0), Handler) + server.response_status = status + server.response_type = content_type + server.response_body = body + server.truncate = truncate + server.received_headers = {} + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return f"http://127.0.0.1:{server.server_address[1]}", server, thread + + +def _proxy(upstream: str): + started = _mlflow_proxy.start(upstream) + assert started is not None + server, base = started + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return base, server, thread + + +def _stop(server: HTTPServer, thread: threading.Thread) -> None: + server.shutdown() + server.server_close() + thread.join(timeout=2) + assert not thread.is_alive() + + +def _post(base: str, *, authorization: str | None = None) -> tuple[int, dict[str, str], bytes]: + headers = {"Content-Type": "application/json"} + if authorization: + headers["Authorization"] = authorization + request = urllib.request.Request( + f"{base}/ai-gateway/mlflow/v1/chat/completions", + data=b'{"stream":true}', + headers=headers, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=5) as response: + return response.status, dict(response.headers.items()), response.read() + except urllib.error.HTTPError as exc: + return exc.code, dict(exc.headers.items()), exc.read() + + +class TestSseRepair: + def test_healthy_stream_body_is_byte_identical(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_WITH_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + try: + status, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert status == 200 + assert body == _STREAM_WITH_FINISH + assert body.count(b"finish_reason") == 1 + + def test_missing_finish_is_injected_before_done(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_NO_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, headers, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body.count(b"finish_reason") == 1 + assert body.index(b"finish_reason") < body.index(b"[DONE]") + assert "Content-Length" not in headers + + def test_data_field_without_space_and_absent_id(self): + stream = b'data:{"choices":[{"delta":{"content":"ok"}}]}\n\ndata:[DONE]\n\n' + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert b"finish_reason" in body + assert b'"id":null' not in body + + def test_transport_truncated_stream_is_not_turned_into_success(self): + stream = b'data: {"id":"c3","choices":[{"delta":{"content":"partial"}}]}\n\n' + upstream, gateway, gateway_thread = _gateway(stream, truncate=True) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body == stream + assert b"finish_reason" not in body + assert b"[DONE]" not in body + + def test_multiline_finish_event_is_not_repaired_twice(self): + stream = ( + b'data: {"id":"c4","choices":[{"delta":{},\n' + b'data: "index":0,"finish_reason":"stop"}]}\n\n' + b"data: [DONE]\n\n" + ) + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body == stream + assert body.count(b"finish_reason") == 1 + + def test_any_choice_finish_reason_suppresses_injection(self): + stream = ( + b'data: {"choices":[{"delta":{}},{"delta":{},"finish_reason":"length"}]}\n\n' + b"data: [DONE]\n\n" + ) + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body == stream + assert body.count(b"finish_reason") == 1 + + @pytest.mark.parametrize( + "stream", + [ + b'data: {"error":{"message":"rate limited"}}\n\n', + b'event: error\ndata: {"message":"rate limited"}\n\n', + b'event:error\ndata: {"message":"rate limited"}\n\n', + ], + ) + def test_explicit_sse_error_is_not_turned_into_success(self, stream): + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body == stream + assert b"finish_reason" not in body + assert b"[DONE]" not in body + + +class TestPassthroughAndFailures: + def test_non_streaming_json_status_headers_and_body_preserved(self): + payload = b'{"choices":[{"message":{"content":"ok"}}]}' + upstream, gateway, gateway_thread = _gateway(payload, content_type="application/json") + base, proxy, proxy_thread = _proxy(upstream) + try: + status, headers, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert status == 200 + assert headers["Content-Type"] == "application/json" + assert headers["Content-Length"] == str(len(payload)) + assert headers["X-Upstream"] == "yes" + assert body == payload + + def test_http_error_is_relayed_without_repair(self): + payload = b'{"error":"rate limited"}' + upstream, gateway, gateway_thread = _gateway( + payload, status=429, content_type="application/json" + ) + base, proxy, proxy_thread = _proxy(upstream) + try: + status, headers, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert status == 429 + assert headers["Content-Type"] == "application/json" + assert body == payload + assert b"finish_reason" not in body + + def test_connection_refused_returns_controlled_502(self): + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + base, proxy, proxy_thread = _proxy(f"http://127.0.0.1:{port}") + try: + status, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + assert status == 502 + assert json.loads(body)["error"] + assert b"finish_reason" not in body + + def test_absolute_request_target_is_rejected_without_forwarding_auth(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_WITH_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + proxy_port = int(base.rsplit(":", 1)[1]) + attacker, attacker_server, attacker_thread = _gateway( + b"captured", content_type="text/plain" + ) + request = ( + f"POST {attacker}/capture HTTP/1.1\r\n" + "Host: ignored\r\n" + "Authorization: Bearer secret-value\r\n" + "Content-Length: 2\r\n" + "Connection: close\r\n\r\n{}" + ).encode() + sock = socket.create_connection(("127.0.0.1", proxy_port), timeout=5) + try: + sock.sendall(request) + response = b"" + while chunk := sock.recv(4096): + response += chunk + finally: + sock.close() + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + _stop(attacker_server, attacker_thread) + assert b" 400 " in response.split(b"\r\n", 1)[0] + assert gateway.received_headers == {} + assert attacker_server.received_headers == {} + + def test_authorization_forwarded_and_hop_by_hop_headers_removed(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_WITH_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + try: + _post(base, authorization="Bearer secret-value") + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + lowered = {key.lower(): value for key, value in gateway.received_headers.items()} + assert lowered["authorization"] == "Bearer secret-value" + # urllib regenerates identity after the client value is stripped, so + # the parseable SSE cannot arrive gzip-compressed. + assert lowered["accept-encoding"] == "identity" + assert lowered["host"].startswith("127.0.0.1:") # regenerated for upstream + + +class TestLifecycle: + def test_shutdown_and_server_close_release_port(self): + started = _mlflow_proxy.start("https://example.com") + assert started is not None + server, _ = started + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + _stop(server, thread) + rebound = socket.socket() + try: + rebound.bind(("127.0.0.1", port)) + finally: + rebound.close() + + def test_repeated_start_stop_uses_fresh_live_servers(self): + ports = [] + for _ in range(3): + started = _mlflow_proxy.start("https://example.com") + assert started is not None + server, _ = started + ports.append(server.server_address[1]) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + _stop(server, thread) + assert all(isinstance(port, int) and port > 0 for port in ports) + + def test_bind_failure_warns_and_degrades_to_direct_gateway(self, monkeypatch): + warnings = [] + monkeypatch.setattr( + _mlflow_proxy, + "_Server", + lambda *args, **kwargs: (_ for _ in ()).throw(OSError("no sockets")), + ) + monkeypatch.setattr(_mlflow_proxy, "print_warning", warnings.append) + + assert _mlflow_proxy.start("https://example.com") is None + assert warnings and "not started" in warnings[0]