stravinsky 0.1.2__py3-none-any.whl → 0.2.38__py3-none-any.whl

mcp_bridge/tools/agent_manager.py

@@ -24,9 +24,10 @@ logger = logging.getLogger(__name__)
 @dataclass
 class AgentTask:
     """Represents a background agent task with full tool access."""
+
     id: str
     prompt: str
-    agent_type: str  # explore, librarian, frontend, etc.
+    agent_type: str  # explore, dewey, frontend, delphi, etc.
     description: str
     status: str  # pending, running, completed, failed, cancelled
     created_at: str
@@ -36,12 +37,14 @@ class AgentTask:
     result: Optional[str] = None
     error: Optional[str] = None
     pid: Optional[int] = None
+    timeout: int = 300  # Default 5 minutes
     progress: Optional[Dict[str, Any]] = None  # tool calls, last update
 
 
 @dataclass
 class AgentProgress:
     """Progress tracking for a running agent."""
+
     tool_calls: int = 0
     last_tool: Optional[str] = None
     last_message: Optional[str] = None
@@ -51,72 +54,79 @@ class AgentProgress:
 class AgentManager:
     """
     Manages background agent execution using Claude Code CLI.
-    
+
     Key features:
     - Spawns agents with full tool access via `claude -p`
     - Tracks task status and progress
     - Persists state to .stravinsky/agents.json
     - Provides notification mechanism for task completion
     """
-    
+
     CLAUDE_CLI = "/opt/homebrew/bin/claude"
-    
+
     def __init__(self, base_dir: Optional[str] = None):
         if base_dir:
             self.base_dir = Path(base_dir)
         else:
             self.base_dir = Path.cwd() / ".stravinsky"
-        
+
         self.agents_dir = self.base_dir / "agents"
         self.state_file = self.base_dir / "agents.json"
-        
+
         self.base_dir.mkdir(parents=True, exist_ok=True)
         self.agents_dir.mkdir(parents=True, exist_ok=True)
-        
+
         if not self.state_file.exists():
             self._save_tasks({})
-        
+
         # In-memory tracking for running processes
         self._processes: Dict[str, subprocess.Popen] = {}
         self._notification_queue: Dict[str, List[AgentTask]] = {}
-    
+        self._lock = threading.RLock()
+
     def _load_tasks(self) -> Dict[str, Any]:
         """Load tasks from persistent storage."""
-        try:
-            with open(self.state_file, "r") as f:
-                return json.load(f)
-        except (json.JSONDecodeError, FileNotFoundError):
-            return {}
-    
+        with self._lock:
+            try:
+                if not self.state_file.exists():
+                    return {}
+                with open(self.state_file, "r") as f:
+                    return json.load(f)
+            except (json.JSONDecodeError, FileNotFoundError):
+                return {}
+
     def _save_tasks(self, tasks: Dict[str, Any]):
         """Save tasks to persistent storage."""
-        with open(self.state_file, "w") as f:
-            json.dump(tasks, f, indent=2)
-    
+        with self._lock:
+            with open(self.state_file, "w") as f:
+                json.dump(tasks, f, indent=2)
+
     def _update_task(self, task_id: str, **kwargs):
         """Update a task's fields."""
-        tasks = self._load_tasks()
-        if task_id in tasks:
-            tasks[task_id].update(kwargs)
-            self._save_tasks(tasks)
-    
+        with self._lock:
+            tasks = self._load_tasks()
+            if task_id in tasks:
+                tasks[task_id].update(kwargs)
+                self._save_tasks(tasks)
+
     def get_task(self, task_id: str) -> Optional[Dict[str, Any]]:
         """Get a task by ID."""
         tasks = self._load_tasks()
         return tasks.get(task_id)
-    
+
     def list_tasks(self, parent_session_id: Optional[str] = None) -> List[Dict[str, Any]]:
         """List all tasks, optionally filtered by parent session."""
         tasks = self._load_tasks()
         task_list = list(tasks.values())
-        
+
         if parent_session_id:
             task_list = [t for t in task_list if t.get("parent_session_id") == parent_session_id]
-        
+
         return task_list
-    
+
     def spawn(
         self,
+        token_store: Any,
         prompt: str,
         agent_type: str = "explore",
         description: str = "",
@@ -124,10 +134,11 @@ class AgentManager:
         system_prompt: Optional[str] = None,
         model: str = "gemini-3-flash",
         thinking_budget: int = 0,
+        timeout: int = 300,
     ) -> str:
         """
         Spawn a new background agent.
-        
+
         Args:
             prompt: The task prompt for the agent
             agent_type: Type of agent (explore, dewey, frontend, delphi)
@@ -135,12 +146,13 @@ class AgentManager:
             parent_session_id: Optional parent session for notifications
             system_prompt: Optional custom system prompt
             model: Model to use (gemini-3-flash, claude, etc.)
-            
+            timeout: Maximum execution time in seconds
+
         Returns:
             Task ID for tracking
         """
         task_id = f"agent_{uuid.uuid4().hex[:8]}"
-        
+
         task = AgentTask(
             id=task_id,
             prompt=prompt,
@@ -149,180 +161,185 @@ class AgentManager:
             status="pending",
             created_at=datetime.now().isoformat(),
             parent_session_id=parent_session_id,
+            timeout=timeout,
         )
-        
+
         # Persist task
-        tasks = self._load_tasks()
-        tasks[task_id] = asdict(task)
-        self._save_tasks(tasks)
-        
+        with self._lock:
+            tasks = self._load_tasks()
+            tasks[task_id] = asdict(task)
+            self._save_tasks(tasks)
+
         # Start background execution
-        self._execute_agent(task_id, prompt, agent_type, system_prompt, model, thinking_budget)
-        
+        self._execute_agent(
+            task_id, token_store, prompt, agent_type, system_prompt, model, thinking_budget, timeout
+        )
+
         return task_id
-    
+
     def _execute_agent(
         self,
         task_id: str,
+        token_store: Any,
         prompt: str,
         agent_type: str,
         system_prompt: Optional[str] = None,
         model: str = "gemini-3-flash",
         thinking_budget: int = 0,
+        timeout: int = 300,
     ):
-        """Execute agent in background thread."""
-        
+        """Execute agent using Claude CLI with full tool access.
+
+        Uses `claude -p` to spawn a background agent with complete tool access,
+        just like oh-my-opencode's Sisyphus implementation.
+        """
+
         def run_agent():
             log_file = self.agents_dir / f"{task_id}.log"
             output_file = self.agents_dir / f"{task_id}.out"
-            
-            self._update_task(
-                task_id,
-                status="running",
-                started_at=datetime.now().isoformat()
-            )
-            
+
+            self._update_task(task_id, status="running", started_at=datetime.now().isoformat())
+
             try:
-                # Route based on model type
-                if model.startswith("gemini"):
-                    # Use Gemini via invoke_gemini
-                    logger.info(f"[AgentManager] Spawning Gemini agent {task_id} with model {model}")
-                    
-                    # Build full prompt with system context
-                    full_prompt = prompt
-                    if system_prompt:
-                        full_prompt = f"{system_prompt}\n\n---\n\n{full_prompt}"
-                    
-                    # Import and call invoke_gemini
-                    import asyncio
-                    from .model_invoke import invoke_gemini
-                    
-                    # Run async in thread
-                    loop = asyncio.new_event_loop()
-                    asyncio.set_event_loop(loop)
-                    try:
-                        result = loop.run_until_complete(
-                            invoke_gemini(
-                                prompt=full_prompt, 
-                                model=model,
-                                thinking_budget=thinking_budget
-                            )
-                        )
-                    finally:
-                        loop.close()
-                    
-                    # Write output
-                    with open(output_file, "w") as f:
-                        f.write(result)
-                    
-                    self._update_task(
-                        task_id,
-                        status="completed",
-                        result=result,
-                        completed_at=datetime.now().isoformat()
-                    )
-                    logger.info(f"[AgentManager] Gemini agent {task_id} completed successfully")
-                    
-                else:
-                    # Use Claude CLI for Claude models
-                    cmd = [
-                        self.CLAUDE_CLI,
-                        "-p",  # Non-interactive print mode
-                    ]
-                    
-                    # Add system prompt if provided
-                    if system_prompt:
-                        cmd.extend(["--system-prompt", system_prompt])
-                    
-                    # Add the prompt
-                    cmd.append(prompt)
-                    
-                    logger.info(f"[AgentManager] Spawning Claude agent {task_id}: {' '.join(cmd[:5])}...")
-                    
-                    # Run Claude CLI
+                # Prepare full prompt with system prompt if provided
+                full_prompt = prompt
+                if system_prompt:
+                    full_prompt = f"{system_prompt}\n\n---\n\n{prompt}"
+
+                logger.info(f"[AgentManager] Spawning Claude CLI agent {task_id} ({agent_type})")
+
+                # Build Claude CLI command with full tool access
+                # Using `claude -p` for non-interactive mode with prompt
+                cmd = [
+                    self.CLAUDE_CLI,
+                    "-p",
+                    full_prompt,
+                    "--output-format",
+                    "text",
+                ]
+
+                # Add model selection if specified
+                if model:
+                    cmd.extend(["--model", model])
+
+                # Add system prompt file if we have one
+                if system_prompt:
+                    system_file = self.agents_dir / f"{task_id}.system"
+                    system_file.write_text(system_prompt)
+                    cmd.extend(["--system-prompt", str(system_file)])
+
+                # Execute Claude CLI as subprocess with full tool access
+                logger.info(f"[AgentManager] Running: {' '.join(cmd[:3])}...")
+
+                with open(log_file, "w") as log_f:
                     process = subprocess.Popen(
                         cmd,
-                        stdout=open(output_file, "w"),
-                        stderr=open(log_file, "w"),
+                        stdout=subprocess.PIPE,
+                        stderr=log_f,
+                        text=True,
                         cwd=str(Path.cwd()),
-                        start_new_session=True,
+                        env={**os.environ, "CLAUDE_CODE_ENTRYPOINT": "stravinsky-agent"},
+                        start_new_session=True,  # Allow process group management
                     )
-                    
-                    self._processes[task_id] = process
-                    self._update_task(task_id, pid=process.pid)
-                    
-                    # Wait for completion
-                    return_code = process.wait()
-                    
-                    # Read output
-                    result = ""
-                    if output_file.exists():
-                        result = output_file.read_text()
-                    
-                    if return_code == 0:
+
+                # Track the process
+                self._processes[task_id] = process
+                self._update_task(task_id, pid=process.pid)
+
+                # Wait for completion with timeout
+                try:
+                    stdout, _ = process.communicate(timeout=timeout)
+                    result = stdout.strip() if stdout else ""
+
+                    if process.returncode == 0:
+                        output_file.write_text(result)
                         self._update_task(
                             task_id,
                             status="completed",
                             result=result,
-                            completed_at=datetime.now().isoformat()
+                            completed_at=datetime.now().isoformat(),
                         )
                         logger.info(f"[AgentManager] Agent {task_id} completed successfully")
                     else:
-                        error_log = ""
+                        error_msg = f"Claude CLI exited with code {process.returncode}"
                         if log_file.exists():
-                            error_log = log_file.read_text()
+                            error_msg += f"\n{log_file.read_text()}"
                         self._update_task(
                             task_id,
                             status="failed",
-                            error=f"Exit code {return_code}: {error_log}",
-                            completed_at=datetime.now().isoformat()
+                            error=error_msg,
+                            completed_at=datetime.now().isoformat(),
                         )
-                        logger.error(f"[AgentManager] Agent {task_id} failed: {error_log[:200]}")
-                    
+                        logger.error(f"[AgentManager] Agent {task_id} failed: {error_msg}")
+
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                    self._update_task(
+                        task_id,
+                        status="failed",
+                        error=f"Agent timed out after {timeout}s",
+                        completed_at=datetime.now().isoformat(),
+                    )
+                    logger.warning(f"[AgentManager] Agent {task_id} timed out")
+
+            except FileNotFoundError:
+                error_msg = f"Claude CLI not found at {self.CLAUDE_CLI}. Install with: npm install -g @anthropic-ai/claude-code"
+                log_file.write_text(error_msg)
+                self._update_task(
+                    task_id,
+                    status="failed",
+                    error=error_msg,
+                    completed_at=datetime.now().isoformat(),
+                )
+                logger.error(f"[AgentManager] {error_msg}")
+
             except Exception as e:
+                error_msg = str(e)
+                log_file.write_text(error_msg)
                 self._update_task(
                     task_id,
                     status="failed",
-                    error=str(e),
-                    completed_at=datetime.now().isoformat()
+                    error=error_msg,
+                    completed_at=datetime.now().isoformat(),
                 )
                 logger.exception(f"[AgentManager] Agent {task_id} exception")
+
             finally:
                 self._processes.pop(task_id, None)
                 self._notify_completion(task_id)
-        
+
         # Run in background thread
         thread = threading.Thread(target=run_agent, daemon=True)
         thread.start()
-    
+
     def _notify_completion(self, task_id: str):
         """Queue notification for parent session."""
         task = self.get_task(task_id)
         if not task:
             return
-        
+
         parent_id = task.get("parent_session_id")
         if parent_id:
             if parent_id not in self._notification_queue:
                 self._notification_queue[parent_id] = []
-            
+
             self._notification_queue[parent_id].append(task)
             logger.info(f"[AgentManager] Queued notification for {parent_id}: task {task_id}")
-    
+
     def get_pending_notifications(self, session_id: str) -> List[Dict[str, Any]]:
         """Get and clear pending notifications for a session."""
         notifications = self._notification_queue.pop(session_id, [])
         return notifications
-    
+
     def cancel(self, task_id: str) -> bool:
         """Cancel a running agent task."""
         task = self.get_task(task_id)
         if not task:
             return False
-        
+
         if task["status"] != "running":
             return False
-        
+
         process = self._processes.get(task_id)
         if process:
             try:
@@ -334,31 +351,56 @@ class AgentManager:
                     process.kill()
                 except:
                     pass
-        
-        self._update_task(
-            task_id,
-            status="cancelled",
-            completed_at=datetime.now().isoformat()
-        )
-        
+
+        self._update_task(task_id, status="cancelled", completed_at=datetime.now().isoformat())
+
         return True
-    
+
+    def stop_all(self, clear_history: bool = False) -> int:
+        """
+        Stop all running agents and optionally clear task history.
+
+        Args:
+            clear_history: If True, also remove completed/failed tasks from history
+
+        Returns:
+            Number of tasks stopped/cleared
+        """
+        tasks = self._load_tasks()
+        stopped_count = 0
+
+        # Stop running tasks
+        for task_id, task in list(tasks.items()):
+            if task.get("status") == "running":
+                self.cancel(task_id)
+                stopped_count += 1
+
+        # Optionally clear history
+        if clear_history:
+            cleared = len(tasks)
+            self._save_tasks({})
+            self._processes.clear()
+            logger.info(f"[AgentManager] Cleared all {cleared} agent tasks")
+            return cleared
+
+        return stopped_count
+
     def get_output(self, task_id: str, block: bool = False, timeout: float = 30.0) -> str:
         """
         Get output from an agent task.
-        
+
         Args:
             task_id: The task ID
             block: If True, wait for completion
             timeout: Max seconds to wait if blocking
-            
+
         Returns:
             Formatted task output/status
         """
         task = self.get_task(task_id)
         if not task:
             return f"Task {task_id} not found."
-        
+
         if block and task["status"] == "running":
             # Poll for completion
             start = datetime.now()
@@ -367,11 +409,11 @@ class AgentManager:
                 if task["status"] != "running":
                     break
                 asyncio.sleep(0.5)
-        
+
         status = task["status"]
         description = task.get("description", "")
         agent_type = task.get("agent_type", "unknown")
-        
+
         if status == "completed":
             result = task.get("result", "(no output)")
             return f"""✅ Agent Task Completed
@@ -382,7 +424,7 @@ class AgentManager:
 
 **Result**:
 {result}"""
-        
+
         elif status == "failed":
             error = task.get("error", "(no error details)")
             return f"""❌ Agent Task Failed
@@ -393,14 +435,14 @@ class AgentManager:
 
 **Error**:
 {error}"""
-        
+
         elif status == "cancelled":
             return f"""⚠️ Agent Task Cancelled
 
 **Task ID**: {task_id}
 **Agent**: {agent_type}
 **Description**: {description}"""
-        
+
         else:  # pending or running
             pid = task.get("pid", "N/A")
             started = task.get("started_at", "N/A")
@@ -413,29 +455,47 @@ class AgentManager:
 **Started**: {started}
 
 Use `agent_output` with block=true to wait for completion."""
-    
+
     def get_progress(self, task_id: str, lines: int = 20) -> str:
         """
         Get real-time progress from a running agent's output.
-        
+
         Args:
             task_id: The task ID
             lines: Number of lines to show from the end
-            
+
         Returns:
             Recent output lines and status
         """
         task = self.get_task(task_id)
         if not task:
             return f"Task {task_id} not found."
-        
+
         output_file = self.agents_dir / f"{task_id}.out"
         log_file = self.agents_dir / f"{task_id}.log"
-        
+
         status = task["status"]
         description = task.get("description", "")
         agent_type = task.get("agent_type", "unknown")
-        
+        pid = task.get("pid")
+
+        # Zombie Detection: If running but process is gone
+        if status == "running" and pid:
+            try:
+                import psutil
+
+                if not psutil.pid_exists(pid):
+                    status = "failed"
+                    self._update_task(
+                        task_id,
+                        status="failed",
+                        error="Agent process died unexpectedly (Zombie detected)",
+                        completed_at=datetime.now().isoformat(),
+                    )
+                    logger.warning(f"[AgentManager] Zombie agent detected: {task_id}")
+            except ImportError:
+                pass
+
         # Read recent output
         output_content = ""
         if output_file.exists():
@@ -447,7 +507,7 @@ Use `agent_output` with block=true to wait for completion."""
                     output_content = "\n".join(recent)
             except Exception:
                 pass
-        
+
         # Check log for errors
         log_content = ""
         if log_file.exists():
@@ -455,7 +515,7 @@ Use `agent_output` with block=true to wait for completion."""
                 log_content = log_file.read_text().strip()
             except Exception:
                 pass
-        
+
         # Status emoji
         status_emoji = {
             "pending": "⏳",
@@ -464,7 +524,7 @@ Use `agent_output` with block=true to wait for completion."""
             "failed": "❌",
             "cancelled": "⚠️",
         }.get(status, "❓")
-        
+
         result = f"""{status_emoji} **Agent Progress**
 
 **Task ID**: {task_id}
@@ -472,56 +532,64 @@ Use `agent_output` with block=true to wait for completion."""
 **Description**: {description}
 **Status**: {status}
 """
-        
+
         if output_content:
             result += f"\n**Recent Output** (last {lines} lines):\n```\n{output_content}\n```"
         elif status == "running":
             result += "\n*Agent is working... no output yet.*"
-        
+
         if log_content and status == "failed":
             # Truncate log if too long
             if len(log_content) > 500:
                 log_content = log_content[:500] + "..."
             result += f"\n\n**Error Log**:\n```\n{log_content}\n```"
-        
+
         return result
 
 
 # Global manager instance
 _manager: Optional[AgentManager] = None
+_manager_lock = threading.Lock()
 
 
 def get_manager() -> AgentManager:
     """Get or create the global AgentManager instance."""
     global _manager
     if _manager is None:
-        _manager = AgentManager()
+        with _manager_lock:
+            # Double-check pattern to avoid race condition
+            if _manager is None:
+                _manager = AgentManager()
     return _manager
 
 
 # Tool interface functions
 
+
 async def agent_spawn(
     prompt: str,
     agent_type: str = "explore",
     description: str = "",
     model: str = "gemini-3-flash",
     thinking_budget: int = 0,
+    timeout: int = 300,
 ) -> str:
     """
     Spawn a background agent.
-    
+
     Args:
         prompt: The task for the agent to perform
         agent_type: Type of agent (explore, dewey, frontend, delphi)
         description: Short description shown in status
         model: Model to use (gemini-3-flash, gemini-2.0-flash, claude)
-        
+        thinking_budget: Reserved reasoning tokens
+        timeout: Execution timeout in seconds
+
     Returns:
         Task ID and instructions
     """
     manager = get_manager()
-    
+
     # Map agent types to system prompts
     system_prompts = {
         "explore": "You are a codebase exploration specialist. Find files, patterns, and answer 'where is X?' questions efficiently.",
@@ -552,25 +620,90 @@ ULTRATHINK MODE (when user says "ULTRATHINK" or "think harder"):
 2. Edge Case Analysis: What could go wrong and how we prevented it
 3. The Code: Optimized, bespoke, production-ready, utilizing existing libraries""",
         "delphi": "You are a strategic advisor. Provide architecture guidance, debugging assistance, and code review.",
+        "document_writer": """You are a Technical Documentation Specialist. Your expertise is creating clear, comprehensive documentation.
+
+DOCUMENT TYPES YOU EXCEL AT:
+- README files with proper structure
+- API documentation with examples
+- Architecture decision records (ADRs)
+- User guides and tutorials
+- Inline code documentation
+
+DOCUMENTATION PRINCIPLES:
+- Audience-first: Know who's reading and what they need
+- Progressive disclosure: Overview → Details → Edge cases
+- Examples over explanations: Show, don't just tell
+- Keep it DRY: Reference rather than repeat
+- Version awareness: Note when behavior differs across versions
+
+RESPONSE FORMAT:
+1. Document type and target audience identified
+2. The documentation, properly formatted in markdown""",
+        "multimodal": """You interpret media files that cannot be read as plain text.
+
+Your job: examine the attached file and extract ONLY what was requested.
+
+CAPABILITIES:
+- PDFs: extract text, structure, tables, data from specific sections
+- Images: describe layouts, UI elements, text, diagrams, charts
+- Diagrams: explain relationships, flows, architecture depicted
+- Screenshots: analyze UI/UX, identify components, extract text
+
+HOW YOU WORK:
+1. Receive a file path and a goal describing what to extract
+2. Read and analyze the file deeply using Gemini's vision capabilities
+3. Return ONLY the relevant extracted information
+4. The main agent never processes the raw file - you save context tokens
+
+RESPONSE RULES:
+- Return extracted information directly, no preamble
+- If info not found, state clearly what's missing
+- Be thorough on the goal, concise on everything else""",
     }
-    
+
     system_prompt = system_prompts.get(agent_type, None)
-    
+
+    # Override model and thinking_budget based on agent type for optimal performance
+    # Per project requirements:
+    # - Gemini tiers are controlled by thinking_budget, NOT model name
+    # - high: 32000, medium: 16000, low: 8000 thinking tokens
+    # - "gemini-3-pro-high" = gemini-3-pro-low + thinking_budget=32000
+    # - "gemini-3-flash" equivalent = gemini-3-pro-low + thinking_budget=0
+    agent_configs = {
+        "stravinsky": {"model": "claude-opus-4-5", "thinking_budget": 0},
+        "frontend": {"model": "gemini-3-pro-low", "thinking_budget": 32000},  # HIGH tier
+        "document_writer": {"model": "gemini-3-pro-low", "thinking_budget": 0},  # Flash equiv
+        "multimodal": {"model": "gemini-3-pro-low", "thinking_budget": 0},  # Flash equiv
+        "explore": {"model": "gemini-3-pro-low", "thinking_budget": 0},  # Flash equiv
+        "delphi": {"model": "gpt-5.2", "thinking_budget": 0},  # GPT-5.2 strategic
+        "dewey": {"model": "gemini-3-pro-low", "thinking_budget": 0},  # Flash equiv
+    }
+    config = agent_configs.get(agent_type, {"model": model, "thinking_budget": thinking_budget})
+    actual_model = config["model"]
+    actual_thinking_budget = config["thinking_budget"]
+
+    # Get token store for authentication
+    from ..auth.token_store import TokenStore
+
+    token_store = TokenStore()
+
     task_id = manager.spawn(
+        token_store=token_store,
         prompt=prompt,
         agent_type=agent_type,
         description=description or prompt[:50],
         system_prompt=system_prompt,
-        model=model,
-        thinking_budget=thinking_budget,
+        model=actual_model,
+        thinking_budget=actual_thinking_budget,
+        timeout=timeout,
     )
-    
+
     return f"""🚀 Background agent spawned successfully.
 
 **Task ID**: {task_id}
 **Agent Type**: {agent_type}
-**Model**: {model}
-**Thinking Budget**: {thinking_budget if thinking_budget > 0 else "N/A"}
+**Model**: {actual_model}
+**Thinking Budget**: {actual_thinking_budget if actual_thinking_budget > 0 else "N/A"}
 **Description**: {description or prompt[:50]}
 
 The agent is now running. Use:
@@ -582,11 +715,11 @@ The agent is now running. Use:
 async def agent_output(task_id: str, block: bool = False) -> str:
     """
     Get output from a background agent task.
-    
+
     Args:
         task_id: The task ID from agent_spawn
         block: If True, wait for the task to complete (up to 30s)
-        
+
     Returns:
         Task status and output
     """
@@ -594,19 +727,55 @@ async def agent_output(task_id: str, block: bool = False) -> str:
     return manager.get_output(task_id, block=block)
 
 
+async def agent_retry(
+    task_id: str,
+    new_prompt: Optional[str] = None,
+    new_timeout: Optional[int] = None,
+) -> str:
+    """
+    Retry a failed or timed-out background agent.
+
+    Args:
+        task_id: The ID of the task to retry
+        new_prompt: Optional refined prompt for the retry
+        new_timeout: Optional new timeout in seconds
+
+    Returns:
+        New Task ID and status
+    """
+    manager = get_manager()
+    task = manager.get_task(task_id)
+
+    if not task:
+        return f"❌ Task {task_id} not found."
+
+    if task["status"] in ["running", "pending"]:
+        return f"⚠️ Task {task_id} is still {task['status']}. Cancel it first if you want to retry."
+
+    prompt = new_prompt or task["prompt"]
+    timeout = new_timeout or task.get("timeout", 300)
+
+    return await agent_spawn(
+        prompt=prompt,
+        agent_type=task["agent_type"],
+        description=f"Retry of {task_id}: {task['description']}",
+        timeout=timeout,
+    )
+
+
 async def agent_cancel(task_id: str) -> str:
     """
     Cancel a running background agent.
-    
+
     Args:
         task_id: The task ID to cancel
-        
+
     Returns:
         Cancellation result
     """
     manager = get_manager()
     success = manager.cancel(task_id)
-    
+
     if success:
         return f"✅ Agent task {task_id} has been cancelled."
     else:
@@ -620,18 +789,18 @@ async def agent_cancel(task_id: str) -> str:
 async def agent_list() -> str:
     """
     List all background agent tasks.
-    
+
     Returns:
         Formatted list of tasks
     """
     manager = get_manager()
     tasks = manager.list_tasks()
-    
+
     if not tasks:
         return "No background agent tasks found."
-    
+
     lines = ["**Background Agent Tasks**", ""]
-    
+
     for t in sorted(tasks, key=lambda x: x.get("created_at", ""), reverse=True):
         status_emoji = {
             "pending": "⏳",
@@ -640,24 +809,24 @@ async def agent_list() -> str:
             "failed": "❌",
             "cancelled": "⚠️",
         }.get(t["status"], "❓")
-        
+
         desc = t.get("description", t.get("prompt", "")[:40])
         lines.append(f"- {status_emoji} [{t['id']}] {t['agent_type']}: {desc}")
-    
+
     return "\n".join(lines)
 
 
 async def agent_progress(task_id: str, lines: int = 20) -> str:
     """
     Get real-time progress from a running background agent.
-    
+
     Shows the most recent output lines from the agent, useful for
     monitoring what the agent is currently doing.
-    
+
     Args:
         task_id: The task ID from agent_spawn
         lines: Number of recent output lines to show (default 20)
-        
+
     Returns:
         Recent agent output and status
     """