Add Codex Azure OpenAI and Microsoft Foundry MPS test coverage

scripts/run_integration.py

@@ -115,6 +115,26 @@ def arguments():
         default="gpt-5-nano",
         help="Model allowed by the OpenAI MPS selected in the configure CUJ.",
     )
+    parser.add_argument(
+        "--codex-azure-provider",
+        default="main.ucode.ci_azure_openai_mps",
+        help="Existing Azure OpenAI MPS (reasoning-capable model) for the Codex azure_openai CUJ.",
+    )
+    parser.add_argument(
+        "--codex-azure-provider-model",
+        default="gpt-5-nano",
+        help="Reasoning-capable model allowed by the Azure OpenAI MPS.",
+    )
+    parser.add_argument(
+        "--codex-foundry-provider",
+        default="main.ucode.ci_foundry_mps",
+        help="Existing Microsoft Foundry MPS (reasoning-capable model) for the Codex foundry CUJ.",
+    )
+    parser.add_argument(
+        "--codex-foundry-provider-model",
+        default="gpt-5-nano",
+        help="Reasoning-capable model allowed by the Microsoft Foundry MPS.",
+    )
     parser.add_argument(
         "--parent-schema",
         default="main.ucode",
@@ -295,6 +315,10 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
             "claude_relayed_provider": args.claude_relayed_provider,
             "codex_provider": args.codex_provider,
             "codex_provider_model": args.codex_provider_model,
+            "codex_azure_provider": args.codex_azure_provider,
+            "codex_azure_provider_model": args.codex_azure_provider_model,
+            "codex_foundry_provider": args.codex_foundry_provider,
+            "codex_foundry_provider_model": args.codex_foundry_provider_model,
             "parent_schema": args.parent_schema,
             "dependencies": args.dependency,
             "workspace": args.workspace,
@@ -536,6 +560,10 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
                 "UG_INTEGRATION_CLAUDE_OAUTH_TOKEN": oauth_token,
                 "UG_INTEGRATION_CODEX_PROVIDER": args.codex_provider,
                 "UG_INTEGRATION_CODEX_PROVIDER_MODEL": args.codex_provider_model,
+                "UG_INTEGRATION_CODEX_AZURE_PROVIDER": args.codex_azure_provider,
+                "UG_INTEGRATION_CODEX_AZURE_PROVIDER_MODEL": args.codex_azure_provider_model,
+                "UG_INTEGRATION_CODEX_FOUNDRY_PROVIDER": args.codex_foundry_provider,
+                "UG_INTEGRATION_CODEX_FOUNDRY_PROVIDER_MODEL": args.codex_foundry_provider_model,
                 "UG_INTEGRATION_PARENT_SCHEMA": args.parent_schema,
                 "UCODE_TEST_WORKSPACE": args.workspace or "",
                 "DATABRICKS_BEARER": bearer,

tests/integration/conftest.py

@@ -116,6 +116,26 @@ def codex_provider_model():
     return os.environ["UG_INTEGRATION_CODEX_PROVIDER_MODEL"]
 
 
+@pytest.fixture(scope="session")
+def codex_azure_provider():
+    return os.environ["UG_INTEGRATION_CODEX_AZURE_PROVIDER"]
+
+
+@pytest.fixture(scope="session")
+def codex_azure_provider_model():
+    return os.environ["UG_INTEGRATION_CODEX_AZURE_PROVIDER_MODEL"]
+
+
+@pytest.fixture(scope="session")
+def codex_foundry_provider():
+    return os.environ["UG_INTEGRATION_CODEX_FOUNDRY_PROVIDER"]
+
+
+@pytest.fixture(scope="session")
+def codex_foundry_provider_model():
+    return os.environ["UG_INTEGRATION_CODEX_FOUNDRY_PROVIDER_MODEL"]
+
+
 @pytest.fixture(scope="session")
 def parent_schema():
     return os.environ["UG_INTEGRATION_PARENT_SCHEMA"]

tests/integration/test_ug_configure_codex.py

@@ -97,3 +97,67 @@ def test_ug_configure_codex_openai_mps(
         tui.exit_normally()
     task.assert_completed(session, "codex")
     session.assert_not_routed()
+
+
+def _run_codex_openai_dialect_mps_cuj(session, workspace, provider, model, label):
+    """Configure Codex through an OpenAI-dialect MPS, then complete a task on one of its models.
+
+    Azure OpenAI and Microsoft Foundry both speak the OpenAI dialect the gateway fronts for Codex,
+    so the CUJ is identical apart from which MPS the picker selects. Codex always sends
+    `reasoning.effort`, so a completed task also proves two things a plain routing check would miss:
+    the provider-type allowlist accepts the MPS's type for Codex, and the backing deployment is
+    reasoning-capable (a non-reasoning model 400s the request's `reasoning.effort`).
+    """
+    task = FileTask(session)
+
+    command = [
+        str(session.binary),
+        "configure",
+        "--workspace",
+        workspace,
+        "--skip-upgrade",
+        "--disable-databricks-ai-tools",
+    ]
+    with ConfigureTerminal(session, "codex", command, f"configure-{label}-provider") as configure:
+        configure.select_agent("Codex")
+        configure.choose("How should Codex get its models?", "External Models")
+        configure.choose("Select a model provider service:", provider)
+        configure.finish(timeout=240)
+    assert session.workspace_state()["provider_services"]["codex"] == provider
+    assert provider in session.run("status").stdout
+
+    command = [str(session.binary), "codex", "--", "--model", model]
+    with AgentTerminal(session, "codex", command, f"{label}-provider-session") as tui:
+        tui.boot()
+        tui.submit(task.prompt)
+        tui.wait_for_task(task, timeout=300)
+        tui.exit_normally()
+    task.assert_completed(session, "codex")
+    session.assert_not_routed()
+
+
+def test_ug_configure_codex_azure_openai_mps(
+    live_session, workspace, codex_azure_provider, codex_azure_provider_model
+):
+    """Scenario: select an Azure OpenAI (`azure_openai`) MPS in ug configure's picker for Codex.
+
+    Expected: ug saves the provider and a real Codex session completes a file task on one of its
+    reasoning-capable models (see _run_codex_openai_dialect_mps_cuj).
+    """
+    _run_codex_openai_dialect_mps_cuj(
+        live_session, workspace, codex_azure_provider, codex_azure_provider_model, "azure"
+    )
+
+
+def test_ug_configure_codex_foundry_mps(
+    live_session, workspace, codex_foundry_provider, codex_foundry_provider_model
+):
+    """Scenario: select a Microsoft Foundry (`microsoft_foundry`) MPS in ug configure's picker.
+
+    Expected: ug saves the provider and a real Codex session completes a file task, exercising the
+    `microsoft_foundry` allowlist entry for Codex. Foundry speaks the same OpenAI dialect as Azure
+    OpenAI, so this reuses the Azure CUJ (see _run_codex_openai_dialect_mps_cuj).
+    """
+    _run_codex_openai_dialect_mps_cuj(
+        live_session, workspace, codex_foundry_provider, codex_foundry_provider_model, "foundry"
+    )

tests/test_e2e.py

@@ -57,6 +57,12 @@
 )
 CI_OPENAI_MPS = "main.ucode.ci_openai_mps"  # api-key OpenAI (for codex)
 CI_OPENAI_MODEL = "gpt-5-nano"
+# Azure OpenAI MPS for codex, with a mixed allowlist: two reasoning-capable models plus one
+# non-reasoning model. Codex always sends `reasoning.effort`, so this single MPS exercises the
+# reasoning gate, model switching, and allowlist rejection (see TestAzureOpenAiMps).
+CI_AZURE_OPENAI_MPS = "main.ucode.ci_azure_openai_mps"
+CI_AZURE_REASONING_MODELS = ("gpt-5-nano", "gpt-5.6-luna")
+CI_AZURE_NON_REASONING_MODEL = "gpt-4.1-mini"
 
 # Claude Code's client-side model pre-flight is flaky for MPS-routed ids and emits this; the
 # launch retries on it before giving up (see TestModelProviderLaunch).
@@ -848,6 +854,84 @@ def test_mps_is_listed_usable_and_nonrelayed(self, e2e_workspace, e2e_token):
         assert not entry.get("relayed")
 
 
+class TestAzureOpenAiMps:
+    """Azure OpenAI MPS for codex, pinned to CI_AZURE_OPENAI_MPS.
+
+    Exercised without the agent binary — direct Responses API calls through the gateway MPS,
+    the same surface codex uses. The allowlist mixes reasoning-capable models (gpt-5 family)
+    with a non-reasoning one (gpt-4.1-mini); codex always sends reasoning.effort, so these probes
+    pin the behavior a real codex launch depends on: any reasoning model routes (so a session can
+    switch between them), a non-reasoning model rejects reasoning.effort, and an off-allowlist
+    model is refused up front.
+    """
+
+    @staticmethod
+    def _responses_request(
+        workspace: str, token: str, model: str, *, reasoning: bool = True
+    ) -> httpx.Response:
+        body: dict = {
+            "model": model,
+            "input": "say hi in 5 words or less",
+            "max_output_tokens": 16,
+        }
+        if reasoning:
+            body["reasoning"] = {"effort": "low"}
+        return httpx.post(
+            f"{build_tool_base_url('codex', workspace)}/responses",
+            headers={
+                "Authorization": f"Bearer {token}",
+                "Databricks-Model-Provider-Service": CI_AZURE_OPENAI_MPS,
+            },
+            json=body,
+            timeout=60,
+        )
+
+    @staticmethod
+    def _skip_if_unusable(resp: httpx.Response) -> None:
+        # Environmental account conditions, not ucode bugs (mirrors TestAnthropicNonRelayMps). The
+        # allowlist 403 ("not in the allowed models list") is a real assertion, not a skip — only a
+        # connection-permission 401/403 or an exhausted account falls outside the test's scope.
+        text = resp.text
+        if resp.status_code in (401, 403) and ("USE CONNECTION" in text or "EXECUTE" in text):
+            pytest.skip(f"no permission on {CI_AZURE_OPENAI_MPS}: {text[:200]}")
+        if "quota" in text.lower():
+            pytest.skip(f"{CI_AZURE_OPENAI_MPS} account is out of quota: {text[:200]}")
+
+    @pytest.mark.parametrize("model", CI_AZURE_REASONING_MODELS)
+    def test_reasoning_model_serves_codex_reasoning_effort(self, model, e2e_workspace, e2e_token):
+        # codex's request shape (reasoning.effort) routes on every reasoning-capable allowed model,
+        # so an in-session /model switch between them stays within the allowlist.
+        resp = self._responses_request(e2e_workspace, e2e_token, model)
+        self._skip_if_unusable(resp)
+        assert resp.status_code == 200, f"model={model} HTTP {resp.status_code}: {resp.text[:300]}"
+
+    def test_non_reasoning_model_rejects_reasoning_effort(self, e2e_workspace, e2e_token):
+        # Parity invariant: a model that would not work with codex natively must not work through the
+        # gateway either. gpt-4.1-mini is allow-listed but non-reasoning, and codex always sends
+        # reasoning.effort, so the gateway must surface the provider's 400 — not paper over it by
+        # stripping the param.
+        resp = self._responses_request(e2e_workspace, e2e_token, CI_AZURE_NON_REASONING_MODEL)
+        self._skip_if_unusable(resp)
+        assert resp.status_code == 400, f"HTTP {resp.status_code}: {resp.text[:300]}"
+        assert "reasoning.effort" in resp.text
+
+    def test_non_reasoning_model_serves_without_reasoning(self, e2e_workspace, e2e_token):
+        # The rejection above is the reasoning parameter, not the model or the allowlist: the same
+        # model serves a plain request.
+        resp = self._responses_request(
+            e2e_workspace, e2e_token, CI_AZURE_NON_REASONING_MODEL, reasoning=False
+        )
+        self._skip_if_unusable(resp)
+        assert resp.status_code == 200, f"HTTP {resp.status_code}: {resp.text[:300]}"
+
+    def test_rejects_model_outside_allowlist(self, e2e_workspace, e2e_token):
+        # The MPS enforces its allowlist up front: a model it doesn't declare is refused.
+        resp = self._responses_request(e2e_workspace, e2e_token, "gpt-4o")
+        self._skip_if_unusable(resp)
+        assert resp.status_code == 403, f"HTTP {resp.status_code}: {resp.text[:300]}"
+        assert "not in the allowed models list" in resp.text
+
+
 class TestGeminiLaunch:
     """Run the real Gemini CLI against every available gemini model."""