diff --git a/scripts/run_integration.py b/scripts/run_integration.py index 9fa73e2b..b7b418e3 100644 --- a/scripts/run_integration.py +++ b/scripts/run_integration.py @@ -115,6 +115,11 @@ def arguments(): default="gpt-5-nano", help="Model allowed by the OpenAI MPS selected in the configure CUJ.", ) + parser.add_argument( + "--parent-schema", + default="main.ucode", + help="Schema containing the dedicated model-discovery Model Services.", + ) parser.add_argument("--python", default=sys.executable, help="Python 3.12+ path or uv version.") parser.add_argument("--dependency", action="append", default=[], metavar="PACKAGE==VERSION") parser.add_argument("--constraints", type=Path, help="Replay a previous dependencies.txt.") @@ -290,6 +295,7 @@ 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, + "parent_schema": args.parent_schema, "dependencies": args.dependency, "workspace": args.workspace, }, @@ -530,6 +536,7 @@ 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_PARENT_SCHEMA": args.parent_schema, "UCODE_TEST_WORKSPACE": args.workspace or "", "DATABRICKS_BEARER": bearer, } diff --git a/tests/AGENTS.md b/tests/AGENTS.md index 6cc382d3..957e1c89 100644 --- a/tests/AGENTS.md +++ b/tests/AGENTS.md @@ -30,9 +30,10 @@ tests. Keep work scoped to the behavior requested by the user. exception is the `managed_fixture` marker: it uses the built-in `UCODE_MANAGED_CONFIG_STUB` hook to inject the admin CodingAgentConfig INPUT so the real `ug configure` path can be exercised across config shapes the live workspace does - not publish. The gateway, agent binaries, ug internals, and ug state stay real; the - config fetch/wire contract stays covered by the un-stubbed `managed` tests; and the - hook must never be used to disable validation or conceal a failure. + not publish and one workspace fetch can be replayed across isolated cases. The gateway, + agent binaries, ug internals, and ug state stay real; the config fetch/wire contract stays + covered by the un-stubbed `managed` tests; and the hook must never be used to disable + validation or conceal a failure. 5. **Real responses and binaries.** Pin requested ug and agent versions. Never substitute a missing binary/service. Reuse explicit e2e workspace/auth settings; never pick a developer's Databricks profile automatically. diff --git a/tests/README.md b/tests/README.md index c3ae140e..957bbfec 100644 --- a/tests/README.md +++ b/tests/README.md @@ -52,18 +52,28 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup | | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch | | `test_ug_configure_managed_codex_catalog_fallback` | Configure from an injected managed response containing a GPT model absent from Codex's bundled catalog | Actionable metadata warning; conservative catalog entry for the unknown model; real Codex prompt on the valid default model | +| `test_case_01_*`, `test_case_02_*` | Launch Claude/Codex after managed configure and from fresh state | Exact admin MPS source; Claude opens its native picker after an MPS-scoped cache refresh, and Codex's app server matches its scoped catalog | +| `test_case_03_*`, `test_case_04_*` | Disable personal discovery, then launch managed Claude/Codex after configure and from fresh state | The managed MPS still supplies each agent's native catalog | +| `test_case_05_*`, `test_case_06_*` | Pass a provider override to managed Claude/Codex after configure and from fresh state | ug rejects before agent startup without changing agent-owned state/files | +| `test_case_07_*`, `test_case_08_*` | Pass a model-location override to managed Claude/Codex after configure and from fresh state | ug rejects before agent startup without changing agent-owned state/files | +| `test_case_09_*`, `test_case_10_*` | Disable discovery and pass a provider override to managed Claude/Codex | ug still rejects both configured and fresh launches | +| `test_case_11_*`, `test_case_12_*` | Disable discovery and pass a model-location override to managed Claude/Codex | ug still rejects both configured and fresh launches | | `test_ug_installed_wheel_exposes_help_and_version` | Invoke freshly installed console command | Package version matches; public help works | | `test_ug_status_in_fresh_home_is_unconfigured` | Request status before configure | Unconfigured status | | `test_ug_auth_without_configuration_explains_how_to_configure` | Request auth before configure | Actionable setup error and nonzero exit | | `test_ug_and_ucode_auth_helpers_emit_only_the_supplied_bearer` | Run both auth helper commands with the public bearer override, with and without forced refresh | Exact token-only stdout, no warnings or ANSI escapes; no workspace authentication or saved state | | `test_ug_and_ucode_web_search_helpers_preserve_mcp_stdio` | Initialize and list tools through both web-search helper commands | Exactly the MCP JSON-RPC responses; no text/ANSI contamination; existing server/tool identities preserved; no model request | -With both agents selected there are **41 live cases** (6 interactive TUI cases), -**3 managed-workspace cases** (marker `managed`, run against a separate workspace that -publishes a CodingAgentConfig), **3 managed-fixture cases** (marker `managed_fixture`, with only -the CodingAgentConfig input injected), and **5 installation checks**. Parametrization varies -argument spelling or routing mode, never hides the agent/provider in the test name. Duplicate boot-only cases -are incorporated into the Databricks configuration TUI journeys. +With both agents selected there are **42 live cases** (6 interactive TUI cases), +**4 managed-workspace cases** (marker `managed`, run against a separate workspace that +publishes a CodingAgentConfig), **28 managed-fixture cases** (marker `managed_fixture`), and +**5 installation checks**. The 12 numbered scenarios fetch the published config once per agent, +replace that agent's static model source with its dedicated test MPS, drop its incompatible static +defaults, and reuse the result across 24 explicit configured/fresh journeys. The other four +collected cases, from three test functions, inject focused model and MCP shapes. Parametrization +varies argument spelling or routing mode, never hides the +agent/provider in the test name. Duplicate boot-only cases are incorporated into the Databricks +configuration TUI journeys. Generated-file cleanup and strict app-server stdout assertions remain enforced. ug no longer runs a post-configure agent probe; the deprecated `--skip-validate` diff --git a/tests/integration/README.md b/tests/integration/README.md index 432a2776..c772bd4c 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -98,6 +98,8 @@ test_ug_codex_commands.py # command help and parser error forwardi test_ug_codex_app_server.py # actual client/server initialize exchange test_ug_configure_claude_lifecycle.py # repeat setup, revert, rejected credentials test_ug_configure_codex_lifecycle.py # repeat setup, revert, rejected credentials +test_ug_claude_managed_model_discovery.py # fetched/reused Claude MPS policy cases +test_ug_codex_managed_model_discovery.py # fetched/reused Codex MPS policy cases test_ug_configure_managed.py # managed workspace: static model list, no agent selector test_ug_configure_managed_models.py # injected model lists: pickers and Codex fallback metadata test_ug_configure_managed_mcp.py # injected managed MCP list @@ -155,11 +157,13 @@ No service is created or modified. A missing service, permission, or OAuth token fails the selected CUJ, rather than skipping it. There are **42 live cases** (including 6 TUI journeys) and **5 installation -checks** with both agents. A separate **3 managed-workspace cases** (one per agent -plus an idempotent re-configure, marker `managed`) run against a workspace that publishes a CodingAgentConfig; see -"Managed-workspace journeys" below. A further **3 `managed_fixture` cases** inject the admin config -locally (via `UCODE_MANAGED_CONFIG_STUB`) to cover shapes the live workspace does not publish; each -differs from the published config in what it asserts so it proves the injected config drove configure. +checks** with both agents. A separate **4 managed-workspace cases** (one per agent, +an idempotent re-configure, and cache reuse within the TTL; marker `managed`) run against a +workspace that publishes a CodingAgentConfig; see "Managed-workspace journeys" below. A further +**28 `managed_fixture` cases** use `UCODE_MANAGED_CONFIG_STUB`: 24 explicit configured/fresh +managed-discovery journeys fetch the published config once per agent module, replace that agent's +static model source with its dedicated MPS, drop its incompatible static defaults, and reuse the +result; four collected cases from three test functions inject focused model/MCP shapes. See the named coverage and gaps matrix in [../README.md](../README.md). @@ -275,12 +279,23 @@ with no agent selector, and each agent's generated config exposes exactly the ad Treat that published CodingAgentConfig as shared CI fixture state. The managed lanes assert its exact model ids and its both-agent enablement, so editing the managed workspace's config (models, -enabled agents, or defaults) breaks these lanes until the constants in `test_ug_configure_managed.py` -are updated to match. Do not change it casually. +enabled agents, or defaults) breaks these lanes until the constants in `utils/constants.py` are +updated to match. Do not change it casually. The `managed_fixture` journeys use `UCODE_MANAGED_CONFIG_STUB` to short-circuit only the -managed-config HTTP read for config shapes that workspace does not publish. In particular, -`test_ug_configure_managed_codex_catalog_fallback` injects the intentionally nonexistent +managed-config HTTP read for config shapes that workspace does not publish. Each managed-discovery +module fetches the workspace's published config once with the CI bearer, replaces that agent's +static model source with its dedicated ca-central-1 MPS, drops its incompatible static defaults, +and reuses the resulting stub across every isolated scenario: + +- Claude Code: `main.default.ci_e2e_anthropic_mps` +- Codex: `main.default.ci_e2e_openai_mps` + +The real agent then performs its native MPS-scoped model discovery. Claude must write a nonempty +gateway-model cache for the exact Anthropic base URL and open its native picker; Codex's app server +must match the generated scoped catalog. The persisted managed input names the exact provider for +both agents, and Claude's generated config carries its exact provider header. The existing +`test_ug_configure_managed_codex_catalog_fallback` also injects the intentionally nonexistent `system.ai.gpt-99`, keeping it out of the real workspace while launching Codex through that workspace on the valid default model `system.ai.gpt-5-6-sol`. With smart routing enabled, it opens the real Codex `/models` picker and requires that injected custom-catalog model to be listed. The @@ -299,7 +314,8 @@ Run it locally the same way, pointing at the managed workspace: ```bash export UCODE_TEST_WORKSPACE=https:// export DATABRICKS_CLIENT_ID= DATABRICKS_CLIENT_SECRET= -python scripts/run_integration.py --claude-version --codex-version -- -m managed +python scripts/run_integration.py --claude-version --codex-version \ + -- -m "managed or managed_fixture" ``` Each job uses fresh consumer dependency resolution. There is no default dependency diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 7c8c7a52..7f14fd53 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -114,3 +114,8 @@ def codex_provider(): @pytest.fixture(scope="session") def codex_provider_model(): return os.environ["UG_INTEGRATION_CODEX_PROVIDER_MODEL"] + + +@pytest.fixture(scope="session") +def parent_schema(): + return os.environ["UG_INTEGRATION_PARENT_SCHEMA"] diff --git a/tests/integration/test_ug_claude_managed_model_discovery.py b/tests/integration/test_ug_claude_managed_model_discovery.py new file mode 100644 index 00000000..8fe84584 --- /dev/null +++ b/tests/integration/test_ug_claude_managed_model_discovery.py @@ -0,0 +1,397 @@ +"""Claude managed-config CUJs for Tests-table cases 1, 3, 5, 7, 9, and 11. + +The admin CodingAgentConfig is fetched once from the managed workspace, its Claude model source is +set to the dedicated test MPS, and the result is reused through ``UCODE_MANAGED_CONFIG_STUB`` in +each isolated session. Normalization, config writers, the gateway, and Claude Code remain real. +""" + +import json +import os +import re + +import pytest +from utils.constants import MANAGED_CLAUDE_PROVIDER_SERVICE +from utils.managed import ( + fetch_managed_config_stub, + is_managed_config_control_plane_cache, + use_managed_config_stub, +) +from utils.terminal import AgentTerminal + +pytestmark = [pytest.mark.managed_fixture, pytest.mark.claude] + + +@pytest.fixture(scope="module") +def _managed_claude_config_stub(workspace, tmp_path_factory): + return fetch_managed_config_stub( + workspace, + os.environ["DATABRICKS_BEARER"], + tmp_path_factory.mktemp("managed-config-claude"), + "managed-config-claude.json", + agent="CODING_AGENT_CLAUDE_CODE", + provider_service=MANAGED_CLAUDE_PROVIDER_SERVICE, + ) + + +@pytest.fixture(autouse=True) +def _managed_claude_config(live_session, _managed_claude_config_stub): + use_managed_config_stub(live_session, _managed_claude_config_stub) + + +def _claude_state_and_agent_files(session): + paths = [session.home / ".claude.json"] + for directory in (session.home / ".ucode", session.home / ".claude"): + if directory.exists(): + paths.extend(path for path in directory.rglob("*") if path.is_file()) + return { + str(path.relative_to(session.home)): path.read_bytes() + for path in paths + if path.is_file() + # A fresh launch must retrieve and cache the control-plane input before it can reject an + # override. Exclude only that expected cache; every agent-owned state/file stays compared. + and not is_managed_config_control_plane_cache(session.home, path) + } + + +def _assert_rejected_before_claude_started(session, result, requested_source, before): + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert requested_source.lower() in output + assert "admin has specified managed" in output + assert MANAGED_CLAUDE_PROVIDER_SERVICE.lower() in output + assert _claude_state_and_agent_files(session) == before + + +def _assert_managed_provider_in_picker(session, workspace, screen): + settings = json.loads((session.home / ".claude" / "ucode-settings.json").read_text()) + headers = (settings.get("env") or {}).get("ANTHROPIC_CUSTOM_HEADERS", "").splitlines() + expected_header = f"Databricks-Model-Provider-Service: {MANAGED_CLAUDE_PROVIDER_SERVICE}" + assert headers.count(expected_header) == 1, settings + # MPS models come from Claude Code's native gateway discovery, not a static managed picker. + assert not {"availableModels", "enforceAvailableModels", "modelPicker"} & settings.keys(), ( + settings + ) + # Native gateway rows deduplicate against built-ins, so exact cached ids need not be rendered. + assert re.search(r"(?m)^\s*(?:[❯›>]\s*)?\d+\.\s+\S", screen), screen + + cache = json.loads((session.home / ".claude/cache/gateway-models.json").read_text()) + assert cache.get("baseUrl") == workspace.rstrip("/") + "/ai-gateway/anthropic", cache + assert isinstance(cache.get("fetchedAt"), int) and cache["fetchedAt"] > 0, cache + cached_models = cache.get("models") + assert isinstance(cached_models, list) and cached_models, cache + cached_ids = [model.get("id") for model in cached_models if isinstance(model, dict)] + assert len(cached_ids) == len(cached_models), cache + assert cached_ids and all(isinstance(model_id, str) and model_id for model_id in cached_ids), ( + cache + ) + + +@pytest.mark.tui +def test_case_01_managed_claude_uses_admin_discovery_after_configure(live_session, workspace): + """Scenario: configure managed Claude, then launch its model picker. + + Expected: the managed model catalog wins after configuration. + """ + session = live_session + result = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in result.stdout, result.stdout + + command = [str(session.binary), "claude"] + with AgentTerminal(session, "claude", command, "case-01-managed") as tui: + tui.boot() + screen = tui.open_model_picker() + tui.exit_normally() + + _assert_managed_provider_in_picker(session, workspace, screen) + + +@pytest.mark.tui +def test_case_01_fresh_managed_claude_uses_admin_discovery(live_session, workspace): + """Scenario: launch managed Claude's model picker from fresh state. + + Expected: the managed model catalog wins without prior configuration. + """ + session = live_session + command = [str(session.binary), "claude", "--workspace", workspace] + with AgentTerminal(session, "claude", command, "case-01-fresh-managed") as tui: + tui.boot() + screen = tui.open_model_picker() + tui.exit_normally() + + _assert_managed_provider_in_picker(session, workspace, screen) + + +@pytest.mark.tui +def test_case_03_managed_claude_ignores_discovery_disable_after_configure(live_session, workspace): + """Scenario: configure managed Claude, disable discovery, then launch its model picker. + + Expected: workspace-managed discovery still supplies the admin's catalog. + """ + session = live_session + result = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in result.stdout, result.stdout + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + command = [str(session.binary), "claude"] + with AgentTerminal(session, "claude", command, "case-03-managed-disabled") as tui: + tui.boot() + screen = tui.open_model_picker() + tui.exit_normally() + + _assert_managed_provider_in_picker(session, workspace, screen) + + +@pytest.mark.tui +def test_case_03_fresh_managed_claude_ignores_discovery_disable(live_session, workspace): + """Scenario: disable discovery and launch managed Claude's model picker from fresh state. + + Expected: workspace-managed discovery still supplies the admin's catalog. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + command = [str(session.binary), "claude", "--workspace", workspace] + with AgentTerminal(session, "claude", command, "case-03-fresh-managed-disabled") as tui: + tui.boot() + screen = tui.open_model_picker() + tui.exit_normally() + + _assert_managed_provider_in_picker(session, workspace, screen) + + +def test_case_05_managed_claude_rejects_provider_override(live_session, workspace, claude_provider): + """Scenario: configure managed Claude, then pass --provider. + + Expected: ug rejects the override without changing agent-owned state/files. + """ + session = live_session + configured_result = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured_result.stdout + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--provider", + claude_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, f"provider {claude_provider}", before) + + +def test_case_05_fresh_managed_claude_rejects_provider_override( + live_session, workspace, claude_provider +): + """Scenario: pass --provider while launching managed Claude from fresh state. + + Expected: ug rejects the override without changing agent-owned state/files; only the + managed-config retrieval cache may be written. + """ + session = live_session + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--workspace", + workspace, + "--provider", + claude_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, f"provider {claude_provider}", before) + + +def test_case_07_managed_claude_rejects_model_location_override( + live_session, workspace, parent_schema +): + """Scenario: configure managed Claude, then pass --model-location. + + Expected: ug rejects the override without changing agent-owned state/files. + """ + session = live_session + configured_result = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured_result.stdout + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, "--model-location", before) + + +def test_case_07_fresh_managed_claude_rejects_model_location_override( + live_session, workspace, parent_schema +): + """Scenario: pass --model-location while launching managed Claude from fresh state. + + Expected: ug rejects the override without changing agent-owned state/files; only the + managed-config retrieval cache may be written. + """ + session = live_session + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, "--model-location", before) + + +def test_case_09_managed_claude_rejects_provider_when_discovery_disabled( + live_session, workspace, claude_provider +): + """Scenario: configure managed Claude, disable discovery, then pass --provider. + + Expected: ug rejects the override without changing agent-owned state/files. + """ + session = live_session + configured_result = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured_result.stdout + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--provider", + claude_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, f"provider {claude_provider}", before) + + +def test_case_09_fresh_managed_claude_rejects_provider_when_discovery_disabled( + live_session, workspace, claude_provider +): + """Scenario: disable discovery and pass --provider to managed Claude from fresh state. + + Expected: ug rejects the override without changing agent-owned state/files; only the + managed-config retrieval cache may be written. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--workspace", + workspace, + "--provider", + claude_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, f"provider {claude_provider}", before) + + +def test_case_11_managed_claude_rejects_model_location_when_discovery_disabled( + live_session, workspace, parent_schema +): + """Scenario: configure managed Claude, disable discovery, then pass --model-location. + + Expected: ug rejects the override without changing agent-owned state/files. + """ + session = live_session + configured_result = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured_result.stdout + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, "--model-location", before) + + +def test_case_11_fresh_managed_claude_rejects_model_location_when_discovery_disabled( + live_session, workspace, parent_schema +): + """Scenario: disable discovery and pass --model-location to managed Claude from fresh state. + + Expected: ug rejects the override without changing agent-owned state/files; only the + managed-config retrieval cache may be written. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _claude_state_and_agent_files(session) + result = session.run( + "claude", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result, "--model-location", before) diff --git a/tests/integration/test_ug_codex_managed_model_discovery.py b/tests/integration/test_ug_codex_managed_model_discovery.py new file mode 100644 index 00000000..54d455dc --- /dev/null +++ b/tests/integration/test_ug_codex_managed_model_discovery.py @@ -0,0 +1,434 @@ +"""Codex managed-config CUJs for Tests-table cases 2, 4, 6, 8, 10, and 12. + +The admin CodingAgentConfig is fetched once from the managed workspace, its Codex model source is +set to the dedicated test MPS, and the result is reused through ``UCODE_MANAGED_CONFIG_STUB`` in +each isolated session. Normalization, config writers, the gateway, and Codex remain real. +""" + +import json +import os +import tomllib + +import pytest +from utils.constants import MANAGED_CODEX_PROVIDER_SERVICE +from utils.managed import ( + fetch_managed_config_stub, + is_managed_config_control_plane_cache, + use_managed_config_stub, +) + +pytestmark = [pytest.mark.managed_fixture, pytest.mark.codex] + + +@pytest.fixture(scope="module") +def _managed_codex_config_stub(workspace, tmp_path_factory): + return fetch_managed_config_stub( + workspace, + os.environ["DATABRICKS_BEARER"], + tmp_path_factory.mktemp("managed-config-codex"), + "managed-config-codex.json", + agent="CODING_AGENT_CODEX", + provider_service=MANAGED_CODEX_PROVIDER_SERVICE, + ) + + +@pytest.fixture(autouse=True) +def _managed_codex_config(live_session, _managed_codex_config_stub): + use_managed_config_stub(live_session, _managed_codex_config_stub) + + +def _codex_state_and_agent_files(session): + paths = [] + for directory in (session.home / ".ucode", session.home / ".codex"): + if directory.exists(): + paths.extend(path for path in directory.rglob("*") if path.is_file()) + return { + str(path.relative_to(session.home)): path.read_bytes() + for path in paths + if path.is_file() + # A fresh launch must retrieve and cache the control-plane input before it can reject an + # override. Exclude only that expected cache; every agent-owned state/file stays compared. + and not is_managed_config_control_plane_cache(session.home, path) + } + + +def _assert_managed_provider_catalog(session, models): + config = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) + # The provider header is a launch-only overlay; neither it nor the scoped catalog is persisted + # in Codex's generated profile. + assert "model_catalog_json" not in config, config + + managed_cache = json.loads((session.home / ".ucode/managed-config.json").read_text()) + raw_config = managed_cache.get("config") + enabled_agents = raw_config.get("enabled_agents") if isinstance(raw_config, dict) else None + assert isinstance(enabled_agents, list), "persisted managed config had no enabled_agents list" + codex_entries = [ + entry + for entry in enabled_agents + if isinstance(entry, dict) and entry.get("agent") == "CODING_AGENT_CODEX" + ] + assert len(codex_entries) == 1, "persisted managed config did not contain one Codex entry" + managed_codex = codex_entries[0].get("config") + assert isinstance(managed_codex, dict), "persisted managed Codex config was not an object" + assert managed_codex.get("models") == { + "model_provider_service": MANAGED_CODEX_PROVIDER_SERVICE + }, "persisted managed Codex config did not select the dedicated MPS" + assert "default_models" not in managed_codex, "static Codex defaults survived the MPS variant" + + catalog_paths = list((session.home / ".ucode").glob("codex-model-catalog-*.json")) + assert len(catalog_paths) == 1, catalog_paths + catalog = json.loads(catalog_paths[0].read_text()) + catalog_ids = [ + model.get("slug") + for model in catalog.get("models", []) + if isinstance(model, dict) and model.get("visibility") == "list" + ] + assert catalog_ids, catalog + assert models == catalog_ids, (models, catalog) + + +def test_case_02_managed_codex_uses_admin_discovery_after_configure(live_session, workspace): + """Scenario: configure managed Codex, then launch its app server. + + Expected: Codex exposes exactly the admin-managed model catalog. + """ + session = live_session + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + + models = session.codex_model_ids(["app-server", "--listen", "stdio://"]) + + _assert_managed_provider_catalog(session, models) + + +def test_case_02_managed_codex_uses_admin_discovery_from_fresh_state(live_session, workspace): + """Scenario: launch managed Codex with --workspace from fresh state. + + Expected: Codex exposes exactly the admin-managed model catalog. + """ + session = live_session + models = session.codex_model_ids( + ["--workspace", workspace, "--", "app-server", "--listen", "stdio://"] + ) + + _assert_managed_provider_catalog(session, models) + + +def test_case_04_managed_codex_ignores_discovery_disable_after_configure(live_session, workspace): + """Scenario: configure managed Codex, disable discovery, then launch its app server. + + Expected: workspace-managed discovery still supplies the admin's catalog. + """ + session = live_session + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + + models = session.codex_model_ids(["app-server", "--listen", "stdio://"]) + + _assert_managed_provider_catalog(session, models) + + +def test_case_04_managed_codex_ignores_discovery_disable_from_fresh_state(live_session, workspace): + """Scenario: disable discovery and launch managed Codex with --workspace from fresh state. + + Expected: workspace-managed discovery still supplies the admin's catalog. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + + models = session.codex_model_ids( + ["--workspace", workspace, "--", "app-server", "--listen", "stdio://"] + ) + + _assert_managed_provider_catalog(session, models) + + +def test_case_06_managed_codex_rejects_provider_override_after_configure( + live_session, workspace, codex_provider +): + """Scenario: configure managed Codex, then pass a --provider override. + + Expected: ug rejects the override without changing agent-owned state or files. + """ + session = live_session + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--provider", + codex_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert f"provider {codex_provider}".lower() in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before + + +def test_case_06_managed_codex_rejects_provider_override_from_fresh_state( + live_session, workspace, codex_provider +): + """Scenario: pass --workspace and a --provider override from fresh state. + + Expected: ug rejects the override without changing agent-owned state or files, apart from the + managed-config retrieval cache. + """ + session = live_session + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--workspace", + workspace, + "--provider", + codex_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert f"provider {codex_provider}".lower() in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before + + +def test_case_08_managed_codex_rejects_model_location_override_after_configure( + live_session, workspace, parent_schema +): + """Scenario: configure managed Codex, then pass a --model-location override. + + Expected: ug rejects the override without changing agent-owned state or files. + """ + session = live_session + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert "--model-location" in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before + + +def test_case_08_managed_codex_rejects_model_location_override_from_fresh_state( + live_session, workspace, parent_schema +): + """Scenario: pass --workspace and a --model-location override from fresh state. + + Expected: ug rejects the override without changing agent-owned state or files, apart from the + managed-config retrieval cache. + """ + session = live_session + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert "--model-location" in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before + + +def test_case_10_managed_codex_rejects_provider_when_discovery_disabled_after_configure( + live_session, workspace, codex_provider +): + """Scenario: configure managed Codex, disable discovery, then pass a --provider override. + + Expected: ug rejects the override without changing agent-owned state or files. + """ + session = live_session + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--provider", + codex_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert f"provider {codex_provider}".lower() in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before + + +def test_case_10_managed_codex_rejects_provider_when_discovery_disabled_from_fresh_state( + live_session, workspace, codex_provider +): + """Scenario: disable discovery and pass --workspace plus --provider from fresh state. + + Expected: ug rejects the override without changing agent-owned state or files, apart from the + managed-config retrieval cache. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--workspace", + workspace, + "--provider", + codex_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert f"provider {codex_provider}".lower() in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before + + +def test_case_12_managed_codex_rejects_model_location_when_discovery_disabled_after_configure( + live_session, workspace, parent_schema +): + """Scenario: configure managed Codex, disable discovery, then pass --model-location. + + Expected: ug rejects the override without changing agent-owned state or files. + """ + session = live_session + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert "--model-location" in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before + + +def test_case_12_managed_codex_rejects_model_location_when_discovery_disabled_from_fresh_state( + live_session, workspace, parent_schema +): + """Scenario: disable discovery and pass --workspace plus --model-location from fresh state. + + Expected: ug rejects the override without changing agent-owned state or files, apart from the + managed-config retrieval cache. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + before = _codex_state_and_agent_files(session) + + result = session.run( + "codex", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert "--model-location" in output + assert "admin has specified managed" in output + assert MANAGED_CODEX_PROVIDER_SERVICE.lower() in output + assert _codex_state_and_agent_files(session) == before diff --git a/tests/integration/test_ug_configure_managed.py b/tests/integration/test_ug_configure_managed.py index 759f2fd7..b1bc8b6b 100644 --- a/tests/integration/test_ug_configure_managed.py +++ b/tests/integration/test_ug_configure_managed.py @@ -11,15 +11,9 @@ import json import pytest +from utils.constants import MANAGED_CLAUDE_MODELS, MANAGED_CODEX_MODELS from utils.terminal import AgentTerminal -MANAGED_CLAUDE_MODELS = [ - "system.ai.claude-opus-4-8", - "system.ai.claude-sonnet-4-6", - "system.ai.claude-haiku-4-5", -] -MANAGED_CODEX_MODEL = "system.ai.gpt-5-6-sol" - @pytest.mark.managed @pytest.mark.claude @@ -67,7 +61,7 @@ def test_ug_configure_managed_codex(live_session, workspace): for model in catalog.get("models", []) if model.get("visibility") == "list" ] - assert listed == [MANAGED_CODEX_MODEL], catalog + assert listed == MANAGED_CODEX_MODELS, catalog with AgentTerminal(session, "codex", [str(session.binary), "codex"], "managed-codex") as tui: tui.boot() @@ -93,7 +87,7 @@ def test_ug_configure_managed_is_idempotent(live_session, workspace): listed = [m.get("slug") for m in catalog.get("models", []) if m.get("visibility") == "list"] runs.append((settings.get("availableModels"), picker, listed)) - expected = (MANAGED_CLAUDE_MODELS, MANAGED_CLAUDE_MODELS, [MANAGED_CODEX_MODEL]) + expected = (MANAGED_CLAUDE_MODELS, MANAGED_CLAUDE_MODELS, MANAGED_CODEX_MODELS) assert runs == [expected, expected], runs diff --git a/tests/integration/utils/constants.py b/tests/integration/utils/constants.py index 0bc7e866..d48ce82b 100644 --- a/tests/integration/utils/constants.py +++ b/tests/integration/utils/constants.py @@ -1,3 +1,14 @@ """Shared constants for the integration CUJs.""" CODEX_TEST_MODEL = "system.ai.gpt-5-4-nano" + +MANAGED_CLAUDE_MODELS = [ + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5", +] +MANAGED_CODEX_MODELS = ["system.ai.gpt-5-6-sol"] + +# Dedicated ca-central-1 services used only by the managed-discovery fixture variants. +MANAGED_CLAUDE_PROVIDER_SERVICE = "main.default.ci_e2e_anthropic_mps" +MANAGED_CODEX_PROVIDER_SERVICE = "main.default.ci_e2e_openai_mps" diff --git a/tests/integration/utils/harness.py b/tests/integration/utils/harness.py index db663082..9644bf4c 100644 --- a/tests/integration/utils/harness.py +++ b/tests/integration/utils/harness.py @@ -192,7 +192,13 @@ def assert_not_routed(self) -> None: for name in ("codex-v2-interposer.log", "claude-v2-pty.log"): assert not (self.home / ".ucode" / name).exists(), f"Unexpected routing: {name}" - def app_server_handshake(self, args: list[str], timeout: int = 120) -> dict: + def app_server_handshake( + self, + args: list[str], + timeout: int = 120, + request: tuple[str, dict] | None = None, + name: str = "app-server", + ) -> dict: """Speak the real Codex stdio protocol and require an initialize response.""" command = [str(self.binary), "codex", *args] messages: queue.Queue = queue.Queue() @@ -229,37 +235,47 @@ def read_diagnostics(): reader.start() stderr_reader.start() try: - proc.stdin.write( - json.dumps( - { - "id": 1, - "method": "initialize", - "params": {"clientInfo": {"name": "ug-integration", "version": "1.0.0"}}, - } + + def send(message): + proc.stdin.write(json.dumps(message) + "\n") + proc.stdin.flush() + + def wait_for_response(request_id, description): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + message = messages.get(timeout=max(0.01, deadline - time.monotonic())) + except queue.Empty: + break + if message is None: + break + assert isinstance(message, dict), message + assert "protocol_error" not in message, ( + "Non-JSON output on the app-server protocol stream: " + str(message) + ) + if message.get("id") == request_id: + assert "error" not in message, message + assert isinstance(message.get("result"), dict), message + return message + raise AssertionError( + f"No app-server {description} response:\n" + "".join(transcript) ) - + "\n" + + send( + { + "id": 1, + "method": "initialize", + "params": {"clientInfo": {"name": "ug-integration", "version": "1.0.0"}}, + } ) - proc.stdin.flush() - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - try: - message = messages.get(timeout=max(0.01, deadline - time.monotonic())) - except queue.Empty: - break - if message is None: - break - assert isinstance(message, dict), message - assert "protocol_error" not in message, ( - "Non-JSON output on the app-server protocol stream: " + str(message) - ) - if isinstance(message, dict) and message.get("id") == 1: - assert "error" not in message, message - assert isinstance(message.get("result"), dict), message - assert message["result"].get("userAgent"), message - proc.stdin.write('{"method":"initialized","params":{}}\n') - proc.stdin.flush() - return message - raise AssertionError("No app-server initialize response:\n" + "".join(transcript)) + initialized = wait_for_response(1, "initialize") + assert initialized["result"].get("userAgent"), initialized + send({"method": "initialized", "params": {}}) + if request is None: + return initialized + method, params = request + send({"id": 2, "method": method, "params": params}) + return wait_for_response(2, method) finally: stop_process(proc) reader.join(timeout=5) @@ -268,5 +284,26 @@ def read_diagnostics(): proc.stdout.close() proc.stderr.close() self.record( - "app-server.json", {"argv": command, "stdout": transcript, "stderr": diagnostics} + f"{name}.json", {"argv": command, "stdout": transcript, "stderr": diagnostics} ) + + def codex_model_ids(self, args: list[str], name: str = "codex-models") -> list[str]: + """Ask the real Codex app-server for the catalog its model picker uses.""" + response = self.app_server_handshake( + args, + request=( + "model/list", + {"cursor": None, "limit": 1000, "includeHidden": False}, + ), + name=f"{name}-app-server", + ) + result = response["result"] + models = result.get("data") + assert isinstance(models, list), response + assert result.get("nextCursor") is None, "Codex model catalog exceeded the test page size" + ids = [model.get("model") for model in models if isinstance(model, dict)] + assert len(ids) == len(models) and all(isinstance(model_id, str) for model_id in ids), ( + response + ) + self.record(f"{name}.json", response) + return ids diff --git a/tests/integration/utils/managed.py b/tests/integration/utils/managed.py index 9c5c6684..9c132223 100644 --- a/tests/integration/utils/managed.py +++ b/tests/integration/utils/managed.py @@ -1,20 +1,79 @@ -"""Shared builders and stub injection for the managed-config integration suites. +"""Shared fetching, builders, and stub injection for managed-config integration suites. -Pure CodingAgentConfig construction plus the ``UCODE_MANAGED_CONFIG_STUB`` filesystem/env mechanics. The configure invocation, launch, and assertions stay visible in each test (tests/AGENTS.md), and this module imports nothing from the ``ucode`` application package (tests/test_integration_contract.py enforces that boundary). """ import json +import urllib.request from pathlib import Path +MANAGED_CONFIGS_PATH = "/api/ai-gateway/v2/coding-agent-configs" + + +def fetch_managed_config_stub( + workspace: str, + token: str, + directory: Path, + filename: str, + *, + agent: str, + provider_service: str, +) -> Path: + """Fetch the workspace config and persist one agent's MPS-backed test variant.""" + request = urllib.request.Request( + workspace.rstrip("/") + MANAGED_CONFIGS_PATH, + headers={"Authorization": f"Bearer {token}", "Accept": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=30) as response: + payload = json.load(response) + + if isinstance(payload, dict): + configs = payload.get("coding_agent_configs") + elif isinstance(payload, list): + configs = payload + else: + configs = None + assert isinstance(configs, list) and configs, "workspace returned no managed CodingAgentConfig" + config = json.loads(json.dumps(configs[0])) + assert isinstance(config, dict), "managed CodingAgentConfig was not an object" + + enabled_agents = config.get("enabled_agents") + assert isinstance(enabled_agents, list), "managed CodingAgentConfig had no enabled_agents list" + matching_agents = [ + entry for entry in enabled_agents if isinstance(entry, dict) and entry.get("agent") == agent + ] + assert len(matching_agents) == 1, f"expected exactly one {agent} entry" + agent_config = matching_agents[0].get("config") + assert isinstance(agent_config, dict), f"{agent} had no managed agent config" + + # Keep the published policy as the fixture base, but exercise the dedicated MPS source for this + # module. The published defaults name Hosted models, so they cannot accompany an MPS oneof. + agent_config["models"] = {"model_provider_service": provider_service} + agent_config.pop("default_models", None) + + stub = directory / filename + stub.write_text(json.dumps(config), encoding="utf-8") + stub.chmod(0o600) + return stub + + +def use_managed_config_stub(session, stub: Path) -> None: + """Point one isolated session at an already-persisted raw CodingAgentConfig.""" + session.env["UCODE_MANAGED_CONFIG_STUB"] = str(stub) + def set_managed_config_stub(session, tmp_path, config: dict) -> None: """Write ``config`` to a file and point ``UCODE_MANAGED_CONFIG_STUB`` at it for this session.""" stub = Path(tmp_path) / "managed-config.json" stub.write_text(json.dumps(config)) - session.env["UCODE_MANAGED_CONFIG_STUB"] = str(stub) + use_managed_config_stub(session, stub) + + +def is_managed_config_control_plane_cache(home: Path, path: Path) -> bool: + """Whether ``path`` is ug's expected fetched-config cache, not agent-owned state.""" + return path == home / ".ucode" / "managed-config.json" def build_coding_agent_config( diff --git a/tests/integration/utils/terminal.py b/tests/integration/utils/terminal.py index b49da7a8..4bf092fb 100644 --- a/tests/integration/utils/terminal.py +++ b/tests/integration/utils/terminal.py @@ -289,6 +289,23 @@ def check_input_and_exit(self): self.wait_for(lambda text: marker not in text, "cleared prompt") self.exit_normally() + def open_model_picker(self): + """Open Claude's real model picker, record it, then return to the prompt.""" + self.submit("/model") + self.wait_for( + lambda text: "Select model" in text and "Switch between Claude models." in text, + "the model picker", + timeout=60, + ) + screen = self.visible + self.actions.append({"reason": "model-picker-visible", "screen": screen}) + self.send("\x1b", "close the model picker") + self.wait_for( + lambda text: "Select model" not in text, + "the prompt after closing the model picker", + ) + return screen + def wait_for_task(self, task, timeout=180): permission_in_progress = False diff --git a/tests/test_integration_contract.py b/tests/test_integration_contract.py index def610a9..f0499ec1 100644 --- a/tests/test_integration_contract.py +++ b/tests/test_integration_contract.py @@ -67,7 +67,7 @@ def test_live_integration_cases_belong_to_exactly_one_ci_agent(): for node in tree.body: if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"): marks = module_marks | _markers(node.decorator_list) - if marks & {"live", "managed"}: + if marks & {"live", "managed", "managed_fixture"}: assert len(marks & {"claude", "codex"}) == 1, node.name