Anto01
a26914bbe8
Dashboard: - Add Studio page with drag-drop model upload and Claude chat - Add intake system for study creation workflow - Improve session manager and context builder - Add intake API routes and frontend components Optimization Engine: - Add CLI module for command-line operations - Add intake module for study preprocessing - Add validation module with gate checks - Improve Zernike extractor documentation - Update spec models with better validation - Enhance solve_simulation robustness Documentation: - Add ATOMIZER_STUDIO.md planning doc - Add ATOMIZER_UX_SYSTEM.md for UX patterns - Update extractor library docs - Add study-readme-generator skill Tools: - Add test scripts for extraction validation - Add Zernike recentering test Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
539 lines
19 KiB
Python
539 lines
19 KiB
Python
"""
|
|
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
|