Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,27 @@ def _managed_pinned_model() -> tuple[Path, str] | None:
return (path, str(env["ANTHROPIC_MODEL"]))


def managed_settings_model_overrides() -> Path | None:
"""Path to enterprise managed settings when they pin a model ucode selects with, else None.

The enterprise scope outranks the ``--settings`` file ucode passes, so a model set there wins
over the one an admin published in the workspace's managed config — and unlike the user and
project scopes it can't be excluded with ``--setting-sources``. Callers surface this as a warning
so a developer whose models don't match their admin's config knows where to look.

Only the keys ucode actually writes count. The ``_NAME`` companions in
:data:`CLAUDE_MANAGED_MODEL_ENV_KEYS` are picker labels that select nothing, so an enterprise
value there can't override anything and warning about it would be noise."""
path = _managed_settings_path()
if path is None or not path.is_file():
return None
env = read_json_safe(path).get("env")
if not isinstance(env, dict):
return None
selecting_keys = (key for key in CLAUDE_MANAGED_MODEL_ENV_KEYS if not key.endswith("_NAME"))
return path if any(env.get(key) for key in selecting_keys) else None


def relayed_proxy_base_url(state: dict) -> str:
"""Loopback base URL for the relayed refresh proxy, allocating a free port
on first call and caching it in state so config and launch agree."""
Expand Down
28 changes: 25 additions & 3 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,16 @@
resolve_pat_token,
run_databricks_login,
)
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.managed_config import (
load_managed_state,
managed_agent_config_enabled,
managed_launch_state,
)
from ucode.managed_resolve import (
managed_default_model,
managed_provider_service,
managed_supplies_models,
)
from ucode.mcp import (
MCP_CLIENTS,
SKILLS_MCP_KIND,
Expand Down Expand Up @@ -1179,6 +1187,10 @@ 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)
# Discovery exists to find models and isn't needed for managed config that already names them.
managed_models_known = managed_agent_config_enabled() and managed_supplies_models(
load_managed_state(state.get("workspace")), tool
)
# 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
Expand All @@ -1188,7 +1200,7 @@ def _launch_tool(
state["workspace"],
profile=state.get("profile"),
tools=[tool],
skip_model_discovery=bool(provider),
skip_model_discovery=bool(provider) or managed_models_known,
skip_preflight=skip_preflight,
)
# An admin-published managed config wins over the developer's own settings. Resolved before
Expand All @@ -1202,6 +1214,16 @@ def _launch_tool(
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")
# The enterprise scope outranks the --settings file ucode writes, so a model pinned
# there quietly beats the admin's — point at the file rather than let the mismatch
# look like a ucode bug.
if tool == "claude":
overrides = claude_agent.managed_settings_model_overrides()
if overrides is not None:
print_warning(
f"Default models are set in your enterprise managed settings at "
f"{overrides}, which may override your admin's managed config."
)
else:
print_note("No managed coding agent config found; using your own settings")
if managed is not None:
Expand Down
45 changes: 36 additions & 9 deletions src/ucode/managed_resolve.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,19 +60,28 @@ def _agent_model_config(managed: dict, tool: str) -> dict[str, object]:
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.
Once the manifest says anything about ``tool``'s models it is the whole allowlist: the
developer's discovered models drop out entirely, so a launch can only reach models the admin
named. Each claude family resolves to its own slot, or to ``default_model`` when that slot is
unset; with neither set the family is left out, so ucode writes no
``ANTHROPIC_DEFAULT_<FAMILY>_MODEL`` and the agent uses its own default rather than a model the
admin never sanctioned. Every other agent stores a flat list, which has no per-key identity —
there the manifest's list replaces the local one outright. Only when the manifest names nothing
for ``tool`` does the developer's own list stand.
"""
manifest_models = _agent_model_config(managed, tool).get("models")
model_config = _agent_model_config(managed, tool)
manifest_models = model_config.get("models")
default_model = _str(model_config.get("default_model"))
if tool == "claude":
local = dict(_as_dict(state.get("claude_models")))
slots: dict[str, str] = {}
for slot, family in _CLAUDE_FAMILY_SLOTS.items():
model = _str(_as_dict(manifest_models).get(slot))
model = _str(_as_dict(manifest_models).get(slot)) or default_model
Comment thread
asujithan marked this conversation as resolved.
Outdated
if model:
local[family] = model
return local or None
slots[family] = model
if slots:
return slots
local = _as_dict(state.get("claude_models"))
return dict(local) if local else None
if isinstance(manifest_models, list):
models = [m for m in (_str(item) for item in manifest_models) if m]
if models:
Expand All @@ -81,6 +90,24 @@ def effective_agent_models(managed: dict, state: dict, tool: str) -> dict | list
return local_list if local_list else None


def managed_supplies_models(managed: dict | None, tool: str) -> bool:
"""True when the managed config already says which models ``tool`` should use.

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``.
"""
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


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"))
Expand Down
41 changes: 41 additions & 0 deletions tests/test_agent_claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -864,3 +864,44 @@ def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch):
assert state.get(claude.SMART_ROUTING_STATE_KEY) is None
assert list(doc["hooks"]) == ["PreToolUse"]
assert doc["hooks"]["PreToolUse"][0]["hooks"][0]["command"] == "user-policy"


class TestManagedSettingsModelOverrides:
"""Enterprise managed settings outrank ucode's --settings, so a model pinned there beats the
one an admin published — worth pointing a developer at the file."""

@staticmethod
def _write(monkeypatch, tmp_path, payload):
path = tmp_path / "managed-settings.json"
path.write_text(json.dumps(payload), encoding="utf-8")
monkeypatch.setattr(claude, "_managed_settings_path", lambda: path)
return path

@pytest.mark.parametrize(
"key",
["ANTHROPIC_MODEL", "ANTHROPIC_DEFAULT_OPUS_MODEL", "ANTHROPIC_DEFAULT_HAIKU_MODEL"],
)
def test_reports_the_path_when_a_model_is_pinned(self, monkeypatch, tmp_path, key):
path = self._write(monkeypatch, tmp_path, {"env": {key: "system.ai.claude-opus-5"}})
assert claude.managed_settings_model_overrides() == path

def test_none_for_name_companions_that_select_nothing(self, monkeypatch, tmp_path):
# The `_NAME` keys are picker labels, so an enterprise value there overrides no model.
self._write(monkeypatch, tmp_path, {"env": {"ANTHROPIC_DEFAULT_OPUS_MODEL_NAME": "Opus 5"}})
assert claude.managed_settings_model_overrides() is None

def test_none_when_no_model_keys_are_set(self, monkeypatch, tmp_path):
self._write(monkeypatch, tmp_path, {"env": {"SOMETHING_ELSE": "1"}})
assert claude.managed_settings_model_overrides() is None

def test_none_when_env_block_is_absent(self, monkeypatch, tmp_path):
self._write(monkeypatch, tmp_path, {"permissions": {}})
assert claude.managed_settings_model_overrides() is None

def test_none_on_platforms_without_managed_settings(self, monkeypatch):
monkeypatch.setattr(claude, "_managed_settings_path", lambda: None)
assert claude.managed_settings_model_overrides() is None

def test_none_when_the_file_does_not_exist(self, monkeypatch, tmp_path):
monkeypatch.setattr(claude, "_managed_settings_path", lambda: tmp_path / "missing.json")
assert claude.managed_settings_model_overrides() is None
82 changes: 78 additions & 4 deletions tests/test_managed_resolve.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
effective_agent_models,
managed_default_model,
managed_provider_service,
managed_supplies_models,
resolve_state,
)
from ucode.state import MANAGED_OVERLAY_KEY
Expand Down Expand Up @@ -58,28 +59,50 @@ def _state(**overrides) -> dict:
class TestClaudeModels:
def test_proto_slots_map_to_families(self):
# The manifest keeps proto spelling (`default_opus_model`); render_overlay reads `opus`.
# `fable` has no slot in this manifest, so it takes the config's `default_model`.
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",
"fable": "system.ai.claude-opus-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.
def test_family_absent_from_manifest_is_dropped(self):
# The manifest is the whole allowlist: a family the admin didn't pin is left unset rather
# than inheriting the developer's model, so a launch can't reach models the admin didn't
# sanction. Nothing is written for it, so the agent uses its own default.
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"}
assert effective_agent_models(managed, state, "claude") == {"opus": "managed-opus"}

def test_unset_families_fall_back_to_the_configs_default_model(self):
managed = {
"enabled_agents": {
"claude": {
"model_config": {
"default_model": "managed-default",
"models": {"default_opus_model": "managed-opus"},
}
}
}
}
state = _state(claude_models={"sonnet": "local-sonnet"})
assert effective_agent_models(managed, state, "claude") == {
"opus": "managed-opus",
"sonnet": "managed-default",
"haiku": "managed-default",
"fable": "managed-default",
}

def test_no_manifest_models_falls_back_to_local(self):
state = _state(claude_models={"sonnet": "local-sonnet"})
Expand Down Expand Up @@ -307,3 +330,54 @@ def test_survives_a_config_with_no_model_list(self):
# 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


class TestManagedSuppliesModels:
"""Whether the config already says which models an agent uses, so discovery can be skipped."""

def test_true_when_a_family_slot_is_pinned(self):
managed = {
"enabled_agents": {
"claude": {
"model_config": {"models": {"default_opus_model": "system.ai.claude-opus-5"}}
}
}
}
assert managed_supplies_models(managed, "claude") is True

def test_true_for_a_default_model(self):
managed = {"enabled_agents": {"codex": {"model_config": {"default_model": "gpt"}}}}
assert managed_supplies_models(managed, "codex") is True

def test_true_for_a_provider(self):
# A provider routes by header and pins no Databricks model, so discovery is moot.
managed = {
"enabled_agents": {
"claude": {"model_config": {"model_provider_service": "main.default.mps"}}
}
}
assert managed_supplies_models(managed, "claude") is True

def test_true_for_a_flat_model_list(self):
managed = {"enabled_agents": {"opencode": {"model_config": {"models": ["a", "b"]}}}}
assert managed_supplies_models(managed, "opencode") is True

def test_false_when_the_config_names_no_models(self):
# Discovery still has to run, or the launch has nothing to pin.
managed = {"enabled_agents": {"claude": {"use_as_global_settings": True}}}
assert managed_supplies_models(managed, "claude") is False

def test_false_for_an_agent_the_config_does_not_cover(self):
assert managed_supplies_models(MANAGED, "gemini") is False

def test_false_for_no_config_at_all(self):
# First launch has no persisted copy yet, so discovery runs exactly as it always did.
assert managed_supplies_models(None, "claude") is False

def test_false_when_slots_are_present_but_blank(self):
managed = {
"enabled_agents": {
"claude": {"model_config": {"models": {"default_opus_model": " "}}}
}
}
assert managed_supplies_models(managed, "claude") is False
Loading