diff --git a/app/repositories/conversations.py b/app/repositories/conversations.py index 24608a1..c9a4df6 100644 --- a/app/repositories/conversations.py +++ b/app/repositories/conversations.py @@ -1,16 +1,24 @@ """Persistence queries for user-owned assistant conversations.""" -from sqlalchemy import func, select +from sqlalchemy import and_, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload -from app.models.entities import Conversation, ConversationMessage +from app.models.entities import Conversation, ConversationMessage, IntegrationConversation class ConversationRepository: def __init__(self, session: AsyncSession) -> None: self.session = session + async def is_external(self, conversation_id: int) -> bool: + return ( + await self.session.scalar( + select(IntegrationConversation.id).where(IntegrationConversation.conversation_id == conversation_id) + ) + is not None + ) + async def list_for_user( self, user_id: int, @@ -39,16 +47,29 @@ class ConversationRepository: offset: int = 0, ) -> tuple[Conversation, list[ConversationMessage], int] | None: conversation = await self.session.scalar( - select(Conversation) - .where( + select(Conversation).where( Conversation.id == conversation_id, Conversation.user_id == user_id, ) ) if conversation is None: return None + return await self.get_message_page(conversation, limit=limit, offset=offset) + + async def get_message_page( + self, + conversation: Conversation, + *, + limit: int = 20, + offset: int = 0, + ) -> tuple[Conversation, list[ConversationMessage], int]: + """Page messages after the caller has verified ownership of the conversation.""" + + conversation_id = conversation.id total = await self.session.scalar( - select(func.count()).select_from(ConversationMessage).where( + select(func.count()) + .select_from(ConversationMessage) + .where( ConversationMessage.conversation_id == conversation_id, ) ) @@ -81,3 +102,52 @@ class ConversationRepository: await self.session.delete(conversation) await self.session.flush() return True + + async def get_withdraw_target_for_user( + self, + conversation_id: int, + message_id: int, + user_id: int, + ) -> tuple[Conversation, ConversationMessage | None] | None: + """Check ownership and load the target in one query without loading history.""" + + row = ( + await self.session.execute( + select(Conversation, ConversationMessage) + .outerjoin( + ConversationMessage, + and_( + ConversationMessage.conversation_id == Conversation.id, + ConversationMessage.id == message_id, + ), + ) + .where( + Conversation.id == conversation_id, + Conversation.user_id == user_id, + ) + ) + ).one_or_none() + return (row[0], row[1]) if row is not None else None + + async def get_next_message( + self, + conversation_id: int, + message: ConversationMessage, + ) -> ConversationMessage | None: + """Load only the message immediately following the target message.""" + + return await self.session.scalar( + select(ConversationMessage) + .where( + ConversationMessage.conversation_id == conversation_id, + or_( + ConversationMessage.created_at > message.created_at, + and_( + ConversationMessage.created_at == message.created_at, + ConversationMessage.id > message.id, + ), + ), + ) + .order_by(ConversationMessage.created_at, ConversationMessage.id) + .limit(1) + ) diff --git a/app/services/conversations.py b/app/services/conversations.py index bb95f9c..b1a1db6 100644 --- a/app/services/conversations.py +++ b/app/services/conversations.py @@ -785,17 +785,18 @@ class ConversationService: ) -> ConversationData: """Delete one user turn and its immediately paired assistant reply.""" - conversation = await self.repository.get_for_user(conversation_id, user.id) - if conversation is None: + target = await self.repository.get_withdraw_target_for_user(conversation_id, message_id, user.id) + if target is None: raise AppException("对话不存在或无权访问", code=404, status_code=404) - message = next((item for item in conversation.messages if item.id == message_id), None) + if await self.repository.is_external(conversation_id): + raise AppException("报告关联会话暂不支持撤回", code=409, status_code=409) + conversation, message = target if message is None or message.role != "user": raise AppException("只能撤回用户消息", code=422, status_code=422) - messages = sorted(conversation.messages, key=lambda item: (item.created_at, item.id)) - index = messages.index(message) ids = [message.id] - if index + 1 < len(messages) and messages[index + 1].role == "assistant": - ids.append(messages[index + 1].id) + next_message = await self.repository.get_next_message(conversation_id, message) + if next_message is not None and next_message.role == "assistant": + ids.append(next_message.id) await self.session.execute( delete(ConversationMessage).where(ConversationMessage.id.in_(ids)) ) @@ -811,9 +812,7 @@ class ConversationService: ) ) await self.session.flush() - page = await self.repository.get_page_for_user(conversation_id, user.id, limit=20, offset=0) - if page is None: - raise AppException("对话读取失败", code=500, status_code=500) + page = await self.repository.get_message_page(conversation, limit=20, offset=0) persisted, page_messages, total = page return self._data(persisted, messages=page_messages, total=total, offset=0) diff --git a/tests/test_conversation_withdraw_repository.py b/tests/test_conversation_withdraw_repository.py new file mode 100644 index 0000000..33cf938 --- /dev/null +++ b/tests/test_conversation_withdraw_repository.py @@ -0,0 +1,156 @@ +"""Exercise real withdrawal SQL and guard against loading an entire long conversation.""" + +from datetime import UTC, datetime, timedelta + +import pytest +from app.core.exceptions import AppException +from app.models.entities import ( + Conversation, + ConversationMessage, + IntegrationAssessment, + IntegrationConversation, + OperationAudit, + UserAccount, +) +from app.services.conversations import ConversationService +from sqlalchemy import JSON, BigInteger, Integer, MetaData, event, inspect, select +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + + +@pytest.fixture +async def withdraw_db(tmp_path): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'withdraw.db'}") + schema = MetaData() + for model in ( + UserAccount, + Conversation, + ConversationMessage, + OperationAudit, + IntegrationAssessment, + IntegrationConversation, + ): + copied = model.__table__.to_metadata(schema) + for column in copied.columns: + if isinstance(column.type, JSONB): + column.type = JSON() + elif isinstance(column.type, BigInteger): + column.type = Integer() + async with engine.begin() as connection: + await connection.run_sync(schema.create_all) + sessions = async_sessionmaker(engine, expire_on_commit=False) + start = datetime(2026, 10, 27, tzinfo=UTC) + async with sessions() as session: + session.add_all([UserAccount(id=i, emp_no=f"u{i}", name=f"用户{i}", role="user") for i in (1, 2)]) + session.add_all([Conversation(id=i, user_id=i, title=f"会话{i}") for i in (1, 2)]) + for message_id, role, minute in ( + (1, "user", 0), + (2, "assistant", 0), + (3, "user", 1), + (4, "user", 2), + (5, "assistant", 2), + (6, "user", 3), + ): + session.add( + ConversationMessage( + id=message_id, + conversation_id=1, + role=role, + content=f"消息{message_id}", + created_at=start + timedelta(minutes=minute), + ) + ) + session.add(ConversationMessage(id=7, conversation_id=2, role="user", content="他人的消息")) + await session.commit() + try: + yield sessions, engine, start + finally: + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("target_id", "deleted"), [(1, [1, 2]), (3, [3]), (6, [6])]) +async def test_withdraw_only_removes_target_and_immediate_assistant(withdraw_db, target_id, deleted): + sessions, _, _ = withdraw_db + async with sessions() as session: + user = await session.get(UserAccount, 1) + result = await ConversationService(session).withdraw_turn(1, target_id, user) + assert result.messages_total == 6 - len(deleted) + assert [message.id for message in result.messages] == [i for i in range(1, 7) if i not in deleted] + await session.commit() + async with sessions() as session: + ids = list(await session.scalars(select(ConversationMessage.id).order_by(ConversationMessage.id))) + assert ids == [i for i in range(1, 8) if i not in deleted] + audit = await session.scalar(select(OperationAudit)) + assert audit.action == "conversation.withdraw" + assert audit.audit_metadata == {"deleted_message_ids": deleted} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("conversation_id", "message_id", "user_id", "status"), + [(1, 1, 2, 404), (999, 1, 1, 404), (1, 999, 1, 422), (1, 2, 1, 422), (1, 7, 1, 422)], +) +async def test_withdraw_rejects_other_users_and_invalid_targets( + withdraw_db, conversation_id, message_id, user_id, status +): + sessions, _, _ = withdraw_db + async with sessions() as session: + user = await session.get(UserAccount, user_id) + with pytest.raises(AppException) as error: + await ConversationService(session).withdraw_turn(conversation_id, message_id, user) + assert error.value.status_code == status + await session.rollback() + async with sessions() as session: + assert len(list(await session.scalars(select(ConversationMessage.id)))) == 7 + assert list(await session.scalars(select(OperationAudit.id))) == [] + + +@pytest.mark.asyncio +async def test_long_conversation_withdraw_loads_only_target_pair_and_latest_page(withdraw_db): + sessions, engine, start = withdraw_db + async with sessions() as session: + session.add_all( + [ + ConversationMessage( + id=message_id, + conversation_id=1, + role="user", + content="较长的历史文字" * 100, + created_at=start + timedelta(minutes=message_id), + ) + for message_id in range(8, 1008) + ] + ) + await session.commit() + async with sessions() as session: + user = await session.get(UserAccount, 1) + statements = [] + + def record_query(_, __, statement, ___, ____, _____): + statements.append(statement) + + event.listen(engine.sync_engine, "before_cursor_execute", record_query) + try: + result = await ConversationService(session).withdraw_turn(1, 1, user) + finally: + event.remove(engine.sync_engine, "before_cursor_execute", record_query) + assert result.messages_total == 1004 + assert [message.id for message in result.messages] == list(range(988, 1008)) + assert result.has_more_messages is True + conversation = await session.get(Conversation, 1) + assert "messages" in inspect(conversation).unloaded + assert sum(isinstance(row, ConversationMessage) for row in session.identity_map.values()) <= 22 + # Ownership + target, next message, count, latest page: no redundant header query. + assert len([sql for sql in statements if sql.lstrip().upper().startswith("SELECT")]) == 5 + + +@pytest.mark.asyncio +async def test_withdraw_rollback_restores_messages_and_drops_audit(withdraw_db): + sessions, _, _ = withdraw_db + async with sessions() as session: + await ConversationService(session).withdraw_turn(1, 1, await session.get(UserAccount, 1)) + await session.rollback() + async with sessions() as session: + assert len(list(await session.scalars(select(ConversationMessage.id)))) == 7 + assert list(await session.scalars(select(OperationAudit.id))) == []