agentpool 2.1.9__py3-none-any.whl → 2.5.0__py3-none-any.whl

agentpool_storage/sql_provider/sql_provider.py

@@ -15,7 +15,12 @@ from agentpool.utils.now import get_now
 from agentpool.utils.parse_time import parse_time_period
 from agentpool_storage.base import StorageProvider
 from agentpool_storage.models import QueryFilters
-from agentpool_storage.sql_provider.models import CommandHistory, Conversation, Message
+from agentpool_storage.sql_provider.models import (
+    CommandHistory,
+    Conversation,
+    Message,
+    Project,
+)
 from agentpool_storage.sql_provider.utils import (
     build_message_query,
     format_conversation,
@@ -33,6 +38,7 @@ if TYPE_CHECKING:
 
     from agentpool.common_types import JsonValue
     from agentpool.messaging import ChatMessage
+    from agentpool.sessions.models import ProjectData
     from agentpool_config.session import SessionQuery
     from agentpool_config.storage import SQLStorageConfig
     from agentpool_storage.models import ConversationData, StatsFilters
@@ -120,10 +126,10 @@ class SQLModelProvider(StorageProvider):
         content: str,
         role: str,
         name: str | None = None,
+        parent_id: str | None = None,
         cost_info: TokenCost | None = None,
         model: str | None = None,
         response_time: float | None = None,
-        forwarded_from: list[str] | None = None,
         provider_name: str | None = None,
         provider_response_id: str | None = None,
         messages: str | None = None,
@@ -138,6 +144,7 @@ class SQLModelProvider(StorageProvider):
             msg = Message(
                 conversation_id=conversation_id,
                 id=message_id,
+                parent_id=parent_id,
                 content=content,
                 role=role,
                 name=name,
@@ -149,7 +156,6 @@ class SQLModelProvider(StorageProvider):
                 input_tokens=cost_info.token_usage.input_tokens if cost_info else None,
                 output_tokens=cost_info.token_usage.output_tokens if cost_info else None,
                 cost=float(cost_info.total_cost) if cost_info else None,
-                forwarded_from=forwarded_from,
                 provider_name=provider_name,
                 provider_response_id=provider_response_id,
                 messages=messages,
@@ -202,6 +208,141 @@ class SQLModelProvider(StorageProvider):
             )
             return result.scalar_one_or_none()
 
+    async def get_conversation_messages(
+        self,
+        conversation_id: str,
+        *,
+        include_ancestors: bool = False,
+    ) -> list[ChatMessage[str]]:
+        """Get all messages for a conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+            include_ancestors: If True, traverse parent_id chain to include
+                messages from ancestor conversations (for forked convos).
+
+        Returns:
+            List of messages ordered by timestamp.
+        """
+        async with AsyncSession(self.engine) as session:
+            # Get messages for this conversation
+            result = await session.execute(
+                select(Message)
+                .where(Message.conversation_id == conversation_id)
+                .order_by(Message.timestamp.asc())  # type: ignore
+            )
+            messages = [to_chat_message(m) for m in result.scalars().all()]
+
+            if not include_ancestors or not messages:
+                return messages
+
+            # Find the first message's parent_id to get ancestor chain
+            first_msg = messages[0]
+            if first_msg.parent_id:
+                ancestors = await self.get_message_ancestry(first_msg.parent_id)
+                return ancestors + messages
+
+            return messages
+
+    async def get_message(
+        self,
+        message_id: str,
+    ) -> ChatMessage[str] | None:
+        """Get a single message by ID."""
+        async with AsyncSession(self.engine) as session:
+            result = await session.execute(select(Message).where(Message.id == message_id))
+            msg = result.scalar_one_or_none()
+            return to_chat_message(msg) if msg else None
+
+    async def get_message_ancestry(
+        self,
+        message_id: str,
+    ) -> list[ChatMessage[str]]:
+        """Get the ancestry chain of a message.
+
+        Traverses parent_id chain to build full history.
+        """
+        ancestors: list[ChatMessage[str]] = []
+        current_id: str | None = message_id
+
+        async with AsyncSession(self.engine) as session:
+            while current_id:
+                result = await session.execute(select(Message).where(Message.id == current_id))
+                msg = result.scalar_one_or_none()
+                if not msg:
+                    break
+                ancestors.append(to_chat_message(msg))
+                current_id = msg.parent_id
+
+        # Reverse to get oldest first
+        ancestors.reverse()
+        return ancestors
+
+    async def fork_conversation(
+        self,
+        *,
+        source_conversation_id: str,
+        new_conversation_id: str,
+        fork_from_message_id: str | None = None,
+        new_agent_name: str | None = None,
+    ) -> str | None:
+        """Fork a conversation at a specific point.
+
+        Creates a new conversation record. The fork point message_id is returned
+        so callers can set it as parent_id for new messages.
+        """
+        async with AsyncSession(self.engine) as session:
+            # Get source conversation
+            result = await session.execute(
+                select(Conversation).where(Conversation.id == source_conversation_id)
+            )
+            source_conv = result.scalar_one_or_none()
+            if not source_conv:
+                msg = f"Source conversation not found: {source_conversation_id}"
+                raise ValueError(msg)
+
+            # Determine fork point
+            fork_point_id: str | None = None
+            if fork_from_message_id:
+                # Verify the message exists and belongs to the source conversation
+                msg_result = await session.execute(
+                    select(Message).where(
+                        Message.id == fork_from_message_id,
+                        Message.conversation_id == source_conversation_id,
+                    )
+                )
+                fork_msg = msg_result.scalar_one_or_none()
+                if not fork_msg:
+                    err = f"Message {fork_from_message_id} not found in conversation"
+                    raise ValueError(err)
+                fork_point_id = fork_from_message_id
+            else:
+                # Fork from the last message
+                msg_result = await session.execute(
+                    select(Message)
+                    .where(Message.conversation_id == source_conversation_id)
+                    .order_by(desc(Message.timestamp))
+                    .limit(1)
+                )
+                last_msg = msg_result.scalar_one_or_none()
+                if last_msg:
+                    fork_point_id = last_msg.id
+
+            # Create new conversation
+            agent_name = new_agent_name or source_conv.agent_name
+            new_conv = Conversation(
+                id=new_conversation_id,
+                agent_name=agent_name,
+                title=f"{source_conv.title or 'Conversation'} (fork)"
+                if source_conv.title
+                else None,
+                start_time=get_now(),
+            )
+            session.add(new_conv)
+            await session.commit()
+
+            return fork_point_id
+
     async def log_command(
         self,
         *,
@@ -442,3 +583,128 @@ class SQLModelProvider(StorageProvider):
         msg_count = len(msg_result.scalars().all())
 
         return conv_count, msg_count
+
+    async def delete_conversation_messages(
+        self,
+        conversation_id: str,
+    ) -> int:
+        """Delete all messages for a conversation."""
+        from sqlalchemy import delete, func
+
+        async with AsyncSession(self.engine) as session:
+            # First count messages to return
+            count_result = await session.execute(
+                select(func.count()).where(Message.conversation_id == conversation_id)
+            )
+            count = count_result.scalar() or 0
+
+            # Then delete
+            await session.execute(
+                delete(Message).where(Message.conversation_id == conversation_id)  # type: ignore[arg-type]
+            )
+            await session.commit()
+            return count
+
+    # Project methods
+
+    def _to_project_data(self, row: Project) -> ProjectData:
+        """Convert database model to ProjectData."""
+        from agentpool.sessions.models import ProjectData
+
+        return ProjectData(
+            project_id=row.project_id,
+            worktree=row.worktree,
+            name=row.name,
+            vcs=row.vcs,
+            config_path=row.config_path,
+            created_at=row.created_at,
+            last_active=row.last_active,
+            settings=row.settings_json or {},
+        )
+
+    def _to_project_model(self, data: ProjectData) -> Project:
+        """Convert ProjectData to database model."""
+        return Project(
+            project_id=data.project_id,
+            worktree=data.worktree,
+            name=data.name,
+            vcs=data.vcs,
+            config_path=data.config_path,
+            created_at=data.created_at,
+            last_active=data.last_active,
+            settings_json=data.settings,
+        )
+
+    async def save_project(self, project: ProjectData) -> None:
+        """Save or update a project."""
+        from sqlalchemy import delete
+
+        async with AsyncSession(self.engine) as session:
+            # Delete existing if present (upsert via delete+insert)
+            stmt = delete(Project).where(Project.project_id == project.project_id)  # type: ignore[arg-type]
+            await session.execute(stmt)
+
+            # Insert new/updated
+            db_project = self._to_project_model(project)
+            session.add(db_project)
+            await session.commit()
+            logger.debug("Saved project", project_id=project.project_id)
+
+    async def get_project(self, project_id: str) -> ProjectData | None:
+        """Get a project by ID."""
+        async with AsyncSession(self.engine) as session:
+            stmt = select(Project).where(Project.project_id == project_id)
+            result = await session.execute(stmt)
+            row = result.scalars().first()
+            return self._to_project_data(row) if row else None
+
+    async def get_project_by_worktree(self, worktree: str) -> ProjectData | None:
+        """Get a project by worktree path."""
+        async with AsyncSession(self.engine) as session:
+            stmt = select(Project).where(Project.worktree == worktree)
+            result = await session.execute(stmt)
+            row = result.scalars().first()
+            return self._to_project_data(row) if row else None
+
+    async def get_project_by_name(self, name: str) -> ProjectData | None:
+        """Get a project by friendly name."""
+        async with AsyncSession(self.engine) as session:
+            stmt = select(Project).where(Project.name == name)
+            result = await session.execute(stmt)
+            row = result.scalars().first()
+            return self._to_project_data(row) if row else None
+
+    async def list_projects(self, limit: int | None = None) -> list[ProjectData]:
+        """List all projects, ordered by last_active descending."""
+        async with AsyncSession(self.engine) as session:
+            stmt = select(Project).order_by(desc(Project.last_active))
+            if limit is not None:
+                stmt = stmt.limit(limit)
+            result = await session.execute(stmt)
+            return [self._to_project_data(row) for row in result.scalars().all()]
+
+    async def delete_project(self, project_id: str) -> bool:
+        """Delete a project."""
+        from sqlalchemy import delete
+
+        async with AsyncSession(self.engine) as session:
+            stmt = delete(Project).where(Project.project_id == project_id)  # type: ignore[arg-type]
+            result = await session.execute(stmt)
+            await session.commit()
+            deleted: bool = result.rowcount > 0  # type: ignore[attr-defined]
+            if deleted:
+                logger.debug("Deleted project", project_id=project_id)
+            return deleted
+
+    async def touch_project(self, project_id: str) -> None:
+        """Update project's last_active timestamp."""
+        from sqlalchemy import update
+
+        async with AsyncSession(self.engine) as session:
+            stmt = (
+                update(Project)
+                .where(Project.project_id == project_id)  # type: ignore[arg-type]
+                .values(last_active=get_now())
+            )
+            await session.execute(stmt)
+            await session.commit()

agentpool_storage/sql_provider/utils.py

@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
+import contextlib
 from datetime import datetime
 from decimal import Decimal
 from typing import TYPE_CHECKING, Any
 
 from pydantic_ai import RunUsage
-from sqlalchemy import JSON, Column, and_, or_
-from sqlalchemy.sql import expression
+from sqlalchemy import Column, and_
 from sqlmodel import select
 
 from agentpool.messaging import ChatMessage, TokenCost
@@ -67,7 +67,6 @@ def to_chat_message(db_message: Message) -> ChatMessage[str]:
         model_name=db_message.model,
         cost_info=cost_info,
         response_time=db_message.response_time,
-        forwarded_from=db_message.forwarded_from or [],
         timestamp=db_message.timestamp,
         provider_name=db_message.provider_name,
         provider_response_id=db_message.provider_response_id,
@@ -101,9 +100,14 @@ def auto_migrate_columns(sync_conn: Any, dialect: Any) -> None:
     from sqlmodel import SQLModel
 
     inspector = inspect(sync_conn)
+    existing_tables = set(inspector.get_table_names())
 
     # For each table in our models
     for table_name, table in SQLModel.metadata.tables.items():
+        # Skip tables that don't exist yet (they'll be created fresh)
+        if table_name not in existing_tables:
+            continue
+
         existing = {col["name"] for col in inspector.get_columns(table_name)}
 
         # For each column in model that doesn't exist in DB
@@ -116,7 +120,9 @@ def auto_migrate_columns(sync_conn: Any, dialect: Any) -> None:
                 sql = (
                     f"ALTER TABLE {table_name} ADD COLUMN {col.name} {type_sql}{nullable}{default}"
                 )
-                sync_conn.execute(text(sql))
+                # Column may already exist (race condition or stale inspector cache)
+                with contextlib.suppress(Exception):
+                    sync_conn.execute(text(sql))
 
 
 def parse_model_info(model: str | None) -> tuple[str | None, str | None]:
@@ -160,15 +166,7 @@ def build_message_query(query: SessionQuery) -> SelectOfScalar[Any]:
     if query.name:
         conditions.append(Message.conversation_id == query.name)
     if query.agents:
-        agent_conditions = [Column("name").in_(query.agents)]
-        if query.include_forwarded:
-            agent_conditions.append(
-                and_(
-                    Column("forwarded_from").isnot(None),
-                    expression.cast(Column("forwarded_from"), JSON).contains(list(query.agents)),  # type: ignore
-                )
-            )
-        conditions.append(or_(*agent_conditions))
+        conditions.append(Column("name").in_(query.agents))
     if query.since and (cutoff := query.get_time_cutoff()):
         conditions.append(Message.timestamp >= cutoff)
     if query.until:
@@ -222,6 +220,7 @@ def format_conversation(
                     },
                     cost=float(msg.cost_info.total_cost) if msg.cost_info else None,
                     response_time=msg.response_time,
+                    parent_id=msg.parent_id,
                 )
                 for msg in chat_messages
             ],

agentpool_storage/zed_provider/__init__.py

@@ -0,0 +1,16 @@
+"""Zed IDE storage provider.
+
+This package implements a read-only storage backend that reads Zed IDE's
+native thread format from ~/.local/share/zed/threads/threads.db.
+
+Zed stores conversations as zstd-compressed JSON in a SQLite database.
+This provider enables importing and analyzing Zed conversations.
+"""
+
+from __future__ import annotations
+
+from agentpool_storage.zed_provider.provider import ZedStorageProvider
+
+__all__ = [
+    "ZedStorageProvider",
+]

agentpool_storage/zed_provider/helpers.py

@@ -0,0 +1,281 @@
+"""Helper functions for Zed storage provider.
+
+Stateless conversion and utility functions for working with Zed format.
+"""
+
+from __future__ import annotations
+
+import base64
+from datetime import datetime
+import io
+import json
+from typing import Any
+
+from pydantic_ai.messages import (
+    BinaryContent,
+    ModelRequest,
+    ModelResponse,
+    TextPart,
+    ThinkingPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from pydantic_ai.usage import RequestUsage
+import zstandard
+
+from agentpool.log import get_logger
+from agentpool.messaging import ChatMessage
+from agentpool.utils.now import get_now
+from agentpool_storage.zed_provider.models import ZedThread
+
+
+logger = get_logger(__name__)
+
+
+def decompress_thread(data: bytes, data_type: str) -> ZedThread:
+    """Decompress and parse thread data.
+
+    Args:
+        data: Compressed thread data
+        data_type: Type of compression ("zstd" or plain)
+
+    Returns:
+        Parsed ZedThread object
+    """
+    if data_type == "zstd":
+        dctx = zstandard.ZstdDecompressor()
+        # Use stream_reader for data without content size in header
+        reader = dctx.stream_reader(io.BytesIO(data))
+        json_data = reader.read()
+    else:
+        json_data = data
+
+    thread_dict = json.loads(json_data)
+    return ZedThread.model_validate(thread_dict)
+
+
+def parse_user_content(  # noqa: PLR0915
+    content_list: list[dict[str, Any]],
+) -> tuple[str, list[str | BinaryContent]]:
+    """Parse user message content blocks.
+
+    Args:
+        content_list: List of content blocks from Zed user message
+
+    Returns:
+        Tuple of (display_text, pydantic_ai_content_list)
+    """
+    display_parts: list[str] = []
+    pydantic_content: list[str | BinaryContent] = []
+
+    for item in content_list:
+        if "Text" in item:
+            text = item["Text"]
+            display_parts.append(text)
+            pydantic_content.append(text)
+
+        elif "Image" in item:
+            image_data = item["Image"]
+            source = image_data.get("source", "")
+            try:
+                binary_data = base64.b64decode(source)
+                # Try to detect image type from magic bytes
+                media_type = "image/png"  # default
+                if binary_data[:3] == b"\xff\xd8\xff":
+                    media_type = "image/jpeg"
+                elif binary_data[:4] == b"\x89PNG":
+                    media_type = "image/png"
+                elif binary_data[:6] in (b"GIF87a", b"GIF89a"):
+                    media_type = "image/gif"
+                elif binary_data[:4] == b"RIFF" and binary_data[8:12] == b"WEBP":
+                    media_type = "image/webp"
+
+                pydantic_content.append(BinaryContent(data=binary_data, media_type=media_type))
+                display_parts.append("[image]")
+            except (ValueError, TypeError, IndexError) as e:
+                logger.warning("Failed to decode image", error=str(e))
+                display_parts.append("[image decode error]")
+
+        elif "Mention" in item:
+            mention = item["Mention"]
+            uri = mention.get("uri", {})
+            content = mention.get("content", "")
+
+            # Format mention based on type
+            if "File" in uri:
+                file_info = uri["File"]
+                path = file_info.get("abs_path", "unknown")
+                formatted = f"[File: {path}]\n{content}"
+            elif "Directory" in uri:
+                dir_info = uri["Directory"]
+                path = dir_info.get("abs_path", "unknown")
+                formatted = f"[Directory: {path}]\n{content}"
+            elif "Symbol" in uri:
+                symbol_info = uri["Symbol"]
+                path = symbol_info.get("abs_path", "unknown")
+                name = symbol_info.get("name", "")
+                formatted = f"[Symbol: {name} in {path}]\n{content}"
+            elif "Selection" in uri:
+                sel_info = uri["Selection"]
+                path = sel_info.get("abs_path", "unknown")
+                formatted = f"[Selection: {path}]\n{content}"
+            elif "Fetch" in uri:
+                fetch_info = uri["Fetch"]
+                url = fetch_info.get("url", "unknown")
+                formatted = f"[Fetched: {url}]\n{content}"
+            else:
+                formatted = content
+
+            display_parts.append(formatted)
+            pydantic_content.append(formatted)
+
+    display_text = "\n".join(display_parts)
+    return display_text, pydantic_content
+
+
+def parse_agent_content(
+    content_list: list[dict[str, Any]],
+) -> tuple[str, list[TextPart | ThinkingPart | ToolCallPart]]:
+    """Parse agent message content blocks.
+
+    Args:
+        content_list: List of content blocks from Zed agent message
+
+    Returns:
+        Tuple of (display_text, pydantic_ai_parts)
+    """
+    display_parts: list[str] = []
+    pydantic_parts: list[TextPart | ThinkingPart | ToolCallPart] = []
+
+    for item in content_list:
+        if "Text" in item:
+            text = item["Text"]
+            display_parts.append(text)
+            pydantic_parts.append(TextPart(content=text))
+
+        elif "Thinking" in item:
+            thinking = item["Thinking"]
+            text = thinking.get("text", "")
+            signature = thinking.get("signature")
+            display_parts.append(f"<thinking>\n{text}\n</thinking>")
+            pydantic_parts.append(ThinkingPart(content=text, signature=signature))
+
+        elif "ToolUse" in item:
+            tool_use = item["ToolUse"]
+            tool_id = tool_use.get("id", "")
+            tool_name = tool_use.get("name", "")
+            tool_input = tool_use.get("input", {})
+            display_parts.append(f"[Tool: {tool_name}]")
+            pydantic_parts.append(
+                ToolCallPart(tool_name=tool_name, args=tool_input, tool_call_id=tool_id)
+            )
+
+    display_text = "\n".join(display_parts)
+    return display_text, pydantic_parts
+
+
+def parse_tool_results(tool_results: dict[str, Any]) -> list[ToolReturnPart]:
+    """Parse tool results into ToolReturnParts.
+
+    Args:
+        tool_results: Dictionary of tool results from Zed agent message
+
+    Returns:
+        List of ToolReturnPart objects
+    """
+    parts: list[ToolReturnPart] = []
+
+    for tool_id, result in tool_results.items():
+        if isinstance(result, dict):
+            tool_name = result.get("tool_name", "")
+            output = result.get("output", "")
+            content = result.get("content", {})
+
+            # Handle output being a dict like {"Text": "..."}
+            if isinstance(output, dict) and "Text" in output:
+                output = output["Text"]
+            # Extract text content if available
+            elif isinstance(content, dict) and "Text" in content:
+                output = content["Text"]
+            elif isinstance(content, str):
+                output = content
+
+            parts.append(
+                ToolReturnPart(tool_name=tool_name, content=output or "", tool_call_id=tool_id)
+            )
+
+    return parts
+
+
+def thread_to_chat_messages(thread: ZedThread, thread_id: str) -> list[ChatMessage[str]]:
+    """Convert a Zed thread to ChatMessages.
+
+    Args:
+        thread: Zed thread object
+        thread_id: Thread identifier
+
+    Returns:
+        List of ChatMessage objects
+    """
+    messages: list[ChatMessage[str]] = []
+    try:
+        updated_at = datetime.fromisoformat(thread.updated_at.replace("Z", "+00:00"))
+    except (ValueError, AttributeError):
+        updated_at = get_now()
+    # Get model info
+    model_name = None
+    if thread.model:
+        model_name = f"{thread.model.provider}:{thread.model.model}"
+
+    for idx, msg in enumerate(thread.messages):
+        if msg == "Resume":
+            continue  # Skip control messages
+        msg_id = f"{thread_id}_{idx}"
+
+        if msg.User is not None:
+            user_msg = msg.User
+            display_text, pydantic_content = parse_user_content(user_msg.content)
+            part = UserPromptPart(content=pydantic_content)
+            model_request = ModelRequest(parts=[part])
+
+            messages.append(
+                ChatMessage[str](
+                    content=display_text,
+                    conversation_id=thread_id,
+                    role="user",
+                    message_id=user_msg.id or msg_id,
+                    name=None,
+                    model_name=None,
+                    timestamp=updated_at,  # Zed doesn't store per-message timestamps
+                    messages=[model_request],
+                )
+            )
+
+        elif msg.Agent is not None:
+            agent_msg = msg.Agent
+            display_text, pydantic_parts = parse_agent_content(agent_msg.content)
+            # Build ModelResponse
+            usage = RequestUsage()
+            model_response = ModelResponse(parts=pydantic_parts, usage=usage, model_name=model_name)
+            # Build tool return parts for next request if there are tool results
+            tool_return_parts = parse_tool_results(agent_msg.tool_results)
+            # Create the messages list - response, then optionally tool returns
+            pydantic_messages: list[ModelResponse | ModelRequest] = [model_response]
+            if tool_return_parts:
+                pydantic_messages.append(ModelRequest(parts=tool_return_parts))
+
+            messages.append(
+                ChatMessage[str](
+                    content=display_text,
+                    conversation_id=thread_id,
+                    role="assistant",
+                    message_id=msg_id,
+                    name="zed",
+                    model_name=model_name,
+                    timestamp=updated_at,
+                    messages=pydantic_messages,
+                )
+            )
+
+    return messages