From be6ebcb08372f7af84454b27136581c335280e71 Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:30:31 +0000 Subject: [PATCH 1/9] Resolve smart-routing presets separately for Claude and Codex --- README.md | 12 +-- src/ucode/agents/__init__.py | 81 +++++++++--------- src/ucode/agents/codex.py | 9 +- src/ucode/cli.py | 50 +++++------ src/ucode/smart_routing/config.py | 108 +++++++++++++++++------- src/ucode/smart_routing/orchestrator.py | 14 +-- src/ucode/smart_routing/session_env.py | 8 +- src/ucode/smart_routing/v2.py | 20 +++-- tests/README.md | 5 ++ tests/integration/README.md | 5 ++ tests/test_agents_init.py | 38 +++++++++ tests/test_cli.py | 51 +++++++++++ tests/test_smart_routing_config.py | 59 ++++++++++++- 13 files changed, 335 insertions(+), 125 deletions(-) diff --git a/README.md b/README.md index 59374ecdf..affff4856 100644 --- a/README.md +++ b/README.md @@ -251,7 +251,9 @@ The generated shell hooks expect Git Bash; PowerShell-only setups are not covere ### Smart Router Orchestrator -Use `SMART_ROUTER_CONFIG_VERSION` at launch to select a smart-routing configuration: +Use `SMART_ROUTER_CONFIG_VERSION` at launch to select a smart-routing configuration. +Each version defines separate Claude Code and Codex flag mappings; the existing versions +use identical mappings for both agents: | Version | Subagent routing | First-prompt routing | Orchestrator | | --- | --- | --- | --- | @@ -282,8 +284,8 @@ For example: SMART_ROUTER_CONFIG_VERSION=subagent_orch_v0 ug claude ``` -The version takes precedence over conflicting legacy flags. Before parsing command options -or running any command callbacks, UG expands it into +The version takes precedence over conflicting legacy flags. Once the agent is selected, +UG expands its mapping into `ENABLE_SMART_ROUTING_V2`, `ENABLE_SMART_ROUTING_SUBAGENT_ONLY`, and `ENABLE_SMART_ROUTER_ORCHESTRATOR` for the launched session. When the version is unset or empty, these legacy flags retain their existing behavior, including @@ -295,10 +297,10 @@ Workspace smart-routing defaults do not rewrite the selected version's flags. Version names require an explicit suffix. Future revisions use new `_v1`, `_v2`, etc. names without changing existing versions. -Version definitions fail validation at module import if any flag in +Version definitions require both Claude and Codex mappings and fail validation at module import if any flag in `SMART_ROUTING_ENV_KEYS` is missing, has a value other than `"0"` or `"1"`, or an unknown flag is present. Register new managed flags in that tuple and -explicitly set them in every version. +explicitly set them for both agents in every version. Use `ug codex` in the same command for Codex. Smart Router Orchestrator assigns bounded work to explorer, researcher, worker, tester, and reviewer roles while the root plans, diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 3091a948f..b8869c5dc 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -27,6 +27,7 @@ from ucode.managed_config import refresh_managed_config from ucode.managed_files import managed_write_batch from ucode.os_compatibility import subprocess_cross_os +from ucode.smart_routing import v2 as smart_routing_v2 from ucode.state import get_provider_service, load_state, save_state from ucode.telemetry import agent_version from ucode.ui import ( @@ -411,45 +412,49 @@ def configure_tool( parent_schema: str | None = None, picker_catalog: AnthropicModelCatalog | None = None, ) -> dict: - result: dict | tuple[dict, str] - if tool == "codex": - result = codex.write_tool_config( - state, model, provider=provider, parent_schema=parent_schema - ) - elif tool == "claude": - # A Model Provider Service or parent schema routes by header and discovers models natively, - # so the usual "model required" guard doesn't apply to either Claude source. - if not model and not provider and not parent_schema: - raise RuntimeError(f"A {tool} model must be selected before configuration.") - result = claude.write_tool_config( - state, - model, - provider=provider, - provider_models=provider_models, - relayed=relayed, - route_root_model=route_root_model, - custom_model=custom_model, - coding_agent_config_defaults=coding_agent_config_defaults, - parent_schema=parent_schema, - picker_catalog=picker_catalog, - ) - else: - # Every tool in this branch needs a model — including gemini under a provider, - # which still pins the service's target model in the URL. - if not model: - raise RuntimeError(f"A {tool} model must be selected before configuration.") - if tool == "gemini": - result = gemini.write_tool_config(state, model, provider=provider) - elif tool == "copilot": - result = copilot.write_tool_config(state, model) - elif tool == "pi": - result = pi.write_tool_config(state, model) + previous = smart_routing_v2.apply_config(agent=tool) + try: + result: dict | tuple[dict, str] + if tool == "codex": + result = codex.write_tool_config( + state, model, provider=provider, parent_schema=parent_schema + ) + elif tool == "claude": + # A Model Provider Service or parent schema routes by header and discovers models natively, + # so the usual "model required" guard doesn't apply to either Claude source. + if not model and not provider and not parent_schema: + raise RuntimeError(f"A {tool} model must be selected before configuration.") + result = claude.write_tool_config( + state, + model, + provider=provider, + provider_models=provider_models, + relayed=relayed, + route_root_model=route_root_model, + custom_model=custom_model, + coding_agent_config_defaults=coding_agent_config_defaults, + parent_schema=parent_schema, + picker_catalog=picker_catalog, + ) else: - result = opencode.write_tool_config(state, model) - # gemini/opencode/copilot/pi return (state, token); codex/claude return state - if isinstance(result, tuple): - return result[0] - return result + # Every tool in this branch needs a model — including gemini under a provider, + # which still pins the service's target model in the URL. + if not model: + raise RuntimeError(f"A {tool} model must be selected before configuration.") + if tool == "gemini": + result = gemini.write_tool_config(state, model, provider=provider) + elif tool == "copilot": + result = copilot.write_tool_config(state, model) + elif tool == "pi": + result = pi.write_tool_config(state, model) + else: + result = opencode.write_tool_config(state, model) + # gemini/opencode/copilot/pi return (state, token); codex/claude return state + if isinstance(result, tuple): + return result[0] + return result + finally: + smart_routing_v2.restore_smart_routing_env(previous) def configured_paths(tool: str, state: dict) -> list[str]: diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 1ff680c14..423b8fe82 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -36,6 +36,7 @@ write_toml_file, ) from ucode.constants import ( + AGENT_CODEX, LOOPBACK_HOST, MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, @@ -206,7 +207,7 @@ def _provider_block( http_headers[MODEL_PROVIDER_SERVICE_HEADER] = provider elif parent_schema: http_headers[MODEL_SERVICE_PARENT_SCHEMA_HEADER] = parent_schema - if smart_routing_v2.smart_routing_enabled(): + if smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): http_headers[SMART_ROUTER_RECIPE_HEADER] = configured_router_name() _apply_managed_headers(http_headers, managed_http_headers) return { @@ -485,7 +486,7 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: prune_key_paths(base, _PROVIDER_HTTP_HEADERS_KEY_PATHS) deep_merge_dict(base, copy.deepcopy(overlay)) # deep_merge can't drop keys, so clear model preferences from an earlier run. - if chosen_model is None and not smart_routing_v2.smart_routing_enabled(): + if chosen_model is None and not smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): for key in ("model", "model_reasoning_effort"): base.pop(key, None) if include_catalog: @@ -814,7 +815,7 @@ def default_model(state: dict) -> str | None: """Return a managed Codex model, or leave selection to Codex.""" if isinstance(state.get("codex_default_model"), str): return state["codex_default_model"] - if smart_routing_v2.smart_routing_enabled(): + if smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): return _smart_routing_config_model(state) return None @@ -842,7 +843,7 @@ def config_precedence_paths() -> tuple[Path, ...]: def clear_model_preferences(state: dict) -> bool: """Remove ucode profile model preferences so Codex selects its default.""" - if smart_routing_v2.smart_routing_enabled(): + if smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): return False if isinstance(state.get("codex_default_model"), str): return False diff --git a/src/ucode/cli.py b/src/ucode/cli.py index d6dbfd15f..ea3a45897 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -46,7 +46,7 @@ from ucode.agents.codex import revert_legacy_shared_config from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import is_dry_run, restore_file, set_dry_run -from ucode.constants import SMART_ROUTING_ENV_KEYS +from ucode.constants import AGENT_CLAUDE, AGENT_CODEX, SMART_ROUTING_ENV_KEYS from ucode.custom_oauth import ( CUSTOM_OAUTH_CLI_ENV_VAR, custom_oauth_cli_enabled, @@ -1387,22 +1387,6 @@ def revert() -> int: class _HelpOrderedGroup(TyperGroup): """Keep top-level help organized across commands and nested Typer apps.""" - def make_context( - self, - info_name: str | None, - args: list[str], - parent: _click.Context | None = None, - **extra: Any, - ) -> _click.Context: - previous = smart_routing_v2.apply_config() - try: - ctx = super().make_context(info_name, args, parent, **extra) - except BaseException: - smart_routing_v2.restore_smart_routing_env(previous) - raise - ctx.call_on_close(lambda: smart_routing_v2.restore_smart_routing_env(previous)) - return ctx - def list_commands(self, ctx: _click.Context) -> list[str]: commands = super().list_commands(ctx) order = {name: index for index, name in enumerate(_HELP_COMMAND_ORDER)} @@ -2296,7 +2280,9 @@ def codex_router_hook_cmd( import json import sys - if not smart_routing_v2.smart_routing_enabled(effective_environment()): + if not smart_routing_v2.smart_routing_enabled( + effective_environment(agent=AGENT_CODEX), agent=AGENT_CODEX + ): return from ucode.smart_routing.codex_routing import ( @@ -2375,7 +2361,9 @@ def claude_router_hook_cmd( import json import sys - if not smart_routing_v2.smart_routing_enabled(effective_environment()): + if not smart_routing_v2.smart_routing_enabled( + effective_environment(agent=AGENT_CLAUDE), agent=AGENT_CLAUDE + ): return from ucode.smart_routing.claude_routing import ( @@ -2492,12 +2480,12 @@ def _auto_configure_tool(tool: str, custom_oauth: CustomOAuthConfig | None = Non @contextmanager -def _smart_routing_v2_flag(enabled: bool | None) -> Iterator[None]: +def _smart_routing_v2_flag(enabled: bool | None, *, agent: str) -> Iterator[None]: """Apply an explicit routing choice without leaking into an embedding process.""" previous = ( - smart_routing_v2.apply_config() + smart_routing_v2.apply_config(agent=agent) if enabled is None - else smart_routing_v2.override_smart_routing(enabled) + else smart_routing_v2.override_smart_routing(enabled, agent=agent) ) try: yield @@ -2820,8 +2808,10 @@ def _launch_tool( parent_schema: str | None = None, custom_oauth: CustomOAuthConfig | None = None, ) -> None: + previous = {} try: tool = normalize_tool(tool_name) + previous = smart_routing_v2.apply_config(agent=tool) if not custom_oauth_cli_enabled(custom_oauth): os.environ.pop(CUSTOM_OAUTH_CLI_ENV_VAR, None) # Before any status print: a stdio-protocol subcommand owns stdout, so @@ -3201,7 +3191,8 @@ def _launch_tool( if managed_smart_routing_enabled and smart_routing_enabled and not smart_routing_v2.smart_routing_enabled() - else None + else None, + agent=tool, ): launch_agent(tool, state, ctx.args, options=launch_options) except RuntimeError as exc: @@ -3210,6 +3201,8 @@ def _launch_tool( except KeyboardInterrupt: print_err("Interrupted.") raise typer.Exit(130) from None + finally: + smart_routing_v2.restore_smart_routing_env(previous) # Launch-only escape hatch for managed/headless launchers (e.g. omnigent) that @@ -3297,10 +3290,9 @@ def default( return set_dry_run(dry_run) try: - with _smart_routing_v2_flag(None): - _launch_managed_default( - ctx, dry_run=dry_run, skip_preflight=skip_preflight, workspace=workspace - ) + _launch_managed_default( + ctx, dry_run=dry_run, skip_preflight=skip_preflight, workspace=workspace + ) except typer.Exit: # `typer.Exit` subclasses RuntimeError, so it has to be re-raised ahead of the handler # below. Otherwise a launch that already reported its own error is followed by @@ -3436,7 +3428,7 @@ def codex_cmd( except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from exc - with _smart_routing_v2_flag(enable_smart_routing_flag): + with _smart_routing_v2_flag(enable_smart_routing_flag, agent=AGENT_CODEX): with _disable_smart_routing_for_subcommand("codex", ctx): _launch_tool( "codex", @@ -3524,7 +3516,7 @@ def claude_cmd( except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from exc - with _smart_routing_v2_flag(enable_smart_routing_flag): + with _smart_routing_v2_flag(enable_smart_routing_flag, agent=AGENT_CLAUDE): with _disable_smart_routing_for_subcommand("claude", ctx): _launch_tool( "claude", diff --git a/src/ucode/smart_routing/config.py b/src/ucode/smart_routing/config.py index 2a7fe9c01..bc371b3e6 100644 --- a/src/ucode/smart_routing/config.py +++ b/src/ucode/smart_routing/config.py @@ -6,6 +6,8 @@ from collections.abc import Mapping, MutableMapping from ucode.constants import ( + AGENT_CLAUDE, + AGENT_CODEX, ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, ENABLE_SMART_ROUTING_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR, @@ -30,68 +32,114 @@ _VERSIONS = { FIRST_PROMPT_AND_SUBAGENT_NO_ORCH_V0: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + AGENT_CLAUDE: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + }, + AGENT_CODEX: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + }, }, SUBAGENT_ONLY_V0: { - ENABLE_SMART_ROUTING_ENV_VAR: "0", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + AGENT_CLAUDE: { + ENABLE_SMART_ROUTING_ENV_VAR: "0", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + }, + AGENT_CODEX: { + ENABLE_SMART_ROUTING_ENV_VAR: "0", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + }, }, SUBAGENT_ONLY_V1: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + AGENT_CLAUDE: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + }, + AGENT_CODEX: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + }, }, SUBAGENT_ORCH_V0: { - ENABLE_SMART_ROUTING_ENV_VAR: "0", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + AGENT_CLAUDE: { + ENABLE_SMART_ROUTING_ENV_VAR: "0", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + }, + AGENT_CODEX: { + ENABLE_SMART_ROUTING_ENV_VAR: "0", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + }, }, SUBAGENT_ORCH_V1: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + AGENT_CLAUDE: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + }, + AGENT_CODEX: { + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", + }, }, } -def _validate_versions(versions: Mapping[str, Mapping[str, str]]) -> None: - """Require every version to explicitly configure the complete managed flag set.""" +def _validate_versions(versions: Mapping[str, Mapping[str, Mapping[str, str]]]) -> None: + """Require both agents and a complete, explicit flag set for each.""" + expected_agents = {AGENT_CLAUDE, AGENT_CODEX} expected = set(SMART_ROUTING_ENV_KEYS) - for version, values in versions.items(): - missing = expected - values.keys() - unexpected = values.keys() - expected - if missing or unexpected: + for version, agents in versions.items(): + if agents.keys() != expected_agents: raise ValueError( f"Invalid smart-routing version {version!r}: " - f"missing env vars {sorted(missing)}; unexpected env vars {sorted(unexpected)}." + f"expected agents {sorted(expected_agents)}, got {sorted(agents)}." ) - for key, value in values.items(): - if value not in ("0", "1"): + for agent, values in agents.items(): + missing = expected - values.keys() + unexpected = values.keys() - expected + if missing or unexpected: raise ValueError( - f"Invalid smart-routing version {version!r}: " - f"{key} must be '0' or '1', got {value!r}." + f"Invalid smart-routing version {version!r} for {agent!r}: " + f"missing env vars {sorted(missing)}; unexpected env vars {sorted(unexpected)}." ) + for key, value in values.items(): + if value not in ("0", "1"): + raise ValueError( + f"Invalid smart-routing version {version!r} for {agent!r}: " + f"{key} must be '0' or '1', got {value!r}." + ) _validate_versions(_VERSIONS) -def resolve_environment(env: Mapping[str, str] | None = None) -> dict[str, str]: +def resolve_environment( + env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE +) -> dict[str, str]: """Expand a version before applying any launch or session-specific overrides.""" resolved = dict(os.environ if env is None else env) version = resolved.pop(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, "").strip() - resolved.update(_VERSIONS.get(version, {})) + resolved.update(_VERSIONS.get(version, {}).get(agent, {})) return resolved -def apply_config(env: MutableMapping[str, str] | None = None) -> dict[str, str | None]: +def apply_config( + env: MutableMapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE +) -> dict[str, str | None]: """Consume the launch selector, returning the values needed to restore its input.""" target = os.environ if env is None else env version = target.get(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, "").strip() - preset = _VERSIONS.get(version) + preset = _VERSIONS.get(version, {}).get(agent) if preset is None: return {} keys = (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) diff --git a/src/ucode/smart_routing/orchestrator.py b/src/ucode/smart_routing/orchestrator.py index a09aecb22..81bfc6f39 100644 --- a/src/ucode/smart_routing/orchestrator.py +++ b/src/ucode/smart_routing/orchestrator.py @@ -31,16 +31,16 @@ ) -def feature_enabled(env: Mapping[str, str] | None = None) -> bool: - source = resolve_environment(env) +def feature_enabled(env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE) -> bool: + source = resolve_environment(env, agent=agent) return source.get(ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR) == "1" -def enabled(env: Mapping[str, str] | None = None) -> bool: +def enabled(env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE) -> bool: from ucode.smart_routing.v2 import smart_routing_enabled source = os.environ if env is None else env - if not feature_enabled(source): + if not feature_enabled(source, agent=agent): return False if source.get("ISAAC_LAUNCH_MODE", "").strip().lower() == "omni": return False @@ -50,7 +50,7 @@ def enabled(env: Mapping[str, str] | None = None) -> bool: return False except (RuntimeError, OSError): return False - return smart_routing_enabled(effective_environment(source)) + return smart_routing_enabled(effective_environment(source, agent=agent), agent=agent) def skill_directory() -> Path: @@ -65,7 +65,7 @@ def add_claude_agents(plugin_dir: Path) -> None: def sync_hooks(doc: dict, *, agent: str) -> None: groups = {} - if feature_enabled(): + if feature_enabled(agent=agent): argv = [sys.executable, "-m", HOOK_MODULE, "--agent", agent] hook = { "type": "command", @@ -90,7 +90,7 @@ def hook_output(payload: object, *, agent: str) -> dict | None: ): return None context = DISABLED_CONTEXT - if enabled(): + if enabled(agent=agent): directory = skill_directory() try: workflow = skills.skill_entrypoint(directory, agent).read_text(encoding="utf-8") diff --git a/src/ucode/smart_routing/session_env.py b/src/ucode/smart_routing/session_env.py index d5f4a467a..f88378bd5 100644 --- a/src/ucode/smart_routing/session_env.py +++ b/src/ucode/smart_routing/session_env.py @@ -10,7 +10,7 @@ from pathlib import Path from ucode.config_io import atomic_write_json -from ucode.constants import SMART_ROUTING_ENV_KEYS +from ucode.constants import AGENT_CLAUDE, SMART_ROUTING_ENV_KEYS from ucode.smart_routing.config import resolve_environment SESSION_ENV_VAR = "UCODE_SESSION_ENV_FILE" @@ -52,9 +52,11 @@ def _read(path: Path) -> dict[str, str]: return _validate(json.loads(path.read_text(encoding="utf-8"))) -def effective_environment(env: Mapping[str, str] | None = None) -> dict[str, str]: +def effective_environment( + env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE +) -> dict[str, str]: """Overlay the latest session controls on the hook process environment.""" - effective = resolve_environment(env) + effective = resolve_environment(env, agent=agent) try: path = session_env_path(effective) except RuntimeError: diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index cc2260612..bb8aa1e58 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -27,6 +27,8 @@ write_text_file, ) from ucode.constants import ( + AGENT_CLAUDE, + AGENT_CODEX, ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, ENABLE_SMART_ROUTING_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR, @@ -87,7 +89,7 @@ class ClaudeRoutingSetupError(RuntimeError): def _prepare_smart_router_session(agent: str) -> Path: skills = [SMART_ROUTER_SKILL] - if orchestrator.feature_enabled(): + if orchestrator.feature_enabled(agent=agent): skills.append(SMART_ROUTER_ORCHESTRATOR_SKILL) for skill in skills: try: @@ -152,9 +154,9 @@ def _model_picker_catalog() -> AnthropicModelCatalog | None: def smart_routing_enabled( - env: MutableMapping[str, str] | None = None, *, default: bool = False + env: MutableMapping[str, str] | None = None, *, default: bool = False, agent: str = AGENT_CLAUDE ) -> bool: - source = resolve_environment(env) + source = resolve_environment(env, agent=agent) values = [ source.get(var) for var in (ENABLE_SMART_ROUTING_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR) ] @@ -165,9 +167,11 @@ def smart_routing_enabled( return default -def first_prompt_routing_enabled(env: MutableMapping[str, str] | None = None) -> bool: +def first_prompt_routing_enabled( + env: MutableMapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE +) -> bool: """Whether the first prompt is routed. Subagent-only wins over the full V2 flag.""" - source = resolve_environment(env) + source = resolve_environment(env, agent=agent) return ( source.get(ENABLE_SMART_ROUTING_ENV_VAR) == "1" and source.get(ENABLE_SUBAGENT_ROUTING_ENV_VAR) != "1" @@ -184,11 +188,13 @@ def enable_smart_routing( def override_smart_routing( enabled: bool, env: MutableMapping[str, str] | None = None, + *, + agent: str = AGENT_CLAUDE, ) -> dict[str, str | None]: """Set an explicit launch-scoped routing choice and return the prior values.""" target = os.environ if env is None else env previous = {var: target.get(var) for var in SMART_ROUTING_ENV_KEYS} - previous.update(apply_config(target)) + previous.update(apply_config(target, agent=agent)) if enabled: target[ENABLE_SMART_ROUTING_ENV_VAR] = "1" else: @@ -682,7 +688,7 @@ def launch_codex( SESSION_PYTHON_ENV_VAR ] config_args = codex_config_args(overlay) - if not first_prompt_routing_enabled(): + if not first_prompt_routing_enabled(agent=AGENT_CODEX): # Subagent-only routing needs neither the app-server nor the interposer: # the hooks ride in the CLI config, so launch the TUI directly. exec_or_spawn([binary, *config_args, *tool_args]) diff --git a/tests/README.md b/tests/README.md index 6421ddd4b..e1c21d992 100644 --- a/tests/README.md +++ b/tests/README.md @@ -35,6 +35,11 @@ CUJ4's Claude preset sessions also use it after verifying delegated file tasks: answer does not prove that a resumed child has stopped running. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. +Agent-specific smart-routing preset component checks cover identical existing mappings, +complete per-agent validation, configure order, bootstrap and managed-default launch +selection, session controls, and restoration after configuration/launch failures. +These checks do not establish live agent coverage. + CUJ3 and CUJ4 also verify Claude's native recovery from the known thinking-display 400: the same payload without display must receive a non-empty 200 for adaptive or enabled thinking. If the next attempt rejects `safeguards`, only its native removal is accepted before the final 200. diff --git a/tests/integration/README.md b/tests/integration/README.md index e81b10ae2..b7a3c1f8d 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -15,6 +15,11 @@ managed-default and explicit-model cases unchanged. One additional test runs the `SMART_ROUTER_CONFIG_VERSION` values, checking both agents' first-prompt routing, orchestrator context in inference input, and completed explicitly requested routed subagents. This is not automatic orchestrator-delegation coverage. No workspace configuration is modified. +Agent-specific smart-routing preset component checks cover identical existing mappings, +complete per-agent validation, configure order, bootstrap and managed-default launch +selection, session controls, and restoration after configuration/launch failures. +These checks do not establish live agent coverage. + CUJ3 and CUJ4 require a non-empty HTTP 200 for inference. Claude 2.1.290's known thinking-display 400 is accepted only when the next inference request with the same system context and session metadata removes `display`, diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index fd234e8b4..bced2e57a 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -7,6 +7,7 @@ import subprocess import sys from contextlib import contextmanager, nullcontext, redirect_stdout +from copy import deepcopy from unittest.mock import MagicMock import pytest @@ -1111,3 +1112,40 @@ def test_appends_os_managed_file_recorded_in_state(self): paths = configured_paths("codex", state) assert str(CODEX_CONFIG_PATH).replace(str(CODEX_CONFIG_PATH.home()), "~", 1) in paths assert "/etc/codex/managed_config.toml" in paths + + +@pytest.mark.parametrize("order", [("claude", "codex"), ("codex", "claude")]) +@pytest.mark.parametrize("fail_write", [False, True]) +def test_configure_dispatch_scopes_agent_presets(monkeypatch, order, fail_write): + from ucode.constants import ( + AGENT_CODEX, + SMART_ROUTER_CONFIG_VERSION_ENV_VAR, + SMART_ROUTING_ENV_KEYS, + ) + from ucode.smart_routing import config + + version = config.SUBAGENT_ORCH_V0 + agents = deepcopy(config._VERSIONS[version]) + agents[AGENT_CODEX] = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") + monkeypatch.setitem(config._VERSIONS, version, agents) + monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) + before = { + key: os.environ.get(key) + for key in (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) + } + for agent in order: + + def write(state, *args, agent=agent, **kwargs): + assert {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} == agents[agent] + assert SMART_ROUTER_CONFIG_VERSION_ENV_VAR not in os.environ + if fail_write: + raise RuntimeError("write failed") + return state + + monkeypatch.setattr(agents_mod._MODULES[agent], "write_tool_config", write) + if fail_write: + with pytest.raises(RuntimeError, match="write failed"): + agents_mod.configure_tool(agent, {}, model="model") + else: + assert agents_mod.configure_tool(agent, {}, model="model") == {} + assert {key: os.environ.get(key) for key in before} == before diff --git a/tests/test_cli.py b/tests/test_cli.py index b95af40d1..c7facdc99 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -11,6 +11,7 @@ import sys import time import tomllib +from copy import deepcopy from importlib import metadata from pathlib import Path from unittest.mock import MagicMock, call, patch @@ -24,8 +25,15 @@ import ucode.cli as cli_mod import ucode.databricks as db_mod from ucode.cli import app +from ucode.constants import ( + AGENT_CLAUDE, + AGENT_CODEX, + SMART_ROUTER_CONFIG_VERSION_ENV_VAR, + SMART_ROUTING_ENV_KEYS, +) from ucode.databricks import GatewayProbe from ucode.managed_config import normalize_managed_config +from ucode.smart_routing import config as routing_config _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -497,6 +505,49 @@ def _launch_policy_patches( } +@pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) +@pytest.mark.parametrize("bare", [False, True]) +@pytest.mark.parametrize("fail_launch", [False, True]) +def test_agent_preset_reaches_bootstrap_and_launch_and_is_restored( + monkeypatch, agent, bare, fail_launch +): + version = routing_config.SUBAGENT_ORCH_V0 + agents = deepcopy(routing_config._VERSIONS[version]) + agents[AGENT_CODEX] = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") + monkeypatch.setitem(routing_config._VERSIONS, version, agents) + monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) + before = { + key: os.environ.get(key) + for key in (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) + } + managed = {"default_agent": agent, "enabled_agents": {agent: {"smart_routing_enabled": True}}} + observed = [] + + def capture_launch(tool, state, args, *, options): + observed.append({key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS}) + assert options.launch_smart_routing == (agent == AGENT_CLAUDE) + if fail_launch: + raise RuntimeError("launch failed") + + with _launch_policy_patches(managed) as calls: + calls["launch"].side_effect = capture_launch + monkeypatch.setattr( + cli_mod, + "ensure_bootstrap_dependencies", + lambda *args, **kwargs: observed.append( + {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} + ), + ) + monkeypatch.setattr(cli_mod, "install_databricks_cli", lambda *args, **kwargs: None) + monkeypatch.setattr( + cli_mod, "refresh_managed_config", lambda *args, **kwargs: (managed, False) + ) + result = runner.invoke(app, [] if bare else [agent]) + assert result.exit_code == (1 if fail_launch else 0), result.output + assert observed == [agents[agent], agents[agent]] + assert {key: os.environ.get(key) for key in before} == before + + class TestSubcommandRouting: @pytest.mark.parametrize("tool", TOOLS) def test_subcommand_calls_correct_tool(self, tool): diff --git a/tests/test_smart_routing_config.py b/tests/test_smart_routing_config.py index c797902a0..bd19685c2 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -2,9 +2,13 @@ from __future__ import annotations +from copy import deepcopy + import pytest from ucode.constants import ( + AGENT_CLAUDE, + AGENT_CODEX, ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR, ENABLE_SMART_ROUTING_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR, @@ -62,6 +66,7 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): assert applied == original +@pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) @pytest.mark.parametrize("ENABLE_SMART_ROUTING_V2", [None, "0", "1"]) @pytest.mark.parametrize("ENABLE_SMART_ROUTING_SUBAGENT_ONLY", [None, "0", "1"]) @pytest.mark.parametrize("ENABLE_SMART_ROUTER_ORCHESTRATOR", [None, "0", "1"]) @@ -77,6 +82,7 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): ], ) def test_smart_routing_config_cartesian_grid( + agent, ENABLE_SMART_ROUTING_V2, ENABLE_SMART_ROUTING_SUBAGENT_ONLY, ENABLE_SMART_ROUTER_ORCHESTRATOR, @@ -101,13 +107,13 @@ def test_smart_routing_config_cartesian_grid( if SMART_ROUTER_CONFIG_VERSION is not None: expected.update(_EXPECTED_PRESETS[SMART_ROUTER_CONFIG_VERSION]) - resolved = config.resolve_environment(source) + resolved = config.resolve_environment(source, agent=agent) assert resolved == expected assert source == original applied = source.copy() - previous = config.apply_config(applied) + previous = config.apply_config(applied, agent=agent) assert applied == expected v2.restore_smart_routing_env(previous, applied) @@ -140,3 +146,52 @@ def test_legacy_preset_equivalence(version, expected_routing): v2.first_prompt_routing_enabled(environment), orchestrator.feature_enabled(environment), ) == expected_routing + + +@pytest.mark.parametrize( + "failure", ["missing_agent", "unknown_agent", "missing_flag", "unknown_flag", "invalid_value"] +) +def test_agent_presets_require_complete_valid_configuration(failure): + agents = deepcopy(config._VERSIONS[config.SUBAGENT_ORCH_V0]) + if failure == "missing_agent": + agents.pop(AGENT_CODEX) + elif failure == "unknown_agent": + agents["other"] = agents[AGENT_CODEX].copy() + elif failure == "missing_flag": + agents[AGENT_CODEX].pop(ENABLE_SMART_ROUTING_ENV_VAR) + elif failure == "unknown_flag": + agents[AGENT_CODEX]["UNKNOWN_FLAG"] = "0" + else: + agents[AGENT_CODEX][ENABLE_SMART_ROUTING_ENV_VAR] = "yes" + with pytest.raises(ValueError, match="Invalid smart-routing version"): + config._validate_versions({"invalid_v0": agents}) + + +def test_existing_agent_presets_are_equal_and_independent(): + for version, expected in _EXPECTED_PRESETS.items(): + agents = config._VERSIONS[version] + assert agents[AGENT_CLAUDE] == agents[AGENT_CODEX] == expected + assert agents[AGENT_CLAUDE] is not agents[AGENT_CODEX] + + +@pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) +def test_agent_resolution_and_session_controls(monkeypatch, agent, tmp_path): + from ucode.smart_routing import session_env + + version = config.SUBAGENT_ORCH_V0 + agents = deepcopy(config._VERSIONS[version]) + agents[AGENT_CODEX] = dict.fromkeys(agents[AGENT_CODEX], "0") + monkeypatch.setitem(config._VERSIONS, version, agents) + source = {SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version} + expected = agents[agent] + assert config.resolve_environment(source, agent=agent) == expected + assert v2.smart_routing_enabled(source, default=True, agent=agent) == (agent == AGENT_CLAUDE) + assert orchestrator.feature_enabled(source, agent=agent) == (agent == AGENT_CLAUDE) + control = tmp_path / "session.json" + control.write_text('{"ENABLE_SMART_ROUTER_ORCHESTRATOR":"0"}') + source[session_env.SESSION_ENV_VAR] = str(control) + assert session_env.effective_environment(source, agent=agent) == { + **expected, + session_env.SESSION_ENV_VAR: str(control), + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + } From 9be35cd314cfa29bda15a3dae12cbd5f301b96c3 Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:39:28 +0000 Subject: [PATCH 2/9] Require explicit routing agents and use shared agent constants --- src/ucode/agents/__init__.py | 43 +++++++-------- src/ucode/agents/claude.py | 3 +- src/ucode/agents/codex.py | 44 ++++++++------- src/ucode/cli.py | 68 ++++++++++++------------ src/ucode/smart_routing/codex_routing.py | 3 +- src/ucode/smart_routing/config.py | 6 +-- src/ucode/smart_routing/orchestrator.py | 10 ++-- src/ucode/smart_routing/session_env.py | 6 +-- src/ucode/smart_routing/v2.py | 28 +++++----- tests/test_claude_smart_routing_v2.py | 13 ++--- tests/test_cli.py | 2 +- tests/test_smart_routing_config.py | 25 ++++++--- 12 files changed, 135 insertions(+), 116 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b8869c5dc..995072cb4 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -16,6 +16,7 @@ from pathlib import Path from ucode.config_io import ToolSpec +from ucode.constants import AGENT_CLAUDE, AGENT_CODEX from ucode.databricks import ( AnthropicModelCatalog, get_databricks_token, @@ -47,8 +48,8 @@ from .args import explicit_model_arg_value as explicit_model_arg_value _MODULES = { - "codex": codex, - "claude": claude, + AGENT_CODEX: codex, + AGENT_CLAUDE: claude, "gemini": gemini, "opencode": opencode, "copilot": copilot, @@ -63,9 +64,9 @@ # rejects it and the model-config paths never see it. The `configure`/MCP flows # handle "cursor" separately as an MCP-only client (see MCP_ONLY_CLIENTS). TOOL_ALIASES = { - "codex": "codex", - "claude": "claude", - "claude-code": "claude", + AGENT_CODEX: AGENT_CODEX, + AGENT_CLAUDE: AGENT_CLAUDE, + "claude-code": AGENT_CLAUDE, "gemini": "gemini", "gemini-cli": "gemini", "opencode": "opencode", @@ -73,18 +74,18 @@ "pi": "pi", } -DEFAULT_TOOL = "codex" +DEFAULT_TOOL = AGENT_CODEX BUNDLE_VERSION = 1 -_MANAGED_SETTINGS_TOOLS = {"claude", "codex"} +_MANAGED_SETTINGS_TOOLS = {AGENT_CLAUDE, AGENT_CODEX} _NATIVE_UPGRADE_COMMANDS = { - "claude": ["claude", "upgrade"], - "codex": ["codex", "update"], + AGENT_CLAUDE: ["claude", "upgrade"], + AGENT_CODEX: ["codex", "update"], } # ucode tool -> `databricks aitools` agent id. gemini/pi aren't supported. AITOOLS_AGENT_TOKENS = { - "claude": "claude-code", - "codex": "codex", + AGENT_CLAUDE: "claude-code", + AGENT_CODEX: "codex", "opencode": "opencode", "copilot": "copilot", } @@ -137,7 +138,7 @@ def _update_installed_tool_binary(tool: str, version: str | None = None) -> bool command = ["npm", "install", "-g", target] print_note(f"Upgrading {spec['display']}...") - if tool == "codex": + if tool == AGENT_CODEX: # Detach potentially incompatible metadata until the next validated refresh. codex.detach_app_model_catalog() try: @@ -230,7 +231,7 @@ def install_tool_binary( print_section("Bootstrap") print_warning(f"`{binary}` was not found. Installing {spec['display']}...") - if tool == "codex": + if tool == AGENT_CODEX: codex.detach_app_model_catalog() try: subprocess_cross_os.run( @@ -307,7 +308,7 @@ def resolve_launch_model( ) -> tuple[dict, str | None]: model = explicit_model or default_model_for_tool(tool, state) # if model is not specified for codex, then launch with harness's default model. - if not model and tool != "codex": + if not model and tool != AGENT_CODEX: raise RuntimeError( f"No models available for {tool}. Run `ucode configure` to set up your workspace." ) @@ -345,7 +346,7 @@ def resolve_provider_models( # Only Claude pins per-family model ids. Codex ignores this map, and gemini resolves # its target through resolve_gemini_provider_model instead — so mapping their targets # through Claude-family logic would be meaningless (see docstring). - if tool != "claude": + if tool != AGENT_CLAUDE: return None, None, relayed return map_claude_family_models(service.get("targets") or []) or None, None, relayed @@ -415,11 +416,11 @@ def configure_tool( previous = smart_routing_v2.apply_config(agent=tool) try: result: dict | tuple[dict, str] - if tool == "codex": + if tool == AGENT_CODEX: result = codex.write_tool_config( state, model, provider=provider, parent_schema=parent_schema ) - elif tool == "claude": + elif tool == AGENT_CLAUDE: # A Model Provider Service or parent schema routes by header and discovers models natively, # so the usual "model required" guard doesn't apply to either Claude source. if not model and not provider and not parent_schema: @@ -488,11 +489,11 @@ def launch( def check_gateway_endpoint(state: dict, tool: str) -> bool: """V2-only: a tool is available iff we discovered models for it.""" - if tool == "claude": + if tool == AGENT_CLAUDE: return bool(state.get("claude_models")) if tool == "opencode": return bool(state.get("opencode_models")) - if tool == "codex": + if tool == AGENT_CODEX: return bool(state.get("codex_models")) if tool == "gemini": return bool(state.get("gemini_models")) @@ -570,8 +571,8 @@ def _configure_one( return configure_tool( tool, state, None, provider=provider, provider_models=provider_models, relayed=relayed ) - if tool == "codex": - return configure_tool("codex", state) + if tool == AGENT_CODEX: + return configure_tool(AGENT_CODEX, state) state, model = resolve_launch_model(tool, state, None) return configure_tool(tool, state, model) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 4eeda776b..2e944c901 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -27,6 +27,7 @@ write_json_file, ) from ucode.constants import ( + AGENT_CLAUDE, LOOPBACK_HOST, MCP_CLEANUP_SCOPES, MCP_USER_SCOPE, @@ -471,7 +472,7 @@ def render_overlay( header_lines.append(f"{MODEL_PROVIDER_SERVICE_HEADER}: {provider}") elif parent_schema: header_lines.append(f"{MODEL_SERVICE_PARENT_SCHEMA_HEADER}: {parent_schema}") - if smart_routing_v2.smart_routing_enabled(): + if smart_routing_v2.smart_routing_enabled(None, default=False, agent=AGENT_CLAUDE): header_lines.append(f"{SMART_ROUTER_RECIPE_HEADER}: {configured_router_name()}") # Relayed: the X-Databricks-AI-Gateway-Token swap header is added per request # by the refresh proxy, not here — a static value would go stale mid-session. diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 423b8fe82..d767f9912 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -199,7 +199,7 @@ def _provider_block( auth_argv = build_custom_auth_token_argv(workspace, custom_oauth) else: auth_argv = build_auth_token_argv(workspace, databricks_profile, use_pat=use_pat) - base_url = build_tool_base_url("codex", workspace) + base_url = build_tool_base_url(AGENT_CODEX, workspace) http_headers = { "User-Agent": f"ucode/{ug_version()} codex/{agent_version('codex')}", } @@ -207,7 +207,7 @@ def _provider_block( http_headers[MODEL_PROVIDER_SERVICE_HEADER] = provider elif parent_schema: http_headers[MODEL_SERVICE_PARENT_SCHEMA_HEADER] = parent_schema - if smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): + if smart_routing_v2.smart_routing_enabled(None, default=False, agent=AGENT_CODEX): http_headers[SMART_ROUTER_RECIPE_HEADER] = configured_router_name() _apply_managed_headers(http_headers, managed_http_headers) return { @@ -391,7 +391,7 @@ def configured_paths(state: dict) -> list[str]: if ( isinstance(static_models, list) and static_models - and not get_provider_service(state, "codex") + and not get_provider_service(state, AGENT_CODEX) ): paths.append(str(CODEX_MODEL_CATALOG_PATH)) return paths @@ -448,7 +448,7 @@ def write_tool_config( profiles[CODEX_PROFILE_NAME].pop(key, None) _set_provider_header(doc, None) write_toml_file(LEGACY_CODEX_CONFIG_PATH, doc) - state = mark_tool_managed(state, "codex", LEGACY_MANAGED_KEYS) + state = mark_tool_managed(state, AGENT_CODEX, LEGACY_MANAGED_KEYS) save_state(state) return state @@ -469,7 +469,7 @@ def write_tool_config( # Back up only a file that predates ucode's management of the tool. A # re-configure would otherwise snapshot ucode's own generated file, and # revert would restore that snapshot instead of deleting the file. - if not is_tool_managed(state, "codex"): + if not is_tool_managed(state, AGENT_CODEX): backup_existing_file(CODEX_CONFIG_PATH, CODEX_BACKUP_PATH) overlay = render_overlay( workspace, @@ -486,7 +486,9 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: prune_key_paths(base, _PROVIDER_HTTP_HEADERS_KEY_PATHS) deep_merge_dict(base, copy.deepcopy(overlay)) # deep_merge can't drop keys, so clear model preferences from an earlier run. - if chosen_model is None and not smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): + if chosen_model is None and not smart_routing_v2.smart_routing_enabled( + None, default=False, agent=AGENT_CODEX + ): for key in ("model", "model_reasoning_effort"): base.pop(key, None) if include_catalog: @@ -513,7 +515,7 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: ) write_toml_file(CODEX_CONFIG_PATH, doc) _reconcile_managed_config(state, lambda base: compose(base, include_catalog=False)) - state = mark_tool_managed(state, "codex", MANAGED_KEYS) + state = mark_tool_managed(state, AGENT_CODEX, MANAGED_KEYS) save_state(state) return state @@ -565,18 +567,18 @@ def managed_config_is_current(state: dict) -> bool: if path is None: return True required_scope = "managed" if managed_writes_allowed() else None - return managed_file_is_verified(state, "codex", path, required_scope=required_scope) + return managed_file_is_verified(state, AGENT_CODEX, path, required_scope=required_scope) def managed_config_status(state: dict) -> tuple[Path | None, str, str]: path = codex_managed_config_path() - status, backup = managed_file_status(state, "codex", path, parser=_parse_managed_config) + status, backup = managed_file_status(state, AGENT_CODEX, path, parser=_parse_managed_config) return path, status, backup def revert_managed_config() -> str: return revert_managed_file( - "codex", + AGENT_CODEX, display="Codex", parser=_parse_managed_config, dumper=tomlkit.dumps, @@ -609,13 +611,13 @@ def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> N desired_doc = compose(existing) conflicts = managed_file_conflicts(managed_before, desired_doc, MANAGED_KEYS) if not managed_writes_allowed() and not conflicts: - mark_managed_file_verified(state, "codex", path, scope="local-compatible") + mark_managed_file_verified(state, AGENT_CODEX, path, scope="local-compatible") return try: reconcile_managed_file( path, tomlkit.dumps(desired_doc), - tool="codex", + tool=AGENT_CODEX, display="Codex", owned_paths=MANAGED_KEYS, parser=_parse_managed_config, @@ -627,9 +629,9 @@ def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> N f"Codex OS-managed settings could not be updated at {path}; continuing with local " f"settings at {CODEX_CONFIG_PATH}." ) - mark_managed_file_verified(state, "codex", path, scope="local-compatible") + mark_managed_file_verified(state, AGENT_CODEX, path, scope="local-compatible") return - mark_managed_file_verified(state, "codex", path) + mark_managed_file_verified(state, AGENT_CODEX, path) MANAGED_MCP_CONFIG_KEY = "mcp_servers" @@ -772,7 +774,7 @@ def reconcile_managed_mcp(state: dict, servers: dict[str, dict]) -> bool: reconcile_managed_file( path, tomlkit.dumps(existing), - tool="codex", + tool=AGENT_CODEX, display="Codex", owned_paths=[[MANAGED_MCP_CONFIG_KEY]], parser=_parse_managed_config, @@ -781,7 +783,9 @@ def reconcile_managed_mcp(state: dict, servers: dict[str, dict]) -> bool: return False # Preserve the scope the model reconcile recorded (e.g. relay-compatible); an MCP-only write only # refreshes the fingerprint, it does not change how the file relates to the model settings. - mark_managed_file_verified(state, "codex", path, scope=managed_file_scope(state, "codex")) + mark_managed_file_verified( + state, AGENT_CODEX, path, scope=managed_file_scope(state, AGENT_CODEX) + ) return True @@ -815,7 +819,7 @@ def default_model(state: dict) -> str | None: """Return a managed Codex model, or leave selection to Codex.""" if isinstance(state.get("codex_default_model"), str): return state["codex_default_model"] - if smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): + if smart_routing_v2.smart_routing_enabled(None, default=False, agent=AGENT_CODEX): return _smart_routing_config_model(state) return None @@ -843,7 +847,7 @@ def config_precedence_paths() -> tuple[Path, ...]: def clear_model_preferences(state: dict) -> bool: """Remove ucode profile model preferences so Codex selects its default.""" - if smart_routing_v2.smart_routing_enabled(agent=AGENT_CODEX): + if smart_routing_v2.smart_routing_enabled(None, default=False, agent=AGENT_CODEX): return False if isinstance(state.get("codex_default_model"), str): return False @@ -856,7 +860,7 @@ def clear_model_preferences(state: dict) -> bool: if changed: # Never snapshot ucode's own generated file here; revert would restore # the snapshot instead of deleting the file. - if not is_tool_managed(state, "codex"): + if not is_tool_managed(state, AGENT_CODEX): backup_existing_file(CODEX_CONFIG_PATH, CODEX_BACKUP_PATH) write_toml_file(CODEX_CONFIG_PATH, doc) return changed @@ -1137,7 +1141,7 @@ def launch( # Launch-scoped admin routing wins over persisted developer configuration. A transient provider # is most specific; otherwise a transient UC parent must suppress a saved provider. provider = transient_provider or ( - None if parent_schema else get_provider_service(state, "codex") + None if parent_schema else get_provider_service(state, AGENT_CODEX) ) if workspace and (provider or parent_schema): _reject_managed_model_catalog() diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ea3a45897..2d9139717 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -770,7 +770,7 @@ def _setup_single_agent( state = states[0] parent_schema = None managed = None - if apply_managed and tool in ("claude", "codex"): + if apply_managed and tool in (AGENT_CLAUDE, AGENT_CODEX): managed, _ = refresh_managed_config(state, force_refresh=True) _reject_disabled_agent(managed, tool) if managed is not None: @@ -783,7 +783,7 @@ def _setup_single_agent( if apply_managed and not managed_enabled_tools(managed or {}): state = _record_unmanaged_configured_agents(state, [tool]) install_databricks_ai_tools_for_agents( - [tool], state, force_refresh=tool not in ("claude", "codex") + [tool], state, force_refresh=tool not in (AGENT_CLAUDE, AGENT_CODEX) ) spec = TOOL_SPECS[tool] provider_summary = "Databricks" if parent_schema else _provider_summary(tool, state) @@ -828,7 +828,7 @@ def _maybe_select_provider_service(tool: str, state: dict) -> dict: No-op for tools other than claude/codex/gemini. Falls back to Databricks when no matching provider services are found or the listing fails. """ - if tool not in ("claude", "codex", "gemini"): + if tool not in (AGENT_CLAUDE, AGENT_CODEX, "gemini"): return state display = TOOL_SPECS[tool]["display"] @@ -960,7 +960,7 @@ def _configure_workspace_command( resolved = resolve_state(managed, state, tool_name) parent_schema = ( managed_unity_catalog_location(managed, tool_name) - if tool_name in ("claude", "codex") + if tool_name in (AGENT_CLAUDE, AGENT_CODEX) and not managed_provider_service(managed, tool_name) else None ) @@ -1121,7 +1121,7 @@ def _status_models(tool: str, state: dict) -> list[str]: static_models = _model_values(state.get(f"{tool}_static_models")) if static_models: models = static_models - elif tool in ("claude", "codex", "gemini", "opencode"): + elif tool in (AGENT_CLAUDE, AGENT_CODEX, "gemini", "opencode"): models = _model_values(state.get(f"{tool}_models")) elif tool == "copilot": models = _model_values(state.get("copilot_models")) or ( @@ -1294,7 +1294,7 @@ def status() -> int: default_model = _status_default_model(tool, effective_state, models) if default_model: rows.append(("Default model", default_model)) - if tool in ("claude", "codex"): + if tool in (AGENT_CLAUDE, AGENT_CODEX): rows.append( ( "Tracing", @@ -2281,7 +2281,7 @@ def codex_router_hook_cmd( import sys if not smart_routing_v2.smart_routing_enabled( - effective_environment(agent=AGENT_CODEX), agent=AGENT_CODEX + effective_environment(agent=AGENT_CODEX), default=False, agent=AGENT_CODEX ): return @@ -2362,7 +2362,7 @@ def claude_router_hook_cmd( import sys if not smart_routing_v2.smart_routing_enabled( - effective_environment(agent=AGENT_CLAUDE), agent=AGENT_CLAUDE + effective_environment(agent=AGENT_CLAUDE), default=False, agent=AGENT_CLAUDE ): return @@ -2476,7 +2476,7 @@ def _auto_configure_tool(tool: str, custom_oauth: CustomOAuthConfig | None = Non ) -CAN_USE_CACHED_CONFIG_AGENTS = frozenset({"claude", "codex"}) +CAN_USE_CACHED_CONFIG_AGENTS = frozenset({AGENT_CLAUDE, AGENT_CODEX}) @contextmanager @@ -2735,9 +2735,9 @@ def _child_owns_stdout(tool: str, tool_args: list[str]) -> bool: """ if "--" in tool_args: tool_args = tool_args[: tool_args.index("--")] - if tool == "claude": + if tool == AGENT_CLAUDE: return any(arg in {"-p", "--print"} for arg in tool_args) - return tool == "codex" and any(arg in {"exec", "e", "app-server"} for arg in tool_args) + return tool == AGENT_CODEX and any(arg in {"exec", "e", "app-server"} for arg in tool_args) def _should_launch_smart_routing( @@ -2756,7 +2756,7 @@ def _smart_routing_launch_shape(tool: str, tool_args: list[str], explicit_prompt """Whether the forwarded arguments represent an interactive launch.""" if not tool_args or explicit_prompt: return True - return tool == "claude" and tool_args[0].startswith("-") + return tool == AGENT_CLAUDE and tool_args[0].startswith("-") def _launch_options( @@ -2823,7 +2823,7 @@ def _launch_tool( # `--model` lands in ctx.args instead of a ucode option. It still determines the effective # launch model and should therefore win in the launch summary. forwarded_model = ( - explicit_model_arg_value(ctx.args) if tool in {"claude", "codex"} else None + explicit_model_arg_value(ctx.args) if tool in {AGENT_CLAUDE, AGENT_CODEX} else None ) # `--model` is exposed by the claude and gemini launch commands. Under a provider it selects # which of the service's targets/tiers to launch on, rather than being rejected — see the @@ -2905,7 +2905,7 @@ def _launch_tool( managed_provider = managed_provider_service(managed or {}, tool) managed_parent_schema = ( managed_unity_catalog_location(managed or {}, tool) - if tool in {"claude", "codex"} and not managed_provider + if tool in {AGENT_CLAUDE, AGENT_CODEX} and not managed_provider else None ) if managed_provider: @@ -2919,13 +2919,13 @@ def _launch_tool( # Unmanaged Claude launches discover gateway models automatically; with no # provider or parent header the gateway defaults to system.ai. Managed # configs opt into discovery by selecting an MPS or Unity Catalog location. - if tool == "claude" and (managed is None or managed_provider or managed_parent_schema): + if tool == AGENT_CLAUDE and (managed is None or managed_provider or managed_parent_schema): os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" # The environment switch remains a developer override; managed config is the workspace # policy equivalent and must take effect before launch options are computed. managed_smart_routing_enabled = _managed_smart_routing_enabled(managed, tool) smart_routing_enabled = smart_routing_v2.smart_routing_enabled( - default=managed_smart_routing_enabled + None, default=managed_smart_routing_enabled, agent=tool ) # Discovery exists to find models and isn't needed for managed config that already names them. managed_models_known = managed_supplies_models(managed, tool) @@ -2986,7 +2986,7 @@ def _launch_tool( relayed = False coding_agent_config_defaults = ( managed_claude_family_models(managed) or {} - if tool == "claude" and managed is not None + if tool == AGENT_CLAUDE and managed is not None else {} ) if provider and tool != "gemini": @@ -3005,7 +3005,7 @@ def _launch_tool( # re-derivation (see resolve_provider_models). Only when the manifest actually selected # this provider for claude, and authored something to pin. if ( - tool == "claude" + tool == AGENT_CLAUDE and managed is not None and provider == managed_provider_service(managed, tool) ): @@ -3015,7 +3015,7 @@ def _launch_tool( coding_agent_config_defaults = authored # Managed defaults choose models without limiting the selected source's catalog. should_fetch_claude_picker_catalog = ( - tool == "claude" + tool == AGENT_CLAUDE and not relayed and ( bool(managed_parent_schema or managed_provider) @@ -3050,19 +3050,19 @@ def _launch_tool( resolved_model = None managed_source_model = ( managed_launch_model(managed or {}, recommendation, tool) - if tool == "claude" and (managed_provider or managed_parent_schema) + if tool == AGENT_CLAUDE and (managed_provider or managed_parent_schema) else None ) - if tool == "claude" and managed_parent_schema: + if tool == AGENT_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: + if tool == AGENT_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 provider and tool == AGENT_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. @@ -3094,16 +3094,16 @@ def _launch_tool( # 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": + if tool == AGENT_CLAUDE: route_root_model = managed_model else: resolved_model = managed_model # An explicit `--model` is the user's own choice and outranks everything above (managed # default, smart-routing pick). Non-claude agents take it as the resolved model, which # Codex keeps an explicit --model in ctx.args and passes it to its CLI verbatim. - if model and tool != "claude": + if model and tool != AGENT_CLAUDE: resolved_model = model - if model and tool == "claude" and not provider: + if model and tool == AGENT_CLAUDE and not provider: route_root_model = None if coding_agent_config_defaults and not state.get("claude_static_models") and not relayed: picker_catalog = claude_agent.default_model_picker_catalog( @@ -3141,7 +3141,7 @@ def _launch_tool( refresh_downloaded_skills_on_launch(state) # Relayed = a Claude subscription: forward the model to Claude Code's own flag, like `-- --model X`. should_forward_relayed_model = ( - tool == "claude" + tool == AGENT_CLAUDE and provider and relayed and relayed_forward_model @@ -3167,10 +3167,10 @@ def _launch_tool( _print_budget_panel(recommendation, tool, managed) # The managed config's MCP servers and skills are both applied at `ug configure`, not here, # so the launch hot path makes no per-launch discovery calls for them. - if tool == "claude": + if tool == AGENT_CLAUDE: if provider: state["_claude_launch_provider"] = provider - elif tool == "codex": + elif tool == AGENT_CODEX: if provider: state["_codex_launch_provider"] = provider elif parent_schema: @@ -3190,7 +3190,7 @@ def _launch_tool( True if managed_smart_routing_enabled and smart_routing_enabled - and not smart_routing_v2.smart_routing_enabled() + and not smart_routing_v2.smart_routing_enabled(None, default=False, agent=tool) else None, agent=tool, ): @@ -3429,9 +3429,9 @@ def codex_cmd( print_err(str(exc)) raise typer.Exit(1) from exc with _smart_routing_v2_flag(enable_smart_routing_flag, agent=AGENT_CODEX): - with _disable_smart_routing_for_subcommand("codex", ctx): + with _disable_smart_routing_for_subcommand(AGENT_CODEX, ctx): _launch_tool( - "codex", + AGENT_CODEX, ctx, provider=provider, refresh=refresh, @@ -3517,9 +3517,9 @@ def claude_cmd( print_err(str(exc)) raise typer.Exit(1) from exc with _smart_routing_v2_flag(enable_smart_routing_flag, agent=AGENT_CLAUDE): - with _disable_smart_routing_for_subcommand("claude", ctx): + with _disable_smart_routing_for_subcommand(AGENT_CLAUDE, ctx): _launch_tool( - "claude", + AGENT_CLAUDE, ctx, provider=provider, model=model, diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index 6d6137d28..70c157b35 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -19,6 +19,7 @@ from typing import Any from ucode.config_io import APP_DIR +from ucode.constants import AGENT_CODEX from ucode.smart_routing import orchestrator, routing from ucode.smart_routing.routing import RoutingDecision @@ -110,7 +111,7 @@ def route_pre_tool_use( def record(payload, task, decision, requested): routing.write_decision_record(DECISIONS_PATH, payload, task, decision, requested) - name_suffix = " [orchestrator on]" if orchestrator.enabled() else "" + name_suffix = " [orchestrator on]" if orchestrator.enabled(agent=AGENT_CODEX) else "" return routing.route_spawn_tool( payload, is_spawn_agent=is_spawn_agent_tool, diff --git a/src/ucode/smart_routing/config.py b/src/ucode/smart_routing/config.py index bc371b3e6..1120f23ad 100644 --- a/src/ucode/smart_routing/config.py +++ b/src/ucode/smart_routing/config.py @@ -123,9 +123,7 @@ def _validate_versions(versions: Mapping[str, Mapping[str, Mapping[str, str]]]) _validate_versions(_VERSIONS) -def resolve_environment( - env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE -) -> dict[str, str]: +def resolve_environment(env: Mapping[str, str] | None = None, *, agent: str) -> dict[str, str]: """Expand a version before applying any launch or session-specific overrides.""" resolved = dict(os.environ if env is None else env) version = resolved.pop(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, "").strip() @@ -134,7 +132,7 @@ def resolve_environment( def apply_config( - env: MutableMapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE + env: MutableMapping[str, str] | None = None, *, agent: str ) -> dict[str, str | None]: """Consume the launch selector, returning the values needed to restore its input.""" target = os.environ if env is None else env diff --git a/src/ucode/smart_routing/orchestrator.py b/src/ucode/smart_routing/orchestrator.py index 81bfc6f39..8bc6afe16 100644 --- a/src/ucode/smart_routing/orchestrator.py +++ b/src/ucode/smart_routing/orchestrator.py @@ -31,12 +31,12 @@ ) -def feature_enabled(env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE) -> bool: +def feature_enabled(env: Mapping[str, str] | None = None, *, agent: str) -> bool: source = resolve_environment(env, agent=agent) return source.get(ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR) == "1" -def enabled(env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE) -> bool: +def enabled(env: Mapping[str, str] | None = None, *, agent: str) -> bool: from ucode.smart_routing.v2 import smart_routing_enabled source = os.environ if env is None else env @@ -50,7 +50,9 @@ def enabled(env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE) return False except (RuntimeError, OSError): return False - return smart_routing_enabled(effective_environment(source, agent=agent), agent=agent) + return smart_routing_enabled( + effective_environment(source, agent=agent), default=False, agent=agent + ) def skill_directory() -> Path: @@ -59,7 +61,7 @@ def skill_directory() -> Path: def add_claude_agents(plugin_dir: Path) -> None: """Load roles alongside the router's exact-model agents, only for this launch.""" - if feature_enabled(): + if feature_enabled(agent=AGENT_CLAUDE): shutil.copytree(skill_directory() / "agents", plugin_dir / "agents", dirs_exist_ok=True) diff --git a/src/ucode/smart_routing/session_env.py b/src/ucode/smart_routing/session_env.py index f88378bd5..1432601ef 100644 --- a/src/ucode/smart_routing/session_env.py +++ b/src/ucode/smart_routing/session_env.py @@ -10,7 +10,7 @@ from pathlib import Path from ucode.config_io import atomic_write_json -from ucode.constants import AGENT_CLAUDE, SMART_ROUTING_ENV_KEYS +from ucode.constants import SMART_ROUTING_ENV_KEYS from ucode.smart_routing.config import resolve_environment SESSION_ENV_VAR = "UCODE_SESSION_ENV_FILE" @@ -52,9 +52,7 @@ def _read(path: Path) -> dict[str, str]: return _validate(json.loads(path.read_text(encoding="utf-8"))) -def effective_environment( - env: Mapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE -) -> dict[str, str]: +def effective_environment(env: Mapping[str, str] | None = None, *, agent: str) -> dict[str, str]: """Overlay the latest session controls on the hook process environment.""" effective = resolve_environment(env, agent=agent) try: diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index bb8aa1e58..0e419d049 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -154,7 +154,7 @@ def _model_picker_catalog() -> AnthropicModelCatalog | None: def smart_routing_enabled( - env: MutableMapping[str, str] | None = None, *, default: bool = False, agent: str = AGENT_CLAUDE + env: MutableMapping[str, str] | None, *, default: bool, agent: str ) -> bool: source = resolve_environment(env, agent=agent) values = [ @@ -168,7 +168,7 @@ def smart_routing_enabled( def first_prompt_routing_enabled( - env: MutableMapping[str, str] | None = None, *, agent: str = AGENT_CLAUDE + env: MutableMapping[str, str] | None = None, *, agent: str ) -> bool: """Whether the first prompt is routed. Subagent-only wins over the full V2 flag.""" source = resolve_environment(env, agent=agent) @@ -179,17 +179,17 @@ def first_prompt_routing_enabled( def enable_smart_routing( - env: MutableMapping[str, str] | None = None, + env: MutableMapping[str, str] | None = None, *, agent: str ) -> dict[str, str | None]: """Set the full smart-routing env var and return the prior value of every routing var.""" - return override_smart_routing(True, env) + return override_smart_routing(True, env, agent=agent) def override_smart_routing( enabled: bool, env: MutableMapping[str, str] | None = None, *, - agent: str = AGENT_CLAUDE, + agent: str, ) -> dict[str, str | None]: """Set an explicit launch-scoped routing choice and return the prior values.""" target = os.environ if env is None else env @@ -366,7 +366,7 @@ def _request_claude_routing_decision( available.setdefault(_claude_router_model_id(model), model) if not available: return None, "Anthropic models endpoint returned no Claude models" - route_options = [(model, "claude") for model in available] + route_options = [(model, AGENT_CLAUDE) for model in available] return routing.select_route( workspace, token, @@ -429,7 +429,7 @@ def route_claude_pre_tool_use( route.routed_model, ) agent_name = claude_routing.SUBAGENT_NOTICE_CONFIG.name(route.tool_input) or "subagent" - if orchestrator.enabled(): + if orchestrator.enabled(agent=AGENT_CLAUDE): agent_name += " [orchestrator on]" routing_message = claude_routing.SUBAGENT_NOTICE_CONFIG.message( route.decision, @@ -530,7 +530,7 @@ def launch_claude( ) model_ids = catalog.model_ids - route_first_prompt = first_prompt_routing_enabled() + route_first_prompt = first_prompt_routing_enabled(agent=AGENT_CLAUDE) # TODO: Restore first-prompt routing on Windows after replacing the Unix-only PTY wrapper: # https://databricks.atlassian.net/browse/AIGTWY-4385 if route_first_prompt and os.name == "nt": @@ -547,7 +547,9 @@ def launch_claude( if not isinstance(env, dict): raise RuntimeError("Claude settings 'env' must be an object for smart routing.") env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None) - env[ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR] = "1" if orchestrator.feature_enabled() else "0" + env[ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR] = ( + "1" if orchestrator.feature_enabled(agent=AGENT_CLAUDE) else "0" + ) if route_first_prompt: env[ENABLE_SMART_ROUTING_ENV_VAR] = "1" else: @@ -563,7 +565,7 @@ def launch_claude( sync_smart_routing_hooks(settings, routing_state, enabled=True) if route_first_prompt: sync_first_prompt_hook(settings, hook_executable) - orchestrator.sync_hooks(settings, agent="claude") + orchestrator.sync_hooks(settings, agent=AGENT_CLAUDE) model_setting = _ClaudeModelSettingGuard(user_settings_path) def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: @@ -583,7 +585,7 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: plugin_dir = launch_dir / "plugin" if route_first_prompt: env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) - session_path = _prepare_smart_router_session("claude") + session_path = _prepare_smart_router_session(AGENT_CLAUDE) env[SESSION_ENV_VAR] = str(session_path) env[SESSION_PYTHON_ENV_VAR] = os.environ[SESSION_PYTHON_ENV_VAR] try: @@ -636,7 +638,7 @@ def _v2_hooks(state: dict, available_models: list[str]) -> dict: "PreToolUse": merge_pre_tool_use_hooks([], state, available_models=available_models), } } - orchestrator.sync_hooks(doc, agent="codex") + orchestrator.sync_hooks(doc, agent=AGENT_CODEX) return doc["hooks"] @@ -680,7 +682,7 @@ def launch_codex( overlay["model_catalog_json"] = str(catalog_path) overlay["hooks"] = _v2_hooks(state, available_models) overlay["features.hooks"] = True - session_env_path = _prepare_smart_router_session("codex") + session_env_path = _prepare_smart_router_session(AGENT_CODEX) # Codex constructs tool subprocess environments through its shell policy. # Pass both the session marker and its launching interpreter through that policy. overlay[f"shell_environment_policy.set.{SESSION_ENV_VAR}"] = str(session_env_path) diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 848386ae7..fac4836b4 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -13,6 +13,7 @@ import pytest from ucode.agents import LaunchOptions, claude +from ucode.constants import AGENT_CLAUDE from ucode.databricks import AnthropicModelCatalog from ucode.smart_routing import claude_hooks, routing, v2 @@ -199,22 +200,22 @@ class TestSmartRoutingEnvVars: def test_either_flag_enables_smart_routing(self, monkeypatch): monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) monkeypatch.delenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, raising=False) - assert not v2.smart_routing_enabled() + assert not v2.smart_routing_enabled(None, default=False, agent=AGENT_CLAUDE) monkeypatch.setenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, "1") - assert v2.smart_routing_enabled() + assert v2.smart_routing_enabled(None, default=False, agent=AGENT_CLAUDE) monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") - assert v2.smart_routing_enabled() + assert v2.smart_routing_enabled(None, default=False, agent=AGENT_CLAUDE) def test_subagent_only_flag_suppresses_first_prompt_routing(self, monkeypatch): monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) monkeypatch.setenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, "1") - assert not v2.first_prompt_routing_enabled() + assert not v2.first_prompt_routing_enabled(agent=AGENT_CLAUDE) # Subagent-only wins when both are set: an ambient full flag cannot # override an explicit subagent-only session. monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") - assert not v2.first_prompt_routing_enabled() + assert not v2.first_prompt_routing_enabled(agent=AGENT_CLAUDE) monkeypatch.delenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR) - assert v2.first_prompt_routing_enabled() + assert v2.first_prompt_routing_enabled(agent=AGENT_CLAUDE) class TestV2Launch: diff --git a/tests/test_cli.py b/tests/test_cli.py index c7facdc99..4f746ae72 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -997,7 +997,7 @@ def test_disable_smart_routing_is_consumed_by_ug(self, tool): patch( "ucode.cli._launch_tool", side_effect=lambda *_args, **_kwargs: routing_during_launch.append( - cli_mod.smart_routing_v2.smart_routing_enabled() + cli_mod.smart_routing_v2.smart_routing_enabled(None, default=False, agent=tool) ), ) as mock_launch, ): diff --git a/tests/test_smart_routing_config.py b/tests/test_smart_routing_config.py index bd19685c2..3563ba0a4 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -60,9 +60,9 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): applied[SMART_ROUTER_CONFIG_VERSION_ENV_VAR] = selector original = applied.copy() - assert config.resolve_environment(applied) == expected + assert config.resolve_environment(applied, agent=AGENT_CLAUDE) == expected assert applied == original - assert config.apply_config(applied) == {} + assert config.apply_config(applied, agent=AGENT_CLAUDE) == {} assert applied == original @@ -135,16 +135,16 @@ def test_legacy_preset_equivalence(version, expected_routing): SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version, "SMART_ROUTER_NAME": legacy_env["SMART_ROUTER_NAME"], } - resolved = config.resolve_environment(preset_env) + resolved = config.resolve_environment(preset_env, agent=AGENT_CLAUDE) applied = preset_env.copy() - config.apply_config(applied) + config.apply_config(applied, agent=AGENT_CLAUDE) assert applied == resolved for environment in (legacy_env, preset_env, resolved, applied): assert environment["SMART_ROUTER_NAME"] == legacy_env["SMART_ROUTER_NAME"] assert ( - v2.smart_routing_enabled(environment), - v2.first_prompt_routing_enabled(environment), - orchestrator.feature_enabled(environment), + v2.smart_routing_enabled(environment, default=False, agent=AGENT_CLAUDE), + v2.first_prompt_routing_enabled(environment, agent=AGENT_CLAUDE), + orchestrator.feature_enabled(environment, agent=AGENT_CLAUDE), ) == expected_routing @@ -195,3 +195,14 @@ def test_agent_resolution_and_session_controls(monkeypatch, agent, tmp_path): session_env.SESSION_ENV_VAR: str(control), ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", } + + +@pytest.mark.parametrize("resolver", [config.resolve_environment, config.apply_config]) +def test_config_resolution_requires_an_explicit_agent(resolver): + with pytest.raises(TypeError, match="agent"): + resolver({}) + + +def test_smart_routing_check_requires_an_explicit_agent(): + with pytest.raises(TypeError, match="agent"): + v2.smart_routing_enabled(None, default=False) From ad7ee48f4f48976c320d74d0fadef3038a7d2f3b Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:41:28 +0000 Subject: [PATCH 3/9] Limit agent constant replacements to routing config changes --- src/ucode/agents/__init__.py | 43 ++++++++++++------------- src/ucode/agents/codex.py | 34 ++++++++++---------- src/ucode/cli.py | 60 +++++++++++++++++------------------ src/ucode/smart_routing/v2.py | 10 +++--- 4 files changed, 72 insertions(+), 75 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 995072cb4..b8869c5dc 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -16,7 +16,6 @@ from pathlib import Path from ucode.config_io import ToolSpec -from ucode.constants import AGENT_CLAUDE, AGENT_CODEX from ucode.databricks import ( AnthropicModelCatalog, get_databricks_token, @@ -48,8 +47,8 @@ from .args import explicit_model_arg_value as explicit_model_arg_value _MODULES = { - AGENT_CODEX: codex, - AGENT_CLAUDE: claude, + "codex": codex, + "claude": claude, "gemini": gemini, "opencode": opencode, "copilot": copilot, @@ -64,9 +63,9 @@ # rejects it and the model-config paths never see it. The `configure`/MCP flows # handle "cursor" separately as an MCP-only client (see MCP_ONLY_CLIENTS). TOOL_ALIASES = { - AGENT_CODEX: AGENT_CODEX, - AGENT_CLAUDE: AGENT_CLAUDE, - "claude-code": AGENT_CLAUDE, + "codex": "codex", + "claude": "claude", + "claude-code": "claude", "gemini": "gemini", "gemini-cli": "gemini", "opencode": "opencode", @@ -74,18 +73,18 @@ "pi": "pi", } -DEFAULT_TOOL = AGENT_CODEX +DEFAULT_TOOL = "codex" BUNDLE_VERSION = 1 -_MANAGED_SETTINGS_TOOLS = {AGENT_CLAUDE, AGENT_CODEX} +_MANAGED_SETTINGS_TOOLS = {"claude", "codex"} _NATIVE_UPGRADE_COMMANDS = { - AGENT_CLAUDE: ["claude", "upgrade"], - AGENT_CODEX: ["codex", "update"], + "claude": ["claude", "upgrade"], + "codex": ["codex", "update"], } # ucode tool -> `databricks aitools` agent id. gemini/pi aren't supported. AITOOLS_AGENT_TOKENS = { - AGENT_CLAUDE: "claude-code", - AGENT_CODEX: "codex", + "claude": "claude-code", + "codex": "codex", "opencode": "opencode", "copilot": "copilot", } @@ -138,7 +137,7 @@ def _update_installed_tool_binary(tool: str, version: str | None = None) -> bool command = ["npm", "install", "-g", target] print_note(f"Upgrading {spec['display']}...") - if tool == AGENT_CODEX: + if tool == "codex": # Detach potentially incompatible metadata until the next validated refresh. codex.detach_app_model_catalog() try: @@ -231,7 +230,7 @@ def install_tool_binary( print_section("Bootstrap") print_warning(f"`{binary}` was not found. Installing {spec['display']}...") - if tool == AGENT_CODEX: + if tool == "codex": codex.detach_app_model_catalog() try: subprocess_cross_os.run( @@ -308,7 +307,7 @@ def resolve_launch_model( ) -> tuple[dict, str | None]: model = explicit_model or default_model_for_tool(tool, state) # if model is not specified for codex, then launch with harness's default model. - if not model and tool != AGENT_CODEX: + if not model and tool != "codex": raise RuntimeError( f"No models available for {tool}. Run `ucode configure` to set up your workspace." ) @@ -346,7 +345,7 @@ def resolve_provider_models( # Only Claude pins per-family model ids. Codex ignores this map, and gemini resolves # its target through resolve_gemini_provider_model instead — so mapping their targets # through Claude-family logic would be meaningless (see docstring). - if tool != AGENT_CLAUDE: + if tool != "claude": return None, None, relayed return map_claude_family_models(service.get("targets") or []) or None, None, relayed @@ -416,11 +415,11 @@ def configure_tool( previous = smart_routing_v2.apply_config(agent=tool) try: result: dict | tuple[dict, str] - if tool == AGENT_CODEX: + if tool == "codex": result = codex.write_tool_config( state, model, provider=provider, parent_schema=parent_schema ) - elif tool == AGENT_CLAUDE: + elif tool == "claude": # A Model Provider Service or parent schema routes by header and discovers models natively, # so the usual "model required" guard doesn't apply to either Claude source. if not model and not provider and not parent_schema: @@ -489,11 +488,11 @@ def launch( def check_gateway_endpoint(state: dict, tool: str) -> bool: """V2-only: a tool is available iff we discovered models for it.""" - if tool == AGENT_CLAUDE: + if tool == "claude": return bool(state.get("claude_models")) if tool == "opencode": return bool(state.get("opencode_models")) - if tool == AGENT_CODEX: + if tool == "codex": return bool(state.get("codex_models")) if tool == "gemini": return bool(state.get("gemini_models")) @@ -571,8 +570,8 @@ def _configure_one( return configure_tool( tool, state, None, provider=provider, provider_models=provider_models, relayed=relayed ) - if tool == AGENT_CODEX: - return configure_tool(AGENT_CODEX, state) + if tool == "codex": + return configure_tool("codex", state) state, model = resolve_launch_model(tool, state, None) return configure_tool(tool, state, model) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index d767f9912..1d591381f 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -199,7 +199,7 @@ def _provider_block( auth_argv = build_custom_auth_token_argv(workspace, custom_oauth) else: auth_argv = build_auth_token_argv(workspace, databricks_profile, use_pat=use_pat) - base_url = build_tool_base_url(AGENT_CODEX, workspace) + base_url = build_tool_base_url("codex", workspace) http_headers = { "User-Agent": f"ucode/{ug_version()} codex/{agent_version('codex')}", } @@ -391,7 +391,7 @@ def configured_paths(state: dict) -> list[str]: if ( isinstance(static_models, list) and static_models - and not get_provider_service(state, AGENT_CODEX) + and not get_provider_service(state, "codex") ): paths.append(str(CODEX_MODEL_CATALOG_PATH)) return paths @@ -448,7 +448,7 @@ def write_tool_config( profiles[CODEX_PROFILE_NAME].pop(key, None) _set_provider_header(doc, None) write_toml_file(LEGACY_CODEX_CONFIG_PATH, doc) - state = mark_tool_managed(state, AGENT_CODEX, LEGACY_MANAGED_KEYS) + state = mark_tool_managed(state, "codex", LEGACY_MANAGED_KEYS) save_state(state) return state @@ -469,7 +469,7 @@ def write_tool_config( # Back up only a file that predates ucode's management of the tool. A # re-configure would otherwise snapshot ucode's own generated file, and # revert would restore that snapshot instead of deleting the file. - if not is_tool_managed(state, AGENT_CODEX): + if not is_tool_managed(state, "codex"): backup_existing_file(CODEX_CONFIG_PATH, CODEX_BACKUP_PATH) overlay = render_overlay( workspace, @@ -515,7 +515,7 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict: ) write_toml_file(CODEX_CONFIG_PATH, doc) _reconcile_managed_config(state, lambda base: compose(base, include_catalog=False)) - state = mark_tool_managed(state, AGENT_CODEX, MANAGED_KEYS) + state = mark_tool_managed(state, "codex", MANAGED_KEYS) save_state(state) return state @@ -567,18 +567,18 @@ def managed_config_is_current(state: dict) -> bool: if path is None: return True required_scope = "managed" if managed_writes_allowed() else None - return managed_file_is_verified(state, AGENT_CODEX, path, required_scope=required_scope) + return managed_file_is_verified(state, "codex", path, required_scope=required_scope) def managed_config_status(state: dict) -> tuple[Path | None, str, str]: path = codex_managed_config_path() - status, backup = managed_file_status(state, AGENT_CODEX, path, parser=_parse_managed_config) + status, backup = managed_file_status(state, "codex", path, parser=_parse_managed_config) return path, status, backup def revert_managed_config() -> str: return revert_managed_file( - AGENT_CODEX, + "codex", display="Codex", parser=_parse_managed_config, dumper=tomlkit.dumps, @@ -611,13 +611,13 @@ def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> N desired_doc = compose(existing) conflicts = managed_file_conflicts(managed_before, desired_doc, MANAGED_KEYS) if not managed_writes_allowed() and not conflicts: - mark_managed_file_verified(state, AGENT_CODEX, path, scope="local-compatible") + mark_managed_file_verified(state, "codex", path, scope="local-compatible") return try: reconcile_managed_file( path, tomlkit.dumps(desired_doc), - tool=AGENT_CODEX, + tool="codex", display="Codex", owned_paths=MANAGED_KEYS, parser=_parse_managed_config, @@ -629,9 +629,9 @@ def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> N f"Codex OS-managed settings could not be updated at {path}; continuing with local " f"settings at {CODEX_CONFIG_PATH}." ) - mark_managed_file_verified(state, AGENT_CODEX, path, scope="local-compatible") + mark_managed_file_verified(state, "codex", path, scope="local-compatible") return - mark_managed_file_verified(state, AGENT_CODEX, path) + mark_managed_file_verified(state, "codex", path) MANAGED_MCP_CONFIG_KEY = "mcp_servers" @@ -774,7 +774,7 @@ def reconcile_managed_mcp(state: dict, servers: dict[str, dict]) -> bool: reconcile_managed_file( path, tomlkit.dumps(existing), - tool=AGENT_CODEX, + tool="codex", display="Codex", owned_paths=[[MANAGED_MCP_CONFIG_KEY]], parser=_parse_managed_config, @@ -783,9 +783,7 @@ def reconcile_managed_mcp(state: dict, servers: dict[str, dict]) -> bool: return False # Preserve the scope the model reconcile recorded (e.g. relay-compatible); an MCP-only write only # refreshes the fingerprint, it does not change how the file relates to the model settings. - mark_managed_file_verified( - state, AGENT_CODEX, path, scope=managed_file_scope(state, AGENT_CODEX) - ) + mark_managed_file_verified(state, "codex", path, scope=managed_file_scope(state, "codex")) return True @@ -860,7 +858,7 @@ def clear_model_preferences(state: dict) -> bool: if changed: # Never snapshot ucode's own generated file here; revert would restore # the snapshot instead of deleting the file. - if not is_tool_managed(state, AGENT_CODEX): + if not is_tool_managed(state, "codex"): backup_existing_file(CODEX_CONFIG_PATH, CODEX_BACKUP_PATH) write_toml_file(CODEX_CONFIG_PATH, doc) return changed @@ -1141,7 +1139,7 @@ def launch( # Launch-scoped admin routing wins over persisted developer configuration. A transient provider # is most specific; otherwise a transient UC parent must suppress a saved provider. provider = transient_provider or ( - None if parent_schema else get_provider_service(state, AGENT_CODEX) + None if parent_schema else get_provider_service(state, "codex") ) if workspace and (provider or parent_schema): _reject_managed_model_catalog() diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 2d9139717..11f53876e 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -770,7 +770,7 @@ def _setup_single_agent( state = states[0] parent_schema = None managed = None - if apply_managed and tool in (AGENT_CLAUDE, AGENT_CODEX): + if apply_managed and tool in ("claude", "codex"): managed, _ = refresh_managed_config(state, force_refresh=True) _reject_disabled_agent(managed, tool) if managed is not None: @@ -783,7 +783,7 @@ def _setup_single_agent( if apply_managed and not managed_enabled_tools(managed or {}): state = _record_unmanaged_configured_agents(state, [tool]) install_databricks_ai_tools_for_agents( - [tool], state, force_refresh=tool not in (AGENT_CLAUDE, AGENT_CODEX) + [tool], state, force_refresh=tool not in ("claude", "codex") ) spec = TOOL_SPECS[tool] provider_summary = "Databricks" if parent_schema else _provider_summary(tool, state) @@ -828,7 +828,7 @@ def _maybe_select_provider_service(tool: str, state: dict) -> dict: No-op for tools other than claude/codex/gemini. Falls back to Databricks when no matching provider services are found or the listing fails. """ - if tool not in (AGENT_CLAUDE, AGENT_CODEX, "gemini"): + if tool not in ("claude", "codex", "gemini"): return state display = TOOL_SPECS[tool]["display"] @@ -960,7 +960,7 @@ def _configure_workspace_command( resolved = resolve_state(managed, state, tool_name) parent_schema = ( managed_unity_catalog_location(managed, tool_name) - if tool_name in (AGENT_CLAUDE, AGENT_CODEX) + if tool_name in ("claude", "codex") and not managed_provider_service(managed, tool_name) else None ) @@ -1121,7 +1121,7 @@ def _status_models(tool: str, state: dict) -> list[str]: static_models = _model_values(state.get(f"{tool}_static_models")) if static_models: models = static_models - elif tool in (AGENT_CLAUDE, AGENT_CODEX, "gemini", "opencode"): + elif tool in ("claude", "codex", "gemini", "opencode"): models = _model_values(state.get(f"{tool}_models")) elif tool == "copilot": models = _model_values(state.get("copilot_models")) or ( @@ -1294,7 +1294,7 @@ def status() -> int: default_model = _status_default_model(tool, effective_state, models) if default_model: rows.append(("Default model", default_model)) - if tool in (AGENT_CLAUDE, AGENT_CODEX): + if tool in ("claude", "codex"): rows.append( ( "Tracing", @@ -2476,7 +2476,7 @@ def _auto_configure_tool(tool: str, custom_oauth: CustomOAuthConfig | None = Non ) -CAN_USE_CACHED_CONFIG_AGENTS = frozenset({AGENT_CLAUDE, AGENT_CODEX}) +CAN_USE_CACHED_CONFIG_AGENTS = frozenset({"claude", "codex"}) @contextmanager @@ -2735,9 +2735,9 @@ def _child_owns_stdout(tool: str, tool_args: list[str]) -> bool: """ if "--" in tool_args: tool_args = tool_args[: tool_args.index("--")] - if tool == AGENT_CLAUDE: + if tool == "claude": return any(arg in {"-p", "--print"} for arg in tool_args) - return tool == AGENT_CODEX and any(arg in {"exec", "e", "app-server"} for arg in tool_args) + return tool == "codex" and any(arg in {"exec", "e", "app-server"} for arg in tool_args) def _should_launch_smart_routing( @@ -2756,7 +2756,7 @@ def _smart_routing_launch_shape(tool: str, tool_args: list[str], explicit_prompt """Whether the forwarded arguments represent an interactive launch.""" if not tool_args or explicit_prompt: return True - return tool == AGENT_CLAUDE and tool_args[0].startswith("-") + return tool == "claude" and tool_args[0].startswith("-") def _launch_options( @@ -2823,7 +2823,7 @@ def _launch_tool( # `--model` lands in ctx.args instead of a ucode option. It still determines the effective # launch model and should therefore win in the launch summary. forwarded_model = ( - explicit_model_arg_value(ctx.args) if tool in {AGENT_CLAUDE, AGENT_CODEX} else None + explicit_model_arg_value(ctx.args) if tool in {"claude", "codex"} else None ) # `--model` is exposed by the claude and gemini launch commands. Under a provider it selects # which of the service's targets/tiers to launch on, rather than being rejected — see the @@ -2905,7 +2905,7 @@ def _launch_tool( managed_provider = managed_provider_service(managed or {}, tool) managed_parent_schema = ( managed_unity_catalog_location(managed or {}, tool) - if tool in {AGENT_CLAUDE, AGENT_CODEX} and not managed_provider + if tool in {"claude", "codex"} and not managed_provider else None ) if managed_provider: @@ -2919,7 +2919,7 @@ def _launch_tool( # Unmanaged Claude launches discover gateway models automatically; with no # provider or parent header the gateway defaults to system.ai. Managed # configs opt into discovery by selecting an MPS or Unity Catalog location. - if tool == AGENT_CLAUDE and (managed is None or managed_provider or managed_parent_schema): + if tool == "claude" and (managed is None or managed_provider or managed_parent_schema): os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" # The environment switch remains a developer override; managed config is the workspace # policy equivalent and must take effect before launch options are computed. @@ -2986,7 +2986,7 @@ def _launch_tool( relayed = False coding_agent_config_defaults = ( managed_claude_family_models(managed) or {} - if tool == AGENT_CLAUDE and managed is not None + if tool == "claude" and managed is not None else {} ) if provider and tool != "gemini": @@ -3005,7 +3005,7 @@ def _launch_tool( # re-derivation (see resolve_provider_models). Only when the manifest actually selected # this provider for claude, and authored something to pin. if ( - tool == AGENT_CLAUDE + tool == "claude" and managed is not None and provider == managed_provider_service(managed, tool) ): @@ -3015,7 +3015,7 @@ def _launch_tool( coding_agent_config_defaults = authored # Managed defaults choose models without limiting the selected source's catalog. should_fetch_claude_picker_catalog = ( - tool == AGENT_CLAUDE + tool == "claude" and not relayed and ( bool(managed_parent_schema or managed_provider) @@ -3050,19 +3050,19 @@ def _launch_tool( resolved_model = None managed_source_model = ( managed_launch_model(managed or {}, recommendation, tool) - if tool == AGENT_CLAUDE and (managed_provider or managed_parent_schema) + if tool == "claude" and (managed_provider or managed_parent_schema) else None ) - if tool == AGENT_CLAUDE and managed_parent_schema: + 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 == AGENT_CLAUDE and managed_provider: + 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 == AGENT_CLAUDE and (provider_launch_model or provider_models): + 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. @@ -3094,16 +3094,16 @@ def _launch_tool( # 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 == AGENT_CLAUDE: + if tool == "claude": route_root_model = managed_model else: resolved_model = managed_model # An explicit `--model` is the user's own choice and outranks everything above (managed # default, smart-routing pick). Non-claude agents take it as the resolved model, which # Codex keeps an explicit --model in ctx.args and passes it to its CLI verbatim. - if model and tool != AGENT_CLAUDE: + if model and tool != "claude": resolved_model = model - if model and tool == AGENT_CLAUDE and not provider: + if model and tool == "claude" and not provider: route_root_model = None if coding_agent_config_defaults and not state.get("claude_static_models") and not relayed: picker_catalog = claude_agent.default_model_picker_catalog( @@ -3141,7 +3141,7 @@ def _launch_tool( refresh_downloaded_skills_on_launch(state) # Relayed = a Claude subscription: forward the model to Claude Code's own flag, like `-- --model X`. should_forward_relayed_model = ( - tool == AGENT_CLAUDE + tool == "claude" and provider and relayed and relayed_forward_model @@ -3167,10 +3167,10 @@ def _launch_tool( _print_budget_panel(recommendation, tool, managed) # The managed config's MCP servers and skills are both applied at `ug configure`, not here, # so the launch hot path makes no per-launch discovery calls for them. - if tool == AGENT_CLAUDE: + if tool == "claude": if provider: state["_claude_launch_provider"] = provider - elif tool == AGENT_CODEX: + elif tool == "codex": if provider: state["_codex_launch_provider"] = provider elif parent_schema: @@ -3429,9 +3429,9 @@ def codex_cmd( print_err(str(exc)) raise typer.Exit(1) from exc with _smart_routing_v2_flag(enable_smart_routing_flag, agent=AGENT_CODEX): - with _disable_smart_routing_for_subcommand(AGENT_CODEX, ctx): + with _disable_smart_routing_for_subcommand("codex", ctx): _launch_tool( - AGENT_CODEX, + "codex", ctx, provider=provider, refresh=refresh, @@ -3517,9 +3517,9 @@ def claude_cmd( print_err(str(exc)) raise typer.Exit(1) from exc with _smart_routing_v2_flag(enable_smart_routing_flag, agent=AGENT_CLAUDE): - with _disable_smart_routing_for_subcommand(AGENT_CLAUDE, ctx): + with _disable_smart_routing_for_subcommand("claude", ctx): _launch_tool( - AGENT_CLAUDE, + "claude", ctx, provider=provider, model=model, diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 0e419d049..b03d4d97b 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -366,7 +366,7 @@ def _request_claude_routing_decision( available.setdefault(_claude_router_model_id(model), model) if not available: return None, "Anthropic models endpoint returned no Claude models" - route_options = [(model, AGENT_CLAUDE) for model in available] + route_options = [(model, "claude") for model in available] return routing.select_route( workspace, token, @@ -565,7 +565,7 @@ def launch_claude( sync_smart_routing_hooks(settings, routing_state, enabled=True) if route_first_prompt: sync_first_prompt_hook(settings, hook_executable) - orchestrator.sync_hooks(settings, agent=AGENT_CLAUDE) + orchestrator.sync_hooks(settings, agent="claude") model_setting = _ClaudeModelSettingGuard(user_settings_path) def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: @@ -585,7 +585,7 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: plugin_dir = launch_dir / "plugin" if route_first_prompt: env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) - session_path = _prepare_smart_router_session(AGENT_CLAUDE) + session_path = _prepare_smart_router_session("claude") env[SESSION_ENV_VAR] = str(session_path) env[SESSION_PYTHON_ENV_VAR] = os.environ[SESSION_PYTHON_ENV_VAR] try: @@ -638,7 +638,7 @@ def _v2_hooks(state: dict, available_models: list[str]) -> dict: "PreToolUse": merge_pre_tool_use_hooks([], state, available_models=available_models), } } - orchestrator.sync_hooks(doc, agent=AGENT_CODEX) + orchestrator.sync_hooks(doc, agent="codex") return doc["hooks"] @@ -682,7 +682,7 @@ def launch_codex( overlay["model_catalog_json"] = str(catalog_path) overlay["hooks"] = _v2_hooks(state, available_models) overlay["features.hooks"] = True - session_env_path = _prepare_smart_router_session(AGENT_CODEX) + session_env_path = _prepare_smart_router_session("codex") # Codex constructs tool subprocess environments through its shell policy. # Pass both the session marker and its launching interpreter through that policy. overlay[f"shell_environment_policy.set.{SESSION_ENV_VAR}"] = str(session_env_path) From b6b4527a8673f1ea6fb0cfaa21584b1cae8c683c Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:44:38 +0000 Subject: [PATCH 4/9] Share routing presets and define only intentional agent overrides --- src/ucode/smart_routing/config.py | 104 ++++++++++++----------------- tests/README.md | 4 +- tests/integration/README.md | 4 +- tests/test_agents_init.py | 10 +-- tests/test_cli.py | 9 ++- tests/test_smart_routing_config.py | 62 +++++++++-------- 6 files changed, 91 insertions(+), 102 deletions(-) diff --git a/src/ucode/smart_routing/config.py b/src/ucode/smart_routing/config.py index 1120f23ad..7aa69b950 100644 --- a/src/ucode/smart_routing/config.py +++ b/src/ucode/smart_routing/config.py @@ -32,79 +32,54 @@ _VERSIONS = { FIRST_PROMPT_AND_SUBAGENT_NO_ORCH_V0: { - AGENT_CLAUDE: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", - }, - AGENT_CODEX: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", - }, + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", }, SUBAGENT_ONLY_V0: { - AGENT_CLAUDE: { - ENABLE_SMART_ROUTING_ENV_VAR: "0", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", - }, - AGENT_CODEX: { - ENABLE_SMART_ROUTING_ENV_VAR: "0", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", - }, + ENABLE_SMART_ROUTING_ENV_VAR: "0", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", }, SUBAGENT_ONLY_V1: { - AGENT_CLAUDE: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", - }, - AGENT_CODEX: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", - }, + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", }, SUBAGENT_ORCH_V0: { - AGENT_CLAUDE: { - ENABLE_SMART_ROUTING_ENV_VAR: "0", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", - }, - AGENT_CODEX: { - ENABLE_SMART_ROUTING_ENV_VAR: "0", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", - }, + ENABLE_SMART_ROUTING_ENV_VAR: "0", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", }, SUBAGENT_ORCH_V1: { - AGENT_CLAUDE: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", - }, - AGENT_CODEX: { - ENABLE_SMART_ROUTING_ENV_VAR: "1", - ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", - }, + ENABLE_SMART_ROUTING_ENV_VAR: "1", + ENABLE_SUBAGENT_ROUTING_ENV_VAR: "1", + ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "1", }, } +# Only intentional agent differences belong here; omitted flags inherit the shared preset. +# Comment each override with the reason that agent differs. +_AGENT_OVERRIDES: dict[str, dict[str, dict[str, str]]] = {} -def _validate_versions(versions: Mapping[str, Mapping[str, Mapping[str, str]]]) -> None: - """Require both agents and a complete, explicit flag set for each.""" - expected_agents = {AGENT_CLAUDE, AGENT_CODEX} + +def _validate_versions( + versions: Mapping[str, Mapping[str, str]], + overrides: Mapping[str, Mapping[str, Mapping[str, str]]], +) -> None: + """Require complete shared presets and valid, sparse agent overrides.""" expected = set(SMART_ROUTING_ENV_KEYS) - for version, agents in versions.items(): - if agents.keys() != expected_agents: + for version, agents in overrides.items(): + if version not in versions or agents.keys() - {AGENT_CLAUDE, AGENT_CODEX}: raise ValueError( - f"Invalid smart-routing version {version!r}: " - f"expected agents {sorted(expected_agents)}, got {sorted(agents)}." + f"Invalid smart-routing version {version!r}: unknown version or agent." ) - for agent, values in agents.items(): + for version, shared in versions.items(): + configs = {"shared": shared} + configs.update( + {agent: {**shared, **values} for agent, values in overrides.get(version, {}).items()} + ) + for agent, values in configs.items(): missing = expected - values.keys() unexpected = values.keys() - expected if missing or unexpected: @@ -120,14 +95,21 @@ def _validate_versions(versions: Mapping[str, Mapping[str, Mapping[str, str]]]) ) -_validate_versions(_VERSIONS) +_validate_versions(_VERSIONS, _AGENT_OVERRIDES) + + +def _version_config(version: str, *, agent: str) -> dict[str, str] | None: + shared = _VERSIONS.get(version) + if shared is None or agent not in (AGENT_CLAUDE, AGENT_CODEX): + return None + return {**shared, **_AGENT_OVERRIDES.get(version, {}).get(agent, {})} def resolve_environment(env: Mapping[str, str] | None = None, *, agent: str) -> dict[str, str]: """Expand a version before applying any launch or session-specific overrides.""" resolved = dict(os.environ if env is None else env) version = resolved.pop(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, "").strip() - resolved.update(_VERSIONS.get(version, {}).get(agent, {})) + resolved.update(_version_config(version, agent=agent) or {}) return resolved @@ -137,7 +119,7 @@ def apply_config( """Consume the launch selector, returning the values needed to restore its input.""" target = os.environ if env is None else env version = target.get(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, "").strip() - preset = _VERSIONS.get(version, {}).get(agent) + preset = _version_config(version, agent=agent) if preset is None: return {} keys = (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) diff --git a/tests/README.md b/tests/README.md index e1c21d992..a3f98e351 100644 --- a/tests/README.md +++ b/tests/README.md @@ -35,8 +35,8 @@ CUJ4's Claude preset sessions also use it after verifying delegated file tasks: answer does not prove that a resumed child has stopped running. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. -Agent-specific smart-routing preset component checks cover identical existing mappings, -complete per-agent validation, configure order, bootstrap and managed-default launch +Agent-specific smart-routing preset component checks cover shared existing presets, sparse agent overrides, +complete flag validation, configure order, bootstrap and managed-default launch selection, session controls, and restoration after configuration/launch failures. These checks do not establish live agent coverage. diff --git a/tests/integration/README.md b/tests/integration/README.md index b7a3c1f8d..68069f814 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -15,8 +15,8 @@ managed-default and explicit-model cases unchanged. One additional test runs the `SMART_ROUTER_CONFIG_VERSION` values, checking both agents' first-prompt routing, orchestrator context in inference input, and completed explicitly requested routed subagents. This is not automatic orchestrator-delegation coverage. No workspace configuration is modified. -Agent-specific smart-routing preset component checks cover identical existing mappings, -complete per-agent validation, configure order, bootstrap and managed-default launch +Agent-specific smart-routing preset component checks cover shared existing presets, sparse agent overrides, +complete flag validation, configure order, bootstrap and managed-default launch selection, session controls, and restoration after configuration/launch failures. These checks do not establish live agent coverage. diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index bced2e57a..e49cc8e55 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -7,7 +7,6 @@ import subprocess import sys from contextlib import contextmanager, nullcontext, redirect_stdout -from copy import deepcopy from unittest.mock import MagicMock import pytest @@ -1125,9 +1124,8 @@ def test_configure_dispatch_scopes_agent_presets(monkeypatch, order, fail_write) from ucode.smart_routing import config version = config.SUBAGENT_ORCH_V0 - agents = deepcopy(config._VERSIONS[version]) - agents[AGENT_CODEX] = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") - monkeypatch.setitem(config._VERSIONS, version, agents) + codex_flags = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") + monkeypatch.setitem(config._AGENT_OVERRIDES, version, {AGENT_CODEX: codex_flags}) monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) before = { key: os.environ.get(key) @@ -1136,7 +1134,9 @@ def test_configure_dispatch_scopes_agent_presets(monkeypatch, order, fail_write) for agent in order: def write(state, *args, agent=agent, **kwargs): - assert {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} == agents[agent] + assert {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} == ( + codex_flags if agent == AGENT_CODEX else config._VERSIONS[version] + ) assert SMART_ROUTER_CONFIG_VERSION_ENV_VAR not in os.environ if fail_write: raise RuntimeError("write failed") diff --git a/tests/test_cli.py b/tests/test_cli.py index 4f746ae72..80d8c314a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -11,7 +11,6 @@ import sys import time import tomllib -from copy import deepcopy from importlib import metadata from pathlib import Path from unittest.mock import MagicMock, call, patch @@ -512,9 +511,8 @@ def test_agent_preset_reaches_bootstrap_and_launch_and_is_restored( monkeypatch, agent, bare, fail_launch ): version = routing_config.SUBAGENT_ORCH_V0 - agents = deepcopy(routing_config._VERSIONS[version]) - agents[AGENT_CODEX] = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") - monkeypatch.setitem(routing_config._VERSIONS, version, agents) + codex_flags = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") + monkeypatch.setitem(routing_config._AGENT_OVERRIDES, version, {AGENT_CODEX: codex_flags}) monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) before = { key: os.environ.get(key) @@ -544,7 +542,8 @@ def capture_launch(tool, state, args, *, options): ) result = runner.invoke(app, [] if bare else [agent]) assert result.exit_code == (1 if fail_launch else 0), result.output - assert observed == [agents[agent], agents[agent]] + expected = codex_flags if agent == AGENT_CODEX else routing_config._VERSIONS[version] + assert observed == [expected, expected] assert {key: os.environ.get(key) for key in before} == before diff --git a/tests/test_smart_routing_config.py b/tests/test_smart_routing_config.py index 3563ba0a4..ddd2099f6 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -2,8 +2,6 @@ from __future__ import annotations -from copy import deepcopy - import pytest from ucode.constants import ( @@ -148,30 +146,32 @@ def test_legacy_preset_equivalence(version, expected_routing): ) == expected_routing -@pytest.mark.parametrize( - "failure", ["missing_agent", "unknown_agent", "missing_flag", "unknown_flag", "invalid_value"] -) -def test_agent_presets_require_complete_valid_configuration(failure): - agents = deepcopy(config._VERSIONS[config.SUBAGENT_ORCH_V0]) - if failure == "missing_agent": - agents.pop(AGENT_CODEX) - elif failure == "unknown_agent": - agents["other"] = agents[AGENT_CODEX].copy() - elif failure == "missing_flag": - agents[AGENT_CODEX].pop(ENABLE_SMART_ROUTING_ENV_VAR) - elif failure == "unknown_flag": - agents[AGENT_CODEX]["UNKNOWN_FLAG"] = "0" - else: - agents[AGENT_CODEX][ENABLE_SMART_ROUTING_ENV_VAR] = "yes" +@pytest.mark.parametrize("scope", ["shared", "override"]) +@pytest.mark.parametrize("failure", ["unknown_flag", "invalid_value"]) +def test_presets_reject_invalid_flags(scope, failure): + shared = config._VERSIONS[config.SUBAGENT_ORCH_V0].copy() + override = {ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0"} + values = shared if scope == "shared" else override + values["UNKNOWN_FLAG" if failure == "unknown_flag" else ENABLE_SMART_ROUTING_ENV_VAR] = "yes" with pytest.raises(ValueError, match="Invalid smart-routing version"): - config._validate_versions({"invalid_v0": agents}) + config._validate_versions({"invalid_v0": shared}, {"invalid_v0": {AGENT_CODEX: override}}) + +def test_shared_presets_require_all_flags(): + shared = config._VERSIONS[config.SUBAGENT_ORCH_V0].copy() + shared.pop(ENABLE_SMART_ROUTING_ENV_VAR) + with pytest.raises(ValueError, match="missing env vars"): + config._validate_versions({"invalid_v0": shared}, {}) -def test_existing_agent_presets_are_equal_and_independent(): - for version, expected in _EXPECTED_PRESETS.items(): - agents = config._VERSIONS[version] - assert agents[AGENT_CLAUDE] == agents[AGENT_CODEX] == expected - assert agents[AGENT_CLAUDE] is not agents[AGENT_CODEX] + +@pytest.mark.parametrize( + "version, agent", [("unknown_v0", AGENT_CODEX), (config.SUBAGENT_ORCH_V0, "other")] +) +def test_overrides_reject_unknown_versions_and_agents(version, agent): + with pytest.raises(ValueError, match="unknown version or agent"): + config._validate_versions( + config._VERSIONS, {version: {agent: {ENABLE_SMART_ROUTING_ENV_VAR: "0"}}} + ) @pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) @@ -179,12 +179,20 @@ def test_agent_resolution_and_session_controls(monkeypatch, agent, tmp_path): from ucode.smart_routing import session_env version = config.SUBAGENT_ORCH_V0 - agents = deepcopy(config._VERSIONS[version]) - agents[AGENT_CODEX] = dict.fromkeys(agents[AGENT_CODEX], "0") - monkeypatch.setitem(config._VERSIONS, version, agents) + # Codex differs only in subagent routing and orchestration; V2 inherits the shared value. + override = {ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0"} + monkeypatch.setitem(config._AGENT_OVERRIDES, version, {AGENT_CODEX: override}) source = {SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version} - expected = agents[agent] + expected = {**_EXPECTED_PRESETS[version], **(override if agent == AGENT_CODEX else {})} + resolved = config.resolve_environment(source, agent=agent) + assert resolved == expected + resolved[ENABLE_SMART_ROUTING_ENV_VAR] = "1" assert config.resolve_environment(source, agent=agent) == expected + applied = source.copy() + previous = config.apply_config(applied, agent=agent) + assert applied == expected + v2.restore_smart_routing_env(previous, applied) + assert applied == source assert v2.smart_routing_enabled(source, default=True, agent=agent) == (agent == AGENT_CLAUDE) assert orchestrator.feature_enabled(source, agent=agent) == (agent == AGENT_CLAUDE) control = tmp_path / "session.json" From bd164d5cb474e3a695356ea715672db9dd1e2e9f Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:47:01 +0000 Subject: [PATCH 5/9] Map agent routing configs to shared base versions --- src/ucode/smart_routing/config.py | 44 +++++++++---------------- tests/README.md | 4 +-- tests/integration/README.md | 4 +-- tests/test_agents_init.py | 6 ++-- tests/test_cli.py | 8 +++-- tests/test_smart_routing_config.py | 52 ++++++++++++++++-------------- 6 files changed, 58 insertions(+), 60 deletions(-) diff --git a/src/ucode/smart_routing/config.py b/src/ucode/smart_routing/config.py index 7aa69b950..076daddfd 100644 --- a/src/ucode/smart_routing/config.py +++ b/src/ucode/smart_routing/config.py @@ -30,7 +30,7 @@ # Route only subagents with V2, subagent-only, and orchestration all enabled. SUBAGENT_ORCH_V1 = "subagent_orch_v1" -_VERSIONS = { +_BASE_VERSIONS = { FIRST_PROMPT_AND_SUBAGENT_NO_ORCH_V0: { ENABLE_SMART_ROUTING_ENV_VAR: "1", ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", @@ -58,28 +58,23 @@ }, } -# Only intentional agent differences belong here; omitted flags inherit the shared preset. -# Comment each override with the reason that agent differs. -_AGENT_OVERRIDES: dict[str, dict[str, dict[str, str]]] = {} +# Agents share the base preset unless explicitly changed; comment any intentional difference. +_VERSIONS = { + version: {AGENT_CLAUDE: flags, AGENT_CODEX: flags} for version, flags in _BASE_VERSIONS.items() +} -def _validate_versions( - versions: Mapping[str, Mapping[str, str]], - overrides: Mapping[str, Mapping[str, Mapping[str, str]]], -) -> None: - """Require complete shared presets and valid, sparse agent overrides.""" +def _validate_versions(versions: Mapping[str, Mapping[str, Mapping[str, str]]]) -> None: + """Require both agents and a complete, explicit flag set for each.""" + expected_agents = {AGENT_CLAUDE, AGENT_CODEX} expected = set(SMART_ROUTING_ENV_KEYS) - for version, agents in overrides.items(): - if version not in versions or agents.keys() - {AGENT_CLAUDE, AGENT_CODEX}: + for version, agents in versions.items(): + if agents.keys() != expected_agents: raise ValueError( - f"Invalid smart-routing version {version!r}: unknown version or agent." + f"Invalid smart-routing version {version!r}: " + f"expected agents {sorted(expected_agents)}, got {sorted(agents)}." ) - for version, shared in versions.items(): - configs = {"shared": shared} - configs.update( - {agent: {**shared, **values} for agent, values in overrides.get(version, {}).items()} - ) - for agent, values in configs.items(): + for agent, values in agents.items(): missing = expected - values.keys() unexpected = values.keys() - expected if missing or unexpected: @@ -95,21 +90,14 @@ def _validate_versions( ) -_validate_versions(_VERSIONS, _AGENT_OVERRIDES) - - -def _version_config(version: str, *, agent: str) -> dict[str, str] | None: - shared = _VERSIONS.get(version) - if shared is None or agent not in (AGENT_CLAUDE, AGENT_CODEX): - return None - return {**shared, **_AGENT_OVERRIDES.get(version, {}).get(agent, {})} +_validate_versions(_VERSIONS) def resolve_environment(env: Mapping[str, str] | None = None, *, agent: str) -> dict[str, str]: """Expand a version before applying any launch or session-specific overrides.""" resolved = dict(os.environ if env is None else env) version = resolved.pop(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, "").strip() - resolved.update(_version_config(version, agent=agent) or {}) + resolved.update(_VERSIONS.get(version, {}).get(agent, {})) return resolved @@ -119,7 +107,7 @@ def apply_config( """Consume the launch selector, returning the values needed to restore its input.""" target = os.environ if env is None else env version = target.get(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, "").strip() - preset = _version_config(version, agent=agent) + preset = _VERSIONS.get(version, {}).get(agent) if preset is None: return {} keys = (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) diff --git a/tests/README.md b/tests/README.md index a3f98e351..8ab3fc76c 100644 --- a/tests/README.md +++ b/tests/README.md @@ -35,8 +35,8 @@ CUJ4's Claude preset sessions also use it after verifying delegated file tasks: answer does not prove that a resumed child has stopped running. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. -Agent-specific smart-routing preset component checks cover shared existing presets, sparse agent overrides, -complete flag validation, configure order, bootstrap and managed-default launch +Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, +complete per-agent validation, configure order, bootstrap and managed-default launch selection, session controls, and restoration after configuration/launch failures. These checks do not establish live agent coverage. diff --git a/tests/integration/README.md b/tests/integration/README.md index 68069f814..a7b39769a 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -15,8 +15,8 @@ managed-default and explicit-model cases unchanged. One additional test runs the `SMART_ROUTER_CONFIG_VERSION` values, checking both agents' first-prompt routing, orchestrator context in inference input, and completed explicitly requested routed subagents. This is not automatic orchestrator-delegation coverage. No workspace configuration is modified. -Agent-specific smart-routing preset component checks cover shared existing presets, sparse agent overrides, -complete flag validation, configure order, bootstrap and managed-default launch +Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, +complete per-agent validation, configure order, bootstrap and managed-default launch selection, session controls, and restoration after configuration/launch failures. These checks do not establish live agent coverage. diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index e49cc8e55..886ab0166 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -1125,7 +1125,9 @@ def test_configure_dispatch_scopes_agent_presets(monkeypatch, order, fail_write) version = config.SUBAGENT_ORCH_V0 codex_flags = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") - monkeypatch.setitem(config._AGENT_OVERRIDES, version, {AGENT_CODEX: codex_flags}) + monkeypatch.setitem( + config._VERSIONS, version, {**config._VERSIONS[version], AGENT_CODEX: codex_flags} + ) monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) before = { key: os.environ.get(key) @@ -1135,7 +1137,7 @@ def test_configure_dispatch_scopes_agent_presets(monkeypatch, order, fail_write) def write(state, *args, agent=agent, **kwargs): assert {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} == ( - codex_flags if agent == AGENT_CODEX else config._VERSIONS[version] + codex_flags if agent == AGENT_CODEX else config._BASE_VERSIONS[version] ) assert SMART_ROUTER_CONFIG_VERSION_ENV_VAR not in os.environ if fail_write: diff --git a/tests/test_cli.py b/tests/test_cli.py index 80d8c314a..39497e890 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -512,7 +512,11 @@ def test_agent_preset_reaches_bootstrap_and_launch_and_is_restored( ): version = routing_config.SUBAGENT_ORCH_V0 codex_flags = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") - monkeypatch.setitem(routing_config._AGENT_OVERRIDES, version, {AGENT_CODEX: codex_flags}) + monkeypatch.setitem( + routing_config._VERSIONS, + version, + {**routing_config._VERSIONS[version], AGENT_CODEX: codex_flags}, + ) monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) before = { key: os.environ.get(key) @@ -542,7 +546,7 @@ def capture_launch(tool, state, args, *, options): ) result = runner.invoke(app, [] if bare else [agent]) assert result.exit_code == (1 if fail_launch else 0), result.output - expected = codex_flags if agent == AGENT_CODEX else routing_config._VERSIONS[version] + expected = codex_flags if agent == AGENT_CODEX else routing_config._BASE_VERSIONS[version] assert observed == [expected, expected] assert {key: os.environ.get(key) for key in before} == before diff --git a/tests/test_smart_routing_config.py b/tests/test_smart_routing_config.py index ddd2099f6..9b4c90ce4 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -146,32 +146,32 @@ def test_legacy_preset_equivalence(version, expected_routing): ) == expected_routing -@pytest.mark.parametrize("scope", ["shared", "override"]) -@pytest.mark.parametrize("failure", ["unknown_flag", "invalid_value"]) -def test_presets_reject_invalid_flags(scope, failure): - shared = config._VERSIONS[config.SUBAGENT_ORCH_V0].copy() - override = {ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0"} - values = shared if scope == "shared" else override - values["UNKNOWN_FLAG" if failure == "unknown_flag" else ENABLE_SMART_ROUTING_ENV_VAR] = "yes" +@pytest.mark.parametrize( + "failure", ["missing_agent", "unknown_agent", "missing_flag", "unknown_flag", "invalid_value"] +) +def test_agent_presets_require_complete_valid_configuration(failure): + agents = { + agent: flags.copy() for agent, flags in config._VERSIONS[config.SUBAGENT_ORCH_V0].items() + } + if failure == "missing_agent": + agents.pop(AGENT_CODEX) + elif failure == "unknown_agent": + agents["other"] = agents[AGENT_CODEX].copy() + elif failure == "missing_flag": + agents[AGENT_CODEX].pop(ENABLE_SMART_ROUTING_ENV_VAR) + elif failure == "unknown_flag": + agents[AGENT_CODEX]["UNKNOWN_FLAG"] = "0" + else: + agents[AGENT_CODEX][ENABLE_SMART_ROUTING_ENV_VAR] = "yes" with pytest.raises(ValueError, match="Invalid smart-routing version"): - config._validate_versions({"invalid_v0": shared}, {"invalid_v0": {AGENT_CODEX: override}}) - + config._validate_versions({"invalid_v0": agents}) -def test_shared_presets_require_all_flags(): - shared = config._VERSIONS[config.SUBAGENT_ORCH_V0].copy() - shared.pop(ENABLE_SMART_ROUTING_ENV_VAR) - with pytest.raises(ValueError, match="missing env vars"): - config._validate_versions({"invalid_v0": shared}, {}) - -@pytest.mark.parametrize( - "version, agent", [("unknown_v0", AGENT_CODEX), (config.SUBAGENT_ORCH_V0, "other")] -) -def test_overrides_reject_unknown_versions_and_agents(version, agent): - with pytest.raises(ValueError, match="unknown version or agent"): - config._validate_versions( - config._VERSIONS, {version: {agent: {ENABLE_SMART_ROUTING_ENV_VAR: "0"}}} - ) +def test_existing_agent_presets_share_the_base(): + for version, expected in _EXPECTED_PRESETS.items(): + agents = config._VERSIONS[version] + assert agents[AGENT_CLAUDE] == agents[AGENT_CODEX] == expected + assert agents[AGENT_CLAUDE] is agents[AGENT_CODEX] is config._BASE_VERSIONS[version] @pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) @@ -181,7 +181,11 @@ def test_agent_resolution_and_session_controls(monkeypatch, agent, tmp_path): version = config.SUBAGENT_ORCH_V0 # Codex differs only in subagent routing and orchestration; V2 inherits the shared value. override = {ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0"} - monkeypatch.setitem(config._AGENT_OVERRIDES, version, {AGENT_CODEX: override}) + monkeypatch.setitem( + config._VERSIONS, + version, + {**config._VERSIONS[version], AGENT_CODEX: {**config._BASE_VERSIONS[version], **override}}, + ) source = {SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version} expected = {**_EXPECTED_PRESETS[version], **(override if agent == AGENT_CODEX else {})} resolved = config.resolve_environment(source, agent=agent) From 186f8e7f432eacd16f26e800d311673a3cdc534b Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:52:07 +0000 Subject: [PATCH 6/9] Remove redundant agent config test coverage --- tests/README.md | 2 +- tests/integration/README.md | 2 +- tests/test_agents_init.py | 6 +++--- tests/test_smart_routing_config.py | 32 +++++------------------------- 4 files changed, 10 insertions(+), 32 deletions(-) diff --git a/tests/README.md b/tests/README.md index 8ab3fc76c..d6cd4044a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -36,7 +36,7 @@ answer does not prove that a resumed child has stopped running. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, -complete per-agent validation, configure order, bootstrap and managed-default launch +complete per-agent validation, configure dispatch, bootstrap and managed-default launch selection, session controls, and restoration after configuration/launch failures. These checks do not establish live agent coverage. diff --git a/tests/integration/README.md b/tests/integration/README.md index a7b39769a..8b5f3212f 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -16,7 +16,7 @@ managed-default and explicit-model cases unchanged. One additional test runs the context in inference input, and completed explicitly requested routed subagents. This is not automatic orchestrator-delegation coverage. No workspace configuration is modified. Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, -complete per-agent validation, configure order, bootstrap and managed-default launch +complete per-agent validation, configure dispatch, bootstrap and managed-default launch selection, session controls, and restoration after configuration/launch failures. These checks do not establish live agent coverage. diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 886ab0166..30e5b1461 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -1113,10 +1113,10 @@ def test_appends_os_managed_file_recorded_in_state(self): assert "/etc/codex/managed_config.toml" in paths -@pytest.mark.parametrize("order", [("claude", "codex"), ("codex", "claude")]) @pytest.mark.parametrize("fail_write", [False, True]) -def test_configure_dispatch_scopes_agent_presets(monkeypatch, order, fail_write): +def test_configure_dispatch_scopes_agent_presets(monkeypatch, fail_write): from ucode.constants import ( + AGENT_CLAUDE, AGENT_CODEX, SMART_ROUTER_CONFIG_VERSION_ENV_VAR, SMART_ROUTING_ENV_KEYS, @@ -1133,7 +1133,7 @@ def test_configure_dispatch_scopes_agent_presets(monkeypatch, order, fail_write) key: os.environ.get(key) for key in (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) } - for agent in order: + for agent in (AGENT_CLAUDE, AGENT_CODEX): def write(state, *args, agent=agent, **kwargs): assert {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} == ( diff --git a/tests/test_smart_routing_config.py b/tests/test_smart_routing_config.py index 9b4c90ce4..06205700c 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -64,7 +64,6 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): assert applied == original -@pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) @pytest.mark.parametrize("ENABLE_SMART_ROUTING_V2", [None, "0", "1"]) @pytest.mark.parametrize("ENABLE_SMART_ROUTING_SUBAGENT_ONLY", [None, "0", "1"]) @pytest.mark.parametrize("ENABLE_SMART_ROUTER_ORCHESTRATOR", [None, "0", "1"]) @@ -80,7 +79,6 @@ def test_absent_or_unknown_selector_is_a_no_op(selector): ], ) def test_smart_routing_config_cartesian_grid( - agent, ENABLE_SMART_ROUTING_V2, ENABLE_SMART_ROUTING_SUBAGENT_ONLY, ENABLE_SMART_ROUTER_ORCHESTRATOR, @@ -105,13 +103,13 @@ def test_smart_routing_config_cartesian_grid( if SMART_ROUTER_CONFIG_VERSION is not None: expected.update(_EXPECTED_PRESETS[SMART_ROUTER_CONFIG_VERSION]) - resolved = config.resolve_environment(source, agent=agent) + resolved = config.resolve_environment(source, agent=AGENT_CLAUDE) assert resolved == expected assert source == original applied = source.copy() - previous = config.apply_config(applied, agent=agent) + previous = config.apply_config(applied, agent=AGENT_CLAUDE) assert applied == expected v2.restore_smart_routing_env(previous, applied) @@ -167,11 +165,10 @@ def test_agent_presets_require_complete_valid_configuration(failure): config._validate_versions({"invalid_v0": agents}) -def test_existing_agent_presets_share_the_base(): +def test_codex_uses_existing_presets(): for version, expected in _EXPECTED_PRESETS.items(): - agents = config._VERSIONS[version] - assert agents[AGENT_CLAUDE] == agents[AGENT_CODEX] == expected - assert agents[AGENT_CLAUDE] is agents[AGENT_CODEX] is config._BASE_VERSIONS[version] + source = {SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version} + assert config.resolve_environment(source, agent=AGENT_CODEX) == expected @pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) @@ -188,15 +185,7 @@ def test_agent_resolution_and_session_controls(monkeypatch, agent, tmp_path): ) source = {SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version} expected = {**_EXPECTED_PRESETS[version], **(override if agent == AGENT_CODEX else {})} - resolved = config.resolve_environment(source, agent=agent) - assert resolved == expected - resolved[ENABLE_SMART_ROUTING_ENV_VAR] = "1" assert config.resolve_environment(source, agent=agent) == expected - applied = source.copy() - previous = config.apply_config(applied, agent=agent) - assert applied == expected - v2.restore_smart_routing_env(previous, applied) - assert applied == source assert v2.smart_routing_enabled(source, default=True, agent=agent) == (agent == AGENT_CLAUDE) assert orchestrator.feature_enabled(source, agent=agent) == (agent == AGENT_CLAUDE) control = tmp_path / "session.json" @@ -207,14 +196,3 @@ def test_agent_resolution_and_session_controls(monkeypatch, agent, tmp_path): session_env.SESSION_ENV_VAR: str(control), ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", } - - -@pytest.mark.parametrize("resolver", [config.resolve_environment, config.apply_config]) -def test_config_resolution_requires_an_explicit_agent(resolver): - with pytest.raises(TypeError, match="agent"): - resolver({}) - - -def test_smart_routing_check_requires_an_explicit_agent(): - with pytest.raises(TypeError, match="agent"): - v2.smart_routing_enabled(None, default=False) From 5561bb6527ed4d3ce16d6fb5cd9754503dffbe22 Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:54:23 +0000 Subject: [PATCH 7/9] Use agent constants in configure dispatch --- src/ucode/agents/__init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b8869c5dc..1f4ced0f8 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -16,6 +16,7 @@ from pathlib import Path from ucode.config_io import ToolSpec +from ucode.constants import AGENT_CLAUDE, AGENT_CODEX from ucode.databricks import ( AnthropicModelCatalog, get_databricks_token, @@ -415,11 +416,11 @@ def configure_tool( previous = smart_routing_v2.apply_config(agent=tool) try: result: dict | tuple[dict, str] - if tool == "codex": + if tool == AGENT_CODEX: result = codex.write_tool_config( state, model, provider=provider, parent_schema=parent_schema ) - elif tool == "claude": + elif tool == AGENT_CLAUDE: # A Model Provider Service or parent schema routes by header and discovers models natively, # so the usual "model required" guard doesn't apply to either Claude source. if not model and not provider and not parent_schema: From d930fbdaacba9703635ec0381bf6550812855f88 Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 21:59:45 +0000 Subject: [PATCH 8/9] Simplify routing tests to direct configuration and boundary checks --- tests/README.md | 4 +- tests/integration/README.md | 4 +- tests/test_agents_init.py | 45 +++++----------------- tests/test_cli.py | 61 ++++++------------------------ tests/test_smart_routing_config.py | 36 ++++++++++-------- 5 files changed, 45 insertions(+), 105 deletions(-) diff --git a/tests/README.md b/tests/README.md index d6cd4044a..569827229 100644 --- a/tests/README.md +++ b/tests/README.md @@ -36,8 +36,8 @@ answer does not prove that a resumed child has stopped running. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, -complete per-agent validation, configure dispatch, bootstrap and managed-default launch -selection, session controls, and restoration after configuration/launch failures. +complete per-agent validation, session precedence, configuration failure restoration, +and preserving the preset selector until managed-default agent selection. These checks do not establish live agent coverage. CUJ3 and CUJ4 also verify Claude's native recovery from the known thinking-display 400: diff --git a/tests/integration/README.md b/tests/integration/README.md index 8b5f3212f..916e5fb7d 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -16,8 +16,8 @@ managed-default and explicit-model cases unchanged. One additional test runs the context in inference input, and completed explicitly requested routed subagents. This is not automatic orchestrator-delegation coverage. No workspace configuration is modified. Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, -complete per-agent validation, configure dispatch, bootstrap and managed-default launch -selection, session controls, and restoration after configuration/launch failures. +complete per-agent validation, session precedence, configuration failure restoration, +and preserving the preset selector until managed-default agent selection. These checks do not establish live agent coverage. CUJ3 and CUJ4 require a non-empty HTTP 200 for inference. Claude 2.1.290's known diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 30e5b1461..d91953529 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -28,6 +28,7 @@ resolve_launch_model, ) from ucode.agents.args import has_explicit_model_arg +from ucode.constants import AGENT_CODEX from ucode.managed_config import ManagedConfigResult from ucode.ui import redirect_output_to_stderr @@ -1113,41 +1114,15 @@ def test_appends_os_managed_file_recorded_in_state(self): assert "/etc/codex/managed_config.toml" in paths -@pytest.mark.parametrize("fail_write", [False, True]) -def test_configure_dispatch_scopes_agent_presets(monkeypatch, fail_write): - from ucode.constants import ( - AGENT_CLAUDE, - AGENT_CODEX, - SMART_ROUTER_CONFIG_VERSION_ENV_VAR, - SMART_ROUTING_ENV_KEYS, +def test_configure_restores_routing_environment_when_writer_fails(monkeypatch): + monkeypatch.setenv("SMART_ROUTER_CONFIG_VERSION", "subagent_orch_v0") + monkeypatch.setenv("ENABLE_SMART_ROUTING_SUBAGENT_ONLY", "0") + before = os.environ.copy() + monkeypatch.setattr( + agents_mod.codex, "write_tool_config", MagicMock(side_effect=RuntimeError("write failed")) ) - from ucode.smart_routing import config - version = config.SUBAGENT_ORCH_V0 - codex_flags = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") - monkeypatch.setitem( - config._VERSIONS, version, {**config._VERSIONS[version], AGENT_CODEX: codex_flags} - ) - monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) - before = { - key: os.environ.get(key) - for key in (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) - } - for agent in (AGENT_CLAUDE, AGENT_CODEX): + with pytest.raises(RuntimeError, match="write failed"): + agents_mod.configure_tool(AGENT_CODEX, {}, model="model") - def write(state, *args, agent=agent, **kwargs): - assert {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} == ( - codex_flags if agent == AGENT_CODEX else config._BASE_VERSIONS[version] - ) - assert SMART_ROUTER_CONFIG_VERSION_ENV_VAR not in os.environ - if fail_write: - raise RuntimeError("write failed") - return state - - monkeypatch.setattr(agents_mod._MODULES[agent], "write_tool_config", write) - if fail_write: - with pytest.raises(RuntimeError, match="write failed"): - agents_mod.configure_tool(agent, {}, model="model") - else: - assert agents_mod.configure_tool(agent, {}, model="model") == {} - assert {key: os.environ.get(key) for key in before} == before + assert os.environ == before diff --git a/tests/test_cli.py b/tests/test_cli.py index 39497e890..e58df6640 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -24,15 +24,8 @@ import ucode.cli as cli_mod import ucode.databricks as db_mod from ucode.cli import app -from ucode.constants import ( - AGENT_CLAUDE, - AGENT_CODEX, - SMART_ROUTER_CONFIG_VERSION_ENV_VAR, - SMART_ROUTING_ENV_KEYS, -) from ucode.databricks import GatewayProbe from ucode.managed_config import normalize_managed_config -from ucode.smart_routing import config as routing_config _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -504,51 +497,19 @@ def _launch_policy_patches( } -@pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) -@pytest.mark.parametrize("bare", [False, True]) -@pytest.mark.parametrize("fail_launch", [False, True]) -def test_agent_preset_reaches_bootstrap_and_launch_and_is_restored( - monkeypatch, agent, bare, fail_launch -): - version = routing_config.SUBAGENT_ORCH_V0 - codex_flags = dict.fromkeys(SMART_ROUTING_ENV_KEYS, "0") - monkeypatch.setitem( - routing_config._VERSIONS, - version, - {**routing_config._VERSIONS[version], AGENT_CODEX: codex_flags}, - ) - monkeypatch.setenv(SMART_ROUTER_CONFIG_VERSION_ENV_VAR, version) - before = { - key: os.environ.get(key) - for key in (*SMART_ROUTING_ENV_KEYS, SMART_ROUTER_CONFIG_VERSION_ENV_VAR) - } - managed = {"default_agent": agent, "enabled_agents": {agent: {"smart_routing_enabled": True}}} - observed = [] +def test_bare_launch_preserves_preset_until_agent_selection(monkeypatch): + monkeypatch.setenv("SMART_ROUTER_CONFIG_VERSION", "subagent_orch_v0") + monkeypatch.setenv("ENABLE_SMART_ROUTING_SUBAGENT_ONLY", "0") - def capture_launch(tool, state, args, *, options): - observed.append({key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS}) - assert options.launch_smart_routing == (agent == AGENT_CLAUDE) - if fail_launch: - raise RuntimeError("launch failed") + def select_and_launch_agent(*args, **kwargs): + assert os.environ["SMART_ROUTER_CONFIG_VERSION"] == "subagent_orch_v0" + assert os.environ["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] == "0" - with _launch_policy_patches(managed) as calls: - calls["launch"].side_effect = capture_launch - monkeypatch.setattr( - cli_mod, - "ensure_bootstrap_dependencies", - lambda *args, **kwargs: observed.append( - {key: os.environ.get(key) for key in SMART_ROUTING_ENV_KEYS} - ), - ) - monkeypatch.setattr(cli_mod, "install_databricks_cli", lambda *args, **kwargs: None) - monkeypatch.setattr( - cli_mod, "refresh_managed_config", lambda *args, **kwargs: (managed, False) - ) - result = runner.invoke(app, [] if bare else [agent]) - assert result.exit_code == (1 if fail_launch else 0), result.output - expected = codex_flags if agent == AGENT_CODEX else routing_config._BASE_VERSIONS[version] - assert observed == [expected, expected] - assert {key: os.environ.get(key) for key in before} == before + with patch("ucode.cli._launch_managed_default", side_effect=select_and_launch_agent) as launch: + result = runner.invoke(app, []) + + assert result.exit_code == 0, result.output + launch.assert_called_once() class TestSubcommandRouting: diff --git a/tests/test_smart_routing_config.py b/tests/test_smart_routing_config.py index 06205700c..4b081402b 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -12,7 +12,7 @@ ENABLE_SUBAGENT_ROUTING_ENV_VAR, SMART_ROUTER_CONFIG_VERSION_ENV_VAR, ) -from ucode.smart_routing import config, orchestrator, v2 +from ucode.smart_routing import config, orchestrator, session_env, v2 _EXPECTED_PRESETS = { config.FIRST_PROMPT_AND_SUBAGENT_NO_ORCH_V0: { @@ -171,28 +171,32 @@ def test_codex_uses_existing_presets(): assert config.resolve_environment(source, agent=AGENT_CODEX) == expected -@pytest.mark.parametrize("agent", [AGENT_CLAUDE, AGENT_CODEX]) -def test_agent_resolution_and_session_controls(monkeypatch, agent, tmp_path): - from ucode.smart_routing import session_env - - version = config.SUBAGENT_ORCH_V0 - # Codex differs only in subagent routing and orchestration; V2 inherits the shared value. - override = {ENABLE_SUBAGENT_ROUTING_ENV_VAR: "0", ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0"} +@pytest.mark.parametrize("agent, enabled", [(AGENT_CLAUDE, "1"), (AGENT_CODEX, "0")]) +def test_agent_config_selection_and_session_precedence(monkeypatch, agent, enabled, tmp_path): monkeypatch.setitem( - config._VERSIONS, - version, - {**config._VERSIONS[version], AGENT_CODEX: {**config._BASE_VERSIONS[version], **override}}, + config._VERSIONS[config.SUBAGENT_ORCH_V0], + AGENT_CODEX, + { + "ENABLE_SMART_ROUTING_V2": "0", + "ENABLE_SMART_ROUTING_SUBAGENT_ONLY": "0", + "ENABLE_SMART_ROUTER_ORCHESTRATOR": "0", + }, ) - source = {SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version} - expected = {**_EXPECTED_PRESETS[version], **(override if agent == AGENT_CODEX else {})} + source = {"SMART_ROUTER_CONFIG_VERSION": "subagent_orch_v0"} + expected = { + "ENABLE_SMART_ROUTING_V2": "0", + "ENABLE_SMART_ROUTING_SUBAGENT_ONLY": enabled, + "ENABLE_SMART_ROUTER_ORCHESTRATOR": enabled, + } assert config.resolve_environment(source, agent=agent) == expected - assert v2.smart_routing_enabled(source, default=True, agent=agent) == (agent == AGENT_CLAUDE) - assert orchestrator.feature_enabled(source, agent=agent) == (agent == AGENT_CLAUDE) + assert v2.smart_routing_enabled(source, default=True, agent=agent) == (enabled == "1") + assert orchestrator.feature_enabled(source, agent=agent) == (enabled == "1") + control = tmp_path / "session.json" control.write_text('{"ENABLE_SMART_ROUTER_ORCHESTRATOR":"0"}') source[session_env.SESSION_ENV_VAR] = str(control) assert session_env.effective_environment(source, agent=agent) == { **expected, session_env.SESSION_ENV_VAR: str(control), - ENABLE_SMART_ROUTER_ORCHESTRATOR_ENV_VAR: "0", + "ENABLE_SMART_ROUTER_ORCHESTRATOR": "0", } From ede2ef0560ea9bb473f799ca059fb5962e47c423 Mon Sep 17 00:00:00 2001 From: "lilly.luo" Date: Fri, 9 Oct 2026 22:06:02 +0000 Subject: [PATCH 9/9] Remove redundant environment mutation around configuration --- src/ucode/agents/__init__.py | 81 +++++++++++++++++------------------- tests/README.md | 4 +- tests/integration/README.md | 4 +- tests/test_agents_init.py | 15 ------- 4 files changed, 42 insertions(+), 62 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 1f4ced0f8..f273e25d7 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -28,7 +28,6 @@ from ucode.managed_config import refresh_managed_config from ucode.managed_files import managed_write_batch from ucode.os_compatibility import subprocess_cross_os -from ucode.smart_routing import v2 as smart_routing_v2 from ucode.state import get_provider_service, load_state, save_state from ucode.telemetry import agent_version from ucode.ui import ( @@ -413,49 +412,45 @@ def configure_tool( parent_schema: str | None = None, picker_catalog: AnthropicModelCatalog | None = None, ) -> dict: - previous = smart_routing_v2.apply_config(agent=tool) - try: - result: dict | tuple[dict, str] - if tool == AGENT_CODEX: - result = codex.write_tool_config( - state, model, provider=provider, parent_schema=parent_schema - ) - elif tool == AGENT_CLAUDE: - # A Model Provider Service or parent schema routes by header and discovers models natively, - # so the usual "model required" guard doesn't apply to either Claude source. - if not model and not provider and not parent_schema: - raise RuntimeError(f"A {tool} model must be selected before configuration.") - result = claude.write_tool_config( - state, - model, - provider=provider, - provider_models=provider_models, - relayed=relayed, - route_root_model=route_root_model, - custom_model=custom_model, - coding_agent_config_defaults=coding_agent_config_defaults, - parent_schema=parent_schema, - picker_catalog=picker_catalog, - ) + result: dict | tuple[dict, str] + if tool == AGENT_CODEX: + result = codex.write_tool_config( + state, model, provider=provider, parent_schema=parent_schema + ) + elif tool == AGENT_CLAUDE: + # A Model Provider Service or parent schema routes by header and discovers models natively, + # so the usual "model required" guard doesn't apply to either Claude source. + if not model and not provider and not parent_schema: + raise RuntimeError(f"A {tool} model must be selected before configuration.") + result = claude.write_tool_config( + state, + model, + provider=provider, + provider_models=provider_models, + relayed=relayed, + route_root_model=route_root_model, + custom_model=custom_model, + coding_agent_config_defaults=coding_agent_config_defaults, + parent_schema=parent_schema, + picker_catalog=picker_catalog, + ) + else: + # Every tool in this branch needs a model — including gemini under a provider, + # which still pins the service's target model in the URL. + if not model: + raise RuntimeError(f"A {tool} model must be selected before configuration.") + if tool == "gemini": + result = gemini.write_tool_config(state, model, provider=provider) + elif tool == "copilot": + result = copilot.write_tool_config(state, model) + elif tool == "pi": + result = pi.write_tool_config(state, model) else: - # Every tool in this branch needs a model — including gemini under a provider, - # which still pins the service's target model in the URL. - if not model: - raise RuntimeError(f"A {tool} model must be selected before configuration.") - if tool == "gemini": - result = gemini.write_tool_config(state, model, provider=provider) - elif tool == "copilot": - result = copilot.write_tool_config(state, model) - elif tool == "pi": - result = pi.write_tool_config(state, model) - else: - result = opencode.write_tool_config(state, model) - # gemini/opencode/copilot/pi return (state, token); codex/claude return state - if isinstance(result, tuple): - return result[0] - return result - finally: - smart_routing_v2.restore_smart_routing_env(previous) + result = opencode.write_tool_config(state, model) + # gemini/opencode/copilot/pi return (state, token); codex/claude return state + if isinstance(result, tuple): + return result[0] + return result def configured_paths(tool: str, state: dict) -> list[str]: diff --git a/tests/README.md b/tests/README.md index 569827229..e4d8d86da 100644 --- a/tests/README.md +++ b/tests/README.md @@ -36,8 +36,8 @@ answer does not prove that a resumed child has stopped running. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, -complete per-agent validation, session precedence, configuration failure restoration, -and preserving the preset selector until managed-default agent selection. +complete per-agent validation, session precedence, and preserving the preset selector +until managed-default agent selection. These checks do not establish live agent coverage. CUJ3 and CUJ4 also verify Claude's native recovery from the known thinking-display 400: diff --git a/tests/integration/README.md b/tests/integration/README.md index 916e5fb7d..85e63db89 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -16,8 +16,8 @@ managed-default and explicit-model cases unchanged. One additional test runs the context in inference input, and completed explicitly requested routed subagents. This is not automatic orchestrator-delegation coverage. No workspace configuration is modified. Agent-specific smart-routing preset component checks cover shared base presets, explicit agent mappings, -complete per-agent validation, session precedence, configuration failure restoration, -and preserving the preset selector until managed-default agent selection. +complete per-agent validation, session precedence, and preserving the preset selector +until managed-default agent selection. These checks do not establish live agent coverage. CUJ3 and CUJ4 require a non-empty HTTP 200 for inference. Claude 2.1.290's known diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index d91953529..fd234e8b4 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -28,7 +28,6 @@ resolve_launch_model, ) from ucode.agents.args import has_explicit_model_arg -from ucode.constants import AGENT_CODEX from ucode.managed_config import ManagedConfigResult from ucode.ui import redirect_output_to_stderr @@ -1112,17 +1111,3 @@ def test_appends_os_managed_file_recorded_in_state(self): paths = configured_paths("codex", state) assert str(CODEX_CONFIG_PATH).replace(str(CODEX_CONFIG_PATH.home()), "~", 1) in paths assert "/etc/codex/managed_config.toml" in paths - - -def test_configure_restores_routing_environment_when_writer_fails(monkeypatch): - monkeypatch.setenv("SMART_ROUTER_CONFIG_VERSION", "subagent_orch_v0") - monkeypatch.setenv("ENABLE_SMART_ROUTING_SUBAGENT_ONLY", "0") - before = os.environ.copy() - monkeypatch.setattr( - agents_mod.codex, "write_tool_config", MagicMock(side_effect=RuntimeError("write failed")) - ) - - with pytest.raises(RuntimeError, match="write failed"): - agents_mod.configure_tool(AGENT_CODEX, {}, model="model") - - assert os.environ == before