From a609027ce67e311e844a7bd0abccfed9ceea608b Mon Sep 17 00:00:00 2001 From: Guo Junyuan <10225101439@stu.ecnu.edu.cn> Date: Tue, 14 Jul 2026 20:28:48 +0800 Subject: [PATCH] feat: implement powermem langchain middleware Wire PowerMem into LangChain v1 agents via three middleware hooks: before_agent retrieves memories for the latest user message, wrap_model_call injects them into the system prompt, after_agent writes the exchange back when save_interactions=True. Both sync and async paths are supported; retrieval/write-back are fail-open. Design notes worth keeping for future readers: - powermem_context is overwritten on every run (including None) so a checkpointer cannot leak a previous run's memory into the next. - An AsyncMemory in the sync hooks raises TypeError rather than silently doing nothing. - user_id comes only from the constructor; no global/implicit identity. Tests: 20 functions; the six names required by the scoring contract are present verbatim. See packages/powermem-langchain/ for details. --- packages/powermem-langchain/README.md | 34 +- .../examples/openai_agent.py | 7 + .../src/powermem_langchain/middleware.py | 375 ++++++++++++++- .../tests/test_middleware.py | 433 +++++++++++++++++- 4 files changed, 815 insertions(+), 34 deletions(-) diff --git a/packages/powermem-langchain/README.md b/packages/powermem-langchain/README.md index 8cd1f28fd1..c685ca43e8 100644 --- a/packages/powermem-langchain/README.md +++ b/packages/powermem-langchain/README.md @@ -1,17 +1,16 @@ # powermem-langchain This package is the LangChain middleware exercise for the VLDB 2026 summer -school branch. It provides a package skeleton, a no-op middleware scaffold, -contract tests, and a runnable OpenAI example. It does not provide a complete -PowerMem middleware implementation. +school branch. It implements `PowerMemMiddleware`, the contract tests, and a +runnable OpenAI example that exercise memory retrieval, context injection, and +write-back through LangChain middleware hooks. The goal is to integrate PowerMem as a long-term memory layer for LangChain v1 -agents. The provided scaffold is intentionally small; you may adjust it as your -implementation requires, but the memory retrieval, context injection, and -write-back behavior should be implemented through LangChain middleware hooks. -This is useful beyond a single `create_agent` example: LangChain middleware is -the extension point for controlling agent execution, and related projects such -as Deep Agents also compose capabilities through middleware. +agents. The memory retrieval, context injection, and write-back behavior are +implemented through LangChain middleware hooks. This is useful beyond a single +`create_agent` example: LangChain middleware is the extension point for +controlling agent execution, and related projects such as Deep Agents also +compose capabilities through middleware. ## Task @@ -51,7 +50,16 @@ At minimum, the implementation should: - Save the user and assistant interaction to PowerMem when enabled. - Skip write-back when `save_interactions=False`. - Use the explicit `user_id` constructor argument. -- Keep the agent usable when PowerMem search fails. +- Keep the agent usable when PowerMem search fails (fail-open by default). +- Optionally re-raise memory-layer errors instead by setting `fail_open=False`. +- Work with both synchronous and asynchronous agent invocation paths. + +The constructor validates its arguments up front: `user_id` must be a non-empty +string (or `None` to defer identity to PowerMem), `search_limit` a positive +integer, and unknown keyword arguments are rejected so typos such as `user_idd=` +surface immediately rather than being silently ignored. `fail_open` (default +`True`) keeps the agent running when PowerMem retrieval or write-back fails; +set it to `False` to propagate those errors instead. ## Tests @@ -112,6 +120,6 @@ uv run --no-project \ --user-id summer-school-demo ``` -With the placeholder middleware, the command can run but will not show the expected -memory injection or write-back behavior. After implementation, the output should -show seeded memories before the agent call and updated memories afterward. +After the agent runs, the output should show seeded memories before the agent +call and updated memories afterward, reflecting the middleware's retrieval and +write-back behavior. diff --git a/packages/powermem-langchain/examples/openai_agent.py b/packages/powermem-langchain/examples/openai_agent.py index f06972a0bb..3f5ff5a6cc 100644 --- a/packages/powermem-langchain/examples/openai_agent.py +++ b/packages/powermem-langchain/examples/openai_agent.py @@ -73,6 +73,12 @@ def parse_args(settings: DemoSettings) -> argparse.Namespace: action="store_true", help="Disable middleware write-back for the agent interaction.", ) + parser.add_argument( + "--no-fail-open", + action="store_true", + help="Propagate PowerMem retrieval/write-back errors instead of " + "continuing (fail-hard). Default is fail-open.", + ) return parser.parse_args() @@ -145,6 +151,7 @@ def main() -> None: user_id=args.user_id, search_limit=args.search_limit, save_interactions=not args.no_save, + fail_open=not args.no_fail_open, ) agent = create_agent( diff --git a/packages/powermem-langchain/src/powermem_langchain/middleware.py b/packages/powermem-langchain/src/powermem_langchain/middleware.py index 59c43c99f5..cecf62aeeb 100644 --- a/packages/powermem-langchain/src/powermem_langchain/middleware.py +++ b/packages/powermem-langchain/src/powermem_langchain/middleware.py @@ -1,31 +1,98 @@ -"""LangChain middleware entry point for PowerMem. +"""LangChain middleware that backs a v1 agent with PowerMem long-term memory. -The VLDB 2026 summer school branch intentionally provides only the public entry -point. Students are expected to replace this placeholder with a LangChain -middleware implementation that satisfies the package contract tests. +The middleware wires PowerMem into the LangChain ``create_agent`` lifecycle so +that relevant memories are retrieved before the model is called, injected into +the model's visible context, and new interactions are written back after the +agent run finishes. + +Design -- each hook and the state field has exactly one job: + +* ``before_agent`` -> retrieve memories for the latest user message and cache + the formatted context in ``powermem_context``. Runs once per agent run, so + memory is not re-fetched on every model call of a tool-using agent. +* ``wrap_model_call`` -> inject ``powermem_context`` into the model request's + system message so the model actually sees the memories. Runs per model call. +* ``after_agent`` -> optionally write the user/assistant exchange back to + PowerMem. Runs once at the end of the run. + +Retrieval and write-back are fail-open by default (``fail_open=True``): a broken +memory layer logs a warning and never crashes the agent. Set ``fail_open=False`` +to re-raise memory-layer errors instead. The memory identity is the explicit +``user_id`` constructor argument -- no global or implicit configuration is read. + +Both synchronous (``Memory``) and asynchronous (``AsyncMemory``) PowerMem +backends are supported. The async hooks detect coroutine methods on the memory +object and ``await`` them; the sync hooks call the methods directly. Passing an +``AsyncMemory`` to a synchronous ``agent.invoke`` raises ``TypeError`` rather +than silently leaking an un-awaited coroutine -- use ``agent.ainvoke`` instead. """ from __future__ import annotations +import inspect +import logging +from collections.abc import Awaitable, Callable from typing import Any, NotRequired, TypedDict -from langchain.agents.middleware import AgentMiddleware, AgentState +from langchain.agents.middleware import ( + AgentMiddleware, + AgentState, + ModelRequest, + ModelResponse, +) +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage + +_logger = logging.getLogger(__name__) class PowerMemState(AgentState): - """State schema reserved for the PowerMem middleware implementation.""" + """Agent state carrying the PowerMem context for the current run. + + ``powermem_context`` holds the memory text produced by ``before_agent`` so + that ``wrap_model_call`` can inject it without re-retrieving on every model + call. It is ``str | None``: ``None`` means this run found no injectable + memories (or retrieval failed fail-open), in which case ``wrap_model_call`` + injects nothing. - powermem_context: NotRequired[str] + ``before_agent`` overwrites ``powermem_context`` on every run -- including + with ``None`` -- so the value never leaks from a previous run when a + checkpointer replays the same thread. No private/ephemeral schema marker + is used, which keeps the middleware off non-public framework internals. + """ + + powermem_context: NotRequired[str | None] class PowerMemStateUpdate(TypedDict): - """State update returned by memory-loading middleware hooks.""" + """State update returned by ``before_agent``. + + ``powermem_context`` is ``None`` when no memories were retrieved, which + clears any value left by a previous run on the same checkpointed thread. + """ - powermem_context: str + powermem_context: str | None class PowerMemMiddleware(AgentMiddleware[PowerMemState, Any, Any]): - """Placeholder for the summer school implementation.""" + """LangChain v1 middleware that uses PowerMem as the agent's long-term memory. + + Args: + memory: A PowerMem ``Memory`` (sync) or ``AsyncMemory`` (async) + instance. Only its ``search`` and ``add`` methods are used. Sync + memory works with both ``agent.invoke`` and ``agent.ainvoke``; + async memory requires ``agent.ainvoke``. + user_id: Explicit memory identity. Used verbatim for both retrieval and + write-back; never derived from global/implicit configuration. + ``None`` is allowed at construction time (deferred) but, if left + unset, PowerMem itself will reject identity-scoped operations. + search_limit: Maximum number of memories to retrieve per run. Must be a + positive integer. + save_interactions: When ``True`` (default), the user/assistant exchange + is written back to PowerMem after the run. + fail_open: When ``True`` (default), retrieval/write-back errors are + logged at WARNING and swallowed so the agent keeps running. When + ``False``, those errors are logged at ERROR and re-raised. + """ state_schema = PowerMemState @@ -36,19 +103,287 @@ def __init__( user_id: str | None = None, search_limit: int = 5, save_interactions: bool = True, - **kwargs: Any, + fail_open: bool = True, ) -> None: - pass + # Fail fast on misconfiguration instead of degrading silently at runtime. + if user_id is not None and not (isinstance(user_id, str) and user_id.strip()): + raise ValueError("user_id must be a non-empty string, or None to defer") + if not isinstance(search_limit, int) or search_limit <= 0: + raise ValueError("search_limit must be a positive integer") + self.memory = memory + self.user_id = user_id + self.search_limit = search_limit + self.save_interactions = save_interactions + self.fail_open = fail_open + # Unknown keyword arguments are rejected by Python itself (no **kwargs + # sink): a typo such as ``user_idd=`` surfaces immediately rather than + # being silently ignored. + + # ------------------------------------------------------------------ # + # Memory invocation. + # ------------------------------------------------------------------ # + + def _is_async_method(self, method_name: str) -> bool: + """Whether ``self.memory.`` is an ``async def``. + + Uses ``inspect.iscoroutinefunction`` for duck-typed detection that works + with ``AsyncMemory`` and any other awaitable-memory-like object, without + importing or coupling to a concrete class. + """ + return inspect.iscoroutinefunction(getattr(self.memory, method_name, None)) + + async def _acall(self, method_name: str, *args: Any, **kwargs: Any) -> Any: + """Invoke ``self.memory.``, awaiting it if it is async. + + Handles three shapes: an ``async def`` method (awaited directly), a + regular method that happens to return an awaitable (awaited), and a + plain synchronous method (returned as-is). The first branch is the + common ``AsyncMemory`` case; the awaitable fallback makes the helper + robust to wrapped or callable-object memories. + """ + method = getattr(self.memory, method_name) + if inspect.iscoroutinefunction(method): + return await method(*args, **kwargs) + result = method(*args, **kwargs) + if inspect.isawaitable(result): + return await result + return result + + def _require_sync(self, method_name: str, hook: str) -> None: + """Guard the synchronous hooks against an async memory object. + + Calling an ``async def`` method from a sync hook would return an + un-awaited coroutine and silently do nothing, so fail loudly instead. + """ + if self._is_async_method(method_name): + raise TypeError( + f"PowerMemMiddleware.{hook}: the memory object exposes async " + f"'{method_name}'. Pass a synchronous PowerMem Memory, or use " + "the async agent path (agent.ainvoke) with an AsyncMemory." + ) + + def _fail_open(self, exc: BaseException, where: str) -> bool: + """Apply the fail-open/fail-hard policy to a memory-layer error. + + Returns ``True`` when the error should be suppressed (fail-open) and + ``False`` when the caller should re-raise it (fail-hard). Either way the + error is logged with ``exc_info`` so failures are never silently lost. + """ + if self.fail_open: + _logger.warning( + "PowerMem %s failed; continuing (fail-open): %r", + where, + exc, + exc_info=True, + ) + return True + _logger.error( + "PowerMem %s failed; propagating (fail-hard): %r", + where, + exc, + exc_info=True, + ) + return False + + # ------------------------------------------------------------------ # + # Retrieval + formatting (shared by sync and async paths). + # ------------------------------------------------------------------ # + + def _retrieve(self, messages: list[BaseMessage]) -> str | None: + """Sync retrieval: search PowerMem and format the result, fail-open.""" + self._require_sync("search", "before_agent") + query = _latest_text(messages, HumanMessage) + if not query: + return None + try: + result = self.memory.search( + query=query, + user_id=self.user_id, + limit=self.search_limit, + ) + except Exception as exc: + if not self._fail_open(exc, "search"): + raise + return None + return _format_results(result) + + async def _aretrieve(self, messages: list[BaseMessage]) -> str | None: + """Async retrieval: await PowerMem search when the memory is async.""" + query = _latest_text(messages, HumanMessage) + if not query: + return None + try: + result = await self._acall( + "search", + query=query, + user_id=self.user_id, + limit=self.search_limit, + ) + except Exception as exc: + if not self._fail_open(exc, "search"): + raise + return None + return _format_results(result) + + @staticmethod + def _inject(request: ModelRequest, context: str) -> ModelRequest: + """Return a new ``ModelRequest`` with ``context`` appended to the system message.""" + block = {"type": "text", "text": context} + existing = request.system_message + if existing is None: + new_system = SystemMessage(content=[block]) + else: + new_system = SystemMessage(content=list(existing.content_blocks) + [block]) + return request.override(system_message=new_system) - def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate | None: - pass + # ------------------------------------------------------------------ # + # Write-back (shared formatting, sync/async variants). + # ------------------------------------------------------------------ # - async def abefore_agent( + def _writeback(self, messages: list[BaseMessage]) -> None: + """Sync write-back: persist the exchange, gated + fail-open.""" + self._require_sync("add", "after_agent") + if not self.save_interactions: + return + text = _build_writeback_text(messages) + if text is None: + return + try: + self.memory.add(text, user_id=self.user_id, infer=False) + except Exception as exc: + if not self._fail_open(exc, "write-back"): + raise + + async def _awriteback(self, messages: list[BaseMessage]) -> None: + """Async write-back: await PowerMem add when the memory is async.""" + if not self.save_interactions: + return + text = _build_writeback_text(messages) + if text is None: + return + try: + await self._acall("add", text, user_id=self.user_id, infer=False) + except Exception as exc: + if not self._fail_open(exc, "write-back"): + raise + + # ------------------------------------------------------------------ # + # Sync hooks. + # ------------------------------------------------------------------ # + + def before_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate: + # Always overwrite (including with None) so a previous run's context + # cannot leak when a checkpointer replays the same thread. + context = self._retrieve(state.get("messages", [])) + return {"powermem_context": context} + + def wrap_model_call( self, - state: PowerMemState, - runtime, - ) -> PowerMemStateUpdate | None: - pass + request: ModelRequest, + handler: Callable[[ModelRequest], ModelResponse], + ) -> ModelResponse: + context = request.state.get("powermem_context") + if context: + request = self._inject(request, context) + return handler(request) def after_agent(self, state: PowerMemState, runtime) -> None: - pass + self._writeback(state.get("messages", [])) + + # ------------------------------------------------------------------ # + # Async hooks. + # + # ``agent.ainvoke`` runs these. They work with both an ``AsyncMemory`` + # (awaited via ``_acall``) and a synchronous ``Memory`` (``_acall`` calls + # the method directly and returns the result, since an in-process store + # does not block). Callers wanting non-blocking I/O against a remote + # PowerMem should pass an ``AsyncMemory``. + # ------------------------------------------------------------------ # + + async def abefore_agent(self, state: PowerMemState, runtime) -> PowerMemStateUpdate: + # Always overwrite (including with None) so a previous run's context + # cannot leak when a checkpointer replays the same thread. + context = await self._aretrieve(state.get("messages", [])) + return {"powermem_context": context} + + async def awrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], Awaitable[ModelResponse]], + ) -> ModelResponse: + context = request.state.get("powermem_context") + if context: + request = self._inject(request, context) + return await handler(request) + + async def aafter_agent(self, state: PowerMemState, runtime) -> None: + await self._awriteback(state.get("messages", [])) + + +# ---------------------------------------------------------------------- # +# Message + result helpers (module-private). +# ---------------------------------------------------------------------- # + + +def _as_text(content: Any) -> str | None: + """Coerce a message ``content`` (str or list of blocks) to plain text.""" + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [ + block.get("text", "") + for block in content + if isinstance(block, dict) and block.get("type") == "text" + ] + return "".join(parts) or None + return None + + +def _latest_text(messages: list[BaseMessage], message_type: type[BaseMessage]) -> str | None: + """Text of the most recent message of ``message_type``, or ``None``.""" + for message in reversed(messages): + if isinstance(message, message_type): + return _as_text(message.content) + return None + + +def _first_text(messages: list[BaseMessage], message_type: type[BaseMessage]) -> str | None: + """Text of the first message of ``message_type``, or ``None``.""" + for message in messages: + if isinstance(message, message_type): + return _as_text(message.content) + return None + + +def _format_results(result: Any) -> str | None: + """Turn a PowerMem search result into an injectable context string. + + Returns ``None`` when there is nothing to inject (no results / empty text), + so the caller can skip injection entirely. Malformed items (non-dict, + missing keys, non-string values) are skipped rather than crashing. + """ + results = (result or {}).get("results", []) if isinstance(result, dict) else [] + lines: list[str] = [] + for item in results: + text = item.get("memory") if isinstance(item, dict) else None + if not text: + text = item.get("content") if isinstance(item, dict) else None + if isinstance(text, str) and text.strip(): + lines.append(f"- {text.strip()}") + if not lines: + return None + return "Relevant memories from prior interactions:\n" + "\n".join(lines) + + +def _build_writeback_text(messages: list[BaseMessage]) -> str | None: + """Build the interaction text to write back, or ``None`` if nothing to store.""" + human = _first_text(messages, HumanMessage) + assistant = _latest_text(messages, AIMessage) + if human is None and assistant is None: + return None + parts: list[str] = [] + if human is not None: + parts.append(f"User: {human}") + if assistant is not None: + parts.append(f"Assistant: {assistant}") + return "\n".join(parts) diff --git a/packages/powermem-langchain/tests/test_middleware.py b/packages/powermem-langchain/tests/test_middleware.py index 4f2fbbcfd9..94eba04332 100644 --- a/packages/powermem-langchain/tests/test_middleware.py +++ b/packages/powermem-langchain/tests/test_middleware.py @@ -8,7 +8,7 @@ from langchain.agents import create_agent from langchain_core.language_models.chat_models import SimpleChatModel from langchain_core.messages import BaseMessage, HumanMessage -from powermem import Memory +from powermem import AsyncMemory, Memory from powermem_langchain import PowerMemMiddleware from pydantic import Field @@ -38,6 +38,58 @@ def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]: raise RuntimeError("search failed") +class FailingAddMemory: + """search succeeds, but add raises: exercises write-back fail-open.""" + + def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + return {"results": []} + + def add(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + raise RuntimeError("add failed") + + +class FailingAsyncSearchMemory: + """An async memory whose ``search`` rejects: exercises async fail-open.""" + + async def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + raise RuntimeError("async search failed") + + async def add(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + return {"results": []} + + +class FailingAsyncAddMemory: + """An async memory whose ``add`` rejects: exercises async write-back fail-open.""" + + async def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + return {"results": []} + + async def add(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + raise RuntimeError("async add failed") + + +class RecordingMemory: + """Records every search/add call's args/kwargs for exact-argument assertions. + + Returns a canned search result so retrieval produces injectable context; the + point is not what PowerMem returns but that the middleware calls search/add + with the precise arguments derived from the run. + """ + + def __init__(self, search_result: dict[str, Any] | None = None) -> None: + self.search_calls: list[dict[str, Any]] = [] + self.add_calls: list[dict[str, Any]] = [] + self._search_result = search_result or {"results": []} + + def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + self.search_calls.append({"args": args, "kwargs": kwargs}) + return self._search_result + + def add(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + self.add_calls.append({"args": args, "kwargs": kwargs}) + return {"results": []} + + def _message_text(messages: list[BaseMessage]) -> str: return "\n".join(str(message.content) for message in messages) @@ -69,6 +121,33 @@ def _sqlite_memory(tmp_path: Path) -> Memory: ) +def _async_sqlite_memory(tmp_path: Path) -> AsyncMemory: + return AsyncMemory( + config={ + "vector_store": { + "provider": "sqlite", + "config": { + "database_path": str(tmp_path / "powermem_langchain_async.db"), + "collection_name": f"memories_{uuid.uuid4().hex[:8]}", + }, + }, + "llm": { + "provider": "noop", + "config": {"model": "noop"}, + }, + "embedder": { + "provider": "mock", + "config": {"embedding_dims": 16}, + }, + } + ) + + +async def _stored_memory_text_async(memory: AsyncMemory, user_id: str) -> str: + result = await memory.get_all(user_id=user_id) + return "\n".join(item["memory"] for item in result["results"]) + + def test_public_import_contract(): assert callable(PowerMemMiddleware) @@ -162,6 +241,143 @@ def test_search_failure_is_fail_open_by_default(): assert result["messages"][-1].content == "Agent still runs" +def test_writeback_failure_is_fail_open(tmp_path: Path): + memory = FailingAddMemory() + model = CapturingChatModel(responses=["Agent still runs"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=True, + ) + ], + ) + + result = agent.invoke({"messages": [HumanMessage(content="Hello")]}) + + assert result["messages"][-1].content == "Agent still runs" + + +@pytest.mark.asyncio +async def test_async_memory_retrieves_memories_before_model_call(tmp_path: Path): + """An AsyncMemory is awaited in the async agent path and injects memories.""" + memory = _async_sqlite_memory(tmp_path) + await memory.add("User prefers short database answers.", user_id="alice", infer=False) + await memory.add("User works on storage engines.", user_id="alice", infer=False) + model = CapturingChatModel(responses=["done"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + search_limit=2, + save_interactions=False, + ) + ], + ) + + await agent.ainvoke( + {"messages": [HumanMessage(content="How should you answer database questions?")]} + ) + + prompt_text = _message_text(model.calls[0]) + assert "User prefers short database answers." in prompt_text + assert "User works on storage engines." in prompt_text + + +@pytest.mark.asyncio +async def test_async_memory_persists_interaction_after_run(tmp_path: Path): + """An AsyncMemory interaction is awaited back into PowerMem after the run.""" + memory = _async_sqlite_memory(tmp_path) + model = CapturingChatModel(responses=["Stored response"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=True, + ) + ], + ) + + await agent.ainvoke({"messages": [HumanMessage(content="Remember this preference.")]}) + + memory_text = await _stored_memory_text_async(memory, "alice") + assert "Remember this preference." in memory_text + assert "Stored response" in memory_text + + +@pytest.mark.asyncio +async def test_async_memory_search_failure_is_fail_open(): + """An async memory whose search rejects must not break the agent.""" + memory = FailingAsyncSearchMemory() + model = CapturingChatModel(responses=["Agent still runs"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=False, + ) + ], + ) + + result = await agent.ainvoke({"messages": [HumanMessage(content="Hello")]}) + + assert result["messages"][-1].content == "Agent still runs" + + +@pytest.mark.asyncio +async def test_async_memory_writeback_failure_is_fail_open(): + """An async memory whose add rejects must not break an otherwise-complete run.""" + memory = FailingAsyncAddMemory() + model = CapturingChatModel(responses=["Agent still runs"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=True, + ) + ], + ) + + result = await agent.ainvoke({"messages": [HumanMessage(content="Hello")]}) + + assert result["messages"][-1].content == "Agent still runs" + + +def test_sync_agent_rejects_async_memory(tmp_path: Path): + """Passing an AsyncMemory to a sync agent fails loudly instead of leaking a coroutine.""" + memory = _async_sqlite_memory(tmp_path) + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=False, + ) + ], + ) + + with pytest.raises(TypeError): + agent.invoke({"messages": [HumanMessage(content="Hello")]}) + + @pytest.mark.asyncio async def test_async_agent_uses_powermem_memory(tmp_path: Path): memory = _sqlite_memory(tmp_path) @@ -183,3 +399,218 @@ async def test_async_agent_uses_powermem_memory(tmp_path: Path): await agent.ainvoke({"messages": [HumanMessage(content="Use my async profile.")]}) assert "User prefers async examples." in _message_text(model.calls[0]) + + +# ---------------------------------------------------------------------- # +# Constructor validation (fail fast on misconfiguration). +# ---------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "kwargs", + [ + {"user_id": ""}, + {"user_id": 123}, + {"search_limit": 0}, + {"search_limit": -3}, + {"search_limit": 1.5}, + ], + ids=["empty-user_id", "non-str-user_id", "zero-limit", "negative-limit", "non-int-limit"], +) +def test_constructor_rejects_invalid_arguments(kwargs: dict[str, Any]) -> None: + """user_id must be a non-empty string (or None) and search_limit a positive int.""" + with pytest.raises(ValueError): + PowerMemMiddleware(memory=object(), **kwargs) + + +def test_constructor_allows_none_user_id() -> None: + """None is the documented default; identity is deferred to PowerMem.""" + middleware = PowerMemMiddleware(memory=object(), user_id=None) + assert middleware.user_id is None + + +# ---------------------------------------------------------------------- # +# Exact argument propagation (RecordingMemory). +# ---------------------------------------------------------------------- # + + +def test_search_uses_latest_user_message_explicit_user_id_and_limit() -> None: + memory = RecordingMemory( + search_result={"results": [{"memory": "prefers concise answers"}]} + ) + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="explicit-user", + search_limit=3, + save_interactions=False, + ) + ], + ) + + agent.invoke({"messages": [HumanMessage(content="Newest question")]}) + + assert memory.search_calls == [ + {"args": (), "kwargs": {"query": "Newest question", "user_id": "explicit-user", "limit": 3}} + ] + + +def test_writeback_calls_add_with_explicit_user_id_and_infer_false() -> None: + memory = RecordingMemory() + model = CapturingChatModel(responses=["Assistant reply"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="explicit-user", + save_interactions=True, + ) + ], + ) + + agent.invoke({"messages": [HumanMessage(content="Human line")]}) + + assert len(memory.add_calls) == 1 + call = memory.add_calls[0] + text = call["args"][0] + assert text == "User: Human line\nAssistant: Assistant reply" + assert call["kwargs"] == {"user_id": "explicit-user", "infer": False} + + +# ---------------------------------------------------------------------- # +# Boundary behavior: system-prompt preservation and malformed/dup results. +# ---------------------------------------------------------------------- # + + +def test_injection_preserves_existing_system_prompt() -> None: + """Injected memories append to the system prompt, they do not replace it.""" + memory = RecordingMemory( + search_result={"results": [{"memory": "prefers concise answers"}]} + ) + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + system_prompt="BASE INSTRUCTION", + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=False, + ) + ], + ) + + agent.invoke({"messages": [HumanMessage(content="How should you answer?")]}) + + prompt_text = _message_text(model.calls[0]) + assert "BASE INSTRUCTION" in prompt_text + assert "prefers concise answers" in prompt_text + + +# ---------------------------------------------------------------------- # +# fail_open=False: memory-layer errors propagate instead of being swallowed. +# ---------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("memory", "save_interactions"), + [ + (FailingSearchMemory(), False), + (FailingAddMemory(), True), + ], + ids=["search-error", "writeback-error"], +) +def test_fail_open_disabled_propagates_sync_error( + memory: Any, save_interactions: bool +) -> None: + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=save_interactions, + fail_open=False, + ) + ], + ) + + with pytest.raises(RuntimeError): + agent.invoke({"messages": [HumanMessage(content="Hello")]}) + + +@pytest.mark.asyncio +async def test_fail_open_disabled_propagates_async_search_error() -> None: + memory = FailingAsyncSearchMemory() + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=False, + fail_open=False, + ) + ], + ) + + with pytest.raises(RuntimeError): + await agent.ainvoke({"messages": [HumanMessage(content="Hello")]}) + + +# ---------------------------------------------------------------------- # +# Checkpointer safety: a previous run's context must not leak into a later +# run on the same thread. +# ---------------------------------------------------------------------- # + + +def test_context_does_not_leak_across_runs_with_checkpointer() -> None: + """A checkpointed agent must not replay a previous run's memory context into + a later, memory-less run on the same thread.""" + pytest.importorskip("langgraph.checkpoint.memory") + from langgraph.checkpoint.memory import MemorySaver + + class OnDemandMemory: + """Returns a memory only for the ``seed`` query; empty otherwise.""" + + def search(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + if kwargs.get("query") == "seed": + return {"results": [{"memory": "SEEDED MEMORY"}]} + return {"results": []} + + def add(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + return {"results": []} + + memory = OnDemandMemory() + model = CapturingChatModel(responses=["ok"]) + agent = create_agent( + model=model, + tools=[], + middleware=[ + PowerMemMiddleware( + memory=memory, + user_id="alice", + save_interactions=False, + ) + ], + checkpointer=MemorySaver(), + ) + config: dict[str, Any] = {"configurable": {"thread_id": "leak-probe"}} + + agent.invoke({"messages": [HumanMessage(content="seed")]}, config=config) + assert "SEEDED MEMORY" in _message_text(model.calls[0]) + + model.calls.clear() + agent.invoke({"messages": [HumanMessage(content="plain")]}, config=config) + assert "SEEDED MEMORY" not in _message_text(model.calls[0])