From 9a51fa7083c2d8e1546c475ddf97993b15346111 Mon Sep 17 00:00:00 2001 From: Toby Date: Tue, 7 Jul 2026 19:39:46 -0500 Subject: [PATCH 1/3] fix: multi-agent isolation + SQL safety at the data layer (#19) Configurable isolation policy, SQL-safety fixes, by-id scoping, and server-derived caller identity. Stacked on #28 (fix/mcp-http-auth). SQL safety (unconditional): - increment_recall_counts: allowlist the interpolated table identifier ({memory_entries, conversations, delegations}); reject anything else. - replace()/remove(): escape LIKE metacharacters (\ % _) + ESCAPE '\', so remove(old_text="%") no longer wipes the whole (agent,target) scope. Isolation policy (HEXUS_MEMORY_ISOLATION, default "shared"): - shared: reads/recall/search span all agents (trusted-fleet default). - strict: reads scoped to the caller's identity. - Cross-agent MUTATIONS (confirm/reject/remove/forget/summarize by id) are always caller-scoped, in both modes. - Resolves the empty-agent_identity asymmetry (item C): empty means "all agents" in shared / "this agent" in strict, consistently. - list_entries now treats agent_identity=None as all-agents, matching search()/count(). By-id scoping (item B): - fetch_full/confirm_entry/reject_entry/summarize_session take an optional agent_identity and filter on it; fetch_full bypasses the id-keyed CCRCache when strict-scoped so it can't leak another agent's row. Threaded through the MCP tools and the hermes handlers. Server-derived identity (item A): - HTTP transport publishes the authenticated X-Hermes-Session-Key as the caller identity (tools.current_caller ContextVar); it overrides the client-supplied agent_identity for writes/mutations, so an authenticated client can no longer act as another agent. stdio / in-process callers keep the prior behavior. Tests: SQL-safety + isolation unit tests (no DB) in test_http_auth.py; DB-backed cross-agent scoping tests in test_mcp_server.py. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 1 + hexus/__init__.py | 7 +- hexus/store.py | 215 +++++++++++++++++++++++++++++---------- mcp_server/server.py | 33 ++++++ mcp_server/tools.py | 119 +++++++++++++++++----- tests/test_http_auth.py | 129 +++++++++++++++++++++++ tests/test_mcp_server.py | 119 ++++++++++++++++++++++ 7 files changed, 543 insertions(+), 80 deletions(-) diff --git a/README.md b/README.md index 15702e4..d14b59b 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,7 @@ When running via Docker or as a standalone MCP server, you can pass the followin * `HEXUS_DB_PASS` - Used by our `compose.yml` to set the Postgres password (and the default DSN's password). * `HEXUS_TRANSPORT` - MCP transport: `"stdio"` (default) or `"http"`. * `HEXUS_AGENT_IDENTITY` - Default agent identity for tool calls that don't supply one (default: `"default"`). +* `HEXUS_MEMORY_ISOLATION` - Multi-agent read isolation: `"shared"` (default) lets any agent recall/search/read every agent's memory — a single shared knowledge base for a trusted fleet; `"strict"` scopes reads to the calling agent's own identity. Cross-agent **mutations** (confirm/reject/remove/forget/summarize by id) are always scoped to the caller in **both** modes. On the HTTP transport the caller's identity is taken server-side from the `X-Hermes-Session-Key` header and overrides any client-supplied `agent_identity`, so an authenticated client cannot act as another agent. * `HEXUS_EMBED_EAGER_LOAD` - Set to `"1"` to pre-load the local embedding model at startup (saves ~1-2s on first use). * `HEXUS_EMBED_DEVICE` - Torch device for the embedder (default: `"cpu"`). * `HEXUS_WEBHOOK_URL` / `HEXUS_WEBHOOK_SECRET` - (Optional) POST a signed webhook on memory writes. diff --git a/hexus/__init__.py b/hexus/__init__.py index 739d43d..674d83c 100644 --- a/hexus/__init__.py +++ b/hexus/__init__.py @@ -1649,7 +1649,7 @@ def _handle_confirm_memory(self, args: Dict[str, Any]) -> str: return tool_error("id must be an integer") try: - success = self._store.confirm_entry(entry_id) + success = self._store.confirm_entry(entry_id, self._agent_identity) return json.dumps({"id": entry_id, "success": success}) except Exception as exc: return json.dumps({"error": f"db: {exc}"}) @@ -1668,7 +1668,7 @@ def _handle_reject_memory(self, args: Dict[str, Any]) -> str: return tool_error("id must be an integer") try: - success = self._store.reject_entry(entry_id) + success = self._store.reject_entry(entry_id, self._agent_identity) return json.dumps({"id": entry_id, "success": success}) except Exception as exc: return json.dumps({"error": f"db: {exc}"}) @@ -1691,6 +1691,7 @@ def _handle_summarize_session(self, args: Dict[str, Any]) -> str: res = self._store.summarize_session( session_id=session_id, limit=limit, + agent_identity=self._agent_identity, ) return json.dumps(res) except Exception as exc: @@ -1710,7 +1711,7 @@ def _handle_headroom_retrieve(self, args: Dict[str, Any]) -> str: return tool_error("id must be an integer") try: - content = self._store.fetch_full(entry_id) + content = self._store.fetch_full(entry_id, self._agent_identity) if content is None: return json.dumps({"id": entry_id, "found": False, "content": None}) return json.dumps({"id": entry_id, "found": True, "content": content}) diff --git a/hexus/store.py b/hexus/store.py index fde8a29..3f106e7 100644 --- a/hexus/store.py +++ b/hexus/store.py @@ -54,6 +54,41 @@ _cross_encoder_model: Any = None _cross_encoder_lock = threading.Lock() +# Tables that carry a `metadata` JSONB column and a sequential `id`; the only +# identifiers ever interpolated into `increment_recall_counts`'s SQL. Anything +# outside this set is rejected rather than interpolated — a table name is an +# identifier, not a bindable parameter, so an allowlist is the safe equivalent +# of parameterization here. +_RECALL_COUNT_TABLES = frozenset({"memory_entries", "conversations", "delegations"}) + + +def _resolve_isolation(value: Optional[str] = None) -> str: + """Resolve the multi-agent isolation policy. + + 'shared' (default): reads/recall/search may span every agent's memory — + convenient for a trusted fleet that wants one shared + knowledge base. Cross-agent *mutations* are still + blocked (see the by-id store methods). + 'strict': reads are scoped to the caller's own agent_identity; + cross-agent reads require an explicit identity. + + Resolution: explicit `value` arg, else HEXUS_MEMORY_ISOLATION env var, + else 'shared'. + """ + raw = value if value is not None else os.environ.get("HEXUS_MEMORY_ISOLATION") + return "strict" if (raw or "").strip().lower() == "strict" else "shared" + + +def _escape_like(text: str) -> str: + r"""Escape LIKE metacharacters so caller-supplied text matches literally. + + Backslash must be escaped first (it is the ESCAPE character), then the two + wildcards `%` and `_`. Callers must pair this with an ``ESCAPE '\'`` clause. + Without it, `remove(old_text="%")` matches — and deletes — every row in the + scope, and `_` silently matches any single character. + """ + return text.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + def get_cross_encoder() -> Any: global _cross_encoder_model @@ -349,6 +384,7 @@ def __init__( max_lifetime: float = 300.0, entity_extractor_enabled: bool = True, entity_extractor_patterns: Optional[Dict[str, str]] = None, + isolation: Optional[str] = None, ): """Open a lazily-initialized, self-draining ConnectionPool. @@ -393,6 +429,11 @@ def __init__( ) self._ccr_cache = CCRCache() + # Multi-agent isolation policy. 'shared' (default) lets reads span + # agents; 'strict' scopes reads to the caller. Cross-agent mutations + # are blocked in both modes (see confirm_entry/reject_entry/fetch_full). + self.isolation = _resolve_isolation(isolation) + # Resolve vector precision configuration precision = os.environ.get("HEXUS_VECTOR_PRECISION", "float32").lower() if precision in ("float16", "fp16", "half"): @@ -904,22 +945,24 @@ def replace( hash_target = compressed if compressed is not None else new_content content_hash = hashlib.sha256(hash_target.encode("utf-8")).digest() + like_pattern = f"%{_escape_like(old_text)}%" + with self._get_pool().connection() as conn: with conn.cursor() as cur: # Find matching rows to update cache cur.execute( - """ + r""" SELECT id FROM memory_entries WHERE agent_identity = %s AND target = %s - AND content LIKE %s + AND content LIKE %s ESCAPE '\' """, - (agent_identity, target, f"%{old_text}%"), + (agent_identity, target, like_pattern), ) matching_ids = [r[0] for r in cur.fetchall()] cur.execute( - """ + r""" UPDATE memory_entries SET content = %s, embedding = %s::vector, @@ -929,7 +972,7 @@ def replace( updated_at = now() WHERE agent_identity = %s AND target = %s - AND content LIKE %s + AND content LIKE %s ESCAPE '\' """, ( new_content, @@ -939,7 +982,7 @@ def replace( content_hash, agent_identity, target, - f"%{old_text}%", + like_pattern, ), ) updated = cur.rowcount @@ -962,16 +1005,18 @@ def remove( Returns the number of rows deleted. """ + like_pattern = f"%{_escape_like(old_text)}%" + with self._get_pool().connection() as conn: with conn.cursor() as cur: cur.execute( - """ + r""" DELETE FROM memory_entries WHERE agent_identity = %s AND target = %s - AND content LIKE %s + AND content LIKE %s ESCAPE '\' """, - (agent_identity, target, f"%{old_text}%"), + (agent_identity, target, like_pattern), ) deleted = cur.rowcount conn.commit() @@ -982,17 +1027,25 @@ def remove( def list_entries( self, *, - agent_identity: str, + agent_identity: Optional[str] = None, target: Optional[str] = None, limit: int = 100, offset: int = 0, ) -> List[Dict[str, Any]]: - """List entries in an agent's scope. If target is None, both stores.""" - params: List[Any] = [agent_identity] - target_clause = "" + """List entries in an agent's scope. + + agent_identity=None/empty → list across ALL agents (matches `search` + and `count`). target=None → both stores. + """ + clauses: List[str] = [] + params: List[Any] = [] + if agent_identity: + clauses.append("agent_identity = %s") + params.append(agent_identity) if target: - target_clause = "AND target = %s" + clauses.append("target = %s") params.append(target) + where = ("WHERE " + " AND ".join(clauses)) if clauses else "" params.append(limit) params.append(offset) @@ -1002,8 +1055,7 @@ def list_entries( f""" SELECT id, agent_identity, target, content, created_at, updated_at, metadata FROM memory_entries - WHERE agent_identity = %s - {target_clause} + {where} ORDER BY updated_at DESC LIMIT %s OFFSET %s @@ -1295,18 +1347,37 @@ def hybrid_search( return rows - def fetch_full(self, memory_id: int) -> Optional[str]: - """Fetch the original full content of a memory entry, checking CCRCache first.""" - cached = self._ccr_cache.get(memory_id) - if cached is not None: - return cached + def fetch_full( + self, memory_id: int, agent_identity: Optional[str] = None + ) -> Optional[str]: + """Fetch the original full content of a memory entry, checking CCRCache first. + + When isolation is 'strict' and an `agent_identity` is supplied, the + read is scoped to that agent — the CCRCache short-circuit is bypassed + too, since the cache is keyed by id alone and would otherwise leak a + row the caller doesn't own. In 'shared' mode reads span agents (the + default trusted-fleet behavior) and the cache is used as before. + """ + scoped = self.isolation == "strict" and agent_identity is not None + + if not scoped: + cached = self._ccr_cache.get(memory_id) + if cached is not None: + return cached with self._get_pool().connection() as conn: with conn.cursor() as cur: - cur.execute( - "SELECT content FROM memory_entries WHERE id = %s", - (memory_id,), - ) + if scoped: + cur.execute( + "SELECT content FROM memory_entries " + "WHERE id = %s AND agent_identity = %s", + (memory_id, agent_identity), + ) + else: + cur.execute( + "SELECT content FROM memory_entries WHERE id = %s", + (memory_id,), + ) row = cur.fetchone() if row: content = row[0] @@ -2030,38 +2101,49 @@ def _apply_min_confidence( filtered.append(r) return filtered - def confirm_entry(self, entry_id: int) -> bool: - """Increment confirm_count in metadata JSONB for the given entry ID.""" - sql = """ - UPDATE memory_entries - SET metadata = jsonb_set( - metadata, - '{confirm_count}', - (COALESCE(metadata->>'confirm_count', '0')::int + 1)::text::jsonb - ) - WHERE id = %s + def confirm_entry( + self, entry_id: int, agent_identity: Optional[str] = None + ) -> bool: + """Increment confirm_count in metadata JSONB for the given entry ID. + + When `agent_identity` is supplied the mutation is scoped to that agent, + so one agent cannot bump another agent's entry by guessing its id. + Cross-agent mutations are blocked regardless of the isolation mode. """ - with self._get_pool().connection() as conn: - with conn.cursor() as cur: - cur.execute(sql, (entry_id,)) - updated = cur.rowcount - conn.commit() - return updated > 0 + return self._bump_entry_count(entry_id, "confirm_count", agent_identity) - def reject_entry(self, entry_id: int) -> bool: - """Increment reject_count in metadata JSONB for the given entry ID.""" - sql = """ + def reject_entry( + self, entry_id: int, agent_identity: Optional[str] = None + ) -> bool: + """Increment reject_count in metadata JSONB for the given entry ID. + + Scoped to `agent_identity` when supplied — see `confirm_entry`. + """ + return self._bump_entry_count(entry_id, "reject_count", agent_identity) + + def _bump_entry_count( + self, entry_id: int, field: str, agent_identity: Optional[str] + ) -> bool: + # `field` is one of the fixed literals passed by confirm/reject above, + # never caller input, so interpolating it into the JSONB path is safe. + if field not in ("confirm_count", "reject_count"): + raise ValueError(f"unknown count field: {field!r}") + sql = f""" UPDATE memory_entries SET metadata = jsonb_set( metadata, - '{reject_count}', - (COALESCE(metadata->>'reject_count', '0')::int + 1)::text::jsonb + '{{{field}}}', + (COALESCE(metadata->>'{field}', '0')::int + 1)::text::jsonb ) WHERE id = %s """ + params: List[Any] = [entry_id] + if agent_identity is not None: + sql += " AND agent_identity = %s\n" + params.append(agent_identity) with self._get_pool().connection() as conn: with conn.cursor() as cur: - cur.execute(sql, (entry_id,)) + cur.execute(sql, tuple(params)) updated = cur.rowcount conn.commit() return updated > 0 @@ -2069,6 +2151,12 @@ def reject_entry(self, entry_id: int) -> bool: def increment_recall_counts(self, table: str, ids: List[int]) -> None: if not ids: return + if table not in _RECALL_COUNT_TABLES: + # `table` is an SQL identifier interpolated below, so it cannot be + # a bound parameter — reject anything outside the known set instead + # of trusting the caller (guards against SQL injection via a + # derived/attacker-influenced table name). + raise ValueError(f"unknown table for recall counts: {table!r}") sql = f""" UPDATE {table} SET metadata = jsonb_set( @@ -2244,26 +2332,45 @@ def summarize_session( *, session_id: str, limit: int = 5, + agent_identity: Optional[str] = None, ) -> Dict[str, Any]: - """Compute the vector centroid of a session's turns and find the K closest turns.""" - count_sql = "SELECT COUNT(*) FROM conversations WHERE session_id = %s" - sql = """ + """Compute the vector centroid of a session's turns and find the K closest turns. + + When `agent_identity` is supplied the session lookup is scoped to that + agent, so a caller cannot summarize another agent's session by id. + """ + scope = "" if agent_identity is None else " AND agent_identity = %s" + count_sql = f"SELECT COUNT(*) FROM conversations WHERE session_id = %s{scope}" + sql = f""" WITH centroid AS ( SELECT AVG(embedding) AS vec FROM conversations - WHERE session_id = %s + WHERE session_id = %s{scope} ) SELECT id, role, content, ts, metadata, 1 - (embedding <=> (SELECT vec FROM centroid)) AS centrality_score FROM conversations - WHERE session_id = %s + WHERE session_id = %s{scope} AND embedding IS NOT NULL ORDER BY embedding <=> (SELECT vec FROM centroid) LIMIT %s """ + count_params: Tuple[Any, ...] = ( + (session_id,) if agent_identity is None else (session_id, agent_identity) + ) + if agent_identity is None: + main_params: Tuple[Any, ...] = (session_id, session_id, limit) + else: + main_params = ( + session_id, + agent_identity, + session_id, + agent_identity, + limit, + ) with self._get_pool().connection() as conn: with conn.cursor() as cur: - cur.execute(count_sql, (session_id,)) + cur.execute(count_sql, count_params) total_turns = cur.fetchone()[0] if total_turns == 0: @@ -2274,7 +2381,7 @@ def summarize_session( } with conn.cursor(row_factory=dict_row) as cur: - cur.execute(sql, (session_id, session_id, limit)) + cur.execute(sql, main_params) rows = list(cur.fetchall()) for r in rows: if r.get("ts") and hasattr(r["ts"], "isoformat"): diff --git a/mcp_server/server.py b/mcp_server/server.py index f68de2c..2392637 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -389,6 +389,34 @@ async def asgi(scope, receive, send): return asgi +def _wrap_with_identity(app): + """Wrap an ASGI app so each HTTP request's `X-Hermes-Session-Key` header + is published as the authoritative caller identity for the duration of the + request (``tools.current_caller``). + + This is what turns `agent_identity` from a client-asserted free choice into + a server-derived value (issue #19 item A): tool functions prefer this + identity over the client's `agent_identity` arg for writes/mutations, so an + authenticated client can no longer act as another agent. Runs inside the + bearer-auth gate, so only authenticated requests set an identity. Non-HTTP + scopes (lifespan/websocket) pass through untouched. + """ + + async def asgi(scope, receive, send): + if scope.get("type") != "http": + await app(scope, receive, send) + return + headers = dict(scope.get("headers") or []) + raw = headers.get(b"x-hermes-session-key", b"").decode("latin-1").strip() + token = tools.current_caller.set(raw or None) + try: + await app(scope, receive, send) + finally: + tools.current_caller.reset(token) + + return asgi + + def _build_server( store: MemoryStore, *, @@ -1253,6 +1281,11 @@ async def metrics(request): # app (rather than adding Starlette middleware) keeps /metrics # covered and leaves the streamable-http lifespan/websocket scopes # untouched. + # Publish the per-request X-Hermes-Session-Key as the caller identity + # (server-derived agent identity — issue #19 item A) before the auth + # gate hands off to the MCP app. + app = _wrap_with_identity(app) + token = os.environ.get("HEXUS_API_TOKEN") if token: logger.info("HTTP transport: bearer-token auth enabled (HEXUS_API_TOKEN).") diff --git a/mcp_server/tools.py b/mcp_server/tools.py index 4f139f4..7cd7628 100644 --- a/mcp_server/tools.py +++ b/mcp_server/tools.py @@ -23,6 +23,7 @@ import logging import os +from contextvars import ContextVar from typing import Any, Dict, List, Optional from hexus.embed import EmbeddingError, embed @@ -101,6 +102,89 @@ def _coerce_agent_identity(args: Dict[str, Any]) -> str: return default_agent_identity() +# --------------------------------------------------------------------------- +# Identity resolution for multi-agent isolation (issue #19). +# +# The *authoritative* caller identity comes from the transport. On the HTTP +# transport, server.py derives it from the authenticated `X-Hermes-Session-Key` +# header and publishes it in the `current_caller` ContextVar (below) for the +# duration of the request. When present it wins over the client-chosen +# `agent_identity` arg for writes/mutations, which stops an authenticated +# client from acting as another agent. When absent (stdio clients, +# direct/in-process calls) we fall back to the arg / env default and preserve +# the pre-#19 behavior. +# --------------------------------------------------------------------------- + +_CALLER_KEY = "_caller_identity" + +# Set per-request by the HTTP transport (server.py) from the authenticated +# `X-Hermes-Session-Key` header. It is a ContextVar rather than a tool arg so +# the ~19 tool wrappers don't each have to thread it through, and — crucially — +# so a client cannot set it: the MCP tool schema has no such parameter, and the +# only writer is the server's request middleware. stdio / in-process callers +# leave it None and keep the pre-#19 behavior. An explicit `_caller_identity` +# in args (in-process callers, tests) still takes precedence over the var. +current_caller: ContextVar[Optional[str]] = ContextVar( + "hexus_current_caller", default=None +) + + +def _caller_identity(args: Dict[str, Any]) -> Optional[str]: + """Server-derived transport identity, or None when none was set. + + Precedence: explicit `_caller_identity` in args (trusted in-process + callers) → the per-request ContextVar (HTTP transport) → None. + """ + c = args.get(_CALLER_KEY) + if isinstance(c, str) and c.strip(): + return c.strip() + v = current_caller.get() + if isinstance(v, str) and v.strip(): + return v.strip() + return None + + +def _write_identity(args: Dict[str, Any]) -> str: + """Scope an agent writes into. Transport identity is authoritative; else + fall back to the explicit arg, then the env default.""" + return _caller_identity(args) or _coerce_agent_identity(args) + + +def _scope_identity(args: Dict[str, Any]) -> Optional[str]: + """Identity to scope a by-id mutation/read to, or None to stay unscoped. + + Transport identity wins; otherwise an explicit non-empty `agent_identity` + arg. Returns None when neither is present so direct/stdio callers keep the + unscoped behavior, while the networked fleet (which always carries a + transport identity) gets cross-agent reads/mutations blocked. + """ + caller = _caller_identity(args) + if caller is not None: + return caller + a = args.get("agent_identity") + if isinstance(a, str) and a.strip(): + return a.strip() + return None + + +def _read_identity(store: MemoryStore, args: Dict[str, Any]) -> Optional[str]: + """Effective agent_identity filter for a read/search, resolving the + empty-identity asymmetry (issue #19 item C) consistently via the store's + isolation policy: + + - explicit non-empty `agent_identity` arg → scope to it (shared mode only; + in strict mode reads are always confined to the caller) + - empty/None → 'strict': the caller's own identity; 'shared': None (search + every agent — the trusted-fleet default). + """ + if getattr(store, "isolation", "shared") == "strict": + return _scope_identity(args) or default_agent_identity() + a = args.get("agent_identity") + if isinstance(a, str) and a.strip(): + return a.strip() + return None + + def _coerce_target(args: Dict[str, Any]) -> Optional[str]: """target ∈ {'memory', 'user', None}. Anything else is rejected. @@ -136,7 +220,7 @@ def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: raise ValueError(f"contents[{i}] must be a non-empty string") target = _coerce_target(args) or "memory" - agent = _coerce_agent_identity(args) + agent = _write_identity(args) doc_type = args.get("doc_type", "memory") source_url = args.get("source_url") metadata_in = args.get("metadata") @@ -229,9 +313,7 @@ def memory_recall(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: if top_k > 100: top_k = 100 - agent = args.get("agent_identity") - if isinstance(agent, str) and agent.strip() == "": - agent = None + agent = _read_identity(store, args) target = _coerce_target(args) min_similarity = float(args.get("min_similarity", 0.0)) @@ -286,9 +368,7 @@ def memory_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: Returns: {"count": N, "rows": [...], "limit": L, "offset": O} """ - agent = args.get("agent_identity") - if agent is None or (isinstance(agent, str) and not agent.strip()): - agent = _coerce_agent_identity(args) + agent = _read_identity(store, args) target = _coerce_target(args) limit = int(args.get("limit", 20)) if limit < 1: @@ -326,7 +406,7 @@ def memory_forget(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: "would_delete_id": entry_id, "hint": "pass confirm=true to actually delete", } - agent = args.get("agent_identity") or _coerce_agent_identity(args) + agent = _write_identity(args) with store._get_pool().connection() as conn: # noqa: SLF001 — admin path with conn.cursor() as cur: cur.execute( @@ -440,7 +520,7 @@ def memory_append_turn(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, An content = args.get("content") if not isinstance(content, str) or not content.strip(): raise ValueError("content must be a non-empty string") - agent = _coerce_agent_identity(args) + agent = _write_identity(args) metadata = args.get("metadata") if metadata is not None and not isinstance(metadata, dict): raise ValueError("metadata must be a dict or None") @@ -471,9 +551,7 @@ def memory_count(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: Returns: {"memory_entries": N, "conversations": M, "agent_identity": ..., "target": ...} """ - agent = args.get("agent_identity") - if agent is None or (isinstance(agent, str) and not agent.strip()): - agent = _coerce_agent_identity(args) + agent = _read_identity(store, args) target = _coerce_target(args) session_id = args.get("session_id") if isinstance(session_id, str) and session_id.strip() == "": @@ -543,9 +621,7 @@ def memory_hybrid_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, vector_weight = float(args.get("vector_weight", 0.7)) text_weight = float(args.get("text_weight", 0.3)) - agent = args.get("agent_identity") - if isinstance(agent, str) and agent.strip() == "": - agent = None + agent = _read_identity(store, args) target = _coerce_target(args) min_similarity = float(args.get("min_similarity", 0.0)) @@ -687,11 +763,7 @@ def memory_record_delegation( result = args.get("result") or "" - agent = args.get("agent_identity") - if isinstance(agent, str) and agent.strip() == "": - agent = "default" - elif not agent: - agent = "default" + agent = _write_identity(args) metadata = args.get("metadata") or {} @@ -874,7 +946,7 @@ def memory_confirm(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: except (TypeError, ValueError): raise ValueError("id must be an integer") - success = store.confirm_entry(entry_id) + success = store.confirm_entry(entry_id, _scope_identity(args)) return {"id": entry_id, "success": success} @@ -888,7 +960,7 @@ def memory_reject(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: except (TypeError, ValueError): raise ValueError("id must be an integer") - success = store.reject_entry(entry_id) + success = store.reject_entry(entry_id, _scope_identity(args)) return {"id": entry_id, "success": success} @@ -906,6 +978,7 @@ def memory_summarize_session( return store.summarize_session( session_id=session_id, limit=limit, + agent_identity=_scope_identity(args), ) @@ -923,7 +996,7 @@ def memory_retrieve(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: except (TypeError, ValueError): raise ValueError("id must be an integer") - content = store.fetch_full(entry_id) + content = store.fetch_full(entry_id, _scope_identity(args)) if content is None: return {"id": entry_id, "found": False, "content": None} return {"id": entry_id, "found": True, "content": content} diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py index 7787feb..6cc2cd7 100644 --- a/tests/test_http_auth.py +++ b/tests/test_http_auth.py @@ -105,3 +105,132 @@ async def send(msg): asyncio.run(app({"type": "lifespan"}, receive, send)) assert downstream.calls == ["lifespan"] # forwarded, no 401 + + +# ----------------------------------------------------------------------- +# Server-derived caller identity (issue #19 item A) +# ----------------------------------------------------------------------- + +from mcp_server.server import _wrap_with_identity # noqa: E402 +from mcp_server import tools # noqa: E402 + + +def _identity_scope(session_key=None): + headers = [] + if session_key is not None: + headers.append((b"x-hermes-session-key", session_key.encode("latin-1"))) + return {"type": "http", "headers": headers, "method": "POST", "path": "/mcp"} + + +def test_identity_header_published_to_contextvar(): + """The X-Hermes-Session-Key header is visible via tools.current_caller for + the duration of the request, and reset afterwards.""" + seen = {} + + async def app(scope, receive, send): + seen["caller"] = tools.current_caller.get() + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + wrapped = _wrap_with_identity(app) + _drive(wrapped, _identity_scope(session_key="marketing")) + assert seen["caller"] == "marketing" + # No leak across requests. + assert tools.current_caller.get() is None + + +def test_identity_absent_header_is_none(): + seen = {} + + async def app(scope, receive, send): + seen["caller"] = tools.current_caller.get() + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + wrapped = _wrap_with_identity(app) + _drive(wrapped, _identity_scope(session_key=None)) + assert seen["caller"] is None + + +def test_identity_contextvar_beats_client_arg_for_writes(): + """When a transport identity is set, it is authoritative for writes even + if the client passes a different agent_identity arg.""" + token = tools.current_caller.set("real-agent") + try: + assert tools._write_identity({"agent_identity": "spoofed"}) == "real-agent" + assert tools._scope_identity({"agent_identity": "spoofed"}) == "real-agent" + finally: + tools.current_caller.reset(token) + + +def test_write_identity_falls_back_to_arg_then_env(monkeypatch): + """Without a transport identity, writes fall back to the arg, then env.""" + assert tools.current_caller.get() is None + assert tools._write_identity({"agent_identity": "sales"}) == "sales" + monkeypatch.setenv("HEXUS_AGENT_IDENTITY", "fallback") + assert tools._write_identity({}) == "fallback" + + +def test_scope_identity_none_when_nothing_known(): + """By-id ops stay unscoped for direct/stdio callers with no identity.""" + assert tools.current_caller.get() is None + assert tools._scope_identity({}) is None + + +class _FakeStore: + def __init__(self, isolation): + self.isolation = isolation + + +def test_read_identity_shared_empty_means_all_agents(): + store = _FakeStore("shared") + assert tools._read_identity(store, {"agent_identity": ""}) is None + # Explicit arg is honored as a filter. + assert tools._read_identity(store, {"agent_identity": "sales"}) == "sales" + + +def test_read_identity_strict_confines_to_caller(monkeypatch): + store = _FakeStore("strict") + monkeypatch.setenv("HEXUS_AGENT_IDENTITY", "me") + # Empty → caller (env default here). + assert tools._read_identity(store, {"agent_identity": ""}) == "me" + # Even an explicit other-agent arg is overridden by the caller identity. + token = tools.current_caller.set("real") + try: + assert tools._read_identity(store, {"agent_identity": "other"}) == "real" + finally: + tools.current_caller.reset(token) + + +# ----------------------------------------------------------------------- +# SQL-safety helpers + isolation policy (issue #19), no DB required +# ----------------------------------------------------------------------- + +from hexus.store import _escape_like, _resolve_isolation # noqa: E402 + + +def test_escape_like_neutralizes_wildcards(): + assert _escape_like("%") == r"\%" + assert _escape_like("_") == r"\_" + assert _escape_like("a_b%c") == r"a\_b\%c" + # Backslash escaped first so it can't double-escape a following wildcard. + assert _escape_like("\\%") == r"\\\%" + # Plain text is untouched. + assert _escape_like("hello world") == "hello world" + + +def test_resolve_isolation_default_shared(monkeypatch): + monkeypatch.delenv("HEXUS_MEMORY_ISOLATION", raising=False) + assert _resolve_isolation() == "shared" + + +def test_resolve_isolation_strict_from_env(monkeypatch): + monkeypatch.setenv("HEXUS_MEMORY_ISOLATION", "STRICT") + assert _resolve_isolation() == "strict" + monkeypatch.setenv("HEXUS_MEMORY_ISOLATION", "shared") + assert _resolve_isolation() == "shared" + # Explicit arg wins over env. + assert _resolve_isolation("strict") == "strict" + # Unknown value falls back to shared. + monkeypatch.setenv("HEXUS_MEMORY_ISOLATION", "banana") + assert _resolve_isolation() == "shared" diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index c9269a2..54319e5 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1361,3 +1361,122 @@ def __exit__(self, exc_type, exc_val, exc_tb): assert res["status"] == "ok" assert "low_confidence" in res assert "cooccurring" in res + + +# ----------------------------------------------------------------------- +# Multi-agent isolation + SQL safety (issue #19) +# ----------------------------------------------------------------------- + +_EMB = [0.1] * 384 + + +def test_increment_recall_counts_rejects_unknown_table(store): + """The table identifier is interpolated into SQL, so anything outside the + known set must be rejected before a query is built (no DB round-trip).""" + import pytest as _pytest + + with _pytest.raises(ValueError): + store.increment_recall_counts( + "memory_entries; DROP TABLE memory_entries; --", [1] + ) + # A known table with no ids is a no-op and must not raise. + store.increment_recall_counts("memory_entries", []) + + +def test_remove_escapes_like_wildcards(store): + """`remove(old_text="%")` must delete only rows containing a literal '%', + not every row in the (agent, target) scope.""" + agent = agent_of(store) + store.add(agent_identity=agent, target="memory", content="100% done", embedding=_EMB) + store.add(agent_identity=agent, target="memory", content="all clear", embedding=_EMB) + + deleted = store.remove(agent_identity=agent, target="memory", old_text="%") + assert deleted == 1 + + remaining = store.list_entries(agent_identity=agent, target="memory", limit=50) + contents = [r["content"] for r in remaining] + assert "all clear" in contents + assert "100% done" not in contents + + +def test_remove_escapes_underscore_wildcard(store): + """'_' must match literally, not 'any single character'.""" + agent = agent_of(store) + store.add(agent_identity=agent, target="memory", content="a_b marker", embedding=_EMB) + store.add(agent_identity=agent, target="memory", content="axb marker", embedding=_EMB) + + deleted = store.remove(agent_identity=agent, target="memory", old_text="a_b") + assert deleted == 1 + remaining = [r["content"] for r in store.list_entries(agent_identity=agent, limit=50)] + assert "axb marker" in remaining + + +def test_confirm_reject_scoped_to_agent(store): + """A caller cannot confirm/reject another agent's entry by id.""" + agent = agent_of(store) + row_id = store.add( + agent_identity=agent, target="memory", content="scoped entry", embedding=_EMB + ) + assert row_id is not None + + # Wrong identity → no row matches → no bump. + assert store.confirm_entry(row_id, "someone-else") is False + assert store.reject_entry(row_id, "someone-else") is False + # Correct identity → bump succeeds. + assert store.confirm_entry(row_id, agent) is True + assert store.reject_entry(row_id, agent) is True + # Unscoped (None) still works — the in-process/stdio compatibility path. + assert store.confirm_entry(row_id, None) is True + + +def test_summarize_session_scoped_to_agent(store): + """summarize_session scoped to the wrong agent sees no turns.""" + from mcp_server import tools + + agent = agent_of(store) + session_id = "iso-session-" + agent + for role, text in (("user", "hello there"), ("assistant", "general kenobi")): + tools.memory_append_turn( + store, + {"session_id": session_id, "agent_identity": agent, "role": role, "content": text}, + ) + + wrong = store.summarize_session(session_id=session_id, agent_identity="not-me") + assert wrong["turn_count"] == 0 + + right = store.summarize_session(session_id=session_id, agent_identity=agent) + assert right["turn_count"] == 2 + + +def test_fetch_full_shared_reads_cross_agent(store): + """Default 'shared' isolation: reads by id span agents (the feature).""" + agent = agent_of(store) + row_id = store.add( + agent_identity=agent, target="memory", content="shared content", embedding=_EMB + ) + # A different identity can still read it in shared mode. + assert store.fetch_full(row_id, "another-agent") == "shared content" + + +def test_fetch_full_strict_scopes_reads(store): + """'strict' isolation: fetch_full is confined to the caller's identity, + and the id-keyed CCRCache cannot leak another agent's row.""" + from hexus.store import MemoryStore + + agent = agent_of(store) + row_id = store.add( + agent_identity=agent, target="memory", content="secret content", embedding=_EMB + ) + + strict = MemoryStore(os.environ["PG_TEST_DSN"], isolation="strict") + try: + assert strict.isolation == "strict" + # Wrong identity → no read, even though the id exists. + assert strict.fetch_full(row_id, "not-the-owner") is None + # Correct identity → read succeeds. + assert strict.fetch_full(row_id, agent) == "secret content" + # And after a legit read populated the cache, a wrong-identity read + # still must not be served from it. + assert strict.fetch_full(row_id, "not-the-owner") is None + finally: + strict.close() From c42f88ca7a9cae1b78dbc0120603fc65920f5de3 Mon Sep 17 00:00:00 2001 From: Toby Date: Sat, 11 Jul 2026 14:30:23 -0500 Subject: [PATCH 2/3] fix: clean SQL building in summarize_session, verify ContextVar propagation warning, update CI packaging check, and bump version to 0.9.2 --- .github/workflows/ci.yml | 2 ++ hexus/store.py | 29 ++++++++++++----------------- mcp_server/server.py | 2 ++ mcp_server/tools.py | 9 +++++++++ pyproject.toml | 2 +- 5 files changed, 26 insertions(+), 18 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 88c36cc..e3193db 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -81,6 +81,8 @@ jobs: "mcp_server/__init__.py", "mcp_server/cli.py", "mcp_server/import_cli.py", + "mcp_server/server.py", + "mcp_server/tools.py", ] missing = [r for r in required if r not in names] assert not missing, f"{whl} is missing: {missing}" diff --git a/hexus/store.py b/hexus/store.py index 3f106e7..f1abcb8 100644 --- a/hexus/store.py +++ b/hexus/store.py @@ -2339,35 +2339,30 @@ def summarize_session( When `agent_identity` is supplied the session lookup is scoped to that agent, so a caller cannot summarize another agent's session by id. """ - scope = "" if agent_identity is None else " AND agent_identity = %s" - count_sql = f"SELECT COUNT(*) FROM conversations WHERE session_id = %s{scope}" + clauses = ["session_id = %s"] + params = [session_id] + if agent_identity is not None: + clauses.append("agent_identity = %s") + params.append(agent_identity) + + where_clause = " AND ".join(clauses) + count_sql = f"SELECT COUNT(*) FROM conversations WHERE {where_clause}" sql = f""" WITH centroid AS ( SELECT AVG(embedding) AS vec FROM conversations - WHERE session_id = %s{scope} + WHERE {where_clause} ) SELECT id, role, content, ts, metadata, 1 - (embedding <=> (SELECT vec FROM centroid)) AS centrality_score FROM conversations - WHERE session_id = %s{scope} + WHERE {where_clause} AND embedding IS NOT NULL ORDER BY embedding <=> (SELECT vec FROM centroid) LIMIT %s """ - count_params: Tuple[Any, ...] = ( - (session_id,) if agent_identity is None else (session_id, agent_identity) - ) - if agent_identity is None: - main_params: Tuple[Any, ...] = (session_id, session_id, limit) - else: - main_params = ( - session_id, - agent_identity, - session_id, - agent_identity, - limit, - ) + count_params = tuple(params) + main_params = tuple(params) + tuple(params) + (limit,) with self._get_pool().connection() as conn: with conn.cursor() as cur: cur.execute(count_sql, count_params) diff --git a/mcp_server/server.py b/mcp_server/server.py index 2392637..2ae097b 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1269,6 +1269,8 @@ def memory_stats() -> Dict[str, Any]: def get_asgi_app_with_metrics(*args, **kwargs): app = _orig_get_asgi_app(*args, **kwargs) from starlette.responses import Response + from . import tools + tools.http_transport_active = True async def metrics(request): return Response(_generate_metrics(store), media_type="text/plain") diff --git a/mcp_server/tools.py b/mcp_server/tools.py index 7cd7628..a2fe351 100644 --- a/mcp_server/tools.py +++ b/mcp_server/tools.py @@ -128,6 +128,9 @@ def _coerce_agent_identity(args: Dict[str, Any]) -> str: "hexus_current_caller", default=None ) +# Flag to indicate if the HTTP transport is active. Used to audit ContextVar propagation. +http_transport_active = False + def _caller_identity(args: Dict[str, Any]) -> Optional[str]: """Server-derived transport identity, or None when none was set. @@ -141,6 +144,12 @@ def _caller_identity(args: Dict[str, Any]) -> Optional[str]: v = current_caller.get() if isinstance(v, str) and v.strip(): return v.strip() + if http_transport_active: + logger.warning( + "current_caller ContextVar is None inside tool invocation on HTTP transport. " + "This suggests ContextVar propagation through FastMCP dispatch has failed, " + "or the request was missing the X-Hermes-Session-Key header." + ) return None diff --git a/pyproject.toml b/pyproject.toml index 092dc1b..94670e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "hexus" -version = "0.9.1" +version = "0.9.2" description = "Hexus: Postgres + hexus memory provider plugin for hermes-agent AND a standalone MCP server. Local BERT (MiniLM-L6-v2) embeddings replace the upstream HTTP embedder; multi-agent storage with per-minion themes, async writer, no LLM in the memory hot path." readme = "README.md" requires-python = ">=3.11" From 37dc24086d61963158f943e0f42a323aded8c313 Mon Sep 17 00:00:00 2001 From: Toby Date: Sat, 11 Jul 2026 14:31:22 -0500 Subject: [PATCH 3/3] style: run ruff format to fix styling check failures --- hexus/store.py | 4 +--- mcp_server/server.py | 1 + tests/test_mcp_server.py | 27 +++++++++++++++++++++------ 3 files changed, 23 insertions(+), 9 deletions(-) diff --git a/hexus/store.py b/hexus/store.py index f1abcb8..be6e29d 100644 --- a/hexus/store.py +++ b/hexus/store.py @@ -2112,9 +2112,7 @@ def confirm_entry( """ return self._bump_entry_count(entry_id, "confirm_count", agent_identity) - def reject_entry( - self, entry_id: int, agent_identity: Optional[str] = None - ) -> bool: + def reject_entry(self, entry_id: int, agent_identity: Optional[str] = None) -> bool: """Increment reject_count in metadata JSONB for the given entry ID. Scoped to `agent_identity` when supplied — see `confirm_entry`. diff --git a/mcp_server/server.py b/mcp_server/server.py index 2ae097b..e357856 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1270,6 +1270,7 @@ def get_asgi_app_with_metrics(*args, **kwargs): app = _orig_get_asgi_app(*args, **kwargs) from starlette.responses import Response from . import tools + tools.http_transport_active = True async def metrics(request): diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 54319e5..cd2630e 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1387,8 +1387,12 @@ def test_remove_escapes_like_wildcards(store): """`remove(old_text="%")` must delete only rows containing a literal '%', not every row in the (agent, target) scope.""" agent = agent_of(store) - store.add(agent_identity=agent, target="memory", content="100% done", embedding=_EMB) - store.add(agent_identity=agent, target="memory", content="all clear", embedding=_EMB) + store.add( + agent_identity=agent, target="memory", content="100% done", embedding=_EMB + ) + store.add( + agent_identity=agent, target="memory", content="all clear", embedding=_EMB + ) deleted = store.remove(agent_identity=agent, target="memory", old_text="%") assert deleted == 1 @@ -1402,12 +1406,18 @@ def test_remove_escapes_like_wildcards(store): def test_remove_escapes_underscore_wildcard(store): """'_' must match literally, not 'any single character'.""" agent = agent_of(store) - store.add(agent_identity=agent, target="memory", content="a_b marker", embedding=_EMB) - store.add(agent_identity=agent, target="memory", content="axb marker", embedding=_EMB) + store.add( + agent_identity=agent, target="memory", content="a_b marker", embedding=_EMB + ) + store.add( + agent_identity=agent, target="memory", content="axb marker", embedding=_EMB + ) deleted = store.remove(agent_identity=agent, target="memory", old_text="a_b") assert deleted == 1 - remaining = [r["content"] for r in store.list_entries(agent_identity=agent, limit=50)] + remaining = [ + r["content"] for r in store.list_entries(agent_identity=agent, limit=50) + ] assert "axb marker" in remaining @@ -1438,7 +1448,12 @@ def test_summarize_session_scoped_to_agent(store): for role, text in (("user", "hello there"), ("assistant", "general kenobi")): tools.memory_append_turn( store, - {"session_id": session_id, "agent_identity": agent, "role": role, "content": text}, + { + "session_id": session_id, + "agent_identity": agent, + "role": role, + "content": text, + }, ) wrong = store.summarize_session(session_id=session_id, agent_identity="not-me")