diff --git a/docs/a2a.md b/docs/a2a.md index a90f2c1e..be3e882c 100644 --- a/docs/a2a.md +++ b/docs/a2a.md @@ -81,3 +81,16 @@ Each conversation is identified by a `context_id`. All messages within the same - Remember what was discussed - Track which documents were already found - Provide contextual follow-up answers + +### Memory Management + +To prevent memory growth, the server uses LRU (Least Recently Used) eviction: + +- Maximum 1000 contexts kept in memory (configurable via `A2A_MAX_CONTEXTS`) +- When limit exceeded, least recently used contexts are automatically evicted +- No periodic cleanup needed - eviction happens on-demand + +Configure via environment variable: +```bash +export A2A_MAX_CONTEXTS=1000 +``` diff --git a/src/haiku/rag/a2a.py b/src/haiku/rag/a2a.py index fff1d7c2..38020e68 100644 --- a/src/haiku/rag/a2a.py +++ b/src/haiku/rag/a2a.py @@ -1,5 +1,6 @@ import logging import uuid +from collections import OrderedDict from contextlib import asynccontextmanager from pathlib import Path @@ -24,9 +25,10 @@ try: Message, TaskIdParams, TaskSendParams, + TaskState, TextPart, ) - from fasta2a.storage import InMemoryStorage # type: ignore + from fasta2a.storage import InMemoryStorage, Storage # type: ignore except ImportError as e: raise ImportError( "A2A support requires the 'a2a' extra. " @@ -148,6 +150,64 @@ def save_message_history(message_history: list[ModelMessage]) -> Message: ) +class LRUMemoryStorage(Storage[list["Message"]]): # type: ignore + """Storage wrapper with LRU eviction for contexts. + + Enforces a maximum context limit using LRU (Least Recently Used) eviction. + """ + + def __init__(self, storage: InMemoryStorage, max_contexts: int): + self.storage = storage + self.max_contexts = max_contexts + # Track context access order (LRU cache) + self.context_order: OrderedDict[str, None] = OrderedDict() + + async def load_context(self, context_id: str) -> list["Message"] | None: + """Load context and update access order.""" + result = await self.storage.load_context(context_id) + if result is not None: + # Move to end (most recently used) + self.context_order.pop(context_id, None) + self.context_order[context_id] = None + return result + + async def update_context(self, context_id: str, context: list["Message"]) -> None: + """Update context and enforce LRU limit.""" + await self.storage.update_context(context_id, context) + # Move to end (most recently used) + self.context_order.pop(context_id, None) + self.context_order[context_id] = None + + # Enforce max contexts limit (LRU eviction) + while len(self.context_order) > self.max_contexts: + # Remove oldest (first item in OrderedDict) + oldest_context_id = next(iter(self.context_order)) + self.context_order.pop(oldest_context_id) + logger.debug( + f"Evicted context {oldest_context_id} (LRU, limit={self.max_contexts})" + ) + + async def load_task(self, task_id: str, history_length: int | None = None): + """Delegate to underlying storage.""" + return await self.storage.load_task(task_id, history_length) + + async def update_task( + self, + task_id: str, + state: TaskState, + new_artifacts: list["Artifact"] | None = None, + new_messages: list["Message"] | None = None, + ): + """Delegate to underlying storage.""" + return await self.storage.update_task( + task_id, state, new_artifacts, new_messages + ) + + async def submit_task(self, context_id: str, message: "Message"): + """Delegate to underlying storage.""" + return await self.storage.submit_task(context_id, message) + + def extract_question_from_task(task_history: list[Message]) -> str | None: """Extract the user's question from task history. @@ -176,7 +236,10 @@ def create_a2a_app(db_path: Path): Returns: A FastA2A ASGI application """ - storage = InMemoryStorage() + base_storage = InMemoryStorage() + storage = LRUMemoryStorage( + storage=base_storage, max_contexts=Config.A2A_MAX_CONTEXTS + ) broker = InMemoryBroker() # Create the agent with native search tool @@ -341,6 +404,7 @@ def create_a2a_app(db_path: Path): # Create FastA2A app with custom worker lifecycle @asynccontextmanager async def lifespan(app): + logger.info(f"Started A2A server (max contexts: {Config.A2A_MAX_CONTEXTS})") async with app.task_manager: async with worker.run(): yield diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 40fcc8cd..b108ac32 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -62,6 +62,10 @@ class AppConfig(BaseModel): # to allow concurrent connections to safely use recent versions. VACUUM_RETENTION_SECONDS: int = 60 + # Maximum number of A2A contexts to keep in memory. When exceeded, least + # recently used contexts will be evicted. Default is 1000. + A2A_MAX_CONTEXTS: int = 1000 + @field_validator("MONITOR_DIRECTORIES", mode="before") @classmethod def parse_monitor_directories(cls, v): diff --git a/tests/test_a2a.py b/tests/test_a2a.py index 9efdf900..3ae9b404 100644 --- a/tests/test_a2a.py +++ b/tests/test_a2a.py @@ -138,6 +138,76 @@ async def test_extract_question_from_task_no_text(): assert question is None +@pytest.mark.asyncio +async def test_lru_memory_storage_lru_eviction(): + """Test that LRUMemoryStorage evicts least recently used contexts.""" + from fasta2a.storage import InMemoryStorage + + from haiku.rag.a2a import LRUMemoryStorage + + base_storage = InMemoryStorage() + storage = LRUMemoryStorage(storage=base_storage, max_contexts=3) + + # Add 3 contexts (at limit) + await storage.update_context("ctx1", []) + await storage.update_context("ctx2", []) + await storage.update_context("ctx3", []) + + # All 3 should be tracked + assert len(storage.context_order) == 3 + assert "ctx1" in storage.context_order + assert "ctx2" in storage.context_order + assert "ctx3" in storage.context_order + + # Add 4th context - should evict ctx1 (oldest) + await storage.update_context("ctx4", []) + assert len(storage.context_order) == 3 + assert "ctx1" not in storage.context_order + assert "ctx2" in storage.context_order + assert "ctx3" in storage.context_order + assert "ctx4" in storage.context_order + + # Access ctx2 (moves it to end) + await storage.load_context("ctx2") + + # Add 5th context - should evict ctx3 (now oldest since ctx2 was accessed) + await storage.update_context("ctx5", []) + assert len(storage.context_order) == 3 + assert "ctx3" not in storage.context_order + assert "ctx2" in storage.context_order # Still present (was accessed) + assert "ctx4" in storage.context_order + assert "ctx5" in storage.context_order + + +@pytest.mark.asyncio +async def test_lru_memory_storage_access_order(): + """Test that accessing contexts updates their order.""" + from fasta2a.storage import InMemoryStorage + + from haiku.rag.a2a import LRUMemoryStorage + + base_storage = InMemoryStorage() + storage = LRUMemoryStorage(storage=base_storage, max_contexts=2) + + # Add 2 contexts + await storage.update_context("ctx1", []) + await storage.update_context("ctx2", []) + + # Order should be: ctx1, ctx2 + assert list(storage.context_order.keys()) == ["ctx1", "ctx2"] + + # Load ctx1 (moves to end) + await storage.load_context("ctx1") + # Order should be: ctx2, ctx1 + assert list(storage.context_order.keys()) == ["ctx2", "ctx1"] + + # Add ctx3 - should evict ctx2 (oldest) + await storage.update_context("ctx3", []) + assert "ctx2" not in storage.context_order + assert "ctx1" in storage.context_order + assert "ctx3" in storage.context_order + + @pytest.mark.asyncio async def test_a2a_app_creation(temp_db_path): """Test that A2A app can be created successfully."""