fix: raise SessionNotFoundError when appending to a missing session

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 954860721

Author: GWeale

src/google/adk/errors/session_not_found_error.py

@@ -21,5 +21,5 @@ class SessionNotFoundError(ValueError):
   Inherits from ValueError (for backward compatibility).
   """
 
-  def __init__(self, message="Session not found."):
+  def __init__(self, message: str = "Session not found.") -> None:
     super().__init__(message)

src/google/adk/integrations/firestore/firestore_session_service.py

@@ -29,6 +29,7 @@
 from typing import Optional
 
 from ...errors.already_exists_error import AlreadyExistsError
+from ...errors.session_not_found_error import SessionNotFoundError
 from ...events.event import Event
 from ...platform import uuid as platform_uuid
 from ...sessions import _session_util
@@ -504,7 +505,7 @@ async def _append_txn(transaction: firestore.AsyncTransaction) -> int:
         # 1. Reads
         session_snap = await session_ref.get(transaction=transaction)
         if not session_snap.exists:
-          raise ValueError(f"Session {session.id} not found.")
+          raise SessionNotFoundError(f"Session {session.id} not found.")
 
         session_doc = session_snap.to_dict() or {}
         if session_doc.get("status") == "DELETING":

src/google/adk/sessions/database_session_service.py

@@ -49,6 +49,7 @@
 
 from . import _session_util
 from ..errors.already_exists_error import AlreadyExistsError
+from ..errors.session_not_found_error import SessionNotFoundError
 from ..events.event import Event
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
@@ -779,7 +780,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
         storage_session_result = await sql_session.execute(storage_session_stmt)
         storage_session = storage_session_result.scalars().one_or_none()
         if storage_session is None:
-          raise ValueError(f"Session {session.id} not found.")
+          raise SessionNotFoundError(f"Session {session.id} not found.")
         storage_update_time = storage_session.get_update_timestamp(
             is_sqlite=is_sqlite, is_postgresql=is_postgresql
         )

src/google/adk/sessions/sqlite_session_service.py

@@ -31,6 +31,7 @@
 
 from . import _session_util
 from ..errors.already_exists_error import AlreadyExistsError
+from ..errors.session_not_found_error import SessionNotFoundError
 from ..events.event import Event
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
@@ -388,7 +389,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
       ) as cursor:
         row = await cursor.fetchone()
         if row is None:
-          raise ValueError(f"Session {session.id} not found.")
+          raise SessionNotFoundError(f"Session {session.id} not found.")
         storage_update_time = row["update_time"]
         if storage_update_time > session.last_update_time:
           raise ValueError(

tests/unittests/integrations/firestore/test_firestore_session_service.py

@@ -20,6 +20,7 @@
 from unittest import mock
 
 from google.adk.errors.already_exists_error import AlreadyExistsError
+from google.adk.errors.session_not_found_error import SessionNotFoundError
 from google.adk.events.event import Event
 from google.adk.events.event import EventActions
 from google.adk.integrations.firestore.firestore_session_service import FirestoreSessionService
@@ -268,6 +269,28 @@ async def test_append_event(mock_firestore_client):
   assert session.last_update_time == event.timestamp
 
 
+@pytest.mark.asyncio
+async def test_append_event_session_not_found(mock_firestore_client):
+  service = FirestoreSessionService(client=mock_firestore_client)
+  session = Session(id="test_session", app_name="test_app", user_id="test_user")
+  event = Event(invocation_id="test_inv", author="user")
+
+  session_doc_snapshot = mock.MagicMock()
+  session_doc_snapshot.exists = False
+
+  root_coll = mock_firestore_client.collection.return_value
+  app_ref = root_coll.document.return_value
+  users_coll = app_ref.collection.return_value
+  user_ref = users_coll.document.return_value
+  sessions_ref = user_ref.collection.return_value
+  session_doc_ref = sessions_ref.document.return_value
+  session_doc_ref.get = mock.AsyncMock(return_value=session_doc_snapshot)
+
+  with mock.patch("google.cloud.firestore.async_transactional", lambda x: x):
+    with pytest.raises(SessionNotFoundError):
+      await service.append_event(session, event)
+
+
 @pytest.mark.asyncio
 async def test_append_event_with_state_delta(mock_firestore_client):
   service = FirestoreSessionService(client=mock_firestore_client)

tests/unittests/sessions/test_session_service.py

@@ -22,6 +22,7 @@
 from unittest import mock
 
 from google.adk.errors.already_exists_error import AlreadyExistsError
+from google.adk.errors.session_not_found_error import SessionNotFoundError
 from google.adk.events.event import Event
 from google.adk.events.event_actions import EventActions
 from google.adk.features import FeatureName
@@ -802,6 +803,32 @@ async def test_session_last_update_time_updates_on_event(session_service):
   assert refreshed_session.last_update_time > original_update_time
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'service_type', [SessionServiceType.DATABASE, SessionServiceType.SQLITE]
+)
+async def test_append_event_to_deleted_session_raises_session_not_found(
+    service_type, tmp_path
+):
+  session_service = get_session_service(service_type, tmp_path)
+  try:
+    app_name = 'my_app'
+    user_id = 'user'
+    session = await session_service.create_session(
+        app_name=app_name, user_id=user_id
+    )
+    await session_service.delete_session(
+        app_name=app_name, user_id=user_id, session_id=session.id
+    )
+
+    event = Event(invocation_id='inv1', author='user')
+    with pytest.raises(SessionNotFoundError):
+      await session_service.append_event(session, event)
+  finally:
+    if isinstance(session_service, DatabaseSessionService):
+      await session_service.close()
+
+
 @pytest.mark.asyncio
 async def test_append_event_to_stale_session():
   session_service = get_session_service(