From 5851bbc7fc4556aaa68a6e2b6c3f6a93a4b93aab Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 12 Aug 2026 15:20:35 +0000 Subject: [PATCH 1/3] Treat FEATURE_DISABLED as no managed config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the coding-agent config feature is gated off server-side, the AI Gateway returns FEATURE_DISABLED. Treat it exactly like NOT_FOUND — the workspace behaves as though no managed config existed — so ucode falls back to its defaults cleanly instead of warning and reapplying a cached config. - get_managed_config: collapse FEATURE_DISABLED to (None, None), the authoritative "no config" that also clears a previously cached config. - get_model_recommendation: collapse FEATURE_DISABLED to (None, None) so the budget check falls back silently to the default model. Co-authored-by: Isaac --- src/ucode/managed_config.py | 18 +++++++++++++++++- tests/test_managed_config.py | 26 ++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 5e9bbcee..f82ca3d8 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -276,6 +276,10 @@ def get_model_recommendation(workspace: str, token: str) -> tuple[dict | None, s """ payload, reason = fetch_model_recommendation(workspace, token) if reason is not None: + # A disabled feature means the workspace behaves as though no config existed: fall back to + # the default model silently instead of warning about a budget we couldn't read. + 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")) @@ -311,7 +315,10 @@ 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): + # A FEATURE_DISABLED means the feature is switched off server-side; either way the workspace + # has no config in effect, so both collapse to the authoritative (None, None) that clears a + # previously cached config rather than a warning that would reapply it. + if _is_not_found(reason) or _is_feature_disabled(reason): return None, None return None, reason if not configs: @@ -328,6 +335,15 @@ def _is_not_found(reason: str) -> bool: return "http 404" in lowered or "not_found" in lowered +def _is_feature_disabled(reason: str) -> bool: + """True when a read failure means the coding-agent config feature is switched off server-side. + + The gateway gates the feature behind a SAFE flag; when it is off, reads fail with a + ``FEATURE_DISABLED`` error whose body carries that ``error_code``. A disabled feature is not a + failure to report — the workspace behaves exactly as though no managed config existed.""" + 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/tests/test_managed_config.py b/tests/test_managed_config.py index 38c8a201..71bd157f 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -178,6 +178,25 @@ 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): + # A FEATURE_DISABLED means the feature is switched off server-side — the workspace behaves as + # though no config existed, so it collapses to (None, None) and clears any cached config. + 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 +511,13 @@ 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): + # A disabled feature reads as no recommendation, not a failure to warn the developer about. + 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"} From 525c086d84982c71b09d63252b0e7ab024af25de Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 12 Aug 2026 18:01:43 +0000 Subject: [PATCH 2/3] Reject `ucode setup` when coding-agent config feature is disabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit setup previously walked the admin through the whole wizard even when the server has the feature gated off: the existing-config read reports "no config" (FEATURE_DISABLED collapses to that for the launch fallback), so setup saw nothing published and offered to create one. Add an up-front `_require_feature_enabled` gate to the interactive setup flow. It lists configs and, on FEATURE_DISABLED, fails with a clear message pointing the admin at the Enhanced Unity AI Gateway Preview — before any prompting, and before the admin check to match the server's own gate order. Other read failures pass through, since the API still enforces the gate at publish time. Co-authored-by: Isaac --- src/ucode/managed_wizard.py | 22 ++++++++++++++++++++++ tests/test_managed_wizard.py | 23 +++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index f5da862b..7ed52c83 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, @@ -39,6 +40,7 @@ update_coding_agent_config, ) from ucode.managed_config import ( + _is_feature_disabled, get_managed_config, load_managed_state, managed_state_workspace, @@ -860,6 +862,25 @@ def _render_summary(workspace: str, manifest: dict) -> None: print_panel("Configuration summary", lines) +def _require_feature_enabled(workspace: str, token: str) -> None: + """Stop unless managed coding-agent config is enabled for this workspace. + + When the server has the feature gated off it rejects every config read and write with + FEATURE_DISABLED, so authoring one can't work. Detect that here — with the same list read the + wizard makes next — and fail with a clear message instead of walking the admin through a setup + that could never publish. Mirrors the server's own order: the feature gate precedes the admin + check. A read that fails for any other reason is allowed through; the API still enforces the gate + at publish time. + """ + 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 +1053,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_wizard.py b/tests/test_managed_wizard.py index cc4b8e50..9d3271a8 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -213,6 +213,29 @@ 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): + # A non-FEATURE_DISABLED read failure must not block setup — the API still gates writes. + 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", From a4e6441383bbdafa566461d19fdf07295ec563bd Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 12 Aug 2026 18:33:12 +0000 Subject: [PATCH 3/3] Bare `ucode`: dedicated message when the feature is disabled On bare `ucode`, a disabled feature previously showed the generic "no managed config found; using your local settings" line, because the launch read collapses FEATURE_DISABLED to an empty result. Re-read on the no-config path to tell the two apart and, when the feature is off, point the developer at a specific agent (`ucode codex` / `ucode claude`). Also promote is_feature_disabled to a public helper (now shared by cli, managed_wizard, and managed_config). Co-authored-by: Isaac --- src/ucode/cli.py | 10 ++++++++++ src/ucode/managed_config.py | 17 +++-------------- src/ucode/managed_wizard.py | 13 ++----------- tests/test_managed_config.py | 3 --- tests/test_managed_wizard.py | 1 - 5 files changed, 15 insertions(+), 29 deletions(-) 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 f82ca3d8..2ff28065 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -276,9 +276,7 @@ def get_model_recommendation(workspace: str, token: str) -> tuple[dict | None, s """ payload, reason = fetch_model_recommendation(workspace, token) if reason is not None: - # A disabled feature means the workspace behaves as though no config existed: fall back to - # the default model silently instead of warning about a budget we couldn't read. - if _is_feature_disabled(reason): + if is_feature_disabled(reason): return None, None return None, reason agent = AGENT_ENUM_TO_TOOL.get(_str(payload.get("recommended_agent")) or "") @@ -314,11 +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. - # A FEATURE_DISABLED means the feature is switched off server-side; either way the workspace - # has no config in effect, so both collapse to the authoritative (None, None) that clears a - # previously cached config rather than a warning that would reapply it. - if _is_not_found(reason) or _is_feature_disabled(reason): + if _is_not_found(reason) or is_feature_disabled(reason): return None, None return None, reason if not configs: @@ -335,12 +329,7 @@ def _is_not_found(reason: str) -> bool: return "http 404" in lowered or "not_found" in lowered -def _is_feature_disabled(reason: str) -> bool: - """True when a read failure means the coding-agent config feature is switched off server-side. - - The gateway gates the feature behind a SAFE flag; when it is off, reads fail with a - ``FEATURE_DISABLED`` error whose body carries that ``error_code``. A disabled feature is not a - failure to report — the workspace behaves exactly as though no managed config existed.""" +def is_feature_disabled(reason: str) -> bool: return "feature_disabled" in reason.lower() diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 7ed52c83..ab7b2bf4 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -40,8 +40,8 @@ update_coding_agent_config, ) from ucode.managed_config import ( - _is_feature_disabled, get_managed_config, + is_feature_disabled, load_managed_state, managed_state_workspace, save_managed_state, @@ -863,18 +863,9 @@ def _render_summary(workspace: str, manifest: dict) -> None: def _require_feature_enabled(workspace: str, token: str) -> None: - """Stop unless managed coding-agent config is enabled for this workspace. - - When the server has the feature gated off it rejects every config read and write with - FEATURE_DISABLED, so authoring one can't work. Detect that here — with the same list read the - wizard makes next — and fail with a clear message instead of walking the admin through a setup - that could never publish. Mirrors the server's own order: the feature gate precedes the admin - check. A read that fails for any other reason is allowed through; the API still enforces the gate - at publish time. - """ 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): + 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." diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 71bd157f..b45f7f7d 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -186,8 +186,6 @@ def test_not_found_is_treated_as_no_config(self, monkeypatch, not_found_reason): ], ) def test_feature_disabled_is_treated_as_no_config(self, monkeypatch, disabled_reason): - # A FEATURE_DISABLED means the feature is switched off server-side — the workspace behaves as - # though no config existed, so it collapses to (None, None) and clears any cached config. monkeypatch.setattr( mc_mod, "fetch_managed_coding_agent_configs", @@ -512,7 +510,6 @@ def test_failed_read_surfaces_the_reason(self, monkeypatch): assert mc_mod.get_model_recommendation("https://w", "tok") == (None, "HTTP 500") def test_feature_disabled_is_no_recommendation(self, monkeypatch): - # A disabled feature reads as no recommendation, not a failure to warn the developer about. self._stub( monkeypatch, {}, reason='HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED"}' ) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 9d3271a8..73eef967 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -227,7 +227,6 @@ def test_feature_enabled_passes(self): wizard._require_feature_enabled(WORKSPACE, "token") # must not raise def test_transient_read_failure_is_allowed_through(self): - # A non-FEATURE_DISABLED read failure must not block setup — the API still gates writes. with patch.object( wizard, "fetch_managed_coding_agent_configs",