Update Google ADK examples to use Restate's ADK extensions (#32)

* use ADK extensions from Restate SDK

* fix session ids and user ids

* Make examples more concise

* Update ADK examples

* Move runner out of the handler

* Cleanup

* Cleanup

* Fix broken link

* Fix broken link

Author: gvdongen

a2a/README.md

@@ -2,7 +2,7 @@
 
 **This integration is work-in-progress.**
 
-These examples use [Restate](https://ai.restate.dev/) to implement the [Agent2Agent (A2A) protocol](https://github.com/google/A2A).
+These examples use [Restate](https://ai.restate.dev/) to implement the [Agent2Agent (A2A) protocol](https://github.com/a2aproject/A2A).
 
 Restate acts as a scalable, resilient task orchestrator that speaks the A2A protocol and gives you:
 - 🔁 **Automatic retries** - Handles LLM API downtime, timeouts, and infrastructure failures

google-adk/example/app/chat.py

@@ -1,53 +1,38 @@
 import restate
+
 from google.adk import Runner
 from google.adk.apps import App
-from google.adk.sessions import Session
-from google.genai import types as genai_types
-
-from app.utils.models import ChatMessage
+from google.genai.types import Content, Part
 from google.adk.agents.llm_agent import Agent
+from restate.ext.adk import RestateSessionService, RestatePlugin
 
-from middleware.restate_plugin import RestatePlugin
-from middleware.restate_session_service import RestateSessionService
-from middleware.restate_utils import restate_overrides
+from app.utils.models import ChatMessage
 
 APP_NAME = "agents"
 
-# AGENT
 agent = Agent(
     model="gemini-2.5-flash",
     name="assistant",
-    description="A helpful assistant that can answer questions.",
     instruction="You are a helpful assistant. Be concise and helpful.",
 )
 
 # Enables retries and recovery for model calls and tool executions
 app = App(name=APP_NAME, root_agent=agent, plugins=[RestatePlugin()])
-session_service = RestateSessionService()
+runner = Runner(app=app, session_service=RestateSessionService())
 
 chat = restate.VirtualObject("Chat")
 
 
-# HANDLER
 @chat.handler()
-async def message(ctx: restate.ObjectContext, req: ChatMessage) -> str:
-    session_id = ctx.key()
-    with restate_overrides(ctx):
-        await session_service.create_session(
-            app_name=APP_NAME, user_id=req.user_id, session_id=session_id
-        )
-
-        runner = Runner(app=app, session_service=session_service)
-        events = runner.run_async(
-            user_id=req.user_id,
-            session_id=session_id,
-            new_message=genai_types.Content(
-                role="user", parts=[genai_types.Part.from_text(text=req.message)]
-            ),
-        )
-        final_response = ""
-        async for event in events:
-            if event.is_final_response() and event.content and event.content.parts:
-                if event.content.parts[0].text:
-                    final_response = event.content.parts[0].text
-        return final_response
+async def message(ctx: restate.ObjectContext, req: ChatMessage) -> str | None:
+    events = runner.run_async(
+        user_id=ctx.key(),
+        session_id=req.session_id,
+        new_message=Content(role="user", parts=[Part.from_text(text=req.message)]),
+    )
+    final_response = None
+    async for event in events:
+        if event.is_final_response() and event.content and event.content.parts:
+            if event.content.parts[0].text:
+                final_response = event.content.parts[0].text
+    return final_response

google-adk/example/app/durable_agent.py

@@ -4,23 +4,19 @@
 from google.adk.apps import App
 from google.adk.sessions import InMemorySessionService
 from google.genai.types import Content, Part
-from app.utils.models import WeatherResponse, WeatherPrompt
-from app.utils.utils import call_weather_api
-from google.adk.tools.tool_context import ToolContext
 from google.adk.agents.llm_agent import Agent
+from restate.ext.adk import RestatePlugin, restate_object_context
 
-from middleware.restate_plugin import RestatePlugin
-from middleware.restate_utils import restate_overrides
+from app.utils.models import WeatherResponse, WeatherPrompt
+from app.utils.utils import call_weather_api
 
 APP_NAME = "agents"
 
 
-# TOOLS
-async def get_weather(tool_context: ToolContext, city: str) -> WeatherResponse:
+async def get_weather(city: str) -> WeatherResponse:
     """Get the current weather for a given city."""
-    restate_context = tool_context.session.state["restate_context"]
     #  Do one or more durable steps using the Restate context
-    return await restate_context.run_typed(
+    return await restate_object_context().run_typed(
         f"Get weather {city}", call_weather_api, city=city
     )
 
@@ -29,9 +25,7 @@ async def get_weather(tool_context: ToolContext, city: str) -> WeatherResponse:
 agent = Agent(
     model="gemini-2.5-flash",
     name="weather_agent",
-    description="Agent that provides weather updates for cities.",
-    instruction="""You are a helpful agent that provides weather updates.
-    Use the get_weather tool to fetch current weather information.""",
+    instruction="You are a helpful agent that provides weather updates.",
     tools=[get_weather],
 )
 
@@ -41,23 +35,22 @@ async def get_weather(tool_context: ToolContext, city: str) -> WeatherResponse:
 
 
 @agent_service.handler()
-async def run(ctx: restate.Context, req: WeatherPrompt) -> str:
-    session_id = str(ctx.uuid())
-    with restate_overrides(ctx):
-        session_service = InMemorySessionService()
-        await session_service.create_session(
-            app_name=APP_NAME, user_id=req.user_id, session_id=session_id
-        )
-        runner = Runner(app=app, session_service=session_service)
-
-        events = runner.run_async(
-            user_id=req.user_id,
-            session_id=session_id,
-            new_message=Content(role="user", parts=[Part.from_text(text=req.message)]),
-        )
-        final_response = ""
-        async for event in events:
-            if event.is_final_response() and event.content and event.content.parts:
-                if event.content.parts[0].text:
-                    final_response = event.content.parts[0].text
-        return final_response
+async def run(_ctx: restate.Context, req: WeatherPrompt) -> str | None:
+    session_service = InMemorySessionService()
+    await session_service.create_session(
+        app_name=APP_NAME, user_id="user-123", session_id=req.session_id
+    )
+
+    runner = Runner(app=app, session_service=session_service)
+    events = runner.run_async(
+        user_id="user-123",
+        session_id=req.session_id,
+        new_message=Content(role="user", parts=[Part.from_text(text=req.message)]),
+    )
+
+    final_response = None
+    async for event in events:
+        if event.is_final_response() and event.content and event.content.parts:
+            if event.content.parts[0].text:
+                final_response = event.content.parts[0].text
+    return final_response

google-adk/example/app/durable_stateful_agent.py

@@ -1,65 +1,49 @@
-from datetime import timedelta
-
 import restate
+
 from google.adk import Runner
 from google.adk.apps import App
 from google.genai.types import Content, Part
-from app.utils.models import WeatherResponse, WeatherPrompt
-from app.utils.utils import call_weather_api
-from google.adk.tools.tool_context import ToolContext
 from google.adk.agents.llm_agent import Agent
+from restate.ext.adk import RestateSessionService, RestatePlugin, restate_object_context
 
-from middleware.restate_plugin import RestatePlugin
-from middleware.restate_session_service import RestateSessionService
-from middleware.restate_utils import restate_overrides
+from app.utils.models import WeatherResponse, WeatherPrompt
+from app.utils.utils import call_weather_api
 
 APP_NAME = "agents"
 
 
-# TOOLS
-async def get_weather(tool_context: ToolContext, city: str) -> WeatherResponse:
+async def get_weather(city: str) -> WeatherResponse:
     """Get the current weather for a given city."""
-    restate_context = tool_context.session.state["restate_context"]
-
     #  call tool wrapped as Restate durable step
-    return await restate_context.run_typed("Get weather", call_weather_api, city=city)
+    return await restate_object_context().run_typed(
+        f"Get weather {city}", call_weather_api, city=city
+    )
 
 
 agent = Agent(
     model="gemini-2.5-flash",
     name="weather_agent",
-    description="Agent that provides weather updates for cities.",
-    instruction="""You are a helpful agent that provides weather updates.
-    Use the get_weather tool to fetch current weather information.""",
+    instruction="You are a helpful agent that provides weather updates.",
     tools=[get_weather],
 )
 
 app = App(name=APP_NAME, root_agent=agent, plugins=[RestatePlugin()])
-session_service = RestateSessionService()
-
+runner = Runner(app=app, session_service=RestateSessionService())
 
 agent_service = restate.VirtualObject("StatefulWeatherAgent")
 
 
-# HANDLER
 @agent_service.handler()
-async def run(ctx: restate.ObjectContext, req: WeatherPrompt) -> str:
-    session_id = ctx.key()
-    with restate_overrides(ctx):
-        # Use Restate session service to persist session state in Restate
-        await session_service.create_session(
-            app_name=APP_NAME, user_id=req.user_id, session_id=session_id
-        )
-        runner = Runner(app=app, session_service=session_service)
-
-        events = runner.run_async(
-            user_id=req.user_id,
-            session_id=session_id,
-            new_message=Content(role="user", parts=[Part.from_text(text=req.message)]),
-        )
-        final_response = ""
-        async for event in events:
-            if event.is_final_response() and event.content and event.content.parts:
-                if event.content.parts[0].text:
-                    final_response = event.content.parts[0].text
-        return final_response
+async def run(ctx: restate.ObjectContext, req: WeatherPrompt) -> str | None:
+    events = runner.run_async(
+        user_id=ctx.key(),
+        session_id=req.session_id,
+        new_message=Content(role="user", parts=[Part.from_text(text=req.message)]),
+    )
+
+    final_response = None
+    async for event in events:
+        if event.is_final_response() and event.content and event.content.parts:
+            if event.content.parts[0].text:
+                final_response = event.content.parts[0].text
+    return final_response

google-adk/example/app/human_approval_agent.py

@@ -2,29 +2,25 @@
 from google.adk import Runner
 from google.adk.agents.llm_agent import Agent
 from google.adk.apps import App
-from google.adk.tools.tool_context import ToolContext
 from google.genai.types import Content, Part
+from restate.ext.adk import RestatePlugin, RestateSessionService, restate_object_context
 
 from app.utils.models import ClaimPrompt, InsuranceClaim
 from app.utils.utils import request_human_review
 
-from middleware.restate_plugin import RestatePlugin
-from middleware.restate_session_service import RestateSessionService
-from middleware.restate_utils import restate_overrides
-
 APP_NAME = "agents"
 
 
 # TOOLS
-async def human_approval(tool_context: ToolContext, claim: InsuranceClaim) -> str:
+async def human_approval(claim: InsuranceClaim) -> str:
     """Ask for human approval for high-value claims."""
-    restate_context = tool_context.session.state["restate_context"]
+    ctx = restate_object_context()
 
     # Create an awakeable for human approval
-    approval_id, approval_promise = restate_context.awakeable(type_hint=str)
+    approval_id, approval_promise = ctx.awakeable(type_hint=str)
 
     # Request human review
-    await restate_context.run_typed(
+    await ctx.run_typed(
         "Request review",
         request_human_review,
         claim=claim,
@@ -39,7 +35,6 @@ async def human_approval(tool_context: ToolContext, claim: InsuranceClaim) -> st
 agent = Agent(
     model="gemini-2.5-flash",
     name="claim_approval_agent",
-    description="Insurance claim evaluation agent that handles human approval workflows.",
     instruction="""You are an insurance claim evaluation agent. Use these rules: 
     - if the amount is more than 1000, ask for human approval using tools; 
     - if the amount is less than 1000, decide by yourself.""",
@@ -48,29 +43,23 @@ async def human_approval(tool_context: ToolContext, claim: InsuranceClaim) -> st
 
 
 app = App(name=APP_NAME, root_agent=agent, plugins=[RestatePlugin()])
-session_service = RestateSessionService()
+runner = Runner(app=app, session_service=RestateSessionService())
 
 agent_service = restate.VirtualObject("HumanClaimApprovalAgent")
 
 
 # HANDLER
 @agent_service.handler()
-async def run(ctx: restate.ObjectContext, req: ClaimPrompt) -> str:
-    session_id = ctx.key()
-    with restate_overrides(ctx):
-        await session_service.create_session(
-            app_name=APP_NAME, user_id=req.user_id, session_id=session_id
-        )
-        runner = Runner(app=app, session_service=session_service)
+async def run(ctx: restate.ObjectContext, req: ClaimPrompt) -> str | None:
+    events = runner.run_async(
+        user_id=ctx.key(),
+        session_id=req.session_id,
+        new_message=Content(role="user", parts=[Part.from_text(text=req.message)]),
+    )
 
-        events = runner.run_async(
-            user_id=req.user_id,
-            session_id=session_id,
-            new_message=Content(role="user", parts=[Part.from_text(text=req.message)]),
-        )
-        final_response = ""
-        async for event in events:
-            if event.is_final_response() and event.content and event.content.parts:
-                if event.content.parts[0].text:
-                    final_response = event.content.parts[0].text
-        return final_response
+    final_response = None
+    async for event in events:
+        if event.is_final_response() and event.content and event.content.parts:
+            if event.content.parts[0].text:
+                final_response = event.content.parts[0].text
+    return final_response