codex: try --profile, fall back on codex's rejection error

src/ucode/agents/codex.py

@@ -4,6 +4,10 @@
 
 import os
 import re
+import signal
+import subprocess
+import sys
+import tempfile
 from pathlib import Path
 
 from ucode.agent_updates import available_npm_package_update
@@ -391,12 +395,49 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]])
     return max(parsed, key=_gpt_version_key)[0]
 
 
+# Stderr prefix codex prints (and exits nonzero) when --profile is passed to a
+# subcommand that doesn't accept it. Stable across codex 0.137 and 0.141.
+_PROFILE_REJECTED_MARKER = "--profile only applies to runtime commands"
+_STDERR_HEAD_LIMIT = 64 * 1024  # bytes of stderr to scan for the marker
+
+
 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])
+    # Spawn codex with --profile first — the TUI and runtime subcommands
+    # (exec/resume/mcp/...) keep ucode's routing, including subcommands 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
+    # exactly that rejection, relaunch without --profile. Any other failure
+    # propagates unchanged, so an unrelated error is never silently retried
+    # without ucode's routing. codex's stderr goes to a temp file — unlike a
+    # pipe it has no backpressure, so codex can't block on it — read once codex
+    # exits (the trade: its stderr surfaces at exit, not live).
+    with tempfile.TemporaryFile() as stderr_file:
+        proc = subprocess.Popen(
+            [binary, "--profile", CODEX_PROFILE_NAME, *tool_args], stderr=stderr_file
+        )
+        try:
+            returncode = proc.wait()
+        except KeyboardInterrupt:
+            # Forward Ctrl-C and let codex set its own exit code (mirrors
+            # exec_or_spawn's Windows path) rather than racing it.
+            proc.send_signal(signal.SIGINT)
+            returncode = proc.wait()
+        stderr_file.seek(0)
+        head = stderr_file.read(_STDERR_HEAD_LIMIT)
+        if returncode != 0 and _PROFILE_REJECTED_MARKER.encode() in head:
+            exec_or_spawn([binary, *tool_args])  # swallow the expected rejection
+            return  # unreachable in production (exec/exit); keeps tests honest
+        # Replay codex's stderr and propagate its exit code.
+        stderr_file.seek(0)
+        for chunk in iter(lambda: stderr_file.read(65536), b""):
+            sys.stderr.write(chunk.decode(errors="replace"))
+        sys.stderr.flush()
+    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,85 @@ 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]]] = []
+    """launch() spawns codex with --profile first, and relaunches without it only
+    on codex's --profile-rejection error. A fake Popen stands in for codex: its
+    wait() writes the given stderr into the temp file launch() reads."""
+
+    @staticmethod
+    def _patch(monkeypatch, *, returncode: int, stderr: bytes = b""):
+        spawned: list[list[str]] = []
+        fallbacks: list[list[str]] = []
+
+        class _Proc:
+            def __init__(self, stderr_file):
+                self._stderr_file = stderr_file
+
+            def wait(self):
+                self._stderr_file.write(stderr)
+                self._stderr_file.flush()
+                return returncode
+
+            def send_signal(self, sig):  # pragma: no cover - interrupt not exercised
+                pass
 
-        def fake_execvp(binary: str, args: list[str]) -> None:
-            exec_calls.append((binary, args))
-            raise RuntimeError("stop")
+        def fake_popen(argv, stderr=None, **kwargs):
+            spawned.append(argv)
+            return _Proc(stderr)
 
+        monkeypatch.setattr(codex.subprocess, "Popen", fake_popen)
+        monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: fallbacks.append(argv))
+        monkeypatch.setattr(codex, "get_databricks_token", lambda workspace, profile=None: "tok")
+        return spawned, fallbacks
+
+    def test_sets_oauth_token_and_spawns_with_profile(self, monkeypatch):
         monkeypatch.delenv("OAUTH_TOKEN", raising=False)
+        spawned, fallbacks = self._patch(monkeypatch, returncode=0)
         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 spawned == [["codex", "--profile", "ucode", "--search"]]
+        assert fallbacks == []
+
+    def test_success_propagates_exit_and_replays_stderr(self, monkeypatch, capsys):
+        spawned, fallbacks = self._patch(monkeypatch, returncode=0, stderr=b"warn: heads up\n")
+        with pytest.raises(SystemExit) as exc:
+            codex.launch({"workspace": WS}, ["exec", "hi"])
+        assert exc.value.code == 0
+        assert spawned == [["codex", "--profile", "ucode", "exec", "hi"]]
+        assert fallbacks == []
+        assert "warn: heads up" in capsys.readouterr().err
+
+    def test_profile_rejection_relaunches_without_profile_and_swallows_it(
+        self, monkeypatch, capsys
+    ):
+        marker = f"Error: {codex._PROFILE_REJECTED_MARKER} and `codex mcp`\n".encode()
+        for args in (["app-server", "--listen", "u"], ["mcp-server"]):
+            spawned, fallbacks = self._patch(monkeypatch, returncode=1, stderr=marker)
+            codex.launch({"workspace": WS}, args)
+            assert spawned == [["codex", "--profile", "ucode", *args]]
+            assert fallbacks == [["codex", *args]]
+            # The expected rejection is noise once we relaunch — swallowed.
+            assert capsys.readouterr().err == ""
+
+    def test_unrelated_failure_does_not_fall_back(self, monkeypatch, capsys):
+        # An unmarked failure must NOT retry without --profile — that would
+        # silently drop ucode's routing.
+        spawned, fallbacks = self._patch(monkeypatch, returncode=3, stderr=b"boom\n")
+        with pytest.raises(SystemExit) as exc:
+            codex.launch({"workspace": WS}, ["exec", "hi"])
+        assert exc.value.code == 3
+        assert fallbacks == []
+        assert "boom" in capsys.readouterr().err
+
+    def test_marker_on_zero_exit_does_not_fall_back(self, monkeypatch, capsys):
+        # The marker only triggers the fallback paired with a nonzero exit.
+        marker = f"{codex._PROFILE_REJECTED_MARKER}\n".encode()
+        spawned, fallbacks = self._patch(monkeypatch, returncode=0, stderr=marker)
+        with pytest.raises(SystemExit) as exc:
+            codex.launch({"workspace": WS}, [])
+        assert exc.value.code == 0
+        assert fallbacks == []