Refine course progression and runtime safety

s15_agent_teams/code.py

@@ -20,7 +20,8 @@ ASCII flow:
                   └──────── MessageBus + typed protocol ┘
 """
 
-import os, subprocess, json, time, random, threading, queue, re
+import atexit, fcntl, os, signal, subprocess, json, time, random, threading, queue, re
+from contextlib import contextmanager
 from pathlib import Path
 from datetime import datetime
 from dataclasses import dataclass, asdict, field
@@ -50,10 +51,54 @@ TASKS_DIR = WORKDIR / ".tasks"
 TASKS_DIR.mkdir(exist_ok=True)
 TASKS_ROOT = TASKS_DIR.resolve()
 task_lock = threading.RLock()
+TASK_LOCK_PATH = TASKS_DIR / ".lock"
+_task_store_state = threading.local()
 
 # owner -> {"task_id": str, "cwd": Path}. A teammate gets one assignment at
 # a time, and every filesystem tool resolves its cwd through this registry.
 teammate_assignments: dict[str, dict[str, object]] = {}
+assignment_versions: dict[str, int] = {}
+
+
+@contextmanager
+def task_store_lock():
+    """Serialize task mutations across threads and host processes."""
+    with task_lock:
+        depth = getattr(_task_store_state, "depth", 0)
+        if depth == 0:
+            handle = TASK_LOCK_PATH.open("a+")
+            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
+            _task_store_state.handle = handle
+        _task_store_state.depth = depth + 1
+        try:
+            yield
+        finally:
+            _task_store_state.depth -= 1
+            if _task_store_state.depth == 0:
+                handle = _task_store_state.handle
+                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+                handle.close()
+                del _task_store_state.handle
+
+
+def advance_assignment_version(owner: str):
+    """Invalidate old approvals without clearing an explicit plan requirement."""
+    with task_lock:
+        assignment_versions[owner] = assignment_versions.get(owner, 0) + 1
+        gates = globals().get("plan_gates")
+        request_ids = globals().get("plan_request_ids")
+        team = globals().get("team_lock")
+        if team is not None:
+            team.acquire()
+        try:
+            if (isinstance(gates, dict) and owner in gates
+                    and gates[owner] != "not_required"):
+                gates[owner] = "required"
+            if isinstance(request_ids, dict):
+                request_ids.pop(owner, None)
+        finally:
+            if team is not None:
+                team.release()
 
 
 @dataclass
@@ -92,8 +137,16 @@ def create_task(subject: str, description: str = "",
 
 
 def save_task(task: Task):
-    with task_lock:
-        _task_path(task.id).write_text(json.dumps(asdict(task), indent=2))
+    with task_store_lock():
+        path = _task_path(task.id)
+        temporary = path.with_name(
+            f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp"
+        )
+        try:
+            temporary.write_text(json.dumps(asdict(task), indent=2))
+            os.replace(temporary, path)
+        finally:
+            temporary.unlink(missing_ok=True)
 
 
 def load_task(task_id: str) -> Task:
@@ -151,12 +204,16 @@ def _incomplete_dependencies(task: Task) -> list[str]:
 
 def claim_task(task_id: str, owner: str = "agent") -> str:
     """Atomically claim one task and bind the owner's filesystem cwd."""
-    with task_lock:
+    with task_store_lock():
         task = load_task(task_id)
         if task.status != "pending":
             return f"Task {task_id} is {task.status}, cannot claim"
         if task.owner:
             return f"Task {task_id} is already owned by {task.owner}"
+        assignment = teammate_assignments.get(owner)
+        if assignment:
+            return (f"Owner {owner} must finish the current work turn for "
+                    f"{assignment['task_id']} before claiming another task")
         current = _owner_in_progress(owner)
         if current:
             return (f"Owner {owner} must complete {current.id} before "
@@ -170,24 +227,31 @@ def claim_task(task_id: str, owner: str = "agent") -> str:
         task.status = "in_progress"
         save_task(task)
         teammate_assignments[owner] = {"task_id": task.id, "cwd": cwd}
+        advance_assignment_version(owner)
     print(f"  \033[36m[claim] {task.subject} → in_progress (owner: {owner})\033[0m")
     return f"Claimed {task.id} ({task.subject})"
 
 
 def complete_task(task_id: str, owner: str = "agent") -> str:
     """Complete an assignment only when the caller owns it."""
-    with task_lock:
+    with task_store_lock():
         task = load_task(task_id)
         if task.status != "in_progress":
             return f"Task {task_id} is {task.status}, cannot complete"
         if task.owner != owner:
             return (f"Task {task_id} is owned by {task.owner}, "
                     f"not {owner}; cannot complete")
+        gate = globals().get("plan_gates", {}).get(owner, "not_required")
+        if gate in {"required", "pending", "rejected"}:
+            return f"Task {task_id} cannot complete while plan status is {gate}"
+        assignment = teammate_assignments.get(owner)
+        if not assignment or assignment.get("task_id") != task.id:
+            cwd, error = task_worktree_cwd(task)
+            if error:
+                return f"Task {task_id} cannot complete: {error}"
+            teammate_assignments[owner] = {"task_id": task.id, "cwd": cwd}
         task.status = "completed"
         save_task(task)
-        assignment = teammate_assignments.get(owner)
-        if assignment and assignment.get("task_id") == task_id:
-            teammate_assignments.pop(owner, None)
         unblocked = [t.subject for t in list_tasks()
                      if t.status == "pending" and t.blockedBy and can_start(t.id)]
     print(f"  \033[32m[complete] {task.subject} ✓\033[0m")
@@ -228,8 +292,8 @@ def _worktree_branch(name: str) -> str:
     return f"wt/{name}"
 
 
-def run_git(args: list[str], cwd: Path | None = None) -> tuple[bool, str]:
-    """Run Git without shell interpolation and return (ok, combined output)."""
+def _run_git(args: list[str], cwd: Path | None = None) -> tuple[bool, str]:
+    """Run Git without shell interpolation and preserve machine output."""
     try:
         result = subprocess.run(
             ["git", *args], cwd=cwd or WORKDIR,
@@ -238,11 +302,17 @@ def run_git(args: list[str], cwd: Path | None = None) -> tuple[bool, str]:
     except (OSError, subprocess.TimeoutExpired) as exc:
         return False, f"{type(exc).__name__}: {exc}"
     output = (result.stdout + result.stderr).strip()
-    return result.returncode == 0, output[:5000] or "(no output)"
+    return result.returncode == 0, output or "(no output)"
+
+
+def run_git(args: list[str], cwd: Path | None = None) -> tuple[bool, str]:
+    """Run Git and bound only the text returned to the model."""
+    ok, output = _run_git(args, cwd)
+    return ok, output[:5000]
 
 
 def _registered_worktrees() -> tuple[dict[Path, dict[str, str]], str | None]:
-    ok, output = run_git(["worktree", "list", "--porcelain"])
+    ok, output = _run_git(["worktree", "list", "--porcelain"])
     if not ok:
         return {}, f"cannot read Git worktree registry: {output}"
     entries: dict[Path, dict[str, str]] = {}
@@ -289,12 +359,17 @@ def task_worktree_cwd(task: Task) -> tuple[Path, str | None]:
 def assignment_cwd(owner: str) -> Path:
     with task_lock:
         assignment = teammate_assignments.get(owner)
-        if not assignment:
-            if _owner_in_progress(owner):
-                raise ValueError(f"Missing assignment metadata for {owner}")
+        task = _owner_in_progress(owner)
+        if task and (not assignment or assignment.get("task_id") != task.id):
+            cwd, error = task_worktree_cwd(task)
+            if error:
+                raise ValueError(error)
+            assignment = {"task_id": task.id, "cwd": cwd}
+            teammate_assignments[owner] = assignment
+        elif not assignment:
             return WORKDIR
         task = load_task(str(assignment["task_id"]))
-        if task.status != "in_progress" or task.owner != owner:
+        if task.status not in {"in_progress", "completed"} or task.owner != owner:
             raise ValueError(f"Assignment for {owner} is no longer active")
         cwd, error = task_worktree_cwd(task)
         if error:
@@ -304,6 +379,22 @@ def assignment_cwd(owner: str) -> Path:
         return cwd
 
 
+def release_completed_assignment(owner: str) -> bool:
+    """Release a completed cwd lease only at a model turn boundary."""
+    with task_lock:
+        assignment = teammate_assignments.get(owner)
+        if not assignment:
+            return False
+        task = load_task(str(assignment["task_id"]))
+        if task.status != "completed" or task.owner != owner:
+            return False
+        teammate_assignments.pop(owner, None)
+        advance_assignment_version(owner)
+        if owner in globals().get("plan_gates", {}):
+            globals()["plan_gates"][owner] = "not_required"
+        return True
+
+
 def release_teammate_assignment(owner: str):
     """Return abandoned teammate work to the task board on thread exit."""
     with task_lock:
@@ -315,6 +406,9 @@ def release_teammate_assignment(owner: str):
                 save_task(task)
         finally:
             teammate_assignments.pop(owner, None)
+            advance_assignment_version(owner)
+            if owner in globals().get("plan_gates", {}):
+                globals()["plan_gates"][owner] = "not_required"
 
 
 def create_worktree(name: str, task_id: str) -> str:
@@ -412,6 +506,19 @@ def remove_worktree(name: str, discard_changes: bool = False) -> str:
         if active:
             return (f"Error: Worktree '{name}' is bound to active task "
                     f"{active[0].id}; complete it before removal")
+        leased = [owner for owner, assignment in teammate_assignments.items()
+                  if Path(assignment["cwd"]).resolve() == path.resolve()]
+        if leased:
+            return (f"Error: Worktree '{name}' is still in use by "
+                    f"{', '.join(sorted(leased))}; wait for the turn to end")
+        with globals().get("background_lock", threading.Lock()):
+            running = [task for task in globals().get("background_tasks", {}).values()
+                       if task.get("status") == "running"
+                       and task.get("cwd")
+                       and Path(task["cwd"]).resolve() == path.resolve()]
+        if running:
+            return (f"Error: Worktree '{name}' has a running background command; "
+                    "wait for it to finish")
 
         ok, status = run_git(
             ["status", "--porcelain", "--ignored"], cwd=path
@@ -452,7 +559,7 @@ PROMPT_SECTIONS = {
              "get_task, create_task, list_tasks, claim_task, complete_task, "
              "schedule_cron, list_crons, cancel_cron, "
              "spawn_teammate, send_message, request_shutdown, "
-             "request_plan, review_plan, create_worktree, remove_worktree.",
+             "request_plan, review_plan, create_worktree.",
     "teams": (
         "When parallel work would help, first propose a small team with clear "
         "responsibilities and wait for the user's confirmation. Do not call "
@@ -461,8 +568,8 @@ PROMPT_SECTIONS = {
         "create a task-bound worktree only when a separate working directory "
         "would prevent conflicting edits. A teammate must complete its current "
         "Task before claiming another. A worktree changes tool default cwd "
-        "only; it is not a sandbox. The remove_worktree tool removes only clean "
-        "checkouts and never discards changes. React to team events delivered by the "
+        "only; it is not a sandbox. Worktree removal stays with the host or "
+        "user. React to team events delivered by the "
         "runtime, and shut teammates down when coordination is complete."
     ),
     "workspace": f"Working directory: {WORKDIR}",
@@ -504,18 +611,78 @@ def safe_path(p: str, cwd: Path | None = None) -> Path:
     return path
 
 
+_shell_processes: set[subprocess.Popen] = set()
+_shell_process_lock = threading.RLock()
+
+
+def _stop_process_group(process: subprocess.Popen):
+    """Stop processes that remain in the command's original process group."""
+    for sig in (signal.SIGTERM, signal.SIGKILL):
+        try:
+            os.killpg(process.pid, sig)
+        except ProcessLookupError:
+            return
+        except OSError:
+            return
+        time.sleep(0.05)
+
+
+def _stop_all_shell_processes():
+    with _shell_process_lock:
+        processes = list(_shell_processes)
+    for process in processes:
+        _stop_process_group(process)
+
+
+def _handle_termination_signal(signum, _frame):
+    _stop_all_shell_processes()
+    raise SystemExit(128 + signum)
+
+
+atexit.register(_stop_all_shell_processes)
+signal.signal(signal.SIGTERM, _handle_termination_signal)
+
+
+def _run_bash_process(command: str, cwd: Path | None = None) -> tuple[str, int | None]:
+    process = None
+    try:
+        process = subprocess.Popen(
+            command, shell=True, cwd=cwd or WORKDIR,
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+            text=True, start_new_session=True,
+        )
+        with _shell_process_lock:
+            _shell_processes.add(process)
+        stdout, stderr = process.communicate(timeout=120)
+        out = (stdout + stderr).strip()
+        return (out[:50000] if out else "(no output)"), process.returncode
+    except subprocess.TimeoutExpired:
+        return "Error: Timeout (120s)", None
+    except OSError as exc:
+        return f"Error: {type(exc).__name__}: {exc}", None
+    finally:
+        if process is not None:
+            _stop_process_group(process)
+            try:
+                process.wait(timeout=0.2)
+            except subprocess.TimeoutExpired:
+                pass
+            with _shell_process_lock:
+                _shell_processes.discard(process)
+
+
+def _format_bash_result(output: str, exit_code: int | None) -> str:
+    if exit_code == 0:
+        return output
+    if exit_code is None:
+        return output
+    return f"Error: command exited with status {exit_code}\n{output}"
+
+
 def run_bash(command: str, run_in_background: bool = False,
              cwd: Path | None = None) -> str:
     # run_in_background is handled by agent_loop dispatch, not here
-    try:
-        r = subprocess.run(command, shell=True, cwd=cwd or WORKDIR,
-                           capture_output=True, text=True, timeout=120)
-        out = (r.stdout + r.stderr).strip()
-        return out[:50000] if out else "(no output)"
-    except subprocess.TimeoutExpired:
-        return "Error: Timeout (120s)"
-    except OSError as exc:
-        return f"Error: {type(exc).__name__}: {exc}"
+    return _format_bash_result(*_run_bash_process(command, cwd))
 
 
 def run_read(path: str, limit: int | None = None,
@@ -539,6 +706,28 @@ def run_write(path: str, content: str, cwd: Path | None = None) -> str:
         return f"Error: {e}"
 
 
+def _agent_cwd() -> tuple[Path | None, str | None]:
+    try:
+        return assignment_cwd("agent"), None
+    except (FileNotFoundError, ValueError) as exc:
+        return None, f"Error: Invalid task assignment: {exc}"
+
+
+def run_agent_bash(command: str, run_in_background: bool = False) -> str:
+    cwd, error = _agent_cwd()
+    return error or run_bash(command, run_in_background, cwd)
+
+
+def run_agent_read(path: str, limit: int | None = None) -> str:
+    cwd, error = _agent_cwd()
+    return error or run_read(path, limit, cwd)
+
+
+def run_agent_write(path: str, content: str) -> str:
+    cwd, error = _agent_cwd()
+    return error or run_write(path, content, cwd)
+
+
 # Task tools
 
 def run_create_task(subject: str, description: str = "",
@@ -613,15 +802,18 @@ def is_slow_operation(tool_name: str, tool_input: dict) -> bool:
 
 def should_run_background(tool_name: str, tool_input: dict) -> bool:
     """Model explicit request takes priority; fallback to heuristic."""
-    if tool_input.get("run_in_background"):
-        return True
-    return is_slow_operation(tool_name, tool_input)
+    return tool_name == "bash" and (
+        tool_input.get("run_in_background") is True
+        or is_slow_operation(tool_name, tool_input)
+    )
 
 
 def execute_tool(block) -> str:
     """Execute a tool call block, return output."""
     handler = {
-        "bash": run_bash, "read_file": run_read, "write_file": run_write,
+        "bash": run_agent_bash,
+        "read_file": run_agent_read,
+        "write_file": run_agent_write,
         "create_task": run_create_task, "list_tasks": run_list_tasks,
         "get_task": run_get_task, "claim_task": run_claim_task,
         "complete_task": run_complete_task,
@@ -633,24 +825,37 @@ def execute_tool(block) -> str:
         "request_plan": run_request_plan,
         "review_plan": run_review_plan,
         "create_worktree": run_create_worktree,
-        "remove_worktree": run_remove_worktree,
     }.get(block.name)
-    if handler:
-        return handler(**block.input)
-    return f"Unknown tool: {block.name}"
+    if not handler:
+        return f"Unknown tool: {block.name}"
+    try:
+        return str(handler(**block.input))
+    except (TypeError, ValueError) as exc:
+        return f"Error: {exc}"
 
 
 def start_background_task(block) -> str:
-    """Run tool in a daemon thread. Returns background task ID."""
+    """Run one bash call in a daemon thread with a fixed dispatch cwd."""
     global _bg_counter
     _bg_counter += 1
     bg_id = f"bg_{_bg_counter:04d}"
     cmd = block.input.get("command", block.name)
+    cwd, cwd_error = _agent_cwd()
 
     def worker():
-        result = execute_tool(block)
+        try:
+            if block.name != "bash":
+                raise ValueError("only bash can run in the background")
+            if cwd_error:
+                raise ValueError(cwd_error.removeprefix("Error: "))
+            output, exit_code = _run_bash_process(str(block.input["command"]), cwd)
+            result = _format_bash_result(output, exit_code)
+            status = "completed" if exit_code == 0 else "failed"
+        except Exception as exc:
+            result = f"Error: {type(exc).__name__}: {exc}"
+            status = "failed"
         with background_lock:
-            background_tasks[bg_id]["status"] = "completed"
+            background_tasks[bg_id]["status"] = status
             background_results[bg_id] = result
 
     with background_lock:
@@ -658,6 +863,7 @@ def start_background_task(block) -> str:
             "tool_use_id": block.id,
             "command": cmd,
             "status": "running",
+            "cwd": str(cwd) if cwd else None,
         }
     threading.Thread(target=worker, daemon=True).start()
     print(f"  \033[33m[background] dispatched {bg_id}: {cmd[:40]}\033[0m")
@@ -665,10 +871,10 @@ def start_background_task(block) -> str:
 
 
 def collect_background_results() -> list[str]:
-    """Collect completed background results as task_notification messages."""
+    """Collect terminal background results as task_notification messages."""
     with background_lock:
         ready_ids = [bid for bid, task in background_tasks.items()
-                     if task["status"] == "completed"]
+                     if task["status"] in {"completed", "failed"}]
     notifications = []
     for bg_id in ready_ids:
         with background_lock:
@@ -678,7 +884,7 @@ def collect_background_results() -> list[str]:
         notifications.append(
             f"<task_notification>\n"
             f"  <task_id>{bg_id}</task_id>\n"
-            f"  <status>completed</status>\n"
+            f"  <status>{task['status']}</status>\n"
             f"  <command>{task['command']}</command>\n"
             f"  <summary>{summary}</summary>\n"
             f"</task_notification>")
@@ -688,10 +894,11 @@ def collect_background_results() -> list[str]:
 
 
 def has_pending_background() -> bool:
-    """Non-destructive: True if any background task has completed and is
+    """Non-destructive: True if any background task is terminal and is
     waiting to be collected. The inbox poller uses this in its wake condition."""
     with background_lock:
-        return any(t["status"] == "completed" for t in background_tasks.values())
+        return any(t["status"] in {"completed", "failed"}
+                   for t in background_tasks.values())
 
 
 # ── Cron Scheduler (from s14, synced) ──
@@ -706,11 +913,12 @@ class CronJob:
     prompt: str      # message to inject when fired
     recurring: bool  # True = recurring, False = one-shot
     durable: bool    # True = persist to disk
+    pending_delivery: bool = False
 
 
 scheduled_jobs: dict[str, CronJob] = {}
 cron_queue: list[CronJob] = []
-cron_lock = threading.Lock()
+cron_lock = threading.RLock()
 _last_fired: dict[str, str] = {}  # job_id → "YYYY-MM-DD HH:MM"
 
 
@@ -811,8 +1019,11 @@ def validate_cron(cron_expr: str) -> str | None:
 
 def save_durable_jobs():
     """Persist durable jobs to .scheduled_tasks.json."""
-    durable = [asdict(j) for j in scheduled_jobs.values() if j.durable]
-    DURABLE_PATH.write_text(json.dumps(durable, indent=2))
+    with cron_lock:
+        durable = [asdict(j) for j in scheduled_jobs.values() if j.durable]
+        temporary = DURABLE_PATH.with_suffix(".json.tmp")
+        temporary.write_text(json.dumps(durable, indent=2))
+        os.replace(temporary, DURABLE_PATH)
 
 
 def load_durable_jobs():
@@ -828,6 +1039,8 @@ def load_durable_jobs():
                 print(f"  \033[31m[cron] skipping invalid job {job.id}: {err}\033[0m")
                 continue
             scheduled_jobs[job.id] = job
+            if job.pending_delivery:
+                cron_queue.append(job)
         valid = [j for j in jobs if j["id"] in scheduled_jobs]
         if valid:
             print(f"  \033[35m[cron] loaded {len(valid)} durable job(s)\033[0m")
@@ -848,8 +1061,8 @@ def schedule_job(cron: str, prompt: str, recurring: bool = True,
     )
     with cron_lock:
         scheduled_jobs[job.id] = job
-    if durable:
-        save_durable_jobs()
+        if durable:
+            save_durable_jobs()
     print(f"  \033[35m[cron register] {job.id} '{cron}' → {prompt[:40]}\033[0m")
     return job
 
@@ -858,14 +1071,28 @@ def cancel_job(job_id: str) -> str:
     """Cancel a cron job."""
     with cron_lock:
         job = scheduled_jobs.pop(job_id, None)
+        cron_queue[:] = [queued for queued in cron_queue if queued.id != job_id]
+        if job and job.durable:
+            save_durable_jobs()
     if not job:
         return f"Job {job_id} not found"
-    if job.durable:
-        save_durable_jobs()
     print(f"  \033[31m[cron cancel] {job_id}\033[0m")
     return f"Cancelled {job_id}"
 
 
+def _enqueue_due_job(job: CronJob):
+    """Persist a one-shot delivery before exposing it through the queue."""
+    if not job.recurring:
+        job.pending_delivery = True
+        try:
+            if job.durable:
+                save_durable_jobs()
+        except Exception:
+            job.pending_delivery = False
+            raise
+    cron_queue.append(job)
+
+
 def cron_scheduler_loop():
     """Independent daemon thread: poll every 1s, fire matching jobs.
     Individual job errors are caught to prevent one bad job from
@@ -878,16 +1105,14 @@ def cron_scheduler_loop():
         with cron_lock:
             for job in list(scheduled_jobs.values()):
                 try:
+                    if job.pending_delivery:
+                        continue
                     if cron_matches(job.cron, now):
                         if _last_fired.get(job.id) != minute_marker:
-                            cron_queue.append(job)
+                            _enqueue_due_job(job)
                             _last_fired[job.id] = minute_marker
                             print(f"  \033[35m[cron fire] {job.id} → "
                                   f"{job.prompt[:40]}\033[0m")
-                        if not job.recurring:
-                            scheduled_jobs.pop(job.id, None)
-                            if job.durable:
-                                save_durable_jobs()
                 except Exception as e:
                     print(f"  \033[31m[cron error] {job.id}: {e}\033[0m")
 
@@ -900,6 +1125,35 @@ def consume_cron_queue() -> list[CronJob]:
     return fired
 
 
+def has_cron_queue() -> bool:
+    with cron_lock:
+        return bool(cron_queue)
+
+
+def acknowledge_cron_jobs(jobs: list[CronJob]):
+    """Remove one-shot jobs after a model call accepts their prompts."""
+    durable_changed = False
+    with cron_lock:
+        for job in jobs:
+            current = scheduled_jobs.get(job.id)
+            if current and not current.recurring and current.pending_delivery:
+                scheduled_jobs.pop(job.id, None)
+                durable_changed = durable_changed or current.durable
+        if durable_changed:
+            save_durable_jobs()
+
+
+def restore_cron_jobs(jobs: list[CronJob]):
+    """Put unacknowledged deliveries back after a failed model call."""
+    with cron_lock:
+        queued_ids = {job.id for job in cron_queue}
+        for job in jobs:
+            current = scheduled_jobs.get(job.id)
+            if current and current.id not in queued_ids:
+                cron_queue.append(current)
+                queued_ids.add(current.id)
+
+
 # Load durable jobs on startup, then start scheduler thread
 load_durable_jobs()
 threading.Thread(target=cron_scheduler_loop, daemon=True).start()
@@ -1023,6 +1277,8 @@ class ProtocolState:
     target: str
     status: str
     payload: str
+    work_version: int | None = None
+    task_id: str | None = None
     created_at: float = field(default_factory=time.time)
 
 
@@ -1098,22 +1354,35 @@ def _last_assistant_text(content) -> str:
     return ""
 
 
+def current_work_identity(owner: str) -> tuple[int, str | None]:
+    with task_lock:
+        assignment = teammate_assignments.get(owner)
+        task_id = str(assignment["task_id"]) if assignment else None
+        return assignment_versions.get(owner, 0), task_id
+
+
 def _teammate_submit_plan(from_name: str, plan: str) -> str:
-    with team_lock:
-        if plan_gates.get(from_name) == "pending":
-            return "A plan is already waiting for review."
-        request_id = new_request_id()
-        pending_requests[request_id] = ProtocolState(
-            request_id=request_id,
-            type="plan_approval",
-            sender=from_name,
-            target="lead",
-            status="pending",
-            payload=plan,
-        )
-        plan_gates[from_name] = "pending"
-        plan_request_ids[from_name] = request_id
-        active_teammates[from_name] = "waiting_approval"
+    with task_lock:
+        assignment = teammate_assignments.get(from_name)
+        task_id = str(assignment["task_id"]) if assignment else None
+        work_version = assignment_versions.get(from_name, 0)
+        with team_lock:
+            if plan_gates.get(from_name) == "pending":
+                return "A plan is already waiting for review."
+            request_id = new_request_id()
+            pending_requests[request_id] = ProtocolState(
+                request_id=request_id,
+                type="plan_approval",
+                sender=from_name,
+                target="lead",
+                status="pending",
+                payload=plan,
+                work_version=work_version,
+                task_id=task_id,
+            )
+            plan_gates[from_name] = "pending"
+            plan_request_ids[from_name] = request_id
+            active_teammates[from_name] = "waiting_approval"
     BUS.send(from_name, "lead", plan, "plan_approval_request",
              {"request_id": request_id})
     return f"Plan submitted ({request_id}). Wait for Lead's decision."
@@ -1133,6 +1402,7 @@ def apply_plan_response(name: str, msg: dict) -> tuple[bool, str]:
     """Apply only the Lead response for this teammate's current plan."""
     metadata = msg.get("metadata", {})
     request_id = metadata.get("request_id", "")
+    work_version, task_id = current_work_identity(name)
     with team_lock:
         state = pending_requests.get(request_id)
         expected_id = plan_request_ids.get(name)
@@ -1144,6 +1414,8 @@ def apply_plan_response(name: str, msg: dict) -> tuple[bool, str]:
             and state.type == "plan_approval"
             and state.sender == name
             and state.target == "lead"
+            and state.work_version == work_version
+            and state.task_id == task_id
             and state.status in {"approved", "rejected"}
             and metadata.get("approve", False)
             == (state.status == "approved")
@@ -1208,7 +1480,7 @@ def scan_unclaimed_tasks() -> list[Task]:
 def claim_next_task(name: str) -> Task | None:
     """Claim the first still-available task, never a second assignment."""
     with task_lock:
-        if _owner_in_progress(name):
+        if teammate_assignments.get(name) or _owner_in_progress(name):
             return None
     for task in scan_unclaimed_tasks():
         result = claim_task(task.id, owner=name)
@@ -1219,7 +1491,8 @@ def claim_next_task(name: str) -> Task | None:
 
 # ── Teammate Thread ──
 
-def spawn_teammate_thread(name: str, role: str, prompt: str) -> str:
+def spawn_teammate_thread(name: str, role: str, prompt: str,
+                          require_plan: bool = False) -> str:
     """Spawn a persistent teammate that alternates between WORK and IDLE."""
     if not is_valid_agent_name(name):
         return ("Invalid teammate name: use 1-64 letters, digits, "
@@ -1231,7 +1504,8 @@ def spawn_teammate_thread(name: str, role: str, prompt: str) -> str:
                for existing in active_teammates):
             return f"Teammate '{name}' already exists"
         active_teammates[name] = "working"
-        plan_gates[name] = "not_required"
+        plan_gates[name] = "required" if require_plan else "not_required"
+        assignment_versions[name] = 1
 
     system = (f"You are '{name}', a {role}. "
               "Use tools to complete assigned work. You can list, claim, and "
@@ -1278,7 +1552,11 @@ def spawn_teammate_thread(name: str, role: str, prompt: str) -> str:
             except FileNotFoundError:
                 return f"Error: Task {task_id} not found"
 
-        messages = [{"role": "user", "content": prompt}]
+        initial_prompt = prompt
+        if require_plan:
+            initial_prompt += ("\n\n[Plan required] Submit a plan and wait for "
+                               "Lead approval before bash or write_file.")
+        messages = [{"role": "user", "content": initial_prompt}]
         sub_tools = [
             {"name": "bash", "description": "Run a shell command.",
              "input_schema": {"type": "object",
@@ -1367,6 +1645,8 @@ def spawn_teammate_thread(name: str, role: str, prompt: str) -> str:
 
         should_stop = False
         while not should_stop:
+            if handle_messages(BUS.read_inbox(name)):
+                break
             with team_lock:
                 active_teammates[name] = "working"
             try:
@@ -1398,6 +1678,7 @@ def spawn_teammate_thread(name: str, role: str, prompt: str) -> str:
                 with team_lock:
                     active_teammates[name] = "waiting_approval"
             else:
+                release_completed_assignment(name)
                 with team_lock:
                     active_teammates[name] = "idle"
                 BUS.send(name, "lead", "Waiting for more work.",
@@ -1462,13 +1743,15 @@ def spawn_teammate_thread(name: str, role: str, prompt: str) -> str:
 
 # ── Lead Team Tools ──
 
-def run_spawn_teammate(name: str, role: str, prompt: str) -> str:
-    return spawn_teammate_thread(name, role, prompt)
+def run_spawn_teammate(name: str, role: str, prompt: str,
+                       require_plan: bool = False) -> str:
+    return spawn_teammate_thread(name, role, prompt, require_plan)
 
 
 def run_send_message(to: str, content: str) -> str:
     if to not in active_teammates:
         return f"Teammate '{to}' is not active"
+    advance_assignment_version(to)
     BUS.send("lead", to, content)
     return f"Sent to {to}"
 
@@ -1502,6 +1785,10 @@ def run_request_plan(teammate: str, task: str) -> str:
 
 def run_review_plan(request_id: str, approve: bool,
                     feedback: str = "") -> str:
+    state = pending_requests.get(request_id)
+    if not state:
+        return f"Request {request_id} not found"
+    work_version, task_id = current_work_identity(state.sender)
     with team_lock:
         state = pending_requests.get(request_id)
         if not state:
@@ -1510,6 +1797,8 @@ def run_review_plan(request_id: str, approve: bool,
             return f"Request {request_id} is not a plan"
         if state.status != "pending":
             return f"Request {request_id} already {state.status}"
+        if (state.work_version != work_version or state.task_id != task_id):
+            return f"Request {request_id} belongs to an earlier assignment"
         if plan_request_ids.get(state.sender) != request_id:
             return f"Request {request_id} is not the current plan"
         state.status = "approved" if approve else "rejected"
@@ -1524,11 +1813,6 @@ def run_create_worktree(name: str, task_id: str) -> str:
     return create_worktree(name, task_id)
 
 
-def run_remove_worktree(name: str) -> str:
-    """Model-facing cleanup never opts into destructive removal."""
-    return remove_worktree(name)
-
-
 # ── Tool Definitions ──
 
 TOOLS = [
@@ -1607,7 +1891,8 @@ TOOLS = [
                               "pattern": "^[A-Za-z0-9_-]{1,64}$",
                           },
                           "role": {"type": "string"},
-                          "prompt": {"type": "string"}},
+                          "prompt": {"type": "string"},
+                          "require_plan": {"type": "boolean"}},
                       "required": ["name", "role", "prompt"]}},
     {"name": "send_message",
      "description": "Send a message to a teammate via MessageBus.",
@@ -1646,18 +1931,6 @@ TOOLS = [
                                      "task_id": {"type": "string"}},
                       "required": ["name", "task_id"],
                       "additionalProperties": False}},
-    {"name": "remove_worktree",
-     "description": "Remove a clean task worktree while retaining its branch.",
-     "input_schema": {"type": "object",
-                      "properties": {
-                          "name": {
-                              "type": "string",
-                              "pattern": ("^(?!.*\\.\\.)[A-Za-z0-9]"
-                                          "[A-Za-z0-9._-]{0,63}$"),
-                              "maxLength": 64,
-                          }},
-                      "required": ["name"],
-                      "additionalProperties": False}},
 ]
 
 
@@ -1690,19 +1963,23 @@ def agent_loop(messages: list, context: dict):
             messages.append({"role": "user",
                              "content": f"[Scheduled] {job.prompt}"})
             print(f"  \033[35m[inject cron] {job.prompt[:50]}\033[0m")
-
         try:
             response = client.messages.create(
                 model=MODEL, system=system, messages=messages,
                 tools=TOOLS, max_tokens=8000)
         except Exception as e:
+            restore_cron_jobs(fired)
             messages.append({"role": "assistant", "content": [
                 {"type": "text",
                  "text": f"[Error] {type(e).__name__}: {e}"}]})
+            release_completed_assignment("agent")
             return
 
+        acknowledge_cron_jobs(fired)
+
         messages.append({"role": "assistant", "content": response.content})
         if response.stop_reason != "tool_use":
+            release_completed_assignment("agent")
             return
 
         results = []
@@ -1761,7 +2038,8 @@ if __name__ == "__main__":
         # so the final message can outlive its registry entry.
         while True:
             time.sleep(1)
-            if BUS.peek("lead") or has_pending_background():
+            if (BUS.peek("lead") or has_pending_background()
+                    or has_cron_queue()):
                 events.put(("wake", None))
 
     threading.Thread(target=input_reader, daemon=True).start()
@@ -1778,17 +2056,18 @@ if __name__ == "__main__":
             history.append({"role": "user", "content": payload})
         else:  # "wake": teammate inbox or background results are ready
             parts = []
+            cron_ready = has_cron_queue()
             inbox = consume_lead_inbox()
             if inbox:
                 parts.append(format_team_events(inbox))
             bg = collect_background_results()
             parts.extend(bg)
-            if not parts:
+            if not parts and not cron_ready:
                 continue  # already drained by an earlier wake (idempotent)
             history.append({"role": "user", "content": "\n".join(parts)})
             print(f"\n\033[33m[wake: {len(inbox)} team events + "
                   f"{len(bg)} background "
-                  f"-> new turn]\033[0m")
+                  f"{1 if cron_ready else 0} cron -> new turn]\033[0m")
 
         # One turn for whichever source woke us.
         agent_loop(history, context)