"""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)