zwarm 3.2.1__py3-none-any.whl → 3.6.0__py3-none-any.whl

zwarm/sessions/manager.py

@@ -1,387 +1,178 @@
 """
 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
+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
-from dataclasses import dataclass, field
+import tomllib
 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)
-
-    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
+    adapter_name = "codex"
+    default_model = "gpt-5.1-codex-mini"
 
-        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
+    # =========================================================================
+    # Codex-specific config handling
+    # =========================================================================
 
-    def _is_output_complete(self, session_id: str, turn: int) -> bool:
+    def _load_codex_config(self) -> dict[str, Any]:
         """
-        Check if output file indicates the task completed.
+        Load codex.toml from state_dir.
 
-        Looks for completion markers like 'turn.completed' or 'task.completed'
-        in the JSONL output. This is more reliable than PID checking.
+        Returns parsed TOML as dict, or empty dict if not found.
         """
-        output_path = self._output_path(session_id, turn)
-        if not output_path.exists():
-            return False
-
+        codex_toml = self.state_dir / "codex.toml"
+        if not codex_toml.exists():
+            return {}
         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
+            with open(codex_toml, "rb") as f:
+                return tomllib.load(f)
         except Exception:
-            pass
+            return {}
 
-        return False
+    def _build_codex_overrides(self, config: dict[str, Any]) -> list[str]:
+        """
+        Convert codex.toml config to -c override flags.
 
-    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}"
+        Handles nested sections like [features] and [sandbox_workspace_write].
+
+        Returns list of ["-c", "key=value", "-c", "key2=value2", ...]
+        """
+        overrides = []
+
+        def add_override(key: str, value: Any):
+            """Add a -c override for a key=value pair."""
+            if isinstance(value, bool):
+                value = "true" if value else "false"
+            overrides.extend(["-c", f"{key}={value}"])
+
+        for key, value in config.items():
+            if isinstance(value, dict):
+                # Nested section like [features] or [sandbox_workspace_write]
+                for subkey, subvalue in value.items():
+                    add_override(f"{key}.{subkey}", subvalue)
             else:
-                session.status = SessionStatus.COMPLETED
-        else:
-            session.status = SessionStatus.FAILED
-            session.error = "No output file found"
+                # Top-level key
+                add_override(key, value)
 
-        self._save_session(session)
+        return overrides
+
+    # =========================================================================
+    # Session lifecycle (Codex-specific implementation)
+    # =========================================================================
 
     def start_session(
         self,
         task: str,
         working_dir: Path | None = None,
-        model: str = "gpt-5.1-codex-mini",
+        model: str | None = None,
         sandbox: str = "workspace-write",
         source: str = "user",
-        adapter: str = "codex",
-    ) -> CodexSession:
+    ) -> Session:
         """
         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
+            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
+
+        Note:
+            Settings are read from .zwarm/codex.toml and passed via -c overrides.
+            Run `zwarm init` to set up the config.
         """
         session_id = str(uuid4())
         working_dir = working_dir or Path.cwd()
         now = datetime.now().isoformat()
 
-        session = CodexSession(
+        # Load codex config from .zwarm/codex.toml
+        codex_config = self._load_codex_config()
+
+        # Get model from config or use default
+        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 = Session(
             id=session_id,
             task=task,
             status=SessionStatus.PENDING,
             working_dir=working_dir,
             created_at=now,
             updated_at=now,
-            model=model,
+            model=effective_model,
             turn=1,
             messages=[SessionMessage(role="user", content=task, timestamp=now)],
             source=source,
-            adapter=adapter,
+            adapter=self.adapter_name,
         )
 
         # Create session directory
         session_dir = self._session_dir(session_id)
         session_dir.mkdir(parents=True, exist_ok=True)
 
-        # Build command
-        cmd = [
-            "codex", "exec",
+        # Build command with -c overrides from codex.toml
+        cmd = ["codex"]
+
+        # Add -c overrides from codex.toml (excluding special keys we handle separately)
+        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))
+
+        # Add exec command and flags
+        cmd.extend([
+            "exec",
             "--json",
-            "--full-auto",
             "--skip-git-repo-check",
-            "--model", model,
-            "-C", str(working_dir.absolute()),
-        ]
-
-        # Add sandbox mode
-        if sandbox == "danger-full-access":
+            "--model",
+            effective_model,
+            "-C",
+            str(working_dir.absolute()),
+        ])
+
+        # Full danger mode bypasses all safety controls
+        if full_danger:
             cmd.append("--dangerously-bypass-approvals-and-sandbox")
-        elif sandbox == "workspace-write":
-            # Default codex behavior
-            pass
 
         cmd.extend(["--", task])
 
@@ -407,7 +198,7 @@ class CodexSessionManager:
         self,
         session_id: str,
         message: str,
-    ) -> CodexSession | None:
+    ) -> Session | None:
         """
         Inject a follow-up message into a completed session.
 
@@ -425,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
@@ -449,18 +238,35 @@ 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))
+        session.messages.append(
+            SessionMessage(role="user", content=message, timestamp=now)
+        )
 
-        # Build command
-        cmd = [
-            "codex", "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"]
+
+        # 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",
-            "--full-auto",
             "--skip-git-repo-check",
-            "--model", session.model,
-            "-C", str(session.working_dir.absolute()),
-            "--", augmented_task,
-        ]
+            "--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)
@@ -480,105 +286,50 @@ 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
+    # =========================================================================
+    # Output parsing (Codex-specific JSONL format)
+    # =========================================================================
 
-        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:
+    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]:
+    def _parse_output(
+        self, output_path: Path
+    ) -> tuple[list[SessionMessage], dict[str, int], str | None]:
         """
         Parse JSONL output from codex exec.
 
@@ -612,11 +363,13 @@ Continue from where you left off, addressing the user's new message."""
                 if item_type == "agent_message":
                     text = item.get("text", "")
                     if text:
-                        messages.append(SessionMessage(
-                            role="assistant",
-                            content=text,
-                            timestamp=datetime.now().isoformat(),
-                        ))
+                        messages.append(
+                            SessionMessage(
+                                role="assistant",
+                                content=text,
+                                timestamp=datetime.now().isoformat(),
+                            )
+                        )
 
                 elif item_type == "reasoning":
                     # Could optionally capture reasoning
@@ -625,19 +378,23 @@ Continue from where you left off, addressing the user's new message."""
                 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},
-                    ))
+                    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]}",
-                        ))
+                        messages.append(
+                            SessionMessage(
+                                role="tool",
+                                content=f"[Output]: {output[:500]}",
+                            )
+                        )
 
             elif event_type == "turn.completed":
                 turn_usage = event.get("usage", {})
@@ -645,14 +402,18 @@ Continue from where you left off, addressing the user's new message."""
                     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)
+                    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]:
+    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.
 
@@ -699,13 +460,16 @@ Continue from where you left off, addressing the user's new message."""
                     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,
-                        })
+                        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", "")
@@ -715,78 +479,61 @@ Continue from where you left off, addressing the user's new message."""
                         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,
-                        })
+                        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,
-                        })
+                        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,
-                        })
+                        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),
-                        })
+                        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