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/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index f577eb39..712f82b8 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -392,9 +392,9 @@ def configure_tool( state, model, provider=provider, parent_schema=parent_schema ) elif tool == "claude": - # A Model Provider Service routes by header and pins no Databricks - # model, so the usual "model required" guard doesn't apply to claude. - if not model and not provider: + # A Model Provider Service or parent schema routes by header and discovers models natively, + # so the usual "model required" guard doesn't apply to either Claude source. + if not model and not provider and not parent_schema: raise RuntimeError(f"A {tool} model must be selected before configuration.") result = claude.write_tool_config( state, @@ -497,12 +497,12 @@ def _availability_failure_detail(tool: str, state: dict) -> str: return " (" + "; ".join(parts) + ")" -def configure_single_tool(tool: str, state: dict) -> dict: +def configure_single_tool(tool: str, state: dict, *, parent_schema: str | None = None) -> dict: """Check availability, configure, and persist state for one tool only.""" - provider = get_provider_service(state, tool) - # A Model Provider Service routes through the same gateway and pins no - # Databricks model, so the per-tool model availability check doesn't apply. - if not provider: + provider = None if parent_schema else get_provider_service(state, tool) + # A Model Provider Service or parent schema routes through the same gateway and pins no + # globally discovered Databricks model, so the availability check doesn't apply. + if not provider and not parent_schema: with spinner(f"Checking {TOOL_SPECS[tool]['display']} availability..."): ok = check_gateway_endpoint(state, tool) if not ok: @@ -511,15 +511,19 @@ def configure_single_tool(tool: str, state: dict) -> dict: f"{TOOL_SPECS[tool]['display']} is not available on this workspace.{detail}" ) with managed_write_batch(_managed_settings_displays([tool])): - state = _configure_one(tool, state, provider) + state = _configure_one(tool, state, provider, parent_schema=parent_schema) available_tools = list(set((state.get("available_tools") or []) + [tool])) state["available_tools"] = available_tools save_state(state) return state -def _configure_one(tool: str, state: dict, provider: str | None) -> dict: +def _configure_one( + tool: str, state: dict, provider: str | None, *, parent_schema: str | None = None +) -> dict: """Write one tool's config, routing through ``provider`` when set.""" + if parent_schema: + return configure_tool(tool, state, parent_schema=parent_schema) if provider: if tool == "gemini": # Gemini pins a concrete target in the URL, so configure must resolve one now — @@ -542,7 +546,11 @@ def _configure_one(tool: str, state: dict, provider: str | None) -> dict: def configure_selected_tools( - state: dict, tools: list[str], *, install_ai_tools: bool = True + state: dict, + tools: list[str], + *, + install_ai_tools: bool = True, + parent_schemas: dict[str, str] | None = None, ) -> dict: """Configure the given tools. Caller is responsible for ensuring each tool is available on the workspace. @@ -553,7 +561,9 @@ def configure_selected_tools( """ with managed_write_batch(_managed_settings_displays(tools)): for tool in tools: - state = _configure_one(tool, state, get_provider_service(state, tool)) + parent_schema = (parent_schemas or {}).get(tool) + provider = None if parent_schema else get_provider_service(state, tool) + state = _configure_one(tool, state, provider, parent_schema=parent_schema) existing = state.get("available_tools") or [] state["available_tools"] = sorted(set(existing) | set(tools)) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 9c927553..ad8f18c8 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -447,7 +447,7 @@ def render_overlay( # provider and Claude Code's own canonical model names are sent verbatim — # pinning a Databricks model id here would mislabel the picker and isn't # routable. - elif claude_models and not provider: + elif claude_models and not provider and not parent_schema: # Picker rows show the raw routable id (e.g. "system.ai.claude-opus-4-8[1m]") # so users can see which gateway-routable model is behind each shortcut. # We deliberately don't set the `_NAME` companion env vars — the raw id @@ -486,7 +486,7 @@ def render_overlay( overlay["permissions"] = {"deny": ["WebSearch"]} keys.append(["permissions", "deny"]) - if static_models and not provider and not relayed: + if static_models and not provider and not parent_schema and not relayed: overlay["availableModels"] = list(static_models) overlay["enforceAvailableModels"] = True overlay["modelPicker"] = { @@ -865,6 +865,14 @@ def write_tool_config( # revert would restore that snapshot instead of deleting the file. if not is_tool_managed(state, "claude"): backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) + previous_keys = ((state.get("managed_configs") or {}).get("claude") or {}).get("keys", []) + # Native discovery must not inherit UG's prior static allow-list. Untracked picker values + # remain user/admin-owned and are preserved by the merge below. + stale_picker_keys = ( + [key for key in CLAUDE_MANAGED_PICKER_KEYS if [key] in previous_keys] + if provider or parent_schema + else [] + ) web_search_model = _resolve_web_search_model(state) # Relayed inference points at a local refresh proxy; its loopback base URL is # recorded in state so launch starts the proxy on the matching port. @@ -889,7 +897,8 @@ def write_tool_config( ) managed_file_keys = list(managed_keys) for path in ( - [["env", key] for key in CLAUDE_MANAGED_MODEL_ENV_KEYS] + [[key] for key in stale_picker_keys] + + [["env", key] for key in CLAUDE_MANAGED_MODEL_ENV_KEYS] + [["env", key] for key in CLAUDE_CONDITIONAL_ENV_KEYS] + [["env", key] for key in CLAUDE_REMOVED_ENV_KEYS] + [["env", key] for key in CLAUDE_OTEL_TRACE_ENV_KEYS] @@ -953,6 +962,8 @@ def _compose( else: target_env[key] = selected_default_model merged = deep_merge_dict(base, overlay_for_merge) + for key in stale_picker_keys: + merged.pop(key, None) overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( existing_custom_headers, overlay_custom_headers @@ -1013,7 +1024,7 @@ def _compose( state, lambda base: _compose( base, - enforce_model_default_hierarchy=provider is None, + enforce_model_default_hierarchy=provider is None and parent_schema is None, managed_settings_snapshots=managed_snapshots, ), managed_file_keys, diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 20bb5abd..8f30d63b 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -917,10 +917,10 @@ def launch( binary = SPEC["binary"] workspace = state.get("workspace") launch_provider = state.get("_codex_launch_provider") - provider = ( + transient_provider = ( launch_provider.strip() if isinstance(launch_provider, str) and launch_provider.strip() - else get_provider_service(state, "codex") + else None ) launch_parent_schema = state.get("_codex_launch_parent_schema") parent_schema = ( @@ -928,6 +928,11 @@ def launch( if isinstance(launch_parent_schema, str) and launch_parent_schema.strip() else None ) + # Launch-scoped admin routing wins over persisted developer configuration. A transient provider + # is most specific; otherwise a transient UC parent must suppress a saved provider. + provider = transient_provider or ( + None if parent_schema else get_provider_service(state, "codex") + ) if workspace and (provider or parent_schema): _reject_managed_model_catalog() token = None diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 83d193f8..a7684e5f 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -91,6 +91,7 @@ managed_provider_family_models, managed_provider_service, managed_supplies_models, + managed_unity_catalog_location, managed_unservable_models, recommended_agent, resolve_state, @@ -775,16 +776,25 @@ def configure_workspace_command( clear_custom_oauth=custom_oauth is None, ) state = states[0] - state = configure_single_tool(tool, state) - # No managed refresh precedes this branch, so read fresh here: `ug configure` never decides - # from a stale cache. - install_databricks_ai_tools_for_agents([tool], state, force_refresh=True) + parent_schema = None + if tool in ("claude", "codex"): + managed, _ = refresh_managed_config(state, force_refresh=True) + _reject_disabled_agent(managed, tool) + if managed is not None: + state = resolve_state(managed, state, tool) + if not managed_provider_service(managed, tool): + parent_schema = managed_unity_catalog_location(managed, tool) + state = configure_single_tool(tool, state, parent_schema=parent_schema) + install_databricks_ai_tools_for_agents( + [tool], state, force_refresh=tool not in ("claude", "codex") + ) spec = TOOL_SPECS[tool] + provider_summary = "Databricks" if parent_schema else _provider_summary(tool, state) console.print( Panel( f"[bold]Workspace:[/bold] [cyan]{state['workspace']}[/cyan]\n" f"[bold]{spec['display']}:[/bold] [green]configured[/green] " - f"[dim](Provider: {_provider_summary(tool, state)})[/dim]", + f"[dim](Provider: {provider_summary})[/dim]", title="Configuration Complete", style="green", expand=False, @@ -812,13 +822,25 @@ def configure_workspace_command( if managed is not None and managed_tools: configured_tools: list[str] = [] for tool_name in managed_tools: - if check_gateway_endpoint(state, tool_name): + resolved = resolve_state(managed, state, tool_name) + parent_schema = ( + managed_unity_catalog_location(managed, tool_name) + if tool_name in ("claude", "codex") + and not managed_provider_service(managed, tool_name) + else None + ) + if ( + get_provider_service(resolved, tool_name) + or parent_schema + or check_gateway_endpoint(resolved, tool_name) + ): if not install_tool_binary(tool_name, strict=False): continue configured = configure_selected_tools( - resolve_state(managed, state, tool_name), + resolved, [tool_name], install_ai_tools=not is_dry_run(), + parent_schemas={tool_name: parent_schema} if parent_schema else None, ) # Each iteration resolves from `state` and persists a copy, so carry the # accumulated available_tools forward — otherwise the last agent's save drops @@ -2337,6 +2359,22 @@ def _launch_tool( raise RuntimeError("--model-location must be `.`.") # Checked before discovery, which can take tens of seconds, so a blocked launch fails fast. _reject_disabled_agent(managed, tool) + managed_provider = managed_provider_service(managed or {}, tool) + managed_parent_schema = ( + managed_unity_catalog_location(managed or {}, tool) + if tool in {"claude", "codex"} and not managed_provider + else None + ) + if managed_provider: + provider = managed_provider + parent_schema = None + elif managed_parent_schema: + # Managed UC discovery supersedes a developer's persisted provider without + # rewriting it; the admin's location exists only for this launch. + provider = None + parent_schema = managed_parent_schema + if tool == "claude" and (managed_provider or managed_parent_schema): + os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" # The environment switch remains a developer override; managed config is the workspace # policy equivalent and must take effect before launch options are computed. managed_smart_routing_enabled = _managed_smart_routing_enabled(managed, tool) @@ -2353,7 +2391,9 @@ def _launch_tool( state["workspace"], profile=state.get("profile"), tools=[tool], - skip_model_discovery=bool(provider) or managed_models_known, + skip_model_discovery=( + bool(provider) or bool(managed_parent_schema) or managed_models_known + ), skip_preflight=skip_preflight, **configure_kwargs, ) @@ -2374,10 +2414,6 @@ def _launch_tool( ) elif not coding_agent_config_feature_disabled: print_note("No managed coding agent config found; using your own settings") - if managed is not None: - managed_provider = managed_provider_service(managed, tool) - if managed_provider: - provider = managed_provider if provider and parent_schema is not None: raise RuntimeError("--provider and --model-location cannot be used together.") # Checked after the managed config settles `provider`: an admin-set provider must trip this @@ -2428,13 +2464,12 @@ def _launch_tool( route_root_model = None managed_model = None relayed_forward_model = None # forwarded to Claude Code's --model for a relayed provider - if provider: + if provider or managed_parent_schema: # Routing through a Model Provider Service pins no Databricks model; - # the agent uses its own canonical model names (header selects the - # provider). Skip model resolution, which would otherwise fail when - # the workspace has no matching Databricks models. + # managed UC discovery likewise lets the agent select from the parent schema. Skip model + # resolution, which would otherwise fail when global discovery found no models. resolved_model = None - if tool == "claude" and (model or provider_models): + if provider and tool == "claude" and (model or provider_models): if relayed: # Resolve against a curated allowlist so the forwarded id is one the gateway # allows; an allow_all relay declares none, so forward as-is. @@ -2445,7 +2480,7 @@ def _launch_tool( ) else: route_root_model = resolve_provider_launch_model(model, provider_models or {}) - if tool == "gemini": + if provider and tool == "gemini": # Gemini is the exception: the request still names a concrete model # in the URL, so pin one of the service's targets (--model or default). resolved_model, gemini_error = resolve_gemini_provider_model(state, provider, model) diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index c249c698..82d1427d 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -190,6 +190,11 @@ def managed_provider_service(managed: dict, tool: str) -> str | None: return _str(_agent_model_config(managed, tool).get("model_provider_service")) +def managed_unity_catalog_location(managed: dict, tool: str) -> str | None: + """Return only the Unity Catalog model location managed config specifies for ``tool``.""" + return _str(_agent_model_config(managed, tool).get("unity_catalog_location")) + + def managed_static_models(managed: dict, tool: str) -> list[str] | None: """The explicit model allow-list (``model_config.model_services``) the config sets for ``tool``. 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 7150f84b..42d62593 100644 --- a/tests/README.md +++ b/tests/README.md @@ -51,6 +51,12 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_configure_claude_repeat_and_revert`, `test_ug_configure_codex_repeat_and_revert` | Configure twice over user settings; complete a task; revert twice | Settings preserved; no bearer in ug state; generated config removed; status unconfigured | | `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_case_01_*`, `test_case_03_*` | Launch managed Claude after configure and from fresh state, with personal discovery enabled and disabled | Claude receives the admin MPS header, caches native discovery results, and opens its real model picker | +| `test_case_05_*`, `test_case_07_*` | Pass a provider or model-location override to managed Claude after configure and from fresh state | ug rejects the override before Claude starts and preserves agent-owned state | +| `test_case_09_*`, `test_case_11_*` | Disable discovery and pass a provider or model-location override to managed Claude | ug still rejects both configured and fresh launches | +| `test_case_02_*`, `test_case_04_*` | Launch managed Codex after configure and from fresh state, with personal discovery enabled and disabled | Codex exposes exactly the admin MPS-scoped catalog | +| `test_case_06_*`, `test_case_08_*` | Pass a provider or model-location override to managed Codex after configure and from fresh state | ug rejects the override before Codex starts and preserves agent-owned state | +| `test_case_10_*`, `test_case_12_*` | Disable discovery and pass a provider or model-location override to managed Codex | ug still rejects both configured and fresh launches | | `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_managed_fixture_claude_model_lifecycle`, `test_managed_fixture_codex_model_lifecycle` | Configure across no config -> static A -> static B -> MPS -> no config (stub-injected, `null` for no-config; MPS via a real provider service) | Each agent's model files reconcile to each static config (removed models pruned); switching to an MPS and a workspace with no managed config both clear ug's managed model settings so no stale list is enforced | | `test_ug_installed_wheel_exposes_help_and_version` | Invoke freshly installed console command | Package version matches; public help works | @@ -60,9 +66,12 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `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), **5 managed-fixture cases** (marker `managed_fixture`, with only -the CodingAgentConfig input injected), and **5 installation checks**. Parametrization varies +**4 managed-workspace cases** (marker `managed`, run against a separate workspace that +publishes a CodingAgentConfig), **34 managed-fixture cases** (marker `managed_fixture`, with only +the CodingAgentConfig input injected), and **5 installation checks**. The 24 numbered scenarios +cover configured and fresh state across the Claude and Codex managed-discovery matrix; ten +existing collected cases cover focused model, MCP, skills, and lifecycle 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. diff --git a/tests/integration/README.md b/tests/integration/README.md index c2bd85f4..8b5b09d9 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -98,8 +98,10 @@ 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_models.py # injected model sources and Codex fallback metadata test_ug_configure_managed_mcp.py # injected managed MCP list test_ug_configure_managed_skills.py # injected managed skills: download, coexist, reconcile away test_ug_configure_managed_lifecycle.py # none -> A -> B -> MPS -> none: reconcile, clear on MPS/no-config @@ -157,19 +159,14 @@ 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 **9 `managed_fixture` cases** inject the admin config -locally (via `UCODE_MANAGED_CONFIG_STUB`) to cover shapes the live workspace does not publish, -including a managed MCP server landing in Codex's OS-managed `[mcp_servers]` (interactive configure) -while the developer's own config stays untouched, and reaching Claude's `/mcp` view via the -user-scope fallback (non-interactive configure); each -differs from the published config in what it asserts so it proves the injected config drove configure. -Two of those are per-agent lifecycle journeys (no config -> static A -> static B -> MPS -> no config): -they assert the generated model files reconcile to each static config, are cleared when switching to -a Model Provider Service, and are cleared when the workspace has no managed config (an explicit -`null` stub reproduces the no-config states; the MPS states use real provider services on the managed -workspace, `main.default.ci_e2e_anthropic_mps` and `main.default.ci_e2e_openai_mps`). +checks** with both agents. A separate **4 managed-workspace cases** (one per agent, an idempotent +re-configure, and a cache-TTL journey; marker `managed`) run against a workspace that publishes a +CodingAgentConfig; see "Managed-workspace journeys" below. A further **34 `managed_fixture` +cases** use `UCODE_MANAGED_CONFIG_STUB`. Twenty-four explicit configured/fresh Claude and Codex +discovery and source-override journeys fetch the published config once per agent, replace that +agent's static source with its dedicated MPS, and reuse the result. Ten existing collected cases +cover focused model, MCP, skills, and lifecycle shapes, including per-agent model reconciliation +and managed skill cleanup. See the named coverage and gaps matrix in [../README.md](../README.md). @@ -289,7 +286,13 @@ enabled agents, or defaults) breaks these lanes until the constants in `test_ug_ 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, +managed-config HTTP read for config shapes that workspace does not publish. The Claude discovery +module fetches the workspace's published config once, replaces Claude's static model source with +`main.default.ci_e2e_anthropic_mps`, drops incompatible static defaults, and reuses that fixture +across all configured/fresh scenarios. The Codex module does the same with +`main.default.ci_e2e_openai_mps`. The tests verify Claude's admin header, native cache and real +model picker, Codex's exact app-server catalog, and both agents' rejection of personal source +overrides. In addition, `test_ug_configure_managed_codex_catalog_fallback` 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 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..fd00dfde --- /dev/null +++ b/tests/integration/test_ug_claude_managed_model_discovery.py @@ -0,0 +1,393 @@ +"""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, before=None): + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert "`--provider` or `--model-location` is not allowed" in output + assert "managed config exists for the workspace" in output + if before is not None: + 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, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Claude. + """ + session = live_session + result = session.run( + "claude", + "--workspace", + workspace, + "--provider", + claude_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result) + + +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, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Claude. + """ + session = live_session + result = session.run( + "claude", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result) + + +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, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Claude. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + result = session.run( + "claude", + "--workspace", + workspace, + "--provider", + claude_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result) + + +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, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Claude. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + result = session.run( + "claude", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_claude_started(session, result) 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..efa50080 --- /dev/null +++ b/tests/integration/test_ug_codex_managed_model_discovery.py @@ -0,0 +1,395 @@ +"""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_rejected_before_codex_started(session, result, before=None): + output = (result.stdout + result.stderr).lower() + assert result.returncode == 1 + assert "`--provider` or `--model-location` is not allowed" in output + assert "managed config exists for the workspace" in output + if before is not None: + assert _codex_state_and_agent_files(session) == before + + +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, + ) + + _assert_rejected_before_codex_started(session, result, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Codex. + """ + session = live_session + result = session.run( + "codex", + "--workspace", + workspace, + "--provider", + codex_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_codex_started(session, result) + + +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, + ) + + _assert_rejected_before_codex_started(session, result, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Codex. + """ + session = live_session + result = session.run( + "codex", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_codex_started(session, result) + + +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, + ) + + _assert_rejected_before_codex_started(session, result, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Codex. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + result = session.run( + "codex", + "--workspace", + workspace, + "--provider", + codex_provider, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_codex_started(session, result) + + +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, + ) + + _assert_rejected_before_codex_started(session, result, 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 may establish the fresh workspace/agent configuration, then rejects the + override before starting Codex. + """ + session = live_session + session.env["UG_ENABLE_MODEL_DISCOVERY"] = "0" + result = session.run( + "codex", + "--workspace", + workspace, + "--model-location", + parent_schema, + "--", + "--version", + ok=False, + timeout=240, + ) + + _assert_rejected_before_codex_started(session, result) diff --git a/tests/integration/utils/constants.py b/tests/integration/utils/constants.py index 0bc7e866..f5c6374c 100644 --- a/tests/integration/utils/constants.py +++ b/tests/integration/utils/constants.py @@ -1,3 +1,6 @@ """Shared constants for the integration CUJs.""" CODEX_TEST_MODEL = "system.ai.gpt-5-4-nano" + +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 60b7c521..4a7891a0 100644 --- a/tests/integration/utils/managed.py +++ b/tests/integration/utils/managed.py @@ -1,4 +1,4 @@ -"""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 @@ -7,8 +7,61 @@ """ 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" + + 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) -> None: """Write ``config`` to a file and point ``UCODE_MANAGED_CONFIG_STUB`` at it for this session. @@ -17,7 +70,12 @@ def set_managed_config_stub(session, tmp_path, config: dict | None) -> None: managed config (the stub's no-config state).""" 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 b2460d74..e86f607f 100644 --- a/tests/integration/utils/terminal.py +++ b/tests/integration/utils/terminal.py @@ -298,6 +298,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_agent_claude.py b/tests/test_agent_claude.py index 1e3810b3..d1f8b10a 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -372,10 +372,9 @@ def test_fable_not_pinned_under_provider(self): def test_provider_adds_routing_header(self): overlay, _ = claude.render_overlay(WS, "s4", provider="main.aarushi.aarushi-claude") - assert ( - "Databricks-Model-Provider-Service: main.aarushi.aarushi-claude" - in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] - ) + headers = overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert "Databricks-Model-Provider-Service: main.aarushi.aarushi-claude" in headers + assert "Databricks-Model-Service-Parent-Schema" not in headers def test_provider_skips_model_pinning(self): models = { @@ -396,11 +395,18 @@ def test_no_provider_header_without_flag(self): assert "Databricks-Model-Provider-Service" not in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] def test_parent_adds_discovery_header(self): - overlay, _ = claude.render_overlay(WS, "s4", parent_schema="main.default") - assert ( - "Databricks-Model-Service-Parent-Schema: main.default" - in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] + overlay, _ = claude.render_overlay( + WS, + "s4", + claude_models={"sonnet": "system.ai.claude-sonnet-4-6"}, + parent_schema="main.default", + static_models=["system.ai.claude-sonnet-4-6"], ) + headers = overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert "Databricks-Model-Service-Parent-Schema: main.default" in headers + assert "Databricks-Model-Provider-Service" not in headers + assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in overlay["env"] + assert "availableModels" not in overlay def test_bedrock_provider_pins_model_ids(self): provider_models = { @@ -961,7 +967,7 @@ def test_managed_file_updates_gateway_settings_without_changing_model_picker(sel self._patch(monkeypatch, private_writes, managed_writes, existing) state = {"workspace": WS, "codex_models": []} - claude.write_tool_config(state, "databricks-claude-sonnet-4") + claude.write_tool_config(state, "databricks-claude-sonnet-4", parent_schema="main.default") written = json.loads(managed_writes[0][1]) assert written["modelPicker"] == picker @@ -1053,6 +1059,56 @@ def test_managed_file_omits_workspace_defaults_for_provider(self, monkeypatch): env = json.loads(managed_writes[0][1])["env"] assert not set(claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS.values()) & env.keys() + def test_managed_file_omits_workspace_defaults_for_parent_schema(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + existing = { + str(FAKE_MANAGED_PATH): { + "env": {"ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-4-8"} + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = { + "workspace": WS, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + } + + claude.write_tool_config(state, None, parent_schema="main.default") + + env = json.loads(managed_writes[0][1])["env"] + assert not set(claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS.values()) & env.keys() + + def test_parent_schema_prunes_previous_static_picker(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + picker = { + "availableModels": ["system.ai.claude-opus-4-8"], + "enforceAvailableModels": True, + "modelPicker": {"replaceBuiltInOptions": True, "options": []}, + "companyPolicy": "keep", + } + existing = { + str(claude.CLAUDE_SETTINGS_PATH): picker, + str(FAKE_MANAGED_PATH): picker, + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = { + "workspace": WS, + "managed_configs": { + "claude": {"keys": [[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS]} + }, + } + + updated = claude.write_tool_config(state, None, parent_schema="main.default") + + for written in (private_writes[0][1], json.loads(managed_writes[0][1])): + assert not set(claude.CLAUDE_MANAGED_PICKER_KEYS) & written.keys() + assert written["companyPolicy"] == "keep" + assert not any( + [key] in updated["managed_configs"]["claude"]["keys"] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + ) + def test_managed_file_keeps_provider_model_pins(self, monkeypatch): private_writes: list = [] managed_writes: list = [] diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 997c01cc..2b16252d 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -913,6 +913,38 @@ def test_parent_discovery_uses_authoritative_catalog(self, tmp_path, monkeypatch ) assert 'Databricks-Model-Service-Parent-Schema = "main.default"' in parent_arg + def test_transient_parent_suppresses_persisted_provider(self, tmp_path, monkeypatch): + launches = self._patch(tmp_path, monkeypatch) + seen = {} + monkeypatch.setattr( + codex, "_model_catalog_path", lambda workspace, scope: tmp_path / "models.json" + ) + monkeypatch.setattr( + codex, + "_fetch_codex_model_catalog", + lambda workspace, token, **kwargs: seen.update(kwargs) or {"models": []}, + ) + + codex.launch( + { + "workspace": WS, + "provider_services": {"codex": "main.default.developer"}, + "_codex_launch_parent_schema": "main.managed", + }, + [], + options=LaunchOptions(), + ) + + assert seen == { + "source": codex.CodexCatalogSource.PARENT_SCHEMA, + "identifier": "main.managed", + } + provider_arg = next( + arg for arg in launches[0] if arg.startswith("model_providers.Databricks=") + ) + assert 'Databricks-Model-Service-Parent-Schema = "main.managed"' in provider_arg + assert "Databricks-Model-Provider-Service" not in provider_arg + def test_parent_discovery_refreshes_when_parent_changes(self, tmp_path, monkeypatch): launches = self._patch(tmp_path, monkeypatch) monkeypatch.setattr(codex, "CODEX_MODEL_CATALOG_PATH", tmp_path / "models.json") diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 6f3abbe7..2384f6a4 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -4,6 +4,7 @@ import subprocess from contextlib import contextmanager +from unittest.mock import MagicMock import pytest @@ -183,6 +184,23 @@ def test_configure_single_tool_does_not_install(self, monkeypatch): agents_mod.configure_single_tool("codex", {"codex_models": ["m"], "profile": "myprof"}) assert captured == {} + def test_managed_parent_skips_global_availability_and_writes_header(self, monkeypatch): + state = {"workspace": "https://x.databricks.com"} + monkeypatch.setattr( + agents_mod, + "check_gateway_endpoint", + lambda *_a: pytest.fail("managed parent must not require global model availability"), + ) + configure = MagicMock(return_value=state) + monkeypatch.setattr(agents_mod, "configure_tool", configure) + monkeypatch.setattr(agents_mod, "save_state", lambda _state: None) + + assert ( + agents_mod.configure_single_tool("claude", state, parent_schema="main.default") is state + ) + + configure.assert_called_once_with("claude", state, parent_schema="main.default") + def test_configure_selected_tools_triggers_install(self, monkeypatch): captured = self._stub_configure(monkeypatch) agents_mod.configure_selected_tools( @@ -739,7 +757,9 @@ def capture_batch(displays): yield monkeypatch.setattr(agents_mod, "managed_write_batch", capture_batch) - monkeypatch.setattr(agents_mod, "_configure_one", lambda tool, state, provider: state) + monkeypatch.setattr( + agents_mod, "_configure_one", lambda tool, state, provider, **kwargs: state + ) monkeypatch.setattr(agents_mod, "save_state", lambda state: None) monkeypatch.setattr(agents_mod, "install_databricks_ai_tools_for_agents", lambda *_: None) diff --git a/tests/test_cli.py b/tests/test_cli.py index 88aedbfd..c8a7bef7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -452,24 +452,33 @@ def _launch_policy_patches( *, persisted_provider: str | None = None, ): + launch_state = dict(MINIMAL_STATE) with ( patch("ucode.cli.ensure_bootstrap_dependencies"), - patch("ucode.cli.load_state", return_value=MINIMAL_STATE), - patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.load_state", return_value=launch_state), + patch("ucode.cli.ensure_provider_state", return_value=launch_state), patch("ucode.cli._fetch_managed_config", return_value=(managed, False)), patch("ucode.cli._fetch_budget_recommendation", return_value=None), - patch("ucode.cli.get_provider_service", return_value=persisted_provider), - patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE), - patch("ucode.cli.resolve_provider_models", return_value=(None, None, False)), + patch("ucode.cli.get_provider_service", return_value=persisted_provider) as get_provider, + patch("ucode.cli.configure_shared_state", return_value=launch_state) as shared, + patch( + "ucode.cli.resolve_provider_models", return_value=(None, None, False) + ) as resolve_provider, patch("ucode.cli.resolve_gemini_provider_model", return_value=("gemini-2.0-flash", None)), patch( "ucode.cli.resolve_launch_model", - return_value=(MINIMAL_STATE, "databricks-claude-sonnet-4"), + return_value=(launch_state, "databricks-claude-sonnet-4"), ), - patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE), + patch("ucode.cli.configure_tool", return_value=launch_state) as configure, patch("ucode.cli.launch_agent") as launch, ): - yield launch + yield { + "get_provider": get_provider, + "shared": shared, + "resolve_provider": resolve_provider, + "configure": configure, + "launch": launch, + } class TestSubcommandRouting: @@ -979,19 +988,162 @@ class TestManagedConfigLaunchSourceGuard: ], ) def test_managed_config_rejects_launch_source_options(self, tool, option, value): - with _launch_policy_patches({}) as launch: + with _launch_policy_patches({}) as calls: result = runner.invoke(app, [tool, option, value]) assert result.exit_code == 1 assert "`--provider` or `--model-location` is not allowed" in _strip_ansi(result.output) - launch.assert_not_called() + calls["launch"].assert_not_called() def test_persisted_provider_is_not_mistaken_for_an_explicit_option(self): - with _launch_policy_patches({}, persisted_provider="main.default.provider") as launch: + with _launch_policy_patches({}, persisted_provider="main.default.provider") as calls: result = runner.invoke(app, ["claude"]) assert result.exit_code == 0, result.output - launch.assert_called_once() + assert calls["get_provider"].call_count == 1 + assert calls["get_provider"].call_args.args[1] == "claude" + calls["resolve_provider"].assert_called_once() + calls["configure"].assert_called_once() + calls["launch"].assert_called_once() + + +class TestManagedClaudeModelDiscovery: + MPS_CONFIG = { + "enabled_agents": { + "claude": { + "model_config": { + "model_provider_service": "main.default.anthropic-mps", + "models": { + "default_sonnet_model": "anthropic.claude-sonnet-4-6", + "default_opus_model": "anthropic.claude-opus-4-8", + "default_haiku_model": "anthropic.claude-haiku-4-5", + "default_fable_model": "anthropic.claude-fable-5-1", + }, + "default_model": "anthropic.claude-sonnet-4-6", + } + } + } + } + UC_CONFIG = { + "enabled_agents": {"claude": {"model_config": {"unity_catalog_location": "main.default"}}} + } + + @staticmethod + def _invoke(monkeypatch, managed): + state = { + **MINIMAL_STATE, + "claude_models": {}, + "provider_services": {"claude": "main.developer.provider"}, + } + shared = MagicMock(return_value=state) + configure = MagicMock(side_effect=lambda _tool, configured, *_args, **_kwargs: configured) + launch = MagicMock() + monkeypatch.setattr(cli_mod, "ensure_bootstrap_dependencies", lambda *_a, **_k: None) + monkeypatch.setattr(cli_mod, "load_state", lambda: state) + monkeypatch.setattr(cli_mod, "ensure_provider_state", lambda *_a: state) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda _state: (managed, False)) + monkeypatch.setattr(cli_mod, "get_provider_service", lambda *_a: "main.developer.provider") + monkeypatch.setattr(cli_mod, "configure_shared_state", shared) + resolve_provider = MagicMock(return_value=(None, None, False)) + monkeypatch.setattr(cli_mod, "resolve_provider_models", resolve_provider) + resolve_model = MagicMock(side_effect=AssertionError("must use native discovery")) + monkeypatch.setattr(cli_mod, "resolve_launch_model", resolve_model) + monkeypatch.setattr(cli_mod, "configure_tool", configure) + monkeypatch.setattr(cli_mod, "launch_agent", launch) + + result = runner.invoke(app, ["claude"]) + return { + "result": result, + "state": state, + "shared": shared, + "configure": configure, + "launch": launch, + "resolve_provider": resolve_provider, + "resolve_model": resolve_model, + } + + @pytest.mark.parametrize( + ("managed", "expected_provider", "expected_parent"), + [ + (MPS_CONFIG, "main.default.anthropic-mps", None), + (UC_CONFIG, None, "main.default"), + ], + ids=["mps", "uc-parent"], + ) + def test_launch_uses_managed_source_and_native_discovery( + self, monkeypatch, managed, expected_provider, expected_parent + ): + calls = self._invoke(monkeypatch, managed) + + assert calls["result"].exit_code == 0, calls["result"].output + assert calls["shared"].call_args.kwargs["skip_model_discovery"] is True + calls["resolve_model"].assert_not_called() + if expected_provider: + assert calls["resolve_provider"].call_args.args[2] == expected_provider + else: + calls["resolve_provider"].assert_not_called() + assert calls["configure"].call_args.kwargs["provider"] == expected_provider + assert calls["configure"].call_args.kwargs["parent_schema"] == expected_parent + assert calls["configure"].call_args.args[1]["provider_services"]["claude"] == ( + expected_provider or "main.developer.provider" + ) + assert os.environ["ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"] == "1" + + +def test_claude_discovery_changes_do_not_break_other_managed_providers(): + managed = { + "enabled_agents": { + "gemini": {"model_config": {"model_provider_service": "main.default.gemini-mps"}} + } + } + with _launch_policy_patches(managed): + result = runner.invoke(app, ["gemini"]) + + assert result.exit_code == 0, result.output + assert "main.default.gemini-mps" in _strip_ansi(result.output) + + +class TestManagedCodexModelSource: + @pytest.mark.parametrize( + ("model_config", "expected_provider", "expected_parent"), + [ + ( + {"model_provider_service": "main.default.managed-mps"}, + "main.default.managed-mps", + None, + ), + ({"unity_catalog_location": "main.managed"}, None, "main.managed"), + ], + ids=["mps", "uc-parent"], + ) + def test_managed_source_overrides_saved_provider( + self, monkeypatch, model_config, expected_provider, expected_parent + ): + monkeypatch.delenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", raising=False) + managed = {"enabled_agents": {"codex": {"model_config": model_config}}} + + with _launch_policy_patches( + managed, + persisted_provider="main.default.developer", + ) as calls: + result = runner.invoke(app, ["codex"]) + + assert result.exit_code == 0, result.output + assert calls["shared"].call_args.kwargs["skip_model_discovery"] is True + if expected_provider: + calls["resolve_provider"].assert_called_once() + else: + calls["resolve_provider"].assert_not_called() + assert calls["configure"].call_args.kwargs["provider"] == expected_provider + assert calls["configure"].call_args.kwargs["parent_schema"] == expected_parent + launch_state = calls["launch"].call_args.args[1] + if expected_provider: + assert launch_state["_codex_launch_provider"] == expected_provider + assert "_codex_launch_parent_schema" not in launch_state + else: + assert launch_state["_codex_launch_parent_schema"] == expected_parent + assert "_codex_launch_provider" not in launch_state + assert "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" not in os.environ class TestClaudeModelFlag: @@ -3001,6 +3153,131 @@ def test_managed_config_applies_all_enabled_and_skips_selection(self, monkeypatc assert installed == ["claude", "codex"] assert configured == ["claude", "codex"] + @pytest.mark.parametrize( + ("model_config", "expected_provider", "expected_parent"), + [ + ( + {"model_provider_service": "main.default.anthropic-mps"}, + "main.default.anthropic-mps", + None, + ), + ({"unity_catalog_location": "main.models"}, None, "main.models"), + ], + ids=["mps", "uc-parent"], + ) + def test_managed_claude_source_configures_without_global_models( + self, monkeypatch, model_config, expected_provider, expected_parent + ): + state = { + **MINIMAL_STATE, + "available_tools": [], + "claude_models": {}, + "provider_services": {"claude": "main.default.developer"}, + } + managed = {"enabled_agents": {"claude": {"model_config": model_config}}} + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s, **_k: (managed, False)) + monkeypatch.setattr( + cli_mod, + "check_gateway_endpoint", + lambda *_a: pytest.fail("managed model sources do not require global models"), + ) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr(cli_mod, "_print_managed_summary", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "_configure_managed_mcp_servers", lambda *_a: None) + configured: list[tuple[dict, dict]] = [] + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda resolved, tools, **kwargs: configured.append((resolved, kwargs)) or resolved, + ) + + assert cli_mod.configure_workspace_command(workspaces=[("https://w.com", None)]) == 0 + + assert len(configured) == 1 + resolved, kwargs = configured[0] + expected_saved_provider = expected_provider or "main.default.developer" + assert cli_mod.get_provider_service(resolved, "claude") == expected_saved_provider + assert kwargs["parent_schemas"] == ( + {"claude": expected_parent} if expected_parent else None + ) + + def test_single_claude_agent_passes_managed_uc_parent_directly(self, monkeypatch): + state = {**MINIMAL_STATE, "provider_services": {"claude": "main.default.developer"}} + managed = { + "enabled_agents": { + "claude": {"model_config": {"unity_catalog_location": "main.models"}} + } + } + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda *a, **k: (managed, False)) + configure = MagicMock(return_value=state) + monkeypatch.setattr(cli_mod, "configure_single_tool", configure) + monkeypatch.setattr(cli_mod, "install_databricks_ai_tools_for_agents", lambda *a, **k: None) + + assert ( + cli_mod.configure_workspace_command(tool="claude", workspaces=[("https://w.com", None)]) + == 0 + ) + configure.assert_called_once_with("claude", state, parent_schema="main.models") + + def test_managed_codex_parent_is_passed_to_generic_configure(self, monkeypatch): + state = { + **MINIMAL_STATE, + "available_tools": [], + "codex_models": [], + "provider_services": {"codex": "main.default.developer"}, + } + managed = { + "enabled_agents": {"codex": {"model_config": {"unity_catalog_location": "main.models"}}} + } + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "refresh_managed_config", lambda s, **_k: (managed, False)) + monkeypatch.setattr( + cli_mod, + "check_gateway_endpoint", + lambda *_a: pytest.fail("managed model sources do not require global models"), + ) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr(cli_mod, "_print_managed_summary", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "_configure_managed_mcp_servers", lambda *_a: None) + configured: list[dict] = [] + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda resolved, tools, **kwargs: configured.append(kwargs) or resolved, + ) + + assert cli_mod.configure_workspace_command(workspaces=[("https://w.com", None)]) == 0 + + assert configured[0]["parent_schemas"] == {"codex": "main.models"} + + def test_single_codex_agent_passes_managed_parent_directly(self, monkeypatch): + state = { + **MINIMAL_STATE, + "provider_services": {"codex": "main.default.developer"}, + } + managed = { + "enabled_agents": {"codex": {"model_config": {"unity_catalog_location": "main.models"}}} + } + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + refresh = MagicMock(return_value=(managed, False)) + monkeypatch.setattr(cli_mod, "refresh_managed_config", refresh) + configure = MagicMock(return_value=state) + monkeypatch.setattr(cli_mod, "configure_single_tool", configure) + install_ai_tools = MagicMock() + monkeypatch.setattr(cli_mod, "install_databricks_ai_tools_for_agents", install_ai_tools) + + result = runner.invoke( + app, ["configure", "--agent", "codex", "--workspace", "https://w.com"] + ) + + assert result.exit_code == 0, result.output + assert "(Provider: Databricks)" in _strip_ansi(result.output) + refresh.assert_called_once_with(state, force_refresh=True) + configure.assert_called_once_with("codex", state, parent_schema="main.models") + install_ai_tools.assert_called_once_with(["codex"], state, force_refresh=False) + def test_managed_config_fails_when_no_enabled_agent_is_available(self, monkeypatch): import ucode.cli as cli_mod diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index 867de657..8fc832b0 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -283,6 +283,7 @@ def test_configure_without_custom_options_resets_custom_oauth(self): state = {"workspace": WS, "available_tools": ["claude"]} with ( patch("ucode.cli._configure_shared_workspace_states", return_value=[state]) as shared, + patch("ucode.cli.refresh_managed_config", return_value=(None, False)), patch("ucode.cli.configure_single_tool", return_value=state), patch("ucode.cli.install_databricks_ai_tools_for_agents"), ): diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index f7a2a2f7..b82a026d 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -20,6 +20,7 @@ managed_state_overrides, managed_static_models, managed_supplies_models, + managed_unity_catalog_location, managed_unservable_models, recommended_agent, resolve_state, @@ -207,6 +208,26 @@ def test_none_for_agent_not_in_manifest(self): assert managed_provider_service(MANAGED, "gemini") is None +class TestManagedUnityCatalogLocation: + def test_returns_the_normalized_agent_location(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"unity_catalog_location": " main.default "}} + } + } + assert managed_unity_catalog_location(managed, "claude") == "main.default" + + def test_ignores_other_agents_and_blank_locations(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"unity_catalog_location": " "}}, + "codex": {"model_config": {"unity_catalog_location": "main.codex"}}, + } + } + assert managed_unity_catalog_location(managed, "claude") is None + assert managed_unity_catalog_location(managed, "gemini") is None + + class TestResolveState: def test_does_not_mutate_input_state(self): # managed-state.json and state.json stay separate files: resolution is per-write and