Atomizer/atomizer-dashboard/backend/api/services/session_manager.py

"""
Session Manager

Manages persistent Claude Code sessions with MCP integration.
Fixed for Windows compatibility - uses subprocess.Popen with ThreadPoolExecutor.
"""

import asyncio
import json
import os
import subprocess
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import AsyncGenerator, Dict, List, Literal, Optional

from .conversation_store import ConversationStore
from .context_builder import ContextBuilder

# Paths
ATOMIZER_ROOT = Path(__file__).parent.parent.parent.parent.parent
MCP_SERVER_PATH = ATOMIZER_ROOT / "mcp-server" / "atomizer-tools"

# Thread pool for subprocess operations (Windows compatible)
_executor = ThreadPoolExecutor(max_workers=4)


@dataclass
class ClaudeSession:
    """Represents an active Claude Code session"""

    session_id: str
    mode: Literal["user", "power"]
    study_id: Optional[str]
    created_at: datetime = field(default_factory=datetime.now)
    last_active: datetime = field(default_factory=datetime.now)

    def is_alive(self) -> bool:
        """Session is always 'alive' - we use stateless CLI calls"""
        return True


class SessionManager:
    """Manages Claude Code sessions with MCP tools"""

    def __init__(self):
        self.sessions: Dict[str, ClaudeSession] = {}
        self.store = ConversationStore()
        self.context_builder = ContextBuilder()
        self._cleanup_task: Optional[asyncio.Task] = None
        self._lock: Optional[asyncio.Lock] = None

    def _get_lock(self) -> asyncio.Lock:
        """Get or create the async lock (must be called from async context)"""
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    async def start(self):
        """Start the session manager"""
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    async def stop(self):
        """Stop the session manager and all sessions"""
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass

        # Clean up temp files
        for session in list(self.sessions.values()):
            self._cleanup_session_files(session.session_id)

    async def create_session(
        self,
        mode: Literal["user", "power"] = "user",
        study_id: Optional[str] = None,
        resume_session_id: Optional[str] = None,
    ) -> ClaudeSession:
        """
        Create or resume a Claude Code session.

        Note: Sessions are now stateless - we don't spawn persistent processes.
        Each message is handled via a one-shot CLI call for Windows compatibility.
        """
        async with self._get_lock():
            # Resume existing session if requested
            if resume_session_id and resume_session_id in self.sessions:
                session = self.sessions[resume_session_id]
                session.last_active = datetime.now()
                self.store.touch_session(session.session_id)
                return session

            session_id = resume_session_id or str(uuid.uuid4())[:8]

            # Create or update session in store
            existing = self.store.get_session(session_id)
            if existing:
                self.store.update_session(session_id, mode=mode, study_id=study_id)
            else:
                self.store.create_session(session_id, mode, study_id)

            # Build MCP config for this session
            mcp_config = self._build_mcp_config(mode)
            mcp_config_path = ATOMIZER_ROOT / f".claude-mcp-{session_id}.json"
            with open(mcp_config_path, "w") as f:
                json.dump(mcp_config, f)

            # Create session object (no subprocess - stateless)
            session = ClaudeSession(
                session_id=session_id,
                mode=mode,
                study_id=study_id,
            )

            self.sessions[session_id] = session
            return session

    async def send_message(
        self,
        session_id: str,
        message: str,
        canvas_state: Optional[Dict] = None,
    ) -> AsyncGenerator[Dict, None]:
        """
        Send a message to a session and stream the response.

        Uses synchronous subprocess.Popen via ThreadPoolExecutor for Windows compatibility.

        Args:
            session_id: The session ID
            message: User message
            canvas_state: Optional canvas state (nodes, edges) from UI
        """
        session = self.sessions.get(session_id)

        if not session:
            yield {"type": "error", "message": "Session not found"}
            return

        session.last_active = datetime.now()

        # Store user message
        self.store.add_message(session_id, "user", message)

        # Build context with conversation history AND canvas state
        history = self.store.get_history(session_id, limit=10)
        full_prompt = self.context_builder.build(
            mode=session.mode,
            study_id=session.study_id,
            conversation_history=history[:-1],
            canvas_state=canvas_state,  # Pass canvas state for context
        )
        full_prompt += f"\n\nUser: {message}\n\nRespond helpfully and concisely:"

        # Build CLI arguments
        cli_args = ["claude", "--print"]

        # Ensure MCP config exists
        mcp_config_path = ATOMIZER_ROOT / f".claude-mcp-{session_id}.json"
        if not mcp_config_path.exists():
            mcp_config = self._build_mcp_config(session.mode)
            with open(mcp_config_path, "w") as f:
                json.dump(mcp_config, f)
        cli_args.extend(["--mcp-config", str(mcp_config_path)])

        if session.mode == "user":
            cli_args.extend([
                "--allowedTools",
                "Read Write(**/STUDY_REPORT.md) Write(**/3_results/*.md) Bash(python:*) mcp__atomizer-tools__*"
            ])
        else:
            cli_args.append("--dangerously-skip-permissions")

        cli_args.append("-")  # Read from stdin

        full_response = ""
        tool_calls: List[Dict] = []

        try:
            loop = asyncio.get_event_loop()

            # Run subprocess in thread pool (Windows compatible)
            def run_claude():
                try:
                    process = subprocess.Popen(
                        cli_args,
                        stdin=subprocess.PIPE,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        cwd=str(ATOMIZER_ROOT),
                        text=True,
                        encoding='utf-8',
                        errors='replace',
                    )
                    stdout, stderr = process.communicate(input=full_prompt, timeout=300)
                    return {
                        "stdout": stdout,
                        "stderr": stderr,
                        "returncode": process.returncode,
                    }
                except subprocess.TimeoutExpired:
                    process.kill()
                    return {"error": "Response timeout (5 minutes)"}
                except FileNotFoundError:
                    return {"error": "Claude CLI not found in PATH. Install with: npm install -g @anthropic-ai/claude-code"}
                except Exception as e:
                    return {"error": str(e)}

            result = await loop.run_in_executor(_executor, run_claude)

            if "error" in result:
                yield {"type": "error", "message": result["error"]}
            else:
                full_response = result["stdout"] or ""

                if full_response:
                    yield {"type": "text", "content": full_response}

                if result["returncode"] != 0 and result["stderr"]:
                    yield {"type": "error", "message": f"CLI error: {result['stderr']}"}

        except Exception as e:
            yield {"type": "error", "message": str(e)}

        # Store assistant response
        if full_response:
            self.store.add_message(
                session_id,
                "assistant",
                full_response.strip(),
                tool_calls=tool_calls if tool_calls else None,
            )

        yield {"type": "done", "tool_calls": tool_calls}

    async def switch_mode(
        self,
        session_id: str,
        new_mode: Literal["user", "power"],
    ) -> ClaudeSession:
        """Switch a session's mode"""
        session = self.sessions.get(session_id)
        if not session:
            raise ValueError(f"Session {session_id} not found")

        session.mode = new_mode
        self.store.update_session(session_id, mode=new_mode)

        # Rebuild MCP config with new mode
        mcp_config = self._build_mcp_config(new_mode)
        mcp_config_path = ATOMIZER_ROOT / f".claude-mcp-{session_id}.json"
        with open(mcp_config_path, "w") as f:
            json.dump(mcp_config, f)

        return session

    async def set_study_context(
        self,
        session_id: str,
        study_id: str,
    ):
        """Update the study context for a session"""
        session = self.sessions.get(session_id)
        if session:
            session.study_id = study_id
            self.store.update_session(session_id, study_id=study_id)

    def get_session(self, session_id: str) -> Optional[ClaudeSession]:
        """Get session by ID"""
        return self.sessions.get(session_id)

    def get_session_info(self, session_id: str) -> Optional[Dict]:
        """Get session info including database record"""
        session = self.sessions.get(session_id)
        if not session:
            return None

        db_record = self.store.get_session(session_id)
        return {
            "session_id": session.session_id,
            "mode": session.mode,
            "study_id": session.study_id,
            "is_alive": session.is_alive(),
            "created_at": session.created_at.isoformat(),
            "last_active": session.last_active.isoformat(),
            "message_count": self.store.get_message_count(session_id),
            **({} if not db_record else {"db_record": db_record}),
        }

    def _build_mcp_config(self, mode: Literal["user", "power"]) -> dict:
        """Build MCP configuration for Claude"""
        return {
            "mcpServers": {
                "atomizer": {
                    "command": "node",
                    "args": [str(MCP_SERVER_PATH / "dist" / "index.js")],
                    "env": {
                        "ATOMIZER_MODE": mode,
                        "ATOMIZER_ROOT": str(ATOMIZER_ROOT),
                    },
                },
            },
        }

    def _cleanup_session_files(self, session_id: str):
        """Clean up temp files for a session"""
        for pattern in [
            f".claude-mcp-{session_id}.json",
            f".claude-prompt-{session_id}.md",
        ]:
            path = ATOMIZER_ROOT / pattern
            if path.exists():
                try:
                    path.unlink()
                except Exception:
                    pass

    async def _cleanup_loop(self):
        """Periodically clean up stale sessions"""
        while True:
            try:
                await asyncio.sleep(300)  # Every 5 minutes

                now = datetime.now()
                stale = [
                    sid
                    for sid, session in list(self.sessions.items())
                    if (now - session.last_active).total_seconds() > 3600
                ]

                for sid in stale:
                    self._cleanup_session_files(sid)
                    self.sessions.pop(sid, None)

                self.store.cleanup_stale_sessions(max_age_hours=24)

            except asyncio.CancelledError:
                break
            except Exception:
                pass


# Global instance
_session_manager: Optional[SessionManager] = None


def get_session_manager() -> SessionManager:
    """Get or create the global session manager instance"""
    global _session_manager
    if _session_manager is None:
        _session_manager = SessionManager()
    return _session_manager