diff --git a/README.md b/README.md index e450bdaa..a51bc212 100644 --- a/README.md +++ b/README.md @@ -356,6 +356,7 @@ control the installation. | File | Tool | |------|------| | `~/.codex/ucode.config.toml` (or legacy `~/.codex/config.toml`) | Codex | +| `~/.codex/ucode-models.json` | Codex static model catalog generated by ug | | `~/.claude/ucode-settings.json` | Claude Code settings generated by ug | | `/etc/claude-code/managed-settings.json` (Linux) or `/Library/Application Support/ClaudeCode/managed-settings.json` (macOS) | Claude Code OS-managed settings | | `/etc/codex/managed_config.toml` | Codex OS-managed settings | @@ -369,6 +370,22 @@ control the installation. Existing files are backed up before being overwritten. `ug revert` restores backups. +### What `ug configure` and `ug` write for Claude Code + +`ug configure` and every `ug` launch apply the workspace's managed config to these files: + +- `~/.claude/ucode-settings.json`: the settings ug generates for Claude Code. Its `env` block sets `ANTHROPIC_BASE_URL` (the gateway), the per family default models `ANTHROPIC_DEFAULT_OPUS_MODEL`, `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_FABLE_MODEL`, `ANTHROPIC_CUSTOM_HEADERS`, any tracing variables, and, when the config uses model discovery, the gateway model discovery flag. When the config pins a static model list, ug also writes the model picker keys `availableModels`, `enforceAvailableModels`, and a `modelPicker` here, so `/model` shows exactly those models. Any ug managed `permissions` and hooks live here too. +- `~/.claude/settings.json`: Claude Code's own user settings. ug records the selected default model here. +- `~/.claude.json`: ug registers the Databricks `web_search` MCP server here when a suitable endpoint is available. +- `/etc/claude-code/managed-settings.json` (Linux) or `/Library/Application Support/ClaudeCode/managed-settings.json` (macOS): the OS enterprise managed settings. ug mirrors the same settings here so a bare `claude`, launched outside ug, still routes through the gateway. Writing this needs local admin rights; without them ug applies only the user level file above. + +### What `ug configure` and `ug` write for Codex + +- `~/.codex/ucode.config.toml`: the `ucode` profile that points Codex at the gateway. It sets `model_provider = "ucode-databricks"`, the `[model_providers.ucode-databricks]` block (gateway `base_url`, `wire_api`, the `auth` command, and headers), the default `model`, and, when the config pins a static model list, `model_catalog_json` pointing at the catalog file below. +- `~/.codex/ucode-models.json`: the model catalog ug generates from a static allow-list. Codex reads it through `model_catalog_json` so `/model` shows exactly those models. ug writes it only for a static list and removes it when the config uses discovery instead. +- `~/.codex/config.toml`: the legacy single file layout for Codex older than 0.134.0. ug writes the same settings under `[profiles.ucode]` here instead. +- `/etc/codex/managed_config.toml`: the OS enterprise managed Codex config. ug mirrors its Codex settings here so a bare `codex` uses the gateway. Writing this needs local admin rights. + ## Documentation diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 5b8e1666..24eb08e3 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -80,6 +80,10 @@ # Retained only to identify and remove state written by the legacy persisted opt-in. SMART_ROUTING_STATE_KEY = smart_routing_v2.LEGACY_STATE_KEY +# Tracks the modelPicker that ucode wrote (for ownership detection on later pruning). +CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY = "claude_managed_model_picker" +# Tracks the availableModels/enforceAvailableModels that ucode wrote (for ownership detection). +CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY = "claude_managed_pruned_picker" def _parse_version(value: str) -> tuple[int, int, int] | None: @@ -161,6 +165,8 @@ def _resolve_web_search_model(state: dict) -> str | None: # Launch-scoped feature flags that ucode may write into Claude settings. These # must be removed again when the corresponding launch flag is absent. CLAUDE_CONDITIONAL_ENV_KEYS = ("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",) +CLAUDE_MANAGED_PICKER_KEYS = ("availableModels", "enforceAvailableModels", "modelPicker") +CLAUDE_PRUNED_PICKER_KEYS = ("availableModels", "enforceAvailableModels") # Env keys ucode used to write but no longer does; stripped from the managed # settings file on every launch so stale values never linger. CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",) @@ -324,6 +330,8 @@ def render_overlay( relayed_base_url: str | None = None, route_root_model: str | None = None, custom_model: str | None = None, + static_models: list[str] | None = None, + model_service_location: str | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -449,9 +457,26 @@ def render_overlay( overlay["permissions"] = {"deny": ["WebSearch"]} keys.append(["permissions", "deny"]) + if static_models and not provider and not relayed: + overlay["availableModels"] = list(static_models) + overlay["enforceAvailableModels"] = True + overlay["modelPicker"] = { + "replaceBuiltInOptions": True, + "options": [{"model": m, "label": _picker_label(m)} for m in static_models], + } + keys += [[key] for key in CLAUDE_MANAGED_PICKER_KEYS] + elif model_service_location and not provider and not relayed: + env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" + keys.append(["env", "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"]) + return overlay, keys +def _picker_label(model: str) -> str: + """A short picker label for a model id — the raw id minus the ``system.ai.`` prefix.""" + return model.removeprefix("system.ai.") + + def _maybe_add_1m_suffix(model: str) -> str: if model.endswith("[1m]"): return model @@ -596,6 +621,8 @@ def write_tool_config( relayed_base_url=relayed_base_url, route_root_model=route_root_model, custom_model=custom_model, + static_models=state.get("claude_static_models"), + model_service_location=state.get("claude_model_service_location"), ) tracing_env_vars = tracing_env(state, "claude") stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None @@ -619,13 +646,16 @@ def write_tool_config( + [["env", key] for key in CLAUDE_TRACING_ENV_KEYS] + [["hooks", "Stop"]] + [["hooks", event] for event in ("PreToolUse", "SessionStart", "SubagentStart")] + + [[key] for key in CLAUDE_MANAGED_PICKER_KEYS] ): if path not in managed_file_keys: managed_file_keys.append(path) # V2 installs routing hooks in a transient per-launch settings file. Persistent settings must # contain no ucode routing hooks; surgically strip legacy ones while preserving user hooks. - def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: + def _compose( + base: dict, *, enforce_model_default_hierarchy: bool, is_managed: bool = False + ) -> dict: base_env = base.get("env") existing_custom_headers = ( base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None @@ -694,21 +724,93 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: # longer writes. for key in CLAUDE_REMOVED_ENV_KEYS: merged_env.pop(key, None) + # Prune availableModels/enforceAvailableModels when ucode stops writing them: the private + # file unconditionally (ucode owns it), the managed file only when they match the ownership + # marker (ucode wrote them) so admin-set values survive. + if not is_managed: + for picker_key in CLAUDE_PRUNED_PICKER_KEYS: + if picker_key not in overlay_for_merge: + merged.pop(picker_key, None) + else: + saved_pruned = state.get(CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY) + if saved_pruned is not None: + current_pruned = {k: merged.get(k) for k in CLAUDE_PRUNED_PICKER_KEYS} + if current_pruned == saved_pruned: + for picker_key in CLAUDE_PRUNED_PICKER_KEYS: + if picker_key not in overlay_for_merge: + merged.pop(picker_key, None) + # Prune ucode-owned modelPicker when transitioning to dynamic config. Only prune if + # it matches the ownership marker (ucode wrote it); preserve admin-authored pickers. + if "modelPicker" not in overlay_for_merge: + saved_picker = state.get(CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY) + current_picker = merged.get("modelPicker") + if saved_picker is not None and current_picker == saved_picker: + merged.pop("modelPicker", None) sync_smart_routing_hooks(merged, state, enabled=False) return merged write_json_file( CLAUDE_SETTINGS_PATH, - _compose(read_json_safe(CLAUDE_SETTINGS_PATH), enforce_model_default_hierarchy=False), + _compose( + read_json_safe(CLAUDE_SETTINGS_PATH), + enforce_model_default_hierarchy=False, + is_managed=False, + ), + ) + + # Prior markers let compose_managed prune keys ucode wrote on a previous run (e.g. a + # static->discovery transition). + prior_picker_marker = state.get(CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY) + prior_pruned_marker = state.get(CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY) + + is_writing_static_picker = ( + "modelPicker" in overlay + or "availableModels" in overlay + or "enforceAvailableModels" in overlay ) - _reconcile_managed_settings( + # Persist this run's markers before the fallible reconcile so a write survives a failed retry; + # discovery runs instead clear them only after reconcile succeeds (below). + if is_writing_static_picker: + if "modelPicker" in overlay: + state[CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY] = overlay["modelPicker"] + else: + state.pop(CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY, None) + if "availableModels" in overlay or "enforceAvailableModels" in overlay: + state[CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY] = { + k: overlay.get(k) for k in CLAUDE_PRUNED_PICKER_KEYS + } + else: + state.pop(CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY, None) + save_state(state) + + def compose_managed(base: dict) -> dict: + # Compose against the prior markers so stale keys prune, then restore the current ones. + saved_picker = state.get(CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY) + saved_pruned = state.get(CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY) + try: + state[CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY] = prior_picker_marker + state[CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY] = prior_pruned_marker + return _compose(base, enforce_model_default_hierarchy=True, is_managed=True) + finally: + state[CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY] = saved_picker + state[CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY] = saved_pruned + + managed_file_was_written = _reconcile_managed_settings( state, - lambda base: _compose(base, enforce_model_default_hierarchy=True), + compose_managed, managed_file_keys, relayed, ) + # Discovery runs clear the markers only after reconcile succeeds, so a failed reconcile leaves + # the prior markers in place for a retry to prune. Only clear if the managed file was actually + # written; if it was skipped/suppressed, keep the markers so a later retry can prune. + if not is_writing_static_picker and managed_file_was_written: + state.pop(CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY, None) + state.pop(CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY, None) + save_state(state) + if web_search_model: web_search_entry = _web_search_mcp_entry( state["workspace"], web_search_model, state.get("profile") @@ -788,18 +890,21 @@ def _reconcile_managed_settings( compose: Callable[[dict], dict], owned_paths: list[list[str]], relayed: bool, -) -> None: +) -> bool: """Reconcile Claude Code's OS-managed settings so a bare ``claude`` uses the gateway. The managed file is root-owned and the highest-precedence scope, so every normal Claude configuration mirrors ucode's settings there. The same compose operation that produced the private file is applied to the existing managed file, preserving unrelated IT-authored keys. - `ug configure` updates gateway-owned fields in this file, but does not generate or modify - the `modelPicker` object; an existing picker is retained by the merge. + `ug configure` generates the `modelPicker` object only when the managed config supplies a + static model list; without one the `/model` surface comes from gateway model discovery instead, + so ucode neither writes nor prunes `modelPicker` and an existing picker is retained by the merge. Relayed launches are skipped: they depend on a per-session loopback refresh proxy that only runs during `ucode claude`, so a bare `claude` could not reach the gateway anyway. + + Returns True if the managed file was written, False otherwise. """ path = _managed_settings_path() if path is None: @@ -807,7 +912,7 @@ def _reconcile_managed_settings( "Machine-wide Claude settings aren't supported on this platform; skipped the managed " "settings." ) - return + return False if path.is_symlink(): raise RuntimeError( f"Refusing to use Claude Code managed settings through symlink {path}. Replace it " @@ -823,7 +928,7 @@ def _reconcile_managed_settings( "created them, run `ucode revert` from an interactive terminal first." ) mark_managed_file_verified(state, "claude", path, scope="relay-compatible") - return + return False current_text = read_managed_file(path) try: @@ -846,9 +951,9 @@ def _reconcile_managed_settings( "your administrator." ) mark_managed_file_verified(state, "claude", path, scope="local-compatible") - return + return False try: - reconcile_managed_file( + result = reconcile_managed_file( path, _dump_managed_settings(desired_settings), tool="claude", @@ -864,8 +969,10 @@ def _reconcile_managed_settings( f"local settings at {CLAUDE_SETTINGS_PATH}." ) mark_managed_file_verified(state, "claude", path, scope="local-compatible") - return + return False mark_managed_file_verified(state, "claude", path) + # Only signal successful write if the file was actually written, not for "unchanged" or "unsupported" + return result in ("created", "written") def _preserve_permission_denies(existing: dict, desired: dict) -> None: @@ -1214,12 +1321,20 @@ def _build_claude_argv( caller_settings = _merge_claude_settings(caller_settings, _load_caller_settings(value)) # ucode wins over the caller for conflicting keys (protects gateway auth); # hooks from both sides survive. - merged = _merge_claude_settings(caller_settings, read_json_safe(CLAUDE_SETTINGS_PATH)) + ucode_settings = read_json_safe(CLAUDE_SETTINGS_PATH) + merged = _merge_claude_settings(caller_settings, ucode_settings) if settings_override is not None: merged = _merge_claude_settings(merged, settings_override) merged_env = merged.get("env") if isinstance(merged_env, dict): - merged_env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None) + # Keep the gateway-discovery flag only if ucode's persistent settings set it (managed + # config); otherwise strip a stale one from the --settings merge. + ucode_env = ucode_settings.get("env") + if not ( + isinstance(ucode_env, dict) + and ucode_env.get("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY") == "1" + ): + merged_env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None) return [ binary, *source_args, diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 0643fdcd..f483496c 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -17,7 +17,9 @@ ToolSpec, backup_existing_file, deep_merge_dict, + is_dry_run, read_toml_safe, + write_json_file, write_toml_file, ) from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_token_argv @@ -57,6 +59,7 @@ CODEX_PROFILE_NAME = "ucode" CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / f"{CODEX_PROFILE_NAME}.config.toml" CODEX_BACKUP_PATH = APP_DIR / "codex-ucode-config.backup.toml" +CODEX_CATALOG_PATH = CODEX_CONFIG_DIR / "ucode-models.json" LEGACY_CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / "config.toml" LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml" CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" @@ -79,10 +82,81 @@ MANAGED_KEYS: list[list[str]] = [ ["model_provider"], ["model"], + ["model_catalog_json"], ["model_providers", CODEX_MODEL_PROVIDER_NAME], ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers"], ] + +def _strip_provider_header(config: dict) -> None: + """Remove the stale Databricks-Model-Provider-Service header when moving off a provider.""" + providers = config.get("model_providers") + if isinstance(providers, dict): + provider_block = providers.get(CODEX_MODEL_PROVIDER_NAME) + if isinstance(provider_block, dict): + headers = provider_block.get("http_headers") + if isinstance(headers, dict): + headers.pop("Databricks-Model-Provider-Service", None) + + +def _fetch_gateway_codex_models(workspace: str, token: str) -> dict | None: + """Fetch the real Codex model catalog from the AI Gateway. + + Returns the response payload `{"models": [, ...]}` on success, or None + if the fetch fails. This is the same shape that Codex's own remote discovery uses. + """ + from ucode.databricks import _http_get_json, build_tool_base_url + + base_url = build_tool_base_url("codex", workspace) + url = f"{base_url}/models" + payload, reason = _http_get_json(url, token) + if payload is None: + print_warning_err(f"Failed to fetch Codex model catalog from {url}: {reason}") + return None + if not isinstance(payload, dict) or "models" not in payload: + print_warning_err( + f"Codex model catalog from {url} has unexpected shape; expected {{'models': [...]}}" + ) + return None + return payload + + +def build_codex_catalog(workspace: str, token: str, allow_list: list[str]) -> dict | None: + """Build the ``model_catalog_json`` payload by fetching and filtering the real gateway catalog. + + Setting ``model_catalog_json`` makes Codex's StaticModelsManager the model source in place of + remote discovery, so ``/model`` lists exactly these entries. Each is a full ModelPreset from + the gateway, preserving the real `base_instructions`, `priority`, and all other fields verbatim. + The gateway's ordering is preserved for the filtered entries. + + Returns the catalog payload on success, an empty dict on success-but-no-matches (fall back to + discovery), or None on fetch failure (preserve existing catalog). + """ + gateway_catalog = _fetch_gateway_codex_models(workspace, token) + if gateway_catalog is None: + # Fetch failed; preserve existing catalog to survive transient network blips. + return None + + models = gateway_catalog.get("models") + if not isinstance(models, list) or not models: + print_warning_err("Codex model catalog is empty; falling back to discovery") + # Fetch succeeded but is empty; fall back to discovery and remove static catalog. + return {} + + allow_set = set(allow_list) + filtered = [m for m in models if isinstance(m, dict) and m.get("slug") in allow_set] + + if not filtered: + print_warning_err( + f"No allow-listed models found in gateway catalog; falling back to discovery. " + f"Allow-list: {allow_list}, gateway models: {[m.get('slug') for m in models if isinstance(m, dict)]}" + ) + # Fetch succeeded but no allow-listed models match; fall back to discovery. + return {} + + return {"models": filtered} + + LEGACY_MANAGED_KEYS: list[list[str]] = [ ["profile"], ["profiles", CODEX_PROFILE_NAME], @@ -185,10 +259,13 @@ def render_overlay( use_pat: bool = False, provider: str | None = None, custom_oauth: CustomOAuthConfig | None = None, + catalog_path: str | None = None, ) -> dict: overlay: dict = {"model_provider": CODEX_MODEL_PROVIDER_NAME} if model: overlay["model"] = model + if catalog_path: + overlay["model_catalog_json"] = catalog_path overlay["model_providers"] = { CODEX_MODEL_PROVIDER_NAME: _provider_block( workspace, databricks_profile, use_pat, provider, custom_oauth @@ -346,11 +423,18 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non ): for key in ("model", "model_reasoning_effort"): profiles[CODEX_PROFILE_NAME].pop(key, None) + # Drop a stale Databricks-Model-Provider-Service header when moving off a provider. + if not provider: + _strip_provider_header(doc) write_toml_file(LEGACY_CODEX_CONFIG_PATH, doc) state = mark_tool_managed(state, "codex", LEGACY_MANAGED_KEYS) save_state(state) return state + static_models = state.get("codex_static_models") + static_models = static_models if isinstance(static_models, list) and static_models else None + catalog_path = str(CODEX_CATALOG_PATH) if static_models and not provider else None + _remove_legacy_ucode_profile() backup_existing_file(CODEX_CONFIG_PATH, CODEX_BACKUP_PATH) overlay = render_overlay( @@ -360,6 +444,7 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non use_pat=bool(state.get("use_pat")), provider=provider, custom_oauth=state.get("custom_oauth"), + catalog_path=catalog_path, ) def compose(base: dict) -> dict: @@ -368,8 +453,30 @@ def compose(base: dict) -> dict: if chosen_model is None and not smart_routing_v2.enabled(): for key in ("model", "model_reasoning_effort"): base.pop(key, None) + if not catalog_path: + base.pop("model_catalog_json", None) + # Drop a stale Databricks-Model-Provider-Service header when moving off a provider. + if not provider: + _strip_provider_header(base) return base + if static_models and not provider: + token = get_databricks_token(workspace, databricks_profile) + catalog = build_codex_catalog(workspace, token, static_models) + if catalog is None: + # Fetch failed; preserve existing catalog to survive transient network blips. + pass + elif catalog: + # Fetch succeeded and has matching models; write the catalog. + write_json_file(CODEX_CATALOG_PATH, catalog) + else: + # Fetch succeeded but nothing matched allow-list; fall back to discovery. + catalog_path = None + if CODEX_CATALOG_PATH.exists() and not is_dry_run(): + CODEX_CATALOG_PATH.unlink() + elif CODEX_CATALOG_PATH.exists() and not is_dry_run(): + CODEX_CATALOG_PATH.unlink() + doc = read_toml_safe(CODEX_CONFIG_PATH) compose(doc) sync_smart_routing_hooks( diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 64cbb3d1..7e861dee 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -84,6 +84,13 @@ def managed_state_overrides(managed: dict, tool: str) -> dict[str, object]: default_model = _str(_agent_model_config(managed, tool).get("default_model")) if default_model: overrides[f"{tool}_default_model"] = default_model + if tool in ("claude", "codex"): + static_models = managed_static_models(managed, tool) + if static_models: + overrides[f"{tool}_static_models"] = static_models + location = managed_model_service_location(managed, tool) + if location: + overrides[f"{tool}_model_service_location"] = location return overrides @@ -114,7 +121,8 @@ def managed_unservable_models(managed: dict, tool: str) -> list[str]: def _manifest_models(managed: dict, tool: str) -> dict | list | None: """The manifest's models for ``tool`` in its own vocabulary, or None when it names none.""" - manifest_models = _agent_model_config(managed, tool).get("models") + model_config = _agent_model_config(managed, tool) + manifest_models = model_config.get("models") if tool == "claude": slots: dict[str, str] = {} for slot, family in _CLAUDE_FAMILY_SLOTS.items(): @@ -122,6 +130,14 @@ def _manifest_models(managed: dict, tool: str) -> dict | list | None: if model: slots[family] = model return slots or None + # For flat-list agents (gemini, opencode, pi, copilot), check the `names` key first + # (from managed static model lists), then fall back to legacy `models` list. + if tool not in ("claude", "codex"): + names = model_config.get("names") + if isinstance(names, list): + listed = [model for model in (_str(item) for item in names) if model] + if listed: + return listed if isinstance(manifest_models, list): listed = [model for model in (_str(item) for item in manifest_models) if model] return listed or None @@ -158,17 +174,23 @@ def managed_supplies_models(managed: dict | None, tool: str) -> bool: Lets the launch path skip Databricks model discovery, whose whole purpose is to find the models the config has now specified. Any of the three counts: a provider (the agent routes by header and - pins no Databricks model), a ``default_model``, or at least one entry in ``models``. + pins no Databricks model), a ``default_model``, or at least one entry in ``models`` (or ``names`` + for flat-list agents). """ model_config = _agent_model_config(managed or {}, tool) if _str(model_config.get("model_provider_service")) or _str(model_config.get("default_model")): return True - models = model_config.get("models") - if isinstance(models, dict): - return any(_str(value) for value in models.values()) - if isinstance(models, list): - return any(_str(item) for item in models) - return False + if tool == "claude" and ( + managed_static_models(managed or {}, tool) + or _str(model_config.get("model_service_location")) + ): + return True + if tool == "codex" and managed_static_models(managed or {}, tool): + return True + # For flat-list agents (gemini, opencode, pi, copilot), check both names (new static lists) + # and models (legacy lists), via _manifest_models which already handles both. + manifest_models = _manifest_models(managed or {}, tool) + return manifest_models is not None def managed_provider_service(managed: dict, tool: str) -> str | None: @@ -176,6 +198,26 @@ def managed_provider_service(managed: dict, tool: str) -> str | None: return _str(_agent_model_config(managed, tool).get("model_provider_service")) +def managed_static_models(managed: dict, tool: str) -> list[str] | None: + """The explicit model allow-list (``models.names``) the config sets for ``tool``, or None. + + Static curation: the launch path writes exactly these into the agent's own picker allow-list + (Claude ``availableModels``/``modelPicker``, Codex ``model_catalog_json``) instead of discovering + the workspace's models. The order is the admin's; empty and non-string entries are dropped.""" + names = _agent_model_config(managed, tool).get("names") + if isinstance(names, list): + listed = [model for model in (_str(item) for item in names) if model] + return listed or None + return None + + +def managed_model_service_location(managed: dict, tool: str) -> str | None: + """The UC catalog/schema (``models.model_service_location``) the config points ``tool`` at for + auto model discovery, or None. The agent discovers from the gateway rather than ucode pinning a + list, so the launch path only turns discovery on for this source.""" + return _str(_agent_model_config(managed, tool).get("model_service_location")) + + def managed_default_model(managed: dict, tool: str) -> str | None: """Return the model the managed config wants ``tool`` to launch on, if it names one. diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 1ce87a8b..6011621a 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy import json import os import shlex @@ -936,6 +937,59 @@ def test_relayed_rejects_invalid_managed_json(self, monkeypatch): assert managed_writes == [] + def test_admin_available_models_preserved_in_managed_file(self, monkeypatch): + # FIX P1-7: Enterprise admin's availableModels in OS-managed file must survive + # when ucode writes no static list. ucode's own keys should still be pruned. + private_writes: list = [] + managed_writes: list = [] + admin_available = ["system.ai.custom-model"] + existing_managed = { + str(FAKE_MANAGED_PATH): { + "availableModels": admin_available, + "enforceAvailableModels": True, + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing_managed) + state = {"workspace": WS, "codex_models": []} + claude.write_tool_config(state, "databricks-claude-sonnet-4") + _, text = managed_writes[0] + written = json.loads(text) + # Admin's availableModels should survive since no ownership marker was set. + assert written["availableModels"] == admin_available + assert written["enforceAvailableModels"] is True + + def test_ucode_available_models_pruned_when_transitioning_to_discovery(self, monkeypatch): + # FIX P1-7: When ucode transitions from static to discovery config, only prune + # the availableModels/enforceAvailableModels that ucode itself wrote. Admin keys stay. + private_writes: list = [] + managed_writes: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + # Step 1: Write with static models to the managed file. + state1 = {"workspace": WS, "claude_static_models": ["system.ai.claude-opus-4-8"]} + claude.write_tool_config(state1, None) + # Now managed file has ucode-written availableModels/enforceAvailableModels. + # Verify the marker was saved. + assert claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY in state1 + # Step 2: Rewrite to managed file in discovery mode (no static list). + # Use the saved state with the ownership marker. + private_writes.clear() + managed_writes.clear() + state2 = { + "workspace": WS, + claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY: state1[ + claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY + ], + "claude_model_service_location": "system.ai", + } + claude.write_tool_config(state2, None) + # Now the ucode-written keys should be pruned from managed file. + _, text = managed_writes[0] + written = json.loads(text) + assert "availableModels" not in written + assert "enforceAvailableModels" not in written + assert written.get("env", {}).get("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY") == "1" + def test_noninteractive_uses_local_settings_when_managed_file_is_compatible(self, monkeypatch): private_writes: list = [] managed_writes: list = [] @@ -1015,6 +1069,44 @@ def deny_managed_write(*args, **kwargs): {"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4" ) + def test_ownership_markers_persist_even_if_managed_reconcile_fails(self, monkeypatch): + # FIX #6: Ownership markers must be saved before the reconcile so they survive + # if reconcile raises, allowing later transitions to recognize stale keys. + private_writes: list = [] + managed_writes: list = [] + state_saves: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + + # Mock save_state to capture when markers are saved, and make reconcile fail. + original_save = claude.save_state + + def capture_save(s): + state_saves.append(copy.deepcopy(s)) + return original_save(s) + + monkeypatch.setattr(claude, "save_state", capture_save) + + def failing_reconcile(*args, **kwargs): + raise RuntimeError("Simulated reconcile failure") + + monkeypatch.setattr(claude, "_reconcile_managed_settings", failing_reconcile) + + state = {"workspace": WS, "claude_static_models": ["system.ai.claude-opus-4-8"]} + + # Reconcile will fail, but markers should still be saved. + with pytest.raises(RuntimeError, match="reconcile failure"): + claude.write_tool_config(state, None) + + # Verify markers were saved before the failure. + assert len(state_saves) > 0 + saved_state = state_saves[0] + assert claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY in saved_state + # The saved marker should have availableModels/enforceAvailableModels. + saved_marker = saved_state[claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY] + assert "availableModels" in saved_marker + assert "enforceAvailableModels" in saved_marker + class TestRegisterWebSearchMcp: def test_skips_registration_when_entry_is_current(self, monkeypatch): @@ -1366,6 +1458,236 @@ def test_prunes_stale_name_companion_keys_from_older_ucode(self, monkeypatch): assert "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME" not in env +class TestMarkerRecordingBeforeManagedReconcile: + """FIX P2-3: Ownership markers must be recorded before _reconcile_managed_settings, + so they survive if that step fails.""" + + def test_ownership_marker_set_before_managed_reconcile_fails(self, monkeypatch): + # FIX P2-3: Markers are recorded in state BEFORE the fallible OS reconcile, + # so they're set even if reconcile fails. On next run, markers persist to help + # recognize stale values. This test verifies markers are placed before reconcile. + monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) + monkeypatch.setattr(claude, "read_json_safe", lambda path: {}) + monkeypatch.setattr(claude, "write_json_file", lambda path, payload: None) + monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) + marker_in_state_at_failure: dict = {} + + def capture_state_on_failure(*args, **kwargs): + # Capture what's in state when reconcile is called (markers should be there). + # We use args[0] since compose is passed as a lambda, but state is in scope. + raise RuntimeError("Simulated OS reconcile failure") + + def failing_reconcile(state, *args, **kwargs): + # Capture state before raising. + marker_in_state_at_failure.update( + { + "marker": state.get(claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY), + } + ) + raise RuntimeError("Simulated OS reconcile failure") + + monkeypatch.setattr(claude, "_reconcile_managed_settings", failing_reconcile) + state = { + "workspace": WS, + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + + with pytest.raises(RuntimeError, match="Simulated OS reconcile failure"): + claude.write_tool_config(state, None) + + # The markers should be in state when reconcile is called (before it fails). + assert marker_in_state_at_failure["marker"] is not None + + +class TestDiscoveryMarkerTiming: + """FIX 2: When discovery run (no static picker), markers must be cleared and persisted AFTER + reconcile succeeds. If reconcile fails, prior markers stay persisted for retry pruning.""" + + def _patch(self, monkeypatch, private_writes, managed_writes, existing=None): + monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) + monkeypatch.setattr( + claude, "read_json_safe", lambda path: existing.get(str(path), {}) if existing else {} + ) + monkeypatch.setattr( + claude, "write_json_file", lambda path, payload: private_writes.append(payload) + ) + monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) + + def reconcile_managed_file(path, desired_text, **kwargs): + managed_writes.append({"path": str(path), "payload": desired_text}) + return "written" + + monkeypatch.setattr(claude, "reconcile_managed_file", reconcile_managed_file) + + def test_discovery_run_with_reconcile_failure_retains_prior_markers(self, monkeypatch): + # FIX 2b: When a DISCOVERY run (no static picker, markers cleared) has a reconcile failure, + # the prior markers must NOT be cleared/persisted before reconcile. This way, a retry can + # still recognize and prune stale keys using the prior marker. + private_writes: list = [] + managed_writes: list = [] + state_saves: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + + original_save = claude.save_state + + def capture_save(s): + state_saves.append(copy.deepcopy(s)) + return original_save(s) + + monkeypatch.setattr(claude, "save_state", capture_save) + + # Capture state at reconcile time to verify markers are still there. + state_at_reconcile: dict = {} + + def failing_reconcile(state, *args, **kwargs): + state_at_reconcile.update(copy.deepcopy(state)) + raise RuntimeError("Simulated reconcile failure") + + monkeypatch.setattr(claude, "_reconcile_managed_settings", failing_reconcile) + + # DISCOVERY run (no static_models, so markers will be cleared AFTER reconcile). + # Start with prior markers from a previous static run. + state = { + "workspace": WS, + claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY: { + "availableModels": ["old-model"], + "enforceAvailableModels": True, + }, + } + + # Reconcile will fail, but prior markers should remain unpersisted before the call. + # The key is: markers were NOT cleared and persisted BEFORE reconcile. + with pytest.raises(RuntimeError, match="reconcile failure"): + claude.write_tool_config(state, None) + + # Verify the prior marker was STILL in state when reconcile was called. + # This proves markers were NOT cleared before reconcile on a DISCOVERY run. + assert claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY in state_at_reconcile + assert state_at_reconcile[claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY][ + "availableModels" + ] == ["old-model"] + + def test_discovery_run_with_reconcile_success_clears_markers(self, monkeypatch): + # FIX 2c & FIX A: When a DISCOVERY run (no static picker) has a SUCCESSFUL reconcile, + # the markers should be cleared and persisted AFTER reconcile (only if the file was written). + private_writes: list = [] + managed_writes: list = [] + state_saves: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + + original_save = claude.save_state + + def capture_save(s): + state_saves.append(copy.deepcopy(s)) + return original_save(s) + + monkeypatch.setattr(claude, "save_state", capture_save) + + # Mock _reconcile_managed_settings to return True (file was written). + monkeypatch.setattr(claude, "_reconcile_managed_settings", lambda *a, **kw: True) + + # DISCOVERY run (no static_models) - start with prior markers from a previous static run. + state = { + "workspace": WS, + claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY: { + "availableModels": ["old-model"], + "enforceAvailableModels": True, + }, + } + + claude.write_tool_config(state, None) + + # Verify markers were cleared and persisted AFTER reconcile. + # The last save_state call should have cleared the markers. + assert len(state_saves) > 0 + final_save = state_saves[-1] + assert claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY not in final_save + assert claude.CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY not in final_save + + def test_static_write_run_persists_markers_before_reconcile_fails(self, monkeypatch): + # FIX 2a: When a WRITE run (static picker) has a reconcile failure, the new markers + # must be persisted BEFORE reconcile so they survive the failure. + private_writes: list = [] + managed_writes: list = [] + state_saves: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + + original_save = claude.save_state + + def capture_save(s): + state_saves.append(copy.deepcopy(s)) + return original_save(s) + + monkeypatch.setattr(claude, "save_state", capture_save) + + def failing_reconcile(*args, **kwargs): + raise RuntimeError("Simulated reconcile failure") + + monkeypatch.setattr(claude, "_reconcile_managed_settings", failing_reconcile) + + # WRITE run (has static_models, so new markers are set). + state = {"workspace": WS, "claude_static_models": ["system.ai.claude-opus-4-8"]} + + with pytest.raises(RuntimeError, match="reconcile failure"): + claude.write_tool_config(state, None) + + # Verify markers were saved BEFORE the failure. + # There should be at least one save_state call before the exception. + assert len(state_saves) > 0 + first_save = state_saves[0] + # The new marker should be in the save (showing it was persisted before reconcile). + assert claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY in first_save + saved_marker = first_save[claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY] + assert "availableModels" in saved_marker + assert "enforceAvailableModels" in saved_marker + + def test_discovery_run_with_suppressed_managed_write_preserves_markers(self, monkeypatch): + # FIX A: When a discovery run (no static picker) has a SUPPRESSED managed write + # (e.g., write suppressed context, or unchanged content), markers should NOT be + # cleared so a later successful reconcile can prune them. + private_writes: list = [] + managed_writes: list = [] + state_saves: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + + original_save = claude.save_state + + def capture_save(s): + state_saves.append(copy.deepcopy(s)) + return original_save(s) + + monkeypatch.setattr(claude, "save_state", capture_save) + + # Mock _reconcile_managed_settings to return False (write was suppressed/skipped) + monkeypatch.setattr(claude, "_reconcile_managed_settings", lambda *a, **kw: False) + + # DISCOVERY run (no static_models) - start with prior markers from a previous static run. + state = { + "workspace": WS, + claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY: { + "availableModels": ["old-model"], + "enforceAvailableModels": True, + }, + } + + claude.write_tool_config(state, None) + + # Verify markers were NOT cleared when managed write was suppressed. + # The last save_state call should have the markers STILL present. + assert len(state_saves) > 0 + final_save = state_saves[-1] + # Markers should be preserved (not cleared) since the managed write was suppressed + assert claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY in final_save + assert final_save[claude.CLAUDE_MANAGED_PRUNED_PICKER_STATE_KEY] == { + "availableModels": ["old-model"], + "enforceAvailableModels": True, + } + + class TestBuildClaudeArgv: def test_no_caller_settings_uses_ucode_file(self, monkeypatch): monkeypatch.setattr(claude, "read_json_safe", lambda p: {"apiKeyHelper": "u"}) @@ -1493,6 +1815,51 @@ def test_malformed_file_json_raises(self, tmp_path, monkeypatch): with pytest.raises(RuntimeError, match="not valid JSON"): claude._build_claude_argv("claude", ["--settings", str(bad_file)]) + def test_managed_discovery_flag_survives_caller_settings(self, monkeypatch): + # FIX P2-2: When managed config enables discovery, preserve the flag even + # when a caller passes --settings. The flag is launched-scoped and should + # survive the merge. + ucode_settings = { + "apiKeyHelper": "ucode-helper", + "env": { + "ANTHROPIC_BASE_URL": "https://gw", + "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1", + }, + } + monkeypatch.setattr(claude, "read_json_safe", lambda p: ucode_settings) + caller = json.dumps({"statusLine": {"type": "command", "command": "sl"}}) + argv = claude._build_claude_argv("claude", ["--settings", caller, "-p", "hi"]) + # Exactly one --settings reaches Claude, and the discovery flag is preserved. + assert argv.count("--settings") == 1 + assert argv[:2] == ["claude", "--settings"] + assert argv[3:] == ["-p", "hi"] + merged = json.loads(argv[2]) + # Discovery flag from managed config should survive. + assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert merged["env"]["ANTHROPIC_BASE_URL"] == "https://gw" + assert merged["statusLine"] == {"type": "command", "command": "sl"} + + def test_stale_discovery_flag_stripped_without_managed_config(self, monkeypatch): + # When managed config does NOT enable discovery, strip any stale flag + # that might come from a --settings merge. + ucode_settings = { + "apiKeyHelper": "ucode-helper", + "env": {"ANTHROPIC_BASE_URL": "https://gw"}, + } + monkeypatch.setattr(claude, "read_json_safe", lambda p: ucode_settings) + caller = json.dumps( + { + "env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1"}, + "statusLine": {"type": "command", "command": "sl"}, + } + ) + argv = claude._build_claude_argv("claude", ["--settings", caller]) + merged = json.loads(argv[2]) + # Stale discovery flag should be stripped since managed config doesn't have it. + assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in merged["env"] + assert merged["env"]["ANTHROPIC_BASE_URL"] == "https://gw" + assert merged["statusLine"] == {"type": "command", "command": "sl"} + class TestClaudeSmartRouting: def _capture_write(self, monkeypatch, existing, written): @@ -1639,3 +2006,101 @@ def test_missing_login_runs_browser_flow(self, monkeypatch): monkeypatch.setattr(claude, "print_success", lambda *a, **kw: None) claude._ensure_subscription_login() assert calls == [[claude.SPEC["binary"], "auth", "login"]] + + +class TestManagedModelPicker: + """A managed static `names` list drives Claude Code's own picker allow-list; a discovery + location instead turns on gateway model discovery.""" + + WS = "https://ws.example.com" + + def test_static_names_write_available_models_and_picker(self): + overlay, keys = claude.render_overlay( + self.WS, None, {}, static_models=["system.ai.claude-opus-4-8", "system.ai.kimi-k3"] + ) + assert overlay["availableModels"] == ["system.ai.claude-opus-4-8", "system.ai.kimi-k3"] + assert overlay["enforceAvailableModels"] is True + assert overlay["modelPicker"] == { + "replaceBuiltInOptions": True, + "options": [ + {"model": "system.ai.claude-opus-4-8", "label": "claude-opus-4-8"}, + {"model": "system.ai.kimi-k3", "label": "kimi-k3"}, + ], + } + for key in claude.CLAUDE_MANAGED_PICKER_KEYS: + assert [key] in keys + + def test_model_service_location_enables_gateway_discovery(self): + overlay, keys = claude.render_overlay(self.WS, None, {}, model_service_location="system.ai") + assert overlay["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert "availableModels" not in overlay + assert ["env", "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] in keys + + def test_provider_suppresses_the_static_picker(self): + overlay, _ = claude.render_overlay( + self.WS, None, {}, provider="cat.sch.mps", static_models=["system.ai.claude-opus-4-8"] + ) + assert "availableModels" not in overlay + assert "modelPicker" not in overlay + + def test_stale_picker_keys_pruned_when_no_static_list(self, tmp_path, monkeypatch): + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text( + json.dumps({"availableModels": ["old"], "enforceAvailableModels": True}) + ) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "backup.json") + monkeypatch.setattr(claude, "_reconcile_managed_settings", lambda *a, **k: None) + monkeypatch.setattr(claude, "save_state", lambda s: None) + claude.write_tool_config({"workspace": self.WS}, None) + doc = json.loads(settings_path.read_text()) + assert "availableModels" not in doc + assert "enforceAvailableModels" not in doc + + def test_static_to_dynamic_transition_prunes_ucode_owned_picker(self, tmp_path, monkeypatch): + settings_path = tmp_path / "ucode-settings.json" + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "backup.json") + monkeypatch.setattr(claude, "_reconcile_managed_settings", lambda *a, **k: None) + saved_state = {} + monkeypatch.setattr(claude, "save_state", lambda s: saved_state.update(s)) + # Step 1: Write with static models, generating a picker and setting the marker. + state1 = {"workspace": self.WS, "claude_static_models": ["system.ai.claude-opus-4-8"]} + claude.write_tool_config(state1, None) + doc1 = json.loads(settings_path.read_text()) + assert "modelPicker" in doc1 + assert "availableModels" in doc1 + assert claude.CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY in saved_state + saved_picker = saved_state[claude.CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY] + # Step 2: Transition to discovery config, reusing same on-disk settings. + state2 = { + "workspace": self.WS, + claude.CLAUDE_MANAGED_MODEL_PICKER_STATE_KEY: saved_picker, + "claude_model_service_location": "system.ai", + } + claude.write_tool_config(state2, None) + doc2 = json.loads(settings_path.read_text()) + # The ucode-owned picker should be pruned; discovery flag should be set. + assert "modelPicker" not in doc2 + assert "availableModels" not in doc2 + assert "enforceAvailableModels" not in doc2 + assert doc2.get("env", {}).get("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY") == "1" + + def test_admin_picker_preserved_under_dynamic_transition(self, tmp_path, monkeypatch): + settings_path = tmp_path / "ucode-settings.json" + admin_picker = { + "replaceBuiltInOptions": True, + "options": [{"model": "system.ai.claude-custom", "label": "custom"}], + } + settings_path.write_text(json.dumps({"modelPicker": admin_picker})) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "backup.json") + monkeypatch.setattr(claude, "_reconcile_managed_settings", lambda *a, **k: None) + monkeypatch.setattr(claude, "save_state", lambda s: None) + # Apply discovery config without any ownership marker set. + state = {"workspace": self.WS, "claude_model_service_location": "system.ai"} + claude.write_tool_config(state, None) + doc = json.loads(settings_path.read_text()) + # Admin picker should be preserved even though we're in discovery mode. + assert doc["modelPicker"] == admin_picker + assert doc.get("env", {}).get("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY") == "1" diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 2ca9ad61..a0ccf0d5 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os from pathlib import Path @@ -414,6 +415,38 @@ def test_legacy_write_preserves_other_profiles_in_shared_config(self, tmp_path, assert doc["profiles"]["other"]["model_provider"] == "keep" assert doc["profiles"]["ucode"]["model_provider"] == "ucode-databricks" + def test_legacy_layout_prunes_stale_provider_header_on_transition(self, tmp_path, monkeypatch): + # FIX #4: Prune the stale routing header in the legacy layout when transitioning + # from provider to static/discovery (no provider). + config_dir = tmp_path / ".codex" + config_dir.mkdir() + legacy_path = config_dir / "config.toml" + # Simulates a prior run with a provider that wrote the header. + legacy_path.write_text( + 'profile = "ucode"\n' + "[profiles.ucode]\n" + 'model_provider = "ucode-databricks"\n' + "[model_providers.ucode-databricks]\n" + 'http_headers = { "Databricks-Model-Provider-Service" = "main.old.provider" }\n', + encoding="utf-8", + ) + profile_path = config_dir / "ucode.config.toml" + backup_path = tmp_path / "codex-ucode-config.backup.toml" + legacy_backup_path = tmp_path / "codex-config.backup.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", backup_path) + monkeypatch.setattr(codex, "LEGACY_CODEX_CONFIG_PATH", legacy_path) + monkeypatch.setattr(codex, "LEGACY_CODEX_BACKUP_PATH", legacy_backup_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.133.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + # Rewrite in legacy layout WITHOUT provider - should prune the stale header. + codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}, provider=None) + + doc = read_toml_safe(legacy_path) + headers = doc["model_providers"]["ucode-databricks"].get("http_headers", {}) + assert "Databricks-Model-Provider-Service" not in headers + class TestCodexLegacyLayoutDetection: def test_new_codex_uses_modern_layout(self, monkeypatch): @@ -841,3 +874,283 @@ def deny_managed_write(*args, **kwargs): with pytest.raises(managed_files.ManagedFileWriteUnavailable, match="sudo denied"): codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) + + +class TestCodexStaticCatalog: + """A managed static `names` list is written as a `model_catalog_json` catalog of full presets.""" + + def test_build_catalog_fetches_and_filters_gateway_models(self, monkeypatch): + # Mock the gateway fetch to return real ModelPreset structures. + gateway_response = { + "models": [ + { + "slug": "system.ai.kimi-k3", + "display_name": "kimi-k3", + "priority": 100, + "base_instructions": "Full operating prompt for kimi...", + "visibility": "list", + "shell_type": "shell_command", + "truncation_policy": {"mode": "tokens", "limit": 10000}, + "supported_reasoning_levels": [ + {"effort": "low", "description": "Fast"}, + {"effort": "high", "description": "Deep"}, + ], + }, + { + "slug": "gpt-5.4", + "display_name": "gpt-5.4", + "priority": 90, + "base_instructions": "Full operating prompt for gpt-5.4...", + "visibility": "list", + "shell_type": "shell_command", + }, + { + "slug": "other-model", + "display_name": "other", + "priority": 50, + "base_instructions": "Should be filtered out...", + "visibility": "hidden", + }, + ] + } + + def mock_fetch(workspace, token): + return gateway_response + + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", mock_fetch) + + catalog = codex.build_codex_catalog(WS, "token", ["system.ai.kimi-k3", "gpt-5.4"]) + assert list(catalog) == ["models"] + assert len(catalog["models"]) == 2 + first, second = catalog["models"] + + # Verify real fields are preserved verbatim from gateway. + assert first["slug"] == "system.ai.kimi-k3" + assert first["display_name"] == "kimi-k3" + assert first["priority"] == 100 + assert first["base_instructions"] == "Full operating prompt for kimi..." + assert first["visibility"] == "list" + assert first["shell_type"] == "shell_command" + assert first["truncation_policy"] == {"mode": "tokens", "limit": 10000} + assert first["supported_reasoning_levels"][0]["effort"] == "low" + + # Gateway order is preserved for filtered entries. + assert second["slug"] == "gpt-5.4" + assert second["priority"] == 90 + + def test_build_catalog_returns_none_on_fetch_failure(self, monkeypatch): + # FIX C: Fetch failure should return None so existing catalog is preserved. + # Mock fetch to return None (failure). + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", lambda w, t: None) + + result = codex.build_codex_catalog(WS, "token", ["system.ai.kimi-k3"]) + assert result is None + + def test_build_catalog_returns_empty_dict_when_no_matching_slugs(self, monkeypatch): + # FIX C: Success with no matches should return {} so catalog falls back to discovery. + # Mock gateway catalog with models not in the allow-list. + gateway_response = { + "models": [ + { + "slug": "other-model", + "display_name": "other", + "priority": 50, + "base_instructions": "prompt", + "visibility": "list", + } + ] + } + + def mock_fetch(workspace, token): + return gateway_response + + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", mock_fetch) + + result = codex.build_codex_catalog(WS, "token", ["system.ai.kimi-k3"]) + assert result == {} + + def test_build_catalog_returns_empty_dict_when_gateway_empty(self, monkeypatch): + # FIX C: When fetch succeeds but returns empty models list, return {} for fallback. + gateway_response = {"models": []} + + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", lambda w, t: gateway_response) + + result = codex.build_codex_catalog(WS, "token", ["system.ai.kimi-k3"]) + assert result == {} + + def test_write_config_emits_catalog_file_and_reference(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + catalog_path = tmp_path / ".codex" / "ucode-models.json" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "CODEX_CATALOG_PATH", catalog_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + monkeypatch.setattr(codex, "get_databricks_token", lambda w, p: "token") + + # Mock the gateway fetch to return real presets. + gateway_response = { + "models": [ + { + "slug": "system.ai.kimi-k3", + "display_name": "kimi-k3", + "priority": 100, + "base_instructions": "kimi prompt", + "visibility": "list", + }, + { + "slug": "gpt-5.4", + "display_name": "gpt-5.4", + "priority": 90, + "base_instructions": "gpt prompt", + "visibility": "list", + }, + ] + } + + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", lambda w, t: gateway_response) + + codex.write_tool_config( + {"workspace": WS, "codex_static_models": ["system.ai.kimi-k3", "gpt-5.4"]} + ) + + doc = read_toml_safe(config_path) + assert doc["model_catalog_json"] == str(catalog_path) + written = json.loads(catalog_path.read_text()) + assert [m["slug"] for m in written["models"]] == ["system.ai.kimi-k3", "gpt-5.4"] + # Verify real fields are preserved from gateway. + assert written["models"][0]["base_instructions"] == "kimi prompt" + + def test_fetch_failure_preserves_existing_catalog(self, tmp_path, monkeypatch): + # FIX C: On fetch failure, preserve existing catalog so transient network blips + # don't destroy a good catalog. + config_path = tmp_path / ".codex" / "ucode.config.toml" + catalog_path = tmp_path / ".codex" / "ucode-models.json" + catalog_path.parent.mkdir(parents=True, exist_ok=True) + existing_catalog = {"models": [{"slug": "old-good-model", "display_name": "Old Good"}]} + catalog_path.write_text(json.dumps(existing_catalog)) + + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "CODEX_CATALOG_PATH", catalog_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + monkeypatch.setattr(codex, "get_databricks_token", lambda w, p: "token") + # Mock fetch to return None (failure). + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", lambda w, t: None) + + codex.write_tool_config({"workspace": WS, "codex_static_models": ["system.ai.kimi-k3"]}) + + # On fetch failure, existing catalog should be preserved (not unlinked). + assert catalog_path.exists() + preserved = json.loads(catalog_path.read_text()) + assert preserved["models"][0]["slug"] == "old-good-model" + # Config should have the catalog reference preserved. + doc = read_toml_safe(config_path) + assert "model_catalog_json" in doc + + def test_catalog_removed_when_success_but_no_matches(self, tmp_path, monkeypatch): + # FIX C: When fetch succeeds but nothing matches the allow-list, remove the catalog + # to fall back to discovery (not on transient network failure). + config_path = tmp_path / ".codex" / "ucode.config.toml" + catalog_path = tmp_path / ".codex" / "ucode-models.json" + catalog_path.parent.mkdir(parents=True, exist_ok=True) + # Pre-populate with a stale catalog. + catalog_path.write_text('{"models": []}') + + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "CODEX_CATALOG_PATH", catalog_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + monkeypatch.setattr(codex, "get_databricks_token", lambda w, p: "token") + + # Mock fetch to succeed but return models not in allow-list. + gateway_response = { + "models": [{"slug": "other-model", "display_name": "Other", "priority": 50}] + } + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", lambda w, t: gateway_response) + + codex.write_tool_config({"workspace": WS, "codex_static_models": ["system.ai.kimi-k3"]}) + + # Catalog should be removed (no matches), allowing discovery to run. + assert not catalog_path.exists() + doc = read_toml_safe(config_path) + assert "model_catalog_json" not in doc + + def test_stale_catalog_removed_when_no_static_list(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + catalog_path = tmp_path / ".codex" / "ucode-models.json" + catalog_path.parent.mkdir(parents=True, exist_ok=True) + catalog_path.write_text('{"models": []}') + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "CODEX_CATALOG_PATH", catalog_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + codex.write_tool_config({"workspace": WS}) + + assert not catalog_path.exists() + assert "model_catalog_json" not in read_toml_safe(config_path) + + def test_provider_suppresses_static_catalog(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + catalog_path = tmp_path / ".codex" / "ucode-models.json" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "CODEX_CATALOG_PATH", catalog_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + codex.write_tool_config( + {"workspace": WS, "codex_static_models": ["system.ai.kimi-k3"]}, + provider="main.default.mps", + ) + + assert not catalog_path.exists() + assert "model_catalog_json" not in read_toml_safe(config_path) + + def test_stale_provider_header_removed_on_transition_to_static(self, tmp_path, monkeypatch): + """P1-4: Stale Databricks-Model-Provider-Service header removed on provider->static transition.""" + config_path = tmp_path / ".codex" / "ucode.config.toml" + catalog_path = tmp_path / ".codex" / "ucode-models.json" + # Pre-populate config with a stale provider header from a prior provider run. + config_path.parent.mkdir(parents=True, exist_ok=True) + config_path.write_text( + 'model_provider = "ucode-databricks"\n' + "[model_providers.ucode-databricks]\n" + 'name = "Gateway"\n' + "[model_providers.ucode-databricks.http_headers]\n" + '"Databricks-Model-Provider-Service" = "main.old.provider"\n', + encoding="utf-8", + ) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "CODEX_CATALOG_PATH", catalog_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + monkeypatch.setattr(codex, "get_databricks_token", lambda w, p: "token") + + # Mock the gateway fetch. + gateway_response = { + "models": [ + { + "slug": "system.ai.kimi-k3", + "display_name": "kimi-k3", + "priority": 100, + "base_instructions": "prompt", + "visibility": "list", + } + ] + } + monkeypatch.setattr(codex, "_fetch_gateway_codex_models", lambda w, t: gateway_response) + + # Transition from provider to static (no provider parameter). + codex.write_tool_config({"workspace": WS, "codex_static_models": ["system.ai.kimi-k3"]}) + + # Verify stale header is removed. + doc = read_toml_safe(config_path) + headers = doc.get("model_providers", {}).get("ucode-databricks", {}).get("http_headers", {}) + assert "Databricks-Model-Provider-Service" not in headers + assert "model_catalog_json" in doc # Static catalog is present. diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index b7c78def..b83ccacb 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -14,8 +14,10 @@ managed_default_model, managed_enabled_tools, managed_launch_model, + managed_model_service_location, managed_provider_service, managed_state_overrides, + managed_static_models, managed_supplies_models, managed_unservable_models, recommended_agent, @@ -401,6 +403,71 @@ def test_false_when_slots_are_present_but_blank(self): } assert managed_supplies_models(managed, "claude") is False + def test_true_for_flat_agent_with_static_names_list(self): + # FIX 1: Flat-list agents (gemini, opencode, pi, copilot) with managed static `names` + # list must return True so discovery is skipped. + managed = { + "enabled_agents": {"gemini": {"model_config": {"names": ["system.ai.gemini-3-flash"]}}} + } + assert managed_supplies_models(managed, "gemini") is True + + def test_true_for_opencode_with_static_names_list(self): + # Verify opencode also recognizes the names list. + managed = { + "enabled_agents": { + "opencode": {"model_config": {"names": ["system.ai.claude-opus-4-8"]}} + } + } + assert managed_supplies_models(managed, "opencode") is True + + def test_true_for_flat_agent_with_legacy_models_list(self): + # Legacy flat `models` list for flat-agents should still work. + managed = { + "enabled_agents": {"gemini": {"model_config": {"models": ["system.ai.gemini-3-flash"]}}} + } + assert managed_supplies_models(managed, "gemini") is True + + def test_false_when_names_list_is_blank(self): + # Empty or whitespace-only names list should not count. + managed = {"enabled_agents": {"gemini": {"model_config": {"names": [" ", ""]}}}} + assert managed_supplies_models(managed, "gemini") is False + + def test_false_for_codex_with_only_model_service_location(self): + # FIX B: Codex does NOT read model_service_location, so location-only config + # should return False to allow discovery to run normally. Claude returns True for + # location-only config because it reads it, but codex must still discover. + managed = { + "enabled_agents": { + "codex": {"model_config": {"model_service_location": "main.catalog.models"}} + } + } + assert managed_supplies_models(managed, "codex") is False + + def test_true_for_claude_with_only_model_service_location(self): + # Claude DOES read model_service_location, so location-only config returns True + # to suppress discovery (gateway discovery will run instead). + managed = { + "enabled_agents": { + "claude": {"model_config": {"model_service_location": "main.catalog.models"}} + } + } + assert managed_supplies_models(managed, "claude") is True + + def test_true_for_codex_with_static_models_and_location(self): + # Codex with static models (names) should return True even if location is also set, + # since static models take precedence. + managed = { + "enabled_agents": { + "codex": { + "model_config": { + "names": ["gpt-4o"], + "model_service_location": "main.catalog.models", + } + } + } + } + assert managed_supplies_models(managed, "codex") is True + class TestManagedStateOverrides: """Each agent reads its models from a different shape, so the manifest has to be translated.""" @@ -598,3 +665,134 @@ def test_default_model_stands_without_a_recommendation(self): def test_none_when_neither_names_a_model(self): assert managed_launch_model({}, None, "pi") is None + + +class TestStaticAndAutoModels: + """The static `names` allow-list and the `model_service_location` auto-discovery source.""" + + @staticmethod + def _managed(tool, model_config): + return {"default_agent": tool, "enabled_agents": {tool: {"model_config": model_config}}} + + def test_static_models_reads_the_names_list(self): + m = self._managed("claude", {"names": ["system.ai.claude-opus-4-8", "system.ai.kimi-k3"]}) + assert managed_static_models(m, "claude") == [ + "system.ai.claude-opus-4-8", + "system.ai.kimi-k3", + ] + + def test_static_models_none_when_absent_or_empty(self): + assert managed_static_models(self._managed("claude", {"names": []}), "claude") is None + assert managed_static_models(self._managed("claude", {}), "claude") is None + + def test_model_service_location_read(self): + m = self._managed("codex", {"model_service_location": "main.agents"}) + assert managed_model_service_location(m, "codex") == "main.agents" + + def test_names_and_location_count_as_supplying_models(self): + assert managed_supplies_models(self._managed("claude", {"names": ["x"]}), "claude") is True + # FIX B: Codex does NOT read model_service_location, so it should return False + # even when location is set. Claude returns True for location-only config. + assert ( + managed_supplies_models( + self._managed("codex", {"model_service_location": "system.ai"}), "codex" + ) + is False + ) + # But claude with location should return True + assert ( + managed_supplies_models( + self._managed("claude", {"model_service_location": "system.ai"}), "claude" + ) + is True + ) + + def test_overrides_layer_static_and_location_for_claude_and_codex(self): + m = self._managed("claude", {"names": ["a", "b"]}) + assert managed_state_overrides(m, "claude")["claude_static_models"] == ["a", "b"] + m2 = self._managed("codex", {"model_service_location": "main.agents"}) + assert managed_state_overrides(m2, "codex")["codex_model_service_location"] == "main.agents" + + def test_resolve_state_layers_static_models_into_state(self): + m = self._managed("claude", {"names": ["system.ai.claude-opus-4-8"]}) + resolved = resolve_state(m, {"workspace": WORKSPACE}, "claude") + assert resolved["claude_static_models"] == ["system.ai.claude-opus-4-8"] + + +class TestFlatListAgentsWithStaticNames: + """Flat-list agents (gemini, opencode, pi, copilot) should read managed static model lists from the `names` key.""" + + @staticmethod + def _managed(tool, model_config): + return {"default_agent": tool, "enabled_agents": {tool: {"model_config": model_config}}} + + @pytest.mark.parametrize("tool", ["gemini", "opencode", "pi", "copilot"]) + def test_flat_list_agents_resolve_managed_names_list(self, tool): + # Managed static model lists (from wire `model_services`) normalize to internal `names` key. + # Flat-list agents should resolve this to their agent-specific state key. + m = self._managed( + tool, + { + "names": [ + "system.ai.claude-opus-4-8", + "system.ai.gemini-3-flash", + "system.ai.kimi-k2-7-code", + ] + }, + ) + overrides = managed_state_overrides(m, tool) + # For opencode, names are bucketed by provider; for others, a flat list + if tool == "opencode": + assert overrides == { + "opencode_models": { + "anthropic": ["system.ai.claude-opus-4-8"], + "gemini": ["system.ai.gemini-3-flash"], + "oss": ["system.ai.kimi-k2-7-code"], + } + } + else: + assert overrides == { + f"{tool}_models": [ + "system.ai.claude-opus-4-8", + "system.ai.gemini-3-flash", + "system.ai.kimi-k2-7-code", + ] + } + + @pytest.mark.parametrize("tool", ["gemini", "pi", "copilot"]) + def test_flat_list_agents_names_list_in_resolved_state(self, tool): + # The resolved state passed to write_tool_config should contain the managed list. + m = self._managed(tool, {"names": ["model-a", "model-b", "model-c"]}) + state = _state() + resolved = resolve_state(m, state, tool) + assert resolved[f"{tool}_models"] == ["model-a", "model-b", "model-c"] + + @pytest.mark.parametrize("tool", ["gemini", "opencode", "pi", "copilot"]) + def test_flat_list_agents_legacy_models_list_fallback(self, tool): + # Backward compat: legacy flat `models` list should still work when `names` is absent. + m = self._managed(tool, {"models": ["legacy-a", "legacy-b"]}) + overrides = managed_state_overrides(m, tool) + if tool == "opencode": + # opencode still gets bucketed (no family classification for "legacy-*" so likely empty) + assert "opencode_models" not in overrides or overrides["opencode_models"] == {} + else: + assert overrides == {f"{tool}_models": ["legacy-a", "legacy-b"]} + + @pytest.mark.parametrize("tool", ["gemini", "opencode", "pi", "copilot"]) + def test_names_takes_precedence_over_models(self, tool): + # When both `names` and `models` are present, `names` should win. + m = self._managed( + tool, + { + "names": ["system.ai.claude-opus-4-8"], + "models": ["system.ai.claude-sonnet-4-6"], + }, + ) + overrides = managed_state_overrides(m, tool) + if tool == "opencode": + # opencode still gets bucketed + assert overrides.get("opencode_models", {}).get("anthropic") == [ + "system.ai.claude-opus-4-8" + ] + else: + assert overrides[f"{tool}_models"] == ["system.ai.claude-opus-4-8"]