Add model discovery integration tests with managed config

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,
             }

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.

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`

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://<managed-workspace>
 export DATABRICKS_CLIENT_ID=<sp-app-id> DATABRICKS_CLIENT_SECRET=<sp-oauth-secret>
-python scripts/run_integration.py --claude-version <v> --codex-version <v> -- -m managed
+python scripts/run_integration.py --claude-version <v> --codex-version <v> \
+  -- -m "managed or managed_fixture"
 ```
 
 Each job uses fresh consumer dependency resolution. There is no default dependency

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"]