diff --git a/.changeset/cold-queens-train.md b/.changeset/cold-queens-train.md new file mode 100644 index 00000000..77ddf0d1 --- /dev/null +++ b/.changeset/cold-queens-train.md @@ -0,0 +1,5 @@ +--- +"llama-agents-server": minor +--- + +Add max_completed history cap to MemoryWorkflowStore in order to control memory consumption diff --git a/packages/llama-agents-server/src/llama_agents/server/_store/memory_workflow_store.py b/packages/llama-agents-server/src/llama_agents/server/_store/memory_workflow_store.py index de8d5af3..5c66cb74 100644 --- a/packages/llama-agents-server/src/llama_agents/server/_store/memory_workflow_store.py +++ b/packages/llama-agents-server/src/llama_agents/server/_store/memory_workflow_store.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import logging import weakref +from collections import deque from collections.abc import AsyncIterator from datetime import datetime, timezone from typing import Any, Dict, List @@ -17,6 +18,7 @@ from .abstract_workflow_store import ( PersistentHandler, StoredEvent, StoredTick, + is_terminal_status, ) logger = logging.getLogger(__name__) @@ -57,7 +59,10 @@ def _matches_query(handler: PersistentHandler, query: HandlerQuery) -> bool: class MemoryWorkflowStore(AbstractWorkflowStore): - def __init__(self) -> None: + def __init__(self, max_completed: int | None = 1000) -> None: + if max_completed is not None and max_completed < 0: + raise ValueError("max_completed must be >= 0 or None") + self.handlers: Dict[str, PersistentHandler] = {} self.events: Dict[str, List[StoredEvent]] = {} self.ticks: Dict[str, List[StoredTick]] = {} @@ -65,6 +70,8 @@ class MemoryWorkflowStore(AbstractWorkflowStore): self._conditions: weakref.WeakValueDictionary[str, asyncio.Condition] = ( weakref.WeakValueDictionary() ) + self.max_completed = max_completed + self._terminal_queue: deque[str] = deque() def create_state_store( self, @@ -99,6 +106,9 @@ class MemoryWorkflowStore(AbstractWorkflowStore): async def update(self, handler: PersistentHandler) -> None: self.handlers[handler.handler_id] = handler + if is_terminal_status(handler.status): + self._terminal_queue.append(handler.handler_id) + self._evict_oldest_completed() async def delete(self, query: HandlerQuery) -> int: to_delete = [ @@ -110,6 +120,33 @@ class MemoryWorkflowStore(AbstractWorkflowStore): del self.handlers[handler_id] return len(to_delete) + def _evict_oldest_completed(self) -> None: + """Remove the oldest completed handlers when the cap is exceeded. + + Uses _terminal_queue (insertion-ordered deque) for O(1) eviction + instead of scanning and sorting all handlers. + """ + if self.max_completed is None: + return + + while len(self._terminal_queue) > self.max_completed: + handler_id = self._terminal_queue.popleft() + handler = self.handlers.get(handler_id) + if handler is None: + # Already removed (e.g. via delete()), skip. + continue + if not is_terminal_status(handler.status): + # Stale terminal-queue entry for a handler_id that was upserted + # into a newer non-terminal row. + continue + + self.handlers.pop(handler_id, None) + run_id = handler.run_id + if run_id is not None: + self.events.pop(run_id, None) + self.ticks.pop(run_id, None) + self.state_stores.pop(run_id, None) + def _get_or_create_condition(self, run_id: str) -> asyncio.Condition: """Get or create a condition for a run_id. diff --git a/packages/llama-agents-server/tests/server/test_memory_workflow_store.py b/packages/llama-agents-server/tests/server/test_memory_workflow_store.py index 47febbcd..08a7f1fd 100644 --- a/packages/llama-agents-server/tests/server/test_memory_workflow_store.py +++ b/packages/llama-agents-server/tests/server/test_memory_workflow_store.py @@ -4,6 +4,7 @@ from __future__ import annotations from datetime import datetime, timezone +from typing import Any import pytest from llama_agents.client.protocol.serializable_events import EventEnvelopeWithMetadata @@ -21,24 +22,69 @@ from workflows.events import ( WorkflowFailedEvent, ) +T0 = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc) -@pytest.mark.asyncio -async def test_update_and_query_returns_inserted_handler() -> None: - store = MemoryWorkflowStore() - handler = PersistentHandler( - handler_id="h1", - workflow_name="wf_a", - status="running", +def _ts(seconds: int) -> datetime: + return datetime(2024, 1, 1, 0, 0, seconds, tzinfo=timezone.utc) + + +def _handler( + handler_id: str = "h1", + workflow_name: str = "wf", + status: Status = "running", + **kwargs: Any, +) -> PersistentHandler: + return PersistentHandler( + handler_id=handler_id, + workflow_name=workflow_name, + status=status, + **kwargs, ) - await store.update(handler) - # Filter by workflow_name list +async def _insert( + store: MemoryWorkflowStore, + handler_id: str = "h1", + workflow_name: str = "wf", + status: Status = "running", + **kwargs: Any, +) -> PersistentHandler: + h = _handler( + handler_id=handler_id, workflow_name=workflow_name, status=status, **kwargs + ) + await store.update(h) + return h + + +async def _query_ids(store: MemoryWorkflowStore, **kwargs: Any) -> set[str]: + result = await store.query(HandlerQuery(**kwargs)) + return {h.handler_id for h in result} + + +def _make_stored_event(event: Event, run_id: str = "run-1") -> StoredEvent: + return StoredEvent( + run_id=run_id, + sequence=0, + timestamp=datetime.now(timezone.utc), + event=EventEnvelopeWithMetadata.from_event(event), + ) + + +@pytest.fixture +def store() -> MemoryWorkflowStore: + return MemoryWorkflowStore() + + +@pytest.mark.asyncio +async def test_update_and_query_returns_inserted_handler( + store: MemoryWorkflowStore, +) -> None: + await _insert(store, handler_id="h1", workflow_name="wf_a") + result = await store.query( HandlerQuery(workflow_name_in=["wf_a"], status_in=["running"]) ) - assert len(result) == 1 found = result[0] assert found.handler_id == "h1" @@ -47,34 +93,17 @@ async def test_update_and_query_returns_inserted_handler() -> None: @pytest.mark.asyncio -async def test_update_on_conflict_overwrites_existing_row() -> None: - store = MemoryWorkflowStore() +async def test_update_on_conflict_overwrites_existing_row( + store: MemoryWorkflowStore, +) -> None: + await _insert(store, handler_id="h2", workflow_name="wf_b") + await _insert(store, handler_id="h2", workflow_name="wf_b", status="completed") - # Initial insert (in-progress) - await store.update( - PersistentHandler( - handler_id="h2", - workflow_name="wf_b", - status="running", - ) + assert ( + await _query_ids(store, workflow_name_in=["wf_b"], status_in=["running"]) + == set() ) - # Update same handler_id (completed) - await store.update( - PersistentHandler( - handler_id="h2", - workflow_name="wf_b", - status="completed", - ) - ) - - # Should not be returned for status=running - result_in_progress = await store.query( - HandlerQuery(workflow_name_in=["wf_b"], status_in=["running"]) - ) - assert result_in_progress == [] - - # Should be returned for status=completed with latest values result_completed = await store.query( HandlerQuery(workflow_name_in=["wf_b"], status_in=["completed"]) ) @@ -86,281 +115,130 @@ async def test_update_on_conflict_overwrites_existing_row() -> None: @pytest.mark.asyncio -async def test_delete_filters_by_query() -> None: - store = MemoryWorkflowStore() - - await store.update( - PersistentHandler( - handler_id="delete-me", - workflow_name="wf_delete", - status="completed", - ) - ) - await store.update( - PersistentHandler( - handler_id="keep-me", - workflow_name="wf_keep", - status="running", - ) +async def test_delete_filters_by_query(store: MemoryWorkflowStore) -> None: + await _insert( + store, handler_id="delete-me", workflow_name="wf_delete", status="completed" ) + await _insert(store, handler_id="keep-me", workflow_name="wf_keep") deleted = await store.delete(HandlerQuery(handler_id_in=["delete-me"])) - assert deleted == 1 - remaining = await store.query(HandlerQuery()) - ids = {handler.handler_id for handler in remaining} - assert ids == {"keep-me"} + assert await _query_ids(store) == {"keep-me"} @pytest.mark.asyncio -async def test_delete_noop_on_empty_filter() -> None: - store = MemoryWorkflowStore() - - await store.update( - PersistentHandler( - handler_id="delete-me", - workflow_name="wf_delete", - status="completed", - ) +async def test_delete_noop_on_empty_filter(store: MemoryWorkflowStore) -> None: + await _insert( + store, handler_id="delete-me", workflow_name="wf_delete", status="completed" ) deleted = await store.delete(HandlerQuery(handler_id_in=[])) - assert deleted == 0 - remaining = await store.query(HandlerQuery()) - assert len(remaining) == 1 - assert remaining[0].handler_id == "delete-me" + assert await _query_ids(store) == {"delete-me"} @pytest.mark.asyncio -async def test_query_filters_by_handler_id_and_empty_lists() -> None: - store = MemoryWorkflowStore() - - # Seed three handlers +async def test_query_filters_by_handler_id_and_empty_lists( + store: MemoryWorkflowStore, +) -> None: for hid, wf in [("h1", "wf_a"), ("h2", "wf_a"), ("h3", "wf_b")]: - await store.update( - PersistentHandler( - handler_id=hid, - workflow_name=wf, - status="running", - ) - ) + await _insert(store, handler_id=hid, workflow_name=wf) - # Filter by specific handler ids - result = await store.query(HandlerQuery(handler_id_in=["h1", "h3"])) - ids = {h.handler_id for h in result} - assert ids == {"h1", "h3"} - - # Empty handler_id list short-circuits to [] - result_empty_ids = await store.query(HandlerQuery(handler_id_in=[])) - assert result_empty_ids == [] - - # Empty workflow_name list short-circuits to [] - result_empty_wf = await store.query(HandlerQuery(workflow_name_in=[])) - assert result_empty_wf == [] - - # No filters returns all - all_rows = await store.query(HandlerQuery()) - assert {h.handler_id for h in all_rows} == {"h1", "h2", "h3"} + assert await _query_ids(store, handler_id_in=["h1", "h3"]) == {"h1", "h3"} + assert await store.query(HandlerQuery(handler_id_in=[])) == [] + assert await store.query(HandlerQuery(workflow_name_in=[])) == [] + assert await _query_ids(store) == {"h1", "h2", "h3"} @pytest.mark.asyncio -async def test_query_filters_by_multiple_statuses() -> None: - store = MemoryWorkflowStore() +async def test_query_filters_by_multiple_statuses(store: MemoryWorkflowStore) -> None: + statuses: list[tuple[str, Status]] = [ + ("h1", "running"), + ("h2", "completed"), + ("h3", "failed"), + ("h4", "cancelled"), + ] + for hid, status in statuses: + await _insert(store, handler_id=hid, status=status) - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - ) - ) - await store.update( - PersistentHandler( - handler_id="h2", - workflow_name="wf", - status="completed", - ) - ) - await store.update( - PersistentHandler( - handler_id="h3", - workflow_name="wf", - status="failed", - ) - ) - await store.update( - PersistentHandler( - handler_id="h4", - workflow_name="wf", - status="cancelled", - ) - ) - - # Query for multiple statuses - result = await store.query(HandlerQuery(status_in=["completed", "failed"])) - ids = {h.handler_id for h in result} - assert ids == {"h2", "h3"} - - # Empty status list returns nothing - result_empty = await store.query(HandlerQuery(status_in=[])) - assert result_empty == [] + assert await _query_ids(store, status_in=["completed", "failed"]) == {"h2", "h3"} + assert await store.query(HandlerQuery(status_in=[])) == [] @pytest.mark.asyncio -async def test_query_filters_by_workflow_name() -> None: - store = MemoryWorkflowStore() +async def test_query_filters_by_workflow_name(store: MemoryWorkflowStore) -> None: + await _insert(store, handler_id="h1", workflow_name="wf_a") + await _insert(store, handler_id="h2", workflow_name="wf_b") + await _insert(store, handler_id="h3", workflow_name="wf_a", status="completed") - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf_a", - status="running", - ) - ) - await store.update( - PersistentHandler( - handler_id="h2", - workflow_name="wf_b", - status="running", - ) - ) - await store.update( - PersistentHandler( - handler_id="h3", - workflow_name="wf_a", - status="completed", - ) - ) - - # Query for specific workflow - result = await store.query(HandlerQuery(workflow_name_in=["wf_a"])) - ids = {h.handler_id for h in result} - assert ids == {"h1", "h3"} - - # Query for multiple workflows - result_multi = await store.query(HandlerQuery(workflow_name_in=["wf_a", "wf_b"])) - ids_multi = {h.handler_id for h in result_multi} - assert ids_multi == {"h1", "h2", "h3"} + assert await _query_ids(store, workflow_name_in=["wf_a"]) == {"h1", "h3"} + assert await _query_ids(store, workflow_name_in=["wf_a", "wf_b"]) == { + "h1", + "h2", + "h3", + } @pytest.mark.asyncio -async def test_query_combines_multiple_filters() -> None: - store = MemoryWorkflowStore() +async def test_query_combines_multiple_filters(store: MemoryWorkflowStore) -> None: + await _insert(store, handler_id="h1", workflow_name="wf_a") + await _insert(store, handler_id="h2", workflow_name="wf_a", status="completed") + await _insert(store, handler_id="h3", workflow_name="wf_b") + await _insert(store, handler_id="h4", workflow_name="wf_b", status="completed") - # Seed multiple handlers with different combinations - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf_a", - status="running", - ) - ) - await store.update( - PersistentHandler( - handler_id="h2", - workflow_name="wf_a", - status="completed", - ) - ) - await store.update( - PersistentHandler( - handler_id="h3", - workflow_name="wf_b", - status="running", - ) - ) - await store.update( - PersistentHandler( - handler_id="h4", - workflow_name="wf_b", - status="completed", - ) - ) - - # Combine workflow and status filters - result = await store.query( - HandlerQuery(workflow_name_in=["wf_a"], status_in=["running"]) - ) - ids = {h.handler_id for h in result} - assert ids == {"h1"} - - # Combine all three filters - result_triple = await store.query( - HandlerQuery( - handler_id_in=["h2", "h4"], - workflow_name_in=["wf_a"], - status_in=["completed"], - ) - ) - ids_triple = {h.handler_id for h in result_triple} - assert ids_triple == {"h2"} + assert await _query_ids( + store, workflow_name_in=["wf_a"], status_in=["running"] + ) == {"h1"} + assert await _query_ids( + store, + handler_id_in=["h2", "h4"], + workflow_name_in=["wf_a"], + status_in=["completed"], + ) == {"h2"} @pytest.mark.asyncio -async def test_delete_removes_multiple_matching_handlers() -> None: - store = MemoryWorkflowStore() - - # Seed multiple handlers +async def test_delete_removes_multiple_matching_handlers( + store: MemoryWorkflowStore, +) -> None: for i in range(5): - await store.update( - PersistentHandler( - handler_id=f"h{i}", - workflow_name="wf", - status="completed" if i % 2 == 0 else "running", - ) + await _insert( + store, handler_id=f"h{i}", status="completed" if i % 2 == 0 else "running" ) - # Delete all completed handlers deleted = await store.delete(HandlerQuery(status_in=["completed"])) - assert deleted == 3 # h0, h2, h4 - - remaining = await store.query(HandlerQuery()) - ids = {h.handler_id for h in remaining} - assert ids == {"h1", "h3"} + assert deleted == 3 + assert await _query_ids(store) == {"h1", "h3"} @pytest.mark.asyncio -async def test_store_handles_all_datetime_fields() -> None: - store = MemoryWorkflowStore() - +async def test_store_handles_all_datetime_fields(store: MemoryWorkflowStore) -> None: now = datetime.now(timezone.utc) - handler = PersistentHandler( - handler_id="h1", - workflow_name="wf", + stop = StopEvent(result={"output": "success"}) + await _insert( + store, status="completed", run_id="run123", error=None, - result=StopEvent(result={"output": "success"}), + result=stop, started_at=now, updated_at=now, completed_at=now, ) - await store.update(handler) - result = await store.query(HandlerQuery(handler_id_in=["h1"])) assert len(result) == 1 found = result[0] assert found.run_id == "run123" - assert found.result == StopEvent(result={"output": "success"}) + assert found.result == stop assert found.started_at == now assert found.updated_at == now assert found.completed_at == now @pytest.mark.asyncio -async def test_store_handles_error_field() -> None: - store = MemoryWorkflowStore() - - handler = PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="failed", - error="Something went wrong", - ) - - await store.update(handler) +async def test_store_handles_error_field(store: MemoryWorkflowStore) -> None: + await _insert(store, status="failed", error="Something went wrong") result = await store.query(HandlerQuery(handler_id_in=["h1"])) assert len(result) == 1 @@ -368,104 +246,64 @@ async def test_store_handles_error_field() -> None: @pytest.mark.asyncio -async def test_empty_store_returns_empty_results() -> None: - store = MemoryWorkflowStore() - - # Query empty store - result = await store.query(HandlerQuery()) - assert result == [] - - # Delete from empty store - deleted = await store.delete(HandlerQuery(handler_id_in=["nonexistent"])) - assert deleted == 0 +async def test_empty_store_returns_empty_results(store: MemoryWorkflowStore) -> None: + assert await store.query(HandlerQuery()) == [] + assert await store.delete(HandlerQuery(handler_id_in=["nonexistent"])) == 0 @pytest.mark.asyncio -async def test_update_handler_status_with_nonexistent_run_id() -> None: - store = MemoryWorkflowStore() - # Should not raise when run_id does not exist +async def test_update_handler_status_with_nonexistent_run_id( + store: MemoryWorkflowStore, +) -> None: await store.update_handler_status("nonexistent-run-id", status="completed") @pytest.mark.asyncio -async def test_update_handler_status_sets_status_and_completed_at() -> None: - store = MemoryWorkflowStore() - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - run_id="run-1", - ) - ) +async def test_update_handler_status_sets_status_and_completed_at( + store: MemoryWorkflowStore, +) -> None: + await _insert(store, run_id="run-1") await store.update_handler_status("run-1", status="completed") result = await store.query(HandlerQuery(run_id_in=["run-1"])) assert len(result) == 1 - handler = result[0] - assert handler.status == "completed" - assert handler.updated_at is not None - assert handler.completed_at is not None + assert result[0].status == "completed" + assert result[0].updated_at is not None + assert result[0].completed_at is not None @pytest.mark.asyncio -async def test_update_handler_status_with_result() -> None: - store = MemoryWorkflowStore() - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - run_id="run-1", - ) - ) +async def test_update_handler_status_with_result(store: MemoryWorkflowStore) -> None: + await _insert(store, run_id="run-1") stop = StopEvent(result={"answer": 42}) await store.update_handler_status("run-1", status="completed", result=stop) result = await store.query(HandlerQuery(run_id_in=["run-1"])) - handler = result[0] - assert handler.status == "completed" - assert handler.result == stop + assert result[0].status == "completed" + assert result[0].result == stop @pytest.mark.asyncio -async def test_update_handler_status_with_error() -> None: - store = MemoryWorkflowStore() - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - run_id="run-1", - ) - ) +async def test_update_handler_status_with_error(store: MemoryWorkflowStore) -> None: + await _insert(store, run_id="run-1") await store.update_handler_status("run-1", status="failed", error="boom") result = await store.query(HandlerQuery(run_id_in=["run-1"])) - handler = result[0] - assert handler.status == "failed" - assert handler.error == "boom" - assert handler.completed_at is not None + assert result[0].status == "failed" + assert result[0].error == "boom" + assert result[0].completed_at is not None @pytest.mark.asyncio -async def test_update_handler_status_idle_since_explicit_none_clears() -> None: - store = MemoryWorkflowStore() +async def test_update_handler_status_idle_since_explicit_none_clears( + store: MemoryWorkflowStore, +) -> None: now = datetime.now(timezone.utc) - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - run_id="run-1", - idle_since=now, - ) - ) + await _insert(store, run_id="run-1", idle_since=now) - # Passing idle_since=None explicitly should clear it await store.update_handler_status("run-1", idle_since=None) result = await store.query(HandlerQuery(run_id_in=["run-1"])) @@ -473,20 +311,12 @@ async def test_update_handler_status_idle_since_explicit_none_clears() -> None: @pytest.mark.asyncio -async def test_update_handler_status_idle_since_unset_preserves() -> None: - store = MemoryWorkflowStore() +async def test_update_handler_status_idle_since_unset_preserves( + store: MemoryWorkflowStore, +) -> None: now = datetime.now(timezone.utc) - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - run_id="run-1", - idle_since=now, - ) - ) + await _insert(store, run_id="run-1", idle_since=now) - # Not passing idle_since at all should preserve the existing value await store.update_handler_status("run-1", status="running") result = await store.query(HandlerQuery(run_id_in=["run-1"])) @@ -494,18 +324,11 @@ async def test_update_handler_status_idle_since_unset_preserves() -> None: @pytest.mark.asyncio -async def test_update_handler_status_non_terminal_does_not_set_completed_at() -> None: - store = MemoryWorkflowStore() - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - run_id="run-1", - ) - ) +async def test_update_handler_status_non_terminal_does_not_set_completed_at( + store: MemoryWorkflowStore, +) -> None: + await _insert(store, run_id="run-1") - # Update status to "running" (non-terminal) should not set completed_at await store.update_handler_status("run-1", status="running") result = await store.query(HandlerQuery(run_id_in=["run-1"])) @@ -515,17 +338,10 @@ async def test_update_handler_status_non_terminal_does_not_set_completed_at() -> @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["completed", "failed", "cancelled"]) async def test_update_handler_status_terminal_sets_completed_at( + store: MemoryWorkflowStore, terminal_status: Status, ) -> None: - store = MemoryWorkflowStore() - await store.update( - PersistentHandler( - handler_id="h1", - workflow_name="wf", - status="running", - run_id="run-1", - ) - ) + await _insert(store, run_id="run-1") await store.update_handler_status("run-1", status=terminal_status) @@ -533,15 +349,6 @@ async def test_update_handler_status_terminal_sets_completed_at( assert result[0].completed_at is not None -def _make_stored_event(event: Event, run_id: str = "run-1") -> StoredEvent: - return StoredEvent( - run_id=run_id, - sequence=0, - timestamp=datetime.now(timezone.utc), - event=EventEnvelopeWithMetadata.from_event(event), - ) - - def test_is_terminal_event_stop_event() -> None: stored = _make_stored_event(StopEvent(result="done")) assert AbstractWorkflowStore._is_terminal_event(stored) is True @@ -553,7 +360,6 @@ def test_is_terminal_event_regular_event() -> None: def test_is_terminal_event_workflow_failed_event() -> None: - # WorkflowFailedEvent extends StopEvent, so it should be terminal event = WorkflowFailedEvent( step_name="my_step", exception_type="ValueError", @@ -567,6 +373,172 @@ def test_is_terminal_event_workflow_failed_event() -> None: def test_is_terminal_event_workflow_cancelled_event() -> None: - # WorkflowCancelledEvent extends StopEvent, so it should be terminal stored = _make_stored_event(WorkflowCancelledEvent()) assert AbstractWorkflowStore._is_terminal_event(stored) is True + + +# --- max_completed history cap tests --- + + +@pytest.mark.asyncio +async def test_max_completed_default_is_1000() -> None: + assert MemoryWorkflowStore().max_completed == 1000 + + +@pytest.mark.asyncio +async def test_max_completed_none_means_unlimited() -> None: + store = MemoryWorkflowStore(max_completed=None) + assert store.max_completed is None + + for i in range(50): + await _insert( + store, + handler_id=f"h{i}", + status="completed", + run_id=f"run-{i}", + completed_at=_ts(i), + ) + + assert len(await store.query(HandlerQuery())) == 50 + + +def test_max_completed_negative_raises_value_error() -> None: + with pytest.raises(ValueError, match="max_completed must be >= 0 or None"): + MemoryWorkflowStore(max_completed=-1) + + +@pytest.mark.asyncio +async def test_max_completed_evicts_oldest_when_exceeded() -> None: + store = MemoryWorkflowStore(max_completed=3) + + for i in range(5): + await _insert( + store, + handler_id=f"h{i}", + status="completed", + run_id=f"run-{i}", + completed_at=_ts(i), + ) + + assert await _query_ids(store) == {"h2", "h3", "h4"} + + +@pytest.mark.asyncio +async def test_max_completed_does_not_evict_running_handlers() -> None: + store = MemoryWorkflowStore(max_completed=2) + + for i in range(3): + await _insert(store, handler_id=f"running-{i}", run_id=f"run-r{i}") + + for i in range(3): + await _insert( + store, + handler_id=f"done-{i}", + status="completed", + run_id=f"run-d{i}", + completed_at=_ts(i), + ) + + assert await _query_ids(store) == { + "running-0", + "running-1", + "running-2", + "done-1", + "done-2", + } + + +@pytest.mark.asyncio +async def test_max_completed_applies_to_all_terminal_statuses() -> None: + store = MemoryWorkflowStore(max_completed=2) + + await _insert( + store, handler_id="h-completed", status="completed", completed_at=_ts(0) + ) + await _insert(store, handler_id="h-failed", status="failed", completed_at=_ts(1)) + await _insert( + store, handler_id="h-cancelled", status="cancelled", completed_at=_ts(2) + ) + + assert await _query_ids(store) == {"h-failed", "h-cancelled"} + + +@pytest.mark.asyncio +async def test_max_completed_cleans_up_events_ticks_and_state() -> None: + store = MemoryWorkflowStore(max_completed=1) + + store.events["run-old"] = [] + store.ticks["run-old"] = [] + store.create_state_store("run-old") + + await _insert( + store, + handler_id="h-old", + status="completed", + run_id="run-old", + completed_at=_ts(0), + ) + assert "run-old" in store.events + assert "run-old" in store.ticks + assert "run-old" in store.state_stores + + await _insert( + store, + handler_id="h-new", + status="completed", + run_id="run-new", + completed_at=_ts(1), + ) + + assert "run-old" not in store.events + assert "run-old" not in store.ticks + assert "run-old" not in store.state_stores + remaining = await store.query(HandlerQuery()) + assert len(remaining) == 1 + assert remaining[0].handler_id == "h-new" + + +@pytest.mark.asyncio +async def test_max_completed_eviction_via_update_handler_status() -> None: + """Eviction triggers when status changes to terminal via update_handler_status.""" + store = MemoryWorkflowStore(max_completed=2) + + for i in range(3): + await _insert(store, handler_id=f"h{i}", run_id=f"run-{i}") + + await store.update_handler_status("run-0", status="completed") + await store.update_handler_status("run-1", status="completed") + await store.update_handler_status("run-2", status="completed") + + ids = await _query_ids(store) + assert len(ids) == 2 + assert "h0" not in ids + assert "h1" in ids + assert "h2" in ids + + +@pytest.mark.asyncio +async def test_max_completed_ignores_stale_terminal_queue_entries() -> None: + store = MemoryWorkflowStore(max_completed=1) + + await _insert( + store, + handler_id="shared", + status="completed", + run_id="run-old", + completed_at=_ts(0), + ) + await _insert(store, handler_id="shared", status="running", run_id="run-active") + + await _insert( + store, + handler_id="done-2", + status="completed", + run_id="run-2", + completed_at=_ts(1), + ) + + handlers = await store.query(HandlerQuery()) + by_id = {handler.handler_id: handler for handler in handlers} + assert set(by_id) == {"shared", "done-2"} + assert by_id["shared"].status == "running"