codex: allow-list --profile by subcommand; omit it where codex rejects it

src/ucode/agents/codex.py

@@ -391,12 +391,84 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]])
     return max(parsed, key=_gpt_version_key)[0]
 
 
+# codex global options that consume the following token as their value. Used to
+# skip past a leading global (e.g. ``--model X``) when locating the subcommand,
+# so it isn't mistaken for the subcommand itself. Boolean flags take no value;
+# ``--flag=value`` forms are self-contained and handled inline.
+_CODEX_GLOBAL_VALUE_FLAGS = frozenset(
+    {
+        "-c",
+        "--config",
+        "-i",
+        "--image",
+        "-m",
+        "--model",
+        "-p",
+        "--profile",
+        "-s",
+        "--sandbox",
+        "-C",
+        "--cd",
+        "-a",
+        "--ask-for-approval",
+        "--enable",
+        "--disable",
+        "--remote",
+        "--remote-auth-token-env",
+        "--local-provider",
+        "--add-dir",
+    }
+)
+
+
+def _codex_subcommand(tool_args: list[str]) -> str | None:
+    """Return the codex subcommand in *tool_args*, skipping global options.
+
+    ``None`` when there is no subcommand (the interactive TUI).
+    """
+    i = 0
+    while i < len(tool_args):
+        arg = tool_args[i]
+        if not arg.startswith("-"):
+            return arg
+        # A bare value flag (``--model X``) consumes the next token too; a
+        # ``--flag=value`` form or a boolean flag consumes only itself.
+        if "=" not in arg and arg in _CODEX_GLOBAL_VALUE_FLAGS:
+            i += 1
+        i += 1
+    return None
+
+
+# codex accepts the global --profile ONLY on these runtime subcommands (plus the
+# bare interactive TUI). Server-family subcommands (app-server, mcp-server,
+# exec-server, remote-control) and utilities reject it: "--profile only applies
+# to runtime commands and codex mcp". Verified against codex 0.137 + 0.141.
+# ALLOW-LIST (not deny-list) so a future codex server subcommand defaults to the
+# safe no-profile path instead of breaking.
+_CODEX_PROFILE_SUBCOMMANDS = frozenset(
+    {"exec", "review", "resume", "archive", "delete", "unarchive", "fork", "mcp", "sandbox"}
+)
+
+
+def _codex_accepts_profile(tool_args: list[str]) -> bool:
+    sub = _codex_subcommand(tool_args)
+    return sub is None or sub in _CODEX_PROFILE_SUBCOMMANDS
+
+
 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])
+    if _codex_accepts_profile(tool_args):
+        exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args])
+    else:
+        # codex rejects the global --profile on these subcommands (app-server,
+        # mcp-server, exec-server, remote-control, ...). They are caller-
+        # configured — e.g. omnigent runs `codex app-server` with its own
+        # CODEX_HOME + config — so ucode simply omits --profile and stays out of
+        # the way instead of injecting a named profile codex won't accept.
+        exec_or_spawn([binary, *tool_args])
 
 
 def validate_cmd(binary: str) -> list[str]:

tests/test_agent_codex.py

@@ -436,3 +436,70 @@ def fake_execvp(binary: str, args: list[str]) -> None:
 
         assert os.environ["OAUTH_TOKEN"] == "fresh-token"
         assert exec_calls == [("codex", ["codex", "--profile", "ucode", "--search"])]
+
+
+class TestCodexSubcommand:
+    def test_none_for_interactive(self):
+        assert codex._codex_subcommand([]) is None
+        assert codex._codex_subcommand(["--search"]) is None
+
+    def test_finds_bare_subcommand(self):
+        assert codex._codex_subcommand(["exec", "hi"]) == "exec"
+        assert codex._codex_subcommand(["app-server"]) == "app-server"
+
+    def test_skips_leading_value_flags(self):
+        assert codex._codex_subcommand(["--model", "gpt-5", "app-server"]) == "app-server"
+        assert codex._codex_subcommand(["-c", "k=v", "exec"]) == "exec"
+
+    def test_handles_equals_form(self):
+        assert codex._codex_subcommand(["--model=gpt-5", "app-server"]) == "app-server"
+
+
+class TestCodexAcceptsProfile:
+    def test_bare_tui_accepts_profile(self):
+        assert codex._codex_accepts_profile([])
+        assert codex._codex_accepts_profile(["--search"])
+
+    def test_runtime_subcommands_accept_profile(self):
+        for sub in codex._CODEX_PROFILE_SUBCOMMANDS:
+            assert codex._codex_accepts_profile([sub])
+
+    def test_server_family_and_utilities_reject_profile(self):
+        for sub in ("app-server", "mcp-server", "exec-server", "remote-control", "login", "debug"):
+            assert not codex._codex_accepts_profile([sub])
+
+
+class TestCodexLaunchProfileRouting:
+    """launch() attaches --profile only where codex accepts it. Subcommands that
+    reject it (the server family) get the caller's args through untouched, with
+    no --profile — the caller configures those (e.g. omnigent via CODEX_HOME)."""
+
+    @staticmethod
+    def _capture(monkeypatch):
+        calls: list[list[str]] = []
+        monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: calls.append(argv))
+        monkeypatch.setattr(codex, "get_databricks_token", lambda workspace, profile=None: "tok")
+        return calls
+
+    def test_bare_keeps_profile(self, monkeypatch):
+        calls = self._capture(monkeypatch)
+        codex.launch({"workspace": WS}, [])
+        assert calls == [["codex", "--profile", "ucode"]]
+
+    def test_exec_keeps_profile(self, monkeypatch):
+        calls = self._capture(monkeypatch)
+        codex.launch({"workspace": WS}, ["exec", "say hi"])
+        assert calls == [["codex", "--profile", "ucode", "exec", "say hi"]]
+
+    def test_runtime_subcommands_keep_profile(self, monkeypatch):
+        for sub in ("mcp", "resume"):
+            calls = self._capture(monkeypatch)
+            codex.launch({"workspace": WS}, [sub])
+            assert calls[0] == ["codex", "--profile", "ucode", sub]
+
+    def test_server_family_omits_profile(self, monkeypatch):
+        for sub in ("app-server", "mcp-server", "exec-server", "remote-control"):
+            calls = self._capture(monkeypatch)
+            codex.launch({"workspace": WS}, [sub, "--listen", "u"])
+            # No --profile injected; the caller's args pass through untouched.
+            assert calls[0] == ["codex", sub, "--listen", "u"]