Yash030/claude-code-proxy

1"""2CLI Session Manager for Multi-Instance Claude CLI Support3 4Manages a pool of CLISession instances, each handling one conversation.5This enables true parallel processing where multiple conversations run6simultaneously in separate CLI processes.7"""8 9import asyncio10import uuid11 12from loguru import logger13 14from .session import CLISession15 16 17class CLISessionManager:18    """19    Manages multiple CLISession instances for parallel conversation processing.20 21    Each new conversation gets its own CLISession with its own subprocess.22    Replies to existing conversations reuse the same CLISession instance.23    """24 25    def __init__(26        self,27        workspace_path: str,28        api_url: str,29        allowed_dirs: list[str] | None = None,30        plans_directory: str | None = None,31        claude_bin: str = "claude",32        *,33        log_raw_cli_diagnostics: bool = False,34        log_messaging_error_details: bool = False,35    ):36        """37        Initialize the session manager.38 39        Args:40            workspace_path: Working directory for CLI processes41            api_url: API URL for the proxy42            allowed_dirs: Directories the CLI is allowed to access43            plans_directory: Directory for Claude Code CLI plan files (passed via --settings)44        """45        self.workspace = workspace_path46        self.api_url = api_url47        self.allowed_dirs = allowed_dirs or []48        self.plans_directory = plans_directory49        self.claude_bin = claude_bin50        self._log_raw_cli_diagnostics = log_raw_cli_diagnostics51        self._log_messaging_error_details = log_messaging_error_details52 53        self._sessions: dict[str, CLISession] = {}54        self._pending_sessions: dict[str, CLISession] = {}55        self._temp_to_real: dict[str, str] = {}56        self._real_to_temp: dict[str, str] = {}57        self._lock = asyncio.Lock()58 59        logger.info("CLISessionManager initialized")60 61    async def get_or_create_session(62        self, session_id: str | None = None63    ) -> tuple[CLISession, str, bool]:64        """65        Get an existing session or create a new one.66 67        Returns:68            Tuple of (CLISession instance, session_id, is_new_session)69        """70        async with self._lock:71            if session_id:72                lookup_id = self._temp_to_real.get(session_id, session_id)73 74                if lookup_id in self._sessions:75                    return self._sessions[lookup_id], lookup_id, False76                if lookup_id in self._pending_sessions:77                    return self._pending_sessions[lookup_id], lookup_id, False78 79            temp_id = session_id if session_id else f"pending_{uuid.uuid4().hex[:8]}"80 81            new_session = CLISession(82                workspace_path=self.workspace,83                api_url=self.api_url,84                allowed_dirs=self.allowed_dirs,85                plans_directory=self.plans_directory,86                claude_bin=self.claude_bin,87                log_raw_cli_diagnostics=self._log_raw_cli_diagnostics,88            )89            self._pending_sessions[temp_id] = new_session90            logger.info(f"Created new session: {temp_id}")91 92            return new_session, temp_id, True93 94    async def register_real_session_id(95        self, temp_id: str, real_session_id: str96    ) -> bool:97        """Register the real session ID from CLI output."""98        async with self._lock:99            if temp_id not in self._pending_sessions:100                logger.warning(f"Temp session {temp_id} not found")101                return False102 103            session = self._pending_sessions.pop(temp_id)104            self._sessions[real_session_id] = session105            self._temp_to_real[temp_id] = real_session_id106            self._real_to_temp[real_session_id] = temp_id107 108            logger.info(f"Registered session: {temp_id} -> {real_session_id}")109            return True110 111    async def remove_session(self, session_id: str) -> bool:112        """Remove a session from the manager."""113        async with self._lock:114            if session_id in self._pending_sessions:115                session = self._pending_sessions.pop(session_id)116                await session.stop()117                return True118 119            if session_id in self._sessions:120                session = self._sessions.pop(session_id)121                await session.stop()122                temp_id = self._real_to_temp.pop(session_id, None)123                if temp_id is not None:124                    self._temp_to_real.pop(temp_id, None)125                return True126 127            return False128 129    async def stop_all(self):130        """Stop all sessions."""131        async with self._lock:132            all_sessions = list(self._sessions.values()) + list(133                self._pending_sessions.values()134            )135            for session in all_sessions:136                try:137                    await session.stop()138                except Exception as e:139                    if self._log_messaging_error_details:140                        logger.error(141                            "Error stopping session: {}: {}",142                            type(e).__name__,143                            e,144                        )145                    else:146                        logger.error(147                            "Error stopping session: exc_type={}",148                            type(e).__name__,149                        )150 151            self._sessions.clear()152            self._pending_sessions.clear()153            self._temp_to_real.clear()154            self._real_to_temp.clear()155            logger.info("All sessions stopped")156 157    def get_stats(self) -> dict:158        """Get session statistics."""159        return {160            "active_sessions": len(self._sessions),161            "pending_sessions": len(self._pending_sessions),162            "busy_count": sum(1 for s in self._sessions.values() if s.is_busy),163        }164