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..f273e25d7 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, @@ -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: 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 1ff680c14..1d591381f 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(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 diff --git a/src/ucode/cli.py b/src/ucode/cli.py index d6dbfd15f..11f53876e 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), 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", 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 2a7fe9c01..076daddfd 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, @@ -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) diff --git a/src/ucode/smart_routing/orchestrator.py b/src/ucode/smart_routing/orchestrator.py index a09aecb22..8bc6afe16 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) -> 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) -> 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,9 @@ 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), default=False, agent=agent + ) def skill_directory() -> Path: @@ -59,13 +61,13 @@ 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) 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 +92,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..1432601ef 100644 --- a/src/ucode/smart_routing/session_env.py +++ b/src/ucode/smart_routing/session_env.py @@ -52,9 +52,9 @@ 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) -> 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..b03d4d97b 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, *, default: bool, agent: str ) -> 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 +) -> 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" @@ -175,20 +179,22 @@ def first_prompt_routing_enabled(env: MutableMapping[str, str] | None = None) -> 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, ) -> 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: @@ -423,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, @@ -524,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": @@ -541,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: @@ -682,7 +690,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..e4d8d86da 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 shared base presets, explicit agent mappings, +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: 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..85e63db89 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 shared base presets, explicit agent mappings, +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 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_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 b95af40d1..e58df6640 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -497,6 +497,21 @@ def _launch_policy_patches( } +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 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 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: @pytest.mark.parametrize("tool", TOOLS) def test_subcommand_calls_correct_tool(self, tool): @@ -946,7 +961,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 c797902a0..4b081402b 100644 --- a/tests/test_smart_routing_config.py +++ b/tests/test_smart_routing_config.py @@ -5,12 +5,14 @@ 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, 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: { @@ -56,9 +58,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 @@ -101,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) + resolved = config.resolve_environment(source, agent=AGENT_CLAUDE) assert resolved == expected assert source == original applied = source.copy() - previous = config.apply_config(applied) + previous = config.apply_config(applied, agent=AGENT_CLAUDE) assert applied == expected v2.restore_smart_routing_env(previous, applied) @@ -129,14 +131,72 @@ 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 + + +@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": agents}) + + +def test_codex_uses_existing_presets(): + for version, expected in _EXPECTED_PRESETS.items(): + source = {SMART_ROUTER_CONFIG_VERSION_ENV_VAR: version} + assert config.resolve_environment(source, agent=AGENT_CODEX) == expected + + +@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[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": "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) == (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": "0", + }