From 6cab5df88d7ddc6f13802d0dae5460fbdf955f9d Mon Sep 17 00:00:00 2001 From: David Liu Date: Tue, 1 Sep 2026 23:23:55 +0000 Subject: [PATCH] Split managed state into draft and published slots `~/.ucode/managed-state.json` held one config per workspace, written both by `refresh_managed_config` (the copy fetched from the workspace on every launch) and by the authoring wizard (the admin's local, unpublished draft). The two clobbered each other: a launch wiped an in-progress draft, and a launch could apply an unpublished draft as if it were published policy. Store both under a versioned per-workspace map with separate `draft` and `published` slots. A v1 file migrates on read, with its original bytes kept at `managed-state.json.pre-v2.bak` on the next write. The map is written through a sibling temp file and renamed into place, because it now holds the admin's draft: nothing can refetch that, so a torn write would lose it outright. `load_managed_state` and `save_managed_state` stay as thin wrappers over the published slot so existing callers keep working; a follow-up moves them over. Co-authored-by: Isaac --- README.md | 1 + src/ucode/managed_config.py | 162 +++++++++++++++++++++++++++-------- tests/test_managed_config.py | 85 +++++++++++++++++- 3 files changed, 211 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index 661ab410..641cb20c 100644 --- a/README.md +++ b/README.md @@ -387,6 +387,7 @@ control the installation. | `~/.pi/agent/models.json` | Pi | | `~/.cursor/mcp.json` | Cursor Agent (MCP servers only) | | `~/.ucode/managed-state.json` | The managed config — authored by `ucode setup` (admins) and refreshed from the workspace on launch | +| `~/.ucode/managed-state.json.pre-v2.bak` | One-time copy of a pre-slots `managed-state.json`, kept when it is first migrated | | `~/.ucode/managed-backups/` | Baseline backups for OS-managed files changed by ucode | Existing files are backed up before being overwritten. `ucode revert` restores backups. diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 7bbf512b..5ce3a114 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -36,10 +36,12 @@ fetch_model_recommendation, get_databricks_token, ) -from ucode.ui import console, print_warning +from ucode.ui import print_warning MANAGED_STATE_PATH = config_io.APP_DIR / "managed-state.json" +MANAGED_STATE_VERSION = 2 + # Opt-in switch while the feature is in bug bash: unset means launches ignore managed configs # entirely and behave exactly as they did before. MANAGED_CONFIG_ENV_VAR = "ENABLE_MANAGED_AGENT_CONFIG" @@ -342,29 +344,96 @@ def _is_permission_denied(reason: str) -> bool: return "http 403" in lowered or "permission_denied" in lowered -def save_managed_state(workspace: str, config: dict) -> None: - """Persist the normalized managed config to ``~/.ucode/managed-state.json`` at mode 0600. +def _migrated_workspaces(data: dict) -> dict[str, dict]: + """Return the ``workspaces`` map from a raw ``managed-state.json`` dict, migrating v1 in-memory. + + v1 stored a single ``{workspace, config}`` slot shared by both the admin's authored draft and the + launch-fetched published copy, so a legacy value's provenance can't be recovered. It is migrated + into the ``published`` slot — the common case is a developer's fetched snapshot, and treating it + as published keeps ``ucode export`` / ``ucode publish`` (now strictly draft-only) from mistaking a + cached fetch for authored work. A one-time backup (see :func:`_backup_legacy_file_once`) makes a + rare unpublished admin draft recoverable. Read-only: never writes. + """ + if data.get("version") == MANAGED_STATE_VERSION and isinstance(data.get("workspaces"), dict): + return {k: v for k, v in cast("dict", data["workspaces"]).items() if isinstance(v, dict)} + workspaces: dict[str, dict] = {} + legacy_ws = data.get("workspace") + if isinstance(legacy_ws, str) and legacy_ws: + legacy_cfg = data.get("config") + workspaces[legacy_ws] = {"published": legacy_cfg if isinstance(legacy_cfg, dict) else {}} + return workspaces + + +def _is_legacy_file(data: dict) -> bool: + """True when ``data`` is a pre-v2 ``{workspace, config}`` file (not the versioned map).""" + return data.get("version") != MANAGED_STATE_VERSION and "workspace" in data + + +def _backup_legacy_file_once() -> None: + """Copy a pre-v2 ``managed-state.json`` to ``.pre-v2.bak`` before it is overwritten. + + Best-effort and idempotent: migration maps the single legacy slot into ``published``, which can't + preserve a rare unpublished admin draft, so the original bytes are kept once for recovery. Written + through a temp file so an interrupted copy cannot leave a truncated backup that the ``exists()`` + guard would then treat as the original.""" + backup = MANAGED_STATE_PATH.with_suffix(MANAGED_STATE_PATH.suffix + ".pre-v2.bak") + if backup.exists() or not MANAGED_STATE_PATH.exists(): + return + if not _is_legacy_file(config_io.read_json_safe(MANAGED_STATE_PATH)): + return + tmp = backup.with_name(backup.name + ".tmp") + try: + tmp.write_bytes(MANAGED_STATE_PATH.read_bytes()) + _restrict_permissions(tmp) + os.replace(tmp, backup) + except OSError: + pass + + +def _save_slot(workspace: str, slot: str, config: dict) -> None: + """Write ``config`` to ``workspace``'s ``published`` or ``draft`` slot, preserving everything else. - The file is org-authored, not developer-editable — 0600 keeps it readable/writable only by the - user (a light guard; hard enforcement / sudo ownership is a separate concern). No-op in dry-run. + Reads the current file, migrates it to the v2 map, sets one slot for one workspace, and writes it + back — so a refresh of the published slot never disturbs an admin's draft (or other workspaces). + 0600 keeps the org-authored file readable/writable only by the user. No-op write in dry-run. - An empty ``config`` records "this workspace has no managed config", which matters because the - file doubles as the fallback when a later read fails: without it, removing a config server-side - would leave the old one on disk to be reapplied after a transient outage. + An empty ``config`` in the ``published`` slot records "this workspace has no managed config", + which matters because that slot doubles as the fallback when a later read fails: without it, + removing a config server-side would leave the old one on disk to be reapplied after an outage. """ - payload = {"workspace": workspace, "config": config} + workspaces = _migrated_workspaces(config_io.read_json_safe(MANAGED_STATE_PATH)) + workspaces[workspace] = {**workspaces.get(workspace, {}), slot: config} + payload = {"version": MANAGED_STATE_VERSION, "workspaces": workspaces} if config_io.is_dry_run(): - # Print rather than write, matching how the agent config writers behave under --dry-run. - console.print( - f"\n[bold]\\[dry run] {MANAGED_STATE_PATH}[/bold]\n{json.dumps(payload, indent=2)}\n" - ) + config_io.write_json_file(MANAGED_STATE_PATH, payload) return - config_io.ensure_parent_dir(MANAGED_STATE_PATH) + _backup_legacy_file_once() + _write_state_atomically(payload) + + +def _write_state_atomically(payload: dict) -> None: + """Write the state map through a sibling temp file and rename it into place. + + The map holds the admin's draft, which nothing can rebuild: a torn write would take the draft + with it, where before v2 an interrupted write only cost a snapshot the next launch refetches. + """ + tmp = MANAGED_STATE_PATH.with_name(MANAGED_STATE_PATH.name + ".tmp") + config_io.write_json_file(tmp, payload) + _restrict_permissions(tmp) try: - MANAGED_STATE_PATH.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + os.replace(tmp, MANAGED_STATE_PATH) except OSError as exc: raise RuntimeError(f"Failed to write managed state file: {MANAGED_STATE_PATH}") from exc - _restrict_permissions(MANAGED_STATE_PATH) + + +def save_published_config(workspace: str, config: dict) -> None: + """Persist the workspace's launch-fetched ``published`` snapshot, leaving any ``draft`` intact.""" + _save_slot(workspace, "published", config) + + +def save_draft_config(workspace: str, config: dict) -> None: + """Persist the admin's locally authored ``draft``, leaving the ``published`` snapshot intact.""" + _save_slot(workspace, "draft", config) def _restrict_permissions(path: Path) -> None: @@ -376,34 +445,55 @@ def _restrict_permissions(path: Path) -> None: pass -def load_managed_state(workspace: str | None) -> dict | None: - """Load the persisted managed config for ``workspace``, or None if absent/mismatched. - - Returns the normalized config dict (the ``config`` field), only when the stored file is for the - same workspace — so a stale file from another workspace is ignored rather than misapplied. - - This is the single local managed config: ``ucode setup`` authors it here, ``ucode publish`` - publishes it, and a launch refreshes it from the workspace. The admin-authored draft and the - pulled copy share one file because the workspace is the source of truth — to keep a draft, - publish it with ``ucode publish``. - """ +def _load_slot(workspace: str | None, slot: str) -> dict | None: + """Return ``workspace``'s ``published`` or ``draft`` config, or None if absent.""" if not workspace: return None - data = config_io.read_json_safe(MANAGED_STATE_PATH) - if data.get("workspace") != workspace: - return None - config = data.get("config") + entry = _migrated_workspaces(config_io.read_json_safe(MANAGED_STATE_PATH)).get(workspace) or {} + config = entry.get(slot) return config if isinstance(config, dict) else None +def load_published_config(workspace: str | None) -> dict | None: + """Load the launch-fetched ``published`` snapshot for ``workspace``, or None if absent. + + This is what the launch path overlays (:func:`ucode.managed_resolve.resolve_state`) and what + ``ucode status`` reports — the admin-defined config a developer actually runs under. A stale + snapshot from a different workspace is ignored rather than misapplied. + """ + return _load_slot(workspace, "published") + + +def load_draft_config(workspace: str | None) -> dict | None: + """Load the admin's locally authored, unpublished ``draft`` for ``workspace``, or None. + + Only ``ucode configure`` (admin authoring) writes this, and only ``ucode export`` / ``ucode + publish`` read it — never the launch path. A developer who has only ever fetched a published + snapshot has no draft, which is why export/publish are draft-only rather than falling back to the + fetched copy: a cached publication is not authored work. + """ + return _load_slot(workspace, "draft") + + def managed_state_workspace() -> str | None: - """The workspace the on-disk managed config was authored/pulled for, or None when there is none. + """The sole workspace recorded in ``managed-state.json``, or None when absent/ambiguous. - Lets a caller that has no workspace in local ucode state (e.g. ``ucode setup --show`` before - ``ucode configure``) still find the manifest on disk and report which workspace it belongs to. + Lets a caller with no workspace in local ucode state still find the config on disk. With several + workspaces recorded the answer is ambiguous, so it returns None and the caller reports that a + workspace must be selected first. """ - workspace = config_io.read_json_safe(MANAGED_STATE_PATH).get("workspace") - return workspace if isinstance(workspace, str) and workspace else None + workspaces = _migrated_workspaces(config_io.read_json_safe(MANAGED_STATE_PATH)) + return next(iter(workspaces)) if len(workspaces) == 1 else None + + +def save_managed_state(workspace: str, config: dict) -> None: + """Deprecated alias for :func:`save_published_config`, kept while callers migrate to the slots.""" + save_published_config(workspace, config) + + +def load_managed_state(workspace: str | None) -> dict | None: + """Deprecated alias for :func:`load_published_config`, kept while callers migrate to the slots.""" + return load_published_config(workspace) def refresh_managed_config(state: dict) -> tuple[dict | None, bool]: diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 05c179fc..f6f3c0a3 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -13,11 +13,15 @@ import ucode.managed_config as mc_mod from ucode.managed_config import ( get_managed_config, + load_draft_config, load_managed_state, + load_published_config, managed_state_workspace, normalize_managed_config, refresh_managed_config, + save_draft_config, save_managed_state, + save_published_config, ) from ucode.managed_setup import serialize_managed_config @@ -246,7 +250,9 @@ def test_workspace_is_none_when_absent(self, _managed_path): def test_dry_run_writes_nothing(self, _managed_path, monkeypatch): # Under --dry-run the config writers print instead of touching disk, so a launch that # dry-runs an admin's authored draft never overwrites it. - monkeypatch.setattr(config_io_mod, "is_dry_run", lambda: True) + # Patch the flag itself rather than `is_dry_run`: the shared JSON writer reads the module + # global directly. + monkeypatch.setattr(config_io_mod, "_dry_run", True) save_managed_state("https://ws.example.com", {"default_agent": "claude"}) assert not _managed_path.exists() @@ -268,6 +274,83 @@ def test_loaded_config_serializes_to_a_json_encodable_payload(self, _managed_pat assert json.loads(json.dumps(serialize_managed_config(loaded))) +class TestDraftPublishedSeparation: + """The draft (admin-authored, unpublished) and published (fetched) slots are kept apart.""" + + @pytest.fixture(autouse=True) + def _managed_path(self, tmp_path, monkeypatch): + path = tmp_path / ".ucode" / "managed-state.json" + monkeypatch.setattr(mc_mod, "MANAGED_STATE_PATH", path) + return path + + def test_slots_are_independent(self, _managed_path): + ws = "https://ws.example.com" + save_draft_config(ws, {"default_agent": "claude"}) + save_published_config(ws, {"default_agent": "codex"}) + assert load_draft_config(ws) == {"default_agent": "claude"} + assert load_published_config(ws) == {"default_agent": "codex"} + + def test_saving_published_preserves_the_draft(self, _managed_path): + ws = "https://ws.example.com" + save_draft_config(ws, {"default_agent": "claude"}) + save_published_config(ws, {"default_agent": "codex"}) + save_published_config(ws, {"default_agent": "gemini"}) + assert load_draft_config(ws) == {"default_agent": "claude"} + + def test_refresh_never_clobbers_a_draft(self, _managed_path, monkeypatch): + ws = "https://ws.example.com" + save_draft_config(ws, {"default_agent": "claude", "enabled_agents": {"claude": {}}}) + monkeypatch.setattr(mc_mod, "get_databricks_token", lambda w, p: "tok") + monkeypatch.setattr( + mc_mod, "get_managed_config", lambda w, tok: ({"default_agent": "codex"}, None) + ) + result, _ = refresh_managed_config({"workspace": ws}) + assert result == {"default_agent": "codex"} + assert load_published_config(ws) == {"default_agent": "codex"} + assert load_draft_config(ws) == { + "default_agent": "claude", + "enabled_agents": {"claude": {}}, + } + + def test_refresh_of_one_workspace_keeps_another_workspaces_draft( + self, _managed_path, monkeypatch + ): + ws_a, ws_b = "https://a.example.com", "https://b.example.com" + save_draft_config(ws_a, {"default_agent": "claude"}) + monkeypatch.setattr(mc_mod, "get_databricks_token", lambda w, p: "tok") + monkeypatch.setattr( + mc_mod, "get_managed_config", lambda w, tok: ({"default_agent": "codex"}, None) + ) + refresh_managed_config({"workspace": ws_b}) + assert load_draft_config(ws_a) == {"default_agent": "claude"} + + def test_a_failed_write_leaves_the_previous_state_intact(self, _managed_path, monkeypatch): + ws = "https://ws.example.com" + save_draft_config(ws, {"default_agent": "claude"}) + + def partial_write(path, payload): + path.write_text('{"version": 2, "workspa', encoding="utf-8") + raise RuntimeError("disk full") + + monkeypatch.setattr(mc_mod.config_io, "write_json_file", partial_write) + with pytest.raises(RuntimeError): + save_published_config(ws, {"default_agent": "codex"}) + assert load_draft_config(ws) == {"default_agent": "claude"} + + def test_legacy_file_migrates_into_the_published_slot_with_a_backup(self, _managed_path): + ws = "https://ws.example.com" + _managed_path.parent.mkdir(parents=True, exist_ok=True) + legacy = {"workspace": ws, "config": {"default_agent": "claude"}} + _managed_path.write_text(json.dumps(legacy), encoding="utf-8") + assert load_published_config(ws) == {"default_agent": "claude"} + assert load_draft_config(ws) is None + backup = _managed_path.with_suffix(_managed_path.suffix + ".pre-v2.bak") + assert not backup.exists() + save_published_config(ws, {"default_agent": "codex"}) + assert backup.exists() + assert json.loads(backup.read_text()) == legacy + + class TestFetchClient: """fetch_managed_coding_agent_configs lives in databricks.py; test its response parsing."""