diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index d0d0380..1afe962 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -62,7 +62,11 @@ def _resolve_web_search_model(state: dict) -> str | None: WEB_SEARCH_MCP_NAME = "web_search" -_CLAUDE_MODEL_RE = re.compile(r"^databricks-claude-(opus|sonnet)-(\d+)-(\d+)(.*)$") +# Matches both the AI Gateway form (`databricks-claude-opus-4-8`) and the UC +# model-services form (`system.ai.claude-opus-4-8`). +_CLAUDE_MODEL_RE = re.compile( + r"^(?:system\.ai\.)?(?:databricks-)?claude-(opus|sonnet)-(\d+)-(\d+)(.*)$" +) # Env keys the MLflow Stop hook reads to route traces. Written into the # settings `env` block alongside the hook itself. diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index e0bb64b..e8d5eb9 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -255,6 +255,10 @@ def _openai_model_id(model: str | None) -> str | None: def _codex_model_id(model: str | None) -> str | None: + # UC model-services ids (`system.ai.gpt-5`) route by name through the + # gateway, so they must be sent verbatim — not rewritten to an OpenAI id. + if model and model.startswith("system.ai."): + return model if model in CODEX_OPENAI_ID_INCOMPATIBLE_MODELS: return model return _openai_model_id(model) @@ -263,7 +267,12 @@ def _codex_model_id(model: str | None) -> str | None: def _parse_gpt(model: str | None) -> tuple[int, int | None, int | None, str] | None: if not model: return None - match = _GPT_RE.fullmatch(model.split("/")[-1]) + # Strip the UC model-services prefix so `system.ai.gpt-5` parses for version + # selection; the original id is preserved by callers that need it verbatim. + tail = model.split("/")[-1] + if tail.startswith("system.ai."): + tail = tail[len("system.ai.") :] + match = _GPT_RE.fullmatch(tail) if not match: return None major, minor, patch, suffix = match.groups() diff --git a/src/ucode/cli.py b/src/ucode/cli.py index c363e22..15973b0 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -33,6 +33,7 @@ discover_claude_models, discover_codex_models, discover_gemini_models, + discover_model_services, ensure_ai_gateway_v2, ensure_databricks_auth, find_profile_name_for_host, @@ -41,6 +42,7 @@ install_databricks_cli, normalize_workspace_url, run_databricks_login, + uc_enabled, ) from ucode.mcp import ( MCP_CLIENTS, @@ -149,6 +151,8 @@ def configure_shared_state( profile: str | None = None, tools: list[str] | None = None, force_login: bool = False, + enable_uc: bool | None = None, + reset_uc: bool = False, ) -> dict: """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state. @@ -158,9 +162,28 @@ def configure_shared_state( ``--profile`` to every CLI invocation so ambiguous `~/.databrickscfg` entries (e.g. DEFAULT and a named profile both pointing at the same host) don't error out. If ``None``, we resolve it from the host after login. + ``enable_uc`` is the resolved CLI flag (`--enable-uc`): when not None + it overrides both the env var and the persisted state. + ``reset_uc`` is True only on the explicit ``ucode configure`` flow. """ workspace = normalize_workspace_url(workspace) - previous_workspace = load_state().get("workspace") + prior_state = load_state() + previous_workspace = prior_state.get("workspace") + # Precedence: explicit CLI flag > env var > (configure: reset to False; + # launch: target workspace's persisted state). Use *target* state on the + # launch path so the flag is sticky per-workspace and doesn't leak + # across workspace switches. + # TODO: when this flips uc_enabled True->False, prune any + # `system.ai.*` MCP services from state["mcp_servers"] (and their + # cross-tool registrations). Today they linger as orphans pointing at + # /ai-gateway/mcp-services/* until the user re-runs `configure mcp` + # or switches workspaces. + if enable_uc is None: + if reset_uc: + enable_uc = uc_enabled(default=False) + else: + target_ws_state = load_full_state().get("workspaces", {}).get(workspace) or {} + enable_uc = uc_enabled(default=bool(target_ws_state.get("uc_enabled"))) fetch_all = tools is None if force_login: run_databricks_login(workspace, profile) @@ -184,19 +207,29 @@ def configure_shared_state( claude_reason: str | None = None gemini_reason: str | None = None codex_reason: str | None = None - with spinner("Fetching available models..."): + claude_models = {} + gemini_models = [] + codex_models = [] + if enable_uc: + # Opt-in: one UC model-services call yields all families as + # `system.ai.` ids, bucketed by name. The single reason is + # shared across the families that were requested. + with spinner("Fetching available models (model services)..."): + ms_claude, ms_codex, ms_gemini, ms_reason = discover_model_services(workspace, token) if want_claude: - claude_models, claude_reason = discover_claude_models(workspace, token) - else: - claude_models = {} + claude_models, claude_reason = ms_claude, ms_reason if want_gemini: - gemini_models, gemini_reason = discover_gemini_models(workspace, token) - else: - gemini_models = [] + gemini_models, gemini_reason = ms_gemini, ms_reason if want_codex: - codex_models, codex_reason = discover_codex_models(workspace, token) - else: - codex_models = [] + codex_models, codex_reason = ms_codex, ms_reason + else: + with spinner("Fetching available models..."): + if want_claude: + claude_models, claude_reason = discover_claude_models(workspace, token) + if want_gemini: + gemini_models, gemini_reason = discover_gemini_models(workspace, token) + if want_codex: + codex_models, codex_reason = discover_codex_models(workspace, token) opencode_models: dict[str, list[str]] = {} if claude_models: opencode_models["anthropic"] = list(claude_models.values()) @@ -210,6 +243,9 @@ def configure_shared_state( state["profile"] = profile else: state.pop("profile", None) + # Persist the resolved flag so subsequent launches stay on the same + # discovery path without the env var or CLI flag being re-passed. + state["uc_enabled"] = enable_uc state["base_urls"] = build_shared_base_urls(workspace) if want_claude: state["claude_models"] = claude_models @@ -239,13 +275,22 @@ def _configure_shared_workspace_states( tools: list[str] | None, *, force_login: bool, + enable_uc: bool | None = None, + reset_uc: bool = False, ) -> list[dict]: if not workspaces: raise RuntimeError("At least one workspace must be provided.") states: list[dict] = [] for workspace, profile in workspaces: states.append( - configure_shared_state(workspace, profile=profile, tools=tools, force_login=force_login) + configure_shared_state( + workspace, + profile=profile, + tools=tools, + force_login=force_login, + enable_uc=enable_uc, + reset_uc=reset_uc, + ) ) return states @@ -256,6 +301,8 @@ def configure_workspace_command( workspaces: list[tuple[str, str | None]] | None = None, *, prompt_optional_updates: bool = True, + enable_uc: bool | None = None, + reset_uc: bool = False, ) -> int: if tool is not None and selected_tools is not None: raise RuntimeError("Use either --agent or --agents, not both.") @@ -263,7 +310,9 @@ def configure_workspace_command( workspace_entries = workspaces or [_prompt_for_configuration(tool)] if tool is not None: - states = _configure_shared_workspace_states(workspace_entries, [tool], force_login=True) + states = _configure_shared_workspace_states( + workspace_entries, [tool], force_login=True, enable_uc=enable_uc, reset_uc=reset_uc + ) state = states[0] state = configure_single_tool(tool, state) spec = TOOL_SPECS[tool] @@ -290,7 +339,13 @@ def configure_workspace_command( raise RuntimeError(f"{spec['display']} validation failed — config reverted.") return 0 - states = _configure_shared_workspace_states(workspace_entries, selected_tools, force_login=True) + states = _configure_shared_workspace_states( + workspace_entries, + selected_tools, + force_login=True, + enable_uc=enable_uc, + reset_uc=reset_uc, + ) state = states[0] save_state(state) @@ -649,6 +704,18 @@ def configure( "'low' prints terse single-line status instead.", ), ] = "normal", + enable_uc: Annotated[ + bool, + typer.Option( + "--enable-uc", + help="Discover models via UC `model-services` (`system.ai.`) and " + "surface curated `system.ai.*` MCP services. Equivalent to setting " + "UCODE_ENABLE_UC=1 for this configure run. The value is persisted so " + "subsequent `ucode ` launches stay on the same discovery path; " + "re-run `ucode configure` without the flag (and without " + "UCODE_ENABLE_UC=1 in the env) to turn UC discovery back off.", + ), + ] = False, ) -> None: """Configure workspace URL and AI Gateway.""" if ctx.invoked_subcommand is not None: @@ -659,6 +726,10 @@ def configure( set_dry_run(dry_run) set_verbosity(verbose) prompt_optional_updates = not skip_upgrade + flag_enable_uc: bool | None = True if enable_uc else None + # Explicit `ucode configure` is a clean slate: when the user omits both + # `--enable-uc` and `UCODE_ENABLE_UC`, persisted `uc_enabled=true` from + # a prior run is reset to false. try: install_databricks_cli() if agent is not None and agents is not None: @@ -673,31 +744,46 @@ def configure( prompt_optional_updates=prompt_optional_updates, ) if workspace_entries is None: - configure_workspace_command(tool) + configure_workspace_command(tool, enable_uc=flag_enable_uc, reset_uc=True) else: - configure_workspace_command(tool, workspaces=workspace_entries) + configure_workspace_command( + tool, + workspaces=workspace_entries, + enable_uc=flag_enable_uc, + reset_uc=True, + ) elif agents is not None: selected_tools = _parse_agents_option(agents) if workspace_entries is None: configure_workspace_command( selected_tools=selected_tools, prompt_optional_updates=prompt_optional_updates, + enable_uc=flag_enable_uc, + reset_uc=True, ) else: configure_workspace_command( selected_tools=selected_tools, workspaces=workspace_entries, prompt_optional_updates=prompt_optional_updates, + enable_uc=flag_enable_uc, + reset_uc=True, ) else: # Tool binaries are installed after the user picks which agents # they want, in configure_workspace_command. if workspace_entries is None: - configure_workspace_command(prompt_optional_updates=prompt_optional_updates) + configure_workspace_command( + prompt_optional_updates=prompt_optional_updates, + enable_uc=flag_enable_uc, + reset_uc=True, + ) else: configure_workspace_command( workspaces=workspace_entries, prompt_optional_updates=prompt_optional_updates, + enable_uc=flag_enable_uc, + reset_uc=True, ) if tracing: # The workspaces were just configured, so enable tracing for them diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 2d45feb..db26583 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -17,7 +17,7 @@ from typing import Literal, cast, overload from urllib import error as urllib_error from urllib import request as urllib_request -from urllib.parse import urlparse +from urllib.parse import urlencode, urlparse from databricks.sql.exc import ServerOperationError @@ -977,6 +977,228 @@ def build_auth_shell_command(workspace: str, profile: str | None = None) -> str: ) +def uc_enabled(default: bool = False) -> bool: + """True when the opt-in UC-securables discovery path is enabled. + + Three input precedences, callers handle the highest one first: + 1. ``ucode configure --enable-uc / --no-enable-uc`` (resolved by the + CLI before this function is called and passed in via ``default``, + since it overrides everything). + 2. ``UCODE_ENABLE_UC=1`` (or true/yes/on) env var. + 3. The value persisted in state (sticky, also passed via ``default``). + + Enabling UC discovery makes ucode: + - resolve models via UC `model-services` as `system.ai.` + instead of the per-family AI Gateway listings + - surface curated `system.ai.*` MCP services in `ucode configure mcp` + """ + raw = os.environ.get("UCODE_ENABLE_UC") + if raw is None or not raw.strip(): + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +# A model-service's `name` is `model-services/system.ai.`; the +# part after the prefix is exactly the model string agents send (no +# `databricks-` infix — that only appears on the inner destination name). +_MODEL_SERVICE_NAME_PREFIX = "model-services/" +# The metastore-scope listing returns services from EVERY schema (e.g. +# `main.user.foo`, `temp.*`, internal DLT schemas). We only want the +# Databricks-managed foundation models under `system.ai`. +_MODEL_SERVICE_REQUIRED_PREFIX = "system.ai." + + +def _model_service_id(service: dict) -> str | None: + """Extract the `system.ai.` id from one model-service entry. + + Returns None for services in any other schema, so user/internal model + services don't leak into the family buckets.""" + name = service.get("name") + if not isinstance(name, str): + return None + name = name.strip() + if name.startswith(_MODEL_SERVICE_NAME_PREFIX): + name = name[len(_MODEL_SERVICE_NAME_PREFIX) :] + if not name.startswith(_MODEL_SERVICE_REQUIRED_PREFIX): + return None + return name or None + + +# The model-services metastore listing is slow and flaky — large pages +# routinely 504 with `Timeout listing model services under metastore`. A small +# page is far more likely to come back, and each page gets a few retries before +# we give up. +_MODEL_SERVICES_PAGE_SIZE = 10 +_MODEL_SERVICES_PAGE_RETRIES = 4 + + +def _get_model_services_page( + url: str, token: str, *, retries: int = _MODEL_SERVICES_PAGE_RETRIES +) -> tuple[dict | list | None, str | None]: + """GET one model-services page, retrying on failure. + + The endpoint frequently 504s under load; a retry usually succeeds. Returns + the same (payload, reason) shape as ``_http_get_json`` — the last attempt's + result when all retries are exhausted.""" + payload: dict | list | None = None + reason: str | None = None + for attempt in range(retries): + payload, reason = _http_get_json(url, token, timeout=30) + if payload is not None: + return payload, None + _debug("model-services page", f"attempt {attempt + 1}/{retries} failed: {reason}") + return payload, reason + + +def list_model_services( + workspace: str, + token: str, + *, + page_size: int = _MODEL_SERVICES_PAGE_SIZE, + max_pages: int = 100, +) -> tuple[list[str], str | None]: + """List all `system.ai.*` model ids via the UC model-services API. + + Pages through ``/api/2.1/unity-catalog/model-services`` (metastore scope) + and returns the de-duplicated, sorted list of ``system.ai.`` + ids. Uses a small page size with per-page retries because the endpoint is + slow and frequently 504s. Returns (ids, reason); reason is None on success, + otherwise it describes why the list is empty (HTTP/network error or no + services). + """ + hostname = workspace_hostname(workspace) + ids: list[str] = [] + page_token: str | None = None + seen_tokens: set[str] = set() + last_reason: str | None = None + for _ in range(max_pages): + params: dict[str, str] = {"page_size": str(page_size)} + if page_token: + params["page_token"] = page_token + url = f"https://{hostname}/api/2.1/unity-catalog/model-services?{urlencode(params)}" + payload, reason = _get_model_services_page(url, token) + if payload is None: + # Surface the failure only if we have nothing yet; a mid-pagination + # blip still returns whatever we collected. + last_reason = reason + break + data = cast(dict, payload) if isinstance(payload, dict) else {} + for service in data.get("model_services", []): + if isinstance(service, dict): + model_id = _model_service_id(service) + if model_id: + ids.append(model_id) + page_token = data.get("next_page_token") or None + if not page_token: + last_reason = None + break + if page_token in seen_tokens: + break + seen_tokens.add(page_token) + + deduped = sorted(set(ids)) + if deduped: + return deduped, None + return [], last_reason or "model-services listing returned no models" + + +def discover_model_services( + workspace: str, token: str +) -> tuple[dict[str, str], list[str], list[str], str | None]: + """Discover models via UC model-services and bucket them by family name. + + Returns (claude_models, codex_models, gemini_models, reason): + + - ``claude_models`` maps ``opus``/``sonnet``/``haiku`` to the newest + matching ``system.ai.claude-*`` id (mirrors ``discover_claude_models``). + - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids. + - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. + + ``reason`` is None on success, else explains why nothing was found. Family + bucketing is by name substring because the model-services API does not + expose per-model API dialects. + """ + ids, reason = list_model_services(workspace, token) + if not ids: + return {}, [], [], reason + + claude_models: dict[str, str] = {} + for family in ("opus", "sonnet", "haiku"): + candidates = sorted( + [m for m in ids if f"claude-{family}-" in m], + reverse=True, + ) + if candidates: + claude_models[family] = candidates[0] + + codex_models = [m for m in ids if "gpt-" in m] + gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) + + if not (claude_models or codex_models or gemini_models): + sample = ", ".join(ids[:5]) + return ( + {}, + [], + [], + ( + "model-services returned model ids but none matched " + f"claude/gpt/gemini families (got: {sample})" + ), + ) + return claude_models, codex_models, gemini_models, None + + +# --- MCP services (parallel to model services) ----------------------------- + + +_MCP_SERVICE_NAME_PREFIX = "mcp-services/" +_MCP_SERVICE_REQUIRED_PREFIX = "system.ai." + + +def _mcp_service_full_name(service: dict) -> str | None: + """Extract `system.ai.` from one mcp-service entry, or None.""" + name = service.get("name") + if not isinstance(name, str): + return None + name = name.strip().removeprefix(_MCP_SERVICE_NAME_PREFIX) + if not name.startswith(_MCP_SERVICE_REQUIRED_PREFIX): + return None + status = ((service.get("config") or {}).get("connection") or {}).get("status") + if status is not None and status != "ACTIVE": + return None + return name + + +def list_mcp_services(workspace: str, token: str) -> tuple[list[str], str | None]: + """List Databricks-curated `system.ai.*` MCP services. + + The listing endpoint requires `?parent=schemas/system.ai`; without it the + request returns 499 with a truncated body (verified against e2-dogfood + 2026-06-11). Returns (full_names, reason). + """ + hostname = workspace_hostname(workspace) + url = ( + f"https://{hostname}/api/2.1/unity-catalog/mcp-services" + f"?{urlencode({'parent': 'schemas/system.ai'})}" + ) + payload, reason = _http_get_json(url, token, timeout=30) + if payload is None: + return [], reason + data = cast(dict, payload) if isinstance(payload, dict) else {} + names = [] + for service in data.get("mcp_services") or []: + if not isinstance(service, dict): + continue + full_name = _mcp_service_full_name(service) + if full_name: + names.append(full_name) + return sorted(set(names)), None if names else (reason or "no `system.ai.*` mcp services found") + + +def build_mcp_service_url(workspace: str, full_name: str) -> str: + return f"{workspace}/ai-gateway/mcp-services/{full_name}" + + def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], str | None]: """Discover Claude families on this workspace's AI Gateway. diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index ba94521..82c6bef 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -26,10 +26,14 @@ from ucode.agents import copilot, opencode from ucode.config_io import restore_file from ucode.databricks import ( + build_mcp_service_url, ensure_databricks_auth, + get_databricks_token, list_databricks_apps, list_databricks_connections, list_genie_spaces, + list_mcp_services, + uc_enabled, workspace_hostname, ) from ucode.state import load_full_state, load_state, save_state @@ -75,6 +79,7 @@ SQL_MCP_VALUE = "managed:sql" GENIE_SPACE_SELECTION_PREFIX = "genie-space:" APP_MCP_SELECTION_PREFIX = "app:" +MCP_SERVICE_SELECTION_PREFIX = "mcp-service:" MCP_ADD_PREFIX = "add:" MCP_CONNECTION_MARKERS = ( "is_mcp", @@ -352,6 +357,15 @@ def discover_external_mcp_connection_names(workspace: str, profile: str | None = return external_mcp_connection_names(list_databricks_connections(workspace, profile)) +def discover_mcp_service_names(workspace: str, profile: str | None = None) -> list[str]: + """Curated `system.ai.*` MCP services. Empty list if discovery fails so + callers can fall back to legacy connection discovery without surfacing + every error to the picker.""" + token = get_databricks_token(workspace, profile) + names, _reason = list_mcp_services(workspace, token) + return names + + def _normalize_workspace_title(text: str) -> str: """Collapse a Databricks workspace title to lowercase alphanumerics joined by single hyphens, trimmed at the edges. Output is safe to use as an MCP @@ -669,6 +683,7 @@ def build_mcp_picker_choices( available_genie_servers: list[dict], available_app_servers: list[dict], original_servers: list[dict], + available_mcp_service_names: list[str] | None = None, ) -> list[questionary.Choice | questionary.Separator]: original_by_name = _servers_by_name(original_servers) known_names = set(original_by_name) @@ -682,6 +697,17 @@ def build_mcp_picker_choices( choices.append(_add_choice(SQL_MCP_VALUE, "Databricks SQL")) displayed_names.add("databricks-sql") + for name in available_mcp_service_names or []: + # Picker shows the dotted UC name; state/agents store the dashed form + # (see resolver). Compare against the dashed form when checking what's + # already registered. + registered_as = name.replace(".", "-") + if registered_as in known_names: + choices.append(_server_choice(registered_as, True, name)) + else: + choices.append(_add_choice(f"{MCP_SERVICE_SELECTION_PREFIX}{name}", name)) + displayed_names.add(registered_as) + for name in available_external_names: if name in known_names: choices.append(_server_choice(name, True, name)) @@ -733,6 +759,7 @@ def prompt_for_mcp_server_choices( available_genie_servers: list[dict], available_app_servers: list[dict], original_servers: list[dict], + available_mcp_service_names: list[str] | None = None, ) -> list[str] | None: selection = _scrolling_checkbox( "MCP:", @@ -741,6 +768,7 @@ def prompt_for_mcp_server_choices( available_genie_servers, available_app_servers, original_servers, + available_mcp_service_names, ), style=_picker_style(), instruction="(space to toggle, enter to save, type to filter)", @@ -791,6 +819,14 @@ def _resolve_mcp_selection( raise RuntimeError("missing external connection name") return server_name, f"{workspace}/api/2.0/mcp/external/{server_name}" + if selection.startswith(MCP_SERVICE_SELECTION_PREFIX): + full_name = selection.removeprefix(MCP_SERVICE_SELECTION_PREFIX) + if not full_name: + raise RuntimeError("missing MCP service name") + # Agent CLIs (claude/codex/gemini) reject dots in registered names. + # URL keeps the UC `..` form; entry name uses dashes. + return full_name.replace(".", "-"), build_mcp_service_url(workspace, full_name) + if selection == SQL_MCP_VALUE: return "databricks-sql", f"{workspace}/api/2.0/mcp/sql" @@ -942,6 +978,18 @@ def configure_mcp_command() -> int: "Databricks apps", lambda: discover_app_mcp_servers(workspace, profile), ) + # Curated `system.ai.*` MCP services live behind a separate UC API and + # are gated on the same UC opt-in that enables model-services discovery + # (env: UCODE_ENABLE_UC, CLI: `ucode configure --enable-uc`, persisted + # in state on configure). + available_mcp_service_names = ( + _discover_mcp_source( + "MCP services", + lambda: discover_mcp_service_names(workspace, profile), + ) + if uc_enabled(default=bool(state.get("uc_enabled"))) + else [] + ) original_mcp_servers: list[dict] = list(state.get("mcp_servers") or []) original_by_name = _servers_by_name(original_mcp_servers) @@ -950,6 +998,7 @@ def configure_mcp_command() -> int: available_genie_mcp_servers, available_app_mcp_servers, original_mcp_servers, + available_mcp_service_names, ) if selections is None: return 0 diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index ea33c63..9888efd 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -41,6 +41,14 @@ def test_does_not_duplicate_1m_suffix(self): overlay, _ = claude.render_overlay(WS, "databricks-claude-opus-4-7[1m]") assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-opus-4-7[1m]" + def test_adds_1m_suffix_for_model_services_name(self): + overlay, _ = claude.render_overlay(WS, "system.ai.claude-opus-4-8") + assert overlay["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + + def test_no_1m_suffix_for_model_services_haiku(self): + overlay, _ = claude.render_overlay(WS, "system.ai.claude-haiku-4-6") + assert overlay["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-haiku-4-6" + def test_sets_anthropic_base_url(self): overlay, _ = claude.render_overlay(WS, "s4") assert overlay["env"]["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic" diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index b84b667..f8d6baf 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -337,6 +337,17 @@ def test_openai_model_id_maps_databricks_naming(self): def test_codex_model_id_preserves_openai_incompatible_models(self): assert codex._codex_model_id("databricks-gpt-5-2-codex") == "databricks-gpt-5-2-codex" assert codex._codex_model_id("databricks-gpt-5-4-nano") == "databricks-gpt-5-4-nano" + + def test_codex_model_id_passes_model_services_id_verbatim(self): + # UC model-services ids route by name, so they must not be rewritten + # to the OpenAI id form. + assert codex._codex_model_id("system.ai.gpt-5") == "system.ai.gpt-5" + assert codex._codex_model_id("system.ai.gpt-5-2-codex") == "system.ai.gpt-5-2-codex" + + def test_default_model_selects_model_services_gpt(self): + models = ["system.ai.gpt-5", "system.ai.gpt-5-5", "system.ai.claude-opus-4-8"] + + assert codex.default_model({"codex_models": models}) == "system.ai.gpt-5-5" assert codex._codex_model_id("databricks-gpt-5-5") == "gpt-5.5" diff --git a/tests/test_cli.py b/tests/test_cli.py index 9809d1a..7e2a98c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -80,8 +80,11 @@ def test_configure_help_lists_agents_flag(self): result = runner.invoke(app, ["configure", "--help"]) assert result.exit_code == 0 output = _strip_ansi(result.output) + # Typer wraps long help text across lines and pads with box-drawing + # characters; collapse whitespace + box chars before substring-matching. + flat = re.sub(r"[│╭╮╯╰─\s]+", " ", output) assert "--agents" in output - assert "comma-separated list of agents" in output + assert "comma-separated list of agents" in flat assert "--workspaces" in output @@ -380,7 +383,9 @@ def test_no_flag_calls_configure_all(self): ): result = runner.invoke(app, ["configure"]) assert result.exit_code == 0, result.output - mock_cfg.assert_called_once_with(prompt_optional_updates=True) + mock_cfg.assert_called_once_with( + prompt_optional_updates=True, enable_uc=None, reset_uc=True + ) def test_agents_flag_calls_configure_with_tools(self): with ( @@ -392,7 +397,10 @@ def test_agents_flag_calls_configure_with_tools(self): assert result.exit_code == 0, result.output mock_install.assert_not_called() mock_cfg.assert_called_once_with( - selected_tools=["claude", "codex"], prompt_optional_updates=True + selected_tools=["claude", "codex"], + prompt_optional_updates=True, + enable_uc=None, + reset_uc=True, ) def test_agents_flag_normalizes_aliases_and_dedupes(self): @@ -404,7 +412,10 @@ def test_agents_flag_normalizes_aliases_and_dedupes(self): result = runner.invoke(app, ["configure", "--agents", " claude-code, codex,claude "]) assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with( - selected_tools=["claude", "codex"], prompt_optional_updates=True + selected_tools=["claude", "codex"], + prompt_optional_updates=True, + enable_uc=None, + reset_uc=True, ) def test_workspaces_flag_calls_configure_with_workspaces(self): @@ -428,6 +439,8 @@ def test_workspaces_flag_calls_configure_with_workspaces(self): ("https://second.databricks.com", None), ], prompt_optional_updates=True, + enable_uc=None, + reset_uc=True, ) def test_agents_and_workspaces_flags_call_configure_with_both(self): @@ -445,6 +458,8 @@ def test_agents_and_workspaces_flags_call_configure_with_both(self): selected_tools=["claude", "codex"], workspaces=[("https://first.com", None)], prompt_optional_updates=True, + enable_uc=None, + reset_uc=True, ) def test_agent_and_workspaces_flags_call_configure_with_both(self): @@ -461,7 +476,9 @@ def test_agent_and_workspaces_flags_call_configure_with_both(self): mock_install.assert_called_once_with( "claude", strict=True, update_existing=True, prompt_optional_updates=True ) - mock_cfg.assert_called_once_with("claude", workspaces=[("https://first.com", None)]) + mock_cfg.assert_called_once_with( + "claude", workspaces=[("https://first.com", None)], enable_uc=None, reset_uc=True + ) def test_agent_flag_calls_configure_with_tool(self): with ( @@ -474,7 +491,7 @@ def test_agent_flag_calls_configure_with_tool(self): mock_install.assert_called_once_with( "claude", strict=True, update_existing=True, prompt_optional_updates=True ) - mock_cfg.assert_called_once_with("claude") + mock_cfg.assert_called_once_with("claude", enable_uc=None, reset_uc=True) def test_skip_upgrade_flag_disables_optional_update_prompt(self): with ( @@ -484,7 +501,9 @@ def test_skip_upgrade_flag_disables_optional_update_prompt(self): ): result = runner.invoke(app, ["configure", "--skip-upgrade"]) assert result.exit_code == 0, result.output - mock_cfg.assert_called_once_with(prompt_optional_updates=False) + mock_cfg.assert_called_once_with( + prompt_optional_updates=False, enable_uc=None, reset_uc=True + ) def test_skip_upgrade_flag_with_agent_skips_optional_update(self): with ( @@ -507,7 +526,10 @@ def test_skip_upgrade_flag_with_agents_forwards_to_configure(self): result = runner.invoke(app, ["configure", "--agents", "claude,codex", "--skip-upgrade"]) assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with( - selected_tools=["claude", "codex"], prompt_optional_updates=False + selected_tools=["claude", "codex"], + prompt_optional_updates=False, + enable_uc=None, + reset_uc=True, ) def test_agent_flag_normalizes_alias(self): @@ -518,7 +540,7 @@ def test_agent_flag_normalizes_alias(self): ): result = runner.invoke(app, ["configure", "--agent", "claude-code"]) assert result.exit_code == 0, result.output - mock_cfg.assert_called_once_with("claude") + mock_cfg.assert_called_once_with("claude", enable_uc=None, reset_uc=True) def test_upgrade_runs_uv_tool_install(self): with patch("subprocess.run") as mock_run: @@ -658,7 +680,14 @@ def test_multiple_workspaces_configure_all_and_use_first(self, monkeypatch): } configured_shared: list[tuple[str, str | None, tuple[str, ...] | None, bool]] = [] - def fake_configure_shared_state(workspace, profile=None, tools=None, force_login=False): + def fake_configure_shared_state( + workspace, + profile=None, + tools=None, + force_login=False, + enable_uc=None, + reset_uc=False, + ): configured_shared.append( (workspace, profile, tuple(tools) if tools is not None else None, force_login) ) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index d3feeba..3a850e7 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -132,6 +132,266 @@ def test_selects_opus_4_8_when_advertised(self, monkeypatch): assert models["opus"] == "databricks-claude-opus-4-8" +def _model_service(model_id: str) -> dict: + """A model-services entry whose `name` strips to `model_id`.""" + return {"name": f"model-services/{model_id}"} + + +class TestUcEnabled: + def test_off_by_default(self, monkeypatch): + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + assert db_mod.uc_enabled() is False + + def test_truthy_values_enable(self, monkeypatch): + for value in ("1", "true", "TRUE", "yes", "on"): + monkeypatch.setenv("UCODE_ENABLE_UC", value) + assert db_mod.uc_enabled() is True + + def test_falsey_values_disable(self, monkeypatch): + # A non-empty, non-truthy value explicitly disables — even over a + # persisted default of True. + for value in ("0", "false", "no"): + monkeypatch.setenv("UCODE_ENABLE_UC", value) + assert db_mod.uc_enabled(default=True) is False + + def test_unset_falls_back_to_default(self, monkeypatch): + # Sticky behavior: when the env var is unset (or blank), the persisted + # default decides. + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + assert db_mod.uc_enabled(default=True) is True + assert db_mod.uc_enabled(default=False) is False + monkeypatch.setenv("UCODE_ENABLE_UC", "") + assert db_mod.uc_enabled(default=True) is True + + def test_env_var_overrides_default(self, monkeypatch): + monkeypatch.setenv("UCODE_ENABLE_UC", "1") + assert db_mod.uc_enabled(default=False) is True + + +class TestDiscoverModelServices: + def test_buckets_families_by_name(self, monkeypatch): + payload = { + "model_services": [ + _model_service("system.ai.claude-opus-4-7"), + _model_service("system.ai.claude-opus-4-8"), + _model_service("system.ai.claude-sonnet-4-6"), + _model_service("system.ai.gpt-5"), + _model_service("system.ai.gemini-2-5-flash"), + _model_service("system.ai.gemini-3-5-flash"), + _model_service("system.ai.llama-4-maverick"), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + # Newest opus wins; sonnet bucketed; haiku absent. + assert claude == { + "opus": "system.ai.claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-4-6", + } + assert codex == ["system.ai.gpt-5"] + # Gemini ordered newest-first via the shared sort key. + assert gemini[0] == "system.ai.gemini-3-5-flash" + # llama is not bucketed into any of the three families. + assert "system.ai.llama-4-maverick" not in codex + gemini + + def test_paginates_via_next_page_token(self, monkeypatch): + pages = { + None: { + "model_services": [_model_service("system.ai.gpt-5")], + "next_page_token": "tok2", + }, + "tok2": { + "model_services": [_model_service("system.ai.claude-opus-4-8")], + }, + } + + def fake_get(url, token, timeout=10): + token_param = None + if "page_token=" in url: + token_param = url.split("page_token=")[1].split("&")[0] + return pages[token_param], None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + claude, codex, _, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert codex == ["system.ai.gpt-5"] + assert claude == {"opus": "system.ai.claude-opus-4-8"} + + def test_http_failure_returns_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (None, "HTTP 500 Server Error") + ) + + claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + + assert (claude, codex, gemini) == ({}, [], []) + assert reason == "HTTP 500 Server Error" + + def test_no_matching_families_reports_sample(self, monkeypatch): + payload = {"model_services": [_model_service("system.ai.llama-4-maverick")]} + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + + assert (claude, codex, gemini) == ({}, [], []) + assert reason is not None and "llama-4-maverick" in reason + + def test_ignores_non_system_ai_schemas(self, monkeypatch): + # The metastore listing returns services from every schema; only + # system.ai.* foundation models should be picked up. + payload = { + "model_services": [ + _model_service("system.ai.gpt-5"), + _model_service("main.svenwb.gpt-5-5"), + _model_service("temp.erni.claude-opus-4-8"), + _model_service("dnasi_agent_cuj.default.dnasi-gpt55-test"), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert codex == ["system.ai.gpt-5"] + assert claude == {} # temp.erni.claude-* must not be bucketed + assert gemini == [] + + def test_retries_page_before_giving_up(self, monkeypatch): + payload = {"model_services": [_model_service("system.ai.gpt-5")]} + calls = {"n": 0} + + def flaky_get(url, token, timeout=10): + calls["n"] += 1 + if calls["n"] < 3: + return None, "HTTP 504 Gateway Timeout" + return payload, None + + monkeypatch.setattr(db_mod, "_http_get_json", flaky_get) + + ids, reason = db_mod.list_model_services(WS, "token") + + assert reason is None + assert ids == ["system.ai.gpt-5"] + assert calls["n"] == 3 # two failures, third succeeds + + +class TestListMcpServices: + def test_accepts_entries_without_connection_status(self, monkeypatch): + payload = { + "mcp_services": [ + { + "name": "mcp-services/system.ai.github", + "config": {"usage_tracking": {"enabled": True}, "tracing": {"enabled": True}}, + }, + { + "name": "mcp-services/system.ai.atlassian", + "config": {}, + }, + { + "name": "mcp-services/system.ai.slack", + }, + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + + names, reason = db_mod.list_mcp_services(WS, "token") + + assert reason is None + assert names == ["system.ai.atlassian", "system.ai.github", "system.ai.slack"] + + def test_accepts_legacy_active_status(self, monkeypatch): + payload = { + "mcp_services": [ + { + "name": "mcp-services/system.ai.github", + "config": {"connection": {"status": "ACTIVE"}}, + }, + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + + names, reason = db_mod.list_mcp_services(WS, "token") + + assert reason is None + assert names == ["system.ai.github"] + + def test_rejects_explicit_non_active_status(self, monkeypatch): + # If the field is present and non-ACTIVE, drop the entry — the + # backing connection is broken and the proxy will fail. + payload = { + "mcp_services": [ + { + "name": "mcp-services/system.ai.github", + "config": {"connection": {"status": "ACTIVE"}}, + }, + { + "name": "mcp-services/system.ai.broken", + "config": {"connection": {"status": "FAILED"}}, + }, + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + + names, _reason = db_mod.list_mcp_services(WS, "token") + + assert names == ["system.ai.github"] + + def test_ignores_non_system_ai_entries(self, monkeypatch): + payload = { + "mcp_services": [ + {"name": "mcp-services/system.ai.github"}, + {"name": "mcp-services/main.svenwb.github_mcp"}, + {"name": "mcp-services/temp.erni.github_mcp"}, + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + + names, _reason = db_mod.list_mcp_services(WS, "token") + + assert names == ["system.ai.github"] + + def test_http_failure_propagates_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, timeout=30: (None, "HTTP 500 Server Error"), + ) + + names, reason = db_mod.list_mcp_services(WS, "token") + + assert names == [] + assert reason == "HTTP 500 Server Error" + + def test_empty_payload_reports_no_results(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: ({"mcp_services": []}, None) + ) + + names, reason = db_mod.list_mcp_services(WS, "token") + + assert names == [] + assert reason and "no `system.ai.*`" in reason + + def _foundation_models_payload(names): return { "endpoints": [ diff --git a/tests/test_e2e_uc.py b/tests/test_e2e_uc.py new file mode 100644 index 0000000..d0f994d --- /dev/null +++ b/tests/test_e2e_uc.py @@ -0,0 +1,216 @@ +"""End-to-end tests for the `--enable-uc` / `UCODE_ENABLE_UC` opt-in. + +Verifies UC-securables discovery (model-services + MCP services) and the +flag-precedence ladder (CLI flag > env var > persisted state) against a +live Databricks workspace. + +Run with: + UCODE_TEST_WORKSPACE=https://your-workspace.databricks.com \ + uv run pytest tests/test_e2e_uc.py -v +""" + +from __future__ import annotations + +import pytest + +from ucode.cli import configure_shared_state +from ucode.databricks import ( + discover_model_services, + list_mcp_services, + uc_enabled, +) +from ucode.state import load_state + + +def _has_uc_models(workspace: str, token: str) -> bool: + claude, codex, gemini, _reason = discover_model_services(workspace, token) + return bool(claude or codex or gemini) + + +def _all_resolved_model_ids(state: dict) -> list[str]: + ids: list[str] = list((state.get("claude_models") or {}).values()) + ids += state.get("codex_models") or [] + ids += state.get("gemini_models") or [] + return ids + + +# --------------------------------------------------------------------------- +# UC discovery primitives — verify the endpoints return only `system.ai.*` +# entries (the per-family/connection filters drop everything else). +# --------------------------------------------------------------------------- + + +class TestDiscoverModelServicesE2E: + def test_returns_only_system_ai_models(self, e2e_workspace, e2e_token): + claude, codex, gemini, reason = discover_model_services(e2e_workspace, e2e_token) + if not (claude or codex or gemini): + pytest.skip(f"No system.ai.* model services on workspace: {reason}") + non_system = sorted( + { + m + for m in _all_resolved_model_ids( + {"claude_models": claude, "codex_models": codex, "gemini_models": gemini} + ) + if not m.startswith("system.ai.") + } + ) + assert not non_system, f"Non-system.ai entries leaked through: {non_system[:5]}" + + +class TestListMcpServicesE2E: + def test_returns_only_system_ai_mcp_services(self, e2e_workspace, e2e_token): + names, reason = list_mcp_services(e2e_workspace, e2e_token) + if not names: + pytest.skip(f"No system.ai.* MCP services on workspace: {reason}") + non_system = sorted({n for n in names if not n.startswith("system.ai.")}) + assert not non_system, f"Non-system.ai entries leaked through: {non_system[:5]}" + + +# --------------------------------------------------------------------------- +# `uc_enabled` precedence — env var alone (no `default` arg). +# --------------------------------------------------------------------------- + + +class TestUcEnabledEnvE2E: + def test_env_on_overrides_default_off(self, monkeypatch): + monkeypatch.setenv("UCODE_ENABLE_UC", "1") + assert uc_enabled(default=False) is True + + def test_env_off_overrides_default_on(self, monkeypatch): + monkeypatch.setenv("UCODE_ENABLE_UC", "0") + assert uc_enabled(default=True) is False + + +# --------------------------------------------------------------------------- +# `configure_shared_state` end-to-end: flag resolution + persistence to +# `state["uc_enabled"]` + which discovery path runs (UC vs legacy). +# --------------------------------------------------------------------------- + + +class TestConfigureSharedStateEnableUcE2E: + """Resolution ladder: CLI flag > env var > persisted state. Each test + runs the full configure path against the live workspace and asserts on + the resolved flag and the actual model namespaces written to state.""" + + def test_explicit_true_persists_and_discovers_system_ai( + self, monkeypatch, e2e_workspace, e2e_token + ): + if not _has_uc_models(e2e_workspace, e2e_token): + pytest.skip("Workspace has no system.ai.* model services.") + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + + state = configure_shared_state(e2e_workspace, force_login=False, enable_uc=True) + assert state["uc_enabled"] is True + assert load_state()["uc_enabled"] is True + ids = _all_resolved_model_ids(state) + assert any(m.startswith("system.ai.") for m in ids), ( + f"Expected at least one system.ai.* model id, got: {ids[:5]}" + ) + + def test_env_off_overrides_persisted_true(self, monkeypatch, e2e_workspace): + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + # Pre-seed the target workspace's state with uc_enabled=True. + configure_shared_state(e2e_workspace, force_login=False, enable_uc=True) + assert load_state()["uc_enabled"] is True + + # Now opt back out via env var, no CLI flag. + monkeypatch.setenv("UCODE_ENABLE_UC", "0") + state = configure_shared_state(e2e_workspace, force_login=False, enable_uc=None) + assert state["uc_enabled"] is False + assert load_state()["uc_enabled"] is False + ids = _all_resolved_model_ids(state) + assert all(not m.startswith("system.ai.") for m in ids), ( + f"Legacy discovery leaked system.ai entries: " + f"{[m for m in ids if m.startswith('system.ai.')][:5]}" + ) + + def test_env_resolves_when_cli_flag_omitted(self, monkeypatch, e2e_workspace, e2e_token): + if not _has_uc_models(e2e_workspace, e2e_token): + pytest.skip("Workspace has no system.ai.* model services.") + monkeypatch.setenv("UCODE_ENABLE_UC", "1") + + state = configure_shared_state(e2e_workspace, force_login=False, enable_uc=None) + assert state["uc_enabled"] is True + ids = _all_resolved_model_ids(state) + assert any(m.startswith("system.ai.") for m in ids) + + def test_plain_configure_resets_persisted_uc_enabled( + self, monkeypatch, e2e_workspace, e2e_token + ): + """`ucode configure` without `--enable-uc` and without + UCODE_ENABLE_UC is a clean slate: a previously-persisted + `uc_enabled=True` is flipped back to False, and discovery returns + to legacy `databricks-*` ids.""" + if not _has_uc_models(e2e_workspace, e2e_token): + pytest.skip("Workspace has no system.ai.* model services.") + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + + # First configure persists `uc_enabled=True`. + configure_shared_state(e2e_workspace, force_login=False, enable_uc=True) + assert load_state()["uc_enabled"] is True + + # Second configure with reset_uc=True (the explicit `ucode configure` + # path) clears the flag. + state = configure_shared_state( + e2e_workspace, force_login=False, enable_uc=None, reset_uc=True + ) + assert state["uc_enabled"] is False + assert load_state()["uc_enabled"] is False + ids = _all_resolved_model_ids(state) + assert all(not m.startswith("system.ai.") for m in ids), ( + f"Reset run still pulled UC ids: {[m for m in ids if m.startswith('system.ai.')][:5]}" + ) + + def test_launch_path_preserves_persisted_uc_enabled( + self, monkeypatch, e2e_workspace, e2e_token + ): + """Launch-time refetches (`ucode `) call configure_shared_state + without `reset_uc`. They must keep an existing persisted True so a + Claude/Codex/Gemini launch right after `--enable-uc` doesn't silently + drop UC discovery.""" + if not _has_uc_models(e2e_workspace, e2e_token): + pytest.skip("Workspace has no system.ai.* model services.") + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + + # User runs `ucode configure --enable-uc`. + configure_shared_state(e2e_workspace, force_login=False, enable_uc=True) + assert load_state()["uc_enabled"] is True + + # User then runs `ucode claude` — same call shape as + # _launch_tool's refetch (no reset_uc, no enable_uc, no env). + state = configure_shared_state(e2e_workspace, force_login=False, enable_uc=None) + assert state["uc_enabled"] is True + assert any(m.startswith("system.ai.") for m in _all_resolved_model_ids(state)) + + def test_default_off_when_no_flag_no_env_no_state(self, monkeypatch, e2e_workspace): + # Fresh state (autouse fixture redirects STATE_PATH per test) plus + # no env var means the flag falls through to its default of False. + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + + state = configure_shared_state(e2e_workspace, force_login=False, enable_uc=None) + assert state["uc_enabled"] is False + ids = _all_resolved_model_ids(state) + assert all(not m.startswith("system.ai.") for m in ids) + + def test_other_workspace_flag_does_not_leak_into_target(self, monkeypatch, e2e_workspace): + """Regression: enabling UC on workspace A must not silently turn it + on for a fresh `ucode configure` on workspace B. The default has to + come from B's own persisted state, not A's (which is whatever + happens to be `current_workspace`).""" + from ucode.state import save_state + + monkeypatch.delenv("UCODE_ENABLE_UC", raising=False) + + other_ws = "https://other-workspace.cloud.databricks.com" + save_state({"workspace": other_ws, "uc_enabled": True}) + + state = configure_shared_state(e2e_workspace, force_login=False, enable_uc=None) + assert state["uc_enabled"] is False, ( + "Cross-workspace leak: another workspace's uc_enabled bled into " + "the target workspace's default." + ) + ids = _all_resolved_model_ids(state) + assert all(not m.startswith("system.ai.") for m in ids), ( + f"Discovery used UC despite per-workspace default being False: " + f"{[m for m in ids if m.startswith('system.ai.')][:5]}" + )