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 <no-reply@databricks.com>

Author: david-siqi-liu

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.

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 ``<path>.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]:

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."""