Anto01
ba0b9a1fae
Backend: - spec.py: New AtomizerSpec REST API endpoints - spec_manager.py: SpecManager service for unified config - interview_engine.py: Study creation interview logic - claude.py: Enhanced Claude API with context - optimization.py: Extended optimization endpoints - context_builder.py, session_manager.py: Improved services Frontend: - Chat components: Enhanced message rendering, tool call cards - Hooks: useClaudeCode, useSpecWebSocket, improved useChat - Pages: Updated Dashboard, Analysis, Insights, Setup, Home - Components: ParallelCoordinatesPlot, ParetoPlot improvements - App.tsx: Route updates for canvas/studio Infrastructure: - vite.config.ts: Build configuration updates - start/stop-dashboard.bat: Script improvements
454 lines
16 KiB
Python
454 lines
16 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:
|
|
# Check if response contains canvas modifications (from MCP tools)
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
modifications = self._extract_canvas_modifications(full_response)
|
|
logger.info(f"[SEND_MSG] Found {len(modifications)} canvas modifications to send")
|
|
|
|
for mod in modifications:
|
|
logger.info(f"[SEND_MSG] Sending canvas_modification: {mod.get('action')} {mod.get('nodeType')}")
|
|
yield {"type": "canvas_modification", "modification": mod}
|
|
|
|
# Always send the text 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 _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
|
|
|
|
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
|