gobby 0.2.5__py3-none-any.whl

gobby/sessions/transcripts/gemini.py

@@ -0,0 +1,195 @@
+"""
+Gemini transcript parser.
+
+Parses JSONL transcript files generated by Gemini CLI.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import UTC, datetime
+from typing import Any
+
+from gobby.sessions.transcripts.base import ParsedMessage, TokenUsage
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiTranscriptParser:
+    """
+    Parses JSONL transcript files from Gemini.
+
+    Implements the TranscriptParser protocol for Gemini's transcript format.
+    Assumes a standard JSONL structure where each line is a message or event.
+    """
+
+    def __init__(self, logger_instance: logging.Logger | None = None):
+        """
+        Initialize GeminiTranscriptParser.
+
+        Args:
+            logger_instance: Optional logger instance.
+        """
+        self.logger = logger_instance or logger
+
+    def extract_last_messages(
+        self, turns: list[dict[str, Any]], num_pairs: int = 2
+    ) -> list[dict[str, Any]]:
+        """
+        Extract last N user<>agent message pairs.
+        """
+        messages: list[dict[str, str]] = []
+        for turn in reversed(turns):
+            # Adapt to generic turn structure
+            # Assumes generic schema: {"role": "...", "content": "..."} or nested in "message"
+            role = turn.get("role") or turn.get("message", {}).get("role")
+            content = turn.get("content") or turn.get("message", {}).get("content")
+
+            if role in ["user", "model", "assistant"]:
+                norm_role = "assistant" if role == "model" else role
+
+                # Handle complex content types if necessary
+                if isinstance(content, list):
+                    content = " ".join(str(part) for part in content)
+
+                messages.insert(0, {"role": norm_role, "content": str(content)})
+                if len(messages) >= num_pairs * 2:
+                    break
+        return messages
+
+    def extract_turns_since_clear(
+        self, turns: list[dict[str, Any]], max_turns: int = 50
+    ) -> list[dict[str, Any]]:
+        """
+        Extract turns since the most recent session boundary.
+        For Gemini, we might look for specific clear events or just return the tail.
+        """
+        # Placeholder: just return last N turns for now until we know the clear signal
+        return turns[-max_turns:] if len(turns) > max_turns else turns
+
+    def is_session_boundary(self, turn: dict[str, Any]) -> bool:
+        """
+        Check if a turn is a session boundary.
+        """
+        # Placeholder for Gemini specific boundary
+        return False
+
+    def parse_line(self, line: str, index: int) -> ParsedMessage | None:
+        """
+        Parse a single line from the transcript JSONL.
+        """
+        if not line.strip():
+            return None
+
+        try:
+            data = json.loads(line)
+        except json.JSONDecodeError:
+            self.logger.warning(f"Invalid JSON at line {index}")
+            return None
+
+        # Extract timestamp
+        timestamp_str = data.get("timestamp") or datetime.now(UTC).isoformat()
+        try:
+            timestamp = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
+        except ValueError:
+            timestamp = datetime.now(UTC)
+
+        # Determine role and content
+        # Check top-level or nested 'message'
+        role = data.get("role")
+        content = data.get("content")
+
+        if not role and "message" in data:
+            msg = data["message"]
+            role = msg.get("role")
+            content = msg.get("content")
+
+        if not role:
+            # Try type field common in other schemas
+            msg_type = data.get("type")
+            if msg_type == "user":
+                role = "user"
+            elif msg_type == "model":
+                role = "assistant"
+
+        # Normalize role
+        if role == "model":
+            role = "assistant"
+
+        if not role:
+            # Maybe a tool result or system event
+            if "tool_result" in data:
+                role = "tool"
+                content = str(data["tool_result"])
+            else:
+                # Unknown or uninteresting line
+                return None
+
+        # Normalize content
+        content_type = "text"
+        tool_name = None
+        tool_input = None
+        tool_result = None
+
+        if isinstance(content, list):
+            # Handle potential rich content
+            text_parts: list[str] = []
+            for part in content:
+                if isinstance(part, str):
+                    text_parts.append(part)
+                elif isinstance(part, dict):
+                    if "text" in part:
+                        text_parts.append(str(part["text"]))
+                    # Check for tool calls
+                    if "functionCall" in part:
+                        content_type = "tool_use"
+                        tool_name = part["functionCall"].get("name")
+                        tool_input = part["functionCall"].get("args")
+            content = " ".join(text_parts)
+        else:
+            content = str(content or "")
+
+        return ParsedMessage(
+            index=index,
+            role=role,
+            content=content,
+            content_type=content_type,
+            tool_name=tool_name,
+            tool_input=tool_input,
+            tool_result=tool_result,
+            timestamp=timestamp,
+            raw_json=data,
+            usage=self._extract_usage(data),
+        )
+
+    def _extract_usage(self, data: dict[str, Any]) -> TokenUsage | None:
+        """Extract token usage from Gemini message data."""
+        # Gemini API standard is usageMetadata
+        usage_data = data.get("usageMetadata")
+
+        if not usage_data:
+            return None
+
+        return TokenUsage(
+            input_tokens=usage_data.get("promptTokenCount", 0),
+            output_tokens=usage_data.get("candidatesTokenCount", 0),
+            # Gemini doesn't always split cache tokens in this view, strictly speaking
+            # but usually total is prompt + candidates
+            total_cost_usd=None,  # Cost calculation not standard in CLI output
+        )
+
+    def parse_lines(self, lines: list[str], start_index: int = 0) -> list[ParsedMessage]:
+        """
+        Parse a list of transcript lines.
+        """
+        parsed_messages = []
+        current_index = start_index
+
+        for line in lines:
+            message = self.parse_line(line, current_index)
+            if message:
+                parsed_messages.append(message)
+                current_index += 1
+
+        return parsed_messages

gobby/storage/__init__.py

@@ -0,0 +1,21 @@
+"""Local storage layer for Gobby daemon."""
+
+from gobby.storage.database import LocalDatabase
+from gobby.storage.inter_session_messages import InterSessionMessageManager
+from gobby.storage.mcp import LocalMCPManager
+from gobby.storage.migrations import run_migrations
+from gobby.storage.projects import LocalProjectManager
+from gobby.storage.sessions import LocalSessionManager
+from gobby.storage.task_dependencies import TaskDependencyManager
+from gobby.storage.tasks import LocalTaskManager
+
+__all__ = [
+    "InterSessionMessageManager",
+    "LocalDatabase",
+    "LocalMCPManager",
+    "LocalProjectManager",
+    "LocalSessionManager",
+    "LocalTaskManager",
+    "TaskDependencyManager",
+    "run_migrations",
+]

gobby/storage/agents.py

@@ -0,0 +1,409 @@
+"""Storage manager for agent runs."""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Any, Literal
+
+from gobby.storage.database import DatabaseProtocol
+
+logger = logging.getLogger(__name__)
+
+AgentRunStatus = Literal["pending", "running", "success", "error", "timeout", "cancelled"]
+
+
+@dataclass
+class AgentRun:
+    """Agent run data model."""
+
+    id: str
+    parent_session_id: str
+    provider: str
+    prompt: str
+    status: AgentRunStatus
+    created_at: str
+    updated_at: str
+    # Optional fields
+    child_session_id: str | None = None
+    workflow_name: str | None = None
+    model: str | None = None
+    result: str | None = None
+    error: str | None = None
+    tool_calls_count: int = 0
+    turns_used: int = 0
+    started_at: str | None = None
+    completed_at: str | None = None
+
+    @classmethod
+    def from_row(cls, row: Any) -> AgentRun:
+        """Create AgentRun from database row."""
+        return cls(
+            id=row["id"],
+            parent_session_id=row["parent_session_id"],
+            child_session_id=row["child_session_id"],
+            workflow_name=row["workflow_name"],
+            provider=row["provider"],
+            model=row["model"],
+            status=row["status"],
+            prompt=row["prompt"],
+            result=row["result"],
+            error=row["error"],
+            tool_calls_count=row["tool_calls_count"] or 0,
+            turns_used=row["turns_used"] or 0,
+            started_at=row["started_at"],
+            completed_at=row["completed_at"],
+            created_at=row["created_at"],
+            updated_at=row["updated_at"],
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "id": self.id,
+            "parent_session_id": self.parent_session_id,
+            "child_session_id": self.child_session_id,
+            "workflow_name": self.workflow_name,
+            "provider": self.provider,
+            "model": self.model,
+            "status": self.status,
+            "prompt": self.prompt,
+            "result": self.result,
+            "error": self.error,
+            "tool_calls_count": self.tool_calls_count,
+            "turns_used": self.turns_used,
+            "started_at": self.started_at,
+            "completed_at": self.completed_at,
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+        }
+
+
+class LocalAgentRunManager:
+    """Manager for agent run storage operations."""
+
+    def __init__(self, db: DatabaseProtocol):
+        """Initialize with database connection."""
+        self.db = db
+
+    def create(
+        self,
+        parent_session_id: str,
+        provider: str,
+        prompt: str,
+        workflow_name: str | None = None,
+        model: str | None = None,
+        child_session_id: str | None = None,
+    ) -> AgentRun:
+        """
+        Create a new agent run.
+
+        Args:
+            parent_session_id: Session that spawned this agent.
+            provider: LLM provider (claude, gemini, etc.)
+            prompt: The prompt given to the agent.
+            workflow_name: Optional workflow being executed.
+            model: Optional model override.
+            child_session_id: Optional child session for the agent.
+
+        Returns:
+            Created AgentRun.
+        """
+        run_id = f"ar-{uuid.uuid4().hex[:12]}"
+        now = datetime.now(UTC).isoformat()
+
+        self.db.execute(
+            """
+            INSERT INTO agent_runs (
+                id, parent_session_id, child_session_id, workflow_name,
+                provider, model, status, prompt, created_at, updated_at
+            )
+            VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?)
+            """,
+            (
+                run_id,
+                parent_session_id,
+                child_session_id,
+                workflow_name,
+                provider,
+                model,
+                prompt,
+                now,
+                now,
+            ),
+        )
+
+        logger.debug(f"Created agent run {run_id} for session {parent_session_id}")
+        agent_run = self.get(run_id)
+        if agent_run is None:
+            raise RuntimeError(f"Failed to retrieve newly created agent run: {run_id}")
+        return agent_run
+
+    def get(self, run_id: str) -> AgentRun | None:
+        """Get agent run by ID."""
+        row = self.db.fetchone("SELECT * FROM agent_runs WHERE id = ?", (run_id,))
+        return AgentRun.from_row(row) if row else None
+
+    def start(self, run_id: str) -> AgentRun | None:
+        """Mark agent run as started."""
+        now = datetime.now(UTC).isoformat()
+        self.db.execute(
+            """
+            UPDATE agent_runs
+            SET status = 'running', started_at = ?, updated_at = ?
+            WHERE id = ?
+            """,
+            (now, now, run_id),
+        )
+        return self.get(run_id)
+
+    def complete(
+        self,
+        run_id: str,
+        result: str,
+        tool_calls_count: int = 0,
+        turns_used: int = 0,
+    ) -> AgentRun | None:
+        """
+        Mark agent run as completed successfully.
+
+        Args:
+            run_id: The agent run ID.
+            result: The agent's output/result.
+            tool_calls_count: Number of tool calls made.
+            turns_used: Number of turns used.
+
+        Returns:
+            Updated AgentRun.
+        """
+        now = datetime.now(UTC).isoformat()
+        self.db.execute(
+            """
+            UPDATE agent_runs
+            SET status = 'success',
+                result = ?,
+                tool_calls_count = ?,
+                turns_used = ?,
+                completed_at = ?,
+                updated_at = ?
+            WHERE id = ?
+            """,
+            (result, tool_calls_count, turns_used, now, now, run_id),
+        )
+        return self.get(run_id)
+
+    def fail(
+        self,
+        run_id: str,
+        error: str,
+        tool_calls_count: int = 0,
+        turns_used: int = 0,
+    ) -> AgentRun | None:
+        """
+        Mark agent run as failed.
+
+        Args:
+            run_id: The agent run ID.
+            error: Error message.
+            tool_calls_count: Number of tool calls made before failure.
+            turns_used: Number of turns used before failure.
+
+        Returns:
+            Updated AgentRun.
+        """
+        now = datetime.now(UTC).isoformat()
+        self.db.execute(
+            """
+            UPDATE agent_runs
+            SET status = 'error',
+                error = ?,
+                tool_calls_count = ?,
+                turns_used = ?,
+                completed_at = ?,
+                updated_at = ?
+            WHERE id = ?
+            """,
+            (error, tool_calls_count, turns_used, now, now, run_id),
+        )
+        return self.get(run_id)
+
+    def timeout(self, run_id: str, turns_used: int = 0) -> AgentRun | None:
+        """Mark agent run as timed out."""
+        now = datetime.now(UTC).isoformat()
+        self.db.execute(
+            """
+            UPDATE agent_runs
+            SET status = 'timeout',
+                error = 'Execution timed out',
+                turns_used = ?,
+                completed_at = ?,
+                updated_at = ?
+            WHERE id = ?
+            """,
+            (turns_used, now, now, run_id),
+        )
+        return self.get(run_id)
+
+    def cancel(self, run_id: str) -> AgentRun | None:
+        """Mark agent run as cancelled."""
+        now = datetime.now(UTC).isoformat()
+        self.db.execute(
+            """
+            UPDATE agent_runs
+            SET status = 'cancelled', completed_at = ?, updated_at = ?
+            WHERE id = ?
+            """,
+            (now, now, run_id),
+        )
+        return self.get(run_id)
+
+    def update_child_session(self, run_id: str, child_session_id: str) -> AgentRun | None:
+        """Update the child session ID for an agent run."""
+        now = datetime.now(UTC).isoformat()
+        self.db.execute(
+            """
+            UPDATE agent_runs
+            SET child_session_id = ?, updated_at = ?
+            WHERE id = ?
+            """,
+            (child_session_id, now, run_id),
+        )
+        return self.get(run_id)
+
+    def list_by_session(
+        self,
+        parent_session_id: str,
+        status: AgentRunStatus | None = None,
+        limit: int = 100,
+    ) -> list[AgentRun]:
+        """
+        List agent runs for a session.
+
+        Args:
+            parent_session_id: The parent session ID.
+            status: Optional status filter.
+            limit: Maximum number of results.
+
+        Returns:
+            List of AgentRun objects.
+        """
+        if status:
+            rows = self.db.fetchall(
+                """
+                SELECT * FROM agent_runs
+                WHERE parent_session_id = ? AND status = ?
+                ORDER BY created_at DESC
+                LIMIT ?
+                """,
+                (parent_session_id, status, limit),
+            )
+        else:
+            rows = self.db.fetchall(
+                """
+                SELECT * FROM agent_runs
+                WHERE parent_session_id = ?
+                ORDER BY created_at DESC
+                LIMIT ?
+                """,
+                (parent_session_id, limit),
+            )
+        return [AgentRun.from_row(row) for row in rows]
+
+    def list_running(self, limit: int = 100) -> list[AgentRun]:
+        """List all currently running agent runs."""
+        rows = self.db.fetchall(
+            """
+            SELECT * FROM agent_runs
+            WHERE status = 'running'
+            ORDER BY started_at ASC
+            LIMIT ?
+            """,
+            (limit,),
+        )
+        return [AgentRun.from_row(row) for row in rows]
+
+    def count_by_session(self, parent_session_id: str) -> dict[str, int]:
+        """
+        Count agent runs by status for a session.
+
+        Args:
+            parent_session_id: The parent session ID.
+
+        Returns:
+            Dict mapping status to count.
+        """
+        rows = self.db.fetchall(
+            """
+            SELECT status, COUNT(*) as count
+            FROM agent_runs
+            WHERE parent_session_id = ?
+            GROUP BY status
+            """,
+            (parent_session_id,),
+        )
+        return {row["status"]: row["count"] for row in rows}
+
+    def delete(self, run_id: str) -> bool:
+        """Delete an agent run."""
+        cursor = self.db.execute("DELETE FROM agent_runs WHERE id = ?", (run_id,))
+        return bool(cursor.rowcount and cursor.rowcount > 0)
+
+    def cleanup_stale_runs(self, timeout_minutes: int = 30) -> int:
+        """
+        Mark stale running agent runs as timed out.
+
+        Args:
+            timeout_minutes: Minutes of inactivity before timeout.
+
+        Returns:
+            Number of runs timed out.
+        """
+        now = datetime.now(UTC).isoformat()
+        cursor = self.db.execute(
+            """
+            UPDATE agent_runs
+            SET status = 'timeout',
+                error = 'Stale run timed out',
+                completed_at = ?,
+                updated_at = ?
+            WHERE status = 'running'
+            AND datetime(started_at) < datetime('now', 'utc', ? || ' minutes')
+            """,
+            (now, now, f"-{timeout_minutes}"),
+        )
+        count = cursor.rowcount or 0
+        if count > 0:
+            logger.info(f"Timed out {count} stale agent runs (>{timeout_minutes}m)")
+        return count
+
+    def cleanup_stale_pending_runs(self, timeout_minutes: int = 60) -> int:
+        """
+        Mark stale pending agent runs as failed.
+
+        Pending runs that never started within the timeout period are marked as errors.
+
+        Args:
+            timeout_minutes: Minutes since creation before marking as failed.
+
+        Returns:
+            Number of runs failed.
+        """
+        now = datetime.now(UTC).isoformat()
+        cursor = self.db.execute(
+            """
+            UPDATE agent_runs
+            SET status = 'error',
+                error = 'Pending run never started',
+                completed_at = ?,
+                updated_at = ?
+            WHERE status = 'pending'
+            AND datetime(created_at) < datetime('now', 'utc', ? || ' minutes')
+            """,
+            (now, now, f"-{timeout_minutes}"),
+        )
+        count = cursor.rowcount or 0
+        if count > 0:
+            logger.info(f"Failed {count} stale pending agent runs (>{timeout_minutes}m)")
+        return count