190 lines
8.1 KiB
Python
190 lines
8.1 KiB
Python
"""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)
|