zwarm 3.4.0__py3-none-any.whl → 3.7.0__py3-none-any.whl

zwarm/sessions/manager.py

@@ -1,189 +1,54 @@
 """
 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
-- Config is read from .zwarm/codex.toml and passed via -c overrides
+This module implements the CodexSessionManager, which handles:
+- Spawning `codex exec --json` subprocesses
+- Parsing Codex's JSONL output format
+- Loading config from .zwarm/codex.toml
+
+Inherits shared functionality from BaseSessionManager.
 """
 
 from __future__ import annotations
 
 import json
-import os
-import signal
 import subprocess
-import time
 import tomllib
-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
 
+from .base import (
+    BaseSessionManager,
+    Session,
+    SessionMessage,
+    SessionStatus,
+)
 
-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"),
-        )
+# Re-export for backwards compatibility
+CodexSession = Session
 
-    @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:
+class CodexSessionManager(BaseSessionManager):
     """
     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)
+    adapter_name = "codex"
+    default_model = "gpt-5.1-codex-mini"
+
+    # =========================================================================
+    # Codex-specific config handling
+    # =========================================================================
 
     def _load_codex_config(self) -> dict[str, Any]:
         """
@@ -227,160 +92,9 @@ class CodexSessionManager:
 
         return overrides
 
-    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)
+    # =========================================================================
+    # Session lifecycle (Codex-specific implementation)
+    # =========================================================================
 
     def start_session(
         self,
@@ -389,8 +103,7 @@ class CodexSessionManager:
         model: str | None = None,
         sandbox: str = "workspace-write",
         source: str = "user",
-        adapter: str = "codex",
-    ) -> CodexSession:
+    ) -> Session:
         """
         Start a new Codex session in the background.
 
@@ -400,7 +113,6 @@ class CodexSessionManager:
             model: Model override (default: from codex.toml or gpt-5.1-codex-mini)
             sandbox: Sandbox mode (ignored if full_danger=true in codex.toml)
             source: Who spawned this session ("user" or "orchestrator:<id>")
-            adapter: Which adapter to use ("codex", "claude_code")
 
         Returns:
             The created session
@@ -417,12 +129,12 @@ class CodexSessionManager:
         codex_config = self._load_codex_config()
 
         # Get model from config or use default
-        effective_model = model or codex_config.get("model", "gpt-5.1-codex-mini")
+        effective_model = model or codex_config.get("model", self.default_model)
 
         # Check if full_danger mode is enabled
         full_danger = codex_config.get("full_danger", False)
 
-        session = CodexSession(
+        session = Session(
             id=session_id,
             task=task,
             status=SessionStatus.PENDING,
@@ -433,7 +145,7 @@ class CodexSessionManager:
             turn=1,
             messages=[SessionMessage(role="user", content=task, timestamp=now)],
             source=source,
-            adapter=adapter,
+            adapter=self.adapter_name,
         )
 
         # Create session directory
@@ -441,7 +153,6 @@ class CodexSessionManager:
         session_dir.mkdir(parents=True, exist_ok=True)
 
         # Build command with -c overrides from codex.toml
-        # This ensures each .zwarm dir has its own codex config
         cmd = ["codex"]
 
         # Add -c overrides from codex.toml (excluding special keys we handle separately)
@@ -487,7 +198,7 @@ class CodexSessionManager:
         self,
         session_id: str,
         message: str,
-    ) -> CodexSession | None:
+    ) -> Session | None:
         """
         Inject a follow-up message into a completed session.
 
@@ -505,8 +216,6 @@ class CodexSessionManager:
             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
@@ -533,25 +242,31 @@ Continue from where you left off, addressing the user's new message."""
             SessionMessage(role="user", content=message, timestamp=now)
         )
 
-        # Build command
-        # Note: --search is a global flag that must come before 'exec'
+        # Build command with -c overrides from codex.toml
+        codex_config = self._load_codex_config()
+        full_danger = codex_config.get("full_danger", False)
+
         cmd = ["codex"]
-        if session.web_search:
-            cmd.append("--search")
-        cmd.extend(
-            [
-                "exec",
-                "--json",
-                "--full-auto",
-                "--skip-git-repo-check",
-                "--model",
-                session.model,
-                "-C",
-                str(session.working_dir.absolute()),
-                "--",
-                augmented_task,
-            ]
-        )
+
+        # Add -c overrides from codex.toml (excluding special keys)
+        config_for_overrides = {k: v for k, v in codex_config.items() if k not in ("model", "full_danger")}
+        cmd.extend(self._build_codex_overrides(config_for_overrides))
+
+        cmd.extend([
+            "exec",
+            "--json",
+            "--skip-git-repo-check",
+            "--model",
+            session.model,
+            "-C",
+            str(session.working_dir.absolute()),
+        ])
+
+        # Full danger mode bypasses all safety controls
+        if full_danger:
+            cmd.append("--dangerously-bypass-approvals-and-sandbox")
+
+        cmd.extend(["--", augmented_task])
 
         # Start process
         output_path = self._output_path(session.id, session.turn)
@@ -571,104 +286,47 @@ Continue from where you left off, addressing the user's new message."""
 
         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
+    # =========================================================================
+    # Output parsing (Codex-specific JSONL format)
+    # =========================================================================
 
-    def delete_session(self, session_id: str) -> bool:
+    def _is_output_complete(self, session_id: str, turn: int) -> bool:
         """
-        Delete a session entirely (removes from disk).
-
-        Kills the process first if still running.
+        Check if output file indicates the task completed.
 
-        Returns True if deleted, False if not found.
+        Looks for completion markers like 'turn.completed' or 'task.completed'
+        in the JSONL output. This is more reliable than PID checking.
         """
-        import shutil
-
-        session = self.get_session(session_id)
-        if not session:
+        output_path = self._output_path(session_id, turn)
+        if not output_path.exists():
             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
+        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 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]:
@@ -879,34 +537,3 @@ Continue from where you left off, addressing the user's new message."""
                         )
 
         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