agentpool 2.1.9__py3-none-any.whl

agentpool/sessions/manager.py

@@ -0,0 +1,302 @@
+"""Session manager for pool-level session lifecycle."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any, Self
+from uuid import uuid4
+
+from agentpool.log import get_logger
+from agentpool.sessions.models import SessionData
+from agentpool.sessions.session import ClientSession
+from agentpool.sessions.store import MemorySessionStore
+
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+    from agentpool.delegation.pool import AgentPool
+    from agentpool.sessions.store import SessionStore
+
+logger = get_logger(__name__)
+
+
+class SessionManager:
+    """Manages session lifecycle at the pool level.
+
+    Handles:
+    - Session creation and initialization
+    - Active session tracking (in-memory)
+    - Session persistence via SessionStore
+    - Session resumption from storage
+    - Cleanup of expired sessions
+    """
+
+    def __init__(
+        self,
+        pool: AgentPool[Any],
+        store: SessionStore | None = None,
+        pool_id: str | None = None,
+    ) -> None:
+        """Initialize session manager.
+
+        Args:
+            pool: Agent pool for agent access
+            store: Session persistence backend (defaults to MemorySessionStore)
+            pool_id: Optional identifier for this pool (for multi-pool setups)
+        """
+        self._pool = pool
+        self._store = store or MemorySessionStore()
+        self._pool_id = pool_id
+        self._active: dict[str, ClientSession] = {}
+        self._lock = asyncio.Lock()
+        logger.debug("Initialized session manager", pool_id=pool_id)
+
+    @property
+    def pool(self) -> AgentPool[Any]:
+        """Get the agent pool."""
+        return self._pool
+
+    @property
+    def store(self) -> SessionStore:
+        """Get the session store."""
+        return self._store
+
+    @property
+    def active_sessions(self) -> dict[str, ClientSession]:
+        """Get currently active sessions (read-only view)."""
+        return dict(self._active)
+
+    async def __aenter__(self) -> Self:
+        """Enter context and initialize store."""
+        await self._store.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit context and clean up all sessions."""
+        # Close all active sessions
+        async with self._lock:
+            sessions = list(self._active.values())
+            self._active.clear()
+
+        for session in sessions:
+            try:
+                await session.close()
+            except Exception:
+                logger.exception("Error closing session", session_id=session.session_id)
+
+        await self._store.__aexit__(exc_type, exc_val, exc_tb)
+        logger.debug("Session manager closed", session_count=len(sessions))
+
+    def generate_session_id(self) -> str:
+        """Generate a unique session ID."""
+        return f"sess_{uuid4().hex[:12]}"
+
+    async def create(
+        self,
+        agent_name: str,
+        *,
+        session_id: str | None = None,
+        cwd: str | None = None,
+        conversation_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        session_class: type[ClientSession] = ClientSession,
+        **session_kwargs: Any,
+    ) -> ClientSession:
+        """Create a new session.
+
+        Args:
+            agent_name: Name of the initial agent
+            session_id: Optional specific session ID (generated if None)
+            cwd: Working directory for the session
+            conversation_id: Optional conversation ID (generated if None)
+            metadata: Optional session metadata
+            session_class: Session class to instantiate (for protocol-specific sessions)
+            **session_kwargs: Additional kwargs passed to session constructor
+
+        Returns:
+            Created session instance
+
+        Raises:
+            ValueError: If session_id already exists
+            KeyError: If agent_name not found in pool
+        """
+        # Validate agent exists
+        if agent_name not in self._pool.all_agents:
+            msg = f"Agent '{agent_name}' not found in pool"
+            raise KeyError(msg)
+
+        async with self._lock:
+            # Generate or validate session ID
+            if session_id is None:
+                session_id = self.generate_session_id()
+            elif session_id in self._active:
+                msg = f"Session '{session_id}' already exists"
+                raise ValueError(msg)
+
+            # Create session data
+            data = SessionData(
+                session_id=session_id,
+                agent_name=agent_name,
+                conversation_id=conversation_id or f"conv_{uuid4().hex[:12]}",
+                pool_id=self._pool_id,
+                cwd=cwd,
+                metadata=metadata or {},
+            )
+
+            # Persist to store
+            await self._store.save(data)
+
+            # Create runtime session
+            session = session_class(
+                data=data,
+                pool=self._pool,
+                manager=self,
+                **session_kwargs,
+            )
+
+            self._active[session_id] = session
+            logger.info(
+                "Created session",
+                session_id=session_id,
+                agent=agent_name,
+            )
+            return session
+
+    async def get(self, session_id: str) -> ClientSession | None:
+        """Get an active session by ID.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            Active session if found, None otherwise
+        """
+        return self._active.get(session_id)
+
+    async def resume(
+        self,
+        session_id: str,
+        session_class: type[ClientSession] = ClientSession,
+        **session_kwargs: Any,
+    ) -> ClientSession | None:
+        """Resume a session from storage.
+
+        Loads session data from store and creates a runtime session.
+
+        Args:
+            session_id: Session identifier
+            session_class: Session class to instantiate
+            **session_kwargs: Additional kwargs passed to session constructor
+
+        Returns:
+            Resumed session if found in store, None otherwise
+        """
+        async with self._lock:
+            # Check if already active
+            if session_id in self._active:
+                return self._active[session_id]
+
+            # Try to load from store
+            data = await self._store.load(session_id)
+            if data is None:
+                return None
+
+            # Validate agent still exists
+            if data.agent_name not in self._pool.all_agents:
+                logger.warning(
+                    "Session agent no longer exists",
+                    session_id=session_id,
+                    agent=data.agent_name,
+                )
+                return None
+
+            # Create runtime session
+            session = session_class(
+                data=data,
+                pool=self._pool,
+                manager=self,
+                **session_kwargs,
+            )
+
+            self._active[session_id] = session
+            logger.info("Resumed session", session_id=session_id)
+            return session
+
+    async def close(self, session_id: str, *, delete: bool = False) -> bool:
+        """Close and optionally delete a session.
+
+        Args:
+            session_id: Session identifier
+            delete: Whether to also delete from store
+
+        Returns:
+            True if session was closed, False if not found
+        """
+        async with self._lock:
+            session = self._active.pop(session_id, None)
+
+        if session:
+            await session.close()
+            if delete:
+                await self._store.delete(session_id)
+            logger.info("Closed session", session_id=session_id, deleted=delete)
+            return True
+
+        # Session not active, but maybe in store
+        if delete:
+            return await self._store.delete(session_id)
+
+        return False
+
+    async def save(self, data: SessionData) -> None:
+        """Save session data to store.
+
+        Args:
+            data: Session data to persist
+        """
+        await self._store.save(data)
+
+    async def list_sessions(
+        self,
+        *,
+        active_only: bool = False,
+        agent_name: str | None = None,
+    ) -> list[str]:
+        """List session IDs.
+
+        Args:
+            active_only: Only return currently active sessions
+            agent_name: Filter by agent name
+
+        Returns:
+            List of session IDs
+        """
+        if active_only:
+            sessions = list(self._active.keys())
+            if agent_name:
+                sessions = [sid for sid, s in self._active.items() if s.agent_name == agent_name]
+            return sessions
+
+        return await self._store.list_sessions(
+            pool_id=self._pool_id,
+            agent_name=agent_name,
+        )
+
+    async def cleanup_expired(self, max_age_hours: int = 24) -> int:
+        """Remove expired sessions from store.
+
+        Does not affect currently active sessions.
+
+        Args:
+            max_age_hours: Maximum session age in hours
+
+        Returns:
+            Number of sessions removed
+        """
+        return await self._store.cleanup_expired(max_age_hours)

agentpool/sessions/models.py

@@ -0,0 +1,71 @@
+"""Session data models."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import Field
+from schemez import Schema
+
+from agentpool.utils.now import get_now
+
+
+class SessionData(Schema):
+    """Persistable session state.
+
+    Contains all information needed to persist and restore a session.
+    Protocol-specific data (ACP capabilities, web cookies, etc.) goes in metadata.
+    """
+
+    session_id: str
+    """Unique session identifier."""
+
+    agent_name: str
+    """Name of the currently active agent."""
+
+    conversation_id: str
+    """Links to conversation in StorageManager."""
+
+    title: str | None = None
+    """AI-generated or user-provided title for the conversation."""
+
+    pool_id: str | None = None
+    """Optional pool/manifest identifier for multi-pool setups."""
+
+    cwd: str | None = None
+    """Working directory for the session."""
+
+    created_at: datetime = Field(default_factory=get_now)
+    """When the session was created."""
+
+    last_active: datetime = Field(default_factory=get_now)
+    """Last activity timestamp."""
+
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    """Protocol-specific or custom metadata.
+
+    Examples:
+        - ACP: client_capabilities, mcp_servers
+        - Web: user_id, auth_token
+        - CLI: terminal_size, color_support
+    """
+
+    def touch(self) -> None:
+        """Update last_active timestamp."""
+        # Note: Schema is frozen by default, so we need to work around that
+        # by using object.__setattr__ or making this field mutable
+        object.__setattr__(self, "last_active", get_now())
+
+    def with_agent(self, agent_name: str) -> SessionData:
+        """Return copy with different agent."""
+        return self.model_copy(update={"agent_name": agent_name, "last_active": get_now()})
+
+    def with_metadata(self, **kwargs: Any) -> SessionData:
+        """Return copy with updated metadata."""
+        new_metadata = {**self.metadata, **kwargs}
+        return self.model_copy(update={"metadata": new_metadata, "last_active": get_now()})
+
+    def with_title(self, title: str) -> SessionData:
+        """Return copy with updated title."""
+        return self.model_copy(update={"title": title, "last_active": get_now()})

agentpool/sessions/session.py

@@ -0,0 +1,239 @@
+"""Protocol-agnostic client session."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any, Self
+
+from agentpool.log import get_logger
+from agentpool.messaging.message_history import MessageHistory
+
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+    from agentpool.agents.agent import Agent
+    from agentpool.delegation.pool import AgentPool
+    from agentpool.messaging import ChatMessage
+    from agentpool.sessions.manager import SessionManager
+    from agentpool.sessions.models import SessionData
+
+logger = get_logger(__name__)
+
+
+class ClientSession:
+    """Protocol-agnostic runtime session.
+
+    Base class for protocol-specific sessions (ACP, Web, etc.).
+    Manages the runtime state of an active session, including:
+    - Reference to session data (persistable state)
+    - Active agent instance
+    - Conversation history (owned by session, not agent)
+    - Pool access for agent switching
+
+    The session owns the conversation history and passes it to the agent
+    on each run. This makes the agent stateless from the session's perspective.
+
+    Subclasses add protocol-specific functionality:
+    - ACPSession: ACP connection, capabilities, slash commands
+    - WebSession: WebSocket handling, HTTP session state
+    """
+
+    def __init__(
+        self,
+        data: SessionData,
+        pool: AgentPool[Any],
+        manager: SessionManager | None = None,
+    ) -> None:
+        """Initialize client session.
+
+        Args:
+            data: Persistable session state
+            pool: Agent pool for agent access and switching
+            manager: Optional session manager for lifecycle operations
+        """
+        self._data = data
+        self._pool = pool
+        self._manager = manager
+        self._agent: Agent[Any, Any] | None = None
+        self._closed = False
+        self._title_generation_triggered = False
+        self._title_task: asyncio.Task[None] | None = None
+        # Session owns conversation history - agent is stateless
+        self._history = MessageHistory()
+
+        logger.debug(
+            "Created client session",
+            session_id=data.session_id,
+            agent=data.agent_name,
+        )
+
+    @property
+    def session_id(self) -> str:
+        """Get session identifier."""
+        return self._data.session_id
+
+    @property
+    def data(self) -> SessionData:
+        """Get session data (persistable state)."""
+        return self._data
+
+    @property
+    def pool(self) -> AgentPool[Any]:
+        """Get agent pool."""
+        return self._pool
+
+    @property
+    def agent(self) -> Agent[Any, Any]:
+        """Get current active agent, creating if needed."""
+        if self._agent is None:
+            self._agent = self._pool.get_agent(self._data.agent_name)
+        return self._agent
+
+    @property
+    def agent_name(self) -> str:
+        """Get current agent name."""
+        return self._data.agent_name
+
+    @property
+    def conversation_id(self) -> str:
+        """Get conversation ID for message storage."""
+        return self._data.conversation_id
+
+    @property
+    def history(self) -> MessageHistory:
+        """Get the session's conversation history."""
+        return self._history
+
+    @property
+    def is_closed(self) -> bool:
+        """Check if session is closed."""
+        return self._closed
+
+    async def __aenter__(self) -> Self:
+        """Enter session context."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit session context."""
+        await self.close()
+
+    async def run(self, prompt: str, **kwargs: Any) -> ChatMessage[Any]:
+        """Run the agent with a prompt, using session's history.
+
+        The session passes its own MessageHistory to the agent, making
+        the agent stateless from the session's perspective. Messages
+        are automatically added to the session's history.
+
+        Args:
+            prompt: User prompt to send to the agent
+            **kwargs: Additional arguments passed to agent.run()
+
+        Returns:
+            The agent's response message
+        """
+        result = await self.agent.run(
+            prompt,
+            message_history=self._history,
+            conversation_id=self.conversation_id,
+            **kwargs,
+        )
+
+        # Trigger title generation after first exchange (fire-and-forget)
+        # Note: Message logging is handled by the agent itself via MessageNode.log_message
+        if not self._title_generation_triggered and self._pool.storage:
+            self._title_generation_triggered = True
+            self._title_task = asyncio.create_task(self._generate_title())
+
+        return result
+
+    async def _generate_title(self) -> None:
+        """Generate conversation title in the background."""
+        if not self._pool.storage:
+            return
+        try:
+            messages = self._history.get_history()
+            if messages:
+                title = await self._pool.storage.generate_conversation_title(
+                    self.conversation_id,
+                    messages,
+                )
+                # Also update SessionData so title is available when listing sessions
+                if title and self._manager:
+                    self._data = self._data.with_title(title)
+                    await self._manager.save(self._data)
+        except Exception:
+            logger.exception(
+                "Failed to generate conversation title",
+                conversation_id=self.conversation_id,
+            )
+
+    async def switch_agent(self, agent_name: str) -> None:
+        """Switch to a different agent.
+
+        The conversation history is preserved across agent switches.
+
+        Args:
+            agent_name: Name of agent to switch to
+
+        Raises:
+            KeyError: If agent not found in pool
+        """
+        if agent_name not in self._pool.all_agents:
+            msg = f"Agent '{agent_name}' not found in pool"
+            raise KeyError(msg)
+
+        self._agent = self._pool.get_agent(agent_name)
+        self._data = self._data.with_agent(agent_name)
+
+        # Persist the change
+        if self._manager:
+            await self._manager.save(self._data)
+
+        logger.info(
+            "Switched agent",
+            session_id=self.session_id,
+            agent=agent_name,
+        )
+
+    async def touch(self) -> None:
+        """Update last_active timestamp and persist."""
+        self._data.touch()
+        if self._manager:
+            await self._manager.save(self._data)
+
+    async def save(self) -> None:
+        """Persist current session state."""
+        if self._manager:
+            await self._manager.save(self._data)
+            logger.debug("Saved session", session_id=self.session_id)
+
+    async def close(self) -> None:
+        """Close the session and clean up resources."""
+        if self._closed:
+            return
+
+        self._closed = True
+        self._agent = None
+        logger.debug("Closed session", session_id=self.session_id)
+
+    def update_metadata(self, **kwargs: Any) -> None:
+        """Update session metadata.
+
+        Args:
+            **kwargs: Key-value pairs to add/update in metadata
+        """
+        self._data = self._data.with_metadata(**kwargs)
+
+    def clear_history(self) -> None:
+        """Clear the session's conversation history."""
+        self._history.clear()
+
+    def get_history_messages(self) -> list[ChatMessage[Any]]:
+        """Get all messages in the session's history."""
+        return self._history.get_history()