Merge branch 'main' into notif-model

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: