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] 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