Merge pull request #224 from MervinPraison/claude/issue-221-20260715-000000

feat: session fork endpoint for Agent Playground (STITCH-011)

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

src/praisonaiui/features/sessions_ext.py

@@ -132,6 +132,14 @@ class _InMemorySessionStore(SessionProtocol):
 _session_metadata: Dict[str, Dict[str, Any]] = {}
 
 
+_MESSAGE_COPY_FIELDS = ("role", "content", "text", "tool_calls", "name", "metadata")
+
+
+def _sanitize_message(message: Dict[str, Any]) -> Dict[str, Any]:
+    """Copy a message for a forked session, stripping ephemeral identifiers."""
+    return {k: message[k] for k in _MESSAGE_COPY_FIELDS if k in message}
+
+
 def _get_metadata(session_id: str) -> Dict[str, Any]:
     """Get session metadata (labels, usage, etc)."""
     if session_id not in _session_metadata:
@@ -166,6 +174,7 @@ class SessionsFeature(BaseFeatureProtocol):
             Route("/api/sessions/{session_id}/context", self._build_context, methods=["POST"]),
             Route("/api/sessions/{session_id}/compact", self._compact, methods=["POST"]),
             Route("/api/sessions/{session_id}/reset", self._reset, methods=["POST"]),
+            Route("/api/sessions/{session_id}/fork", self._fork, methods=["POST"]),
             Route("/api/sessions/{session_id}/preview", self._preview, methods=["GET"]),
             Route("/api/sessions/{session_id}/labels", self._labels, methods=["GET"]),
             Route("/api/sessions/{session_id}/labels", self._set_labels, methods=["POST"]),
@@ -519,6 +528,89 @@ class SessionsFeature(BaseFeatureProtocol):
             }
         )
 
+    async def _fork(self, request: Request) -> JSONResponse:
+        """POST /api/sessions/{id}/fork — Branch a session into a new child.
+
+        Body (all optional):
+            up_to_index: int   — keep messages[0..up_to_index] inclusive
+            title_suffix: str  — appended to the parent title (default " (fork)")
+            agent_id: str      — override recorded on the child metadata
+            labels: list[str]  — labels applied to the child
+        """
+        source_id = request.path_params["session_id"]
+
+        if source_id.startswith("agent:"):
+            return JSONResponse(
+                {"error": "Cannot fork a virtual agent session"}, status_code=400
+            )
+
+        content_type = request.headers.get("content-type") or ""
+        body = await request.json() if content_type.startswith("application/json") else {}
+
+        try:
+            from praisonaiui.server import get_datastore
+
+            ds = get_datastore()
+        except Exception as e:
+            logger.warning(f"fork: datastore unavailable: {e}")
+            return JSONResponse({"error": "Datastore unavailable"}, status_code=503)
+
+        source = await ds.get_session(source_id)
+        if source is None:
+            return JSONResponse({"error": "Source session not found"}, status_code=404)
+
+        messages = await ds.get_messages(source_id)
+        total = len(messages)
+
+        up_to_index = body.get("up_to_index")
+        if up_to_index is not None:
+            if not isinstance(up_to_index, int) or up_to_index < 0 or up_to_index >= max(total, 1):
+                return JSONResponse(
+                    {"error": "up_to_index out of range", "total_messages": total},
+                    status_code=400,
+                )
+            messages = messages[: up_to_index + 1]
+
+        child = await ds.create_session()
+        child_id = child.get("id") or child.get("session_id")
+
+        for msg in messages:
+            await ds.add_message(child_id, _sanitize_message(msg))
+
+        suffix = body.get("title_suffix", " (fork)")
+        parent_title = source.get("title") or ""
+        child_meta = {
+            "parent_session_id": source_id,
+            "forked_at": time.time(),
+            "fork_point_index": up_to_index,
+            "source": "playground-fork",
+        }
+        if body.get("agent_id"):
+            child_meta["agent_id"] = body["agent_id"]
+
+        try:
+            await ds.update_session(
+                child_id,
+                title=f"{parent_title}{suffix}",
+                metadata=child_meta,
+            )
+        except Exception as e:
+            logger.warning(f"fork: update_session metadata failed: {e}")
+
+        labels = body.get("labels")
+        if isinstance(labels, list) and labels:
+            _get_metadata(child_id)["_labels"] = [str(x) for x in labels]
+
+        return JSONResponse(
+            {
+                "session_id": child_id,
+                "parent_id": source_id,
+                "forked_at_index": up_to_index,
+                "message_count": len(messages),
+                "created_at": time.time(),
+            }
+        )
+
     async def _labels(self, request: Request) -> JSONResponse:
         sid = request.path_params["session_id"]
         meta = _get_metadata(sid)

tests/unit/test_session_fork.py

@@ -0,0 +1,173 @@
+"""Unit tests for the session fork endpoint (POST /api/sessions/{id}/fork)."""
+
+import asyncio
+
+import pytest
+from starlette.testclient import TestClient
+
+from praisonaiui.datastore import MemoryDataStore
+from praisonaiui.features.sessions_ext import _sanitize_message
+
+
+@pytest.fixture
+def store():
+    return MemoryDataStore()
+
+
+@pytest.fixture
+def client(store):
+    import praisonaiui.server as server
+
+    original = server._datastore
+    server._datastore = store
+    try:
+        yield TestClient(server.create_app())
+    finally:
+        server._datastore = original
+
+
+def _seed(store, session_id, count):
+    async def _run():
+        await store.create_session(session_id)
+        for i in range(count):
+            role = "user" if i % 2 == 0 else "assistant"
+            await store.add_message(session_id, {"role": role, "content": f"msg {i}"})
+
+    asyncio.run(_run())
+
+
+def _messages(store, session_id):
+    return asyncio.run(store.get_messages(session_id))
+
+
+def _session(store, session_id):
+    return asyncio.run(store.get_session(session_id))
+
+
+def test_sanitize_message_strips_ephemeral_ids():
+    msg = {
+        "id": "ephemeral-1",
+        "role": "user",
+        "content": "hi",
+        "tool_calls": [{"name": "x"}],
+        "timestamp": 123,
+    }
+    out = _sanitize_message(msg)
+    assert "id" not in out
+    assert "timestamp" not in out
+    assert out["role"] == "user"
+    assert out["content"] == "hi"
+    assert out["tool_calls"] == [{"name": "x"}]
+
+
+def test_fork_copies_all_messages(store, client):
+    _seed(store, "parent", 4)
+
+    resp = client.post("/api/sessions/parent/fork", json={})
+    assert resp.status_code == 200
+    body = resp.json()
+
+    assert body["parent_id"] == "parent"
+    assert body["message_count"] == 4
+    child_id = body["session_id"]
+    assert child_id != "parent"
+
+    child_messages = _messages(store, child_id)
+    assert len(child_messages) == 4
+    assert child_messages[0]["content"] == "msg 0"
+
+
+def test_fork_up_to_index_truncates(store, client):
+    _seed(store, "parent", 10)
+
+    resp = client.post("/api/sessions/parent/fork", json={"up_to_index": 3})
+    assert resp.status_code == 200
+    body = resp.json()
+
+    assert body["forked_at_index"] == 3
+    assert body["message_count"] == 4
+
+    child_messages = _messages(store, body["session_id"])
+    assert len(child_messages) == 4
+    assert child_messages[-1]["content"] == "msg 3"
+
+
+def test_fork_records_parent_metadata(store, client):
+    _seed(store, "parent", 2)
+
+    resp = client.post(
+        "/api/sessions/parent/fork",
+        json={"agent_id": "support_bot", "labels": ["playground-fork"]},
+    )
+    assert resp.status_code == 200
+    child_id = resp.json()["session_id"]
+
+    child = _session(store, child_id)
+    meta = child.get("metadata", {})
+    assert meta.get("parent_session_id") == "parent"
+    assert meta.get("agent_id") == "support_bot"
+    assert meta.get("source") == "playground-fork"
+
+
+def test_fork_missing_source_returns_404(client):
+    resp = client.post("/api/sessions/does-not-exist/fork", json={})
+    assert resp.status_code == 404
+
+
+def test_fork_index_out_of_range_returns_400(store, client):
+    _seed(store, "parent", 3)
+
+    resp = client.post("/api/sessions/parent/fork", json={"up_to_index": 99})
+    assert resp.status_code == 400
+    assert resp.json()["total_messages"] == 3
+
+
+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