From 0497fc578cef4c397077152914aa655b9602b63d Mon Sep 17 00:00:00 2001
From: Shangjie Chen <deanchen@google.com>
Date: Mon, 19 May 2025 08:46:19 -0700
Subject: [PATCH] fix: fix return type of get_session.

PiperOrigin-RevId: 760638917
---
 src/google/adk/sessions/base_session_service.py      |  2 +-
 src/google/adk/sessions/in_memory_session_service.py | 10 ++++++----
 src/google/adk/sessions/vertex_ai_session_service.py |  2 +-
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/google/adk/sessions/base_session_service.py b/src/google/adk/sessions/base_session_service.py
index ed1540a..25e46ba 100644
--- a/src/google/adk/sessions/base_session_service.py
+++ b/src/google/adk/sessions/base_session_service.py
@@ -99,7 +99,7 @@ class BaseSessionService(abc.ABC):
     session.events.append(event)
     return event
 
-  def __update_session_state(self, session: Session, event: Event):
+  def __update_session_state(self, session: Session, event: Event) -> None:
     """Updates the session state based on the event."""
     if not event.actions or not event.actions.state_delta:
       return
diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py
index 9eb72c6..282e8e5 100644
--- a/src/google/adk/sessions/in_memory_session_service.py
+++ b/src/google/adk/sessions/in_memory_session_service.py
@@ -113,7 +113,7 @@ class InMemorySessionService(BaseSessionService):
       user_id: str,
       session_id: str,
       config: Optional[GetSessionConfig] = None,
-  ) -> Session:
+  ) -> Optional[Session]:
     return self._get_session_impl(
         app_name=app_name,
         user_id=user_id,
@@ -128,7 +128,7 @@ class InMemorySessionService(BaseSessionService):
       user_id: str,
       session_id: str,
       config: Optional[GetSessionConfig] = None,
-  ) -> Session:
+  ) -> Optional[Session]:
     logger.warning('Deprecated. Please migrate to the async method.')
     return self._get_session_impl(
         app_name=app_name,
@@ -144,7 +144,7 @@ class InMemorySessionService(BaseSessionService):
       user_id: str,
       session_id: str,
       config: Optional[GetSessionConfig] = None,
-  ) -> Session:
+  ) -> Optional[Session]:
     if app_name not in self.sessions:
       return None
     if user_id not in self.sessions[app_name]:
@@ -171,7 +171,9 @@ class InMemorySessionService(BaseSessionService):
 
     return self._merge_state(app_name, user_id, copied_session)
 
-  def _merge_state(self, app_name: str, user_id: str, copied_session: Session):
+  def _merge_state(
+      self, app_name: str, user_id: str, copied_session: Session
+  ) -> Session:
     # Merge app state
     if app_name in self.app_state:
       for key in self.app_state[app_name].keys():
diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py
index c56b23a..a6d9053 100644
--- a/src/google/adk/sessions/vertex_ai_session_service.py
+++ b/src/google/adk/sessions/vertex_ai_session_service.py
@@ -119,7 +119,7 @@ class VertexAiSessionService(BaseSessionService):
       user_id: str,
       session_id: str,
       config: Optional[GetSessionConfig] = None,
-  ) -> Session:
+  ) -> Optional[Session]:
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
 
     # Get session resource