222 lines
11 KiB
Python
222 lines
11 KiB
Python
|
import asyncio
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import subprocess
|
|||
|
|
import sys
|
|||
|
|
from pathlib import Path
|
|||
|
|
from contextlib import aclosing, asynccontextmanager
|
|||
|
|
from dataclasses import replace
|
|||
|
|
from uuid import uuid4
|
|||
|
|
|
|||
|
|
import pytest
|
|||
|
|
|
|||
|
|
from src.agent.graph import build_graph
|
|||
|
|
from src.backends.chat import ChatBackendConflictError, ChatInput, FormUpdate
|
|||
|
|
from src.backends.langgraph import LangGraphBackend
|
|||
|
|
from src.services.app_service import AccidentAppService
|
|||
|
|
from src.storage.runtime import StorageSettings, open_storage
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run(coroutine):
|
|||
|
|
# psycopg 在 Windows 要求 SelectorEventLoop；Linux 使用标准 runner。
|
|||
|
|
factory = asyncio.SelectorEventLoop if os.name == 'nt' else None
|
|||
|
|
with asyncio.Runner(loop_factory=factory) as runner:
|
|||
|
|
return runner.run(coroutine)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.fixture(params=['sqlite', 'postgresql'])
|
|||
|
|
def settings(request, tmp_path):
|
|||
|
|
if request.param == 'postgresql':
|
|||
|
|
dsn = os.getenv('AGENT_TEST_POSTGRES_DSN')
|
|||
|
|
if not dsn:
|
|||
|
|
pytest.skip('Set AGENT_TEST_POSTGRES_DSN to an isolated PostgreSQL database')
|
|||
|
|
return StorageSettings(kind='postgresql', postgres_dsn=dsn, auto_migrate=True, pool_size=2)
|
|||
|
|
return StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'agent.sqlite3'), auto_migrate=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RecordingLLM:
|
|||
|
|
def __init__(self):
|
|||
|
|
self.histories = []
|
|||
|
|
|
|||
|
|
async def complete(self, system, user, *, history=None):
|
|||
|
|
if '【字段提取】' in system:
|
|||
|
|
return '{"has_injury":false,"vehicle_count":0}'
|
|||
|
|
self.histories.append(history or [])
|
|||
|
|
return '<state>1002</state>已记录'
|
|||
|
|
|
|||
|
|
async def stream_tokens(self, system, user, *, history=None):
|
|||
|
|
self.histories.append(history or [])
|
|||
|
|
yield '<state>1002</state>'
|
|||
|
|
yield '已记录'
|
|||
|
|
|
|||
|
|
|
|||
|
|
@asynccontextmanager
|
|||
|
|
async def service_for(settings, llm=None):
|
|||
|
|
async with open_storage(settings) as storage:
|
|||
|
|
backend = LangGraphBackend(build_graph(storage.checkpointer), llm=llm or RecordingLLM())
|
|||
|
|
yield AccidentAppService(backend, storage.repository, locks=storage.locks), storage
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def session_row(storage, session):
|
|||
|
|
async with storage.repository.database.transaction() as conn:
|
|||
|
|
return await conn.one('SELECT * FROM agent_sessions WHERE session_id=?', (session,))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_restart_keeps_fields_phase_and_history(settings):
|
|||
|
|
async def scenario():
|
|||
|
|
session = uuid4().hex
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
await service.set_info(session, 'hphm1', 'test-plate')
|
|||
|
|
first = await service.chat(ChatInput(session, '新对话'))
|
|||
|
|
assert first.content.startswith('<state>1001</state>')
|
|||
|
|
await service.chat(ChatInput(session, '继续办理'))
|
|||
|
|
values = await storage.repository.load(session)
|
|||
|
|
assert values['ywrysw'] is False
|
|||
|
|
assert type(values['jdcsl']) is int and values['jdcsl'] == 0
|
|||
|
|
llm = RecordingLLM()
|
|||
|
|
async with service_for(replace(settings, auto_migrate=False), llm) as (service, storage):
|
|||
|
|
assert json.loads(await service.get_info(session, 'hphm1')) == 'test-plate'
|
|||
|
|
result = await service.chat(ChatInput(session, '重启后的新输入'))
|
|||
|
|
assert result.content.startswith('<state>1002</state>')
|
|||
|
|
assert any(m['content'] == '继续办理' for m in llm.histories[0])
|
|||
|
|
row = await session_row(storage, session)
|
|||
|
|
graph = json.loads(row['graph_state'])
|
|||
|
|
assert graph['turn_count'] == 3
|
|||
|
|
assert graph['dialogue'][-2]['content'] == '重启后的新输入'
|
|||
|
|
assert graph['state_version'] == row['version']
|
|||
|
|
run(scenario())
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.parametrize('stream', [False, True])
|
|||
|
|
def test_idempotency_replay_after_restart_and_conflict(settings, stream):
|
|||
|
|
async def scenario():
|
|||
|
|
session = uuid4().hex
|
|||
|
|
request = ChatInput(session, '继续办理', request_key='stable-key', need_form_update=True)
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
await service.chat(ChatInput(session, '新对话'))
|
|||
|
|
if stream:
|
|||
|
|
first = [event async for event in service.stream(request)]
|
|||
|
|
assert any(isinstance(event, FormUpdate) for event in first)
|
|||
|
|
else:
|
|||
|
|
first = await service.chat(request)
|
|||
|
|
before = await session_row(storage, session)
|
|||
|
|
llm = RecordingLLM()
|
|||
|
|
async with service_for(settings, llm) as (service, storage):
|
|||
|
|
result = await service.chat(request)
|
|||
|
|
assert result.content == '<state>1002</state>已记录'
|
|||
|
|
assert llm.histories == []
|
|||
|
|
assert (await session_row(storage, session))['version'] == before['version']
|
|||
|
|
with pytest.raises(ChatBackendConflictError):
|
|||
|
|
await service.chat(replace(request, text='different'))
|
|||
|
|
run(scenario())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_checkpoint_ahead_of_business_commit_is_discarded(settings):
|
|||
|
|
async def scenario():
|
|||
|
|
session = uuid4().hex
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
await service.chat(ChatInput(session, '新对话'))
|
|||
|
|
original = storage.repository.commit_turn
|
|||
|
|
async def fail(*args): raise ConnectionError('business commit failed')
|
|||
|
|
async def unavailable(*args): raise ConnectionError('journal cleanup failed')
|
|||
|
|
storage.repository.commit_turn = fail
|
|||
|
|
storage.repository.abort_turn = unavailable
|
|||
|
|
with pytest.raises(ConnectionError):
|
|||
|
|
await service.chat(ChatInput(session, '未提交的输入', request_key='retry'))
|
|||
|
|
storage.repository.commit_turn = original
|
|||
|
|
llm = RecordingLLM()
|
|||
|
|
async with service_for(settings, llm) as (service, storage):
|
|||
|
|
await service.chat(ChatInput(session, '未提交的输入', request_key='retry'))
|
|||
|
|
assert all(m['content'] != '未提交的输入' for m in llm.histories[0])
|
|||
|
|
assert json.loads((await session_row(storage, session))['graph_state'])['turn_count'] == 2
|
|||
|
|
run(scenario())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_committed_business_state_restores_missing_checkpoint(settings):
|
|||
|
|
async def scenario():
|
|||
|
|
session = uuid4().hex
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
await service.chat(ChatInput(session, '新对话'))
|
|||
|
|
await service.chat(ChatInput(session, '继续办理'))
|
|||
|
|
await storage.checkpointer.adelete_thread(session)
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
reply = await service.chat(ChatInput(session, '再次输入'))
|
|||
|
|
assert reply.content.startswith('<state>1002</state>')
|
|||
|
|
assert json.loads((await session_row(storage, session))['graph_state'])['turn_count'] == 3
|
|||
|
|
run(scenario())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_disconnect_after_commit_can_replay_without_advancing(settings):
|
|||
|
|
async def scenario():
|
|||
|
|
session = uuid4().hex
|
|||
|
|
request = ChatInput(session, '继续办理', request_key='disconnect')
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
await service.chat(ChatInput(session, '新对话'))
|
|||
|
|
initial = await session_row(storage, session)
|
|||
|
|
async with aclosing(service.stream(request)) as events:
|
|||
|
|
await anext(events)
|
|||
|
|
async with asyncio.timeout(5):
|
|||
|
|
while (await session_row(storage, session))['version'] == initial['version']:
|
|||
|
|
await asyncio.sleep(0.01)
|
|||
|
|
# 此时图及业务提交已完成，但消费者没有读到回合末尾。
|
|||
|
|
before = await session_row(storage, session)
|
|||
|
|
llm = RecordingLLM()
|
|||
|
|
async with service_for(settings, llm) as (service, storage):
|
|||
|
|
await service.chat(request)
|
|||
|
|
assert llm.histories == []
|
|||
|
|
assert (await session_row(storage, session))['version'] == before['version']
|
|||
|
|
run(scenario())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_delete_then_restart_and_stale_commit_cannot_revive_data(settings):
|
|||
|
|
async def scenario():
|
|||
|
|
session = uuid4().hex
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
await service.set_info(session, 'hphm1', 'old-plate')
|
|||
|
|
await service.chat(ChatInput(session, '新对话', request_key='old-key'))
|
|||
|
|
ticket = await storage.repository.begin_turn(session, None, 'pending')
|
|||
|
|
await service.delete_session(session, service._backend.reset_checkpoint)
|
|||
|
|
with pytest.raises(ChatBackendConflictError):
|
|||
|
|
await storage.repository.commit_turn(ticket, {'hphm1':'revived'}, {})
|
|||
|
|
assert await storage.checkpointer.aget_tuple({'configurable': {'thread_id':session}}) is None
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
assert json.loads(await service.get_info(session, 'hphm1')) == ''
|
|||
|
|
reply = await service.chat(ChatInput(session, '新会话', request_key='old-key'))
|
|||
|
|
assert reply.content.startswith('<state>1001</state>')
|
|||
|
|
assert json.loads((await session_row(storage, session))['graph_state'])['turn_count'] == 1
|
|||
|
|
run(scenario())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_business_version_conflict_rejects_stale_turn(settings):
|
|||
|
|
async def scenario():
|
|||
|
|
session = uuid4().hex
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
ticket = await storage.repository.begin_turn(session, None, 'pending')
|
|||
|
|
await storage.repository.apply_update(session, {'jdcsl':2})
|
|||
|
|
with pytest.raises(ChatBackendConflictError):
|
|||
|
|
await storage.repository.commit_turn(ticket, {'jdcsl':1}, {})
|
|||
|
|
assert (await storage.repository.load(session))['jdcsl'] == 2
|
|||
|
|
run(scenario())
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.parametrize('boundary', ['before', 'after'])
|
|||
|
|
def test_abrupt_process_exit_at_commit_boundary(settings, boundary):
|
|||
|
|
session = uuid4().hex
|
|||
|
|
async def initialize():
|
|||
|
|
async with service_for(settings) as (service, storage):
|
|||
|
|
await service.chat(ChatInput(session, '新对话'))
|
|||
|
|
run(initialize())
|
|||
|
|
worker = Path(__file__).with_name('crash_worker.py')
|
|||
|
|
result = subprocess.run([sys.executable, str(worker), settings.kind, settings.sqlite_path, session, boundary],
|
|||
|
|
capture_output=True, timeout=45)
|
|||
|
|
assert result.returncode == (23 if boundary=='before' else 24), result.stderr.decode(errors='replace')
|
|||
|
|
async def recover():
|
|||
|
|
llm = RecordingLLM()
|
|||
|
|
async with service_for(replace(settings,auto_migrate=False),llm) as (service,storage):
|
|||
|
|
reply=await service.chat(ChatInput(session,'继续办理',request_key='crash-key'))
|
|||
|
|
assert reply.content.startswith('<state>1002</state>')
|
|||
|
|
assert bool(llm.histories) is (boundary=='before')
|
|||
|
|
row=await session_row(storage,session)
|
|||
|
|
assert json.loads(row['graph_state'])['turn_count']==2
|
|||
|
|
run(recover())
|