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..fb2564e5 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. @@ -1636,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``. @@ -1665,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 @@ -1758,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: @@ -1796,7 +2014,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. @@ -1822,8 +2040,16 @@ 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) - - codex_models = sorted([m for m in ids if "gpt-" in m], key=model_version_sort_key) + 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) oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)] @@ -2856,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, ) @@ -2898,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( @@ -2922,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 {}, ( @@ -3362,7 +3586,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..8ab459a8 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) @@ -224,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`.""" @@ -247,6 +266,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 +352,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 +375,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. @@ -338,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") @@ -2498,6 +2624,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"), @@ -2550,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() 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