diff --git a/src/va_agent/session.py b/src/va_agent/session.py index 706d7de..462dbea 100644 --- a/src/va_agent/session.py +++ b/src/va_agent/session.py @@ -6,6 +6,7 @@ import asyncio import logging import time import uuid +from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, override from google.adk.errors.already_exists_error import AlreadyExistsError @@ -102,6 +103,24 @@ class FirestoreSessionService(BaseSessionService): def _events_col(self, app_name: str, user_id: str, session_id: str) -> Any: return self._session_ref(app_name, user_id, session_id).collection("events") + @staticmethod + def _timestamp_to_float(value: Any, default: float = 0.0) -> float: + if value is None: + return default + if isinstance(value, (int, float)): + return float(value) + if hasattr(value, "timestamp"): + try: + return float(value.timestamp()) + except ( + TypeError, + ValueError, + OSError, + OverflowError, + ) as exc: # pragma: no cover + logger.debug("Failed to convert timestamp %r: %s", value, exc) + return default + # ------------------------------------------------------------------ # State helpers # ------------------------------------------------------------------ @@ -171,7 +190,7 @@ class FirestoreSessionService(BaseSessionService): ) ) - now = time.time() + now = datetime.now(UTC) write_coros.append( self._session_ref(app_name, user_id, session_id).set( { @@ -196,7 +215,7 @@ class FirestoreSessionService(BaseSessionService): user_id=user_id, id=session_id, state=merged, - last_update_time=now, + last_update_time=now.timestamp(), ) @override @@ -283,7 +302,9 @@ class FirestoreSessionService(BaseSessionService): id=session_id, state=merged, events=events, - last_update_time=session_data.get("last_update_time", 0.0), + last_update_time=self._timestamp_to_float( + session_data.get("last_update_time"), 0.0 + ), ) @override @@ -326,7 +347,9 @@ class FirestoreSessionService(BaseSessionService): id=data["session_id"], state=merged, events=[], - last_update_time=data.get("last_update_time", 0.0), + last_update_time=self._timestamp_to_float( + data.get("last_update_time"), 0.0 + ), ) ) @@ -366,6 +389,8 @@ class FirestoreSessionService(BaseSessionService): # Persist state deltas session_ref = self._session_ref(app_name, user_id, session_id) + last_update_dt = datetime.fromtimestamp(event.timestamp, UTC) + if event.actions and event.actions.state_delta: state_deltas = _session_util.extract_state_delta(event.actions.state_delta) @@ -386,16 +411,16 @@ class FirestoreSessionService(BaseSessionService): FieldPath("state", k).to_api_repr(): v for k, v in state_deltas["session"].items() } - field_updates["last_update_time"] = event.timestamp + field_updates["last_update_time"] = last_update_dt write_coros.append(session_ref.update(field_updates)) else: write_coros.append( - session_ref.update({"last_update_time": event.timestamp}) + session_ref.update({"last_update_time": last_update_dt}) ) await asyncio.gather(*write_coros) else: - await session_ref.update({"last_update_time": event.timestamp}) + await session_ref.update({"last_update_time": last_update_dt}) # Log token usage if event.usage_metadata: