diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 84977294..afc09b3e 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -31,7 +31,6 @@ from sqlalchemy.ext.asyncio import async_sessionmaker from sqlalchemy.ext.asyncio import AsyncEngine from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory from sqlalchemy.ext.asyncio import create_async_engine -from sqlalchemy.inspection import inspect from sqlalchemy.pool import StaticPool from typing_extensions import override from tzlocal import get_localzone @@ -275,22 +274,29 @@ class DatabaseSessionService(BaseSessionService): storage_user_state.state = storage_user_state.state | user_state_delta # Store the session + now = datetime.now(timezone.utc) + is_sqlite = self.db_engine.dialect.name == "sqlite" + if is_sqlite: + now = now.replace(tzinfo=None) + storage_session = schema.StorageSession( app_name=app_name, user_id=user_id, id=session_id, state=session_state, + create_time=now, + update_time=now, ) sql_session.add(storage_session) await sql_session.commit() - await sql_session.refresh(storage_session) - # Merge states for response merged_state = _merge_state( storage_app_state.state, storage_user_state.state, session_state ) - session = storage_session.to_session(state=merged_state) + session = storage_session.to_session( + state=merged_state, is_sqlite=is_sqlite + ) return session @override @@ -350,7 +356,10 @@ class DatabaseSessionService(BaseSessionService): # Convert storage session to session events = [e.to_event() for e in reversed(storage_events)] - session = storage_session.to_session(state=merged_state, events=events) + is_sqlite = self.db_engine.dialect.name == "sqlite" + session = storage_session.to_session( + state=merged_state, events=events, is_sqlite=is_sqlite + ) return session @override @@ -393,11 +402,14 @@ class DatabaseSessionService(BaseSessionService): user_states_map[storage_user_state.user_id] = storage_user_state.state sessions = [] + is_sqlite = self.db_engine.dialect.name == "sqlite" for storage_session in results: session_state = storage_session.state user_state = user_states_map.get(storage_session.user_id, {}) merged_state = _merge_state(app_state, user_state, session_state) - sessions.append(storage_session.to_session(state=merged_state)) + sessions.append( + storage_session.to_session(state=merged_state, is_sqlite=is_sqlite) + ) return ListSessionsResponse(sessions=sessions) @override @@ -433,15 +445,6 @@ class DatabaseSessionService(BaseSessionService): schema.StorageSession, (session.app_name, session.user_id, session.id) ) - if storage_session.update_timestamp_tz > session.last_update_time: - raise ValueError( - "The last_update_time provided in the session object" - f" {datetime.fromtimestamp(session.last_update_time):'%Y-%m-%d %H:%M:%S'} is" - " earlier than the update_time in the storage_session" - f" {datetime.fromtimestamp(storage_session.update_timestamp_tz):'%Y-%m-%d %H:%M:%S'}." - " Please check if it is a stale session." - ) - # Fetch states from storage storage_app_state = await sql_session.get( schema.StorageAppState, (session.app_name) @@ -450,6 +453,29 @@ class DatabaseSessionService(BaseSessionService): schema.StorageUserState, (session.app_name, session.user_id) ) + is_sqlite = self.db_engine.dialect.name == "sqlite" + if ( + storage_session.get_update_timestamp(is_sqlite) + > session.last_update_time + ): + # Reload the session from storage if it has been updated since it was + # loaded. + app_state = storage_app_state.state if storage_app_state else {} + user_state = storage_user_state.state if storage_user_state else {} + session_state = storage_session.state + session.state = _merge_state(app_state, user_state, session_state) + + stmt = ( + select(schema.StorageEvent) + .filter(schema.StorageEvent.app_name == session.app_name) + .filter(schema.StorageEvent.session_id == session.id) + .filter(schema.StorageEvent.user_id == session.user_id) + .order_by(schema.StorageEvent.timestamp.asc()) + ) + result = await sql_session.stream_scalars(stmt) + storage_events = [e async for e in result] + session.events = [e.to_event() for e in storage_events] + # Extract state delta if event.actions and event.actions.state_delta: state_deltas = _session_util.extract_state_delta( @@ -466,7 +492,7 @@ class DatabaseSessionService(BaseSessionService): if session_state_delta: storage_session.state = storage_session.state | session_state_delta - if storage_session._dialect_name == "sqlite": + if is_sqlite: update_time = datetime.fromtimestamp( event.timestamp, timezone.utc ).replace(tzinfo=None) @@ -476,10 +502,9 @@ class DatabaseSessionService(BaseSessionService): sql_session.add(schema.StorageEvent.from_event(session, event)) await sql_session.commit() - await sql_session.refresh(storage_session) # Update timestamp with commit time - session.last_update_time = storage_session.update_timestamp_tz + session.last_update_time = storage_session.get_update_timestamp(is_sqlite) # Also update the in-memory session await super().append_event(session=session, event=event) diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index bbdd66a2..8603ed81 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -41,7 +41,6 @@ from sqlalchemy import func from sqlalchemy import Text from sqlalchemy.dialects import mysql from sqlalchemy.ext.mutable import MutableDict -from sqlalchemy.inspection import inspect from sqlalchemy.orm import DeclarativeBase from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column @@ -131,15 +130,9 @@ class StorageSession(Base): def __repr__(self): return f"" - @property - def _dialect_name(self) -> Optional[str]: - session = inspect(self).session - return session.bind.dialect.name if session else None - - @property - def update_timestamp_tz(self) -> datetime: + def get_update_timestamp(self, is_sqlite: bool) -> float: """Returns the time zone aware update timestamp.""" - if self._dialect_name == "sqlite": + if is_sqlite: # SQLite does not support timezone. SQLAlchemy returns a naive datetime # object without timezone information. We need to convert it to UTC # manually. @@ -150,6 +143,7 @@ class StorageSession(Base): self, state: dict[str, Any] | None = None, events: list[Event] | None = None, + is_sqlite: bool = False, ) -> Session: """Converts the storage session to a session object.""" if state is None: @@ -163,7 +157,7 @@ class StorageSession(Base): id=self.id, state=state, events=events, - last_update_time=self.update_timestamp_tz, + last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite), ) diff --git a/src/google/adk/sessions/schemas/v1.py b/src/google/adk/sessions/schemas/v1.py index 34b84dc1..a7c429b2 100644 --- a/src/google/adk/sessions/schemas/v1.py +++ b/src/google/adk/sessions/schemas/v1.py @@ -32,7 +32,6 @@ import uuid from sqlalchemy import ForeignKeyConstraint from sqlalchemy import func from sqlalchemy.ext.mutable import MutableDict -from sqlalchemy.inspection import inspect from sqlalchemy.orm import DeclarativeBase from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column @@ -106,15 +105,9 @@ class StorageSession(Base): def __repr__(self): return f"" - @property - def _dialect_name(self) -> Optional[str]: - session = inspect(self).session - return session.bind.dialect.name if session else None - - @property - def update_timestamp_tz(self) -> datetime: + def get_update_timestamp(self, is_sqlite: bool) -> float: """Returns the time zone aware update timestamp.""" - if self._dialect_name == "sqlite": + if is_sqlite: # SQLite does not support timezone. SQLAlchemy returns a naive datetime # object without timezone information. We need to convert it to UTC # manually. @@ -125,6 +118,7 @@ class StorageSession(Base): self, state: dict[str, Any] | None = None, events: list[Event] | None = None, + is_sqlite: bool = False, ) -> Session: """Converts the storage session to a session object.""" if state is None: @@ -138,7 +132,7 @@ class StorageSession(Base): id=self.id, state=state, events=events, - last_update_time=self.update_timestamp_tz, + last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite), ) diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index c4c86575..30eb1567 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -504,26 +504,65 @@ async def test_session_last_update_time_updates_on_event(session_service): @pytest.mark.asyncio -async def test_get_session_with_config(session_service): +async def test_append_event_to_stale_session(): + session_service = get_session_service( + service_type=SessionServiceType.DATABASE + ) + app_name = 'my_app' user_id = 'user' + current_time = datetime.now().astimezone(timezone.utc).timestamp() - session = await session_service.create_session( + original_session = await session_service.create_session( app_name=app_name, user_id=user_id ) - original_update_time = session.last_update_time - - event = Event(invocation_id='invocation', author='user') - await session_service.append_event(session=session, event=event) - - assert session.last_update_time >= event.timestamp - - refreshed_session = await session_service.get_session( - app_name=app_name, user_id=user_id, session_id=session.id + event1 = Event( + invocation_id='inv1', + author='user', + timestamp=current_time + 1, + actions=EventActions(state_delta={'sk1': 'v1'}), ) - assert refreshed_session is not None - assert refreshed_session.last_update_time >= event.timestamp - assert refreshed_session.last_update_time > original_update_time + await session_service.append_event(original_session, event1) + + updated_session = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=original_session.id + ) + event2 = Event( + invocation_id='inv2', + author='user', + timestamp=current_time + 2, + actions=EventActions(state_delta={'sk2': 'v2'}), + ) + await session_service.append_event(updated_session, event2) + + # original_session is now stale + assert original_session.last_update_time < updated_session.last_update_time + assert len(original_session.events) == 1 + assert 'sk2' not in original_session.state + + # Appending another event to stale original_session + event3 = Event( + invocation_id='inv3', + author='user', + timestamp=current_time + 3, + actions=EventActions(state_delta={'sk3': 'v3'}), + ) + await session_service.append_event(original_session, event3) + + # If we fetch session from DB, it should contain all 3 events and all state + # changes. + session_final = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=original_session.id + ) + assert len(session_final.events) == 3 + assert session_final.state.get('sk1') == 'v1' + assert session_final.state.get('sk2') == 'v2' + assert session_final.state.get('sk3') == 'v3' + assert [e.invocation_id for e in session_final.events] == [ + 'inv1', + 'inv2', + 'inv3', + ] @pytest.mark.asyncio