From 9642d7aa664c587982d3694ddfd50843e4e9d1f0 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Wed, 16 Sep 2026 21:01:54 +0000 Subject: [PATCH] Make Claude picker ownership crash-safe --- src/ucode/agents/__init__.py | 12 +- src/ucode/agents/claude.py | 1071 +++++++++++++++++++-- src/ucode/agents/codex.py | 19 +- src/ucode/cli.py | 14 +- src/ucode/config_io.py | 18 +- src/ucode/managed_files.py | 272 +++++- tests/conftest.py | 4 + tests/test_agent_claude.py | 1720 +++++++++++++++++++++++++++++++++- tests/test_agent_codex.py | 3 +- tests/test_agents_init.py | 31 + tests/test_cli.py | 119 ++- tests/test_config_io.py | 36 + tests/test_managed_files.py | 401 ++++++++ 13 files changed, 3535 insertions(+), 185 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 0ed36e066..28fa0e2b2 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -25,7 +25,8 @@ ) from ucode.managed_config import refresh_managed_config from ucode.managed_files import managed_write_batch -from ucode.state import get_provider_service, load_state, save_state +from ucode.state import get_provider_service, load_state +from ucode.state import save_state as save_state from ucode.telemetry import agent_version from ucode.ui import ( print_err, @@ -519,11 +520,9 @@ def configure_single_tool(tool: str, state: dict) -> dict: raise RuntimeError( f"{TOOL_SPECS[tool]['display']} is not available on this workspace.{detail}" ) + state["available_tools"] = list(set((state.get("available_tools") or []) + [tool])) with managed_write_batch(_managed_settings_displays([tool])): state = _configure_one(tool, state, provider) - available_tools = list(set((state.get("available_tools") or []) + [tool])) - state["available_tools"] = available_tools - save_state(state) return state @@ -562,11 +561,10 @@ def configure_selected_tools( """ with managed_write_batch(_managed_settings_displays(tools)): for tool in tools: + existing = state.get("available_tools") or [] + state["available_tools"] = sorted(set(existing) | {tool}) state = _configure_one(tool, state, get_provider_service(state, tool)) - existing = state.get("available_tools") or [] - state["available_tools"] = sorted(set(existing) | set(tools)) - save_state(state) if install_ai_tools: install_databricks_ai_tools_for_agents(tools, state) return state diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 1e1cd358c..951be3771 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -3,17 +3,20 @@ from __future__ import annotations import copy +import hashlib import json import os import re +import secrets import shutil import signal import socket import subprocess import threading -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from pathlib import Path -from typing import cast +from typing import Protocol, cast from ucode import gateway_proxy from ucode.config_io import ( @@ -21,6 +24,7 @@ ToolSpec, backup_existing_file, deep_merge_dict, + is_dry_run, read_json_safe, write_json_file, ) @@ -48,11 +52,14 @@ managed_file_conflicts, managed_file_is_verified, managed_file_status, + managed_last_applied_paths, managed_writes_allowed, mark_managed_file_verified, read_managed_file, reconcile_managed_file, + restore_unchanged_managed_paths, revert_managed_file, + write_private_json_file, ) from ucode.mcp_oauth import CLAUDE_CODE_OAUTH_CLIENT_ID, MCP_OAUTH_CALLBACK_PORT from ucode.smart_routing import v2 as smart_routing_v2 @@ -60,13 +67,29 @@ remove_smart_routing_hooks, sync_smart_routing_hooks, ) -from ucode.state import MANAGED_OVERLAY_KEY, is_tool_managed, mark_tool_managed, save_state +from ucode.state import ( + MANAGED_OVERLAY_KEY, + developer_state_from_resolved, + is_tool_managed, + load_full_state, + load_state, + mark_tool_managed, + save_state, +) from ucode.telemetry import agent_version, ug_version from ucode.tracing import tracing_env from ucode.ui import print_note, print_success, print_warning from .args import LaunchOptions, has_explicit_model_arg + +class _WindowsFileLockApi(Protocol): + LK_LOCK: int + LK_UNLCK: int + + def locking(self, fd: int, mode: int, nbytes: int) -> None: ... + + GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" # If set, Claude Code launches in headless mode instead of the interactive login flow. CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR = "CLAUDE_CODE_OAUTH_TOKEN" @@ -76,7 +99,10 @@ # The default model is stored in Claude's default user settings, not the ucode settings. CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" +CLAUDE_PICKER_MANAGEMENT_PATH = APP_DIR / "claude-picker-management.json" +CLAUDE_PICKER_MANAGEMENT_VERSION = 2 WEB_SEARCH_MCP_STATE_KEY = "claude_web_search_mcp" +WEB_SEARCH_MCP_GENERATION_KEY = "claude_web_search_generation" MINIMUM_CLAUDE_VERSION = (2, 1, 248) MINIMUM_CLAUDE_VERSION_TEXT = "2.1.248" @@ -88,6 +114,11 @@ "backup_path": CLAUDE_BACKUP_PATH, } +_MISSING_PICKER_VALUE = object() +_picker_thread_lock = threading.RLock() +_picker_lock_state = threading.local() +_web_search_registration_thread_lock = threading.Lock() + # Retained only to identify and remove state written by the legacy persisted opt-in. SMART_ROUTING_STATE_KEY = smart_routing_v2.LEGACY_STATE_KEY @@ -203,7 +234,6 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: # settings file on every launch so stale values never linger. CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",) CLAUDE_MANAGED_PICKER_KEYS = ("availableModels", "enforceAvailableModels", "modelPicker") -CLAUDE_PRUNED_PICKER_KEYS = ("availableModels", "enforceAvailableModels") ANTHROPIC_CUSTOM_HEADERS_ENV_KEY = "ANTHROPIC_CUSTOM_HEADERS" CLAUDE_MANAGED_CUSTOM_HEADER_NAMES = frozenset( { @@ -292,12 +322,85 @@ def managed_settings_status(state: dict) -> tuple[Path | None, str, str]: def revert_managed_settings() -> str: - return revert_managed_file( - "claude", - display="Claude Code", - parser=_parse_managed_settings, - dumper=_dump_managed_settings, - ) + path = _managed_settings_path() + if path is None: + return revert_managed_file( + "claude", + display="Claude Code", + parser=_parse_managed_settings, + dumper=_dump_managed_settings, + ) + with _picker_management_lock(): + current_text = read_managed_file(path) + current = _parse_managed_settings(current_text) if current_text is not None else {} + lease, recovery, phase = _recover_picker_transition("managed", path, current) + if phase == "applying" and recovery in {"target", "drift"}: + raise RuntimeError( + "Cannot safely revert Claude Code managed settings while a picker update is " + "incomplete. Re-run configuration to repair its metadata, then run `ucode " + "revert` again." + ) + if not lease and recovery is None: + lease = _legacy_managed_picker_management(path, current) + + transitioned = copy.deepcopy(current) + if lease: + _transition_private_picker(transitioned, {}, lease) + target = _picker_group(transitioned) + journaled = bool(lease or recovery is not None) + if journaled: + _begin_picker_transition( + "managed", + path, + lease, + {}, + _picker_group(current), + target, + phase="reverting", + ) + acquisition_lease = [ + { + "path": [key], + "baseline_exists": entry["original_exists"], + **( + {"baseline": copy.deepcopy(entry["original"])} + if entry["original_exists"] + else {} + ), + "applied": copy.deepcopy(entry["last_applied"]), + } + for key, entry in lease.items() + ] + + def mark_revert_written() -> None: + written_text = read_managed_file(path) + written = _parse_managed_settings(written_text) if written_text is not None else {} + if not _picker_group_matches(written, target): + raise RuntimeError(f"Could not verify restored Claude picker settings at {path}.") + _mark_picker_revert_written("managed", path) + + result = revert_managed_file( + "claude", + display="Claude Code", + parser=_parse_managed_settings, + dumper=_dump_managed_settings, + acquisition_lease=acquisition_lease, + excluded_owned_paths=[[key] for key in CLAUDE_MANAGED_PICKER_KEYS], + before_backup_delete=mark_revert_written if journaled else None, + ) + written_text = read_managed_file(path) + written = _parse_managed_settings(written_text) if written_text is not None else {} + if journaled and not _picker_group_matches(written, target): + raise RuntimeError(f"Could not verify restored Claude picker settings at {path}.") + if journaled: + _save_managed_picker_management(path, {}) + return result + + +def revert_settings(state: dict) -> tuple[str, bool]: + """Revert both Claude settings scopes as one serialized transaction.""" + with _picker_management_lock(): + return revert_managed_settings(), revert_private_settings(state) def _managed_relayed_conflicts(path: Path) -> list[str]: @@ -491,7 +594,7 @@ def render_overlay( overlay["permissions"] = {"deny": ["WebSearch"]} keys.append(["permissions", "deny"]) - if static_models and not provider and not relayed: + if static_models and not provider and not parent_schema and not relayed: overlay["availableModels"] = list(static_models) overlay["enforceAvailableModels"] = True overlay["modelPicker"] = { @@ -516,6 +619,693 @@ def _picker_label(model: str) -> str: return model.removeprefix("system.ai.") +@contextmanager +def _picker_management_lock() -> Iterator[None]: + """Serialize sidecar and settings transitions that share picker ownership.""" + if is_dry_run(): + yield + return + + with _picker_thread_lock: + depth = getattr(_picker_lock_state, "depth", 0) + if depth: + _picker_lock_state.depth = depth + 1 + try: + yield + finally: + _picker_lock_state.depth = depth + return + _picker_lock_state.depth = 1 + try: + with _picker_process_lock(): + yield + finally: + _picker_lock_state.depth = 0 + + +@contextmanager +def _picker_process_lock() -> Iterator[None]: + lock_path = CLAUDE_PICKER_MANAGEMENT_PATH.with_name( + f"{CLAUDE_PICKER_MANAGEMENT_PATH.name}.lock" + ) + with _process_file_lock(lock_path): + yield + + +@contextmanager +def _process_file_lock(lock_path: Path) -> Iterator[None]: + """Hold one cross-platform advisory byte lock at an explicit private path.""" + try: + lock_path.parent.mkdir(parents=True, exist_ok=True) + lock_file = lock_path.open("a+b") + except OSError as exc: + raise RuntimeError(f"Cannot lock Claude picker metadata at {lock_path}: {exc}") from exc + with lock_file: + try: + os.chmod(lock_path, 0o600) + _lock_picker_file(lock_file) + except OSError as exc: + raise RuntimeError(f"Cannot lock Claude picker metadata at {lock_path}: {exc}") from exc + try: + yield + finally: + _unlock_picker_file(lock_file) + + +@contextmanager +def _web_search_registration_lock() -> Iterator[None]: + lock_path = CLAUDE_PICKER_MANAGEMENT_PATH.with_name("claude-web-search-registration.lock") + with _web_search_registration_thread_lock: + with _process_file_lock(lock_path): + yield + + +def _lock_picker_file(lock_file) -> None: + if current_os() is OS.WINDOWS: + import msvcrt + + windows_lock = cast("_WindowsFileLockApi", msvcrt) + lock_file.seek(0, os.SEEK_END) + if lock_file.tell() == 0: + lock_file.write(b"\0") + lock_file.flush() + lock_file.seek(0) + windows_lock.locking(lock_file.fileno(), windows_lock.LK_LOCK, 1) + return + import fcntl + + fcntl.flock(lock_file, fcntl.LOCK_EX) + + +def _unlock_picker_file(lock_file) -> None: + if current_os() is OS.WINDOWS: + import msvcrt + + windows_lock = cast("_WindowsFileLockApi", msvcrt) + lock_file.seek(0) + windows_lock.locking(lock_file.fileno(), windows_lock.LK_UNLCK, 1) + return + import fcntl + + fcntl.flock(lock_file, fcntl.LOCK_UN) + + +def _valid_picker_entries(entries: object) -> bool: + if not isinstance(entries, dict) or set(entries) != set(CLAUDE_MANAGED_PICKER_KEYS): + return False + for entry in entries.values(): + if not isinstance(entry, dict) or not isinstance(entry.get("original_exists"), bool): + return False + expected = {"original_exists", "last_applied"} + if entry["original_exists"]: + expected.add("original") + if set(entry) != expected: + return False + return True + + +def _picker_group(settings: dict) -> dict[str, dict]: + group: dict[str, dict] = {} + for key in CLAUDE_MANAGED_PICKER_KEYS: + exists = key in settings + entry: dict = {"exists": exists} + if exists: + entry["value"] = copy.deepcopy(settings[key]) + group[key] = entry + return group + + +def _valid_picker_group(group: object) -> bool: + if not isinstance(group, dict) or set(group) != set(CLAUDE_MANAGED_PICKER_KEYS): + return False + for entry in group.values(): + if not isinstance(entry, dict) or not isinstance(entry.get("exists"), bool): + return False + if set(entry) != ({"exists", "value"} if entry["exists"] else {"exists"}): + return False + return True + + +def _picker_group_matches(settings: dict, group: dict[str, dict]) -> bool: + return _picker_group(settings) == group + + +def _private_document_sha256(settings: dict) -> str: + encoded = json.dumps(settings, sort_keys=True, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _read_private_json_object(path: Path) -> dict: + """Read private Claude transaction input strictly; only a missing file means empty.""" + try: + text = path.read_text(encoding="utf-8") + except FileNotFoundError: + return {} + except OSError as exc: + raise RuntimeError(f"Cannot read Claude settings at {path}: {exc}") from exc + try: + settings = json.loads(text) + except json.JSONDecodeError as exc: + raise RuntimeError(f"Cannot parse Claude settings at {path}: {exc}") from exc + if not isinstance(settings, dict): + raise RuntimeError(f"Claude settings at {path} must contain a JSON object.") + return settings + + +def _load_picker_management() -> dict[str, dict]: + """Load per-scope acquisition leases for ucode's Claude picker fields.""" + path = CLAUDE_PICKER_MANAGEMENT_PATH + if path.is_symlink(): + raise RuntimeError(f"Refusing to read symlinked Claude picker metadata at {path}.") + try: + text = path.read_text(encoding="utf-8") + except FileNotFoundError: + return {} + except OSError as exc: + raise RuntimeError(f"Cannot read Claude picker metadata at {path}: {exc}") from exc + try: + metadata = json.loads(text) + except json.JSONDecodeError as exc: + raise RuntimeError(f"Cannot parse Claude picker metadata at {path}: {exc}") from exc + if ( + not isinstance(metadata, dict) + or metadata.get("version") != CLAUDE_PICKER_MANAGEMENT_VERSION + or not isinstance(metadata.get("leases"), dict) + ): + raise RuntimeError(f"Invalid Claude picker metadata at {path}.") + + leases = metadata["leases"] + if not leases or not set(leases).issubset({"private", "managed"}): + raise RuntimeError(f"Invalid Claude picker metadata at {path}.") + pending_keys = {"phase", "prior", "intended", "before", "target"} + document_proof_keys = { + "before_document_exists", + "before_document_sha256", + "target_document_exists", + "target_document_sha256", + } + for scope, lease in leases.items(): + if not isinstance(lease, dict) or not isinstance(lease.get("path"), str): + raise RuntimeError(f"Invalid Claude picker metadata at {path}.") + pending = lease.get("pending") + if pending is None: + if set(lease) != {"path", "keys"} or not _valid_picker_entries(lease.get("keys")): + raise RuntimeError(f"Invalid Claude picker metadata at {path}.") + continue + if ( + set(lease) != {"path", "pending"} + or not isinstance(pending, dict) + or frozenset(pending) + not in {frozenset(pending_keys), frozenset(pending_keys | document_proof_keys)} + or pending.get("phase") not in {"applying", "reverting", "revert_written"} + or (pending.get("prior") is not None and not _valid_picker_entries(pending["prior"])) + or ( + pending.get("intended") is not None + and not _valid_picker_entries(pending["intended"]) + ) + or not _valid_picker_group(pending.get("before")) + or not _valid_picker_group(pending.get("target")) + or ( + document_proof_keys.issubset(pending) + and ( + scope != "private" + or not isinstance(pending["before_document_exists"], bool) + or not isinstance(pending["target_document_exists"], bool) + or not isinstance(pending["before_document_sha256"], str) + or re.fullmatch(r"[0-9a-f]{64}", pending["before_document_sha256"]) is None + or not isinstance(pending["target_document_sha256"], str) + or re.fullmatch(r"[0-9a-f]{64}", pending["target_document_sha256"]) is None + ) + ) + ): + raise RuntimeError(f"Invalid Claude picker metadata at {path}.") + return copy.deepcopy(leases) + + +def _save_picker_management(leases: dict[str, dict]) -> None: + path = CLAUDE_PICKER_MANAGEMENT_PATH + if leases: + write_private_json_file( + path, + { + "version": CLAUDE_PICKER_MANAGEMENT_VERSION, + "leases": leases, + }, + ) + return + if is_dry_run(): + return + if path.is_symlink(): + raise RuntimeError(f"Refusing to remove symlinked Claude picker metadata at {path}.") + try: + path.unlink(missing_ok=True) + except OSError as exc: + raise RuntimeError(f"Cannot remove Claude picker metadata at {path}: {exc}") from exc + + +def _save_picker_lease(scope: str, path: Path, entries: dict[str, dict]) -> None: + if entries and not _valid_picker_entries(entries): + raise RuntimeError("Invalid Claude picker lease.") + leases = _load_picker_management() + if entries: + leases[scope] = {"path": str(path), "keys": copy.deepcopy(entries)} + else: + leases.pop(scope, None) + _save_picker_management(leases) + + +def _begin_picker_transition( + scope: str, + path: Path, + previous: dict[str, dict], + intended: dict[str, dict], + before: dict[str, dict], + target: dict[str, dict], + *, + phase: str = "applying", + document_transition: tuple[bool, dict, bool, dict] | None = None, +) -> None: + if ( + (previous and not _valid_picker_entries(previous)) + or (intended and not _valid_picker_entries(intended)) + or not _valid_picker_group(before) + or not _valid_picker_group(target) + or phase not in {"applying", "reverting"} + ): + raise RuntimeError("Invalid Claude picker transition.") + pending = { + "phase": phase, + "prior": copy.deepcopy(previous) or None, + "intended": copy.deepcopy(intended) or None, + "before": copy.deepcopy(before), + "target": copy.deepcopy(target), + } + if document_transition is not None: + before_exists, before_document, target_exists, target_document = document_transition + pending.update( + { + "before_document_exists": before_exists, + "before_document_sha256": _private_document_sha256(before_document), + "target_document_exists": target_exists, + "target_document_sha256": _private_document_sha256(target_document), + } + ) + leases = _load_picker_management() + leases[scope] = { + "path": str(path), + "pending": pending, + } + _save_picker_management(leases) + + +def _recover_picker_transition( + scope: str, path: Path, settings: dict +) -> tuple[dict[str, dict], str | None, str | None]: + """Resolve an interrupted transition without discarding its retry proof.""" + leases = _load_picker_management() + lease = leases.get(scope) + if lease is None: + return {}, None, None + if lease["path"] != str(path): + raise RuntimeError( + f"Claude picker metadata for {scope} targets {lease['path']}, not {path}." + ) + pending = lease.get("pending") + if pending is None: + return copy.deepcopy(lease["keys"]), None, None + + current_exists = path.exists() + document_sha256 = _private_document_sha256(settings) + has_document_proof = "target_document_sha256" in pending + target_matches = ( + current_exists == pending["target_document_exists"] + and document_sha256 == pending["target_document_sha256"] + if has_document_proof + else _picker_group_matches(settings, pending["target"]) + ) + before_matches = ( + current_exists == pending["before_document_exists"] + and document_sha256 == pending["before_document_sha256"] + if has_document_proof + else _picker_group_matches(settings, pending["before"]) + ) + if target_matches: + recovered = pending["intended"] + recovery = "target" + elif pending["phase"] == "revert_written": + # Once the revert target was durably verified, every later non-target value is external + # post-revert drift, even if it exactly recreates the old pre-revert settings. + recovered = pending["intended"] + recovery = "drift" + elif before_matches: + recovered = pending["prior"] + recovery = "before" + else: + # The complete group is neither the pre-write nor intended state. Treat it as external + # drift after the attempted transition; the next transition will preserve or rebase it. + recovered = pending["intended"] + recovery = "drift" + return copy.deepcopy(recovered or {}), recovery, pending["phase"] + + +def _save_private_picker_management(entries: dict[str, dict]) -> None: + _save_picker_lease("private", CLAUDE_SETTINGS_PATH, entries) + + +def _save_managed_picker_management(path: Path, entries: dict[str, dict]) -> None: + _save_picker_lease("managed", path, entries) + + +def _mark_picker_revert_written(scope: str, path: Path) -> None: + """Durably record that one revert target was written and fully verified.""" + leases = _load_picker_management() + lease = leases.get(scope) + if not isinstance(lease, dict): + raise RuntimeError(f"Claude {scope} revert journal is not ready to commit.") + pending = lease.get("pending") + if ( + not isinstance(pending, dict) + or lease.get("path") != str(path) + or pending.get("phase") != "reverting" + ): + raise RuntimeError(f"Claude {scope} revert journal is not ready to commit.") + pending["phase"] = "revert_written" + _save_picker_management(leases) + + +def _apply_picker_group(settings: dict, group: dict[str, dict]) -> None: + for key, entry in group.items(): + if entry["exists"]: + settings[key] = copy.deepcopy(entry["value"]) + else: + settings.pop(key, None) + + +def _state_records_legacy_picker_ownership(state: dict) -> bool: + states = [state] + workspaces = load_full_state().get("workspaces") + if isinstance(workspaces, dict): + states.extend(entry for entry in workspaces.values() if isinstance(entry, dict)) + for candidate in states: + managed_configs = candidate.get("managed_configs") + claude_management = ( + managed_configs.get("claude") if isinstance(managed_configs, dict) else None + ) + paths = claude_management.get("keys") if isinstance(claude_management, dict) else None + if isinstance(paths, list) and any([key] in paths for key in CLAUDE_MANAGED_PICKER_KEYS): + return True + return False + + +def _legacy_private_picker_management(settings: dict) -> dict[str, dict]: + """Bootstrap pre-sidecar ownership only from an all-three managed snapshot proof.""" + managed_path = _managed_settings_path() + picker_paths = [[key] for key in CLAUDE_MANAGED_PICKER_KEYS] + if managed_path is None: + return {} + last_managed, recorded_paths = managed_last_applied_paths( + "claude", managed_path, picker_paths, parser=_parse_managed_settings + ) + if {tuple(path) for path in recorded_paths} != {tuple(path) for path in picker_paths}: + return {} + private_matches_last = all( + settings.get(key, _MISSING_PICKER_VALUE) == last_managed.get(key, _MISSING_PICKER_VALUE) + for key in CLAUDE_MANAGED_PICKER_KEYS + ) + if not private_matches_last: + return {} + original = _read_private_json_object(CLAUDE_BACKUP_PATH) + entries: dict[str, dict] = {} + for key in CLAUDE_MANAGED_PICKER_KEYS: + if key not in settings: + continue + entry = { + "original_exists": key in original, + "last_applied": copy.deepcopy(settings[key]), + } + if key in original: + entry["original"] = copy.deepcopy(original[key]) + entries[key] = entry + return entries if _valid_picker_entries(entries) else {} + + +def _legacy_managed_picker_management(path: Path, settings: dict) -> dict[str, dict]: + """Bootstrap one all-three lease from integrity-checked legacy manifest snapshots.""" + picker_paths = [[key] for key in CLAUDE_MANAGED_PICKER_KEYS] + applied = copy.deepcopy(settings) + restored, restored_paths = restore_unchanged_managed_paths( + "claude", + path, + copy.deepcopy(settings), + picker_paths, + parser=_parse_managed_settings, + ) + if {tuple(candidate) for candidate in restored_paths} != { + tuple(candidate) for candidate in picker_paths + }: + return {} + if not all(key in applied for key in CLAUDE_MANAGED_PICKER_KEYS): + return {} + entries: dict[str, dict] = {} + for key in CLAUDE_MANAGED_PICKER_KEYS: + entry = { + "original_exists": key in restored, + "last_applied": copy.deepcopy(applied[key]), + } + if key in restored: + entry["original"] = copy.deepcopy(restored[key]) + entries[key] = entry + return entries + + +def _transition_private_picker( + settings: dict, overlay: dict, previous: dict[str, dict] +) -> dict[str, dict]: + """Restore or rebase private picker fields before applying the current overlay.""" + static_picker_active = all(key in overlay for key in CLAUDE_MANAGED_PICKER_KEYS) + if not static_picker_active: + unchanged = bool(previous) and all( + settings.get(key, _MISSING_PICKER_VALUE) == entry["last_applied"] + for key, entry in previous.items() + ) + if unchanged: + for key, entry in previous.items(): + if entry["original_exists"]: + settings[key] = copy.deepcopy(entry["original"]) + else: + settings.pop(key, None) + return {} + + current_management: dict[str, dict] = {} + previous_unchanged = bool(previous) and all( + settings.get(key, _MISSING_PICKER_VALUE) == entry["last_applied"] + for key, entry in previous.items() + ) + for key in CLAUDE_MANAGED_PICKER_KEYS: + current = settings.get(key, _MISSING_PICKER_VALUE) + previous_entry = previous.get(key) + if previous_unchanged and previous_entry is not None: + original_exists = previous_entry["original_exists"] + original = previous_entry.get("original") + else: + original_exists = current is not _MISSING_PICKER_VALUE + original = current + entry = { + "original_exists": original_exists, + "last_applied": copy.deepcopy(overlay[key]), + } + if original_exists: + entry["original"] = copy.deepcopy(original) + current_management[key] = entry + return current_management + + +def _reconcile_private_settings( + state: dict, overlay: dict, compose: Callable[[dict], dict] +) -> None: + """Apply private settings with a crash-safe picker lease; caller holds the lock.""" + settings = _read_private_json_object(CLAUDE_SETTINGS_PATH) + settings_existed = CLAUDE_SETTINGS_PATH.exists() + settings_before = copy.deepcopy(settings) + before = _picker_group(settings) + previous, recovery, phase = _recover_picker_transition( + "private", CLAUDE_SETTINGS_PATH, settings + ) + if phase == "reverting" and recovery == "drift": + raise RuntimeError( + "Cannot safely configure Claude settings while an interrupted private revert has " + "unverified external changes. Restore the file to its pre-revert state or remove " + f"{CLAUDE_SETTINGS_PATH}, then run `ug revert` again." + ) + legacy_management = ( + not previous and recovery is None and _state_records_legacy_picker_ownership(state) + ) + + # Back up only a file that predates ucode's management of the tool. A re-configure would + # otherwise snapshot ucode's generated file, and revert would restore that snapshot. + if not is_tool_managed(state, "claude") and not previous and not legacy_management: + backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) + elif phase == "revert_written" and recovery == "drift": + # The old backup may still exist if the verified revert crashed before cleanup. Replace + # it atomically so this complete post-revert document becomes the next lease baseline. + if settings_existed: + write_json_file(CLAUDE_BACKUP_PATH, settings_before) + else: + _remove_private_backup() + elif ( + recovery == "target" + and phase in {"reverting", "revert_written"} + and not CLAUDE_BACKUP_PATH.exists() + ): + # The prior revert restored and verified this complete document before its backup was + # removed, but failed to clear the journal. Treat it as the baseline of this new lease. + backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) + if legacy_management: + previous = _legacy_private_picker_management(settings) + + intended = _transition_private_picker(settings, overlay, previous) + desired = compose(settings) + target = _picker_group(desired) + if not previous and intended and before == target: + # Matching IT policy was not changed by ucode, so it must not become leased merely + # because the requested static policy happens to have the same values. + intended = {} + transition = bool(recovery is not None or previous or intended) + if transition: + _begin_picker_transition( + "private", + CLAUDE_SETTINGS_PATH, + previous, + intended, + before, + target, + document_transition=(settings_existed, settings_before, True, desired), + ) + write_json_file(CLAUDE_SETTINGS_PATH, desired) + if is_dry_run(): + return + written = _read_private_json_object(CLAUDE_SETTINGS_PATH) + if written != desired: + raise RuntimeError(f"Could not verify Claude settings at {CLAUDE_SETTINGS_PATH}.") + if transition: + _save_private_picker_management(intended) + + +def _remove_private_backup() -> None: + try: + CLAUDE_BACKUP_PATH.unlink(missing_ok=True) + except OSError as exc: + raise RuntimeError( + f"Failed to remove Claude settings backup at {CLAUDE_BACKUP_PATH}" + ) from exc + + +def revert_private_settings(state: dict) -> bool: + """Restore private settings atomically while honoring the picker acquisition lease.""" + with _picker_management_lock(): + current = _read_private_json_object(CLAUDE_SETTINGS_PATH) + lease, recovery, phase = _recover_picker_transition( + "private", CLAUDE_SETTINGS_PATH, current + ) + legacy_claim = _state_records_legacy_picker_ownership(state) + + if phase in {"reverting", "revert_written"} and recovery == "target": + _remove_private_backup() + _save_private_picker_management({}) + return True + if phase == "reverting" and recovery == "drift": + raise RuntimeError( + "Cannot safely resume Claude settings revert because the file changed after the " + "revert was journaled but before its target was verified. Restore the file to " + f"its pre-revert state or remove {CLAUDE_SETTINGS_PATH}, then retry." + ) + if not lease and recovery is None and legacy_claim: + lease = _legacy_private_picker_management(current) + + backup_exists = CLAUDE_BACKUP_PATH.exists() + managed_configs = state.get("managed_configs") + state_managed = isinstance(managed_configs, dict) and bool(managed_configs.get("claude")) + if not (backup_exists or state_managed or lease or recovery is not None or legacy_claim): + return False + + retry_preserves_current = phase == "revert_written" and recovery == "drift" + if retry_preserves_current: + desired = copy.deepcopy(current) + elif backup_exists: + desired = _read_private_json_object(CLAUDE_BACKUP_PATH) + else: + desired = {} + preserve_current_group = phase == "reverting" and recovery == "drift" + if lease: + transitioned = copy.deepcopy(current) + _transition_private_picker(transitioned, {}, lease) + _apply_picker_group(desired, _picker_group(transitioned)) + elif legacy_claim or preserve_current_group: + _apply_picker_group(desired, _picker_group(current)) + + desired_exists = ( + CLAUDE_SETTINGS_PATH.exists() + if retry_preserves_current + else backup_exists or bool(desired) + ) + before = _picker_group(current) + target = _picker_group(desired if desired_exists else {}) + _begin_picker_transition( + "private", + CLAUDE_SETTINGS_PATH, + lease, + {}, + before, + target, + phase="reverting", + document_transition=( + CLAUDE_SETTINGS_PATH.exists(), + current, + desired_exists, + desired if desired_exists else {}, + ), + ) + + if desired_exists: + write_json_file(CLAUDE_SETTINGS_PATH, desired) + if _read_private_json_object(CLAUDE_SETTINGS_PATH) != desired: + raise RuntimeError( + f"Could not verify restored Claude settings at {CLAUDE_SETTINGS_PATH}." + ) + else: + try: + CLAUDE_SETTINGS_PATH.unlink(missing_ok=True) + except OSError as exc: + raise RuntimeError( + f"Failed to remove Claude settings at {CLAUDE_SETTINGS_PATH}" + ) from exc + if CLAUDE_SETTINGS_PATH.exists(): + raise RuntimeError( + f"Could not verify removal of Claude settings at {CLAUDE_SETTINGS_PATH}." + ) + + _mark_picker_revert_written("private", CLAUDE_SETTINGS_PATH) + _remove_private_backup() + _save_private_picker_management({}) + return True + + +def private_settings_are_globally_managed(state: dict | None = None) -> bool: + """Whether global metadata records picker fields written into ucode's Claude settings.""" + with _picker_management_lock(): + settings = _read_private_json_object(CLAUDE_SETTINGS_PATH) + lease, _recovery, _phase = _recover_picker_transition( + "private", CLAUDE_SETTINGS_PATH, settings + ) + return bool(lease) or _state_records_legacy_picker_ownership(state or {}) + + +def clear_private_picker_management() -> None: + """Forget private picker ownership after the corresponding settings file was reverted.""" + with _picker_management_lock(): + _save_private_picker_management({}) + + def _maybe_add_1m_suffix(model: str) -> str: if model.endswith("[1m]"): return model @@ -695,6 +1485,58 @@ def _web_search_mcp_is_current(state: dict, entry: dict) -> bool: return isinstance(servers, dict) and servers.get(WEB_SEARCH_MCP_NAME) == entry +def _claude_state_generation(state: dict) -> tuple: + """Return persisted developer fields that identify one Claude settings generation.""" + managed_configs = state.get("managed_configs") + claude_management = managed_configs.get("claude") if isinstance(managed_configs, dict) else None + fingerprints = state.get("managed_file_fingerprints") + claude_fingerprint = fingerprints.get("claude") if isinstance(fingerprints, dict) else None + return ( + state.get("workspace"), + state.get("profile"), + state.get("web_search_model"), + state.get("codex_models"), + state.get("claude_static_models"), + state.get("provider_services"), + claude_management, + claude_fingerprint, + ) + + +def _register_web_search_for_current_generation(state: dict, entry: dict) -> dict: + """Register only if this generation is still current, then CAS its cache.""" + generation = state.get(WEB_SEARCH_MCP_GENERATION_KEY) + if not isinstance(generation, str) or not generation: + return load_state() + persisted_generation = _claude_state_generation(developer_state_from_resolved(state)) + + with _web_search_registration_lock(): + with _picker_management_lock(): + latest = load_state() + if ( + latest.get(WEB_SEARCH_MCP_GENERATION_KEY) != generation + or _claude_state_generation(latest) != persisted_generation + ): + return latest + already_current = _web_search_mcp_is_current(latest, entry) + + model = _resolve_web_search_model(state) + registration_success = bool(model) and ( + already_current + or _register_web_search_mcp(state["workspace"], model, state.get("profile")) + ) + with _picker_management_lock(): + latest = load_state() + if ( + registration_success + and latest.get(WEB_SEARCH_MCP_GENERATION_KEY) == generation + and _claude_state_generation(latest) == persisted_generation + ): + latest[WEB_SEARCH_MCP_STATE_KEY] = entry + save_state(latest) + return latest + + def _unregister_web_search_mcp() -> None: """Remove the web_search MCP server from all scopes. Used by revert.""" for scope in MCP_CLEANUP_SCOPES: @@ -732,12 +1574,12 @@ def write_tool_config( coding_agent_config_defaults: dict[str, str] | None = None, parent_schema: str | None = None, ) -> dict: - # Back up only a file that predates ucode's management of the tool. A - # re-configure would otherwise snapshot ucode's own generated file, and - # revert would restore that snapshot instead of deleting the file. - if not is_tool_managed(state, "claude"): - backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) web_search_model = _resolve_web_search_model(state) + web_search_entry = ( + _web_search_mcp_entry(state["workspace"], web_search_model, state.get("profile")) + if web_search_model + else None + ) # Relayed inference points at a local refresh proxy; its loopback base URL is # recorded in state so launch starts the proxy on the matching port. relayed_base_url = relayed_proxy_base_url(state) if relayed else None @@ -773,7 +1615,8 @@ def write_tool_config( "to install the Claude Stop hook — traces won't be emitted. Re-run " "`ucode configure tracing`." ) - managed_file_keys = list(managed_keys) + picker_paths = [[key] for key in CLAUDE_MANAGED_PICKER_KEYS] + managed_file_keys = [path for path in managed_keys if path not in picker_paths] for path in ( [["env", key] for key in CLAUDE_MANAGED_MODEL_ENV_KEYS] + [["env", key] for key in CLAUDE_CONDITIONAL_ENV_KEYS] @@ -862,44 +1705,42 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: sync_smart_routing_hooks(merged, state, enabled=False) return merged - write_json_file( - CLAUDE_SETTINGS_PATH, - _compose(read_json_safe(CLAUDE_SETTINGS_PATH), enforce_model_default_hierarchy=False), - ) - - _reconcile_managed_settings( - state, - lambda base: _compose( - base, - enforce_model_default_hierarchy=provider is None and parent_schema is None, - ), - managed_file_keys, - relayed, - ) - - if web_search_model: - web_search_entry = _web_search_mcp_entry( - state["workspace"], web_search_model, state.get("profile") + with _picker_management_lock(): + _reconcile_private_settings( + state, + overlay, + lambda base: _compose(base, enforce_model_default_hierarchy=False), ) - if not _web_search_mcp_is_current(state, web_search_entry): - # Registration runs multiple `claude mcp` subprocesses and can take several seconds. - registration_success = _register_web_search_mcp( - state["workspace"], web_search_model, state.get("profile") - ) - if registration_success: - state[WEB_SEARCH_MCP_STATE_KEY] = web_search_entry - else: - state.pop(WEB_SEARCH_MCP_STATE_KEY, None) + _reconcile_managed_settings( + state, + lambda base: _compose( + base, + enforce_model_default_hierarchy=provider is None and parent_schema is None, + ), + managed_file_keys, + relayed, + static_picker_active=all(key in overlay for key in CLAUDE_MANAGED_PICKER_KEYS), + ) + if web_search_entry is None: + state.pop(WEB_SEARCH_MCP_STATE_KEY, None) + state.pop(WEB_SEARCH_MCP_GENERATION_KEY, None) + else: + # This opaque token is saved atomically with the settings files. Unlike resolved + # managed values, it survives save_state's managed-overlay stripping and therefore + # remains a stable CAS identity for the deferred MCP registration. + state[WEB_SEARCH_MCP_GENERATION_KEY] = secrets.token_hex(16) + # Persist the state generation under the same lock as both settings scopes. Registration + # runs after releasing the lock and uses a generation check before recording its cache. + if relayed: + state["claude_relayed"] = True + else: + state.pop("claude_relayed", None) + state.pop("relayed_proxy_port", None) + state = mark_tool_managed(state, "claude", managed_keys) + save_state(state) - # Persist relayed mode + proxy port so launch() wires the refresh proxy and - # subscription login; cleared on a non-relayed launch. - if relayed: - state["claude_relayed"] = True - else: - state.pop("claude_relayed", None) - state.pop("relayed_proxy_port", None) - state = mark_tool_managed(state, "claude", managed_keys) - save_state(state) + if web_search_entry is not None: + state = _register_web_search_for_current_generation(state, web_search_entry) return state @@ -956,6 +1797,8 @@ def _reconcile_managed_settings( compose: Callable[[dict], dict], owned_paths: list[list[str]], relayed: bool, + *, + static_picker_active: bool, ) -> None: """Reconcile Claude Code's OS-managed settings so a bare ``claude`` uses the gateway. @@ -963,8 +1806,8 @@ def _reconcile_managed_settings( configuration mirrors ucode's settings there. The same compose operation that produced the private file is applied to the existing managed file, preserving unrelated IT-authored keys. - `ug configure` updates gateway-owned fields in this file, but does not generate or modify - the `modelPicker` object; an existing picker is retained by the merge. + An existing IT-authored picker is retained by the merge. Picker fields that ucode recorded as + managed are removed when the active model source switches away from a static model list. Relayed launches are skipped: they depend on a per-session loopback refresh proxy that only runs during `ucode claude`, so a bare `claude` could not reach the gateway anyway. @@ -993,47 +1836,87 @@ def _reconcile_managed_settings( mark_managed_file_verified(state, "claude", path, scope="relay-compatible") return - current_text = read_managed_file(path) - try: - existing = _parse_managed_settings(current_text) if current_text is not None else {} - except RuntimeError as exc: - raise RuntimeError( - f"Cannot safely update Claude Code managed settings at {path}: {exc}. " - "ucode did not modify the file. Repair it or contact your administrator." - ) from exc - managed_before = copy.deepcopy(existing) - desired_settings = compose(existing) - _preserve_permission_denies(managed_before, desired_settings) - if not managed_writes_allowed(): - conflicts = managed_file_conflicts(managed_before, desired_settings, owned_paths) - if conflicts: + with _picker_management_lock(): + current_text = read_managed_file(path) + try: + existing = _parse_managed_settings(current_text) if current_text is not None else {} + managed_before = copy.deepcopy(existing) + before = _picker_group(existing) + previous, recovery, _phase = _recover_picker_transition("managed", path, existing) + if not previous and recovery is None: + previous = _legacy_managed_picker_management(path, existing) + picker_overlay: dict = {} + if static_picker_active: + composed = compose(copy.deepcopy(existing)) + picker_overlay = { + key: copy.deepcopy(composed[key]) for key in CLAUDE_MANAGED_PICKER_KEYS + } + intended = _transition_private_picker(existing, picker_overlay, previous) + desired_settings = compose(existing) + target = _picker_group(desired_settings) + if not previous and intended and before == target: + intended = {} + except (KeyError, RuntimeError) as exc: raise RuntimeError( - "Claude Code configuration cannot be applied non-interactively because " - f"OS-managed settings at {path} override ucode values: {', '.join(conflicts)}. " - "Run `ucode configure --agent claude` from an interactive terminal or contact " - "your administrator." + f"Cannot safely update Claude Code managed settings at {path}: {exc}. " + "ucode did not modify the file. Repair it or contact your administrator." + ) from exc + + transition = bool(recovery is not None or previous or intended) + conflict_paths = list(owned_paths) + if transition or before != target: + for picker_path in [[key] for key in CLAUDE_MANAGED_PICKER_KEYS]: + if picker_path not in conflict_paths: + conflict_paths.append(picker_path) + _preserve_permission_denies(managed_before, desired_settings) + if not managed_writes_allowed(): + conflicts = managed_file_conflicts(managed_before, desired_settings, conflict_paths) + if conflicts: + raise RuntimeError( + "Claude Code configuration cannot be applied non-interactively because " + f"OS-managed settings at {path} override ucode values: {', '.join(conflicts)}. " + "Run `ucode configure --agent claude` from an interactive terminal or contact " + "your administrator." + ) + mark_managed_file_verified(state, "claude", path, scope="local-compatible") + return + + picker_paths = [[key] for key in CLAUDE_MANAGED_PICKER_KEYS] + if transition: + _begin_picker_transition("managed", path, previous, intended, before, target) + try: + reconcile_managed_file( + path, + _dump_managed_settings(desired_settings), + tool="claude", + display="Claude Code", + owned_paths=owned_paths, + conditional_owned_paths=picker_paths, + repair_last_applied=bool(previous or recovery is not None), ) - mark_managed_file_verified(state, "claude", path, scope="local-compatible") - return - try: - reconcile_managed_file( - path, - _dump_managed_settings(desired_settings), - tool="claude", - display="Claude Code", - owned_paths=owned_paths, - ) - except ManagedFileWriteUnavailable: - conflicts = managed_file_conflicts(managed_before, desired_settings, owned_paths) - if conflicts: - raise - print_warning( - f"Claude Code OS-managed settings could not be updated at {path}; continuing with " - f"local settings at {CLAUDE_SETTINGS_PATH}." - ) - mark_managed_file_verified(state, "claude", path, scope="local-compatible") - return - mark_managed_file_verified(state, "claude", path) + except ManagedFileWriteUnavailable: + conflicts = managed_file_conflicts(managed_before, desired_settings, conflict_paths) + if conflicts: + raise + print_warning( + f"Claude Code OS-managed settings could not be updated at {path}; continuing with " + f"local settings at {CLAUDE_SETTINGS_PATH}." + ) + mark_managed_file_verified(state, "claude", path, scope="local-compatible") + return + + if is_dry_run(): + return + written_text = read_managed_file(path) + try: + written = _parse_managed_settings(written_text) if written_text is not None else {} + except RuntimeError as exc: + raise RuntimeError(f"Could not verify Claude picker settings at {path}: {exc}") from exc + if not _picker_group_matches(written, target): + raise RuntimeError(f"Could not verify Claude picker settings at {path}.") + if transition: + _save_managed_picker_management(path, intended) + mark_managed_file_verified(state, "claude", path) def _preserve_permission_denies(existing: dict, desired: dict) -> None: diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 572d98989..b24272162 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -6,7 +6,6 @@ import hashlib import os import re -import tempfile from collections.abc import Callable from pathlib import Path @@ -667,24 +666,10 @@ def _model_catalog_path(workspace: str, scope: str) -> Path: def _write_model_catalog(path: Path, catalog: dict) -> None: - temp_path = None try: - path.parent.mkdir(parents=True, exist_ok=True) - fd, raw_temp_path = tempfile.mkstemp( - prefix=f".{path.name}.", suffix=".tmp", dir=path.parent - ) - os.close(fd) - temp_path = Path(raw_temp_path) - write_json_file(temp_path, catalog) - os.replace(temp_path, path) - except OSError as exc: + write_json_file(path, catalog) + except RuntimeError as exc: raise RuntimeError(f"Could not write Codex model catalog at {path}.") from exc - finally: - if temp_path is not None: - try: - temp_path.unlink(missing_ok=True) - except OSError: - pass def _launch_token(state: dict, workspace: str) -> str: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 29c1a91ce..ebb4f4ee9 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -982,11 +982,9 @@ def _state_with_model_location(state: dict, tool: str, location: str | None) -> def _configure_location_backed_tool(state: dict, tool: str, location: str) -> dict: """Configure one agent at ``location`` and persist only ordinary developer state.""" - state = configure_tool(tool, state, parent_schema=location) existing = state.get("available_tools") or [] state["available_tools"] = sorted(set(existing) | {tool}) - save_state(state) - return state + return configure_tool(tool, state, parent_schema=location) def _configure_tools_with_model_location( @@ -1152,12 +1150,16 @@ def revert() -> int: state = load_state() managed_configs = state.get("managed_configs") or {} mcp_results = revert_mcp_configs(state) - claude_managed_result = claude_agent.revert_managed_settings() + claude_managed_result, claude_private_result = claude_agent.revert_settings(state) codex_managed_result = codex_agent.revert_managed_config() results: dict[str, bool] = { - tool: restore_file( - spec["config_path"], spec["backup_path"], bool(managed_configs.get(tool)) + tool: ( + claude_private_result + if tool == "claude" + else restore_file( + spec["config_path"], spec["backup_path"], bool(managed_configs.get(tool)) + ) ) for tool, spec in TOOL_SPECS.items() } diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index 3444abd95..b4e0d8163 100644 --- a/src/ucode/config_io.py +++ b/src/ucode/config_io.py @@ -109,10 +109,26 @@ def write_json_file(path: Path, payload: dict) -> None: console.print(f"\n[bold]\\[dry run] {path}[/bold]\n{content}") return ensure_parent_dir(path) + tmp_path: Path | None = None try: - path.write_text(content, encoding="utf-8") + with tempfile.NamedTemporaryFile( + mode="w", dir=path.parent, delete=False, encoding="utf-8" + ) as tmp: + tmp_path = Path(tmp.name) + tmp.write(content) + tmp.flush() + os.fsync(tmp.fileno()) + os.chmod(tmp_path, 0o600) + os.replace(tmp_path, path) + tmp_path = None except OSError as exc: raise RuntimeError(f"Failed to write config file: {path}") from exc + finally: + if tmp_path is not None: + try: + tmp_path.unlink(missing_ok=True) + except OSError: + pass def deep_merge_dict(base: dict, overlay: dict) -> dict: diff --git a/src/ucode/managed_files.py b/src/ucode/managed_files.py index 0bffc81f5..5dd4ff5af 100644 --- a/src/ucode/managed_files.py +++ b/src/ucode/managed_files.py @@ -136,6 +136,13 @@ def managed_writes_allowed() -> bool: return sys.stdin.isatty() +def write_private_json_file(path: Path, payload: dict) -> None: + """Atomically write owner-only JSON metadata on the same filesystem as managed backups.""" + if is_dry_run(): + return + _write_private_file(path, json.dumps(payload, indent=2) + "\n") + + @contextmanager def managed_write_batch(displays: list[str]) -> Iterator[None]: """Group setup messaging for agents configured in one command.""" @@ -182,6 +189,112 @@ def managed_file_conflicts( return conflicts +def managed_last_applied_paths( + tool: str, + path: Path, + candidate_paths: list[list[str]], + *, + parser: ManagedParser, +) -> tuple[dict, list[list[str]]]: + """Return the last-applied document and candidate paths recorded as ucode-owned. + + Ownership of an OS-managed file is machine-global, while workspace state is not. Consult the + integrity-checked snapshot so a new workspace can identify values written by an earlier one. + """ + entry = _manifest_files(_load_manifest()).get(tool) + if not isinstance(entry, dict) or entry.get("path") != str(path): + return {}, [] + last_text = _snapshot_text(entry, "last_applied_file") + if last_text is None: + return {}, [] + try: + last = parser(last_text) + except Exception as exc: # noqa: BLE001 + raise RuntimeError( + f"Cannot safely inspect the last managed settings written for {tool} at {path}: {exc}" + ) from exc + if not isinstance(last, dict): + raise RuntimeError( + f"Cannot safely inspect the last managed settings written for {tool} at {path}: " + "the snapshot is not an object." + ) + + owned_paths = entry.get("owned_paths") + recorded = { + tuple(owned_path) + for raw_path in (owned_paths if isinstance(owned_paths, list) else []) + if (owned_path := _owned_path(raw_path)) is not None + } + candidates: list[list[str]] = [] + for raw_path in candidate_paths: + candidate = _owned_path(raw_path) + if candidate is None or tuple(candidate) not in recorded: + continue + candidates.append(candidate) + return deepcopy(last), candidates + + +def unchanged_managed_paths( + tool: str, + path: Path, + current: dict, + candidate_paths: list[list[str]], + *, + parser: ManagedParser, +) -> list[list[str]]: + """Return recorded candidates whose current values still match ucode's last write.""" + last, recorded_candidates = managed_last_applied_paths( + tool, path, candidate_paths, parser=parser + ) + unchanged: list[list[str]] = [] + for candidate in recorded_candidates: + if _path_value(current, candidate) == _path_value(last, candidate): + unchanged.append(candidate) + return unchanged + + +def restore_unchanged_managed_paths( + tool: str, + path: Path, + current: dict, + candidate_paths: list[list[str]], + *, + parser: ManagedParser, +) -> tuple[dict, list[list[str]]]: + """Restore unchanged owned candidates to their pre-ucode values. + + Returns a copy of ``current`` plus the paths restored. Values changed since ucode's last write + remain untouched and are omitted so the next successful reconcile can relinquish ownership. + """ + unchanged = unchanged_managed_paths(tool, path, current, candidate_paths, parser=parser) + if not unchanged: + return deepcopy(current), [] + entry = _manifest_files(_load_manifest()).get(tool) + if not isinstance(entry, dict) or entry.get("path") != str(path): + return deepcopy(current), [] + original_text = _original_text(entry) + try: + original = parser(original_text) if original_text is not None else {} + except Exception as exc: # noqa: BLE001 + raise RuntimeError( + f"Cannot safely inspect the original managed settings for {tool} at {path}: {exc}" + ) from exc + if not isinstance(original, dict): + raise RuntimeError( + f"Cannot safely inspect the original managed settings for {tool} at {path}: " + "the snapshot is not an object." + ) + + restored = deepcopy(current) + for owned_path in unchanged: + original_value = _path_value(original, owned_path) + if original_value is _MISSING: + _delete_path_value(restored, owned_path) + else: + _set_path_value(restored, owned_path, original_value) + return restored, unchanged + + def managed_file_status( state: dict, tool: str, @@ -233,11 +346,15 @@ def reconcile_managed_file( tool: str, display: str, owned_paths: list[list[str]], + conditional_owned_paths: list[list[str]] | None = None, + repair_last_applied: bool = False, ) -> str: """Back up, atomically write, and verify one OS-managed settings file. The first pre-ucode contents are retained until ``ucode revert``. Subsequent writes update only - the last-applied snapshot used for drift-safe three-way restoration. + the last-applied snapshot used for drift-safe three-way restoration. Conditional paths are + relinquished when omitted from ``owned_paths`` so externally changed values cannot become + ucode-owned merely because they appear in the next whole-file snapshot. """ if not managed_files_supported(): print_warning( @@ -256,6 +373,18 @@ def reconcile_managed_file( ) current_text = read_managed_file(path) if current_text == desired_text: + # A prior write may have replaced the settings and then failed while recording its + # last-applied snapshot. Only an active caller-owned transaction proves that retry; the + # same-path backup alone must not claim an unrelated exact-match file. + entry = _manifest_files(_load_manifest()).get(tool) + if repair_last_applied and isinstance(entry, dict) and entry.get("path") == str(path): + _record_last_applied( + tool, + path, + desired_text, + owned_paths, + conditional_owned_paths=conditional_owned_paths or [], + ) return "unchanged" if is_dry_run(): console.print(f"\n[bold]\\[dry run] {path} (via sudo)[/bold]\n{desired_text}") @@ -298,7 +427,13 @@ def reconcile_managed_file( f"{display} managed settings changed concurrently at {path}. ucode will not overwrite " "the newer policy; run the command again or contact your administrator." ) - _record_last_applied(tool, path, desired_text, owned_paths) + _record_last_applied( + tool, + path, + desired_text, + owned_paths, + conditional_owned_paths=conditional_owned_paths or [], + ) if not _managed_write_batch: print_success(f"Settings configured for {display}") return "created" if created else "written" @@ -310,6 +445,9 @@ def revert_managed_file( display: str, parser: ManagedParser, dumper: ManagedDumper, + acquisition_lease: list[dict] | None = None, + excluded_owned_paths: list[list[str]] | None = None, + before_backup_delete: Callable[[], None] | None = None, ) -> str: """Restore one managed file from its baseline while preserving later external edits.""" manifest = _load_manifest() @@ -327,23 +465,45 @@ def revert_managed_file( original_text = _original_text(entry) last_text = _snapshot_text(entry, "last_applied_file") - if current_text == last_text: - desired_text = original_text - elif current_text is None or last_text is None: + lease_drift = False + if last_text is None and current_text != original_text: + raise RuntimeError( + f"Cannot safely revert {display} managed settings at {path}: the last-applied " + f"snapshot is missing. The backup was retained under {MANAGED_BACKUP_DIR}. " + "Re-run configuration to repair its metadata, then run `ucode revert` again." + ) + if current_text is None or last_text is None: desired_text = current_text else: try: current_doc = parser(current_text) original_doc = parser(original_text) if original_text is not None else {} last_doc = parser(last_text) + if not all(isinstance(doc, dict) for doc in (current_doc, original_doc, last_doc)): + raise ValueError("managed settings snapshots must be objects") except Exception as exc: # noqa: BLE001 raise RuntimeError( f"Cannot safely revert {display} managed settings at {path}: {exc}" ) from exc owned_paths = entry.get("owned_paths") paths = owned_paths if isinstance(owned_paths, list) else [] - reverted = _three_way_revert(current_doc, original_doc, last_doc, paths) - desired_text = dumper(reverted) + lease_paths = _acquisition_lease_paths(acquisition_lease or []) + excluded_paths = { + tuple(excluded_path) + for raw_path in (excluded_owned_paths or []) + if (excluded_path := _owned_path(raw_path)) is not None + } + generic_paths = [ + raw_path + for raw_path in paths + if (path_value := _owned_path(raw_path)) is None + or tuple(path_value) not in lease_paths | excluded_paths + ] + reverted = _three_way_revert(current_doc, original_doc, last_doc, generic_paths) + reverted, lease_drift = _apply_acquisition_lease( + current_doc, reverted, acquisition_lease or [] + ) + desired_text = None if original_text is None and not reverted else dumper(reverted) if desired_text != current_text: if not managed_writes_allowed(): @@ -368,10 +528,12 @@ def revert_managed_file( f"retained under {MANAGED_BACKUP_DIR}." ) + if before_backup_delete is not None: + before_backup_delete() _delete_backup(tool, manifest, entry) if original_text is None and desired_text is None: return "removed" - if current_text != last_text: + if current_text != last_text or lease_drift: return "ucode entries removed; external changes preserved" return "restored" @@ -436,9 +598,9 @@ def _backup_filename(tool: str, path: Path) -> str: return f"{tool}-managed-settings.backup{suffix}" -def _last_applied_filename(tool: str, path: Path) -> str: +def _last_applied_filename(tool: str, path: Path, content_sha256: str) -> str: suffix = path.suffix or ".txt" - return f"{tool}-managed-settings.last-applied{suffix}" + return f"{tool}-managed-settings.last-applied-{content_sha256}{suffix}" def _ensure_backup(tool: str, path: Path, current_text: str | None) -> bool: @@ -471,22 +633,53 @@ def _ensure_backup(tool: str, path: Path, current_text: str | None) -> bool: def _record_last_applied( - tool: str, path: Path, desired_text: str, owned_paths: list[list[str]] + tool: str, + path: Path, + desired_text: str, + owned_paths: list[list[str]], + *, + conditional_owned_paths: list[list[str]], ) -> None: manifest = _load_manifest() entry = _manifest_files(manifest).get(tool) if not isinstance(entry, dict): raise RuntimeError(f"Missing managed-settings backup metadata for {tool}.") - last_file = _last_applied_filename(tool, path) + content_sha256 = _sha256(desired_text) + last_file = _last_applied_filename(tool, path, content_sha256) + previous_last_file = entry.get("last_applied_file") _write_private_file(MANAGED_BACKUP_DIR / last_file, desired_text) entry["last_applied_file"] = last_file - entry["last_applied_sha256"] = _sha256(desired_text) + entry["last_applied_sha256"] = content_sha256 known_paths = entry.get("owned_paths") if isinstance(entry.get("owned_paths"), list) else [] + active = { + tuple(owned_path) + for raw_path in owned_paths + if (owned_path := _owned_path(raw_path)) is not None + } + conditional = { + tuple(owned_path) + for raw_path in conditional_owned_paths + if (owned_path := _owned_path(raw_path)) is not None + } + known_paths = [ + raw_path + for raw_path in known_paths + if (known_path := _owned_path(raw_path)) is None + or tuple(known_path) not in conditional + or tuple(known_path) in active + ] for owned_path in owned_paths: if owned_path not in known_paths: known_paths.append(list(owned_path)) entry["owned_paths"] = known_paths _write_manifest(manifest) + if isinstance(previous_last_file, str) and previous_last_file != last_file: + try: + _snapshot_path(previous_last_file).unlink(missing_ok=True) + except OSError: + # The manifest now points at the new immutable generation. A stale unreferenced + # snapshot is harmless and can be removed by a later revert or manual cleanup. + pass def _snapshot_text(entry: dict, key: str) -> str | None: @@ -520,16 +713,22 @@ def _backup_label(tool: str) -> str: def _delete_backup(tool: str, manifest: dict, entry: dict) -> None: - for key in ("backup_file", "last_applied_file"): - filename = entry.get(key) - if isinstance(filename, str): - try: - _snapshot_path(filename).unlink(missing_ok=True) - except OSError as exc: - raise RuntimeError(f"Could not remove managed-settings backup: {exc}") from exc + snapshots = [ + _snapshot_path(filename) + for key in ("backup_file", "last_applied_file") + if isinstance((filename := entry.get(key)), str) + ] _manifest_files(manifest).pop(tool, None) _write_manifest(manifest) + # Once the manifest no longer references these generations, cleanup is best-effort. A stale + # private snapshot is harmless; unlinking first would make a failed manifest switch unretryable. + for snapshot in snapshots: + try: + snapshot.unlink(missing_ok=True) + except OSError: + pass + def _snapshot_path(filename: str) -> Path: if Path(filename).name != filename: @@ -579,6 +778,39 @@ def _owned_path(value: object) -> list[str] | None: return cast(list[str], value) +def _acquisition_lease_paths(lease: list[dict]) -> set[tuple[str, ...]]: + paths: set[tuple[str, ...]] = set() + for entry in lease: + path = _owned_path(entry.get("path")) if isinstance(entry, dict) else None + if ( + path is None + or not isinstance(entry.get("baseline_exists"), bool) + or "applied" not in entry + or (entry["baseline_exists"] and "baseline" not in entry) + ): + raise RuntimeError("Invalid managed-settings acquisition lease.") + paths.add(tuple(path)) + return paths + + +def _apply_acquisition_lease(current: dict, reverted: dict, lease: list[dict]) -> tuple[dict, bool]: + if not lease: + return reverted, False + _acquisition_lease_paths(lease) + unchanged = all( + _path_value(current, cast(list[str], entry["path"])) == entry["applied"] for entry in lease + ) + if not unchanged: + return reverted, True + for entry in lease: + path = cast(list[str], entry["path"]) + if entry["baseline_exists"]: + _set_path_value(reverted, path, entry["baseline"]) + else: + _delete_path_value(reverted, path) + return reverted, False + + def _three_way_revert(current: dict, original: dict, last: dict, paths: list) -> dict: reverted = deepcopy(current) for raw_path in paths: diff --git a/tests/conftest.py b/tests/conftest.py index a1a7419ff..6f700637d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -33,6 +33,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch): import ucode.managed_config as managed_config_mod import ucode.managed_files as managed_files_mod import ucode.state as state_mod + from ucode.agents import claude as claude_mod from ucode.agents import codex as codex_mod state_dir = tmp_path / ".ucode" @@ -49,6 +50,9 @@ def _isolate_ucode_state(tmp_path, monkeypatch): monkeypatch.setattr( managed_files_mod, "MANAGED_BACKUP_MANIFEST_PATH", backup_dir / "manifest.json" ) + monkeypatch.setattr( + claude_mod, "CLAUDE_PICKER_MANAGEMENT_PATH", state_dir / "claude-picker-management.json" + ) monkeypatch.setattr(codex_mod, "codex_managed_config_path", lambda: None) def reject_privileged_write(path, _desired_text): diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index af3531f74..29da28930 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -6,7 +6,11 @@ import os import shlex import subprocess +import sys +import threading +import time from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, Mock import pytest @@ -14,7 +18,7 @@ from ucode import managed_files from ucode.agents import LaunchOptions, claude from ucode.smart_routing import claude_routing, v2 -from ucode.state import MANAGED_OVERLAY_KEY +from ucode.state import MANAGED_OVERLAY_KEY, developer_state_from_resolved WS = "https://example.databricks.com" # A connection MCP proxy argv, used by the Claude MCP-registration helper tests. @@ -28,6 +32,24 @@ def _proxy_argv() -> list[str]: return build_mcp_proxy_argv(GH_URL, WS, "p") +def _patch_private_json_store(monkeypatch, initial: dict, on_write=None) -> dict[str, dict]: + """Mock the strict private settings reader and its matching atomic writer.""" + store = {str(claude.CLAUDE_SETTINGS_PATH): json.loads(json.dumps(initial))} + + def read(path): + return json.loads(json.dumps(store.get(str(path), {}))) + + def write(path, payload): + copied = json.loads(json.dumps(payload)) + store[str(path)] = copied + if on_write is not None: + on_write(path, copied) + + monkeypatch.setattr(claude, "_read_private_json_object", read) + monkeypatch.setattr(claude, "write_json_file", write) + return store + + @pytest.fixture(autouse=True) def _avoid_real_managed_settings(monkeypatch): monkeypatch.setattr(claude, "_managed_settings_path", lambda: None) @@ -44,6 +66,52 @@ def test_display(self): assert claude.SPEC["display"] == "Claude Code" +def test_picker_process_lock_uses_msvcrt_on_windows(tmp_path, monkeypatch): + metadata_path = tmp_path / "claude-picker-management.json" + calls: list[tuple[int, int]] = [] + fake_msvcrt = SimpleNamespace( + LK_LOCK=11, + LK_UNLCK=12, + locking=lambda _fd, mode, size: calls.append((mode, size)), + ) + monkeypatch.setattr(claude, "CLAUDE_PICKER_MANAGEMENT_PATH", metadata_path) + monkeypatch.setattr(claude, "current_os", lambda: claude.OS.WINDOWS) + monkeypatch.setitem(sys.modules, "msvcrt", fake_msvcrt) + + with claude._picker_process_lock(): + assert metadata_path.with_name(f"{metadata_path.name}.lock").read_bytes() == b"\0" + + assert calls == [(fake_msvcrt.LK_LOCK, 1), (fake_msvcrt.LK_UNLCK, 1)] + + +def test_process_file_lock_blocks_a_second_process(tmp_path): + lock_path = tmp_path / "registration.lock" + ready_path = tmp_path / "child-ready" + acquired_path = tmp_path / "child-acquired" + script = "\n".join( + [ + "from pathlib import Path", + "from ucode.agents.claude import _process_file_lock", + f"lock_path = Path({str(lock_path)!r})", + f"Path({str(ready_path)!r}).write_text('ready')", + "with _process_file_lock(lock_path):", + f" Path({str(acquired_path)!r}).write_text('acquired')", + ] + ) + + with claude._process_file_lock(lock_path): + child = subprocess.Popen([sys.executable, "-c", script]) + deadline = time.monotonic() + 5 + while not ready_path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + assert ready_path.exists() + time.sleep(0.1) + assert not acquired_path.exists() + + assert child.wait(timeout=5) == 0 + assert acquired_path.read_text() == "acquired" + + class TestMinimumVersion: @pytest.mark.parametrize("version", ["2.1.248", "2.1.250", "3.0.0"]) def test_supported_version(self, monkeypatch, version): @@ -701,9 +769,15 @@ def test_has_max_turns(self): class TestWriteToolConfigMcpRegistration: def _common_patches(self, monkeypatch, calls): monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - monkeypatch.setattr(claude, "read_json_safe", lambda path: {}) - monkeypatch.setattr(claude, "write_json_file", lambda path, payload: None) - monkeypatch.setattr(claude, "save_state", lambda state: None) + _patch_private_json_store(monkeypatch, {}) + persisted: dict = {} + + def save(state): + persisted.clear() + persisted.update(json.loads(json.dumps(state))) + + monkeypatch.setattr(claude, "save_state", save) + monkeypatch.setattr(claude, "load_state", lambda: json.loads(json.dumps(persisted))) monkeypatch.setattr( claude, "_register_web_search_mcp", @@ -741,9 +815,8 @@ class TestWriteToolConfigStripsRemovedEnvKeys: def _patch(self, monkeypatch, existing, written): monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - monkeypatch.setattr(claude, "read_json_safe", lambda path: existing) - monkeypatch.setattr( - claude, "write_json_file", lambda path, payload: written.append(payload) + _patch_private_json_store( + monkeypatch, existing, lambda _path, payload: written.append(payload) ) monkeypatch.setattr(claude, "save_state", lambda state: None) monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) @@ -814,17 +887,16 @@ class TestWriteToolConfigManagedSettings: def _patch(self, monkeypatch, private_writes, managed_writes, existing_by_path=None): existing_by_path = existing_by_path or {} monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - # Deep-copy the seeded existing content so the compose step can't mutate the fixture. - monkeypatch.setattr( - claude, - "read_json_safe", - lambda path: json.loads(json.dumps(existing_by_path.get(str(path), {}))), - ) - monkeypatch.setattr( - claude, - "write_json_file", - lambda path, payload: private_writes.append((str(path), payload)), - ) + + def fake_write_private(path, payload): + existing_by_path[str(path)] = payload + private_writes.append((str(path), payload)) + + private_initial = existing_by_path.get(str(claude.CLAUDE_SETTINGS_PATH), {}) + private_store = _patch_private_json_store(monkeypatch, private_initial, fake_write_private) + existing_by_path[str(claude.CLAUDE_SETTINGS_PATH)] = private_store[ + str(claude.CLAUDE_SETTINGS_PATH) + ] monkeypatch.setattr(claude, "save_state", lambda state: None) monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) @@ -840,6 +912,7 @@ def _patch(self, monkeypatch, private_writes, managed_writes, existing_by_path=N monkeypatch.setattr(claude, "mark_managed_file_verified", lambda *a, **kw: None) def fake_write_managed(path, text, **kwargs): + existing_by_path[str(path)] = json.loads(text) managed_writes.append((str(path), text)) return "written" @@ -941,6 +1014,135 @@ def test_managed_file_updates_gateway_settings_without_changing_model_picker(sel assert written["modelPicker"] == picker assert written["env"]["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic" + @pytest.mark.parametrize( + "source_kwargs", + [ + {"provider": "main.default.anthropic"}, + {"parent_schema": "main.managed_models"}, + ], + ids=["provider", "model-location"], + ) + def test_native_discovery_removes_previously_managed_static_picker( + self, monkeypatch, source_kwargs + ): + private_writes: list = [] + managed_writes: list = [] + stale_picker = { + "availableModels": ["system.ai.claude-opus-4-8"], + "enforceAvailableModels": True, + "modelPicker": { + "replaceBuiltInOptions": True, + "options": [ + { + "model": "system.ai.claude-opus-4-8", + "label": "claude-opus-4-8", + } + ], + }, + } + existing_settings = { + **stale_picker, + "companyPolicy": {"keep": True}, + "env": {"MY_OWN": "keep"}, + } + existing = { + str(claude.CLAUDE_SETTINGS_PATH): existing_settings, + str(FAKE_MANAGED_PATH): existing_settings, + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + + def restore_managed_picker(tool, path, current, candidate_paths, **kwargs): + for candidate_path in candidate_paths: + current.pop(candidate_path[0], None) + return current, candidate_paths + + monkeypatch.setattr(claude, "restore_unchanged_managed_paths", restore_managed_picker) + monkeypatch.setattr( + claude, + "managed_last_applied_paths", + lambda tool, path, candidate_paths, **kwargs: ( + existing_settings, + candidate_paths, + ), + ) + state = { + "workspace": WS, + "codex_models": [], + "claude_static_models": stale_picker["availableModels"], + "managed_configs": { + "claude": {"keys": [[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS]} + }, + } + + result = claude.write_tool_config(state, None, **source_kwargs) + + written_settings = [private_writes[0][1], json.loads(managed_writes[0][1])] + for written in written_settings: + assert not set(claude.CLAUDE_MANAGED_PICKER_KEYS) & written.keys() + assert written["companyPolicy"] == {"keep": True} + assert written["env"]["MY_OWN"] == "keep" + assert not any( + [key] in result["managed_configs"]["claude"]["keys"] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + ) + + def test_model_location_preserves_unowned_picker_settings(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + picker_settings = { + "availableModels": ["enterprise-model"], + "enforceAvailableModels": True, + "modelPicker": { + "replaceBuiltInOptions": True, + "options": [{"model": "enterprise-model", "label": "Enterprise"}], + }, + } + existing = { + str(claude.CLAUDE_SETTINGS_PATH): picker_settings, + str(FAKE_MANAGED_PATH): picker_settings, + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = { + "workspace": WS, + "codex_models": [], + "managed_configs": {"claude": {"keys": [["env", "ANTHROPIC_BASE_URL"]]}}, + } + + claude.write_tool_config(state, None, parent_schema="main.managed_models") + + assert { + key: private_writes[0][1][key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS + } == picker_settings + managed = json.loads(managed_writes[0][1]) + assert {key: managed[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == picker_settings + + def test_managed_file_preserves_picker_without_global_ownership_proof(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + picker_settings = { + "availableModels": ["enterprise-model"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "enterprise-model"}]}, + } + self._patch( + monkeypatch, + private_writes, + managed_writes, + {str(FAKE_MANAGED_PATH): picker_settings}, + ) + state = { + "workspace": WS, + "codex_models": [], + "managed_configs": { + "claude": {"keys": [[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS]} + }, + } + + claude.write_tool_config(state, None, parent_schema="main.managed_models") + + managed = json.loads(managed_writes[0][1]) + assert {key: managed[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == picker_settings + def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch): private_writes: list = [] managed_writes: list = [] @@ -1214,7 +1416,7 @@ def test_static_models_written_to_picker(self, monkeypatch): "codex_models": [], "claude_static_models": static_models, } - claude.write_tool_config(state, "system.ai.claude-opus-4-8") + result = claude.write_tool_config(state, "system.ai.claude-opus-4-8") # Managed file should have the picker. assert len(managed_writes) > 0 managed_content = json.loads(managed_writes[0][1]) @@ -1222,6 +1424,10 @@ def test_static_models_written_to_picker(self, monkeypatch): assert managed_content["enforceAvailableModels"] is True assert "modelPicker" in managed_content assert len(managed_content["modelPicker"]["options"]) == 2 + assert all( + [key] in result["managed_configs"]["claude"]["keys"] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + ) def test_static_models_not_written_when_absent(self, monkeypatch): # When claude_static_models is not in state, picker fields are not written. @@ -1237,6 +1443,1262 @@ def test_static_models_not_written_when_absent(self, monkeypatch): assert "modelPicker" not in managed_content +class TestPickerOwnershipAcrossWorkspaces: + @staticmethod + def _patch_files(monkeypatch, tmp_path): + private_path = tmp_path / "ucode-settings.json" + managed_path = tmp_path / "managed-settings.json" + backup_path = tmp_path / "ucode-settings.backup.json" + metadata_path = tmp_path / "claude-picker-management.json" + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", private_path) + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", backup_path) + monkeypatch.setattr(claude, "CLAUDE_PICKER_MANAGEMENT_PATH", metadata_path) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed_path) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + monkeypatch.setattr(managed_files, "managed_writes_allowed", lambda: True) + monkeypatch.setattr(managed_files, "managed_files_supported", lambda: True) + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda path, text: path.write_text(text, encoding="utf-8"), + ) + monkeypatch.setattr(claude, "save_state", lambda state: None) + monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) + return private_path, managed_path, backup_path, metadata_path + + def test_scoped_configuration_without_picker_lease_preserves_it_policy( + self, monkeypatch, tmp_path + ): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + picker = { + "availableModels": ["enterprise"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(picker), encoding="utf-8") + + claude.write_tool_config( + {"workspace": "https://workspace.example.com", "codex_models": []}, + None, + parent_schema="main.models", + ) + + for path in (private_path, managed_path): + settings = json.loads(path.read_text()) + assert {key: settings[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == picker + assert not metadata_path.exists() + manifest = json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text()) + assert not any( + [key] in manifest["files"]["claude"]["owned_paths"] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + ) + + def test_legacy_private_backup_without_managed_proof_preserves_current_picker( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + backup_picker = { + "availableModels": ["enterprise-one"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise-one"}]}, + } + current_picker = { + "availableModels": ["enterprise-two"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "enterprise-two"}]}, + } + private_path.write_text(json.dumps(current_picker), encoding="utf-8") + managed_path.write_text(json.dumps(current_picker), encoding="utf-8") + backup_path.write_text(json.dumps(backup_picker), encoding="utf-8") + monkeypatch.setattr( + claude, + "load_full_state", + lambda: { + "workspaces": { + "https://workspace-a.example.com": { + "managed_configs": { + "claude": {"keys": [[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS]} + } + } + } + }, + ) + + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.models", + ) + + written = json.loads(private_path.read_text()) + assert {key: written[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == current_picker + assert json.loads(backup_path.read_text()) == backup_picker + assert not metadata_path.exists() + + @pytest.mark.parametrize( + ("contents", "message"), + [("{", "Cannot parse Claude settings"), ("[]", "must contain a JSON object")], + ids=["invalid", "non-object"], + ) + def test_invalid_private_settings_fail_before_any_mutation( + self, monkeypatch, tmp_path, contents, message + ): + private_path, managed_path, backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + private_path.write_text(contents, encoding="utf-8") + managed_path.write_text('{"companyPolicy": "keep"}', encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "managed_configs": {"claude": {"keys": [["env", "ANTHROPIC_BASE_URL"]]}}, + } + + with pytest.raises(RuntimeError, match=message): + claude.write_tool_config(state, "system.ai.claude-opus-4-8") + + assert private_path.read_text() == contents + assert json.loads(managed_path.read_text()) == {"companyPolicy": "keep"} + assert not backup_path.exists() + assert not metadata_path.exists() + + def test_private_revert_restores_latest_acquisition_baseline(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + picker_one = { + "availableModels": ["enterprise-one"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise-one"}]}, + } + picker_two = { + "availableModels": ["enterprise-two"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "enterprise-two"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(picker_one), encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + static_state = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + } + claude.write_tool_config(static_state, static_models[0]) + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.models", + ) + current = json.loads(private_path.read_text()) + current.update(picker_two) + private_path.write_text(json.dumps(current), encoding="utf-8") + claude.write_tool_config(static_state, static_models[0]) + + assert claude.revert_private_settings(static_state) is True + + restored = json.loads(private_path.read_text()) + assert {key: restored[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == picker_two + + def test_private_revert_preserves_whole_drifted_picker_group(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "availableModels": ["enterprise-one"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise-one"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + state = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + } + claude.write_tool_config(state, static_models[0]) + drift = { + "availableModels": ["enterprise-drift"], + "enforceAvailableModels": True, + "modelPicker": { + "replaceBuiltInOptions": False, + "options": [{"model": "enterprise-drift", "label": "Enterprise"}], + }, + } + current = json.loads(private_path.read_text()) + current.update(drift) + private_path.write_text(json.dumps(current), encoding="utf-8") + + assert claude.revert_private_settings(state) is True + + restored = json.loads(private_path.read_text()) + assert {key: restored[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == drift + + def test_private_revert_retry_after_lease_clear_failure_keeps_restored_target( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + save_lease = claude._save_private_picker_management + failed = False + + def fail_first_clear(entries): + nonlocal failed + if not entries and not failed: + failed = True + raise RuntimeError("injected private lease clear failure") + save_lease(entries) + + monkeypatch.setattr(claude, "_save_private_picker_management", fail_first_clear) + with pytest.raises(RuntimeError, match="injected private lease clear failure"): + claude.revert_private_settings(state) + + assert json.loads(private_path.read_text()) == baseline + assert not backup_path.exists() + assert ( + json.loads(metadata_path.read_text())["leases"]["private"]["pending"]["phase"] + == "revert_written" + ) + + assert claude.revert_private_settings(state) is True + assert json.loads(private_path.read_text()) == baseline + assert set(json.loads(metadata_path.read_text())["leases"]) == {"managed"} + + def test_private_revert_without_picker_lease_is_journaled_for_retry( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + original = {"companyPolicy": "original"} + private_path.write_text(json.dumps(original), encoding="utf-8") + managed_path.write_text("{}", encoding="utf-8") + state = {"workspace": "https://workspace.example.com", "codex_models": []} + claude.write_tool_config(state, "system.ai.claude-opus-4-8") + assert backup_path.exists() + save_lease = claude._save_private_picker_management + failed = False + + def fail_first_clear(entries): + nonlocal failed + if not entries and not failed: + failed = True + raise RuntimeError("injected private lease clear failure") + save_lease(entries) + + monkeypatch.setattr(claude, "_save_private_picker_management", fail_first_clear) + with pytest.raises(RuntimeError, match="injected private lease clear failure"): + claude.revert_private_settings(state) + + assert json.loads(private_path.read_text()) == original + assert not backup_path.exists() + pending = json.loads(metadata_path.read_text())["leases"]["private"]["pending"] + assert pending["phase"] == "revert_written" + assert "target_document_sha256" in pending + + externally_edited = json.loads(private_path.read_text()) + externally_edited["external"] = "preserve" + private_path.write_text(json.dumps(externally_edited), encoding="utf-8") + + assert claude.revert_private_settings(state) is True + assert json.loads(private_path.read_text()) == { + **original, + "external": "preserve", + } + assert not metadata_path.exists() + + def test_reconfigure_after_failed_revert_reacquires_complete_private_baseline( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + original = { + "companyPolicy": {"preserve": True}, + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(original), encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + save_lease = claude._save_private_picker_management + failed = False + + def fail_first_clear(entries): + nonlocal failed + if not entries and not failed: + failed = True + raise RuntimeError("injected private lease clear failure") + save_lease(entries) + + monkeypatch.setattr(claude, "_save_private_picker_management", fail_first_clear) + with pytest.raises(RuntimeError, match="injected private lease clear failure"): + claude.revert_private_settings(state) + assert json.loads(private_path.read_text()) == original + assert not backup_path.exists() + + claude.write_tool_config(state, state["claude_static_models"][0]) + assert backup_path.exists() + assert claude.revert_private_settings(state) is True + assert json.loads(private_path.read_text()) == original + + def test_reconfigure_replaces_stale_backup_after_verified_revert_drift( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + original = { + "companyPolicy": {"preserve": True}, + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(original), encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + remove_backup = claude._remove_private_backup + failed = False + + def fail_first_backup_cleanup(): + nonlocal failed + if not failed: + failed = True + raise RuntimeError("injected crash before backup cleanup") + remove_backup() + + monkeypatch.setattr(claude, "_remove_private_backup", fail_first_backup_cleanup) + with pytest.raises(RuntimeError, match="injected crash before backup cleanup"): + claude.revert_private_settings(state) + assert json.loads(private_path.read_text()) == original + assert json.loads(backup_path.read_text()) == original + + drifted = {**original, "external": "preserve this edit"} + private_path.write_text(json.dumps(drifted), encoding="utf-8") + claude.write_tool_config(state, state["claude_static_models"][0]) + assert json.loads(backup_path.read_text()) == drifted + + assert claude.revert_private_settings(state) is True + assert json.loads(private_path.read_text()) == drifted + + def test_verified_revert_exact_before_recreation_is_postwrite_drift( + self, monkeypatch, tmp_path + ): + private_path, managed_path, _backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + original = { + "companyPolicy": {"preserve": True}, + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(original), encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + exact_before = json.loads(private_path.read_text()) + save_lease = claude._save_private_picker_management + failed = False + + def fail_first_clear(entries): + nonlocal failed + if not entries and not failed: + failed = True + raise RuntimeError("injected private lease clear failure") + save_lease(entries) + + monkeypatch.setattr(claude, "_save_private_picker_management", fail_first_clear) + with pytest.raises(RuntimeError, match="injected private lease clear failure"): + claude.revert_private_settings(state) + + private_path.write_text(json.dumps(exact_before), encoding="utf-8") + assert claude.revert_private_settings(state) is True + assert json.loads(private_path.read_text()) == exact_before + + def test_prewrite_revert_drift_does_not_succeed_with_ucode_settings( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + managed_path.write_text("{}", encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + assert not backup_path.exists() + begin_transition = claude._begin_picker_transition + + def crash_after_journal(*args, **kwargs): + begin_transition(*args, **kwargs) + if kwargs.get("phase") == "reverting": + raise RuntimeError("injected crash after revert journal") + + monkeypatch.setattr(claude, "_begin_picker_transition", crash_after_journal) + with pytest.raises(RuntimeError, match="injected crash after revert journal"): + claude.revert_private_settings(state) + settings = json.loads(private_path.read_text()) + settings["external"] = "edited" + private_path.write_text(json.dumps(settings), encoding="utf-8") + monkeypatch.setattr(claude, "_begin_picker_transition", begin_transition) + + with pytest.raises(RuntimeError, match="before its target was verified"): + claude.revert_private_settings(state) + + remaining = json.loads(private_path.read_text()) + assert remaining["external"] == "edited" + assert "apiKeyHelper" in remaining + assert not backup_path.exists() + + def test_reconfigure_does_not_capture_unverified_prewrite_revert_drift( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + managed_path.write_text("{}", encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + begin_transition = claude._begin_picker_transition + + def crash_after_journal(*args, **kwargs): + begin_transition(*args, **kwargs) + if kwargs.get("phase") == "reverting": + raise RuntimeError("injected crash after revert journal") + + monkeypatch.setattr(claude, "_begin_picker_transition", crash_after_journal) + with pytest.raises(RuntimeError, match="injected crash after revert journal"): + claude.revert_private_settings(state) + settings = json.loads(private_path.read_text()) + settings["external"] = "edited" + private_path.write_text(json.dumps(settings), encoding="utf-8") + monkeypatch.setattr(claude, "_begin_picker_transition", begin_transition) + + with pytest.raises(RuntimeError, match="unverified external changes"): + claude.write_tool_config(state, state["claude_static_models"][0]) + + assert not backup_path.exists() + assert json.loads(private_path.read_text())["apiKeyHelper"] == settings["apiKeyHelper"] + + def test_managed_revert_retry_after_lease_clear_failure_clears_stale_lease( + self, monkeypatch, tmp_path + ): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + save_lease = claude._save_managed_picker_management + failed = False + + def fail_first_clear(path, entries): + nonlocal failed + if not entries and not failed: + failed = True + raise RuntimeError("injected managed lease clear failure") + save_lease(path, entries) + + monkeypatch.setattr(claude, "_save_managed_picker_management", fail_first_clear) + with pytest.raises(RuntimeError, match="injected managed lease clear failure"): + claude.revert_managed_settings() + + assert { + key: json.loads(managed_path.read_text())[key] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + } == baseline + assert ( + "claude" + not in json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text())["files"] + ) + + assert claude.revert_managed_settings() == "unchanged" + assert set(json.loads(metadata_path.read_text())["leases"]) == {"private"} + + def test_managed_verified_revert_preserves_recreated_before_picker(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "companyPolicy": "preserve", + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + applied = json.loads(managed_path.read_text()) + applied_picker = {key: applied[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} + save_lease = claude._save_managed_picker_management + failed = False + + def fail_first_clear(path, entries): + nonlocal failed + if not entries and not failed: + failed = True + raise RuntimeError("injected managed lease clear failure") + save_lease(path, entries) + + monkeypatch.setattr(claude, "_save_managed_picker_management", fail_first_clear) + with pytest.raises(RuntimeError, match="injected managed lease clear failure"): + claude.revert_managed_settings() + assert ( + json.loads(metadata_path.read_text())["leases"]["managed"]["pending"]["phase"] + == "revert_written" + ) + + recreated = json.loads(managed_path.read_text()) + recreated.update(applied_picker) + managed_path.write_text(json.dumps(recreated), encoding="utf-8") + + assert claude.revert_managed_settings() == "unchanged" + written = json.loads(managed_path.read_text()) + assert {key: written[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == applied_picker + assert written["companyPolicy"] == "preserve" + assert set(json.loads(metadata_path.read_text())["leases"]) == {"private"} + + def test_managed_marker_failure_retains_manifest_for_exact_before_retry( + self, monkeypatch, tmp_path + ): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "companyPolicy": "preserve", + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + state = { + "workspace": "https://workspace.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + claude.write_tool_config(state, state["claude_static_models"][0]) + applied = json.loads(managed_path.read_text()) + applied_picker = {key: applied[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} + mark_written = claude._mark_picker_revert_written + failed = False + + def fail_first_managed_marker(scope, path): + nonlocal failed + if scope == "managed" and not failed: + failed = True + raise RuntimeError("injected managed marker failure") + mark_written(scope, path) + + monkeypatch.setattr(claude, "_mark_picker_revert_written", fail_first_managed_marker) + with pytest.raises(RuntimeError, match="injected managed marker failure"): + claude.revert_managed_settings() + + manifest = json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text()) + assert "claude" in manifest["files"] + assert ( + json.loads(metadata_path.read_text())["leases"]["managed"]["pending"]["phase"] + == "reverting" + ) + recreated = json.loads(managed_path.read_text()) + recreated.update(applied_picker) + managed_path.write_text(json.dumps(recreated), encoding="utf-8") + + assert ( + claude.revert_managed_settings() == "ucode entries removed; external changes preserved" + ) + written = json.loads(managed_path.read_text()) + assert {key: written[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == { + key: baseline[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS + } + assert written["companyPolicy"] == "preserve" + assert ( + "claude" + not in json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text())["files"] + ) + assert set(json.loads(metadata_path.read_text())["leases"]) == {"private"} + + @pytest.mark.parametrize( + "source_kwargs", + [ + {"provider": "main.default.anthropic"}, + {"parent_schema": "main.managed_models"}, + ], + ids=["provider", "model-location"], + ) + def test_fresh_workspace_restores_preexisting_picker( + self, monkeypatch, tmp_path, source_kwargs + ): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + original_picker = { + "availableModels": ["user-model"], + "enforceAvailableModels": False, + "modelPicker": { + "replaceBuiltInOptions": False, + "options": [{"model": "user-model", "label": "User"}], + }, + } + original = {**original_picker, "companyPolicy": {"keep": True}} + private_path.write_text(json.dumps(original), encoding="utf-8") + managed_path.write_text(json.dumps(original), encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + workspace_a = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + } + + claude.write_tool_config(workspace_a, static_models[0]) + assert json.loads(private_path.read_text())["availableModels"] == static_models + assert json.loads(managed_path.read_text())["availableModels"] == static_models + + workspace_b = {"workspace": "https://workspace-b.example.com", "codex_models": []} + claude.write_tool_config(workspace_b, None, **source_kwargs) + + for path in (private_path, managed_path): + written = json.loads(path.read_text()) + assert {key: written[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == ( + original_picker + ) + assert written["companyPolicy"] == {"keep": True} + assert not metadata_path.exists() + manifest = json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text()) + owned_paths = manifest["files"]["claude"]["owned_paths"] + assert not any([key] in owned_paths for key in claude.CLAUDE_MANAGED_PICKER_KEYS) + + def test_reacquisition_restores_each_scopes_new_picker_baseline(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + picker_one = { + "availableModels": ["enterprise-one"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise-one"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps({**picker_one, "companyPolicy": "keep"}), encoding="utf-8") + + static_models = ["system.ai.claude-opus-4-8"] + static_state = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + } + scoped_state = {"workspace": "https://workspace-b.example.com", "codex_models": []} + claude.write_tool_config(dict(static_state), static_models[0]) + metadata = json.loads(metadata_path.read_text()) + assert set(metadata["leases"]) == {"private", "managed"} + for lease in metadata["leases"].values(): + assert set(lease["keys"]) == set(claude.CLAUDE_MANAGED_PICKER_KEYS) + manifest = json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text()) + assert not any( + [key] in manifest["files"]["claude"]["owned_paths"] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + ) + claude.write_tool_config(dict(scoped_state), None, parent_schema="main.models") + + picker_two = { + "availableModels": ["enterprise-two"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "enterprise-two"}]}, + } + for path in (private_path, managed_path): + settings = json.loads(path.read_text()) + settings.update(picker_two) + path.write_text(json.dumps(settings), encoding="utf-8") + + claude.write_tool_config(dict(static_state), static_models[0]) + claude.write_tool_config(dict(scoped_state), None, provider="main.default.anthropic") + + for path in (private_path, managed_path): + settings = json.loads(path.read_text()) + assert {key: settings[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == picker_two + assert settings["companyPolicy"] == "keep" + assert not metadata_path.exists() + + def test_private_static_update_failure_keeps_committed_baseline(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + state_a = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + state_b = { + "workspace": "https://workspace-b.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-sonnet-4-6"], + } + claude.write_tool_config(dict(state_a), state_a["claude_static_models"][0]) + write_json = claude.write_json_file + + def fail_private_write(path, payload): + if path == private_path: + raise RuntimeError("injected private write failure") + write_json(path, payload) + + monkeypatch.setattr(claude, "write_json_file", fail_private_write) + with pytest.raises(RuntimeError, match="injected private write failure"): + claude.write_tool_config(dict(state_b), state_b["claude_static_models"][0]) + assert ( + json.loads(private_path.read_text())["availableModels"] + == state_a["claude_static_models"] + ) + + monkeypatch.setattr(claude, "write_json_file", write_json) + claude.write_tool_config(dict(state_b), state_b["claude_static_models"][0]) + claude.write_tool_config( + {"workspace": "https://workspace-c.example.com", "codex_models": []}, + None, + parent_schema="main.models", + ) + settings = json.loads(private_path.read_text()) + assert {key: settings[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == baseline + + def test_managed_static_update_failure_keeps_committed_baseline(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + state_a = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-opus-4-8"], + } + state_b = { + "workspace": "https://workspace-b.example.com", + "codex_models": [], + "claude_static_models": ["system.ai.claude-sonnet-4-6"], + } + claude.write_tool_config(dict(state_a), state_a["claude_static_models"][0]) + replace = managed_files._sudo_replace + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda path, text: (_ for _ in ()).throw(PermissionError("injected managed failure")), + ) + with pytest.raises(managed_files.ManagedFileWriteUnavailable): + claude.write_tool_config(dict(state_b), state_b["claude_static_models"][0]) + assert ( + json.loads(managed_path.read_text())["availableModels"] + == state_a["claude_static_models"] + ) + + monkeypatch.setattr(managed_files, "_sudo_replace", replace) + claude.write_tool_config(dict(state_b), state_b["claude_static_models"][0]) + claude.write_tool_config( + {"workspace": "https://workspace-c.example.com", "codex_models": []}, + None, + parent_schema="main.models", + ) + settings = json.loads(managed_path.read_text()) + assert {key: settings[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == baseline + + def test_managed_release_repairs_snapshot_before_clearing_pending_lease( + self, monkeypatch, tmp_path + ): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "availableModels": ["enterprise"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + claude.write_tool_config( + { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + }, + static_models[0], + ) + scoped_state = {"workspace": "https://workspace-b.example.com", "codex_models": []} + record_last_applied = managed_files._record_last_applied + + def fail_record(*args, **kwargs): + raise RuntimeError("injected managed metadata failure") + + monkeypatch.setattr(managed_files, "_record_last_applied", fail_record) + with pytest.raises(RuntimeError, match="injected managed metadata failure"): + claude.write_tool_config(dict(scoped_state), None, parent_schema="main.models") + assert "pending" in json.loads(metadata_path.read_text())["leases"]["managed"] + assert { + key: json.loads(managed_path.read_text())[key] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + } == baseline + + monkeypatch.setattr(managed_files, "_record_last_applied", record_last_applied) + claude.write_tool_config(dict(scoped_state), None, parent_schema="main.models") + + assert not metadata_path.exists() + manifest = json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text()) + entry = manifest["files"]["claude"] + assert not any([key] in entry["owned_paths"] for key in claude.CLAUDE_MANAGED_PICKER_KEYS) + assert (managed_files.MANAGED_BACKUP_DIR / entry["last_applied_file"]).read_text() == ( + managed_path.read_text() + ) + + def test_managed_revert_uses_latest_acquisition_baseline(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + picker_one = { + "availableModels": ["enterprise-one"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise-one"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(picker_one), encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + static_state = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + } + scoped_state = {"workspace": "https://workspace-b.example.com", "codex_models": []} + claude.write_tool_config(dict(static_state), static_models[0]) + claude.write_tool_config(dict(scoped_state), None, parent_schema="main.models") + + picker_two = { + "availableModels": ["enterprise-two"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "enterprise-two"}]}, + } + managed = json.loads(managed_path.read_text()) + managed.update(picker_two) + managed_path.write_text(json.dumps(managed), encoding="utf-8") + claude.write_tool_config(dict(static_state), static_models[0]) + + assert claude.revert_managed_settings() == "restored" + restored = json.loads(managed_path.read_text()) + assert {key: restored[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == picker_two + metadata = json.loads(metadata_path.read_text()) + assert set(metadata["leases"]) == {"private"} + + def test_managed_revert_preserves_whole_drifted_picker_group(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + baseline = { + "availableModels": ["enterprise-one"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "enterprise-one"}]}, + } + for path in (private_path, managed_path): + path.write_text(json.dumps(baseline), encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + claude.write_tool_config( + { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + }, + static_models[0], + ) + drift = { + "availableModels": ["enterprise-drift"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "enterprise-drift"}]}, + } + managed = json.loads(managed_path.read_text()) + managed.update(drift) + managed_path.write_text(json.dumps(managed), encoding="utf-8") + + assert ( + claude.revert_managed_settings() == "ucode entries removed; external changes preserved" + ) + restored = json.loads(managed_path.read_text()) + assert {key: restored[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == drift + metadata = json.loads(metadata_path.read_text()) + assert set(metadata["leases"]) == {"private"} + + def test_fresh_workspace_migrates_legacy_picker_ownership(self, monkeypatch, tmp_path): + private_path, managed_path, backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + original_picker = { + "availableModels": ["user-model"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "user-model"}]}, + } + stale_picker = { + "availableModels": ["system.ai.claude-opus-4-8"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "system.ai.claude-opus-4-8"}]}, + } + managed_path.write_text(json.dumps(original_picker), encoding="utf-8") + managed_files.reconcile_managed_file( + managed_path, + json.dumps(stale_picker), + tool="claude", + display="Claude Code", + owned_paths=[[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS], + ) + private_path.write_text(json.dumps(stale_picker), encoding="utf-8") + backup_path.write_text(json.dumps(original_picker), encoding="utf-8") + monkeypatch.setattr( + claude, + "load_full_state", + lambda: { + "workspaces": { + "https://workspace-a.example.com": { + "managed_configs": { + "claude": {"keys": [[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS]} + } + } + } + }, + ) + + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.managed_models", + ) + + written = json.loads(private_path.read_text()) + assert {key: written[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == original_picker + assert json.loads(backup_path.read_text()) == original_picker + assert not metadata_path.exists() + + def test_legacy_migration_preserves_private_picker_changed_after_ucode( + self, monkeypatch, tmp_path + ): + private_path, managed_path, backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + stale_picker = { + "availableModels": ["system.ai.claude-opus-4-8"], + "enforceAvailableModels": True, + "modelPicker": {"options": [{"model": "system.ai.claude-opus-4-8"}]}, + } + managed_path.write_text("{}", encoding="utf-8") + managed_files.reconcile_managed_file( + managed_path, + json.dumps(stale_picker), + tool="claude", + display="Claude Code", + owned_paths=[[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS], + ) + edited_picker = { + "availableModels": ["user-edited-model"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "user-edited-model"}]}, + } + private_path.write_text(json.dumps(edited_picker), encoding="utf-8") + backup_path.write_text("{}", encoding="utf-8") + monkeypatch.setattr( + claude, + "load_full_state", + lambda: { + "workspaces": { + "https://workspace-a.example.com": { + "managed_configs": { + "claude": {"keys": [[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS]} + } + } + } + }, + ) + + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.managed_models", + ) + + written_private = json.loads(private_path.read_text()) + assert { + key: written_private[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS + } == edited_picker + assert ( + not set(claude.CLAUDE_MANAGED_PICKER_KEYS) & json.loads(managed_path.read_text()).keys() + ) + assert not metadata_path.exists() + + def test_fresh_workspace_preserves_post_ucode_picker_edits(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + private_path.write_text('{"companyPolicy": {"keep": true}}', encoding="utf-8") + managed_path.write_text('{"companyPolicy": {"keep": true}}', encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + claude.write_tool_config( + { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + }, + static_models[0], + ) + edited_picker = { + "availableModels": ["enterprise-model"], + "enforceAvailableModels": True, + "modelPicker": { + "replaceBuiltInOptions": True, + "options": [{"model": "enterprise-model", "label": "Enterprise"}], + }, + } + for path in (private_path, managed_path): + settings = json.loads(path.read_text()) + settings.update(edited_picker) + path.write_text(json.dumps(settings), encoding="utf-8") + + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.managed_models", + ) + claude.write_tool_config( + {"workspace": "https://workspace-c.example.com", "codex_models": []}, + None, + provider="main.default.anthropic", + ) + + for path in (private_path, managed_path): + written = json.loads(path.read_text()) + assert {key: written[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == edited_picker + assert written["companyPolicy"] == {"keep": True} + assert not metadata_path.exists() + manifest = json.loads(managed_files.MANAGED_BACKUP_MANIFEST_PATH.read_text()) + owned_paths = manifest["files"]["claude"]["owned_paths"] + assert not any([key] in owned_paths for key in claude.CLAUDE_MANAGED_PICKER_KEYS) + + def test_required_managed_picker_cleanup_cannot_fall_back_noninteractively( + self, monkeypatch, tmp_path + ): + private_path, managed_path, _backup_path, _metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + static_models = ["system.ai.claude-opus-4-8"] + claude.write_tool_config( + { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + }, + static_models[0], + ) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: False) + + with pytest.raises(RuntimeError, match="cannot be applied non-interactively"): + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.managed_models", + ) + + assert "availableModels" not in json.loads(private_path.read_text()) + assert json.loads(managed_path.read_text())["availableModels"] == static_models + + def test_static_picker_sidecar_failure_is_retry_safe(self, monkeypatch, tmp_path): + private_path, _managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + user_picker = { + "availableModels": ["user-model"], + "enforceAvailableModels": False, + "modelPicker": {"options": [{"model": "user-model"}]}, + } + private_path.write_text(json.dumps(user_picker), encoding="utf-8") + static_models = ["system.ai.claude-opus-4-8"] + static_state = { + "workspace": "https://workspace-a.example.com", + "codex_models": [], + "claude_static_models": static_models, + } + begin_picker_transition = claude._begin_picker_transition + attempts = 0 + + def fail_first_sidecar_write(*args, **kwargs): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("injected sidecar failure") + begin_picker_transition(*args, **kwargs) + + monkeypatch.setattr(claude, "_begin_picker_transition", fail_first_sidecar_write) + + with pytest.raises(RuntimeError, match="injected sidecar failure"): + claude.write_tool_config(dict(static_state), static_models[0]) + assert { + key: json.loads(private_path.read_text())[key] + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + } == user_picker + + claude.write_tool_config(dict(static_state), static_models[0]) + assert json.loads(private_path.read_text())["availableModels"] == static_models + assert metadata_path.exists() + + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.managed_models", + ) + written = json.loads(private_path.read_text()) + assert {key: written[key] for key in claude.CLAUDE_MANAGED_PICKER_KEYS} == user_picker + assert not metadata_path.exists() + + def test_picker_sidecar_supports_each_scope_and_both(self, monkeypatch, tmp_path): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + entries = { + key: {"original_exists": False, "last_applied": f"applied-{key}"} + for key in claude.CLAUDE_MANAGED_PICKER_KEYS + } + + claude._save_private_picker_management(entries) + assert set(json.loads(metadata_path.read_text())["leases"]) == {"private"} + claude._save_managed_picker_management(managed_path, entries) + assert set(json.loads(metadata_path.read_text())["leases"]) == {"private", "managed"} + claude._save_private_picker_management({}) + assert set(json.loads(metadata_path.read_text())["leases"]) == {"managed"} + claude._save_managed_picker_management(managed_path, {}) + assert not metadata_path.exists() + + metadata_path.write_text( + json.dumps({"version": claude.CLAUDE_PICKER_MANAGEMENT_VERSION, "leases": {}}), + encoding="utf-8", + ) + with pytest.raises(RuntimeError, match="Invalid Claude picker metadata"): + claude._load_picker_management() + + def test_malformed_private_picker_metadata_fails_without_changing_settings( + self, monkeypatch, tmp_path + ): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + private_path.write_text('{"user": "keep"}', encoding="utf-8") + managed_path.write_text('{"companyPolicy": "keep"}', encoding="utf-8") + metadata_path.write_text("{", encoding="utf-8") + + with pytest.raises(RuntimeError, match="Cannot parse Claude picker metadata"): + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.managed_models", + ) + + assert json.loads(private_path.read_text()) == {"user": "keep"} + assert json.loads(managed_path.read_text()) == {"companyPolicy": "keep"} + + @pytest.mark.parametrize("managed_keys", [[], ["availableModels"]], ids=["empty", "partial"]) + def test_incomplete_private_picker_metadata_fails_without_changing_settings( + self, monkeypatch, tmp_path, managed_keys + ): + private_path, managed_path, _backup_path, metadata_path = self._patch_files( + monkeypatch, tmp_path + ) + private_path.write_text('{"user": "keep"}', encoding="utf-8") + managed_path.write_text('{"companyPolicy": "keep"}', encoding="utf-8") + metadata_path.write_text( + json.dumps( + { + "version": claude.CLAUDE_PICKER_MANAGEMENT_VERSION, + "leases": { + "private": { + "path": str(private_path), + "keys": { + key: {"original_exists": False, "last_applied": []} + for key in managed_keys + }, + } + }, + } + ), + encoding="utf-8", + ) + + with pytest.raises(RuntimeError, match="Invalid Claude picker metadata"): + claude.write_tool_config( + {"workspace": "https://workspace-b.example.com", "codex_models": []}, + None, + parent_schema="main.managed_models", + ) + + assert json.loads(private_path.read_text()) == {"user": "keep"} + assert json.loads(managed_path.read_text()) == {"companyPolicy": "keep"} + + class TestAddClaudeMcpServer: def test_registers_stdio_proxy_command(self, monkeypatch): calls: list[dict] = [] @@ -1451,21 +2913,31 @@ def test_write_tool_config_completes_when_mcp_registration_fails(self, monkeypat # block the rest of `ucode claude` setup (state save, managed-key # marking, etc.) from completing. monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - monkeypatch.setattr(claude, "read_json_safe", lambda path: {}) - monkeypatch.setattr(claude, "write_json_file", lambda path, payload: None) + _patch_private_json_store(monkeypatch, {}) saved: list[dict] = [] monkeypatch.setattr(claude, "save_state", lambda state: saved.append(state)) + monkeypatch.setattr( + claude, + "load_state", + lambda: json.loads(json.dumps(saved[-1])) if saved else {}, + ) monkeypatch.setattr(claude, "remove_claude_mcp_server", lambda name, scope: False) + attempts = 0 + def boom(name, entry, scope=claude.MCP_USER_SCOPE): + nonlocal attempts + attempts += 1 raise RuntimeError("Failed to add MCP server 'web_search' via claude CLI.") monkeypatch.setattr(claude, "add_claude_mcp_server", boom) state = {"workspace": WS, "codex_models": ["databricks-gpt-5"]} result = claude.write_tool_config(state, "databricks-claude-sonnet-4") + result = claude.write_tool_config(result, "databricks-claude-sonnet-4") assert saved, "save_state should still be called when MCP registration fails" assert result["workspace"] == WS + assert attempts == 2 class TestClaudeLaunch: @@ -1684,13 +3156,12 @@ class TestWriteToolConfigPrunesStaleModelEnv: def _patch(self, monkeypatch, existing_settings): monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - monkeypatch.setattr(claude, "read_json_safe", lambda path: existing_settings) written: dict = {} def fake_write(path, payload): written["payload"] = payload - monkeypatch.setattr(claude, "write_json_file", fake_write) + _patch_private_json_store(monkeypatch, existing_settings, fake_write) monkeypatch.setattr(claude, "save_state", lambda state: None) monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) return written @@ -1883,9 +3354,8 @@ def test_malformed_file_json_raises(self, tmp_path, monkeypatch): class TestClaudeSmartRouting: def _capture_write(self, monkeypatch, existing, written): monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - monkeypatch.setattr(claude, "read_json_safe", lambda path: existing) - monkeypatch.setattr( - claude, "write_json_file", lambda path, payload: written.append(payload) + _patch_private_json_store( + monkeypatch, existing, lambda _path, payload: written.append(payload) ) monkeypatch.setattr(claude, "save_state", lambda state: None) monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) @@ -2062,3 +3532,201 @@ def test_reconfigure_does_not_back_up_generated_file(self, tmp_path, monkeypatch claude.write_tool_config(state, "databricks-claude-sonnet-4") assert not (tmp_path / "backup.json").exists() + + +def test_write_tool_config_serializes_both_scopes_and_state_save(monkeypatch): + calls: list[str] = [] + calls_lock = threading.Lock() + b_started = threading.Event() + errors: list[BaseException] = [] + + def record(label: str) -> None: + with calls_lock: + calls.append(label) + + def private(state, _overlay, _compose): + label = state["workspace"] + record(f"{label}-private") + if label == "A": + assert b_started.wait(timeout=5) + + def managed(state, *_args, **_kwargs): + record(f"{state['workspace']}-managed") + + monkeypatch.setattr(claude, "_reconcile_private_settings", private) + monkeypatch.setattr(claude, "_reconcile_managed_settings", managed) + monkeypatch.setattr(claude, "save_state", lambda state: record(f"{state['workspace']}-save")) + + def configure(label: str) -> None: + try: + if label == "B": + b_started.set() + claude.write_tool_config( + {"workspace": label, "codex_models": []}, + "system.ai.claude-opus-4-8", + ) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + thread_a = threading.Thread(target=configure, args=("A",)) + thread_b = threading.Thread(target=configure, args=("B",)) + thread_a.start() + thread_b.start() + thread_a.join(timeout=5) + thread_b.join(timeout=5) + + assert not thread_a.is_alive() + assert not thread_b.is_alive() + assert errors == [] + assert calls == [ + "A-private", + "A-managed", + "A-save", + "B-private", + "B-managed", + "B-save", + ] + + +def test_web_search_registration_converges_to_latest_generation(monkeypatch): + persisted: dict = {} + persisted_lock = threading.Lock() + b_saved = threading.Event() + release_a = threading.Event() + picker_released = threading.Event() + registrations: list[tuple[str, str | None]] = [] + errors: list[BaseException] = [] + + monkeypatch.setattr(claude, "_reconcile_private_settings", lambda *_args, **_kwargs: None) + monkeypatch.setattr(claude, "_reconcile_managed_settings", lambda *_args, **_kwargs: None) + monkeypatch.setattr(claude, "_web_search_mcp_is_current", lambda *_args: False) + + def save(state): + with persisted_lock: + persisted.clear() + persisted.update(json.loads(json.dumps(state))) + if state.get("web_search_model") == "model-b": + b_saved.set() + + def load(): + with persisted_lock: + return json.loads(json.dumps(persisted)) + + def register(_workspace, model, profile=None): + registrations.append((model, profile)) + if model == "model-a": + + def probe_lock(): + with claude._picker_management_lock(): + picker_released.set() + + probe = threading.Thread(target=probe_lock) + probe.start() + assert picker_released.wait(timeout=5) + probe.join(timeout=5) + assert release_a.wait(timeout=5) + return True + + monkeypatch.setattr(claude, "save_state", save) + monkeypatch.setattr(claude, "load_state", load) + monkeypatch.setattr(claude, "_register_web_search_mcp", register) + + def configure(model: str, profile: str) -> None: + try: + claude.write_tool_config( + { + "workspace": WS, + "profile": profile, + "web_search_model": model, + "codex_models": [f"fallback-{model}"], + }, + "system.ai.claude-opus-4-8", + ) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + thread_a = threading.Thread(target=configure, args=("model-a", "profile-a")) + thread_b = threading.Thread(target=configure, args=("model-b", "profile-b")) + thread_a.start() + assert picker_released.wait(timeout=5) + thread_b.start() + assert b_saved.wait(timeout=5) + release_a.set() + thread_a.join(timeout=5) + thread_b.join(timeout=5) + + assert not thread_a.is_alive() + assert not thread_b.is_alive() + assert errors == [] + assert registrations == [("model-a", "profile-a"), ("model-b", "profile-b")] + latest = load() + assert latest["web_search_model"] == "model-b" + assert latest["profile"] == "profile-b" + assert latest[claude.WEB_SEARCH_MCP_STATE_KEY] == claude._web_search_mcp_entry( + WS, "model-b", "profile-b" + ) + + +def test_managed_web_search_registration_uses_persisted_generation(monkeypatch): + persisted: dict = {} + registrations: list[tuple[str, str | None]] = [] + + monkeypatch.setattr(claude, "_reconcile_private_settings", lambda *_args, **_kwargs: None) + monkeypatch.setattr(claude, "_reconcile_managed_settings", lambda *_args, **_kwargs: None) + monkeypatch.setattr(claude, "_web_search_mcp_is_current", lambda *_args: False) + + def save(state): + persisted.clear() + persisted.update(json.loads(json.dumps(developer_state_from_resolved(state)))) + + monkeypatch.setattr(claude, "save_state", save) + monkeypatch.setattr(claude, "load_state", lambda: json.loads(json.dumps(persisted))) + monkeypatch.setattr( + claude, + "_register_web_search_mcp", + lambda _workspace, model, profile=None: registrations.append((model, profile)) or True, + ) + + claude.write_tool_config( + { + "workspace": WS, + "profile": "managed-profile", + "web_search_model": "managed-model", + "codex_models": ["managed-fallback"], + MANAGED_OVERLAY_KEY: { + "profile": "developer-profile", + "web_search_model": None, + "codex_models": [], + }, + }, + "system.ai.claude-opus-4-8", + ) + + expected = claude._web_search_mcp_entry(WS, "managed-model", "managed-profile") + assert registrations == [("managed-model", "managed-profile")] + assert persisted[claude.WEB_SEARCH_MCP_STATE_KEY] == expected + assert persisted[claude.WEB_SEARCH_MCP_GENERATION_KEY] + + +def test_web_search_registration_rejects_token_preserving_state_change(monkeypatch): + entry = claude._web_search_mcp_entry(WS, "model-a", "profile-a") + state = { + "workspace": WS, + "profile": "profile-a", + "web_search_model": "model-a", + claude.WEB_SEARCH_MCP_GENERATION_KEY: "same-token", + } + latest = {**state, "profile": "profile-b"} + registrations: list[str] = [] + + monkeypatch.setattr(claude, "load_state", lambda: dict(latest)) + monkeypatch.setattr( + claude, + "_register_web_search_mcp", + lambda *_args, **_kwargs: registrations.append("registered") or True, + ) + + result = claude._register_web_search_for_current_generation(state, entry) + + assert result == latest + assert registrations == [] diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index a95a1fe61..d2c283407 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -1041,9 +1041,10 @@ def test_catalog_write_reports_path_on_failure(self, tmp_path, monkeypatch): path = tmp_path / "models.json" monkeypatch.setattr(codex.os, "replace", lambda *args: (_ for _ in ()).throw(OSError())) - with pytest.raises(RuntimeError, match=str(path)): + with pytest.raises(RuntimeError, match=str(path)) as exc_info: codex._write_model_catalog(path, {"models": [{"slug": "gpt-mps"}]}) + assert str(path) in str(exc_info.value.__cause__) assert list(tmp_path.glob(".models.json.*.tmp")) == [] def test_catalog_cleanup_does_not_mask_write_failure(self, tmp_path, monkeypatch): diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 4333bffe1..1734f5893 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -3,6 +3,7 @@ from __future__ import annotations import subprocess +import threading from contextlib import contextmanager import pytest @@ -774,6 +775,36 @@ def test_empty_selection_preserves_existing(self, monkeypatch): result = configure_selected_tools(state, []) assert result["available_tools"] == ["codex"] + def test_concurrent_newer_state_is_not_overwritten_after_claude_returns(self, monkeypatch): + saved: list[dict] = [] + saved_lock = threading.Lock() + + def save(state): + with saved_lock: + saved.append(dict(state)) + + def configure(_tool, state, _provider): + assert state["available_tools"] == ["claude"] + save(state) + writer = threading.Thread( + target=save, + args=({**state, "generation": "newer"},), + ) + writer.start() + writer.join(timeout=5) + return state + + monkeypatch.setattr(agents_mod, "_configure_one", configure) + monkeypatch.setattr(agents_mod, "save_state", save) + monkeypatch.setattr(agents_mod, "install_databricks_ai_tools_for_agents", lambda *_: None) + + configure_selected_tools( + {"workspace": "https://x.databricks.com", "generation": "original"}, + ["claude"], + ) + + assert saved[-1]["generation"] == "newer" + class TestConfiguredPaths: def test_claude_reports_its_settings_file_home_abbreviated(self): diff --git a/tests/test_cli.py b/tests/test_cli.py index db245d396..849952350 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -8,6 +8,7 @@ import os import re import subprocess +import threading import time import tomllib from importlib import metadata @@ -2705,6 +2706,10 @@ def test_reverts_mcp_configs_before_clearing_state(self): ), ), patch("ucode.cli.clear_state", side_effect=lambda: cleared.append(True)), + patch( + "ucode.cli.claude_agent.revert_settings", + return_value=("unchanged", False), + ), ): result = runner.invoke(app, ["revert"]) @@ -2713,6 +2718,44 @@ def test_reverts_mcp_configs_before_clearing_state(self): assert cleared == [True] assert "Claude Code MCP config: restored" in result.output + def test_global_picker_ownership_removes_claude_file_for_fresh_workspace(self): + with ( + patch("ucode.cli.load_state", return_value={"workspace": MINIMAL_STATE["workspace"]}), + patch("ucode.cli.restore_file", return_value=False) as mock_restore, + patch("ucode.cli.revert_mcp_configs", return_value={}), + patch("ucode.cli.clear_state"), + patch( + "ucode.cli.claude_agent.revert_settings", + return_value=("unchanged", True), + ), + patch("ucode.cli.codex_agent.revert_managed_config", return_value="unchanged"), + ): + result = runner.invoke(app, ["revert"]) + + assert result.exit_code == 0, result.output + assert all( + call.args[0] != cli_mod.TOOL_SPECS["claude"]["config_path"] + for call in mock_restore.call_args_list + ) + assert "Claude Code config: restored" in result.output + + def test_global_picker_metadata_survives_when_claude_file_was_not_reverted(self): + with ( + patch("ucode.cli.load_state", return_value={"workspace": MINIMAL_STATE["workspace"]}), + patch("ucode.cli.restore_file", return_value=False), + patch("ucode.cli.revert_mcp_configs", return_value={}), + patch("ucode.cli.clear_state"), + patch( + "ucode.cli.claude_agent.revert_settings", + return_value=("unchanged", False), + ), + patch("ucode.cli.codex_agent.revert_managed_config", return_value="unchanged"), + ): + result = runner.invoke(app, ["revert"]) + + assert result.exit_code == 0, result.output + assert "Claude Code config: unchanged" in result.output + class TestDoctorCommand: def test_invokes_doctor(self): @@ -2929,14 +2972,15 @@ def test_model_location_autoconfigure_is_launch_scoped(self, tool): "codex_models": [], } saved_states = [] + + def configure_and_save(_tool, current, **_kwargs): + saved_states.append(json.loads(json.dumps(current))) + return current + with ( patch("ucode.cli.load_state", return_value=existing_state), patch("ucode.cli.configure_shared_state", return_value=configured_state) as mock_shared, - patch("ucode.cli.configure_tool", return_value=configured_state) as mock_configure, - patch( - "ucode.cli.save_state", - side_effect=lambda state: saved_states.append(json.loads(json.dumps(state))), - ), + patch("ucode.cli.configure_tool", side_effect=configure_and_save) as mock_configure, ): cli_mod._auto_configure_tool(tool, model_location="main.models") @@ -3001,6 +3045,11 @@ def test_true_first_run_configures_managed_location_without_persisting_policy(se } } persisted: list[dict] = [] + + def configure_and_save(_tool, current, *_args, **_kwargs): + cli_mod.save_state(current) + return current + with ( patch("ucode.cli.ensure_bootstrap_dependencies"), patch("ucode.cli.load_state", return_value={}), @@ -3019,9 +3068,7 @@ def test_true_first_run_configures_managed_location_without_persisting_policy(se patch("ucode.cli.ensure_provider_state", return_value=configured_state), patch("ucode.cli._fetch_budget_recommendation", return_value=None), patch("ucode.cli.resolve_launch_model") as mock_resolve, - patch( - "ucode.cli.configure_tool", side_effect=lambda *args, **kwargs: args[1] - ) as mock_configure, + patch("ucode.cli.configure_tool", side_effect=configure_and_save) as mock_configure, patch( "ucode.cli.save_state", side_effect=lambda state: persisted.append( @@ -3156,6 +3203,39 @@ def test_cursor_launch_uses_unity_gateway_branding(): class TestConfigureAgentFlag: + def test_model_location_does_not_overwrite_newer_state_after_agent_returns(self): + saved: list[dict] = [] + saved_lock = threading.Lock() + + def save(state): + with saved_lock: + saved.append(dict(state)) + + def configure(tool, state, **kwargs): + assert tool == "claude" + assert kwargs == {"parent_schema": "main.models"} + assert state["available_tools"] == ["claude"] + save(state) + writer = threading.Thread( + target=save, + args=({**state, "generation": "newer"},), + ) + writer.start() + writer.join(timeout=5) + return state + + with ( + patch("ucode.cli.configure_tool", side_effect=configure), + patch("ucode.cli.save_state", side_effect=save), + ): + cli_mod._configure_location_backed_tool( + {"workspace": MINIMAL_STATE["workspace"], "generation": "original"}, + "claude", + "main.models", + ) + + assert saved[-1]["generation"] == "newer" + def test_help_lists_model_location(self): result = runner.invoke(app, ["configure", "--help"]) @@ -3363,6 +3443,11 @@ def test_model_location_configures_fallback_tools_with_fresh_managed_gemini_sour managed = {"enabled_agents": enabled_agents} state = {**MINIMAL_STATE, "available_tools": []} saved_states: list[dict] = [] + + def write_and_save(_tool, current, **_kwargs): + saved_states.append(json.loads(json.dumps(current))) + return current + with ( patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), @@ -3373,9 +3458,7 @@ def test_model_location_configures_fallback_tools_with_fresh_managed_gemini_sour patch("ucode.cli.refresh_managed_config", return_value=(managed, False)), patch("ucode.cli.check_gateway_endpoint", return_value=True) as mock_available, patch("ucode.cli.resolve_state", wraps=cli_mod.resolve_state) as mock_resolve, - patch( - "ucode.cli.configure_tool", side_effect=lambda tool, current, **kwargs: current - ) as mock_write, + patch("ucode.cli.configure_tool", side_effect=write_and_save) as mock_write, patch( "ucode.cli.configure_selected_tools", side_effect=lambda current, *args, **kwargs: current, @@ -3431,6 +3514,10 @@ def test_later_fallback_save_does_not_persist_earlier_managed_agent_overlay(self def configure_managed(current, *args, **kwargs): return {**current, "available_tools": ["claude"]} + def write_and_save(_tool, current, **_kwargs): + saved_states.append(json.loads(json.dumps(current))) + return current + with ( patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), @@ -3442,7 +3529,7 @@ def configure_managed(current, *args, **kwargs): patch( "ucode.cli.configure_selected_tools", side_effect=configure_managed ) as mock_managed, - patch("ucode.cli.configure_tool", side_effect=lambda tool, current, **kwargs: current), + patch("ucode.cli.configure_tool", side_effect=write_and_save), patch( "ucode.cli.save_state", side_effect=lambda current: saved_states.append(json.loads(json.dumps(current))), @@ -3492,6 +3579,11 @@ def test_managed_gemini_provider_does_not_restore_fallback_agent_providers(self) }, } saved_states: list[dict] = [] + + def write_and_save(_tool, current, **_kwargs): + saved_states.append(json.loads(json.dumps(current))) + return current + with ( patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), @@ -3503,7 +3595,7 @@ def test_managed_gemini_provider_does_not_restore_fallback_agent_providers(self) "ucode.cli.configure_selected_tools", side_effect=lambda current, *args, **kwargs: current, ) as mock_managed, - patch("ucode.cli.configure_tool", side_effect=lambda tool, current, **kwargs: current), + patch("ucode.cli.configure_tool", side_effect=write_and_save), patch( "ucode.cli.save_state", side_effect=lambda current: saved_states.append(json.loads(json.dumps(current))), @@ -4089,6 +4181,7 @@ def test_model_location_persists_only_after_each_agent_succeeds(self, monkeypatc def configure(tool, candidate, **kwargs): if tool == "codex": raise RuntimeError("codex write failed") + cli_mod.save_state(candidate) return candidate monkeypatch.setattr(cli_mod, "configure_tool", configure) diff --git a/tests/test_config_io.py b/tests/test_config_io.py index 9ee1c58d7..91d14c9ae 100644 --- a/tests/test_config_io.py +++ b/tests/test_config_io.py @@ -165,6 +165,42 @@ def test_write_json_file(self, tmp_path): write_json_file(p, {"a": 1}) data = json.loads(p.read_text()) assert data == {"a": 1} + assert p.stat().st_mode & 0o777 == 0o600 + + def test_write_json_file_replace_failure_preserves_original(self, tmp_path, monkeypatch): + p = tmp_path / "out.json" + original = b'{"original": true}\n' + p.write_bytes(original) + original_entries = set(tmp_path.iterdir()) + + def fail_replace(source, target): + assert target == p + raise OSError("injected replace failure") + + monkeypatch.setattr(config_io.os, "replace", fail_replace) + + with pytest.raises(RuntimeError, match="Failed to write config file"): + write_json_file(p, {"replacement": True}) + + assert p.read_bytes() == original + assert set(tmp_path.iterdir()) == original_entries + + def test_write_json_file_fsync_failure_removes_temp_file(self, tmp_path, monkeypatch): + p = tmp_path / "out.json" + original = b'{"original": true}\n' + p.write_bytes(original) + original_entries = set(tmp_path.iterdir()) + + def fail_fsync(_fd): + raise OSError("injected fsync failure") + + monkeypatch.setattr(config_io.os, "fsync", fail_fsync) + + with pytest.raises(RuntimeError, match=f"Failed to write config file: {p}"): + write_json_file(p, {"replacement": True}) + + assert p.read_bytes() == original + assert set(tmp_path.iterdir()) == original_entries def test_write_json_file_dry_run_no_write(self, tmp_path): set_dry_run(True) diff --git a/tests/test_managed_files.py b/tests/test_managed_files.py index aa047dffa..9ffe94ccc 100644 --- a/tests/test_managed_files.py +++ b/tests/test_managed_files.py @@ -91,6 +91,14 @@ def run(command, **kwargs): class TestManagedFileLifecycle: + def test_private_json_metadata_is_atomic_and_owner_only(self, tmp_path, backup_dir): + path = tmp_path / "metadata.json" + + managed_files.write_private_json_file(path, {"version": 1}) + + assert json.loads(path.read_text()) == {"version": 1} + assert path.stat().st_mode & 0o777 == 0o600 + def test_dry_run_does_not_write_or_backup(self, tmp_path, backup_dir, monkeypatch): path = tmp_path / "managed.json" config_io.set_dry_run(True) @@ -228,6 +236,193 @@ def test_unchanged_file_never_creates_backup(self, tmp_path, backup_dir, monkeyp assert result == "unchanged" assert not backup_dir.exists() + def test_unchanged_retry_repairs_last_applied_after_metadata_failure( + self, tmp_path, backup_dir, monkeypatch + ): + path = tmp_path / "managed.json" + path.write_text('{"picker": "enterprise"}\n', encoding="utf-8") + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + record_last_applied = managed_files._record_last_applied + + def fail_record(*args, **kwargs): + raise RuntimeError("injected metadata failure") + + monkeypatch.setattr(managed_files, "_record_last_applied", fail_record) + + with pytest.raises(RuntimeError, match="injected metadata failure"): + managed_files.reconcile_managed_file( + path, + '{"picker": "ucode"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["picker"]], + conditional_owned_paths=[["picker"]], + ) + assert json.loads(path.read_text()) == {"picker": "ucode"} + assert (backup_dir / "manifest.json").exists() + with pytest.raises(RuntimeError, match="last-applied snapshot is missing"): + managed_files.revert_managed_file( + "claude", + display="Claude Code", + parser=json.loads, + dumper=lambda doc: json.dumps(doc) + "\n", + ) + assert (backup_dir / "manifest.json").exists() + + monkeypatch.setattr(managed_files, "_record_last_applied", record_last_applied) + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda *args: pytest.fail("unchanged retry must not rewrite settings"), + ) + assert ( + managed_files.reconcile_managed_file( + path, + '{"picker": "ucode"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["picker"]], + conditional_owned_paths=[["picker"]], + repair_last_applied=True, + ) + == "unchanged" + ) + + restored, restored_paths = managed_files.restore_unchanged_managed_paths( + "claude", + path, + {"picker": "ucode"}, + [["picker"]], + parser=json.loads, + ) + assert restored == {"picker": "enterprise"} + assert restored_paths == [["picker"]] + + def test_unchanged_file_with_backup_is_not_claimed_without_repair_proof( + self, tmp_path, backup_dir, monkeypatch + ): + path = tmp_path / "managed.json" + path.write_text('{"picker": "ucode"}\n', encoding="utf-8") + managed_files._ensure_backup("claude", path, '{"picker": "enterprise"}\n') + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda *args: pytest.fail("unchanged settings must not be rewritten"), + ) + + result = managed_files.reconcile_managed_file( + path, + '{"picker": "ucode"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["picker"]], + ) + + assert result == "unchanged" + entry = json.loads((backup_dir / "manifest.json").read_text())["files"]["claude"] + assert "last_applied_file" not in entry + assert entry["owned_paths"] == [] + + def test_failed_manifest_switch_keeps_previous_snapshot_readable( + self, tmp_path, backup_dir, monkeypatch + ): + path = tmp_path / "managed.json" + path.write_text('{"value": "original"}\n', encoding="utf-8") + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + managed_files.reconcile_managed_file( + path, + '{"value": "one"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["value"]], + ) + first_manifest = json.loads((backup_dir / "manifest.json").read_text()) + first_file = first_manifest["files"]["claude"]["last_applied_file"] + first_text = (backup_dir / first_file).read_text() + write_manifest = managed_files._write_manifest + monkeypatch.setattr( + managed_files, + "_write_manifest", + lambda manifest: (_ for _ in ()).throw(RuntimeError("injected manifest failure")), + ) + + with pytest.raises(RuntimeError, match="injected manifest failure"): + managed_files.reconcile_managed_file( + path, + '{"value": "two"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["value"]], + ) + + persisted = json.loads((backup_dir / "manifest.json").read_text()) + assert persisted["files"]["claude"]["last_applied_file"] == first_file + assert (backup_dir / first_file).read_text() == first_text + + monkeypatch.setattr(managed_files, "_write_manifest", write_manifest) + assert ( + managed_files.reconcile_managed_file( + path, + '{"value": "two"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["value"]], + repair_last_applied=True, + ) + == "unchanged" + ) + repaired = json.loads((backup_dir / "manifest.json").read_text()) + second_file = repaired["files"]["claude"]["last_applied_file"] + assert second_file != first_file + assert (backup_dir / second_file).read_text() == '{"value": "two"}\n' + + def test_old_snapshot_cleanup_failure_is_harmless(self, tmp_path, backup_dir, monkeypatch): + path = tmp_path / "managed.json" + path.write_text('{"value": "original"}\n', encoding="utf-8") + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + managed_files.reconcile_managed_file( + path, + '{"value": "one"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["value"]], + ) + first_manifest = json.loads((backup_dir / "manifest.json").read_text()) + first_file = first_manifest["files"]["claude"]["last_applied_file"] + first_path = backup_dir / first_file + unlink = type(first_path).unlink + + def fail_old_snapshot(candidate, *args, **kwargs): + if candidate == first_path: + raise OSError("injected cleanup failure") + return unlink(candidate, *args, **kwargs) + + monkeypatch.setattr(type(first_path), "unlink", fail_old_snapshot) + + managed_files.reconcile_managed_file( + path, + '{"value": "two"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["value"]], + ) + + persisted = json.loads((backup_dir / "manifest.json").read_text()) + assert persisted["files"]["claude"]["last_applied_file"] != first_file + assert first_path.exists() + def test_verified_check_uses_fingerprint(self, tmp_path): path = tmp_path / "managed.json" path.write_text("current", encoding="utf-8") @@ -238,6 +433,94 @@ def test_verified_check_uses_fingerprint(self, tmp_path): path.write_text("changed-content", encoding="utf-8") assert managed_files.managed_file_is_verified(state, "claude", path) is False + def test_unchanged_managed_paths_match_integrity_checked_snapshot( + self, tmp_path, backup_dir, monkeypatch + ): + path = tmp_path / "managed.json" + path.write_text('{"picker": "enterprise", "other": 1}\n', encoding="utf-8") + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + managed_files.reconcile_managed_file( + path, + '{"picker": "ucode", "other": 1}\n', + tool="claude", + display="Claude Code", + owned_paths=[["picker"], ["other"]], + ) + + assert managed_files.unchanged_managed_paths( + "claude", + path, + {"picker": "ucode", "other": 1}, + [["picker"], ["unowned"]], + parser=json.loads, + ) == [["picker"]] + restored, restored_paths = managed_files.restore_unchanged_managed_paths( + "claude", + path, + {"picker": "ucode", "other": 1}, + [["picker"]], + parser=json.loads, + ) + assert restored == {"picker": "enterprise", "other": 1} + assert restored_paths == [["picker"]] + + path.write_text('{"picker": "enterprise", "other": 1}\n', encoding="utf-8") + assert ( + managed_files.unchanged_managed_paths( + "claude", + path, + {"picker": "enterprise", "other": 1}, + [["picker"]], + parser=json.loads, + ) + == [] + ) + + def test_conditional_ownership_is_relinquished_after_external_change( + self, tmp_path, backup_dir, monkeypatch + ): + path = tmp_path / "managed.json" + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + managed_files.reconcile_managed_file( + path, + '{"picker": "ucode", "other": 1}\n', + tool="claude", + display="Claude Code", + owned_paths=[["picker"], ["other"]], + conditional_owned_paths=[["picker"]], + ) + path.write_text('{"picker": "enterprise", "other": 1}\n', encoding="utf-8") + + managed_files.reconcile_managed_file( + path, + '{"picker": "enterprise", "other": 2}\n', + tool="claude", + display="Claude Code", + owned_paths=[["other"]], + conditional_owned_paths=[["picker"]], + ) + + manifest = json.loads((backup_dir / "manifest.json").read_text()) + assert ["picker"] not in manifest["files"]["claude"]["owned_paths"] + assert ( + managed_files.unchanged_managed_paths( + "claude", + path, + {"picker": "enterprise", "other": 2}, + [["picker"]], + parser=json.loads, + ) + == [] + ) + def test_revert_restores_exact_original(self, tmp_path, backup_dir, monkeypatch): path = tmp_path / "managed.json" path.write_text('{"enterprise": true}\n', encoding="utf-8") @@ -265,6 +548,96 @@ def test_revert_restores_exact_original(self, tmp_path, backup_dir, monkeypatch) assert path.read_text() == '{"enterprise": true}\n' assert json.loads((backup_dir / "manifest.json").read_text())["files"] == {} + def test_revert_manifest_switch_failure_keeps_snapshots_for_retry( + self, tmp_path, backup_dir, monkeypatch + ): + path = tmp_path / "managed.json" + path.write_text('{"value": "original"}\n', encoding="utf-8") + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + managed_files.reconcile_managed_file( + path, + '{"value": "ucode"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["value"]], + ) + persisted = json.loads((backup_dir / "manifest.json").read_text()) + entry = persisted["files"]["claude"] + referenced = [ + backup_dir / entry[key] for key in ("backup_file", "last_applied_file") if key in entry + ] + write_manifest = managed_files._write_manifest + monkeypatch.setattr( + managed_files, + "_write_manifest", + lambda manifest: (_ for _ in ()).throw(RuntimeError("injected manifest failure")), + ) + + with pytest.raises(RuntimeError, match="injected manifest failure"): + managed_files.revert_managed_file( + "claude", + display="Claude Code", + parser=json.loads, + dumper=lambda doc: json.dumps(doc) + "\n", + ) + + assert "claude" in json.loads((backup_dir / "manifest.json").read_text())["files"] + assert all(snapshot.exists() for snapshot in referenced) + + monkeypatch.setattr(managed_files, "_write_manifest", write_manifest) + assert ( + managed_files.revert_managed_file( + "claude", + display="Claude Code", + parser=json.loads, + dumper=lambda doc: json.dumps(doc) + "\n", + ) + == "ucode entries removed; external changes preserved" + ) + assert json.loads((backup_dir / "manifest.json").read_text())["files"] == {} + + def test_revert_snapshot_gc_failure_is_harmless(self, tmp_path, backup_dir, monkeypatch): + path = tmp_path / "managed.json" + path.write_text('{"value": "original"}\n', encoding="utf-8") + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + managed_files.reconcile_managed_file( + path, + '{"value": "ucode"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["value"]], + ) + entry = json.loads((backup_dir / "manifest.json").read_text())["files"]["claude"] + stale_snapshot = backup_dir / entry["last_applied_file"] + unlink = type(stale_snapshot).unlink + + def fail_snapshot_gc(candidate, *args, **kwargs): + if candidate == stale_snapshot: + raise OSError("injected snapshot GC failure") + return unlink(candidate, *args, **kwargs) + + monkeypatch.setattr(type(stale_snapshot), "unlink", fail_snapshot_gc) + + assert ( + managed_files.revert_managed_file( + "claude", + display="Claude Code", + parser=json.loads, + dumper=lambda doc: json.dumps(doc) + "\n", + ) + == "restored" + ) + assert json.loads((backup_dir / "manifest.json").read_text())["files"] == {} + assert stale_snapshot.exists() + def test_revert_removes_file_created_by_ucode(self, tmp_path, backup_dir, monkeypatch): path = tmp_path / "managed.json" monkeypatch.setattr( @@ -321,6 +694,34 @@ def test_revert_preserves_external_changes(self, tmp_path, backup_dir, monkeypat assert result == "ucode entries removed; external changes preserved" assert json.loads(path.read_text()) == {"enterprise": "new-policy", "new": True} + def test_revert_exact_last_snapshot_preserves_unowned_paths( + self, tmp_path, backup_dir, monkeypatch + ): + path = tmp_path / "managed.json" + path.write_text('{"picker": "enterprise", "owned": "old"}\n', encoding="utf-8") + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + managed_files.reconcile_managed_file( + path, + '{"picker": "relinquished", "owned": "new"}\n', + tool="claude", + display="Claude Code", + owned_paths=[["owned"]], + ) + + result = managed_files.revert_managed_file( + "claude", + display="Claude Code", + parser=json.loads, + dumper=lambda doc: json.dumps(doc, sort_keys=True) + "\n", + ) + + assert result == "restored" + assert json.loads(path.read_text()) == {"picker": "relinquished", "owned": "old"} + def test_reconcile_retries_exact_mdm_restore_once(self, tmp_path, backup_dir, monkeypatch): path = tmp_path / "managed.json" path.write_text('{"enterprise": true}\n', encoding="utf-8")