From 18fb1e5114d9666d8d586345a8ec4fd3156d0960 Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Sun, 14 Jun 2026 17:52:27 +0000 Subject: [PATCH] Make UC model discovery the default; remove --enable-uc ucode configure only found the opus Claude model on workspaces where the legacy /ai-gateway/anthropic/v1/models listing dropped sonnet/haiku once system.ai endpoints existed. UC model-services discovery found all three but was hidden behind --enable-uc (and had no fallback). Now discovery is UC-first and best-effort: try model-services first, fall back to the legacy per-family listing per family when UC returns nothing. The --enable-uc flag, UCODE_ENABLE_UC env var, and uc_enabled helper are removed, and curated system.ai.* MCP services are surfaced by default. The model-services endpoint still 499s without a bounded page_size, so pagination (page_size=100, follow next_page_token, light per-page retry) is retained. Co-authored-by: Isaac --- src/ucode/cli.py | 90 ++++------------------ src/ucode/databricks.py | 47 ++++-------- src/ucode/mcp.py | 18 ++--- tests/test_cli.py | 83 +++++++++++++-------- tests/test_databricks.py | 50 +++++-------- tests/test_e2e_uc.py | 157 ++++++--------------------------------- tests/test_mcp.py | 4 + 7 files changed, 128 insertions(+), 321 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 76aa47a..0d60916 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -46,7 +46,6 @@ normalize_workspace_url, resolve_pat_token, run_databricks_login, - uc_enabled, ) from ucode.mcp import ( MCP_CLIENTS, @@ -195,8 +194,6 @@ 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, use_pat: bool | None = None, ) -> dict: """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state. @@ -211,28 +208,10 @@ 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) 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"))) if use_pat is None: use_pat = bool(prior_state.get("use_pat")) and previous_workspace == workspace fetch_all = tools is None @@ -278,25 +257,23 @@ def configure_shared_state( 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) + # UC-first, best-effort: one UC model-services call yields all families as + # `system.ai.` ids, bucketed by name. If a family comes back + # empty (workspace without UC model-services, or the listing failed), fall + # back to the per-family AI Gateway listing for that family only. + with spinner("Fetching available models..."): + ms_claude, ms_codex, ms_gemini, ms_reason = discover_model_services(workspace, token) if want_claude: claude_models, claude_reason = ms_claude, ms_reason + if not claude_models: + claude_models, claude_reason = discover_claude_models(workspace, token) if want_gemini: gemini_models, gemini_reason = ms_gemini, ms_reason + if not gemini_models: + gemini_models, gemini_reason = discover_gemini_models(workspace, token) if want_codex: 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: + if not codex_models: codex_models, codex_reason = discover_codex_models(workspace, token) opencode_models: dict[str, list[str]] = {} if claude_models: @@ -311,9 +288,8 @@ 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 + # UC discovery is now always-on; drop any flag persisted by older versions. + state.pop("uc_enabled", None) # Persist the auth mode so launches rebuild the same (PAT-based) agent # auth command; an explicit re-configure without --use-pat clears it. if use_pat: @@ -349,8 +325,6 @@ def _configure_shared_workspace_states( tools: list[str] | None, *, force_login: bool, - enable_uc: bool | None = None, - reset_uc: bool = False, use_pat: bool = False, ) -> list[dict]: if not workspaces: @@ -363,8 +337,6 @@ def _configure_shared_workspace_states( profile=profile, tools=tools, force_login=force_login, - enable_uc=enable_uc, - reset_uc=reset_uc, use_pat=use_pat, ) ) @@ -377,8 +349,6 @@ 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, use_pat: bool = False, skip_validate: bool = False, ) -> int: @@ -392,8 +362,6 @@ def configure_workspace_command( workspace_entries, [tool], force_login=True, - enable_uc=enable_uc, - reset_uc=reset_uc, use_pat=use_pat, ) state = states[0] @@ -429,8 +397,6 @@ def configure_workspace_command( workspace_entries, selected_tools, force_login=True, - enable_uc=enable_uc, - reset_uc=reset_uc, use_pat=use_pat, ) state = states[0] @@ -827,18 +793,6 @@ 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: @@ -849,10 +803,6 @@ 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: @@ -883,15 +833,11 @@ def configure( prompt_optional_updates=prompt_optional_updates, ) if workspace_entries is None: - configure_workspace_command( - tool, enable_uc=flag_enable_uc, reset_uc=True, **skip_kwargs - ) + configure_workspace_command(tool, **skip_kwargs) else: configure_workspace_command( tool, workspaces=workspace_entries, - enable_uc=flag_enable_uc, - reset_uc=True, **skip_kwargs, ) elif agents is not None: @@ -900,8 +846,6 @@ def configure( configure_workspace_command( selected_tools=selected_tools, prompt_optional_updates=prompt_optional_updates, - enable_uc=flag_enable_uc, - reset_uc=True, **skip_kwargs, ) else: @@ -909,8 +853,6 @@ def configure( selected_tools=selected_tools, workspaces=workspace_entries, prompt_optional_updates=prompt_optional_updates, - enable_uc=flag_enable_uc, - reset_uc=True, **skip_kwargs, ) else: @@ -919,16 +861,12 @@ def configure( if workspace_entries is None: configure_workspace_command( prompt_optional_updates=prompt_optional_updates, - enable_uc=flag_enable_uc, - reset_uc=True, **skip_kwargs, ) else: configure_workspace_command( workspaces=workspace_entries, prompt_optional_updates=prompt_optional_updates, - enable_uc=flag_enable_uc, - reset_uc=True, **skip_kwargs, ) if tracing: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b35c07b..574d906 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1050,27 +1050,6 @@ def build_auth_shell_command( ) -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). @@ -1097,11 +1076,12 @@ def _model_service_id(service: dict) -> str | 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 +# The model-services metastore listing REQUIRES a bounded `page_size`: +# unparameterized or large-page requests (verified against +# eng-ml-agent-platform.staging 2026-06-14) return `HTTP 499` with an empty +# body, while pages of 10–100 come back reliably. A page can still 499 +# intermittently under load, so each gets a few retries before we give up. +_MODEL_SERVICES_PAGE_SIZE = 100 _MODEL_SERVICES_PAGE_RETRIES = 4 @@ -1110,9 +1090,9 @@ def _get_model_services_page( ) -> 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.""" + The endpoint intermittently 499/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): @@ -1133,11 +1113,10 @@ def list_model_services( """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). + with a bounded ``page_size`` (the endpoint 499s without one) and returns the + de-duplicated, sorted list of ``system.ai.`` ids. 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] = [] diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 18fe57d..e2912a2 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -35,7 +35,6 @@ list_databricks_connections, list_genie_spaces, list_mcp_services, - uc_enabled, workspace_hostname, ) from ucode.state import load_full_state, load_state, save_state @@ -990,17 +989,12 @@ 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 [] + # Curated `system.ai.*` MCP services live behind a separate UC API. Like + # the other sources this is best-effort — `_discover_mcp_source` swallows + # failures and returns [] so workspaces without them just see nothing extra. + available_mcp_service_names = _discover_mcp_source( + "MCP services", + lambda: discover_mcp_service_names(workspace, profile), ) original_mcp_servers: list[dict] = list(state.get("mcp_servers") or []) diff --git a/tests/test_cli.py b/tests/test_cli.py index 9a4fe5d..cf156bc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -383,9 +383,7 @@ 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, enable_uc=None, reset_uc=True - ) + mock_cfg.assert_called_once_with(prompt_optional_updates=True) def test_agents_flag_calls_configure_with_tools(self): with ( @@ -399,8 +397,6 @@ def test_agents_flag_calls_configure_with_tools(self): mock_cfg.assert_called_once_with( selected_tools=["claude", "codex"], prompt_optional_updates=True, - enable_uc=None, - reset_uc=True, ) def test_agents_flag_normalizes_aliases_and_dedupes(self): @@ -414,8 +410,6 @@ def test_agents_flag_normalizes_aliases_and_dedupes(self): mock_cfg.assert_called_once_with( selected_tools=["claude", "codex"], prompt_optional_updates=True, - enable_uc=None, - reset_uc=True, ) def test_workspaces_flag_calls_configure_with_workspaces(self): @@ -439,8 +433,6 @@ 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): @@ -458,8 +450,6 @@ 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): @@ -476,9 +466,7 @@ 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)], enable_uc=None, reset_uc=True - ) + mock_cfg.assert_called_once_with("claude", workspaces=[("https://first.com", None)]) def test_agent_flag_calls_configure_with_tool(self): with ( @@ -491,7 +479,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", enable_uc=None, reset_uc=True) + mock_cfg.assert_called_once_with("claude") def test_skip_upgrade_flag_disables_optional_update_prompt(self): with ( @@ -501,9 +489,7 @@ 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, enable_uc=None, reset_uc=True - ) + mock_cfg.assert_called_once_with(prompt_optional_updates=False) def test_skip_upgrade_flag_with_agent_skips_optional_update(self): with ( @@ -528,8 +514,6 @@ def test_skip_upgrade_flag_with_agents_forwards_to_configure(self): mock_cfg.assert_called_once_with( selected_tools=["claude", "codex"], prompt_optional_updates=False, - enable_uc=None, - reset_uc=True, ) def test_agent_flag_normalizes_alias(self): @@ -540,7 +524,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", enable_uc=None, reset_uc=True) + mock_cfg.assert_called_once_with("claude") def test_upgrade_runs_uv_tool_install(self): with patch("subprocess.run") as mock_run: @@ -685,8 +669,6 @@ def fake_configure_shared_state( profile=None, tools=None, force_login=False, - enable_uc=None, - reset_uc=False, use_pat=False, ): configured_shared.append( @@ -788,8 +770,6 @@ def test_profiles_flag_resolves_workspaces(self): mock_cfg.assert_called_once_with( workspaces=[("https://first.databricks.com", "DEFAULT")], prompt_optional_updates=True, - enable_uc=None, - reset_uc=True, ) def test_profiles_flag_with_agents(self): @@ -807,8 +787,6 @@ def test_profiles_flag_with_agents(self): selected_tools=["claude", "codex"], workspaces=[("https://first.databricks.com", "DEFAULT")], prompt_optional_updates=True, - enable_uc=None, - reset_uc=True, ) def test_profiles_flag_with_agent(self): @@ -823,8 +801,6 @@ def test_profiles_flag_with_agent(self): mock_cfg.assert_called_once_with( "claude", workspaces=[("https://first.databricks.com", "DEFAULT")], - enable_uc=None, - reset_uc=True, ) def test_use_pat_and_skip_validate_are_forwarded(self): @@ -851,8 +827,6 @@ def test_use_pat_and_skip_validate_are_forwarded(self): selected_tools=["claude", "codex"], workspaces=[("https://first.databricks.com", "DEFAULT")], prompt_optional_updates=True, - enable_uc=None, - reset_uc=True, use_pat=True, skip_validate=True, ) @@ -926,6 +900,7 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], None)) monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) @@ -990,6 +965,51 @@ def test_reconfigure_without_flag_clears_use_pat(self, monkeypatch): assert logins == [(self.WS, "DEFAULT")] assert "use_pat" not in state + def test_uc_models_used_without_legacy_fallback(self, monkeypatch): + # When model-services returns models, they're used and the legacy + # per-family discovery is never consulted. + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({"opus": "system.ai.claude-opus-4-8"}, ["system.ai.gpt-5"], [], None), + ) + legacy_called: list[str] = [] + monkeypatch.setattr( + cli_mod, + "discover_claude_models", + lambda w, t: legacy_called.append("claude") or ({}, None), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["claude_models"] == {"opus": "system.ai.claude-opus-4-8"} + assert state["codex_models"] == ["system.ai.gpt-5"] + assert legacy_called == [] + assert "uc_enabled" not in state + + def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): + # No UC model-services: each family falls back to the legacy listing. + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, "discover_model_services", lambda w, t: ({}, [], [], "no model services") + ) + monkeypatch.setattr( + cli_mod, + "discover_claude_models", + lambda w, t: ( + {"opus": "databricks-claude-opus-4-8", "sonnet": "databricks-claude-sonnet-4-6"}, + None, + ), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["claude_models"] == { + "opus": "databricks-claude-opus-4-8", + "sonnet": "databricks-claude-sonnet-4-6", + } + class TestConfigureSkipValidate: def test_skip_validate_skips_agent_validation(self, monkeypatch): @@ -1050,6 +1070,7 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], None)) monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 5be51f0..65f05d8 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -137,37 +137,6 @@ def _model_service(model_id: str) -> dict: 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 = { @@ -267,6 +236,23 @@ def test_ignores_non_system_ai_schemas(self, monkeypatch): assert claude == {} # temp.erni.claude-* must not be bucketed assert gemini == [] + def test_requests_bounded_page_size(self, monkeypatch): + # The endpoint 499s without a bounded page_size, so every request must + # carry one. + urls: list[str] = [] + + def fake_get(url, token, timeout=10): + urls.append(url) + return {"model_services": [_model_service("system.ai.gpt-5")]}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + ids, reason = db_mod.list_model_services(WS, "token") + + assert ids == ["system.ai.gpt-5"] + assert reason is None + assert all("page_size=" in u for u in urls) + def test_retries_page_before_giving_up(self, monkeypatch): payload = {"model_services": [_model_service("system.ai.gpt-5")]} calls = {"n": 0} @@ -274,7 +260,7 @@ def test_retries_page_before_giving_up(self, monkeypatch): def flaky_get(url, token, timeout=10): calls["n"] += 1 if calls["n"] < 3: - return None, "HTTP 504 Gateway Timeout" + return None, "HTTP 499 Unknown" return payload, None monkeypatch.setattr(db_mod, "_http_get_json", flaky_get) diff --git a/tests/test_e2e_uc.py b/tests/test_e2e_uc.py index d0f994d..72fc2f1 100644 --- a/tests/test_e2e_uc.py +++ b/tests/test_e2e_uc.py @@ -1,8 +1,9 @@ -"""End-to-end tests for the `--enable-uc` / `UCODE_ENABLE_UC` opt-in. +"""End-to-end tests for UC-securables discovery (now always-on). -Verifies UC-securables discovery (model-services + MCP services) and the -flag-precedence ladder (CLI flag > env var > persisted state) against a -live Databricks workspace. +Verifies that `configure_shared_state` discovers models via UC model-services +(`system.ai.*`) by default, falls back to the legacy per-family AI Gateway +listings when UC model-services are absent, and surfaces only `system.ai.*` +entries from the UC primitives. Run with: UCODE_TEST_WORKSPACE=https://your-workspace.databricks.com \ @@ -17,7 +18,6 @@ from ucode.databricks import ( discover_model_services, list_mcp_services, - uc_enabled, ) from ucode.state import load_state @@ -67,150 +67,35 @@ def test_returns_only_system_ai_mcp_services(self, e2e_workspace, e2e_token): # --------------------------------------------------------------------------- -# `uc_enabled` precedence — env var alone (no `default` arg). +# `configure_shared_state` end-to-end: UC discovery is the default, with a +# best-effort fallback to the legacy `databricks-*` listings per family. # --------------------------------------------------------------------------- -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 - ): +class TestConfigureSharedStateE2E: + def test_default_discovers_system_ai(self, monkeypatch, e2e_workspace, e2e_token): + """No flag, no env: a workspace with UC model-services resolves + `system.ai.*` ids and never persists a `uc_enabled` flag.""" 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 + state = configure_shared_state(e2e_workspace, force_login=False) + assert "uc_enabled" not in load_state() 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) + def test_falls_back_to_legacy_when_no_uc_models(self, monkeypatch, e2e_workspace, e2e_token): + """A workspace without UC model-services must still configure, via the + legacy per-family AI Gateway listings (`databricks-*` ids).""" + if _has_uc_models(e2e_workspace, e2e_token): + pytest.skip("Workspace has system.ai.* model services; fallback not exercised.") - # 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." - ) + state = configure_shared_state(e2e_workspace, force_login=False) ids = _all_resolved_model_ids(state) + assert ids, "Fallback discovery returned no models at all." assert all(not m.startswith("system.ai.") for m in ids), ( - f"Discovery used UC despite per-workspace default being False: " + f"Fallback unexpectedly returned UC ids: " f"{[m for m in ids if m.startswith('system.ai.')][:5]}" ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 333bb32..a0b4eba 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -408,6 +408,10 @@ def _patch_mcp_choices(monkeypatch, *values: str) -> None: "prompt_for_mcp_server_choices", lambda *args, **kwargs: list(values), ) + # Curated system.ai.* MCP-services discovery now always runs; stub it so + # configure_mcp_command tests don't shell out to the `databricks` CLI. + # Tests that exercise it override this after calling the helper. + monkeypatch.setattr(mcp, "discover_mcp_service_names", lambda workspace, profile=None: []) class TestConfigureMcpCommand: