codex: retry without --profile when the first attempt fails fast

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]:

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 == []