Session lifecycle and robustness

src/autopilot_loop/cli.py

@@ -38,6 +38,24 @@ def _generate_task_id():
     return uuid.uuid4().hex[:8]
 
 
+def _detect_autopilot_branch():
+    """If the current git branch matches autopilot/*, return the branch name.
+
+    Returns None if not on an autopilot branch or git is unavailable.
+    """
+    try:
+        result = subprocess.run(
+            ["git", "branch", "--show-current"],
+            capture_output=True, text=True, check=True,
+        )
+        branch = result.stdout.strip()
+        if branch.startswith("autopilot/"):
+            return branch
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        pass
+    return None
+
+
 def cmd_start(args):
     """Start a new autopilot task."""
     config = load_config({
@@ -57,7 +75,14 @@ def cmd_start(args):
         sys.exit(1)
 
     task_id = _generate_task_id()
-    branch = config["branch_pattern"].format(task_id=task_id)
+
+    # Detect if we're already on an autopilot branch
+    existing_branch = _detect_autopilot_branch()
+    if existing_branch:
+        branch = existing_branch
+        logger.info("Detected existing autopilot branch: %s", branch)
+    else:
+        branch = config["branch_pattern"].format(task_id=task_id)
 
     create_task(
         task_id=task_id,
@@ -71,6 +96,9 @@ def cmd_start(args):
     from autopilot_loop.persistence import update_task
     update_task(task_id, branch=branch)
 
+    if existing_branch:
+        update_task(task_id, existing_branch=1)
+
     if args.dry_run:
         print("DRY RUN — would start task %s" % task_id)
         print("  Branch: %s" % branch)
@@ -392,12 +420,74 @@ def cmd_stop(args):
     except (subprocess.CalledProcessError, FileNotFoundError):
         print("No tmux session found for task %s" % task_id)
 
-    # Update task state
+    # Update task state — save the current state before marking STOPPED
     task = get_task(task_id)
-    if task and task["state"] not in ("COMPLETE", "FAILED"):
+    if task and task["state"] not in ("COMPLETE", "FAILED", "STOPPED"):
         from autopilot_loop.persistence import update_task
-        update_task(task_id, state="FAILED")
-        print("✓ Task %s marked as FAILED" % task_id)
+        update_task(task_id, pre_stop_state=task["state"], state="STOPPED")
+        print("✓ Task %s marked as STOPPED" % task_id)
+
+
+def cmd_restart(args):
+    """Restart a stopped task from its current phase."""
+    task_id = args.task_id
+    task = get_task(task_id)
+
+    if not task:
+        print("Error: task %s not found" % task_id, file=sys.stderr)
+        sys.exit(1)
+
+    if task["state"] != "STOPPED":
+        print("Error: task %s is in state %s, only STOPPED tasks can be restarted" % (task_id, task["state"]),
+              file=sys.stderr)
+        sys.exit(1)
+
+    config = load_config({"model": task["model"]})
+
+    # Determine the phase to restart from. For states that are mid-action
+    # (e.g. FIX, IMPLEMENT), restart from the beginning of that phase.
+    # For waiting states, restart from the state that triggered the wait.
+    restart_state = task["state"]
+    stopped_state = task.get("pre_stop_state") or "INIT"
+
+    # Map waiting/verification states back to their action states
+    _RESTART_STATE_MAP = {
+        "VERIFY_PUSH": "FIX" if task.get("task_mode") != "ci" else "FIX_CI",
+        "WAIT_REVIEW": "REQUEST_REVIEW",
+        "WAIT_CI": "FIX_CI",
+        "VERIFY_PR": "IMPLEMENT",
+    }
+    restart_state = _RESTART_STATE_MAP.get(stopped_state, stopped_state)
+
+    from autopilot_loop.persistence import update_task
+    update_task(task_id, state=restart_state)
+
+    # Launch in tmux
+    sessions_dir = get_sessions_dir(task_id)
+    log_file = os.path.join(sessions_dir, "orchestrator.log")
+    tmux_session = "autopilot-%s" % task_id
+    run_cmd = "autopilot _run --task-id %s 2>&1 | tee -a %s" % (task_id, log_file)
+
+    try:
+        subprocess.run(
+            ["tmux", "new-session", "-d", "-s", tmux_session, run_cmd],
+            check=True,
+        )
+    except FileNotFoundError:
+        logger.warning("tmux not found, running in foreground")
+        cmd_run(argparse.Namespace(task_id=task_id))
+        return
+    except subprocess.CalledProcessError as e:
+        print("Error: failed to create tmux session: %s" % e, file=sys.stderr)
+        sys.exit(1)
+
+    print("✓ Restarting task %s from state %s" % (task_id, restart_state))
+    print("✓ Running in tmux session: %s" % tmux_session)
+    print()
+    print("  To check progress:  autopilot status")
+    print("  To view logs:       autopilot logs --session %s" % task_id)
+    print("  To attach to tmux:  tmux attach -t %s" % tmux_session)
+    print("  To stop:            autopilot stop %s" % task_id)
 
 
 def main():
@@ -433,6 +523,10 @@ def main():
     p_stop = subparsers.add_parser("stop", help="Stop a running task")
     p_stop.add_argument("task_id", type=str, help="Task ID to stop")
 
+    # restart
+    p_restart = subparsers.add_parser("restart", help="Restart a stopped task")
+    p_restart.add_argument("task_id", type=str, help="Task ID to restart")
+
     # fix-ci
     p_fixci = subparsers.add_parser("fix-ci", help="Fix CI failures on an existing PR")
     p_fixci.add_argument("--pr", type=int, required=True, help="PR number")
@@ -457,6 +551,8 @@ def main():
         cmd_logs(args)
     elif args.command == "stop":
         cmd_stop(args)
+    elif args.command == "restart":
+        cmd_restart(args)
     elif args.command == "fix-ci":
         cmd_fix_ci(args)
     elif args.command == "_run":

src/autopilot_loop/codespace.py

@@ -1,24 +1,52 @@
 """Codespace idle timeout management via gh api."""
 
+import json
 import logging
 import os
 import subprocess
 
 logger = logging.getLogger(__name__)
 
 
+def get_idle_timeout():
+    """Get the current idle timeout for the codespace, or None if not in a codespace."""
+    codespace_name = os.environ.get("CODESPACE_NAME")
+    if not codespace_name:
+        return None
+
+    try:
+        result = subprocess.run(
+            [
+                "gh", "api",
+                "/user/codespaces/%s" % codespace_name,
+                "--jq", ".idle_timeout_minutes",
+            ],
+            capture_output=True, text=True, check=True,
+        )
+        return int(result.stdout.strip())
+    except (FileNotFoundError, subprocess.CalledProcessError, ValueError):
+        return None
+
+
 def set_idle_timeout(minutes=120):
     """Set the idle timeout for the current codespace.
 
     Uses: gh api -X PATCH "/user/codespaces/$CODESPACE_NAME" -F idle_timeout_minutes=<value>
 
+    Skips the update if the current timeout is already >= the desired value.
     Non-fatal — logs a warning if it fails (e.g., not in a codespace).
     """
     codespace_name = os.environ.get("CODESPACE_NAME")
     if not codespace_name:
         logger.debug("Not in a codespace (CODESPACE_NAME not set), skipping idle timeout")
         return
 
+    # Check current timeout before updating
+    current = get_idle_timeout()
+    if current is not None and current >= minutes:
+        logger.info("Codespace idle timeout already set to %dm (>= %dm), skipping", current, minutes)
+        return
+
     try:
         subprocess.run(
             [

src/autopilot_loop/orchestrator.py

@@ -41,6 +41,7 @@
     fix_prompt,
     format_ci_annotations_for_prompt,
     format_review_for_prompt,
+    implement_on_existing_branch_prompt,
     implement_prompt,
     plan_and_implement_prompt,
 )
@@ -63,8 +64,12 @@
     "RESOLVE_COMMENTS",
     "COMPLETE",
     "FAILED",
+    "STOPPED",
 ]
 
+# Terminal states — the orchestrator stops when reaching any of these
+TERMINAL_STATES = frozenset({"COMPLETE", "FAILED", "STOPPED"})
+
 
 class BaseOrchestrator:
     """Shared infrastructure for state machine orchestrators."""
@@ -81,11 +86,11 @@ def _get_handlers(self):
         raise NotImplementedError
 
     def run(self):
-        """Run the state machine until COMPLETE or FAILED."""
+        """Run the state machine until a terminal state (COMPLETE, FAILED, or STOPPED)."""
         state = self.task["state"]
         logger.info("[%s] Starting orchestrator from state: %s", self.task_id, state)
 
-        while state not in ("COMPLETE", "FAILED"):
+        while state not in TERMINAL_STATES:
             logger.info("[%s] %s", self.task_id, state)
             try:
                 state = self._transition(state)
@@ -253,11 +258,20 @@ def _retry_fix_state(self):
     def _do_implement(self):
         """Run copilot agent with implement prompt."""
         branch = self.task["branch"]
-        prompt = implement_prompt(
-            task_description=self.task["prompt"],
-            branch_name=branch,
-            custom_instructions=self.config.get("custom_instructions", ""),
-        )
+
+        # Use existing-branch prompt if the branch already exists remotely
+        if self.task.get("existing_branch"):
+            prompt = implement_on_existing_branch_prompt(
+                task_description=self.task["prompt"],
+                branch_name=branch,
+                custom_instructions=self.config.get("custom_instructions", ""),
+            )
+        else:
+            prompt = implement_prompt(
+                task_description=self.task["prompt"],
+                branch_name=branch,
+                custom_instructions=self.config.get("custom_instructions", ""),
+            )
 
         result = self._run_agent_with_retry("IMPLEMENT", prompt, "implement")
         if result is None:

src/autopilot_loop/persistence.py

@@ -28,7 +28,7 @@
 
 # Bump this when the schema changes. Additive changes (new nullable columns)
 # are handled by _migrate(). Breaking changes trigger a DB recreate.
-SCHEMA_VERSION = 3
+SCHEMA_VERSION = 4
 
 SCHEMA = """
 CREATE TABLE IF NOT EXISTS schema_meta (
@@ -50,6 +50,8 @@
     last_review_id INTEGER,
     task_mode TEXT NOT NULL DEFAULT 'review',
     ci_check_names TEXT,
+    pre_stop_state TEXT,
+    existing_branch INTEGER NOT NULL DEFAULT 0,
     created_at REAL NOT NULL,
     updated_at REAL NOT NULL
 );
@@ -81,6 +83,8 @@
     (2, "tasks", "last_review_id", "INTEGER"),
     (3, "tasks", "task_mode", "TEXT NOT NULL DEFAULT 'review'"),
     (3, "tasks", "ci_check_names", "TEXT"),
+    (4, "tasks", "pre_stop_state", "TEXT"),
+    (4, "tasks", "existing_branch", "INTEGER NOT NULL DEFAULT 0"),
 ]
 
 
@@ -183,7 +187,7 @@ def get_task(task_id):
 _TASK_COLUMNS = frozenset({
     "prompt", "state", "pr_number", "branch", "iteration",
     "max_iterations", "plan_mode", "dry_run", "model", "last_review_id",
-    "task_mode", "ci_check_names", "updated_at",
+    "task_mode", "ci_check_names", "pre_stop_state", "existing_branch", "updated_at",
 })
 
 

src/autopilot_loop/prompts.py

@@ -6,6 +6,7 @@
 
 __all__ = [
     "implement_prompt",
+    "implement_on_existing_branch_prompt",
     "plan_and_implement_prompt",
     "fix_prompt",
     "fix_ci_prompt",
@@ -50,6 +51,44 @@ def implement_prompt(task_description, branch_name, custom_instructions=""):
     return "\n".join(parts)
 
 
+def implement_on_existing_branch_prompt(task_description, branch_name, custom_instructions=""):
+    """Prompt for the IMPLEMENT phase on an existing branch.
+
+    Used when the user starts a new task on a branch that already exists
+    (e.g. autopilot/<task-id>). The agent should NOT create a new branch
+    and should commit directly on the current branch.
+    """
+    parts = []
+
+    if custom_instructions:
+        parts.append(custom_instructions.strip())
+        parts.append("")
+
+    parts.append("## Task")
+    parts.append(task_description.strip())
+    parts.append("")
+    parts.append("## Instructions")
+    parts.append(
+        "IMPORTANT: You are on the existing branch `%s`. Do NOT create a new branch.\n"
+        "Work directly on this branch.\n\n"
+        "1. Implement the task described above.\n"
+        "2. Run any relevant tests to verify your changes work.\n"
+        "3. Stage and commit your changes with a proper, descriptive commit message\n"
+        "   that explains what was changed and why. Do NOT use generic messages.\n"
+        "4. If a PR already exists for this branch, push your changes. If no PR exists,\n"
+        "   create a draft pull request using `gh pr create --draft`. Use the repo's\n"
+        "   PR template (check `.github/PULL_REQUEST_TEMPLATE.md` or similar) to\n"
+        "   structure the PR body. Write a clear title and fill in the template sections.\n"
+        "5. Push the branch: `git push origin %s`\n"
+        "6. After pushing, review your own changes by running `git diff HEAD~1` and\n"
+        "   examining the output. If you find any issues (bugs, missing tests, style\n"
+        "   problems), fix them, commit with a descriptive message, and push again.\n"
+        % (branch_name, branch_name)
+    )
+
+    return "\n".join(parts)
+
+
 def plan_and_implement_prompt(task_description, branch_name, custom_instructions=""):
     """Prompt for the PLAN_AND_IMPLEMENT phase.