Bind tool calls to the active engine session

engine.py

@@ -64,9 +64,6 @@ def __init__(self, config: LCMConfig | None = None,
         # genuinely new messages (appended after compaction) get ingested.
         self._ingest_cursor: int = 0
 
-        # Wire tool handlers
-        lcm_tools.set_engine(self)
-
         # State required by ContextEngine ABC and run_agent.py compatibility
         self.model = ""
         self.base_url = ""
@@ -263,7 +260,7 @@ def handle_tool_call(self, name: str, args: Dict[str, Any], **kwargs) -> str:
         }
         handler = handlers.get(name)
         if handler:
-            return handler(args)
+            return handler(args, engine=self)
         return json.dumps({"error": f"Unknown LCM tool: {name}"})
 
     def get_status(self) -> Dict[str, Any]:

tests/test_lcm_engine.py

@@ -5,6 +5,7 @@
 
 from agent.context_engine import ContextEngine
 from hermes_lcm.config import LCMConfig
+from hermes_lcm.dag import SummaryNode
 from hermes_lcm.engine import LCMEngine
 
 
@@ -387,3 +388,62 @@ def test_handle_describe_overview(self, engine):
     def test_handle_unknown_tool(self, engine):
         result = json.loads(engine.handle_tool_call("unknown_tool", {}))
         assert "error" in result
+
+    def test_tool_dispatch_is_bound_to_engine_instance(self, tmp_path):
+        config_a = LCMConfig(database_path=str(tmp_path / "a.db"))
+        config_b = LCMConfig(database_path=str(tmp_path / "b.db"))
+
+        engine_a = LCMEngine(config=config_a)
+        engine_a._session_id = "session-a"
+        engine_b = LCMEngine(config=config_b)
+        engine_b._session_id = "session-b"
+
+        engine_a._store.append("session-a", {"role": "user", "content": "alpha project"})
+        engine_b._store.append("session-b", {"role": "user", "content": "beta project"})
+
+        result_a = json.loads(engine_a.handle_tool_call("lcm_grep", {"query": "alpha"}))
+        result_b = json.loads(engine_b.handle_tool_call("lcm_grep", {"query": "beta"}))
+
+        assert result_a["total_results"] == 1
+        assert result_b["total_results"] == 1
+        assert "alpha" in result_a["results"][0]["snippet"]
+        assert "beta" in result_b["results"][0]["snippet"]
+
+    def test_describe_and_expand_are_session_scoped(self, engine):
+        node_id = engine._dag.add_node(
+            SummaryNode(
+                session_id="session-a",
+                depth=0,
+                summary="secret summary",
+                token_count=10,
+                source_token_count=20,
+                source_ids=[],
+                source_type="messages",
+                created_at=0,
+            )
+        )
+
+        engine._session_id = "session-b"
+
+        describe = json.loads(engine.handle_tool_call("lcm_describe", {"node_id": node_id}))
+        expand = json.loads(engine.handle_tool_call("lcm_expand", {"node_id": node_id}))
+
+        assert "error" in describe
+        assert "error" in expand
+
+    def test_describe_overview_includes_sparse_high_depth_nodes(self, engine):
+        engine._dag.add_node(
+            SummaryNode(
+                session_id="test-session",
+                depth=2,
+                summary="durable summary",
+                token_count=100,
+                source_token_count=500,
+                source_ids=[],
+                source_type="messages",
+                created_at=0,
+            )
+        )
+
+        overview = json.loads(engine.handle_tool_call("lcm_describe", {}))
+        assert "d2" in overview["depths"]

tools.py

@@ -1,5 +1,7 @@
 """Tool handlers for LCM — the code that runs when the LLM calls each tool."""
 
+from __future__ import annotations
+
 import json
 import logging
 from typing import Any, Dict, TYPE_CHECKING
@@ -9,172 +11,180 @@
 
 logger = logging.getLogger(__name__)
 
-# The engine instance is set by the engine on init
-_engine: "LCMEngine | None" = None
+
+def _require_engine(kwargs: Dict[str, Any]) -> "LCMEngine | None":
+    engine = kwargs.get("engine")
+    return engine if engine is not None else None
 
 
-def set_engine(engine: "LCMEngine") -> None:
-    global _engine
-    _engine = engine
+def _get_session_node(engine: "LCMEngine", node_id: int):
+    node = engine._dag.get_node(node_id)
+    if node is None or node.session_id != engine._session_id:
+        return None
+    return node
 
 
 def lcm_grep(args: Dict[str, Any], **kwargs) -> str:
-    """Search across the full DAG and raw messages."""
-    if _engine is None:
+    """Search across the full DAG and raw messages for the current session."""
+    engine = _require_engine(kwargs)
+    if engine is None:
         return json.dumps({"error": "LCM engine not initialized"})
 
     query = args.get("query", "").strip()
     if not query:
         return json.dumps({"error": "No query provided"})
 
     limit = args.get("limit", 10)
-    session_id = _engine._session_id
-
+    session_id = engine._session_id
     results = []
 
-    # Search raw messages
     try:
-        msg_hits = _engine._store.search(query, session_id=session_id, limit=limit)
+        msg_hits = engine._store.search(query, session_id=session_id, limit=limit)
         for hit in msg_hits:
-            results.append({
-                "type": "message",
-                "depth": "raw",
-                "store_id": hit["store_id"],
-                "role": hit["role"],
-                "snippet": hit.get("snippet", hit.get("content", "")[:200]),
-            })
-    except Exception as e:
-        logger.debug("Message search failed: %s", e)
-
-    # Search summary nodes
+            results.append(
+                {
+                    "type": "message",
+                    "depth": "raw",
+                    "store_id": hit["store_id"],
+                    "role": hit["role"],
+                    "snippet": hit.get("snippet", hit.get("content", "")[:200]),
+                }
+            )
+    except Exception as exc:
+        logger.debug("Message search failed: %s", exc)
+
     try:
-        node_hits = _engine._dag.search(query, session_id=session_id, limit=limit)
+        node_hits = engine._dag.search(query, session_id=session_id, limit=limit)
         for node in node_hits:
-            results.append({
-                "type": "summary",
-                "depth": f"d{node.depth}",
-                "node_id": node.node_id,
-                "snippet": node.summary[:300],
-                "token_count": node.token_count,
-                "expand_hint": node.expand_hint,
-            })
-    except Exception as e:
-        logger.debug("Node search failed: %s", e)
-
-    # Sort by relevance (raw messages first, then by depth ascending)
-    results.sort(key=lambda r: (0 if r["type"] == "message" else 1, r.get("depth", "")))
-
-    return json.dumps({
-        "query": query,
-        "total_results": len(results),
-        "results": results[:limit],
-    })
+            results.append(
+                {
+                    "type": "summary",
+                    "depth": f"d{node.depth}",
+                    "node_id": node.node_id,
+                    "snippet": node.summary[:300],
+                    "token_count": node.token_count,
+                    "expand_hint": node.expand_hint,
+                }
+            )
+    except Exception as exc:
+        logger.debug("Node search failed: %s", exc)
+
+    results.sort(key=lambda result: (0 if result["type"] == "message" else 1, result.get("depth", "")))
+    return json.dumps({"query": query, "total_results": len(results), "results": results[:limit]})
 
 
 def lcm_describe(args: Dict[str, Any], **kwargs) -> str:
     """Inspect a node's subtree or get session DAG overview."""
-    if _engine is None:
+    engine = _require_engine(kwargs)
+    if engine is None:
         return json.dumps({"error": "LCM engine not initialized"})
 
     node_id = args.get("node_id")
-    session_id = _engine._session_id
+    session_id = engine._session_id
 
     if node_id is not None:
-        info = _engine._dag.describe_subtree(node_id)
+        node = _get_session_node(engine, node_id)
+        if node is None:
+            return json.dumps({"error": f"Node {node_id} not found in current session"})
+        info = engine._dag.describe_subtree(node_id)
         return json.dumps(info)
 
-    # Session overview: count nodes at each depth
+    all_nodes = engine._dag.get_session_nodes(session_id)
     overview = {
         "session_id": session_id,
-        "store_message_count": _engine._store.get_session_count(session_id),
+        "store_message_count": engine._store.get_session_count(session_id),
         "depths": {},
     }
 
-    for depth in range(10):  # check up to d9
-        count = _engine._dag.count_at_depth(session_id, depth)
-        if count == 0 and depth > 0:
-            break
-        if count > 0:
-            nodes = _engine._dag.get_session_nodes(session_id, depth=depth, limit=100)
-            overview["depths"][f"d{depth}"] = {
-                "count": count,
-                "total_tokens": sum(n.token_count for n in nodes),
-                "total_source_tokens": sum(n.source_token_count for n in nodes),
-                "nodes": [
-                    {
-                        "node_id": n.node_id,
-                        "token_count": n.token_count,
-                        "expand_hint": n.expand_hint,
-                    }
-                    for n in nodes[:20]  # cap at 20 per depth for display
-                ],
-            }
+    for depth in sorted({node.depth for node in all_nodes}):
+        nodes = [node for node in all_nodes if node.depth == depth]
+        overview["depths"][f"d{depth}"] = {
+            "count": len(nodes),
+            "total_tokens": sum(node.token_count for node in nodes),
+            "total_source_tokens": sum(node.source_token_count for node in nodes),
+            "nodes": [
+                {
+                    "node_id": node.node_id,
+                    "token_count": node.token_count,
+                    "expand_hint": node.expand_hint,
+                }
+                for node in nodes[:20]
+            ],
+        }
 
     return json.dumps(overview)
 
 
 def lcm_expand(args: Dict[str, Any], **kwargs) -> str:
     """Expand a summary node to its source content."""
-    if _engine is None:
+    engine = _require_engine(kwargs)
+    if engine is None:
         return json.dumps({"error": "LCM engine not initialized"})
 
     node_id = args.get("node_id")
     if node_id is None:
         return json.dumps({"error": "node_id is required"})
 
-    max_tokens = args.get("max_tokens", 4000)
+    node = _get_session_node(engine, node_id)
+    if node is None:
+        return json.dumps({"error": f"Node {node_id} not found in current session"})
 
-    node = _engine._dag.get_node(node_id)
-    if not node:
-        return json.dumps({"error": f"Node {node_id} not found"})
+    max_tokens = args.get("max_tokens", 4000)
 
     if node.source_type == "messages":
-        # Expand to raw messages
-        messages = []
         from .tokens import count_tokens
+
+        messages = []
         budget_used = 0
-        for sid in node.source_ids:
-            stored = _engine._store.get(sid)
-            if not stored:
+        for store_id in node.source_ids:
+            stored = engine._store.get(store_id)
+            if not stored or stored.get("session_id") != engine._session_id:
                 continue
             content = stored.get("content", "")
             msg_tokens = count_tokens(content)
             if budget_used + msg_tokens > max_tokens and messages:
-                messages.append({
-                    "note": f"Truncated — {len(node.source_ids) - len(messages)} more messages available",
-                })
+                messages.append(
+                    {
+                        "note": f"Truncated — {len(node.source_ids) - len(messages)} more messages available",
+                    }
+                )
                 break
-            messages.append({
-                "store_id": stored["store_id"],
-                "role": stored["role"],
-                "content": content[:2000] if len(content) > 2000 else content,
-            })
-            budget_used += msg_tokens
-
-        return json.dumps({
-            "node_id": node_id,
-            "depth": node.depth,
-            "source_type": "messages",
-            "expanded": messages,
-        })
-
-    elif node.source_type == "nodes":
-        # Expand to child summaries
-        children = _engine._dag.get_source_nodes(node)
-        return json.dumps({
-            "node_id": node_id,
-            "depth": node.depth,
-            "source_type": "nodes",
-            "expanded": [
+            messages.append(
                 {
-                    "node_id": c.node_id,
-                    "depth": c.depth,
-                    "summary": c.summary[:1000],
-                    "token_count": c.token_count,
-                    "expand_hint": c.expand_hint,
+                    "store_id": stored["store_id"],
+                    "role": stored["role"],
+                    "content": content[:2000] if len(content) > 2000 else content,
                 }
-                for c in children
-            ],
-        })
+            )
+            budget_used += msg_tokens
+
+        return json.dumps(
+            {
+                "node_id": node_id,
+                "depth": node.depth,
+                "source_type": "messages",
+                "expanded": messages,
+            }
+        )
+
+    if node.source_type == "nodes":
+        children = [child for child in engine._dag.get_source_nodes(node) if child.session_id == engine._session_id]
+        return json.dumps(
+            {
+                "node_id": node_id,
+                "depth": node.depth,
+                "source_type": "nodes",
+                "expanded": [
+                    {
+                        "node_id": child.node_id,
+                        "depth": child.depth,
+                        "summary": child.summary[:1000],
+                        "token_count": child.token_count,
+                        "expand_hint": child.expand_hint,
+                    }
+                    for child in children
+                ],
+            }
+        )
 
     return json.dumps({"error": f"Unknown source_type: {node.source_type}"})