From 42bcaa8b0da1d11d989272f777eca45a66acdd6b Mon Sep 17 00:00:00 2001 From: Tien Le Date: Wed, 5 Aug 2026 15:32:21 +0000 Subject: [PATCH 1/6] Add PATCH/DELETE transport and coding-agent-config CRUD clients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The write-side API plumbing for `ucode apply` (next change). No CLI wiring and no interactive flow: transport helpers plus three clients, mirroring how the read client landed. `databricks.py` had `_http_get_json`/`_http_post_json` but no PATCH or DELETE. Rather than a third and fourth near-copy of the same 40 lines of error handling, the body-sending path is factored into `_http_send_json(method, ...)` and the three verbs become thin wrappers. `_http_post_json`'s behavior is unchanged. DELETE needed one real difference: its success response is `google.protobuf.Empty`, which arrives as `{}` or an empty body depending on the gateway. An empty body would otherwise be reported as "response was not valid JSON", so `allow_empty_body` treats it as success. `delete_coding_agent_config` returns only a reason — there is no payload worth handing back. Clients for the three admin RPCs, all workspace-admin gated server-side: - `create_coding_agent_config` — POST to the collection. v0 allows one config per workspace, so this returns ALREADY_EXISTS when one exists. - `update_coding_agent_config` — PATCH the resource. Preferred over delete-then-create, which has a window where the workspace has *no* managed config: if the create failed, every developer would lose their config until someone re-ran the command. The server applies the mask inside a single entity-store update, so a failed write leaves the old config intact. - `delete_coding_agent_config` — DELETE by resource name. `MANAGED_CONFIG_UPDATE_MASK_PATHS` is every field ucode's manifest can set. The server requires a non-empty mask and rejects paths outside its mutable set; this is that set minus what ucode doesn't author — `budget_id` (deprecated for `budget_policy.budget_id`, and rejected on write) and `default_options`/`tiers` (the legacy model-only shape). Sending every path ucode owns, not just the populated ones, is what lets a re-run *clear* a field the admin removed: the server merges per path, so an omitted path leaves the old value in place. `_coding_agent_config_url` joins on the API root rather than the collection URL, since the resource name already carries the `coding-agent-configs/` segment and would otherwise be duplicated. Tests: 15 cases. The mask is checked against `serialize_managed_config`'s actual output rather than a restated list, so adding a manifest field fails the test instead of shipping a mask that cannot clear it. Mutation-verified three ways: dropping `allow_empty_body`, dropping the `update_mask`, and dropping one mask path each fail a specific test. Co-authored-by: Isaac --- src/ucode/databricks.py | 160 ++++++++++++++++++++++++++---- tests/test_databricks.py | 204 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 345 insertions(+), 19 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index e873831..1816fcd 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -252,28 +252,37 @@ def _http_get_json( return None, f"network error: {exc}" -def _http_post_json( - url: str, token: str, payload: dict, *, timeout: int = 10 +def _http_send_json( + method: str, + url: str, + token: str, + payload: dict | None, + *, + timeout: int = 10, + allow_empty_body: bool = False, ) -> tuple[dict | list | None, str | None]: - """POST a JSON body to an endpoint. Returns (payload, None) on success, - (None, reason) on failure. Mirrors `_http_get_json`.""" - body_bytes = json.dumps(payload).encode("utf-8") - request = urllib_request.Request( - url, - data=body_bytes, - method="POST", - headers={ - "Authorization": f"Bearer {token}", - "Accept": "application/json", - "Content-Type": "application/json", - }, - ) + """Send a request that may carry a JSON body, and decode a JSON response. + + Shared by `_http_post_json`, `_http_patch_json`, and `_http_delete` — the three differ only in + verb, whether they send a body, and whether an empty response is success. Returns + ``(payload, None)`` on success and ``(None, reason)`` on failure, like `_http_get_json`. + + ``allow_empty_body`` is for DELETE, whose success response is ``google.protobuf.Empty`` — an + empty body there is the expected result, not a decode failure. + """ + body_bytes = json.dumps(payload).encode("utf-8") if payload is not None else None + headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"} + if body_bytes is not None: + headers["Content-Type"] = "application/json" + request = urllib_request.Request(url, data=body_bytes, method=method, headers=headers) try: with urllib_request.urlopen(request, timeout=timeout) as response: body = response.read().decode("utf-8") - _debug(f"POST {url}", f"HTTP {response.status}, {len(body)} bytes") + _debug(f"{method} {url}", f"HTTP {response.status}, {len(body)} bytes") if _debug_enabled(): _debug("body", body[:4000]) + if allow_empty_body and not body.strip(): + return None, None try: return json.loads(body), None except json.JSONDecodeError as exc: @@ -284,7 +293,7 @@ def _http_post_json( body = exc.read().decode("utf-8", errors="replace") if exc.fp else "" except Exception: body = "" - _debug(f"POST {url}", f"HTTP {exc.code} {exc.reason}") + _debug(f"{method} {url}", f"HTTP {exc.code} {exc.reason}") if _debug_enabled() and body: _debug("body", body[:4000]) reason = f"HTTP {exc.code} {exc.reason}" @@ -293,15 +302,43 @@ def _http_post_json( reason = f"{reason}: {body_excerpt}" return None, reason except urllib_error.URLError as exc: - _debug(f"POST {url}", f"URLError: {exc.reason}") + _debug(f"{method} {url}", f"URLError: {exc.reason}") return None, f"network error: {exc.reason}" except OSError as exc: # See `_http_get_json`: a bare socket timeout is an OSError, not a # URLError, and would otherwise escape the caller's error handling. - _debug(f"POST {url}", f"OSError: {exc}") + _debug(f"{method} {url}", f"OSError: {exc}") return None, f"network error: {exc}" +def _http_post_json( + url: str, token: str, payload: dict, *, timeout: int = 10 +) -> tuple[dict | list | None, str | None]: + """POST a JSON body to an endpoint. Returns (payload, None) on success, + (None, reason) on failure. Mirrors `_http_get_json`.""" + return _http_send_json("POST", url, token, payload, timeout=timeout) + + +def _http_patch_json( + url: str, token: str, payload: dict, *, timeout: int = 10 +) -> tuple[dict | list | None, str | None]: + """PATCH a JSON body to an endpoint. Returns (payload, None) on success, + (None, reason) on failure.""" + return _http_send_json("PATCH", url, token, payload, timeout=timeout) + + +def _http_delete( + url: str, token: str, *, timeout: int = 10 +) -> tuple[dict | list | None, str | None]: + """DELETE a resource. Returns (payload, None) on success, (None, reason) on failure. + + A successful delete returns ``google.protobuf.Empty``, which serializes as ``{}`` or an empty + body depending on the gateway, so both count as success and yield ``(None, None)``. Callers + should test ``reason`` rather than the payload. + """ + return _http_send_json("DELETE", url, token, None, timeout=timeout, allow_empty_body=True) + + def _http_get_bytes(url: str, token: str, *, timeout: int = 10) -> tuple[bytes | None, str | None]: """GET raw bytes. Returns (body, None) on success, (None, reason) on failure. @@ -1575,6 +1612,91 @@ def fetch_model_recommendation(workspace: str, token: str) -> tuple[dict, str | return payload, None +# Every field ucode's manifest can set, as `update_mask` paths for a PATCH. The server rejects a +# missing or empty mask, and rejects paths outside its own mutable set — this is that set minus the +# fields ucode doesn't author: `budget_id` (deprecated in favour of `budget_policy.budget_id`, and +# rejected on write) and `default_options`/`tiers` (the legacy model-only shape superseded by +# `enabled_agents`/`budget_policy`). Sending every path ucode owns, rather than only the ones +# currently populated, is what lets a re-run *clear* a field the admin removed: the server merges +# per path, so an omitted path leaves the old value in place. +MANAGED_CONFIG_UPDATE_MASK_PATHS: tuple[str, ...] = ( + "display_name", + "default_agent", + "enabled_agents", + "mcp_servers", + "skills", + "tracing", + "budget_policy", +) + + +def _coding_agent_config_url(workspace: str, name: str | None = None) -> str: + """The collection URL, or one config's resource URL when ``name`` is given. + + ``name`` is the server-assigned resource name (``coding-agent-configs/{id}``), which the Get and + Update paths template directly, so it is appended as-is rather than rebuilt from an id. + """ + hostname = workspace_hostname(workspace) + base = f"https://{hostname}{_CODING_AGENT_CONFIGS_API_PATH}" + if name is None: + return base + # The resource name already carries the collection segment, so join on the API root. + root = base.rsplit("/coding-agent-configs", 1)[0] + return f"{root}/{name.strip().strip('/')}" + + +def create_coding_agent_config( + workspace: str, token: str, config: dict +) -> tuple[dict | None, str | None]: + """Create the workspace's managed CodingAgentConfig. + + v0 allows at most one config per workspace, so this fails with ALREADY_EXISTS when one is + already defined; callers should update that one instead of creating a second. + """ + url = _coding_agent_config_url(workspace) + payload, reason = _http_post_json(url, token, config, timeout=30) + if reason is not None: + return None, reason + if not isinstance(payload, dict): + return None, "coding-agent-config create returned an unexpected response shape" + return payload, None + + +def update_coding_agent_config( + workspace: str, + token: str, + name: str, + config: dict, + *, + update_mask: tuple[str, ...] = MANAGED_CONFIG_UPDATE_MASK_PATHS, +) -> tuple[dict | None, str | None]: + """Update an existing managed CodingAgentConfig in place. + + Preferred over delete-then-create: the server applies the mask inside a single entity-store + update, so the workspace is never left without a config if the write fails partway. ``name`` + identifies the config and is echoed in the body, which is what the API's path template expects. + """ + url = _coding_agent_config_url(workspace, name) + body = {**config, "name": name, "update_mask": {"paths": list(update_mask)}} + payload, reason = _http_patch_json(url, token, body, timeout=30) + if reason is not None: + return None, reason + if not isinstance(payload, dict): + return None, "coding-agent-config update returned an unexpected response shape" + return payload, None + + +def delete_coding_agent_config(workspace: str, token: str, name: str) -> str | None: + """Delete a managed CodingAgentConfig by resource name. Returns None on success, else a reason. + + Returns only the failure reason: a successful delete responds with ``Empty``, so there is no + payload worth handing back. + """ + url = _coding_agent_config_url(workspace, name) + _, reason = _http_delete(url, token, timeout=30) + return reason + + # --- MCP services (parallel to model services) ----------------------------- diff --git a/tests/test_databricks.py b/tests/test_databricks.py index bfdf3e3..41bc12b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2286,3 +2286,207 @@ def test_false_when_the_payload_names_no_groups(self, monkeypatch, payload): # A well-formed `Me` for a user in no groups omits `groups` entirely. self._stub(monkeypatch, payload) assert db_mod.is_workspace_admin("https://w", "tok") is False + + +class TestCodingAgentConfigUrls: + def test_collection_url(self): + assert db_mod._coding_agent_config_url(WS) == f"{WS}/api/ai-gateway/v2/coding-agent-configs" + + def test_resource_url_appends_the_server_assigned_name(self): + # The API templates Get/Update/Delete on `{name=coding-agent-configs/*}`, so the resource + # name already carries the collection segment and must not be duplicated. + url = db_mod._coding_agent_config_url(WS, "coding-agent-configs/abc123") + assert url == f"{WS}/api/ai-gateway/v2/coding-agent-configs/abc123" + + def test_stray_slashes_are_tolerated(self): + url = db_mod._coding_agent_config_url(WS, "/coding-agent-configs/abc123/") + assert url == f"{WS}/api/ai-gateway/v2/coding-agent-configs/abc123" + + +class TestHttpDelete: + """A successful delete returns `google.protobuf.Empty`, so an empty body is success.""" + + @staticmethod + def _empty_response(body: str = ""): + from unittest.mock import MagicMock + + response = MagicMock() + response.__enter__ = lambda s: s + response.__exit__ = MagicMock(return_value=False) + response.read.return_value = body.encode("utf-8") + response.status = 200 + return response + + def test_empty_body_is_success_not_a_decode_error(self, monkeypatch): + # Without `allow_empty_body` this would fail with "response was not valid JSON". + monkeypatch.setattr( + db_mod.urllib_request, "urlopen", lambda request, timeout=None: self._empty_response() + ) + payload, reason = db_mod._http_delete(f"{WS}/api/anything", "tok") + assert reason is None + assert payload is None + + def test_empty_json_object_is_also_success(self, monkeypatch): + monkeypatch.setattr( + db_mod.urllib_request, + "urlopen", + lambda request, timeout=None: self._empty_response("{}"), + ) + payload, reason = db_mod._http_delete(f"{WS}/api/anything", "tok") + assert reason is None + assert payload == {} + + def test_uses_the_delete_verb_and_sends_no_body(self, monkeypatch): + seen = {} + + def capture(request, timeout=None): + seen["method"] = request.get_method() + seen["data"] = request.data + return self._empty_response() + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", capture) + db_mod._http_delete(f"{WS}/api/anything", "tok") + assert seen["method"] == "DELETE" + assert seen["data"] is None + + def test_http_error_surfaces_the_body(self, monkeypatch): + import io + from unittest.mock import MagicMock + from urllib.error import HTTPError + + body = '{"error_code":"PERMISSION_DENIED","message":"admin required"}' + + def raise_http_error(request, timeout=None): + raise HTTPError( + url="", code=403, msg="Forbidden", hdrs=MagicMock(), fp=io.BytesIO(body.encode()) + ) + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", raise_http_error) + _, reason = db_mod._http_delete(f"{WS}/api/anything", "tok") + assert reason is not None + assert "403" in reason + assert "PERMISSION_DENIED" in reason + + +class TestHttpPatchJson: + def test_uses_the_patch_verb_and_sends_the_body(self, monkeypatch): + from unittest.mock import MagicMock + + seen = {} + + def capture(request, timeout=None): + seen["method"] = request.get_method() + seen["data"] = request.data + seen["content_type"] = request.get_header("Content-type") + response = MagicMock() + response.__enter__ = lambda s: s + response.__exit__ = MagicMock(return_value=False) + response.read.return_value = b'{"name":"coding-agent-configs/x"}' + response.status = 200 + return response + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", capture) + payload, reason = db_mod._http_patch_json(f"{WS}/api/anything", "tok", {"k": "v"}) + assert reason is None + assert payload == {"name": "coding-agent-configs/x"} + assert seen["method"] == "PATCH" + assert json.loads(seen["data"]) == {"k": "v"} + assert seen["content_type"] == "application/json" + + +class TestCodingAgentConfigCrudClients: + CONFIG = {"default_agent": "CODING_AGENT_CLAUDE_CODE"} + + def test_create_posts_the_config_to_the_collection(self, monkeypatch): + seen = {} + + def fake_post(url, token, payload, *, timeout=10): + seen.update(url=url, payload=payload) + return {"name": "coding-agent-configs/new"}, None + + monkeypatch.setattr(db_mod, "_http_post_json", fake_post) + config, reason = db_mod.create_coding_agent_config(WS, "tok", self.CONFIG) + assert reason is None + assert config == {"name": "coding-agent-configs/new"} + assert seen["url"] == f"{WS}/api/ai-gateway/v2/coding-agent-configs" + assert seen["payload"] == self.CONFIG + + def test_create_surfaces_the_failure_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda *a, **k: (None, 'HTTP 400: {"error_code":"ALREADY_EXISTS"}'), + ) + config, reason = db_mod.create_coding_agent_config(WS, "tok", self.CONFIG) + assert config is None + assert "ALREADY_EXISTS" in reason + + def test_update_patches_the_resource_with_a_mask(self, monkeypatch): + seen = {} + + def fake_patch(url, token, payload, *, timeout=10): + seen.update(url=url, payload=payload) + return {"name": "coding-agent-configs/abc"}, None + + monkeypatch.setattr(db_mod, "_http_patch_json", fake_patch) + config, reason = db_mod.update_coding_agent_config( + WS, "tok", "coding-agent-configs/abc", self.CONFIG + ) + assert reason is None + assert config == {"name": "coding-agent-configs/abc"} + assert seen["url"] == f"{WS}/api/ai-gateway/v2/coding-agent-configs/abc" + # The server rejects a missing or empty mask, and needs `name` in the body for its path + # template — both must be present alongside the config's own fields. + assert seen["payload"]["name"] == "coding-agent-configs/abc" + assert seen["payload"]["update_mask"]["paths"] + assert seen["payload"]["default_agent"] == "CODING_AGENT_CLAUDE_CODE" + + def test_update_mask_never_names_a_field_the_server_rejects(self): + # The server's mutable set is the upper bound; `budget_id` is in it but deprecated and + # rejected on write, so ucode must not name it. `default_options`/`tiers` are the legacy + # model-only shape ucode never authors. + assert "budget_id" not in db_mod.MANAGED_CONFIG_UPDATE_MASK_PATHS + assert "default_options" not in db_mod.MANAGED_CONFIG_UPDATE_MASK_PATHS + assert "tiers" not in db_mod.MANAGED_CONFIG_UPDATE_MASK_PATHS + + def test_update_mask_covers_every_field_the_manifest_can_set(self): + # A path ucode omits is a field a re-run silently cannot clear, since the server merges per + # path. Derive the expectation from the serializer rather than restating it, so adding a + # manifest field fails here instead of shipping a mask that can't clear it. + from ucode.managed_setup import serialize_managed_config + + emitted = set( + serialize_managed_config( + { + "display_name": "org config", + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-5"}} + }, + "mcp_servers": [{"name": "databricks-sql", "type": "sql"}], + "skills": {"names": ["main.default"]}, + "tracing_table": "main.default.traces", + "budget_policy": { + "budget_id": "11111111-1111-1111-1111-111111111111", + "tiers": [], + }, + } + ) + ) + assert emitted == set(db_mod.MANAGED_CONFIG_UPDATE_MASK_PATHS) + + def test_delete_returns_only_a_reason(self, monkeypatch): + seen = {} + + def fake_delete(url, token, *, timeout=10): + seen["url"] = url + return None, None + + monkeypatch.setattr(db_mod, "_http_delete", fake_delete) + assert db_mod.delete_coding_agent_config(WS, "tok", "coding-agent-configs/abc") is None + assert seen["url"] == f"{WS}/api/ai-gateway/v2/coding-agent-configs/abc" + + def test_delete_surfaces_the_failure_reason(self, monkeypatch): + monkeypatch.setattr(db_mod, "_http_delete", lambda *a, **k: (None, "HTTP 404 Not Found")) + reason = db_mod.delete_coding_agent_config(WS, "tok", "coding-agent-configs/abc") + assert reason == "HTTP 404 Not Found" From 9852bcf4245479909cdf9329512b17ead066c67f Mon Sep 17 00:00:00 2001 From: Tien Le Date: Wed, 5 Aug 2026 15:36:40 +0000 Subject: [PATCH 2/6] setup: require a UUID budget id, and index tiers the way the server does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three gaps found by reading the server-side validation this manifest is written for (universe #2365441), all reachable through `--from-file` even though the wizard can't produce them. `budget_policy.budget_id` must parse as a UUID. The handler requires it, and until now ucode only checked non-empty — so a hand-written manifest carrying `"budget_id": "eng-budget"` passed local validation and failed at the API with an INVALID_PARAMETER_VALUE. Local pre-flight exists precisely to spend the round trip on real problems. The message names `budget_configuration_id` so an admin knows where to get a valid one. Tier positions are now reported 0-based. The server indexes with `zipWithIndex`, so ucode's `tiers[1]` and the API's `tiers[0]` described the same tier — an admin reconciling the two messages would be looking at the wrong one. Added a test that the deprecated top-level `CodingAgentConfig.budget_id` (field 3) is never emitted, even when a hand-written manifest sets it. The serializer already only writes `budget_policy.budget_id`; the handler rejects the top-level field, so this pins behavior that is currently correct by construction rather than by intent. Test fixtures used short placeholders (`"b"`, `"budget-1"`) where a real `budget_configuration_id` would be, so those are now UUIDs — 20 occurrences across the two files. `list_workspace_budgets` only ever returns real ones, so the fixtures were describing input the wizard can't produce. Tests: +6. Mutation-verified: dropping the UUID check fails four cases, reverting to 1-based indices fails `test_tier_positions_are_reported_zero_based`, and emitting the top-level `budget_id` fails `test_a_manifest_carrying_a_top_level_budget_id_still_omits_it`. Co-authored-by: Isaac --- src/ucode/managed_setup.py | 23 +++++++++++-- tests/test_managed_setup.py | 67 +++++++++++++++++++++++++++++++----- tests/test_managed_wizard.py | 26 ++++++++------ 3 files changed, 93 insertions(+), 23 deletions(-) diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index 2ef8373..b94ee7a 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -25,6 +25,7 @@ import json import os +import uuid from pathlib import Path from typing import cast @@ -532,14 +533,30 @@ def _agent_model_ids(agent_config: dict) -> set[str]: def _validate_budget_policy(budget_policy: dict, enabled_agents: dict[str, dict]) -> list[str]: - """Validate a ``budget_policy`` against the agents the manifest enables.""" + """Validate a ``budget_policy`` against the agents the manifest enables. + + Tier positions are reported 0-based to match the server's own messages, which index with + ``zipWithIndex`` — an admin comparing the two error sources should see the same number. + """ errors: list[str] = [] - if not budget_policy.get("budget_id"): + budget_id = budget_policy.get("budget_id") + if not budget_id: errors.append("budget_policy.budget_id is required.") + else: + # The server requires a parseable UUID here. The wizard can only offer real + # `budget_configuration_id`s, but `--from-file` and hand-edited manifests can carry + # anything, and catching it locally beats an INVALID_PARAMETER_VALUE round-trip. + try: + uuid.UUID(str(budget_id)) + except ValueError: + errors.append( + f"budget_policy.budget_id must be a UUID (got '{budget_id}'). Use the " + "budget_configuration_id from the workspace's AI Gateway budgets." + ) percentages: list[float] = [] tiers = budget_policy.get("tiers") - for index, tier in enumerate(tiers if isinstance(tiers, list) else [], start=1): + for index, tier in enumerate(tiers if isinstance(tiers, list) else []): if not isinstance(tier, dict): errors.append(f"budget_policy.tiers[{index}] must be an object.") continue diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index 8a08200..8bac8c8 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -36,6 +36,10 @@ WORKSPACE = "https://ws.example.com" +# The server requires `budget_policy.budget_id` to parse as a UUID, so fixtures that aren't +# *testing* that rule need a real one. +BUDGET_ID = "11111111-1111-1111-1111-111111111111" + # A workspace state shaped like `configure_shared_state` produces. STATE = { "workspace": WORKSPACE, @@ -256,6 +260,19 @@ def test_budget_tiers_keep_fractions(self): assert [tier["spending_percentage"] for tier in tiers] == [0.8, 1.0] assert tiers[1]["default_agent"] == "CODING_AGENT_OPENCODE" + def test_the_deprecated_top_level_budget_id_is_never_emitted(self): + # `CodingAgentConfig.budget_id` (field 3) is deprecated in favour of + # `budget_policy.budget_id`, and the CRUD handler rejects a write that sets it. The budget + # id must appear only under the policy. + payload = serialize_managed_config(_full_manifest()) + assert "budget_id" not in payload + assert payload["budget_policy"]["budget_id"] == "c6563b45-df9a-4b19-afb2-d42dc2b52576" + + def test_a_manifest_carrying_a_top_level_budget_id_still_omits_it(self): + # A hand-written `--from-file` manifest could set it; the serializer must not pass it on. + payload = serialize_managed_config({**_full_manifest(), "budget_id": BUDGET_ID}) + assert "budget_id" not in payload + def test_unknown_agent_is_dropped(self): payload = serialize_managed_config( { @@ -633,7 +650,7 @@ def test_tier_percentage_must_be_a_fraction(self, pct): manifest = { **_minimal_manifest(), "budget_policy": { - "budget_id": "b", + "budget_id": BUDGET_ID, "tiers": [ { "spending_percentage": pct, @@ -654,7 +671,7 @@ def test_tier_percentages_must_be_unique(self): } manifest = { **_minimal_manifest(), - "budget_policy": {"budget_id": "b", "tiers": [tier, dict(tier)]}, + "budget_policy": {"budget_id": BUDGET_ID, "tiers": [tier, dict(tier)]}, } errors = validate_manifest(manifest, STATE) assert any("must be unique" in e for e in errors) @@ -663,7 +680,7 @@ def test_tier_agent_must_be_enabled(self): manifest = { **_minimal_manifest(), "budget_policy": { - "budget_id": "b", + "budget_id": BUDGET_ID, "tiers": [ { "spending_percentage": 0.5, @@ -680,7 +697,7 @@ def test_tier_needs_a_default_model(self): manifest = { **_minimal_manifest(), "budget_policy": { - "budget_id": "b", + "budget_id": BUDGET_ID, "tiers": [{"spending_percentage": 0.5, "default_agent": "claude"}], }, } @@ -701,7 +718,7 @@ def test_tier_model_must_be_one_the_agent_has(self): } }, "budget_policy": { - "budget_id": "b", + "budget_id": BUDGET_ID, "tiers": [ { "spending_percentage": 0.8, @@ -726,7 +743,7 @@ def test_tier_model_from_the_agents_list_is_accepted(self): } }, "budget_policy": { - "budget_id": "b", + "budget_id": BUDGET_ID, "tiers": [ { "spending_percentage": 0.8, @@ -750,7 +767,7 @@ def test_tier_model_matching_a_claude_family_slot_is_accepted(self): } }, "budget_policy": { - "budget_id": "b", + "budget_id": BUDGET_ID, "tiers": [ { "spending_percentage": 0.8, @@ -775,7 +792,7 @@ def test_tier_model_check_skipped_when_the_agent_lists_nothing(self): } }, "budget_policy": { - "budget_id": "b", + "budget_id": BUDGET_ID, "tiers": [ { "spending_percentage": 0.8, @@ -788,9 +805,41 @@ def test_tier_model_check_skipped_when_the_agent_lists_nothing(self): assert validate_manifest(manifest, STATE) == [] def test_budget_policy_alone_still_requires_a_default_agent(self): - errors = validate_manifest({"budget_policy": {"budget_id": "b"}}) + errors = validate_manifest({"budget_policy": {"budget_id": BUDGET_ID}}) assert any("default_agent is required" in e for e in errors) + @pytest.mark.parametrize("bad_id", ["not-a-uuid", "b", "1111", "11111111-1111-1111-1111"]) + def test_budget_id_must_be_a_uuid(self, bad_id): + # The server requires a parseable UUID. The wizard can only offer real ids, but + # `--from-file` can carry anything, and rejecting it here beats a round-trip failure. + manifest = { + **_minimal_manifest(), + "budget_policy": {"budget_id": bad_id, "tiers": []}, + } + errors = validate_manifest(manifest, STATE) + assert any("must be a UUID" in e for e in errors), errors + + def test_a_real_uuid_is_accepted(self): + manifest = { + **_minimal_manifest(), + "budget_policy": {"budget_id": "c6563b45-df9a-4b19-afb2-d42dc2b52576", "tiers": []}, + } + assert validate_manifest(manifest, STATE) == [] + + def test_tier_positions_are_reported_zero_based(self): + # The server indexes tiers with `zipWithIndex`, so an admin comparing ucode's message with + # the API's must see the same number for the same tier. + manifest = { + **_minimal_manifest(), + "budget_policy": { + "budget_id": BUDGET_ID, + "tiers": [{"spending_percentage": 0.5, "default_agent": "claude"}], + }, + } + errors = validate_manifest(manifest, STATE) + assert any("tiers[0]" in e for e in errors), errors + assert not any("tiers[1]" in e for e in errors), errors + def test_errors_accumulate(self): manifest = { "default_agent": "codex", diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 126d630..5366936 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -24,6 +24,10 @@ WORKSPACE = "https://ws.example.com" +# `list_workspace_budgets` returns real `budget_configuration_id`s, and validation requires a +# parseable UUID, so the fixtures use one rather than a readable placeholder. +BUDGET_ID = "c6563b45-df9a-4b19-afb2-d42dc2b52576" + STATE = { "workspace": WORKSPACE, "claude_models": { @@ -885,14 +889,14 @@ def test_no_budgets_warns_and_yields_none(self): assert warn.called def test_percentages_are_stored_as_fractions(self): - budgets = [{"id": "budget-1", "display_name": "eng"}] + budgets = [{"id": BUDGET_ID, "display_name": "eng"}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, "prompt_for_selection", - side_effect=["budget-1", "claude", "system.ai.claude-opus-4-8"], + side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], ), patch.object(wizard, "prompt_for_text", return_value="tiered"), # prompt_for_percentage already converts; it returns the fraction. @@ -900,7 +904,7 @@ def test_percentages_are_stored_as_fractions(self): ): policy = wizard._prompt_budget_policy(WORKSPACE, "token", CLAUDE_ONLY, STATE) assert policy is not None - assert policy["budget_id"] == "budget-1" + assert policy["budget_id"] == BUDGET_ID assert policy["tiers"] == [ { "spending_percentage": 0.8, @@ -920,14 +924,14 @@ def test_offers_only_the_models_the_agent_was_configured_with(self): } } } - budgets = [{"id": "budget-1", "display_name": "eng"}] + budgets = [{"id": BUDGET_ID, "display_name": "eng"}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, "prompt_for_selection", - side_effect=["budget-1", "pi", "system.ai.kimi-k2-6"], + side_effect=[BUDGET_ID, "pi", "system.ai.kimi-k2-6"], ) as select, patch.object(wizard, "prompt_for_text", return_value="tiered"), patch.object(wizard, "prompt_for_percentage", return_value=0.8), @@ -949,14 +953,14 @@ def test_claude_family_slots_are_flattened_for_the_picker(self): } } } - budgets = [{"id": "budget-1", "display_name": "eng"}] + budgets = [{"id": BUDGET_ID, "display_name": "eng"}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, "prompt_for_selection", - side_effect=["budget-1", "claude", "system.ai.claude-opus-4-8"], + side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], ) as select, patch.object(wizard, "prompt_for_text", return_value="tiered"), patch.object(wizard, "prompt_for_percentage", return_value=0.8), @@ -968,14 +972,14 @@ def test_claude_family_slots_are_flattened_for_the_picker(self): def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): # An agent configured through a provider service has no enumerable list; better to offer the # catalog than nothing at all. - budgets = [{"id": "budget-1", "display_name": "eng"}] + budgets = [{"id": BUDGET_ID, "display_name": "eng"}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, "prompt_for_selection", - side_effect=["budget-1", "gemini", "system.ai.gemini-3-flash"], + side_effect=[BUDGET_ID, "gemini", "system.ai.gemini-3-flash"], ) as select, patch.object(wizard, "prompt_for_text", return_value="tiered"), patch.object(wizard, "prompt_for_percentage", return_value=0.8), @@ -985,14 +989,14 @@ def test_falls_back_to_the_catalog_when_an_agent_lists_nothing(self): assert offered == ["system.ai.gemini-3-flash"] def test_authored_policy_validates(self): - budgets = [{"id": "budget-1", "display_name": "eng"}] + budgets = [{"id": BUDGET_ID, "display_name": "eng"}] with ( patch.object(wizard, "prompt_yes_no_default", side_effect=[True, False]), patch.object(wizard, "list_workspace_budgets", return_value=(budgets, None)), patch.object( wizard, "prompt_for_selection", - side_effect=["budget-1", "claude", "system.ai.claude-opus-4-8"], + side_effect=[BUDGET_ID, "claude", "system.ai.claude-opus-4-8"], ), patch.object(wizard, "prompt_for_text", return_value="tiered"), patch.object(wizard, "prompt_for_percentage", return_value=0.8), From d0f0ae6f44c4e730b80b7803a284a6860f9997ef Mon Sep 17 00:00:00 2001 From: Tien Le Date: Wed, 5 Aug 2026 19:28:14 +0000 Subject: [PATCH 3/6] setup: drop the agent -> oneof-variant identity map MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback on #267: `_AGENT_MODEL_CONFIG_VARIANT` mapped each agent to its `AgentModelConfig` oneof key, but the proto's field names are ucode's tool names verbatim — claude, codex, opencode, pi, gemini, copilot — so every entry mapped a name to itself. One use site, so the tool now serves as the key directly. The dict's only other effect was a KeyError on an unknown agent, which was already unreachable: `serialize_managed_config` filters to `tool in AGENT_TOOL_TO_ENUM` before calling this, and the `AGENT_TOOL_TO_ENUM[tool]` lookup two lines down would raise first anyway. No new test. The variant keys are already covered — hard-coding the wrong one fails `test_codex_model_config_has_no_model_list` and `test_flat_list_agents_use_repeated_models`, and the round-trip through `normalize_managed_config` asserts the alignment for every agent. The other half of that review comment — validating a model against its agent's dialect, so a manifest can't pin a GPT id for Claude Code — is deliberately not here. It turned out to need a decision rather than a patch: the agent -> families mapping already exists twice (`agents._TOOL_DISCOVERY_SOURCES` and `managed_setup._AGENT_MODEL_FAMILIES`) and the two disagree for codex, copilot, opencode, and pi. Picking one requires checking each agent's own writer, and the likely outcome is that this module's opencode entry is wrong — it lists `codex` while `build_opencode_base_urls` serves no OpenAI route — which would be a picker bug in the wizard, not a refactor. Landing that separately. Co-authored-by: Isaac --- src/ucode/managed_setup.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index b94ee7a..759ee9c 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -48,17 +48,6 @@ AGENT_TOOL_TO_ENUM: dict[str, str] = {tool: enum for enum, tool in AGENT_ENUM_TO_TOOL.items()} MCP_TAG_TO_TYPE_ENUM: dict[str, str] = {tag: enum for enum, tag in MCP_TYPE_ENUM_TO_TAG.items()} -# `AgentModelConfig` oneof variant key per agent. The server rejects a config whose variant doesn't -# match its agent (`validateAgentModelConfig`), so this mapping is not cosmetic. -_AGENT_MODEL_CONFIG_VARIANT: dict[str, str] = { - "claude": "claude", - "codex": "codex", - "opencode": "opencode", - "pi": "pi", - "gemini": "gemini", - "copilot": "copilot", -} - # Agents whose model config carries a flat `models` list. Claude instead uses per-family slots # (`ClaudeDefaultModels`), and Codex has no model list at all — it selects exactly one model. _FLAT_MODEL_LIST_AGENTS = frozenset({"opencode", "pi", "gemini", "copilot"}) @@ -250,8 +239,11 @@ def _enabled_agent_payload(tool: str, agent_config: dict) -> dict: if isinstance(model_config, dict): body = _model_config_payload(tool, model_config) if body: - variant = _AGENT_MODEL_CONFIG_VARIANT[tool] - config["model_config"] = {variant: body} + # The `AgentModelConfig` oneof field names are ucode's tool names verbatim (claude, + # codex, opencode, pi, gemini, copilot), so the tool doubles as the variant key. The + # server rejects a variant that doesn't match its agent (`validateAgentModelConfig`), + # and the round-trip through `normalize_managed_config` pins that alignment in tests. + config["model_config"] = {tool: body} entry: dict = {"agent": AGENT_TOOL_TO_ENUM[tool]} if config: From ae99ee5716d33348b44e875b55458d60f5c01456 Mon Sep 17 00:00:00 2001 From: Tien Le <44335153+tt-le@users.noreply.github.com> Date: Fri, 7 Aug 2026 02:33:38 -0400 Subject: [PATCH 4/6] Add `ucode apply`: publish the authored managed config (#271) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The publish step for the manifest `ucode setup` authors. Validates, shows what would change, confirms, then writes it to the workspace via the clients added in the parent commit. Updates in place rather than replacing. When the workspace already has a config, `apply` PATCHes it using its resource name (read back from the existing-config GET, which `normalize_managed_config` preserves). Delete-then-create was the original plan — v0's Create returns ALREADY_EXISTS — but it has a window where the workspace has *no* managed config, and if the create failed there every developer would silently fall back to their own settings until someone re-ran the command. The server applies the update mask inside a single entity-store update, so a failed PATCH leaves the current config intact. It is still a whole-manifest write: every path ucode owns is sent, so a field the admin dropped on a re-run is cleared rather than left behind. Refuses to publish when it cannot tell whether a config already exists. A failed existence check used to be the one case where "just try the create" would either duplicate or silently overwrite an admin's work, so an unreadable check is a hard error naming the reason rather than a warning. `_explain_publish_failure` maps the failures an admin will actually hit. FEATURE_DISABLED is the likely first experience — the CRUD flag is off by default — so it names `codingAgentConfigCrudEnabled` instead of printing an HTTP 400. INVALID_PARAMETER_VALUE is passed through verbatim: the server names the offending field, which beats any paraphrase. `--yes` skips the confirmation for CI; `--dry-run` validates and previews without writing. The admin gate and validation both run before anything is sent, so an invalid manifest or a non-admin costs no round trip. README documents the publish step and drops the "no partial update yet" caveat, which the PATCH path makes untrue. Verified against eng-ml-inference.staging: `apply --dry-run` authenticated, verified admin, rendered the summary, detected the existing config, and chose the update path without writing. Tests: 19 cases. Mutation-verified three ways — always-create instead of PATCH, publishing despite an unreadable existence check, and publishing an invalid manifest each fail a specific test. Also covers that `typer.Exit(0)` isn't caught by the command's own RuntimeError handler, the same trap `setup` hit. Co-authored-by: Isaac --- README.md | 21 +++- src/ucode/cli.py | 30 ++++- src/ucode/managed_wizard.py | 124 ++++++++++++++++++- tests/test_managed_wizard.py | 222 +++++++++++++++++++++++++++++++++++ 4 files changed, 393 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index b9dd630..7d6fd7a 100644 --- a/README.md +++ b/README.md @@ -200,8 +200,23 @@ ucode setup --dry-run ucode setup --from-file ./managed-settings.json ``` -Publishing replaces the workspace's config outright — there is no partial update yet, so anything -skipped in a re-run is dropped. +Once the manifest looks right, publish it: + +```bash +# Validate, show what would change, and ask before publishing. +ucode apply + +# Preview without publishing. +ucode apply --dry-run + +# Publish without the confirmation prompt (for CI). +ucode apply --yes +``` + +`apply` updates the workspace's existing config in place rather than replacing it, so a failed +publish leaves the current config intact. It is still a whole-manifest write: every field ucode +authors is sent, so anything skipped in a re-run is cleared rather than carried over. Developers +pick the new config up on their next ucode run. --- @@ -230,6 +245,8 @@ skipped in a re-run is dropped. | `ucode setup` | Author the workspace's managed coding config (workspace admins only) | | `ucode setup show` | Print the authored config and the payload `ucode apply` would publish | | `ucode setup --from-file ` | Load a hand-written managed config instead of running the prompts | +| `ucode apply` | Publish the authored managed config to the workspace (workspace admins only) | +| `ucode apply --yes` | Publish without the confirmation prompt | ## Managed Local Files diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a465338..2a079ff 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -77,7 +77,7 @@ recommended_agent, resolve_state, ) -from ucode.managed_wizard import setup_command, show_command +from ucode.managed_wizard import apply_command, setup_command, show_command from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, @@ -2339,6 +2339,34 @@ def setup_show_cmd() -> None: raise typer.Exit(code) +@app.command("apply") +def apply_cmd( + yes: Annotated[ + bool, + typer.Option("--yes", "-y", help="Publish without the confirmation prompt."), + ] = False, + dry_run: Annotated[ + bool, + typer.Option("--dry-run", help="Validate and preview without publishing."), + ] = False, +) -> None: + """Publish this workspace's managed coding config (workspace admins only).""" + set_dry_run(dry_run) + # See the `setup` callback: `typer.Exit` subclasses RuntimeError, so it must be raised after + # the try block or the handler below would report a successful exit as an error. + try: + install_databricks_cli() + code = apply_command(yes=yes) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + except KeyboardInterrupt: + print_err("Interrupted.") + raise typer.Exit(130) from None + if code: + raise typer.Exit(code) + + @app.command("status") def status_cmd() -> None: """Show current workspace, tool configs, and saved model selections.""" diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index 9cf2655..e99a059 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -15,10 +15,13 @@ import json from pathlib import Path +from typing import cast from ucode.agents import TOOL_SPECS, check_gateway_endpoint +from ucode.config_io import is_dry_run from ucode.databricks import ( ANTHROPIC_FAMILIES, + create_coding_agent_config, discover_claude_models_unbucketed, ensure_databricks_auth, get_databricks_token, @@ -28,6 +31,7 @@ list_model_provider_services, list_workspace_budgets, service_usable_for_tool, + update_coding_agent_config, ) from ucode.managed_config import get_managed_config from ucode.managed_setup import ( @@ -897,4 +901,122 @@ def show_command() -> int: return 0 -__all__ = ["setup_command", "setup_from_file", "show_command"] +# Server-side failures an admin is actually likely to hit, mapped to something they can act on. The +# raw reasons are `HTTP : ` strings from the transport, and the body carries the +# API's `error_code`, so matching on that is more robust than on status codes alone. +def _explain_publish_failure(reason: str) -> str: + lowered = reason.lower() + if "feature_disabled" in lowered: + return ( + "Managed coding-agent configs aren't enabled on this workspace yet. Ask your Databricks " + "contact to enable the `codingAgentConfigCrudEnabled` flag for it, then re-run " + "`ucode apply`." + ) + if "permission_denied" in lowered or "http 403" in lowered: + return ( + "Publishing a managed config requires workspace admin. Your account can read the " + "workspace but not author its coding config." + ) + if "already_exists" in lowered: + return ( + "This workspace already has a managed config, but ucode couldn't read it to update in " + "place. Run `ucode apply` again — if it keeps failing, the existing config may need to " + "be deleted by hand." + ) + if "invalid_parameter_value" in lowered: + # The server names the offending field; passing it through beats paraphrasing. + return f"The workspace rejected the config: {reason}" + return f"Could not publish the managed config: {reason}" + + +def apply_command(*, yes: bool = False) -> int: + """Publish the authored manifest to the workspace. + + Updates the existing config in place when there is one, rather than deleting and recreating it: + a failed recreate would leave the workspace with no managed config at all, and every developer + would silently fall back to their own settings. Returns a process exit code. + """ + from ucode.cli import _prompt_for_configuration + + print_section("ucode apply") + + state = load_state() + workspace = state.get("workspace") + profile = state.get("profile") + if not workspace: + workspace, profile = _prompt_for_configuration() + + manifest = load_managed_settings(workspace) + if manifest is None: + raise RuntimeError( + "No managed config has been authored for this workspace. Run `ucode setup` first " + "(or `ucode setup --from-file `)." + ) + + errors = validate_manifest(manifest, state) + if errors: + print_err("The authored config is not valid, so it was not published:") + for error in errors: + print_note(error) + print_note("Re-run `ucode setup` to fix it, or edit ~/.ucode/managed-settings.json.") + return 1 + + ensure_databricks_auth(workspace, profile) + token = get_databricks_token(workspace, profile) + _require_admin(workspace, token) + + payload = serialize_managed_config(manifest) + _render_summary(workspace, manifest) + + # Read before writing: the resource name tells us whether to create or update, and shows the + # admin what they are about to overwrite. + with spinner("Checking for an existing managed config..."): + existing, reason = get_managed_config(workspace, token) + if reason is not None: + raise RuntimeError( + f"Could not check whether {workspace} already has a managed config: {reason}. " + "Refusing to publish without knowing, since that could overwrite a config silently." + ) + + existing_name = (existing or {}).get("name") + if existing is not None and not isinstance(existing_name, str): + raise RuntimeError( + "This workspace has a managed config but the API didn't return its resource name, so " + "ucode can't update it in place. Delete it in the workspace and re-run `ucode apply`." + ) + + console.print() + if existing is None: + print_note(f"This will create a new managed config on {workspace}.") + else: + agents = ", ".join((existing.get("enabled_agents") or {}).keys()) or "no agents" + print_warning( + f"This will replace the config already published on {workspace} (currently: {agents}). " + "Every developer picks the new one up on their next ucode run." + ) + if not yes and not prompt_yes_no_default("Publish this config?", default=False): + print_note("Nothing was published.") + return 1 + + if is_dry_run(): + print_success("Dry run: the config was validated but not published.") + return 0 + + if existing is None: + with spinner("Publishing the managed config..."): + published, publish_reason = create_coding_agent_config(workspace, token, payload) + else: + with spinner("Updating the managed config..."): + published, publish_reason = update_coding_agent_config( + workspace, token, cast("str", existing_name), payload + ) + if publish_reason is not None: + raise RuntimeError(_explain_publish_failure(publish_reason)) + + name = (published or {}).get("name") or existing_name or "coding-agent-configs/?" + print_success(f"Published {name} to {workspace}") + print_note("Developers pick this up on their next ucode run.") + return 0 + + +__all__ = ["apply_command", "setup_command", "setup_from_file", "show_command"] diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 5366936..b8b8d77 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -14,6 +14,7 @@ import typer.main from typer.testing import CliRunner +import ucode.cli as cli_mod import ucode.config_io as config_io_mod import ucode.managed_setup as managed_setup_mod import ucode.managed_wizard as wizard @@ -1320,6 +1321,202 @@ def fake_sel(prompt, options, **kwargs): assert any("model" in p for p in searchable_prompts), searchable_prompts +class TestApplyCommand: + MANIFEST = { + "default_agent": "claude", + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} + }, + } + + @staticmethod + def _patches(**overrides): + """The network/auth boundary `apply_command` sits behind, with per-test overrides.""" + defaults = { + "load_state": lambda: {"workspace": WORKSPACE, "profile": "p", **STATE}, + "ensure_databricks_auth": lambda *a, **k: None, + "get_databricks_token": lambda *a, **k: "tok", + "is_workspace_admin": lambda *a, **k: True, + "get_managed_config": lambda *a, **k: (None, None), + "create_coding_agent_config": lambda *a, **k: ( + {"name": "coding-agent-configs/new"}, + None, + ), + "update_coding_agent_config": lambda *a, **k: ( + {"name": "coding-agent-configs/old"}, + None, + ), + "prompt_yes_no_default": lambda *a, **k: True, + } + defaults.update(overrides) + return [patch.object(wizard, name, value) for name, value in defaults.items()] + + def _run(self, *, yes=False, **overrides): + import contextlib + + with contextlib.ExitStack() as stack: + for p in self._patches(**overrides): + stack.enter_context(p) + return wizard.apply_command(yes=yes) + + def test_unauthored_config_is_an_actionable_error(self): + with patch.object(wizard, "load_state", return_value={"workspace": WORKSPACE}): + with pytest.raises(RuntimeError, match="ucode setup"): + wizard.apply_command() + + def test_creates_when_no_config_exists(self): + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + created = {} + + def fake_create(workspace, token, payload): + created.update(workspace=workspace, payload=payload) + return {"name": "coding-agent-configs/new"}, None + + assert self._run(create_coding_agent_config=fake_create) == 0 + assert created["workspace"] == WORKSPACE + # What goes over the wire is proto-JSON, not ucode's manifest shape. + assert created["payload"]["default_agent"] == "CODING_AGENT_CLAUDE_CODE" + + def test_updates_in_place_when_a_config_exists(self): + # Delete-then-create would leave the workspace with no config if the create failed, so an + # existing config must be PATCHed rather than replaced. + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + existing = {"name": "coding-agent-configs/abc", "enabled_agents": {"codex": {}}} + updated = {} + created = {"called": False} + + def fake_update(workspace, token, name, payload): + updated.update(name=name, payload=payload) + return {"name": name}, None + + def fake_create(*a, **k): + created["called"] = True + return {}, None + + assert ( + self._run( + get_managed_config=lambda *a, **k: (existing, None), + update_coding_agent_config=fake_update, + create_coding_agent_config=fake_create, + ) + == 0 + ) + assert updated["name"] == "coding-agent-configs/abc" + assert created["called"] is False + + def test_invalid_manifest_is_not_published(self): + # `default_agent` names an agent that isn't enabled. + managed_setup_mod.save_managed_settings( + WORKSPACE, {"default_agent": "codex", "enabled_agents": {"claude": {}}} + ) + created = {"called": False} + + def fake_create(*a, **k): + created["called"] = True + return {}, None + + assert self._run(create_coding_agent_config=fake_create) == 1 + assert created["called"] is False + + def test_declining_the_prompt_publishes_nothing(self): + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + created = {"called": False} + + def fake_create(*a, **k): + created["called"] = True + return {}, None + + code = self._run( + prompt_yes_no_default=lambda *a, **k: False, create_coding_agent_config=fake_create + ) + assert code == 1 + assert created["called"] is False + + def test_yes_skips_the_prompt(self): + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + + def refuse(*a, **k): + raise AssertionError("--yes must not prompt") + + assert self._run(yes=True, prompt_yes_no_default=refuse) == 0 + + def test_non_admin_is_rejected_before_publishing(self): + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + created = {"called": False} + + def fake_create(*a, **k): + created["called"] = True + return {}, None + + with pytest.raises(RuntimeError, match="not an admin"): + self._run( + is_workspace_admin=lambda *a, **k: False, create_coding_agent_config=fake_create + ) + assert created["called"] is False + + def test_unreadable_existing_config_refuses_to_publish(self): + # Publishing without knowing whether a config exists risks silently overwriting one. + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + created = {"called": False} + + def fake_create(*a, **k): + created["called"] = True + return {}, None + + with pytest.raises(RuntimeError, match="Refusing to publish"): + self._run( + get_managed_config=lambda *a, **k: (None, "HTTP 500 Server Error"), + create_coding_agent_config=fake_create, + ) + assert created["called"] is False + + def test_existing_config_without_a_resource_name_is_an_error(self): + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + with pytest.raises(RuntimeError, match="resource name"): + self._run(get_managed_config=lambda *a, **k: ({"enabled_agents": {}}, None)) + + def test_dry_run_validates_without_publishing(self, monkeypatch): + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + monkeypatch.setattr(config_io_mod, "_dry_run", True) + created = {"called": False} + + def fake_create(*a, **k): + created["called"] = True + return {}, None + + assert self._run(create_coding_agent_config=fake_create) == 0 + assert created["called"] is False + + +class TestPublishFailureMessages: + """The server's error codes, turned into something an admin can act on.""" + + def test_feature_disabled_names_the_flag(self): + message = wizard._explain_publish_failure( + 'HTTP 400 Bad Request: {"error_code":"FEATURE_DISABLED","message":"..."}' + ) + assert "codingAgentConfigCrudEnabled" in message + + def test_permission_denied_says_admin_is_required(self): + message = wizard._explain_publish_failure( + 'HTTP 403 Forbidden: {"error_code":"PERMISSION_DENIED"}' + ) + assert "workspace admin" in message + + def test_invalid_parameter_value_is_passed_through_verbatim(self): + # The server names the offending field, which is more useful than any paraphrase. + reason = ( + 'HTTP 400 Bad Request: {"error_code":"INVALID_PARAMETER_VALUE",' + '"message":"budget_policy.tiers[0].spending_percentage must be between 0 and 1"}' + ) + message = wizard._explain_publish_failure(reason) + assert "budget_policy.tiers[0].spending_percentage" in message + + def test_unknown_failure_still_surfaces_the_reason(self): + message = wizard._explain_publish_failure("network error: timed out") + assert "timed out" in message + + class TestCliWiring: def test_setup_is_registered(self): result = runner.invoke(app, ["--help"]) @@ -1341,6 +1538,31 @@ def test_setup_show_is_registered(self): assert result.exit_code == 0 assert "show" in result.output + def test_apply_is_registered(self): + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0 + assert "apply" in result.output + + def test_apply_declares_yes_and_dry_run(self): + # Asserted on the declared options rather than rendered help, which Rich ellipsizes at + # narrow terminal widths (see test_setup_help_lists_from_file). + command = typer.main.get_command(app).commands["apply"] # type: ignore[attr-defined] + declared = {opt for param in command.params for opt in param.opts} + assert {"--yes", "--dry-run"} <= declared + + def test_apply_error_exits_nonzero_with_a_message(self): + with patch.object(cli_mod, "apply_command", side_effect=RuntimeError("no config authored")): + result = runner.invoke(app, ["apply"]) + assert result.exit_code == 1 + + def test_successful_apply_exits_zero(self): + # Same trap as `setup`: `typer.Exit` subclasses RuntimeError, so raising it inside the + # command's try block would report success as "ERROR 0". + with patch.object(cli_mod, "apply_command", return_value=0): + result = runner.invoke(app, ["apply"]) + assert result.exit_code == 0 + assert "ERROR" not in result.output + def test_successful_setup_exits_zero(self): # `typer.Exit` subclasses RuntimeError, so a success code must not be caught and reported # as an error by the command's own RuntimeError handler. From 9f1c69ba3d27c7d5c1f7ec797bf64431437b0e9d Mon Sep 17 00:00:00 2001 From: Tien Le Date: Fri, 7 Aug 2026 13:43:11 +0000 Subject: [PATCH 5/6] apply: validate against the workspace's full model listing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode apply` rejected a model `ucode setup` had just offered: claude: model 'system.ai.claude-opus-4-8' is not available on this workspace. `state["claude_models"]` holds only the newest id per family, because the launch path pins one model per family alias. The wizard deliberately offers the older versions too — pinning `default_opus_model` to a known-good `claude-opus-4-8` is a normal thing for an admin to want — and stashes the full listing on `state["all_claude_models"]` so validation recognizes them. That stash is never persisted. `setup` saves the manifest, not the state, so a separate `apply` process starts from a fresh `load_state()` without it and validates against the narrow per-family inventory. The failure lands at the very end of the flow, naming a model the wizard itself had listed a moment earlier. `apply` now re-fetches the listing instead of trusting what `setup` left behind, which also covers a hand-edited or `--from-file` manifest authored on another machine. Best-effort: a failed listing leaves validation on the narrower inventory rather than blocking a publish on a transient API error. `ensure_databricks_auth` moves above validation since the listing needs a token; nothing is written until well after. Aarushi predicted this failure, and its exact message, reviewing the invariant on `_claude_candidates` in #268. The fix there documented and tested the invariant within one process — both paths satisfied it — and so missed that `setup` and `apply` are two processes. Tests: +2, covering the older-version publish and that a failed fetch still publishes. Mutation-verified: validating against bare `state` reproduces the original error. Co-authored-by: Isaac --- src/ucode/managed_wizard.py | 41 ++++++++++++++++++++++-- tests/test_managed_wizard.py | 60 ++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 2 deletions(-) diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index e99a059..da68ba4 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -929,6 +929,40 @@ def _explain_publish_failure(reason: str) -> str: return f"Could not publish the managed config: {reason}" +def _with_claude_inventory(state: dict, workspace: str, profile: str | None) -> dict: + """``state`` plus the full Claude listing, for validating a manifest against the workspace. + + ``state["claude_models"]`` holds only the newest id per family (the launch path pins one model + per family alias), but `ucode setup` deliberately offers the older versions too — pinning + ``default_opus_model`` to a known-good ``claude-opus-4-8`` is a normal thing for an admin to + want. Validating against ``claude_models`` alone therefore rejected a model the wizard itself + had just offered: + + claude: model 'system.ai.claude-opus-4-8' is not available on this workspace. + + The wizard stashes the full listing on ``state["all_claude_models"]`` mid-run, but that is never + persisted — `setup` saves the manifest, not the state — so a separate `ucode apply` process + starts from a fresh ``load_state()`` without it. Re-fetching here makes the check independent of + what the wizard happened to leave behind, which also covers a hand-edited or ``--from-file`` + manifest authored on another machine. + + Best-effort: a failed listing returns ``state`` untouched, leaving validation on the narrower + inventory rather than blocking a publish on a transient API error. + """ + if isinstance(state.get("all_claude_models"), list) and state["all_claude_models"]: + return state + try: + token = get_databricks_token(workspace, profile) + all_claude, _ = discover_claude_models_unbucketed(workspace, token) + except (RuntimeError, OSError): + # OSError covers a missing `databricks` binary: `get_databricks_token` shells out, so a + # machine without the CLI on PATH raises FileNotFoundError rather than RuntimeError. + return state + if not all_claude: + return state + return {**state, "all_claude_models": all_claude} + + def apply_command(*, yes: bool = False) -> int: """Publish the authored manifest to the workspace. @@ -953,7 +987,11 @@ def apply_command(*, yes: bool = False) -> int: "(or `ucode setup --from-file `)." ) - errors = validate_manifest(manifest, state) + # Auth first: validating a Claude manifest needs the workspace's full model listing, and that + # listing needs a token. Nothing is written until well below this point. + ensure_databricks_auth(workspace, profile) + + errors = validate_manifest(manifest, _with_claude_inventory(state, workspace, profile)) if errors: print_err("The authored config is not valid, so it was not published:") for error in errors: @@ -961,7 +999,6 @@ def apply_command(*, yes: bool = False) -> int: print_note("Re-run `ucode setup` to fix it, or edit ~/.ucode/managed-settings.json.") return 1 - ensure_databricks_auth(workspace, profile) token = get_databricks_token(workspace, profile) _require_admin(workspace, token) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index b8b8d77..27ebe69 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -1418,6 +1418,66 @@ def fake_create(*a, **k): assert self._run(create_coding_agent_config=fake_create) == 1 assert created["called"] is False + def test_an_older_family_version_the_wizard_offered_still_publishes(self): + # `setup` offers every version of a Claude family, but `claude_models` keeps only the newest + # per family and the wizard's `all_claude_models` stash is never persisted — so a separate + # `apply` process used to reject a model it had just offered: + # claude: model 'system.ai.claude-opus-4-1' is not available on this workspace. + # `apply` re-fetches the full listing rather than trusting what `setup` left in state. + managed_setup_mod.save_managed_settings( + WORKSPACE, + { + "default_agent": "claude", + "enabled_agents": { + "claude": { + "model_config": { + "default_model": "system.ai.claude-opus-4-1", + "models": {"default_opus_model": "system.ai.claude-opus-4-1"}, + } + } + }, + }, + ) + published: dict = {} + + def fake_create(workspace, token, payload): + published["payload"] = payload + return {"name": "coding-agent-configs/new"}, None + + # State carries only the newest Opus, as a fresh `load_state()` would. + narrow = {"workspace": WORKSPACE, "profile": "p", "claude_models": {"opus": "newest"}} + assert ( + self._run( + load_state=lambda: dict(narrow), + discover_claude_models_unbucketed=lambda *a, **k: ( + ["system.ai.claude-opus-4-1", "newest"], + None, + ), + create_coding_agent_config=fake_create, + ) + == 0 + ) + assert published, "the manifest should have been published" + + def test_a_failed_inventory_fetch_does_not_block_publishing(self): + # The re-fetch is best-effort: a transient listing failure must not turn into a refusal to + # publish a manifest that validates against what state already knows. + managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) + published: dict = {} + + def fake_create(workspace, token, payload): + published["payload"] = payload + return {"name": "coding-agent-configs/new"}, None + + assert ( + self._run( + discover_claude_models_unbucketed=lambda *a, **k: ([], "HTTP 500"), + create_coding_agent_config=fake_create, + ) + == 0 + ) + assert published + def test_declining_the_prompt_publishes_nothing(self): managed_setup_mod.save_managed_settings(WORKSPACE, self.MANIFEST) created = {"called": False} From 695791b723a391f79dfea08ca5e7bae3fd3ae831 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Fri, 7 Aug 2026 13:47:41 +0000 Subject: [PATCH 6/6] apply: send update_mask as a query param, not in the body MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PATCH was rejected by the server it was written for: HTTP 400 INVALID_PARAMETER_VALUE: Field 'update_mask' is required and must contain at least one subfield with a non-default value! Two mistakes, both visible in the RPC's HTTP binding (universe ai-gateway-api/api/proto/service.proto:319): patch: "/ai-gateway/v2/{coding_agent_config.name=coding-agent-configs/*}" body: "coding_agent_config" `body: "coding_agent_config"` means the config *is* the entire request body, so there is nowhere in it for a sibling `update_mask` — the one we nested was parsed as an unknown config field, leaving the mask genuinely absent. It belongs in the query string. And `update_mask` is a `google.protobuf.FieldMask`, whose JSON and query form is a single comma-separated string, not the `{"paths": [...]}` object we sent. `name` stays in the body: the path template reads it from the config. The mask contents were right, and the 15 tests covering them all passed — they asserted on `payload["update_mask"]["paths"]`, which is exactly the shape the server rejects. A unit test that mirrors the client's own assumption cannot catch a wire-format error; this was only ever going to surface against a real workspace. Tests: the update test now asserts the mask arrives in the query string, as one comma-separated FieldMask, and is absent from the body. Mutation-verified: restoring the nested-object form fails it. Co-authored-by: Isaac --- src/ucode/databricks.py | 16 ++++++++++++++-- tests/test_databricks.py | 15 +++++++++++---- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 1816fcd..e902e34 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1675,9 +1675,21 @@ def update_coding_agent_config( Preferred over delete-then-create: the server applies the mask inside a single entity-store update, so the workspace is never left without a config if the write fails partway. ``name`` identifies the config and is echoed in the body, which is what the API's path template expects. + + ``update_mask`` goes in the query string, not the body. The RPC's HTTP binding is + ``patch: "…/{coding_agent_config.name=coding-agent-configs/*}"`` with ``body: + "coding_agent_config"`` — the config *is* the whole body, so a mask nested inside it is parsed + as an unknown config field and the server reports the mask as missing: + + Field 'update_mask' is required and must contain at least one subfield with a non-default + value! + + It is also a ``google.protobuf.FieldMask``, whose JSON/query form is one comma-separated string + rather than a ``{"paths": [...]}`` object. """ - url = _coding_agent_config_url(workspace, name) - body = {**config, "name": name, "update_mask": {"paths": list(update_mask)}} + query = urlencode({"update_mask": ",".join(update_mask)}) + url = f"{_coding_agent_config_url(workspace, name)}?{query}" + body = {**config, "name": name} payload, reason = _http_patch_json(url, token, body, timeout=30) if reason is not None: return None, reason diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 41bc12b..02f15ce 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -5,6 +5,7 @@ import json import os import subprocess +from urllib.parse import parse_qs import pytest @@ -2434,11 +2435,17 @@ def fake_patch(url, token, payload, *, timeout=10): ) assert reason is None assert config == {"name": "coding-agent-configs/abc"} - assert seen["url"] == f"{WS}/api/ai-gateway/v2/coding-agent-configs/abc" - # The server rejects a missing or empty mask, and needs `name` in the body for its path - # template — both must be present alongside the config's own fields. + # The mask rides in the query string: the RPC binds `body: "coding_agent_config"`, so the + # config is the whole body and a mask nested inside it is read as an unknown config field — + # the server then reports the mask as missing. A FieldMask's JSON form is one + # comma-separated string, not a `{"paths": [...]}` object. + url, _, query = seen["url"].partition("?") + assert url == f"{WS}/api/ai-gateway/v2/coding-agent-configs/abc" + mask = parse_qs(query)["update_mask"][0].split(",") + assert mask == list(db_mod.MANAGED_CONFIG_UPDATE_MASK_PATHS) + assert "update_mask" not in seen["payload"] + # `name` still goes in the body: the API's path template reads it from the config. assert seen["payload"]["name"] == "coding-agent-configs/abc" - assert seen["payload"]["update_mask"]["paths"] assert seen["payload"]["default_agent"] == "CODING_AGENT_CLAUDE_CODE" def test_update_mask_never_names_a_field_the_server_rejects(self):