Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
3 changes: 3 additions & 0 deletions src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -707,6 +707,9 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict:
else:
target_env[key] = selected_default_model
merged = deep_merge_dict(base, overlay_for_merge)
if enforce_model_default_hierarchy and "modelPicker" in base:
# Claude's managed modelPicker is administrator-owned; ucode must leave it intact.
merged["modelPicker"] = copy.deepcopy(base["modelPicker"])
overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY]
merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers(
existing_custom_headers, overlay_custom_headers
Expand Down
45 changes: 44 additions & 1 deletion src/ucode/smart_routing/v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from ucode.config_io import APP_DIR, read_json_safe, read_toml_safe, write_json_file
from ucode.constants import LOOPBACK_HOST
from ucode.databricks import (
AnthropicModelCatalog,
build_auth_token_argv,
get_databricks_token,
list_anthropic_model_catalog,
Expand Down Expand Up @@ -58,6 +59,47 @@
_ANTHROPIC_AIGW_MODEL_RE = re.compile(r"^anthropic-aigw-[0-9a-fA-F]{8}-(.+)$")


def _model_picker_catalog() -> AnthropicModelCatalog | None:
"""Read model-picker rows using the managed-settings then ucode-settings waterfall.

A managed picker is authoritative for smart routing: its rows are the models the
administrator exposed, so there is no need to query the gateway catalog first.
"""
try:
from ucode.agents.claude import (
CLAUDE_SETTINGS_PATH,
CLAUDE_USER_SETTINGS_PATH,
_managed_settings_path,
)

# Hierarchy: managed settings, CLI-supplied settings (ucode-settings.json), local user
# settings, based on the modelPicker scope documented at https://code.claude.com/docs/en/settings-reference#modelpicker.
paths = [_managed_settings_path(), CLAUDE_SETTINGS_PATH, CLAUDE_USER_SETTINGS_PATH]
except (ImportError, OSError):
return None
for path in paths:
if path is None or not path.is_file():
continue
settings = read_json_safe(path)
picker_settings = settings.get("modelPicker") if isinstance(settings, dict) else None
picker = picker_settings.get("options") if isinstance(picker_settings, dict) else None
if not isinstance(picker, list):
continue
model_ids: list[str] = []
seen: set[str] = set()
for row in picker:
if not isinstance(row, dict) or not isinstance(row.get("model"), str):
continue
model_id = row["model"].strip()
if not model_id or model_id in seen:
continue
seen.add(model_id)
model_ids.append(model_id)
if model_ids:
return AnthropicModelCatalog(model_ids, {})
return None


def enabled() -> bool:
return os.environ.get(ENV_VAR) == "1"

Expand Down Expand Up @@ -348,7 +390,8 @@ def launch_claude(
os.environ[OAUTH_TOKEN_ENV_VAR] = token
os.environ[GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1"
os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
catalog = list_anthropic_model_catalog(workspace, token)
# modelPicker takes priority over model discovery.
catalog = _model_picker_catalog() or list_anthropic_model_catalog(workspace, token)
if not catalog.model_ids:
raise RuntimeError(
catalog.error_msg or "Anthropic models endpoint returned no Claude models"
Expand Down
19 changes: 19 additions & 0 deletions tests/test_agent_claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -762,6 +762,25 @@ def test_managed_file_preserves_other_keys(self, monkeypatch):
assert written["env"]["ANTHROPIC_BASE_URL"]
assert written["apiKeyHelper"]

def test_managed_file_preserves_model_picker(self, monkeypatch):
private_writes: list = []
managed_writes: list = []
picker = {
"replaceBuiltInOptions": True,
"options": [
{"model": "system.ai.claude-opus-4-8"},
{"model": "system.ai.glm-5-2"},
],
}
existing = {str(FAKE_MANAGED_PATH): {"modelPicker": picker}}
self._patch(monkeypatch, private_writes, managed_writes, existing)
state = {"workspace": WS, "codex_models": []}

claude.write_tool_config(state, "databricks-claude-sonnet-4")

written = json.loads(managed_writes[0][1])
assert written["modelPicker"] == picker

def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch):
private_writes: list = []
managed_writes: list = []
Expand Down
61 changes: 61 additions & 0 deletions tests/test_claude_smart_routing_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,67 @@
from ucode.smart_routing import claude_hooks, claude_pty, routing, v2


class TestManagedModelPicker:
def test_reads_model_ids_from_managed_picker(self, tmp_path, monkeypatch):
path = tmp_path / "managed-settings.json"
path.write_text(
json.dumps(
{
"modelPicker": {
"options": [
{"model": "system.ai.claude-opus-4-8", "label": "Opus"},
{"model": "system.ai.claude-sonnet-5", "label": "Sonnet"},
]
}
}
)
)
monkeypatch.setattr(claude, "_managed_settings_path", lambda: path)

catalog = v2._model_picker_catalog()

assert catalog is not None
assert catalog.model_ids == ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"]
assert catalog.model_id_to_display_name == {}

def test_ignores_empty_or_missing_picker(self, tmp_path, monkeypatch):
path = tmp_path / "managed-settings.json"
path.write_text(json.dumps({"env": {}}))
monkeypatch.setattr(claude, "_managed_settings_path", lambda: path)
monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "ucode-settings.json")

assert v2._model_picker_catalog() is None

def test_falls_back_to_ucode_settings_picker(self, tmp_path, monkeypatch):
managed = tmp_path / "managed-settings.json"
managed.write_text(json.dumps({"env": {}}))
ucode_settings = tmp_path / "ucode-settings.json"
ucode_settings.write_text(
json.dumps({"modelPicker": {"options": [{"model": "system.ai.claude-opus-5"}]}})
)
monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed)
monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings)

catalog = v2._model_picker_catalog()

assert catalog is not None
assert catalog.model_ids == ["system.ai.claude-opus-5"]

def test_falls_back_to_user_settings_picker(self, tmp_path, monkeypatch):
user_settings = tmp_path / "settings.json"
user_settings.write_text(
json.dumps({"modelPicker": {"options": [{"model": "system.ai.claude-sonnet-5"}]}})
)
monkeypatch.setattr(claude, "_managed_settings_path", lambda: tmp_path / "missing-managed")
monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "missing-ucode")
monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user_settings)

catalog = v2._model_picker_catalog()

assert catalog is not None
assert catalog.model_ids == ["system.ai.claude-sonnet-5"]


class TestDirectModelCommand:
@pytest.mark.parametrize(
"name",
Expand Down
Loading