feat(api, sdk): finalization of the history deletion feature in the sdk

Signed-off-by: Jan Jeliga <jeliga.jan@gmail.com>

Author: jjeliga

apps/agentstack-sdk-py/src/agentstack_sdk/server/context.py

@@ -3,6 +3,7 @@
 
 
 from collections.abc import AsyncIterator
+from uuid import UUID
 
 import janus
 from a2a.types import Artifact, Message, MessageSendConfiguration, Task
@@ -36,6 +37,11 @@ async def load_history(self) -> AsyncIterator[Message | Artifact]:
         async for item in self._store.load_history():
             yield item
 
+    async def delete_history_from_id(self, from_id: UUID) -> None:
+        if not self._store:
+            raise RuntimeError("Context store is not initialized")
+        await self._store.delete_history_from_id(from_id)
+
     def yield_sync(self, value: RunYield) -> RunYieldResume:
         self._yield_queue.sync_q.put(value)
         return self._yield_resume_queue.sync_q.get()

apps/agentstack-sdk-py/src/agentstack_sdk/server/store/context_store.py

@@ -6,6 +6,7 @@
 import abc
 from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Protocol
+from uuid import UUID
 
 from a2a.types import Artifact, Message
 
@@ -19,6 +20,8 @@ async def load_history(self) -> AsyncIterator[Message | Artifact]:
 
     async def store(self, data: Message | Artifact) -> None: ...
 
+    async def delete_history_from_id(self, from_id: UUID) -> None: ...
+
 
 class ContextStore(abc.ABC):
     def modify_dependencies(self, dependencies: dict[str, Depends]) -> None:

apps/agentstack-sdk-py/src/agentstack_sdk/server/store/memory_context_store.py

@@ -3,6 +3,7 @@
 
 from collections.abc import AsyncIterator
 from datetime import timedelta
+from uuid import UUID
 
 from a2a.types import Artifact, Message
 from cachetools import TTLCache
@@ -23,6 +24,15 @@ async def load_history(self) -> AsyncIterator[Message | Artifact]:
     async def store(self, data: Message | Artifact) -> None:
         self._history.append(data.model_copy(deep=True))
 
+    async def delete_history_from_id(self, from_id: UUID) -> None:
+        # Does not allow to delete from an artifact onwards
+        index = next(
+            (i for i, item in enumerate(self._history) if isinstance(item, Message) and item.message_id == from_id),
+            None,
+        )
+        if index is not None:
+            self._history = self._history[:index]
+
 
 class InMemoryContextStore(ContextStore):
     def __init__(self, max_contexts: int = 1000, context_ttl: timedelta = timedelta(hours=1)):

apps/agentstack-sdk-py/src/agentstack_sdk/server/store/platform_context_store.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections.abc import AsyncIterator
+from uuid import UUID
 
 from a2a.types import Artifact, Message
 
@@ -46,3 +47,7 @@ async def load_history(self) -> AsyncIterator[Message | Artifact]:
     async def store(self, data: Message | Artifact) -> None:
         async with self._platform_extension.use_client():
             await Context.add_history_item(self._context_id, data=data)
+
+    async def delete_history_from_id(self, from_id: UUID) -> None:
+        async with self._platform_extension.use_client():
+            await Context.delete_history_from_id(self._context_id, from_id=from_id)

apps/agentstack-sdk-py/tests/e2e/test_history.py

@@ -28,6 +28,17 @@ async def get_final_task_from_stream(stream: AsyncIterator[ClientEvent | Message
     return final_task
 
 
+async def send_message_get_response(
+    client: Client, content: str, context_id: str | None = None
+) -> tuple[list[str], str]:
+    message = create_text_message_object(content=content)
+    if context_id is not None:
+        message.context_id = context_id
+    final_task = await get_final_task_from_stream(client.send_message(message))
+    agent_messages = [msg.parts[0].root.text for msg in final_task.history or []]
+    return agent_messages, final_task.context_id
+
+
 @pytest.fixture
 async def history_agent(create_server_with_agent) -> AsyncGenerator[tuple[Server, Client]]:
     """Agent that tests context.store.load_history() functionality."""
@@ -44,25 +55,40 @@ async def history_agent(input: Message, context: RunContext) -> AsyncGenerator[R
         yield server, client
 
 
+@pytest.fixture
+async def history_deleting_agent(create_server_with_agent) -> AsyncGenerator[tuple[Server, Client]]:
+    """Agent that tests context.store.load_history() functionality."""
+    context_store = InMemoryContextStore()
+
+    async def history_agent(input: Message, context: RunContext) -> AsyncGenerator[RunYield, None]:
+        await context.store(input)
+        n_messages = 0
+        async for message in context.load_history():
+            n_messages += 1
+            if n_messages == 1:
+                delete_id = message.message_id
+            if n_messages > 3:
+                await context.delete_history_from_id(delete_id)
+
+        async for message in context.load_history():
+            message.role = Role.agent
+            yield message
+
+    async with create_server_with_agent(history_agent, context_store=context_store) as (server, client):
+        yield server, client
+
+
 async def test_agent_history(history_agent):
     """Test that history starts empty."""
     _, client = history_agent
-    message = create_text_message_object(content="first message")
 
-    final_task = await get_final_task_from_stream(client.send_message(message))
-    agent_messages = [msg.parts[0].root.text for msg in final_task.history]
+    agent_messages, context_id = await send_message_get_response(client, "first message")
     assert agent_messages == ["first message"]
 
-    message = create_text_message_object(content="second message")
-    message.context_id = final_task.context_id
-    final_task = await get_final_task_from_stream(client.send_message(message))
-    agent_messages = [msg.parts[0].root.text for msg in final_task.history]
+    agent_messages, context_id = await send_message_get_response(client, "second message", context_id=context_id)
     assert agent_messages == ["first message", "first message", "second message"]
 
-    message = create_text_message_object(content="third message")
-    message.context_id = final_task.context_id
-    final_task = await get_final_task_from_stream(client.send_message(message))
-    agent_messages = [msg.parts[0].root.text for msg in final_task.history]
+    agent_messages, context_id = await send_message_get_response(client, "third message", context_id=context_id)
     assert agent_messages == [
         # first run
         "first message",
@@ -75,3 +101,23 @@ async def test_agent_history(history_agent):
         "second message",
         "third message",
     ]
+
+
+async def test_agent_deleting_history(history_deleting_agent):
+    """Test that history starts empty."""
+    _, client = history_deleting_agent
+
+    agent_messages, context_id = await send_message_get_response(client, "first message")
+    assert agent_messages == ["first message"]
+
+    agent_messages, context_id = await send_message_get_response(client, "second message", context_id=context_id)
+    assert agent_messages == ["first message", "second message"]
+
+    agent_messages, context_id = await send_message_get_response(client, "third message", context_id=context_id)
+    assert agent_messages == ["first message", "second message", "third message"]
+
+    agent_messages, context_id = await send_message_get_response(client, "delete message", context_id=context_id)
+    assert agent_messages == []
+
+    agent_messages, context_id = await send_message_get_response(client, "first message")
+    assert agent_messages == ["first message"]

apps/agentstack-server/src/agentstack_server/infrastructure/persistence/repositories/context.py

@@ -214,7 +214,7 @@ async def list_history(
         )
 
     async def delete_history_from_id(self, *, context_id: UUID, from_id: UUID) -> int:
-        """Delete all history items from a specific item onwards (inclusive) for a context"""
+        """Delete all history items from a specific item onwards (inclusive) in given context"""
         # First, get the created_at timestamp of the item to delete from
         query_item = select(context_history_table.c.created_at).where(
             context_history_table.c.id == from_id,

docs/development/agent-integration/multi-turn.mdx

@@ -300,14 +300,15 @@ async def edit_message_in_context(run_context: RunContext, message_id: str, new_
     # Step 1: Delete from this message onwards
     await context.delete_history_from_id(from_id=message_id)
 
-    # Step 2: Add the corrected message
+    # Step 2: Create the corrected message
     corrected_message = Message(
         message_id=str(uuid.uuid4()),
-        parts=[Part(root=TextPart(text=new_text))],
+        parts=[Part(TextPart(text=new_text))],
         role=Role.user,
-        context_id=context.id,
         metadata=metadata,
     )
+
+    # Step 3: Store the corrected message
     await run_context.store(data=corrected_message)
 ```