Skip to content
Merged
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
35 changes: 23 additions & 12 deletions src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -582,6 +582,7 @@ def _enforce_model_default_hierarchy(
ucode_defaults: dict[str, str],
ucode_last_written_defaults: dict[str, str],
enforced_models: list[str] | None,
add_1m_suffix: bool = True,
) -> str | None:
"""Resolve one Claude family's managed-file default model.

Expand All @@ -598,7 +599,7 @@ def _enforce_model_default_hierarchy(
selected = ucode_defaults.get(family)
if selected is None:
return None
if family in ("opus", "sonnet"):
if add_1m_suffix and family in ("opus", "sonnet"):
selected = _maybe_add_1m_suffix(selected)
if enforced_models is not None and selected.split("[", 1)[0] not in enforced_models:
return None
Expand Down Expand Up @@ -919,6 +920,7 @@ def write_tool_config(
otel_tracing=bool(state.get("claude_otel_tracing")),
picker_catalog=picker_catalog,
)
source_scoped_defaults = bool((provider or parent_schema) and coding_agent_config_defaults)
# Native discovery must not inherit UG's prior static allow-list. Keep a replacement picker
# written by this launch, and remove only previously owned picker keys that no longer apply.
stale_picker_keys = [
Expand Down Expand Up @@ -957,15 +959,21 @@ def _compose(
settings_file_env = base_env if isinstance(base_env, dict) else {}
target_env = overlay_for_merge["env"]
configured_defaults = coding_agent_config_defaults or {}
settings_file_existing_defaults = {
family: model
for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items()
if isinstance((model := settings_file_env.get(key)), str)
}
managed_overlay = state.get(MANAGED_OVERLAY_KEY, {})
ucode_defaults = (
managed_overlay.get("claude_models") or state.get("claude_models") or {}
)
if source_scoped_defaults:
# The managed map is complete policy for this source: omitted families must not
# inherit targets from local settings or live discovery.
settings_file_existing_defaults = {}
ucode_defaults = {}
else:
settings_file_existing_defaults = {
family: model
for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items()
if isinstance((model := settings_file_env.get(key)), str)
}
managed_overlay = state.get(MANAGED_OVERLAY_KEY, {})
ucode_defaults = (
managed_overlay.get("claude_models") or state.get("claude_models") or {}
)

enforced_models = overlay_for_merge.get("availableModels")
last_applied_env = {}
Expand All @@ -987,6 +995,7 @@ def _compose(
ucode_defaults=ucode_defaults,
ucode_last_written_defaults=ucode_last_written_defaults,
enforced_models=enforced_models,
add_1m_suffix=provider is None,
)
if selected_default_model is None:
target_env.pop(key, None)
Expand Down Expand Up @@ -1046,7 +1055,7 @@ def _compose(
CLAUDE_SETTINGS_PATH,
_compose(
read_json_safe(CLAUDE_SETTINGS_PATH),
enforce_model_default_hierarchy=False,
enforce_model_default_hierarchy=source_scoped_defaults,
managed_settings_snapshots=None,
),
)
Expand All @@ -1055,7 +1064,9 @@ def _compose(
state,
lambda base: _compose(
base,
enforce_model_default_hierarchy=provider is None and parent_schema is None,
enforce_model_default_hierarchy=(
source_scoped_defaults or (provider is None and parent_schema is None)
),
managed_settings_snapshots=managed_snapshots,
),
managed_file_keys,
Expand Down
24 changes: 20 additions & 4 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -2715,17 +2715,33 @@ def _launch_tool(
# managed UC discovery likewise lets the agent select from the parent schema. Skip model
# resolution, which would otherwise fail when global discovery found no models.
resolved_model = None
if provider and tool == "claude" and (model or provider_models):
managed_source_model = (
managed_launch_model(managed or {}, recommendation, tool)
if tool == "claude" and (managed_provider or managed_parent_schema)
else None
)
if tool == "claude" and managed_parent_schema:
# Native discovery supplies the catalog, but the managed policy still controls
# which model Claude starts on.
route_root_model = managed_source_model
provider_launch_model = model
if tool == "claude" and managed_provider:
# A CLI model still wins, followed by the budget recommendation and the managed
# default. Unmanaged providers retain their existing target-selection behavior.
provider_launch_model = provider_launch_model or managed_source_model
if provider and tool == "claude" and (provider_launch_model or provider_models):
if relayed:
# Resolve against a curated allowlist so the forwarded id is one the gateway
# allows; an allow_all relay declares none, so forward as-is.
relayed_forward_model = (
resolve_provider_launch_model(model, provider_models)
resolve_provider_launch_model(provider_launch_model, provider_models)
if provider_models
else model
else provider_launch_model
)
else:
route_root_model = resolve_provider_launch_model(model, provider_models or {})
route_root_model = resolve_provider_launch_model(
provider_launch_model, provider_models or {}
)
if provider and tool == "gemini":
# Gemini is the exception: the request still names a concrete model
# in the URL, so pin one of the service's targets (--model or default).
Expand Down
2 changes: 1 addition & 1 deletion src/ucode/managed_resolve.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,7 +249,7 @@ def managed_provider_family_models(managed: dict) -> dict[str, str] | None:

config = _agent_model_config(managed, "claude")
slots: dict[str, str] = {}
raw_slots = _as_dict(config.get("models"))
raw_slots = _as_dict(config.get("default_models_by_model_family"))
for slot, family in _CLAUDE_FAMILY_SLOTS.items():
model = _str(raw_slots.get(slot))
if model:
Expand Down
1 change: 1 addition & 0 deletions tests/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`.
| `test_case_06_*`, `test_case_08_*` | Pass a provider or model-location override to managed Codex after configure and from fresh state | ug rejects the override before Codex starts and preserves agent-owned state |
| `test_case_10_*`, `test_case_12_*` | Disable discovery and pass a provider or model-location override to managed Codex | ug still rejects both configured and fresh launches |
| `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_mps_defaults_accompany_discovery`, `test_managed_fixture_claude_parent_schema_defaults_accompany_discovery` | Launch Claude from injected managed defaults with MPS and Unity Catalog discovery | Both generated settings files retain every admin-authored default alongside the source header; only UC Opus/Sonnet family ids gain `[1m]` |
| `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 |
Expand Down
3 changes: 2 additions & 1 deletion tests/integration/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,8 @@ cases** use `UCODE_MANAGED_CONFIG_STUB`. Twenty-four explicit configured/fresh C
discovery and source-override journeys fetch the published config once per agent, replace that
agent's static source with its dedicated MPS, and reuse the result. Twelve existing collected cases
cover focused model, MCP, skills, and lifecycle shapes, including per-agent model reconciliation
and managed skill cleanup.
and managed skill cleanup. The two Claude default-model cases launch with injected MPS and Unity
Catalog sources and verify both generated settings files retain all admin-authored family defaults.
See the named coverage and gaps matrix in
[../README.md](../README.md).

Expand Down
111 changes: 110 additions & 1 deletion tests/integration/test_ug_configure_managed_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,15 @@
import json

import pytest
from utils.constants import MANAGED_CLAUDE_PROVIDER_SERVICE
from utils.evidence import FileTask
from utils.managed import (
build_claude_agent_config,
build_codex_agent_config,
build_coding_agent_config,
set_managed_config_stub,
)
from utils.terminal import AgentTerminal
from utils.terminal import AgentTerminal, TerminalProcess

CLAUDE_OPUS = "system.ai.claude-opus-4-8"
# A real ca-central model absent from the live published config: its presence in the picker can
Expand All @@ -24,6 +25,12 @@
LIVE_ONLY = "haiku-4-5" # published live, but not in the injected list below
CODEX_DEFAULT = "system.ai.gpt-5-6-sol"
CODEX_WITHOUT_BUNDLED_METADATA = "system.ai.gpt-99"
MANAGED_CLAUDE_DEFAULT_ENV_KEYS = {
"default_fable_model": "ANTHROPIC_DEFAULT_FABLE_MODEL",
"default_opus_model": "ANTHROPIC_DEFAULT_OPUS_MODEL",
"default_sonnet_model": "ANTHROPIC_DEFAULT_SONNET_MODEL",
"default_haiku_model": "ANTHROPIC_DEFAULT_HAIKU_MODEL",
}

SMART_ROUTING_BANNER = "Using Unity Gateway Smart Router."
CLAUDE_SMART_ROUTING_MODELS = [
Expand All @@ -44,6 +51,108 @@
]


@pytest.mark.managed_fixture
@pytest.mark.claude
def test_managed_fixture_claude_mps_defaults_accompany_discovery(live_session, workspace, tmp_path):
"""Scenario: launch Claude with managed defaults and MPS discovery.

Expected: the installed ug launch writes the MPS header and every admin-authored default to
both Claude settings files without changing the model ids. This settings reconciliation check
does not claim model inference.
"""
session = live_session
defaults = {
"default_model": "anthropic.claude-sonnet-5",
"default_fable_model": "anthropic.claude-fable-5-1",
"default_opus_model": "anthropic.claude-opus-5",
"default_sonnet_model": "anthropic.claude-sonnet-5",
"default_haiku_model": "anthropic.claude-haiku-4-5",
}
config = build_coding_agent_config(
"CODING_AGENT_CLAUDE_CODE",
{
"agent": "CODING_AGENT_CLAUDE_CODE",
"config": {
"models": {"model_provider_service": MANAGED_CLAUDE_PROVIDER_SERVICE},
"default_models": defaults,
},
},
)
set_managed_config_stub(session, tmp_path, config)
result = session.run("configure", "--workspace", workspace, "--skip-upgrade", timeout=240)
assert "Select coding agents to configure:" not in result.stdout, result.stdout

command = [str(session.binary), "claude", "--", "--version"]
with TerminalProcess(session, "claude", command, "managed-defaults-mps") as terminal:
terminal.finish(timeout=240)

private_settings = json.loads((session.home / ".claude" / "ucode-settings.json").read_text())
os_managed_settings = json.loads(
session.run("/etc/claude-code/managed-settings.json", binary="cat", timeout=30).stdout
)
for settings in (private_settings, os_managed_settings):
env = settings.get("env") or {}
expected_header = f"Databricks-Model-Provider-Service: {MANAGED_CLAUDE_PROVIDER_SERVICE}"
assert expected_header in env.get("ANTHROPIC_CUSTOM_HEADERS", "").splitlines(), settings
assert env.get("ANTHROPIC_MODEL") == defaults["default_model"], settings
for config_key, env_key in MANAGED_CLAUDE_DEFAULT_ENV_KEYS.items():
assert env.get(env_key) == defaults[config_key], settings


@pytest.mark.managed_fixture
@pytest.mark.claude
def test_managed_fixture_claude_parent_schema_defaults_accompany_discovery(
live_session, workspace, tmp_path
):
"""Scenario: launch Claude with managed defaults and Unity Catalog discovery.

Expected: the installed ug launch writes the parent-schema header and every admin-authored
default to both Claude settings files, adding ``[1m]`` only to Opus and Sonnet family defaults.
This settings reconciliation check does not claim model inference.
"""
session = live_session
parent_schema = "system.ai"
defaults = {
"default_model": f"{parent_schema}.claude-sonnet-5",
"default_fable_model": f"{parent_schema}.claude-fable-5-1",
"default_opus_model": f"{parent_schema}.claude-opus-5",
"default_sonnet_model": f"{parent_schema}.claude-sonnet-5",
"default_haiku_model": f"{parent_schema}.claude-haiku-4-5",
}
config = build_coding_agent_config(
"CODING_AGENT_CLAUDE_CODE",
{
"agent": "CODING_AGENT_CLAUDE_CODE",
"config": {
"models": {"unity_catalog_location": parent_schema},
"default_models": defaults,
},
},
)
set_managed_config_stub(session, tmp_path, config)
result = session.run("configure", "--workspace", workspace, "--skip-upgrade", timeout=240)
assert "Select coding agents to configure:" not in result.stdout, result.stdout

command = [str(session.binary), "claude", "--", "--version"]
with TerminalProcess(session, "claude", command, "managed-defaults-parent-schema") as terminal:
terminal.finish(timeout=240)

private_settings = json.loads((session.home / ".claude" / "ucode-settings.json").read_text())
os_managed_settings = json.loads(
session.run("/etc/claude-code/managed-settings.json", binary="cat", timeout=30).stdout
)
for settings in (private_settings, os_managed_settings):
env = settings.get("env") or {}
expected_header = f"Databricks-Model-Service-Parent-Schema: {parent_schema}"
assert expected_header in env.get("ANTHROPIC_CUSTOM_HEADERS", "").splitlines(), settings
assert env.get("ANTHROPIC_MODEL") == defaults["default_model"], settings
for config_key, env_key in MANAGED_CLAUDE_DEFAULT_ENV_KEYS.items():
expected = defaults[config_key]
if config_key in {"default_opus_model", "default_sonnet_model"}:
expected += "[1m]"
assert env.get(env_key) == expected, settings


@pytest.mark.managed_fixture
@pytest.mark.claude
def test_managed_fixture_claude_model_picker_reflects_the_config(live_session, workspace, tmp_path):
Expand Down
Loading