diff --git a/core/database.py b/core/database.py index a9ad90b8b..250fafbf4 100644 --- a/core/database.py +++ b/core/database.py @@ -251,6 +251,16 @@ class Session(TimestampMixin, Base): 'crew_member_id': self.crew_member_id, } + +class GroupChatState(TimestampMixin, Base): + __tablename__ = "group_chat_states" + + parent_session_id = Column(String, ForeignKey("sessions.id", ondelete="CASCADE"), primary_key=True, index=True) + owner = Column(String, nullable=True, index=True) + mode = Column(String, nullable=False, default="parallel") + state = Column(JSON, nullable=False, default=dict) + + class ChatMessage(Base): """ SQLAlchemy model for ChatMessage table. diff --git a/routes/session_routes.py b/routes/session_routes.py index dc29a64e4..12fd63b19 100644 --- a/routes/session_routes.py +++ b/routes/session_routes.py @@ -10,8 +10,8 @@ import logging from core.session_manager import SessionManager from core.models import ChatMessage from src.request_models import SessionResponse -from core.database import Session as DbSession, SessionLocal, Document, GalleryImage, utcnow_naive -from src.auth_helpers import effective_user, _auth_disabled, owner_filter +from core.database import Session as DbSession, GroupChatState, SessionLocal, Document, GalleryImage, utcnow_naive +from src.auth_helpers import get_current_user, effective_user, _auth_disabled, owner_filter from src.session_image_cleanup import _generated_image_path_for_cleanup, session_image_refs from src.session_actions import is_session_recently_active from src.upload_handler import reserve_message_upload_references @@ -173,6 +173,105 @@ def _persist_session_headers(session_id: str, headers: dict | None) -> None: db.close() +_GROUP_CHAT_MODES = {"parallel", "round-robin"} +_GROUP_CHAT_MAX_PARTICIPANTS = 8 + + +def _group_state_str(value, max_len: int = 4096) -> str | None: + if value is None: + return None + text = str(value) + return text[:max_len] + + +def _normalize_group_character(value) -> dict | None: + if not isinstance(value, dict): + return None + normalized = {} + character_id = _group_state_str(value.get("characterId"), 256) + character_name = _group_state_str(value.get("characterName"), 1024) + character_prompt = _group_state_str(value.get("characterPrompt"), 50000) + if character_id: + normalized["characterId"] = character_id + if character_name: + normalized["characterName"] = character_name + if character_prompt: + normalized["characterPrompt"] = character_prompt + return normalized or None + + +def _normalize_group_model(value) -> dict: + if not isinstance(value, dict): + raise HTTPException(400, "Invalid group model") + mid = _group_state_str(value.get("mid"), 1024) + url = _group_state_str(value.get("url"), 2048) + if not mid or not url: + raise HTTPException(400, "Group model is missing its model id or endpoint URL") + + normalized = { + "mid": mid, + "url": url, + "display": _group_state_str(value.get("display") or mid, 1024), + } + endpoint_id = _group_state_str(value.get("endpointId"), 256) + endpoint_name = _group_state_str(value.get("epName"), 1024) + group_name = _group_state_str(value.get("_groupName"), 1024) + character = _normalize_group_character(value.get("character")) + if endpoint_id: + normalized["endpointId"] = endpoint_id + if endpoint_name: + normalized["epName"] = endpoint_name + if group_name: + normalized["_groupName"] = group_name + if character: + normalized["character"] = character + return normalized + + +def _normalize_group_state(raw_state, parent_session_id: str) -> dict: + if isinstance(raw_state, dict) and isinstance(raw_state.get("group_state"), dict): + raw_state = raw_state["group_state"] + if not isinstance(raw_state, dict): + raise HTTPException(400, "Invalid group chat state") + + mode = raw_state.get("mode") + if mode not in _GROUP_CHAT_MODES: + mode = "parallel" + + raw_models = raw_state.get("models") + if not isinstance(raw_models, list): + raise HTTPException(400, "Group chat state is missing models") + models = [ + _normalize_group_model(model) + for model in raw_models[:_GROUP_CHAT_MAX_PARTICIPANTS] + ] + if len(models) < 2: + raise HTTPException(400, "Group chat state must include at least two models") + + raw_participants = raw_state.get("participantSessions", []) + if not isinstance(raw_participants, list): + raise HTTPException(400, "Invalid group participant session list") + participant_sessions = [] + for value in raw_participants[:len(models)]: + participant_sessions.append(_group_state_str(value, 256) if value else None) + while len(participant_sessions) < len(models): + participant_sessions.append(None) + + try: + round_robin_idx = max(0, int(raw_state.get("roundRobinIdx", 0))) + except (TypeError, ValueError): + round_robin_idx = 0 + + return { + "active": True, + "mode": mode, + "models": models, + "participantSessions": participant_sessions, + "parentSessionId": parent_session_id, + "roundRobinIdx": round_robin_idx, + } + + _HIDDEN_SYSTEM_SESSION_NAMES = { "[Task] Chat Sessions Tidy", "[Task] Documents Tidy", @@ -568,6 +667,67 @@ def setup_session_routes( session_manager.save_sessions() return {"ok": True, "count": len(messages)} + @router.put("/session/{sid}/group_state") + async def save_group_state(request: Request, sid: str): + """Persist the group-chat participant/model state for the parent session.""" + _verify_session_owner(request, sid) + state = _normalize_group_state(await request.json(), sid) + user = effective_user(request) + participant_ids = { + participant_id + for participant_id in state["participantSessions"] + if participant_id + } + + db = SessionLocal() + try: + if participant_ids: + rows = db.query(DbSession.id, DbSession.owner).filter(DbSession.id.in_(participant_ids)).all() + owner_by_id = {row.id: row.owner for row in rows} + if set(owner_by_id) != participant_ids: + raise HTTPException(400, "Group participant session not found") + if user and any(owner != user for owner in owner_by_id.values()): + raise HTTPException(404, "Group participant session not found") + + group_state = db.query(GroupChatState).filter(GroupChatState.parent_session_id == sid).first() + if group_state is None: + group_state = GroupChatState(parent_session_id=sid) + db.add(group_state) + group_state.owner = user + group_state.mode = state["mode"] + group_state.state = state + group_state.updated_at = utcnow_naive() + db.commit() + return {"ok": True, "group_state": state} + except HTTPException: + db.rollback() + raise + except Exception: + db.rollback() + raise + finally: + db.close() + + @router.get("/session/{sid}/group_state") + def get_group_state(request: Request, sid: str): + """Return persisted group-chat state for a parent session, if one exists.""" + _verify_session_owner(request, sid) + user = effective_user(request) + db = SessionLocal() + try: + group_state = db.query(GroupChatState).filter(GroupChatState.parent_session_id == sid).first() + if group_state is None: + return {"ok": False, "group_state": None} + if user and group_state.owner and group_state.owner != user: + raise HTTPException(404, "Group chat state not found") + try: + state = _normalize_group_state(group_state.state, sid) + except HTTPException: + return {"ok": False, "group_state": None} + return {"ok": True, "group_state": state} + finally: + db.close() + @router.post("/session/{sid}/delete") def delete_session_beacon(request: Request, sid: str): """Delete session via POST (for navigator.sendBeacon on page close).""" diff --git a/static/js/group.js b/static/js/group.js index 2d0ab5b3e..732ebf3dd 100644 --- a/static/js/group.js +++ b/static/js/group.js @@ -675,6 +675,7 @@ export async function startGroup(models, parentSessionId) { } _saveState(); + await _saveStateToServer(); // Now select the session so the UI switches to it. if (_parentSessionId && window.sessionModule) { @@ -811,6 +812,7 @@ async function _sendRoundRobin(msg, box) { // Order is randomized per-message now, so _roundRobinIdx no longer drives // turn order; left in state for backward compat only. _saveState(); + _saveStateToServer().catch(() => {}); } /** After parallel responses, inject each model's response into all other sessions. */ @@ -975,32 +977,75 @@ async function _streamToHolder(modelIdx, sessionId, msg, holderEl, abortCtrl) { // ── State Persistence ──────────────────────────────── +function _buildState() { + return { + active: _active, + mode: _mode, + models: _models, + participantSessions: _participantSessions, + parentSessionId: _parentSessionId, + roundRobinIdx: _roundRobinIdx, + }; +} + +function _applyState(s, sessionId) { + if (!s || !s.active || s.parentSessionId !== sessionId) return false; + const models = Array.isArray(s.models) ? s.models : []; + if (models.length < 2) return false; + _active = true; + _mode = s.mode || 'parallel'; + _models = models; + _participantSessions = Array.isArray(s.participantSessions) ? s.participantSessions : []; + while (_participantSessions.length < _models.length) _participantSessions.push(null); + _parentSessionId = s.parentSessionId; + _roundRobinIdx = s.roundRobinIdx || 0; + return true; +} + function _saveState() { try { - localStorage.setItem(GROUP_STATE_KEY, JSON.stringify({ - active: _active, - mode: _mode, - models: _models, - participantSessions: _participantSessions, - parentSessionId: _parentSessionId, - roundRobinIdx: _roundRobinIdx, - })); + localStorage.setItem(GROUP_STATE_KEY, JSON.stringify(_buildState())); } catch (e) {} } -export function restoreState(sessionId) { +async function _saveStateToServer() { + if (!_parentSessionId || !_active || _models.length < 2) return; try { - const s = JSON.parse(localStorage.getItem(GROUP_STATE_KEY) || 'null'); - if (s && s.active && s.parentSessionId === sessionId) { - _active = true; - _mode = s.mode || 'parallel'; - _models = s.models || []; - _participantSessions = s.participantSessions || []; - _parentSessionId = s.parentSessionId; - _roundRobinIdx = s.roundRobinIdx || 0; + const res = await fetch(`${API_BASE}/api/session/${_parentSessionId}/group_state`, { + method: 'PUT', + credentials: 'same-origin', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(_buildState()), + }); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + } catch (e) { + console.warn('[group] Failed to persist group state:', e); + } +} + +export async function restoreState(sessionId) { + try { + const s = Storage.getJSON(GROUP_STATE_KEY, null); + if (_applyState(s, sessionId)) { + _saveState(); return true; } } catch (e) {} + + try { + const res = await fetch(`${API_BASE}/api/session/${sessionId}/group_state`, { + credentials: 'same-origin', + }); + if (!res.ok) return false; + const data = await res.json(); + const s = data.group_state || data.groupState; + if (data.ok && _applyState(s, sessionId)) { + _saveState(); + return true; + } + } catch (e) { + console.warn('[group] Failed to restore group state:', e); + } return false; } diff --git a/static/js/sessions.js b/static/js/sessions.js index cf59d478c..4e6ba44d1 100644 --- a/static/js/sessions.js +++ b/static/js/sessions.js @@ -2063,7 +2063,12 @@ export async function selectSession(id, { keepSidebar = false, showLoading = tru window.chatModule.checkPendingResearch(id); } // Restore group chat state if this is a group session - if (window.groupModule && window.groupModule.restoreState && window.groupModule.restoreState(id)) { + let restoredGroupState = false; + if (window.groupModule && window.groupModule.restoreState) { + restoredGroupState = await window.groupModule.restoreState(id); + if (navToken !== _sessionNavToken || currentSessionId !== id) return; + } + if (restoredGroupState) { if (window._syncGroupIndicator) window._syncGroupIndicator(true); // Hide model picker for group sessions const _mpw = document.getElementById('model-picker-wrap'); diff --git a/tests/test_group_chat_state_routes.py b/tests/test_group_chat_state_routes.py new file mode 100644 index 000000000..4db5f677e --- /dev/null +++ b/tests/test_group_chat_state_routes.py @@ -0,0 +1,171 @@ +import asyncio +import sys +import tempfile +import types +import uuid +from unittest.mock import MagicMock + +import pytest +from fastapi import HTTPException +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import NullPool + +import core.database as cdb +from core.database import GroupChatState, Session as DbSession + +_TMPDB = tempfile.NamedTemporaryFile(suffix=".db", delete=False) +_ENGINE = create_engine( + f"sqlite:///{_TMPDB.name}", + connect_args={"check_same_thread": False}, + poolclass=NullPool, +) +cdb.Base.metadata.create_all(_ENGINE) +_TS = sessionmaker(bind=_ENGINE, autoflush=False, autocommit=False) + + +class _JsonRequest: + def __init__(self, payload=None): + self._payload = payload or {} + + async def json(self): + return self._payload + + +def _stub_multipart_if_missing(monkeypatch): + try: + import python_multipart # noqa: F401 + return + except ImportError: + pass + stub = types.ModuleType("python_multipart") + stub.__version__ = "0.0.20" + monkeypatch.setitem(sys.modules, "python_multipart", stub) + + +def _reset_db(): + db = _TS() + try: + db.query(GroupChatState).delete() + db.query(DbSession).delete() + db.commit() + finally: + db.close() + + +def _add_session(session_id, owner="alice", name="session"): + db = _TS() + try: + db.add(DbSession( + id=session_id, + owner=owner, + name=name, + endpoint_url="http://localhost:11434", + model="llama3", + archived=False, + )) + db.commit() + finally: + db.close() + + +def _endpoint(router, path, method): + return next( + r.endpoint for r in reversed(router.routes) + if getattr(r, "path", "") == path and method in getattr(r, "methods", set()) + ) + + +def _routes(monkeypatch): + import routes.session_routes as sr + + _stub_multipart_if_missing(monkeypatch) + monkeypatch.setattr(sr, "SessionLocal", _TS) + monkeypatch.setattr(sr, "effective_user", lambda request: "alice") + return sr.setup_session_routes(MagicMock(), {}) + + +def _group_payload(parent_id, participant_ids): + return { + "active": True, + "mode": "round-robin", + "models": [ + { + "mid": "llama3", + "display": "Llama 3", + "url": "http://localhost:11434", + "endpointId": "local", + "_groupName": "Athena", + "character": { + "characterId": "athena", + "characterName": "Athena", + "characterPrompt": "Offer wise counsel.", + }, + }, + { + "mid": "mistral", + "display": "Mistral", + "url": "http://localhost:11434", + "_groupName": "Mistral", + }, + ], + "participantSessions": participant_ids, + "parentSessionId": "client-side-placeholder", + "roundRobinIdx": 3, + } + + +def test_group_chat_state_round_trips_personas_and_participant_sessions(monkeypatch): + _reset_db() + router = _routes(monkeypatch) + parent_id = str(uuid.uuid4()) + participant_ids = [str(uuid.uuid4()), str(uuid.uuid4())] + _add_session(parent_id, name="parent") + for session_id in participant_ids: + _add_session(session_id, name="participant") + + save = _endpoint(router, "/api/session/{sid}/group_state", "PUT") + get = _endpoint(router, "/api/session/{sid}/group_state", "GET") + + saved = asyncio.run(save(request=_JsonRequest(_group_payload(parent_id, participant_ids)), sid=parent_id)) + restored = get(request=MagicMock(), sid=parent_id) + + assert saved["ok"] is True + assert restored["ok"] is True + state = restored["group_state"] + assert state["parentSessionId"] == parent_id + assert state["mode"] == "round-robin" + assert state["participantSessions"] == participant_ids + assert state["models"][0]["_groupName"] == "Athena" + assert state["models"][0]["character"]["characterPrompt"] == "Offer wise counsel." + + +def test_group_chat_state_rejects_participants_from_other_users(monkeypatch): + _reset_db() + router = _routes(monkeypatch) + parent_id = str(uuid.uuid4()) + alice_participant = str(uuid.uuid4()) + bob_participant = str(uuid.uuid4()) + _add_session(parent_id, name="parent") + _add_session(alice_participant, name="alice participant") + _add_session(bob_participant, owner="bob", name="bob participant") + + save = _endpoint(router, "/api/session/{sid}/group_state", "PUT") + + with pytest.raises(HTTPException) as exc: + asyncio.run(save( + request=_JsonRequest(_group_payload(parent_id, [alice_participant, bob_participant])), + sid=parent_id, + )) + assert exc.value.status_code == 404 + + +def test_missing_group_chat_state_returns_empty_result(monkeypatch): + _reset_db() + router = _routes(monkeypatch) + parent_id = str(uuid.uuid4()) + _add_session(parent_id, name="parent") + + get = _endpoint(router, "/api/session/{sid}/group_state", "GET") + + assert get(request=MagicMock(), sid=parent_id) == {"ok": False, "group_state": None} diff --git a/tests/test_group_chat_storage.py b/tests/test_group_chat_storage.py index 2bd422311..063251e61 100644 --- a/tests/test_group_chat_storage.py +++ b/tests/test_group_chat_storage.py @@ -5,9 +5,20 @@ SOURCE = ( Path(__file__).resolve().parent.parent / "static" / "js" / "group.js" ).read_text(encoding="utf-8") +SESSIONS_SOURCE = ( + Path(__file__).resolve().parent.parent / "static" / "js" / "sessions.js" +).read_text(encoding="utf-8") + def test_group_session_sidebar_cache_uses_safe_json_loader(): assert "import Storage from './storage.js';" in SOURCE assert "Storage.getJSON('odysseus-group-sessions', [])" in SOURCE assert "Array.isArray(storedGroupSessions)" in SOURCE assert "JSON.parse(localStorage.getItem('odysseus-group-sessions')" not in SOURCE + + +def test_group_state_persists_to_server_and_restores_async(): + assert "async function _saveStateToServer()" in SOURCE + assert "/group_state" in SOURCE + assert "export async function restoreState(sessionId)" in SOURCE + assert "await window.groupModule.restoreState(id)" in SESSIONS_SOURCE