diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 2a5c657..2bd4627 100644 --- a/src/ucode/agents/codex.py +++ b/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]: diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 9a68e03..453666a 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,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 == []