steerdev 0.4.27__py3-none-any.whl

steerdev_agent/daemon.py

@@ -0,0 +1,675 @@
+"""Persistent agent daemon mode.
+
+Runs indefinitely, polling for commands from the API and falling back
+to the task queue when idle. Each command/task execution creates a session
+with full event streaming, reusing the existing Runner infrastructure.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import signal
+import socket
+from pathlib import Path
+
+import httpx
+from loguru import logger
+from rich.console import Console
+from rich.panel import Panel
+
+from steerdev_agent.api.commands import CommandResponse, CommandsClient
+from steerdev_agent.api.events import EventsClient
+from steerdev_agent.api.sessions import SessionCreateRequest, SessionsClient
+from steerdev_agent.api.tasks import TasksClient
+from steerdev_agent.config.models import DaemonConfig, ExecutorConfig
+from steerdev_agent.executor import ExecutorFactory
+from steerdev_agent.executor.base import EventType
+from steerdev_agent.executor.claude import ClaudeExecutorError
+from steerdev_agent.prompt.builder import PromptBuilder, PromptContext, TaskContext
+
+console = Console()
+
+
+class DaemonError(Exception):
+    """Error during daemon execution."""
+
+
+class Daemon:
+    """Persistent agent that polls for commands and executes them.
+
+    Lifecycle:
+    1. Register agent (get/create agent_id via API)
+    2. Recover stale commands from previous crash
+    3. Start heartbeat background task
+    4. Enter poll loop:
+       a. Poll /commands/next
+       b. If command found -> execute it
+       c. If no command AND auto_fetch_tasks -> fetch next task
+       d. If nothing -> sleep(idle_poll_interval)
+       e. Check for shutdown signal
+    5. On shutdown: send offline heartbeat, exit
+    """
+
+    def __init__(
+        self,
+        project_id: str,
+        agent_name: str,
+        working_directory: str | Path | None = None,
+        api_key: str | None = None,
+        agent_type: str = "claude",
+        model: str | None = None,
+        max_turns: int | None = None,
+        daemon_config: DaemonConfig | None = None,
+        executor_config: ExecutorConfig | None = None,
+    ) -> None:
+        self.project_id = project_id
+        self.agent_name = agent_name
+        self.working_directory = Path(working_directory or Path.cwd())
+        self._api_key = api_key
+        self.agent_type = agent_type
+        self.model = model
+        self.max_turns = max_turns
+        self._daemon_config = daemon_config or DaemonConfig()
+        self._executor_config = executor_config or ExecutorConfig()
+
+        # State
+        self._agent_id: str | None = None
+        self._shutdown_event = asyncio.Event()
+        self._is_busy = False
+        self._consecutive_errors = 0
+        self._commands_executed = 0
+        self._commands_succeeded = 0
+        self._commands_failed = 0
+        self._tasks_executed = 0
+
+        # Components
+        self._commands_client: CommandsClient | None = None
+        self._sessions_client: SessionsClient | None = None
+        self._prompt_builder = PromptBuilder()
+        self._heartbeat_task: asyncio.Task[None] | None = None
+
+    async def start(self) -> None:
+        """Start the daemon loop."""
+        self._setup_signal_handlers()
+
+        console.print(
+            Panel(
+                f"[bold blue]SteerDev Daemon[/bold blue]\n"
+                f"Project: {self.project_id}\n"
+                f"Agent: {self.agent_name}\n"
+                f"Working Dir: {self.working_directory}\n"
+                f"Poll Interval: {self._daemon_config.poll_interval_seconds}s\n"
+                f"Auto-fetch Tasks: {self._daemon_config.auto_fetch_tasks}",
+                title="Starting Daemon",
+            )
+        )
+
+        try:
+            # Step 1: Register agent
+            self._agent_id = await self._register_agent()
+            if not self._agent_id:
+                raise DaemonError("Failed to register agent")
+
+            logger.info(f"Agent registered: {self._agent_id}")
+
+            # Initialize clients
+            self._commands_client = CommandsClient(
+                agent_id=self._agent_id,
+                api_key=self._api_key,
+            )
+            self._sessions_client = SessionsClient(api_key=self._api_key)
+
+            # Step 2: Recover stale commands
+            await self._recover_stale_commands()
+
+            # Step 3: Start heartbeat
+            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+
+            # Step 4: Main poll loop
+            await self._poll_loop()
+
+        except DaemonError as e:
+            logger.error(f"Daemon error: {e}")
+            console.print(f"[red]Daemon error: {e}[/red]")
+        except Exception as e:
+            logger.exception(f"Unexpected daemon error: {e}")
+            console.print(f"[red]Unexpected error: {e}[/red]")
+        finally:
+            # Step 5: Shutdown
+            await self._shutdown()
+
+    async def _register_agent(self) -> str | None:
+        """Register this agent via the API and return the agent_id."""
+        try:
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                from steerdev_agent.api.client import get_api_endpoint, get_api_key
+
+                api_key = self._api_key or get_api_key()
+                api_base = get_api_endpoint()
+
+                response = await client.post(
+                    f"{api_base}/agents",
+                    headers={
+                        "Authorization": f"Bearer {api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "name": self.agent_name,
+                        "application": self.agent_type,
+                        "hostname": socket.gethostname(),
+                    },
+                )
+
+                if response.status_code in (200, 201):
+                    data = response.json()
+                    return data["id"]
+
+                logger.error(f"Failed to register agent: {response.status_code} - {response.text}")
+                return None
+        except httpx.RequestError as e:
+            logger.error(f"Request error registering agent: {e}")
+            return None
+
+    async def _recover_stale_commands(self) -> None:
+        """Mark any claimed/executing commands from a previous crash as failed."""
+        if not self._commands_client:
+            return
+
+        for stale_status in ("claimed", "executing"):
+            result = await self._commands_client.list_commands(status=stale_status)  # type: ignore[arg-type]
+            if not result:
+                continue
+
+            for cmd in result.commands:
+                logger.warning(f"Recovering stale command {cmd.id} (was {cmd.status})")
+                await self._commands_client.mark_failed(
+                    cmd.id,
+                    error=f"Recovered on daemon restart (was {cmd.status})",
+                )
+
+    async def _poll_loop(self) -> None:
+        """Main loop: poll for commands, execute them, or fall back to tasks."""
+        logger.info("Entering poll loop")
+        console.print("[green]Daemon ready, polling for commands...[/green]")
+
+        while not self._shutdown_event.is_set():
+            try:
+                executed = await self._poll_once()
+
+                if executed:
+                    self._consecutive_errors = 0
+                    interval = self._daemon_config.poll_interval_seconds
+                else:
+                    interval = self._daemon_config.idle_poll_interval_seconds
+
+            except Exception as e:
+                self._consecutive_errors += 1
+                logger.error(f"Poll error ({self._consecutive_errors}): {e}")
+
+                if self._consecutive_errors >= self._daemon_config.max_consecutive_errors:
+                    logger.error("Max consecutive errors reached, shutting down")
+                    console.print("[red]Max consecutive errors reached, shutting down[/red]")
+                    break
+
+                # Exponential backoff: 2^n seconds, capped at 60s
+                backoff = min(2**self._consecutive_errors, 60)
+                interval = backoff
+
+            # Sleep with cancellation support
+            try:
+                await asyncio.wait_for(
+                    self._shutdown_event.wait(),
+                    timeout=interval,
+                )
+                # If we get here, shutdown was signaled
+                break
+            except TimeoutError:
+                # Normal timeout, continue polling
+                pass
+
+    async def _poll_once(self) -> bool:
+        """Poll for one command or task. Returns True if something was executed."""
+        assert self._commands_client is not None
+
+        # Try command queue first
+        command = await self._commands_client.poll_next()
+
+        if command:
+            # Handle control commands specially
+            if command.command_type == "control":
+                return await self._handle_control_command(command)
+
+            await self._execute_command(command)
+            return True
+
+        # Fall back to task queue if enabled
+        if self._daemon_config.auto_fetch_tasks:
+            executed = await self._execute_task_fallback()
+            if executed:
+                return True
+
+        return False
+
+    async def _handle_control_command(self, command: CommandResponse) -> bool:
+        """Handle control commands (shutdown, etc.)."""
+        assert self._commands_client is not None
+
+        if command.control_action == "shutdown":
+            logger.info("Received shutdown command")
+            console.print("[yellow]Received shutdown command[/yellow]")
+            await self._commands_client.mark_completed(command.id, result="Shutdown acknowledged")
+            self._shutdown_event.set()
+            return True
+
+        logger.warning(f"Unknown control action: {command.control_action}")
+        await self._commands_client.mark_failed(
+            command.id,
+            error=f"Unknown control action: {command.control_action}",
+        )
+        return True
+
+    async def _execute_command(self, command: CommandResponse) -> None:
+        """Execute a single command, creating a session and streaming events."""
+        assert self._commands_client is not None
+        assert self._sessions_client is not None
+
+        self._is_busy = True
+        self._commands_executed += 1
+
+        cmd_desc = (
+            f"task={command.task_id}"
+            if command.command_type == "task"
+            else f"prompt={command.prompt[:50]}..."
+            if command.prompt and len(command.prompt) > 50
+            else f"prompt={command.prompt}"
+            if command.prompt
+            else command.command_type
+        )
+        console.print(f"\n[bold cyan]Executing command: {cmd_desc}[/bold cyan]")
+        console.print(f"[dim]Command ID: {command.id}[/dim]")
+
+        # Build prompt based on command type
+        prompt = await self._build_prompt_for_command(command)
+        if not prompt:
+            await self._commands_client.mark_failed(command.id, error="Failed to build prompt")
+            self._commands_failed += 1
+            self._is_busy = False
+            return
+
+        # Create session
+        session = await self._sessions_client.create_session(
+            SessionCreateRequest(
+                project_id=self.project_id,
+                task_id=command.task_id,
+                agent_name=self.agent_name,
+                agent_type=self.agent_type,
+                prompt=prompt,
+                working_directory=str(command.working_directory or self.working_directory),
+            )
+        )
+
+        if not session:
+            await self._commands_client.mark_failed(command.id, error="Failed to create session")
+            self._commands_failed += 1
+            self._is_busy = False
+            return
+
+        # Mark command as executing with session_id
+        await self._commands_client.mark_executing(command.id, session.id)
+        await self._sessions_client.mark_running(session.id)
+
+        # Execute
+        events_client = EventsClient(session_id=session.id, api_key=self._api_key)
+        await events_client.start()
+        events_sent = 0
+
+        try:
+            executor = ExecutorFactory.create(
+                config=self._executor_config,
+                working_directory=str(command.working_directory or self.working_directory),
+                model=self.model,
+                max_turns=self.max_turns,
+            )
+
+            # Start with timeout
+            await executor.start(prompt)
+            console.print("[dim]Agent started, streaming output...[/dim]")
+
+            async def _run_execution() -> int:
+                nonlocal events_sent
+                async for event in executor.stream_events():
+                    await events_client.add_event(
+                        event_type=event.event_type.value,
+                        data=event.data,
+                        raw_json=event.raw_json,
+                        timestamp=event.timestamp,
+                    )
+                    events_sent += 1
+
+                    if event.event_type == EventType.ASSISTANT:
+                        msg = event.data.get("message", {})
+                        content = msg.get("content", "")
+                        if isinstance(content, str) and content:
+                            preview = content[:100] + "..." if len(content) > 100 else content
+                            console.print(f"[cyan]Assistant:[/cyan] {preview}")
+
+                    # Check for cancellation
+                    if self._shutdown_event.is_set():
+                        logger.info("Shutdown during execution, stopping executor")
+                        await executor.stop()
+                        return -1
+
+                return await executor.wait()
+
+            exit_code = await asyncio.wait_for(
+                _run_execution(),
+                timeout=self._daemon_config.command_timeout_seconds,
+            )
+
+            agent_session_id = executor.session_id
+
+            if exit_code != 0:
+                stderr = await executor.get_stderr()
+                error_msg = stderr.strip() if stderr else f"Process exited with code {exit_code}"
+                await self._commands_client.mark_failed(command.id, error=error_msg)
+                await self._sessions_client.mark_failed(
+                    session.id,
+                    metadata={"error": error_msg, "exit_code": exit_code},
+                )
+                self._commands_failed += 1
+                console.print(f"[red]Command failed: {error_msg}[/red]")
+            else:
+                await self._commands_client.mark_completed(
+                    command.id,
+                    result=f"Completed successfully. Events: {events_sent}",
+                )
+                await self._sessions_client.mark_completed(
+                    session.id,
+                    agent_session_id=agent_session_id,
+                )
+                self._commands_succeeded += 1
+                console.print("[green]Command completed successfully[/green]")
+
+        except TimeoutError:
+            error_msg = f"Command timed out after {self._daemon_config.command_timeout_seconds}s"
+            logger.error(error_msg)
+            await self._commands_client.mark_failed(command.id, error=error_msg)
+            await self._sessions_client.mark_failed(session.id, metadata={"error": error_msg})
+            self._commands_failed += 1
+            console.print(f"[red]{error_msg}[/red]")
+
+        except ClaudeExecutorError as e:
+            error_msg = str(e)
+            logger.error(f"Executor error: {error_msg}")
+            await self._commands_client.mark_failed(command.id, error=error_msg)
+            await self._sessions_client.mark_failed(session.id, metadata={"error": error_msg})
+            self._commands_failed += 1
+
+        except Exception as e:
+            error_msg = str(e)
+            logger.exception(f"Unexpected error executing command: {error_msg}")
+            await self._commands_client.mark_failed(command.id, error=error_msg)
+            await self._sessions_client.mark_failed(session.id, metadata={"error": error_msg})
+            self._commands_failed += 1
+
+        finally:
+            await events_client.close()
+            self._is_busy = False
+
+    async def _build_prompt_for_command(self, command: CommandResponse) -> str | None:
+        """Build a prompt string from a command."""
+        if command.command_type == "prompt":
+            return command.prompt
+
+        if command.command_type == "task":
+            if not command.task_id:
+                return None
+
+            # Fetch task details
+            with TasksClient(api_key=self._api_key) as client:
+                task = client.get_task(command.task_id)
+                if not task:
+                    logger.error(f"Task not found: {command.task_id}")
+                    return None
+
+            context = PromptContext(
+                task=TaskContext(
+                    id=task.get("id", ""),
+                    title=task.get("title", "Unknown"),
+                    prompt=task.get("prompt", ""),
+                    status=task.get("status", "unstarted"),
+                    priority=task.get("priority", 3),
+                    working_directory=task.get("working_directory"),
+                ),
+            )
+            return self._prompt_builder.build(context)
+
+        return None
+
+    async def _execute_task_fallback(self) -> bool:
+        """Fetch and execute the next task from the task queue.
+
+        Returns True if a task was executed, False if none available.
+        """
+        with TasksClient(api_key=self._api_key) as client:
+            task = client.get_next_task(project_id=self.project_id)
+
+        if not task:
+            return False
+
+        task_id = task.get("id", "")
+        task_title = task.get("title", "Unknown")
+
+        console.print(f"\n[bold cyan]Auto-fetched task: {task_title}[/bold cyan]")
+        console.print(f"[dim]Task ID: {task_id}[/dim]")
+
+        self._is_busy = True
+        self._tasks_executed += 1
+
+        # Build prompt
+        context = PromptContext(
+            task=TaskContext(
+                id=task_id,
+                title=task_title,
+                prompt=task.get("prompt", ""),
+                status=task.get("status", "unstarted"),
+                priority=task.get("priority", 3),
+                working_directory=task.get("working_directory"),
+            ),
+        )
+        prompt = self._prompt_builder.build(context)
+
+        # Create session
+        assert self._sessions_client is not None
+        session = await self._sessions_client.create_session(
+            SessionCreateRequest(
+                project_id=self.project_id,
+                task_id=task_id,
+                agent_name=self.agent_name,
+                agent_type=self.agent_type,
+                prompt=prompt,
+                working_directory=str(self.working_directory),
+            )
+        )
+
+        if not session:
+            logger.error("Failed to create session for task fallback")
+            self._is_busy = False
+            return True  # We tried, count as executed
+
+        await self._sessions_client.mark_running(session.id)
+
+        events_client = EventsClient(session_id=session.id, api_key=self._api_key)
+        await events_client.start()
+
+        try:
+            executor = ExecutorFactory.create(
+                config=self._executor_config,
+                working_directory=str(self.working_directory),
+                model=self.model,
+                max_turns=self.max_turns,
+            )
+
+            await executor.start(prompt)
+
+            async for event in executor.stream_events():
+                await events_client.add_event(
+                    event_type=event.event_type.value,
+                    data=event.data,
+                    raw_json=event.raw_json,
+                    timestamp=event.timestamp,
+                )
+
+                if event.event_type == EventType.ASSISTANT:
+                    msg = event.data.get("message", {})
+                    content = msg.get("content", "")
+                    if isinstance(content, str) and content:
+                        preview = content[:100] + "..." if len(content) > 100 else content
+                        console.print(f"[cyan]Assistant:[/cyan] {preview}")
+
+            exit_code = await executor.wait()
+            agent_session_id = executor.session_id
+
+            if exit_code != 0:
+                stderr = await executor.get_stderr()
+                error_msg = stderr.strip() if stderr else f"Process exited with code {exit_code}"
+                await self._sessions_client.mark_failed(
+                    session.id,
+                    metadata={"error": error_msg, "exit_code": exit_code},
+                )
+                console.print(f"[red]Task failed: {error_msg}[/red]")
+            else:
+                await self._sessions_client.mark_completed(
+                    session.id,
+                    agent_session_id=agent_session_id,
+                )
+                console.print("[green]Task completed successfully[/green]")
+
+        except Exception as e:
+            logger.exception(f"Error executing task fallback: {e}")
+            await self._sessions_client.mark_failed(
+                session.id,
+                metadata={"error": str(e)},
+            )
+        finally:
+            await events_client.close()
+            self._is_busy = False
+
+        return True
+
+    async def _heartbeat_loop(self) -> None:
+        """Background task: send heartbeat every N seconds."""
+        while not self._shutdown_event.is_set():
+            try:
+                await self._send_heartbeat()
+            except Exception as e:
+                logger.warning(f"Heartbeat error: {e}")
+
+            try:
+                await asyncio.wait_for(
+                    self._shutdown_event.wait(),
+                    timeout=self._daemon_config.heartbeat_interval_seconds,
+                )
+                break  # Shutdown signaled
+            except TimeoutError:
+                pass
+
+    async def _send_heartbeat(self) -> None:
+        """Send a heartbeat to update agent status."""
+        try:
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                from steerdev_agent.api.client import get_api_endpoint, get_api_key
+
+                api_key = self._api_key or get_api_key()
+                api_base = get_api_endpoint()
+
+                # Re-register to update heartbeat (getOrCreateAgent updates last_heartbeat_at)
+                await client.post(
+                    f"{api_base}/agents",
+                    headers={
+                        "Authorization": f"Bearer {api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "name": self.agent_name,
+                        "application": self.agent_type,
+                        "hostname": socket.gethostname(),
+                    },
+                )
+        except httpx.RequestError as e:
+            logger.debug(f"Heartbeat request error: {e}")
+
+    def _setup_signal_handlers(self) -> None:
+        """Set up SIGINT/SIGTERM handlers for graceful shutdown."""
+        loop = asyncio.get_running_loop()
+
+        def _signal_handler() -> None:
+            logger.info("Received shutdown signal")
+            console.print("\n[yellow]Received shutdown signal, finishing current work...[/yellow]")
+            self._shutdown_event.set()
+
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            loop.add_signal_handler(sig, _signal_handler)
+
+    async def _shutdown(self) -> None:
+        """Clean up resources and send offline heartbeat."""
+        console.print("[dim]Shutting down daemon...[/dim]")
+
+        # Cancel heartbeat task
+        if self._heartbeat_task and not self._heartbeat_task.done():
+            self._heartbeat_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._heartbeat_task
+
+        # Send offline heartbeat (best effort)
+        with contextlib.suppress(Exception):
+            await self._send_heartbeat()
+
+        # Close clients
+        if self._commands_client:
+            await self._commands_client.close()
+        if self._sessions_client:
+            await self._sessions_client.close()
+
+        # Print summary
+        console.print(
+            Panel(
+                f"[bold]Daemon Summary[/bold]\n"
+                f"Commands Executed: {self._commands_executed}\n"
+                f"Commands Succeeded: {self._commands_succeeded}\n"
+                f"Commands Failed: {self._commands_failed}\n"
+                f"Tasks Auto-fetched: {self._tasks_executed}",
+                title="Shutdown Complete",
+                border_style="yellow",
+            )
+        )
+
+
+async def run_daemon(
+    project_id: str,
+    agent_name: str,
+    working_directory: str | None = None,
+    api_key: str | None = None,
+    agent_type: str = "claude",
+    model: str | None = None,
+    max_turns: int | None = None,
+    daemon_config: DaemonConfig | None = None,
+    executor_config: ExecutorConfig | None = None,
+) -> None:
+    """Run the daemon.
+
+    Convenience function for running the daemon with minimal setup.
+    """
+    daemon = Daemon(
+        project_id=project_id,
+        agent_name=agent_name,
+        working_directory=working_directory,
+        api_key=api_key,
+        agent_type=agent_type,
+        model=model,
+        max_turns=max_turns,
+        daemon_config=daemon_config,
+        executor_config=executor_config,
+    )
+    await daemon.start()