diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index f0f10429..69dc261b 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -525,7 +525,7 @@ def launch( def _launch_smart_routing(state: dict, tool_args: list[str]) -> None: - """Launch the Codex TUI through the smart-routing interposer.""" + """Launch Codex with smart-routing configuration.""" clear_model_preferences(state) binary = SPEC["binary"] version_text = agent_version(binary) @@ -543,6 +543,15 @@ def _launch_smart_routing(state: dict, tool_args: list[str]) -> None: or (codex_model_id(models[0]) if models else None) or APP_SERVER_SMART_ROUTING_STARTING_MODEL ) + if tool_args[:1] == ["app-server"]: + config_args, _ = smart_routing_v2.prepare_codex_routing( + state, + start_model=start_model, + render_overlay=render_overlay, + ) + exec_or_spawn([binary, "app-server", *config_args, *tool_args[1:]]) + return # unreachable in production (exec replaces the process) + smart_routing_v2.launch_codex( state, tool_args, diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0ebd20ec..50c095b8 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1924,6 +1924,11 @@ def _should_launch_smart_routing( ) -> bool: if model is not None or has_explicit_model_arg(tool_args): return False + if tool == "codex" and ( + tool_args[:1] == ["app-server"] + or any(arg == "--remote" or arg.startswith("--remote=") for arg in tool_args) + ): + return True if not tool_args or explicit_prompt: return True return tool == "claude" and tool_args[0].startswith("-") diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index e8ce72c6..f3c874c5 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -486,14 +486,13 @@ def _v2_pre_tool_use_hooks(state: dict, available_models: list[str]) -> list[dic ) -def launch_codex( +def prepare_codex_routing( state: dict, - tool_args: list[str], *, - binary: str, - start_model: str | None, + start_model: str, render_overlay: Callable[..., dict], -) -> NoReturn: +) -> tuple[list[str], list[str]]: + """Prepare the CLI config and model catalog shared by Codex routing modes.""" workspace = state.get("workspace") if not workspace: raise RuntimeError( @@ -521,25 +520,76 @@ def launch_codex( overlay["hooks"] = { "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), } - config_args = codex_config_args(overlay) - app_port = _free_port() - app_server_url = _loopback_websocket_url(app_port) - - # Preserve the user's normal CODEX_HOME (including MCP servers, skills, and - # preferences) and layer only ucode's gateway settings at CLI precedence. - app_server = subprocess.Popen( - [binary, "app-server", *config_args, "--listen", app_server_url], - env=os.environ.copy(), - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, + return codex_config_args(overlay), available_models + + +def _take_remote_arg(tool_args: list[str]) -> tuple[str | None, list[str]]: + """Remove Codex's remote option while preserving all other argument ordering.""" + remote_url = None + remaining: list[str] = [] + index = 0 + while index < len(tool_args): + arg = tool_args[index] + if arg == "--": + remaining.extend(tool_args[index:]) + break + if arg == "--remote" and index + 1 < len(tool_args) and tool_args[index + 1] != "--": + remote_url = tool_args[index + 1] + index += 2 + continue + if arg.startswith("--remote="): + remote_url = arg.partition("=")[2] + index += 1 + continue + remaining.append(arg) + index += 1 + return remote_url, remaining + + +def launch_codex( + state: dict, + tool_args: list[str], + *, + binary: str, + start_model: str, + render_overlay: Callable[..., dict], +) -> NoReturn: + workspace = state.get("workspace") + if not isinstance(workspace, str) or not workspace: + raise RuntimeError( + "Smart routing v2 needs a configured workspace; run `ucode configure codex` first." + ) + profile = state.get("profile") + config_args, available_models = prepare_codex_routing( + state, + start_model=start_model, + render_overlay=render_overlay, ) + external_app_server_url, tui_args = _take_remote_arg(tool_args) + + # With no caller-owned remote, preserve the user's normal CODEX_HOME + # (including MCP servers, skills, and preferences) and layer only ucode's + # gateway settings at CLI precedence on the app-server we create. + app_server = None stop_interposer = None try: - if not _wait_for_app_server(app_port, timeout=APP_SERVER_READY_TIMEOUT_SECONDS): - raise RuntimeError( - "Codex app-server did not become ready for smart routing v2; check workspace auth." + if external_app_server_url: + app_server_url = external_app_server_url + else: + app_port = _free_port() + app_server_url = _loopback_websocket_url(app_port) + app_server = subprocess.Popen( + [binary, "app-server", *config_args, "--listen", app_server_url], + env=os.environ.copy(), + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, ) + if not _wait_for_app_server(app_port, timeout=APP_SERVER_READY_TIMEOUT_SECONDS): + raise RuntimeError( + "Codex app-server did not become ready for smart routing v2; " + "check workspace auth." + ) tui_port, stop_interposer = codex_interposer.start_interposer_thread( LOOPBACK_HOST, app_server_url, @@ -550,7 +600,7 @@ def launch_codex( log_path=CODEX_INTERPOSER_LOG, ) tui_url = _loopback_websocket_url(tui_port) - tui = subprocess.Popen([binary, "--remote", tui_url, "--model", start_model, *tool_args]) + tui = subprocess.Popen([binary, "--remote", tui_url, "--model", start_model, *tui_args]) try: returncode = tui.wait() except KeyboardInterrupt: @@ -559,9 +609,10 @@ def launch_codex( finally: if stop_interposer is not None: stop_interposer() - app_server.terminate() - try: - app_server.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT_SECONDS) - except Exception: # noqa: BLE001 - app_server.kill() + if app_server is not None: + app_server.terminate() + try: + app_server.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT_SECONDS) + except Exception: # noqa: BLE001 + app_server.kill() sys.exit(returncode) diff --git a/tests/test_cli.py b/tests/test_cli.py index 11e4aa09..fe625ffd 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -522,6 +522,26 @@ def test_codex_and_claude_share_smart_routing_policy( assert options.launch_smart_routing is expected + @pytest.mark.parametrize( + "tool_args", + [ + ["app-server", "--listen", "ws://127.0.0.1:7107"], + ["--remote", "ws://127.0.0.1:7107"], + ["--remote=ws://127.0.0.1:7107"], + ], + ) + def test_codex_options_enable_smart_routing_for_openui_launch_shapes(self, tool_args): + options = cli_mod._launch_options( + "codex", + tool_args, + smart_routing_enabled=True, + explicit_prompt=False, + model=None, + provider=None, + ) + + assert options.launch_smart_routing is True + @pytest.mark.parametrize( ("tool_args", "expected"), [ diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index fbc61e7d..9aec22ab 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -122,6 +122,58 @@ def launch_v2(state, tool_args, **kwargs): assert calls[0]["start_model"] == "gpt-5.6-luna" + def test_codex_app_server_uses_requested_listener_instead_of_smart_routing_stack( + self, monkeypatch + ): + prepared = [] + launches = [] + monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.148.0") + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setattr(codex, "default_model", lambda state: "gpt-start") + monkeypatch.setattr( + v2, + "launch_codex", + lambda *args, **kwargs: pytest.fail("app-server entered the managed TUI stack"), + ) + + def prepare(state, **kwargs): + prepared.append((state, kwargs)) + return ["--config", 'model_provider="ucode-databricks"'], [] + + monkeypatch.setattr(v2, "prepare_codex_routing", prepare) + monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) + state = {"workspace": WS} + requested_args = [ + "app-server", + "--listen", + "ws://127.0.0.1:7107", + "-c", + "mcp_servers.openui.enabled=true", + ] + + codex.launch( + state, + requested_args, + options=LaunchOptions(launch_smart_routing=True), + ) + + assert prepared == [ + ( + state, + {"start_model": "gpt-start", "render_overlay": codex.render_overlay}, + ) + ] + assert launches == [ + [ + "codex", + "app-server", + "--config", + 'model_provider="ucode-databricks"', + *requested_args[1:], + ] + ] + def test_owns_app_server_interposer_and_tui_lifecycle(self, monkeypatch): processes = [] interposer_args = {} @@ -228,6 +280,71 @@ def start_interposer(*args, **kwargs): assert stopped == [True] assert processes[0].terminated is True + def test_reuses_external_app_server_and_starts_only_interposer_and_tui(self, monkeypatch): + processes = [] + interposer_targets = [] + stopped = [] + external_url = "ws://127.0.0.1:7107" + + class FakeProcess: + def __init__(self, argv, **kwargs): + self.argv = argv + self.kwargs = kwargs + processes.append(self) + + def wait(self, timeout=None): + assert timeout is None + return 9 + + def send_signal(self, _signal): + raise AssertionError("test does not interrupt the TUI") + + monkeypatch.setattr(v2.subprocess, "Popen", FakeProcess) + monkeypatch.setattr(v2, "get_databricks_token", lambda workspace, profile: "token") + monkeypatch.setattr( + v2, + "_free_port", + lambda: pytest.fail("external remote must not allocate an app-server port"), + ) + monkeypatch.setattr( + v2, + "_wait_for_app_server", + lambda *args, **kwargs: pytest.fail("external app-server is owned by the caller"), + ) + + def start_interposer(host, target, **kwargs): + interposer_targets.append((host, target, kwargs)) + return 41002, lambda: stopped.append(True) + + monkeypatch.setattr(codex_interposer, "start_interposer_thread", start_interposer) + + with pytest.raises(SystemExit) as exc: + v2.launch_codex( + {"workspace": WS, "codex_models": ["system.ai.gpt-5-6-sol"]}, + ["--remote", external_url, "--search", "--", "keep the prompt intact"], + binary="codex", + start_model="gpt-start", + render_overlay=codex.render_overlay, + ) + + assert exc.value.code == 9 + assert len(processes) == 1 + assert processes[0].argv == [ + "codex", + "--remote", + "ws://127.0.0.1:41002", + "--model", + "gpt-start", + "--search", + "--", + "keep the prompt intact", + ] + assert ( + sum(arg == "--remote" or arg.startswith("--remote=") for arg in processes[0].argv) == 1 + ) + assert interposer_targets[0][0:2] == (v2.LOOPBACK_HOST, external_url) + assert stopped == [True] + def test_v2_pre_tool_hook_preserves_user_hooks(self, tmp_path, monkeypatch): codex_home = tmp_path / ".codex" codex_home.mkdir()