feat: replace custom ContextStore with A2A-native task history

agents/canvas/src/canvas/agent.py

@@ -50,7 +50,6 @@ async def canvas_agent(
         yield "Can't run without a LLM."
         return
 
-    await context.store(message)
     edit_request = await canvas.parse_canvas_edit_request(message=message)
 
     user_text_content = _get_text(message)
@@ -132,12 +131,6 @@ async def canvas_agent(
                 parts=[TextPart(text=content_delta)],
             )
 
-    final_artifact = AgentArtifact(
-        artifact_id=artifact.artifact_id,
-        name=artifact.name,
-        parts=[TextPart(text=buffer)],
-    )
-    await context.store(final_artifact)
 
 
 def serve():

agents/chat/src/chat/agent.py

@@ -24,19 +24,18 @@
     PlatformApiExtensionServer,
     PlatformApiExtensionSpec,
 )
-from kagenti_adk.a2a.types import AgentArtifact, AgentMessage
+from kagenti_adk.a2a.types import AgentArtifact
 from kagenti_adk.server import Server
 from kagenti_adk.server.context import RunContext
 from kagenti_adk.server.middleware.platform_auth_backend import PlatformAuthBackend
-from kagenti_adk.server.store.platform_context_store import PlatformContextStore
 from beeai_framework.adapters.agentstack.backend.chat import AgentStackChatModel
 from beeai_framework.agents.requirement import RequirementAgent
 from beeai_framework.agents.requirement.events import (
     RequirementAgentFinalAnswerEvent,
     RequirementAgentSuccessEvent,
 )
 from beeai_framework.agents.requirement.utils._tool import FinalAnswerTool
-from beeai_framework.backend import AssistantMessage, ChatModelParameters
+from beeai_framework.backend import ChatModelParameters
 from beeai_framework.errors import FrameworkError
 from beeai_framework.middleware.trajectory import GlobalTrajectoryMiddleware
 from beeai_framework.tools import AnyTool, Tool
@@ -45,7 +44,6 @@
 from beeai_framework.tools.weather import OpenMeteoTool
 from openinference.instrumentation.beeai import BeeAIInstrumentor
 
-from chat.helpers.citations import extract_citations
 from chat.helpers.trajectory import TrajectoryContent
 from chat.tools.files.file_creator import FileCreatorTool, FileCreatorToolOutput
 from chat.tools.files.file_reader import FileReaderTool
@@ -149,8 +147,6 @@ async def chat(
     _p: Annotated[PlatformApiExtensionServer, PlatformApiExtensionSpec()],
 ):
     """Agent with memory and access to web search, Wikipedia, and weather."""
-    await context.store(input)
-
     # Send initial trajectory
     yield trajectory.trajectory_metadata(title="Starting", content="Received your request")
 
@@ -220,7 +216,6 @@ async def chat(
         middlewares=[GlobalTrajectoryMiddleware(included=[Tool])],
     )
 
-    final_answer: AssistantMessage | None = None
     new_messages = [to_framework_message(item, extracted_files) for item in history]
 
     try:
@@ -244,8 +239,6 @@ async def chat(
                 case RequirementAgentFinalAnswerEvent(delta=delta):
                     yield delta
                 case RequirementAgentSuccessEvent(state=state):
-                    final_answer = state.answer
-
                     last_step = state.steps[-1]
                     if last_step.tool and last_step.tool.name == FinalAnswerTool.name:  # internal tool
                         continue
@@ -259,24 +252,14 @@ async def chat(
                         group_id=last_step.id,
                     )
                     yield metadata
-                    await context.store(AgentMessage(metadata=metadata))
 
                     if isinstance(last_step.output, FileCreatorToolOutput):
                         for file_info in last_step.output.result.files:
                             part = file_info.file.to_part()
                             part.filename = file_info.display_filename
                             artifact = AgentArtifact(name=file_info.display_filename, parts=[part])
                             yield artifact
-                            await context.store(artifact)
-
-        if final_answer:
-            citations, clean_text = extract_citations(final_answer.text)
 
-            message = AgentMessage(
-                text=clean_text,
-                metadata=(citation.citation_metadata(citations=citations) if citations else None),
-            )
-            await context.store(message)
     except FrameworkError as err:
         raise RuntimeError(err.explain())
 
@@ -287,7 +270,6 @@ def serve():
             host=os.getenv("HOST", "127.0.0.1"),
             port=int(os.getenv("PORT", 8000)),
             configure_telemetry=True,
-            context_store=PlatformContextStore(),
             auth_backend=PlatformAuthBackend(),
         )
     except KeyboardInterrupt:

agents/deepagents_content_builder/src/content_builder/agent.py

@@ -91,7 +91,6 @@ async def content_builder_agent(
         return
 
     started_at = datetime.now(timezone.utc)
-    await context.store(data=message)
 
     subagents: list[SubAgent] = []
     for sub_agent in AVAILABLE_SUBAGENTS:
@@ -140,7 +139,6 @@ async def content_builder_agent(
                             title=data["name"], content=json.dumps(obj=data["args"])
                         )
                         yield tool_call_metadata
-                        await context.store(data=AgentMessage(metadata=tool_call_metadata))
                     tool_calls.clear()
 
                 elif last_msg.tool_call_chunks:
@@ -151,12 +149,10 @@ async def content_builder_agent(
                             tool_calls[tc_id]["args"] += tc.get("args") or ""
                 elif last_msg.text:
                     yield AgentMessage(text=last_msg.text)
-                    await context.store(AgentMessage(text=last_msg.text))
 
             elif isinstance(last_msg, ToolMessage) and last_msg.name and last_msg.text:
                 tool_message_metadata = trajectory.trajectory_metadata(title=last_msg.name, content=last_msg.text)
                 yield tool_message_metadata
-                await context.store(data=AgentMessage(metadata=tool_message_metadata))
 
     updated_files = await agent_stack_backend.alist(order_by="created_at", order="asc", created_after=started_at)
     for updated_file in updated_files:

agents/rag/src/rag/agent.py

@@ -28,7 +28,6 @@
 from kagenti_adk.server import Server
 from kagenti_adk.server.context import RunContext
 from kagenti_adk.server.middleware.platform_auth_backend import PlatformAuthBackend
-from kagenti_adk.server.store.platform_context_store import PlatformContextStore
 from beeai_framework.adapters.agentstack.backend.chat import AgentStackChatModel
 from beeai_framework.agents.requirement import RequirementAgent
 from beeai_framework.agents.requirement.utils._tool import FinalAnswerTool
@@ -116,7 +115,6 @@ async def rag(
     _: Annotated[PlatformApiExtensionServer, PlatformApiExtensionSpec()],
 ):
     """RAG agent that retrieves and generates text based on user queries."""
-    await context.store(input)
     llm, embedding = _get_clients(llm_ext, embedding_ext)
 
     history = [m async for m in context.load_history()]
@@ -181,7 +179,6 @@ async def rag(
                 phase="end",
             ).metadata(trajectory)
             yield vector_store_create_metadata
-            await context.store(AgentMessage(metadata=vector_store_create_metadata))
 
         tools.append(cast(Tool, VectorSearchTool(vector_store_id=vector_store_id, embedding_function=embedding)))
         async for item in embed_all_files(
@@ -300,7 +297,6 @@ async def handle_tool_success(event, meta):
             metadata=(citation.citation_metadata(citations=citations) if citations else None),
         )
         yield message
-        await context.store(message)
 
 
 def _get_clients(
@@ -331,7 +327,6 @@ def serve():
             host=os.getenv("HOST", "127.0.0.1"),
             port=int(os.getenv("PORT", 8000)),
             configure_telemetry=True,
-            context_store=PlatformContextStore(),
             auth_backend=PlatformAuthBackend(),
         )
     except KeyboardInterrupt:

apps/adk-py/examples/canvas_ui_code_agent.py

@@ -84,8 +84,6 @@ async def artifacts_agent(
 ):
     """Works with artifacts"""
 
-    await context.store(input)
-
     canvas_edit_request = await canvas.parse_canvas_edit_request(message=input)
 
     if canvas_edit_request:
@@ -126,7 +124,6 @@ async def artifacts_agent(
     if pre_text := response[: match.start()]:
         message = AgentMessage(text=pre_text)
         yield message
-        await context.store(message)
 
     await asyncio.sleep(1)
 
@@ -153,7 +150,6 @@ async def artifacts_agent(
         name=artifact_name,
         parts=[TextPart(text=code_content)],
     )
-    await context.store(artifact)
 
     # Send first chunk with artifact_id to establish the artifact
     first_artifact = AgentArtifact(
@@ -171,13 +167,11 @@ async def artifacts_agent(
             parts=[TextPart(text=chunk)],
         )
         yield chunk_artifact
-        await context.store(chunk_artifact)
         await asyncio.sleep(0.3)
 
     if post_text := response[match.end() :]:
         message = AgentMessage(text=post_text)
         yield message
-        await context.store(message)
 
 
 if __name__ == "__main__":

apps/adk-py/examples/canvas_ui_test_agent.py

@@ -67,8 +67,6 @@ async def artifacts_agent(
 ):
     """Works with artifacts"""
 
-    await context.store(input)
-
     canvas_edit_request = await canvas.parse_canvas_edit_request(message=input)
 
     if canvas_edit_request:
@@ -106,7 +104,6 @@ async def artifacts_agent(
     if pre_text := response[: match.start()].strip():
         message = AgentMessage(text=pre_text)
         yield message
-        await context.store(message)
 
     await asyncio.sleep(1)
 
@@ -137,7 +134,6 @@ async def artifacts_agent(
         name=artifact_name,
         parts=[TextPart(text=recipe_content)],
     )
-    await context.store(artifact)
 
     # Send first chunk with artifact_id to establish the artifact
     first_artifact = AgentArtifact(
@@ -155,13 +151,11 @@ async def artifacts_agent(
             parts=[TextPart(text=chunk)],
         )
         yield chunk_artifact
-        await context.store(chunk_artifact)
         await asyncio.sleep(0.3)
 
     if post_text := response[match.end() :]:
         message = AgentMessage(text=post_text)
         yield message
-        await context.store(message)
 
 
 if __name__ == "__main__":

apps/adk-py/examples/citation_agent.py

@@ -12,7 +12,6 @@
 from kagenti_adk.a2a.types import AgentMessage
 from kagenti_adk.server import Server
 from kagenti_adk.server.context import RunContext
-from kagenti_adk.server.store.platform_context_store import PlatformContextStore
 
 server = Server()
 
@@ -27,9 +26,6 @@ async def example_agent(
 ):
     """Agent that demonstrates citation extension usage"""
 
-    # Store the current message in the context store
-    await context.store(input)
-
     # Simulate researching multiple sources
     research_text = """Based on recent research, artificial intelligence has made significant progress in natural
 language processing. Studies show that transformer models have revolutionized the field, and
@@ -60,12 +56,11 @@ async def example_agent(
         metadata=citation.citation_metadata(citations=citations),
     )
     yield message
-    await context.store(message)
 
 
 def run():
     server.run(
-        host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", 8000)), context_store=PlatformContextStore()
+        host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", 8000))
     )
 
 

apps/adk-py/examples/citation_agent_artifact.py

@@ -12,7 +12,6 @@
 from kagenti_adk.a2a.types import AgentArtifact
 from kagenti_adk.server import Server
 from kagenti_adk.server.context import RunContext
-from kagenti_adk.server.store.platform_context_store import PlatformContextStore
 
 server = Server()
 
@@ -27,9 +26,6 @@ async def example_agent(
 ):
     """Agent that demonstrates citation extension usage"""
 
-    # Store the current message in the context store
-    await context.store(input)
-
     # Simulate researching multiple sources
     research_text = """Based on recent research, artificial intelligence has made significant progress in natural
 language processing. Studies show that transformer models have revolutionized the field, and
@@ -62,12 +58,10 @@ async def example_agent(
     )
     yield artifact
 
-    await context.store(artifact)
-
 
 def run():
     server.run(
-        host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", 8002)), context_store=PlatformContextStore()
+        host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", 8002))
     )
 
 

apps/adk-py/examples/history.py

@@ -19,9 +19,6 @@
 async def example_agent(input: Message, context: RunContext):
     """Agent that demonstrates conversation history access"""
 
-    # Store the current message in the context store
-    await context.store(input)
-
     # Get the current user message
     current_message = get_message_text(input)
     print(f"Current message: {current_message}")
@@ -36,9 +33,6 @@ async def example_agent(input: Message, context: RunContext):
     message = AgentMessage(text=f"Hello! I can see we have {len(history)} messages in our conversation.")
     yield message
 
-    # Store the message in the context store
-    await context.store(message)
-
 
 def run():
     server.run(host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", 8000)))

apps/adk-py/examples/history_framework.py

@@ -18,7 +18,6 @@
 from kagenti_adk.a2a.types import AgentMessage
 from kagenti_adk.server import Server
 from kagenti_adk.server.context import RunContext
-from kagenti_adk.server.store.platform_context_store import PlatformContextStore
 
 server = Server()
 
@@ -44,8 +43,6 @@ async def multi_turn_chat_agent(
     llm: Annotated[LLMServiceExtensionServer, LLMServiceExtensionSpec.single_demand()],
 ):
     """Multi-turn chat agent with conversation memory and LLM integration"""
-    await context.store(input)
-
     # Load conversation history
     history = [message async for message in context.load_history() if isinstance(message, Message) and message.parts]
 
@@ -81,14 +78,12 @@ async def multi_turn_chat_agent(
                 response = AgentMessage(text=step.input["response"])
 
                 yield response
-                await context.store(response)
 
 
 def run():
     server.run(
         host=os.getenv("HOST", "127.0.0.1"),
         port=int(os.getenv("PORT", "8000")),
-        context_store=PlatformContextStore(),  # Enable persistent storage
     )