From b21108a40700b13dbcd027bb7cc372ff5cf6a529 Mon Sep 17 00:00:00 2001 From: Agent Zero Date: Sun, 23 Aug 2026 19:36:30 +0200 Subject: [PATCH] fix(memory): block turns when recall times out --- plugins/_memory/AGENTS.md | 1 + .../_91_recall_wait.py | 22 ++- tests/test_memory_recall_wait.py | 178 ++++++++++++++++++ 3 files changed, 199 insertions(+), 2 deletions(-) create mode 100644 tests/test_memory_recall_wait.py diff --git a/plugins/_memory/AGENTS.md b/plugins/_memory/AGENTS.md index ed8b56d118..17ced65ff8 100644 --- a/plugins/_memory/AGENTS.md +++ b/plugins/_memory/AGENTS.md @@ -18,6 +18,7 @@ - Preserve embedding metadata needed to rebuild indexes safely. - `memory_load` accepts numeric `threshold` and `limit` values as native numbers or numeric strings and coerces them before vector search. - Avoid storing transient action-history noise as durable memory. +- When enabled automatic recall times out, record a concise visible error and block that agent turn; never generate a response without the required recalled-memory context. Preserve normal cancellation propagation. ## Work Guidance diff --git a/plugins/_memory/extensions/python/message_loop_prompts_after/_91_recall_wait.py b/plugins/_memory/extensions/python/message_loop_prompts_after/_91_recall_wait.py index 6076847ca9..e201686858 100644 --- a/plugins/_memory/extensions/python/message_loop_prompts_after/_91_recall_wait.py +++ b/plugins/_memory/extensions/python/message_loop_prompts_after/_91_recall_wait.py @@ -1,4 +1,7 @@ +import asyncio + from helpers.extension import Extension +from helpers.errors import HandledException from agent import LoopData from plugins._memory.extensions.python.message_loop_prompts_after._50_recall_memories import DATA_NAME_TASK as DATA_NAME_TASK_MEMORIES, DATA_NAME_ITER as DATA_NAME_ITER_MEMORIES from helpers import plugins @@ -16,7 +19,7 @@ async def execute(self, loop_data: LoopData = LoopData(), **kwargs): task = self.agent.get_data(DATA_NAME_TASK_MEMORIES) iter = self.agent.get_data(DATA_NAME_ITER_MEMORIES) or 0 - if task and not task.done(): + if task: # if memory recall is set to delayed mode, do not await on the iteration it was called if set["memory_recall_delayed"]: @@ -27,4 +30,19 @@ async def execute(self, loop_data: LoopData = LoopData(), **kwargs): return # otherwise await the task - await task + try: + await task + except asyncio.TimeoutError as error: + self.agent.context.log.log( + type="error", + heading="Memory recall timed out", + content=( + "No response was generated because required memory recall did " + "not finish within 30 seconds. Retry the request after the " + "memory service recovers." + ), + ) + self.agent.set_data(DATA_NAME_TASK_MEMORIES, None) + raise HandledException( + "Required memory recall timed out; agent response blocked." + ) from error diff --git a/tests/test_memory_recall_wait.py b/tests/test_memory_recall_wait.py new file mode 100644 index 0000000000..7b612b1b33 --- /dev/null +++ b/tests/test_memory_recall_wait.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import asyncio +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +TARGET_PATH = ( + PROJECT_ROOT + / "plugins" + / "_memory" + / "extensions" + / "python" + / "message_loop_prompts_after" + / "_91_recall_wait.py" +) +DATA_NAME_TASK = "_recall_memories_task" +DATA_NAME_ITER = "_recall_memories_iter" + + +def load_recall_wait_module(): + extension_module = types.ModuleType("helpers.extension") + + class Extension: + def __init__(self, agent=None, **_kwargs) -> None: + self.agent = agent + + extension_module.Extension = Extension + + agent_module = types.ModuleType("agent") + agent_module.LoopData = object + + recall_memories_module = types.ModuleType( + "plugins._memory.extensions.python.message_loop_prompts_after._50_recall_memories" + ) + recall_memories_module.DATA_NAME_TASK = DATA_NAME_TASK + recall_memories_module.DATA_NAME_ITER = DATA_NAME_ITER + + plugins_module = types.ModuleType("helpers.plugins") + plugins_module.get_plugin_config = lambda *_args, **_kwargs: { + "memory_recall_delayed": False, + } + errors_module = types.ModuleType("helpers.errors") + + class HandledException(Exception): + pass + + errors_module.HandledException = HandledException + helpers_module = types.ModuleType("helpers") + helpers_module.plugins = plugins_module + + modules = { + "agent": agent_module, + "helpers": helpers_module, + "helpers.extension": extension_module, + "helpers.errors": errors_module, + "helpers.plugins": plugins_module, + "plugins._memory.extensions.python.message_loop_prompts_after._50_recall_memories": recall_memories_module, + } + with patch.dict(sys.modules, modules): + spec = importlib.util.spec_from_file_location("test_recall_wait", TARGET_PATH) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class FakeLog: + def __init__(self) -> None: + self.entries: list[dict] = [] + + def log(self, **entry) -> None: + self.entries.append(entry) + + +class FakeAgent: + def __init__(self, task: asyncio.Task) -> None: + self.data = {DATA_NAME_TASK: task} + self.context = SimpleNamespace(log=FakeLog()) + + def get_data(self, key): + return self.data.get(key) + + def set_data(self, key, value) -> None: + self.data[key] = value + + +class RecallWaitTests(unittest.TestCase): + def setUp(self) -> None: + self.recall_wait = load_recall_wait_module() + + def test_timeout_logs_visible_block_and_prevents_agent_turn(self) -> None: + async def run() -> FakeAgent: + async def timed_out_recall() -> None: + await asyncio.sleep(0) + raise asyncio.TimeoutError + + task = asyncio.create_task(timed_out_recall()) + agent = FakeAgent(task) + + with self.assertRaises(self.recall_wait.HandledException) as exc: + await self.recall_wait.RecallWait(agent).execute( + SimpleNamespace(iteration=0) + ) + self.assertEqual( + str(exc.exception), + "Required memory recall timed out; agent response blocked.", + ) + return agent + + agent = asyncio.run(run()) + + assert agent.data[DATA_NAME_TASK] is None + assert agent.context.log.entries == [ + { + "type": "error", + "heading": "Memory recall timed out", + "content": ( + "No response was generated because required memory recall did " + "not finish within 30 seconds. Retry the request after the " + "memory service recovers." + ), + } + ] + + def test_completed_timeout_task_still_blocks_agent_turn(self) -> None: + async def run() -> FakeAgent: + async def timed_out_recall() -> None: + raise asyncio.TimeoutError + + task = asyncio.create_task(timed_out_recall()) + try: + await task + except asyncio.TimeoutError: + pass + self.assertTrue(task.done()) + + agent = FakeAgent(task) + with self.assertRaises(self.recall_wait.HandledException): + await self.recall_wait.RecallWait(agent).execute( + SimpleNamespace(iteration=0) + ) + return agent + + agent = asyncio.run(run()) + + assert agent.data[DATA_NAME_TASK] is None + assert agent.context.log.entries[0]["type"] == "error" + + def test_cancellation_still_propagates(self) -> None: + async def run() -> FakeAgent: + async def cancelled_recall() -> None: + await asyncio.sleep(0) + raise asyncio.CancelledError + + task = asyncio.create_task(cancelled_recall()) + agent = FakeAgent(task) + + with self.assertRaises(asyncio.CancelledError): + await self.recall_wait.RecallWait(agent).execute( + SimpleNamespace(iteration=0) + ) + return agent + + agent = asyncio.run(run()) + + assert agent.data[DATA_NAME_TASK] is not None + assert agent.context.log.entries == [] + + +if __name__ == "__main__": + unittest.main()