codex: allow-list --profile by subcommand; deliver config via -c where --profile is rejected

ucode's codex launch always ran `codex --profile ucode <args>`, but codex only
accepts the global `--profile` on runtime subcommands (the interactive TUI,
exec/review/resume/archive/delete/unarchive/fork/mcp/sandbox). Server-family
subcommands (app-server, mcp-server, exec-server, remote-control) and utilities
reject it up front:

  Error: --profile only applies to runtime commands and `codex mcp`: ...

so e.g. `ucode codex app-server` was functionally broken.

Allow-list the subcommands that accept `--profile` (a deny-list would break the
next codex server subcommand by default). Where `--profile` is rejected, deliver
the same provider block ucode writes as codex root `-c` overrides instead (root
`model_provider`/`model` keys + one inline `model_providers.<name>` table,
mirroring omnigent), so those subcommands still route through the Databricks AI
Gateway without a named profile.

Co-authored-by: Isaac

Author: Edwinhe03

src/ucode/agents/codex.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 import os
 import re
 from pathlib import Path
@@ -21,7 +22,7 @@
     get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
-from ucode.state import mark_tool_managed, save_state
+from ucode.state import get_provider_service, mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 
 CODEX_CONFIG_DIR = Path.home() / ".codex"
@@ -391,12 +392,125 @@ 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 _toml_key(k: str) -> str:
+    return k if re.fullmatch(r"[A-Za-z0-9_-]+", k) else json.dumps(k)
+
+
+def _toml_value(v) -> str:
+    if isinstance(v, bool):
+        return "true" if v else "false"
+    if isinstance(v, str):
+        return json.dumps(v)  # JSON basic string == TOML basic string
+    if isinstance(v, int):
+        return str(v)
+    if isinstance(v, list):
+        return "[" + ",".join(_toml_value(x) for x in v) + "]"
+    if isinstance(v, dict):
+        return "{" + ",".join(f"{_toml_key(k)}={_toml_value(x)}" for k, x in v.items()) + "}"
+    raise TypeError(f"unsupported TOML value for codex -c override: {v!r}")
+
+
+def _codex_config_overrides(state: dict) -> list[str]:
+    """Provider config as codex root -c overrides, for subcommands that reject
+    --profile (they cannot select a named profile, so deliver the provider as the
+    default/root config instead)."""
+    provider = get_provider_service(state, "codex")
+    block = _provider_block(
+        state["workspace"],
+        state.get("profile"),
+        use_pat=bool(state.get("use_pat")),
+        provider=provider,
+    )
+    entries = [f"model_provider={_toml_value(CODEX_MODEL_PROVIDER_NAME)}"]
+    # pinned model, if ucode has one written for this workspace
+    model = None
+    doc = read_toml_safe(CODEX_CONFIG_PATH)
+    if isinstance(doc, dict):
+        model = doc.get("model")
+    if model:
+        entries.append(f"model={_toml_value(model)}")
+    entries.append(f"model_providers.{CODEX_MODEL_PROVIDER_NAME}={_toml_value(block)}")
+    args: list[str] = []
+    for e in entries:
+        args += ["-c", e]
+    return args
+
+
 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 --profile here and offers no named-profile selection for
+        # these subcommands; deliver the same provider config as root -c overrides.
+        exec_or_spawn([binary, *_codex_config_overrides(state), *tool_args])
 
 
 def validate_cmd(binary: str) -> list[str]:

tests/test_agent_codex.py

@@ -436,3 +436,115 @@ 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, and otherwise
+    delivers the provider config as root -c overrides (server-family subcommands
+    reject --profile)."""
+
+    @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")
+        # Keep provider-block rendering off the network/subprocess in tests.
+        monkeypatch.setattr(codex, "agent_version", lambda binary: "0.141.0")
+        monkeypatch.setattr(codex, "ucode_version", lambda: "0.1.0")
+        return calls
+
+    def test_bare_keeps_profile_no_overrides(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_without_overrides(self, monkeypatch):
+        for sub in ("mcp", "resume"):
+            calls = self._capture(monkeypatch)
+            codex.launch({"workspace": WS}, [sub])
+            argv = calls[0]
+            assert argv[:3] == ["codex", "--profile", "ucode"]
+            assert not any(a.startswith("model_provider=") for a in argv)
+
+    def test_server_family_uses_config_overrides_not_profile(self, monkeypatch, tmp_path):
+        monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", tmp_path / "ucode.config.toml")
+        for sub in ("app-server", "mcp-server", "exec-server"):
+            calls = self._capture(monkeypatch)
+            codex.launch({"workspace": WS}, [sub, "--listen", "u"])
+            argv = calls[0]
+            assert "--profile" not in argv
+            assert argv[0] == "codex"
+            assert "-c" in argv
+            joined = " ".join(argv)
+            assert "model_provider=" in joined
+            assert "model_providers.ucode-databricks=" in joined
+            # The caller's own args are preserved at the tail, untouched.
+            assert argv[-3:] == [sub, "--listen", "u"]
+
+    def test_config_overrides_pin_model_when_written(self, monkeypatch, tmp_path):
+        config = tmp_path / "ucode.config.toml"
+        config.write_text('model = "gpt-5"\n', encoding="utf-8")
+        monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config)
+        calls = self._capture(monkeypatch)
+        codex.launch({"workspace": WS}, ["app-server"])
+        joined = " ".join(calls[0])
+        assert 'model="gpt-5"' in joined
+
+
+class TestCodexConfigOverrides:
+    def test_serializes_provider_block_as_c_overrides(self, monkeypatch, tmp_path):
+        monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", tmp_path / "missing.toml")
+        monkeypatch.setattr(codex, "agent_version", lambda binary: "0.141.0")
+        monkeypatch.setattr(codex, "ucode_version", lambda: "0.1.0")
+        args = codex._codex_config_overrides({"workspace": WS, "profile": "DEFAULT"})
+        # Every override is delivered as a `-c key=value` pair.
+        assert args[0::2] == ["-c", "-c"]
+        assert args[1].startswith("model_provider=")
+        assert args[3].startswith("model_providers.ucode-databricks=")
+
+    def test_toml_value_quotes_strings_and_serializes_nested(self):
+        assert codex._toml_value("hi") == '"hi"'
+        assert codex._toml_value(5000) == "5000"
+        assert codex._toml_value(True) == "true"
+        assert codex._toml_value(["a", "b"]) == '["a","b"]'
+        assert codex._toml_value({"k": "v"}) == '{k="v"}'
+
+    def test_toml_key_quotes_non_bare_keys(self):
+        assert codex._toml_key("User-Agent") == "User-Agent"
+        assert codex._toml_key("has space") == '"has space"'