diff --git a/alembic/versions/0008_session_conversation_id.py b/alembic/versions/0008_session_conversation_id.py index a28a6a34d..20cf679d8 100644 --- a/alembic/versions/0008_session_conversation_id.py +++ b/alembic/versions/0008_session_conversation_id.py @@ -23,5 +23,6 @@ def upgrade() -> None: def downgrade() -> None: - op.drop_index("ix_session_ownership_conversation_id", table_name="session_ownership") - op.drop_column("session_ownership", "conversation_id") + with op.batch_alter_table("session_ownership") as batch: + batch.drop_index("ix_session_ownership_conversation_id") + batch.drop_column("conversation_id") diff --git a/docs/docs/core-components/chat.mdx b/docs/docs/core-components/chat.mdx index 91255de7c..7ed92b7c5 100644 --- a/docs/docs/core-components/chat.mdx +++ b/docs/docs/core-components/chat.mdx @@ -34,11 +34,19 @@ These IDs can differ when Langflow returns a new response ID. Both IDs are checked against the authenticated user's session ownership; interruption recovery does not bypass those checks. Response aliases are durably associated with their canonical conversation, including across backend restarts. +After a backend restart, continuing with an owned response alias reuses the canonical sidebar thread and preserves its saved title and message count. +Successful message bodies are still retrieved from Langflow; the backend does not reconstruct them from metadata. Deleting a conversation revokes its canonical ID and all recorded aliases and attempts to remove every associated Langflow session. A deleted alias cannot be used to resume the conversation. +Bulk deletion also removes the caller's own legacy local conversation records that have no ownership row. +Those orphan records do not authorize upstream session deletion or bypass the ownership checks on chat requests. +Orphan cleanup is bulk-only at the API: single-conversation delete endpoints return `404` for records without an ownership row. IDs are sent only after ownership and its conversation association have been persisted. In `files` mode this requires a successful file write; in `db` mode it requires a committed database write. +File updates are serialized and persisted on a worker thread, so durable writes do not block the event loop or overwrite concurrent ownership changes. +File writes sync both the file contents and the parent directory after atomic replacement before ownership is acknowledged. +If directory sync fails after replacement, file-backed authorization is suspended and further file mutations require a successful reload and directory sync, preventing updates from overwriting ownership changes with stale cached data. In `hybrid` mode the database is preferred, with a successful durable JSON write as fallback if the database is unavailable. If neither configured destination can persist the claim, the stream returns an error without exposing a new response ID; an in-memory cache entry is not sufficient. diff --git a/src/agent.py b/src/agent.py index 8172039af..39ac4d0cb 100644 --- a/src/agent.py +++ b/src/agent.py @@ -161,6 +161,8 @@ async def store_conversation_thread(user_id: str, response_id: str, conversation # 2. Store only essential metadata to disk (simplified JSON) messages = conversation_state.get("messages", []) + # A restarted thread has durable metadata, but not successful message bodies. + persisted_metadata = conversation_state.get("_persisted_metadata", {}) first_user_msg = next((msg for msg in messages if msg.get("role") == "user"), None) title = "New Chat" if first_user_msg: @@ -170,17 +172,21 @@ async def store_conversation_thread(user_id: str, response_id: str, conversation user_assistant_messages = [msg for msg in messages if msg.get("role") in ["user", "assistant"]] metadata_only = { "response_id": response_id, - "title": title, + "title": persisted_metadata.get("title") or title, "endpoint": "langflow", "created_at": conversation_state.get("created_at"), "last_activity": conversation_state.get("last_activity"), - "previous_response_id": conversation_state.get("previous_response_id"), + "previous_response_id": persisted_metadata.get( + "previous_response_id", conversation_state.get("previous_response_id") + ), "filter_id": conversation_state.get("filter_id"), - "total_messages": len(user_assistant_messages), + "total_messages": (persisted_metadata.get("total_messages") or 0) + + len(user_assistant_messages), } # Snapshot error-thread turns: Langflow history is incomplete after fresh-session retries. - if any(msg.get("error") for msg in user_assistant_messages): - metadata_only["messages"] = [ + persisted_messages = persisted_metadata.get("messages", []) + if persisted_messages or any(msg.get("error") for msg in user_assistant_messages): + metadata_only["messages"] = persisted_messages + [ { "role": msg.get("role"), "content": msg.get("content"), @@ -910,10 +916,37 @@ async def async_langflow_chat_stream( thread_id, ) previous_loaded_from_memory = _conversation_thread_in_memory(user_id, thread_id) - conversation_state = get_conversation_thread(user_id, thread_id) + owned_session_ids = [] + persisted_metadata = {} + if thread_id and not previous_loaded_from_memory: + from services.session_ownership_service import session_ownership_service + owned_session_ids = await session_ownership_service.get_conversation_sessions( + user_id, thread_id + ) + if owned_session_ids: + thread_id = owned_session_ids[0] + previous_loaded_from_memory = _conversation_thread_in_memory(user_id, thread_id) + if not previous_loaded_from_memory: + persisted_metadata = await conversation_persistence.get_conversation_thread( + user_id, thread_id + ) from datetime import datetime + conversation_state = get_conversation_thread(user_id, thread_id) + if persisted_metadata: + conversation_state["_persisted_metadata"] = persisted_metadata + restored_created_at = persisted_metadata.get("created_at") + if isinstance(restored_created_at, str): + try: + restored_created_at = datetime.fromisoformat(restored_created_at) + except ValueError: + restored_created_at = None + if isinstance(restored_created_at, datetime): + conversation_state["created_at"] = restored_created_at + conversation_state["filter_id"] = persisted_metadata.get("filter_id") + conversation_state["response_ids"] = list(owned_session_ids) + conversation_state["messages"].append( {"role": "user", "content": prompt, "timestamp": datetime.now()} ) @@ -930,7 +963,8 @@ async def async_langflow_chat_stream( usage_data = None collected_chunks = [] # Store all chunks for function call data error_occurred = False - store_id = None + # A durable owned group can safely retain its canonical sidebar ID on restart. + store_id = thread_id if owned_session_ids else None claimed_ids: set[str] = set() turn_finished = False @@ -1127,9 +1161,14 @@ async def persist_stream_state(*session_ids, mint_if_missing=False): await upstream.aclose() -async def delete_user_conversation(user_id: str, response_id: str) -> bool: +async def delete_user_conversation( + user_id: str, response_id: str, *, session_ids: list[str] | None = None +) -> bool: """Delete a conversation for a user from both memory and persistent storage. + ``session_ids`` is the caller's already-authorized ownership snapshot, when + provided. An empty group permits only user-scoped local orphan cleanup. + Returns: True — conversation was found and deleted from at least one store. False — conversation did not exist in any store (confirmed not-found). @@ -1140,13 +1179,15 @@ async def delete_user_conversation(user_id: str, response_id: str) -> bool: """ from services.session_ownership_service import session_ownership_service - session_ids = await session_ownership_service.get_conversation_sessions(user_id, response_id) - if not session_ids: - return False - canonical_id = session_ids[0] - # Revoke ALL durable aliases before removing the local index. Failures must - # reach the caller, not turn into a successful delete with reusable IDs. - await session_ownership_service.release_conversation_sessions(user_id, session_ids) + if session_ids is None: + session_ids = await session_ownership_service.get_conversation_sessions( + user_id, response_id + ) + canonical_id = session_ids[0] if session_ids else response_id + if session_ids: + # Revoke only the snapshot used for upstream cleanup. Never revoke a late + # alias whose upstream deletion was not attempted. + await session_ownership_service.release_conversation_sessions(user_id, session_ids) deleted = await conversation_persistence.delete_conversation_thread(user_id, canonical_id) if canonical_id in active_conversations.get(user_id, {}): del active_conversations[user_id][canonical_id] diff --git a/src/db/repositories/session_ownership_repo.py b/src/db/repositories/session_ownership_repo.py index 44b6bbd52..64ef9d42c 100644 --- a/src/db/repositories/session_ownership_repo.py +++ b/src/db/repositories/session_ownership_repo.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime from sqlalchemy import select +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlmodel import col @@ -27,11 +28,6 @@ async def claim( canonical = await self.get(canonical_id) if canonical is None or canonical.user_id != user_id: raise PermissionError("Conversation is not owned by this user") - if existing is not None and ( - existing.user_id != user_id - or (existing.conversation_id is not None and existing.conversation_id != canonical_id) - ): - raise PermissionError("Session ownership or conversation conflict") if existing is None: row = SessionOwnership( response_id=response_id, @@ -40,9 +36,40 @@ async def claim( created_at=now, last_accessed=now, ) - self.session.add(row) - await self.session.flush() - return row + try: + # begin_nested flushes pending caller work before establishing + # its savepoint, so add only this claim inside the savepoint. + async with self.session.begin_nested(): + self.session.add(row) + await self.session.flush() + except IntegrityError as exc: + # Only primary-key/unique violations can be an INSERT race. + # Never hide another constraint failure behind a matching row. + if ( + getattr(exc.orig, "sqlstate", None) != "23505" + and getattr(exc.orig, "sqlite_errorname", None) + != "SQLITE_CONSTRAINT_PRIMARYKEY" + ): + raise + # The savepoint has rolled back only this insert; preserve the + # caller's transaction while reading the winning claim afresh. + existing = await self.session.get( + SessionOwnership, response_id, populate_existing=True + ) + if existing is None: + raise + if canonical_id != response_id: + canonical = await self.session.get( + SessionOwnership, canonical_id, populate_existing=True + ) + if canonical is None or canonical.user_id != user_id: + raise PermissionError("Conversation is not owned by this user") from None + else: + return row + if existing.user_id != user_id or ( + (existing.conversation_id or response_id) != canonical_id + ): + raise PermissionError("Session ownership or conversation conflict") existing.conversation_id = canonical_id existing.last_accessed = now self.session.add(existing) diff --git a/src/services/chat_service.py b/src/services/chat_service.py index 7b139cc48..d4df2d0e8 100644 --- a/src/services/chat_service.py +++ b/src/services/chat_service.py @@ -688,11 +688,12 @@ async def get_chat_history(self, user_id: str): if first_user_msg else "New chat" ) + persisted_metadata = conversation_state.get("_persisted_metadata", {}) conversations.append( { "response_id": response_id, - "title": title, + "title": persisted_metadata.get("title") or title, "endpoint": "chat", "messages": messages, "created_at": conversation_state.get("created_at").isoformat() @@ -703,7 +704,8 @@ async def get_chat_history(self, user_id: str): else None, "previous_response_id": conversation_state.get("previous_response_id"), "filter_id": conversation_state.get("filter_id"), - "total_messages": len(messages), + "total_messages": (persisted_metadata.get("total_messages") or 0) + + len(messages), "source": "in_memory", } ) @@ -948,7 +950,9 @@ async def delete_session(self, user_id: str, session_id: str): ] langflow_deleted = bool(upstream_results) and all(upstream_results) - local_deleted = await delete_user_conversation(user_id, session_id) + local_deleted = await delete_user_conversation( + user_id, session_id, session_ids=session_ids + ) if not local_deleted: return { @@ -981,7 +985,7 @@ async def delete_sessions(self, user_id: str, session_ids: list[str]) -> dict[st for session_id in dict.fromkeys(session_ids): try: owner = await session_ownership_service.get_session_owner(session_id) - if owner != user_id: + if owner is not None and owner != user_id: logger.warning("bulk_delete: %s not owned by %s", session_id, user_id) failed.append(session_id) continue diff --git a/src/services/session_ownership_service.py b/src/services/session_ownership_service.py index 6ecb2cae1..e9c3277ca 100644 --- a/src/services/session_ownership_service.py +++ b/src/services/session_ownership_service.py @@ -14,8 +14,10 @@ from __future__ import annotations +import asyncio import json import os +import stat import tempfile from collections.abc import Callable from datetime import UTC, datetime @@ -32,6 +34,10 @@ logger = get_logger(__name__) +class _IndeterminateOwnershipWrite(OSError): + """The replacement is visible, but its directory entry was not synced.""" + + class SessionOwnershipService: """Tracks which user owns which session.""" @@ -39,6 +45,8 @@ def __init__(self, session_factory: Callable | None = None): self.ownership_file = get_data_file("session_ownership.json") os.makedirs(os.path.dirname(self.ownership_file), exist_ok=True) self._session_factory = session_factory + self._file_lock = asyncio.Lock() + self._file_state_indeterminate = False # JSON cache — eagerly loaded so `files` and `hybrid` modes # behave like the legacy implementation. self.ownership_data: dict[str, dict[str, Any]] = self._load_ownership_data() @@ -58,22 +66,75 @@ def _load_ownership_data(self) -> dict[str, dict[str, Any]]: return {} def _save_ownership_data(self, data: dict[str, dict[str, Any]]) -> None: - """Publish a complete file before changing the authorization cache.""" + """Atomically replace and sync file + directory on the caller's worker thread.""" temporary = None try: + try: + mode = stat.S_IMODE(os.stat(self.ownership_file).st_mode) + except FileNotFoundError: + mode = None with tempfile.NamedTemporaryFile( mode="w", dir=os.path.dirname(self.ownership_file), delete=False ) as f: temporary = f.name + if mode is not None: + os.fchmod(f.fileno(), mode) json.dump(data, f, indent=2) f.flush() os.fsync(f.fileno()) os.replace(temporary, self.ownership_file) - self.ownership_data = data + temporary = None + # Persist the replacement directory entry before publishing ownership. + try: + directory_fd = os.open(os.path.dirname(self.ownership_file), os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + except OSError as exc: + raise _IndeterminateOwnershipWrite(str(exc)) from exc finally: if temporary is not None and os.path.exists(temporary): os.unlink(temporary) + def _read_and_sync_ownership_data(self) -> dict[str, dict[str, Any]]: + """Strictly reload an uncertain replacement and make its name durable.""" + with open(self.ownership_file) as f: + data = json.load(f) + directory_fd = os.open(os.path.dirname(self.ownership_file), os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + return data + + async def _reconcile_ownership_data(self) -> None: + """Reconcile before mutation, under ``_file_lock``; failures stay closed.""" + if self._file_state_indeterminate: + data = await asyncio.to_thread(self._read_and_sync_ownership_data) + self.ownership_data = data + self._file_state_indeterminate = False + + async def _persist_ownership_data(self, data: dict[str, dict[str, Any]]) -> None: + """Persist then publish while the caller holds ``_file_lock``.""" + write = asyncio.create_task(asyncio.to_thread(self._save_ownership_data, data)) + cancelled = False + try: + while not write.done(): + try: + await asyncio.shield(write) + except asyncio.CancelledError: + # Keep the lock until the worker finishes, even on cancellation. + cancelled = True + write.result() + except _IndeterminateOwnershipWrite: + # A failed sync cannot undo an already completed replacement. + self._file_state_indeterminate = True + raise + self.ownership_data = data + if cancelled: + raise asyncio.CancelledError + # ------------------------------------------------------------------ # Public async API # ------------------------------------------------------------------ @@ -88,8 +149,15 @@ async def claim_session( An unsuccessful write never adds an authorization entry to the cache. """ canonical_id = conversation_id or session_id - if file_writes_enabled() and self._file_claim_conflicts(user_id, session_id, canonical_id): - return False + if file_writes_enabled(): + async with self._file_lock: + try: + await self._reconcile_ownership_data() + except Exception as exc: + logger.error("File ownership reconciliation failed", error=str(exc)) + return False + if self._file_claim_conflicts(user_id, session_id, canonical_id): + return False durable = False if db_writes_enabled(): try: @@ -97,30 +165,36 @@ async def claim_session( except PermissionError: return False if file_writes_enabled(): - # Another claim may have published while the DB write was awaiting. - if self._file_claim_conflicts(user_id, session_id, canonical_id): - return False - now = datetime.now(UTC).isoformat() - existing = self.ownership_data.get(session_id, {}) - if canonical_id != session_id and not durable: - canonical = self.ownership_data.get(canonical_id) - if not canonical or canonical.get("user_id") != user_id: + async with self._file_lock: + try: + await self._reconcile_ownership_data() + except Exception as exc: + logger.error("File ownership reconciliation failed", error=str(exc)) + return durable + # Recheck after DB work and after any pending file publication. + if self._file_claim_conflicts(user_id, session_id, canonical_id): return False - updated = { - **self.ownership_data, - session_id: { - **existing, - "user_id": user_id, - "conversation_id": canonical_id, - "created_at": existing.get("created_at", now), - "last_accessed": now, - }, - } - try: - self._save_ownership_data(updated) - durable = True - except Exception as exc: # noqa: BLE001 - logger.error("File claim_session failed", error=str(exc)) + now = datetime.now(UTC).isoformat() + existing = self.ownership_data.get(session_id, {}) + if canonical_id != session_id and not durable: + canonical = self.ownership_data.get(canonical_id) + if not canonical or canonical.get("user_id") != user_id: + return False + updated = { + **self.ownership_data, + session_id: { + **existing, + "user_id": user_id, + "conversation_id": canonical_id, + "created_at": existing.get("created_at", now), + "last_accessed": now, + }, + } + try: + await self._persist_ownership_data(updated) + durable = True + except Exception as exc: # noqa: BLE001 + logger.error("File claim_session failed", error=str(exc)) return durable async def get_session_owner(self, session_id: str) -> str | None: @@ -132,6 +206,8 @@ async def get_session_owner(self, session_id: str) -> str | None: if mode == "db": return None # files or hybrid-with-no-db-row → fall back to JSON + if self._file_state_indeterminate: + return None data = self.ownership_data.get(session_id) return data.get("user_id") if data else None @@ -143,10 +219,14 @@ async def get_user_sessions(self, user_id: str) -> list[str]: return db_sessions # hybrid: union with JSON-only entries json_sessions = [ - sid for sid, data in self.ownership_data.items() if data.get("user_id") == user_id + sid + for sid, data in self.ownership_data.items() + if not self._file_state_indeterminate and data.get("user_id") == user_id ] return list(dict.fromkeys(db_sessions + json_sessions)) # files mode + if self._file_state_indeterminate: + return [] return [sid for sid, data in self.ownership_data.items() if data.get("user_id") == user_id] async def is_session_owned_by_user(self, session_id: str, user_id: str) -> bool: @@ -163,6 +243,8 @@ async def get_conversation_sessions(self, user_id: str, session_id: str) -> list durable aliases authorized while reporting successful removal. """ mode = get_storage_mode() + if mode != "db" and self._file_state_indeterminate: + raise RuntimeError("File ownership state requires reconciliation") files = self.ownership_data if mode != "db" else {} data = files.get(session_id, {}) canonical_id = ( @@ -222,11 +304,13 @@ async def release_conversation_sessions(self, user_id: str, session_ids: list[st hybrid groups can contain JSON-only aliases absent from a DB rollback. """ if file_writes_enabled(): - updated = dict(self.ownership_data) - for session_id in session_ids: - if updated.get(session_id, {}).get("user_id") == user_id: - del updated[session_id] - self._save_ownership_data(updated) + async with self._file_lock: + await self._reconcile_ownership_data() + updated = dict(self.ownership_data) + for session_id in session_ids: + if updated.get(session_id, {}).get("user_id") == user_id: + del updated[session_id] + await self._persist_ownership_data(updated) if db_writes_enabled(): from db.repositories import SessionOwnershipRepo @@ -244,16 +328,20 @@ async def release_conversation_sessions(self, user_id: str, session_ids: list[st async def release_session(self, user_id: str, session_id: str) -> bool: released = False - if file_writes_enabled() and session_id in self.ownership_data: - if self.ownership_data[session_id].get("user_id") == user_id: - updated = dict(self.ownership_data) - del updated[session_id] - self._save_ownership_data(updated) - released = True - else: - logger.warning( - f"User {user_id} tried to release session {session_id} they don't own (json)" - ) + if file_writes_enabled(): + async with self._file_lock: + await self._reconcile_ownership_data() + existing = self.ownership_data.get(session_id) + if existing is not None: + if existing.get("user_id") == user_id: + updated = dict(self.ownership_data) + del updated[session_id] + await self._persist_ownership_data(updated) + released = True + else: + logger.warning( + f"User {user_id} tried to release session {session_id} they don't own (json)" + ) if db_writes_enabled(): db_released = await self._db_release(session_id, user_id) @@ -319,7 +407,7 @@ def _file_claim_conflicts(self, user_id: str, session_id: str, canonical_id: str existing and ( existing.get("user_id") != user_id - or (existing.get("conversation_id") or canonical_id) != canonical_id + or (existing.get("conversation_id") or session_id) != canonical_id ) or canonical and canonical.get("user_id") != user_id @@ -350,10 +438,10 @@ async def _db_claim(self, user_id: str, session_id: str, conversation_id: str) - return True except PermissionError: raise - except IntegrityError as exc: - # A competing insert may now own the ID. Never turn a uniqueness - # conflict into permission to publish a JSON fallback for this user. - raise PermissionError("Concurrent session ownership conflict") from exc + except IntegrityError: + # The repository recovers only a confirmed, compatible INSERT + # winner. Any other integrity failure must not authorize fallback. + raise except Exception as exc: # noqa: BLE001 logger.error("DB claim_session failed", error=str(exc)) return False diff --git a/tests/unit/db/repositories/test_session_ownership_repo_savepoint.py b/tests/unit/db/repositories/test_session_ownership_repo_savepoint.py new file mode 100644 index 000000000..e84bc108a --- /dev/null +++ b/tests/unit/db/repositories/test_session_ownership_repo_savepoint.py @@ -0,0 +1,115 @@ +"""Claim-race recovery must not discard the caller's transaction.""" + +import sys +from pathlib import Path + +import pytest +import pytest_asyncio +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlmodel import SQLModel + +ROOT = Path(__file__).resolve().parents[4] +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +import db.models # noqa: E402,F401 +from db.models import SessionOwnership # noqa: E402 +from db.repositories.session_ownership_repo import SessionOwnershipRepo # noqa: E402 + + +@pytest_asyncio.fixture +async def session_factory(): + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(SQLModel.metadata.create_all) + yield async_sessionmaker(engine, expire_on_commit=False) + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "winner_user,winner_canonical,accepted", + [("alice", "canonical", True), ("bob", "canonical", False), ("alice", "other", False)], +) +async def test_alias_race_preserves_pending_canonical_claim( + monkeypatch, session_factory, winner_user, winner_canonical, accepted +): + async with session_factory() as session: + session.add( + SessionOwnership( + response_id="alias", user_id=winner_user, conversation_id=winner_canonical + ) + ) + await session.commit() + + async with session_factory() as session: + repo = SessionOwnershipRepo(session) + canonical = await repo.claim("canonical", "alice") + # Keep additional unflushed work pending as begin_nested is entered. + canonical.last_accessed = canonical.created_at + original_get = repo.get + stale_read = False + + async def stale_get(response_id): + nonlocal stale_read + if response_id == "alias" and not stale_read: + stale_read = True + return None + return await original_get(response_id) + + monkeypatch.setattr(repo, "get", stale_get) + if accepted: + row = await repo.claim("alias", "alice", "canonical") + assert row.conversation_id == "canonical" + else: + with pytest.raises(PermissionError, match="ownership or conversation conflict"): + await repo.claim("alias", "alice", "canonical") + assert session.in_transaction() + assert (await original_get("canonical")).user_id == "alice" + await session.commit() + + async with session_factory() as session: + repo = SessionOwnershipRepo(session) + assert (await repo.get("canonical")).user_id == "alice" + winner = await repo.get("alias") + assert (winner.user_id, winner.conversation_id) == (winner_user, winner_canonical) + + +@pytest.mark.asyncio +async def test_nonunique_constraint_is_not_recovered_as_matching_winner( + monkeypatch, session_factory +): + async with session_factory() as session: + repo = SessionOwnershipRepo(session) + await repo.claim("canonical", "alice") + await repo.claim("alias", "alice", "canonical") + await session.execute( + text( + "CREATE TRIGGER reject_alias BEFORE INSERT ON session_ownership " + "WHEN NEW.response_id = 'alias' BEGIN " + "SELECT RAISE(ABORT, 'unrelated ownership constraint'); END" + ) + ) + await session.commit() + + async with session_factory() as session: + repo = SessionOwnershipRepo(session) + await repo.claim("prior", "alice") + original_get = repo.get + + async def stale_get(response_id): + if response_id == "alias": + return None + return await original_get(response_id) + + monkeypatch.setattr(repo, "get", stale_get) + with pytest.raises(IntegrityError, match="unrelated ownership constraint"): + await repo.claim("alias", "alice", "canonical") + assert session.is_active + await session.commit() + + async with session_factory() as session: + assert (await SessionOwnershipRepo(session).get("prior")).user_id == "alice" diff --git a/tests/unit/services/test_chat_session_lifecycle.py b/tests/unit/services/test_chat_session_lifecycle.py new file mode 100644 index 000000000..f3750fbac --- /dev/null +++ b/tests/unit/services/test_chat_session_lifecycle.py @@ -0,0 +1,765 @@ +"""Durable chat ownership, continuation, cancellation, and deletion regressions.""" + +from __future__ import annotations + +import asyncio +import json +from datetime import datetime +from pathlib import Path +from types import SimpleNamespace + +import anyio +import httpx +import pytest +import pytest_asyncio +from fastapi import HTTPException +from openai import AsyncOpenAI +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlmodel import SQLModel + +import agent as agent_module +from api.chat import ChatBody, _assert_owns, delete_session_endpoint, langflow_endpoint +from services.chat_service import ChatService +from services.conversation_persistence_service import ConversationPersistenceService +from services.session_ownership_service import SessionOwnershipService + + +async def _collect_error_chunks(stream) -> list[dict]: + return [json.loads(raw.decode("utf-8")) async for raw in stream] + + +@pytest_asyncio.fixture(params=["db", "hybrid", "files"]) +async def durable_stream_stores(request, monkeypatch, tmp_path): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", request.param) + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'chat.db'}") + async with engine.begin() as connection: + await connection.run_sync(SQLModel.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + monkeypatch.setattr( + "services.session_ownership_service.get_data_file", + lambda name: str(tmp_path / name), + ) + ownership = SessionOwnershipService(session_factory=factory) + persistence = ConversationPersistenceService( + storage_file=str(tmp_path / "conversations.json"), session_factory=factory + ) + monkeypatch.setattr("services.session_ownership_service.session_ownership_service", ownership) + monkeypatch.setattr(agent_module, "conversation_persistence", persistence) + monkeypatch.setattr(agent_module, "active_conversations", {}) + yield ownership, persistence, factory + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("abort", ["close", "cancel"]) +async def test_aborted_langflow_stream_keeps_owned_id_and_partial_turn( + monkeypatch, durable_stream_stores, abort +): + ownership, persistence, factory = durable_stream_stores + waiting = asyncio.Event() + upstream_closed = asyncio.Event() + calls = [] + + async def upstream(*_args, **kwargs): + calls.append(kwargs.get("previous_response_id")) + try: + yield b'{"id":"resp-first","delta":{"content":"Partial answer"}}\n' + waiting.set() + await asyncio.Event().wait() + finally: + upstream_closed.set() + + monkeypatch.setattr(agent_module, "async_stream", upstream) + stream = agent_module.async_langflow_chat_stream(object(), "flow", "First question", "alice") + chunk = json.loads(await anext(stream)) + # The ID must already be durable when it first crosses the generator boundary. + await _assert_owns(chunk["id"], "alice") + assert (await persistence.get_conversation_thread("alice", chunk["id"]))["total_messages"] >= 1 + if abort == "close": + await stream.aclose() + else: + pending = asyncio.create_task(anext(stream)) + await waiting.wait() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert upstream_closed.is_set() + # Reload ownership rather than relying on the live service's cache. + reloaded = SessionOwnershipService(session_factory=factory) + assert await reloaded.get_session_owner(chunk["id"]) == "alice" + with pytest.raises(HTTPException) as forbidden: + await _assert_owns(chunk["id"], "mallory") + assert forbidden.value.status_code == 403 + with pytest.raises(HTTPException) as missing: + await _assert_owns("unknown-session", "alice") + assert missing.value.status_code == 404 + + async def continuation(*_args, **kwargs): + calls.append(kwargs.get("previous_response_id")) + yield b'{"id":"resp-next","response_id":"resp-next-alias","delta":{"content":"Next"}}\n' + + monkeypatch.setattr(agent_module, "async_stream", continuation) + await _assert_owns(chunk["id"], "alice") + followup = agent_module.async_langflow_chat_stream( + object(), + "flow", + "Follow-up", + "alice", + previous_response_id=chunk["id"], + conversation_id=chunk["id"], + ) + next_chunk = json.loads(await anext(followup)) + await _assert_owns(next_chunk["id"], "alice") + await _assert_owns(next_chunk["response_id"], "alice") + assert await _collect_error_chunks(followup) == [] + assert calls == [None, "resp-first"] # Never replay the interrupted prompt. + conversations = await persistence.get_user_conversations("alice") + assert list(conversations) == ["resp-first"] # Stable sidebar id, changing response ids. + messages = agent_module.active_conversations["alice"]["resp-first"]["messages"] + assert [(m["role"], m["content"]) for m in messages if m["role"] != "system"] == [ + ("user", "First question"), + ("assistant", "Partial answer"), + ("user", "Follow-up"), + ("assistant", "Next"), + ] + assert conversations["resp-first"]["total_messages"] == 4 + + # SDK callers may only carry the changing response ID, not conversation_id. + await _collect_error_chunks( + agent_module.async_langflow_chat_stream( + object(), + "flow", + "Third question", + "alice", + previous_response_id=next_chunk["response_id"], + ) + ) + assert list(await persistence.get_user_conversations("alice")) == ["resp-first"] + assert calls == [None, "resp-first", "resp-next-alias"] + assert [m["content"] for m in messages if m["role"] == "user"] == [ + "First question", + "Follow-up", + "Third question", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("claim_failure", ["unavailable", "foreign"]) +async def test_stream_does_not_expose_unowned_response_id( + monkeypatch, durable_stream_stores, claim_failure +): + ownership, persistence, _ = durable_stream_stores + if claim_failure == "foreign": + await ownership.claim_session("mallory", "resp-unclaimable") + else: + + async def unavailable(*_args, **_kwargs): + return False + + monkeypatch.setattr(ownership, "claim_session", unavailable) + + async def upstream(*_args, **_kwargs): + yield b'{"id":"resp-unclaimable","delta":{"content":"Not exposed"}}\n' + + monkeypatch.setattr(agent_module, "async_stream", upstream) + chunks = await _collect_error_chunks( + agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") + ) + assert len(chunks) == 1 + assert chunks[0]["status"] == "failed" + assert "id" not in chunks[0] and "response_id" not in chunks[0] + assert await persistence.get_user_conversations("alice") == {} + assert await ownership.get_session_owner("resp-unclaimable") == ( + "mallory" if claim_failure == "foreign" else None + ) + + +@pytest.mark.asyncio +async def test_local_asgi_disconnect_smoke_preserves_followup(monkeypatch, durable_stream_stores): + """Real endpoint → generator → OpenAI SSE adapter → durable stores, locally.""" + _, persistence, factory = durable_stream_stores + disconnected = anyio.Event() + upstream_closed = anyio.Event() + calls = [] + + class ProviderStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield ( + b"event: response.created\n" + b'data: {"type":"response.created","response":{"id":"resp-http"}}\n\n' + ) + yield ( + b"event: response.output_text.delta\n" + b'data: {"type":"response.output_text.delta","id":"resp-http",' + b'"delta":{"content":"Partial HTTP answer"}}\n\n' + ) + await anyio.sleep_forever() + + async def aclose(self): + upstream_closed.set() + + async def provider(request): + payload = json.loads(request.content) + calls.append(payload) + if len(calls) == 1: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=ProviderStream() + ) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b"event: response.output_text.delta\n" + b'data: {"type":"response.output_text.delta","id":"resp-http-next",' + b'"delta":{"content":"Follow-up HTTP answer"}}\n\n' + b"event: response.completed\n" + b'data: {"type":"response.completed","response":{"id":"resp-http-next",' + b'"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7}}}\n\n' + ), + ) + + async with AsyncOpenAI( + api_key="local-test-only", + base_url="https://provider.invalid/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(provider)), + max_retries=0, + ) as client: + + class EndpointChatService: + async def langflow_chat(self, prompt, _subject, _jwt, **kwargs): + return agent_module.async_langflow_chat_stream( + client, + "flow", + prompt, + kwargs["storage_user_id"], + previous_response_id=kwargs["previous_response_id"], + conversation_id=kwargs["conversation_id"], + ) + + user = SimpleNamespace( + user_id="oauth-alice", + db_user_id="alice", + jwt_token="local-test-only", + name="Alice", + email="alice@example.invalid", + ) + service = EndpointChatService() + response = await langflow_endpoint( + ChatBody(prompt="First HTTP question", stream=True), + chat_service=service, + session_manager=None, + user=user, + ) + emitted = [] + + async def send(message): + if message["type"] == "http.response.body" and message.get("body"): + chunk = json.loads(message["body"]) + emitted.append(chunk) + response_id = chunk.get("id") or (chunk.get("response") or {}).get("id") + if response_id: + await _assert_owns(response_id, "alice") + if chunk.get("delta", {}).get("content"): + disconnected.set() + + async def receive(): + await disconnected.wait() + return {"type": "http.disconnect"} + + # Starlette's real disconnect listener cancels the body iterator using an + # AnyIO cancel scope, rather than manually throwing into our generator. + with anyio.fail_after(5): + await response({"type": "http", "asgi": {"spec_version": "2.3"}}, receive, send) + assert upstream_closed.is_set() + assert len(emitted) == 2 + metadata = await persistence.get_conversation_thread("alice", "resp-http") + assert metadata["total_messages"] == 2 + + followup = await langflow_endpoint( + ChatBody( + prompt="Next HTTP question", + stream=True, + previous_response_id="resp-http", + conversation_id="resp-http", + ), + chat_service=service, + session_manager=None, + user=user, + ) + chunks = await _collect_error_chunks(followup.body_iterator) + await _assert_owns("resp-http-next", "alice") + assert chunks[-1]["response"]["usage"]["total_tokens"] == 7 + assert calls[1]["previous_response_id"] == "resp-http" + assert [call["input"] for call in calls] == ["First HTTP question", "Next HTTP question"] + conversations = await persistence.get_user_conversations("alice") + assert list(conversations) == ["resp-http"] + assert conversations["resp-http"]["total_messages"] == 4 + messages = agent_module.active_conversations["alice"]["resp-http"]["messages"] + assert [m["content"] for m in messages if m["role"] != "system"] == [ + "First HTTP question", + "Partial HTTP answer", + "Next HTTP question", + "Follow-up HTTP answer", + ] + + # Restart before canonical deletion: aliases must come from durable state. + reloaded = SessionOwnershipService(session_factory=factory) + monkeypatch.setattr( + "services.session_ownership_service.session_ownership_service", reloaded + ) + monkeypatch.setattr(agent_module, "active_conversations", {}) + upstream_deletes = [] + + async def monitor_delete(method, path): + assert method == "DELETE" + upstream_deletes.append(path.rsplit("/", 1)[-1]) + return httpx.Response(204) + + monkeypatch.setattr("services.chat_service.clients.langflow_request", monitor_delete) + deletion = await delete_session_endpoint( + session_id="resp-http", chat_service=ChatService(), user=user + ) + assert deletion.status_code == 200 + assert set(upstream_deletes) == {"resp-http", "resp-http-next"} + fresh = SessionOwnershipService(session_factory=factory) + monkeypatch.setattr("services.session_ownership_service.session_ownership_service", fresh) + with pytest.raises(HTTPException) as missing: + await langflow_endpoint( + ChatBody( + prompt="Cannot resurrect", stream=True, previous_response_id="resp-http-next" + ), + chat_service=service, + session_manager=None, + user=user, + ) + assert missing.value.status_code == 404 + assert len(calls) == 2 + + # Actual endpoint + SSE adapter with both persistent destinations broken. + ownership_path = Path(fresh.ownership_file) + if ownership_path.exists(): + ownership_path.unlink() + ownership_path.mkdir() + + def unavailable_database(): + raise OSError("local smoke database unavailable") + + monkeypatch.setattr(fresh, "_resolve_session_factory", lambda: unavailable_database) + failed = await langflow_endpoint( + ChatBody(prompt="Fail closed", stream=True), + chat_service=service, + session_manager=None, + user=user, + ) + failed_chunks = await _collect_error_chunks(failed.body_iterator) + assert len(failed_chunks) == 1 and failed_chunks[0]["status"] == "failed" + assert "id" not in failed_chunks[0] and "response_id" not in failed_chunks[0] + assert ( + await SessionOwnershipService(session_factory=factory).get_session_owner( + "resp-http-next" + ) + is None + ) + print( + "LOCAL SMOKE: disconnect/follow-up retained 4 messages; restart + " + "canonical deletion removed both upstream sessions and alias authorization; " + "persistent failures emitted no client ID" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delta", [None, {"content": None}]) +async def test_metadata_only_chunk_claims_id_before_exposure( + monkeypatch, durable_stream_stores, delta +): + async def upstream(*_args, **_kwargs): + yield (json.dumps({"id": "resp-metadata", "delta": delta}) + "\n").encode() + + monkeypatch.setattr(agent_module, "async_stream", upstream) + stream = agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") + chunk = json.loads(await anext(stream)) + await _assert_owns(chunk["id"], "alice") + await stream.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delete_id", ["resp-first", "resp-second"]) +async def test_delete_canonical_revokes_aliases_after_reload( + monkeypatch, durable_stream_stores, delete_id +): + ownership, persistence, factory = durable_stream_stores + + async def upstream(*_args, **kwargs): + response_id = "resp-second" if kwargs.get("previous_response_id") else "resp-first" + yield (json.dumps({"id": response_id, "delta": {"content": "Answer"}}) + "\n").encode() + + monkeypatch.setattr(agent_module, "async_stream", upstream) + for previous in (None, "resp-first"): + await _collect_error_chunks( + agent_module.async_langflow_chat_stream( + object(), "flow", "Question", "alice", previous_response_id=previous + ) + ) + await ownership.claim_session("mallory", "resp-foreign") + await ownership.claim_session("alice", "resp-unrelated") + assert list(await persistence.get_user_conversations("alice")) == ["resp-first"] + assert await agent_module.delete_user_conversation("mallory", "resp-first") is False + assert agent_module.active_conversations["alice"]["resp-first"]["messages"] + # Restart: neither ownership cache nor the in-memory response_ids can help. + reloaded = SessionOwnershipService(session_factory=factory) + monkeypatch.setattr("services.session_ownership_service.session_ownership_service", reloaded) + monkeypatch.setattr(agent_module, "active_conversations", {}) + persistence = ConversationPersistenceService( + storage_file=persistence.storage_file, session_factory=factory + ) + monkeypatch.setattr(agent_module, "conversation_persistence", persistence) + assert await agent_module.delete_user_conversation("alice", delete_id) is True + fresh = SessionOwnershipService(session_factory=factory) + monkeypatch.setattr("services.session_ownership_service.session_ownership_service", fresh) + for response_id in ("resp-first", "resp-second"): + with pytest.raises(HTTPException) as missing: + await _assert_owns(response_id, "alice") + assert missing.value.status_code == 404 + assert await fresh.get_session_owner("resp-foreign") == "mallory" + assert await fresh.get_session_owner("resp-unrelated") == "alice" + assert await persistence.get_user_conversations("alice") == {} + + +@pytest.mark.asyncio +async def test_persistent_claim_failure_never_exposes_cached_id(monkeypatch, durable_stream_stores): + ownership, persistence, factory = durable_stream_stores + # A real IO failure, not a stubbed claim: cache echo must not prove durability. + Path(ownership.ownership_file).mkdir() + + def unavailable_database(): + raise OSError("local database unavailable") + + monkeypatch.setattr(ownership, "_resolve_session_factory", lambda: unavailable_database) + + async def upstream(*_args, **_kwargs): + yield b'{"id":"resp-file-failure","delta":{"content":"Not exposed"}}\n' + + monkeypatch.setattr(agent_module, "async_stream", upstream) + chunks = await _collect_error_chunks( + agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") + ) + assert len(chunks) == 1 + assert chunks[0]["status"] == "failed" + assert "id" not in chunks[0] and "response_id" not in chunks[0] + assert await ownership.get_session_owner("resp-file-failure") is None + assert ( + await SessionOwnershipService(session_factory=factory).get_session_owner( + "resp-file-failure" + ) + is None + ) + assert await persistence.get_user_conversations("alice") == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("durable_stream_stores", ["hybrid", "files"], indirect=True) +async def test_failed_file_revocation_is_not_reported_as_success( + monkeypatch, durable_stream_stores +): + ownership, _, _ = durable_stream_stores + + async def upstream(*_args, **_kwargs): + yield b'{"id":"resp-delete-failure","delta":{"content":"Answer"}}\n' + + monkeypatch.setattr(agent_module, "async_stream", upstream) + await _collect_error_chunks( + agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") + ) + ownership_path = Path(ownership.ownership_file) + ownership_path.unlink() + ownership_path.mkdir() + + async def monitor_delete(_self, _session_id): + return True + + monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) + result = await ChatService().delete_session("alice", "resp-delete-failure") + assert result["success"] is False + await _assert_owns("resp-delete-failure", "alice") + assert "resp-delete-failure" in agent_module.active_conversations["alice"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("durable_stream_stores", ["hybrid"], indirect=True) +async def test_json_only_alias_cleanup_precedes_failed_db_revocation( + monkeypatch, durable_stream_stores +): + from db.repositories import SessionOwnershipRepo + + ownership, _, factory = durable_stream_stores + + async def upstream(*_args, **kwargs): + response_id = "alias" if kwargs.get("previous_response_id") else "canonical" + yield (json.dumps({"id": response_id, "delta": {"content": "Answer"}}) + "\n").encode() + + monkeypatch.setattr(agent_module, "async_stream", upstream) + await _collect_error_chunks( + agent_module.async_langflow_chat_stream(object(), "flow", "First", "alice") + ) + + def unavailable_database(): + raise OSError("database temporarily unavailable") + + ownership._session_factory = unavailable_database + await _collect_error_chunks( + agent_module.async_langflow_chat_stream( + object(), "flow", "Next", "alice", previous_response_id="canonical" + ) + ) + ownership._session_factory = factory + async with factory() as session: + assert await SessionOwnershipRepo(session).get("alias") is None + assert await ownership.get_session_owner("alias") == "alice" + original = SessionOwnershipRepo.release + + async def failed_release(*_args): + raise OSError("database revocation failed") + + upstream_deletes = [] + + async def monitor_delete(_self, session_id): + upstream_deletes.append(session_id) + return True + + monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) + monkeypatch.setattr(SessionOwnershipRepo, "release", failed_release) + result = await ChatService().delete_session("alice", "canonical") + assert result["success"] is False + assert set(upstream_deletes) == {"canonical", "alias"} + # Even though JSON-only ownership was revoked, its upstream cleanup was not + # lost. The surviving canonical DB row can be retried after a restart. + reloaded = SessionOwnershipService(session_factory=factory) + monkeypatch.setattr("services.session_ownership_service.session_ownership_service", reloaded) + monkeypatch.setattr(SessionOwnershipRepo, "release", original) + assert (await ChatService().delete_session("alice", "canonical"))["success"] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("storage", ["persisted", "memory"]) +@pytest.mark.parametrize("operation", ["direct", "single", "bulk"]) +async def test_delete_own_orphan_without_upstream_authorization( + monkeypatch, durable_stream_stores, storage, operation +): + ownership, persistence, _ = durable_stream_stores + state = {"messages": [], "created_at": datetime.now()} + if storage == "persisted": + await persistence.store_conversation_thread("alice", "orphan", state) + else: + agent_module.active_conversations["alice"] = {"orphan": state} + assert await ownership.get_session_owner("orphan") is None + upstream_deletes = [] + + async def monitor_delete(_self, session_id): + upstream_deletes.append(session_id) + return True + + monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) + if operation == "direct": + assert await agent_module.delete_user_conversation("alice", "orphan") is True + elif operation == "single": + assert (await ChatService().delete_session("alice", "orphan"))["success"] is True + else: + assert await ChatService().delete_sessions("alice", ["orphan"]) == { + "deleted": ["orphan"], + "failed": [], + } + assert await persistence.get_conversation_thread("alice", "orphan") == {} + assert "orphan" not in agent_module.active_conversations.get("alice", {}) + assert upstream_deletes == [] + with pytest.raises(HTTPException) as missing: + await _assert_owns("orphan", "alice") + assert missing.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_delete_all_removes_persisted_orphans_only_for_caller( + monkeypatch, durable_stream_stores +): + _, persistence, _ = durable_stream_stores + state = {"messages": [], "created_at": datetime.now()} + await persistence.store_conversation_thread("alice", "orphan", state) + await persistence.store_conversation_thread("mallory", "foreign", state) + upstream_deletes = [] + + async def monitor_delete(_self, session_id): + upstream_deletes.append(session_id) + return True + + monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) + assert await ChatService().delete_all_user_sessions("alice") == { + "success": True, + "deleted_count": 1, + "total_count": 1, + } + assert await persistence.get_user_conversations("alice") == {} + assert await persistence.get_conversation_thread("mallory", "foreign") + assert upstream_deletes == [] + + +@pytest.mark.asyncio +async def test_missing_and_foreign_deletes_never_touch_upstream(monkeypatch, durable_stream_stores): + ownership, persistence, _ = durable_stream_stores + state = {"messages": [], "created_at": datetime.now()} + await persistence.store_conversation_thread("mallory", "foreign", state) + await ownership.claim_session("mallory", "foreign") + upstream_deletes = [] + + async def monitor_delete(_self, session_id): + upstream_deletes.append(session_id) + return True + + monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) + service = ChatService() + for session_id in ("missing", "foreign"): + assert await agent_module.delete_user_conversation("alice", session_id) is False + assert (await service.delete_session("alice", session_id))["not_found"] is True + assert await service.delete_sessions("alice", ["missing", "foreign"]) == { + "deleted": [], + "failed": ["missing", "foreign"], + } + assert await persistence.get_conversation_thread("mallory", "foreign") + assert await ownership.get_session_owner("foreign") == "mallory" + assert upstream_deletes == [] + user = SimpleNamespace(db_user_id="alice", user_id="oauth-alice") + for session_id, status in (("missing", 404), ("foreign", 403)): + with pytest.raises(HTTPException) as denied: + await delete_session_endpoint(session_id=session_id, chat_service=service, user=user) + assert denied.value.status_code == status + + +@pytest.mark.asyncio +async def test_delete_uses_one_owned_snapshot_for_upstream_and_revocation( + monkeypatch, durable_stream_stores +): + ownership, _, _ = durable_stream_stores + + async def upstream(*_args, **_kwargs): + yield b'{"id":"canonical","response_id":"alias","delta":{"content":"Answer"}}\n' + + monkeypatch.setattr(agent_module, "async_stream", upstream) + await _collect_error_chunks( + agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") + ) + resolve = ownership.get_conversation_sessions + resolutions = [] + + async def resolve_once(user_id, session_id): + resolutions.append(session_id) + assert len(resolutions) == 1, "Deletion resolved a different ownership snapshot" + return await resolve(user_id, session_id) + + upstream_deletes = [] + + async def monitor_delete(_self, session_id): + upstream_deletes.append(session_id) + if session_id == "canonical": + assert await ownership.claim_session("alice", "late-alias", "canonical") + return True + + monkeypatch.setattr(ownership, "get_conversation_sessions", resolve_once) + monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) + result = await ChatService().delete_session("alice", "alias") + assert result["success"] is True + assert resolutions == ["alias"] + assert set(upstream_deletes) == {"canonical", "alias"} + assert await ownership.get_session_owner("canonical") is None + assert await ownership.get_session_owner("alias") is None + # A claim arriving after the snapshot is not revoked without upstream cleanup. + assert await ownership.get_session_owner("late-alias") == "alice" + + +@pytest.mark.asyncio +async def test_restart_alias_only_continuation_keeps_durable_canonical( + monkeypatch, durable_stream_stores +): + _, persistence, factory = durable_stream_stores + calls = [] + + async def upstream(*_args, **kwargs): + calls.append(kwargs.get("previous_response_id")) + response_id = ("canonical", "alias", "after-restart")[len(calls) - 1] + yield (json.dumps({"id": response_id, "delta": {"content": "Answer"}}) + "\n").encode() + + monkeypatch.setattr(agent_module, "async_stream", upstream) + for previous in (None, "canonical"): + await _collect_error_chunks( + agent_module.async_langflow_chat_stream( + object(), + "flow", + "First question" if previous is None else "Follow-up", + "alice", + previous_response_id=previous, + filter_id="knowledge-filter" if previous is None else None, + ) + ) + before_restart = await persistence.get_conversation_thread("alice", "canonical") + assert before_restart["total_messages"] == 4 + reloaded = SessionOwnershipService(session_factory=factory) + monkeypatch.setattr("services.session_ownership_service.session_ownership_service", reloaded) + monkeypatch.setattr(agent_module, "active_conversations", {}) + reloaded_persistence = ConversationPersistenceService( + storage_file=persistence.storage_file, session_factory=factory + ) + monkeypatch.setattr(agent_module, "conversation_persistence", reloaded_persistence) + await _assert_owns("alias", "alice") + await _collect_error_chunks( + agent_module.async_langflow_chat_stream( + object(), "flow", "After restart", "alice", previous_response_id="alias" + ) + ) + assert calls == [None, "canonical", "alias"] + assert list(await reloaded_persistence.get_user_conversations("alice")) == ["canonical"] + after_restart = await reloaded_persistence.get_conversation_thread("alice", "canonical") + assert after_restart["title"] == before_restart["title"] == "First question" + assert after_restart["created_at"] == before_restart["created_at"] + assert after_restart["filter_id"] == "knowledge-filter" + assert after_restart["total_messages"] == 6 + history = await ChatService().get_chat_history("alice") + restored = next(c for c in history["conversations"] if c["response_id"] == "canonical") + assert restored["created_at"] == before_restart["created_at"] + assert restored["title"] == "First question" + assert restored["total_messages"] == 6 + fresh = SessionOwnershipService(session_factory=factory) + assert set(await fresh.get_conversation_sessions("alice", "after-restart")) == { + "canonical", + "alias", + "after-restart", + } + assert set(agent_module.active_conversations["alice"]) == {"canonical"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("durable_stream_stores", ["files"], indirect=True) +@pytest.mark.parametrize("created_at", [None, "not-a-date", 123]) +async def test_restart_invalid_creation_time_keeps_default( + monkeypatch, durable_stream_stores, created_at +): + ownership, persistence, _ = durable_stream_stores + await ownership.claim_session("alice", "canonical") + await ownership.claim_session("alice", "alias", conversation_id="canonical") + await persistence.store_conversation_thread( + "alice", + "canonical", + {"title": "Saved title", "created_at": created_at, "total_messages": 4}, + ) + + async def upstream(*_args, **_kwargs): + yield b'{"id":"next","delta":{"content":"Answer"}}\n' + + monkeypatch.setattr(agent_module, "async_stream", upstream) + started = datetime.now() + await _collect_error_chunks( + agent_module.async_langflow_chat_stream( + object(), "flow", "Follow-up", "alice", previous_response_id="alias" + ) + ) + history = await ChatService().get_chat_history("alice") + restored = next(c for c in history["conversations"] if c["response_id"] == "canonical") + assert started <= datetime.fromisoformat(restored["created_at"]) <= datetime.now() + assert restored["title"] == "Saved title" + assert restored["total_messages"] == 6 diff --git a/tests/unit/services/test_session_ownership_service_modes.py b/tests/unit/services/test_session_ownership_service_modes.py index 10f1ec2a1..af0f64034 100644 --- a/tests/unit/services/test_session_ownership_service_modes.py +++ b/tests/unit/services/test_session_ownership_service_modes.py @@ -1,7 +1,11 @@ """SessionOwnershipService — hybrid / db / files mode coverage.""" +import asyncio import json +import os +import stat import sys +import threading from pathlib import Path import pytest @@ -314,9 +318,12 @@ async def test_hybrid_losing_db_insert_cannot_publish_fallback_owner( # Model a competing claimant whose SELECT happened just before the winner's # insert. The ensuing flush hits the real DB primary-key constraint. original = SessionOwnershipRepo.get + stale_read = False async def stale_get(repo, response_id): - if response_id == "contended": + nonlocal stale_read + if response_id == "contended" and not stale_read: + stale_read = True return None return await original(repo, response_id) @@ -334,8 +341,6 @@ async def stale_get(repo, response_id): async def test_hybrid_rechecks_cache_after_awaiting_database( monkeypatch, tmp_data_dir, session_factory ): - import asyncio - monkeypatch.setenv("OPENRAG_STORAGE_MODE", "hybrid") svc = _svc(session_factory) waiting = asyncio.Event() @@ -356,3 +361,349 @@ async def delayed_claim(user_id, session_id, conversation_id): winner_published.set() assert await losing_claim is False assert json.loads(tmp_data_dir.read_text())["contended"]["user_id"] == "alice" + + +@pytest.mark.asyncio +async def test_file_fsync_keeps_event_loop_responsive(monkeypatch, tmp_data_dir): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "files") + svc = _svc(None) + heartbeat = threading.Event() + responsive = [] + original_fsync = os.fsync + + def delayed_fsync(fd): + responsive.append(heartbeat.wait(timeout=1)) + original_fsync(fd) + + monkeypatch.setattr(os, "fsync", delayed_fsync) + claim = asyncio.create_task(svc.claim_session("alice", "canonical")) + await asyncio.sleep(0) + heartbeat.set() + assert await claim is True + assert responsive and all(responsive), "fsync blocked the event-loop heartbeat" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mutation", ["claim", "release", "release_group"]) +async def test_concurrent_file_mutations_keep_every_committed_change( + monkeypatch, tmp_data_dir, mutation +): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "files") + svc = _svc(None) + assert await svc.claim_session("alice", "existing") is True + assert await svc.claim_session("alice", "victim") is True + entered = threading.Event() + finish = threading.Event() + original_fsync = os.fsync + + def pause_first_fsync(fd): + if not entered.is_set(): + entered.set() + assert finish.wait(timeout=2), "file persistence blocked the event loop" + original_fsync(fd) + + monkeypatch.setattr(os, "fsync", pause_first_fsync) + first = asyncio.create_task(svc.claim_session("alice", "first")) + try: + assert await asyncio.to_thread(entered.wait, 2) + if mutation == "claim": + operation = svc.claim_session("alice", "second") + elif mutation == "release": + operation = svc.release_session("alice", "victim") + else: + operation = svc.release_conversation_sessions("alice", ["victim"]) + second = asyncio.create_task(operation) + await asyncio.sleep(0) + finally: + finish.set() + first_result, _ = await asyncio.gather(first, second) + assert first_result is True + expected = ( + {"existing", "first", "victim", "second"} + if mutation == "claim" + else { + "existing", + "first", + } + ) + assert set(svc.ownership_data) == expected + assert set(json.loads(tmp_data_dir.read_text())) == expected + assert set(await _svc(None).get_user_sessions("alice")) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mutation", ["claim", "release", "release_group"]) +async def test_failed_file_mutation_preserves_durable_ownership( + monkeypatch, tmp_data_dir, mutation +): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "files") + svc = _svc(None) + assert await svc.claim_session("alice", "canonical") is True + before = tmp_data_dir.read_bytes() + + def fail_replace(*args): + raise OSError("atomic replacement unavailable") + + monkeypatch.setattr(os, "replace", fail_replace) + if mutation == "claim": + assert await svc.claim_session("alice", "alias", conversation_id="canonical") is False + else: + with pytest.raises(OSError, match="atomic replacement unavailable"): + if mutation == "release": + await svc.release_session("alice", "canonical") + else: + await svc.release_conversation_sessions("alice", ["canonical"]) + assert tmp_data_dir.read_bytes() == before + assert svc.ownership_data == json.loads(before) + assert await _svc(None).get_user_sessions("alice") == ["canonical"] + assert list(tmp_data_dir.parent.iterdir()) == [tmp_data_dir] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_fails", [False, True]) +async def test_file_claim_waits_for_directory_durability_before_publishing( + monkeypatch, tmp_data_dir, sync_fails +): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "files") + svc = _svc(None) + entered = threading.Event() + finish = threading.Event() + directory_fds = [] + original_fsync = os.fsync + + def pause_directory_fsync(fd): + if stat.S_ISDIR(os.fstat(fd).st_mode): + directory_fds.append(fd) + entered.set() + assert finish.wait(timeout=2), "directory sync blocked the event loop" + if sync_fails: + raise OSError("directory sync unavailable") + original_fsync(fd) + + monkeypatch.setattr(os, "fsync", pause_directory_fsync) + claim = asyncio.create_task(svc.claim_session("alice", "canonical")) + try: + assert await asyncio.to_thread(entered.wait, 2) + # Replacement alone is not a durable claim and must not authorize it. + assert json.loads(tmp_data_dir.read_text())["canonical"]["user_id"] == "alice" + assert not claim.done() + assert await svc.get_session_owner("canonical") is None + finally: + finish.set() + result = await claim + assert result is not sync_fails + assert await svc.get_session_owner("canonical") == (None if sync_fails else "alice") + assert len(directory_fds) == 1 + with pytest.raises(OSError): + os.fstat(directory_fds[0]) + assert list(tmp_data_dir.parent.iterdir()) == [tmp_data_dir] + + +@pytest.mark.asyncio +async def test_atomic_file_replacement_preserves_existing_permissions(monkeypatch, tmp_data_dir): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "files") + svc = _svc(None) + assert await svc.claim_session("alice", "canonical") is True + tmp_data_dir.chmod(0o640) + assert await svc.claim_session("alice", "alias", conversation_id="canonical") is True + assert stat.S_IMODE(tmp_data_dir.stat().st_mode) == 0o640 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["db", "hybrid"]) +@pytest.mark.parametrize( + ("winner_user", "winner_canonical", "accepted"), + [ + ("alice", "canonical", True), + ("bob", "foreign", False), + ("alice", "other", False), + ], +) +async def test_real_database_claim_race_checks_winning_owner_and_conversation( + monkeypatch, tmp_data_dir, tmp_path, mode, winner_user, winner_canonical, accepted +): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'ownership-race.db'}") + async with engine.begin() as connection: + await connection.run_sync(SQLModel.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "db") + setup = _svc(factory) + for user, canonical in (("alice", "canonical"), ("alice", "other"), ("bob", "foreign")): + assert await setup.claim_session(user, canonical) is True + selected_missing = asyncio.Event() + winner_committed = asyncio.Event() + original_get = SessionOwnershipRepo.get + + async def pause_after_missing_read(repo, response_id): + row = await original_get(repo, response_id) + if response_id == "contended" and row is None and not selected_missing.is_set(): + selected_missing.set() + await winner_committed.wait() + return row + + monkeypatch.setattr(SessionOwnershipRepo, "get", pause_after_missing_read) + loser = _svc(factory) + monkeypatch.setenv("OPENRAG_STORAGE_MODE", mode) + pending = asyncio.create_task( + loser.claim_session("alice", "contended", conversation_id="canonical") + ) + try: + await asyncio.wait_for(selected_missing.wait(), timeout=2) + # The loser already read no row; this insert commits on another real + # connection before its stale INSERT hits the primary-key constraint. + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "db") + assert ( + await _svc(factory).claim_session( + winner_user, "contended", conversation_id=winner_canonical + ) + is True + ) + monkeypatch.setenv("OPENRAG_STORAGE_MODE", mode) + finally: + winner_committed.set() + try: + assert await pending is accepted + fresh = _svc(factory) + assert await fresh.get_session_owner("contended") == winner_user + if accepted: + assert set(await fresh.get_conversation_sessions("alice", "canonical")) == { + "canonical", + "contended", + } + else: + assert "contended" not in loser.ownership_data + if mode == "hybrid": + assert tmp_data_dir.exists() is accepted + finally: + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["db", "hybrid"]) +async def test_unrelated_database_integrity_error_is_not_an_owned_race( + monkeypatch, tmp_data_dir, session_factory, mode +): + from sqlalchemy.exc import IntegrityError + + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "db") + assert await _svc(session_factory).claim_session("alice", "canonical") is True + monkeypatch.setenv("OPENRAG_STORAGE_MODE", mode) + svc = _svc(session_factory) + + async def unrelated_failure(*args): + raise IntegrityError("UPDATE session_ownership", {}, ValueError("unrelated constraint")) + + monkeypatch.setattr(SessionOwnershipRepo, "claim", unrelated_failure) + with pytest.raises(IntegrityError, match="unrelated constraint"): + await svc.claim_session("alice", "canonical") + assert svc.ownership_data == {} + assert not tmp_data_dir.exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["db", "hybrid"]) +async def test_nonunique_insert_integrity_error_cannot_reuse_compatible_owner( + monkeypatch, tmp_data_dir, session_factory, mode +): + from sqlalchemy import text + from sqlalchemy.exc import IntegrityError + + monkeypatch.setenv("OPENRAG_STORAGE_MODE", "db") + assert await _svc(session_factory).claim_session("alice", "canonical") is True + async with session_factory() as session: + await session.execute( + text( + "CREATE TRIGGER reject_ownership_insert BEFORE INSERT ON session_ownership " + "BEGIN SELECT RAISE(ABORT, 'unrelated ownership constraint'); END" + ) + ) + await session.commit() + original_get = SessionOwnershipRepo.get + stale_read = False + + async def stale_once(repo, response_id): + nonlocal stale_read + if response_id == "canonical" and not stale_read: + stale_read = True + return None + return await original_get(repo, response_id) + + monkeypatch.setattr(SessionOwnershipRepo, "get", stale_once) + monkeypatch.setenv("OPENRAG_STORAGE_MODE", mode) + svc = _svc(session_factory) + with pytest.raises(IntegrityError, match="unrelated ownership constraint"): + await svc.claim_session("alice", "canonical") + assert svc.ownership_data == {} + assert not tmp_data_dir.exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["db", "hybrid", "files"]) +async def test_legacy_null_association_remains_its_own_canonical_conversation( + monkeypatch, tmp_data_dir, session_factory, mode +): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", mode) + svc = _svc(session_factory) + assert await svc.claim_session("alice", "legacy") is True + assert await svc.claim_session("alice", "other") is True + if mode != "files": + async with session_factory() as session: + row = await SessionOwnershipRepo(session).get("legacy") + row.conversation_id = None + await session.commit() + if mode != "db": + data = json.loads(tmp_data_dir.read_text()) + data["legacy"]["conversation_id"] = None + tmp_data_dir.write_text(json.dumps(data)) + fresh = _svc(session_factory) + assert await fresh.claim_session("alice", "legacy", conversation_id="other") is False + assert await fresh.get_conversation_sessions("alice", "legacy") == ["legacy"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["files", "hybrid"]) +@pytest.mark.parametrize("operation", ["claim", "release", "release_group"]) +async def test_directory_sync_failure_reconciles_before_next_mutation( + monkeypatch, tmp_data_dir, session_factory, mode, operation +): + monkeypatch.setenv("OPENRAG_STORAGE_MODE", mode) + svc = _svc(session_factory) + if mode == "hybrid": + # Exercise the JSON fallback without a successful database write. + monkeypatch.setattr(svc, "_resolve_session_factory", lambda: None) + assert await svc.claim_session("alice", "existing") + real_fsync = os.fsync + + def failed_directory_sync(fd): + if stat.S_ISDIR(os.fstat(fd).st_mode): + raise OSError("directory sync unavailable") + real_fsync(fd) + + monkeypatch.setattr(os, "fsync", failed_directory_sync) + if operation == "claim": + assert await svc.claim_session("alice", "uncertain") is False + else: + with pytest.raises(OSError, match="directory sync unavailable"): + if operation == "release": + await svc.release_session("alice", "existing") + else: + await svc.release_conversation_sessions("alice", ["existing"]) + replaced = json.loads(tmp_data_dir.read_text()) + assert ("uncertain" in replaced) == (operation == "claim") + assert ("existing" in replaced) == (operation == "claim") + # Neither authorization nor subsequent mutations may use the stale cache. + assert await svc.get_session_owner("existing") is None + assert await svc.get_user_sessions("alice") == [] + assert await svc.claim_session("alice", "later") is False + with pytest.raises(OSError, match="directory sync unavailable"): + await svc.release_session("alice", "existing") + assert json.loads(tmp_data_dir.read_text()) == replaced + + monkeypatch.setattr(os, "fsync", real_fsync) + assert await svc.claim_session("alice", "later") is True + recovered = _svc(session_factory) + if mode == "hybrid": + monkeypatch.setattr(recovered, "_resolve_session_factory", lambda: None) + expected = {"existing", "uncertain", "later"} if operation == "claim" else {"later"} + assert set(await svc.get_user_sessions("alice")) == expected + assert set(await recovered.get_user_sessions("alice")) == expected diff --git a/tests/unit/test_chat_service_bulk_delete.py b/tests/unit/test_chat_service_bulk_delete.py index 81783308b..33b44fa8b 100644 --- a/tests/unit/test_chat_service_bulk_delete.py +++ b/tests/unit/test_chat_service_bulk_delete.py @@ -16,22 +16,6 @@ def _patch_owner(monkeypatch, owners): ) -@pytest.mark.asyncio -async def test_bulk_delete_mixed_ownership(monkeypatch): - svc = _make_service() - # alice owns s1; s2 is bob's; s3 has no owner - _patch_owner(monkeypatch, {"s1": "alice", "s2": "bob", "s3": None}) - svc.delete_session = AsyncMock(return_value={"success": True, "error": None}) - - result = await svc.delete_sessions("alice", ["s1", "s2", "s3"]) - - assert result["deleted"] == ["s1"] - assert sorted(result["failed"]) == ["s2", "s3"] - # only the owned id ever reaches delete_session - called = {c.args[1] for c in svc.delete_session.call_args_list} - assert called == {"s1"} - - @pytest.mark.asyncio async def test_bulk_delete_all_succeed(monkeypatch): svc = _make_service() diff --git a/tests/unit/test_langflow_stream_provider_errors.py b/tests/unit/test_langflow_stream_provider_errors.py index faa4ea1ba..be6f0bff5 100644 --- a/tests/unit/test_langflow_stream_provider_errors.py +++ b/tests/unit/test_langflow_stream_provider_errors.py @@ -1,31 +1,13 @@ -"""Langflow stream lifecycle and provider-error regressions. - -Only the upstream provider is replaced; cancellation regressions use the real -conversation store and ownership guard, including durable SQLite ownership. -""" +"""Langflow provider-error formatting and stream error regressions.""" from __future__ import annotations -import asyncio import json from collections.abc import AsyncIterator -from pathlib import Path -from types import SimpleNamespace -import anyio -import httpx import pytest -import pytest_asyncio -from fastapi import HTTPException -from openai import AsyncOpenAI -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine -from sqlmodel import SQLModel import agent as agent_module -from api.chat import ChatBody, _assert_owns, delete_session_endpoint, langflow_endpoint -from services.chat_service import ChatService -from services.conversation_persistence_service import ConversationPersistenceService -from services.session_ownership_service import SessionOwnershipService def test_format_provider_error_message_extracts_embedded_json(): @@ -119,6 +101,9 @@ async def claim_session(self, user_id, session_id, conversation_id=None): async def get_session_owner(self, session_id): return next((user_id for user_id, sid in claimed if sid == session_id), None) + async def get_conversation_sessions(self, user_id, session_id): + return [] + monkeypatch.setattr(agent_module, "store_conversation_thread", _store) monkeypatch.setattr( "services.session_ownership_service.session_ownership_service", @@ -468,514 +453,3 @@ async def raise_provider_error(*_args, **_kwargs) -> AsyncIterator[bytes]: assert not stored_id.startswith("error_") assert chunks[0].get("id") == stored_id assert stored_id in agent_module.active_conversations.get(user_id, {}) - - -@pytest_asyncio.fixture(params=["db", "hybrid", "files"]) -async def durable_stream_stores(request, monkeypatch, tmp_path): - monkeypatch.setenv("OPENRAG_STORAGE_MODE", request.param) - engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'chat.db'}") - async with engine.begin() as connection: - await connection.run_sync(SQLModel.metadata.create_all) - factory = async_sessionmaker(engine, expire_on_commit=False) - monkeypatch.setattr( - "services.session_ownership_service.get_data_file", - lambda name: str(tmp_path / name), - ) - ownership = SessionOwnershipService(session_factory=factory) - persistence = ConversationPersistenceService( - storage_file=str(tmp_path / "conversations.json"), session_factory=factory - ) - monkeypatch.setattr("services.session_ownership_service.session_ownership_service", ownership) - monkeypatch.setattr(agent_module, "conversation_persistence", persistence) - monkeypatch.setattr(agent_module, "active_conversations", {}) - yield ownership, persistence, factory - await engine.dispose() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("abort", ["close", "cancel"]) -async def test_aborted_langflow_stream_keeps_owned_id_and_partial_turn( - monkeypatch, durable_stream_stores, abort -): - ownership, persistence, factory = durable_stream_stores - waiting = asyncio.Event() - upstream_closed = asyncio.Event() - calls = [] - - async def upstream(*_args, **kwargs): - calls.append(kwargs.get("previous_response_id")) - try: - yield b'{"id":"resp-first","delta":{"content":"Partial answer"}}\n' - waiting.set() - await asyncio.Event().wait() - finally: - upstream_closed.set() - - monkeypatch.setattr(agent_module, "async_stream", upstream) - stream = agent_module.async_langflow_chat_stream(object(), "flow", "First question", "alice") - chunk = json.loads(await anext(stream)) - # The ID must already be durable when it first crosses the generator boundary. - await _assert_owns(chunk["id"], "alice") - assert (await persistence.get_conversation_thread("alice", chunk["id"]))["total_messages"] >= 1 - if abort == "close": - await stream.aclose() - else: - pending = asyncio.create_task(anext(stream)) - await waiting.wait() - pending.cancel() - with pytest.raises(asyncio.CancelledError): - await pending - assert upstream_closed.is_set() - # Reload ownership rather than relying on the live service's cache. - reloaded = SessionOwnershipService(session_factory=factory) - assert await reloaded.get_session_owner(chunk["id"]) == "alice" - with pytest.raises(HTTPException) as forbidden: - await _assert_owns(chunk["id"], "mallory") - assert forbidden.value.status_code == 403 - with pytest.raises(HTTPException) as missing: - await _assert_owns("unknown-session", "alice") - assert missing.value.status_code == 404 - - async def continuation(*_args, **kwargs): - calls.append(kwargs.get("previous_response_id")) - yield b'{"id":"resp-next","response_id":"resp-next-alias","delta":{"content":"Next"}}\n' - - monkeypatch.setattr(agent_module, "async_stream", continuation) - await _assert_owns(chunk["id"], "alice") - followup = agent_module.async_langflow_chat_stream( - object(), - "flow", - "Follow-up", - "alice", - previous_response_id=chunk["id"], - conversation_id=chunk["id"], - ) - next_chunk = json.loads(await anext(followup)) - await _assert_owns(next_chunk["id"], "alice") - await _assert_owns(next_chunk["response_id"], "alice") - assert await _collect_error_chunks(followup) == [] - assert calls == [None, "resp-first"] # Never replay the interrupted prompt. - conversations = await persistence.get_user_conversations("alice") - assert list(conversations) == ["resp-first"] # Stable sidebar id, changing response ids. - messages = agent_module.active_conversations["alice"]["resp-first"]["messages"] - assert [(m["role"], m["content"]) for m in messages if m["role"] != "system"] == [ - ("user", "First question"), - ("assistant", "Partial answer"), - ("user", "Follow-up"), - ("assistant", "Next"), - ] - assert conversations["resp-first"]["total_messages"] == 4 - - # SDK callers may only carry the changing response ID, not conversation_id. - await _collect_error_chunks( - agent_module.async_langflow_chat_stream( - object(), - "flow", - "Third question", - "alice", - previous_response_id=next_chunk["response_id"], - ) - ) - assert list(await persistence.get_user_conversations("alice")) == ["resp-first"] - assert calls == [None, "resp-first", "resp-next-alias"] - assert [m["content"] for m in messages if m["role"] == "user"] == [ - "First question", - "Follow-up", - "Third question", - ] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("claim_failure", ["unavailable", "foreign"]) -async def test_stream_does_not_expose_unowned_response_id( - monkeypatch, durable_stream_stores, claim_failure -): - ownership, persistence, _ = durable_stream_stores - if claim_failure == "foreign": - await ownership.claim_session("mallory", "resp-unclaimable") - else: - - async def unavailable(*_args, **_kwargs): - return False - - monkeypatch.setattr(ownership, "claim_session", unavailable) - - async def upstream(*_args, **_kwargs): - yield b'{"id":"resp-unclaimable","delta":{"content":"Not exposed"}}\n' - - monkeypatch.setattr(agent_module, "async_stream", upstream) - chunks = await _collect_error_chunks( - agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") - ) - assert len(chunks) == 1 - assert chunks[0]["status"] == "failed" - assert "id" not in chunks[0] and "response_id" not in chunks[0] - assert await persistence.get_user_conversations("alice") == {} - assert await ownership.get_session_owner("resp-unclaimable") == ( - "mallory" if claim_failure == "foreign" else None - ) - - -@pytest.mark.asyncio -async def test_local_asgi_disconnect_smoke_preserves_followup(monkeypatch, durable_stream_stores): - """Real endpoint → generator → OpenAI SSE adapter → durable stores, locally.""" - _, persistence, factory = durable_stream_stores - disconnected = anyio.Event() - upstream_closed = anyio.Event() - calls = [] - - class ProviderStream(httpx.AsyncByteStream): - async def __aiter__(self): - yield ( - b"event: response.created\n" - b'data: {"type":"response.created","response":{"id":"resp-http"}}\n\n' - ) - yield ( - b"event: response.output_text.delta\n" - b'data: {"type":"response.output_text.delta","id":"resp-http",' - b'"delta":{"content":"Partial HTTP answer"}}\n\n' - ) - await anyio.sleep_forever() - - async def aclose(self): - upstream_closed.set() - - async def provider(request): - payload = json.loads(request.content) - calls.append(payload) - if len(calls) == 1: - return httpx.Response( - 200, headers={"content-type": "text/event-stream"}, stream=ProviderStream() - ) - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - content=( - b"event: response.output_text.delta\n" - b'data: {"type":"response.output_text.delta","id":"resp-http-next",' - b'"delta":{"content":"Follow-up HTTP answer"}}\n\n' - b"event: response.completed\n" - b'data: {"type":"response.completed","response":{"id":"resp-http-next",' - b'"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7}}}\n\n' - ), - ) - - async with AsyncOpenAI( - api_key="local-test-only", - base_url="https://provider.invalid/v1", - http_client=httpx.AsyncClient(transport=httpx.MockTransport(provider)), - max_retries=0, - ) as client: - - class EndpointChatService: - async def langflow_chat(self, prompt, _subject, _jwt, **kwargs): - return agent_module.async_langflow_chat_stream( - client, - "flow", - prompt, - kwargs["storage_user_id"], - previous_response_id=kwargs["previous_response_id"], - conversation_id=kwargs["conversation_id"], - ) - - user = SimpleNamespace( - user_id="oauth-alice", - db_user_id="alice", - jwt_token="local-test-only", - name="Alice", - email="alice@example.invalid", - ) - service = EndpointChatService() - response = await langflow_endpoint( - ChatBody(prompt="First HTTP question", stream=True), - chat_service=service, - session_manager=None, - user=user, - ) - emitted = [] - - async def send(message): - if message["type"] == "http.response.body" and message.get("body"): - chunk = json.loads(message["body"]) - emitted.append(chunk) - response_id = chunk.get("id") or (chunk.get("response") or {}).get("id") - if response_id: - await _assert_owns(response_id, "alice") - if chunk.get("delta", {}).get("content"): - disconnected.set() - - async def receive(): - await disconnected.wait() - return {"type": "http.disconnect"} - - # Starlette's real disconnect listener cancels the body iterator using an - # AnyIO cancel scope, rather than manually throwing into our generator. - with anyio.fail_after(5): - await response({"type": "http", "asgi": {"spec_version": "2.3"}}, receive, send) - assert upstream_closed.is_set() - assert len(emitted) == 2 - metadata = await persistence.get_conversation_thread("alice", "resp-http") - assert metadata["total_messages"] == 2 - - followup = await langflow_endpoint( - ChatBody( - prompt="Next HTTP question", - stream=True, - previous_response_id="resp-http", - conversation_id="resp-http", - ), - chat_service=service, - session_manager=None, - user=user, - ) - chunks = await _collect_error_chunks(followup.body_iterator) - await _assert_owns("resp-http-next", "alice") - assert chunks[-1]["response"]["usage"]["total_tokens"] == 7 - assert calls[1]["previous_response_id"] == "resp-http" - assert [call["input"] for call in calls] == ["First HTTP question", "Next HTTP question"] - conversations = await persistence.get_user_conversations("alice") - assert list(conversations) == ["resp-http"] - assert conversations["resp-http"]["total_messages"] == 4 - messages = agent_module.active_conversations["alice"]["resp-http"]["messages"] - assert [m["content"] for m in messages if m["role"] != "system"] == [ - "First HTTP question", - "Partial HTTP answer", - "Next HTTP question", - "Follow-up HTTP answer", - ] - - # Restart before canonical deletion: aliases must come from durable state. - reloaded = SessionOwnershipService(session_factory=factory) - monkeypatch.setattr( - "services.session_ownership_service.session_ownership_service", reloaded - ) - monkeypatch.setattr(agent_module, "active_conversations", {}) - upstream_deletes = [] - - async def monitor_delete(method, path): - assert method == "DELETE" - upstream_deletes.append(path.rsplit("/", 1)[-1]) - return httpx.Response(204) - - monkeypatch.setattr("services.chat_service.clients.langflow_request", monitor_delete) - deletion = await delete_session_endpoint( - session_id="resp-http", chat_service=ChatService(), user=user - ) - assert deletion.status_code == 200 - assert set(upstream_deletes) == {"resp-http", "resp-http-next"} - fresh = SessionOwnershipService(session_factory=factory) - monkeypatch.setattr("services.session_ownership_service.session_ownership_service", fresh) - with pytest.raises(HTTPException) as missing: - await langflow_endpoint( - ChatBody( - prompt="Cannot resurrect", stream=True, previous_response_id="resp-http-next" - ), - chat_service=service, - session_manager=None, - user=user, - ) - assert missing.value.status_code == 404 - assert len(calls) == 2 - - # Actual endpoint + SSE adapter with both persistent destinations broken. - ownership_path = Path(fresh.ownership_file) - if ownership_path.exists(): - ownership_path.unlink() - ownership_path.mkdir() - - def unavailable_database(): - raise OSError("local smoke database unavailable") - - monkeypatch.setattr(fresh, "_resolve_session_factory", lambda: unavailable_database) - failed = await langflow_endpoint( - ChatBody(prompt="Fail closed", stream=True), - chat_service=service, - session_manager=None, - user=user, - ) - failed_chunks = await _collect_error_chunks(failed.body_iterator) - assert len(failed_chunks) == 1 and failed_chunks[0]["status"] == "failed" - assert "id" not in failed_chunks[0] and "response_id" not in failed_chunks[0] - assert ( - await SessionOwnershipService(session_factory=factory).get_session_owner( - "resp-http-next" - ) - is None - ) - print( - "LOCAL SMOKE: disconnect/follow-up retained 4 messages; restart + " - "canonical deletion removed both upstream sessions and alias authorization; " - "persistent failures emitted no client ID" - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("delta", [None, {"content": None}]) -async def test_metadata_only_chunk_claims_id_before_exposure( - monkeypatch, durable_stream_stores, delta -): - async def upstream(*_args, **_kwargs): - yield (json.dumps({"id": "resp-metadata", "delta": delta}) + "\n").encode() - - monkeypatch.setattr(agent_module, "async_stream", upstream) - stream = agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") - chunk = json.loads(await anext(stream)) - await _assert_owns(chunk["id"], "alice") - await stream.aclose() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("delete_id", ["resp-first", "resp-second"]) -async def test_delete_canonical_revokes_aliases_after_reload( - monkeypatch, durable_stream_stores, delete_id -): - ownership, persistence, factory = durable_stream_stores - - async def upstream(*_args, **kwargs): - response_id = "resp-second" if kwargs.get("previous_response_id") else "resp-first" - yield (json.dumps({"id": response_id, "delta": {"content": "Answer"}}) + "\n").encode() - - monkeypatch.setattr(agent_module, "async_stream", upstream) - for previous in (None, "resp-first"): - await _collect_error_chunks( - agent_module.async_langflow_chat_stream( - object(), "flow", "Question", "alice", previous_response_id=previous - ) - ) - await ownership.claim_session("mallory", "resp-foreign") - await ownership.claim_session("alice", "resp-unrelated") - assert list(await persistence.get_user_conversations("alice")) == ["resp-first"] - assert await agent_module.delete_user_conversation("mallory", "resp-first") is False - assert agent_module.active_conversations["alice"]["resp-first"]["messages"] - # Restart: neither ownership cache nor the in-memory response_ids can help. - reloaded = SessionOwnershipService(session_factory=factory) - monkeypatch.setattr("services.session_ownership_service.session_ownership_service", reloaded) - monkeypatch.setattr(agent_module, "active_conversations", {}) - persistence = ConversationPersistenceService( - storage_file=persistence.storage_file, session_factory=factory - ) - monkeypatch.setattr(agent_module, "conversation_persistence", persistence) - assert await agent_module.delete_user_conversation("alice", delete_id) is True - fresh = SessionOwnershipService(session_factory=factory) - monkeypatch.setattr("services.session_ownership_service.session_ownership_service", fresh) - for response_id in ("resp-first", "resp-second"): - with pytest.raises(HTTPException) as missing: - await _assert_owns(response_id, "alice") - assert missing.value.status_code == 404 - assert await fresh.get_session_owner("resp-foreign") == "mallory" - assert await fresh.get_session_owner("resp-unrelated") == "alice" - assert await persistence.get_user_conversations("alice") == {} - - -@pytest.mark.asyncio -async def test_persistent_claim_failure_never_exposes_cached_id(monkeypatch, durable_stream_stores): - ownership, persistence, factory = durable_stream_stores - # A real IO failure, not a stubbed claim: cache echo must not prove durability. - Path(ownership.ownership_file).mkdir() - - def unavailable_database(): - raise OSError("local database unavailable") - - monkeypatch.setattr(ownership, "_resolve_session_factory", lambda: unavailable_database) - - async def upstream(*_args, **_kwargs): - yield b'{"id":"resp-file-failure","delta":{"content":"Not exposed"}}\n' - - monkeypatch.setattr(agent_module, "async_stream", upstream) - chunks = await _collect_error_chunks( - agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") - ) - assert len(chunks) == 1 - assert chunks[0]["status"] == "failed" - assert "id" not in chunks[0] and "response_id" not in chunks[0] - assert await ownership.get_session_owner("resp-file-failure") is None - assert ( - await SessionOwnershipService(session_factory=factory).get_session_owner( - "resp-file-failure" - ) - is None - ) - assert await persistence.get_user_conversations("alice") == {} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("durable_stream_stores", ["hybrid", "files"], indirect=True) -async def test_failed_file_revocation_is_not_reported_as_success( - monkeypatch, durable_stream_stores -): - ownership, _, _ = durable_stream_stores - - async def upstream(*_args, **_kwargs): - yield b'{"id":"resp-delete-failure","delta":{"content":"Answer"}}\n' - - monkeypatch.setattr(agent_module, "async_stream", upstream) - await _collect_error_chunks( - agent_module.async_langflow_chat_stream(object(), "flow", "Question", "alice") - ) - ownership_path = Path(ownership.ownership_file) - ownership_path.unlink() - ownership_path.mkdir() - - async def monitor_delete(_self, _session_id): - return True - - monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) - result = await ChatService().delete_session("alice", "resp-delete-failure") - assert result["success"] is False - await _assert_owns("resp-delete-failure", "alice") - assert "resp-delete-failure" in agent_module.active_conversations["alice"] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("durable_stream_stores", ["hybrid"], indirect=True) -async def test_json_only_alias_cleanup_precedes_failed_db_revocation( - monkeypatch, durable_stream_stores -): - from db.repositories import SessionOwnershipRepo - - ownership, _, factory = durable_stream_stores - - async def upstream(*_args, **kwargs): - response_id = "alias" if kwargs.get("previous_response_id") else "canonical" - yield (json.dumps({"id": response_id, "delta": {"content": "Answer"}}) + "\n").encode() - - monkeypatch.setattr(agent_module, "async_stream", upstream) - await _collect_error_chunks( - agent_module.async_langflow_chat_stream(object(), "flow", "First", "alice") - ) - - def unavailable_database(): - raise OSError("database temporarily unavailable") - - ownership._session_factory = unavailable_database - await _collect_error_chunks( - agent_module.async_langflow_chat_stream( - object(), "flow", "Next", "alice", previous_response_id="canonical" - ) - ) - ownership._session_factory = factory - async with factory() as session: - assert await SessionOwnershipRepo(session).get("alias") is None - assert await ownership.get_session_owner("alias") == "alice" - original = SessionOwnershipRepo.release - - async def failed_release(*_args): - raise OSError("database revocation failed") - - upstream_deletes = [] - - async def monitor_delete(_self, session_id): - upstream_deletes.append(session_id) - return True - - monkeypatch.setattr(ChatService, "_delete_langflow_session", monitor_delete) - monkeypatch.setattr(SessionOwnershipRepo, "release", failed_release) - result = await ChatService().delete_session("alice", "canonical") - assert result["success"] is False - assert set(upstream_deletes) == {"canonical", "alias"} - # Even though JSON-only ownership was revoked, its upstream cleanup was not - # lost. The surviving canonical DB row can be retried after a restart. - reloaded = SessionOwnershipService(session_factory=factory) - monkeypatch.setattr("services.session_ownership_service.session_ownership_service", reloaded) - monkeypatch.setattr(SessionOwnershipRepo, "release", original) - assert (await ChatService().delete_session("alice", "canonical"))["success"] is True