fix(storage): finish SQLite transactions before propagating cancellation

docs/agent-storage.md

@@ -11,6 +11,10 @@
 | postgresql | 多 worker、多实例 | AsyncPostgresSaver | 进程内锁 + PostgreSQL session advisory lock |
 
 SQLite 实例锁由操作系统管理，第二个实例打开同一数据库会失败，进程退出自动释放。
+业务事务的 BEGIN、提交与异常路径均在回滚保护范围内；请求取消时，回滚必须
+完成后才能释放连接锁。SQLite checkpoint saver 保留官方表结构，并在本地
+写入/删除操作完成后再传播取消，异常时回滚，避免残留事务锁住业务连接。
+这段取消保护只覆盖数据库操作，不覆盖模型生成。
 SQLite 文件必须放在本地持久化磁盘，不作为多主机共享文件使用。PostgreSQL 的
 锁连接与业务事务连接使用独立连接池；LLM 运行期间不持有业务行事务。
 
@@ -139,3 +143,5 @@ conda run -n py3.11 python -m pytest test/storage/ -q
 没有该变量时跳过 PostgreSQL 用例，不把跳过当作 PostgreSQL 已验证。
 测试使用 fake LLM，不消耗模型额度，覆盖：重启续聊、四接口、两种数据类型、
 幂等重放、版本冲突、checkpoint 缺失/超前、提交前后进程直接退出、数据库锁和删除恢复。
+`test_sqlite_cancellation.py` 另外覆盖 BEGIN、回滚、checkpoint 写入/删除中的
+取消、重复取消及后续请求可用性，验证不依赖服务重启释放残留事务。

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),
+        )

test/storage/test_sqlite_cancellation.py

@@ -0,0 +1,189 @@
+"""SQLite 请求取消后，连接必须先退出事务再交给下一次请求。"""
+
+import asyncio
+from contextlib import aclosing, suppress
+
+import aiosqlite
+import anyio
+import pytest
+
+from src.agent.graph import build_graph
+from src.backends.chat import ChatInput
+from src.backends.langgraph import LangGraphBackend
+from src.services.app_service import AccidentAppService
+from src.storage.database import SqliteDatabase
+from src.storage.runtime import StorageSettings, open_storage
+
+
+@pytest.mark.asyncio
+async def test_cancel_during_begin_does_not_leave_transaction(tmp_path, monkeypatch):
+    async with aiosqlite.connect(tmp_path/'begin.sqlite3', isolation_level=None) as raw:
+        database = SqliteDatabase(raw)
+        started = asyncio.Event()
+        execute = raw.execute
+        async def paused_begin(sql, params=()):
+            cursor = await execute(sql, params)
+            if sql == 'BEGIN IMMEDIATE':
+                started.set()
+                await asyncio.Future()
+            return cursor
+        monkeypatch.setattr(raw, 'execute', paused_begin)
+        async def request():
+            async with database.transaction():
+                pass
+        task = asyncio.create_task(request())
+        await asyncio.wait_for(started.wait(), 3)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        monkeypatch.setattr(raw, 'execute', execute)
+        assert not raw.in_transaction
+        async with database.transaction() as connection:
+            await connection.execute('CREATE TABLE recovered (value TEXT)')
+
+
+@pytest.mark.asyncio
+async def test_asgi_cancel_scope_does_not_interrupt_rollback(tmp_path):
+    async with aiosqlite.connect(tmp_path/'scope.sqlite3', isolation_level=None) as raw:
+        database = SqliteDatabase(raw)
+        with anyio.CancelScope() as scope:
+            async with database.transaction() as connection:
+                await connection.execute('CREATE TABLE uncommitted (value TEXT)')
+                scope.cancel()
+                await anyio.sleep(0)
+        # 用同一 worker 队列作为屏障，确认已排队的 SQL 已完成。
+        await raw.execute('SELECT 1')
+        assert not raw.in_transaction
+        async with database.transaction() as connection:
+            assert await connection.one("SELECT name FROM sqlite_master WHERE name='uncommitted'") is None
+        async with aiosqlite.connect(tmp_path/'scope.sqlite3', isolation_level=None, timeout=0.1) as other:
+            await other.execute('CREATE TABLE unlocked (value TEXT)')
+
+
+@pytest.mark.asyncio
+async def test_repeated_task_cancel_waits_for_rollback_before_unlock(tmp_path, monkeypatch):
+    async with aiosqlite.connect(tmp_path/'repeat.sqlite3', isolation_level=None) as raw:
+        database = SqliteDatabase(raw)
+        entered = asyncio.Event()
+        rolling_back = asyncio.Event()
+        release = asyncio.Event()
+        rollback = raw.rollback
+        async def paused_rollback():
+            rolling_back.set()
+            await release.wait()
+            await rollback()
+        monkeypatch.setattr(raw, 'rollback', paused_rollback)
+        async def request():
+            async with database.transaction():
+                entered.set()
+                await asyncio.Future()
+        task = asyncio.create_task(request())
+        try:
+            await asyncio.wait_for(entered.wait(), 3)
+            task.cancel()
+            await asyncio.wait_for(rolling_back.wait(), 3)
+            task.cancel()
+            await asyncio.sleep(0)
+            assert database.lock.locked()
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            assert not raw.in_transaction
+            assert not database.lock.locked()
+        finally:
+            release.set()
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('operation', ['put', 'writes', 'delete'])
+async def test_checkpoint_cancel_finishes_transaction_before_releasing_lock(tmp_path, monkeypatch, operation):
+    from langgraph.checkpoint.base import empty_checkpoint
+
+    settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'checkpoint.sqlite3'), auto_migrate=True)
+    async with open_storage(settings) as storage:
+        saver = storage.checkpointer
+        config = {'configurable': {'thread_id': 'cancel-check', 'checkpoint_ns': ''}}
+        config = await saver.aput(config, empty_checkpoint(), {}, {})
+        committing = asyncio.Event()
+        release = asyncio.Event()
+        commit = saver.conn.commit
+        async def paused_commit():
+            committing.set()
+            await release.wait()
+            await commit()
+        monkeypatch.setattr(saver.conn, 'commit', paused_commit)
+        if operation == 'put':
+            call = saver.aput(config, empty_checkpoint(), {}, {})
+        elif operation == 'writes':
+            call = saver.aput_writes(config, [('channel', 'value')], 'task')
+        else:
+            call = saver.adelete_thread('cancel-check')
+        task = asyncio.create_task(call)
+        try:
+            await asyncio.wait_for(committing.wait(), 3)
+            task.cancel()
+            await asyncio.sleep(0)
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            assert not saver.conn.in_transaction
+            assert not saver.lock.locked()
+            # 业务连接可以立即写入同一个数据库。
+            await storage.repository.apply_update('next-request', {'hphm1': 'test'})
+        finally:
+            release.set()
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
+
+
+@pytest.mark.asyncio
+async def test_checkpoint_write_failure_rolls_back_before_business_write(tmp_path, monkeypatch):
+    from langgraph.checkpoint.base import empty_checkpoint
+
+    settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'failure.sqlite3'), auto_migrate=True)
+    async with open_storage(settings) as storage:
+        saver = storage.checkpointer
+        config = {'configurable': {'thread_id': 'failed-write', 'checkpoint_ns': ''}}
+        commit = saver.conn.commit
+        async def fail_commit():
+            raise OSError('commit unavailable')
+        monkeypatch.setattr(saver.conn, 'commit', fail_commit)
+        with pytest.raises(OSError):
+            await saver.aput(config, empty_checkpoint(), {}, {})
+        assert not saver.conn.in_transaction
+        monkeypatch.setattr(saver.conn, 'commit', commit)
+        await storage.repository.apply_update('healthy', {'hphm1': 'test'})
+
+
+@pytest.mark.asyncio
+async def test_asgi_stream_cancel_keeps_snapshot_and_next_requests_work(tmp_path):
+    class LLM:
+        async def complete(self, *args, **kwargs):
+            return '{}'
+        async def stream_tokens(self, *args, **kwargs):
+            yield '<state>1002</state>partial'
+            await asyncio.Future()
+
+    settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'stream.sqlite3'), auto_migrate=True)
+    async with open_storage(settings) as storage:
+        backend = LangGraphBackend(build_graph(storage.checkpointer), llm=LLM())
+        service = AccidentAppService(backend, storage.repository, locks=storage.locks)
+        await service.chat(ChatInput('session', 'start'))
+        with anyio.CancelScope() as scope:
+            async with aclosing(service.stream(ChatInput('session', 'cancel-me'))) as events:
+                await anext(events)
+                scope.cancel()
+                await anyio.sleep(0)
+        assert not storage.repository.database.connection.in_transaction
+        assert not storage.checkpointer.conn.in_transaction
+        async with storage.repository.database.transaction() as conn:
+            row = await conn.one('SELECT version FROM agent_sessions WHERE session_id=?', ('session',))
+            assert row['version'] == 1
+        assert not service._locks._entries
+        assert await service.set_info('other', 'hphm1', 'working') == ('200', '')
+        await service.delete_session('session', backend.reset_checkpoint)
+        await service.delete_session('other', backend.reset_checkpoint)

test/storage/test_storage_config_and_locks.py

@@ -6,6 +6,7 @@ from uuid import uuid4
 import pytest
 
 from src.core.config import Config
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from src.storage.runtime import StorageSettings, open_storage
 from .test_persistence import run
 
@@ -36,7 +37,7 @@ def test_sqlite_requires_migration_and_does_not_fall_back(tmp_path):
                 pytest.fail('Uninitialized durable storage must not start')
         async with open_storage(replace(config,auto_migrate=True)): pass
         async with open_storage(config) as runtime:
-            assert runtime.checkpointer.__class__.__name__=='AsyncSqliteSaver'
+            assert isinstance(runtime.checkpointer, AsyncSqliteSaver)
     run(scenario())