From 53fb77b673d66c5640fbe238b96cdfe335b99629 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Wed, 5 Aug 2026 15:32:21 +0000 Subject: [PATCH 1/3] 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 758bb25..ea32872 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. @@ -1531,6 +1568,91 @@ def fetch_managed_coding_agent_configs(workspace: str, token: str) -> tuple[list return [c for c in configs if isinstance(c, dict)], 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 1b142cf..87934ed 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2105,3 +2105,207 @@ def test_a_later_caller_cannot_corrupt_the_cache(self, monkeypatch): hit.pop() again, _ = db_mod.list_model_provider_services(WS, "tok") assert [s["name"] for s in again] == ["main.j.ant", "main.j.oai"] + + +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 0294f47cbc42cb40c6b2780a89c32330e705e850 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Wed, 5 Aug 2026 15:36:40 +0000 Subject: [PATCH 2/3] 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 b6a1e09..d6b2f6b 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": { @@ -883,14 +887,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. @@ -898,7 +902,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, @@ -918,14 +922,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), @@ -947,14 +951,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), @@ -966,14 +970,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), @@ -983,14 +987,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 64d82585007356716a59a3261e87271e6c60bbf3 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Wed, 5 Aug 2026 19:28:14 +0000 Subject: [PATCH 3/3] 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: