diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ffaf0ec9..f573abbe 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -44,6 +44,7 @@ ensure_ai_gateway_v2, ensure_databricks_auth, ensure_pat_bearer, + fetch_managed_coding_agent_configs, find_profile_name_for_host, get_databricks_profiles, get_databricks_token, @@ -63,6 +64,7 @@ ) from ucode.managed_config import ( get_model_recommendation, + is_feature_disabled, load_managed_state, managed_agent_config_enabled, refresh_managed_config, @@ -1853,6 +1855,14 @@ def _launch_managed_default( "dry-run. Run `ucode` without --dry-run to pull your workspace's config first." ) return + token = get_databricks_token(current, state.get("profile")) + _, reason = fetch_managed_coding_agent_configs(current, token) + if reason is not None and is_feature_disabled(reason): + print_warning( + "Managed coding agent config is not enabled for your workspace. Please launch " + "ucode with a specific agent, ex. `ucode codex` or `ucode claude`." + ) + return _print_no_managed_config_guidance(current, state.get("profile")) return # The budget tier can move the org to a cheaper agent, so it outranks the config's diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 5e9bbcee..2ff28065 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -276,6 +276,8 @@ def get_model_recommendation(workspace: str, token: str) -> tuple[dict | None, s """ payload, reason = fetch_model_recommendation(workspace, token) if reason is not None: + if is_feature_disabled(reason): + return None, None return None, reason agent = AGENT_ENUM_TO_TOOL.get(_str(payload.get("recommended_agent")) or "") model = _str(payload.get("recommended_model")) @@ -310,8 +312,7 @@ def get_managed_config(workspace: str, token: str) -> tuple[dict | None, str | N """ configs, reason = fetch_managed_coding_agent_configs(workspace, token) if reason is not None: - # A NOT_FOUND means the admin hasn't defined a config for this workspace — not a failure. - if _is_not_found(reason): + if _is_not_found(reason) or is_feature_disabled(reason): return None, None return None, reason if not configs: @@ -328,6 +329,10 @@ def _is_not_found(reason: str) -> bool: return "http 404" in lowered or "not_found" in lowered +def is_feature_disabled(reason: str) -> bool: + return "feature_disabled" in reason.lower() + + def _is_permission_denied(reason: str) -> bool: """True when the read was refused rather than answering whether a config exists. diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index f5da862b..ab7b2bf4 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -28,6 +28,7 @@ delete_coding_agent_config, discover_claude_models_unbucketed, ensure_databricks_auth, + fetch_managed_coding_agent_configs, get_databricks_token, has_cached_model_provider_services, is_model_provider_feature_unavailable, @@ -40,6 +41,7 @@ ) from ucode.managed_config import ( get_managed_config, + is_feature_disabled, load_managed_state, managed_state_workspace, save_managed_state, @@ -860,6 +862,16 @@ def _render_summary(workspace: str, manifest: dict) -> None: print_panel("Configuration summary", lines) +def _require_feature_enabled(workspace: str, token: str) -> None: + with spinner("Checking whether managed coding-agent config is enabled..."): + _, reason = fetch_managed_coding_agent_configs(workspace, token) + if reason is not None and is_feature_disabled(reason): + raise RuntimeError( + f"Managed coding-agent config is not enabled for {workspace}. Reach out to your " + "account admin to enable the Enhanced Unity AI Gateway Preview." + ) + + def _require_admin(workspace: str, token: str) -> None: """Stop unless the caller is a workspace admin. @@ -1032,6 +1044,7 @@ def setup_command(from_file: str | None = None) -> int: ensure_databricks_auth(workspace, profile, quiet=True) token = get_databricks_token(workspace, profile) + _require_feature_enabled(workspace, token) _require_admin(workspace, token) if not _handle_existing_config(workspace, token): return 0 diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 38c8a201..b45f7f7d 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -178,6 +178,23 @@ def test_not_found_is_treated_as_no_config(self, monkeypatch, not_found_reason): assert cfg is None assert reason is None + @pytest.mark.parametrize( + "disabled_reason", + [ + 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED","message":"..."}', + 'HTTP 403 Forbidden: {"error_code":"FEATURE_DISABLED"}', + ], + ) + def test_feature_disabled_is_treated_as_no_config(self, monkeypatch, disabled_reason): + monkeypatch.setattr( + mc_mod, + "fetch_managed_coding_agent_configs", + lambda ws, tok: ([], disabled_reason), + ) + cfg, reason = get_managed_config("https://ws", "tok") + assert cfg is None + assert reason is None + class TestPersistence: @pytest.fixture(autouse=True) @@ -492,6 +509,12 @@ def test_failed_read_surfaces_the_reason(self, monkeypatch): self._stub(monkeypatch, {}, reason="HTTP 500") assert mc_mod.get_model_recommendation("https://w", "tok") == (None, "HTTP 500") + def test_feature_disabled_is_no_recommendation(self, monkeypatch): + self._stub( + monkeypatch, {}, reason='HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED"}' + ) + assert mc_mod.get_model_recommendation("https://w", "tok") == (None, None) + def test_unparseable_decimals_become_none(self, monkeypatch): self._stub( monkeypatch, {"recommended_agent": "CODING_AGENT_PI", "current_spend": "not-a-number"} diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index cc4b8e50..73eef967 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -213,6 +213,28 @@ def test_unverifiable_check_warns_and_continues(self): assert warn.called +class TestFeatureGate: + def test_feature_disabled_is_rejected(self): + disabled = 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED"}' + with patch.object( + wizard, "fetch_managed_coding_agent_configs", return_value=([], disabled) + ): + with pytest.raises(RuntimeError, match="Enhanced Unity AI Gateway Preview"): + wizard._require_feature_enabled(WORKSPACE, "token") + + def test_feature_enabled_passes(self): + with patch.object(wizard, "fetch_managed_coding_agent_configs", return_value=([], None)): + wizard._require_feature_enabled(WORKSPACE, "token") # must not raise + + def test_transient_read_failure_is_allowed_through(self): + with patch.object( + wizard, + "fetch_managed_coding_agent_configs", + return_value=([], "HTTP 500 Server Error"), + ): + wizard._require_feature_enabled(WORKSPACE, "token") # must not raise + + class TestExistingConfigHandling: RICH_CONFIG = { "name": "coding-agent-configs/abc",