From e6a96ed135a7f4b299b4b1945feac5cf307b210b Mon Sep 17 00:00:00 2001 From: Eric Wang Date: Mon, 26 Oct 2026 16:50:57 +0800 Subject: [PATCH] fix(storage): finish SQLite transactions before propagating cancellation --- docs/agent-storage.md | 6 + src/api/endpoints.py | 11 +- src/storage/cancellation.py | 33 +++ src/storage/database.py | 7 +- src/storage/runtime.py | 4 +- src/storage/sqlite_saver.py | 41 ++++ test/storage/test_sqlite_cancellation.py | 189 ++++++++++++++++++ test/storage/test_storage_config_and_locks.py | 3 +- 8 files changed, 288 insertions(+), 6 deletions(-) create mode 100644 src/storage/cancellation.py create mode 100644 src/storage/sqlite_saver.py create mode 100644 test/storage/test_sqlite_cancellation.py diff --git a/docs/agent-storage.md b/docs/agent-storage.md index b591e7a..8ef99ad 100644 --- a/docs/agent-storage.md +++ b/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 写入/删除中的 +取消、重复取消及后续请求可用性，验证不依赖服务重启释放残留事务。 diff --git a/src/api/endpoints.py b/src/api/endpoints.py index 314375b..c0cc50e 100644 --- a/src/api/endpoints.py +++ b/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} " diff --git a/src/storage/cancellation.py b/src/storage/cancellation.py new file mode 100644 index 0000000..f2105dc --- /dev/null +++ b/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 diff --git a/src/storage/database.py b/src/storage/database.py index 89d5829..0909a28 100644 --- a/src/storage/database.py +++ b/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 diff --git a/src/storage/runtime.py b/src/storage/runtime.py index 857c519..9d7fd2b 100644 --- a/src/storage/runtime.py +++ b/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: diff --git a/src/storage/sqlite_saver.py b/src/storage/sqlite_saver.py new file mode 100644 index 0000000..34d5845 --- /dev/null +++ b/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), + ) diff --git a/test/storage/test_sqlite_cancellation.py b/test/storage/test_sqlite_cancellation.py new file mode 100644 index 0000000..54179a0 --- /dev/null +++ b/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 '1002partial' + 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) diff --git a/test/storage/test_storage_config_and_locks.py b/test/storage/test_storage_config_and_locks.py index 279533b..868ebd0 100644 --- a/test/storage/test_storage_config_and_locks.py +++ b/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())