Resolve smart-routing presets separately for Claude and Codex

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,

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,
@@ -412,11 +413,11 @@ def configure_tool(
     picker_catalog: AnthropicModelCatalog | None = None,
 ) -> dict:
     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:

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.

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(None, default=False, agent=AGENT_CODEX):
         http_headers[SMART_ROUTER_RECIPE_HEADER] = configured_router_name()
     _apply_managed_headers(http_headers, managed_http_headers)
     return {
@@ -485,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():
+        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:
@@ -814,7 +817,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(None, default=False, agent=AGENT_CODEX):
         return _smart_routing_config_model(state)
     return None
 
@@ -842,7 +845,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(None, default=False, agent=AGENT_CODEX):
         return False
     if isinstance(state.get("codex_default_model"), str):
         return False

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), default=False, 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), default=False, 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
@@ -2935,7 +2925,7 @@ def _launch_tool(
         # 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)
@@ -3200,8 +3190,9 @@ def _launch_tool(
             True
             if managed_smart_routing_enabled
             and smart_routing_enabled
-            and not smart_routing_v2.smart_routing_enabled()
-            else None
+            and not smart_routing_v2.smart_routing_enabled(None, default=False, agent=tool)
+            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",

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,

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,
@@ -28,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",
@@ -56,42 +58,56 @@
     },
 }
 
+# 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]]) -> 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) -> 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
+) -> 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)