diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 2a5c657..bf4675e 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -4,6 +4,9 @@ import os import re +import subprocess +import sys +import time from pathlib import Path from ucode.agent_updates import available_npm_package_update @@ -391,12 +394,48 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]]) return max(parsed, key=_gpt_version_key)[0] +# codex rejects the global --profile on subcommands that don't accept it +# (app-server, mcp-server, ...) with a CLI *parse-time* error — before it touches +# auth, the gateway, or the network — so the rejection exits almost instantly. +# We use that to decide when to retry without --profile (see launch()). This +# window is well above codex's ~0.15s cold-start floor and far below the seconds +# any real session needs to connect and then fail, so it never catches a genuine +# failure. Its exit code (1) is indistinguishable from an ordinary failure, so +# elapsed time is the signal we key on rather than stderr text. +_PROFILE_REJECTED_MAX_SECONDS = 3.0 + + def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) - exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args]) + # Run codex with --profile first — the TUI and runtime subcommands + # (exec/resume/mcp/...) keep ucode's Databricks routing, including any added + # by future codex versions. codex rejects the global --profile on + # server-family subcommands (app-server, mcp-server, ...), which are + # caller-configured anyway (e.g. omnigent runs `codex app-server` with its + # own CODEX_HOME); on that rejection we relaunch without --profile. + # + # The retry is gated on the attempt failing *fast*: the rejection is a + # parse-time error (~0.15s), whereas a session that actually starts can only + # fail after a network round-trip (seconds). Without that gate a genuinely + # failing `codex exec` would be silently re-run without --profile — i.e. on + # the user's own OpenAI login instead of the Databricks gateway (ucode writes + # a *named-profile* file, so no --profile means no ucode routing). stdio is + # inherited (no capture), so Ctrl-C reaches codex directly and the resulting + # KeyboardInterrupt propagates past the retry check — quitting an interactive + # session is never mistaken for a --profile rejection. + started = time.monotonic() + returncode = subprocess.run([binary, "--profile", CODEX_PROFILE_NAME, *tool_args]).returncode + if returncode != 0 and time.monotonic() - started < _PROFILE_REJECTED_MAX_SECONDS: + # Fast failure: most likely codex rejected --profile on this subcommand. + # Relaunch without it, handing over the terminal. (A fast failure for + # any other reason — e.g. a bad flag — just re-fails the same way here, + # with no ucode routing to lose since the subcommand had none.) + exec_or_spawn([binary, *tool_args]) + return # unreachable in production (exec replaces the process) + sys.exit(returncode) def validate_cmd(binary: str) -> list[str]: diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 9a68e03..a3cddb8 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -4,6 +4,8 @@ import os +import pytest + from ucode.agents import codex from ucode.config_io import read_toml_safe @@ -416,23 +418,74 @@ def test_skips_git_repo_check(self): class TestCodexLaunch: - def test_sets_oauth_token_and_ucode_profile_before_exec(self, monkeypatch): - exec_calls: list[tuple[str, list[str]]] = [] - - def fake_execvp(binary: str, args: list[str]) -> None: - exec_calls.append((binary, args)) - raise RuntimeError("stop") - + """launch() runs codex with --profile first and relaunches without it only + when that attempt fails *fast* — codex's --profile rejection is a parse-time + error, so a fast nonzero exit means the subcommand didn't accept --profile. + A slow failure is a real session error and is propagated unchanged.""" + + @staticmethod + def _patch(monkeypatch, *, returncode: int, elapsed: float): + """Stub subprocess.run to return `returncode` and make launch() perceive + `elapsed` seconds between its two time.monotonic() reads.""" + runs: list[list[str]] = [] + fallbacks: list[list[str]] = [] + + def fake_run(argv, **kwargs): + runs.append(argv) + return codex.subprocess.CompletedProcess(argv, returncode) + + # launch() reads time.monotonic() once before run and once after. + clock = iter([100.0, 100.0 + elapsed]) + monkeypatch.setattr(codex.subprocess, "run", fake_run) + monkeypatch.setattr(codex.time, "monotonic", lambda: next(clock)) + monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: fallbacks.append(argv)) + monkeypatch.setattr(codex, "get_databricks_token", lambda workspace, profile=None: "tok") + return runs, fallbacks + + def test_sets_oauth_token_and_runs_with_profile(self, monkeypatch): monkeypatch.delenv("OAUTH_TOKEN", raising=False) + runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.5) monkeypatch.setattr( codex, "get_databricks_token", lambda workspace, profile=None: "fresh-token" ) - monkeypatch.setattr(os, "execvp", fake_execvp) - - try: + with pytest.raises(SystemExit) as exc: codex.launch({"workspace": WS}, ["--search"]) - except RuntimeError as exc: - assert str(exc) == "stop" - + assert exc.value.code == 0 assert os.environ["OAUTH_TOKEN"] == "fresh-token" - assert exec_calls == [("codex", ["codex", "--profile", "ucode", "--search"])] + assert runs == [["codex", "--profile", "ucode", "--search"]] + assert fallbacks == [] + + def test_success_propagates_exit_without_retry(self, monkeypatch): + runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.2) + with pytest.raises(SystemExit) as exc: + codex.launch({"workspace": WS}, ["exec", "hi"]) + assert exc.value.code == 0 + assert runs == [["codex", "--profile", "ucode", "exec", "hi"]] + assert fallbacks == [] + + def test_fast_failure_relaunches_without_profile(self, monkeypatch): + # codex rejects --profile on server-family subcommands at parse time — + # a fast nonzero exit → relaunch without --profile. + for args in (["app-server", "--listen", "u"], ["mcp-server"]): + runs, fallbacks = self._patch(monkeypatch, returncode=1, elapsed=0.15) + codex.launch({"workspace": WS}, args) + assert runs == [["codex", "--profile", "ucode", *args]] + assert fallbacks == [["codex", *args]] + + def test_slow_failure_does_not_retry(self, monkeypatch): + # A session that started and then failed (seconds) must NOT be re-run + # without --profile — that would silently drop ucode's Databricks routing + # (relaunching the user's prompt on their own OpenAI login). + runs, fallbacks = self._patch(monkeypatch, returncode=1, elapsed=8.0) + with pytest.raises(SystemExit) as exc: + codex.launch({"workspace": WS}, ["exec", "hi"]) + assert exc.value.code == 1 + assert fallbacks == [] + + def test_fast_success_does_not_retry(self, monkeypatch): + # The retry is gated on a *nonzero* exit; a fast clean exit just returns. + runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.15) + with pytest.raises(SystemExit) as exc: + codex.launch({"workspace": WS}, []) + assert exc.value.code == 0 + assert fallbacks == []