marm-mcp

"""Complete WebSocket handlers for all MCP protocol methods.""" from fastapi import WebSocket from typing import Dict, Any import sqlite3 import logging from datetime import datetime, timezone # Setup logging for security error tracking logger = logging.getLogger(__name__) from core.memory import memory from core.events import events from core.websocket_manager import websocket_manager as ws_manager from core.response_limiter import MCPResponseLimiter from utils.helpers import read_protocol_file # ===== MEMORY HANDLERS ===== async def handle_smart_recall(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle smart recall requests via WebSocket""" try: query = message.get("params", {}).get("query", "") session_name = message.get("params", {}).get("session_name", ws_manager.get_client_session(client_id)) search_all = message.get("params", {}).get("search_all", False) # Perform the search results = await memory.recall_similar(query, session_name if not search_all else None) response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "memories": results, "count": len(results) } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_contextual_log(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle contextual log requests via WebSocket""" try: params = message.get("params", {}) content = params.get("content", "") session_name = params.get("session_name", ws_manager.get_client_session(client_id)) if not content: raise ValueError("content parameter is required") # Store memory with contextual processing memory_id = await memory.store_memory(content, session_name) response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success", "message": f"🧠 Memory stored with contextual processing", "memory_id": memory_id, "session_name": session_name, "content_length": len(content) } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) # ===== SESSION HANDLERS ===== async def handle_start(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle MARM start requests via WebSocket""" try: params = message.get("params", {}) session_name = params.get("session_name", f"ws_session_{client_id}") with memory.get_connection() as conn: conn.execute(''' INSERT OR REPLACE INTO sessions (session_name, marm_active, last_accessed) VALUES (?, TRUE, ?) ''', (session_name, datetime.now(timezone.utc).isoformat())) conn.commit() # Read the current protocol from file protocol_content = await read_protocol_file() await events.emit('marm_started', {'session': session_name}) response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success", "message": f"🚀 MARM protocol activated for session '{session_name}'", "session_name": session_name, "marm_active": True, "protocol_content": protocol_content, "instructions": "The complete MARM protocol documentation has been loaded and is available for reference." } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_refresh(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle MARM refresh requests via WebSocket""" try: params = message.get("params", {}) session_name = params.get("session_name", ws_manager.get_client_session(client_id)) with memory.get_connection() as conn: conn.execute(''' UPDATE sessions SET last_accessed = ? WHERE session_name = ? ''', (datetime.now(timezone.utc).isoformat(), session_name)) conn.commit() # Read the current protocol from file to reaffirm adherence protocol_content = await read_protocol_file() await events.emit('marm_refreshed', {'session': session_name}) response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success", "message": f"🔄 MARM session '{session_name}' refreshed - protocol adherence reaffirmed", "session_name": session_name, "protocol_content": protocol_content, "instructions": "Protocol documentation refreshed. Please review the current MARM protocol specifications above." } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) # ===== LOGGING HANDLERS ===== async def handle_log_entry(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle log entry requests via WebSocket""" try: params = message.get("params", {}) session_name = params.get("session_name", ws_manager.get_client_session(client_id)) entry = params.get("entry", "") if not entry: raise ValueError("Entry parameter is required") # Log the entry memory_id = await memory.store_memory(entry, session_name) response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "id": memory_id, "session_name": session_name, "timestamp": "current_time" } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_log_session(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle log session requests via WebSocket""" try: params = message.get("params", {}) session_name = params.get("session_name", f"ws_session_{client_id}") # Create the session if it doesn't exist with memory.get_connection() as conn: conn.execute(''' INSERT OR IGNORE INTO sessions (session_name, marm_active, last_accessed) VALUES (?, FALSE, ?) ''', (session_name, datetime.now(timezone.utc).isoformat())) conn.commit() await events.emit('log_session_created', {'session': session_name}) response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success", "message": f"📝 Log session '{session_name}' created and activated", "session_name": session_name, "instructions": "Session is ready for logging entries. Use log_entry to add content." } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_log_show(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle log show requests via WebSocket""" try: params = message.get("params", {}) session_name = params.get("session_name") limit = params.get("limit", 20) with memory.get_connection() as conn: if session_name: # Show specific session cursor = conn.execute(''' SELECT entry_date, topic, summary, full_entry FROM log_entries WHERE session_name = ? ORDER BY entry_date DESC LIMIT ? ''', (session_name, limit)) entries = cursor.fetchall() response_data = { "status": "success", "session_name": session_name, "entries": [{"date": e[0], "topic": e[1], "summary": e[2], "content": e[3]} for e in entries], "entry_count": len(entries) } else: # Show all sessions cursor = conn.execute(''' SELECT DISTINCT session_name, COUNT(*) as entry_count, MAX(entry_date) as last_entry FROM log_entries GROUP BY session_name ORDER BY last_entry DESC LIMIT ? ''', (limit,)) sessions = cursor.fetchall() response_data = { "status": "success", "sessions": [{"name": s[0], "entry_count": s[1], "last_entry": s[2]} for s in sessions], "session_count": len(sessions) } response = { "jsonrpc": "2.0", "id": message.get("id"), "result": response_data } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_log_delete(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle log delete requests via WebSocket""" try: params = message.get("params", {}) session_name = params.get("session_name") entry_id = params.get("entry_id") if not session_name: raise ValueError("session_name parameter is required") with memory.get_connection() as conn: if entry_id: # Delete specific entry cursor = conn.execute(''' DELETE FROM log_entries WHERE session_name = ? AND id = ? ''', (session_name, entry_id)) deleted = cursor.rowcount else: # Delete entire session cursor = conn.execute(''' DELETE FROM log_entries WHERE session_name = ? ''', (session_name,)) deleted = cursor.rowcount # Also delete session record conn.execute(''' DELETE FROM sessions WHERE session_name = ? ''', (session_name,)) conn.commit() response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success" if deleted > 0 else "not_found", "message": f"🗑️ Deleted {deleted} entries" if deleted > 0 else "No entries found to delete", "deleted_count": deleted } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) # ===== NOTEBOOK HANDLERS ===== async def handle_notebook_add(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle notebook add requests via WebSocket""" try: params = message.get("params", {}) name = params.get("name", "") data = params.get("data", "") if not name or not data: raise ValueError("Both name and data parameters are required") # Generate embedding if available embedding_bytes = None if memory.encoder: try: embedding = memory.encoder.encode(data) embedding_bytes = embedding.tobytes() except Exception as e: print(f"Failed to generate embedding: {e}") # Add to notebook using same logic as HTTP endpoint with memory.get_connection() as conn: conn.execute(''' INSERT OR REPLACE INTO notebook_entries (name, data, embedding, updated_at) VALUES (?, ?, ?, ?) ''', (name, data, embedding_bytes, datetime.now(timezone.utc).isoformat())) conn.commit() # Emit event await events.emit('notebook_entry_added', { 'name': name, 'data': data }) response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success", "message": f"📓 Notebook entry '{name}' added", "name": name } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_notebook_use(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle notebook use requests via WebSocket""" try: params = message.get("params", {}) name = params.get("name", "") # Updated from 'names' to 'name' for consistency if not name: raise ValueError("name parameter is required") names = [n.strip() for n in name.split(',')] activated_entries = [] with memory.get_connection() as conn: for entry_name in names: cursor = conn.execute('SELECT name, data FROM notebook_entries WHERE name = ?', (entry_name,)) result = cursor.fetchone() if result: activated_entries.append({"name": result[0], "data": result[1]}) # Update active list in memory memory.active_notebook_entries = activated_entries response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success", "message": f"🔧 Activated {len(activated_entries)} notebook entries", "activated_entries": [e["name"] for e in activated_entries], "entries": activated_entries } } await ws_manager.send_personal_message(response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_notebook_show(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle notebook show requests via WebSocket""" try: with memory.get_connection() as conn: cursor = conn.execute('SELECT name, data, created_at, updated_at FROM notebook_entries ORDER BY updated_at DESC') entries = [{"name": row[0], "preview": row[1][:100] + "..." if len(row[1]) > 100 else row[1], "created_at": row[2], "updated_at": row[3]} for row in cursor.fetchall()] response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success", "entries": entries, "total_count": len(entries)}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id) async def handle_notebook_delete(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle notebook delete requests via WebSocket""" try: name = message.get("params", {}).get("name", "") if not name: raise ValueError("name parameter is required") with memory.get_connection() as conn: cursor = conn.execute('DELETE FROM notebook_entries WHERE name = ?', (name,)) deleted = cursor.rowcount conn.commit() response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success" if deleted > 0 else "not_found", "deleted": deleted > 0}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id) async def handle_notebook_clear(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle notebook clear requests via WebSocket""" try: memory.active_notebook_entries = [] response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success", "message": "🧹 Active notebook entries cleared", "active_count": 0}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id) async def handle_notebook_status(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle notebook status requests via WebSocket""" try: active_names = [entry["name"] for entry in memory.active_notebook_entries] response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success", "active_entries": active_names, "active_count": len(active_names)}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id) # ===== REASONING HANDLERS ===== async def handle_summary(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle summary requests via WebSocket""" try: params = message.get("params", {}) session_name = params.get("session_name", ws_manager.get_client_session(client_id)) limit = params.get("limit", 50) # Generate summary using same logic as HTTP endpoint with memory.get_connection() as conn: # Get total count first cursor = conn.execute(''' SELECT COUNT(*) FROM log_entries WHERE session_name = ? ''', (session_name,)) total_entries = cursor.fetchone()[0] # Get limited entries for summary cursor = conn.execute(''' SELECT entry_date, topic, summary, full_entry FROM log_entries WHERE session_name = ? ORDER BY entry_date DESC LIMIT ? ''', (session_name, limit)) entries = cursor.fetchall() if not entries: response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "empty", "message": f"No entries found in session '{session_name}'" } } await ws_manager.send_personal_message(response, client_id) return # Build base response metadata base_response = { "status": "success", "session_name": session_name, "entry_count": len(entries), "total_entries": total_entries } # Build summary with size monitoring summary_lines = [f"# MARM Session Summary: {session_name}"] summary_lines.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M UTC')}") summary_lines.append("") if total_entries > len(entries): summary_lines.append(f"*Showing {len(entries)} most recent entries out of {total_entries} total*") summary_lines.append("") # Add entries with progressive truncation if needed included_entries = [] current_summary_lines = summary_lines.copy() for entry in entries: # Truncate long summaries to prevent size explosion entry_summary = entry[2] if len(entry_summary) > 200: entry_summary = entry_summary[:197] + "..." entry_line = f"**{entry[0]}** [{entry[1]}]: {entry_summary}" test_lines = current_summary_lines + [entry_line] # Test response size with this entry added test_summary = "\n".join(test_lines) test_response = base_response.copy() test_response["summary"] = test_summary response_size = MCPResponseLimiter.estimate_response_size(test_response) if response_size > MCPResponseLimiter.CONTENT_LIMIT: # Can't fit this entry, stop here break # Entry fits, add it current_summary_lines.append(entry_line) included_entries.append(entry) summary_text = "\n".join(current_summary_lines) # Final response with truncation notice if needed final_response = { "jsonrpc": "2.0", "id": message.get("id"), "result": { "status": "success", "session_name": session_name, "summary": summary_text, "entry_count": len(included_entries), "total_entries": total_entries } } # Add truncation notice if we couldn't fit all entries if len(included_entries) < len(entries): final_response["result"]["_mcp_truncated"] = True final_response["result"]["_truncation_reason"] = "Summary limited to 1MB for MCP compliance" final_response["result"]["_entries_shown"] = len(included_entries) final_response["result"]["_entries_available"] = len(entries) await ws_manager.send_personal_message(final_response, client_id) except Exception as e: response = { "jsonrpc": "2.0", "id": message.get("id"), "error": { "code": -32603, "message": "Request processing failed" } } await ws_manager.send_personal_message(response, client_id) async def handle_context_bridge(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle context bridge requests via WebSocket""" try: params = message.get("params", {}) new_topic = params.get("new_topic", "") session_name = params.get("session_name", ws_manager.get_client_session(client_id)) bridge_text = f"# Context Bridge: {new_topic}\nSession: {session_name}\n\nReady to proceed with focused work" response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success", "bridge_text": bridge_text, "session_name": session_name}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id) # ===== SYSTEM HANDLERS ===== async def handle_current_context(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle current context requests via WebSocket""" try: current_time = datetime.now(timezone.utc).isoformat() response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success", "current_datetime": current_time}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id) async def handle_system_info(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle system info requests via WebSocket""" try: response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success", "message": "MARM MCP Server - WebSocket Protocol", "version": "2.2.6-beta"}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id) async def handle_reload_docs(websocket: WebSocket, client_id: str, message: Dict[str, Any]): """Handle reload docs requests via WebSocket""" try: response = {"jsonrpc": "2.0", "id": message.get("id"), "result": {"status": "success", "message": "📚 Documentation reloaded"}} await ws_manager.send_personal_message(response, client_id) except Exception as e: await ws_manager.send_personal_message({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32603, "message": "Request processing failed"}}, client_id)