Anto01
ac5e9b4054
## Documentation Updates - DASHBOARD.md: Updated to V3.0 with Canvas V3 features, file browser, introspection - DASHBOARD_IMPLEMENTATION_STATUS.md: Marked Canvas V3 features as COMPLETE - CANVAS.md: New comprehensive guide for Canvas Builder V3 with all features - CLAUDE.md: Added dashboard quick reference and Canvas V3 features ## Canvas V3 Features Documented - File Browser: Browse studies directory for model files - Model Introspection: Auto-discover expressions, solver type, dependencies - One-Click Add: Add expressions as design variables instantly - Claude Bug Fixes: WebSocket reconnection, SQL errors resolved - Health Check: /api/health endpoint for monitoring ## Backend Services - NX introspection service with expression discovery - File browser API with type filtering - Claude session management improvements - Context builder enhancements ## Frontend Components - FileBrowser: Modal for file selection with search - IntrospectionPanel: View discovered model information - ExpressionSelector: Dropdown for design variable configuration - Improved chat hooks with reconnection logic ## Plan Documents - Added RALPH_LOOP_CANVAS_V2/V3 implementation records - Added ATOMIZER_DASHBOARD_V2_MASTER_PLAN - Added investigation and sync documentation Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
358 lines
12 KiB
Python
358 lines
12 KiB
Python
"""
|
|
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
|