From 7c80be686ba75adfa6c23b7aa5000ae1b2bf79f1 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 25 May 2026 21:59:45 +0000 Subject: [PATCH 1/4] Harden runtime: providers, transport, security, metrics, IO Fixes from a multi-agent system review (isolated, non-overlapping scopes): - mem0: drop instance-method lru_cache (leak + stale backend on reconfigure); attribute post-add fallback to the actual write instead of an arbitrary row - metrics: bound error-type/model label cardinality with __other__ overflow - atomic_io: fsync parent directory after os.replace for crash durability - reconcile: partition claims by scope so cross-user matches aren't false conflicts - clients: per-target error isolation; malformed client config no longer aborts the batch or clobbers an unreadable file - scope_registry: enable WAL + busy_timeout; surface status-write failures - api/oauth: reject PLAIN PKCE (S256 only); exact-hostname redirect_uri check; rate-limit the authorize endpoint - http_client/transport: scope the start lock to spawn only; verify owner by pid; map 401/403 to an actionable auth error https://claude.ai/code/session_01BKjtun7hVwxh6Hv3YLdrww --- agentmemory/api.py | 38 ++++++- agentmemory/clients.py | 142 ++++++++++++++++---------- agentmemory/oauth.py | 4 +- agentmemory/providers/mem0.py | 116 ++++++++++++++++++--- agentmemory/runtime/atomic_io.py | 23 +++++ agentmemory/runtime/http_client.py | 39 +++++-- agentmemory/runtime/metrics.py | 42 +++++++- agentmemory/runtime/reconcile.py | 75 ++++++++------ agentmemory/runtime/scope_registry.py | 32 +++++- agentmemory/runtime/transport.py | 5 + tests/test_agentmemory_clients.py | 46 +++++++++ tests/test_agentmemory_http_client.py | 52 ++++++++++ tests/test_agentmemory_reconcile.py | 24 +++++ tests/test_agentmemory_transport.py | 7 ++ tests/test_mem0_provider.py | 44 ++++++++ tests/test_scope_registry.py | 15 +++ 16 files changed, 584 insertions(+), 120 deletions(-) diff --git a/agentmemory/api.py b/agentmemory/api.py index 63b4ced..0c8251f 100644 --- a/agentmemory/api.py +++ b/agentmemory/api.py @@ -140,6 +140,35 @@ def _ui_disabled() -> bool: return os.environ.get("AGENTMEMORY_DISABLE_UI", "").strip() in {"1", "true", "yes"} +_LOOPBACK_HOSTS = {"127.0.0.1", "::1", "localhost"} + + +def _redirect_uri_allowed(redirect_uri: str) -> bool: + """Validate an OAuth redirect_uri. + + Parses the URI and matches the hostname exactly so that prefix-spoofing + tricks (e.g. ``http://127.0.0.1.evil.com``, ``http://localhost@evil.com``) + are rejected. https is allowed with any host; http is only allowed for + loopback hosts. Any userinfo (``@``) is rejected outright. + """ + if not redirect_uri: + return False + try: + parsed = urlparse(redirect_uri) + except ValueError: + return False + # Reject userinfo tricks like http://127.0.0.1@evil.com + if parsed.username is not None or parsed.password is not None or "@" in (parsed.netloc or ""): + return False + scheme = (parsed.scheme or "").lower() + hostname = (parsed.hostname or "").lower() + if scheme == "https": + return bool(hostname) + if scheme == "http": + return hostname in _LOOPBACK_HOSTS + return False + + class Handler(BaseHTTPRequestHandler): server_version = "AgentMemory/1.0" @@ -330,10 +359,15 @@ def _p(name: str, default: str = "") -> str: scope = _p("scope") or None resource = _p("resource") or None + # Throttle the unauthenticated authorize endpoint (keyed by client_id) + # so it cannot be flooded to grow the in-memory auth-code store. + if not self._require_rate_limit(f"oauth-authorize:{given_client_id or 'anonymous'}"): + return + if given_client_id != expected_client_id: self._send(400, {"error": "invalid_client"}) return - if not (redirect_uri.startswith("https://") or redirect_uri.startswith("http://127.0.0.1") or redirect_uri.startswith("http://localhost")): + if not _redirect_uri_allowed(redirect_uri): self._send(400, {"error": "invalid_request", "error_description": "redirect_uri must be https or loopback"}) return if response_type != "code": @@ -342,7 +376,7 @@ def _p(name: str, default: str = "") -> str: if not code_challenge: self._send(400, {"error": "invalid_request", "error_description": "code_challenge required"}) return - if code_challenge_method.upper() not in {"S256", "PLAIN"}: + if code_challenge_method.upper() != "S256": self._send(400, {"error": "invalid_request", "error_description": "unsupported code_challenge_method"}) return diff --git a/agentmemory/clients.py b/agentmemory/clients.py index fbb1330..e28d8b3 100644 --- a/agentmemory/clients.py +++ b/agentmemory/clients.py @@ -151,10 +151,17 @@ def backup_file(path: Path, backup_dir: Path) -> None: shutil.copy2(path, backup_dir / path.name) +class ConfigParseError(Exception): + """Raised when an existing client config file cannot be parsed as JSON.""" + + def load_json(path: Path, default: dict[str, Any]) -> dict[str, Any]: if not path.exists(): return json.loads(json.dumps(default)) - return json.loads(path.read_text(encoding="utf-8")) + try: + return json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ConfigParseError(f"existing config at {path} is not valid JSON: {exc}") from exc def write_json(path: Path, payload: dict[str, Any]) -> None: @@ -565,20 +572,47 @@ def disconnect_cline(backup_dir: Path) -> dict[str, Any]: return {"target": "cline", "status": "skipped", "reason": "not detected"} +def isolated(target: str, fn: "Any", *args: "Any", **kwargs: "Any") -> dict[str, Any]: + """Run a single target operation, converting any failure into an error result. + + Keeps batch operations resilient: one corrupt config, permission error, + missing directory, or subprocess failure becomes an error entry for that + target instead of aborting the whole batch with a traceback. + """ + try: + return fn(*args, **kwargs) + except ConfigParseError as exc: + return { + "target": target, + "status": "error", + "health": "error", + "reason": str(exc), + "details": str(exc), + } + except Exception as exc: # noqa: BLE001 - isolate any per-target failure + return { + "target": target, + "status": "error", + "health": "error", + "reason": f"{type(exc).__name__}: {exc}", + "details": f"{type(exc).__name__}: {exc}", + } + + def connect_all() -> dict[str, Any]: timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") backup_dir = BACKUP_ROOT / timestamp results = [ - connect_codex(), - connect_claude_code(), - connect_claude_desktop(backup_dir), - connect_gemini_cli(), - connect_qwen_cli(), - connect_cursor(backup_dir), - connect_vscode_copilot(backup_dir), - connect_roo_code(backup_dir), - connect_kilocode(backup_dir), - connect_cline(backup_dir), + isolated("codex", connect_codex), + isolated("claude-code", connect_claude_code), + isolated("claude-desktop", connect_claude_desktop, backup_dir), + isolated("gemini-cli", connect_gemini_cli), + isolated("qwen-cli", connect_qwen_cli), + isolated("cursor", connect_cursor, backup_dir), + isolated("copilot-vscode", connect_vscode_copilot, backup_dir), + isolated("roo-code", connect_roo_code, backup_dir), + isolated("kilocode", connect_kilocode, backup_dir), + isolated("cline", connect_cline, backup_dir), ] return { "server_name": SERVER_NAME, @@ -592,16 +626,16 @@ def disconnect_all() -> dict[str, Any]: timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") backup_dir = BACKUP_ROOT / timestamp results = [ - disconnect_codex(), - disconnect_claude_code(), - disconnect_claude_desktop(backup_dir), - disconnect_gemini_cli(), - disconnect_qwen_cli(), - disconnect_cursor(backup_dir), - disconnect_vscode_copilot(backup_dir), - disconnect_roo_code(backup_dir), - disconnect_kilocode(backup_dir), - disconnect_cline(backup_dir), + isolated("codex", disconnect_codex), + isolated("claude-code", disconnect_claude_code), + isolated("claude-desktop", disconnect_claude_desktop, backup_dir), + isolated("gemini-cli", disconnect_gemini_cli), + isolated("qwen-cli", disconnect_qwen_cli), + isolated("cursor", disconnect_cursor, backup_dir), + isolated("copilot-vscode", disconnect_vscode_copilot, backup_dir), + isolated("roo-code", disconnect_roo_code, backup_dir), + isolated("kilocode", disconnect_kilocode, backup_dir), + isolated("cline", disconnect_cline, backup_dir), ] return { "server_name": SERVER_NAME, @@ -618,19 +652,19 @@ def status_all() -> dict[str, Any]: cline_vscode_mcp = cline_vscode_mcp_path() cline_cursor_mcp = cline_cursor_mcp_path() results = [ - cli_status("codex", "& codex.ps1 mcp list"), - cli_status("claude-code", "& claude mcp list"), - config_status(claude_desktop_config, "mcpServers", "claude-desktop"), - cli_status("gemini-cli", "& gemini.ps1 mcp list"), - cli_status("qwen-cli", "& qwen.ps1 mcp list"), - config_status(CURSOR_MCP, "mcpServers", "cursor"), - config_status(vscode_mcp, "servers", "copilot-vscode"), - config_status(roo_mcp, "mcpServers", "roo-code"), - config_status(kilo_mcp, "mcpServers", "kilocode"), - config_status(cline_vscode_mcp, "mcpServers", "cline"), + isolated("codex", cli_status, "codex", "& codex.ps1 mcp list"), + isolated("claude-code", cli_status, "claude-code", "& claude mcp list"), + isolated("claude-desktop", config_status, claude_desktop_config, "mcpServers", "claude-desktop"), + isolated("gemini-cli", cli_status, "gemini-cli", "& gemini.ps1 mcp list"), + isolated("qwen-cli", cli_status, "qwen-cli", "& qwen.ps1 mcp list"), + isolated("cursor", config_status, CURSOR_MCP, "mcpServers", "cursor"), + isolated("copilot-vscode", config_status, vscode_mcp, "servers", "copilot-vscode"), + isolated("roo-code", config_status, roo_mcp, "mcpServers", "roo-code"), + isolated("kilocode", config_status, kilo_mcp, "mcpServers", "kilocode"), + isolated("cline", config_status, cline_vscode_mcp, "mcpServers", "cline"), ] if not cline_vscode_mcp.exists() and cline_cursor_mcp.exists(): - results[-1] = config_status(cline_cursor_mcp, "mcpServers", "cline") + results[-1] = isolated("cline", config_status, cline_cursor_mcp, "mcpServers", "cline") return { "server_name": SERVER_NAME, "results": results, @@ -645,19 +679,19 @@ def console_status_all() -> dict[str, Any]: cline_vscode_mcp = cline_vscode_mcp_path() cline_cursor_mcp = cline_cursor_mcp_path() results = [ - text_config_status(CODEX_CONFIG, "codex"), - text_config_status(CLAUDE_CODE_CONFIG, "claude-code"), - config_status(claude_desktop_config, "mcpServers", "claude-desktop"), - text_config_status(GEMINI_SETTINGS, "gemini-cli"), - text_config_status(QWEN_SETTINGS, "qwen-cli"), - config_status(CURSOR_MCP, "mcpServers", "cursor"), - config_status(vscode_mcp, "servers", "copilot-vscode"), - config_status(roo_mcp, "mcpServers", "roo-code"), - config_status(kilo_mcp, "mcpServers", "kilocode"), - config_status(cline_vscode_mcp, "mcpServers", "cline"), + isolated("codex", text_config_status, CODEX_CONFIG, "codex"), + isolated("claude-code", text_config_status, CLAUDE_CODE_CONFIG, "claude-code"), + isolated("claude-desktop", config_status, claude_desktop_config, "mcpServers", "claude-desktop"), + isolated("gemini-cli", text_config_status, GEMINI_SETTINGS, "gemini-cli"), + isolated("qwen-cli", text_config_status, QWEN_SETTINGS, "qwen-cli"), + isolated("cursor", config_status, CURSOR_MCP, "mcpServers", "cursor"), + isolated("copilot-vscode", config_status, vscode_mcp, "servers", "copilot-vscode"), + isolated("roo-code", config_status, roo_mcp, "mcpServers", "roo-code"), + isolated("kilocode", config_status, kilo_mcp, "mcpServers", "kilocode"), + isolated("cline", config_status, cline_vscode_mcp, "mcpServers", "cline"), ] if not cline_vscode_mcp.exists() and cline_cursor_mcp.exists(): - results[-1] = config_status(cline_cursor_mcp, "mcpServers", "cline") + results[-1] = isolated("cline", config_status, cline_cursor_mcp, "mcpServers", "cline") return { "server_name": SERVER_NAME, "results": results, @@ -672,19 +706,19 @@ def doctor_all() -> dict[str, Any]: cline_vscode_mcp = cline_vscode_mcp_path() cline_cursor_mcp = cline_cursor_mcp_path() results = [ - cli_doctor("codex", "codex.ps1", "& codex.ps1 mcp list"), - cli_doctor("claude-code", "claude", "& claude mcp list"), - config_doctor(claude_desktop_config, "mcpServers", "claude-desktop"), - cli_doctor("gemini-cli", "gemini.ps1", "& gemini.ps1 mcp list"), - cli_doctor("qwen-cli", "qwen.ps1", "& qwen.ps1 mcp list"), - config_doctor(CURSOR_MCP, "mcpServers", "cursor"), - config_doctor(vscode_mcp, "servers", "copilot-vscode"), - config_doctor(roo_mcp, "mcpServers", "roo-code"), - config_doctor(kilo_mcp, "mcpServers", "kilocode"), - config_doctor(cline_vscode_mcp, "mcpServers", "cline"), + isolated("codex", cli_doctor, "codex", "codex.ps1", "& codex.ps1 mcp list"), + isolated("claude-code", cli_doctor, "claude-code", "claude", "& claude mcp list"), + isolated("claude-desktop", config_doctor, claude_desktop_config, "mcpServers", "claude-desktop"), + isolated("gemini-cli", cli_doctor, "gemini-cli", "gemini.ps1", "& gemini.ps1 mcp list"), + isolated("qwen-cli", cli_doctor, "qwen-cli", "qwen.ps1", "& qwen.ps1 mcp list"), + isolated("cursor", config_doctor, CURSOR_MCP, "mcpServers", "cursor"), + isolated("copilot-vscode", config_doctor, vscode_mcp, "servers", "copilot-vscode"), + isolated("roo-code", config_doctor, roo_mcp, "mcpServers", "roo-code"), + isolated("kilocode", config_doctor, kilo_mcp, "mcpServers", "kilocode"), + isolated("cline", config_doctor, cline_vscode_mcp, "mcpServers", "cline"), ] if not cline_vscode_mcp.exists() and cline_cursor_mcp.exists(): - results[-1] = config_doctor(cline_cursor_mcp, "mcpServers", "cline") + results[-1] = isolated("cline", config_doctor, cline_cursor_mcp, "mcpServers", "cline") return { "server_name": SERVER_NAME, "local_server": local_server_doctor(), diff --git a/agentmemory/oauth.py b/agentmemory/oauth.py index bacd5ee..37fd4a8 100644 --- a/agentmemory/oauth.py +++ b/agentmemory/oauth.py @@ -97,8 +97,8 @@ def consume_auth_code( def _verify_pkce(challenge: str, method: str, verifier: str) -> bool: if not verifier: return False - if method == "PLAIN": - return hmac.compare_digest(challenge, verifier) + # Only S256 is supported. PLAIN is rejected because the challenge equals + # the verifier, making intercepted auth codes fully replayable. if method == "S256": digest = hashlib.sha256(verifier.encode("ascii")).digest() expected = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") diff --git a/agentmemory/providers/mem0.py b/agentmemory/providers/mem0.py index e099af5..487e54d 100644 --- a/agentmemory/providers/mem0.py +++ b/agentmemory/providers/mem0.py @@ -2,7 +2,6 @@ import os import sys -from functools import lru_cache import importlib.metadata import pickle from pathlib import Path @@ -243,6 +242,12 @@ def apply_cli_configuration(cls, *, provider_config: dict[str, Any], args) -> bo def __init__(self, *, runtime_config: dict[str, Any], provider_config: dict[str, Any]) -> None: super().__init__(runtime_config=runtime_config, provider_config=provider_config) self._memory_lock = Lock() + # Per-instance caches. Lifetime is tied to this provider instance, so a + # fresh Mem0Provider (built by config.get_provider after clear_caches on + # reconfigure / key rotation) always rebuilds Memory and re-reads the + # API key. Never promote these to module/global caches. + self._memory: Memory | None = None + self._api_key: str | None = None def capabilities(self) -> ProviderCapabilities: return { @@ -286,20 +291,23 @@ def provider_contract(self) -> ProviderContract: } def clear_caches(self) -> None: - self._get_openrouter_api_key.cache_clear() - self._load_memory.cache_clear() + self._memory = None + self._api_key = None - @lru_cache(maxsize=1) def _get_openrouter_api_key(self) -> str: + if self._api_key is not None: + return self._api_key api_key = os.environ.get("OPENROUTER_API_KEY") if not is_configured_api_key(api_key): raise ProviderConfigurationError( "OPENROUTER_API_KEY is not set. Put it in the shell or in AgentMemory/.env" ) + self._api_key = api_key return api_key - @lru_cache(maxsize=1) def _load_memory(self) -> Memory: + if self._memory is not None: + return self._memory api_key = self._get_openrouter_api_key() _install_openai_usage_capture() @@ -315,7 +323,8 @@ def _load_memory(self) -> Memory: embedder_config["api_key"] = api_key embedder["config"] = embedder_config config["embedder"] = embedder - return Memory.from_config(config) + self._memory = Memory.from_config(config) + return self._memory def _normalize_record(self, payload: dict[str, Any], *, include_score: bool = False) -> MemoryRecord: if not isinstance(payload, dict): @@ -420,44 +429,101 @@ def _record_needs_hydration(self, record: MemoryRecord) -> bool: ) ) + @staticmethod + def _added_text(messages: Any) -> str: + """Best-effort flatten of the add() input into the text mem0 would + store, so the fallback can attribute a listed record to this write.""" + if isinstance(messages, str): + return messages.strip() + if isinstance(messages, dict): + return str(messages.get("content", "")).strip() + if isinstance(messages, list): + parts = [ + str(item.get("content", "")) if isinstance(item, dict) else str(item) + for item in messages + ] + return "\n".join(part for part in parts if part).strip() + return "" + def _fallback_added_record( self, *, memory: Memory, + added_id: str | None = None, + added_text: str = "", user_id=None, agent_id=None, run_id=None, metadata=None, memory_type=None, ) -> MemoryRecord: + # mem0 didn't hand back a proper record. If it surfaced an id for the + # just-saved write, fetch that exact record rather than guessing. + if added_id: + try: + return self._normalize_one_record(memory.get(added_id)) + except Exception: + pass + payload = memory.get_all(user_id=user_id, agent_id=agent_id, run_id=run_id, limit=25, filters=None) records = self._normalize_records(payload) + matched = records if metadata: - filtered = [ + matched = [ record - for record in records + for record in matched if all(record.get("metadata", {}).get(key) == value for key, value in metadata.items()) ] - if filtered: - records = filtered if memory_type is not None: - typed = [ + matched = [ record - for record in records + for record in matched if record.get("memory_type") == memory_type or record.get("metadata", {}).get("memory_type") == memory_type ] - if typed: - records = typed - if not records: - raise ProviderValidationError("Mem0 add succeeded without returning a usable record, and fallback lookup found nothing.") - return records[0] + # Text is the strongest signal that a listed row is the write we just + # made; require it (when known) so we never return a pre-existing row. + if added_text: + text_matched = [ + record for record in matched if str(record.get("memory", "")).strip() == added_text + ] + if text_matched: + matched = text_matched + else: + matched = [] + + # Only trust a listed record when something discriminating narrowed it + # down (text/metadata/memory_type) AND it resolved to a single row; + # otherwise the "first row" is an arbitrary pre-existing record. + had_discriminator = bool(added_text or metadata or memory_type is not None) + if matched and (len(matched) == 1 or had_discriminator): + return matched[0] + + # No confident match. If mem0 surfaced a real id for the write (but the + # earlier get() failed transiently), build a minimal record from the + # known write rather than syncing an arbitrary pre-existing row's id. + if added_id: + return self._normalize_record( + { + "id": added_id, + "memory": added_text, + "metadata": dict(metadata or {}), + "user_id": user_id, + "agent_id": agent_id, + "run_id": run_id, + "memory_type": memory_type, + } + ) + # Without a confident match or a real id we must not invent / borrow an + # id; surfacing the failure is safer than syncing the wrong record. + raise ProviderValidationError("Mem0 add succeeded without returning a usable record, and fallback lookup found nothing.") def _normalize_add_result( self, payload: Any, *, memory: Memory, + added_text: str = "", user_id=None, agent_id=None, run_id=None, @@ -483,8 +549,20 @@ def _normalize_add_result( except Exception: return record return record + # mem0 sometimes wraps the saved id in a non-record envelope (e.g. + # an empty results list alongside raw_saved_record_id). Surface that + # id to the fallback so it can fetch / attribute the exact write. + added_id = None + if isinstance(payload, dict): + for key in ("raw_saved_record_id", "memory_id", "id"): + candidate = payload.get(key) + if isinstance(candidate, str) and candidate: + added_id = candidate + break return self._fallback_added_record( memory=memory, + added_id=added_id, + added_text=added_text, user_id=user_id, agent_id=agent_id, run_id=run_id, @@ -604,6 +682,9 @@ def _iter_scope_payloads(self) -> list[dict[str, Any]]: if not isinstance(blob, (bytes, bytearray)): continue try: + # Trust assumption: the qdrant store is a local, first-party + # file owned by this runtime. pickle.loads here is an RCE vector + # only if an attacker can write to that local store. point = pickle.loads(blob) except Exception: continue @@ -716,6 +797,7 @@ def add_memory(self, *, messages, user_id=None, agent_id=None, run_id=None, meta record = self._normalize_add_result( payload, memory=memory, + added_text=self._added_text(messages), user_id=user_id, agent_id=agent_id, run_id=run_id, diff --git a/agentmemory/runtime/atomic_io.py b/agentmemory/runtime/atomic_io.py index 521e787..30cbe03 100644 --- a/agentmemory/runtime/atomic_io.py +++ b/agentmemory/runtime/atomic_io.py @@ -7,6 +7,28 @@ from typing import Any +def _fsync_dir(directory: Path) -> None: + """Best-effort fsync of a directory to make a rename durable. + + On POSIX, the directory entry created/updated by os.replace is only + guaranteed durable after the containing directory is itself fsynced. + Opening a directory for fsync is not supported on Windows and may fail + on some filesystems, so any OSError is swallowed cleanly. + """ + if os.name != "posix": + return + try: + fd = os.open(directory, os.O_RDONLY) + except OSError: + return + try: + os.fsync(fd) + except OSError: + pass + finally: + os.close(fd) + + def atomic_write_text(path: Path, content: str, *, encoding: str = "utf-8") -> None: """Write text to `path` atomically. @@ -35,6 +57,7 @@ def atomic_write_text(path: Path, content: str, *, encoding: str = "utf-8") -> N os.fsync(temp_file.fileno()) os.replace(temp_path, path) temp_path = None + _fsync_dir(path.parent) finally: if temp_path is not None: Path(temp_path).unlink(missing_ok=True) diff --git a/agentmemory/runtime/http_client.py b/agentmemory/runtime/http_client.py index 05c5066..1ecc400 100644 --- a/agentmemory/runtime/http_client.py +++ b/agentmemory/runtime/http_client.py @@ -12,7 +12,7 @@ from urllib.request import Request, urlopen from agentmemory.platform import launcher_command, launcher_path -from agentmemory.runtime.config import BASE_DIR, active_provider_runtime_policy, clear_caches, current_api_host, current_api_port +from agentmemory.runtime.config import BASE_DIR, active_provider_runtime_policy, clear_caches, current_api_host, current_api_port, read_api_pid from agentmemory.runtime.transport import error_class_for_type from agentmemory.providers.base import ( ProviderError, @@ -31,6 +31,22 @@ import fcntl +def _is_owner_process() -> bool: + # The owner env flag is inherited by any child the API server spawns, so the + # flag alone is not sufficient: an inheriting child would touch the embedded + # backend directly and violate the single-owner lock. Confirm this process is + # actually the recorded API owner by matching its pid against the pid file. + if os.environ.get(OWNER_ENV) != "1": + return False + recorded_pid = read_api_pid() + if recorded_pid is None: + # Startup window: the real server sets the flag and writes its pid early, + # but before the pid file exists we trust the flag so we don't break the + # legitimate owner during its own boot. + return True + return recorded_pid == os.getpid() + + def should_proxy_to_api() -> bool: transport_mode = active_provider_runtime_policy()["transport_mode"] if transport_mode == "remote_only": @@ -38,7 +54,7 @@ def should_proxy_to_api() -> bool: "Provider transport mode 'remote_only' requires a supported remote transport implementation; " "local direct execution is not available." ) - return transport_mode == "owner_process_proxy" and os.environ.get(OWNER_ENV) != "1" + return transport_mode == "owner_process_proxy" and not _is_owner_process() def api_base_url() -> str: @@ -76,6 +92,11 @@ def _request(method: str, path: str, payload: dict[str, Any] | None = None) -> A message = str(exc) error_type = "" error_type_cls = error_class_for_type(error_type, status_code=exc.code) + if exc.code in (401, 403) or error_type == "AuthRequired": + message = ( + f"AgentMemory API at {api_base_url()} requires authentication ({exc.code}). " + "Set AGENTMEMORY_API_TOKEN to the owner process token and retry." + ) raise error_type_cls(message) from exc except URLError as exc: raise ProviderUnavailableError(f"AgentMemory API is not reachable at {api_base_url()}. Start it with `agentmemory start-api`.") from exc @@ -117,6 +138,10 @@ def ensure_api_running() -> None: if api_is_healthy(): return + # Hold the cross-process lock ONLY around the double-checked health re-check + # and the spawn trigger. The up-to-20s health-wait poll runs OUTSIDE the lock + # so parallel clients don't serialize behind it — they each poll independently + # until the shared deadline. with _api_start_lock(): if api_is_healthy(): return @@ -134,11 +159,11 @@ def ensure_api_running() -> None: ) clear_caches() - deadline = time.time() + API_START_TIMEOUT_SECONDS - while time.time() < deadline: - if api_is_healthy(): - return - time.sleep(0.5) + deadline = time.time() + API_START_TIMEOUT_SECONDS + while time.time() < deadline: + if api_is_healthy(): + return + time.sleep(0.5) raise ProviderUnavailableError(f"AgentMemory API did not become ready at {api_base_url()} within {API_START_TIMEOUT_SECONDS:.0f}s") diff --git a/agentmemory/runtime/metrics.py b/agentmemory/runtime/metrics.py index 5309fe5..9d5e320 100644 --- a/agentmemory/runtime/metrics.py +++ b/agentmemory/runtime/metrics.py @@ -31,6 +31,17 @@ ) +# Cardinality caps. `error_type` (from exc.__class__.__name__) and `model` +# (from raw provider responses) are attacker/provider-controlled and would +# otherwise grow the registry — and the resulting Prometheus series count — +# without bound. Once a keyspace hits its cap, further *new* distinct values +# are folded into a catch-all bucket so totals stay accurate while series +# count stays bounded. Existing keys always keep counting. +_OTHER_LABEL = "__other__" +_MAX_ERROR_TYPES_PER_OP = 50 +_MAX_MODELS = 100 + + # OpenRouter pricing (USD per 1M tokens) for the shipped defaults. Numbers # are approximate — update when pricing changes. A model absent from this # table contributes 0 to cost totals but still counts tokens. @@ -111,7 +122,15 @@ def record_operation(self, *, name: str, status: str, duration_seconds: float, e if status == "ok": self._op_ok[name] = self._op_ok.get(name, 0) + 1 else: - key = (name, error_type or "Unknown") + error_type = error_type or "Unknown" + key = (name, error_type) + if key not in self._op_err: + # Cap distinct error_types per operation; fold overflow + # into a catch-all so a flood of novel exception classes + # cannot blow up cardinality. + distinct = sum(1 for n, _t in self._op_err if n == name) + if distinct >= _MAX_ERROR_TYPES_PER_OP: + key = (name, _OTHER_LABEL) self._op_err[key] = self._op_err.get(key, 0) + 1 histogram = self._op_latency.setdefault(name, _Histogram()) histogram.observe(max(duration_seconds, 0.0)) @@ -125,6 +144,16 @@ def record_llm_usage(self, *, model: str, prompt_tokens: int, completion_tokens: return normalized = _normalize_model_key(model) with self._lock: + if ( + normalized not in self._model_tokens_prompt + and normalized not in self._model_tokens_completion + ): + # Cap distinct models; fold overflow into a catch-all so a + # provider returning unbounded model ids cannot blow up + # cardinality. Tokens still accrue, just under __other__. + distinct = len(set(self._model_tokens_prompt) | set(self._model_tokens_completion)) + if distinct >= _MAX_MODELS: + normalized = _OTHER_LABEL self._model_tokens_prompt[normalized] = self._model_tokens_prompt.get(normalized, 0) + max(int(prompt_tokens), 0) self._model_tokens_completion[normalized] = self._model_tokens_completion.get(normalized, 0) + max(int(completion_tokens), 0) @@ -245,7 +274,7 @@ def prometheus_text(self) -> str: cumulative = 0 for idx, boundary in enumerate(_LATENCY_BUCKETS_SECONDS): cumulative += histogram.buckets[idx] - le_label = "+Inf" if boundary == float("inf") else repr(boundary) + le_label = "+Inf" if boundary == float("inf") else _format_float(boundary) lines.append( f'agentmemory_operation_latency_seconds_bucket{{operation="{_esc(name)}",le="{le_label}"}} {cumulative}' ) @@ -300,6 +329,15 @@ def _esc(value: str) -> str: return value.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n") +def _format_float(value: float) -> str: + """Render a bucket boundary as a plain numeric string (no trailing ".0" + for integral values), suitable for a Prometheus `le` label. + """ + if value == int(value): + return str(int(value)) + return repr(value) + + _REGISTRY = _MetricsRegistry() diff --git a/agentmemory/runtime/reconcile.py b/agentmemory/runtime/reconcile.py index bb7f507..2959654 100644 --- a/agentmemory/runtime/reconcile.py +++ b/agentmemory/runtime/reconcile.py @@ -117,45 +117,52 @@ def _conflict_reason(left_claim: dict[str, Any], right_claim: dict[str, Any]) -> return None +def _scope_key(record: MemoryRecord) -> tuple[Any, Any, Any]: + return (record.get("user_id"), record.get("agent_id"), record.get("run_id")) + + def find_conflicts(records: list[MemoryRecord]) -> list[dict[str, Any]]: - claimed: list[tuple[MemoryRecord, dict[str, Any]]] = [] - for record in records: + # Partition records by scope identity so claims from different + # users/agents/runs are never compared against each other. + by_scope: dict[tuple[Any, Any, Any], list[tuple[int, MemoryRecord, dict[str, Any]]]] = {} + for index, record in enumerate(records): claim = _claim_for_record(record) if claim is not None: - claimed.append((record, claim)) + by_scope.setdefault(_scope_key(record), []).append((index, record, claim)) conflicts: list[dict[str, Any]] = [] seen_pairs: set[tuple[str, str]] = set() - for left_index, (left_record, left_claim) in enumerate(claimed): - for right_record, right_claim in claimed[left_index + 1 :]: - reason = _conflict_reason(left_claim, right_claim) - if reason is None: - continue - left_id = str(left_record.get("id") or left_index) - right_id = str(right_record.get("id") or len(seen_pairs)) - pair_key = tuple(sorted((left_id, right_id))) - if pair_key in seen_pairs: - continue - seen_pairs.add(pair_key) - conflicts.append( - { - "reason": reason[0], - "confidence": reason[1], - "subject": left_claim["subject"], - "predicate": left_claim["predicate"], - "left_claim": { - "value": left_claim["value"], - "polarity": left_claim["polarity"], - "source": left_claim["source"], - }, - "right_claim": { - "value": right_claim["value"], - "polarity": right_claim["polarity"], - "source": right_claim["source"], - }, - "left": _record_summary(left_record), - "right": _record_summary(right_record), - } - ) + for claimed in by_scope.values(): + for offset, (left_index, left_record, left_claim) in enumerate(claimed): + for right_index, right_record, right_claim in claimed[offset + 1 :]: + reason = _conflict_reason(left_claim, right_claim) + if reason is None: + continue + left_id = str(left_record.get("id") or f"idx:{left_index}") + right_id = str(right_record.get("id") or f"idx:{right_index}") + pair_key = tuple(sorted((left_id, right_id))) + if pair_key in seen_pairs: + continue + seen_pairs.add(pair_key) + conflicts.append( + { + "reason": reason[0], + "confidence": reason[1], + "subject": left_claim["subject"], + "predicate": left_claim["predicate"], + "left_claim": { + "value": left_claim["value"], + "polarity": left_claim["polarity"], + "source": left_claim["source"], + }, + "right_claim": { + "value": right_claim["value"], + "polarity": right_claim["polarity"], + "source": right_claim["source"], + }, + "left": _record_summary(left_record), + "right": _record_summary(right_record), + } + ) return sorted(conflicts, key=lambda item: (-float(item["confidence"]), str(item["subject"]), str(item["predicate"]))) diff --git a/agentmemory/runtime/scope_registry.py b/agentmemory/runtime/scope_registry.py index 9f5302e..e52f615 100644 --- a/agentmemory/runtime/scope_registry.py +++ b/agentmemory/runtime/scope_registry.py @@ -4,6 +4,7 @@ import base64 from datetime import datetime, timezone import json +import logging import os import sqlite3 from pathlib import Path @@ -18,6 +19,9 @@ import fcntl +_logger = logging.getLogger(__name__) + + _SCHEMA = """ CREATE TABLE IF NOT EXISTS scope_registry ( provider TEXT NOT NULL, @@ -102,7 +106,13 @@ def _connect(runtime_dir: str) -> Iterator[sqlite3.Connection]: path = registry_path(runtime_dir) path.parent.mkdir(parents=True, exist_ok=True) connection = sqlite3.connect(path, timeout=10.0) - connection.execute("PRAGMA busy_timeout = 10000") + # Enable WAL so readers and writers don't block each other across the + # multiple processes (CLI + API server + providers) that sync on every op. + # journal_mode returns a row, so consume it; WAL is persisted per-DB once set. + connection.execute("PRAGMA journal_mode=WAL").fetchone() + connection.execute("PRAGMA busy_timeout=5000") + # NORMAL is durable and safe under WAL while avoiding fsync-per-commit stalls. + connection.execute("PRAGMA synchronous=NORMAL") connection.execute(_SCHEMA) for statement in _INDEXES: connection.execute(statement) @@ -292,7 +302,25 @@ def mark_sync_failed( } ) providers[provider_name] = current - _write_status_document(runtime_dir, payload) + try: + _write_status_document(runtime_dir, payload) + except Exception: + # The status write can fail under the same condition that broke the + # registry write (e.g. disk full). If it is swallowed by callers, + # needs_rebuild never gets persisted and the registry drifts with no + # signal, so emit a loud, observable record before propagating. + _logger.error( + "scope_registry: failed to persist needs_rebuild status for " + "provider=%s operation=%s memory_id=%s; registry may drift " + "undetected (original error: %s: %s)", + provider_name, + operation, + memory_id, + error.__class__.__name__, + error, + exc_info=True, + ) + raise def clear_sync_failure(provider_name: str, runtime_dir: str, *, rebuilt: bool = False) -> None: diff --git a/agentmemory/runtime/transport.py b/agentmemory/runtime/transport.py index 15d3a1d..d9e1ff8 100644 --- a/agentmemory/runtime/transport.py +++ b/agentmemory/runtime/transport.py @@ -180,6 +180,11 @@ def provider_error_payload(exc: ProviderError) -> dict[str, Any]: def error_class_for_type(error_type: str, *, status_code: int) -> type[ProviderError]: if error_type in ERROR_TYPE_MAP: return ERROR_TYPE_MAP[error_type] + # Authentication/authorization failures (401/403) are not validation errors; + # an explicit "AuthRequired" error_type signals the same. Surface them as an + # availability error so the message can tell the caller to set a token. + if status_code in (401, 403) or error_type == "AuthRequired": + return ProviderUnavailableError return ProviderUnavailableError if status_code >= 500 else ProviderValidationError diff --git a/tests/test_agentmemory_clients.py b/tests/test_agentmemory_clients.py index a053252..5322a24 100644 --- a/tests/test_agentmemory_clients.py +++ b/tests/test_agentmemory_clients.py @@ -124,6 +124,52 @@ def test_linux_claude_desktop_path_prefers_existing_lowercase_fallback(self) -> ): self.assertEqual(agentmemory_clients.claude_desktop_config_path(), fallback) + def test_load_json_raises_config_parse_error_on_malformed_json(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "mcp.json" + path.write_text('{"mcpServers": {},}', encoding="utf-8") # trailing comma + with self.assertRaises(agentmemory_clients.ConfigParseError): + agentmemory_clients.load_json(path, {"mcpServers": {}}) + + def test_isolated_corrupt_json_yields_error_entry_without_clobbering(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "mcp.json" + original = '{"mcpServers": {"other": {}},}' # trailing comma -> invalid + path.write_text(original, encoding="utf-8") + backup_dir = Path(tmp) / "backups" + result = agentmemory_clients.isolated( + "cursor", + agentmemory_clients.merge_server_json, + path, + "mcpServers", + agentmemory_clients.SERVER_NAME, + agentmemory_clients.stdio_server_config(), + backup_dir, + ) + self.assertEqual(result["target"], "cursor") + self.assertEqual(result["status"], "error") + self.assertIn("not valid JSON", result["reason"]) + # The unparseable file must not be overwritten or backed up. + self.assertEqual(path.read_text(encoding="utf-8"), original) + self.assertFalse(backup_dir.exists()) + + def test_status_all_isolates_single_corrupt_config(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + corrupt = Path(tmp) / "cursor" / "mcp.json" + corrupt.parent.mkdir(parents=True, exist_ok=True) + corrupt.write_text('{"mcpServers": {},}', encoding="utf-8") + completed = agentmemory_clients.subprocess.CompletedProcess( + args=["pwsh"], returncode=1, stdout="", stderr="" + ) + with ( + mock.patch.object(agentmemory_clients, "CURSOR_MCP", corrupt), + mock.patch.object(agentmemory_clients, "run_pwsh", return_value=completed), + ): + payload = agentmemory_clients.status_all() # must not raise + cursor = next(r for r in payload["results"] if r["target"] == "cursor") + self.assertEqual(cursor["health"], "error") + self.assertIn("not valid JSON", cursor["details"]) + def test_client_path_overrides_take_precedence(self) -> None: with tempfile.TemporaryDirectory() as tmp: custom = Path(tmp) / "custom" / "claude.json" diff --git a/tests/test_agentmemory_http_client.py b/tests/test_agentmemory_http_client.py index e533192..7559f38 100644 --- a/tests/test_agentmemory_http_client.py +++ b/tests/test_agentmemory_http_client.py @@ -388,6 +388,58 @@ def fake_urlopen(*args, **kwargs): finally: agentmemory_http_client.urlopen = original_urlopen # type: ignore[assignment] + def test_http_error_401_maps_to_actionable_authentication_error(self) -> None: + original_urlopen = agentmemory_http_client.urlopen + try: + class FakeHttpError(agentmemory_http_client.HTTPError): + def __init__(self): + super().__init__( + url="http://127.0.0.1:8765/add", + code=401, + msg="Unauthorized", + hdrs=None, + fp=None, + ) + + def read(self): + return b'{"error":"unauthorized"}' + + def fake_urlopen(*args, **kwargs): + raise FakeHttpError() + + agentmemory_http_client.urlopen = fake_urlopen # type: ignore[assignment] + + with self.assertRaises(ProviderUnavailableError) as ctx: + agentmemory_http_client._request("POST", "/add", {"messages": "hi"}) + self.assertIn("AGENTMEMORY_API_TOKEN", str(ctx.exception)) + self.assertIn("401", str(ctx.exception)) + finally: + agentmemory_http_client.urlopen = original_urlopen # type: ignore[assignment] + + def test_owner_with_mismatched_pid_proxies_as_inheriting_child(self) -> None: + original_runtime_policy = agentmemory_http_client.active_provider_runtime_policy + original_read_api_pid = agentmemory_http_client.read_api_pid + original_owner = os.environ.get(agentmemory_http_client.OWNER_ENV) + try: + agentmemory_http_client.active_provider_runtime_policy = lambda: {"transport_mode": "owner_process_proxy"} # type: ignore[assignment] + os.environ[agentmemory_http_client.OWNER_ENV] = "1" + # Recorded owner pid does not match this process -> we are an inheriting child. + agentmemory_http_client.read_api_pid = lambda: os.getpid() + 1 # type: ignore[assignment] + self.assertTrue(agentmemory_http_client.should_proxy_to_api()) + # Matching pid -> we ARE the owner, do not proxy. + agentmemory_http_client.read_api_pid = lambda: os.getpid() # type: ignore[assignment] + self.assertFalse(agentmemory_http_client.should_proxy_to_api()) + # Startup window (no pid recorded yet) -> trust the flag, do not proxy. + agentmemory_http_client.read_api_pid = lambda: None # type: ignore[assignment] + self.assertFalse(agentmemory_http_client.should_proxy_to_api()) + finally: + agentmemory_http_client.active_provider_runtime_policy = original_runtime_policy # type: ignore[assignment] + agentmemory_http_client.read_api_pid = original_read_api_pid # type: ignore[assignment] + if original_owner is None: + os.environ.pop(agentmemory_http_client.OWNER_ENV, None) + else: + os.environ[agentmemory_http_client.OWNER_ENV] = original_owner + def test_http_error_type_maps_capability_error(self) -> None: original_urlopen = agentmemory_http_client.urlopen try: diff --git a/tests/test_agentmemory_reconcile.py b/tests/test_agentmemory_reconcile.py index a14ccf3..10c55ec 100644 --- a/tests/test_agentmemory_reconcile.py +++ b/tests/test_agentmemory_reconcile.py @@ -30,6 +30,30 @@ def test_find_conflicts_detects_different_values_for_same_claim(self) -> None: self.assertEqual(conflicts[0]["subject"], "user") self.assertEqual(conflicts[0]["predicate"], "prefers") + def test_find_conflicts_does_not_cross_scopes(self) -> None: + records = [ + {"id": "a", "memory": "alice is happy", "metadata": {}, "user_id": "u1"}, + {"id": "b", "memory": "alice is not happy", "metadata": {}, "user_id": "u2"}, + ] + + conflicts = reconcile.find_conflicts(records) + + self.assertEqual(conflicts, []) + + def test_find_conflicts_within_single_scope_still_flagged(self) -> None: + records = [ + {"id": "a", "memory": "alice is happy", "metadata": {}, "user_id": "u1"}, + {"id": "b", "memory": "alice is not happy", "metadata": {}, "user_id": "u1"}, + {"id": "c", "memory": "alice is not happy", "metadata": {}, "user_id": "u2"}, + ] + + conflicts = reconcile.find_conflicts(records) + + self.assertEqual(len(conflicts), 1) + self.assertEqual(conflicts[0]["reason"], "opposite_polarity") + self.assertEqual(conflicts[0]["left"]["user_id"], "u1") + self.assertEqual(conflicts[0]["right"]["user_id"], "u1") + def test_find_conflicts_can_use_structured_claim_metadata(self) -> None: records = [ { diff --git a/tests/test_agentmemory_transport.py b/tests/test_agentmemory_transport.py index 534e8bb..010c370 100644 --- a/tests/test_agentmemory_transport.py +++ b/tests/test_agentmemory_transport.py @@ -203,6 +203,13 @@ def test_error_class_for_type_uses_status_code_fallback(self) -> None: self.assertIs(error_class_for_type("", status_code=503), ProviderUnavailableError) self.assertIs(error_class_for_type("", status_code=400), ProviderValidationError) + def test_error_class_for_type_maps_auth_failures_to_unavailable(self) -> None: + self.assertIs(error_class_for_type("", status_code=401), ProviderUnavailableError) + self.assertIs(error_class_for_type("", status_code=403), ProviderUnavailableError) + self.assertIs(error_class_for_type("AuthRequired", status_code=401), ProviderUnavailableError) + # A 401 must NOT fall through to the generic <500 validation default. + self.assertIsNot(error_class_for_type("", status_code=401), ProviderValidationError) + def test_capability_summary_renders_human_readable_flags(self) -> None: summary = capability_summary( { diff --git a/tests/test_mem0_provider.py b/tests/test_mem0_provider.py index 38b0e0f..855bd5b 100644 --- a/tests/test_mem0_provider.py +++ b/tests/test_mem0_provider.py @@ -137,6 +137,24 @@ def add(self, messages, *, user_id=None, agent_id=None, run_id=None, metadata=No return {"results": [], "raw_saved_record_id": record["id"]} +class FakeMem0NoIdAddResultBackend(FakeMem0Backend): + """Persists the write but returns an add() result with no usable id or + record, forcing the get_all-based fallback to attribute the just-added + write among any pre-existing rows in the same scope.""" + + def add(self, messages, *, user_id=None, agent_id=None, run_id=None, metadata=None, infer=True, memory_type=None): + super().add( + messages, + user_id=user_id, + agent_id=agent_id, + run_id=run_id, + metadata=metadata, + infer=infer, + memory_type=memory_type, + ) + return {"results": []} + + class FakeMem0MessageResultBackend(FakeMem0Backend): def update(self, memory_id, data, *, metadata=None): super().update(memory_id, data, metadata=metadata) @@ -213,6 +231,12 @@ def __init__(self, *, runtime_config: dict[str, object], provider_config: dict[s self._fake_memory = FakeMem0EmptyAddResultBackend() +class NoIdAddResultMem0Provider(HarnessMem0Provider): + def __init__(self, *, runtime_config: dict[str, object], provider_config: dict[str, object]) -> None: + Mem0Provider.__init__(self, runtime_config=runtime_config, provider_config=provider_config) + self._fake_memory = FakeMem0NoIdAddResultBackend() + + class MessageResultMem0Provider(HarnessMem0Provider): def __init__(self, *, runtime_config: dict[str, object], provider_config: dict[str, object]) -> None: Mem0Provider.__init__(self, runtime_config=runtime_config, provider_config=provider_config) @@ -323,6 +347,26 @@ def test_add_falls_back_when_mem0_returns_empty_results_wrapper(self) -> None: self.assertEqual(created["metadata"]["topic"], "ops") self.assertEqual(created["provider"], "mem0") + def test_add_fallback_does_not_return_preexisting_record_without_id(self) -> None: + provider = NoIdAddResultMem0Provider( + runtime_config={"runtime_dir": self.temp_dir.name}, + provider_config=Mem0Provider.default_provider_config(runtime_dir=self.temp_dir.name), + ) + existing = provider.add_memory( + messages=[{"role": "user", "content": "pre-existing note"}], + user_id="default", + ) + + created = provider.add_memory( + messages=[{"role": "user", "content": "brand new note"}], + user_id="default", + ) + + # The fallback must attribute the just-added write by its text, never + # borrow the pre-existing row that get_all happens to list first. + self.assertEqual(created["memory"], "brand new note") + self.assertNotEqual(created["id"], existing["id"]) + def test_update_uses_get_fallback_for_message_only_success_response(self) -> None: provider = MessageResultMem0Provider( runtime_config={"runtime_dir": self.temp_dir.name}, diff --git a/tests/test_scope_registry.py b/tests/test_scope_registry.py index 56f2027..895d328 100644 --- a/tests/test_scope_registry.py +++ b/tests/test_scope_registry.py @@ -1,4 +1,5 @@ import multiprocessing +import sqlite3 import tempfile import unittest from datetime import datetime, timezone @@ -145,6 +146,20 @@ def test_list_expired_memory_ids_uses_registry_expiry_index(self) -> None: self.assertEqual(expired, ["expired"]) + def test_wal_mode_enabled_after_write(self) -> None: + scope_registry.upsert_record( + "mem0", + {"id": "1", "memory": "a", "provider": "mem0", "user_id": "user-a"}, + self.runtime_dir, + ) + db_path = scope_registry.registry_path(self.runtime_dir) + connection = sqlite3.connect(db_path) + try: + mode = connection.execute("PRAGMA journal_mode").fetchone()[0] + finally: + connection.close() + self.assertEqual(mode.lower(), "wal") + def test_concurrent_upserts_preserve_all_records(self) -> None: processes = [ multiprocessing.Process(target=_worker_upsert, args=(self.runtime_dir, "localjson", index)) From daaee3cbabf1f852ecf214cfa12b966192178ac4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 25 May 2026 22:02:25 +0000 Subject: [PATCH 2/4] Unify HTTP input contract and schema validation across surfaces - operation_adapters: HTTP `add` now accepts the documented `text` field and synthesizes `messages` like the MCP/CLI adapters (fixes KeyError on the published contract) - HTTP query/body sources now validate against each operation's input schema (limit/enum/type/unexpected-field checks), matching MCP behavior - extract the shared validator into runtime/schema_validation.py; mcp.py reuses it with no behavior change; lazy OPERATIONS import avoids a cycle https://claude.ai/code/session_01BKjtun7hVwxh6Hv3YLdrww --- agentmemory/mcp.py | 52 +------------- agentmemory/runtime/operation_adapters.py | 62 ++++++++++++++--- agentmemory/runtime/schema_validation.py | 72 ++++++++++++++++++++ tests/test_agentmemory_operation_adapters.py | 32 +++++++++ 4 files changed, 158 insertions(+), 60 deletions(-) create mode 100644 agentmemory/runtime/schema_validation.py diff --git a/agentmemory/mcp.py b/agentmemory/mcp.py index 45c2f3a..edc3a12 100644 --- a/agentmemory/mcp.py +++ b/agentmemory/mcp.py @@ -5,6 +5,7 @@ from agentmemory.runtime.operation_adapters import mcp_operation_source from agentmemory.runtime.operations import OPERATIONS_BY_MCP_NAME, mcp_tools +from agentmemory.runtime.schema_validation import validate_arguments from agentmemory.runtime.transport import mcp_result, provider_error_payload from agentmemory.providers.base import ProviderError, ProviderValidationError @@ -42,57 +43,6 @@ def handle_initialize(request_id: Any, params: dict[str, Any]) -> dict[str, Any] return success(request_id, result) -def _schema_type_matches(value: Any, expected_type: str) -> bool: - if expected_type == "object": - return isinstance(value, dict) - if expected_type == "array": - return isinstance(value, list) - if expected_type == "string": - return isinstance(value, str) - if expected_type == "integer": - return isinstance(value, int) and not isinstance(value, bool) - if expected_type == "number": - return isinstance(value, (int, float)) and not isinstance(value, bool) - if expected_type == "boolean": - return isinstance(value, bool) - if expected_type == "null": - return value is None - return True - - -def validate_arguments(schema: dict[str, Any], arguments: Any) -> dict[str, Any]: - if not isinstance(arguments, dict): - raise ProviderValidationError("MCP tool arguments must be an object.") - - if schema.get("type") == "object": - properties = schema.get("properties") or {} - required = schema.get("required") or [] - for field in required: - if field not in arguments: - raise ProviderValidationError(f"Missing required argument: {field}") - - if schema.get("additionalProperties") is False: - extra = sorted(key for key in arguments if key not in properties) - if extra: - raise ProviderValidationError(f"Unexpected argument: {extra[0]}") - - for field, value in arguments.items(): - field_schema = properties.get(field) - if not isinstance(field_schema, dict): - continue - expected_type = field_schema.get("type") - if isinstance(expected_type, str) and not _schema_type_matches(value, expected_type): - raise ProviderValidationError(f"Argument '{field}' must be {expected_type}.") - allowed_values = field_schema.get("enum") - if isinstance(allowed_values, list) and value not in allowed_values: - raise ProviderValidationError(f"Argument '{field}' must be one of: {', '.join(map(str, allowed_values))}.") - minimum = field_schema.get("minimum") - if isinstance(minimum, (int, float)) and isinstance(value, (int, float)) and not isinstance(value, bool) and value < minimum: - raise ProviderValidationError(f"Argument '{field}' must be >= {minimum}.") - - return arguments - - def handle_call(spec: Any, name: str, arguments: Any) -> dict[str, Any]: validated_arguments = validate_arguments(spec.input_schema, arguments) return mcp_result(spec.execute(mcp_operation_source(name, validated_arguments))) diff --git a/agentmemory/runtime/operation_adapters.py b/agentmemory/runtime/operation_adapters.py index aad559d..81d33ae 100644 --- a/agentmemory/runtime/operation_adapters.py +++ b/agentmemory/runtime/operation_adapters.py @@ -4,6 +4,44 @@ from typing import Any, Callable from agentmemory.providers.base import ProviderValidationError +from agentmemory.runtime.schema_validation import validate_arguments + + +def _input_schema_for(operation_name: str) -> dict[str, Any] | None: + # Imported lazily to keep adapters importable even if the operations + # registry (which pulls in providers/config) is mid-initialization; there + # is no static import cycle today, but a lazy lookup keeps it robust. + from agentmemory.runtime.operations import OPERATIONS + + spec = OPERATIONS.get(operation_name) + return spec.input_schema if spec is not None else None + + +def _validate_http_payload(operation_name: str, payload: dict[str, Any]) -> dict[str, Any]: + """Validate a client-supplied request body against the op schema. + + Mirrors MCP: the raw payload is validated as-is so explicit nulls for + typed optional fields are rejected exactly as MCP rejects them. + """ + schema = _input_schema_for(operation_name) + if schema is not None: + validate_arguments(schema, payload) + return payload + + +def _validate_http_source(operation_name: str, source: dict[str, Any]) -> dict[str, Any]: + """Validate an adapter-built source (from query/path params). + + Query/path adapters synthesize optional fields as explicit None to keep a + stable source shape; the schema treats those fields as simply absent (they + are not in `required`). Validate against the present values only so a None + placeholder is not mistaken for a bad type, then return the full source. + """ + schema = _input_schema_for(operation_name) + if schema is not None: + present = {key: value for key, value in source.items() if value is not None} + validate_arguments(schema, present) + return source def operation_name_for_mcp_tool(tool_name: str) -> str: @@ -130,24 +168,30 @@ def http_operation_source( if operation_name == "health": return {} if operation_name == "add": - return dict(payload) + # Validate the RAW payload (which carries `text`, not `messages`) against + # the schema exactly as MCP does, THEN synthesize the `messages` list the + # `_execute_add` handler reads. Order matters: validation must see `text`. + _validate_http_payload(operation_name, payload) + source = dict(payload) + source["messages"] = [{"role": "user", "content": payload["text"]}] + return source if operation_name == "list_scopes": - return { + return _validate_http_source(operation_name, { "limit": int((query_params.get("limit") or ["200"])[0]), "kind": (query_params.get("kind") or [None])[0], "query": (query_params.get("query") or [None])[0], - } + }) if operation_name == "list_scopes_page": - return { + return _validate_http_source(operation_name, { "limit": int((query_params.get("limit") or ["200"])[0]), "cursor": (query_params.get("cursor") or [None])[0], "kind": (query_params.get("kind") or [None])[0], "query": (query_params.get("query") or [None])[0], - } + }) if operation_name in {"search", "search_page"}: - return dict(payload) + return _validate_http_payload(operation_name, dict(payload)) if operation_name == "update": - return dict(payload) + return _validate_http_payload(operation_name, dict(payload)) if operation_name in {"list", "list_page"}: filters = None filters_param = (query_params.get("filters") or [None])[0] @@ -156,14 +200,14 @@ def http_operation_source( filters = json.loads(filters_param) except json.JSONDecodeError as exc: raise ProviderValidationError(f"Invalid JSON: {exc.msg}") from exc - return { + return _validate_http_source(operation_name, { "user_id": (query_params.get("user_id") or [None])[0], "agent_id": (query_params.get("agent_id") or [None])[0], "run_id": (query_params.get("run_id") or [None])[0], "limit": int((query_params.get("limit") or ["100"])[0]), **({"cursor": (query_params.get("cursor") or [None])[0]} if operation_name == "list_page" else {}), "filters": filters, - } + }) if operation_name in {"get", "delete"}: return {"memory_id": path_params["memory_id"]} raise ProviderValidationError(f"Unsupported HTTP operation: {operation_name}") diff --git a/agentmemory/runtime/schema_validation.py b/agentmemory/runtime/schema_validation.py new file mode 100644 index 0000000..07530dd --- /dev/null +++ b/agentmemory/runtime/schema_validation.py @@ -0,0 +1,72 @@ +"""Shared JSON-schema-style argument validation. + +Extracted from ``agentmemory.mcp`` so every surface (MCP, HTTP) can validate +raw operation input against an operation's ``input_schema`` identically. The +validation intentionally mirrors the subset of JSON Schema the operation +schemas use: ``type``, ``required``, ``enum``, ``minimum``, and +``additionalProperties: false``. +""" + +from __future__ import annotations + +from typing import Any + +from agentmemory.providers.base import ProviderValidationError + + +def schema_type_matches(value: Any, expected_type: str) -> bool: + if expected_type == "object": + return isinstance(value, dict) + if expected_type == "array": + return isinstance(value, list) + if expected_type == "string": + return isinstance(value, str) + if expected_type == "integer": + return isinstance(value, int) and not isinstance(value, bool) + if expected_type == "number": + return isinstance(value, (int, float)) and not isinstance(value, bool) + if expected_type == "boolean": + return isinstance(value, bool) + if expected_type == "null": + return value is None + return True + + +def validate_arguments(schema: dict[str, Any], arguments: Any) -> dict[str, Any]: + if not isinstance(arguments, dict): + raise ProviderValidationError("MCP tool arguments must be an object.") + + if schema.get("type") == "object": + properties = schema.get("properties") or {} + required = schema.get("required") or [] + for field in required: + if field not in arguments: + raise ProviderValidationError(f"Missing required argument: {field}") + + if schema.get("additionalProperties") is False: + extra = sorted(key for key in arguments if key not in properties) + if extra: + raise ProviderValidationError(f"Unexpected argument: {extra[0]}") + + for field, value in arguments.items(): + field_schema = properties.get(field) + if not isinstance(field_schema, dict): + continue + expected_type = field_schema.get("type") + if isinstance(expected_type, str) and not schema_type_matches(value, expected_type): + raise ProviderValidationError(f"Argument '{field}' must be {expected_type}.") + allowed_values = field_schema.get("enum") + if isinstance(allowed_values, list) and value not in allowed_values: + raise ProviderValidationError( + f"Argument '{field}' must be one of: {', '.join(map(str, allowed_values))}." + ) + minimum = field_schema.get("minimum") + if ( + isinstance(minimum, (int, float)) + and isinstance(value, (int, float)) + and not isinstance(value, bool) + and value < minimum + ): + raise ProviderValidationError(f"Argument '{field}' must be >= {minimum}.") + + return arguments diff --git a/tests/test_agentmemory_operation_adapters.py b/tests/test_agentmemory_operation_adapters.py index 9b05111..fdac226 100644 --- a/tests/test_agentmemory_operation_adapters.py +++ b/tests/test_agentmemory_operation_adapters.py @@ -162,6 +162,38 @@ def test_http_operation_source_rejects_invalid_filters_json(self) -> None: with self.assertRaises(ProviderValidationError): http_operation_source("list", query_params={"filters": ["{"]}) + def test_http_operation_source_add_synthesizes_messages_from_text(self) -> None: + source = http_operation_source( + "add", + payload={"text": "hello", "user_id": "u1", "memory_type": "preference"}, + ) + + self.assertEqual(source["messages"], [{"role": "user", "content": "hello"}]) + self.assertEqual(source["user_id"], "u1") + self.assertEqual(source["memory_type"], "preference") + # Matches MCP/CLI: raw `text` is preserved alongside synthesized messages. + self.assertEqual(source["text"], "hello") + + def test_http_operation_source_add_requires_text(self) -> None: + with self.assertRaises(ProviderValidationError): + http_operation_source("add", payload={"user_id": "u1"}) + + def test_http_operation_source_add_rejects_unexpected_field(self) -> None: + with self.assertRaises(ProviderValidationError): + http_operation_source("add", payload={"text": "hi", "messages": []}) + + def test_http_operation_source_rejects_invalid_limit(self) -> None: + with self.assertRaises(ProviderValidationError): + http_operation_source("list", query_params={"limit": ["0"]}) + + def test_http_operation_source_search_rejects_invalid_limit(self) -> None: + with self.assertRaises(ProviderValidationError): + http_operation_source("search", payload={"query": "q", "limit": 0}) + + def test_http_operation_source_search_rejects_bad_enum_for_list_scopes_kind(self) -> None: + with self.assertRaises(ProviderValidationError): + http_operation_source("list_scopes", query_params={"kind": ["bogus"]}) + if __name__ == "__main__": unittest.main() From e8c4c1d4e094e315274bc3372e9208712c8d6b57 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 27 May 2026 02:47:50 +0000 Subject: [PATCH 3/4] Make rerank-coercion cert test hermetic (pin transport to direct) The test mocks a localjson-like direct provider and expects the dispatcher to reach the mocked memory_search, but it never pinned the transport. should_proxy_ to_api() then read the ambient active provider (mem0 by default), so the call leaked into the owner-process proxy and failed (rc=2) wherever no API is running -- e.g. the provider-certification CI job, which runs with no config. Mock should_proxy_to_api to False so the test exercises the in-process path it intends. https://claude.ai/code/session_01BKjtun7hVwxh6Hv3YLdrww --- tests/test_provider_contract_v1.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_provider_contract_v1.py b/tests/test_provider_contract_v1.py index 23380aa..364c885 100644 --- a/tests/test_provider_contract_v1.py +++ b/tests/test_provider_contract_v1.py @@ -264,6 +264,7 @@ def test_cli_search_coerces_unsupported_rerank_before_provider_call(self) -> Non original_capabilities = agentmemory_operations.active_provider_capabilities original_provider_name = agentmemory_operations.active_provider_name original_memory_search = agentmemory_operations.memory_search + original_should_proxy = agentmemory_operations.should_proxy_to_api captured: dict = {} def fake_memory_search(**kwargs): @@ -293,6 +294,12 @@ def fake_memory_search(**kwargs): } agentmemory_operations.active_provider_name = lambda: "localjson" # type: ignore[assignment] agentmemory_operations.memory_search = fake_memory_search # type: ignore[assignment] + # The mocked provider is a direct (localjson-like) backend, so the + # dispatcher must run in-process and reach fake_memory_search. Pin + # the transport to direct; otherwise should_proxy_to_api() reads the + # ambient active provider (mem0 by default) and the call leaks into + # the owner-process proxy, failing wherever no API is running. + agentmemory_operations.should_proxy_to_api = lambda: False # type: ignore[assignment] rc = agentmemory_cli.main() finally: sys.argv = original_argv @@ -301,6 +308,7 @@ def fake_memory_search(**kwargs): agentmemory_operations.active_provider_capabilities = original_capabilities # type: ignore[assignment] agentmemory_operations.active_provider_name = original_provider_name # type: ignore[assignment] agentmemory_operations.memory_search = original_memory_search # type: ignore[assignment] + agentmemory_operations.should_proxy_to_api = original_should_proxy # type: ignore[assignment] self.assertEqual(rc, 0) self.assertNotIn("does not support rerank", stderr_buffer.getvalue()) From 3a900f82b53364a165382654c6f53330589854e7 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 27 May 2026 03:00:05 +0000 Subject: [PATCH 4/4] Fix pre-existing CI failures: harness import + portable launcher test These surfaced because this is the first PR to run the CI/certification workflows (they trigger on pull_request and push:main, but the default branch is master), not from the hardening changes. - provider-certification: test_localjson_provider (and mempalace, claude_memory) imported the harness only via `from tests.provider_contract_harness`, which fails under the cert runner's discover when `tests` isn't an importable package (PEP 660 editable install exposes only `agentmemory`). Mirror the robust top-level-first fallback already used by test_mem0_provider. - ubuntu/windows test job: test_text_config_status_detects_configured_launcher hardcoded an absolute launcher path that only matches the author's install root, reading as stale everywhere else. Build it from expected_launcher_path() so the assertion is host-portable. https://claude.ai/code/session_01BKjtun7hVwxh6Hv3YLdrww --- tests/test_agentmemory_clients.py | 10 +++++++++- tests/test_claude_memory_provider.py | 5 ++++- tests/test_localjson_provider.py | 5 ++++- tests/test_mempalace_provider.py | 5 ++++- 4 files changed, 21 insertions(+), 4 deletions(-) diff --git a/tests/test_agentmemory_clients.py b/tests/test_agentmemory_clients.py index 5322a24..1385568 100644 --- a/tests/test_agentmemory_clients.py +++ b/tests/test_agentmemory_clients.py @@ -47,7 +47,15 @@ def test_config_status_detects_stale_launcher(self) -> None: def test_text_config_status_detects_configured_launcher(self) -> None: with tempfile.TemporaryDirectory() as tmp: path = Path(tmp) / "settings.json" - path.write_text('{"agentmemory":{"args":["-File","O:/user files/Projects/tools/AgentMemory/scripts/run-agentmemory-mcp.ps1"]}}', encoding="utf-8") + # Build the config from the current machine's expected launcher so + # the assertion is portable: a hardcoded absolute path only matches + # the author's install root and reads as stale on every other host + # (Linux/Windows CI included). + launcher = agentmemory_clients.expected_launcher_path() + path.write_text( + json.dumps({"agentmemory": {"args": ["-File", launcher]}}), + encoding="utf-8", + ) payload = agentmemory_clients.text_config_status(path, "cli-client") self.assertTrue(payload["configured"]) self.assertEqual(payload["health"], "configured") diff --git a/tests/test_claude_memory_provider.py b/tests/test_claude_memory_provider.py index f02f819..e67df88 100644 --- a/tests/test_claude_memory_provider.py +++ b/tests/test_claude_memory_provider.py @@ -5,7 +5,10 @@ from pathlib import Path from agentmemory.providers.claude_memory import ClaudeMemoryProvider -from tests.provider_contract_harness import ProviderContractHarness +try: + from provider_contract_harness import ProviderContractHarness +except ModuleNotFoundError: # pragma: no cover + from tests.provider_contract_harness import ProviderContractHarness class ClaudeMemoryProviderHarnessTests(ProviderContractHarness, unittest.TestCase): diff --git a/tests/test_localjson_provider.py b/tests/test_localjson_provider.py index 1485690..7910c78 100644 --- a/tests/test_localjson_provider.py +++ b/tests/test_localjson_provider.py @@ -7,7 +7,10 @@ from agentmemory.providers.localjson import LocalJsonProvider from agentmemory.runtime import scope_registry -from tests.provider_contract_harness import ProviderContractHarness +try: + from provider_contract_harness import ProviderContractHarness +except ModuleNotFoundError: # pragma: no cover + from tests.provider_contract_harness import ProviderContractHarness def _add_localjson_records(runtime_dir: str, storage_path: str, count: int, prefix: str) -> None: diff --git a/tests/test_mempalace_provider.py b/tests/test_mempalace_provider.py index 4ca38de..d709eb6 100644 --- a/tests/test_mempalace_provider.py +++ b/tests/test_mempalace_provider.py @@ -7,7 +7,10 @@ from agentmemory.providers.base import MemoryNotFoundError, ProviderCapabilityError from agentmemory.providers.mempalace import MemPalaceProvider -from tests.provider_contract_harness import ProviderContractHarness +try: + from provider_contract_harness import ProviderContractHarness +except ModuleNotFoundError: # pragma: no cover + from tests.provider_contract_harness import ProviderContractHarness class FakePalaceNotFoundError(RuntimeError):