""" Session Manager Manages persistent Claude Code sessions with direct file editing. Fixed for Windows compatibility - uses subprocess.Popen with ThreadPoolExecutor. Strategy: Claude edits atomizer_spec.json directly using Edit/Write tools (no MCP dependency for reliability). """ import asyncio import hashlib 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) import logging logger = logging.getLogger(__name__) @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. Claude edits atomizer_spec.json directly using Edit/Write tools (no MCP). 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) # Get spec path and hash BEFORE Claude runs (to detect changes) spec_path = self._get_spec_path(session.study_id) if session.study_id else None spec_hash_before = self._get_file_hash(spec_path) if spec_path else None # 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, spec_path=str(spec_path) if spec_path else None, # Tell Claude where the spec is ) full_prompt += f"\n\nUser: {message}\n\nRespond helpfully and concisely:" # Build CLI arguments - NO MCP for reliability cli_args = ["claude", "--print"] if session.mode == "user": # User mode: limited tools cli_args.extend( [ "--allowedTools", "Read Bash(python:*)", ] ) else: # Power mode: full access to edit files cli_args.append("--dangerously-skip-permissions") cli_args.append("-") # Read from stdin full_response = "" tool_calls: List[Dict] = [] process: Optional[subprocess.Popen] = None try: loop = asyncio.get_event_loop() # Run subprocess in thread pool (Windows compatible) def run_claude(): nonlocal process 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: if process: 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: # Always send the text response first yield {"type": "text", "content": full_response} if result["returncode"] != 0 and result["stderr"]: logger.warning(f"[SEND_MSG] CLI stderr: {result['stderr']}") except Exception as e: logger.error(f"[SEND_MSG] Exception: {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, ) # Check if spec was modified by comparing hashes if spec_path and session.mode == "power" and session.study_id: spec_hash_after = self._get_file_hash(spec_path) if spec_hash_before != spec_hash_after: logger.info(f"[SEND_MSG] Spec file was modified! Sending update.") spec_update = await self._check_spec_updated(session.study_id) if spec_update: yield { "type": "spec_updated", "spec": spec_update, "tool": "direct_edit", "reason": "Claude modified spec file directly", } yield {"type": "done", "tool_calls": tool_calls} def _get_spec_path(self, study_id: str) -> Optional[Path]: """Get the atomizer_spec.json path for a study.""" if not study_id: return None if study_id.startswith("draft_"): spec_path = ATOMIZER_ROOT / "studies" / "_inbox" / study_id / "atomizer_spec.json" else: spec_path = ATOMIZER_ROOT / "studies" / study_id / "atomizer_spec.json" if not spec_path.exists(): spec_path = ATOMIZER_ROOT / "studies" / study_id / "1_setup" / "atomizer_spec.json" return spec_path if spec_path.exists() else None def _get_file_hash(self, path: Optional[Path]) -> Optional[str]: """Get MD5 hash of a file for change detection.""" if not path or not path.exists(): return None try: with open(path, "rb") as f: return hashlib.md5(f.read()).hexdigest() except Exception: return None 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 _extract_canvas_modifications(self, response: str) -> List[Dict]: """ Extract canvas modification objects from Claude's response. MCP tools like canvas_add_node return JSON with a 'modification' field. This method finds and extracts those modifications so the frontend can apply them. """ import re import logging logger = logging.getLogger(__name__) modifications = [] # Debug: log what we're searching logger.info(f"[CANVAS_MOD] Searching response ({len(response)} chars) for modifications") # Check if "modification" even exists in the response if '"modification"' not in response: logger.info("[CANVAS_MOD] No 'modification' key found in response") return modifications try: # Method 1: Look for JSON in code fences code_block_pattern = r"```(?:json)?\s*([\s\S]*?)```" for match in re.finditer(code_block_pattern, response): block_content = match.group(1).strip() try: obj = json.loads(block_content) if isinstance(obj, dict) and "modification" in obj: logger.info( f"[CANVAS_MOD] Found modification in code fence: {obj['modification']}" ) modifications.append(obj["modification"]) except json.JSONDecodeError: continue # Method 2: Find JSON objects using proper brace matching # This handles nested objects correctly i = 0 while i < len(response): if response[i] == "{": # Found a potential JSON start, find matching close brace_count = 1 j = i + 1 in_string = False escape_next = False while j < len(response) and brace_count > 0: char = response[j] if escape_next: escape_next = False elif char == "\\": escape_next = True elif char == '"' and not escape_next: in_string = not in_string elif not in_string: if char == "{": brace_count += 1 elif char == "}": brace_count -= 1 j += 1 if brace_count == 0: potential_json = response[i:j] try: obj = json.loads(potential_json) if isinstance(obj, dict) and "modification" in obj: mod = obj["modification"] # Avoid duplicates if mod not in modifications: logger.info( f"[CANVAS_MOD] Found inline modification: action={mod.get('action')}, nodeType={mod.get('nodeType')}" ) modifications.append(mod) except json.JSONDecodeError as e: # Not valid JSON, skip pass i = j else: i += 1 except Exception as e: logger.error(f"[CANVAS_MOD] Error extracting modifications: {e}") logger.info(f"[CANVAS_MOD] Extracted {len(modifications)} modification(s)") return modifications async def _check_spec_updated(self, study_id: str) -> Optional[Dict]: """ Check if the atomizer_spec.json was modified and return the updated spec. For drafts in _inbox/, we check the spec file directly. """ import logging logger = logging.getLogger(__name__) try: # Determine spec path based on study_id if study_id.startswith("draft_"): spec_path = ATOMIZER_ROOT / "studies" / "_inbox" / study_id / "atomizer_spec.json" else: # Regular study path spec_path = ATOMIZER_ROOT / "studies" / study_id / "atomizer_spec.json" if not spec_path.exists(): spec_path = ( ATOMIZER_ROOT / "studies" / study_id / "1_setup" / "atomizer_spec.json" ) if not spec_path.exists(): logger.debug(f"[SPEC_CHECK] Spec not found at {spec_path}") return None # Read and return the spec with open(spec_path, "r", encoding="utf-8") as f: spec = json.load(f) logger.info(f"[SPEC_CHECK] Loaded spec from {spec_path}") return spec except Exception as e: logger.error(f"[SPEC_CHECK] Error checking spec: {e}") return None 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