feat: Add remote session support (ActivitySessionService)

Author: marcusmotill

src/google/adk/integrations/temporal.py

@@ -16,6 +16,7 @@
 
 import functools
 import inspect
+import datetime
 from typing import Any, AsyncGenerator, Callable, Optional, List
 
 from temporalio import workflow, activity
@@ -153,3 +154,186 @@ def default_activities(cls) -> List[Callable]:
         Useful for registering activities with the Temporal Worker.
         """
         return [generate_content_activity]
+
+
+class SessionServiceActivities:
+    """Wraps a BaseSessionService to expose its methods as Temporal Activities."""
+    
+    def __init__(self, session_service: Any):
+        # We type hint as Any to avoid circular imports or strict dependency on BaseSessionService 
+        # definition availability at this module level if not necessary, 
+        # but logically it is a BaseSessionService.
+        self.session_service = session_service
+
+    @activity.defn(name="create_session")
+    async def create_session(
+        self,
+        app_name: str,
+        user_id: str,
+        state: Optional[dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+        extra_kwargs: Optional[dict[str, Any]] = None,
+    ) -> Any:
+        # We return Any (dict/model) that Temporal can serialize.
+        # BaseSessionService returns a Session object.
+        return await self.session_service.create_session(
+            app_name=app_name,
+            user_id=user_id,
+            state=state,
+            session_id=session_id,
+            **(extra_kwargs or {})
+        )
+
+    @activity.defn(name="get_session")
+    async def get_session(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        extra_kwargs: Optional[dict[str, Any]] = None,
+    ) -> Optional[Any]:
+        # Note: 'config' argument in get_session might need Pydantic serialization support
+        # We assume kwargs handles simple args or properly serialized objects.
+        return await self.session_service.get_session(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            **(extra_kwargs or {})
+        )
+
+    @activity.defn(name="list_sessions")
+    async def list_sessions(
+        self,
+        app_name: str,
+        user_id: Optional[str] = None
+    ) -> Any:
+        return await self.session_service.list_sessions(
+            app_name=app_name,
+            user_id=user_id
+        )
+
+    @activity.defn(name="delete_session")
+    async def delete_session(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str
+    ) -> None:
+        await self.session_service.delete_session(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id
+        )
+
+    @activity.defn(name="append_event")
+    async def append_event(self, session: Any, event: Any) -> Any:
+        return await self.session_service.append_event(session, event)
+
+    def get_activities(self) -> List[Callable]:
+        """Returns the list of activities to register."""
+        return [
+            self.create_session,
+            self.get_session,
+            self.list_sessions,
+            self.delete_session,
+            self.append_event
+        ]
+
+
+from google.adk.sessions import BaseSessionService, Session
+from google.adk.sessions.base_session_service import ListSessionsResponse
+from google.adk.events import Event
+
+
+class ActivitySessionService(BaseSessionService):
+    """A SessionService that delegates all calls to Temporal Activities.
+    
+    This ensures determinism within a Workflow by offloading the actual
+    session I/O (which might be non-deterministic or remote) to the Worker.
+    """
+    
+    def __init__(
+        self,
+        activity_options: Optional[dict[str, Any]] = None
+    ):
+        """Initializes the ActivitySessionService.
+
+        Args:
+            activity_options: Default options for activity execution (e.g. timeouts).
+                              Defaults to schedule_to_close_timeout=datetime.timedelta(seconds=30) if not provided.
+        """
+        self.activity_options = activity_options or {
+            "schedule_to_close_timeout": datetime.timedelta(seconds=30)
+        }
+
+    async def create_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        state: Optional[dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Session:
+        result = await workflow.execute_activity(
+            "create_session",
+            args=[app_name, user_id, state, session_id, kwargs],
+            **self.activity_options
+        )
+        if isinstance(result, dict):
+            return Session.model_validate(result)
+        return result
+
+    async def get_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        config: Optional[Any] = None,
+    ) -> Optional[Session]:
+        kwargs = {"config": config} if config else {}
+        result = await workflow.execute_activity(
+            "get_session",
+            args=[app_name, user_id, session_id, kwargs],
+            **self.activity_options
+        )
+        if result and isinstance(result, dict):
+            return Session.model_validate(result)
+        return result
+
+    async def list_sessions(
+        self,
+        *,
+        app_name: str,
+        user_id: Optional[str] = None
+    ) -> ListSessionsResponse:
+        result = await workflow.execute_activity(
+            "list_sessions",
+            args=[app_name, user_id],
+            **self.activity_options
+        )
+        if isinstance(result, dict):
+            return ListSessionsResponse.model_validate(result)
+        return result
+
+    async def delete_session(
+        self, *, app_name: str, user_id: str, session_id: str
+    ) -> None:
+        await workflow.execute_activity(
+            "delete_session",
+            args=[app_name, user_id, session_id],
+            **self.activity_options
+        )
+
+    async def append_event(self, session: Session, event: Event) -> Event:
+        # Note: We might need to serialize session/event to dicts if passing them as args causes issues?
+        # Usually PydanticConverter handles serialization of args fine.
+        result = await workflow.execute_activity(
+            "append_event",
+            args=[session, event],
+            **self.activity_options
+        )
+        if isinstance(result, dict):
+            return Event.model_validate(result)
+        return result

src/google/adk/sessions/in_memory_session_service.py

@@ -22,7 +22,6 @@
 
 from typing_extensions import override
 
-from google.adk import runtime
 from . import _session_util
 from ..errors.already_exists_error import AlreadyExistsError
 from ..events.event import Event
@@ -109,14 +108,14 @@ def _create_session_impl(
     session_id = (
         session_id.strip()
         if session_id and session_id.strip()
-        else runtime.new_uuid()
+        else str(uuid.uuid4())
     )
     session = Session(
         app_name=app_name,
         user_id=user_id,
         id=session_id,
         state=session_state or {},
-        last_update_time=runtime.get_time(),
+        last_update_time=time.time(),
     )
 
     if app_name not in self.sessions:

tests/integration/manual_test_temporal_integration.py

@@ -39,7 +39,7 @@
 from google.adk.sessions import InMemorySessionService
 from google.adk.utils.context_utils import Aclosing
 from google.adk.events import Event
-from google.adk.integrations.temporal import activity_as_tool, TemporalModel, generate_content_activity
+from google.adk.integrations.temporal import activity_as_tool, TemporalModel, generate_content_activity, ActivitySessionService, SessionServiceActivities
 
 # Required Environment Variables for this test:
 # - GOOGLE_CLOUD_PROJECT
@@ -102,9 +102,13 @@ async def run(self, prompt: str) -> Event | None:
             tools=[weather_tool]
         )
 
-        # 3. Create Session (uses runtime.new_uuid() -> workflow.uuid4())
-        session_service = InMemorySessionService()
-        logger.info("Create session.")
+        # 3. Create Session (uses ActivitySessionService to delegate to Worker)
+        # This ensures session I/O is deterministic (replay-safe)
+        session_service = ActivitySessionService()
+        logger.info("Create session (via Activity).")
+        # Note: server-generated IDs (or worker-generated) are safe because they come from Activity result.
+        # But we can also pass a deterministic ID if the service supports it.
+        # Here we let the service (Activity) handle it.
         session = await session_service.create_session(app_name="test_app", user_id="test")
 
         logger.info(f"Session created with ID: {session.id}")
@@ -194,7 +198,8 @@ async def run(self, prompt: str) -> str:
         )
 
         # 5. Execute
-        session_service = InMemorySessionService()
+        # Use ActivitySessionService to delegate session I/O to activities (replay-safe)
+        session_service = ActivitySessionService()
         session = await session_service.create_session(app_name="multi_agent_app", user_id="user_MULTI")
 
         runner = Runner(
@@ -287,10 +292,15 @@ async def test_temporalio_integration():
     except RuntimeError:
         pytest.skip("Could not connect to Temporal server. Is it running?")
 
+    # Register InMemorySessionService as activities (simulating remote/worker-side service)
+    # This is key: The state lives in THIS InMemorySessionService instance in the Worker.
+    real_session_service = InMemorySessionService()
+    session_activities = SessionServiceActivities(real_session_service)
+
     async with Worker(
         client,
         workflows=[WeatherAgent, MultiAgentWorkflow],
-        activities=TemporalModel.default_activities() + [
+        activities=TemporalModel.default_activities() + session_activities.get_activities() + [
             get_weather, 
             coordinator_think, 
             tool_agent_think,