diff --git a/hooks/shared.py b/hooks/shared.py index e3c3720f..69fe310a 100644 --- a/hooks/shared.py +++ b/hooks/shared.py @@ -5,7 +5,6 @@ import asyncio import concurrent.futures import logging -import threading from typing import Any from lifecycle.consolidation import ConsolidationEngine @@ -15,29 +14,26 @@ logger = logging.getLogger(__name__) -# Lock to prevent concurrent SQLite access from run_async -_db_lock = threading.Lock() - -# Module-level executor to avoid creating a new pool on every run_async call -_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4) +# Module-level executor — max_workers=1 serializes DB access without blocking the event loop +_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) def run_async(coro): """Run async coroutine from sync context, handling running event loops. - Uses threading.Lock to prevent concurrent SQLite access. + Uses ThreadPoolExecutor(max_workers=1) to serialize DB access + without blocking the main event loop. """ - with _db_lock: - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = None - - if loop and loop.is_running(): - future = _executor.submit(asyncio.run, coro) - return future.result() - else: - return asyncio.run(coro) + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop and loop.is_running(): + future = _executor.submit(asyncio.run, coro) + return future.result(timeout=30) + else: + return asyncio.run(coro) def forgetting_ritual(ctx: dict[str, Any]) -> dict[str, Any]: diff --git a/mcp_server/server.py b/mcp_server/server.py index f84aab18..81bbaf07 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -65,6 +65,7 @@ def __init__(self): @asynccontextmanager async def lifespan(server: FastMCP): from shared.migrations import migration_manager + from lifecycle.importance_scheduler import importance_scheduler result = await migration_manager.migrate() import logging @@ -75,9 +76,33 @@ async def lifespan(server: FastMCP): ctx = AppContext() backup_cron.start() + importance_scheduler.start() + + # Periodic maintenance tasks + async def _periodic_tasks(): + from lifecycle.forgetting import ForgettingSystem + + forgetting = ForgettingSystem() + + while True: + try: + await asyncio.sleep(900) # 15 minutes + await forgetting.cleanup() + except asyncio.CancelledError: + break + except Exception as e: + logging.getLogger(__name__).error("Periodic task error: %s", e) + + periodic_task = asyncio.create_task(_periodic_tasks()) try: yield ctx finally: + periodic_task.cancel() + try: + await periodic_task + except asyncio.CancelledError: + pass + importance_scheduler.stop() backup_cron.stop() diff --git a/mcp_server/tools_layer.py b/mcp_server/tools_layer.py index e9dcf5ef..3368fccb 100644 --- a/mcp_server/tools_layer.py +++ b/mcp_server/tools_layer.py @@ -52,6 +52,27 @@ def _get_hooks(app, layer: str): return app.user_hooks +_VALID_LAYERS = ("user", "agent") + + +def _validate_layer(layer: str) -> str: + """Validate and normalize layer parameter.""" + if layer not in _VALID_LAYERS: + raise ValueError(f"Invalid layer: {layer!r}. Must be one of {_VALID_LAYERS}") + return layer + + +def _fire_hook(hook_name: str, layer: str, context: dict) -> dict: + """Fire a hook safely — logs errors but never breaks the tool.""" + from hooks.registry import hook_registry + + try: + return hook_registry.fire(hook_name, layer, context) + except Exception as e: + logger.warning("Hook %s failed: %s", hook_name, e) + return {"error": str(e)} + + async def _check_rate_limit(app, user_id: str) -> dict | None: """Check rate limit. Returns error dict if exceeded, None if ok.""" try: @@ -125,6 +146,7 @@ async def memory_remember( importance: Importance score 0.0-1.0 (default 0.5) """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_remember") @@ -133,12 +155,9 @@ async def memory_remember( return rate_limit hooks = _get_hooks(app, layer) - try: - gate = hooks._importance_gate({"text": value}) - if gate.get("bypass"): - return RememberResult(status="skipped", reason="below_importance_threshold").dict() - except Exception as e: - logger.warning("Hook _importance_gate failed: %s", e) + gate = _fire_hook("importance_gate", layer, {"text": value, "key": key, "importance": importance}) + if gate.get("results") and any(r.get("bypass") for r in gate["results"] if isinstance(r, dict)): + return RememberResult(status="skipped", reason="below_importance_threshold").dict() mem = _get_memory(app, layer, user_id) @@ -157,14 +176,33 @@ async def memory_remember( ) # Invalidate context cache _context_cache.pop(_get_cache_key(layer, user_id), None) + + # Fire post-save hooks + _fire_hook("message_received", layer, {"text": value, "key": key, "user_id": user_id}) + _fire_hook("emotion_trigger", layer, {"text": value, "user_id": user_id}) + if "error" in key.lower(): + _fire_hook("error_occurred", layer, {"key": key, "value": value, "user_id": user_id}) + elif "decision" in key.lower(): + _fire_hook("decision_made", layer, {"key": key, "value": value, "user_id": user_id}) + elif "correction" in key.lower(): + _fire_hook("self_correction", layer, {"key": key, "value": value, "user_id": user_id}) + return RememberResult(status="ok", entry_id=entry_id, graph_node_id=node_id).dict() entry_id = await mem.remember(key, value, importance) - should_save, emotion_reason, emotion_weight = app.emotion_trigger.should_save(value) - if should_save: - await mem.l3.save(user_id, "%s=%s" % (key, value[:50]), emotion_weight, [emotion_reason]) + + # Fire emotion trigger hook (with direct fallback) + emotion_result = _fire_hook("emotion_trigger", layer, {"text": value, "user_id": user_id}) + if not emotion_result.get("results"): + should_save, emotion_reason, emotion_weight = app.emotion_trigger.should_save(value) + if should_save: + await mem.l3.save(user_id, "%s=%s" % (key, value[:50]), emotion_weight, [emotion_reason]) _context_cache.pop(_get_cache_key(layer, user_id), None) + + # Fire post-save hooks + _fire_hook("message_received", layer, {"text": value, "key": key, "user_id": user_id}) + return RememberResult(status="ok", entry_id=entry_id).dict() @@ -184,15 +222,21 @@ async def memory_recall( limit: Max results (default 10) """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_recall") + _fire_hook("retrieval_router", layer, {"query": query, "user_id": user_id, "limit": limit}) + cached = _get_recall_cache(query, user_id, layer, limit) if cached is not None: return {**RecallResult(results=cached, count=len(cached)).dict(), "cached": True} results = await _get_memory(app, layer, user_id).recall(query, limit) _set_recall_cache(query, user_id, layer, limit, results) + + _fire_hook("auto_context", layer, {"query": query, "results_count": len(results), "user_id": user_id}) + return RecallResult(results=results, count=len(results)).dict() @@ -210,6 +254,7 @@ async def memory_forget( key: Fact key to delete """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_forget") @@ -234,6 +279,7 @@ async def memory_session_start( user_id: User identifier """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_session_start") @@ -242,6 +288,9 @@ async def memory_session_start( return rate_limit session_id = await _get_memory(app, layer, user_id).l2.create_session(user_id) + + _fire_hook("message_received", layer, {"text": "session_started", "session_id": session_id, "user_id": user_id}) + return SessionResult(session_id=session_id).dict() @@ -261,6 +310,7 @@ async def memory_session_end( summary: Session summary """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_session_end") @@ -269,6 +319,10 @@ async def memory_session_end( return rate_limit await _get_memory(app, layer, user_id).l2.close_session(session_id, summary) + + _fire_hook("consolidation", layer, {"trigger": "session_end", "session_id": session_id, "user_id": user_id}) + _fire_hook("state_delta", layer, {"trigger": "session_end", "session_id": session_id, "summary": summary, "user_id": user_id}) + return SessionResult(status="ok").dict() @@ -290,6 +344,7 @@ async def memory_episode_save( tags: Tags (e.g. ["greeting", "decision", "error"]) """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_episode_save") @@ -299,6 +354,12 @@ async def memory_episode_save( episode_id = await _get_memory(app, layer, user_id).l3.save(user_id, summary, weight, tags) _context_cache.pop(_get_cache_key(layer, user_id), None) + + # Fire post-save hooks + _fire_hook("emotion_trigger", layer, {"summary": summary, "emotional_weight": weight, "user_id": user_id}) + _fire_hook("state_delta", layer, {"summary": summary, "tags": tags, "user_id": user_id}) + _fire_hook("consolidation", layer, {"trigger": "episode_save", "user_id": user_id}) + return EpisodeResult(episode_id=episode_id).dict() @@ -318,8 +379,12 @@ async def memory_episode_recall( limit: Max results (default 10) """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_episode_recall") + + _fire_hook("retrieval_router", layer, {"query": tag or "episodes", "user_id": user_id, "limit": limit}) + mem = _get_memory(app, layer, user_id) if tag: episodes = await mem.l3.search_by_tag(user_id, tag, limit) @@ -346,6 +411,7 @@ async def memory_graph_add( tags: Tags """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_graph_add") @@ -355,6 +421,18 @@ async def memory_graph_add( node_id = await _get_graph(app, layer).add_node(user_id, content, node_type, tags) _context_cache.pop(_get_cache_key(layer, user_id), None) + + # Fire graph-specific hooks + hook_map = { + "error_analysis": "error_occurred", + "decision_log": "decision_made", + "personality": "personality_shift", + "emotion": "emotion_context", + } + hook_name = hook_map.get(node_type) + if hook_name: + _fire_hook(hook_name, layer, {"node_type": node_type, "content": content, "user_id": user_id}) + return GraphNodeResult(node_id=node_id).dict() @@ -376,8 +454,12 @@ async def memory_graph_query( limit: Max results """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_graph_query") + + _fire_hook("retrieval_router", layer, {"query": tag or node_type, "user_id": user_id, "limit": limit}) + graph = _get_graph(app, layer) if tag: nodes = await graph.query_by_tag(user_id, tag, limit) @@ -402,6 +484,7 @@ async def memory_session_list( limit: Max results (default 10) """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_session_list") sessions = await _get_memory(app, layer, user_id).l2.get_recent_sessions(user_id, limit) @@ -435,6 +518,7 @@ async def memory_episode_list( offset: Pagination offset """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_episode_list") episodes = await _get_memory(app, layer, user_id).l3.get_episodes(user_id, limit, offset) @@ -458,6 +542,7 @@ async def memory_episode_get( episode_id: Episode database ID """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_episode_get") mem = _get_memory(app, layer, user_id) @@ -493,6 +578,7 @@ async def memory_graph_nodes( limit: Max results (default 20) """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_graph_nodes") graph = _get_graph(app, layer) @@ -524,6 +610,7 @@ async def memory_graph_edges( limit: Max results (default 20) """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_graph_edges") graph = _get_graph(app, layer) @@ -577,6 +664,7 @@ async def memory_stats( user_id: User identifier """ app = _get_ctx(ctx) + layer = _validate_layer(layer) metrics.inc("tool_calls") metrics.inc("tool_stats") mem = _get_memory(app, layer, user_id) @@ -663,6 +751,8 @@ async def memory_context_inject( metrics.inc("tool_calls") metrics.inc("tool_context_inject") + _fire_hook("auto_context", layer, {"query": "context_inject", "user_id": user_id}) + cache_key = _get_cache_key(layer, user_id) cached = _get_cached(cache_key) if cached is not None: