Jack1808/Claude_Code

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    ):32        """33        Initialize the session manager.34 35        Args:36            workspace_path: Working directory for CLI processes37            api_url: API URL for the proxy38            allowed_dirs: Directories the CLI is allowed to access39            plans_directory: Directory for Claude Code CLI plan files (passed via --settings)40        """41        self.workspace = workspace_path42        self.api_url = api_url43        self.allowed_dirs = allowed_dirs or []44        self.plans_directory = plans_directory45 46        self._sessions: dict[str, CLISession] = {}47        self._pending_sessions: dict[str, CLISession] = {}48        self._temp_to_real: dict[str, str] = {}49        self._real_to_temp: dict[str, str] = {}50        self._lock = asyncio.Lock()51 52        logger.info("CLISessionManager initialized")53 54    async def get_or_create_session(55        self, session_id: str | None = None56    ) -> tuple[CLISession, str, bool]:57        """58        Get an existing session or create a new one.59 60        Returns:61            Tuple of (CLISession instance, session_id, is_new_session)62        """63        async with self._lock:64            if session_id:65                lookup_id = self._temp_to_real.get(session_id, session_id)66 67                if lookup_id in self._sessions:68                    return self._sessions[lookup_id], lookup_id, False69                if lookup_id in self._pending_sessions:70                    return self._pending_sessions[lookup_id], lookup_id, False71 72            temp_id = session_id if session_id else f"pending_{uuid.uuid4().hex[:8]}"73 74            new_session = CLISession(75                workspace_path=self.workspace,76                api_url=self.api_url,77                allowed_dirs=self.allowed_dirs,78                plans_directory=self.plans_directory,79            )80            self._pending_sessions[temp_id] = new_session81            logger.info(f"Created new session: {temp_id}")82 83            return new_session, temp_id, True84 85    async def register_real_session_id(86        self, temp_id: str, real_session_id: str87    ) -> bool:88        """Register the real session ID from CLI output."""89        async with self._lock:90            if temp_id not in self._pending_sessions:91                logger.warning(f"Temp session {temp_id} not found")92                return False93 94            session = self._pending_sessions.pop(temp_id)95            self._sessions[real_session_id] = session96            self._temp_to_real[temp_id] = real_session_id97            self._real_to_temp[real_session_id] = temp_id98 99            logger.info(f"Registered session: {temp_id} -> {real_session_id}")100            return True101 102    async def remove_session(self, session_id: str) -> bool:103        """Remove a session from the manager."""104        async with self._lock:105            if session_id in self._pending_sessions:106                session = self._pending_sessions.pop(session_id)107                await session.stop()108                return True109 110            if session_id in self._sessions:111                session = self._sessions.pop(session_id)112                await session.stop()113                temp_id = self._real_to_temp.pop(session_id, None)114                if temp_id is not None:115                    self._temp_to_real.pop(temp_id, None)116                return True117 118            return False119 120    async def stop_all(self):121        """Stop all sessions."""122        async with self._lock:123            all_sessions = list(self._sessions.values()) + list(124                self._pending_sessions.values()125            )126            for session in all_sessions:127                try:128                    await session.stop()129                except Exception as e:130                    logger.error(f"Error stopping session: {e}")131 132            self._sessions.clear()133            self._pending_sessions.clear()134            self._temp_to_real.clear()135            self._real_to_temp.clear()136            logger.info("All sessions stopped")137 138    def get_stats(self) -> dict:139        """Get session statistics."""140        return {141            "active_sessions": len(self._sessions),142            "pending_sessions": len(self._pending_sessions),143            "busy_count": sum(1 for s in self._sessions.values() if s.is_busy),144        }145