Skip to content

Commit a3e4ad3

Browse files
GWealecopybara-github
authored andcommitted
fix: Update session last update time when appending events
The `database_session_service` now updates the `update_time` of a session to the event's timestamp when an event is appended Close #2721 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 834994070
1 parent 848fdbe commit a3e4ad3

File tree

3 files changed

+94
-7
lines changed

3 files changed

+94
-7
lines changed

src/google/adk/sessions/database_session_service.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -722,6 +722,11 @@ async def append_event(self, session: Session, event: Event) -> Event:
722722
if session_state_delta:
723723
storage_session.state = storage_session.state | session_state_delta
724724

725+
if storage_session._dialect_name == "sqlite":
726+
update_time = datetime.utcfromtimestamp(event.timestamp)
727+
else:
728+
update_time = datetime.fromtimestamp(event.timestamp)
729+
storage_session.update_time = update_time
725730
sql_session.add(StorageEvent.from_event(session, event))
726731

727732
await sql_session.commit()

src/google/adk/sessions/sqlite_session_service.py

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
323323

324324
# Trim temp state before persisting
325325
event = self._trim_temp_delta_state(event)
326-
now = time.time()
326+
event_timestamp = event.timestamp
327327

328328
async with self._get_db_connection() as db:
329329
# Check for stale session
@@ -355,11 +355,15 @@ async def append_event(self, session: Session, event: Event) -> Event:
355355

356356
if app_state_delta:
357357
await self._upsert_app_state(
358-
db, session.app_name, app_state_delta, now
358+
db, session.app_name, app_state_delta, event_timestamp
359359
)
360360
if user_state_delta:
361361
await self._upsert_user_state(
362-
db, session.app_name, session.user_id, user_state_delta, now
362+
db,
363+
session.app_name,
364+
session.user_id,
365+
user_state_delta,
366+
event_timestamp,
363367
)
364368
if session_state_delta:
365369
await self._update_session_state_in_db(
@@ -368,7 +372,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
368372
session.user_id,
369373
session.id,
370374
session_state_delta,
371-
now,
375+
event_timestamp,
372376
)
373377
has_session_state_delta = True
374378

@@ -392,12 +396,17 @@ async def append_event(self, session: Session, event: Event) -> Event:
392396
await db.execute(
393397
"UPDATE sessions SET update_time=? WHERE app_name=? AND user_id=?"
394398
" AND id=?",
395-
(now, session.app_name, session.user_id, session.id),
399+
(
400+
event_timestamp,
401+
session.app_name,
402+
session.user_id,
403+
session.id,
404+
),
396405
)
397406
await db.commit()
398407

399-
# Update timestamp with commit time
400-
session.last_update_time = now
408+
# Update timestamp based on event time
409+
session.last_update_time = event_timestamp
401410

402411
# Also update the in-memory session
403412
await super().append_event(session=session, event=event)

tests/unittests/sessions/test_session_service.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -531,6 +531,79 @@ async def test_append_event_complete(service_type, tmp_path):
531531
)
532532

533533

534+
@pytest.mark.asyncio
535+
@pytest.mark.parametrize(
536+
'service_type',
537+
[
538+
SessionServiceType.IN_MEMORY,
539+
SessionServiceType.DATABASE,
540+
SessionServiceType.SQLITE,
541+
],
542+
)
543+
async def test_session_last_update_time_updates_on_event(
544+
service_type, tmp_path
545+
):
546+
session_service = get_session_service(service_type, tmp_path)
547+
app_name = 'my_app'
548+
user_id = 'user'
549+
550+
session = await session_service.create_session(
551+
app_name=app_name, user_id=user_id
552+
)
553+
original_update_time = session.last_update_time
554+
555+
event_timestamp = original_update_time + 10
556+
event = Event(
557+
invocation_id='invocation',
558+
author='user',
559+
timestamp=event_timestamp,
560+
)
561+
await session_service.append_event(session=session, event=event)
562+
563+
assert session.last_update_time == pytest.approx(event_timestamp, abs=1e-6)
564+
565+
refreshed_session = await session_service.get_session(
566+
app_name=app_name, user_id=user_id, session_id=session.id
567+
)
568+
assert refreshed_session is not None
569+
assert refreshed_session.last_update_time == pytest.approx(
570+
event_timestamp, abs=1e-6
571+
)
572+
assert refreshed_session.last_update_time > original_update_time
573+
574+
575+
@pytest.mark.asyncio
576+
@pytest.mark.parametrize(
577+
'service_type',
578+
[
579+
SessionServiceType.IN_MEMORY,
580+
SessionServiceType.DATABASE,
581+
SessionServiceType.SQLITE,
582+
],
583+
)
584+
async def test_get_session_with_config(service_type):
585+
session_service = get_session_service(service_type)
586+
app_name = 'my_app'
587+
user_id = 'user'
588+
589+
session = await session_service.create_session(
590+
app_name=app_name, user_id=user_id
591+
)
592+
original_update_time = session.last_update_time
593+
594+
event = Event(invocation_id='invocation', author='user')
595+
await session_service.append_event(session=session, event=event)
596+
597+
assert session.last_update_time >= event.timestamp
598+
599+
refreshed_session = await session_service.get_session(
600+
app_name=app_name, user_id=user_id, session_id=session.id
601+
)
602+
assert refreshed_session is not None
603+
assert refreshed_session.last_update_time >= event.timestamp
604+
assert refreshed_session.last_update_time > original_update_time
605+
606+
534607
@pytest.mark.asyncio
535608
@pytest.mark.parametrize(
536609
'service_type',

0 commit comments

Comments
 (0)