fix: persist fork lineage metadata on default datastores

The session fork endpoint recorded parent_session_id/agent_id/source via
ds.update_session(metadata=...), but the shipped default stores dropped it:
SDKFileDataStore.update_session only handled title/platform/icon and
JSONFileDataStore had no override at all (base no-op). Fork lineage only
survived on the volatile MemoryDataStore used in tests, silently failing in
production.

- SDKFileDataStore: persist custom session metadata in the sentinel meta
  message and surface it via get_session().
- JSONFileDataStore: add update_session override persisting arbitrary kwargs
  (also fixes title updates being lost).
- Add regression test asserting fork metadata persists on the default stores.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: Mervin Praison <[email protected]>

src/praisonaiui/datastore.py

@@ -476,6 +476,18 @@ class JSONFileDataStore(BaseDataStore):
 
         return await asyncio.to_thread(_get)
 
+    async def update_session(self, session_id: str, **kwargs: Any) -> None:
+        def _update() -> None:
+            path = self._session_path(session_id)
+            data = self._read_session(path)
+            if data is None:
+                return
+            data.update(kwargs)
+            data["updated_at"] = datetime.now(timezone.utc).isoformat()
+            self._write_session(path, data)
+
+        await asyncio.to_thread(_update)
+
     def _read_feedback(self) -> list[dict[str, Any]]:
         """Read feedback from JSON file."""
         try:

src/praisonaiui/datastore_sdk.py

@@ -33,6 +33,8 @@ logger = logging.getLogger(__name__)
 # Sentinel key used to identify the AIUI metadata message inside SDK sessions
 _META_KEY = "_aiui_meta"
 _TITLE_KEY = "_aiui_title"
+# Bucket holding arbitrary session-level metadata (e.g. fork lineage)
+_CUSTOM_META_KEY = "_aiui_custom"
 
 
 class SDKFileDataStore(BaseDataStore):
@@ -116,6 +118,10 @@ class SDKFileDataStore(BaseDataStore):
             result["platform"] = meta["platform"]
         if meta.get("icon"):
             result["icon"] = meta["icon"]
+        # Surface any custom session-level metadata (e.g. fork lineage)
+        custom = meta.get(_CUSTOM_META_KEY)
+        if isinstance(custom, dict) and custom:
+            result["metadata"] = custom
         return result
 
     async def create_session(self, session_id: Optional[str] = None) -> dict[str, Any]:
@@ -183,6 +189,8 @@ class SDKFileDataStore(BaseDataStore):
             await self._persist_meta(session_id, "platform", kwargs["platform"])
         if "icon" in kwargs:
             await self._persist_meta(session_id, "icon", kwargs["icon"])
+        if "metadata" in kwargs and isinstance(kwargs["metadata"], dict):
+            await self._persist_meta(session_id, _CUSTOM_META_KEY, kwargs["metadata"])
 
     async def close(self) -> None:
         pass

tests/unit/test_session_fork.py

@@ -125,3 +125,49 @@ def test_fork_index_out_of_range_returns_400(store, client):
 def test_fork_virtual_agent_session_blocked(client):
     resp = client.post("/api/sessions/agent:support/fork", json={})
     assert resp.status_code == 400
+
+
+def _persistent_stores(tmp_path):
+    """Build the persistent datastores that ship as the default (SDK + JSON)."""
+    stores = []
+    from praisonaiui.datastore import JSONFileDataStore
+
+    stores.append(JSONFileDataStore(data_dir=str(tmp_path / "json")))
+    try:
+        from praisonaiui.datastore_sdk import SDKFileDataStore
+
+        stores.append(SDKFileDataStore(session_dir=str(tmp_path / "sdk")))
+    except (ImportError, Exception):
+        pass
+    return stores
+
+
+def test_fork_metadata_persists_on_default_stores(tmp_path):
+    """Fork lineage must survive on the persistent stores, not just MemoryDataStore.
+
+    Regression: SDKFileDataStore/JSONFileDataStore (the shipped defaults) previously
+    dropped session-level ``metadata`` in ``update_session``, so parent lineage was
+    silently lost outside of tests running on MemoryDataStore.
+    """
+    import praisonaiui.server as server
+
+    for ds in _persistent_stores(tmp_path):
+        original = server._datastore
+        server._datastore = ds
+        try:
+            client = TestClient(server.create_app())
+            _seed(ds, "parent", 2)
+            resp = client.post(
+                "/api/sessions/parent/fork",
+                json={"agent_id": "support_bot"},
+            )
+            assert resp.status_code == 200, type(ds).__name__
+            child_id = resp.json()["session_id"]
+
+            child = asyncio.run(ds.get_session(child_id))
+            meta = child.get("metadata", {})
+            assert meta.get("parent_session_id") == "parent", type(ds).__name__
+            assert meta.get("agent_id") == "support_bot", type(ds).__name__
+            assert meta.get("source") == "playground-fork", type(ds).__name__
+        finally:
+            server._datastore = original