From 4ad71743bc03295faf789cb61041b3db2a35eae4 Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Tue, 4 Aug 2026 23:12:33 +0000 Subject: [PATCH 1/3] Resolve managed coding-agent config over local state at launch --- src/ucode/cli.py | 50 ++++-- src/ucode/managed_config.py | 88 +++++++++-- src/ucode/managed_resolve.py | 115 ++++++++++++++ src/ucode/state.py | 34 +++- tests/test_managed_config.py | 157 ++++++++++++++++++- tests/test_managed_resolve.py | 284 ++++++++++++++++++++++++++++++++++ 6 files changed, 696 insertions(+), 32 deletions(-) create mode 100644 src/ucode/managed_resolve.py create mode 100644 tests/test_managed_resolve.py diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 97f8ae7..3bc14c6 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -55,6 +55,8 @@ resolve_pat_token, run_databricks_login, ) +from ucode.managed_config import managed_launch_state +from ucode.managed_resolve import managed_provider_service from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -1170,6 +1172,9 @@ def _launch_tool( if needs_auto_configure: _auto_configure_tool(tool) state = ensure_provider_state(tool) + # Remembered before the fallback below collapses the two cases: a managed config may not + # silently override a provider the user typed on the command line (it errors instead). + explicit_provider = provider # An explicit --provider overrides the persisted choice; otherwise fall # back to whatever `ucode configure` saved for this tool. provider = provider or get_provider_service(state, tool) @@ -1179,17 +1184,6 @@ def _launch_tool( f"{TOOL_SPECS[tool]['display']} smart routing cannot be enabled with " "--provider. Launch without a Model Provider Service and try again." ) - # Validate the provider service before launching — it must exist, be a - # provider type this tool can route to (e.g. claude can't use an OpenAI - # or Foundry service), and, for Bedrock, expose Claude models to pin. - # Surfaces a clear error up front instead of a cryptic gateway failure - # mid-session. For a Bedrock service this also returns the model ids. - provider_models = None - relayed = False - if provider: - provider_models, error, relayed = resolve_provider_models(tool, state, provider) - if error: - raise RuntimeError(error) # Re-fetch model lists on every launch so newly-added Databricks # endpoints show up without a manual `ucode configure` (and so that # tools like pi which read multiple model bundles never run on @@ -1202,6 +1196,38 @@ def _launch_tool( skip_model_discovery=bool(provider), skip_preflight=skip_preflight, ) + # An admin-published managed config wins over the developer's own settings. Resolved before + # the provider and model are settled below, so each is decided once against the values that + # will actually be written — the two state files are never merged on disk. + state, managed = managed_launch_state(state, tool) + if managed is not None: + managed_provider = managed_provider_service(managed, tool) + if explicit_provider and managed_provider and managed_provider != explicit_provider: + # An explicit --provider that disagrees with the admin's is a hard error rather + # than a silent override: the user asked for something the managed config forbids, + # and quietly routing them elsewhere would hide it. + raise RuntimeError( + f"You cannot launch {TOOL_SPECS[tool]['display']} with provider " + f"{explicit_provider} because your admin has specified managed provider " + f"{managed_provider}." + ) + if managed_provider: + provider = managed_provider + # Validate the provider service before launching — it must exist, be a + # provider type this tool can route to (e.g. claude can't use an OpenAI + # or Foundry service), and, for Bedrock, expose Claude models to pin. + provider_models = None + relayed = False + if provider: + provider_models, error, relayed = resolve_provider_models(tool, state, provider) + if error: + if managed is not None and provider == managed_provider_service(managed, tool): + # Clear error if the admin has Unity Catalog grants the developer doesn't. + raise RuntimeError( + f"Your admin's managed config specifies provider {provider} for " + f"{TOOL_SPECS[tool]['display']}, which can't be used: {error}" + ) + raise RuntimeError(error) if routing_agent is not None and enable_smart_routing_flag: state = routing_agent.enable_smart_routing(state) # The router's per-launch pick for the root session. Codex pins it as the @@ -1241,6 +1267,8 @@ def _launch_tool( route_root_model=route_root_model, ) print_section(f"ucode with {TOOL_SPECS[tool]['display']}") + if managed is not None: + print_kv("Config", "workspace-managed") if provider: print_kv("Provider", provider) elif route_root_model: diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index ec39805..9e49f63 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -4,11 +4,13 @@ (non-admin) and ``ucode`` applies it locally. This module owns the developer-read half: - fetching the raw manifest (via :func:`ucode.databricks.fetch_managed_coding_agent_configs`), -- normalizing the proto-JSON into a stable internal dict keyed by ucode's own tool names, and -- persisting it to ``~/.ucode/managed-state.json`` (0600) so launches can reconcile against it. +- normalizing the proto-JSON into a stable internal dict keyed by ucode's own tool names, +- persisting it to ``~/.ucode/managed-state.json`` (0600), and +- re-reading it on each launch, falling back to the persisted copy when the read fails. -Reconciliation against the local ``state.json`` and applying the manifest to agents live in later -changes; this module deliberately stops at "read + normalize + persist". +:func:`managed_launch_state` is the launch path's entry point: it refreshes the manifest and hands +back the state to configure the agent with. Deciding *which* value wins for a given key is +:mod:`ucode.managed_resolve`'s job, kept separate so that logic stays pure and I/O-free. """ from __future__ import annotations @@ -19,7 +21,9 @@ from typing import cast import ucode.config_io as config_io -from ucode.databricks import fetch_managed_coding_agent_configs +from ucode.databricks import fetch_managed_coding_agent_configs, get_databricks_token +from ucode.managed_resolve import resolve_state +from ucode.ui import print_warning MANAGED_STATE_PATH = config_io.APP_DIR / "managed-state.json" @@ -271,6 +275,10 @@ def save_managed_state(workspace: str, config: dict) -> None: The file is org-authored, not developer-editable — 0600 keeps it readable/writable only by the user (a light guard; hard enforcement / sudo ownership is a separate concern). No-op in dry-run. + + An empty ``config`` records "this workspace has no managed config", which matters because the + file doubles as the fallback when a later read fails: without it, removing a config server-side + would leave the old one on disk to be reapplied after a transient outage. """ if config_io.is_dry_run(): return @@ -307,11 +315,67 @@ def load_managed_state(workspace: str | None) -> dict | None: return config if isinstance(config, dict) else None -def delete_managed_state() -> None: - """Remove the managed-state file, if any. No-op in dry-run.""" - if config_io.is_dry_run(): - return +def refresh_managed_config(state: dict) -> dict | None: + """Fetch the workspace's managed config and persist it, returning the normalized manifest. + + Runs on every launch so a developer picks up an admin's edits without re-running + ``ucode configure``. Returns None when the workspace has no managed config — the normal case for + a workspace whose admin hasn't published one. + + A failed fetch never blocks the launch: an unreachable control plane shouldn't stop someone from + coding. Instead it falls back to the last config persisted for this workspace, so the admin's + most recent known policy still applies; only when there is no persisted config either does the + launch fall through to the developer's own settings. + """ + workspace = state.get("workspace") + if not workspace: + return None try: - MANAGED_STATE_PATH.unlink(missing_ok=True) - except OSError as exc: - raise RuntimeError(f"Failed to remove managed state file: {MANAGED_STATE_PATH}") from exc + token = get_databricks_token(workspace, state.get("profile")) + except RuntimeError as exc: + return _persisted_fallback(workspace, str(exc)) + managed, reason = get_managed_config(workspace, token) + if reason is not None: + return _persisted_fallback(workspace, reason) + if managed is None: + # Record that this workspace has no config, rather than leaving an earlier one on disk: + # the file doubles as the fallback above, so a removed policy would otherwise come back + # into force after the next transient outage. + save_managed_state(workspace, {}) + return None + save_managed_state(workspace, managed) + return managed + + +def _persisted_fallback(workspace: str, reason: str) -> dict | None: + """Return the last persisted config for ``workspace`` after a failed fetch, warning either way. + + Distinguishes the two outcomes in the warning: continuing on a possibly-stale admin config is + materially different from continuing on the developer's own settings. + """ + # An empty persisted config means the last successful read found none, so there is no admin + # policy to fall back to — treat it the same as having no file at all. + persisted = load_managed_state(workspace) + if persisted: + print_warning( + f"Could not read your workspace's managed config ({reason}); " + "using the last one saved for this workspace." + ) + return persisted + print_warning( + f"Could not read your workspace's managed config ({reason}); using your local settings." + ) + return None + + +def managed_launch_state(state: dict, tool: str) -> tuple[dict, dict | None]: + """Return ``(state, managed)`` for launching ``tool`` under any managed config. + + The returned state has the manifest's models and provider layered over the developer's own — + managed wins per key — so the settings file written from it reflects the admin's choices. When + the workspace has no managed config the state is handed back untouched. + """ + managed = refresh_managed_config(state) + if managed is None: + return state, None + return resolve_state(managed, state, tool), managed diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py new file mode 100644 index 0000000..dec85cf --- /dev/null +++ b/src/ucode/managed_resolve.py @@ -0,0 +1,115 @@ +"""Resolve the effective agent settings from the managed config plus local ucode state. + +The admin-authored manifest (``~/.ucode/managed-state.json``, written by +:mod:`ucode.managed_config`) and the developer's own ucode state (``~/.ucode/state.json``) stay +separate files — they are never merged on disk. Instead this module resolves them *per key* at +config-write time: whatever the manifest specifies wins, and anything it leaves unset falls back to +the developer's ucode state. The resolved view is what gets rendered into the agent config files +(e.g. ``~/.claude/ucode-settings.json``), so managed settings take precedence for every ``ucode`` +command without either file being rewritten. + +Only settings the developer set *through* ucode participate in the fallback. Settings they wrote by +hand outside ucode (``~/.claude/settings.json``, etc.) are not read here — Claude Code merges +those scopes itself at launch, underneath the file ucode passes via ``--settings``. + +Everything here is pure: no I/O, no mutation of the inputs. Fetching and persisting the manifest, +and handing the resolved state to the agent config writers, live in :mod:`ucode.managed_config`. +""" + +from __future__ import annotations + +from typing import cast + +from ucode.state import MANAGED_OVERLAY_KEY + +# Proto model-config slot -> the family key `claude.py`'s render_overlay reads. The manifest keeps +# the proto spelling (`default_opus_model`), while ucode state and render_overlay both key claude +# models by bare family (`opus`), so the two have to be bridged before the settings file is written. +_CLAUDE_FAMILY_SLOTS = { + "default_opus_model": "opus", + "default_sonnet_model": "sonnet", + "default_haiku_model": "haiku", + "default_fable_model": "fable", +} + + +def _as_dict(value: object) -> dict[str, object]: + """Return ``value`` as a ``dict[str, object]`` when it is a dict, else an empty dict.""" + return cast("dict[str, object]", value) if isinstance(value, dict) else {} + + +def _str(value: object) -> str | None: + """Return a non-empty stripped string, or None.""" + if isinstance(value, str): + stripped = value.strip() + return stripped or None + return None + + +def _agent_entry(managed: dict, tool: str) -> dict[str, object]: + """Return the manifest's config for ``tool``, or an empty dict when it isn't enabled.""" + enabled = _as_dict(_as_dict(managed).get("enabled_agents")) + return _as_dict(enabled.get(tool)) + + +def _agent_model_config(managed: dict, tool: str) -> dict[str, object]: + """Return the manifest's normalized ``model_config`` for ``tool``, if any.""" + return _as_dict(_agent_entry(managed, tool).get("model_config")) + + +def effective_agent_models(managed: dict, state: dict, tool: str) -> dict | list | None: + """Resolve ``tool``'s model list/slots from the manifest, falling back to ucode state. + + Claude keys its models by family (``opus``/``sonnet``/``haiku``/``fable``) and resolves per + family, so a family the manifest omits keeps the developer's value. Every other agent stores a + flat list, which has no per-key identity — there the manifest's list replaces the local one + outright, or the local one stands when the manifest specifies none. + """ + manifest_models = _agent_model_config(managed, tool).get("models") + if tool == "claude": + local = dict(_as_dict(state.get("claude_models"))) + for slot, family in _CLAUDE_FAMILY_SLOTS.items(): + model = _str(_as_dict(manifest_models).get(slot)) + if model: + local[family] = model + return local or None + if isinstance(manifest_models, list): + models = [m for m in (_str(item) for item in manifest_models) if m] + if models: + return models + local_list = state.get(f"{tool}_models") + return local_list if local_list else None + + +def managed_provider_service(managed: dict, tool: str) -> str | None: + """Return only the provider the managed config specifies for ``tool``, ignoring local state.""" + return _str(_agent_model_config(managed, tool).get("model_provider_service")) + + +def resolve_state(managed: dict, state: dict, tool: str) -> dict: + """Return a copy of ``state`` with ``tool``'s managed values layered on top. + + ``write_tool_config`` reads its models and provider out of the state dict it is handed, so + handing it this resolved copy is what makes managed settings win. Each key the managed config + displaces is recorded under :data:`~ucode.state.MANAGED_OVERLAY_KEY` with the developer's own + value (None when they had none), which ``save_state`` swaps back before writing — so the admin's + settings reach the generated agent config files without ``state.json`` losing what the developer + configured. The two files are never merged on disk. + """ + resolved = dict(state) + overlay: dict[str, object] = {} + models = effective_agent_models(managed, state, tool) + models_key = f"{tool}_models" + if models is not None and models != state.get(models_key): + overlay[models_key] = state.get(models_key) + resolved[models_key] = models + provider = managed_provider_service(managed, tool) + if provider: + providers = dict(_as_dict(state.get("provider_services"))) + if providers.get(tool) != provider: + overlay["provider_services"] = state.get("provider_services") + providers[tool] = provider + resolved["provider_services"] = providers + if overlay: + resolved[MANAGED_OVERLAY_KEY] = overlay + return resolved diff --git a/src/ucode/state.py b/src/ucode/state.py index 7e53d75..0031bc4 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -13,6 +13,10 @@ STATE_PATH = APP_DIR / "state.json" STATE_VERSION = 3 +# Transient key holding the developer's own values for whatever a managed config layered over them. +# Present only in memory: the layered values render the agent settings files, while `save_state` +# restores what's under it so `state.json` keeps recording the developer's own configuration. +MANAGED_OVERLAY_KEY = "_managed_overlay" AUTH_COMMAND_TIMEOUT_MS = 5000 AUTH_REFRESH_INTERVAL_MS = 900_000 @@ -42,14 +46,21 @@ def load_state() -> dict: def save_state(state: dict) -> None: - """Save workspace state back into the per-workspace structure.""" + """Save workspace state back into the per-workspace structure. + + Values a managed config layered over the developer's own are stripped first (see + ``MANAGED_OVERLAY_KEY``), so an admin-published config takes effect through the generated agent + settings files without overwriting what the developer configured for themselves. Read + non-destructively: a launch can save more than once from the same dict (e.g. the relayed proxy + rewriting its port), and every one of those writes must restore the developer's values. + """ if is_dry_run(): return full = load_full_state() workspace = state.get("workspace") or full.get("current_workspace") if workspace: full["current_workspace"] = workspace - full["workspaces"][workspace] = hydrate_state(state) + full["workspaces"][workspace] = hydrate_state(_without_managed_overlay(state)) try: APP_DIR.mkdir(parents=True, exist_ok=True) STATE_PATH.write_text(json.dumps(full, indent=2), encoding="utf-8") @@ -57,6 +68,25 @@ def save_state(state: dict) -> None: raise RuntimeError(f"Failed to write state file: {STATE_PATH}") from exc +def _without_managed_overlay(state: dict) -> dict: + """Return ``state`` with managed-config values swapped back for the developer's own. + + Returns a new dict and leaves ``state`` untouched, so the caller keeps the layered values it + needs for rendering and repeated saves stay idempotent. + """ + overlay = state.get(MANAGED_OVERLAY_KEY) + if not isinstance(overlay, dict): + return state + persisted = {key: value for key, value in state.items() if key != MANAGED_OVERLAY_KEY} + for key, value in overlay.items(): + # A key the developer never set is dropped rather than persisted as None. + if value is None: + persisted.pop(key, None) + else: + persisted[key] = value + return persisted + + def set_current_workspace(workspace: str | None) -> None: """Set ``current_workspace`` without touching the per-workspace blocks. diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 6229a44..346b3ba 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -12,7 +12,9 @@ from ucode.managed_config import ( get_managed_config, load_managed_state, + managed_launch_state, normalize_managed_config, + refresh_managed_config, save_managed_state, ) @@ -203,14 +205,12 @@ def test_load_missing_returns_none(self, _managed_path): def test_load_none_workspace_returns_none(self, _managed_path): assert load_managed_state(None) is None - def test_delete_removes_file(self, _managed_path): + def test_empty_config_overwrites_a_previous_one(self, _managed_path): + # Saving an empty config is how "the admin removed it" is recorded: the stored config must + # be replaced, not left behind for the read-failure fallback to reapply. save_managed_state("https://ws.example.com", {"default_agent": "claude"}) - assert _managed_path.exists() - mc_mod.delete_managed_state() - assert not _managed_path.exists() - - def test_delete_missing_is_noop(self, _managed_path): - mc_mod.delete_managed_state() # should not raise + save_managed_state("https://ws.example.com", {}) + assert load_managed_state("https://ws.example.com") == {} class TestFetchClient: @@ -247,3 +247,146 @@ def test_http_failure_surfaces_reason(self, monkeypatch): configs, reason = db_mod.fetch_managed_coding_agent_configs("https://ws", "tok") assert configs == [] assert reason == "HTTP 403 Forbidden" + + +WORKSPACE = "https://ws.example.com" + +# A normalized managed config, as `normalize_managed_config` produces it. +MANAGED = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-5", + "models": {"default_opus_model": "system.ai.claude-opus-5"}, + } + } + }, +} + + +def _state(**overrides) -> dict: + state = {"workspace": WORKSPACE, "managed_configs": {"claude": {"keys": []}}} + state.update(overrides) + return state + + +class TestRefreshManagedConfig: + """The per-launch re-read, so an admin's edits land without re-running `ucode configure`.""" + + @pytest.fixture(autouse=True) + def _stub_token(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_databricks_token", lambda ws, profile: "tok") + + def test_persists_and_returns_the_manifest(self, monkeypatch): + saved: list[tuple] = [] + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (MANAGED, None)) + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: saved.append((ws, cfg))) + assert refresh_managed_config(_state()) == MANAGED + assert saved == [(WORKSPACE, MANAGED)] + + def test_no_managed_config_returns_none(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + assert refresh_managed_config(_state()) is None + + def test_read_failure_falls_back_to_the_persisted_config(self, monkeypatch): + # The admin's last known policy beats no policy, so a failed fetch reuses what we saved. + warnings: list[str] = [] + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + assert refresh_managed_config(_state()) == MANAGED + assert "HTTP 500" in warnings[0] + assert "last one saved" in warnings[0] + + def test_read_failure_without_persisted_config_uses_local_settings(self, monkeypatch): + warnings: list[str] = [] + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + assert refresh_managed_config(_state()) is None + assert "your local settings" in warnings[0] + + def test_auth_failure_falls_back_to_the_persisted_config(self, monkeypatch): + warnings: list[str] = [] + + def boom(ws, profile): + raise RuntimeError("no token") + + monkeypatch.setattr(mc_mod, "get_databricks_token", boom) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + assert refresh_managed_config(_state()) == MANAGED + assert "no token" in warnings[0] + + def test_auth_failure_without_persisted_config_uses_local_settings(self, monkeypatch): + warnings: list[str] = [] + + def boom(ws, profile): + raise RuntimeError("no token") + + monkeypatch.setattr(mc_mod, "get_databricks_token", boom) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + assert refresh_managed_config(_state()) is None + assert "your local settings" in warnings[0] + + def test_no_config_on_the_server_does_not_use_a_stale_persisted_file(self, monkeypatch): + # A successful read saying "no config" means the admin removed it — that's authoritative, + # so a previously persisted file must not resurrect the old policy. + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setattr( + mc_mod, "load_managed_state", lambda ws: pytest.fail("must not fall back") + ) + assert refresh_managed_config(_state()) is None + + def test_no_config_on_the_server_clears_the_persisted_one(self, monkeypatch): + # Without this, removing the config server-side would leave the old one on disk and the next + # failed read would put a dead policy back into force. + saved: list[tuple] = [] + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: saved.append((ws, cfg))) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + assert refresh_managed_config(_state()) is None + assert saved == [(WORKSPACE, {})] + + def test_empty_persisted_config_is_not_treated_as_a_fallback(self, monkeypatch): + # The empty marker means "no admin policy", so a later failed read must fall through to the + # developer's own settings rather than reporting a managed config. + warnings: list[str] = [] + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: {}) + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + assert refresh_managed_config(_state()) is None + assert "your local settings" in warnings[0] + + def test_no_workspace_is_a_noop(self, monkeypatch): + monkeypatch.setattr( + mc_mod, "get_managed_config", lambda ws, tok: pytest.fail("should not fetch") + ) + assert refresh_managed_config({}) is None + + +class TestManagedLaunchState: + @pytest.fixture(autouse=True) + def _stub_token(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + + def test_layers_managed_models_when_a_config_exists(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (MANAGED, None)) + state = _state(claude_models={"opus": "local-opus"}) + resolved, managed = managed_launch_state(state, "claude") + assert managed == MANAGED + assert resolved["claude_models"]["opus"] == "system.ai.claude-opus-5" + # The developer's own state is untouched — precedence is resolved in memory. + assert state["claude_models"]["opus"] == "local-opus" + + def test_state_untouched_when_no_managed_config(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) + state = _state(claude_models={"opus": "local-opus"}) + resolved, managed = managed_launch_state(state, "claude") + assert managed is None + assert resolved is state diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py new file mode 100644 index 0000000..e328d96 --- /dev/null +++ b/tests/test_managed_resolve.py @@ -0,0 +1,284 @@ +"""Tests for managed_resolve.py / managed_apply.py — resolving and writing managed agent settings.""" + +from __future__ import annotations + +import json + +import pytest + +import ucode.agents.claude as claude +import ucode.config_io as config_io +import ucode.state as state_mod +from ucode.managed_resolve import ( + effective_agent_models, + managed_provider_service, + resolve_state, +) +from ucode.state import MANAGED_OVERLAY_KEY + +WORKSPACE = "https://ws.example.com" + +# A normalized managed config, as `managed_config.normalize_managed_config` produces it. +MANAGED = { + "name": "coding-agent-configs/abc-123", + "default_agent": "claude", + "enabled_agents": { + "claude": { + "use_as_global_settings": True, + "model_config": { + "default_model": "system.ai.claude-opus-5", + "models": { + "default_opus_model": "system.ai.claude-opus-5", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + "default_haiku_model": "system.ai.claude-haiku-4-5", + }, + }, + }, + "codex": { + "model_config": { + "default_model": "databricks-gpt-5-3-codex", + "models": ["databricks-gpt-5-3-codex", "databricks-gpt-5-2-codex"], + } + }, + }, + "budget_policy": {"display_name": "paved-path", "tiers": []}, +} + + +def _state(**overrides) -> dict: + state = { + "workspace": WORKSPACE, + "managed_configs": {"claude": {"keys": []}, "codex": {"keys": []}}, + } + state.update(overrides) + return state + + +class TestClaudeModels: + def test_proto_slots_map_to_families(self): + # The manifest keeps proto spelling (`default_opus_model`); render_overlay reads `opus`. + models = effective_agent_models(MANAGED, _state(), "claude") + assert models == { + "opus": "system.ai.claude-opus-5", + "sonnet": "system.ai.claude-sonnet-4-6", + "haiku": "system.ai.claude-haiku-4-5", + } + + def test_manifest_wins_over_local_per_family(self): + state = _state(claude_models={"opus": "system.ai.claude-opus-4-8"}) + models = effective_agent_models(MANAGED, state, "claude") + assert models["opus"] == "system.ai.claude-opus-5" + + def test_family_absent_from_manifest_keeps_local_value(self): + # Claude resolves per family, so a family the admin didn't pin keeps the developer's choice. + managed = { + "enabled_agents": { + "claude": {"model_config": {"models": {"default_opus_model": "managed-opus"}}} + } + } + state = _state(claude_models={"opus": "local-opus", "fable": "local-fable"}) + models = effective_agent_models(managed, state, "claude") + assert models == {"opus": "managed-opus", "fable": "local-fable"} + + def test_no_manifest_models_falls_back_to_local(self): + state = _state(claude_models={"sonnet": "local-sonnet"}) + assert effective_agent_models({}, state, "claude") == {"sonnet": "local-sonnet"} + + def test_none_when_neither_side_has_models(self): + assert effective_agent_models({}, _state(), "claude") is None + + +class TestListModels: + def test_manifest_list_replaces_local(self): + # A flat list has no per-key identity to merge on, so the manifest's list wins outright. + state = _state(codex_models=["local-codex"]) + assert effective_agent_models(MANAGED, state, "codex") == [ + "databricks-gpt-5-3-codex", + "databricks-gpt-5-2-codex", + ] + + def test_local_list_stands_when_manifest_silent(self): + state = _state(codex_models=["local-codex"]) + assert effective_agent_models({}, state, "codex") == ["local-codex"] + + def test_blank_entries_dropped(self): + managed = {"enabled_agents": {"codex": {"model_config": {"models": [" ", "real", ""]}}}} + assert effective_agent_models(managed, _state(), "codex") == ["real"] + + +class TestManagedProviderService: + """The manifest-only read: needed to attribute a provider to the admin, not to local state.""" + + def test_returns_the_manifest_provider(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"model_provider_service": "main.default.managed"}} + } + } + assert managed_provider_service(managed, "claude") == "main.default.managed" + + def test_ignores_locally_persisted_provider(self): + # No fallback to local state: otherwise a developer's own provider would be misreported as + # the admin's when rejecting a conflicting --provider. + assert managed_provider_service({}, "claude") is None + + def test_none_for_agent_not_in_manifest(self): + assert managed_provider_service(MANAGED, "gemini") is None + + +class TestResolveState: + def test_does_not_mutate_input_state(self): + # managed-state.json and state.json stay separate files: resolution is per-write and + # in-memory, so the developer's own state must come back untouched. + state = _state(claude_models={"opus": "local-opus"}) + before = json.dumps(state, sort_keys=True) + resolve_state(MANAGED, state, "claude") + assert json.dumps(state, sort_keys=True) == before + + def test_layers_managed_models_onto_copy(self): + resolved = resolve_state(MANAGED, _state(), "claude") + assert resolved["claude_models"]["opus"] == "system.ai.claude-opus-5" + + def test_preserves_unrelated_state_keys(self): + resolved = resolve_state(MANAGED, _state(profile="my-profile"), "claude") + assert resolved["profile"] == "my-profile" + assert resolved["workspace"] == WORKSPACE + + def test_layers_provider_without_dropping_other_tools(self): + managed = { + "enabled_agents": { + "codex": {"model_config": {"model_provider_service": "main.default.managed"}} + } + } + state = _state(provider_services={"claude": "main.default.keep"}) + resolved = resolve_state(managed, state, "codex") + assert resolved["provider_services"] == { + "claude": "main.default.keep", + "codex": "main.default.managed", + } + + +class TestStateFileIsNotRewritten: + """The managed config must win by precedence, not by overwriting the developer's state file. + + managed-state.json and state.json stay separate on disk: resolution happens in memory and only + the generated agent settings file reflects it. These tests deliberately let the real + ``save_state`` run against a temp ``state.json`` — stubbing it out is what let this regress, + because the overwrite happens inside ``write_tool_config``, one layer below the resolver. + """ + + @pytest.fixture + def real_state_file(self, tmp_path, monkeypatch): + """Redirect state.json and the claude settings file into tmp_path, unstubbed.""" + monkeypatch.setattr(config_io, "APP_DIR", tmp_path) + monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "ucode-settings.json") + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "backup.json") + # Seed a developer whose own opus choice differs from the manifest's. + state_mod.save_state( + { + "workspace": WORKSPACE, + "managed_configs": {"claude": {"keys": []}}, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + } + ) + return tmp_path + + @staticmethod + def _persisted_claude_models(tmp_path) -> dict: + full = json.loads((tmp_path / "state.json").read_text()) + return full["workspaces"][WORKSPACE].get("claude_models") or {} + + def test_developers_state_file_keeps_their_own_model(self, real_state_file): + # The developer picked opus-4-8; the manifest says opus-5. After configuring under the + # managed config, state.json must still say opus-4-8 — the admin's value belongs only in + # the generated settings file, so removing the managed config restores their own choice. + assert self._persisted_claude_models(real_state_file)["opus"] == "system.ai.claude-opus-4-8" + + resolved_state = resolve_state(MANAGED, state_mod.load_state(), "claude") + claude.write_tool_config(resolved_state, None) + + assert self._persisted_claude_models(real_state_file)["opus"] == "system.ai.claude-opus-4-8" + + def test_settings_file_gets_the_managed_model(self, real_state_file): + # The other half of the contract: precedence must actually reach the generated file. + resolved_state = resolve_state(MANAGED, state_mod.load_state(), "claude") + claude.write_tool_config(resolved_state, None) + + env = json.loads((real_state_file / "ucode-settings.json").read_text())["env"] + assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"].startswith("system.ai.claude-opus-5") + + def test_overlay_bookkeeping_never_lands_on_disk(self, real_state_file): + resolved_state = resolve_state(MANAGED, state_mod.load_state(), "claude") + claude.write_tool_config(resolved_state, None) + + raw = (real_state_file / "state.json").read_text() + assert MANAGED_OVERLAY_KEY not in raw + + def test_repeated_saves_still_restore_the_developers_value(self, real_state_file): + # A launch can save twice from the same dict (the relayed proxy rewrites its port after + # configure), so the swap-back must be idempotent rather than consuming the overlay. + resolved_state = resolve_state(MANAGED, state_mod.load_state(), "claude") + state_mod.save_state(resolved_state) + state_mod.save_state(resolved_state) + + assert self._persisted_claude_models(real_state_file)["opus"] == "system.ai.claude-opus-4-8" + # The in-memory dict still carries the managed value for rendering. + assert resolved_state["claude_models"]["opus"] == "system.ai.claude-opus-5" + + def test_managed_provider_does_not_overwrite_the_developers_provider( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr(config_io, "APP_DIR", tmp_path) + monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") + state_mod.save_state( + {"workspace": WORKSPACE, "provider_services": {"claude": "main.default.mine"}} + ) + managed = { + "enabled_agents": { + "claude": {"model_config": {"model_provider_service": "main.default.admin"}} + } + } + resolved_state = resolve_state(managed, state_mod.load_state(), "claude") + state_mod.save_state(resolved_state) + + full = json.loads((tmp_path / "state.json").read_text()) + assert full["workspaces"][WORKSPACE]["provider_services"] == {"claude": "main.default.mine"} + assert resolved_state["provider_services"]["claude"] == "main.default.admin" + + def test_developer_with_no_prior_value_is_not_given_one(self, tmp_path, monkeypatch): + # The developer never configured claude models; the manifest supplies them for this launch + # only, so state.json must not gain a key recording the admin's choice as theirs. + monkeypatch.setattr(config_io, "APP_DIR", tmp_path) + monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") + state_mod.save_state({"workspace": WORKSPACE}) + + resolved_state = resolve_state(MANAGED, state_mod.load_state(), "claude") + state_mod.save_state(resolved_state) + + full = json.loads((tmp_path / "state.json").read_text()) + assert not full["workspaces"][WORKSPACE].get("claude_models") + + @pytest.mark.parametrize( + ("tool", "models_key", "managed_models"), + [ + ("codex", "codex_models", ["managed-codex"]), + ("gemini", "gemini_models", ["managed-gemini"]), + ], + ) + def test_other_agents_state_is_also_preserved( + self, tmp_path, monkeypatch, tool, models_key, managed_models + ): + # Every agent's write_tool_config calls save_state, so the swap-back has to hold for all of + # them — not just claude. + monkeypatch.setattr(config_io, "APP_DIR", tmp_path) + monkeypatch.setattr(state_mod, "STATE_PATH", tmp_path / "state.json") + state_mod.save_state({"workspace": WORKSPACE, models_key: ["mine"]}) + managed = {"enabled_agents": {tool: {"model_config": {"models": managed_models}}}} + + resolved_state = resolve_state(managed, state_mod.load_state(), tool) + state_mod.save_state(resolved_state) + + full = json.loads((tmp_path / "state.json").read_text()) + assert full["workspaces"][WORKSPACE][models_key] == ["mine"] + assert resolved_state[models_key] == managed_models From 9cf54e36c78df344a6a4817308b16fb9b9a4de78 Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Tue, 4 Aug 2026 23:38:06 +0000 Subject: [PATCH 2/3] stub managed config lookup in launch tests --- tests/test_cli.py | 35 ++++++++++++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index 32bbc75..eb0cad4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -126,6 +126,10 @@ def _patch_launch(tool: str): "ucode.cli.configure_tool", return_value=MINIMAL_STATE, ), + patch( + "ucode.cli.managed_launch_state", + side_effect=lambda state, tool: (state, None), + ), patch("ucode.cli.launch_agent"), ] @@ -141,7 +145,8 @@ def test_subcommand_calls_correct_tool(self, tool): patches[3], patches[4], patches[5], - patches[6] as mock_launch, + patches[6], + patches[7] as mock_launch, ): result = runner.invoke(app, [tool]) assert result.exit_code == 0, result.output @@ -165,6 +170,7 @@ def test_workspace_flag_sets_current_workspace(self): patches[4], patches[5], patches[6], + patches[7], patch("ucode.cli.set_current_workspace") as mock_set, ): result = runner.invoke( @@ -185,6 +191,7 @@ def test_no_workspace_flag_leaves_current_workspace(self): patches[4], patches[5], patches[6], + patches[7], patch("ucode.cli.set_current_workspace") as mock_set, ): result = runner.invoke(app, ["claude"]) @@ -246,6 +253,10 @@ def test_enabled_codex_launch_uses_routed_root_model(self): return_value=(decision, None), ), patch("ucode.cli.configure_tool", return_value=state) as mock_configure, + patch( + "ucode.cli.managed_launch_state", + side_effect=lambda state, tool: (state, None), + ), patch("ucode.cli.launch_agent"), ): result = runner.invoke(app, ["codex"]) @@ -712,6 +723,10 @@ def test_triggers_when_no_workspace(self): return_value=(configured_state, "databricks-claude-sonnet-4"), ), patch("ucode.cli.configure_tool", return_value=configured_state), + patch( + "ucode.cli.managed_launch_state", + side_effect=lambda state, tool: (state, None), + ), patch("ucode.cli.launch_agent"), ): result = runner.invoke(app, ["claude"]) @@ -736,6 +751,10 @@ def test_triggers_when_tool_not_in_available_tools(self): return_value=(MINIMAL_STATE, "databricks-claude-sonnet-4"), ), patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE), + patch( + "ucode.cli.managed_launch_state", + side_effect=lambda state, tool: (state, None), + ), patch("ucode.cli.launch_agent"), ): result = runner.invoke(app, ["claude"]) @@ -759,6 +778,10 @@ def test_skipped_when_already_configured(self): return_value=(MINIMAL_STATE, "databricks-claude-sonnet-4"), ), patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE), + patch( + "ucode.cli.managed_launch_state", + side_effect=lambda state, tool: (state, None), + ), patch("ucode.cli.launch_agent"), ): runner.invoke(app, ["claude"]) @@ -787,7 +810,8 @@ def test_extra_args_forwarded(self, tool, extra_args): patches[3], patches[4], patches[5], - patches[6] as mock_launch, + patches[6], + patches[7] as mock_launch, ): result = runner.invoke(app, [tool, *extra_args]) assert result.exit_code == 0, result.output @@ -803,7 +827,8 @@ def test_no_extra_args_passes_empty_list(self): patches[3], patches[4], patches[5], - patches[6] as mock_launch, + patches[6], + patches[7] as mock_launch, ): runner.invoke(app, ["claude"]) forwarded = mock_launch.call_args[0][2] @@ -2053,6 +2078,10 @@ def _patches(cfg): return_value=(MINIMAL_STATE, "databricks-claude-sonnet-4"), ), patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE), + patch( + "ucode.cli.managed_launch_state", + side_effect=lambda state, tool: (state, None), + ), patch("ucode.cli.launch_agent"), ] From 9d195ec322f91c2e11458bd9b7686755de413570 Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Wed, 5 Aug 2026 16:14:17 +0000 Subject: [PATCH 3/3] gate managed config behind env var --- src/ucode/cli.py | 42 ++++++++++--- src/ucode/managed_config.py | 109 ++++++++++++++++++++++++++------ src/ucode/managed_resolve.py | 10 +++ tests/test_managed_config.py | 113 ++++++++++++++++++++++++++++++---- tests/test_managed_resolve.py | 25 ++++++++ 5 files changed, 257 insertions(+), 42 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3bc14c6..57c081e 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -55,8 +55,8 @@ resolve_pat_token, run_databricks_login, ) -from ucode.managed_config import managed_launch_state -from ucode.managed_resolve import managed_provider_service +from ucode.managed_config import managed_agent_config_enabled, managed_launch_state +from ucode.managed_resolve import managed_default_model, managed_provider_service from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -1179,11 +1179,6 @@ def _launch_tool( # back to whatever `ucode configure` saved for this tool. provider = provider or get_provider_service(state, tool) routing_agent = _ROUTING_AGENTS.get(tool) - if routing_agent is not None and enable_smart_routing_flag and provider: - raise RuntimeError( - f"{TOOL_SPECS[tool]['display']} smart routing cannot be enabled with " - "--provider. Launch without a Model Provider Service and try again." - ) # Re-fetch model lists on every launch so newly-added Databricks # endpoints show up without a manual `ucode configure` (and so that # tools like pi which read multiple model bundles never run on @@ -1199,7 +1194,16 @@ def _launch_tool( # An admin-published managed config wins over the developer's own settings. Resolved before # the provider and model are settled below, so each is decided once against the values that # will actually be written — the two state files are never merged on disk. - state, managed = managed_launch_state(state, tool) + managed = None + if managed_agent_config_enabled(): + # The spinner covers the read; the outcome is printed after it so the line survives + # (a spinner erases itself, leaving nothing to explain which settings won). + with spinner("Checking for a managed coding agent config..."): + state, managed = managed_launch_state(state, tool, skip_preflight=skip_preflight) + if managed is not None: + print_success("Applied your workspace's managed coding agent config") + else: + print_note("No managed coding agent config found; using your own settings") if managed is not None: managed_provider = managed_provider_service(managed, tool) if explicit_provider and managed_provider and managed_provider != explicit_provider: @@ -1213,6 +1217,13 @@ def _launch_tool( ) if managed_provider: provider = managed_provider + # Checked after the managed config settles `provider`: an admin-set provider must trip this + # guard too, or routing would be persisted as on while a provider is active. + if routing_agent is not None and enable_smart_routing_flag and provider: + raise RuntimeError( + f"{TOOL_SPECS[tool]['display']} smart routing cannot be enabled with " + "--provider. Launch without a Model Provider Service and try again." + ) # Validate the provider service before launching — it must exist, be a # provider type this tool can route to (e.g. claude can't use an OpenAI # or Foundry service), and, for Bedrock, expose Claude models to pin. @@ -1240,7 +1251,12 @@ def _launch_tool( # the workspace has no matching Databricks models. resolved_model = None else: - state, resolved_model = resolve_launch_model(tool, state, None) + # A managed default_model is the model the admin wants sessions to start on, so it goes + # in as the explicit model rather than being applied afterwards: for codex the proto has + # no model list at all, so passing it here is the only way a launch succeeds when the + # workspace's own discovery turned up nothing. + managed_model = managed_default_model(managed, tool) if managed is not None else None + state, resolved_model = resolve_launch_model(tool, state, managed_model) if routing_agent is not None and routing_agent.smart_routing_enabled(state): display = TOOL_SPECS[tool]["display"] with spinner(f"Selecting a {display} model with smart routing..."): @@ -1257,6 +1273,14 @@ def _launch_tool( print_warning( f"Smart routing was unavailable ({routing_error}); using {resolved_model}." ) + # The admin's model outranks a smart-routing pick too. Claude only launches on it when + # pinned as ANTHROPIC_MODEL (route_root_model); other agents take `resolved_model`, + # which already holds it from resolve_launch_model above. + if managed_model: + if tool == "claude": + route_root_model = managed_model + else: + resolved_model = managed_model state = configure_tool( tool, state, diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 9e49f63..99a4034 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -27,6 +27,10 @@ MANAGED_STATE_PATH = config_io.APP_DIR / "managed-state.json" +# Opt-in switch while the feature is in bug bash: unset means launches ignore managed configs +# entirely and behave exactly as they did before. +MANAGED_CONFIG_ENV_VAR = "ENABLE_MANAGED_AGENT_CONFIG" + # Shown to a developer when their workspace has no admin-defined managed config yet — the normal # case, not an error. Kept here so the CLI (which surfaces it) uses one consistent message. NO_MANAGED_CONFIG_MESSAGE = "No coding-agent config has been set up by your workspace admin yet." @@ -241,12 +245,15 @@ def get_managed_config(workspace: str, token: str) -> tuple[dict | None, str | N Returns ``(config, reason)``: - ``(config, None)`` — the normalized manifest for the workspace's single config; - - ``(None, None)`` — no managed config is defined for the workspace (not an error); - - ``(None, reason)`` — the read failed; ``reason`` says why. + - ``(None, None)`` — the workspace definitively has no managed config (not an error); + - ``(None, reason)`` — the read didn't settle the question; ``reason`` says why. - "No config defined" arrives two ways depending on the backend: an empty listing (HTTP 200 - with no configs) or a NOT_FOUND (HTTP 404). Both are the normal, non-error case for a workspace - whose admin hasn't set one up, so both collapse to ``(None, None)``. + The distinction matters to callers that cache: only ``(None, None)`` is authoritative enough to + clear a previously stored config. "No config defined" arrives two ways depending on the backend + — an empty listing (HTTP 200 with no configs) or a NOT_FOUND — and both collapse to + ``(None, None)``. Anything else, including a PERMISSION_DENIED, leaves the question unanswered + and is surfaced as a failure: an admin may have published a config the developer can't read, + which they need to know about rather than silently launch without. v0 stores at most one config per workspace, so the first entry is the workspace's config. """ @@ -262,7 +269,7 @@ def get_managed_config(workspace: str, token: str) -> tuple[dict | None, str | N def _is_not_found(reason: str) -> bool: - """True when a read failure reason indicates the config simply doesn't exist yet. + """True when a read failure reason means the workspace definitively has no managed config. ``_http_get_json`` formats failures as ``HTTP [: ]``; a NOT_FOUND surfaces as an ``HTTP 404`` there (and the API's error body carries ``NOT_FOUND``).""" @@ -270,6 +277,17 @@ def _is_not_found(reason: str) -> bool: return "http 404" in lowered or "not_found" in lowered +def _is_permission_denied(reason: str) -> bool: + """True when the read was refused rather than answering whether a config exists. + + The read is meant to be available to any workspace user, so a refusal means the workspace's + managed config isn't readable by this developer — worth telling them about, since an admin may + have published a config that silently isn't reaching them. It settles nothing about whether one + exists, so a cached config is left in place rather than cleared.""" + lowered = reason.lower() + return "http 403" in lowered or "permission_denied" in lowered + + def save_managed_state(workspace: str, config: dict) -> None: """Persist the normalized managed config to ``~/.ucode/managed-state.json`` at mode 0600. @@ -336,7 +354,9 @@ def refresh_managed_config(state: dict) -> dict | None: return _persisted_fallback(workspace, str(exc)) managed, reason = get_managed_config(workspace, token) if reason is not None: - return _persisted_fallback(workspace, reason) + # A refused read leaves the cached config alone: it says nothing about whether the admin's + # config still exists, unlike a successful "no config" answer below. + return _persisted_fallback(workspace, reason, refused=_is_permission_denied(reason)) if managed is None: # Record that this workspace has no config, rather than leaving an earlier one on disk: # the file doubles as the fallback above, so a removed policy would otherwise come back @@ -347,35 +367,84 @@ def refresh_managed_config(state: dict) -> dict | None: return managed -def _persisted_fallback(workspace: str, reason: str) -> dict | None: - """Return the last persisted config for ``workspace`` after a failed fetch, warning either way. +def _persisted_fallback(workspace: str, reason: str, *, refused: bool = False) -> dict | None: + """Return the last persisted config for ``workspace`` after a failed fetch. - Distinguishes the two outcomes in the warning: continuing on a possibly-stale admin config is - materially different from continuing on the developer's own settings. + Warns only when there is a config to fall back on, because then the launch proceeds on an admin + policy that may be out of date. With nothing persisted there is no managed config in play at + all, so staying quiet keeps someone with (say) an expired session from being told about a + feature they don't use — including when the read was ``refused``, since a refusal is no evidence + that a config exists. """ # An empty persisted config means the last successful read found none, so there is no admin # policy to fall back to — treat it the same as having no file at all. persisted = load_managed_state(workspace) - if persisted: + if not persisted: + return None + summary = _summarize_read_failure(reason) + if refused: print_warning( - f"Could not read your workspace's managed config ({reason}); " + f"Your workspace's managed config is not readable by you ({summary}); using the last " + "one saved for this workspace. Ask an admin to grant access." + ) + else: + print_warning( + f"Could not read your workspace's managed config ({summary}); " "using the last one saved for this workspace." ) - return persisted - print_warning( - f"Could not read your workspace's managed config ({reason}); using your local settings." - ) - return None + return persisted -def managed_launch_state(state: dict, tool: str) -> tuple[dict, dict | None]: +def _summarize_read_failure(reason: str) -> str: + """Condense a read failure into one short line fit for a terminal warning. + + ``_http_get_json`` appends the raw response body, which for a gateway error is a multi-line JSON + blob (error_code, message, request_id, trace ids). Surface just the status and the API's own + message; the full text is still available under ``UCODE_DEBUG=1``. + """ + status, _, body = reason.partition(": ") + body = body.strip() + if body.startswith("{"): + try: + parsed = json.loads(body) + except json.JSONDecodeError: + parsed = None + if isinstance(parsed, dict): + message = _str(parsed.get("message")) or _str(parsed.get("error_code")) + if message: + return f"{status.strip()}: {message}" + return status.strip() + condensed = " ".join(reason.split()) + return condensed if len(condensed) <= 160 else condensed[:157] + "..." + + +def managed_agent_config_enabled() -> bool: + """True when managed coding-agent configs are switched on for this run. + + Opt-in while the feature is being bug-bashed: without the env var set, launches behave exactly + as they did before and never read the workspace's config.""" + return os.environ.get(MANAGED_CONFIG_ENV_VAR, "").strip().lower() in ("1", "true", "yes") + + +def managed_launch_state( + state: dict, tool: str, *, skip_preflight: bool = False +) -> tuple[dict, dict | None]: """Return ``(state, managed)`` for launching ``tool`` under any managed config. The returned state has the manifest's models and provider layered over the developer's own — managed wins per key — so the settings file written from it reflects the admin's choices. When the workspace has no managed config the state is handed back untouched. + + ``skip_preflight`` mirrors the launch flag: managed/headless launchers pass it to avoid + per-launch network calls, so the config is read from the last persisted copy instead of being + re-fetched (which means it can be arbitrarily stale until a normal launch refreshes it). """ - managed = refresh_managed_config(state) + if not managed_agent_config_enabled(): + return state, None + if skip_preflight: + managed = load_managed_state(state.get("workspace")) or None + else: + managed = refresh_managed_config(state) if managed is None: return state, None return resolve_state(managed, state, tool), managed diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index dec85cf..a40c419 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -86,6 +86,16 @@ def managed_provider_service(managed: dict, tool: str) -> str | None: return _str(_agent_model_config(managed, tool).get("model_provider_service")) +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. + + Distinct from the family slots :func:`effective_agent_models` resolves: those set what each + family shortcut maps to, while this is the model the session actually starts on. The launch path + pins it explicitly, so the admin's choice holds even for agents that would otherwise pick their + own default.""" + return _str(_agent_model_config(managed, tool).get("default_model")) + + def resolve_state(managed: dict, state: dict, tool: str) -> dict: """Return a copy of ``state`` with ``tool``'s managed values layered on top. diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 346b3ba..2adbafd 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -300,13 +300,15 @@ def test_read_failure_falls_back_to_the_persisted_config(self, monkeypatch): assert "HTTP 500" in warnings[0] assert "last one saved" in warnings[0] - def test_read_failure_without_persisted_config_uses_local_settings(self, monkeypatch): - warnings: list[str] = [] + def test_read_failure_without_persisted_config_is_silent(self, monkeypatch): + # Nothing persisted means no managed config is in play, so an expired session shouldn't + # produce a warning about a feature this developer doesn't use. monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) - monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + monkeypatch.setattr( + mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") + ) assert refresh_managed_config(_state()) is None - assert "your local settings" in warnings[0] def test_auth_failure_falls_back_to_the_persisted_config(self, monkeypatch): warnings: list[str] = [] @@ -320,17 +322,41 @@ def boom(ws, profile): assert refresh_managed_config(_state()) == MANAGED assert "no token" in warnings[0] - def test_auth_failure_without_persisted_config_uses_local_settings(self, monkeypatch): - warnings: list[str] = [] - + def test_auth_failure_without_persisted_config_is_silent(self, monkeypatch): def boom(ws, profile): raise RuntimeError("no token") monkeypatch.setattr(mc_mod, "get_databricks_token", boom) monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) - monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + monkeypatch.setattr( + mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") + ) assert refresh_managed_config(_state()) is None - assert "your local settings" in warnings[0] + + def test_permission_denied_without_cache_is_silent(self, monkeypatch): + # A refusal is no evidence a config exists, so with nothing cached there is no managed + # config in play and warning would be a false positive. + denied = 'HTTP 403 Forbidden: {"error_code":"PERMISSION_DENIED"}' + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, denied)) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + monkeypatch.setattr( + mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") + ) + assert refresh_managed_config(_state()) is None + + def test_permission_denied_warns_and_keeps_the_cached_config(self, monkeypatch): + # A refused read is worth surfacing: an admin may have published a config that isn't + # reaching this developer. It says nothing about whether one exists, so the cache stands. + warnings: list[str] = [] + denied = 'HTTP 403 Forbidden: {"error_code":"PERMISSION_DENIED"}' + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, denied)) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + monkeypatch.setattr( + mc_mod, "save_managed_state", lambda ws, cfg: pytest.fail("must not clear the cache") + ) + assert refresh_managed_config(_state()) == MANAGED + assert "not readable by you" in warnings[0] def test_no_config_on_the_server_does_not_use_a_stale_persisted_file(self, monkeypatch): # A successful read saying "no config" means the admin removed it — that's authoritative, @@ -353,14 +379,14 @@ def test_no_config_on_the_server_clears_the_persisted_one(self, monkeypatch): assert saved == [(WORKSPACE, {})] def test_empty_persisted_config_is_not_treated_as_a_fallback(self, monkeypatch): - # The empty marker means "no admin policy", so a later failed read must fall through to the + # The empty marker means "no admin policy", so a later failed read falls through to the # developer's own settings rather than reporting a managed config. - warnings: list[str] = [] monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, "HTTP 500")) monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: {}) - monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + monkeypatch.setattr( + mc_mod, "print_warning", lambda msg: pytest.fail(f"should not warn: {msg}") + ) assert refresh_managed_config(_state()) is None - assert "your local settings" in warnings[0] def test_no_workspace_is_a_noop(self, monkeypatch): monkeypatch.setattr( @@ -374,6 +400,7 @@ class TestManagedLaunchState: def _stub_token(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_databricks_token", lambda ws, profile: "tok") monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setenv(mc_mod.MANAGED_CONFIG_ENV_VAR, "1") def test_layers_managed_models_when_a_config_exists(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (MANAGED, None)) @@ -390,3 +417,63 @@ def test_state_untouched_when_no_managed_config(self, monkeypatch): resolved, managed = managed_launch_state(state, "claude") assert managed is None assert resolved is state + + @pytest.mark.parametrize("env_value", [None, "", "0", "off", "no"]) + def test_disabled_does_nothing_at_all(self, monkeypatch, env_value): + """While the feature is opt-in, a disabled launch must behave exactly as it did before. + + Every side effect the managed path can have is trip-wired, so this fails if any future + change reaches the network, the cache, or the developer's state without the env var set. + """ + if env_value is None: + monkeypatch.delenv(mc_mod.MANAGED_CONFIG_ENV_VAR, raising=False) + else: + monkeypatch.setenv(mc_mod.MANAGED_CONFIG_ENV_VAR, env_value) + for name in ( + "get_databricks_token", + "fetch_managed_coding_agent_configs", + "get_managed_config", + "load_managed_state", + "save_managed_state", + "resolve_state", + "print_warning", + ): + monkeypatch.setattr( + mc_mod, + name, + lambda *a, called=name, **k: pytest.fail(f"{called} must not run when disabled"), + ) + + assert mc_mod.managed_agent_config_enabled() is False + state = _state(claude_models={"opus": "local-opus"}) + resolved, managed = managed_launch_state(state, "claude") + assert managed is None + # Same object back, so nothing downstream can see a layered value. + assert resolved is state + assert state["claude_models"] == {"opus": "local-opus"} + + @pytest.mark.parametrize("env_value", ["1", "true", "TRUE", "yes"]) + def test_enabled_values(self, monkeypatch, env_value): + monkeypatch.setenv(mc_mod.MANAGED_CONFIG_ENV_VAR, env_value) + assert mc_mod.managed_agent_config_enabled() is True + + def test_skip_preflight_reads_the_cache_without_fetching(self, monkeypatch): + # Headless launchers pass --skip-preflight to avoid per-launch network calls, so the config + # comes from the last persisted copy rather than a fresh read. + monkeypatch.setattr( + mc_mod, "get_managed_config", lambda ws, tok: pytest.fail("should not fetch") + ) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) + resolved, managed = managed_launch_state(_state(), "claude", skip_preflight=True) + assert managed == MANAGED + assert resolved["claude_models"]["opus"] == "system.ai.claude-opus-5" + + def test_skip_preflight_with_no_cache_is_a_noop(self, monkeypatch): + monkeypatch.setattr( + mc_mod, "get_managed_config", lambda ws, tok: pytest.fail("should not fetch") + ) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + state = _state() + resolved, managed = managed_launch_state(state, "claude", skip_preflight=True) + assert managed is None + assert resolved is state diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index e328d96..810556a 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -11,6 +11,7 @@ import ucode.state as state_mod from ucode.managed_resolve import ( effective_agent_models, + managed_default_model, managed_provider_service, resolve_state, ) @@ -282,3 +283,27 @@ def test_other_agents_state_is_also_preserved( full = json.loads((tmp_path / "state.json").read_text()) assert full["workspaces"][WORKSPACE][models_key] == ["mine"] assert resolved_state[models_key] == managed_models + + +class TestManagedDefaultModel: + """The model a launch starts on, which is separate from the family slots.""" + + def test_returns_the_manifest_default_model(self): + assert managed_default_model(MANAGED, "claude") == "system.ai.claude-opus-5" + + def test_none_when_the_manifest_names_no_default(self): + managed = {"enabled_agents": {"claude": {"model_config": {"models": {}}}}} + assert managed_default_model(managed, "claude") is None + + def test_none_for_agent_not_in_manifest(self): + assert managed_default_model({}, "codex") is None + + def test_survives_a_config_with_no_model_list(self): + # CodexModelConfig has no `models` field at all, so default_model is the only model an + # admin can set — it has to be usable on its own or a codex launch can't honor the config. + managed = {"enabled_agents": {"codex": {"model_config": {"default_model": "admin-codex"}}}} + state = {"workspace": WORKSPACE, "managed_configs": {"codex": {"keys": []}}} + assert managed_default_model(managed, "codex") == "admin-codex" + # Nothing lands in the model list, so the launch path must pass the default model into + # resolve_launch_model rather than relying on state having one. + assert resolve_state(managed, state, "codex").get("codex_models") is None