fix(storage): finish SQLite transactions before propagating cancellation

src/api/endpoints.py

@@ -22,6 +22,7 @@ from ..core.fastgpt_client import (
 from ..core.config import Config
 from ..utils.text_chunker import SentenceTextChunker, SentenceTextChunkerConfig
 from loguru import logger
+import asyncio
 import json
 import re
 import time
@@ -412,8 +413,16 @@ async def chat(
                 )
                 yield create_sse_event("done", {"status": "completed"})
                 
+            except asyncio.CancelledError:
+                logger.warning(
+                    "Chat stream cancelled "
+                    f"sessionId={json_data['sessionId']} "
+                    f"duration_ms={(time.perf_counter() - stream_started_at) * 1000:.1f} "
+                    f"text_delta_count={text_delta_count}"
+                )
+                raise
             except Exception as e:
-                logger.error(
+                logger.opt(exception=True).error(
                     "Chat stream failed "
                     f"sessionId={json_data['sessionId']} "
                     f"duration_ms={(time.perf_counter() - stream_started_at) * 1000:.1f} "

src/storage/cancellation.py

@@ -0,0 +1,33 @@
+"""在连接交回调用方前完成已提交给 SQLite worker 的操作。"""
+
+import asyncio
+from collections.abc import Awaitable
+from typing import TypeVar
+
+import anyio
+
+T = TypeVar('T')
+
+
+async def finish_sqlite_operation(operation: Awaitable[T]) -> T:
+    """屏蔽 ASGI 的持续取消与重复 task.cancel；完成后仍向调用方传播取消。
+
+    aiosqlite 取消等待并不会撤销 worker 队列里的 SQL。必须等待其完成，
+    避免连接锁已释放但底层仍持事务，或下一请求与回滚交错。
+    """
+    cancelled = None
+    with anyio.CancelScope(shield=True):
+        task = asyncio.ensure_future(operation)
+        while True:
+            try:
+                result = await asyncio.shield(task)
+                break
+            except asyncio.CancelledError as exc:
+                if task.done():
+                    # 操作本身取消时结束；其他异常也必须被取出。
+                    task.result()
+                    raise
+                cancelled = exc
+    if cancelled is not None:
+        raise cancelled
+    return result

src/storage/database.py

@@ -3,6 +3,8 @@ import os
 from contextlib import asynccontextmanager
 from pathlib import Path
 
+from .cancellation import finish_sqlite_operation
+
 
 class SqliteInstanceLock:
     """拒绝第二个进程/实例打开同一 SQLite 文件；锁由操作系统随退出释放。"""
@@ -73,12 +75,13 @@ class SqliteDatabase:
     async def transaction(self):
         # 不跨 LLM 调用持有事务；仅保护一次业务提交。
         async with self.lock:
-            await self.connection.execute('BEGIN IMMEDIATE')
             try:
+                # BEGIN 也可能已经在 worker 执行，而等待它的请求被取消。
+                await self.connection.execute('BEGIN IMMEDIATE')
                 yield Connection(self.connection)
                 await self.connection.commit()
             except BaseException:
-                await self.connection.rollback()
+                await finish_sqlite_operation(self.connection.rollback())
                 raise
 
 

src/storage/runtime.py

@@ -43,7 +43,7 @@ async def open_storage(settings: StorageSettings):
             return
         if settings.kind == 'sqlite':
             import aiosqlite
-            from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+            from .sqlite_saver import CancellationSafeSqliteSaver
 
             if not settings.sqlite_path.strip() or settings.sqlite_path == ':memory:':
                 raise ValueError('SQLite persistence requires a file path')
@@ -56,7 +56,7 @@ async def open_storage(settings: StorageSettings):
             await business.execute('PRAGMA busy_timeout=5000')
             checkpoint_connection = await stack.enter_async_context(aiosqlite.connect(path))
             await checkpoint_connection.execute('PRAGMA busy_timeout=5000')
-            checkpointer = AsyncSqliteSaver(checkpoint_connection)
+            checkpointer = CancellationSafeSqliteSaver(checkpoint_connection)
             database = SqliteDatabase(business)
             locks = SessionLocks()
         else:

src/storage/sqlite_saver.py

@@ -0,0 +1,41 @@
+"""保留官方 SQLite saver 格式，补充请求取消时的事务收尾。"""
+
+import asyncio
+
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+
+from .cancellation import finish_sqlite_operation
+
+
+class CancellationSafeSqliteSaver(AsyncSqliteSaver):
+    def __init__(self, connection):
+        super().__init__(connection)
+        self._writes_lock = asyncio.Lock()
+
+    async def _finish_write(self, operation):
+        async def write():
+            # 写操作失败后，先回滚再允许后续写入；官方内部仍使用 self.lock。
+            async with self._writes_lock:
+                try:
+                    return await operation()
+                except BaseException:
+                    async with self.lock:
+                        await self.conn.rollback()
+                    raise
+        # SQLite 本地操作完成后再传播取消，不跨模型调用屏蔽取消。
+        return await finish_sqlite_operation(write())
+
+    async def aput(self, config, checkpoint, metadata, new_versions):
+        return await self._finish_write(
+            lambda: super(CancellationSafeSqliteSaver, self).aput(config, checkpoint, metadata, new_versions),
+        )
+
+    async def aput_writes(self, config, writes, task_id, task_path=''):
+        return await self._finish_write(
+            lambda: super(CancellationSafeSqliteSaver, self).aput_writes(config, writes, task_id, task_path),
+        )
+
+    async def adelete_thread(self, thread_id):
+        return await self._finish_write(
+            lambda: super(CancellationSafeSqliteSaver, self).adelete_thread(thread_id),
+        )