feat: add session fork endpoint for Agent Playground (fixes #221)

POST /api/sessions/{id}/fork branches a chat session into a new child,
copying messages up to an optional up_to_index and recording parent
metadata. Backend primitive for the Agent Playground (STITCH-011).

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

Co-Authored-By: Claude <[email protected]>

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,127 @@
+"""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