From a5a4075293267442fad50264b7a1ef28437b81be Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 9 Sep 2026 21:39:56 +0000 Subject: [PATCH] [AIGTWY-4573] Consume v2 coding-agent config + version-gate the launch apply Prepare the ug client for the v2 CodingAgentConfig the AI Gateway emits (AIGTWY-4572) without regressing today's path, and stop every launch from re-applying (and re-prompting for) the config. - normalize_managed_config reads the v2 shape: enabled_agents stays a repeated list of {agent enum, config}, and the inner AgentConfig gained v2 fields (an AgentModels source of names / model_service_location / provider, default_model, default_alias_models, http_headers). The v2 fields win over the deprecated custom_headers / model_config oneof the proto keeps. The spend policy is read from spend_tiers with budget_policy as a legacy fallback; models.provider is read with model_provider_service as a legacy fallback. The config's update_time is captured too. - spec_version forward-compat gate (build supports up to 1): a newer or malformed spec_version is refused as an unresolved read, so the launch keeps its last-known-good cache. - UCODE_MANAGED_CONFIG_STUB reads a local JSON config so the client can be exercised against v2 before the server emits it. See examples/managed-config.stub.json. - Gate the launch apply on update_time, not a time TTL. refresh_managed_config always fetches; a launch re-applies the CLI Managed Configuration (the only step that writes the OS-managed files and can prompt for a password) only when update_time is newer than the applied watermark, or on ug --refresh. An unchanged launch skips the apply and never prompts; a changed launch re-applies all enabled agents in one batched prompt and records the new watermark. ug configure always applies and records the watermark. recommendModel is untouched (still per launch). The static models.names allow-list and model_service_location are parsed but not yet applied to each agent's /model picker; that lands in the stacked follow-up. Co-authored-by: Isaac --- examples/managed-config.stub.json | 55 ++ src/ucode/agents/__init__.py | 146 ++++- src/ucode/cli.py | 182 +++++-- src/ucode/databricks.py | 6 +- src/ucode/managed_config.py | 391 +++++++++++--- src/ucode/managed_export.py | 32 +- src/ucode/managed_files.py | 24 + src/ucode/managed_resolve.py | 2 +- src/ucode/managed_setup.py | 120 +++-- src/ucode/state.py | 21 + tests/conftest.py | 10 +- tests/test_agent_pi.py | 2 +- tests/test_agents_init.py | 35 +- tests/test_cli.py | 852 ++++++++++++++++++++++++++++-- tests/test_custom_oauth.py | 1 + tests/test_e2e.py | 6 +- tests/test_managed_config.py | 396 ++++++++++++-- tests/test_managed_export.py | 130 +++-- tests/test_managed_resolve.py | 4 +- tests/test_managed_setup.py | 117 ++-- tests/test_state.py | 24 + 21 files changed, 2174 insertions(+), 382 deletions(-) create mode 100644 examples/managed-config.stub.json diff --git a/examples/managed-config.stub.json b/examples/managed-config.stub.json new file mode 100644 index 00000000..0a1256e8 --- /dev/null +++ b/examples/managed-config.stub.json @@ -0,0 +1,55 @@ +{ + "spec_version": 1, + "update_time": "2026-09-11T00:00:00Z", + "retrieved_time": "2026-09-11T00:00:00Z", + "default_agent": "CODING_AGENT_CLAUDE_CODE", + "enabled_agents": [ + { + "agent": "CODING_AGENT_CLAUDE_CODE", + "config": { + "models": { + "model_services": [ + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5" + ] + }, + "default_models": { + "default_model": "system.ai.claude-opus-4-8", + "default_opus_model": "system.ai.claude-opus-4-8", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + "default_haiku_model": "system.ai.claude-haiku-4-5" + }, + "smart_routing": { "enabled": false }, + "http_headers": {}, + "tracing": { "enabled": true } + } + }, + { + "agent": "CODING_AGENT_CODEX", + "config": { + "models": { "model_provider_service": "main.default.openai-mps" }, + "default_models": { "default_model": "gpt-5.4" }, + "tracing": { "enabled": false } + } + } + ], + "mcp_servers": { + "names": ["system.ai.github"], + "tags": ["MCP_SERVER_TYPE_UC_SERVICE"] + }, + "skills": { + "names": ["system.ai.pdf-extraction"], + "tags": [] + }, + "spend_tiers": { + "budget_id": "00000000-0000-0000-0000-000000000000", + "tiers": [ + { + "spending_percentage": 0.9, + "recommended_agent": "CODING_AGENT_CODEX", + "recommended_model": "gpt-5.4" + } + ] + } +} diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b12d8d86..f5b079a1 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -26,7 +26,18 @@ resolve_provider_service, ) from ucode.managed_files import managed_write_batch -from ucode.state import get_provider_service, load_state, save_state +from ucode.managed_resolve import ( + managed_provider_service, + managed_supplies_models, + resolve_state, +) +from ucode.state import ( + _without_managed_overlay, + get_provider_service, + load_state, + save_state, + set_provider_service, +) from ucode.telemetry import agent_version from ucode.ui import ( console, @@ -472,24 +483,40 @@ def launch( _MODULES[tool].launch(state, tool_args, options=options) -def check_gateway_endpoint(state: dict, tool: str) -> bool: - """V2-only: a tool is available iff we discovered models for it.""" +def check_gateway_endpoint(state: dict, tool: str, managed: dict | None = None) -> bool: + """V2-only: a tool is available iff we discovered models for it or the managed config supplies them. + + FIX 2: An agent whose models come only from the managed config (no discovered models) + must count as available. Check both discovered models and managed-supplied models. + """ + # Check discovered models if tool == "claude": - return bool(state.get("claude_models")) - if tool == "opencode": - return bool(state.get("opencode_models")) - if tool == "codex": - return bool(state.get("codex_models")) - if tool == "gemini": - return bool(state.get("gemini_models")) - if tool == "copilot": - return bool(state.get("claude_models")) or bool(state.get("codex_models")) - if tool == "pi": - return ( + discovered = bool(state.get("claude_models")) + elif tool == "opencode": + discovered = bool(state.get("opencode_models")) + elif tool == "codex": + discovered = bool(state.get("codex_models")) + elif tool == "gemini": + discovered = bool(state.get("gemini_models")) + elif tool == "copilot": + discovered = bool(state.get("claude_models")) or bool(state.get("codex_models")) + elif tool == "pi": + discovered = ( bool(state.get("claude_models")) or bool(state.get("codex_models")) or bool(state.get("gemini_models")) ) + else: + return False + + # If discovered models exist, the tool is available + if discovered: + return True + + # If managed config supplies models, the tool is available + if managed and managed_supplies_models(managed, tool): + return True + return False @@ -514,14 +541,41 @@ def _availability_failure_detail(tool: str, state: dict) -> str: return " (" + "; ".join(parts) + ")" -def configure_single_tool(tool: str, state: dict) -> dict: - """Check availability, configure, and persist state for one tool only.""" +def resolve_managed_for_tool(managed: dict | None, state: dict, tool: str) -> dict: + """State with the managed config applied for ``tool`` and a persisted provider cleared when a + managed static-list or discovery-location source displaces it, so availability and validation + see the effective config rather than the developer's own settings.""" + if managed is None: + return state + resolved = resolve_state(managed, state, tool) + if managed_supplies_models(managed, tool) and not managed_provider_service(managed, tool): + resolved = set_provider_service(resolved, tool, None) + return resolved + + +def configure_single_tool(tool: str, state: dict, managed: dict | None = None) -> dict: + """Check availability, configure, and persist state for one tool only. + + If managed config is provided, it is applied to the state (its settings take + precedence) and the provider precedence rules are enforced (FIX 2). + """ + # Apply managed config before resolving provider, so admin settings win + if managed is not None: + state = resolve_state(managed, state, tool) + provider = get_provider_service(state, tool) + + # When the managed config names its own model source without a provider, clear any persisted + # provider so the managed source drives the picker/catalog. + if managed is not None and managed_supplies_models(managed, tool): + if not managed_provider_service(managed, tool): + provider = None + # 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: with spinner(f"Checking {TOOL_SPECS[tool]['display']} availability..."): - ok = check_gateway_endpoint(state, tool) + ok = check_gateway_endpoint(state, tool, managed=managed) if not ok: detail = _availability_failure_detail(tool, state) raise RuntimeError( @@ -559,19 +613,48 @@ 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, managed: dict | None = None ) -> dict: """Configure the given tools. Caller is responsible for ensuring each tool is available on the workspace. Merges newly-configured tools into state['available_tools'] rather than replacing it, so a previously-configured tool the user didn't pick this - run is preserved. + run is preserved. If managed config is provided, it is applied to each tool + (its settings take precedence) and the provider precedence rules are enforced. """ + # FIX A: Resolve each tool from a clean (overlay-free) state that accumulates + # configuration writes but never accumulates overlays. This ensures each tool's + # managed overlay is independent and doesn't leak into persisted state. + developer_state = state with managed_write_batch(_managed_settings_displays(tools)): for tool in tools: - state = _configure_one(tool, state, get_provider_service(state, tool)) - + # Apply managed config before resolving provider, so admin settings win. + # Resolve from the accumulated state (without overlay) to preserve each + # tool's configuration writes, but use fresh overlays per tool. + tool_state = developer_state + if managed is not None: + tool_state = resolve_state(managed, developer_state, tool) + + provider = get_provider_service(tool_state, tool) + + # When the managed config names its own model source without a provider, clear any + # persisted provider so the managed source drives the picker/catalog. + if managed is not None and managed_supplies_models(managed, tool): + if not managed_provider_service(managed, tool): + provider = None + + # Configure this tool + tool_configured = _configure_one(tool, tool_state, provider) + # Strip the overlay from this tool's result: the managed values reached the config + # file through resolution above, and now the developer's original state is restored + # for persistence. Non-overlay modifications (auth, etc.) are kept. + if managed is not None: + developer_state = _without_managed_overlay(tool_configured) + else: + developer_state = tool_configured + + state = developer_state existing = state.get("available_tools") or [] state["available_tools"] = sorted(set(existing) | set(tools)) save_state(state) @@ -623,20 +706,27 @@ def ensure_provider_state(tool: str) -> dict: return state -def validate_tool(tool: str) -> tuple[bool, str]: - """Invoke a tool with a simple prompt to verify it works. Returns (ok, error_msg).""" +def validate_tool(tool: str, state: dict | None = None) -> tuple[bool, str]: + """Invoke a tool with a simple prompt to verify it works. Returns (ok, error_msg). + + FIX 1: When ``state`` is provided (post-configure validation), use it as-is since it + holds the managed-resolved values that were written to the agent config. When ``state`` + is None (other paths), load the persisted state. + """ + if state is None: + state = load_state() spec = TOOL_SPECS[tool] binary = spec["binary"] module = _MODULES[tool] # Some configs (e.g. claude relayed) can't be probed with a live message — # the proxy + subscription login only exist at launch. Trust the written config. - if hasattr(module, "skip_validation") and module.skip_validation(load_state()): + if hasattr(module, "skip_validation") and module.skip_validation(state): return True, "" cmd = module.validate_cmd(binary) env = None if hasattr(module, "validate_env"): try: - env = module.validate_env(load_state()) + env = module.validate_env(state) except RuntimeError: env = None try: @@ -685,7 +775,7 @@ def provider_permission_error(tool: str, state: dict, err: str) -> str: return err -def validate_all_tools(state: dict) -> None: +def validate_all_tools(state: dict, managed_config: dict | None = None) -> None: from rich.panel import Panel # local to avoid bumping module-level deps from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH @@ -710,7 +800,9 @@ def validate_all_tools(state: dict) -> None: if tool not in available_tools: continue with spinner(f"Validating {spec['display']}..."): - ok, err = validate_tool(tool) + ok, err = validate_tool( + tool, state=resolve_managed_for_tool(managed_config, state, tool) + ) results.append((tool, ok)) if ok: print_success(f"{spec['display']} is working") diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 5466f2da..7732916e 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -32,6 +32,7 @@ provider_permission_error, resolve_gemini_provider_model, resolve_launch_model, + resolve_managed_for_tool, resolve_provider_models, validate_all_tools, validate_tool, @@ -77,8 +78,11 @@ ManagedConfigResult, get_model_recommendation, load_managed_state, + managed_config_is_newer, + managed_update_time, refresh_managed_config, ) +from ucode.managed_files import suppressed_managed_writes from ucode.managed_resolve import ( managed_claude_family_models, managed_default_model, @@ -115,12 +119,15 @@ from ucode.smart_routing import v2 as smart_routing_v2 from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, ROUTE_FIRST_PROMPT_EVENT from ucode.state import ( + MANAGED_OVERLAY_KEY, STATE_PATH, clear_state, + get_applied_managed_update_time, get_provider_service, load_full_state, load_state, save_state, + set_applied_managed_update_time, set_current_workspace, set_provider_service, ) @@ -260,7 +267,7 @@ def _print_managed_summary_abridged(managed: dict, state: dict, tool: str | None def _confirm_managed_config_applied(managed: dict, workspace: str) -> None: - print_success("A managed config is published for your workspace — you're all set.") + print_success("A CLI Managed Configuration is published for your workspace; you're all set.") _print_managed_summary(managed, {"workspace": workspace}, tool=None) print_note("Run `ug` to launch with your managed settings.") @@ -552,11 +559,11 @@ def configure_shared_state( profile = find_profile_name_for_host(workspace) if profile: state["profile"] = profile - with spinner("Verifying Unity AI Gateway..."): + with spinner("Verifying Unity Gateway..."): token = get_databricks_token(workspace, profile) model_service_probe = probe_unity_gateway_capabilities(workspace, token) if model_service_probe.resource_available: - print_success("Unity AI Gateway connected") + print_success("Unity Gateway connected") else: print_warning(f"Model service: {model_service_probe.detail}") @@ -794,7 +801,13 @@ def configure_workspace_command( clear_custom_oauth=custom_oauth is None, ) state = states[0] - state = configure_single_tool(tool, state) + # Fetch and apply managed config so configure respects the admin's policy + managed, _ = _fetch_managed_config(state) + state = configure_single_tool(tool, state, managed=managed) + # Record the applied watermark so a later unchanged launch doesn't re-apply and re-prompt. + if managed is not None: + state = set_applied_managed_update_time(state, managed_update_time(managed)) + save_state(state) install_databricks_ai_tools_for_agents([tool], state) spec = TOOL_SPECS[tool] console.print( @@ -810,14 +823,17 @@ def configure_workspace_command( if skip_validate: print_note(f"Skipping {spec['display']} validation (--skip-validate).") return 0 + # Validate against the effective managed-resolved config (with a displaced provider + # cleared), not the overlay-stripped persisted state. + validate_state = resolve_managed_for_tool(managed, state, tool) with spinner(f"Validating {spec['display']}..."): - ok, err = validate_tool(tool) + ok, err = validate_tool(tool, state=validate_state) if ok: print_success(f"{spec['display']} is working") else: print_err(f"{spec['display']}: {provider_permission_error(tool, state, err)}") - managed = bool(state.get("managed_configs", {}).get(tool)) - restore_file(spec["config_path"], spec["backup_path"], managed) + managed_flag = bool(state.get("managed_configs", {}).get(tool)) + restore_file(spec["config_path"], spec["backup_path"], managed_flag) available_tools = [t for t in (state.get("available_tools") or []) if t != tool] state["available_tools"] = available_tools save_state(state) @@ -835,13 +851,27 @@ def configure_workspace_command( clear_custom_oauth=custom_oauth is None, ) state = states[0] + # Fetch managed config early so it can be passed to configure functions + managed, _ = _fetch_managed_config(state) save_state(state) + # A managed config's enabled_agents is an allowlist, enforced at launch by + # _reject_disabled_agent. With no explicit --agents, honor it here too: configure exactly the + # enabled agents rather than prompting across every workspace-available one. Passing managed to + # the availability check also lets an agent with managed-only models (none discovered) count. + managed_enabled = managed_enabled_tools(managed or {}) + auto_managed = selected_tools is None and bool(managed_enabled) + available_on_workspace: list[str] = [] - tools_to_check = selected_tools or list(TOOL_SPECS) + if selected_tools is not None: + tools_to_check = selected_tools + elif auto_managed: + tools_to_check = managed_enabled + else: + tools_to_check = list(TOOL_SPECS) for tool_name in tools_to_check: with spinner(f"Checking {TOOL_SPECS[tool_name]['display']} availability..."): - if check_gateway_endpoint(state, tool_name): + if check_gateway_endpoint(state, tool_name, managed=managed): available_on_workspace.append(tool_name) if not available_on_workspace: @@ -849,7 +879,19 @@ def configure_workspace_command( _print_discovery_diagnostics(state) return 1 - if selected_tools is None: + if auto_managed: + unavailable_enabled = [t for t in managed_enabled if t not in available_on_workspace] + if unavailable_enabled: + _print_discovery_diagnostics(state) + displays = ", ".join(TOOL_SPECS[t]["display"] for t in unavailable_enabled) + print_warning(f"Managed config enables agent(s) not available here: {displays}.") + picked = available_on_workspace + print_note( + "Configuring the agents your CLI Managed Configuration enables: " + + ", ".join(TOOL_SPECS[t]["display"] for t in picked) + + "." + ) + elif selected_tools is None: picked = prompt_for_tools([(t, TOOL_SPECS[t]["display"]) for t in available_on_workspace]) else: unavailable_tools = [ @@ -880,23 +922,42 @@ def configure_workspace_command( prompt_optional_updates=prompt_optional_updates, ) - # Offer the provider picker for the chosen claude/codex tools only on the - # interactive path (no --agents); otherwise stay on the Databricks path. - if offer_provider: + # Offer the provider picker for the chosen claude/codex tools only on the interactive path (no + # --agents) and only when the workspace is not fully managed; otherwise stay on the Databricks + # path. A fully-managed configure is non-interactive by design, and any tool whose model source + # the managed config already dictates would have a picked provider overridden at launch anyway. + if offer_provider and not auto_managed: for tool_name in picked: + if managed is not None and managed_supplies_models(managed, tool_name): + continue state = _maybe_select_provider_service(tool_name, state) if offer_optional_setup: - state = configure_selected_tools(state, picked, install_ai_tools=False) + state = configure_selected_tools(state, picked, install_ai_tools=False, managed=managed) else: - state = configure_selected_tools(state, picked) + state = configure_selected_tools(state, picked, managed=managed) + # `ug configure` is the explicit sync: record the applied CLI Managed Configuration watermark so + # later launches skip re-applying (and re-prompting for the OS write) until the admin next edits. + if managed is not None: + state = set_applied_managed_update_time(state, managed_update_time(managed)) + save_state(state) summary_lines = [f"[bold]Workspace:[/bold] [cyan]{state['workspace']}[/cyan]"] + if managed is not None: + enabled_names = ( + ", ".join(TOOL_SPECS[t]["display"] for t in managed_enabled_tools(managed)) or "none" + ) + update_time = managed_update_time(managed) + stamp = f", updated {update_time}" if update_time else "" + summary_lines.append( + f"[bold]CLI Managed Configuration:[/bold] applied{stamp}; enables {enabled_names}" + ) for tool_name in picked: spec = TOOL_SPECS[tool_name] summary_lines.append( f"[bold]{spec['display']}:[/bold] [green]configured[/green] " - f"[dim](Provider: {_provider_summary(tool_name, state)})[/dim]" + f"[dim](Provider: {_provider_summary(tool_name, state)})[/dim]\n" + f" [dim]wrote {spec['config_path']}[/dim]" ) console.print( Panel( @@ -913,7 +974,7 @@ def configure_workspace_command( # Limit validation to just-configured tools so we don't re-validate # previously-configured tools the user didn't touch this run. validate_state = {**state, "available_tools": picked} - validate_all_tools(validate_state) + validate_all_tools(validate_state, managed) if offer_optional_setup and not is_dry_run(): _configure_optional_setup(state, picked) return 0 @@ -1838,7 +1899,7 @@ def _reject_disabled_agent(managed: dict | None, tool: str) -> None: if enabled and tool not in enabled: names = ", ".join(TOOL_SPECS[name]["display"] for name in enabled) raise RuntimeError( - f"Your workspace's managed config doesn't enable {TOOL_SPECS[tool]['display']}. " + f"Your CLI Managed Configuration doesn't enable {TOOL_SPECS[tool]['display']}. " f"Enabled: {names}." ) @@ -1847,7 +1908,9 @@ def _fetch_managed_config(state: dict) -> ManagedConfigResult: """The workspace's managed config for this launch, plus whether the feature is disabled. ``ManagedConfigResult(None, True)`` when the workspace has the feature disabled server-side; - ``ManagedConfigResult(None, False)`` when the feature is on but no config is published. + ``ManagedConfigResult(None, False)`` when the feature is on but no config is published. Always + hits the control plane; the caller decides whether the fetched config is newer than what was + last applied. """ with spinner("Loading..."): return refresh_managed_config(state) @@ -2110,6 +2173,27 @@ def _launch_tool( skip_preflight=skip_preflight, **configure_kwargs, ) + # Version-gate the apply: re-apply the CLI Managed Configuration (and re-write the OS-managed + # files, the only step that can prompt for a password) only when it actually changed since it + # was last applied here. An unchanged launch skips this and never prompts; `ug configure` and + # `--refresh` always re-apply. The launched tool's own config write below runs under + # suppression, so the OS write is owned by this apply-all path and by `ug configure`. + applied_ut = get_applied_managed_update_time(existing) + if managed is not None and (refresh or managed_config_is_newer(managed, applied_ut)): + # Apply to every enabled agent (not just the launched one) so the OS-managed files carry + # the admin's policy. This must run even on a needs_auto_configure launch: the earlier + # _auto_configure_tool wrote the tool's files without a managed config (it hadn't been + # fetched yet), so skipping here would leave the OS-managed file unmanaged while the + # watermark below recorded an apply that never happened. + if applied_ut is not None: + print_note("The CLI Managed Configuration was updated; re-applying it.") + enabled = managed_enabled_tools(managed) + if enabled: + state = configure_selected_tools( + state, enabled, install_ai_tools=False, managed=managed + ) + state = set_applied_managed_update_time(state, managed_update_time(managed)) + save_state(state) # An admin-published managed config wins over the developer's own settings. Layered on after # `configure_shared_state`, whose returned state it overrides, and before the provider and # model are settled below — the two state files are never merged on disk. @@ -2119,15 +2203,15 @@ def _launch_tool( _note_recommended_agent(recommendation, tool) if managed is not None: state = resolve_state(managed, state, tool) - print_success("Applied your workspace's managed coding agent config") + print_success("Applied the CLI Managed Configuration") unservable = managed_unservable_models(managed, tool) if unservable: print_warning( - f"Your workspace's managed config lists no {TOOL_SPECS[tool]['display']}-servable " + f"Your CLI Managed Configuration lists no {TOOL_SPECS[tool]['display']}-servable " f"models ({', '.join(unservable)}); using your discovered models instead." ) elif not coding_agent_config_feature_disabled: - print_note("No managed coding agent config found; using your own settings") + print_note("No CLI Managed Configuration found; using your own settings") if managed is not None: managed_provider = managed_provider_service(managed, tool) if explicit_provider and managed_provider and managed_provider != explicit_provider: @@ -2141,6 +2225,26 @@ def _launch_tool( ) if managed_provider: provider = managed_provider + elif managed_supplies_models(managed, tool): + # The managed config names its own model source, so an explicit --provider conflicts + # with it and a persisted provider is cleared to let the managed source drive the + # picker/catalog. + if explicit_provider: + raise RuntimeError( + f"You cannot launch {TOOL_SPECS[tool]['display']} with provider " + f"{explicit_provider} because your admin's managed config specifies its " + f"own model source." + ) + provider = None + # Clear the provider for this launch so every agent honors the managed source, + # including ones (e.g. Gemini) that re-read get_provider_service. Record the + # developer's own provider in the managed overlay so save_state restores it: the + # clear is launch-scoped and the developer's saved provider survives as a fallback + # if the managed policy later disappears. + overlay = dict(state.get(MANAGED_OVERLAY_KEY) or {}) + overlay.setdefault("provider_services", state.get("provider_services")) + state = set_provider_service(state, tool, None) + state[MANAGED_OVERLAY_KEY] = overlay # Checked after the managed config settles `provider`: an admin-set provider must trip this # guard too, or routing would be persisted as on while a provider is active. if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and provider: @@ -2233,18 +2337,22 @@ def _launch_tool( # Codex keeps an explicit --model in ctx.args and passes it to its CLI verbatim. if model and tool != "claude": resolved_model = model - state = configure_tool( - tool, - state, - resolved_model, - provider=provider, - provider_models=provider_models, - relayed=relayed, - route_root_model=route_root_model, - # Claude's explicit model is launch-scoped and is passed through LaunchOptions below. - custom_model=None, - coding_agent_config_defaults=coding_agent_config_defaults, - ) + # The OS-managed (sudo) write is owned by the version-gated apply-all above and by `ug + # configure`; a launch only refreshes the user-level config and computes launch params, so + # suppress the OS write here to keep an unchanged launch prompt-free. + with suppressed_managed_writes(): + state = configure_tool( + tool, + state, + resolved_model, + provider=provider, + provider_models=provider_models, + relayed=relayed, + route_root_model=route_root_model, + # Claude's explicit model is launch-scoped and is passed through LaunchOptions below. + custom_model=None, + coding_agent_config_defaults=coding_agent_config_defaults, + ) # Relayed = a Claude subscription: forward the model to Claude Code's own flag, like `-- --model X`. should_forward_relayed_model = ( tool == "claude" @@ -2466,7 +2574,7 @@ def _launch_managed_default( ) if not isinstance(tool, str) or not tool: raise RuntimeError( - "Your workspace's managed config names no agent to launch. Ask an admin to set a " + "Your CLI Managed Configuration names no agent to launch. Ask an admin to set a " "default agent, or run `ug ` directly." ) _print_managed_summary(managed, state, tool, abridged=True) @@ -2481,9 +2589,9 @@ def _launch_managed_default( def _print_no_managed_config_guidance() -> None: - """Point the developer at per-user configure when no managed config is published.""" + """Point the developer at per-user configure when no CLI Managed Configuration is published.""" print_note( - "No managed coding agent config is published for this workspace. Run `ug configure` to " + "No CLI Managed Configuration is published for this workspace. Run `ug configure` to " "set up your coding agents, then launch one with `ug ` (for example `ug claude`)." ) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a4ffdd9b..cd294f4a 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1909,7 +1909,7 @@ def fetch_external_model_prices(workspace: str, token: str) -> tuple[list[dict], "mcp_servers", "skills", "tracing", - "budget_policy", + "spend_tiers", ) @@ -2892,7 +2892,7 @@ def _raise_ai_gateway_scope_failure(workspace: str, reason: str) -> NoReturn: def _raise_model_service_permission_failure(workspace: str, model_service_reason: str) -> NoReturn: raise RuntimeError( - "Databricks Unity AI Gateway model service access could not be verified on " + "Databricks Unity Gateway model service access could not be verified on " f"{workspace} ({model_service_reason}). Listing Unity Catalog model services requires " "USE CATALOG on `system`, and USE SCHEMA and EXECUTE on `system.ai`." ) @@ -2915,7 +2915,7 @@ def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe _raise_model_service_permission_failure(workspace, reason) raise RuntimeError( - "Databricks Unity AI Gateway is not enabled on this workspace: model services " + "Databricks Unity Gateway is not enabled on this workspace: model services " f"({reason}) are not available. See {AI_GATEWAY_DOCS_URL}" ) diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index cd2c5912..3b7b1d80 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -2,17 +2,19 @@ An org admin authors a ``CodingAgentConfig`` through the Databricks AI Gateway; developers read it (non-admin) and ``ucode`` applies it locally. This module owns the fetch/normalize side and the one -local file, ``~/.ucode/managed-state.json`` (0600), that both roles share: +local file, ``~/.ucode/managed-configuration.json`` (0600), that both roles share: - fetching the raw manifest (via :func:`ucode.databricks.fetch_managed_coding_agent_configs`), - normalizing the proto-JSON into a stable internal dict keyed by ucode's own tool names, -- persisting it via :func:`save_managed_state` / :func:`load_managed_state`, which the launch path - uses to pull the published copy into the local file, and +- persisting it via :func:`save_managed_state` / :func:`load_managed_state` — the admin-write side + (``managed_setup`` / ``managed_wizard``) authors the manifest here, and the launch path pulls the + published copy back into the same file, and - re-reading it on each launch, falling back to the persisted copy when the read fails. -The workspace is the source of truth: an admin authors the ``CodingAgentConfig`` through the AI -Gateway API or UI, and each launch pulls the published copy into ``managed-state.json``. ``ucode`` -only reads and applies it; it never authors or publishes. +There is deliberately one file, not a separate authored ``managed-settings.json``: the workspace is +the source of truth, so an authored draft and the pulled copy are the same shape and coexist in +``managed-configuration.json``. ``ucode setup`` authors the draft; ``ucode publish`` publishes it; a launch +then pulls the published copy back into the same file. :func:`refresh_managed_config` is the launch path's entry point. It is called before model discovery, because the manifest decides whether that discovery is needed at all; the launch path then hands the @@ -25,6 +27,7 @@ import json import os +from datetime import datetime from pathlib import Path from typing import NamedTuple, cast @@ -36,7 +39,7 @@ ) from ucode.ui import console, print_warning -MANAGED_STATE_PATH = config_io.APP_DIR / "managed-state.json" +MANAGED_CONFIGURATION_PATH = config_io.APP_DIR / "managed-configuration.json" # Shown to a developer when their workspace has no admin-defined managed config yet — the normal # case, not an error. Kept here so the CLI (which surfaces it) uses one consistent message. @@ -55,6 +58,14 @@ "CODING_AGENT_OPENCODE": "opencode", } +_AGENT_ENUM_PREFIX = "CODING_AGENT_" +AGENT_NAME_TO_TOOL: dict[str, str] = { + enum[len(_AGENT_ENUM_PREFIX) :].lower(): tool for enum, tool in AGENT_ENUM_TO_TOOL.items() +} + +MAX_SPEC_VERSION = 1 + + # McpServerType proto enum -> ucode's short type tag. Mirrors the selection prefixes in ``mcp.py``; # the actual name->URL resolution happens there when the manifest is applied (a later change). MCP_TYPE_ENUM_TO_TAG: dict[str, str] = { @@ -147,33 +158,106 @@ def _normalize_model_config(model_config: object) -> dict | None: return result or None +def _normalize_agent_config(config: object) -> dict: + """Normalize an ``AgentConfig`` (the inner per-agent config) into the internal shape. + + The current agent-config fields win over the deprecated predecessors the proto keeps: ``http_headers`` supersedes + ``custom_headers``, and the ``models`` / ``default_model`` / ``default_models`` triple + supersedes the ``model_config`` oneof. A config carrying only the old fields still + normalizes, so the transition until the server stops emitting deprecated fields is covered. + """ + config_in = _as_dict(config) + agent_config: dict = {} + headers = _clean_str_dict(config_in.get("http_headers")) or _clean_str_dict( + config_in.get("custom_headers") + ) + if headers: + agent_config["custom_headers"] = headers + tracing_table = _tracing_table(config_in.get("tracing")) or _tracing_table( + config_in.get("tracing_config") + ) + if tracing_table: + agent_config["tracing_table"] = tracing_table + model_config = _normalize_agent_models(config_in) or _normalize_model_config( + config_in.get("model_config") + ) + if model_config is not None: + agent_config["model_config"] = model_config + return agent_config + + def _normalize_enabled_agent(entry: object) -> tuple[str, dict] | None: - """Normalize one ``EnabledAgent`` into ``(tool, agent_config)``, or None if unusable. + """Normalize one repeated ``EnabledAgent`` (``{agent, config}``) into ``(tool, agent_config)``. - Drops entries whose agent enum is unset/unknown to this ucode build. + This is the wire shape the server uses: ``enabled_agents`` stays a repeated + list keyed by the ``agent`` enum (proto map keys can't be enums), and the inner ``config`` + carries the agent-specific fields. Drops entries whose agent is unset/unknown to this ucode build. """ entry_dict = _as_dict(entry) if not entry_dict: return None - tool = AGENT_ENUM_TO_TOOL.get(_str(entry_dict.get("agent")) or "") + tool = _resolve_agent_tool(entry_dict.get("agent")) if tool is None: return None - config_in = _as_dict(entry_dict.get("config")) - agent_config: dict = {} - headers = config_in.get("custom_headers") - if isinstance(headers, dict): - clean = { - k: v for k, v in _as_dict(headers).items() if isinstance(k, str) and isinstance(v, str) - } - if clean: - agent_config["custom_headers"] = clean - tracing_table = _tracing_table(config_in.get("tracing_config")) - if tracing_table: - agent_config["tracing_table"] = tracing_table - model_config = _normalize_model_config(config_in.get("model_config")) - if model_config is not None: - agent_config["model_config"] = model_config - return tool, agent_config + return tool, _normalize_agent_config(entry_dict.get("config")) + + +def _resolve_agent_tool(key: object) -> str | None: + """Map an agent reference to a ucode tool name, accepting either spelling. + + The server may send agent references as either proto enum (``CODING_AGENT_CLAUDE_CODE``) or + by name (``claude_code``). Both resolve to the same tool, or None when this build doesn't + know the agent. + """ + name = _str(key) + if name is None: + return None + return AGENT_ENUM_TO_TOOL.get(name) or AGENT_NAME_TO_TOOL.get(name) + + +def _normalize_agent_models(agent: dict[str, object]) -> dict | None: + """Normalize an ``AgentConfig``'s model fields into the internal ``model_config`` shape. + + Reads the ``default_models`` map (with keys ``default_model``, ``default_opus_model``, + etc.), and the ``models`` object with its alternatives (``model_provider_service``, + ``unity_catalog_location``, or ``model_services``). The server does not enforce exactly-one under + ``models``, so all present forms are carried and each consumer picks its precedence. Returns None + when the config carries none of them, so the caller can fall back to the deprecated ``model_config`` + oneof. + """ + result: dict = {} + default_models = _as_dict(agent.get("default_models")) + overall_default = _str(default_models.get("default_model")) + if overall_default: + result["default_model"] = overall_default + models_dict = _as_dict(agent.get("models")) + provider = _str(models_dict.get("model_provider_service")) + if provider: + result["model_provider_service"] = provider + location = _str(models_dict.get("unity_catalog_location")) + if location: + result["model_service_location"] = location + model_services = _str_list(models_dict.get("model_services")) + if model_services: + result["names"] = model_services + slots = { + slot: model + for slot in ( + "default_opus_model", + "default_sonnet_model", + "default_haiku_model", + "default_fable_model", + ) + if (model := _str(default_models.get(slot))) + } + if slots: + result["models"] = slots + return result or None + + +def _clean_str_dict(value: object) -> dict[str, str]: + """Keep only the string->string entries of ``value`` (a headers map), or an empty dict.""" + return {k: v for k, v in _as_dict(value).items() if isinstance(k, str) and isinstance(v, str)} def _tracing_table(tracing: object) -> str | None: @@ -182,16 +266,33 @@ def _tracing_table(tracing: object) -> str | None: def _normalize_mcp_servers(value: object) -> list[dict]: - if not isinstance(value, list): - return [] - out: list[dict] = [] - for entry in value: - entry_dict = _as_dict(entry) - name = _str(entry_dict.get("name")) - tag = MCP_TYPE_ENUM_TO_TAG.get(_str(entry_dict.get("type")) or "") - if name and tag: - out.append({"name": name, "type": tag}) - return out + """Normalize ``mcp_servers`` wire format into internal ``list[dict]`` with ``name`` and ``type``. + + Accepts both the new wire format (``{names: [...], tags: [...]}`` parallel arrays) and the old + repeated-list format (for backward compat). + """ + value_dict = _as_dict(value) + if value_dict: + names = value_dict.get("names") + tags = value_dict.get("tags") + if isinstance(names, list) and isinstance(tags, list) and len(names) == len(tags): + out: list[dict] = [] + for name, tag in zip(names, tags, strict=True): + name_str = _str(name) + tag_str = MCP_TYPE_ENUM_TO_TAG.get(_str(tag) or "") + if name_str and tag_str: + out.append({"name": name_str, "type": tag_str}) + return out + if isinstance(value, list): + out = [] + for entry in value: + entry_dict = _as_dict(entry) + name = _str(entry_dict.get("name")) + tag = MCP_TYPE_ENUM_TO_TAG.get(_str(entry_dict.get("type")) or "") + if name and tag: + out.append({"name": name, "type": tag}) + return out + return [] def _normalize_budget_policy(value: object) -> dict | None: @@ -213,10 +314,12 @@ def _normalize_budget_policy(value: object) -> dict | None: if not isinstance(pct, (int, float)) or isinstance(pct, bool): continue tier_out: dict = {"spending_percentage": float(pct)} - agent = AGENT_ENUM_TO_TOOL.get(_str(tier_dict.get("default_agent")) or "") + agent = _resolve_agent_tool( + tier_dict.get("recommended_agent") or tier_dict.get("default_agent") + ) if agent: tier_out["default_agent"] = agent - model = _str(tier_dict.get("default_model")) + model = _str(tier_dict.get("recommended_model")) or _str(tier_dict.get("default_model")) if model: tier_out["default_model"] = model tiers.append(tier_out) @@ -239,33 +342,86 @@ def normalize_managed_config(raw: dict) -> dict: display_name = _str(raw.get("display_name")) if display_name: result["display_name"] = display_name - default_agent = AGENT_ENUM_TO_TOOL.get(_str(raw.get("default_agent")) or "") + default_agent = _resolve_agent_tool(raw.get("default_agent")) if default_agent: result["default_agent"] = default_agent - enabled_agents: dict[str, dict] = {} - raw_agents = raw.get("enabled_agents") - for entry in raw_agents if isinstance(raw_agents, list) else []: - normalized = _normalize_enabled_agent(entry) - if normalized is not None: - tool, agent_config = normalized - enabled_agents[tool] = agent_config + update_time = _str(raw.get("update_time")) + if update_time: + result["update_time"] = update_time + enabled_agents = _normalize_enabled_agents(raw.get("enabled_agents")) if enabled_agents: result["enabled_agents"] = enabled_agents mcp_servers = _normalize_mcp_servers(raw.get("mcp_servers")) if mcp_servers: result["mcp_servers"] = mcp_servers - skill_names = _str_list(_as_dict(raw.get("skills")).get("names")) + skills_obj = _as_dict(raw.get("skills")) + skill_names = _str_list(skills_obj.get("names")) if skill_names: result["skills"] = {"names": skill_names} tracing_table = _tracing_table(raw.get("tracing")) if tracing_table: result["tracing_table"] = tracing_table - budget_policy = _normalize_budget_policy(raw.get("budget_policy")) - if budget_policy is not None: - result["budget_policy"] = budget_policy + spend_tiers = _normalize_budget_policy(raw.get("spend_tiers") or raw.get("budget_policy")) + if spend_tiers is not None: + result["budget_policy"] = spend_tiers return result +def managed_update_time(managed: dict | None) -> str | None: + """The config's server-side ``update_time`` (RFC-3339), or None when absent. + + This is the version watermark: it advances only when an admin edits the workspace config, so a + launch compares it against the last-applied value to decide whether to re-apply. The GET also + returns ``retrieved_time``, which changes on every read and must never be used for this. + """ + return _str(_as_dict(managed).get("update_time")) + + +def _parse_update_time(value: str | None) -> datetime | None: + if not value: + return None + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +def managed_config_is_newer(fetched: dict | None, applied_update_time: str | None) -> bool: + """True when ``fetched`` is a newer version than the last one applied locally. + + A fetched config whose ``update_time`` is missing or unparseable is treated as newer, so a launch + re-applies it rather than trusting possibly-stale local settings; no previously-applied watermark + also counts as newer (the first apply). + """ + fetched_ut = _parse_update_time(managed_update_time(fetched)) + applied_ut = _parse_update_time(applied_update_time) + if fetched_ut is None or applied_ut is None: + return True + return fetched_ut > applied_ut + + +def _normalize_enabled_agents(raw_agents: object) -> dict[str, dict]: + """Normalize ``enabled_agents`` into ``{tool: agent_config}``. + + The server sends a repeated ``EnabledAgent`` list, each entry carrying its own + ``agent`` enum plus a per-agent ``config``. A map keyed by agent name is also accepted defensively. + Either way the result keys by ucode tool name, dropping agents this build doesn't recognize. + """ + enabled_agents: dict[str, dict] = {} + if isinstance(raw_agents, list): + for entry in raw_agents: + normalized = _normalize_enabled_agent(entry) + if normalized is not None: + tool, agent_config = normalized + enabled_agents[tool] = agent_config + elif isinstance(raw_agents, dict): + for key, agent in raw_agents.items(): + tool = _resolve_agent_tool(key) + if tool is not None: + enabled_agents[tool] = _normalize_agent_config(agent) + return enabled_agents + + def _decimal(value: object) -> float | None: """Parse one of the API's decimal-string money fields, or None when absent/unparseable.""" text = _str(value) @@ -318,7 +474,15 @@ def get_managed_config(workspace: str, token: str) -> FetchedManagedConfig: about rather than silently launch without. v0 stores at most one config per workspace, so the first entry is the workspace's config. + + ``UCODE_MANAGED_CONFIG_STUB`` short-circuits the HTTP read: when it names a readable JSON file, + that file's single CodingAgentConfig is used verbatim. It exists so this client can be exercised + against the managed-config shape before the server emits it (AIGTWY-4572); unset in normal use. See + ``examples/managed-config.stub.json`` for a sample. """ + stub = _stub_config() + if stub is not None: + return _gate_config(stub) configs, reason = fetch_managed_coding_agent_configs(workspace, token) if reason is not None: if _is_feature_disabled(reason): @@ -329,7 +493,46 @@ def get_managed_config(workspace: str, token: str) -> FetchedManagedConfig: return FetchedManagedConfig(None, reason) if not configs: return FetchedManagedConfig(None, None) - return FetchedManagedConfig(normalize_managed_config(configs[0]), None) + return _gate_config(configs[0]) + + +def _stub_config() -> dict | None: + """The stub CodingAgentConfig named by ``UCODE_MANAGED_CONFIG_STUB``, or None when unset/bad.""" + path = os.environ.get("UCODE_MANAGED_CONFIG_STUB") + if not path: + return None + try: + raw = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + print_warning(f"UCODE_MANAGED_CONFIG_STUB could not be read ({exc}); ignoring it.") + return None + return raw if isinstance(raw, dict) else None + + +def _gate_config(raw: dict) -> FetchedManagedConfig: + """Apply the ``spec_version`` forward-compat gate and return the raw config unchanged. + + A config declaring a ``spec_version`` newer than this build understands is refused as an + unresolved read (``reason`` set), so the launch path falls back to the last-known-good cache and + never blocks — the same treatment as any read this build can't act on. On a clean read the raw + config is returned verbatim; normalization happens later, at the read/return boundaries, so the + persisted file stays byte-identical to what the gateway returned. + """ + spec = raw.get("spec_version") + if spec is not None: + if isinstance(spec, bool) or not isinstance(spec, int): + return FetchedManagedConfig( + None, + f"Your CLI Managed Configuration has an unrecognized spec_version ({spec!r}); " + "update Unity Gateway with `ug upgrade`.", + ) + if spec > MAX_SPEC_VERSION: + return FetchedManagedConfig( + None, + f"Your CLI Managed Configuration needs a newer Unity Gateway (spec_version {spec}; " + f"this build supports up to {MAX_SPEC_VERSION}). Run `ug upgrade`.", + ) + return FetchedManagedConfig(raw, None) def _is_not_found(reason: str) -> bool: @@ -352,29 +555,44 @@ def _is_permission_denied(reason: str) -> bool: return "http 403" in lowered or "permission_denied" in lowered +def _is_unsupported_spec(reason: str) -> bool: + """True when the read failed because the config's ``spec_version`` is newer than this build. + + Unlike a transient read failure, this is proof a policy exists, so it is surfaced even with no + cached config to fall back on. + """ + return "spec_version" in reason.lower() + + def save_managed_state(workspace: str, config: dict) -> None: - """Persist the normalized managed config to ``~/.ucode/managed-state.json`` at mode 0600. + """Persist the raw managed config to ``~/.ucode/managed-configuration.json`` at mode 0600. - The file is org-authored, not developer-editable — 0600 keeps it readable/writable only by the - user (a light guard; hard enforcement / sudo ownership is a separate concern). No-op in dry-run. + ``config`` is stored verbatim as the gateway returned it (byte-identical to the GET), so the file + is inspectable and ``ug export`` can dump it unchanged; normalization into ucode's internal shape + happens on read (:func:`load_managed_state`), not here. The file is org-authored, not + developer-editable — 0600 keeps it readable/writable only by the user. No-op in dry-run. An empty ``config`` records "this workspace has no managed config", which matters because the file doubles as the fallback when a later read fails: without it, removing a config server-side would leave the old one on disk to be reapplied after a transient outage. """ - payload = {"workspace": workspace, "config": config} + payload: dict = {"workspace": workspace, "config": config} if config_io.is_dry_run(): # Print rather than write, matching how the agent config writers behave under --dry-run. console.print( - f"\n[bold]\\[dry run] {MANAGED_STATE_PATH}[/bold]\n{json.dumps(payload, indent=2)}\n" + f"\n[bold]\\[dry run] {MANAGED_CONFIGURATION_PATH}[/bold]\n{json.dumps(payload, indent=2)}\n" ) return - config_io.ensure_parent_dir(MANAGED_STATE_PATH) + config_io.ensure_parent_dir(MANAGED_CONFIGURATION_PATH) try: - MANAGED_STATE_PATH.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + MANAGED_CONFIGURATION_PATH.write_text( + json.dumps(payload, indent=2) + "\n", encoding="utf-8" + ) except OSError as exc: - raise RuntimeError(f"Failed to write managed state file: {MANAGED_STATE_PATH}") from exc - _restrict_permissions(MANAGED_STATE_PATH) + raise RuntimeError( + f"Failed to write managed state file: {MANAGED_CONFIGURATION_PATH}" + ) from exc + _restrict_permissions(MANAGED_CONFIGURATION_PATH) def _restrict_permissions(path: Path) -> None: @@ -387,19 +605,30 @@ def _restrict_permissions(path: Path) -> None: def load_managed_state(workspace: str | None) -> dict | None: - """Load the persisted managed config for ``workspace``, or None if absent/mismatched. + """Load the persisted managed config for ``workspace`` normalized into ucode's internal shape. + + The file stores the raw gateway config; this reads it and returns + :func:`normalize_managed_config` of it, so every consumer keeps working against the normalized + shape. Returns None when there is no file for this workspace (a stale file from another workspace + is ignored rather than misapplied). A stored empty config normalizes to an empty dict, which + callers already treat as "no config". + """ + raw = load_managed_configuration(workspace) + if raw is None: + return None + return normalize_managed_config(raw) + - Returns the normalized config dict (the ``config`` field), only when the stored file is for the - same workspace — so a stale file from another workspace is ignored rather than misapplied. +def load_managed_configuration(workspace: str | None) -> dict | None: + """Return the raw managed config persisted for ``workspace`` (verbatim as the gateway returned + it), or None if absent or stored for a different workspace. - This is the single local managed config: ``ucode setup`` authors it here, ``ucode publish`` - publishes it, and a launch refreshes it from the workspace. The admin-authored draft and the - pulled copy share one file because the workspace is the source of truth — to keep a draft, - publish it with ``ucode publish``. + Unlike :func:`load_managed_state` this does not normalize: it is the exact CodingAgentConfig, for + ``ug export`` and for inspecting the on-disk file. """ if not workspace: return None - data = config_io.read_json_safe(MANAGED_STATE_PATH) + data = config_io.read_json_safe(MANAGED_CONFIGURATION_PATH) if data.get("workspace") != workspace: return None config = data.get("config") @@ -412,16 +641,19 @@ def managed_state_workspace() -> str | None: Lets a caller that has no workspace in local ucode state (e.g. ``ucode setup --show`` before ``ucode configure``) still find the manifest on disk and report which workspace it belongs to. """ - workspace = config_io.read_json_safe(MANAGED_STATE_PATH).get("workspace") + workspace = config_io.read_json_safe(MANAGED_CONFIGURATION_PATH).get("workspace") return workspace if isinstance(workspace, str) and workspace else None def refresh_managed_config(state: dict) -> ManagedConfigResult: - """Fetch the workspace's managed config and persist it as a :class:`ManagedConfigResult`. + """Fetch the workspace's managed config fresh and persist it as a :class:`ManagedConfigResult`. Runs on every launch so a developer picks up an admin's edits without re-running - ``ucode configure``. The manifest is None when the workspace has no managed config — the normal - case for a workspace whose admin hasn't published one. + ``ucode configure``. It always hits the control plane; whether the fetched config is *newer* than + what was last applied — and so whether the launch re-applies the settings — is the caller's + decision, via :func:`managed_config_is_newer` against the persisted applied watermark. The + manifest is None when the workspace has no managed config — the normal case for a workspace whose + admin hasn't published one. A failed fetch never blocks the launch: an unreachable control plane shouldn't stop someone from coding. Instead it falls back to the last config persisted for this workspace, so the admin's @@ -443,21 +675,22 @@ def refresh_managed_config(state: dict) -> ManagedConfigResult: token = get_databricks_token(workspace, state.get("profile")) except RuntimeError as exc: return ManagedConfigResult(_persisted_fallback(workspace, str(exc)), False) - managed, reason = get_managed_config(workspace, token) + raw, reason = get_managed_config(workspace, token) if reason is not None: if _is_feature_disabled(reason): save_managed_state(workspace, {}) return ManagedConfigResult(None, True) fallback = _persisted_fallback(workspace, reason, refused=_is_permission_denied(reason)) return ManagedConfigResult(fallback, False) - if managed is None: + if raw is None: # Record that this workspace has no config, rather than leaving an earlier one on disk: # the file doubles as the fallback above, so a removed policy would otherwise come back # into force after the next transient outage. save_managed_state(workspace, {}) return ManagedConfigResult(None, False) - save_managed_state(workspace, managed) - return ManagedConfigResult(managed, False) + # Persist the raw config verbatim; hand callers the normalized manifest they expect. + save_managed_state(workspace, raw) + return ManagedConfigResult(normalize_managed_config(raw), False) def _is_feature_disabled(reason: str) -> bool: @@ -477,16 +710,18 @@ def _persisted_fallback(workspace: str, reason: str, *, refused: bool = False) - # policy to fall back to — treat it the same as having no file at all. persisted = load_managed_state(workspace) if not persisted: + if _is_unsupported_spec(reason): + print_warning(reason) return None summary = _summarize_read_failure(reason) if refused: print_warning( - f"Your workspace's managed config is not readable by you ({summary}); using the last " + f"Your CLI Managed Configuration is not readable by you ({summary}); using the last " "one saved for this workspace. Ask an admin to grant access." ) else: print_warning( - f"Could not read your workspace's managed config ({summary}); " + f"Could not read your CLI Managed Configuration ({summary}); " "using the last one saved for this workspace." ) return persisted diff --git a/src/ucode/managed_export.py b/src/ucode/managed_export.py index 48b60325..e1e68226 100644 --- a/src/ucode/managed_export.py +++ b/src/ucode/managed_export.py @@ -18,11 +18,25 @@ import tempfile from pathlib import Path -from ucode.managed_config import load_managed_state, managed_state_workspace -from ucode.managed_setup import serialize_managed_config, validate_manifest +from ucode.managed_config import ( + load_managed_configuration, + managed_state_workspace, + normalize_managed_config, +) +from ucode.managed_setup import validate_manifest from ucode.state import load_state -_SERVER_OWNED_FIELDS = ("name",) +# Server-assigned metadata that is part of the stored raw config but is not authored config and +# must not appear in a portable export (it would be rejected or ignored on re-import). +_SERVER_OWNED_FIELDS = ( + "name", + "workspace_id", + "create_time", + "update_time", + "retrieved_time", + "created_user_id", + "updated_user_id", +) EXPORT_SPEC_VERSION = 1 @@ -36,19 +50,17 @@ def build_export_payload() -> dict: actionable message when no config is authored locally or the config fails structural validation. """ workspace = load_state().get("workspace") or managed_state_workspace() - manifest = load_managed_state(workspace) - if not manifest: + config = load_managed_configuration(workspace) + if not config: raise RuntimeError( - "No managed coding-agent config found locally. Run `ug` against a workspace that " + "No CLI Managed Configuration found locally. Run `ug` against a workspace that " "publishes one, then re-run `ug export`." ) - errors = validate_manifest(manifest, None) + errors = validate_manifest(normalize_managed_config(config), None) if errors: detail = "\n".join(f" - {error}" for error in errors) raise RuntimeError(f"The managed config is not valid, so it was not exported:\n{detail}") - config = serialize_managed_config(manifest) - for field in _SERVER_OWNED_FIELDS: - config.pop(field, None) + config = {key: value for key, value in config.items() if key not in _SERVER_OWNED_FIELDS} return {"workspace": workspace, "spec_version": EXPORT_SPEC_VERSION, **config} diff --git a/src/ucode/managed_files.py b/src/ucode/managed_files.py index 0bffc81f..b1648f90 100644 --- a/src/ucode/managed_files.py +++ b/src/ucode/managed_files.py @@ -31,6 +31,10 @@ _MISSING = object() _managed_write_batch: tuple[str, ...] = () _managed_write_notice_shown = False +# When set, reconcile_managed_file skips the OS-managed (sudo) write entirely. The launch path uses +# this so a plain launch never re-writes /etc — that write is owned by `ug configure` and by the +# launch-time apply that runs only when the CLI Managed Configuration actually changed. +_managed_writes_suppressed = False ManagedParser = Callable[[str], dict] ManagedDumper = Callable[[dict], str] @@ -154,6 +158,24 @@ def managed_write_batch(displays: list[str]) -> Iterator[None]: _managed_write_notice_shown = previous_notice +@contextmanager +def suppressed_managed_writes() -> Iterator[None]: + """Within this context, :func:`reconcile_managed_file` skips the OS-managed (sudo) write. + + Used by the launch path so a launch that isn't re-applying the CLI Managed Configuration never + touches the root-owned /etc file (and never prompts for a password); the user-level config is + still written by the caller. + """ + global _managed_writes_suppressed + + prev = _managed_writes_suppressed + _managed_writes_suppressed = True + try: + yield + finally: + _managed_writes_suppressed = prev + + def _print_managed_write_permission(display: str) -> None: global _managed_write_notice_shown @@ -239,6 +261,8 @@ def reconcile_managed_file( The first pre-ucode contents are retained until ``ucode revert``. Subsequent writes update only the last-applied snapshot used for drift-safe three-way restoration. """ + if _managed_writes_suppressed: + return "unchanged" if not managed_files_supported(): print_warning( f"{display}: OS-managed settings aren't supported on this platform; skipped {path}." diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 4752ff90..64cbb3d1 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -1,6 +1,6 @@ """Resolve the effective agent settings from the managed config plus local ucode state. -The managed config (``~/.ucode/managed-state.json``, published by an admin through the AI Gateway +The managed config (``~/.ucode/managed-configuration.json``, published by an admin through the AI Gateway and refreshed from the workspace at launch through :mod:`ucode.managed_config`) and the developer's own ucode state (``~/.ucode/state.json``) stay separate files — they are never merged on disk. Instead this module resolves them *per key* at config-write time: whatever the manifest specifies wins, and diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index 3c52f909..4cfdf0b3 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -14,7 +14,7 @@ once. Local persistence is not duplicated here: the authored manifest is saved to and loaded from the one -local file, ``~/.ucode/managed-state.json``, via :func:`ucode.managed_config.save_managed_state` and +local file, ``~/.ucode/managed-configuration.json``, via :func:`ucode.managed_config.save_managed_state` and :func:`ucode.managed_config.load_managed_state` — the same file the launch path pulls into. The interactive wizard that calls these helpers, and the publish step, live in @@ -182,36 +182,50 @@ def claude_model_slots(models: list[str]) -> dict[str, str]: def _model_config_payload(tool: str, model_config: dict) -> dict: - """Build one ``AgentModelConfig`` oneof variant body for ``tool``. + """Build one agent's ``models`` object and ``default_models`` map for the wire. - Shapes per the proto: claude gets `models` as a `ClaudeDefaultModels` slot object, codex gets - no model list at all, and the rest get a flat repeated `models`. + The current wire shape has two top-level per-agent fields: ``models`` (the source, + as model_provider_service / unity_catalog_location / model_services) and + ``default_models`` (a flat map with the overall default_model plus family slots). + For claude, the internal has a `models` dict of family slots; for flat-list agents, + it has a `names` list. """ - body: dict = {} + models_obj: dict = {} mps = model_config.get("model_provider_service") if isinstance(mps, str) and mps: - body["model_provider_service"] = mps + models_obj["model_provider_service"] = mps + location = model_config.get("model_service_location") + if isinstance(location, str) and location: + models_obj["unity_catalog_location"] = location + + names = model_config.get("names") + if isinstance(names, list): + model_list = [m for m in names if isinstance(m, str) and m] + if model_list: + models_obj["model_services"] = model_list + + default_models_map: dict = {} default_model = model_config.get("default_model") if isinstance(default_model, str) and default_model: - body["default_model"] = default_model + default_models_map["default_model"] = default_model models = model_config.get("models") - if tool == "claude": - if isinstance(models, dict): - slots = { - slot: value - for slot, value in models.items() - if isinstance(slot, str) and isinstance(value, str) and value - } - if slots: - body["models"] = slots - elif tool in _FLAT_MODEL_LIST_AGENTS: - if isinstance(models, list): - model_list = [m for m in models if isinstance(m, str) and m] - if model_list: - body["models"] = model_list - # codex intentionally carries no model list — CodexModelConfig has only - # model_provider_service + default_model. + if isinstance(models, dict): + for slot in ( + "default_opus_model", + "default_sonnet_model", + "default_haiku_model", + "default_fable_model", + ): + model = models.get(slot) + if isinstance(model, str) and model: + default_models_map[slot] = model + + body: dict = {} + if models_obj: + body["models"] = models_obj + if default_models_map: + body["default_models"] = default_models_map return body @@ -222,19 +236,15 @@ def _enabled_agent_payload(tool: str, agent_config: dict) -> dict: if isinstance(headers, dict): clean = {k: v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)} if clean: - config["custom_headers"] = clean + config["http_headers"] = clean tracing_table = agent_config.get("tracing_table") if isinstance(tracing_table, str) and tracing_table: - config["tracing_config"] = {"table": tracing_table} + config["tracing"] = {"enabled": True} model_config = agent_config.get("model_config") if isinstance(model_config, dict): - body = _model_config_payload(tool, model_config) - if body: - # The `AgentModelConfig` oneof field names are ucode's tool names verbatim (claude, - # codex, opencode, pi, gemini, copilot), so the tool doubles as the variant key. The - # server rejects a variant that doesn't match its agent (`validateAgentModelConfig`), - # and the round-trip through `normalize_managed_config` pins that alignment in tests. - config["model_config"] = {tool: body} + payload = _model_config_payload(tool, model_config) + if payload: + config.update(payload) entry: dict = {"agent": AGENT_TOOL_TO_ENUM[tool]} if config: @@ -243,7 +253,7 @@ def _enabled_agent_payload(tool: str, agent_config: dict) -> dict: def _budget_policy_payload(budget_policy: dict) -> dict: - """Build the ``BudgetPolicy`` body, dropping tiers that name an unknown agent. + """Build the ``spend_tiers`` body, dropping tiers that name an unknown agent. ``spending_percentage`` is passed through as-is: it is a fraction in [0, 1] both in ucode's manifest and in the proto (the server validates that range). Callers prompting an admin in @@ -268,10 +278,10 @@ def _budget_policy_payload(budget_policy: dict) -> dict: tier_payload: dict = {"spending_percentage": float(pct)} agent_enum = AGENT_TOOL_TO_ENUM.get(str(tier.get("default_agent") or "")) if agent_enum: - tier_payload["default_agent"] = agent_enum + tier_payload["recommended_agent"] = agent_enum default_model = tier.get("default_model") if isinstance(default_model, str) and default_model: - tier_payload["default_model"] = default_model + tier_payload["recommended_model"] = default_model tiers.append(tier_payload) if tiers: payload["tiers"] = tiers @@ -282,13 +292,13 @@ def serialize_managed_config(manifest: dict) -> dict: """Serialize ucode's internal manifest into a proto-JSON ``CodingAgentConfig``. The exact inverse of :func:`ucode.managed_config.normalize_managed_config`: tool names become - ``CODING_AGENT_*`` enums, MCP type tags become ``MCP_SERVER_TYPE_*``, and each agent's model - config is wrapped in its matching ``AgentModelConfig`` oneof variant. Agents and MCP types this - build doesn't recognize are dropped, mirroring the read side. + ``CODING_AGENT_*`` enums, internal ``models`` slots become ``default_models`` map keys, and + MCP type tags become ``MCP_SERVER_TYPE_*`` enums. Agents and MCP types this build doesn't + recognize are dropped, mirroring the read side. - Output-only proto fields (``workspace_id``, timestamps, user ids) are never emitted. ``name`` is - carried through when present so an update path can address an existing resource; ``ucode publish`` - omits it on create and lets the server assign one. + Output-only proto fields (``workspace_id``, ``retrieved_time``, user ids) are never emitted. + ``name`` is carried through when present so an update path can address an existing resource; + ``ucode publish`` omits it on create and lets the server assign one. """ payload: dict = {} @@ -315,16 +325,19 @@ def serialize_managed_config(manifest: dict) -> dict: mcp_servers = manifest.get("mcp_servers") if isinstance(mcp_servers, list): - servers: list[dict] = [] + names: list[str] = [] + tags: list[str] = [] for server in mcp_servers: if not isinstance(server, dict): continue server_name = server.get("name") - type_enum = MCP_TAG_TO_TYPE_ENUM.get(str(server.get("type") or "")) + type_tag = str(server.get("type") or "") + type_enum = MCP_TAG_TO_TYPE_ENUM.get(type_tag) if isinstance(server_name, str) and server_name and type_enum: - servers.append({"name": server_name, "type": type_enum}) - if servers: - payload["mcp_servers"] = servers + names.append(server_name) + tags.append(type_enum) + if names: + payload["mcp_servers"] = {"names": names, "tags": tags} skills = manifest.get("skills") if isinstance(skills, dict): @@ -332,17 +345,17 @@ def serialize_managed_config(manifest: dict) -> dict: if isinstance(names, list): skill_names = [n for n in names if isinstance(n, str) and n] if skill_names: - payload["skills"] = {"names": skill_names} + payload["skills"] = {"names": skill_names, "tags": []} tracing_table = manifest.get("tracing_table") if isinstance(tracing_table, str) and tracing_table: - payload["tracing"] = {"table": tracing_table} + payload["tracing"] = {"enabled": True} budget_policy = manifest.get("budget_policy") if isinstance(budget_policy, dict): policy = _budget_policy_payload(budget_policy) if policy: - payload["budget_policy"] = policy + payload["spend_tiers"] = policy return payload @@ -502,9 +515,9 @@ def validate_manifest(manifest: dict, state: dict | None = None) -> list[str]: def _agent_model_ids(agent_config: dict) -> set[str]: """Every model id an agent is configured with — its list plus its default. - Claude's ``models`` is a family-slot dict and the others' a flat list; codex has no list at all, - only ``default_model``. Returns an empty set when nothing is configured, which callers treat as - "can't check" rather than "nothing is allowed". + Claude's ``models`` is a family-slot dict; flat-list agents use ``names`` (a list); + codex has no list at all, only ``default_model``. Returns an empty set when nothing + is configured, which callers treat as "can't check" rather than "nothing is allowed". """ model_config = agent_config.get("model_config") if not isinstance(model_config, dict): @@ -515,6 +528,9 @@ def _agent_model_ids(agent_config: dict) -> set[str]: ids.update(v for v in raw.values() if isinstance(v, str) and v) elif isinstance(raw, list): ids.update(m for m in raw if isinstance(m, str) and m) + names = model_config.get("names") + if isinstance(names, list): + ids.update(m for m in names if isinstance(m, str) and m) default_model = model_config.get("default_model") if isinstance(default_model, str) and default_model: ids.add(default_model) diff --git a/src/ucode/state.py b/src/ucode/state.py index c0285553..df41cd15 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -285,3 +285,24 @@ def set_provider_service(state: dict, tool: str, full_name: str | None) -> dict: else: state.pop("provider_services", None) return state + + +# The CLI Managed Configuration's ``update_time`` last applied to this workspace's agents. A launch +# compares a freshly fetched config against it to decide whether to re-apply, so it is written only +# after an apply succeeds — never on a plain fetch. +APPLIED_MANAGED_UPDATE_TIME_KEY = "applied_managed_update_time" + + +def get_applied_managed_update_time(state: dict) -> str | None: + """The ``update_time`` of the CLI Managed Configuration last applied to this workspace, if any.""" + value = state.get(APPLIED_MANAGED_UPDATE_TIME_KEY) + return value if isinstance(value, str) and value else None + + +def set_applied_managed_update_time(state: dict, update_time: str | None) -> dict: + """Record (or clear) the applied CLI Managed Configuration watermark for this workspace.""" + if update_time: + state[APPLIED_MANAGED_UPDATE_TIME_KEY] = update_time + else: + state.pop(APPLIED_MANAGED_UPDATE_TIME_KEY, None) + return state diff --git a/tests/conftest.py b/tests/conftest.py index 9d861f23..77022660 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -35,9 +35,11 @@ def _isolate_ucode_state(tmp_path, monkeypatch): state_dir.mkdir() monkeypatch.setattr(state_mod, "STATE_PATH", state_dir / "state.json") monkeypatch.setattr(config_io_mod, "APP_DIR", state_dir) - # MANAGED_STATE_PATH is bound from APP_DIR at import, so patching APP_DIR alone doesn't move it; - # rebind it or save_managed_state writes to the developer's real ~/.ucode/managed-state.json. - monkeypatch.setattr(managed_config_mod, "MANAGED_STATE_PATH", state_dir / "managed-state.json") + # MANAGED_CONFIGURATION_PATH is bound from APP_DIR at import, so patching APP_DIR alone doesn't move it; + # rebind it or save_managed_state writes to the developer's real ~/.ucode/managed-configuration.json. + monkeypatch.setattr( + managed_config_mod, "MANAGED_CONFIGURATION_PATH", state_dir / "managed-configuration.json" + ) backup_dir = state_dir / "managed-backups" monkeypatch.setattr(managed_files_mod, "MANAGED_BACKUP_DIR", backup_dir) monkeypatch.setattr( @@ -54,6 +56,8 @@ def reject_privileged_write(path, _desired_text): monkeypatch.setattr(managed_files_mod, "_sudo_replace", reject_privileged_write) monkeypatch.delenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", raising=False) monkeypatch.delenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", raising=False) + # A developer's ambient managed-config stub would otherwise short-circuit every fetch in the suite. + monkeypatch.delenv("UCODE_MANAGED_CONFIG_STUB", raising=False) # The model-services listing is memoized for the life of the process, so without this a cached # result would leak into the next test and make a stubbed listing look like it was never called. databricks_mod.clear_model_services_cache() diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 416c9769..08e5033c 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -440,7 +440,7 @@ def test_failed_pi_validation_rolls_back_settings(self, tmp_path, monkeypatch): monkeypatch.setitem( agents_mod.TOOL_SPECS["pi"], "backup_path", tmp_path / "models.backup.json" ) - monkeypatch.setattr(agents_mod, "validate_tool", lambda tool: (False, "boom")) + monkeypatch.setattr(agents_mod, "validate_tool", lambda tool, **kwargs: (False, "boom")) monkeypatch.setattr(agents_mod, "save_state", lambda s: None) monkeypatch.setattr(agents_mod, "spinner", lambda *_a, **_kw: nullcontext()) diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index a59e40a1..14404255 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -247,6 +247,39 @@ def test_pi_available_with_gemini(self): def test_pi_unavailable_when_no_models(self): assert check_gateway_endpoint({}, "pi") is False + def test_managed_static_list_makes_undiscovered_tool_available(self): + # A managed config can name a tool's models even when discovery found none for it, so a + # single-agent configure must count that as available rather than erroring out. + managed = {"enabled_agents": {"codex": {"model_config": {"models": ["system.ai.gpt-5"]}}}} + assert check_gateway_endpoint({}, "codex", managed=managed) is True + + def test_managed_without_models_leaves_undiscovered_tool_unavailable(self): + managed = {"enabled_agents": {"codex": {"model_config": {}}}} + assert check_gateway_endpoint({}, "codex", managed=managed) is False + + +class TestResolveManagedForTool: + def test_none_managed_returns_state_unchanged(self): + state = {"provider_services": {"codex": "main.x.svc"}} + assert agents_mod.resolve_managed_for_tool(None, state, "codex") is state + + def test_static_list_clears_persisted_provider(self): + # A managed static list with no provider should displace the developer's own provider so + # availability and validation see the managed models, not a stale routed provider. + managed = {"enabled_agents": {"codex": {"model_config": {"models": ["system.ai.gpt-5"]}}}} + state = {"provider_services": {"codex": "main.x.svc"}} + resolved = agents_mod.resolve_managed_for_tool(managed, state, "codex") + assert resolved.get("codex_models") == ["system.ai.gpt-5"] + assert "codex" not in (resolved.get("provider_services") or {}) + + def test_managed_provider_service_is_kept(self): + managed = { + "enabled_agents": {"codex": {"model_config": {"model_provider_service": "main.m.svc"}}} + } + state = {"provider_services": {"codex": "main.x.svc"}} + resolved = agents_mod.resolve_managed_for_tool(managed, state, "codex") + assert resolved["provider_services"]["codex"] == "main.m.svc" + class TestDefaultModelForTool: def test_codex_returns_none_without_a_configured_model(self): @@ -798,7 +831,7 @@ class TestValidateAllToolsVerbosity: def _run(self, monkeypatch, capsys): from contextlib import nullcontext - monkeypatch.setattr(agents_mod, "validate_tool", lambda tool: (True, "")) + monkeypatch.setattr(agents_mod, "validate_tool", lambda tool, **kwargs: (True, "")) monkeypatch.setattr(agents_mod, "save_state", lambda s: None) monkeypatch.setattr(agents_mod, "spinner", lambda *_a, **_kw: nullcontext()) agents_mod.validate_all_tools({"available_tools": ["codex"], "managed_configs": {}}) diff --git a/tests/test_cli.py b/tests/test_cli.py index cd2d3f63..dd2ee90a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -915,7 +915,7 @@ def _provider_launch(monkeypatch, argv, provider_models, relayed=False): monkeypatch.setattr(cli_mod, "load_state", lambda: MINIMAL_STATE) monkeypatch.setattr(cli_mod, "ensure_provider_state", lambda t: MINIMAL_STATE) monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: MINIMAL_STATE) - monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s, **kwargs: (None, False)) monkeypatch.setattr(cli_mod, "_fetch_budget_recommendation", lambda s, m: None) mock_launch = MagicMock() monkeypatch.setattr(cli_mod, "launch_agent", mock_launch) @@ -1054,7 +1054,7 @@ def _launch(monkeypatch, resolve_provider_models): monkeypatch.setattr("ucode.cli.load_state", lambda: state) monkeypatch.setattr("ucode.cli.ensure_provider_state", lambda t: state) monkeypatch.setattr("ucode.cli.configure_shared_state", lambda *a, **k: state) - monkeypatch.setattr("ucode.cli._fetch_managed_config", lambda s: (None, False)) + monkeypatch.setattr("ucode.cli._fetch_managed_config", lambda s, **kwargs: (None, False)) monkeypatch.setattr("ucode.cli.resolve_provider_models", resolve_provider_models) monkeypatch.setattr("ucode.cli.configure_tool", lambda *a, **k: state) monkeypatch.setattr( @@ -1667,6 +1667,9 @@ def _launch(self, monkeypatch, *, managed): patch("ucode.cli.ensure_provider_state", return_value=state), patch("ucode.cli.configure_shared_state", return_value=state), patch("ucode.cli.configure_tool", return_value=state), + # First launch with no applied watermark re-applies all enabled agents; stub it so the + # test exercises the skills path, not a real configure that shells out to databricks. + patch("ucode.cli.configure_selected_tools", return_value=state), patch("ucode.cli.get_databricks_token", return_value="tok"), patch("ucode.cli._fetch_managed_config", return_value=(managed, False)), patch("ucode.cli.apply_managed_mcp_servers", return_value=[]), @@ -2411,7 +2414,7 @@ def test_selected_tools_skip_picker(self, monkeypatch): ) monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *args, **kwargs: state) monkeypatch.setattr( - cli_mod, "check_gateway_endpoint", lambda state, tool: tool in {"claude", "codex"} + cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: tool in {"claude", "codex"} ) monkeypatch.setattr( cli_mod, @@ -2430,25 +2433,147 @@ def test_selected_tools_skip_picker(self, monkeypatch): monkeypatch.setattr( cli_mod, "configure_selected_tools", - lambda state, tools: configured.append(tools) or {**state, "available_tools": tools}, + lambda state, tools, **kwargs: ( + configured.append(tools) or {**state, "available_tools": tools} + ), ) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) assert cli_mod.configure_workspace_command(selected_tools=["claude", "codex"]) == 0 assert install_calls == ["claude", "codex"] assert configured == [["claude", "codex"]] + def test_managed_config_suppresses_provider_picker(self, monkeypatch): + # When the managed config dictates a tool's models, the "Databricks Hosted vs External" + # provider picker must not be shown; a picked provider would only be overridden at launch. + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": []} + monkeypatch.setattr( + cli_mod, "_prompt_for_configuration", lambda tool=None: ("https://example.com", None) + ) + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr( + cli_mod, "prompt_for_tools", lambda available: pytest.fail("no agent picker") + ) + picked_for: list[str] = [] + monkeypatch.setattr( + cli_mod, "_maybe_select_provider_service", lambda tool, s: picked_for.append(tool) or s + ) + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda s, tools, **k: {**s, "available_tools": tools}, + ) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda s, *a, **k: None) + monkeypatch.setattr( + cli_mod, + "_fetch_managed_config", + lambda s: ( + { + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}, + "codex": {"model_config": {"default_model": "system.ai.gpt-5"}}, + } + }, + False, + ), + ) + + assert cli_mod.configure_workspace_command() == 0 + assert picked_for == [] + + def test_managed_enabled_agents_without_models_still_skip_provider_picker(self, monkeypatch): + # A managed config that only enables agents (no model source) is still fully managed: + # `ucode configure` (no --agents) stays non-interactive and defaults to Databricks Hosted, + # never showing the "Databricks Hosted vs External" provider picker. + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": []} + monkeypatch.setattr( + cli_mod, "_prompt_for_configuration", lambda tool=None: ("https://example.com", None) + ) + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr( + cli_mod, "prompt_for_tools", lambda available: pytest.fail("no agent picker") + ) + monkeypatch.setattr( + cli_mod, + "_maybe_select_provider_service", + lambda tool, s: pytest.fail("no provider picker for a managed workspace"), + ) + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda s, tools, **k: {**s, "available_tools": tools}, + ) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda s, *a, **k: None) + monkeypatch.setattr( + cli_mod, + "_fetch_managed_config", + lambda s: ({"enabled_agents": {"claude": {}, "codex": {}}}, False), + ) + + assert cli_mod.configure_workspace_command() == 0 + + def test_managed_enabled_agents_skip_picker(self, monkeypatch): + # A managed config's enabled_agents is an allowlist: `ucode configure` (no --agents) must + # configure exactly those, never prompt across every workspace-available agent. + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": []} + monkeypatch.setattr( + cli_mod, "_prompt_for_configuration", lambda tool=None: ("https://example.com", None) + ) + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + # Every agent looks available, so only the allowlist should narrow the selection. + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) + monkeypatch.setattr(cli_mod, "_maybe_select_provider_service", lambda tool, s: s) + monkeypatch.setattr( + cli_mod, + "prompt_for_tools", + lambda available: pytest.fail( + "prompt_for_tools should not be called for a managed config" + ), + ) + configured: list[list[str]] = [] + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda state, tools, **kwargs: ( + configured.append(tools) or {**state, "available_tools": tools} + ), + ) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) + monkeypatch.setattr( + cli_mod, + "_fetch_managed_config", + lambda s: ({"enabled_agents": {"claude": {}, "codex": {}}}, False), + ) + + assert cli_mod.configure_workspace_command() == 0 + assert configured == [["claude", "codex"]] + def test_provider_picker_gated_by_interactive_path(self, monkeypatch): import ucode.cli as cli_mod state = {**MINIMAL_STATE, "available_tools": []} monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) - monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: t == "claude") + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t, **kw: t == "claude") monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) monkeypatch.setattr( - cli_mod, "configure_selected_tools", lambda s, tools: {**s, "available_tools": tools} + cli_mod, + "configure_selected_tools", + lambda s, tools, **kwargs: {**s, "available_tools": tools}, ) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda s: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda s, *a, **k: None) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) picked_for: list[str] = [] monkeypatch.setattr( cli_mod, @@ -2480,12 +2605,17 @@ def test_unavailable_selected_tool_errors_before_configure(self, monkeypatch): lambda tool=None: ("https://example.com", None), ) monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *args, **kwargs: state) - monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: tool == "claude") + monkeypatch.setattr( + cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: tool == "claude" + ) monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *args, **kwargs: None) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) monkeypatch.setattr( cli_mod, "configure_selected_tools", - lambda state, tools: pytest.fail("configure_selected_tools should not be called"), + lambda state, tools, **kwargs: pytest.fail( + "configure_selected_tools should not be called" + ), ) with pytest.raises(RuntimeError, match="Codex"): @@ -2496,8 +2626,11 @@ def test_strict_error_mentions_skip_unavailable(self, monkeypatch): state = {**MINIMAL_STATE, "available_tools": []} monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) - monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: tool == "claude") + monkeypatch.setattr( + cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: tool == "claude" + ) monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda *a, **k: (None, False)) with pytest.raises(RuntimeError, match="--skip-unavailable"): cli_mod.configure_workspace_command( @@ -2512,7 +2645,7 @@ def test_skip_unavailable_configures_available_subset(self, monkeypatch): state = {**MINIMAL_STATE, "available_tools": []} monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) monkeypatch.setattr( - cli_mod, "check_gateway_endpoint", lambda state, tool: tool in {"claude", "pi"} + cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: tool in {"claude", "pi"} ) installed: list[str] = [] monkeypatch.setattr( @@ -2524,9 +2657,12 @@ def test_skip_unavailable_configures_available_subset(self, monkeypatch): monkeypatch.setattr( cli_mod, "configure_selected_tools", - lambda state, tools: configured.append(tools) or {**state, "available_tools": tools}, + lambda state, tools, **kwargs: ( + configured.append(tools) or {**state, "available_tools": tools} + ), ) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) warnings: list[str] = [] monkeypatch.setattr(cli_mod, "print_warning", lambda msg: warnings.append(msg)) @@ -2548,11 +2684,14 @@ def test_skip_unavailable_still_fails_when_none_available(self, monkeypatch): state = {**MINIMAL_STATE, "available_tools": []} monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) - monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: False) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: False) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) monkeypatch.setattr( cli_mod, "configure_selected_tools", - lambda state, tools: pytest.fail("configure_selected_tools should not be called"), + lambda state, tools, **kwargs: pytest.fail( + "configure_selected_tools should not be called" + ), ) assert ( @@ -2593,16 +2732,17 @@ def fake_configure_shared_state( monkeypatch.setattr(cli_mod, "configure_shared_state", fake_configure_shared_state) monkeypatch.setattr(cli_mod, "save_state", lambda state: (None, False)) - monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: True) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: True) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: ["claude"]) monkeypatch.setattr(cli_mod, "_maybe_select_provider_service", lambda tool, state: state) monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *args, **kwargs: True) monkeypatch.setattr( cli_mod, "configure_selected_tools", - lambda state, tools: {**state, "available_tools": tools}, + lambda state, tools, **kwargs: {**state, "available_tools": tools}, ) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) assert cli_mod.configure_workspace_command() == 0 assert captured["profile"] == "picked-profile" @@ -2635,19 +2775,20 @@ def fake_configure_shared_state( configured_tools: list[tuple[str, list[str]]] = [] monkeypatch.setattr(cli_mod, "configure_shared_state", fake_configure_shared_state) monkeypatch.setattr(cli_mod, "save_state", lambda state: saved.append(state["workspace"])) - monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: True) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool, **kw: True) monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: ["codex"]) monkeypatch.setattr(cli_mod, "_maybe_select_provider_service", lambda tool, state: state) monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *args, **kwargs: True) monkeypatch.setattr( cli_mod, "configure_selected_tools", - lambda state, tools: ( + lambda state, tools, **kwargs: ( configured_tools.append((state["workspace"], tools)) or {**state, "available_tools": tools} ), ) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) assert ( cli_mod.configure_workspace_command( @@ -2925,7 +3066,7 @@ def test_happy_path_prints_success_without_model_service_detail(self, monkeypatc cli_mod.configure_shared_state(self.WS, profile="DEFAULT") output = _strip_ansi(capsys.readouterr().out) - assert "Unity AI Gateway connected" in output + assert "Unity Gateway connected" in output assert "Model service:" not in output @pytest.mark.parametrize( @@ -2964,7 +3105,7 @@ def test_prints_warning_when_model_service_not_detected( output = " ".join(_strip_ansi(capsys.readouterr().out).split()) assert f"Model service: {expected_model_service}" in output - assert "Unity AI Gateway connected" not in output + assert "Unity Gateway connected" not in output assert "(Legacy) endpoints:" not in output assert "V2" not in output assert "V3" not in output @@ -3002,7 +3143,7 @@ def test_local_gateway_probe_failures_do_not_print_success( cli_mod.configure_shared_state(self.WS, profile="DEFAULT") output = _strip_ansi(capsys.readouterr().out) - assert "Unity AI Gateway connected" not in output + assert "Unity Gateway connected" not in output message = str(excinfo.value) assert "v2" not in message.lower() assert "v3" not in message.lower() @@ -3245,15 +3386,16 @@ def test_skip_validate_skips_agent_validation(self, monkeypatch): state = {**MINIMAL_STATE, "workspace": "https://first.com"} monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) monkeypatch.setattr(cli_mod, "save_state", lambda s: None) - monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t: True) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda s, t, **kw: True) monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) monkeypatch.setattr( cli_mod, "configure_selected_tools", - lambda s, tools: {**s, "available_tools": tools}, + lambda s, tools, **kwargs: {**s, "available_tools": tools}, ) validated: list = [] - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda s: validated.append(s)) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda s, *a, **k: validated.append(s)) result = cli_mod.configure_workspace_command( selected_tools=["codex"], @@ -3271,7 +3413,8 @@ def test_single_tool_validation_is_optional(self, monkeypatch, skip_validate, to state = {**MINIMAL_STATE, "workspace": "https://first.com"} monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) - monkeypatch.setattr(cli_mod, "configure_single_tool", lambda t, s: s) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (None, False)) + monkeypatch.setattr(cli_mod, "configure_single_tool", lambda t, s, **kwargs: s) installed: list = [] monkeypatch.setattr( cli_mod, @@ -3279,7 +3422,9 @@ def test_single_tool_validation_is_optional(self, monkeypatch, skip_validate, to lambda tools, s: installed.append(tools), ) validated: list = [] - monkeypatch.setattr(cli_mod, "validate_tool", lambda t: validated.append(t) or (True, "")) + monkeypatch.setattr( + cli_mod, "validate_tool", lambda t, **kwargs: validated.append(t) or (True, "") + ) result = cli_mod.configure_workspace_command( tool, @@ -3549,12 +3694,15 @@ def _fetch(state): def test_fetches_fresh_when_enabled(self, monkeypatch): monkeypatch.setattr( - "ucode.cli.refresh_managed_config", lambda state: ({"enabled_agents": {}}, False) + "ucode.cli.refresh_managed_config", + lambda state, **kwargs: ({"enabled_agents": {}}, False), ) assert self._fetch({"workspace": "https://w"}) == ({"enabled_agents": {}}, False) def test_feature_disabled_returns_none_and_the_flag(self, monkeypatch): - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, True)) + monkeypatch.setattr( + "ucode.cli.refresh_managed_config", lambda state, **kwargs: (None, True) + ) assert self._fetch({"workspace": "https://w"}) == (None, True) @@ -3571,7 +3719,9 @@ def test_a_removed_model_list_no_longer_skips_discovery(self, monkeypatch): } fresh = {"enabled_agents": {"claude": {"model_config": {}}}} monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: stale_cache) - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (fresh, False)) + monkeypatch.setattr( + "ucode.cli.refresh_managed_config", lambda state, **kwargs: (fresh, False) + ) state = dict(MINIMAL_STATE) with ( @@ -3582,6 +3732,9 @@ def test_a_removed_model_list_no_longer_skips_discovery(self, monkeypatch): patch("ucode.cli.ensure_provider_state", return_value=state), patch("ucode.cli.configure_shared_state", return_value=state) as mock_shared, patch("ucode.cli.configure_tool", return_value=state), + # First launch with no applied watermark re-applies all enabled agents; stub it so the + # test exercises the discovery decision, not a real configure. + patch("ucode.cli.configure_selected_tools", return_value=state), patch("ucode.cli.launch_agent"), ): result = runner.invoke(app, ["claude"]) @@ -3610,9 +3763,13 @@ def _run( monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) if coding_agent_config_feature_disabled: - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (None, True)) + monkeypatch.setattr( + "ucode.cli.refresh_managed_config", lambda state, **kwargs: (None, True) + ) else: - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (managed, False)) + monkeypatch.setattr( + "ucode.cli.refresh_managed_config", lambda state, **kwargs: (managed, False) + ) monkeypatch.setattr("ucode.cli.load_managed_state", lambda ws: cached) monkeypatch.setattr( @@ -3732,7 +3889,9 @@ def test_skip_preflight_still_resolves_an_agent_from_the_managed_config(self, mo "default_agent": "claude", "enabled_agents": {"claude": {"model_config": {"default_model": "m"}}}, } - monkeypatch.setattr("ucode.cli.refresh_managed_config", lambda state: (managed, False)) + monkeypatch.setattr( + "ucode.cli.refresh_managed_config", lambda state, **kwargs: (managed, False) + ) monkeypatch.setattr("ucode.cli._fetch_budget_recommendation", lambda state, m: None) monkeypatch.setattr("ucode.cli._print_managed_summary", lambda *a, **k: None) seen: dict = {} @@ -3776,6 +3935,9 @@ def fake_recommendation(workspace, token): patch("ucode.cli.ensure_provider_state", return_value=state), patch("ucode.cli.configure_shared_state", return_value=state), patch("ucode.cli.configure_tool", return_value=state) as cfg, + # A first launch with no applied watermark re-applies all enabled agents; stub that + # apply-all so these tests exercise the per-launch model/budget path, not a real configure. + patch("ucode.cli.configure_selected_tools", return_value=state), patch("ucode.cli.get_databricks_token", return_value="tok"), patch("ucode.cli._fetch_managed_config", return_value=(managed, False)), patch("ucode.cli.launch_agent"), @@ -3865,6 +4027,9 @@ def test_a_token_failure_does_not_block_the_launch(self, monkeypatch): patch("ucode.cli.ensure_provider_state", return_value=state), patch("ucode.cli.configure_shared_state", return_value=state), patch("ucode.cli.configure_tool", return_value=state), + # First launch with no applied watermark re-applies all enabled agents; stub it so the + # test exercises the budget path, not a real configure. + patch("ucode.cli.configure_selected_tools", return_value=state), patch("ucode.cli.get_databricks_token", side_effect=RuntimeError("token expired")), patch( "ucode.cli._fetch_managed_config", @@ -3974,3 +4139,618 @@ def test_still_logs_in_when_nothing_external_is_set(self, monkeypatch): monkeypatch.delenv("DATABRICKS_BEARER_COMMAND", raising=False) assert self._run(monkeypatch) == ["https://ws.cloud.databricks.com"] + + +class TestConfigureAppliesManagedConfig: + """FIX 1: ug configure must fetch and apply managed config, like the launch path does.""" + + def test_configure_single_tool_applies_managed_static_models(self, monkeypatch): + """ug configure --agent claude applies managed config with static model list.""" + import ucode.agents as agents_mod + + state = dict(MINIMAL_STATE) + # Managed config with static model list + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "models": {"default_sonnet_model": "databricks-claude-sonnet-4"} + } + } + } + } + + configure_one_calls = [] + + def mock_configure_one(tool, state, provider): + configure_one_calls.append({"tool": tool, "state": state, "provider": provider}) + return state + + monkeypatch.setattr(agents_mod, "_configure_one", mock_configure_one) + monkeypatch.setattr(agents_mod, "managed_write_batch", contextlib.nullcontext) + monkeypatch.setattr(agents_mod, "save_state", lambda s: None) + monkeypatch.setattr(agents_mod, "check_gateway_endpoint", lambda s, t, **kw: True) + + # Call configure_single_tool with managed config + agents_mod.configure_single_tool("claude", state, managed=managed) + + # Check that _configure_one was called with the resolved state + assert len(configure_one_calls) == 1 + call = configure_one_calls[0] + # The managed config should have applied the sonnet model + assert call["state"].get("claude_models", {}).get("sonnet") == "databricks-claude-sonnet-4" + + def test_configure_clears_persisted_provider_when_managed_supplies_models_without_provider( + self, monkeypatch + ): + """FIX 2: managed config with static models clears persisted provider.""" + import ucode.agents as agents_mod + + state = dict(MINIMAL_STATE) + # Developer has a persisted provider + state["provider_services"] = {"claude": "some.provider.svc"} + + # Managed config with static models but no provider + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "models": {"default_sonnet_model": "databricks-claude-sonnet-4"} + } + } + } + } + + configure_one_calls = [] + + def mock_configure_one(tool, state, provider): + configure_one_calls.append({"tool": tool, "state": state, "provider": provider}) + return state + + monkeypatch.setattr(agents_mod, "_configure_one", mock_configure_one) + monkeypatch.setattr(agents_mod, "managed_write_batch", contextlib.nullcontext) + monkeypatch.setattr(agents_mod, "save_state", lambda s: None) + monkeypatch.setattr(agents_mod, "check_gateway_endpoint", lambda s, t, **kw: True) + + # Call configure_single_tool with managed config + agents_mod.configure_single_tool("claude", state, managed=managed) + + # The provider passed to _configure_one should be None (cleared) because managed config + # supplies models without a provider + assert len(configure_one_calls) == 1 + call = configure_one_calls[0] + # Provider should be None (cleared) to let managed models drive the picker + assert call["provider"] is None + + +class TestManagedProviderPrecedence: + """FIX 2: provider precedence logic when managed config supplies static/location source.""" + + def test_explicit_provider_conflicts_with_managed_static_models_in_launch(self, monkeypatch): + """Test that explicit --provider with managed static models raises error in launch logic.""" + from ucode.managed_resolve import managed_provider_service, managed_supplies_models + + # This test verifies the logic in the launch path that detects the conflict + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "models": {"default_sonnet_model": "databricks-claude-sonnet-4"} + } + } + } + } + + # Managed config supplies models + assert managed_supplies_models(managed, "claude") + # Managed config does NOT supply a provider + assert managed_provider_service(managed, "claude") is None + + def test_managed_provider_with_explicit_provider_validates_conflict(self, monkeypatch): + """Test that explicit --provider matching managed provider is OK, but mismatch errors.""" + from ucode.managed_resolve import managed_provider_service, managed_supplies_models + + managed_with_provider = { + "enabled_agents": { + "claude": {"model_config": {"model_provider_service": "admin.provider.svc"}} + } + } + + managed_provider = managed_provider_service(managed_with_provider, "claude") + assert managed_provider == "admin.provider.svc" + # Managed supplies models through the provider + assert managed_supplies_models(managed_with_provider, "claude") + + def test_managed_static_models_without_persisted_provider_is_clean(self, monkeypatch): + """When managed specifies static models and no persisted provider, provider=None works.""" + from ucode.managed_resolve import managed_provider_service, managed_supplies_models + + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "models": {"default_sonnet_model": "databricks-claude-sonnet-4"} + } + } + } + } + state = dict(MINIMAL_STATE) + + # Managed supplies models but not through a provider + assert managed_supplies_models(managed, "claude") + assert managed_provider_service(managed, "claude") is None + # And there's no persisted provider either + assert "provider_services" not in state or state["provider_services"].get("claude") is None + + def test_managed_clear_of_provider_is_launch_scoped(self): + # A managed static-list source clears the provider for the launch (so every agent honors + # it), but the developer's saved provider is recorded in the managed overlay and restored + # by save_state, surviving as a fallback if the managed policy later disappears. + from ucode.state import ( + MANAGED_OVERLAY_KEY, + _without_managed_overlay, + set_provider_service, + ) + + state: dict = {"provider_services": {"gemini": "dev.provider.svc"}} + overlay = dict(state.get(MANAGED_OVERLAY_KEY) or {}) + overlay.setdefault("provider_services", state.get("provider_services")) + state = set_provider_service(state, "gemini", None) + state[MANAGED_OVERLAY_KEY] = overlay + + # During the launch the provider reads as cleared. + assert (state.get("provider_services") or {}).get("gemini") is None + # But the persisted state keeps the developer's own provider. + assert _without_managed_overlay(state)["provider_services"] == { + "gemini": "dev.provider.svc" + } + + +class TestMultiAgentManagedConfigRegressions: + """Tests for BUG A and BUG B: managed-overlay leakage and provider re-reading.""" + + def test_configure_selected_tools_multiple_agents_no_managed_overlay_leakage(self, monkeypatch): + """BUG A: multi-agent `ug configure` corrupts persisted state via managed-overlay leakage. + + When configuring multiple tools (claude, codex) with a managed config, + each tool's managed values should reach the config file but NOT persist + in state.json. Tool N's overlay should not replace tool N-1's overlay + while tool N-1's managed values remain in the accumulated state. + """ + import ucode.agents as agents_mod + + state = dict(MINIMAL_STATE) + # Developer has previously configured both tools with developer-chosen models + state["claude_models"] = {"sonnet": "dev-claude-sonnet"} + state["codex_models"] = ["dev-codex-mini"] + + # Managed config specifies different models for both tools + managed = { + "enabled_agents": { + "claude": { + "model_config": {"models": {"default_sonnet_model": "managed-claude-sonnet"}} + }, + "codex": {"model_config": {"models": ["managed-codex-mini"]}}, + } + } + + configure_one_calls = [] + + def mock_configure_one(tool, state, provider): + # Record what state was passed to _configure_one for each tool + configure_one_calls.append( + { + "tool": tool, + "state_snapshot": dict(state), # Capture the state passed to config + "provider": provider, + } + ) + return state + + monkeypatch.setattr(agents_mod, "_configure_one", mock_configure_one) + monkeypatch.setattr(agents_mod, "managed_write_batch", contextlib.nullcontext) + monkeypatch.setattr(agents_mod, "save_state", lambda s: None) + monkeypatch.setattr(agents_mod, "check_gateway_endpoint", lambda s, t, **kw: True) + + # Call configure_selected_tools with managed config + result_state = agents_mod.configure_selected_tools( + state, ["claude", "codex"], managed=managed + ) + + # Verify that claude was configured with managed model + claude_call = next((c for c in configure_one_calls if c["tool"] == "claude"), None) + assert claude_call is not None + assert ( + claude_call["state_snapshot"].get("claude_models", {}).get("sonnet") + == "managed-claude-sonnet" + ), "Claude should be configured with managed model" + + # Verify that codex was configured with managed model + codex_call = next((c for c in configure_one_calls if c["tool"] == "codex"), None) + assert codex_call is not None + assert codex_call["state_snapshot"].get("codex_models") == ["managed-codex-mini"], ( + "Codex should be configured with managed model" + ) + + # FIX A TEST: Verify that the final persisted state has developer's original models, + # not the managed ones. This proves that managed values were used for configuration + # but not persisted. + assert result_state.get("claude_models", {}).get("sonnet") == "dev-claude-sonnet", ( + "Final state should have developer's original claude model, not managed model" + ) + assert result_state.get("codex_models") == ["dev-codex-mini"], ( + "Final state should have developer's original codex model, not managed model" + ) + + # Verify no managed overlay marker remains + assert "_managed_overlay" not in result_state, "No overlay should remain in final state" + + def test_gemini_launch_respects_managed_provider_clearing(self, monkeypatch): + """BUG B: when managed config clears a provider, Gemini's launch/refresh should see it. + + When a managed config supplies its own models without a provider, the + launch path clears the provider (sets it to None). Gemini should launch + with provider=None, not re-reading the stale persisted provider. + """ + import subprocess + + import ucode.agents.gemini as gemini_mod + from ucode.agents.args import LaunchOptions + + state = dict(MINIMAL_STATE) + # Developer has persisted a provider for gemini + state["provider_services"] = {"gemini": "stale.provider.svc"} + state["gemini_models"] = ["gemini-2.0-flash"] + + # Simulate managed config that cleared the provider + # (this would be set by the launch path's FIX B: set_provider_service call) + cleared_state = dict(state) + providers = dict(cleared_state.get("provider_services") or {}) + providers.pop("gemini", None) + if providers: + cleared_state["provider_services"] = providers + else: + cleared_state.pop("provider_services", None) + + # Mock the token refresh and config write + write_config_calls = [] + + def mock_write_tool_config(state, model, force_refresh=False, provider=None): + write_config_calls.append( + { + "model": model, + "provider": provider, # Capture the provider passed to write + "state_has_provider": bool(state.get("provider_services", {}).get("gemini")), + } + ) + return (state, "mock-token") + + def mock_get_provider_service(state, tool): + # This simulates the actual behavior after FIX B + providers = state.get("provider_services", {}) + return providers.get(tool) if isinstance(providers, dict) else None + + mock_process = MagicMock() + mock_process.wait.return_value = 0 + + monkeypatch.setattr(gemini_mod, "write_tool_config", mock_write_tool_config) + monkeypatch.setattr(gemini_mod, "get_provider_service", mock_get_provider_service) + monkeypatch.setattr(gemini_mod, "build_runtime_env", lambda *a, **kw: {}) + monkeypatch.setattr(subprocess, "Popen", lambda *a, **kw: mock_process) + + # Launch with the cleared state (simulating what the launch path does after FIX B) + try: + gemini_mod.launch(cleared_state, [], options=LaunchOptions()) + except SystemExit: + pass # launch() raises SystemExit, that's fine for this test + + # FIX B TEST: Verify that write_tool_config was called with provider=None, + # not the stale persisted provider. This proves that Gemini correctly used + # the cleared provider from state, not re-reading the stale value. + assert len(write_config_calls) > 0, "write_tool_config should be called during launch" + config_call = write_config_calls[0] + assert config_call["provider"] is None, ( + f"Gemini should launch with provider=None (managed cleared it), not {config_call['provider']}" + ) + + +class TestFix1PostConfigureValidation: + """FIX 1: Post-configure validation uses managed-resolved config, not stripped state.""" + + def test_single_tool_configure_validates_with_managed_provider(self, monkeypatch): + """When managing gemini through provider B while dev provider A is persisted, + validation must use managed provider B, not the persisted stripped state provider A. + """ + import ucode.agents as agents_mod + import ucode.agents.gemini as gemini_mod + + # Developer has persisted provider A for gemini + state = dict(MINIMAL_STATE) + state["provider_services"] = {"gemini": "dev-provider-a"} + state["gemini_models"] = ["gemini-2.0-flash"] + + # Managed config routes gemini through provider B instead + managed = { + "enabled_agents": { + "gemini": { + "model_config": { + "model_provider_service": "managed-provider-b", + "models": ["managed-gemini-2.0-flash"], + } + } + } + } + + # Track what provider validate_env receives + validate_env_calls = [] + + def mock_validate_env(validate_state): + provider = validate_state.get("provider_services", {}).get("gemini") + validate_env_calls.append({"provider": provider}) + # Build env with the provider it received + return { + "GEMINI_MODEL": "managed-gemini-2.0-flash", + "GOOGLE_GEMINI_BASE_URL": "https://example.databricks.com/ai-gateway/gemini", + "GEMINI_API_KEY_AUTH_MECHANISM": "bearer", + "GEMINI_API_KEY": "token", + "GEMINI_CLI_CUSTOM_HEADERS": f"User-Agent:ucode/test,Databricks-Model-Provider-Service:{provider}" + if provider + else "User-Agent:ucode/test", + "OAUTH_TOKEN": "token", + } + + # Mock subprocess to make validation pass + mock_result = MagicMock() + mock_result.returncode = 0 + + monkeypatch.setattr(gemini_mod, "validate_env", mock_validate_env) + monkeypatch.setattr("subprocess.run", lambda *a, **kw: mock_result) + + # Mock configure_single_tool to apply managed and return resolved state + from ucode.managed_resolve import resolve_state + + def mock_configure_single_tool(tool, state, managed=None): + if managed: + # Apply managed config (resolves state) + state = resolve_state(managed, state, tool) + # Simulate config write and save (which strips overlay) + # The returned state has the overlay stripped + from ucode.state import _without_managed_overlay + + return _without_managed_overlay(state) + + monkeypatch.setattr(agents_mod, "configure_single_tool", mock_configure_single_tool) + monkeypatch.setattr( + agents_mod, "install_databricks_ai_tools_for_agents", lambda *a, **kw: None + ) + + # Simulate the configure_workspace_command path for single tool + from ucode.managed_resolve import resolve_state as resolve_state_impl + + # After configure_single_tool runs, state is stripped of overlay + configured_state = mock_configure_single_tool("gemini", state, managed=managed) + # The configured_state should have stripped the managed provider back to dev provider + assert configured_state.get("provider_services", {}).get("gemini") == "dev-provider-a", ( + "Configured state should be stripped back to developer's provider" + ) + + # Now simulate the validation path (FIX 1): re-resolve the state with managed config + validate_state = resolve_state_impl(managed, configured_state, "gemini") + + # Call validate_env with the re-resolved state (as the fix does) + from ucode.agents.gemini import validate_env + + try: + validate_env(validate_state) + except Exception: + pass # May fail due to missing deps, we only care it received the right state + + # FIX 1 VERIFICATION: validate_env should have seen the managed provider B, + # not the stripped developer provider A + assert len(validate_env_calls) > 0, "validate_env should be called" + assert validate_env_calls[0]["provider"] == "managed-provider-b", ( + f"Validation should use managed provider B, not {validate_env_calls[0]['provider']}" + ) + + def test_validate_tool_accepts_pre_resolved_state(self, monkeypatch): + """validate_tool() accepts an optional state parameter so it validates + with managed-resolved values instead of freshly loaded persisted state.""" + import ucode.agents as agents_mod + + state_with_managed = { + "workspace": "https://example.databricks.com", + "profile": None, + "provider_services": {"gemini": "managed-provider"}, + } + + validate_env_calls = [] + + def mock_validate_env(state): + validate_env_calls.append(state.get("provider_services", {}).get("gemini")) + return {} + + def mock_validate_cmd(binary): + return [binary, "--help"] + + mock_result = MagicMock() + mock_result.returncode = 0 + + # Mock the gemini module + mock_gemini = MagicMock() + mock_gemini.validate_env = mock_validate_env + mock_gemini.validate_cmd = mock_validate_cmd + mock_gemini.skip_validation = MagicMock(return_value=False) + + monkeypatch.setattr(agents_mod, "_MODULES", {"gemini": mock_gemini}) + monkeypatch.setattr("ucode.agents.__init__.TOOL_SPECS", {"gemini": {"binary": "gemini"}}) + monkeypatch.setattr("subprocess.run", lambda *a, **kw: mock_result) + + # Call validate_tool with explicit state (FIX 1) + ok, err = agents_mod.validate_tool("gemini", state=state_with_managed) + + # Verify validate_env received the passed state, not a freshly loaded one + assert len(validate_env_calls) > 0, "validate_env should be called" + assert validate_env_calls[0] == "managed-provider", ( + "validate_env should receive the pre-resolved state with managed provider" + ) + + +class TestFix2ManagedOnlyAgentAvailability: + """FIX 2: Agents with managed-only models (no discovered) must count as available.""" + + def test_check_gateway_endpoint_with_managed_only_models(self): + """An agent with no discovered models but managed models should be available.""" + import ucode.agents as agents_mod + + # State with NO discovered Claude models + state = dict(MINIMAL_STATE) + state["claude_models"] = {} # Empty, no discovered models + + # Managed config provides claude models + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "models": { + "default_sonnet_model": "managed-claude-sonnet", + } + } + } + } + } + + # Without managed config, claude should be unavailable + available_without_managed = agents_mod.check_gateway_endpoint(state, "claude", managed=None) + assert not available_without_managed, "Claude should be unavailable without managed models" + + # With managed config, claude should be available (FIX 2) + available_with_managed = agents_mod.check_gateway_endpoint(state, "claude", managed=managed) + assert available_with_managed, ( + "Claude should be available when managed config supplies models" + ) + + def test_check_gateway_endpoint_prefers_discovered_over_managed(self): + """When both discovered and managed models exist, availability is True.""" + import ucode.agents as agents_mod + + # State with discovered Claude models + state = dict(MINIMAL_STATE) + state["claude_models"] = {"sonnet": "databricks-claude-sonnet"} + + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "models": { + "default_sonnet_model": "managed-claude-sonnet", + } + } + } + } + } + + # With discovered models, should be available + available = agents_mod.check_gateway_endpoint(state, "claude", managed=managed) + assert available, "Claude should be available with discovered models" + + def test_unavailable_agent_with_no_discovery_and_no_managed(self): + """A truly unavailable agent (no discovery AND no managed) stays unavailable.""" + import ucode.agents as agents_mod + + # State with NO discovered Claude models + state = dict(MINIMAL_STATE) + state["claude_models"] = {} # Empty + + # Managed config does NOT provide claude models + managed = { + "enabled_agents": { + "gemini": {"model_config": {}} # Only gemini, no claude + } + } + + # Without discovered or managed models, claude should be unavailable + available = agents_mod.check_gateway_endpoint(state, "claude", managed=managed) + assert not available, ( + "Claude should be unavailable when both discovery and managed provide nothing" + ) + + +class TestLaunchVersionGate: + """A launch re-applies the CLI Managed Configuration (the only step that prompts for the sudo OS + write) only when its update_time is newer than the applied watermark; an unchanged launch skips + the apply-all entirely.""" + + _MANAGED = { + "update_time": "2026-09-11T16:09:31.820Z", + "enabled_agents": {"codex": {"model_config": {"default_model": "system.ai.gpt-5"}}}, + } + + def _launch(self, monkeypatch, applied_update_time): + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": ["codex"]} + if applied_update_time is not None: + state = {**state, "applied_managed_update_time": applied_update_time} + 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 t: state) + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (self._MANAGED, False)) + monkeypatch.setattr(cli_mod, "_fetch_budget_recommendation", lambda s, m: None) + monkeypatch.setattr(cli_mod, "save_state", lambda s: None) + monkeypatch.setattr(cli_mod, "launch_agent", MagicMock()) + monkeypatch.setattr(cli_mod, "configure_tool", MagicMock(return_value=state)) + apply_all = MagicMock(return_value=state) + monkeypatch.setattr(cli_mod, "configure_selected_tools", apply_all) + result = runner.invoke(app, ["codex"]) + return result, apply_all + + def test_unchanged_config_skips_the_apply(self, monkeypatch): + result, apply_all = self._launch( + monkeypatch, applied_update_time="2026-09-11T16:09:31.820Z" + ) + assert result.exit_code == 0, result.output + apply_all.assert_not_called() + + def test_changed_config_reapplies_all_enabled_agents(self, monkeypatch): + result, apply_all = self._launch( + monkeypatch, applied_update_time="2026-09-11T15:00:00.000Z" + ) + assert result.exit_code == 0, result.output + apply_all.assert_called_once() + assert apply_all.call_args.args[1] == ["codex"] + + def test_first_apply_with_no_watermark_reapplies(self, monkeypatch): + result, apply_all = self._launch(monkeypatch, applied_update_time=None) + assert result.exit_code == 0, result.output + apply_all.assert_called_once() + + def test_auto_configure_launch_still_applies_managed(self, monkeypatch): + # Regression: a first launch of an unconfigured tool (needs_auto_configure) must still apply + # the managed config to every enabled agent, not merely stamp the watermark. Otherwise the + # OS-managed file is left unmanaged while the watermark records an apply that never happened. + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": []} + monkeypatch.setattr(cli_mod, "ensure_bootstrap_dependencies", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "_auto_configure_tool", lambda *a, **k: None) + monkeypatch.setattr(cli_mod, "load_state", lambda: state) + monkeypatch.setattr(cli_mod, "ensure_provider_state", lambda t: state) + monkeypatch.setattr(cli_mod, "configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "_fetch_managed_config", lambda s: (self._MANAGED, False)) + monkeypatch.setattr(cli_mod, "_fetch_budget_recommendation", lambda s, m: None) + saved: list[dict] = [] + monkeypatch.setattr(cli_mod, "save_state", lambda s: saved.append(dict(s))) + monkeypatch.setattr(cli_mod, "launch_agent", MagicMock()) + monkeypatch.setattr(cli_mod, "configure_tool", MagicMock(return_value=state)) + apply_all = MagicMock(return_value=state) + monkeypatch.setattr(cli_mod, "configure_selected_tools", apply_all) + + result = runner.invoke(app, ["codex"]) + + assert result.exit_code == 0, result.output + apply_all.assert_called_once() + assert apply_all.call_args.args[1] == ["codex"] + assert any( + s.get("applied_managed_update_time") == self._MANAGED["update_time"] for s in saved + ) diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index dd00e04d..e470be0d 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -249,6 +249,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._fetch_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_e2e.py b/tests/test_e2e.py index 831e6603..b8e20ae8 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -320,7 +320,7 @@ def test_only_picks_codex_writes_only_codex_config(self, tmp_path, monkeypatch, # Skip binary install + post-config validation; we're testing the # selection plumbing, not the agent binaries themselves. monkeypatch.setattr(cli_mod, "install_tool_binary", lambda tool, **kwargs: True) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) # Answer the provider picker; "databricks" keeps the Databricks path. monkeypatch.setattr(cli_mod, "prompt_for_selection", lambda prompt, options: "databricks") @@ -353,7 +353,7 @@ def test_rerun_with_different_pick_preserves_previous( cli_mod, "_prompt_for_configuration", lambda tool=None: (e2e_workspace, None) ) monkeypatch.setattr(cli_mod, "install_tool_binary", lambda tool, **kwargs: True) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) # Answer the provider picker; "databricks" keeps the Databricks path. monkeypatch.setattr(cli_mod, "prompt_for_selection", lambda prompt, options: "databricks") @@ -394,7 +394,7 @@ def test_empty_pick_returns_zero_and_writes_nothing(self, tmp_path, monkeypatch, "install_tool_binary", lambda tool, **kwargs: install_calls.append(tool) or True, ) - monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state, *a, **k: None) rc = cli_mod.configure_workspace_command() assert rc == 0 diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 796e0653..94f81c62 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -13,16 +13,19 @@ import ucode.managed_config as mc_mod from ucode.managed_config import ( get_managed_config, + load_managed_configuration, load_managed_state, + managed_config_is_newer, managed_state_workspace, + managed_update_time, normalize_managed_config, refresh_managed_config, save_managed_state, ) from ucode.managed_setup import serialize_managed_config -# A representative raw CodingAgentConfig proto-JSON manifest (mirrors what the API returns). -RAW_MANIFEST = { +# A representative raw CodingAgentConfig proto-JSON manifest using deprecated model_config oneof. +RAW_MANIFEST_DEPRECATED = { "name": "coding-agent-configs/abc-123", "workspace_id": 1653573648247579, "default_agent": "CODING_AGENT_CLAUDE_CODE", @@ -81,43 +84,43 @@ } -class TestNormalize: +class TestNormalizeDeprecated: def test_full_manifest_maps_enums_to_tool_names(self): - cfg = normalize_managed_config(RAW_MANIFEST) + cfg = normalize_managed_config(RAW_MANIFEST_DEPRECATED) assert cfg["name"] == "coding-agent-configs/abc-123" assert cfg["default_agent"] == "claude" assert set(cfg["enabled_agents"]) == {"claude", "opencode"} def test_claude_agent_config_fields(self): - claude = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["claude"] + claude = normalize_managed_config(RAW_MANIFEST_DEPRECATED)["enabled_agents"]["claude"] assert claude["custom_headers"] == {"x-databricks-workspace": "eng-ml-inference"} assert claude["tracing_table"] == "main.default.ucode_traces" assert claude["model_config"]["default_model"] == "system.ai.claude-opus-4-8" assert claude["model_config"]["models"]["default_opus_model"] == "system.ai.claude-opus-4-8" def test_opencode_model_list_is_flat(self): - opencode = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["opencode"] + opencode = normalize_managed_config(RAW_MANIFEST_DEPRECATED)["enabled_agents"]["opencode"] assert opencode["model_config"]["models"] == [ "system.ai.claude-opus-4-8", "system.ai.kimi-k2-7-code", ] def test_mcp_servers_map_type_enums_to_tags(self): - mcp = normalize_managed_config(RAW_MANIFEST)["mcp_servers"] + mcp = normalize_managed_config(RAW_MANIFEST_DEPRECATED)["mcp_servers"] assert mcp == [ {"name": "system.ai.github", "type": "mcp-service"}, {"name": "some-space-id", "type": "genie-space"}, ] def test_skills_and_tracing_and_budget(self): - cfg = normalize_managed_config(RAW_MANIFEST) + cfg = normalize_managed_config(RAW_MANIFEST_DEPRECATED) assert cfg["skills"] == {"names": ["system.ai.pdf-extraction"]} assert cfg["tracing_table"] == "main.default.ucode_traces" assert cfg["budget_policy"]["budget_id"] == "c6563b45-df9a-4b19-afb2-d42dc2b52576" assert cfg["budget_policy"]["tiers"][1]["default_agent"] == "opencode" def test_reads_top_level_display_name(self): - cfg = normalize_managed_config({**RAW_MANIFEST, "display_name": "paved-path"}) + cfg = normalize_managed_config({**RAW_MANIFEST_DEPRECATED, "display_name": "paved-path"}) assert cfg["display_name"] == "paved-path" def test_display_name_survives_the_serialize_round_trip(self): @@ -138,14 +141,149 @@ def test_empty_manifest_yields_empty_dict(self): assert normalize_managed_config({}) == {} +# A CodingAgentConfig in the current agent-config wire shape as emitted by ai-gateway-api. +# The wire format uses: default_models map (not separate default_model + default_alias_models), +# model field names (model_services, unity_catalog_location), tier field names +# (recommended_agent, recommended_model), and mcp_servers/skills as parallel arrays. +RAW_MANIFEST = { + "spec_version": 1, + "retrieved_time": "2026-09-09T22:00:00Z", + "default_agent": "CODING_AGENT_CLAUDE_CODE", + "enabled_agents": [ + { + "agent": "CODING_AGENT_CLAUDE_CODE", + "config": { + "models": { + "model_services": [ + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5", + ], + }, + "default_models": { + "default_model": "system.ai.claude-opus-4-8", + "default_opus_model": "system.ai.claude-opus-4-8", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + "default_haiku_model": "system.ai.claude-haiku-4-5", + }, + "smart_routing": {"enabled": True}, + "http_headers": {"x-databricks-workspace": "eng-ml-inference"}, + }, + }, + { + "agent": "CODING_AGENT_CODEX", + "config": { + "models": {"model_provider_service": "main.default.openai-mps"}, + "default_models": {"default_model": "gpt-5.4"}, + }, + }, + ], + "mcp_servers": { + "names": ["system.ai.github", "some-space-id"], + "tags": ["MCP_SERVER_TYPE_UC_SERVICE", "MCP_SERVER_TYPE_GENIE"], + }, + "skills": {"names": ["system.ai.pdf-extraction"], "tags": []}, + "tracing": {"enabled": True}, + "spend_tiers": { + "budget_id": "c6563b45-df9a-4b19-afb2-d42dc2b52576", + "tiers": [ + { + "spending_percentage": 0.9, + "recommended_agent": "CODING_AGENT_CODEX", + "recommended_model": "gpt-5.4", + }, + ], + }, +} + + +class TestNormalize: + def test_maps_agents_to_tool_names(self): + cfg = normalize_managed_config(RAW_MANIFEST) + assert cfg["default_agent"] == "claude" + assert set(cfg["enabled_agents"]) == {"claude", "codex"} + + def test_claude_alias_models_map_to_family_slots(self): + claude = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["claude"] + assert claude["model_config"]["models"] == { + "default_opus_model": "system.ai.claude-opus-4-8", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + "default_haiku_model": "system.ai.claude-haiku-4-5", + } + assert claude["model_config"]["default_model"] == "system.ai.claude-opus-4-8" + + def test_http_headers_map_to_custom_headers(self): + claude = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["claude"] + assert claude["custom_headers"] == {"x-databricks-workspace": "eng-ml-inference"} + + def test_static_names_are_carried(self): + claude = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["claude"] + assert claude["model_config"]["names"] == [ + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5", + ] + + def test_unity_catalog_location_is_carried(self): + raw = { + "spec_version": 1, + "enabled_agents": [ + { + "agent": "CODING_AGENT_CLAUDE_CODE", + "config": {"models": {"unity_catalog_location": "main.agents"}}, + } + ], + } + claude = normalize_managed_config(raw)["enabled_agents"]["claude"] + assert claude["model_config"]["model_service_location"] == "main.agents" + + def test_codex_provider_service(self): + codex = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["codex"] + assert codex["model_config"]["model_provider_service"] == "main.default.openai-mps" + assert codex["model_config"]["default_model"] == "gpt-5.4" + + def test_budget_policy_carries_budget_id_and_tiers(self): + cfg = normalize_managed_config(RAW_MANIFEST) + assert cfg["budget_policy"]["budget_id"] == "c6563b45-df9a-4b19-afb2-d42dc2b52576" + assert cfg["budget_policy"]["tiers"][0]["default_agent"] == "codex" + + def test_tracing_enabled_carries_no_table(self): + # Current config `tracing.enabled` has no table FQN, so nothing lands in the (table-shaped) internal key. + assert "tracing_table" not in normalize_managed_config(RAW_MANIFEST) + + def test_spec_version_not_carried_into_internal_manifest(self): + # Kept out so the serialize/normalize round trip (which never sees spec_version) is unaffected. + assert "spec_version" not in normalize_managed_config(RAW_MANIFEST) + + def test_legacy_agent_config_fields_still_read(self): + # An entry with only the deprecated custom_headers + model_config oneof still normalizes. + raw = { + "enabled_agents": [ + { + "agent": "CODING_AGENT_CLAUDE_CODE", + "config": { + "custom_headers": {"x-h": "v"}, + "model_config": {"claude": {"default_model": "system.ai.claude-opus-4-8"}}, + }, + } + ] + } + claude = normalize_managed_config(raw)["enabled_agents"]["claude"] + assert claude["custom_headers"] == {"x-h": "v"} + assert claude["model_config"]["default_model"] == "system.ai.claude-opus-4-8" + + class TestGetManagedConfig: - def test_returns_normalized_first_config(self, monkeypatch): + def test_returns_the_first_config_raw(self, monkeypatch): + # get_managed_config returns the config verbatim now; normalization happens on read. monkeypatch.setattr( - mc_mod, "fetch_managed_coding_agent_configs", lambda ws, tok: ([RAW_MANIFEST], None) + mc_mod, + "fetch_managed_coding_agent_configs", + lambda ws, tok: ([RAW_MANIFEST_DEPRECATED], None), ) cfg, reason = get_managed_config("https://ws", "tok") assert reason is None - assert cfg["default_agent"] == "claude" + assert cfg == RAW_MANIFEST_DEPRECATED def test_no_config_is_not_an_error(self, monkeypatch): monkeypatch.setattr( @@ -197,19 +335,106 @@ def test_feature_disabled_is_not_swallowed_as_not_found(self, monkeypatch): assert cfg is None assert reason_out == reason + def test_returns_raw_wire_shape(self, monkeypatch): + # The raw wire config is returned unchanged (enabled_agents stays a list of {agent, config}); + # normalize_managed_config's translation is covered by its own tests. + monkeypatch.setattr( + mc_mod, "fetch_managed_coding_agent_configs", lambda ws, tok: ([RAW_MANIFEST], None) + ) + cfg, reason = get_managed_config("https://ws", "tok") + assert reason is None + assert cfg == RAW_MANIFEST + + def test_spec_version_newer_than_supported_is_refused(self, monkeypatch): + raw = {**RAW_MANIFEST, "spec_version": mc_mod.MAX_SPEC_VERSION + 1} + monkeypatch.setattr( + mc_mod, "fetch_managed_coding_agent_configs", lambda ws, tok: ([raw], None) + ) + cfg, reason = get_managed_config("https://ws", "tok") + # Refused (reason set) rather than misread, so the launch path keeps the last-known-good + # cache instead of applying a config it can't parse. + assert cfg is None + assert reason is not None and "spec_version" in reason + + @pytest.mark.parametrize("bad_spec", ["2", 2.0, True]) + def test_malformed_spec_version_is_refused(self, monkeypatch, bad_spec): + raw = {**RAW_MANIFEST, "spec_version": bad_spec} + monkeypatch.setattr( + mc_mod, "fetch_managed_coding_agent_configs", lambda ws, tok: ([raw], None) + ) + cfg, reason = get_managed_config("https://ws", "tok") + assert cfg is None + assert reason is not None and "spec_version" in reason + + +class TestManagedConfigStub: + def test_stub_short_circuits_the_http_read(self, tmp_path, monkeypatch): + stub = tmp_path / "managed.json" + stub.write_text(json.dumps(RAW_MANIFEST), encoding="utf-8") + monkeypatch.setenv("UCODE_MANAGED_CONFIG_STUB", str(stub)) + + def _fail(ws, tok): + raise AssertionError("stub set: the HTTP read must not run") + + monkeypatch.setattr(mc_mod, "fetch_managed_coding_agent_configs", _fail) + cfg, reason = get_managed_config("https://ws", "tok") + assert reason is None + assert cfg == RAW_MANIFEST + + def test_stub_applies_the_spec_version_gate(self, tmp_path, monkeypatch): + stub = tmp_path / "managed.json" + stub.write_text( + json.dumps({**RAW_MANIFEST, "spec_version": mc_mod.MAX_SPEC_VERSION + 1}), + encoding="utf-8", + ) + monkeypatch.setenv("UCODE_MANAGED_CONFIG_STUB", str(stub)) + cfg, reason = get_managed_config("https://ws", "tok") + assert cfg is None + assert reason is not None and "spec_version" in reason + + def test_unreadable_stub_falls_through_to_the_http_read(self, tmp_path, monkeypatch): + monkeypatch.setenv("UCODE_MANAGED_CONFIG_STUB", str(tmp_path / "missing.json")) + monkeypatch.setattr( + mc_mod, "fetch_managed_coding_agent_configs", lambda ws, tok: ([RAW_MANIFEST], None) + ) + cfg, reason = get_managed_config("https://ws", "tok") + assert reason is None + assert cfg == RAW_MANIFEST + + +class TestUnsupportedSpecFallback: + def test_unsupported_spec_warns_on_cold_launch(self, monkeypatch): + # No cache + a too-new spec_version is proof a policy exists, so surface it rather than + # silently treating the workspace as having no managed config. + warnings: list = [] + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + reason = "This workspace's managed config needs a newer Unity Gateway (spec_version 2)." + assert mc_mod._persisted_fallback("https://ws", reason) is None + assert warnings and "spec_version" in warnings[0] + + def test_transient_failure_stays_silent_on_cold_launch(self, monkeypatch): + warnings: list = [] + monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) + monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) + assert mc_mod._persisted_fallback("https://ws", "HTTP 503 Service Unavailable") is None + assert warnings == [] + class TestPersistence: @pytest.fixture(autouse=True) def _managed_path(self, tmp_path, monkeypatch): - path = tmp_path / ".ucode" / "managed-state.json" - monkeypatch.setattr(mc_mod, "MANAGED_STATE_PATH", path) + path = tmp_path / ".ucode" / "managed-configuration.json" + monkeypatch.setattr(mc_mod, "MANAGED_CONFIGURATION_PATH", path) return path - def test_save_then_load_round_trips(self, _managed_path): - cfg = normalize_managed_config(RAW_MANIFEST) - save_managed_state("https://ws.example.com", cfg) - loaded = load_managed_state("https://ws.example.com") - assert loaded == cfg + def test_save_stores_raw_and_load_normalizes(self, _managed_path): + # The file holds the raw config verbatim; load_managed_state normalizes on read. + save_managed_state("https://ws.example.com", RAW_MANIFEST) + assert load_managed_configuration("https://ws.example.com") == RAW_MANIFEST + assert load_managed_state("https://ws.example.com") == normalize_managed_config( + RAW_MANIFEST + ) def test_saved_file_is_0600(self, _managed_path): save_managed_state("https://ws.example.com", {"default_agent": "claude"}) @@ -261,7 +486,7 @@ def test_corrupt_file_reads_as_absent(self, _managed_path): def test_loaded_config_serializes_to_a_json_encodable_payload(self, _managed_path): # `ucode publish` POSTs the serialized config, so a manifest that survives a disk round-trip # must still serialize to something json.dumps accepts with no custom encoder. - cfg = normalize_managed_config(RAW_MANIFEST) + cfg = normalize_managed_config(RAW_MANIFEST_DEPRECATED) save_managed_state("https://ws.example.com", cfg) loaded = load_managed_state("https://ws.example.com") assert loaded is not None @@ -272,7 +497,7 @@ class TestFetchClient: """fetch_managed_coding_agent_configs lives in databricks.py; test its response parsing.""" def test_extracts_configs_list(self, monkeypatch): - payload = {"coding_agent_configs": [RAW_MANIFEST]} + payload = {"coding_agent_configs": [RAW_MANIFEST_DEPRECATED]} monkeypatch.setattr( db_mod, "_http_get_json", @@ -333,16 +558,19 @@ class TestRefreshManagedConfig: def _stub_token(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_databricks_token", lambda ws, profile: "tok") - def test_persists_and_returns_the_manifest(self, monkeypatch): + def test_persists_raw_and_returns_the_normalized_manifest(self, monkeypatch): saved: list[tuple] = [] - monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (MANAGED, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: saved.append((ws, cfg))) - assert refresh_managed_config(_state()) == (MANAGED, False) - assert saved == [(WORKSPACE, MANAGED)] + monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (RAW_MANIFEST, None)) + monkeypatch.setattr( + mc_mod, "save_managed_state", lambda ws, cfg, **kwargs: saved.append((ws, cfg)) + ) + # The raw config is persisted verbatim; the caller gets the normalized manifest. + assert refresh_managed_config(_state()) == (normalize_managed_config(RAW_MANIFEST), False) + assert saved == [(WORKSPACE, RAW_MANIFEST)] def test_no_managed_config_returns_none(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg, **kwargs: None) result, _ = refresh_managed_config(_state()) assert result is None @@ -412,7 +640,9 @@ def test_permission_denied_warns_and_keeps_the_cached_config(self, monkeypatch): monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) monkeypatch.setattr(mc_mod, "print_warning", lambda msg: warnings.append(msg)) monkeypatch.setattr( - mc_mod, "save_managed_state", lambda ws, cfg: pytest.fail("must not clear the cache") + mc_mod, + "save_managed_state", + lambda ws, cfg, **kwargs: pytest.fail("must not clear the cache"), ) assert refresh_managed_config(_state()) == (MANAGED, False) assert "not readable by you" in warnings[0] @@ -421,7 +651,7 @@ def test_no_config_on_the_server_does_not_use_a_stale_persisted_file(self, monke # A successful read saying "no config" means the admin removed it — that's authoritative, # so a previously persisted file must not resurrect the old policy. monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg, **kwargs: None) monkeypatch.setattr( mc_mod, "load_managed_state", lambda ws: pytest.fail("must not fall back") ) @@ -433,7 +663,9 @@ def test_no_config_on_the_server_clears_the_persisted_one(self, monkeypatch): # failed read would put a dead policy back into force. saved: list[tuple] = [] monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: saved.append((ws, cfg))) + monkeypatch.setattr( + mc_mod, "save_managed_state", lambda ws, cfg, **kwargs: saved.append((ws, cfg)) + ) monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: None) result, _ = refresh_managed_config(_state()) assert result is None @@ -477,7 +709,9 @@ def test_feature_disabled_ignores_a_cached_config_and_sets_the_flag(self, monkey reason = 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED"}' monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, reason)) monkeypatch.setattr(mc_mod, "load_managed_state", lambda ws: MANAGED) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: saved.append((ws, cfg))) + monkeypatch.setattr( + mc_mod, "save_managed_state", lambda ws, cfg, **kwargs: saved.append((ws, cfg)) + ) monkeypatch.setattr( mc_mod, "print_warning", @@ -500,13 +734,111 @@ def test_transient_failure_does_not_set_the_flag(self, monkeypatch): def test_successful_no_config_clears_the_flag(self, monkeypatch): monkeypatch.setattr(mc_mod, "get_managed_config", lambda ws, tok: (None, None)) - monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg: None) + monkeypatch.setattr(mc_mod, "save_managed_state", lambda ws, cfg, **kwargs: None) state = _state() result, flag = refresh_managed_config(state) assert result is None assert flag is False +class TestRefreshAlwaysFetches: + """The launch-time refresh always hits the control plane; the 30-minute TTL is gone. + + Whether to re-apply the fetched config is decided separately by the caller via + ``managed_config_is_newer`` against the persisted applied watermark, so refresh never short- + circuits on a cached copy. + """ + + @pytest.fixture(autouse=True) + def _stub_token(self, monkeypatch): + monkeypatch.setattr(mc_mod, "get_databricks_token", lambda ws, profile: "tok") + + @staticmethod + def _counting_fetch(monkeypatch, result=(RAW_MANIFEST, None)): + calls = {"n": 0} + + def fetch(ws, tok): + calls["n"] += 1 + return result + + monkeypatch.setattr(mc_mod, "get_managed_config", fetch) + return calls + + def test_fetches_even_with_a_persisted_config(self, monkeypatch): + # A previously-persisted config no longer short-circuits: every launch re-reads the workspace. + save_managed_state(WORKSPACE, RAW_MANIFEST) + calls = self._counting_fetch(monkeypatch) + result, flag = refresh_managed_config(_state()) + assert result == normalize_managed_config(RAW_MANIFEST) + assert flag is False + assert calls["n"] == 1 + + def test_persists_the_fetched_config_raw_without_a_timestamp(self, monkeypatch): + self._counting_fetch(monkeypatch) + refresh_managed_config(_state()) + # The on-disk payload is the raw config verbatim and carries no retrieved_at field. + stored = json.loads(mc_mod.MANAGED_CONFIGURATION_PATH.read_text(encoding="utf-8")) + assert "retrieved_at" not in stored + assert stored["config"] == RAW_MANIFEST + + def test_first_launch_with_no_cache_fetches(self, monkeypatch): + calls = self._counting_fetch(monkeypatch) + refresh_managed_config(_state()) + assert calls["n"] == 1 + + +class TestManagedUpdateTime: + def test_reads_top_level_update_time(self): + assert managed_update_time({"update_time": "2026-09-11T16:09:31.820Z"}) == ( + "2026-09-11T16:09:31.820Z" + ) + + def test_none_when_absent_or_not_a_dict(self): + assert managed_update_time({}) is None + assert managed_update_time(None) is None + + def test_normalize_captures_update_time(self): + raw = { + "default_agent": "CODING_AGENT_CLAUDE_CODE", + "update_time": "2026-09-11T16:09:31.820Z", + } + assert normalize_managed_config(raw)["update_time"] == "2026-09-11T16:09:31.820Z" + + +class TestManagedConfigIsNewer: + OLDER = {"update_time": "2026-09-11T16:00:00.000Z"} + NEWER = {"update_time": "2026-09-11T16:09:31.820Z"} + + def test_true_when_fetched_is_newer(self): + assert managed_config_is_newer(self.NEWER, self.OLDER["update_time"]) is True + + def test_false_when_equal(self): + assert managed_config_is_newer(self.NEWER, self.NEWER["update_time"]) is False + + def test_false_when_fetched_is_older(self): + assert managed_config_is_newer(self.OLDER, self.NEWER["update_time"]) is False + + def test_true_when_no_prior_watermark(self): + # First apply: nothing applied yet, so any fetched config counts as newer. + assert managed_config_is_newer(self.NEWER, None) is True + + def test_true_when_fetched_has_no_update_time(self): + # A config without a parseable update_time is treated as newer so a launch re-applies it + # rather than trusting possibly-stale local settings. + assert managed_config_is_newer({}, self.NEWER["update_time"]) is True + + def test_true_when_watermark_unparseable(self): + assert managed_config_is_newer(self.NEWER, "not-a-timestamp") is True + + def test_handles_z_and_offset_suffixes_equivalently(self): + assert ( + managed_config_is_newer( + {"update_time": "2026-09-11T16:00:00+00:00"}, "2026-09-11T16:00:00Z" + ) + is False + ) + + class TestGetModelRecommendation: """The budget recommendation read. Every response field is optional server-side.""" diff --git a/tests/test_managed_export.py b/tests/test_managed_export.py index b11b3b07..8e6e219f 100644 --- a/tests/test_managed_export.py +++ b/tests/test_managed_export.py @@ -1,10 +1,10 @@ """Tests for `ucode export` and its :mod:`ucode.managed_export` backing module. -`export` is read-only and offline: it serializes the local managed config to the external -proto-JSON `CodingAgentConfig` that `ucode publish -f ` consumes. These focus on the parts -that must not regress — a clean machine-readable stdout stream, byte-identical file output, atomic -replacement that never truncates on failure, exclusion of server-owned fields, and the absence of -any auth/admin/network call. +`export` is read-only and offline: it dumps the locally stored managed config (the raw +`CodingAgentConfig` the gateway returned) as portable JSON, minus server-owned metadata. These +focus on the parts that must not regress — a clean machine-readable stdout stream, byte-identical +file output, atomic replacement that never truncates on failure, exclusion of server-owned fields, +and the absence of any auth/admin/network call. """ from __future__ import annotations @@ -22,7 +22,7 @@ import ucode.managed_export as export_mod from ucode.cli import app from ucode.managed_config import normalize_managed_config -from ucode.managed_setup import serialize_managed_config, validate_manifest +from ucode.managed_setup import validate_manifest runner = CliRunner() @@ -30,84 +30,104 @@ _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") -FULL_MANIFEST = { +# The raw CodingAgentConfig as the gateway returns it (what managed-configuration.json stores). +# Carries server-owned metadata (name, update_time, ...) that export must strip. +FULL_RAW_CONFIG = { "name": "coding-agent-configs/abc123", - "default_agent": "claude", - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}, - "codex": {"model_config": {"default_model": "system.ai.gpt-5-6"}}, - }, - "mcp_servers": [{"name": "system.ai.slack", "type": "mcp-service"}], + "spec_version": 1, + "update_time": "2026-09-11T16:09:31.820Z", + "retrieved_time": "2026-09-11T19:22:49.491Z", + "created_user_id": 74233111714547, + "default_agent": "CODING_AGENT_CLAUDE_CODE", + "enabled_agents": [ + { + "agent": "CODING_AGENT_CLAUDE_CODE", + "config": {"default_models": {"default_model": "system.ai.claude-opus-4-8"}}, + }, + { + "agent": "CODING_AGENT_CODEX", + "config": {"default_models": {"default_model": "system.ai.gpt-5-6"}}, + }, + ], + "mcp_servers": {"names": ["system.ai.slack"], "tags": ["MCP_SERVER_TYPE_UC_SERVICE"]}, "skills": {"names": ["main.default"]}, } +# Server-assigned metadata export strips; the rest of FULL_RAW_CONFIG is emitted verbatim. +_STRIPPED = {"name", "update_time", "retrieved_time", "created_user_id"} +_EXPORTED_CONFIG = {k: v for k, v in FULL_RAW_CONFIG.items() if k not in _STRIPPED} + +# A raw config that normalizes to something structurally invalid (claude enabled, no model). +INVALID_RAW_CONFIG = {"enabled_agents": [{"agent": "CODING_AGENT_CLAUDE_CODE", "config": {}}]} + @pytest.fixture(autouse=True) def _isolate_settings(tmp_path, monkeypatch): """Point the managed-config file at a tmp dir so no test touches the real ~/.ucode.""" monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - monkeypatch.setattr(managed_config_mod, "MANAGED_STATE_PATH", tmp_path / "managed-state.json") + monkeypatch.setattr( + managed_config_mod, "MANAGED_CONFIGURATION_PATH", tmp_path / "managed-configuration.json" + ) monkeypatch.setattr(config_io_mod, "_dry_run", False) @contextlib.contextmanager -def _with_manifest(manifest: dict | None, workspace: str | None = WORKSPACE): - """Patch the module's local reads so a test controls the source config without disk or network.""" +def _with_config(config: dict | None, workspace: str | None = WORKSPACE): + """Patch the module's local reads so a test controls the raw source config without disk/network.""" with ( patch.object(export_mod, "load_state", return_value={"workspace": workspace}), - patch.object(export_mod, "load_managed_state", return_value=manifest), + patch.object(export_mod, "load_managed_configuration", return_value=config), ): yield class TestBuildPayload: - def test_excludes_server_owned_resource_name(self): - with _with_manifest(FULL_MANIFEST): + def test_excludes_server_owned_metadata(self): + with _with_config(FULL_RAW_CONFIG): payload = export_mod.build_export_payload() - assert "name" not in payload + for field in _STRIPPED: + assert field not in payload + # The rest of the raw config passes through verbatim (no translation). assert payload["default_agent"] == "CODING_AGENT_CLAUDE_CODE" - assert payload["mcp_servers"] == [ - {"name": "system.ai.slack", "type": "MCP_SERVER_TYPE_UC_SERVICE"} - ] + assert payload["mcp_servers"] == { + "names": ["system.ai.slack"], + "tags": ["MCP_SERVER_TYPE_UC_SERVICE"], + } + assert payload["enabled_agents"] == FULL_RAW_CONFIG["enabled_agents"] def test_envelope_workspace_first_then_spec_version(self): - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): payload = export_mod.build_export_payload() assert list(payload)[:2] == ["workspace", "spec_version"] assert payload["workspace"] == WORKSPACE assert payload["spec_version"] == 1 - def test_matches_serialize_minus_name_under_envelope(self): - config = serialize_managed_config(FULL_MANIFEST) - config.pop("name", None) - expected = {"workspace": WORKSPACE, "spec_version": 1, **config} - with _with_manifest(FULL_MANIFEST): + def test_payload_is_raw_config_minus_server_fields(self): + expected = {"workspace": WORKSPACE, "spec_version": 1, **_EXPORTED_CONFIG} + with _with_config(FULL_RAW_CONFIG): assert export_mod.build_export_payload() == expected - def test_config_roundtrips_through_parser_and_validator(self): - with _with_manifest(FULL_MANIFEST): + def test_exported_config_still_parses_and_validates(self): + with _with_config(FULL_RAW_CONFIG): payload = export_mod.build_export_payload() config = {k: v for k, v in payload.items() if k not in ("workspace", "spec_version")} - reparsed = normalize_managed_config(config) - assert validate_manifest(reparsed, None) == [] - assert serialize_managed_config(reparsed) == config + assert validate_manifest(normalize_managed_config(config), None) == [] def test_no_config_is_actionable(self): with ( patch.object(export_mod, "load_state", return_value={}), - patch.object(export_mod, "load_managed_state", return_value=None), + patch.object(export_mod, "load_managed_configuration", return_value=None), ): - with pytest.raises(RuntimeError, match="No managed coding-agent config found"): + with pytest.raises(RuntimeError, match="No CLI Managed Configuration found"): export_mod.build_export_payload() def test_invalid_config_is_rejected(self): - invalid = {"enabled_agents": {"claude": {}}} - with _with_manifest(invalid): + with _with_config(INVALID_RAW_CONFIG): with pytest.raises(RuntimeError, match="not valid"): export_mod.build_export_payload() def test_falls_back_to_managed_state_workspace(self): - managed_config_mod.save_managed_state(WORKSPACE, FULL_MANIFEST) + managed_config_mod.save_managed_state(WORKSPACE, FULL_RAW_CONFIG) with patch.object(export_mod, "load_state", return_value={}): payload = export_mod.build_export_payload() assert payload["default_agent"] == "CODING_AGENT_CLAUDE_CODE" @@ -115,7 +135,7 @@ def test_falls_back_to_managed_state_workspace(self): class TestExportCommandStdout: def test_emits_valid_json_with_single_trailing_newline(self, capsys): - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): export_mod.export_command() captured = capsys.readouterr() out = captured.out @@ -125,23 +145,23 @@ def test_emits_valid_json_with_single_trailing_newline(self, capsys): assert captured.err == "" def test_stdout_has_no_rich_or_human_output(self, capsys): - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): export_mod.export_command() out = capsys.readouterr().out assert _ANSI_RE.search(out) is None - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): payload = export_mod.build_export_payload() assert json.loads(out) == payload class TestExportCommandFile: def test_file_bytes_identical_to_stdout_and_stdout_empty(self, capsys, tmp_path): - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): export_mod.export_command() stdout_bytes = capsys.readouterr().out dest = tmp_path / "config.json" - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): export_mod.export_command(file_path=str(dest)) captured = capsys.readouterr() assert captured.out == "" @@ -149,7 +169,7 @@ def test_file_bytes_identical_to_stdout_and_stdout_empty(self, capsys, tmp_path) def test_expands_user_home_in_output_path(self, capsys, tmp_path, monkeypatch): monkeypatch.setenv("HOME", str(tmp_path)) - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): export_mod.export_command(file_path="~/config.json") capsys.readouterr() assert (tmp_path / "config.json").exists() @@ -157,7 +177,7 @@ def test_expands_user_home_in_output_path(self, capsys, tmp_path, monkeypatch): def test_replaces_existing_destination(self, tmp_path): dest = tmp_path / "config.json" dest.write_text("stale contents", encoding="utf-8") - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): export_mod.export_command(file_path=str(dest)) assert json.loads(dest.read_text(encoding="utf-8"))["default_agent"] == ( "CODING_AGENT_CLAUDE_CODE" @@ -166,16 +186,14 @@ def test_replaces_existing_destination(self, tmp_path): def test_invalid_config_leaves_existing_destination_unchanged(self, tmp_path): dest = tmp_path / "config.json" dest.write_text("original", encoding="utf-8") - invalid = {"enabled_agents": {"claude": {}}} - with _with_manifest(invalid): + with _with_config(INVALID_RAW_CONFIG): with pytest.raises(RuntimeError): export_mod.export_command(file_path=str(dest)) assert dest.read_text(encoding="utf-8") == "original" def test_invalid_config_does_not_create_destination(self, tmp_path): dest = tmp_path / "config.json" - invalid = {"enabled_agents": {"claude": {}}} - with _with_manifest(invalid): + with _with_config(INVALID_RAW_CONFIG): with pytest.raises(RuntimeError): export_mod.export_command(file_path=str(dest)) assert not dest.exists() @@ -183,7 +201,7 @@ def test_invalid_config_does_not_create_destination(self, tmp_path): def test_missing_parent_directory_fails_without_creating_it(self, tmp_path): missing_parent = tmp_path / "nope" dest = missing_parent / "config.json" - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): with pytest.raises(RuntimeError, match="parent directory does not exist"): export_mod.export_command(file_path=str(dest)) assert not missing_parent.exists() @@ -191,7 +209,7 @@ def test_missing_parent_directory_fails_without_creating_it(self, tmp_path): def test_write_failure_is_actionable_and_leaves_no_temp_file(self, tmp_path): dest = tmp_path / "adir" dest.mkdir() - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): with pytest.raises(RuntimeError, match="Failed to write"): export_mod.export_command(file_path=str(dest)) leftovers = [p.name for p in tmp_path.iterdir() if p.name.startswith(".ucode-export-")] @@ -201,7 +219,7 @@ def test_write_failure_is_actionable_and_leaves_no_temp_file(self, tmp_path): class TestNoAuthOrAdmin: def test_no_admin_or_token_lookup_occurs(self, capsys): with ( - _with_manifest(FULL_MANIFEST), + _with_config(FULL_RAW_CONFIG), patch("ucode.databricks.is_workspace_admin") as admin, patch("ucode.databricks.get_databricks_token") as token, ): @@ -214,7 +232,7 @@ def test_admin_and_non_admin_produce_identical_output(self, capsys): outputs = [] for admin_value in (True, False): with ( - _with_manifest(FULL_MANIFEST), + _with_config(FULL_RAW_CONFIG), patch("ucode.databricks.is_workspace_admin", MagicMock(return_value=admin_value)), ): export_mod.export_command() @@ -237,7 +255,7 @@ def test_help_documents_command_and_flags(self): def test_long_and_short_file_flags_both_write_the_file(self, tmp_path): for flag in ("--file", "-f"): dest = tmp_path / f"cfg{flag.strip('-')}.json" - with _with_manifest(FULL_MANIFEST): + with _with_config(FULL_RAW_CONFIG): result = runner.invoke(app, ["export", flag, str(dest)]) assert result.exit_code == 0 assert json.loads(dest.read_text(encoding="utf-8"))["default_agent"] == ( @@ -247,7 +265,7 @@ def test_long_and_short_file_flags_both_write_the_file(self, tmp_path): def test_no_config_exits_nonzero(self): with ( patch.object(export_mod, "load_state", return_value={}), - patch.object(export_mod, "load_managed_state", return_value=None), + patch.object(export_mod, "load_managed_configuration", return_value=None), ): result = runner.invoke(app, ["export"]) assert result.exit_code == 1 diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index 96d137fb..b7c78def 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -152,7 +152,7 @@ def test_none_for_agent_not_in_manifest(self): class TestResolveState: def test_does_not_mutate_input_state(self): - # managed-state.json and state.json stay separate files: resolution is per-write and + # managed-configuration.json and state.json stay separate files: resolution is per-write and # in-memory, so the developer's own state must come back untouched. state = _state(claude_models={"opus": "local-opus"}) before = json.dumps(state, sort_keys=True) @@ -185,7 +185,7 @@ def test_layers_provider_without_dropping_other_tools(self): class TestStateFileIsNotRewritten: """The managed config must win by precedence, not by overwriting the developer's state file. - managed-state.json and state.json stay separate on disk: resolution happens in memory and only + managed-configuration.json and state.json stay separate on disk: resolution happens in memory and only the generated agent settings file reflects it. These tests deliberately let the real ``save_state`` run against a temp ``state.json`` — stubbing it out is what let this regress, because the overwrite happens inside ``write_tool_config``, one layer below the resolver. diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index 140bc3d9..f394da64 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -59,13 +59,19 @@ def _minimal_manifest() -> dict: def _full_manifest() -> dict: - """A manifest exercising every field the read side normalizes.""" + """A manifest exercising the fields the current wire shape round-trips. + + The serialize side is being retired, so it inverts only the current wire shape. Two internal + forms are deliberately outside the round-trip and excluded here: a workspace ``tracing_table`` + (the current wire tracing is a client on/off with no table), and the legacy flat-list ``models`` + key (flat-list agents use ``names`` now). ``test_round_trip_boundary_current_wire_shape_only`` + asserts those boundaries explicitly. + """ return { "default_agent": "claude", "enabled_agents": { "claude": { "custom_headers": {"x-databricks-workspace": "eng-ml-inference"}, - "tracing_table": "main.default.claude-traces", "model_config": { "default_model": "system.ai.claude-opus-4-8", "models": { @@ -80,7 +86,7 @@ def _full_manifest() -> dict: "opencode": { "model_config": { "default_model": "system.ai.claude-opus-4-8", - "models": ["system.ai.claude-opus-4-8", "system.ai.kimi-k2-6"], + "names": ["system.ai.claude-opus-4-8", "system.ai.kimi-k2-6"], }, }, }, @@ -89,7 +95,6 @@ def _full_manifest() -> dict: {"name": "genie-space-id", "type": "genie-space"}, ], "skills": {"names": ["system.ai.pdf-extraction"]}, - "tracing_table": "main.default.ucode-traces", "budget_policy": { "display_name": "eng-tiered-routing", "budget_id": "c6563b45-df9a-4b19-afb2-d42dc2b52576", @@ -127,25 +132,47 @@ def test_inversion_is_lossless(self): class TestRoundTrip: - """serialize -> normalize must be the identity on a ucode-native manifest.""" + """serialize -> normalize is the identity on a current-wire-shape ucode-native manifest. + + The serialize/publish authoring path is being retired, so the inverse is only claimed over the + current wire shape. See test_round_trip_boundary_current_wire_shape_only for what is out. + """ def test_full_manifest_round_trips(self): manifest = _full_manifest() assert normalize_managed_config(serialize_managed_config(manifest)) == manifest + def test_round_trip_boundary_current_wire_shape_only(self): + # Explicitly document the inverse's boundary rather than hide it by omission: the retiring + # serialize side does not carry a workspace tracing table (the current wire tracing is a + # client on/off) or the legacy flat-list `models` key (flat-list agents use `names` now), + # so a manifest built from those forms does not round-trip. + manifest = { + "tracing_table": "main.default.traces", + "enabled_agents": { + "opencode": { + "model_config": {"models": ["system.ai.a", "system.ai.b"]}, + }, + }, + } + round_tripped = normalize_managed_config(serialize_managed_config(manifest)) + assert "tracing_table" not in round_tripped + assert round_tripped != manifest + def test_minimal_manifest_round_trips(self): manifest = _minimal_manifest() assert normalize_managed_config(serialize_managed_config(manifest)) == manifest def test_every_known_agent_round_trips(self): # Each agent's oneof variant must survive a round trip, including the flat-list agents and - # codex (which has no model list at all). + # codex (which has no model list at all). Claude uses 'models' (a dict of slots); flat-list + # agents use 'names' (a list); codex uses neither. for tool in AGENT_TOOL_TO_ENUM: model_config: dict = {"default_model": "system.ai.some-model"} if tool == "claude": model_config["models"] = {"default_opus_model": "system.ai.claude-opus-4-8"} elif tool != "codex": - model_config["models"] = ["system.ai.some-model"] + model_config["names"] = ["system.ai.some-model"] manifest = { "default_agent": tool, "enabled_agents": {tool: {"model_config": model_config}}, @@ -171,15 +198,14 @@ def test_claude_model_config_uses_family_slots(self): for entry in payload["enabled_agents"] if entry["agent"] == "CODING_AGENT_CLAUDE_CODE" ) - variant = claude["config"]["model_config"] - assert set(variant) == {"claude"} - assert variant["claude"]["models"] == { + assert claude["config"]["default_models"] == { + "default_model": "system.ai.claude-opus-4-8", "default_opus_model": "system.ai.claude-opus-4-8", "default_sonnet_model": "system.ai.claude-sonnet-4-6", } def test_codex_model_config_has_no_model_list(self): - # CodexModelConfig carries only model_provider_service + default_model. + # Codex carries only default_models, no model_services list. manifest = { "default_agent": "codex", "enabled_agents": { @@ -193,9 +219,9 @@ def test_codex_model_config_has_no_model_list(self): }, } payload = serialize_managed_config(manifest) - variant = payload["enabled_agents"][0]["config"]["model_config"]["codex"] - assert "models" not in variant - assert variant["default_model"] == "system.ai.gpt-5-6" + config = payload["enabled_agents"][0]["config"] + assert "models" not in config + assert config["default_models"]["default_model"] == "system.ai.gpt-5-6" def test_flat_list_agents_use_repeated_models(self): payload = serialize_managed_config(_full_manifest()) @@ -204,8 +230,10 @@ def test_flat_list_agents_use_repeated_models(self): for entry in payload["enabled_agents"] if entry["agent"] == "CODING_AGENT_OPENCODE" ) - variant = opencode["config"]["model_config"]["opencode"] - assert variant["models"] == ["system.ai.claude-opus-4-8", "system.ai.kimi-k2-6"] + assert opencode["config"]["models"]["model_services"] == [ + "system.ai.claude-opus-4-8", + "system.ai.kimi-k2-6", + ] def test_model_provider_service_is_carried_through(self): manifest = { @@ -220,43 +248,46 @@ def test_model_provider_service_is_carried_through(self): }, } payload = serialize_managed_config(manifest) - variant = payload["enabled_agents"][0]["config"]["model_config"]["claude"] - assert variant["model_provider_service"] == "main.default.anthropic-mps" + config = payload["enabled_agents"][0]["config"] + assert config["models"]["model_provider_service"] == "main.default.anthropic-mps" def test_mcp_types_map_to_proto_enums(self): payload = serialize_managed_config(_full_manifest()) - assert payload["mcp_servers"] == [ - {"name": "system.ai.github", "type": "MCP_SERVER_TYPE_UC_SERVICE"}, - {"name": "genie-space-id", "type": "MCP_SERVER_TYPE_GENIE"}, - ] - - def test_tracing_becomes_a_table_object(self): - payload = serialize_managed_config(_full_manifest()) - assert payload["tracing"] == {"table": "main.default.ucode-traces"} + assert payload["mcp_servers"] == { + "names": ["system.ai.github", "genie-space-id"], + "tags": ["MCP_SERVER_TYPE_UC_SERVICE", "MCP_SERVER_TYPE_GENIE"], + } - def test_per_agent_tracing_override(self): - payload = serialize_managed_config(_full_manifest()) + def test_per_agent_tracing_enabled(self): + manifest = { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "tracing_table": "main.default.claude-traces", + "model_config": {"default_model": "system.ai.claude-opus-4-8"}, + } + }, + } + payload = serialize_managed_config(manifest) claude = next( entry for entry in payload["enabled_agents"] if entry["agent"] == "CODING_AGENT_CLAUDE_CODE" ) - assert claude["config"]["tracing_config"] == {"table": "main.default.claude-traces"} + assert claude["config"]["tracing"] == {"enabled": True} def test_budget_tiers_keep_fractions(self): # The server validates 0 <= spending_percentage <= 1, so these stay fractions. payload = serialize_managed_config(_full_manifest()) - tiers = payload["budget_policy"]["tiers"] + tiers = payload["spend_tiers"]["tiers"] assert [tier["spending_percentage"] for tier in tiers] == [0.8, 1.0] - assert tiers[1]["default_agent"] == "CODING_AGENT_OPENCODE" + assert tiers[1]["recommended_agent"] == "CODING_AGENT_OPENCODE" - def test_the_deprecated_top_level_budget_id_is_never_emitted(self): - # `CodingAgentConfig.budget_id` (field 3) is deprecated in favour of - # `budget_policy.budget_id`, and the CRUD handler rejects a write that sets it. The budget - # id must appear only under the policy. + def test_budget_id_appears_only_under_spend_tiers(self): + # The budget id must appear under spend_tiers, not at the top level. payload = serialize_managed_config(_full_manifest()) assert "budget_id" not in payload - assert payload["budget_policy"]["budget_id"] == "c6563b45-df9a-4b19-afb2-d42dc2b52576" + assert payload["spend_tiers"]["budget_id"] == "c6563b45-df9a-4b19-afb2-d42dc2b52576" def test_a_manifest_carrying_a_top_level_budget_id_still_omits_it(self): # A hand-written `--from-file` manifest could set it; the serializer must not pass it on. @@ -281,7 +312,10 @@ def test_unknown_mcp_type_is_dropped(self): payload = serialize_managed_config( {"mcp_servers": [{"name": "a", "type": "not-a-type"}, {"name": "b", "type": "sql"}]} ) - assert payload["mcp_servers"] == [{"name": "b", "type": "MCP_SERVER_TYPE_DATABRICKS_SQL"}] + assert payload["mcp_servers"] == { + "names": ["b"], + "tags": ["MCP_SERVER_TYPE_DATABRICKS_SQL"], + } def test_empty_manifest_serializes_to_empty_payload(self): assert serialize_managed_config({}) == {} @@ -396,7 +430,7 @@ def test_first_model_wins_within_a_family(self): def test_unidentifiable_models_are_skipped(self): assert claude_model_slots(["system.ai.gpt-5-6"]) == {} - def test_slots_serialize_into_the_claude_variant(self): + def test_slots_serialize_into_the_default_models_map(self): manifest = { "default_agent": "claude", "enabled_agents": { @@ -409,8 +443,11 @@ def test_slots_serialize_into_the_claude_variant(self): }, } payload = serialize_managed_config(manifest) - variant = payload["enabled_agents"][0]["config"]["model_config"]["claude"] - assert variant["models"] == {"default_opus_model": "system.ai.claude-opus-4-8"} + config = payload["enabled_agents"][0]["config"] + assert config["default_models"] == { + "default_model": "system.ai.claude-opus-4-8", + "default_opus_model": "system.ai.claude-opus-4-8", + } class TestClaudeFamilyCandidates: diff --git a/tests/test_state.py b/tests/test_state.py index e8ad10fd..05159e49 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -12,12 +12,14 @@ STATE_VERSION, build_agent_state, clear_state, + get_applied_managed_update_time, get_provider_service, hydrate_state, load_full_state, load_state, mark_tool_managed, save_state, + set_applied_managed_update_time, set_provider_service, ) @@ -184,6 +186,28 @@ def test_clearing_one_tool_keeps_the_other(self): assert get_provider_service(state, "codex") == "main.a.openai" +class TestAppliedManagedUpdateTime: + def test_get_returns_none_when_unset(self): + assert get_applied_managed_update_time({}) is None + assert get_applied_managed_update_time({"applied_managed_update_time": ""}) is None + + def test_set_and_get_roundtrip(self): + state = set_applied_managed_update_time({}, "2026-09-11T16:09:31.820Z") + assert get_applied_managed_update_time(state) == "2026-09-11T16:09:31.820Z" + + def test_set_none_clears_the_key(self): + state = set_applied_managed_update_time({}, "2026-09-11T16:09:31.820Z") + state = set_applied_managed_update_time(state, None) + assert get_applied_managed_update_time(state) is None + assert "applied_managed_update_time" not in state + + def test_survives_save_load_roundtrip(self): + state = {"workspace": FAKE_WS, "profile": None} + state = set_applied_managed_update_time(state, "2026-09-11T16:09:31.820Z") + save_state(state) + assert get_applied_managed_update_time(load_state()) == "2026-09-11T16:09:31.820Z" + + # --------------------------------------------------------------------------- # hydrate_state # ---------------------------------------------------------------------------