diff --git a/core/session_manager.py b/core/session_manager.py index fae6fe4..6a884f8 100644 --- a/core/session_manager.py +++ b/core/session_manager.py @@ -202,6 +202,15 @@ class SessionManager: """Persist a single message to the database.""" db = SessionLocal() try: + db_session = db.query(DbSession).filter(DbSession.id == session_id).first() + if db_session is None: + # A stream/tool callback can outlive a session delete. Do not + # create a chat_messages row with no parent session; also drop + # any stale cached session so later writes fail closed too. + self.sessions.pop(session_id, None) + logger.warning("Dropping message for deleted session %s", session_id) + return + msg_id = str(uuid.uuid4()) msg_time = datetime.utcnow() if message.metadata is None: @@ -223,15 +232,13 @@ class SessionManager: ) db.add(db_message) - db_session = db.query(DbSession).filter(DbSession.id == session_id).first() - if db_session: - db_session.message_count = len(self.sessions.get(session_id, {}).history) if session_id in self.sessions else 0 - _now = datetime.now(timezone.utc) - db_session.last_accessed = _now - # Clean "last conversation" timestamp — only bumped here on a - # real message persist, so it powers an accurate "Last active" - # sort that ignores renames / model swaps / mere opens. - db_session.last_message_at = _now + db_session.message_count = len(self.sessions.get(session_id, {}).history) if session_id in self.sessions else 0 + _now = datetime.now(timezone.utc) + db_session.last_accessed = _now + # Clean "last conversation" timestamp — only bumped here on a + # real message persist, so it powers an accurate "Last active" + # sort that ignores renames / model swaps / mere opens. + db_session.last_message_at = _now db.commit() diff --git a/tests/test_session_manager_persist_guard.py b/tests/test_session_manager_persist_guard.py new file mode 100644 index 0000000..cd15c0e --- /dev/null +++ b/tests/test_session_manager_persist_guard.py @@ -0,0 +1,52 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from core.models import ChatMessage +from core.session_manager import SessionManager +import core.session_manager as SM + + +def _manager_with(sessions): + manager = SessionManager.__new__(SessionManager) + manager.sessions = dict(sessions) + return manager + + +def _session_local(parent_row): + db = MagicMock() + db.query.return_value.filter.return_value.first.return_value = parent_row + return MagicMock(return_value=db), db + + +def test_persist_message_drops_write_when_parent_session_is_gone(monkeypatch): + session_local, db = _session_local(None) + monkeypatch.setattr(SM, "SessionLocal", session_local) + + manager = _manager_with({"deleted": SimpleNamespace(history=[])}) + message = ChatMessage("assistant", "late token") + + manager._persist_message("deleted", message) + + assert "deleted" not in manager.sessions + db.add.assert_not_called() + db.commit.assert_not_called() + db.rollback.assert_not_called() + + +def test_persist_message_still_writes_when_parent_session_exists(monkeypatch): + parent = SimpleNamespace(message_count=0, last_accessed=None, last_message_at=None) + session_local, db = _session_local(parent) + monkeypatch.setattr(SM, "SessionLocal", session_local) + + message = ChatMessage("user", "hello") + manager = _manager_with({"sid": SimpleNamespace(history=[message])}) + + manager._persist_message("sid", message) + + db.add.assert_called_once() + db.commit.assert_called_once() + assert parent.message_count == 1 + assert parent.last_accessed is not None + assert parent.last_message_at is not None + assert message.metadata["_db_id"] + assert message.metadata["timestamp"].endswith("Z")