[AIGTWY-4710] Keep managed models across config transitions (#709)

Keeps a coding agent's managed model settings in sync with the managed
config across workspace transitions (AIGTWY-4710). Switching between
managed configs reconciles each agent's model settings to the current
config, and configuring a workspace with no managed config clears ug's
managed model settings so an unmanaged workspace never enforces a stale
list inherited from a previous managed workspace.

Lifecycle exercised, per agent (functional in-process tests plus
end-to-end integration journeys):
- no managed config: no picker (Claude) / no catalog (Codex).
- static config A: the picker/catalog lists exactly A's models.
- static config B: reconciles to B; models A listed that B omits are
pruned.
- model discovery via a Model Provider Service: the picker/catalog is
cleared, since the `Databricks-Model-Provider-Service` header routes and
there is no static list.
- no managed config again: the picker/catalog is cleared, not inherited.

Changes:
- Claude clears the managed picker (`availableModels`,
`enforceAvailableModels`, `modelPicker`) whenever ug is not writing a
static list this run (an unmanaged workspace or a provider switch). It
always clears the picker in its own `ucode-settings.json`. In the OS
`managed-settings.json` the picker is shared with the administrator, so
ug treats it as one unit and reverts it to its pre-ucode baseline only
when the live picker is unchanged since ug last wrote it (a three-way
merge of baseline, ug's last write, and the live file). That way ug
removes only what it introduced and never an administrator's own picker,
a picker they added after ug cleared its own, or one they later edited
(the enforce flag included).
- Codex clears its generated catalog and `model_catalog_json` pointer on
an unmanaged configure or a provider switch, so an unmanaged workspace
uses its own discovery. A managed static config writes them; the
launch-time `model` pin follows Codex's existing clear-when-unmanaged
behavior.
- Prettifies the Claude `/model` picker labels: the enforced allow-list
still shows one row per model (it can include multiple versions per
family and non-Claude models), but each label is now title-cased with a
dotted version and known acronyms upper-cased
(`system.ai.claude-haiku-4-5` shows as `Claude Haiku 4.5`,
`system.ai.glm-5-3` as `GLM 5.3`) instead of the raw id. The routable id
still shows on the right of each row.
- The `UCODE_MANAGED_CONFIG_STUB` test hook accepts an explicit JSON
`null` to reproduce a workspace with no managed config, distinct from an
unset or unreadable stub (which still falls through to the real HTTP
read).
- Fixes a pre-existing bug where a managed config's per-family default
models were dropped. `managed_resolve._manifest_models` read a stale
`models` key instead of the normalized `default_models_by_model_family`,
so `managed_claude_family_models` returned nothing and no
`ANTHROPIC_DEFAULT_<family>_MODEL` was written (only the overall default
and the picker survived, on separate paths). Unit fixtures that
hand-built the stale shape masked it; the fix corrects them and adds a
regression test that builds through `normalize_managed_config` so the
shape can't drift again.
- Reconciles the per-family default env keys the same three-way way as
the picker so a workspace switch does not carry stale defaults over.
Each `ANTHROPIC_DEFAULT_<family>_MODEL` resolves as config, then an
administrator-authored value (one that differs from ug's last write,
preserved verbatim), then discovery; a value ug itself wrote is
re-derived from the new workspace rather than preserved, and any family
default whose model is not in the enforced `availableModels` list is
dropped (e.g. a discovered `fable` default when the admin's list omits
fable).
- Adds functional lifecycle tests (both agents) covering the full
sequence above, including the per-family default env keys landing (opus
and sonnet carry ug's `[1m]` context tag), per-agent `managed_fixture`
integration journeys running it through `ug configure` against the
managed e2e workspace, and coverage-matrix updates. The integration MPS
states point at real provider services on that workspace:
`main.default.ci_e2e_anthropic_mps` (Claude) and
`main.default.ci_e2e_openai_mps` (Codex).

What Claude Code writes for models, by managed model source:
- static (`model_services`): sets `availableModels`,
`enforceAvailableModels`, and `modelPicker` in
`~/.claude/ucode-settings.json` and the OS `managed-settings.json`. The
managed default is applied at launch as `env.ANTHROPIC_MODEL`.
- model discovery via MPS (`model_provider_service`): writes no model
list and clears ug's static picker; the
`Databricks-Model-Provider-Service` header routes and Claude shows its
own model names. A Bedrock-backed provider pins
`ANTHROPIC_DEFAULT_{OPUS,SONNET,HAIKU}_MODEL` from the manifest's family
slots; an Anthropic provider pins nothing.
- model discovery via `unity_catalog_location`: not consumed on main
(schema only; no fields written).

What Codex writes for models, by managed model source:
- static (`model_services`): writes the catalog file
`~/.ucode/codex-model-catalog.json` and points `ucode.config.toml`'s
`model_catalog_json` at it. The managed default is pinned as `model` in
`ucode.config.toml`.
- model discovery via MPS (`model_provider_service`): writes no catalog
and clears any static catalog; the `Databricks-Model-Provider-Service`
header routes and Codex discovers models at launch.
- model discovery via `unity_catalog_location`: not consumed on main
(schema only; no fields written).

Known limitations (out of scope here):
- Clearing on an unmanaged configure also resets in the rare case where
the same workspace's managed config genuinely disappeared, since ug does
not distinguish that from an intentional switch to an unmanaged
workspace.
- On a non-interactive configure (no sudo, `managed_writes_allowed()`
False), ug cannot rewrite the root-owned OS `managed-settings.json`, so
a stale picker there is not cleared in that path; the private
`ucode-settings.json` still is. The interactive flow clears both.
- If ug's managed-file snapshots (its pre-ucode baseline or last write)
are unreadable, ug leaves the managed-file picker in place rather than
risk clobbering an admin's entry.

This pull request and its description were written by Isaac.

Co-authored-by: Isaac <no-reply@databricks.com>

Author: david-siqi-liu

src/ucode/agents/claude.py

@@ -46,10 +46,12 @@
 from ucode.launcher import exec_or_spawn
 from ucode.managed_files import (
     OS,
+    ManagedFileSnapshots,
     ManagedFileWriteUnavailable,
     current_os,
     managed_file_conflicts,
     managed_file_is_verified,
+    managed_file_snapshots,
     managed_file_status,
     managed_writes_allowed,
     mark_managed_file_verified,
@@ -198,7 +200,6 @@ def _otel_trace_env(workspace: str) -> dict[str, str]:
 # settings file on every launch so stale values never linger.
 CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",)
 CLAUDE_MANAGED_PICKER_KEYS = ("availableModels", "enforceAvailableModels", "modelPicker")
-CLAUDE_PRUNED_PICKER_KEYS = ("availableModels", "enforceAvailableModels")
 ANTHROPIC_CUSTOM_HEADERS_ENV_KEY = "ANTHROPIC_CUSTOM_HEADERS"
 CLAUDE_MANAGED_CUSTOM_HEADER_NAMES = frozenset(
     {
@@ -495,9 +496,27 @@ def render_overlay(
     return overlay, keys
 
 
+_MODEL_LABEL_ACRONYMS = frozenset({"glm", "gpt"})
+
+
 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.")
+    """A human-friendly picker label for a model id (e.g. ``system.ai.claude-haiku-4-5`` ->
+    ``Claude Haiku 4.5``): keep the vendor and name words title-cased, uppercase known acronyms,
+    and join a run of numeric segments into a dotted version."""
+    stem = model.removeprefix("system.ai.")
+    parts: list[str] = []
+    version: list[str] = []
+    for token in stem.split("-"):
+        if token.isdigit():
+            version.append(token)
+            continue
+        if version:
+            parts.append(".".join(version))
+            version = []
+        parts.append(token.upper() if token in _MODEL_LABEL_ACRONYMS else token.title())
+    if version:
+        parts.append(".".join(version))
+    return " ".join(parts) if parts else stem
 
 
 def _maybe_add_1m_suffix(model: str) -> str:
@@ -522,24 +541,29 @@ def _enforce_model_default_hierarchy(
     coding_agent_config_defaults: dict[str, str],
     settings_file_existing_defaults: dict[str, str],
     ucode_defaults: dict[str, str],
+    ucode_last_written_defaults: dict[str, str],
+    enforced_models: list[str] | None,
 ) -> str | None:
-    """Apply managed-file model precedence for one Claude family."""
-    coding_agent_config_default_model = coding_agent_config_defaults.get(family)
-    settings_file_existing_default_model = settings_file_existing_defaults.get(family)
-    ucode_default_model = ucode_defaults.get(family)
-
-    if coding_agent_config_default_model is not None:
-        selected_default_model = coding_agent_config_default_model
-    elif settings_file_existing_default_model is not None:
-        return settings_file_existing_default_model
-    else:
-        selected_default_model = ucode_default_model
+    """Resolve one Claude family's managed-file default model.
 
-    if selected_default_model is None:
+    An existing managed-file default ucode did not write itself (it differs from ucode's last write)
+    is an administrator's, so it is preserved verbatim. Otherwise the value is ucode's own or unset,
+    so ucode re-derives it from the coding-agent config, then discovery, resetting a value carried
+    over from a previous workspace, and drops the result when an enforced model list excludes it.
+    """
+    selected = coding_agent_config_defaults.get(family)
+    if selected is None:
+        existing = settings_file_existing_defaults.get(family)
+        if existing is not None and existing != ucode_last_written_defaults.get(family):
+            return existing
+        selected = ucode_defaults.get(family)
+    if selected is None:
         return None
     if family in ("opus", "sonnet"):
-        return _maybe_add_1m_suffix(selected_default_model)
-    return selected_default_model
+        selected = _maybe_add_1m_suffix(selected)
+    if enforced_models is not None and selected.split("[", 1)[0] not in enforced_models:
+        return None
+    return selected
 
 
 def add_claude_mcp_server(
@@ -757,7 +781,12 @@ def write_tool_config(
 
     # 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,
+        managed_settings_snapshots: ManagedFileSnapshots | None,
+    ) -> dict:
         base_env = base.get("env")
         existing_custom_headers = (
             base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None
@@ -778,12 +807,26 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict:
                 managed_overlay.get("claude_models") or state.get("claude_models") or {}
             )
 
+            enforced_models = overlay_for_merge.get("availableModels")
+            last_applied_env = {}
+            if (
+                managed_settings_snapshots is not None
+                and managed_settings_snapshots.last_applied_by_ug
+            ):
+                last_applied_env = managed_settings_snapshots.last_applied_by_ug.get("env") or {}
+            ucode_last_written_defaults = {
+                family: last_applied_env[key]
+                for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items()
+                if isinstance(last_applied_env.get(key), str)
+            }
             for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items():
                 selected_default_model = _enforce_model_default_hierarchy(
                     family,
                     coding_agent_config_defaults=configured_defaults,
                     settings_file_existing_defaults=settings_file_existing_defaults,
                     ucode_defaults=ucode_defaults,
+                    ucode_last_written_defaults=ucode_last_written_defaults,
+                    enforced_models=enforced_models,
                 )
                 if selected_default_model is None:
                     target_env.pop(key, None)
@@ -816,19 +859,43 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict:
             # longer writes.
             for key in CLAUDE_REMOVED_ENV_KEYS:
                 merged_env.pop(key, None)
+        if not any(key in overlay_for_merge for key in CLAUDE_MANAGED_PICKER_KEYS):
+            if managed_settings_snapshots is None:
+                for key in CLAUDE_MANAGED_PICKER_KEYS:
+                    merged.pop(key, None)
+            elif managed_settings_snapshots.last_applied_by_ug is not None:
+                last_applied = managed_settings_snapshots.last_applied_by_ug
+                live_picker = [merged.get(key) for key in CLAUDE_MANAGED_PICKER_KEYS]
+                ucode_picker = [last_applied.get(key) for key in CLAUDE_MANAGED_PICKER_KEYS]
+                if live_picker == ucode_picker:
+                    baseline = managed_settings_snapshots.original_before_ug or {}
+                    for key in CLAUDE_MANAGED_PICKER_KEYS:
+                        if key in baseline:
+                            merged[key] = baseline[key]
+                        else:
+                            merged.pop(key, None)
         if "otelHeadersHelper" not in overlay_for_merge:
             merged.pop("otelHeadersHelper", None)
         sync_smart_routing_hooks(merged, state, enabled=False)
         return merged
 
+    managed_snapshots = managed_file_snapshots("claude", _parse_managed_settings)
     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,
+            managed_settings_snapshots=None,
+        ),
     )
 
     _reconcile_managed_settings(
         state,
-        lambda base: _compose(base, enforce_model_default_hierarchy=provider is None),
+        lambda base: _compose(
+            base,
+            enforce_model_default_hierarchy=provider is None,
+            managed_settings_snapshots=managed_snapshots,
+        ),
         managed_file_keys,
         relayed,
     )
@@ -919,8 +986,9 @@ def _reconcile_managed_settings(
     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` updates gateway-owned fields in this file. It writes the picker
+    (`availableModels`/`modelPicker`) for a static managed list and removes the picker keys it
+    previously wrote when it no longer manages one, leaving an administrator's own picker untouched.
 
     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.

src/ucode/managed_config.py

@@ -511,6 +511,13 @@ def get_model_recommendation(workspace: str, token: str) -> tuple[dict | None, s
     }, None
 
 
+class _StubUnset:
+    """Sentinel type: the stub env var is unset or unreadable, so fall through to the real read."""
+
+
+_STUB_UNSET = _StubUnset()
+
+
 def get_managed_config(workspace: str, token: str) -> FetchedManagedConfig:
     """Fetch and normalize the workspace's managed config.
 
@@ -529,11 +536,14 @@ def get_managed_config(workspace: str, token: str) -> FetchedManagedConfig:
     v0 stores at most one config per workspace, so the first entry is the workspace's config.
 
     ``UCODE_MANAGED_CONFIG_STUB`` short-circuits the HTTP read: when it names a readable JSON file,
-    that file's single CodingAgentConfig is used verbatim. It exists so this client can be exercised
-    against the managed-config shape before the server emits it (AIGTWY-4572); unset in normal use.
+    that file's single CodingAgentConfig is used verbatim, or an explicit ``null`` stands in for a
+    workspace with no managed config. It exists so this client can be exercised against the
+    managed-config shape before the server emits it (AIGTWY-4572); unset in normal use.
     """
     stub = _stub_config()
-    if stub is not None:
+    if not isinstance(stub, _StubUnset):
+        if stub is None:
+            return FetchedManagedConfig(None, None)
         return _gate_config(stub)
     configs, reason = fetch_managed_coding_agent_configs(workspace, token)
     if reason is not None:
@@ -548,17 +558,22 @@ def get_managed_config(workspace: str, token: str) -> FetchedManagedConfig:
     return _gate_config(configs[0])
 
 
-def _stub_config() -> dict | None:
-    """The stub CodingAgentConfig named by ``UCODE_MANAGED_CONFIG_STUB``, or None when unset/bad."""
+def _stub_config() -> dict | None | _StubUnset:
+    """The stub named by ``UCODE_MANAGED_CONFIG_STUB``: the CodingAgentConfig dict it holds, ``None``
+    when it is an explicit JSON ``null`` (a workspace with no managed config), or ``_STUB_UNSET``
+    when the var is unset, the file cannot be read, or its content is not a config dict (fall
+    through to the real read)."""
     path = os.environ.get("UCODE_MANAGED_CONFIG_STUB")
     if not path:
-        return None
+        return _STUB_UNSET
     try:
         raw = json.loads(Path(path).read_text(encoding="utf-8"))
     except (OSError, UnicodeError, json.JSONDecodeError) as exc:
         print_warning(f"UCODE_MANAGED_CONFIG_STUB could not be read ({exc}); ignoring it.")
+        return _STUB_UNSET
+    if raw is None:
         return None
-    return raw if isinstance(raw, dict) else None
+    return raw if isinstance(raw, dict) else _STUB_UNSET
 
 
 def _gate_config(raw: dict) -> FetchedManagedConfig:

src/ucode/managed_files.py

@@ -16,6 +16,7 @@
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from copy import deepcopy
+from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 from typing import Any, cast
@@ -168,6 +169,45 @@ def _print_managed_write_permission(display: str) -> None:
     _managed_write_notice_shown = True
 
 
+@dataclass
+class ManagedFileSnapshots:
+    """The pre-ucode baseline and ucode's last write for a managed file, for a three-way merge."""
+
+    original_before_ug: dict | None
+    last_applied_by_ug: dict | None
+
+
+def managed_file_snapshots(tool: str, parser: ManagedParser) -> ManagedFileSnapshots:
+    """Return the parsed baseline and last-applied snapshots of ``tool``'s managed file.
+
+    ``original_before_ug`` is the pre-ucode baseline; ``last_applied_by_ug`` is what ucode last
+    wrote. A caller reverting a value ucode owns compares the live value against
+    ``last_applied_by_ug`` (unchanged since ucode wrote it) and restores ``original_before_ug`` (the
+    value before ucode, if any) as a three-way merge, so it removes only what ucode itself introduced
+    and never an administrator's own or edited entry. Either field is None when its snapshot is absent
+    or cannot be read/parsed, so callers fall back to keeping the live value.
+    """
+
+    def _parse(text: str | None) -> dict | None:
+        if text is None:
+            return None
+        try:
+            return parser(text)
+        except Exception:  # noqa: BLE001 - an unparseable snapshot just means "unknown".
+            return None
+
+    try:
+        entry = _manifest_files(_load_manifest()).get(tool)
+        if not isinstance(entry, dict):
+            return ManagedFileSnapshots(None, None)
+        return ManagedFileSnapshots(
+            _parse(_snapshot_text(entry, "backup_file")),
+            _parse(_snapshot_text(entry, "last_applied_file")),
+        )
+    except RuntimeError:
+        return ManagedFileSnapshots(None, None)
+
+
 def managed_file_conflicts(
     existing: dict, desired: dict, owned_paths: list[list[str]]
 ) -> list[str]:

src/ucode/managed_resolve.py

@@ -126,14 +126,16 @@ 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)
     if tool == "claude":
+        family_slots = _as_dict(model_config.get("default_models_by_model_family"))
         slots: dict[str, str] = {}
         for slot, family in _CLAUDE_FAMILY_SLOTS.items():
-            model = _str(_as_dict(manifest_models).get(slot))
+            model = _str(family_slots.get(slot))
             if model:
                 slots[family] = model
         return slots or None
+    manifest_models = model_config.get("models")
     if isinstance(manifest_models, list):
         listed = [model for model in (_str(item) for item in manifest_models) if model]
         return listed or None

tests/README.md

@@ -52,6 +52,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`.
 | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup |
 | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch |
 | `test_ug_configure_managed_codex_catalog_fallback` | Configure from an injected managed response containing a GPT model absent from Codex's bundled catalog | Actionable metadata warning; conservative catalog entry for the unknown model; real Codex prompt on the valid default model |
+| `test_managed_fixture_claude_model_lifecycle`, `test_managed_fixture_codex_model_lifecycle` | Configure across no config -> static A -> static B -> MPS -> no config (stub-injected, `null` for no-config; MPS via a real provider service) | Each agent's model files reconcile to each static config (removed models pruned); switching to an MPS and a workspace with no managed config both clear ug's managed model settings so no stale list is enforced |
 | `test_ug_installed_wheel_exposes_help_and_version` | Invoke freshly installed console command | Package version matches; public help works |
 | `test_ug_status_in_fresh_home_is_unconfigured` | Request status before configure | Unconfigured status |
 | `test_ug_auth_without_configuration_explains_how_to_configure` | Request auth before configure | Actionable setup error and nonzero exit |
@@ -60,7 +61,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`.
 
 With both agents selected there are **41 live cases** (6 interactive TUI cases),
 **3 managed-workspace cases** (marker `managed`, run against a separate workspace that
-publishes a CodingAgentConfig), **3 managed-fixture cases** (marker `managed_fixture`, with only
+publishes a CodingAgentConfig), **5 managed-fixture cases** (marker `managed_fixture`, with only
 the CodingAgentConfig input injected), and **5 installation checks**. Parametrization varies
 argument spelling or routing mode, never hides the agent/provider in the test name. Duplicate boot-only cases
 are incorporated into the Databricks configuration TUI journeys.

tests/integration/README.md

@@ -101,6 +101,7 @@ test_ug_configure_codex_lifecycle.py    # repeat setup, revert, rejected credent
 test_ug_configure_managed.py            # managed workspace: static model list, no agent selector
 test_ug_configure_managed_models.py     # injected model lists: pickers and Codex fallback metadata
 test_ug_configure_managed_mcp.py        # injected managed MCP list
+test_ug_configure_managed_lifecycle.py  # none -> A -> B -> MPS -> none: reconcile, clear on MPS/no-config
 test_installation.py                   # fresh installed package
 utils/                                # process/terminal/evidence helpers and Docker files
 ```
@@ -157,9 +158,14 @@ fails the selected CUJ, rather than skipping it.
 There are **42 live cases** (including 6 TUI journeys) and **5 installation
 checks** with both agents. A separate **3 managed-workspace cases** (one per agent
 plus an idempotent re-configure, marker `managed`) run against a workspace that publishes a CodingAgentConfig; see
-"Managed-workspace journeys" below. A further **3 `managed_fixture` cases** inject the admin config
+"Managed-workspace journeys" below. A further **5 `managed_fixture` cases** inject the admin config
 locally (via `UCODE_MANAGED_CONFIG_STUB`) to cover shapes the live workspace does not publish; each
 differs from the published config in what it asserts so it proves the injected config drove configure.
+Two of those are per-agent lifecycle journeys (no config -> static A -> static B -> MPS -> no config):
+they assert the generated model files reconcile to each static config, are cleared when switching to
+a Model Provider Service, and are cleared when the workspace has no managed config (an explicit
+`null` stub reproduces the no-config states; the MPS states use real provider services on the managed
+workspace, `main.default.ci_e2e_anthropic_mps` and `main.default.ci_e2e_openai_mps`).
 See the named coverage and gaps matrix in
 [../README.md](../README.md).