优化会话撤回查询性能

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)
+            select(Conversation).where(
-            .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)
+        )

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)
+        target = await self.repository.get_withdraw_target_for_user(conversation_id, message_id, user.id)
-        if conversation is None:
+        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":
+        next_message = await self.repository.get_next_message(conversation_id, message)
-            ids.append(messages[index + 1].id)
+        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)
+        page = await self.repository.get_message_page(conversation, limit=20, offset=0)
-        if page is None:
-            raise AppException("对话读取失败", code=500, status_code=500)
         persisted, page_messages, total = page
         return self._data(persisted, messages=page_messages, total=total, offset=0)
 

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))) == []