session_manager.py

"""
Session manager for Claude Agent SDK sessions

This module manages execution sessions for Claude agents.
Provides session creation, state management, message tracking, and cleanup functionality.
"""

import asyncio
import logging
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient, Message

from models import MessageInfo, SessionStatus

# Setup logging for debugging
logger = logging.getLogger(__name__)


class SessionInfo:
    """Class for managing information about sessions

    Holds session ID, status, configuration, message history, results, etc.
    """

    def __init__(self, session_id: str, options: ClaudeAgentOptions, prompt: str = ""):
        self.session_id = session_id
        self.status = SessionStatus.PENDING
        self.options = options
        self.prompt = prompt
        self.created_at = datetime.now()
        self.client: Optional[ClaudeSDKClient] = None
        self.messages: List[MessageInfo] = []
        self.result: Optional[Dict[str, Any]] = None
        self.error: Optional[str] = None
        self.start_time: Optional[datetime] = None
        self.end_time: Optional[datetime] = None
        self.task: Optional[asyncio.Task] = None
        # Session ID actually generated by Claude Agent SDK (for resume)
        self.claude_session_id: Optional[str] = None

    def add_message(self, message: Message) -> None:
        """Add a message to the session

        Args:
            message: Message object to add
        """
        message_info = MessageInfo(
            type=message.__class__.__name__,
            content=self._serialize_message(message),
            timestamp=datetime.now().isoformat(),
        )
        self.messages.append(message_info)

    def _serialize_message(self, message: Message) -> Any:
        """Serialize message to JSON-compatible format

        Args:
            message: Message object to serialize

        Returns:
            Message data in JSON-compatible format
        """
        try:
            if hasattr(message, "__dict__"):
                result = {}
                for key, value in message.__dict__.items():
                    if key.startswith("_"):
                        continue
                    try:
                        result[key] = self._serialize_value(value)
                    except Exception as e:
                        logger.debug(f"Error serializing key '{key}': {e}")
                        result[key] = f"<serialization_error: {str(e)}>"
                return result
            return str(message)
        except Exception as e:
            logger.error(f"Error in _serialize_message: {e}")
            return f"<message_serialization_error: {str(e)}>"

    def _serialize_value(self, value: Any) -> Any:
        """Serialize a value to a JSON-compatible format"""
        try:
            if isinstance(value, (str, int, float, bool, type(None))):
                return value
            elif isinstance(value, list):
                return [self._serialize_value(v) for v in value]
            elif isinstance(value, dict):
                return {k: self._serialize_value(v) for k, v in value.items()}
            elif hasattr(value, "__dict__"):
                result = {}
                for key, val in value.__dict__.items():
                    if key.startswith("_"):
                        continue
                    try:
                        result[key] = self._serialize_value(val)
                    except Exception as e:
                        logger.debug(f"Error serializing nested key '{key}': {e}")
                        result[key] = f"<nested_serialization_error: {str(e)}>"
                return result
            else:
                return str(value)
        except Exception as e:
            logger.error(f"Error in _serialize_value: {e}")
            return f"<value_serialization_error: {str(e)}>"

    def get_duration_ms(self) -> Optional[int]:
        """Get session execution time in milliseconds

        Returns:
            Execution time (milliseconds), or None if not executed
        """
        if self.start_time is None:
            return None
        end = self.end_time or datetime.now()
        return int((end - self.start_time).total_seconds() * 1000)


class SessionManager:
    """
    Manager class for Claude Agent SDK sessions

    Provides session creation, execution, state management, cancellation, and cleanup functionality.
    Uses asynchronous locks to ensure thread-safe operations.
    """

    def __init__(self):
        self.sessions: Dict[str, SessionInfo] = {}
        self._lock = asyncio.Lock()

    def generate_session_id(self) -> str:
        """Generate a unique session ID

        Returns:
            Unique session ID in UUID format
        """
        return str(uuid.uuid4())

    async def create_session(
        self,
        prompt: str,
        options: ClaudeAgentOptions,
        resume_session_id: Optional[str] = None,
    ) -> SessionInfo:
        """Create a new session and start the agent

        Args:
            prompt: Prompt to send to the agent
            options: Claude Agent SDK configuration options
            resume_session_id: Session ID to resume (optional)

        Returns:
            Created or resumed session information
        """
        if resume_session_id:
            # Resume existing session
            return await self.resume_session(resume_session_id, prompt, options)
        else:
            # Create new session
            session_id = self.generate_session_id()
            session = SessionInfo(session_id, options, prompt)

            async with self._lock:
                self.sessions[session_id] = session
                logger.info(
                    f"Created session {session_id}, total sessions: {len(self.sessions)}"
                )

            # Start agent in background
            session.task = asyncio.create_task(self._run_agent(session, prompt))

            return session

    async def resume_session(
        self, session_id: str, prompt: str, options: ClaudeAgentOptions
    ) -> SessionInfo:
        """Resume existing session using Claude Agent SDK resume feature

        If there is an existing session with a recorded Claude session ID,
        uses Claude Agent SDK resume feature to continue the conversation.

        Args:
            session_id: Session ID to resume
            prompt: Prompt to send to the agent
            options: Claude Agent SDK configuration options

        Returns:
            Resumed or newly created session information
        """
        async with self._lock:
            existing_session = self.sessions.get(session_id)

            if existing_session and existing_session.claude_session_id:
                # If existing session exists and has Claude session ID
                if existing_session.status in [SessionStatus.RUNNING]:
                    # Return error if currently running
                    raise ValueError(f"Session {session_id} is currently running")

                # Resume session using Claude Agent SDK resume feature
                logger.info(
                    f"Resuming session {session_id} with Claude session ID {existing_session.claude_session_id}"
                )

                # Set resume parameter
                options.resume = existing_session.claude_session_id

                # Update session information
                existing_session.status = SessionStatus.PENDING
                existing_session.prompt = (
                    f"{existing_session.prompt}\n\n--- Session Resume ---\n{prompt}"
                )
                existing_session.options = options
                existing_session.error = None
                existing_session.start_time = None
                existing_session.end_time = None

                # Cancel existing task if present
                if existing_session.task and not existing_session.task.done():
                    existing_session.task.cancel()

                session = existing_session
            else:
                # Error if no existing session or no Claude session ID
                if existing_session:
                    raise ValueError(
                        f"Session {session_id} has no recorded Claude session ID. Cannot resume."
                    )
                else:
                    raise ValueError(f"Session {session_id} not found.")

            logger.info(
                f"Session {session_id} ready for resume with Claude session ID, total sessions: {len(self.sessions)}"
            )

        # Start agent in background
        session.task = asyncio.create_task(self._run_agent(session, prompt))

        return session

    async def _run_agent(self, session: SessionInfo, prompt: str) -> None:
        """Run agent in background

        Creates Claude SDK client, sends prompt,
        receives messages and manages session.

        Args:
            session: Information about the session to execute
            prompt: Prompt to send to the agent
        """
        try:
            session.status = SessionStatus.RUNNING
            session.start_time = datetime.now()

            # Create client and connect
            session.client = ClaudeSDKClient(options=session.options)
            await session.client.connect()

            # Send query
            await session.client.query(prompt, session_id=session.session_id)

            # Receive all messages
            async for message in session.client.receive_messages():
                session.add_message(message)

                # Log detailed message information for debugging
                logger.debug(f"Received message type: {type(message).__name__}")
                logger.debug(
                    f"Message attributes: {[attr for attr in dir(message) if not attr.startswith('_')]}"
                )
                if hasattr(message, "subtype"):
                    logger.debug(f"Message subtype: {message.subtype}")
                if hasattr(message, "is_error"):
                    logger.debug(f"Message is_error: {message.is_error}")
                if hasattr(message, "content"):
                    logger.debug(f"Message content type: {type(message.content)}")

                # Log full message structure for detailed debugging
                logger.debug(f"Full message: {message}")

                # Early capture of Claude session ID (if available from first message)
                if (
                    not session.claude_session_id
                    and hasattr(message, "session_id")
                    and message.session_id
                ):
                    session.claude_session_id = message.session_id
                    logger.debug(
                        f"Early capture of Claude session ID: {session.claude_session_id}"
                    )

                # Check if it's an error message
                if hasattr(message, "is_error") and message.is_error:
                    session.status = SessionStatus.ERROR
                    if hasattr(message, "result"):
                        session.error = message.result
                    else:
                        session.error = "Unknown error occurred"
                    logger.warning("Setting status to ERROR")
                    break

                # Check various completion indicators
                is_final = False

                # Check final_result subtype
                if hasattr(message, "subtype") and message.subtype == "final_result":
                    logger.debug("Found final_result subtype")
                    is_final = True

                # Check other completion indicators
                if hasattr(message, "type") and message.type in [
                    "result",
                    "final",
                    "completion",
                ]:
                    logger.debug(f"Found completion type: {message.type}")
                    is_final = True

                # Check message with turn count and execution time info (likely final message)
                if hasattr(message, "num_turns") and hasattr(message, "duration_ms"):
                    logger.debug("Found message with turn/duration info (likely final)")
                    is_final = True

                if is_final:
                    session.status = SessionStatus.COMPLETED
                    logger.info("Setting status to COMPLETED")
                    # Save result information
                    session.result = {
                        "session_id": (
                            message.session_id
                            if hasattr(message, "session_id")
                            else None
                        ),
                        "num_turns": (
                            message.num_turns if hasattr(message, "num_turns") else None
                        ),
                        "duration_ms": (
                            message.duration_ms
                            if hasattr(message, "duration_ms")
                            else None
                        ),
                        "total_cost_usd": (
                            message.total_cost_usd
                            if hasattr(message, "total_cost_usd")
                            else None
                        ),
                        "usage": session._serialize_value(message.usage)
                        if hasattr(message, "usage")
                        else None,
                    }

                    # Save actual session ID generated by Claude Agent SDK (for resume)
                    if hasattr(message, "session_id") and message.session_id:
                        session.claude_session_id = message.session_id
                        logger.info(
                            f"Saved Claude session ID for resume: {session.claude_session_id}"
                        )

                    break

        except asyncio.CancelledError:
            session.status = SessionStatus.CANCELLED
            session.error = "Session was cancelled"
            logger.info("Session cancelled")
        except Exception as e:
            session.status = SessionStatus.ERROR
            session.error = str(e)
            logger.error(f"Exception in _run_agent: {e}")
            logger.error(f"Exception type: {type(e)}")
            import traceback

            logger.error(f"Full traceback: {traceback.format_exc()}")
        finally:
            session.end_time = datetime.now()
            logger.info(f"Session ended with status: {session.status}")
            if session.client:
                try:
                    await session.client.disconnect()
                except Exception:
                    pass

    async def get_session(self, session_id: str) -> Optional[SessionInfo]:
        """Get session by session ID

        Args:
            session_id: ID of the session to retrieve

        Returns:
            Session information, or None if not found
        """
        async with self._lock:
            return self.sessions.get(session_id)

    async def cancel_session(self, session_id: str) -> bool:
        """Cancel a running session

        Args:
            session_id: ID of the session to cancel

        Returns:
            True if cancellation was successful, False otherwise
        """
        async with self._lock:
            session = self.sessions.get(session_id)
            if not session:
                return False

            if session.status == SessionStatus.RUNNING:
                # Interrupt client
                if session.client:
                    try:
                        await session.client.interrupt()
                    except Exception:
                        pass

                # Cancel task
                if session.task:
                    session.task.cancel()

                session.status = SessionStatus.CANCELLED
                session.end_time = datetime.now()
                return True

            return False

    async def list_sessions(self) -> List[str]:
        """List all session IDs

        Returns:
            List of session IDs
        """
        async with self._lock:
            session_ids = list(self.sessions.keys())
            logger.debug(
                f"list_sessions called, found {len(session_ids)} sessions: {session_ids}"
            )
            return session_ids

    async def get_all_sessions(self) -> List[dict]:
        """Get all sessions with detailed information

        Returns:
            List of dictionaries containing session detailed information
        """
        async with self._lock:
            sessions_data = []
            for session_id, session in self.sessions.items():
                session_data = {
                    "session_id": session_id,
                    "status": session.status.value,
                    "prompt": session.prompt,
                    "created_at": session.created_at.isoformat(),
                    "start_time": session.start_time.isoformat()
                    if session.start_time
                    else None,
                    "end_time": session.end_time.isoformat()
                    if session.end_time
                    else None,
                    "error": session.error,
                    "result": session.result,
                    "message_count": len(session.messages) if session.messages else 0,
                    "claude_session_id": session.claude_session_id,
                }
                sessions_data.append(session_data)

            logger.debug(f"get_all_sessions returning {len(sessions_data)} sessions")
            return sessions_data

    async def cleanup_old_sessions(self, max_age_hours: int = 24) -> int:
        """Clean up old sessions

        Removes sessions whose end time is older than the specified time.

        Args:
            max_age_hours: Maximum age of sessions to keep (default: 24 hours)

        Returns:
            Number of removed sessions
        """
        cutoff = datetime.now().timestamp() - (max_age_hours * 3600)
        removed = 0

        async with self._lock:
            to_remove = []
            for session_id, session in self.sessions.items():
                if session.end_time and session.end_time.timestamp() < cutoff:
                    to_remove.append(session_id)

            for session_id in to_remove:
                del self.sessions[session_id]
                removed += 1

        return removed

    async def delete_session(
        self, session_id: str
    ) -> tuple[bool, Optional[str], Optional[SessionStatus]]:
        """Delete a specific session by session ID

        Deletes session from memory. Cannot delete running sessions.
        Running sessions must be cancelled first using cancel_session().

        Args:
            session_id: ID of the session to delete

        Returns:
            Tuple of (success: bool, error_message: Optional[str], status_before_deletion: Optional[SessionStatus])
            - success: True if deletion was successful, False otherwise
            - error_message: Error message if deletion failed, None otherwise
            - status_before_deletion: Session status before deletion, None if session not found
        """
        async with self._lock:
            session = self.sessions.get(session_id)

            if not session:
                return (False, f"Session {session_id} not found", None)

            # Check if session is in a deletable state
            if session.status == SessionStatus.RUNNING:
                return (
                    False,
                    f"Cannot delete running session {session_id}. Please cancel it first.",
                    session.status,
                )

            # Store status before deletion for response
            status_before = session.status

            # Delete session from memory
            del self.sessions[session_id]
            logger.info(
                f"Deleted session {session_id} with status {status_before}, remaining sessions: {len(self.sessions)}"
            )

            return (True, None, status_before)