zwarm 2.3.5__py3-none-any.whl

zwarm/sessions/manager.py

@@ -0,0 +1,792 @@
+"""
+Codex Session Manager - Background process management for Codex agents.
+
+Architecture:
+- Each session runs `codex exec --json` in a background subprocess
+- Output is streamed to .zwarm/sessions/<session_id>/output.jsonl
+- Session metadata stored in meta.json
+- Can inject follow-up messages by starting new turns with context
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import signal
+import subprocess
+import time
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+
+class SessionStatus(str, Enum):
+    """Status of a codex session."""
+    PENDING = "pending"      # Created but not started
+    RUNNING = "running"      # Process is running
+    COMPLETED = "completed"  # Process exited successfully
+    FAILED = "failed"        # Process exited with error
+    KILLED = "killed"        # Manually killed
+
+
+@dataclass
+class SessionMessage:
+    """A message in a session's history."""
+    role: str  # "user", "assistant", "system", "tool"
+    content: str
+    timestamp: str = ""
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        return {
+            "role": self.role,
+            "content": self.content,
+            "timestamp": self.timestamp,
+            "metadata": self.metadata,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "SessionMessage":
+        return cls(
+            role=data.get("role", "unknown"),
+            content=data.get("content", ""),
+            timestamp=data.get("timestamp", ""),
+            metadata=data.get("metadata", {}),
+        )
+
+
+@dataclass
+class CodexSession:
+    """A managed Codex session."""
+    id: str
+    task: str
+    status: SessionStatus
+    working_dir: Path
+    created_at: str
+    updated_at: str
+    pid: int | None = None
+    exit_code: int | None = None
+    model: str = "gpt-5.1-codex-mini"
+    turn: int = 1
+    messages: list[SessionMessage] = field(default_factory=list)
+    token_usage: dict[str, int] = field(default_factory=dict)
+    error: str | None = None
+    # Source tracking: "user" for direct spawns, "orchestrator:<instance_id>" for delegated
+    source: str = "user"
+    # Adapter used: "codex", "claude_code", etc.
+    adapter: str = "codex"
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "task": self.task,
+            "status": self.status.value,
+            "working_dir": str(self.working_dir),
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+            "pid": self.pid,
+            "exit_code": self.exit_code,
+            "model": self.model,
+            "turn": self.turn,
+            "messages": [m.to_dict() for m in self.messages],
+            "token_usage": self.token_usage,
+            "error": self.error,
+            "source": self.source,
+            "adapter": self.adapter,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "CodexSession":
+        return cls(
+            id=data["id"],
+            task=data["task"],
+            status=SessionStatus(data["status"]),
+            working_dir=Path(data["working_dir"]),
+            created_at=data["created_at"],
+            updated_at=data["updated_at"],
+            pid=data.get("pid"),
+            exit_code=data.get("exit_code"),
+            model=data.get("model", "gpt-5.1-codex-mini"),
+            turn=data.get("turn", 1),
+            messages=[SessionMessage.from_dict(m) for m in data.get("messages", [])],
+            token_usage=data.get("token_usage", {}),
+            error=data.get("error"),
+            source=data.get("source", "user"),
+            adapter=data.get("adapter", "codex"),
+        )
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the session process is still running."""
+        if self.pid is None:
+            return False
+        try:
+            os.kill(self.pid, 0)  # Signal 0 just checks if process exists
+            return True
+        except OSError:
+            return False
+
+    @property
+    def short_id(self) -> str:
+        """Get first 8 chars of ID for display."""
+        return self.id[:8]
+
+    @property
+    def runtime(self) -> str:
+        """Get human-readable runtime."""
+        created = datetime.fromisoformat(self.created_at)
+        now = datetime.now()
+        delta = now - created
+
+        if delta.total_seconds() < 60:
+            return f"{int(delta.total_seconds())}s"
+        elif delta.total_seconds() < 3600:
+            return f"{int(delta.total_seconds() / 60)}m"
+        else:
+            return f"{delta.total_seconds() / 3600:.1f}h"
+
+    @property
+    def source_display(self) -> str:
+        """Get short display string for source."""
+        if self.source == "user":
+            return "you"
+        elif self.source.startswith("orchestrator:"):
+            # Extract instance ID and shorten it
+            instance_id = self.source.split(":", 1)[1]
+            return f"orch:{instance_id[:4]}"
+        else:
+            return self.source[:8]
+
+
+class CodexSessionManager:
+    """
+    Manages background Codex sessions.
+
+    Sessions are stored in:
+    .zwarm/sessions/<session_id>/
+        meta.json      - Session metadata
+        output.jsonl   - Raw JSONL output from codex exec
+        turns/
+            turn_1.jsonl
+            turn_2.jsonl
+            ...
+    """
+
+    def __init__(self, state_dir: Path | str = ".zwarm"):
+        self.state_dir = Path(state_dir)
+        self.sessions_dir = self.state_dir / "sessions"
+        self.sessions_dir.mkdir(parents=True, exist_ok=True)
+
+    def _session_dir(self, session_id: str) -> Path:
+        """Get the directory for a session."""
+        return self.sessions_dir / session_id
+
+    def _meta_path(self, session_id: str) -> Path:
+        """Get the metadata file path for a session."""
+        return self._session_dir(session_id) / "meta.json"
+
+    def _output_path(self, session_id: str, turn: int = 1) -> Path:
+        """Get the output file path for a session turn."""
+        session_dir = self._session_dir(session_id)
+        turns_dir = session_dir / "turns"
+        turns_dir.mkdir(parents=True, exist_ok=True)
+        return turns_dir / f"turn_{turn}.jsonl"
+
+    def _save_session(self, session: CodexSession) -> None:
+        """Save session metadata."""
+        session.updated_at = datetime.now().isoformat()
+        meta_path = self._meta_path(session.id)
+        meta_path.parent.mkdir(parents=True, exist_ok=True)
+        meta_path.write_text(json.dumps(session.to_dict(), indent=2))
+
+    def _load_session(self, session_id: str) -> CodexSession | None:
+        """Load session from disk."""
+        meta_path = self._meta_path(session_id)
+        if not meta_path.exists():
+            return None
+        try:
+            data = json.loads(meta_path.read_text())
+            return CodexSession.from_dict(data)
+        except (json.JSONDecodeError, KeyError) as e:
+            print(f"Error loading session {session_id}: {e}")
+            return None
+
+    def list_sessions(self, status: SessionStatus | None = None) -> list[CodexSession]:
+        """List all sessions, optionally filtered by status."""
+        sessions = []
+        if not self.sessions_dir.exists():
+            return sessions
+
+        for session_dir in self.sessions_dir.iterdir():
+            if not session_dir.is_dir():
+                continue
+            session = self._load_session(session_dir.name)
+            if session:
+                # Update status if process died OR output indicates completion
+                # (output check is more reliable than PID check due to PID reuse)
+                if session.status == SessionStatus.RUNNING:
+                    if self._is_output_complete(session.id, session.turn) or not session.is_running:
+                        self._update_session_status(session)
+
+                if status is None or session.status == status:
+                    sessions.append(session)
+
+        # Sort by created_at descending (newest first)
+        sessions.sort(key=lambda s: s.created_at, reverse=True)
+        return sessions
+
+    def get_session(self, session_id: str) -> CodexSession | None:
+        """Get a session by ID (supports partial ID matching)."""
+        # Try exact match first
+        session = self._load_session(session_id)
+        if session:
+            if session.status == SessionStatus.RUNNING:
+                if self._is_output_complete(session.id, session.turn) or not session.is_running:
+                    self._update_session_status(session)
+            return session
+
+        # Try partial match
+        for session_dir in self.sessions_dir.iterdir():
+            if session_dir.name.startswith(session_id):
+                session = self._load_session(session_dir.name)
+                if session:
+                    if session.status == SessionStatus.RUNNING:
+                        if self._is_output_complete(session.id, session.turn) or not session.is_running:
+                            self._update_session_status(session)
+                    return session
+
+        return None
+
+    def _is_output_complete(self, session_id: str, turn: int) -> bool:
+        """
+        Check if output file indicates the task completed.
+
+        Looks for completion markers like 'turn.completed' or 'task.completed'
+        in the JSONL output. This is more reliable than PID checking.
+        """
+        output_path = self._output_path(session_id, turn)
+        if not output_path.exists():
+            return False
+
+        try:
+            content = output_path.read_text()
+            for line in content.strip().split("\n"):
+                if not line.strip():
+                    continue
+                try:
+                    event = json.loads(line)
+                    event_type = event.get("type", "")
+                    # Check for any completion marker
+                    if event_type in ("turn.completed", "task.completed", "completed", "done"):
+                        return True
+                    # Also check for error as a form of completion
+                    if event_type == "error":
+                        return True
+                except json.JSONDecodeError:
+                    continue
+        except Exception:
+            pass
+
+        return False
+
+    def _update_session_status(self, session: CodexSession) -> None:
+        """Update session status after process completion."""
+        # Parse output to determine status
+        output_path = self._output_path(session.id, session.turn)
+        if output_path.exists():
+            messages, usage, error = self._parse_output(output_path)
+            session.messages = messages
+            session.token_usage = usage
+
+            # Check if we got actual assistant responses
+            has_response = any(m.role == "assistant" for m in messages)
+
+            if error and not has_response:
+                # Only mark as failed if we have an error AND no response
+                session.status = SessionStatus.FAILED
+                session.error = error
+            elif error and has_response:
+                # Got response but also an error (e.g., network disconnect at end)
+                # Treat as completed but note the error
+                session.status = SessionStatus.COMPLETED
+                session.error = f"Completed with error: {error}"
+            else:
+                session.status = SessionStatus.COMPLETED
+        else:
+            session.status = SessionStatus.FAILED
+            session.error = "No output file found"
+
+        self._save_session(session)
+
+    def start_session(
+        self,
+        task: str,
+        working_dir: Path | None = None,
+        model: str = "gpt-5.1-codex-mini",
+        sandbox: str = "workspace-write",
+        source: str = "user",
+        adapter: str = "codex",
+    ) -> CodexSession:
+        """
+        Start a new Codex session in the background.
+
+        Args:
+            task: The task description
+            working_dir: Working directory for codex (default: cwd)
+            model: Model to use
+            sandbox: Sandbox mode
+            source: Who spawned this session ("user" or "orchestrator:<id>")
+            adapter: Which adapter to use ("codex", "claude_code")
+
+        Returns:
+            The created session
+        """
+        session_id = str(uuid4())
+        working_dir = working_dir or Path.cwd()
+        now = datetime.now().isoformat()
+
+        session = CodexSession(
+            id=session_id,
+            task=task,
+            status=SessionStatus.PENDING,
+            working_dir=working_dir,
+            created_at=now,
+            updated_at=now,
+            model=model,
+            turn=1,
+            messages=[SessionMessage(role="user", content=task, timestamp=now)],
+            source=source,
+            adapter=adapter,
+        )
+
+        # Create session directory
+        session_dir = self._session_dir(session_id)
+        session_dir.mkdir(parents=True, exist_ok=True)
+
+        # Build command
+        cmd = [
+            "codex", "exec",
+            "--json",
+            "--full-auto",
+            "--skip-git-repo-check",
+            "--model", model,
+            "-C", str(working_dir.absolute()),
+        ]
+
+        # Add sandbox mode
+        if sandbox == "danger-full-access":
+            cmd.append("--dangerously-bypass-approvals-and-sandbox")
+        elif sandbox == "workspace-write":
+            # Default codex behavior
+            pass
+
+        cmd.extend(["--", task])
+
+        # Start process with output redirected to file
+        output_path = self._output_path(session_id, 1)
+        output_file = open(output_path, "w")
+
+        proc = subprocess.Popen(
+            cmd,
+            cwd=working_dir,
+            stdout=output_file,
+            stderr=subprocess.STDOUT,
+            start_new_session=True,  # Detach from parent process group
+        )
+
+        session.pid = proc.pid
+        session.status = SessionStatus.RUNNING
+        self._save_session(session)
+
+        return session
+
+    def inject_message(
+        self,
+        session_id: str,
+        message: str,
+    ) -> CodexSession | None:
+        """
+        Inject a follow-up message into a completed session.
+
+        This starts a new turn with the conversation context.
+
+        Args:
+            session_id: Session to continue
+            message: The follow-up message
+
+        Returns:
+            Updated session or None if not found/not ready
+        """
+        session = self.get_session(session_id)
+        if not session:
+            return None
+
+        if session.status == SessionStatus.RUNNING:
+            # Can't inject while running - would need to implement
+            # a more complex IPC mechanism for that
+            return None
+
+        # Build context from previous messages
+        context_parts = []
+        for msg in session.messages:
+            if msg.role == "user":
+                context_parts.append(f"USER: {msg.content}")
+            elif msg.role == "assistant":
+                context_parts.append(f"ASSISTANT: {msg.content}")
+
+        # Create augmented prompt with context
+        augmented_task = f"""Continue the following conversation:
+
+{chr(10).join(context_parts)}
+
+USER: {message}
+
+Continue from where you left off, addressing the user's new message."""
+
+        # Start new turn
+        session.turn += 1
+        now = datetime.now().isoformat()
+        session.messages.append(SessionMessage(role="user", content=message, timestamp=now))
+
+        # Build command
+        cmd = [
+            "codex", "exec",
+            "--json",
+            "--full-auto",
+            "--skip-git-repo-check",
+            "--model", session.model,
+            "-C", str(session.working_dir.absolute()),
+            "--", augmented_task,
+        ]
+
+        # Start process
+        output_path = self._output_path(session.id, session.turn)
+        output_file = open(output_path, "w")
+
+        proc = subprocess.Popen(
+            cmd,
+            cwd=session.working_dir,
+            stdout=output_file,
+            stderr=subprocess.STDOUT,
+            start_new_session=True,
+        )
+
+        session.pid = proc.pid
+        session.status = SessionStatus.RUNNING
+        self._save_session(session)
+
+        return session
+
+    def kill_session(self, session_id: str, delete: bool = False) -> bool:
+        """
+        Kill a running session.
+
+        Args:
+            session_id: Session to kill
+            delete: If True, also delete session data entirely
+
+        Returns True if killed, False if not found or not running.
+        """
+        session = self.get_session(session_id)
+        if not session:
+            return False
+
+        if session.pid and session.is_running:
+            try:
+                # Kill the entire process group
+                os.killpg(os.getpgid(session.pid), signal.SIGTERM)
+                time.sleep(0.5)
+
+                # Force kill if still running
+                if session.is_running:
+                    os.killpg(os.getpgid(session.pid), signal.SIGKILL)
+            except (OSError, ProcessLookupError):
+                pass
+
+        if delete:
+            return self.delete_session(session.id)
+
+        session.status = SessionStatus.KILLED
+        session.error = "Manually killed"
+        self._save_session(session)
+        return True
+
+    def delete_session(self, session_id: str) -> bool:
+        """
+        Delete a session entirely (removes from disk).
+
+        Kills the process first if still running.
+
+        Returns True if deleted, False if not found.
+        """
+        import shutil
+
+        session = self.get_session(session_id)
+        if not session:
+            return False
+
+        # Kill if running
+        if session.pid and session.is_running:
+            try:
+                os.killpg(os.getpgid(session.pid), signal.SIGTERM)
+                time.sleep(0.3)
+                if session.is_running:
+                    os.killpg(os.getpgid(session.pid), signal.SIGKILL)
+            except (OSError, ProcessLookupError):
+                pass
+
+        # Remove session directory
+        session_dir = self._session_dir(session.id)
+        if session_dir.exists():
+            shutil.rmtree(session_dir)
+            return True
+
+        return False
+
+    def get_output(self, session_id: str, turn: int | None = None) -> str:
+        """Get raw JSONL output for a session."""
+        session = self.get_session(session_id)
+        if not session:
+            return ""
+
+        if turn is None:
+            turn = session.turn
+
+        output_path = self._output_path(session.id, turn)
+        if not output_path.exists():
+            return ""
+
+        return output_path.read_text()
+
+    def get_messages(self, session_id: str) -> list[SessionMessage]:
+        """Get parsed messages for a session across all turns."""
+        session = self.get_session(session_id)
+        if not session:
+            return []
+
+        all_messages = []
+
+        # Get messages from each turn
+        for turn in range(1, session.turn + 1):
+            output_path = self._output_path(session.id, turn)
+            if output_path.exists():
+                messages, _, _ = self._parse_output(output_path)
+                all_messages.extend(messages)
+
+        return all_messages
+
+    def _parse_output(self, output_path: Path) -> tuple[list[SessionMessage], dict[str, int], str | None]:
+        """
+        Parse JSONL output from codex exec.
+
+        Returns:
+            (messages, token_usage, error)
+        """
+        messages: list[SessionMessage] = []
+        usage: dict[str, int] = {}
+        error: str | None = None
+
+        if not output_path.exists():
+            return messages, usage, "Output file not found"
+
+        content = output_path.read_text()
+
+        for line in content.strip().split("\n"):
+            if not line.strip():
+                continue
+
+            try:
+                event = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+
+            event_type = event.get("type", "")
+
+            if event_type == "item.completed":
+                item = event.get("item", {})
+                item_type = item.get("type", "")
+
+                if item_type == "agent_message":
+                    text = item.get("text", "")
+                    if text:
+                        messages.append(SessionMessage(
+                            role="assistant",
+                            content=text,
+                            timestamp=datetime.now().isoformat(),
+                        ))
+
+                elif item_type == "reasoning":
+                    # Could optionally capture reasoning
+                    pass
+
+                elif item_type == "function_call":
+                    # Track tool calls
+                    func_name = item.get("name", "unknown")
+                    messages.append(SessionMessage(
+                        role="tool",
+                        content=f"[Calling: {func_name}]",
+                        metadata={"function": func_name},
+                    ))
+
+                elif item_type == "function_call_output":
+                    output = item.get("output", "")
+                    if output and len(output) < 500:
+                        messages.append(SessionMessage(
+                            role="tool",
+                            content=f"[Output]: {output[:500]}",
+                        ))
+
+            elif event_type == "turn.completed":
+                turn_usage = event.get("usage", {})
+                for key, value in turn_usage.items():
+                    usage[key] = usage.get(key, 0) + value
+                # Compute total_tokens if not present
+                if "total_tokens" not in usage:
+                    usage["total_tokens"] = usage.get("input_tokens", 0) + usage.get("output_tokens", 0)
+
+            elif event_type == "error":
+                error = event.get("message", str(event))
+
+        return messages, usage, error
+
+    def get_trajectory(self, session_id: str, full: bool = False, max_output_len: int = 200) -> list[dict]:
+        """
+        Get the full trajectory of a session - all steps in order.
+
+        Args:
+            session_id: Session to get trajectory for
+            full: If True, include full untruncated content
+            max_output_len: Max length for outputs when full=False
+
+        Returns a list of step dicts with type, summary, and details.
+        This shows the "broad strokes" of what the agent did.
+        """
+        if full:
+            max_output_len = 999999  # Effectively unlimited
+        session = self.get_session(session_id)
+        if not session:
+            return []
+
+        trajectory = []
+
+        for turn in range(1, session.turn + 1):
+            output_path = self._output_path(session.id, turn)
+            if not output_path.exists():
+                continue
+
+            content = output_path.read_text()
+            step_num = 0
+
+            for line in content.strip().split("\n"):
+                if not line.strip():
+                    continue
+
+                try:
+                    event = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+
+                event_type = event.get("type", "")
+
+                if event_type == "item.completed":
+                    item = event.get("item", {})
+                    item_type = item.get("type", "")
+                    step_num += 1
+
+                    if item_type == "reasoning":
+                        text = item.get("text", "")
+                        summary_len = max_output_len if full else 100
+                        trajectory.append({
+                            "turn": turn,
+                            "step": step_num,
+                            "type": "reasoning",
+                            "summary": text[:summary_len] + ("..." if len(text) > summary_len else ""),
+                            "full_text": text if full else None,
+                        })
+
+                    elif item_type == "command_execution":
+                        cmd = item.get("command", "")
+                        output = item.get("aggregated_output", "")
+                        exit_code = item.get("exit_code")
+                        # Truncate output
+                        output_preview = output[:max_output_len]
+                        if len(output) > max_output_len:
+                            output_preview += "..."
+                        trajectory.append({
+                            "turn": turn,
+                            "step": step_num,
+                            "type": "command",
+                            "command": cmd,
+                            "output": output_preview.strip(),
+                            "exit_code": exit_code,
+                        })
+
+                    elif item_type == "function_call":
+                        func_name = item.get("name", "unknown")
+                        args = item.get("arguments", {})
+                        args_str = str(args)
+                        args_len = max_output_len if full else 100
+                        trajectory.append({
+                            "turn": turn,
+                            "step": step_num,
+                            "type": "tool_call",
+                            "tool": func_name,
+                            "args_preview": args_str[:args_len] + ("..." if len(args_str) > args_len else ""),
+                            "full_args": args if full else None,
+                        })
+
+                    elif item_type == "function_call_output":
+                        output = item.get("output", "")
+                        output_preview = output[:max_output_len]
+                        if len(output) > max_output_len:
+                            output_preview += "..."
+                        trajectory.append({
+                            "turn": turn,
+                            "step": step_num,
+                            "type": "tool_output",
+                            "output": output_preview,
+                        })
+
+                    elif item_type == "agent_message":
+                        text = item.get("text", "")
+                        summary_len = max_output_len if full else 200
+                        trajectory.append({
+                            "turn": turn,
+                            "step": step_num,
+                            "type": "message",
+                            "summary": text[:summary_len] + ("..." if len(text) > summary_len else ""),
+                            "full_text": text if full else None,
+                            "full_length": len(text),
+                        })
+
+        return trajectory
+
+    def cleanup_completed(self, keep_days: int = 7) -> int:
+        """
+        Remove old completed/failed/killed sessions.
+
+        Args:
+            keep_days: Keep sessions newer than this many days
+
+        Returns:
+            Number of sessions cleaned up
+        """
+        import shutil
+        from datetime import timedelta
+
+        cutoff = datetime.now() - timedelta(days=keep_days)
+        cleaned = 0
+
+        for session in self.list_sessions():
+            if session.status in (SessionStatus.COMPLETED, SessionStatus.FAILED, SessionStatus.KILLED):
+                created = datetime.fromisoformat(session.created_at)
+                if created < cutoff:
+                    session_dir = self._session_dir(session.id)
+                    if session_dir.exists():
+                        shutil.rmtree(session_dir)
+                        cleaned += 1
+
+        return cleaned