From 699e90035d8e6c27c0537ea1648fe766f7acae83 Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 18:22:39 -0600 Subject: [PATCH 1/8] fix(controller): lease writable workspaces --- .../prompts/orchestrator/system_prompt.md | 4 +- zenith/src/zenith_harness/cli.py | 36 ++++- zenith/src/zenith_harness/controller.py | 55 ++++++- zenith/src/zenith_harness/server.py | 82 ++++++++-- zenith/src/zenith_harness/storage.py | 148 ++++++++++++++++++ zenith/tests/test_cli.py | 84 ++++++++++ zenith/tests/test_server.py | 136 ++++++++++++++++ zenith/tests/test_storage.py | 93 ++++++++++- 8 files changed, 619 insertions(+), 19 deletions(-) diff --git a/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md b/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md index 22f6589..8acd560 100644 --- a/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md +++ b/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md @@ -25,6 +25,7 @@ Drive the mission only through the orchestrator runtime tools: - Quiescent `mission_running`: when no runnable work remains and evidence supports closure, call `end_mission`. - `attention_needed`: read every open item and its evidence, decide exactly once per item with `decide_attention`, then call `advance_project`. - User-requested or justified mission cancellation: call `abort_project(project_id, reason)`. +- Intentional handoff to another root orchestrator: stop all mutation, then call `release_project(project_id)`. Never share one writable workspace between root orchestrators. - `done`, `failed`, `aborted`: terminal states. Do not continue mission work unless the user starts a new scope or explicitly asks for forensic inspection. `submit_plan` and `decide_attention` persist state; they do not dispatch work. `advance_project` dispatches workers, validators, merge work, and gate evaluation according to runtime state. `end_mission` requests runtime closure and terminal review; call it only after planning, validation, gates, evidence review, and open attention handling support closure. @@ -361,8 +362,9 @@ Orchestrator tools: - `decide_attention(project_id, decisions)`: resolve every open attention item with exactly one decision, then return to runtime flow. Call `advance_project` afterward. - `end_mission(project_id)`: request runtime closure and terminal review only after work is quiescent and evidence supports closure. - `abort_project(project_id, reason)`: terminal cancellation with a recorded reason. +- `release_project(project_id)`: explicitly release persistent workspace ownership so another root orchestrator can take over. Read-only inspection by other sessions remains available before release. -Every orchestrator tool returns an envelope with `projectId`, `state`, `projectRoot`, `harnessRoot`, and `dag`. Some lifecycle/decision tools intentionally return `dag=null`; `submit_plan` and `advance_project` return a compact frontier view; `inspect_project` returns the full task-list view. Trust the envelope, runtime files, and returned paths over session memory. +Lifecycle and inspection tools return an envelope with `projectId`, `state`, `projectRoot`, `harnessRoot`, and `dag`. Some lifecycle/decision tools intentionally return `dag=null`; `submit_plan` and `advance_project` return a compact frontier view; `inspect_project` returns the full task-list view. `release_project` returns a release acknowledgement. Trust the envelope, runtime files, and returned paths over session memory. `submit_plan` accepts a `TaskList`: diff --git a/zenith/src/zenith_harness/cli.py b/zenith/src/zenith_harness/cli.py index 095918e..e72d99c 100644 --- a/zenith/src/zenith_harness/cli.py +++ b/zenith/src/zenith_harness/cli.py @@ -250,9 +250,17 @@ def inspect_tasks_cmd(project_id: str, mission_id: str | None) -> None: @cli.command("abort-project") @click.argument("project_id") @click.option("--reason", required=True) -def abort_project_cmd(project_id: str, reason: str) -> None: +@click.option( + "--owner-id", + default=None, + help=( + "Stable controller/session id. Defaults to ZENITH_CONTROLLER_ID, " + "CODEX_THREAD_ID, or a Claude session id." + ), +) +def abort_project_cmd(project_id: str, reason: str, owner_id: str | None) -> None: """Mark a project Aborted (CLI-side: preserves tasks.json + attempts/).""" - from .controller import ProjectController + from .controller import ProjectController, ToolError from .dispatcher import MockDispatcher, MockTerminalReviewer from .models import TerminalReviewHandoff, WorkHandoff @@ -262,7 +270,29 @@ def abort_project_cmd(project_id: str, reason: str) -> None: MockDispatcher(lambda r: WorkHandoff(node_id=r.task.id, done=False, report="aborted")), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - env = controller.abort_project(project_id, reason) + resolved_owner = owner_id or next( + ( + os.environ[name] + for name in ( + "ZENITH_CONTROLLER_ID", + "CODEX_THREAD_ID", + "CLAUDE_CODE_SESSION_ID", + "CLAUDE_SESSION_ID", + ) + if os.environ.get(name) + ), + None, + ) + if resolved_owner is None: + raise click.ClickException( + "abort-project requires --owner-id or ZENITH_CONTROLLER_ID " + "(CODEX_THREAD_ID and Claude session ids are also accepted)" + ) + try: + controller.claim_project(project_id, resolved_owner) + env = controller.abort_project(project_id, reason) + except ToolError as exc: + raise click.ClickException(str(exc)) from exc click.echo(f"Aborted {project_id}: state={env.state.state}") diff --git a/zenith/src/zenith_harness/controller.py b/zenith/src/zenith_harness/controller.py index f7aeeeb..0bd6ab7 100644 --- a/zenith/src/zenith_harness/controller.py +++ b/zenith/src/zenith_harness/controller.py @@ -1,4 +1,4 @@ -"""ProjectController — routes the 7 orchestrator MCP tools. +"""ProjectController -- routes the 8 orchestrator MCP tools. See `specs/task_list/PRODUCT.md`. The controller owns: - envelope construction @@ -31,7 +31,7 @@ TaskListPatch, TaskStateFile, ) -from .storage import ProjectStore +from .storage import ProjectStore, WorkspaceLeaseConflict from .task_list_patch import apply_patch from .task_validation import ( ValidationError, @@ -71,10 +71,29 @@ def __init__( # Tool methods # ------------------------------------------------------------------ - def start_project(self, brief: str, workspace_dir: str) -> Envelope: + def start_project( + self, brief: str, workspace_dir: str, owner_id: str | None = None + ) -> Envelope: if not brief.strip(): raise ToolError("invalid_brief", "brief is empty") - record = self.store.create_project(brief, workspace_dir) + project_id = self.store.generate_project_id(brief) + if owner_id is not None: + try: + self.store.claim_workspace_lease_for_workspace( + workspace_dir, project_id, owner_id + ) + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc + try: + record = self.store.create_project( + brief, workspace_dir, project_id=project_id + ) + except Exception: + if owner_id is not None: + self.store.release_workspace_lease_for_workspace( + workspace_dir, project_id, owner_id + ) + raise mission_id = self.store.generate_mission_id(1) record.current_mission_id = mission_id self.store.save_project(record) @@ -83,6 +102,34 @@ def start_project(self, brief: str, workspace_dir: str) -> Envelope: ) return self._build_envelope(record.id, dag_mode="none") + def claim_project(self, project_id: str, owner_id: str) -> None: + try: + self.store.claim_workspace_lease(project_id, owner_id) + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc + + def release_project(self, project_id: str, owner_id: str) -> None: + record = self.store.load_project(project_id) + state = self.store.load_state(project_id) + mission_id = self._current_mission_id(record, state) + if mission_id is not None: + task_state = self.store.load_task_state(project_id, mission_id) + running = sorted( + task_id + for task_id, entry in task_state.tasks.items() + if entry.status == "running" + ) + if running: + raise ToolError( + "workspace_busy", + "cannot release controller ownership while tasks are running: " + + ", ".join(running), + ) + try: + self.store.release_workspace_lease(project_id, owner_id) + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc + def submit_plan(self, project_id: str, task_list: TaskList) -> Envelope: state = self._require_state(project_id) if not isinstance(state, MissionPlanning): diff --git a/zenith/src/zenith_harness/server.py b/zenith/src/zenith_harness/server.py index 19c37eb..2223f89 100644 --- a/zenith/src/zenith_harness/server.py +++ b/zenith/src/zenith_harness/server.py @@ -9,6 +9,7 @@ import asyncio import logging import os +from uuid import uuid4 from typing import Annotated, Any from fastmcp import Context, FastMCP @@ -39,7 +40,7 @@ def create_orchestrator_server( config: HarnessConfig, controller: ProjectController | None = None, ) -> FastMCP: - """7 orchestrator tools, registered on a stdio MCP server.""" + """8 orchestrator tools, registered on a stdio MCP server.""" if controller is None: from .dispatcher import MockDispatcher, MockTerminalReviewer @@ -55,8 +56,9 @@ def create_orchestrator_server( name="zenith", instructions=( "Mission orchestration harness. Mode: orchestrator. " - "7 tools: start_project, submit_plan, advance_project, " + "8 tools: start_project, submit_plan, advance_project, " "end_mission, decide_attention, inspect_project, abort_project. " + "release_project performs an explicit controller handoff. " "Lifecycle: plan with submit_plan, run with advance_project, " "request closure with end_mission, resolve attention with decide_attention, " "then call advance_project again." @@ -100,15 +102,29 @@ def create_terminal_reviewer_server() -> FastMCP: def _register_orchestrator_tools(mcp: FastMCP, controller: ProjectController) -> None: - # Per-project lock around mutating controller calls. The thread hop in each - # tool prevents event-loop blocking, but two same-project tool calls could - # otherwise race on disk state (attention, attempts, task-state, tasks). - # docs/v5/07-runtime-architecture.md §9 declares concurrent same-project - # operations undefined behavior; this lock serializes them defensively - # without requiring host coordination. `inspect_project` is read-only and - # stays uncontended. + # The in-process lock serializes overlapping calls in this server. The + # persistent workspace lease below it prevents a second server process or + # root orchestrator from mutating the same workspace. `inspect_project` is + # read-only and stays uncontended so recovery audits remain possible. project_locks: dict[str, asyncio.Lock] = {} locks_guard = asyncio.Lock() + server_instance_owner = f"server:{uuid4()}" + explicit_owner = ( + os.environ.get("ZENITH_CONTROLLER_ID") + or os.environ.get("CODEX_THREAD_ID") + or os.environ.get("CLAUDE_CODE_SESSION_ID") + or os.environ.get("CLAUDE_SESSION_ID") + ) + + def _controller_owner(ctx: Context | None) -> str: + if explicit_owner: + return explicit_owner + if ctx is not None and ctx.request_context is not None: + return ctx.session_id + return server_instance_owner + + def _claim(project_id: str, owner_id: str) -> None: + controller.claim_project(project_id, owner_id) async def _project_lock(project_id: str) -> asyncio.Lock: async with locks_guard: @@ -132,11 +148,17 @@ async def start_project( workspace_dir: Annotated[ str, Field(description="Absolute path to the user's workspace.") ], + ctx: Context | None = None, ) -> dict[str, Any]: # No per-project lock: project_id does not exist until the call returns. try: return _to_payload( - await asyncio.to_thread(controller.start_project, brief, workspace_dir) + await asyncio.to_thread( + controller.start_project, + brief, + workspace_dir, + _controller_owner(ctx), + ) ) except ToolError as exc: return _to_payload(exc) @@ -162,9 +184,12 @@ async def submit_plan( ) ), ], + ctx: Context | None = None, ) -> dict[str, Any]: async with await _project_lock(project_id): try: + owner_id = _controller_owner(ctx) + await asyncio.to_thread(_claim, project_id, owner_id) return _to_payload( await asyncio.to_thread(controller.submit_plan, project_id, task_list) ) @@ -192,6 +217,8 @@ async def advance_project( ) -> dict[str, Any]: async with await _project_lock(project_id): try: + owner_id = _controller_owner(ctx) + await asyncio.to_thread(_claim, project_id, owner_id) return _to_payload( await asyncio.to_thread(controller.advance_project, project_id, max_steps) ) @@ -211,9 +238,12 @@ async def advance_project( ) async def end_mission( project_id: Annotated[str, Field(description="Project id.")], + ctx: Context | None = None, ) -> dict[str, Any]: async with await _project_lock(project_id): try: + owner_id = _controller_owner(ctx) + await asyncio.to_thread(_claim, project_id, owner_id) return _to_payload( await asyncio.to_thread(controller.end_mission, project_id) ) @@ -236,9 +266,12 @@ async def decide_attention( decisions: Annotated[ list[Decision], Field(description="One Decision per open attention item.") ], + ctx: Context | None = None, ) -> dict[str, Any]: async with await _project_lock(project_id): try: + owner_id = _controller_owner(ctx) + await asyncio.to_thread(_claim, project_id, owner_id) return _to_payload( await asyncio.to_thread(controller.decide_attention, project_id, decisions) ) @@ -272,15 +305,44 @@ async def inspect_project( async def abort_project( project_id: Annotated[str, Field(description="Project id.")], reason: Annotated[str, Field(description="Why we are aborting.")], + ctx: Context | None = None, ) -> dict[str, Any]: async with await _project_lock(project_id): try: + owner_id = _controller_owner(ctx) + await asyncio.to_thread(_claim, project_id, owner_id) return _to_payload( await asyncio.to_thread(controller.abort_project, project_id, reason) ) except ToolError as exc: return _to_payload(exc) + @mcp.tool( + name="release_project", + description=( + "Release this controller's persistent workspace lease for an explicit " + "handoff to another root orchestrator. Do not call while a mutating " + "tool is running. Read-only inspection never requires ownership." + ), + ) + async def release_project( + project_id: Annotated[str, Field(description="Project id.")], + ctx: Context | None = None, + ) -> dict[str, Any]: + async with await _project_lock(project_id): + try: + owner_id = _controller_owner(ctx) + await asyncio.to_thread( + controller.release_project, project_id, owner_id + ) + return { + "released": True, + "projectId": project_id, + "ownerId": owner_id, + } + except ToolError as exc: + return _to_payload(exc) + # --------------------------------------------------------------------------- # Worker tool diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index da1d9c0..0a5373f 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -16,9 +16,12 @@ from __future__ import annotations import json +import hashlib import os import re import shutil +import socket +import time from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path @@ -116,6 +119,20 @@ class AttemptRecord: path: Path +@dataclass(frozen=True) +class WorkspaceLease: + owner_id: str + project_id: str + workspace_dir: str + claimed_at: str + last_seen_at: str + path: Path + + +class WorkspaceLeaseConflict(RuntimeError): + """A different controller already owns the writable workspace.""" + + # --------------------------------------------------------------------------- # ProjectStore # --------------------------------------------------------------------------- @@ -157,6 +174,137 @@ def mission_runtime_dir(self, project_id: str, mission_id: str) -> Path: attempts/*.json).""" return self.zenith_runtime_dir(project_id) / "missions" / mission_id + def workspace_lease_path(self, project_id: str) -> Path: + workspace = self.workspace_dir(project_id).expanduser().resolve() + return self._workspace_lease_path(workspace) + + def _workspace_lease_path(self, workspace: str | Path) -> Path: + workspace = Path(workspace).expanduser().resolve() + key = hashlib.sha256(str(workspace).encode("utf-8")).hexdigest()[:24] + return self.config.harness_home / "leases" / "workspaces" / f"{key}.lease" / "owner.json" + + def claim_workspace_lease( + self, project_id: str, owner_id: str + ) -> WorkspaceLease: + return self.claim_workspace_lease_for_workspace( + self.workspace_dir(project_id), project_id, owner_id + ) + + def claim_workspace_lease_for_workspace( + self, + workspace_dir: str | Path, + project_id: str, + owner_id: str, + ) -> WorkspaceLease: + """Claim persistent write ownership for a project's exact workspace. + + The directory creation is the atomic compare-and-set. Ownership persists + across MCP calls and processes until the owning controller explicitly + releases it. Different workspaces remain independently claimable. + """ + if not owner_id.strip(): + raise ValueError("owner_id is empty") + workspace = Path(workspace_dir).expanduser().resolve() + lease_path = self._workspace_lease_path(workspace) + lease_dir = lease_path.parent + lease_dir.parent.mkdir(parents=True, exist_ok=True) + created = False + try: + lease_dir.mkdir() + created = True + except FileExistsError: + pass + + if not created: + current = self._load_workspace_lease_record(lease_path) + if current is None: + for _ in range(5): + time.sleep(0.05) + current = self._load_workspace_lease_record(lease_path) + if current is not None: + break + if current is None: + raise WorkspaceLeaseConflict( + f"workspace lease is incomplete: {lease_path}" + ) + if ( + current.owner_id != owner_id + or current.project_id != project_id + ): + raise WorkspaceLeaseConflict( + "workspace already owned by " + f"owner={current.owner_id} project={current.project_id}; " + f"requesting owner={owner_id} project={project_id}; " + f"lease={lease_path}" + ) + claimed_at = current.claimed_at + else: + claimed_at = utc_now_iso() + + now = utc_now_iso() + payload = { + "schema_version": 1, + "owner_id": owner_id, + "project_id": project_id, + "workspace_dir": str(workspace), + "host": socket.gethostname(), + "claimed_at": claimed_at, + "last_seen_at": now, + } + atomic_write_json(lease_path, payload) + return WorkspaceLease( + owner_id=owner_id, + project_id=project_id, + workspace_dir=str(workspace), + claimed_at=claimed_at, + last_seen_at=now, + path=lease_path, + ) + + def release_workspace_lease(self, project_id: str, owner_id: str) -> None: + self.release_workspace_lease_for_workspace( + self.workspace_dir(project_id), project_id, owner_id + ) + + def release_workspace_lease_for_workspace( + self, + workspace_dir: str | Path, + project_id: str, + owner_id: str, + ) -> None: + lease_path = self._workspace_lease_path(workspace_dir) + if not lease_path.parent.exists(): + return + current = self._load_workspace_lease_record(lease_path) + if current is None: + raise WorkspaceLeaseConflict( + f"workspace lease is incomplete: {lease_path}" + ) + if current.owner_id != owner_id or current.project_id != project_id: + raise WorkspaceLeaseConflict( + "workspace already owned by " + f"owner={current.owner_id} project={current.project_id}; " + f"requesting release owner={owner_id} project={project_id}; " + f"lease={lease_path}" + ) + lease_path.unlink() + lease_path.parent.rmdir() + + @staticmethod + def _load_workspace_lease_record(path: Path) -> WorkspaceLease | None: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + return WorkspaceLease( + owner_id=str(payload["owner_id"]), + project_id=str(payload["project_id"]), + workspace_dir=str(payload["workspace_dir"]), + claimed_at=str(payload["claimed_at"]), + last_seen_at=str(payload["last_seen_at"]), + path=path, + ) + except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError): + return None + # ------------------------------------------------------------------ # Project lifecycle # ------------------------------------------------------------------ diff --git a/zenith/tests/test_cli.py b/zenith/tests/test_cli.py index 87e640e..e8fa6e0 100644 --- a/zenith/tests/test_cli.py +++ b/zenith/tests/test_cli.py @@ -9,6 +9,10 @@ from click.testing import CliRunner from zenith_harness.cli import cli +from zenith_harness.config import HarnessConfig +from zenith_harness.controller import ProjectController +from zenith_harness.dispatcher import MockDispatcher, MockTerminalReviewer +from zenith_harness.models import TerminalReviewHandoff, WorkHandoff @pytest.fixture @@ -35,6 +39,86 @@ def _expected_mcp_server_args() -> list[str]: ] +def _create_owned_project(workspace: Path, owner_id: str) -> str: + config = HarnessConfig.discover() + controller = ProjectController( + config, + MockDispatcher( + lambda request: WorkHandoff( + node_id=request.task.id, done=False, report="unused" + ) + ), + MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), + ) + envelope = controller.start_project("Owned.", str(workspace), owner_id) + return envelope.projectId + + +class TestAbortProjectOwnership: + def test_requires_stable_controller_identity( + self, + runner: CliRunner, + workspace: Path, + env: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + for name in ( + "ZENITH_CONTROLLER_ID", + "CODEX_THREAD_ID", + "CLAUDE_CODE_SESSION_ID", + "CLAUDE_SESSION_ID", + ): + monkeypatch.delenv(name, raising=False) + project_id = _create_owned_project(workspace, "owner-a") + + result = runner.invoke( + cli, ["abort-project", project_id, "--reason", "no owner"] + ) + + assert result.exit_code != 0 + assert "--owner-id or ZENITH_CONTROLLER_ID" in result.output + + def test_non_owner_cannot_abort( + self, runner: CliRunner, workspace: Path, env: dict[str, str] + ) -> None: + project_id = _create_owned_project(workspace, "owner-a") + + result = runner.invoke( + cli, + [ + "abort-project", + project_id, + "--reason", + "wrong owner", + "--owner-id", + "owner-b", + ], + ) + + assert result.exit_code != 0 + assert "workspace_owned" in result.output + + def test_owner_can_abort( + self, runner: CliRunner, workspace: Path, env: dict[str, str] + ) -> None: + project_id = _create_owned_project(workspace, "owner-a") + + result = runner.invoke( + cli, + [ + "abort-project", + project_id, + "--reason", + "intentional", + "--owner-id", + "owner-a", + ], + ) + + assert result.exit_code == 0, result.output + assert f"Aborted {project_id}" in result.output + + class TestInit: def test_stages_host_agent_surface_only( self, runner: CliRunner, workspace: Path, env: dict[str, str] diff --git a/zenith/tests/test_server.py b/zenith/tests/test_server.py index 4b5533c..b12e8d3 100644 --- a/zenith/tests/test_server.py +++ b/zenith/tests/test_server.py @@ -67,6 +67,7 @@ async def test_orchestrator_tools_registered(config: HarnessConfig) -> None: "decide_attention", "inspect_project", "abort_project", + "release_project", } @@ -226,6 +227,141 @@ def responder(req): await server.call_tool("inspect_project", {"project_id": pid}) +@pytest.mark.asyncio +async def test_second_server_cannot_mutate_owned_workspace( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + first = create_orchestrator_server(config) + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-b") + second = create_orchestrator_server(config) + started = await first.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + + blocked = await second.call_tool( + "abort_project", {"project_id": pid, "reason": "should not mutate"} + ) + assert blocked.structured_content["error"] == "workspace_owned" + assert "owner=" in blocked.structured_content["message"] + + +@pytest.mark.asyncio +async def test_second_server_can_inspect_owned_workspace( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + first = create_orchestrator_server(config) + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-b") + second = create_orchestrator_server(config) + started = await first.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + + inspected = await second.call_tool("inspect_project", {"project_id": pid}) + assert inspected.structured_content["projectId"] == pid + assert "error" not in inspected.structured_content + + +@pytest.mark.asyncio +async def test_explicit_release_allows_controller_handoff( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + first = create_orchestrator_server(config) + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-b") + second = create_orchestrator_server(config) + started = await first.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + + released = await first.call_tool("release_project", {"project_id": pid}) + assert released.structured_content["released"] is True + taken_over = await second.call_tool( + "abort_project", {"project_id": pid, "reason": "intentional handoff"} + ) + assert taken_over.structured_content["state"]["state"] == "aborted" + + +@pytest.mark.asyncio +async def test_non_owner_cannot_release_project( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + first = create_orchestrator_server(config) + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-b") + second = create_orchestrator_server(config) + started = await first.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + + blocked = await second.call_tool("release_project", {"project_id": pid}) + assert blocked.structured_content["error"] == "workspace_owned" + + +@pytest.mark.asyncio +async def test_owner_cannot_release_project_while_task_is_running( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + server = create_orchestrator_server(config) + started = await server.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + store = ProjectStore(config) + contract_dir = store.ensure_contract_dir(pid, "mission-001") + (contract_dir / "VAL-001.md").write_text("# VAL-001\n") + await server.call_tool( + "submit_plan", + { + "project_id": pid, + "task_list": { + "tasks": [ + { + "id": "w1", + "type": "work", + "body": "do", + "targets": ["VAL-001"], + "skill": "s", + "depends_on": [], + } + ] + }, + }, + ) + task_state = store.load_task_state(pid, "mission-001") + task_state.set_status("w1", "running") + store.save_task_state(pid, "mission-001", task_state) + + blocked = await server.call_tool("release_project", {"project_id": pid}) + assert blocked.structured_content["error"] == "workspace_busy" + assert "w1" in blocked.structured_content["message"] + + +@pytest.mark.asyncio +async def test_second_server_cannot_start_another_project_in_owned_workspace( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + first = create_orchestrator_server(config) + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-b") + second = create_orchestrator_server(config) + await first.call_tool( + "start_project", {"brief": "First.", "workspace_dir": str(workspace)} + ) + + blocked = await second.call_tool( + "start_project", {"brief": "Second.", "workspace_dir": str(workspace)} + ) + assert blocked.structured_content["error"] == "workspace_owned" + assert len(ProjectStore(config).list_projects()) == 1 + + # --------------------------------------------------------------------------- # Regression: dispatcher that calls asyncio.run() must not crash the MCP # event loop. Reproduces: diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index fdaf380..7b331c1 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -1,6 +1,8 @@ """Storage layer tests. See specs/memory_v2/PRODUCT.md for layout.""" from __future__ import annotations +import json +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import pytest @@ -17,7 +19,12 @@ ValidationItem, WorkHandoff, ) -from zenith_harness.storage import ProjectStore, slugify, utc_now_filesafe +from zenith_harness.storage import ( + ProjectStore, + WorkspaceLeaseConflict, + slugify, + utc_now_filesafe, +) @pytest.fixture @@ -418,3 +425,87 @@ def test_writes_closeout( text = path.read_text() assert "status: done" in text and "Everything shipped." in text assert path == store.mission_dir("p1", "mission-001") / "closeout.md" + + +class TestWorkspaceLease: + def test_same_owner_can_reenter( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + first = store.claim_workspace_lease("p1", "owner-a") + second = store.claim_workspace_lease("p1", "owner-a") + assert second.owner_id == first.owner_id + assert second.claimed_at == first.claimed_at + + def test_different_owner_is_blocked( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + store.claim_workspace_lease("p1", "owner-a") + with pytest.raises(WorkspaceLeaseConflict, match="owner-a"): + store.claim_workspace_lease("p1", "owner-b") + + def test_same_owner_cannot_claim_second_project_for_workspace( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("one", workspace, project_id="p1") + store.create_project("two", workspace, project_id="p2") + store.claim_workspace_lease("p1", "owner-a") + with pytest.raises(WorkspaceLeaseConflict, match="p1"): + store.claim_workspace_lease("p2", "owner-a") + + def test_different_workspaces_allow_parallel_owners( + self, store: ProjectStore, tmp_path: Path + ) -> None: + one = tmp_path / "one" + two = tmp_path / "two" + one.mkdir() + two.mkdir() + store.create_project("one", one, project_id="p1") + store.create_project("two", two, project_id="p2") + assert store.claim_workspace_lease("p1", "owner-a").owner_id == "owner-a" + assert store.claim_workspace_lease("p2", "owner-b").owner_id == "owner-b" + + def test_non_owner_cannot_release( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + store.claim_workspace_lease("p1", "owner-a") + with pytest.raises(WorkspaceLeaseConflict, match="owner-a"): + store.release_workspace_lease("p1", "owner-b") + + def test_release_allows_explicit_handoff( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + store.claim_workspace_lease("p1", "owner-a") + store.release_workspace_lease("p1", "owner-a") + assert store.claim_workspace_lease("p1", "owner-b").owner_id == "owner-b" + + def test_record_contains_forensic_identity( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + payload = json.loads(lease.path.read_text(encoding="utf-8")) + assert payload["owner_id"] == "owner-a" + assert payload["project_id"] == "p1" + assert payload["workspace_dir"] == str(workspace.resolve()) + assert payload["claimed_at"] + assert payload["last_seen_at"] + + def test_atomic_claim_has_exactly_one_winner( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + + def claim(owner: str) -> str: + try: + return store.claim_workspace_lease("p1", owner).owner_id + except WorkspaceLeaseConflict: + return "blocked" + + with ThreadPoolExecutor(max_workers=2) as pool: + results = list(pool.map(claim, ["owner-a", "owner-b"])) + assert results.count("blocked") == 1 + assert len({value for value in results if value != "blocked"}) == 1 From 657a8c344e29ee231cbce75e0bf48a9ca52d4d7c Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 19:32:53 -0600 Subject: [PATCH 2/8] fix(controller): recover interrupted workspace leases --- zenith/src/zenith_harness/controller.py | 7 ++ zenith/src/zenith_harness/server.py | 23 ++-- zenith/src/zenith_harness/storage.py | 140 +++++++++++++++++------- zenith/tests/test_server.py | 21 ++++ zenith/tests/test_storage.py | 55 ++++++++++ 5 files changed, 201 insertions(+), 45 deletions(-) diff --git a/zenith/src/zenith_harness/controller.py b/zenith/src/zenith_harness/controller.py index 0bd6ab7..37dd1f6 100644 --- a/zenith/src/zenith_harness/controller.py +++ b/zenith/src/zenith_harness/controller.py @@ -11,6 +11,7 @@ from __future__ import annotations from dataclasses import dataclass +from pathlib import Path from .config import HarnessConfig from .coordinator import MissionCoordinator @@ -76,6 +77,12 @@ def start_project( ) -> Envelope: if not brief.strip(): raise ToolError("invalid_brief", "brief is empty") + requested_workspace = Path(workspace_dir).expanduser() + if not requested_workspace.is_absolute() or not requested_workspace.is_dir(): + raise ToolError( + "invalid_workspace", + "workspace_dir must be an existing absolute directory", + ) project_id = self.store.generate_project_id(brief) if owner_id is not None: try: diff --git a/zenith/src/zenith_harness/server.py b/zenith/src/zenith_harness/server.py index 2223f89..a9784fc 100644 --- a/zenith/src/zenith_harness/server.py +++ b/zenith/src/zenith_harness/server.py @@ -109,18 +109,25 @@ def _register_orchestrator_tools(mcp: FastMCP, controller: ProjectController) -> project_locks: dict[str, asyncio.Lock] = {} locks_guard = asyncio.Lock() server_instance_owner = f"server:{uuid4()}" - explicit_owner = ( - os.environ.get("ZENITH_CONTROLLER_ID") - or os.environ.get("CODEX_THREAD_ID") - or os.environ.get("CLAUDE_CODE_SESSION_ID") - or os.environ.get("CLAUDE_SESSION_ID") + explicit_owner = next( + ( + value.strip() + for value in ( + os.environ.get("ZENITH_CONTROLLER_ID"), + os.environ.get("CODEX_THREAD_ID"), + os.environ.get("CLAUDE_CODE_SESSION_ID"), + os.environ.get("CLAUDE_SESSION_ID"), + ) + if value and value.strip() + ), + None, ) - def _controller_owner(ctx: Context | None) -> str: + def _controller_owner(_ctx: Context | None) -> str: if explicit_owner: return explicit_owner - if ctx is not None and ctx.request_context is not None: - return ctx.session_id + # FastMCP session ids are scoped to one server and can repeat across + # processes. The server UUID is globally distinct for the fallback path. return server_instance_owner def _claim(project_id: str, owner_id: str) -> None: diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index 0a5373f..7fbaba6 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -21,6 +21,7 @@ import re import shutil import socket +import tempfile import time from dataclasses import dataclass from datetime import UTC, datetime @@ -202,44 +203,55 @@ def claim_workspace_lease_for_workspace( across MCP calls and processes until the owning controller explicitly releases it. Different workspaces remain independently claimable. """ - if not owner_id.strip(): + owner_id = owner_id.strip() + if not owner_id: raise ValueError("owner_id is empty") - workspace = Path(workspace_dir).expanduser().resolve() + requested_workspace = Path(workspace_dir).expanduser() + if not requested_workspace.is_absolute(): + raise ValueError("workspace_dir must be absolute") + workspace = requested_workspace.resolve() + if not workspace.is_dir(): + raise FileNotFoundError(f"workspace_dir does not exist: {workspace}") lease_path = self._workspace_lease_path(workspace) lease_dir = lease_path.parent lease_dir.parent.mkdir(parents=True, exist_ok=True) - created = False - try: - lease_dir.mkdir() - created = True - except FileExistsError: - pass - - if not created: - current = self._load_workspace_lease_record(lease_path) - if current is None: - for _ in range(5): - time.sleep(0.05) - current = self._load_workspace_lease_record(lease_path) - if current is not None: - break - if current is None: - raise WorkspaceLeaseConflict( - f"workspace lease is incomplete: {lease_path}" - ) - if ( - current.owner_id != owner_id - or current.project_id != project_id - ): - raise WorkspaceLeaseConflict( - "workspace already owned by " - f"owner={current.owner_id} project={current.project_id}; " - f"requesting owner={owner_id} project={project_id}; " - f"lease={lease_path}" - ) - claimed_at = current.claimed_at - else: - claimed_at = utc_now_iso() + claimed_at = "" + claimed = False + for _attempt in range(4): + try: + lease_dir.mkdir() + claimed_at = utc_now_iso() + claimed = True + break + except FileExistsError: + current = self._load_workspace_lease_record(lease_path) + if current is None: + for _ in range(20): + time.sleep(0.05) + current = self._load_workspace_lease_record(lease_path) + if current is not None: + break + if current is None: + if self._recover_incomplete_workspace_lease(lease_dir): + continue + continue + if ( + current.owner_id != owner_id + or current.project_id != project_id + ): + raise WorkspaceLeaseConflict( + "workspace already owned by " + f"owner={current.owner_id} project={current.project_id}; " + f"requesting owner={owner_id} project={project_id}; " + f"lease={lease_path}" + ) + claimed_at = current.claimed_at + claimed = True + break + if not claimed: + raise WorkspaceLeaseConflict( + f"workspace lease changed repeatedly during recovery: {lease_path}" + ) now = utc_now_iso() payload = { @@ -251,7 +263,7 @@ def claim_workspace_lease_for_workspace( "claimed_at": claimed_at, "last_seen_at": now, } - atomic_write_json(lease_path, payload) + self._write_workspace_lease_record(lease_path, payload) return WorkspaceLease( owner_id=owner_id, project_id=project_id, @@ -287,8 +299,62 @@ def release_workspace_lease_for_workspace( f"requesting release owner={owner_id} project={project_id}; " f"lease={lease_path}" ) - lease_path.unlink() - lease_path.parent.rmdir() + released_dir = lease_path.parent.with_name( + f"{lease_path.parent.name}.released.{os.getpid()}.{time.time_ns()}" + ) + try: + os.replace(lease_path.parent, released_dir) + except FileNotFoundError as exc: + raise WorkspaceLeaseConflict( + f"workspace lease changed during release: {lease_path}" + ) from exc + shutil.rmtree(released_dir) + + def _recover_incomplete_workspace_lease(self, lease_dir: Path) -> bool: + """Fence and remove an ownerless claim directory after bounded polling.""" + quarantine = lease_dir.with_name( + f"{lease_dir.name}.incomplete.{os.getpid()}.{time.time_ns()}" + ) + try: + os.replace(lease_dir, quarantine) + except FileNotFoundError: + return False + + moved_path = quarantine / "owner.json" + moved = self._load_workspace_lease_record(moved_path) + if moved is not None: + try: + os.replace(quarantine, lease_dir) + except OSError as exc: + raise WorkspaceLeaseConflict( + f"concurrent workspace lease repair preserved at {quarantine}" + ) from exc + raise WorkspaceLeaseConflict( + "workspace already owned by " + f"owner={moved.owner_id} project={moved.project_id}; " + f"lease={lease_dir / 'owner.json'}" + ) + + shutil.rmtree(quarantine) + return True + + @staticmethod + def _write_workspace_lease_record(path: Path, payload: object) -> None: + """Publish one lease record without shared temp-file names.""" + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw_tmp = tempfile.mkstemp( + prefix="owner.json.tmp.", dir=path.parent, text=True + ) + tmp_path = Path(raw_tmp) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, ensure_ascii=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, path) + finally: + tmp_path.unlink(missing_ok=True) @staticmethod def _load_workspace_lease_record(path: Path) -> WorkspaceLease | None: diff --git a/zenith/tests/test_server.py b/zenith/tests/test_server.py index b12e8d3..f8d770a 100644 --- a/zenith/tests/test_server.py +++ b/zenith/tests/test_server.py @@ -362,6 +362,27 @@ async def test_second_server_cannot_start_another_project_in_owned_workspace( assert len(ProjectStore(config).list_projects()) == 1 +@pytest.mark.asyncio +async def test_whitespace_controller_ids_fall_back_to_unique_server_owners( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", " ") + monkeypatch.delenv("CODEX_THREAD_ID", raising=False) + monkeypatch.delenv("CLAUDE_CODE_SESSION_ID", raising=False) + monkeypatch.delenv("CLAUDE_SESSION_ID", raising=False) + first = create_orchestrator_server(config) + second = create_orchestrator_server(config) + started = await first.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + + blocked = await second.call_tool( + "abort_project", {"project_id": pid, "reason": "must not share blank id"} + ) + assert blocked.structured_content["error"] == "workspace_owned" + + # --------------------------------------------------------------------------- # Regression: dispatcher that calls asyncio.run() must not crash the MCP # event loop. Reproduces: diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index 7b331c1..f3043fb 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -509,3 +509,58 @@ def claim(owner: str) -> str: results = list(pool.map(claim, ["owner-a", "owner-b"])) assert results.count("blocked") == 1 assert len({value for value in results if value != "blocked"}) == 1 + + def test_ownerless_crashed_claim_is_fenced_and_recovered( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + lease.path.unlink() + + recovered = store.claim_workspace_lease("p1", "owner-b") + assert recovered.owner_id == "owner-b" + + def test_malformed_crashed_claim_is_fenced_and_recovered( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + lease.path.write_text("{partial", encoding="utf-8") + + recovered = store.claim_workspace_lease("p1", "owner-b") + assert recovered.owner_id == "owner-b" + + def test_release_removes_abandoned_temp_files_atomically( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + (lease.path.parent / "owner.json.tmp.crashed").write_text( + "partial", encoding="utf-8" + ) + + store.release_workspace_lease("p1", "owner-a") + assert store.claim_workspace_lease("p1", "owner-b").owner_id == "owner-b" + + def test_concurrent_same_owner_refreshes_use_unique_temp_files( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + store.claim_workspace_lease("p1", "owner-a") + + with ThreadPoolExecutor(max_workers=4) as pool: + results = list( + pool.map( + lambda _: store.claim_workspace_lease("p1", "owner-a").owner_id, + range(8), + ) + ) + assert results == ["owner-a"] * 8 + + def test_claim_requires_existing_absolute_workspace( + self, store: ProjectStore + ) -> None: + with pytest.raises(ValueError, match="absolute"): + store.claim_workspace_lease_for_workspace( + "relative-workspace", "p1", "owner-a" + ) From 9235ec17d49284f13fe3abbeafa194e86cd4b174 Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 19:54:12 -0600 Subject: [PATCH 3/8] fix(controller): fence stale lease publishers --- zenith/src/zenith_harness/storage.py | 97 ++++++++++++++++++++++++---- zenith/tests/test_storage.py | 41 ++++++++++++ 2 files changed, 127 insertions(+), 11 deletions(-) diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index 7fbaba6..86b1e8a 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -21,7 +21,6 @@ import re import shutil import socket -import tempfile import time from dataclasses import dataclass from datetime import UTC, datetime @@ -217,9 +216,13 @@ def claim_workspace_lease_for_workspace( lease_dir.parent.mkdir(parents=True, exist_ok=True) claimed_at = "" claimed = False + lease_dir_fd: int | None = None for _attempt in range(4): try: lease_dir.mkdir() + lease_dir_fd = os.open( + lease_dir, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) + ) claimed_at = utc_now_iso() claimed = True break @@ -245,10 +248,32 @@ def claim_workspace_lease_for_workspace( f"requesting owner={owner_id} project={project_id}; " f"lease={lease_path}" ) + lease_dir_fd = os.open( + lease_dir, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) + ) + opened = self._load_workspace_lease_record_from_dir_fd( + lease_dir_fd, lease_path + ) + if opened is None: + os.close(lease_dir_fd) + lease_dir_fd = None + continue + if ( + opened.owner_id != owner_id + or opened.project_id != project_id + ): + os.close(lease_dir_fd) + lease_dir_fd = None + raise WorkspaceLeaseConflict( + "workspace ownership changed during refresh to " + f"owner={opened.owner_id} project={opened.project_id}; " + f"requesting owner={owner_id} project={project_id}; " + f"lease={lease_path}" + ) claimed_at = current.claimed_at claimed = True break - if not claimed: + if not claimed or lease_dir_fd is None: raise WorkspaceLeaseConflict( f"workspace lease changed repeatedly during recovery: {lease_path}" ) @@ -263,7 +288,25 @@ def claim_workspace_lease_for_workspace( "claimed_at": claimed_at, "last_seen_at": now, } - self._write_workspace_lease_record(lease_path, payload) + try: + self._write_workspace_lease_record(lease_dir_fd, payload) + opened_stat = os.fstat(lease_dir_fd) + live_stat = lease_dir.stat() + if (opened_stat.st_dev, opened_stat.st_ino) != ( + live_stat.st_dev, + live_stat.st_ino, + ): + raise WorkspaceLeaseConflict( + f"workspace lease changed during claim: {lease_path}" + ) + except WorkspaceLeaseConflict: + raise + except OSError as exc: + raise WorkspaceLeaseConflict( + f"workspace lease changed during claim: {lease_path}" + ) from exc + finally: + os.close(lease_dir_fd) return WorkspaceLease( owner_id=owner_id, project_id=project_id, @@ -339,22 +382,54 @@ def _recover_incomplete_workspace_lease(self, lease_dir: Path) -> bool: return True @staticmethod - def _write_workspace_lease_record(path: Path, payload: object) -> None: - """Publish one lease record without shared temp-file names.""" - path.parent.mkdir(parents=True, exist_ok=True) - fd, raw_tmp = tempfile.mkstemp( - prefix="owner.json.tmp.", dir=path.parent, text=True + def _write_workspace_lease_record(dir_fd: int, payload: object) -> None: + """Publish one lease record within the exact directory inode claimed.""" + temp_name = ( + f"owner.json.tmp.{os.getpid()}.{time.time_ns()}.{id(payload)}" + ) + fd = os.open( + temp_name, + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + 0o600, + dir_fd=dir_fd, ) - tmp_path = Path(raw_tmp) try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump(payload, handle, indent=2, ensure_ascii=False) handle.write("\n") handle.flush() os.fsync(handle.fileno()) - os.replace(tmp_path, path) + os.replace( + temp_name, + "owner.json", + src_dir_fd=dir_fd, + dst_dir_fd=dir_fd, + ) + os.fsync(dir_fd) finally: - tmp_path.unlink(missing_ok=True) + try: + os.unlink(temp_name, dir_fd=dir_fd) + except FileNotFoundError: + pass + + @staticmethod + def _load_workspace_lease_record_from_dir_fd( + dir_fd: int, path: Path + ) -> WorkspaceLease | None: + try: + fd = os.open("owner.json", os.O_RDONLY, dir_fd=dir_fd) + with os.fdopen(fd, "r", encoding="utf-8") as handle: + payload = json.load(handle) + return WorkspaceLease( + owner_id=str(payload["owner_id"]), + project_id=str(payload["project_id"]), + workspace_dir=str(payload["workspace_dir"]), + claimed_at=str(payload["claimed_at"]), + last_seen_at=str(payload["last_seen_at"]), + path=path, + ) + except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError): + return None @staticmethod def _load_workspace_lease_record(path: Path) -> WorkspaceLease | None: diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index f3043fb..cd5dd09 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -4,6 +4,7 @@ import json from concurrent.futures import ThreadPoolExecutor from pathlib import Path +from threading import Event import pytest @@ -510,6 +511,46 @@ def claim(owner: str) -> str: assert results.count("blocked") == 1 assert len({value for value in results if value != "blocked"}) == 1 + def test_recovered_claim_cannot_be_overwritten_by_stale_publisher( + self, + store: ProjectStore, + workspace: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + store.create_project("brief", workspace, project_id="p1") + first_ready = Event() + allow_first_publish = Event() + original_write = store._write_workspace_lease_record + + def pause_first_writer(dir_fd: int, payload: object) -> None: + assert isinstance(payload, dict) + if payload["owner_id"] == "owner-a": + first_ready.set() + assert allow_first_publish.wait(timeout=5) + original_write(dir_fd, payload) + + monkeypatch.setattr( + store, "_write_workspace_lease_record", pause_first_writer + ) + + def claim_first() -> str: + try: + return store.claim_workspace_lease("p1", "owner-a").owner_id + except WorkspaceLeaseConflict: + return "blocked" + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(claim_first) + assert first_ready.wait(timeout=2) + second = pool.submit( + lambda: store.claim_workspace_lease("p1", "owner-b").owner_id + ) + assert second.result(timeout=5) == "owner-b" + allow_first_publish.set() + assert first.result(timeout=5) == "blocked" + + assert store.claim_workspace_lease("p1", "owner-b").owner_id == "owner-b" + def test_ownerless_crashed_claim_is_fenced_and_recovered( self, store: ProjectStore, workspace: Path ) -> None: From 4097e77f3eb416d875c5e8148a2af62aeba48ac5 Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 20:51:06 -0600 Subject: [PATCH 4/8] fix(controller): recover dead workspace leases --- .../prompts/orchestrator/system_prompt.md | 6 +- zenith/src/zenith_harness/cli.py | 2 +- zenith/src/zenith_harness/controller.py | 83 +++++- zenith/src/zenith_harness/server.py | 72 +++++- zenith/src/zenith_harness/storage.py | 241 +++++++++++++++++- zenith/tests/test_server.py | 115 ++++++++- zenith/tests/test_storage.py | 59 ++++- 7 files changed, 539 insertions(+), 39 deletions(-) diff --git a/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md b/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md index 8acd560..c33669c 100644 --- a/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md +++ b/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md @@ -26,6 +26,7 @@ Drive the mission only through the orchestrator runtime tools: - `attention_needed`: read every open item and its evidence, decide exactly once per item with `decide_attention`, then call `advance_project`. - User-requested or justified mission cancellation: call `abort_project(project_id, reason)`. - Intentional handoff to another root orchestrator: stop all mutation, then call `release_project(project_id)`. Never share one writable workspace between root orchestrators. +- Controller crash with a stale lease: call `recover_workspace_lease(workspace_dir, reason)` only after stopping mutation. Recovery succeeds only when the recorded same-host controller PID is dead and writes a durable audit record. - `done`, `failed`, `aborted`: terminal states. Do not continue mission work unless the user starts a new scope or explicitly asks for forensic inspection. `submit_plan` and `decide_attention` persist state; they do not dispatch work. `advance_project` dispatches workers, validators, merge work, and gate evaluation according to runtime state. `end_mission` requests runtime closure and terminal review; call it only after planning, validation, gates, evidence review, and open attention handling support closure. @@ -362,9 +363,10 @@ Orchestrator tools: - `decide_attention(project_id, decisions)`: resolve every open attention item with exactly one decision, then return to runtime flow. Call `advance_project` afterward. - `end_mission(project_id)`: request runtime closure and terminal review only after work is quiescent and evidence supports closure. - `abort_project(project_id, reason)`: terminal cancellation with a recorded reason. -- `release_project(project_id)`: explicitly release persistent workspace ownership so another root orchestrator can take over. Read-only inspection by other sessions remains available before release. +- `release_project(project_id, force?, reason?)`: explicitly release persistent workspace ownership so another root orchestrator can take over. Read-only inspection by other sessions remains available before release. `force=true` is only for stuck running-task records and requires an audit reason. +- `recover_workspace_lease(workspace_dir, reason)`: remove a stale or orphan lease only when the recorded same-host controller process is provably dead. It does not claim the workspace; start or claim the project after recovery. -Lifecycle and inspection tools return an envelope with `projectId`, `state`, `projectRoot`, `harnessRoot`, and `dag`. Some lifecycle/decision tools intentionally return `dag=null`; `submit_plan` and `advance_project` return a compact frontier view; `inspect_project` returns the full task-list view. `release_project` returns a release acknowledgement. Trust the envelope, runtime files, and returned paths over session memory. +Lifecycle and inspection tools return an envelope with `projectId`, `state`, `projectRoot`, `harnessRoot`, and `dag`. Some lifecycle/decision tools intentionally return `dag=null`; `submit_plan` and `advance_project` return a compact frontier view; `inspect_project` returns the full task-list view. Lease release and recovery tools return acknowledgements. Trust the envelope, runtime files, and returned paths over session memory. `submit_plan` accepts a `TaskList`: diff --git a/zenith/src/zenith_harness/cli.py b/zenith/src/zenith_harness/cli.py index e72d99c..4ccf674 100644 --- a/zenith/src/zenith_harness/cli.py +++ b/zenith/src/zenith_harness/cli.py @@ -290,7 +290,7 @@ def abort_project_cmd(project_id: str, reason: str, owner_id: str | None) -> Non ) try: controller.claim_project(project_id, resolved_owner) - env = controller.abort_project(project_id, reason) + env = controller.abort_project(project_id, reason, resolved_owner) except ToolError as exc: raise click.ClickException(str(exc)) from exc click.echo(f"Aborted {project_id}: state={env.state.state}") diff --git a/zenith/src/zenith_harness/controller.py b/zenith/src/zenith_harness/controller.py index 37dd1f6..4f525dd 100644 --- a/zenith/src/zenith_harness/controller.py +++ b/zenith/src/zenith_harness/controller.py @@ -32,7 +32,7 @@ TaskListPatch, TaskStateFile, ) -from .storage import ProjectStore, WorkspaceLeaseConflict +from .storage import ProjectStore, WorkspaceLease, WorkspaceLeaseConflict from .task_list_patch import apply_patch from .task_validation import ( ValidationError, @@ -89,6 +89,8 @@ def start_project( self.store.claim_workspace_lease_for_workspace( workspace_dir, project_id, owner_id ) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc except WorkspaceLeaseConflict as exc: raise ToolError("workspace_owned", str(exc)) from exc try: @@ -112,13 +114,41 @@ def start_project( def claim_project(self, project_id: str, owner_id: str) -> None: try: self.store.claim_workspace_lease(project_id, owner_id) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc + except FileNotFoundError as exc: + raise ToolError("invalid_workspace", str(exc)) from exc except WorkspaceLeaseConflict as exc: raise ToolError("workspace_owned", str(exc)) from exc - def release_project(self, project_id: str, owner_id: str) -> None: + def release_project( + self, + project_id: str, + owner_id: str, + *, + force: bool = False, + reason: str | None = None, + ) -> None: + if force and (not reason or not reason.strip()): + raise ToolError( + "recovery_reason_required", + "forced release requires a non-empty reason", + ) + if reason is not None and len(reason.strip()) > 500: + raise ToolError( + "invalid_recovery", + "recovery reason must be at most 500 characters", + ) + try: + self.store.claim_workspace_lease(project_id, owner_id) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc record = self.store.load_project(project_id) state = self.store.load_state(project_id) mission_id = self._current_mission_id(record, state) + forced_audit: dict[str, object] | None = None if mission_id is not None: task_state = self.store.load_task_state(project_id, mission_id) running = sorted( @@ -127,15 +157,47 @@ def release_project(self, project_id: str, owner_id: str) -> None: if entry.status == "running" ) if running: - raise ToolError( - "workspace_busy", - "cannot release controller ownership while tasks are running: " - + ", ".join(running), + if not force: + raise ToolError( + "workspace_busy", + "cannot release controller ownership while tasks are running: " + + ", ".join(running), + ) + assert reason is not None + forced_audit = { + "project_id": project_id, + "owner_id": owner_id, + "running_task_ids": running, + "reason": reason.strip(), + } + self.store.append_workspace_lease_audit( + {"action": "forced_running_release_authorized", **forced_audit} ) try: self.store.release_workspace_lease(project_id, owner_id) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc except WorkspaceLeaseConflict as exc: raise ToolError("workspace_owned", str(exc)) from exc + if forced_audit is not None: + self.store.append_workspace_lease_audit( + {"action": "forced_running_release_completed", **forced_audit} + ) + + def recover_workspace_lease( + self, workspace_dir: str, owner_id: str, reason: str + ) -> WorkspaceLease: + try: + return self.store.recover_workspace_lease_for_workspace( + workspace_dir, owner_id, reason + ) + except ValueError as exc: + code = "invalid_owner" if "owner_id" in str(exc) else "invalid_recovery" + raise ToolError(code, str(exc)) from exc + except FileNotFoundError as exc: + raise ToolError("invalid_workspace", str(exc)) from exc + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_recovery_blocked", str(exc)) from exc def submit_plan(self, project_id: str, task_list: TaskList) -> Envelope: state = self._require_state(project_id) @@ -238,7 +300,9 @@ def decide_attention( def inspect_project(self, project_id: str) -> Envelope: return self._build_envelope(project_id, dag_mode="full") - def abort_project(self, project_id: str, reason: str) -> Envelope: + def abort_project( + self, project_id: str, reason: str, owner_id: str | None = None + ) -> Envelope: record = self.store.load_project(project_id) state = self.store.load_state(project_id) mid = self._current_mission_id(record, state) @@ -251,7 +315,10 @@ def abort_project(self, project_id: str, reason: str) -> Envelope: pass self.store.clear_attention(project_id) self.store.save_state(project_id, Aborted(reason=reason)) - return self._build_envelope(project_id, dag_mode="none") + envelope = self._build_envelope(project_id, dag_mode="none") + if owner_id is not None: + self.release_project(project_id, owner_id) + return envelope # ------------------------------------------------------------------ # Decision pipeline diff --git a/zenith/src/zenith_harness/server.py b/zenith/src/zenith_harness/server.py index a9784fc..a9be51e 100644 --- a/zenith/src/zenith_harness/server.py +++ b/zenith/src/zenith_harness/server.py @@ -40,7 +40,7 @@ def create_orchestrator_server( config: HarnessConfig, controller: ProjectController | None = None, ) -> FastMCP: - """8 orchestrator tools, registered on a stdio MCP server.""" + """9 orchestrator tools, registered on a stdio MCP server.""" if controller is None: from .dispatcher import MockDispatcher, MockTerminalReviewer @@ -56,9 +56,10 @@ def create_orchestrator_server( name="zenith", instructions=( "Mission orchestration harness. Mode: orchestrator. " - "8 tools: start_project, submit_plan, advance_project, " + "9 tools: start_project, submit_plan, advance_project, " "end_mission, decide_attention, inspect_project, abort_project. " - "release_project performs an explicit controller handoff. " + "release_project performs an explicit controller handoff; " + "recover_workspace_lease is a dead-controller recovery path. " "Lifecycle: plan with submit_plan, run with advance_project, " "request closure with end_mission, resolve attention with decide_attention, " "then call advance_project again." @@ -126,8 +127,9 @@ def _register_orchestrator_tools(mcp: FastMCP, controller: ProjectController) -> def _controller_owner(_ctx: Context | None) -> str: if explicit_owner: return explicit_owner - # FastMCP session ids are scoped to one server and can repeat across - # processes. The server UUID is globally distinct for the fallback path. + # Existing init configs do not inject a dynamic task id. A process-unique + # fallback still provides exclusive ownership; the PID-bound recovery tool + # makes a crashed fallback owner safely recoverable. return server_instance_owner def _claim(project_id: str, owner_id: str) -> None: @@ -319,7 +321,9 @@ async def abort_project( owner_id = _controller_owner(ctx) await asyncio.to_thread(_claim, project_id, owner_id) return _to_payload( - await asyncio.to_thread(controller.abort_project, project_id, reason) + await asyncio.to_thread( + controller.abort_project, project_id, reason, owner_id + ) ) except ToolError as exc: return _to_payload(exc) @@ -334,18 +338,70 @@ async def abort_project( ) async def release_project( project_id: Annotated[str, Field(description="Project id.")], + force: Annotated[ + bool, + Field( + default=False, + description="Allow release with stuck running task records.", + ), + ] = False, + reason: Annotated[ + str | None, + Field( + default=None, + description="Required audit reason when force is true.", + ), + ] = None, ctx: Context | None = None, ) -> dict[str, Any]: async with await _project_lock(project_id): try: owner_id = _controller_owner(ctx) await asyncio.to_thread( - controller.release_project, project_id, owner_id + controller.release_project, + project_id, + owner_id, + force=force, + reason=reason, ) return { "released": True, "projectId": project_id, - "ownerId": owner_id, + "forced": force, + } + except ToolError as exc: + return _to_payload(exc) + + @mcp.tool( + name="recover_workspace_lease", + description=( + "Break a stale workspace lease only when the recorded controller PID " + "is provably dead on this host. Requires an audit reason. Use after a " + "controller crash, including before a project record was created." + ), + ) + async def recover_workspace_lease( + workspace_dir: Annotated[ + str, Field(description="Existing absolute workspace path.") + ], + reason: Annotated[ + str, Field(description="Why dead-controller recovery is required.") + ], + ctx: Context | None = None, + ) -> dict[str, Any]: + async with await _project_lock(f"workspace:{workspace_dir}"): + try: + owner_id = _controller_owner(ctx) + recovered = await asyncio.to_thread( + controller.recover_workspace_lease, + workspace_dir, + owner_id, + reason, + ) + return { + "recovered": True, + "workspaceDir": recovered.workspace_dir, + "previousProjectId": recovered.project_id, } except ToolError as exc: return _to_payload(exc) diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index 86b1e8a..cd22da1 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -126,6 +126,8 @@ class WorkspaceLease: workspace_dir: str claimed_at: str last_seen_at: str + host: str + pid: int path: Path @@ -133,6 +135,32 @@ class WorkspaceLeaseConflict(RuntimeError): """A different controller already owns the writable workspace.""" +_OWNER_ID_RE = re.compile(r"^[A-Za-z0-9._:@-]{1,200}$") + + +def _normalize_owner_id(owner_id: str) -> str: + owner_id = owner_id.strip() + if not owner_id: + raise ValueError("owner_id is empty") + if len(owner_id) > 200: + raise ValueError("owner_id must be at most 200 characters") + if not _OWNER_ID_RE.fullmatch(owner_id): + raise ValueError("owner_id contains unsupported characters") + return owner_id + + +def _pid_is_alive(pid: int) -> bool: + if pid <= 0: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + # --------------------------------------------------------------------------- # ProjectStore # --------------------------------------------------------------------------- @@ -202,9 +230,7 @@ def claim_workspace_lease_for_workspace( across MCP calls and processes until the owning controller explicitly releases it. Different workspaces remain independently claimable. """ - owner_id = owner_id.strip() - if not owner_id: - raise ValueError("owner_id is empty") + owner_id = _normalize_owner_id(owner_id) requested_workspace = Path(workspace_dir).expanduser() if not requested_workspace.is_absolute(): raise ValueError("workspace_dir must be absolute") @@ -224,6 +250,23 @@ def claim_workspace_lease_for_workspace( lease_dir, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) ) claimed_at = utc_now_iso() + claim_payload = { + "schema_version": 1, + "owner_id": owner_id, + "project_id": project_id, + "workspace_dir": str(workspace), + "host": socket.gethostname(), + "pid": os.getpid(), + "claimed_at": claimed_at, + "last_seen_at": claimed_at, + } + try: + self._write_workspace_claim_record(lease_dir_fd, claim_payload) + except OSError: + os.close(lease_dir_fd) + lease_dir_fd = None + shutil.rmtree(lease_dir, ignore_errors=True) + raise claimed = True break except FileExistsError: @@ -235,7 +278,29 @@ def claim_workspace_lease_for_workspace( if current is not None: break if current is None: - if self._recover_incomplete_workspace_lease(lease_dir): + incomplete_marker = self._load_workspace_lease_record( + lease_dir / "claim.json" + ) + if incomplete_marker is None: + raise WorkspaceLeaseConflict( + "workspace lease is incomplete and has no recoverable " + f"claim identity: {lease_path}" + ) + if incomplete_marker.host != socket.gethostname(): + raise WorkspaceLeaseConflict( + "workspace lease has an incomplete claim on another " + f"host ({incomplete_marker.host}): {lease_path}" + ) + if _pid_is_alive(incomplete_marker.pid): + raise WorkspaceLeaseConflict( + "workspace has incomplete active claim by " + f"owner={incomplete_marker.owner_id} " + f"project={incomplete_marker.project_id} " + f"pid={incomplete_marker.pid}; lease={lease_path}" + ) + if self._recover_incomplete_workspace_lease( + lease_dir, incomplete_marker + ): continue continue if ( @@ -248,6 +313,16 @@ def claim_workspace_lease_for_workspace( f"requesting owner={owner_id} project={project_id}; " f"lease={lease_path}" ) + if current.host != socket.gethostname(): + raise WorkspaceLeaseConflict( + "same owner id is active on another host " + f"({current.host}); lease={lease_path}" + ) + if current.pid != os.getpid() and _pid_is_alive(current.pid): + raise WorkspaceLeaseConflict( + "same owner id is active in another controller process " + f"pid={current.pid}; lease={lease_path}" + ) lease_dir_fd = os.open( lease_dir, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) ) @@ -285,10 +360,12 @@ def claim_workspace_lease_for_workspace( "project_id": project_id, "workspace_dir": str(workspace), "host": socket.gethostname(), + "pid": os.getpid(), "claimed_at": claimed_at, "last_seen_at": now, } try: + self._write_workspace_claim_record(lease_dir_fd, payload) self._write_workspace_lease_record(lease_dir_fd, payload) opened_stat = os.fstat(lease_dir_fd) live_stat = lease_dir.stat() @@ -313,6 +390,8 @@ def claim_workspace_lease_for_workspace( workspace_dir=str(workspace), claimed_at=claimed_at, last_seen_at=now, + host=socket.gethostname(), + pid=os.getpid(), path=lease_path, ) @@ -327,6 +406,7 @@ def release_workspace_lease_for_workspace( project_id: str, owner_id: str, ) -> None: + owner_id = _normalize_owner_id(owner_id) lease_path = self._workspace_lease_path(workspace_dir) if not lease_path.parent.exists(): return @@ -353,7 +433,109 @@ def release_workspace_lease_for_workspace( ) from exc shutil.rmtree(released_dir) - def _recover_incomplete_workspace_lease(self, lease_dir: Path) -> bool: + def recover_workspace_lease_for_workspace( + self, + workspace_dir: str | Path, + requesting_owner_id: str, + reason: str, + ) -> WorkspaceLease: + """Release a lease only when its recorded same-host controller is dead.""" + requesting_owner_id = _normalize_owner_id(requesting_owner_id) + reason = reason.strip() + if not reason: + raise ValueError("recovery reason is empty") + if len(reason) > 500: + raise ValueError("recovery reason must be at most 500 characters") + requested_workspace = Path(workspace_dir).expanduser() + if not requested_workspace.is_absolute(): + raise ValueError("workspace_dir must be absolute") + workspace = requested_workspace.resolve() + if not workspace.is_dir(): + raise FileNotFoundError(f"workspace_dir does not exist: {workspace}") + lease_path = self._workspace_lease_path(workspace) + current = self._load_workspace_lease_record(lease_path) + if current is None: + current = self._load_workspace_lease_record( + lease_path.parent / "claim.json" + ) + if current is None: + raise WorkspaceLeaseConflict( + f"workspace lease has no provable owner identity: {lease_path}" + ) + if current.host != socket.gethostname(): + raise WorkspaceLeaseConflict( + "cannot prove controller death on another host: " + f"{current.host}; lease={lease_path}" + ) + if _pid_is_alive(current.pid): + raise WorkspaceLeaseConflict( + "cannot recover live workspace controller " + f"pid={current.pid}; lease={lease_path}" + ) + + recovery_audit = { + "workspace_dir": str(workspace), + "project_id": current.project_id, + "previous_owner_id": current.owner_id, + "previous_host": current.host, + "previous_pid": current.pid, + "requesting_owner_id": requesting_owner_id, + "reason": reason, + } + self.append_workspace_lease_audit( + {"action": "dead_controller_recovery_authorized", **recovery_audit} + ) + + quarantine = lease_path.parent.with_name( + f"{lease_path.parent.name}.recovering.{os.getpid()}.{time.time_ns()}" + ) + try: + os.replace(lease_path.parent, quarantine) + except FileNotFoundError as exc: + raise WorkspaceLeaseConflict( + f"workspace lease changed during recovery: {lease_path}" + ) from exc + moved = self._load_workspace_lease_record(quarantine / "owner.json") + if moved is None: + moved = self._load_workspace_lease_record(quarantine / "claim.json") + if moved is None or self._lease_identity(moved) != self._lease_identity(current): + self._restore_quarantined_workspace_lease(quarantine, lease_path.parent) + raise WorkspaceLeaseConflict( + f"workspace lease changed during recovery: {lease_path}" + ) + if moved.host != socket.gethostname() or _pid_is_alive(moved.pid): + self._restore_quarantined_workspace_lease(quarantine, lease_path.parent) + raise WorkspaceLeaseConflict( + f"workspace controller revived during recovery: {lease_path}" + ) + self.append_workspace_lease_audit( + {"action": "dead_controller_recovery_completed", **recovery_audit} + ) + shutil.rmtree(quarantine) + return moved + + def append_workspace_lease_audit(self, payload: dict[str, object]) -> None: + audit_path = self.config.harness_home / "leases" / "recovery-log.jsonl" + audit_path.parent.mkdir(parents=True, exist_ok=True) + line = { + "schema_version": 1, + "recorded_at": utc_now_iso(), + **payload, + } + encoded = json.dumps(line, ensure_ascii=False, sort_keys=True) + "\n" + fd = os.open(audit_path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600) + try: + remaining = memoryview(encoded.encode("utf-8")) + while remaining: + written = os.write(fd, remaining) + remaining = remaining[written:] + os.fsync(fd) + finally: + os.close(fd) + + def _recover_incomplete_workspace_lease( + self, lease_dir: Path, expected: WorkspaceLease + ) -> bool: """Fence and remove an ownerless claim directory after bounded polling.""" quarantine = lease_dir.with_name( f"{lease_dir.name}.incomplete.{os.getpid()}.{time.time_ns()}" @@ -378,15 +560,52 @@ def _recover_incomplete_workspace_lease(self, lease_dir: Path) -> bool: f"lease={lease_dir / 'owner.json'}" ) + moved_marker = self._load_workspace_lease_record( + quarantine / "claim.json" + ) + if ( + moved_marker is None + or self._lease_identity(moved_marker) != self._lease_identity(expected) + or moved_marker.host != socket.gethostname() + or _pid_is_alive(moved_marker.pid) + ): + self._restore_quarantined_workspace_lease(quarantine, lease_dir) + raise WorkspaceLeaseConflict( + f"workspace incomplete claim changed during recovery: {lease_dir}" + ) + shutil.rmtree(quarantine) return True + @staticmethod + def _lease_identity(lease: WorkspaceLease) -> tuple[str, str, str, int]: + return (lease.owner_id, lease.project_id, lease.host, lease.pid) + + @staticmethod + def _restore_quarantined_workspace_lease( + quarantine: Path, lease_dir: Path + ) -> None: + try: + os.replace(quarantine, lease_dir) + except OSError as exc: + raise WorkspaceLeaseConflict( + f"workspace lease quarantine preserved at {quarantine}" + ) from exc + + @staticmethod + def _write_workspace_claim_record(dir_fd: int, payload: object) -> None: + ProjectStore._write_workspace_json_record(dir_fd, "claim.json", payload) + @staticmethod def _write_workspace_lease_record(dir_fd: int, payload: object) -> None: """Publish one lease record within the exact directory inode claimed.""" - temp_name = ( - f"owner.json.tmp.{os.getpid()}.{time.time_ns()}.{id(payload)}" - ) + ProjectStore._write_workspace_json_record(dir_fd, "owner.json", payload) + + @staticmethod + def _write_workspace_json_record( + dir_fd: int, filename: str, payload: object + ) -> None: + temp_name = f"{filename}.tmp.{os.getpid()}.{time.time_ns()}.{id(payload)}" fd = os.open( temp_name, os.O_WRONLY | os.O_CREAT | os.O_EXCL, @@ -401,7 +620,7 @@ def _write_workspace_lease_record(dir_fd: int, payload: object) -> None: os.fsync(handle.fileno()) os.replace( temp_name, - "owner.json", + filename, src_dir_fd=dir_fd, dst_dir_fd=dir_fd, ) @@ -426,6 +645,8 @@ def _load_workspace_lease_record_from_dir_fd( workspace_dir=str(payload["workspace_dir"]), claimed_at=str(payload["claimed_at"]), last_seen_at=str(payload["last_seen_at"]), + host=str(payload.get("host", "")), + pid=int(payload.get("pid", 0)), path=path, ) except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError): @@ -441,6 +662,8 @@ def _load_workspace_lease_record(path: Path) -> WorkspaceLease | None: workspace_dir=str(payload["workspace_dir"]), claimed_at=str(payload["claimed_at"]), last_seen_at=str(payload["last_seen_at"]), + host=str(payload.get("host", "")), + pid=int(payload.get("pid", 0)), path=path, ) except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError): diff --git a/zenith/tests/test_server.py b/zenith/tests/test_server.py index f8d770a..1baa930 100644 --- a/zenith/tests/test_server.py +++ b/zenith/tests/test_server.py @@ -68,6 +68,7 @@ async def test_orchestrator_tools_registered(config: HarnessConfig) -> None: "inspect_project", "abort_project", "release_project", + "recover_workspace_lease", } @@ -304,7 +305,7 @@ async def test_non_owner_cannot_release_project( @pytest.mark.asyncio -async def test_owner_cannot_release_project_while_task_is_running( +async def test_owner_must_explicitly_force_release_while_task_is_running( config: HarnessConfig, workspace: Path, monkeypatch ) -> None: monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") @@ -342,6 +343,26 @@ async def test_owner_cannot_release_project_while_task_is_running( assert blocked.structured_content["error"] == "workspace_busy" assert "w1" in blocked.structured_content["message"] + missing_reason = await server.call_tool( + "release_project", {"project_id": pid, "force": True} + ) + assert missing_reason.structured_content["error"] == "recovery_reason_required" + + released = await server.call_tool( + "release_project", + { + "project_id": pid, + "force": True, + "reason": "worker process crashed", + }, + ) + assert released.structured_content["released"] is True + audit_path = config.harness_home / "leases" / "recovery-log.jsonl" + audit = json.loads(audit_path.read_text(encoding="utf-8").splitlines()[-1]) + assert audit["action"] == "forced_running_release_completed" + assert audit["running_task_ids"] == ["w1"] + assert audit["reason"] == "worker process crashed" + @pytest.mark.asyncio async def test_second_server_cannot_start_another_project_in_owned_workspace( @@ -363,7 +384,7 @@ async def test_second_server_cannot_start_another_project_in_owned_workspace( @pytest.mark.asyncio -async def test_whitespace_controller_ids_fall_back_to_unique_server_owners( +async def test_missing_controller_identity_uses_unique_process_owners( config: HarnessConfig, workspace: Path, monkeypatch ) -> None: monkeypatch.setenv("ZENITH_CONTROLLER_ID", " ") @@ -383,6 +404,96 @@ async def test_whitespace_controller_ids_fall_back_to_unique_server_owners( assert blocked.structured_content["error"] == "workspace_owned" +@pytest.mark.asyncio +async def test_invalid_controller_identity_returns_structured_error( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner with spaces") + server = create_orchestrator_server(config) + + started = await server.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + + assert started.structured_content["error"] == "invalid_owner" + + +@pytest.mark.asyncio +async def test_abort_releases_workspace_for_next_controller( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + first = create_orchestrator_server(config) + started = await first.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + aborted = await first.call_tool( + "abort_project", {"project_id": pid, "reason": "intentional"} + ) + assert aborted.structured_content["state"]["state"] == "aborted" + + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-b") + second = create_orchestrator_server(config) + restarted = await second.call_tool( + "start_project", {"brief": "Next.", "workspace_dir": str(workspace)} + ) + assert "error" not in restarted.structured_content + + +@pytest.mark.asyncio +async def test_dead_orphan_lease_can_be_recovered_without_project_record( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + store = ProjectStore(config) + lease = store.claim_workspace_lease_for_workspace( + workspace, "ghost-project", "dead-owner" + ) + for name in ("claim.json", "owner.json"): + path = lease.path.parent / name + payload = json.loads(path.read_text(encoding="utf-8")) + payload["pid"] = 999_999_999 + path.write_text(json.dumps(payload), encoding="utf-8") + + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "recovery-owner") + server = create_orchestrator_server(config) + recovered = await server.call_tool( + "recover_workspace_lease", + { + "workspace_dir": str(workspace), + "reason": "controller crashed before project creation", + }, + ) + + assert recovered.structured_content["recovered"] is True + audit_path = config.harness_home / "leases" / "recovery-log.jsonl" + audit = json.loads(audit_path.read_text(encoding="utf-8").splitlines()[-1]) + assert audit["action"] == "dead_controller_recovery_completed" + assert audit["project_id"] == "ghost-project" + started = await server.call_tool( + "start_project", {"brief": "Recovered.", "workspace_dir": str(workspace)} + ) + assert "error" not in started.structured_content + + +@pytest.mark.asyncio +async def test_live_workspace_lease_cannot_be_recovered( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + store = ProjectStore(config) + store.claim_workspace_lease_for_workspace(workspace, "p1", "live-owner") + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "recovery-owner") + server = create_orchestrator_server(config) + + recovered = await server.call_tool( + "recover_workspace_lease", + {"workspace_dir": str(workspace), "reason": "must remain fenced"}, + ) + + assert recovered.structured_content["error"] == "workspace_recovery_blocked" + assert "live workspace controller" in recovered.structured_content["message"] + + # --------------------------------------------------------------------------- # Regression: dispatcher that calls asyncio.run() must not crash the MCP # event loop. Reproduces: diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index cd5dd09..4946fd7 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -494,6 +494,13 @@ def test_record_contains_forensic_identity( assert payload["workspace_dir"] == str(workspace.resolve()) assert payload["claimed_at"] assert payload["last_seen_at"] + assert payload["host"] + assert payload["pid"] > 0 + marker = json.loads( + (lease.path.parent / "claim.json").read_text(encoding="utf-8") + ) + assert marker["owner_id"] == "owner-a" + assert marker["pid"] == payload["pid"] def test_atomic_claim_has_exactly_one_winner( self, store: ProjectStore, workspace: Path @@ -511,7 +518,7 @@ def claim(owner: str) -> str: assert results.count("blocked") == 1 assert len({value for value in results if value != "blocked"}) == 1 - def test_recovered_claim_cannot_be_overwritten_by_stale_publisher( + def test_live_incomplete_claim_cannot_be_overwritten_by_second_owner( self, store: ProjectStore, workspace: Path, @@ -539,38 +546,72 @@ def claim_first() -> str: except WorkspaceLeaseConflict: return "blocked" + def claim_second() -> str: + try: + return store.claim_workspace_lease("p1", "owner-b").owner_id + except WorkspaceLeaseConflict: + return "blocked" + with ThreadPoolExecutor(max_workers=2) as pool: first = pool.submit(claim_first) assert first_ready.wait(timeout=2) - second = pool.submit( - lambda: store.claim_workspace_lease("p1", "owner-b").owner_id - ) - assert second.result(timeout=5) == "owner-b" + second = pool.submit(claim_second) + assert second.result(timeout=5) == "blocked" allow_first_publish.set() - assert first.result(timeout=5) == "blocked" + assert first.result(timeout=5) == "owner-a" - assert store.claim_workspace_lease("p1", "owner-b").owner_id == "owner-b" + assert store.claim_workspace_lease("p1", "owner-a").owner_id == "owner-a" - def test_ownerless_crashed_claim_is_fenced_and_recovered( + def test_ownerless_dead_claim_is_fenced_and_recovered( self, store: ProjectStore, workspace: Path ) -> None: store.create_project("brief", workspace, project_id="p1") lease = store.claim_workspace_lease("p1", "owner-a") lease.path.unlink() + claim_path = lease.path.parent / "claim.json" + claim = json.loads(claim_path.read_text(encoding="utf-8")) + claim["pid"] = 999_999_999 + claim_path.write_text(json.dumps(claim), encoding="utf-8") recovered = store.claim_workspace_lease("p1", "owner-b") assert recovered.owner_id == "owner-b" - def test_malformed_crashed_claim_is_fenced_and_recovered( + def test_malformed_dead_claim_is_fenced_and_recovered( self, store: ProjectStore, workspace: Path ) -> None: store.create_project("brief", workspace, project_id="p1") lease = store.claim_workspace_lease("p1", "owner-a") lease.path.write_text("{partial", encoding="utf-8") + claim_path = lease.path.parent / "claim.json" + claim = json.loads(claim_path.read_text(encoding="utf-8")) + claim["pid"] = 999_999_999 + claim_path.write_text(json.dumps(claim), encoding="utf-8") recovered = store.claim_workspace_lease("p1", "owner-b") assert recovered.owner_id == "owner-b" + def test_malformed_live_claim_remains_fenced( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + lease.path.write_text("{partial", encoding="utf-8") + + with pytest.raises(WorkspaceLeaseConflict, match="incomplete active claim"): + store.claim_workspace_lease("p1", "owner-b") + + def test_owner_id_is_bounded_and_safe( + self, store: ProjectStore, workspace: Path + ) -> None: + with pytest.raises(ValueError, match="unsupported"): + store.claim_workspace_lease_for_workspace( + workspace, "p1", "owner with spaces" + ) + with pytest.raises(ValueError, match="at most 200"): + store.claim_workspace_lease_for_workspace( + workspace, "p1", "x" * 201 + ) + def test_release_removes_abandoned_temp_files_atomically( self, store: ProjectStore, workspace: Path ) -> None: From b732a1f14817c9311f09bc7c9bf859c910eee734 Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 21:08:32 -0600 Subject: [PATCH 5/8] fix(controller): fence unsafe lease handoffs --- .../prompts/orchestrator/system_prompt.md | 2 +- zenith/src/zenith_harness/controller.py | 57 ++++++++----------- zenith/src/zenith_harness/server.py | 25 ++------ zenith/src/zenith_harness/storage.py | 16 ++++-- zenith/tests/test_server.py | 26 +++------ zenith/tests/test_storage.py | 14 +++++ 6 files changed, 64 insertions(+), 76 deletions(-) diff --git a/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md b/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md index c33669c..316b9c1 100644 --- a/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md +++ b/zenith/src/zenith_harness/bundled/prompts/orchestrator/system_prompt.md @@ -363,7 +363,7 @@ Orchestrator tools: - `decide_attention(project_id, decisions)`: resolve every open attention item with exactly one decision, then return to runtime flow. Call `advance_project` afterward. - `end_mission(project_id)`: request runtime closure and terminal review only after work is quiescent and evidence supports closure. - `abort_project(project_id, reason)`: terminal cancellation with a recorded reason. -- `release_project(project_id, force?, reason?)`: explicitly release persistent workspace ownership so another root orchestrator can take over. Read-only inspection by other sessions remains available before release. `force=true` is only for stuck running-task records and requires an audit reason. +- `release_project(project_id)`: explicitly release persistent workspace ownership so another root orchestrator can take over. Read-only inspection by other sessions remains available before release. Running work always blocks release; use audited dead-controller recovery only after the owning process is proven dead. - `recover_workspace_lease(workspace_dir, reason)`: remove a stale or orphan lease only when the recorded same-host controller process is provably dead. It does not claim the workspace; start or claim the project after recovery. Lifecycle and inspection tools return an envelope with `projectId`, `state`, `projectRoot`, `harnessRoot`, and `dag`. Some lifecycle/decision tools intentionally return `dag=null`; `submit_plan` and `advance_project` return a compact frontier view; `inspect_project` returns the full task-list view. Lease release and recovery tools return acknowledgements. Trust the envelope, runtime files, and returned paths over session memory. diff --git a/zenith/src/zenith_harness/controller.py b/zenith/src/zenith_harness/controller.py index 4f525dd..542e5de 100644 --- a/zenith/src/zenith_harness/controller.py +++ b/zenith/src/zenith_harness/controller.py @@ -125,20 +125,7 @@ def release_project( self, project_id: str, owner_id: str, - *, - force: bool = False, - reason: str | None = None, ) -> None: - if force and (not reason or not reason.strip()): - raise ToolError( - "recovery_reason_required", - "forced release requires a non-empty reason", - ) - if reason is not None and len(reason.strip()) > 500: - raise ToolError( - "invalid_recovery", - "recovery reason must be at most 500 characters", - ) try: self.store.claim_workspace_lease(project_id, owner_id) except ValueError as exc: @@ -148,7 +135,6 @@ def release_project( record = self.store.load_project(project_id) state = self.store.load_state(project_id) mission_id = self._current_mission_id(record, state) - forced_audit: dict[str, object] | None = None if mission_id is not None: task_state = self.store.load_task_state(project_id, mission_id) running = sorted( @@ -157,21 +143,10 @@ def release_project( if entry.status == "running" ) if running: - if not force: - raise ToolError( - "workspace_busy", - "cannot release controller ownership while tasks are running: " - + ", ".join(running), - ) - assert reason is not None - forced_audit = { - "project_id": project_id, - "owner_id": owner_id, - "running_task_ids": running, - "reason": reason.strip(), - } - self.store.append_workspace_lease_audit( - {"action": "forced_running_release_authorized", **forced_audit} + raise ToolError( + "workspace_busy", + "cannot release controller ownership while tasks are running: " + + ", ".join(running), ) try: self.store.release_workspace_lease(project_id, owner_id) @@ -179,10 +154,6 @@ def release_project( raise ToolError("invalid_owner", str(exc)) from exc except WorkspaceLeaseConflict as exc: raise ToolError("workspace_owned", str(exc)) from exc - if forced_audit is not None: - self.store.append_workspace_lease_audit( - {"action": "forced_running_release_completed", **forced_audit} - ) def recover_workspace_lease( self, workspace_dir: str, owner_id: str, reason: str @@ -306,6 +277,19 @@ def abort_project( record = self.store.load_project(project_id) state = self.store.load_state(project_id) mid = self._current_mission_id(record, state) + if owner_id is not None and mid is not None: + task_state = self.store.load_task_state(project_id, mid) + running = sorted( + task_id + for task_id, entry in task_state.tasks.items() + if entry.status == "running" + ) + if running: + raise ToolError( + "workspace_busy", + "cannot abort and release ownership while tasks are running: " + + ", ".join(running), + ) if mid: try: self.store.seal_mission( @@ -317,7 +301,12 @@ def abort_project( self.store.save_state(project_id, Aborted(reason=reason)) envelope = self._build_envelope(project_id, dag_mode="none") if owner_id is not None: - self.release_project(project_id, owner_id) + try: + self.store.release_workspace_lease(project_id, owner_id) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc return envelope # ------------------------------------------------------------------ diff --git a/zenith/src/zenith_harness/server.py b/zenith/src/zenith_harness/server.py index a9be51e..a1da79b 100644 --- a/zenith/src/zenith_harness/server.py +++ b/zenith/src/zenith_harness/server.py @@ -109,7 +109,7 @@ def _register_orchestrator_tools(mcp: FastMCP, controller: ProjectController) -> # read-only and stays uncontended so recovery audits remain possible. project_locks: dict[str, asyncio.Lock] = {} locks_guard = asyncio.Lock() - server_instance_owner = f"server:{uuid4()}" + server_instance_owner: tuple[int, str] | None = None explicit_owner = next( ( value.strip() @@ -125,12 +125,16 @@ def _register_orchestrator_tools(mcp: FastMCP, controller: ProjectController) -> ) def _controller_owner(_ctx: Context | None) -> str: + nonlocal server_instance_owner if explicit_owner: return explicit_owner # Existing init configs do not inject a dynamic task id. A process-unique # fallback still provides exclusive ownership; the PID-bound recovery tool # makes a crashed fallback owner safely recoverable. - return server_instance_owner + process_id = os.getpid() + if server_instance_owner is None or server_instance_owner[0] != process_id: + server_instance_owner = (process_id, f"server:{process_id}:{uuid4()}") + return server_instance_owner[1] def _claim(project_id: str, owner_id: str) -> None: controller.claim_project(project_id, owner_id) @@ -338,20 +342,6 @@ async def abort_project( ) async def release_project( project_id: Annotated[str, Field(description="Project id.")], - force: Annotated[ - bool, - Field( - default=False, - description="Allow release with stuck running task records.", - ), - ] = False, - reason: Annotated[ - str | None, - Field( - default=None, - description="Required audit reason when force is true.", - ), - ] = None, ctx: Context | None = None, ) -> dict[str, Any]: async with await _project_lock(project_id): @@ -361,13 +351,10 @@ async def release_project( controller.release_project, project_id, owner_id, - force=force, - reason=reason, ) return { "released": True, "projectId": project_id, - "forced": force, } except ToolError as exc: return _to_payload(exc) diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index cd22da1..cd6729c 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -15,8 +15,9 @@ """ from __future__ import annotations -import json +import fcntl import hashlib +import json import os import re import shutil @@ -318,10 +319,15 @@ def claim_workspace_lease_for_workspace( "same owner id is active on another host " f"({current.host}); lease={lease_path}" ) - if current.pid != os.getpid() and _pid_is_alive(current.pid): + if current.pid != os.getpid(): + if _pid_is_alive(current.pid): + raise WorkspaceLeaseConflict( + "same owner id is active in another controller process " + f"pid={current.pid}; lease={lease_path}" + ) raise WorkspaceLeaseConflict( - "same owner id is active in another controller process " - f"pid={current.pid}; lease={lease_path}" + "same owner id belongs to a dead controller process; " + f"explicit recovery is required; lease={lease_path}" ) lease_dir_fd = os.open( lease_dir, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) @@ -525,12 +531,14 @@ def append_workspace_lease_audit(self, payload: dict[str, object]) -> None: encoded = json.dumps(line, ensure_ascii=False, sort_keys=True) + "\n" fd = os.open(audit_path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600) try: + fcntl.flock(fd, fcntl.LOCK_EX) remaining = memoryview(encoded.encode("utf-8")) while remaining: written = os.write(fd, remaining) remaining = remaining[written:] os.fsync(fd) finally: + fcntl.flock(fd, fcntl.LOCK_UN) os.close(fd) def _recover_incomplete_workspace_lease( diff --git a/zenith/tests/test_server.py b/zenith/tests/test_server.py index 1baa930..f099f54 100644 --- a/zenith/tests/test_server.py +++ b/zenith/tests/test_server.py @@ -305,7 +305,7 @@ async def test_non_owner_cannot_release_project( @pytest.mark.asyncio -async def test_owner_must_explicitly_force_release_while_task_is_running( +async def test_owner_cannot_release_project_while_task_is_running( config: HarnessConfig, workspace: Path, monkeypatch ) -> None: monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") @@ -343,25 +343,15 @@ async def test_owner_must_explicitly_force_release_while_task_is_running( assert blocked.structured_content["error"] == "workspace_busy" assert "w1" in blocked.structured_content["message"] - missing_reason = await server.call_tool( - "release_project", {"project_id": pid, "force": True} - ) - assert missing_reason.structured_content["error"] == "recovery_reason_required" + inspected = await server.call_tool("inspect_project", {"project_id": pid}) + assert inspected.structured_content["state"]["state"] == "mission_running" - released = await server.call_tool( - "release_project", - { - "project_id": pid, - "force": True, - "reason": "worker process crashed", - }, + blocked_abort = await server.call_tool( + "abort_project", {"project_id": pid, "reason": "must stay fenced"} ) - assert released.structured_content["released"] is True - audit_path = config.harness_home / "leases" / "recovery-log.jsonl" - audit = json.loads(audit_path.read_text(encoding="utf-8").splitlines()[-1]) - assert audit["action"] == "forced_running_release_completed" - assert audit["running_task_ids"] == ["w1"] - assert audit["reason"] == "worker process crashed" + assert blocked_abort.structured_content["error"] == "workspace_busy" + inspected = await server.call_tool("inspect_project", {"project_id": pid}) + assert inspected.structured_content["state"]["state"] == "mission_running" @pytest.mark.asyncio diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index 4946fd7..c9c2e61 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -612,6 +612,20 @@ def test_owner_id_is_bounded_and_safe( workspace, "p1", "x" * 201 ) + def test_dead_process_cannot_silently_reuse_same_owner_id( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + for name in ("claim.json", "owner.json"): + path = lease.path.parent / name + payload = json.loads(path.read_text(encoding="utf-8")) + payload["pid"] = 999_999_999 + path.write_text(json.dumps(payload), encoding="utf-8") + + with pytest.raises(WorkspaceLeaseConflict, match="explicit recovery"): + store.claim_workspace_lease("p1", "owner-a") + def test_release_removes_abandoned_temp_files_atomically( self, store: ProjectStore, workspace: Path ) -> None: From 53e3c727ee9229f882fb7b8b52761733296d0008 Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 21:29:05 -0600 Subject: [PATCH 6/8] fix(controller): require audited ownership recovery --- zenith/src/zenith_harness/controller.py | 23 ++++++----- zenith/src/zenith_harness/storage.py | 39 ++++++++++++++++--- zenith/tests/test_coordinator.py | 21 +++++++++- zenith/tests/test_server.py | 52 +++++++++++++++++++++++++ zenith/tests/test_storage.py | 14 ++++++- 5 files changed, 131 insertions(+), 18 deletions(-) diff --git a/zenith/src/zenith_harness/controller.py b/zenith/src/zenith_harness/controller.py index 542e5de..17389e9 100644 --- a/zenith/src/zenith_harness/controller.py +++ b/zenith/src/zenith_harness/controller.py @@ -272,12 +272,18 @@ def inspect_project(self, project_id: str) -> Envelope: return self._build_envelope(project_id, dag_mode="full") def abort_project( - self, project_id: str, reason: str, owner_id: str | None = None + self, project_id: str, reason: str, owner_id: str ) -> Envelope: + try: + self.store.claim_workspace_lease(project_id, owner_id) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc record = self.store.load_project(project_id) state = self.store.load_state(project_id) mid = self._current_mission_id(record, state) - if owner_id is not None and mid is not None: + if mid is not None: task_state = self.store.load_task_state(project_id, mid) running = sorted( task_id @@ -300,13 +306,12 @@ def abort_project( self.store.clear_attention(project_id) self.store.save_state(project_id, Aborted(reason=reason)) envelope = self._build_envelope(project_id, dag_mode="none") - if owner_id is not None: - try: - self.store.release_workspace_lease(project_id, owner_id) - except ValueError as exc: - raise ToolError("invalid_owner", str(exc)) from exc - except WorkspaceLeaseConflict as exc: - raise ToolError("workspace_owned", str(exc)) from exc + try: + self.store.release_workspace_lease(project_id, owner_id) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc return envelope # ------------------------------------------------------------------ diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index cd6729c..6b130c2 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -299,11 +299,10 @@ def claim_workspace_lease_for_workspace( f"project={incomplete_marker.project_id} " f"pid={incomplete_marker.pid}; lease={lease_path}" ) - if self._recover_incomplete_workspace_lease( - lease_dir, incomplete_marker - ): - continue - continue + raise WorkspaceLeaseConflict( + "workspace has an incomplete dead-controller claim; " + f"explicit recovery is required; lease={lease_path}" + ) if ( current.owner_id != owner_id or current.project_id != project_id @@ -428,6 +427,11 @@ def release_workspace_lease_for_workspace( f"requesting release owner={owner_id} project={project_id}; " f"lease={lease_path}" ) + if current.host != socket.gethostname() or current.pid != os.getpid(): + raise WorkspaceLeaseConflict( + "workspace release requires the exact live controller process; " + f"lease={lease_path}" + ) released_dir = lease_path.parent.with_name( f"{lease_path.parent.name}.released.{os.getpid()}.{time.time_ns()}" ) @@ -478,6 +482,13 @@ def recover_workspace_lease_for_workspace( "cannot recover live workspace controller " f"pid={current.pid}; lease={lease_path}" ) + running = self._running_task_ids_for_project(current.project_id) + if running: + raise WorkspaceLeaseConflict( + "cannot recover while the project has running task records: " + + ", ".join(running) + + f"; lease={lease_path}" + ) recovery_audit = { "workspace_dir": str(workspace), @@ -520,6 +531,24 @@ def recover_workspace_lease_for_workspace( shutil.rmtree(quarantine) return moved + def _running_task_ids_for_project(self, project_id: str) -> list[str]: + try: + record = self.load_project(project_id) + except FileNotFoundError: + return [] + mission_id = record.current_mission_id + if mission_id is None: + return [] + try: + state = self.load_task_state(project_id, mission_id) + except FileNotFoundError: + return [] + return sorted( + task_id + for task_id, entry in state.tasks.items() + if entry.status == "running" + ) + def append_workspace_lease_audit(self, payload: dict[str, object]) -> None: audit_path = self.config.harness_home / "leases" / "recovery-log.jsonl" audit_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/zenith/tests/test_coordinator.py b/zenith/tests/test_coordinator.py index 76e88ff..9d24d22 100644 --- a/zenith/tests/test_coordinator.py +++ b/zenith/tests/test_coordinator.py @@ -649,12 +649,29 @@ def test_aborts(self, config: HarnessConfig, workspace: Path) -> None: config, MockDispatcher(lambda r: WorkHandoff(node_id=r.task.id, done=True, report="")), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - pid = _seed_project(controller, workspace) + started = controller.start_project("Brief.", str(workspace), "owner-a") + pid = started.projectId + contract_dir = controller.store.ensure_contract_dir(pid, "mission-001") + (contract_dir / "VAL-001.md").write_text("# VAL-001\n\nStatement body.\n") controller.submit_plan(pid, _simple_tl()) - env = controller.abort_project(pid, "user requested") + env = controller.abort_project(pid, "user requested", "owner-a") assert env.state.state == "aborted" assert env.dag is None + def test_non_owner_cannot_abort_state( + self, config: HarnessConfig, workspace: Path + ) -> None: + controller = ProjectController( + config, MockDispatcher(lambda r: WorkHandoff(node_id=r.task.id, done=True, report="")), + MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), + ) + started = controller.start_project("Brief.", str(workspace), "owner-a") + + with pytest.raises(ToolError, match="workspace already owned"): + controller.abort_project(started.projectId, "hostile", "owner-b") + + assert controller.inspect_project(started.projectId).state.state == "mission_planning" + class TestResume: def test_resume_picks_up_attempt_landed_while_down( diff --git a/zenith/tests/test_server.py b/zenith/tests/test_server.py index f099f54..b64cd53 100644 --- a/zenith/tests/test_server.py +++ b/zenith/tests/test_server.py @@ -484,6 +484,58 @@ async def test_live_workspace_lease_cannot_be_recovered( assert "live workspace controller" in recovered.structured_content["message"] +@pytest.mark.asyncio +async def test_dead_controller_recovery_refuses_running_worker_records( + config: HarnessConfig, workspace: Path, monkeypatch +) -> None: + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-a") + server = create_orchestrator_server(config) + started = await server.call_tool( + "start_project", {"brief": "Owned.", "workspace_dir": str(workspace)} + ) + pid = started.structured_content["projectId"] + store = ProjectStore(config) + contract_dir = store.ensure_contract_dir(pid, "mission-001") + (contract_dir / "VAL-001.md").write_text("# VAL-001\n") + await server.call_tool( + "submit_plan", + { + "project_id": pid, + "task_list": { + "tasks": [ + { + "id": "w1", + "type": "work", + "body": "do", + "targets": ["VAL-001"], + "skill": "s", + "depends_on": [], + } + ] + }, + }, + ) + task_state = store.load_task_state(pid, "mission-001") + task_state.set_status("w1", "running") + store.save_task_state(pid, "mission-001", task_state) + lease_path = store.workspace_lease_path(pid) + for name in ("claim.json", "owner.json"): + path = lease_path.parent / name + payload = json.loads(path.read_text(encoding="utf-8")) + payload["pid"] = 999_999_999 + path.write_text(json.dumps(payload), encoding="utf-8") + + monkeypatch.setenv("ZENITH_CONTROLLER_ID", "owner-b") + recovery_server = create_orchestrator_server(config) + recovered = await recovery_server.call_tool( + "recover_workspace_lease", + {"workspace_dir": str(workspace), "reason": "controller crashed"}, + ) + + assert recovered.structured_content["error"] == "workspace_recovery_blocked" + assert "running task records" in recovered.structured_content["message"] + + # --------------------------------------------------------------------------- # Regression: dispatcher that calls asyncio.run() must not crash the MCP # event loop. Reproduces: diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index c9c2e61..0ad22cd 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -562,7 +562,7 @@ def claim_second() -> str: assert store.claim_workspace_lease("p1", "owner-a").owner_id == "owner-a" - def test_ownerless_dead_claim_is_fenced_and_recovered( + def test_ownerless_dead_claim_requires_audited_recovery( self, store: ProjectStore, workspace: Path ) -> None: store.create_project("brief", workspace, project_id="p1") @@ -573,10 +573,15 @@ def test_ownerless_dead_claim_is_fenced_and_recovered( claim["pid"] = 999_999_999 claim_path.write_text(json.dumps(claim), encoding="utf-8") + with pytest.raises(WorkspaceLeaseConflict, match="explicit recovery"): + store.claim_workspace_lease("p1", "owner-b") + store.recover_workspace_lease_for_workspace( + workspace, "owner-b", "owner publication crashed" + ) recovered = store.claim_workspace_lease("p1", "owner-b") assert recovered.owner_id == "owner-b" - def test_malformed_dead_claim_is_fenced_and_recovered( + def test_malformed_dead_claim_requires_audited_recovery( self, store: ProjectStore, workspace: Path ) -> None: store.create_project("brief", workspace, project_id="p1") @@ -587,6 +592,11 @@ def test_malformed_dead_claim_is_fenced_and_recovered( claim["pid"] = 999_999_999 claim_path.write_text(json.dumps(claim), encoding="utf-8") + with pytest.raises(WorkspaceLeaseConflict, match="explicit recovery"): + store.claim_workspace_lease("p1", "owner-b") + store.recover_workspace_lease_for_workspace( + workspace, "owner-b", "owner publication was malformed" + ) recovered = store.claim_workspace_lease("p1", "owner-b") assert recovered.owner_id == "owner-b" From 691d98872d0b4615419bacc588f5f5177b3dd253 Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 21:44:49 -0600 Subject: [PATCH 7/8] fix(controller): bind recovery to runtime identity --- zenith/src/zenith_harness/controller.py | 73 +++++++++++---- zenith/src/zenith_harness/server.py | 8 +- zenith/src/zenith_harness/storage.py | 108 ++++++++++++++++++++-- zenith/tests/test_coordinator.py | 6 +- zenith/tests/test_coordinator_parallel.py | 16 ++-- zenith/tests/test_runnable_selection.py | 2 +- zenith/tests/test_smoke_parallel_acp.py | 2 +- zenith/tests/test_smoke_real_acp.py | 4 +- zenith/tests/test_storage.py | 39 ++++++++ zenith/tests/test_terminal_review.py | 2 +- 10 files changed, 219 insertions(+), 41 deletions(-) diff --git a/zenith/src/zenith_harness/controller.py b/zenith/src/zenith_harness/controller.py index 17389e9..3d5c8e7 100644 --- a/zenith/src/zenith_harness/controller.py +++ b/zenith/src/zenith_harness/controller.py @@ -73,7 +73,7 @@ def __init__( # ------------------------------------------------------------------ def start_project( - self, brief: str, workspace_dir: str, owner_id: str | None = None + self, brief: str, workspace_dir: str, owner_id: str ) -> Envelope: if not brief.strip(): raise ToolError("invalid_brief", "brief is empty") @@ -84,24 +84,22 @@ def start_project( "workspace_dir must be an existing absolute directory", ) project_id = self.store.generate_project_id(brief) - if owner_id is not None: - try: - self.store.claim_workspace_lease_for_workspace( - workspace_dir, project_id, owner_id - ) - except ValueError as exc: - raise ToolError("invalid_owner", str(exc)) from exc - except WorkspaceLeaseConflict as exc: - raise ToolError("workspace_owned", str(exc)) from exc + try: + self.store.claim_workspace_lease_for_workspace( + workspace_dir, project_id, owner_id + ) + except ValueError as exc: + raise ToolError("invalid_owner", str(exc)) from exc + except WorkspaceLeaseConflict as exc: + raise ToolError("workspace_owned", str(exc)) from exc try: record = self.store.create_project( brief, workspace_dir, project_id=project_id ) except Exception: - if owner_id is not None: - self.store.release_workspace_lease_for_workspace( - workspace_dir, project_id, owner_id - ) + self.store.release_workspace_lease_for_workspace( + workspace_dir, project_id, owner_id + ) raise mission_id = self.store.generate_mission_id(1) record.current_mission_id = mission_id @@ -130,6 +128,8 @@ def release_project( self.store.claim_workspace_lease(project_id, owner_id) except ValueError as exc: raise ToolError("invalid_owner", str(exc)) from exc + except FileNotFoundError as exc: + raise ToolError("invalid_workspace", str(exc)) from exc except WorkspaceLeaseConflict as exc: raise ToolError("workspace_owned", str(exc)) from exc record = self.store.load_project(project_id) @@ -152,12 +152,14 @@ def release_project( self.store.release_workspace_lease(project_id, owner_id) except ValueError as exc: raise ToolError("invalid_owner", str(exc)) from exc + except FileNotFoundError as exc: + raise ToolError("invalid_workspace", str(exc)) from exc except WorkspaceLeaseConflict as exc: raise ToolError("workspace_owned", str(exc)) from exc def recover_workspace_lease( self, workspace_dir: str, owner_id: str, reason: str - ) -> WorkspaceLease: + ) -> WorkspaceLease | None: try: return self.store.recover_workspace_lease_for_workspace( workspace_dir, owner_id, reason @@ -278,6 +280,8 @@ def abort_project( self.store.claim_workspace_lease(project_id, owner_id) except ValueError as exc: raise ToolError("invalid_owner", str(exc)) from exc + except FileNotFoundError as exc: + raise ToolError("invalid_workspace", str(exc)) from exc except WorkspaceLeaseConflict as exc: raise ToolError("workspace_owned", str(exc)) from exc record = self.store.load_project(project_id) @@ -309,9 +313,44 @@ def abort_project( try: self.store.release_workspace_lease(project_id, owner_id) except ValueError as exc: - raise ToolError("invalid_owner", str(exc)) from exc + self.store.append_workspace_lease_audit( + { + "action": "abort_state_committed_lease_release_failed", + "project_id": project_id, + "owner_id": owner_id, + "error": str(exc), + } + ) + raise ToolError( + "abort_succeeded_lease_release_failed", + f"project state is aborted but lease release failed: {exc}", + ) from exc + except FileNotFoundError as exc: + self.store.append_workspace_lease_audit( + { + "action": "abort_state_committed_lease_release_failed", + "project_id": project_id, + "owner_id": owner_id, + "error": str(exc), + } + ) + raise ToolError( + "abort_succeeded_lease_release_failed", + f"project state is aborted but lease release failed: {exc}", + ) from exc except WorkspaceLeaseConflict as exc: - raise ToolError("workspace_owned", str(exc)) from exc + self.store.append_workspace_lease_audit( + { + "action": "abort_state_committed_lease_release_failed", + "project_id": project_id, + "owner_id": owner_id, + "error": str(exc), + } + ) + raise ToolError( + "abort_succeeded_lease_release_failed", + f"project state is aborted but lease release failed: {exc}", + ) from exc return envelope # ------------------------------------------------------------------ diff --git a/zenith/src/zenith_harness/server.py b/zenith/src/zenith_harness/server.py index a1da79b..7e84dd6 100644 --- a/zenith/src/zenith_harness/server.py +++ b/zenith/src/zenith_harness/server.py @@ -387,8 +387,12 @@ async def recover_workspace_lease( ) return { "recovered": True, - "workspaceDir": recovered.workspace_dir, - "previousProjectId": recovered.project_id, + "workspaceDir": ( + recovered.workspace_dir if recovered is not None else workspace_dir + ), + "previousProjectId": ( + recovered.project_id if recovered is not None else None + ), } except ToolError as exc: return _to_payload(exc) diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index 6b130c2..8d4a83b 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -19,10 +19,13 @@ import hashlib import json import os +import platform import re import shutil import socket +import subprocess import time +import uuid from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path @@ -129,6 +132,7 @@ class WorkspaceLease: last_seen_at: str host: str pid: int + runtime_id: str path: Path @@ -162,6 +166,30 @@ def _pid_is_alive(pid: int) -> bool: return True +def _runtime_identity() -> str: + """Bind recovery to this machine boot and PID namespace when available.""" + parts = [platform.system(), socket.gethostname(), str(uuid.getnode())] + if platform.system() == "Linux": + try: + parts.append(Path("/proc/sys/kernel/random/boot_id").read_text().strip()) + parts.append(str(Path("/proc/self/ns/pid").stat().st_ino)) + except OSError: + return "" + elif platform.system() == "Darwin": + try: + boot = subprocess.check_output( + ["/usr/sbin/sysctl", "-n", "kern.boottime"], + text=True, + timeout=2, + ).strip() + except (OSError, subprocess.SubprocessError): + return "" + parts.append(boot) + else: + return "" + return hashlib.sha256("\0".join(parts).encode("utf-8")).hexdigest() + + # --------------------------------------------------------------------------- # ProjectStore # --------------------------------------------------------------------------- @@ -258,6 +286,7 @@ def claim_workspace_lease_for_workspace( "workspace_dir": str(workspace), "host": socket.gethostname(), "pid": os.getpid(), + "runtime_id": _runtime_identity(), "claimed_at": claimed_at, "last_seen_at": claimed_at, } @@ -366,6 +395,7 @@ def claim_workspace_lease_for_workspace( "workspace_dir": str(workspace), "host": socket.gethostname(), "pid": os.getpid(), + "runtime_id": _runtime_identity(), "claimed_at": claimed_at, "last_seen_at": now, } @@ -397,6 +427,7 @@ def claim_workspace_lease_for_workspace( last_seen_at=now, host=socket.gethostname(), pid=os.getpid(), + runtime_id=_runtime_identity(), path=lease_path, ) @@ -448,7 +479,7 @@ def recover_workspace_lease_for_workspace( workspace_dir: str | Path, requesting_owner_id: str, reason: str, - ) -> WorkspaceLease: + ) -> WorkspaceLease | None: """Release a lease only when its recorded same-host controller is dead.""" requesting_owner_id = _normalize_owner_id(requesting_owner_id) reason = reason.strip() @@ -469,14 +500,31 @@ def recover_workspace_lease_for_workspace( lease_path.parent / "claim.json" ) if current is None: - raise WorkspaceLeaseConflict( - f"workspace lease has no provable owner identity: {lease_path}" + for _ in range(20): + time.sleep(0.05) + current = self._load_workspace_lease_record(lease_path) + if current is None: + current = self._load_workspace_lease_record( + lease_path.parent / "claim.json" + ) + if current is not None: + break + if current is None: + self._recover_identityless_workspace_lease( + lease_path, requesting_owner_id, reason ) + return None if current.host != socket.gethostname(): raise WorkspaceLeaseConflict( "cannot prove controller death on another host: " f"{current.host}; lease={lease_path}" ) + runtime_id = _runtime_identity() + if not runtime_id or not current.runtime_id or current.runtime_id != runtime_id: + raise WorkspaceLeaseConflict( + "cannot prove controller death across a different boot or PID namespace; " + f"lease={lease_path}" + ) if _pid_is_alive(current.pid): raise WorkspaceLeaseConflict( "cannot recover live workspace controller " @@ -496,6 +544,7 @@ def recover_workspace_lease_for_workspace( "previous_owner_id": current.owner_id, "previous_host": current.host, "previous_pid": current.pid, + "previous_runtime_id": current.runtime_id, "requesting_owner_id": requesting_owner_id, "reason": reason, } @@ -520,7 +569,11 @@ def recover_workspace_lease_for_workspace( raise WorkspaceLeaseConflict( f"workspace lease changed during recovery: {lease_path}" ) - if moved.host != socket.gethostname() or _pid_is_alive(moved.pid): + if ( + moved.host != socket.gethostname() + or moved.runtime_id != runtime_id + or _pid_is_alive(moved.pid) + ): self._restore_quarantined_workspace_lease(quarantine, lease_path.parent) raise WorkspaceLeaseConflict( f"workspace controller revived during recovery: {lease_path}" @@ -531,6 +584,41 @@ def recover_workspace_lease_for_workspace( shutil.rmtree(quarantine) return moved + def _recover_identityless_workspace_lease( + self, lease_path: Path, requesting_owner_id: str, reason: str + ) -> None: + audit = { + "workspace_dir": str(lease_path), + "requesting_owner_id": requesting_owner_id, + "reason": reason, + } + self.append_workspace_lease_audit( + {"action": "identityless_recovery_authorized", **audit} + ) + quarantine = lease_path.parent.with_name( + f"{lease_path.parent.name}.identityless.{os.getpid()}.{time.time_ns()}" + ) + try: + os.replace(lease_path.parent, quarantine) + except FileNotFoundError as exc: + raise WorkspaceLeaseConflict( + f"workspace lease changed during identityless recovery: {lease_path}" + ) from exc + if self._load_workspace_lease_record( + quarantine / "owner.json" + ) is not None or self._load_workspace_lease_record( + quarantine / "claim.json" + ) is not None: + self._restore_quarantined_workspace_lease(quarantine, lease_path.parent) + raise WorkspaceLeaseConflict( + f"workspace identity appeared during recovery: {lease_path}" + ) + self.append_workspace_lease_audit( + {"action": "identityless_recovery_completed", **audit} + ) + shutil.rmtree(quarantine) + return None + def _running_task_ids_for_project(self, project_id: str) -> list[str]: try: record = self.load_project(project_id) @@ -615,8 +703,14 @@ def _recover_incomplete_workspace_lease( return True @staticmethod - def _lease_identity(lease: WorkspaceLease) -> tuple[str, str, str, int]: - return (lease.owner_id, lease.project_id, lease.host, lease.pid) + def _lease_identity(lease: WorkspaceLease) -> tuple[str, str, str, int, str]: + return ( + lease.owner_id, + lease.project_id, + lease.host, + lease.pid, + lease.runtime_id, + ) @staticmethod def _restore_quarantined_workspace_lease( @@ -684,6 +778,7 @@ def _load_workspace_lease_record_from_dir_fd( last_seen_at=str(payload["last_seen_at"]), host=str(payload.get("host", "")), pid=int(payload.get("pid", 0)), + runtime_id=str(payload.get("runtime_id", "")), path=path, ) except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError): @@ -701,6 +796,7 @@ def _load_workspace_lease_record(path: Path) -> WorkspaceLease | None: last_seen_at=str(payload["last_seen_at"]), host=str(payload.get("host", "")), pid=int(payload.get("pid", 0)), + runtime_id=str(payload.get("runtime_id", "")), path=path, ) except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError): diff --git a/zenith/tests/test_coordinator.py b/zenith/tests/test_coordinator.py index 9d24d22..f1d6691 100644 --- a/zenith/tests/test_coordinator.py +++ b/zenith/tests/test_coordinator.py @@ -94,7 +94,7 @@ def _seed_project( mission_id: str = "mission-001", assertion: str = "VAL-001", ) -> str: - controller.start_project(brief, str(workspace)) + controller.start_project(brief, str(workspace), "test-owner") pid = controller.store.list_projects()[0].id contract_dir = controller.store.ensure_contract_dir(pid, mission_id) (contract_dir / f"{assertion}.md").write_text( @@ -395,7 +395,7 @@ def test_contract_less_plan_rejected( MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) # Seed the project but author no contract assertion files. - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id with pytest.raises(ToolError) as exc: @@ -605,7 +605,7 @@ def responder(req: DispatchRequest) -> NodeHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - start = controller.start_project("Brief.", str(workspace)) + start = controller.start_project("Brief.", str(workspace), "test-owner") assert start.dag is None pid = start.projectId contract_dir = controller.store.ensure_contract_dir(pid, "mission-001") diff --git a/zenith/tests/test_coordinator_parallel.py b/zenith/tests/test_coordinator_parallel.py index a290e22..c354c23 100644 --- a/zenith/tests/test_coordinator_parallel.py +++ b/zenith/tests/test_coordinator_parallel.py @@ -70,7 +70,7 @@ def responder(req: DispatchRequest) -> WorkHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "VAL-A") _write_contract(controller.store, pid, "mission-001", "VAL-B") @@ -102,7 +102,7 @@ def responder(req: DispatchRequest) -> WorkHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "EXP-A") controller.submit_plan( @@ -162,7 +162,7 @@ def responder(req: DispatchRequest) -> WorkHandoff | ValidateHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "VAL-A") controller.submit_plan( @@ -246,7 +246,7 @@ def responder(req: DispatchRequest) -> WorkHandoff | ValidateHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "VAL-A") controller.submit_plan( @@ -320,7 +320,7 @@ def responder(req: DispatchRequest) -> WorkHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "VAL-A") _write_contract(controller.store, pid, "mission-001", "VAL-B") @@ -355,7 +355,7 @@ def responder(req: DispatchRequest) -> WorkHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "VAL-A") _write_contract(controller.store, pid, "mission-001", "VAL-B") @@ -378,7 +378,7 @@ def test_parallel_non_git_workspace_runs_in_workspace( MockDispatcher(lambda r: WorkHandoff(node_id=r.task.id, done=True, report="")), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "VAL-A") controller.submit_plan(pid, TaskList(tasks=[_task("a", "VAL-A")])) @@ -397,7 +397,7 @@ def test_reconcile_applies_multiple_running_attempts( MockDispatcher(lambda r: WorkHandoff(node_id=r.task.id, done=True, report="unused")), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id _write_contract(controller.store, pid, "mission-001", "VAL-A") _write_contract(controller.store, pid, "mission-001", "VAL-B") diff --git a/zenith/tests/test_runnable_selection.py b/zenith/tests/test_runnable_selection.py index a52e8fa..23347c6 100644 --- a/zenith/tests/test_runnable_selection.py +++ b/zenith/tests/test_runnable_selection.py @@ -180,7 +180,7 @@ def responder(req: DispatchRequest) -> NodeHandoff: MockDispatcher(responder), MockTerminalReviewer(TerminalReviewHandoff(done=True, report="")), ) - controller.start_project("Brief.", str(workspace)) + controller.start_project("Brief.", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id contract_dir = controller.store.ensure_contract_dir(pid, "mission-001") (contract_dir / "VAL-A.md").write_text("# VAL-A\n") diff --git a/zenith/tests/test_smoke_parallel_acp.py b/zenith/tests/test_smoke_parallel_acp.py index 3d012d6..0ad30e3 100644 --- a/zenith/tests/test_smoke_parallel_acp.py +++ b/zenith/tests/test_smoke_parallel_acp.py @@ -154,7 +154,7 @@ def _run_parallel_workspace_acp( dispatcher = ACPNodeDispatcher(config) controller = ProjectController(config, dispatcher, _CleanReviewer()) - controller.start_project("parallel smoke", str(workspace)) + controller.start_project("parallel smoke", str(workspace), "test-owner") store = ProjectStore(config) pid = store.list_projects()[0].id diff --git a/zenith/tests/test_smoke_real_acp.py b/zenith/tests/test_smoke_real_acp.py index 66b65b6..ae1f195 100644 --- a/zenith/tests/test_smoke_real_acp.py +++ b/zenith/tests/test_smoke_real_acp.py @@ -252,7 +252,7 @@ def test_smoke_hello_mission( ) # 1) start_project - start_env = controller.start_project("smoke test brief", str(workspace)) + start_env = controller.start_project("smoke test brief", str(workspace), "test-owner") pid = ProjectStore(config).list_projects()[0].id assert start_env.state.state == "mission_planning" @@ -376,7 +376,7 @@ def _run_with_real_terminal_reviewer( reviewer = ACPTerminalReviewer(config) controller = ProjectController(config, dispatcher, reviewer) - controller.start_project("smoke test brief", str(workspace)) + controller.start_project("smoke test brief", str(workspace), "test-owner") pid = ProjectStore(config).list_projects()[0].id _write_contract(workspace, "mission-001") controller.submit_plan(pid, _build_task_list()) diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index 0ad22cd..1a7ff06 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -496,6 +496,7 @@ def test_record_contains_forensic_identity( assert payload["last_seen_at"] assert payload["host"] assert payload["pid"] > 0 + assert payload["runtime_id"] marker = json.loads( (lease.path.parent / "claim.json").read_text(encoding="utf-8") ) @@ -636,6 +637,44 @@ def test_dead_process_cannot_silently_reuse_same_owner_id( with pytest.raises(WorkspaceLeaseConflict, match="explicit recovery"): store.claim_workspace_lease("p1", "owner-a") + def test_identityless_lease_has_audited_break_glass_recovery( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease_path = store.workspace_lease_path("p1") + lease_path.parent.mkdir(parents=True) + + recovered = store.recover_workspace_lease_for_workspace( + workspace, "operator", "claim process died before identity publication" + ) + + assert recovered is None + assert not lease_path.parent.exists() + audit_path = store.config.harness_home / "leases" / "recovery-log.jsonl" + audit = [ + json.loads(line) + for line in audit_path.read_text(encoding="utf-8").splitlines() + ] + assert audit[-2]["action"] == "identityless_recovery_authorized" + assert audit[-1]["action"] == "identityless_recovery_completed" + + def test_recovery_refuses_different_boot_or_pid_namespace( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + for name in ("claim.json", "owner.json"): + path = lease.path.parent / name + payload = json.loads(path.read_text(encoding="utf-8")) + payload["pid"] = 999_999_999 + payload["runtime_id"] = "different-runtime" + path.write_text(json.dumps(payload), encoding="utf-8") + + with pytest.raises(WorkspaceLeaseConflict, match="boot or PID namespace"): + store.recover_workspace_lease_for_workspace( + workspace, "operator", "controller appears dead" + ) + def test_release_removes_abandoned_temp_files_atomically( self, store: ProjectStore, workspace: Path ) -> None: diff --git a/zenith/tests/test_terminal_review.py b/zenith/tests/test_terminal_review.py index 4c19166..f7588f5 100644 --- a/zenith/tests/test_terminal_review.py +++ b/zenith/tests/test_terminal_review.py @@ -72,7 +72,7 @@ def _start_and_seed_contract( controller: ProjectController, workspace: Path ) -> str: """Run start_project, then seed mission-001 contract VAL-001 inside the bucket.""" - controller.start_project("brief", str(workspace)) + controller.start_project("brief", str(workspace), "test-owner") pid = controller.store.list_projects()[0].id contract_dir = controller.store.ensure_contract_dir(pid, "mission-001") (contract_dir / "VAL-001.md").write_text("# VAL-001\n\nstatement\n") From af99cd820e45641191bca5242c33d06f4994c2ba Mon Sep 17 00:00:00 2001 From: Dfskid Date: Mon, 20 Jul 2026 21:58:52 -0600 Subject: [PATCH 8/8] fix: bind leases to process birth identity --- zenith/src/zenith_harness/storage.py | 132 ++++++++++++++++----------- zenith/tests/test_storage.py | 64 ++++++++++--- 2 files changed, 134 insertions(+), 62 deletions(-) diff --git a/zenith/src/zenith_harness/storage.py b/zenith/src/zenith_harness/storage.py index 8d4a83b..440355a 100644 --- a/zenith/src/zenith_harness/storage.py +++ b/zenith/src/zenith_harness/storage.py @@ -132,6 +132,7 @@ class WorkspaceLease: last_seen_at: str host: str pid: int + process_start_id: str runtime_id: str path: Path @@ -166,6 +167,41 @@ def _pid_is_alive(pid: int) -> bool: return True +def _process_start_id(pid: int) -> str: + """Return a boot-local process birth identity, or empty when unprovable.""" + if pid <= 0: + return "" + system = platform.system() + if system == "Linux": + try: + stat = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8") + after_comm = stat.rsplit(")", 1)[1].strip().split() + return f"linux:{after_comm[19]}" + except (OSError, IndexError): + return "" + if system == "Darwin": + try: + started = subprocess.check_output( + ["/bin/ps", "-o", "lstart=", "-p", str(pid)], + text=True, + timeout=2, + ).strip() + except (OSError, subprocess.SubprocessError): + return "" + return f"darwin:{started}" if started else "" + return "" + + +def _process_identity_is_live(pid: int, expected_start_id: str) -> bool: + """Treat an unprovable identity as live; only a proven mismatch is stale.""" + if not _pid_is_alive(pid): + return False + actual_start_id = _process_start_id(pid) + if not actual_start_id or not expected_start_id: + return True + return actual_start_id == expected_start_id + + def _runtime_identity() -> str: """Bind recovery to this machine boot and PID namespace when available.""" parts = [platform.system(), socket.gethostname(), str(uuid.getnode())] @@ -266,6 +302,13 @@ def claim_workspace_lease_for_workspace( workspace = requested_workspace.resolve() if not workspace.is_dir(): raise FileNotFoundError(f"workspace_dir does not exist: {workspace}") + runtime_id = _runtime_identity() + process_start_id = _process_start_id(os.getpid()) + if not runtime_id or not process_start_id: + raise WorkspaceLeaseConflict( + "cannot claim workspace without a provable runtime identity and " + "process start identity" + ) lease_path = self._workspace_lease_path(workspace) lease_dir = lease_path.parent lease_dir.parent.mkdir(parents=True, exist_ok=True) @@ -280,13 +323,14 @@ def claim_workspace_lease_for_workspace( ) claimed_at = utc_now_iso() claim_payload = { - "schema_version": 1, + "schema_version": 2, "owner_id": owner_id, "project_id": project_id, "workspace_dir": str(workspace), "host": socket.gethostname(), "pid": os.getpid(), - "runtime_id": _runtime_identity(), + "process_start_id": process_start_id, + "runtime_id": runtime_id, "claimed_at": claimed_at, "last_seen_at": claimed_at, } @@ -321,7 +365,9 @@ def claim_workspace_lease_for_workspace( "workspace lease has an incomplete claim on another " f"host ({incomplete_marker.host}): {lease_path}" ) - if _pid_is_alive(incomplete_marker.pid): + if _process_identity_is_live( + incomplete_marker.pid, incomplete_marker.process_start_id + ): raise WorkspaceLeaseConflict( "workspace has incomplete active claim by " f"owner={incomplete_marker.owner_id} " @@ -347,8 +393,13 @@ def claim_workspace_lease_for_workspace( "same owner id is active on another host " f"({current.host}); lease={lease_path}" ) - if current.pid != os.getpid(): - if _pid_is_alive(current.pid): + if ( + current.pid != os.getpid() + or current.process_start_id != process_start_id + ): + if _process_identity_is_live( + current.pid, current.process_start_id + ): raise WorkspaceLeaseConflict( "same owner id is active in another controller process " f"pid={current.pid}; lease={lease_path}" @@ -389,13 +440,14 @@ def claim_workspace_lease_for_workspace( now = utc_now_iso() payload = { - "schema_version": 1, + "schema_version": 2, "owner_id": owner_id, "project_id": project_id, "workspace_dir": str(workspace), "host": socket.gethostname(), "pid": os.getpid(), - "runtime_id": _runtime_identity(), + "process_start_id": process_start_id, + "runtime_id": runtime_id, "claimed_at": claimed_at, "last_seen_at": now, } @@ -427,7 +479,8 @@ def claim_workspace_lease_for_workspace( last_seen_at=now, host=socket.gethostname(), pid=os.getpid(), - runtime_id=_runtime_identity(), + process_start_id=process_start_id, + runtime_id=runtime_id, path=lease_path, ) @@ -458,7 +511,11 @@ def release_workspace_lease_for_workspace( f"requesting release owner={owner_id} project={project_id}; " f"lease={lease_path}" ) - if current.host != socket.gethostname() or current.pid != os.getpid(): + if ( + current.host != socket.gethostname() + or current.pid != os.getpid() + or current.process_start_id != _process_start_id(os.getpid()) + ): raise WorkspaceLeaseConflict( "workspace release requires the exact live controller process; " f"lease={lease_path}" @@ -510,10 +567,10 @@ def recover_workspace_lease_for_workspace( if current is not None: break if current is None: - self._recover_identityless_workspace_lease( - lease_path, requesting_owner_id, reason + raise WorkspaceLeaseConflict( + "workspace lease has no recoverable claim identity; manual " + f"forensic intervention is required; lease={lease_path}" ) - return None if current.host != socket.gethostname(): raise WorkspaceLeaseConflict( "cannot prove controller death on another host: " @@ -525,7 +582,7 @@ def recover_workspace_lease_for_workspace( "cannot prove controller death across a different boot or PID namespace; " f"lease={lease_path}" ) - if _pid_is_alive(current.pid): + if _process_identity_is_live(current.pid, current.process_start_id): raise WorkspaceLeaseConflict( "cannot recover live workspace controller " f"pid={current.pid}; lease={lease_path}" @@ -544,6 +601,7 @@ def recover_workspace_lease_for_workspace( "previous_owner_id": current.owner_id, "previous_host": current.host, "previous_pid": current.pid, + "previous_process_start_id": current.process_start_id, "previous_runtime_id": current.runtime_id, "requesting_owner_id": requesting_owner_id, "reason": reason, @@ -572,7 +630,7 @@ def recover_workspace_lease_for_workspace( if ( moved.host != socket.gethostname() or moved.runtime_id != runtime_id - or _pid_is_alive(moved.pid) + or _process_identity_is_live(moved.pid, moved.process_start_id) ): self._restore_quarantined_workspace_lease(quarantine, lease_path.parent) raise WorkspaceLeaseConflict( @@ -584,41 +642,6 @@ def recover_workspace_lease_for_workspace( shutil.rmtree(quarantine) return moved - def _recover_identityless_workspace_lease( - self, lease_path: Path, requesting_owner_id: str, reason: str - ) -> None: - audit = { - "workspace_dir": str(lease_path), - "requesting_owner_id": requesting_owner_id, - "reason": reason, - } - self.append_workspace_lease_audit( - {"action": "identityless_recovery_authorized", **audit} - ) - quarantine = lease_path.parent.with_name( - f"{lease_path.parent.name}.identityless.{os.getpid()}.{time.time_ns()}" - ) - try: - os.replace(lease_path.parent, quarantine) - except FileNotFoundError as exc: - raise WorkspaceLeaseConflict( - f"workspace lease changed during identityless recovery: {lease_path}" - ) from exc - if self._load_workspace_lease_record( - quarantine / "owner.json" - ) is not None or self._load_workspace_lease_record( - quarantine / "claim.json" - ) is not None: - self._restore_quarantined_workspace_lease(quarantine, lease_path.parent) - raise WorkspaceLeaseConflict( - f"workspace identity appeared during recovery: {lease_path}" - ) - self.append_workspace_lease_audit( - {"action": "identityless_recovery_completed", **audit} - ) - shutil.rmtree(quarantine) - return None - def _running_task_ids_for_project(self, project_id: str) -> list[str]: try: record = self.load_project(project_id) @@ -692,7 +715,9 @@ def _recover_incomplete_workspace_lease( moved_marker is None or self._lease_identity(moved_marker) != self._lease_identity(expected) or moved_marker.host != socket.gethostname() - or _pid_is_alive(moved_marker.pid) + or _process_identity_is_live( + moved_marker.pid, moved_marker.process_start_id + ) ): self._restore_quarantined_workspace_lease(quarantine, lease_dir) raise WorkspaceLeaseConflict( @@ -703,12 +728,15 @@ def _recover_incomplete_workspace_lease( return True @staticmethod - def _lease_identity(lease: WorkspaceLease) -> tuple[str, str, str, int, str]: + def _lease_identity( + lease: WorkspaceLease, + ) -> tuple[str, str, str, int, str, str]: return ( lease.owner_id, lease.project_id, lease.host, lease.pid, + lease.process_start_id, lease.runtime_id, ) @@ -778,6 +806,7 @@ def _load_workspace_lease_record_from_dir_fd( last_seen_at=str(payload["last_seen_at"]), host=str(payload.get("host", "")), pid=int(payload.get("pid", 0)), + process_start_id=str(payload.get("process_start_id", "")), runtime_id=str(payload.get("runtime_id", "")), path=path, ) @@ -796,6 +825,7 @@ def _load_workspace_lease_record(path: Path) -> WorkspaceLease | None: last_seen_at=str(payload["last_seen_at"]), host=str(payload.get("host", "")), pid=int(payload.get("pid", 0)), + process_start_id=str(payload.get("process_start_id", "")), runtime_id=str(payload.get("runtime_id", "")), path=path, ) diff --git a/zenith/tests/test_storage.py b/zenith/tests/test_storage.py index 1a7ff06..cc23a67 100644 --- a/zenith/tests/test_storage.py +++ b/zenith/tests/test_storage.py @@ -497,6 +497,7 @@ def test_record_contains_forensic_identity( assert payload["host"] assert payload["pid"] > 0 assert payload["runtime_id"] + assert payload["process_start_id"] marker = json.loads( (lease.path.parent / "claim.json").read_text(encoding="utf-8") ) @@ -637,26 +638,67 @@ def test_dead_process_cannot_silently_reuse_same_owner_id( with pytest.raises(WorkspaceLeaseConflict, match="explicit recovery"): store.claim_workspace_lease("p1", "owner-a") - def test_identityless_lease_has_audited_break_glass_recovery( + def test_identityless_lease_remains_fenced_without_creator_identity( self, store: ProjectStore, workspace: Path ) -> None: store.create_project("brief", workspace, project_id="p1") lease_path = store.workspace_lease_path("p1") lease_path.parent.mkdir(parents=True) + with pytest.raises( + WorkspaceLeaseConflict, match="has no recoverable claim identity" + ): + store.recover_workspace_lease_for_workspace( + workspace, "operator", "claim process died before identity publication" + ) + + assert lease_path.parent.exists() + + def test_claim_fails_before_publish_when_runtime_identity_is_unavailable( + self, + store: ProjectStore, + workspace: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + store.create_project("brief", workspace, project_id="p1") + monkeypatch.setattr("zenith_harness.storage._runtime_identity", lambda: "") + + with pytest.raises(WorkspaceLeaseConflict, match="runtime identity"): + store.claim_workspace_lease("p1", "owner-a") + + assert not store.workspace_lease_path("p1").parent.exists() + + def test_recovery_distinguishes_reused_live_pid_by_process_start( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + for name in ("claim.json", "owner.json"): + path = lease.path.parent / name + payload = json.loads(path.read_text(encoding="utf-8")) + payload["process_start_id"] = "recycled-process-start" + path.write_text(json.dumps(payload), encoding="utf-8") + recovered = store.recover_workspace_lease_for_workspace( - workspace, "operator", "claim process died before identity publication" + workspace, "operator", "recorded PID was reused" ) - assert recovered is None - assert not lease_path.parent.exists() - audit_path = store.config.harness_home / "leases" / "recovery-log.jsonl" - audit = [ - json.loads(line) - for line in audit_path.read_text(encoding="utf-8").splitlines() - ] - assert audit[-2]["action"] == "identityless_recovery_authorized" - assert audit[-1]["action"] == "identityless_recovery_completed" + assert recovered is not None + assert not lease.path.parent.exists() + + def test_reused_pid_cannot_refresh_the_previous_process_lease( + self, store: ProjectStore, workspace: Path + ) -> None: + store.create_project("brief", workspace, project_id="p1") + lease = store.claim_workspace_lease("p1", "owner-a") + for name in ("claim.json", "owner.json"): + path = lease.path.parent / name + payload = json.loads(path.read_text(encoding="utf-8")) + payload["process_start_id"] = "recycled-process-start" + path.write_text(json.dumps(payload), encoding="utf-8") + + with pytest.raises(WorkspaceLeaseConflict, match="explicit recovery"): + store.claim_workspace_lease("p1", "owner-a") def test_recovery_refuses_different_boot_or_pid_namespace( self, store: ProjectStore, workspace: Path