fix(core): session store - update a session

resolves #99

Author: asaf

packages/core/src/flux0_core/sessions.py

@@ -19,7 +19,7 @@
 from flux0_core.users import UserId
 
 SessionId = NewType("SessionId", str)
-SessionMode: TypeAlias = Literal["auto", "manual"]
+SessionMode: TypeAlias = Literal["auto", "manual", "record", "replay"]
 
 
 ConsumerId: TypeAlias = Literal["client"]
@@ -44,6 +44,7 @@ class SessionUpdateParams(TypedDict, total=False):
     mode: SessionMode
     title: Optional[str]
     consumption_offsets: Mapping[ConsumerId, int]
+    metadata: Optional[Mapping[str, JSONSerializable]]
 
 
 EventId = NewType("EventId", str)

packages/core/src/flux0_core/storage/nanodb_memory.py

@@ -24,7 +24,7 @@
 from flux0_core.users import User, UserId, UserStore, UserUpdateParams
 from flux0_nanodb.api import DocumentCollection, DocumentDatabase
 from flux0_nanodb.query import And, Comparison, QueryFilter
-from flux0_nanodb.types import DocumentID, DocumentVersion
+from flux0_nanodb.types import DocumentID, DocumentVersion, JSONPatchOperation
 
 
 #############
@@ -421,7 +421,15 @@ async def update_session(
         session_id: SessionId,
         params: SessionUpdateParams,
     ) -> Session:
-        raise NotImplementedError
+        update_data = {k: v for k, v in params.items() if v is not None}
+        patch: List[JSONPatchOperation] = [
+            {"op": "replace", "path": f"/{k}", "value": v} for k, v in update_data.items()
+        ]
+        await self._session_col.update_one(Comparison(path="id", op="$eq", value=session_id), patch)
+        updated = await self.read_session(session_id)
+        if not updated:
+            raise ValueError(f"Session not found: {session_id}")
+        return updated
 
     @override
     async def list_sessions(

packages/core/tests/storage/test_nanodb_memory.py

@@ -2,7 +2,7 @@
 
 import pytest
 from flux0_core.agents import AgentId, AgentStore, AgentType
-from flux0_core.sessions import SessionStore, StatusEventData
+from flux0_core.sessions import SessionStore, SessionUpdateParams, StatusEventData
 from flux0_core.storage.nanodb_memory import (
     AgentDocumentStore,
     SessionDocumentStore,
@@ -108,7 +108,11 @@ async def test_session_crud(session_store: SessionStore) -> None:
     #
     rs = await session_store.read_session(s.id)
     assert rs == s
-    # TODO update
+    rs = await session_store.update_session(s.id, SessionUpdateParams(title="new title"))
+    assert rs.title == "new title"
+    rs = await session_store.read_session(s.id)
+    assert rs is not None
+    assert rs.title == "new title"
     # delete
     #
     ok = await session_store.delete_session(s.id)