feat(storage): add selectable SQLite and PostgreSQL persistence
- Select one storage backend per instance via AGENT_STORAGE, retaining memory as the backward-compatible default. Manage the selected async repository, official checkpointer and connection pools in lifespan; fail startup on unavailable or uninitialized storage without fallback. - Add SQLite and PostgreSQL adapters, schema initialization CLI, sample PostgreSQL compose configuration and a Windows-compatible serve entry. SQLite enforces single-instance file ownership; PostgreSQL uses a dedicated pool for session advisory locks plus version/epoch fencing. - Make repository IO and commit callbacks awaitable. Persist business fields, a complete committed graph snapshot and turn result/version in one business transaction before success. Rebuild execution checkpoints from that snapshot, discarding uncommitted progress after interruption. - Support optional Idempotency-Key on durable chat requests. Replay committed results and return 409 on changed request content; delete removes business/history data and replay records while retaining only an epoch/version tombstone to reject stale commits. - Document migration, database switching boundaries, backup/deployment expectations and crash recovery. Ignore generated database files. Validation: 315 tests passed with real SQLite and isolated PostgreSQL, including abrupt child-process exits before/after business commit; both backends passed CLI migration and actual HTTP server restart smoke. Existing API tests were rechecked after isolating them from user storage. An initial Windows CreateProcess permission failure did not recur in the full diagnostic run; failure assertions were not skipped or relaxed.
This commit is contained in:
Eric Wang
@@ -0,0 +1 @@
|
||||
"""SQLite/PostgreSQL durability integration tests."""
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Integration-test child: exit abruptly at a durable turn's commit boundary."""
|
||||
import asyncio
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
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.runtime import StorageSettings, open_storage
|
||||
|
||||
|
||||
class LLM:
|
||||
async def complete(self, system, user, *, history=None):
|
||||
if '【字段提取】' in system:
|
||||
return '{"has_injury":false}'
|
||||
return '<state>1002</state>crash-boundary-reply'
|
||||
|
||||
|
||||
async def main():
|
||||
kind, path, session, boundary = sys.argv[1:]
|
||||
settings = StorageSettings(kind=kind, sqlite_path=path,
|
||||
postgres_dsn=os.getenv('AGENT_TEST_POSTGRES_DSN',''), pool_size=2)
|
||||
async with open_storage(settings) as storage:
|
||||
service = AccidentAppService(LangGraphBackend(build_graph(storage.checkpointer), llm=LLM()),
|
||||
storage.repository, locks=storage.locks)
|
||||
original = storage.repository.commit_turn
|
||||
async def crash(*args):
|
||||
if boundary=='before':
|
||||
os._exit(23)
|
||||
await original(*args)
|
||||
os._exit(24)
|
||||
storage.repository.commit_turn = crash
|
||||
await service.chat(ChatInput(session,'继续办理',request_key='crash-key'))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
factory=asyncio.SelectorEventLoop if os.name=='nt' else None
|
||||
with asyncio.Runner(loop_factory=factory) as runner:
|
||||
runner.run(main())
|
||||
@@ -0,0 +1,54 @@
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import replace
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
|
||||
from src.api.endpoints import router
|
||||
from src.core import fastgpt_client as wiring
|
||||
from src.core.config import Config
|
||||
from src.storage.runtime import open_storage
|
||||
from .test_persistence import RecordingLLM, run, settings
|
||||
|
||||
|
||||
def test_http_restart_replay_headers_and_switching(settings, monkeypatch):
|
||||
import src.agent.llm as llm_module
|
||||
monkeypatch.setattr(Config, 'AGENT_BACKEND', 'langgraph')
|
||||
monkeypatch.setattr(Config, 'LLM_PROVIDER', 'fake')
|
||||
monkeypatch.setattr(Config, 'AGENT_STORAGE', settings.kind)
|
||||
monkeypatch.setattr(Config, 'AGENT_SQLITE_PATH', settings.sqlite_path)
|
||||
monkeypatch.setattr(Config, 'AGENT_POSTGRES_DSN', settings.postgres_dsn)
|
||||
monkeypatch.setattr(Config, 'AGENT_DB_POOL_SIZE', settings.pool_size)
|
||||
monkeypatch.setattr(Config, 'AGENT_DB_AUTO_MIGRATE', True)
|
||||
monkeypatch.setattr(llm_module, 'build_llm_service', lambda: RecordingLLM())
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
@asynccontextmanager
|
||||
async def client():
|
||||
async with wiring.lifespan(app):
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as value:
|
||||
yield value
|
||||
async def scenario():
|
||||
import uuid
|
||||
payload = {'sessionId': uuid.uuid4().hex, 'timeStamp':'20261022120000'}
|
||||
request = {**payload, 'text':'新对话'}
|
||||
async with client() as http:
|
||||
assert (await http.post('/set_info', json={**payload,'key':'hphm1','value':'http-plate'})).json()['code']=='200'
|
||||
first = (await http.post('/chat', json=request, headers={'Idempotency-Key':'http-key'})).json()
|
||||
assert first['nextStageCode']=='1001'
|
||||
monkeypatch.setattr(Config, 'AGENT_DB_AUTO_MIGRATE', False)
|
||||
async with client() as http:
|
||||
assert (await http.post('/chat', json=request, headers={'Idempotency-Key':'http-key'})).json()==first
|
||||
conflict = await http.post('/chat', json={**request,'text':'different'}, headers={'Idempotency-Key':'http-key'})
|
||||
assert conflict.json()['code']=='409'
|
||||
stream = await http.post('/chat?stream=true', json={**payload,'text':'继续办理','needFormUpdate':False})
|
||||
assert 'event: done' in stream.text
|
||||
assert 'event: formUpdate' not in stream.text
|
||||
assert '"0"' == (await http.post('/get_info', json={**payload,'key':'jdcsl'})).json()['value']
|
||||
assert (await http.request('DELETE','/delete_session',json=payload)).json()['code']=='200'
|
||||
async with client() as http:
|
||||
assert (await http.post('/get_info',json={**payload,'key':'hphm1'})).json()['value']=='""'
|
||||
assert (await http.post('/chat',json=request,headers={'Idempotency-Key':'http-key'})).json()['nextStageCode']=='1001'
|
||||
run(scenario())
|
||||
@@ -0,0 +1,221 @@
|
||||
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())
|
||||
@@ -0,0 +1,88 @@
|
||||
import asyncio
|
||||
import os
|
||||
from dataclasses import replace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from src.core.config import Config
|
||||
from src.storage.runtime import StorageSettings, open_storage
|
||||
from .test_persistence import run
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def postgres_settings():
|
||||
dsn = os.getenv('AGENT_TEST_POSTGRES_DSN')
|
||||
if not dsn:
|
||||
pytest.skip('Set AGENT_TEST_POSTGRES_DSN to run PostgreSQL integration tests')
|
||||
return StorageSettings(kind='postgresql',postgres_dsn=dsn,auto_migrate=True,pool_size=2)
|
||||
|
||||
|
||||
def test_sqlite_rejects_second_instance_and_releases_guard(tmp_path):
|
||||
async def scenario():
|
||||
config = StorageSettings(kind='sqlite',sqlite_path=str(tmp_path/'single.db'),auto_migrate=True)
|
||||
async with open_storage(config):
|
||||
with pytest.raises(RuntimeError, match='active instance'):
|
||||
async with open_storage(config): pass
|
||||
async with open_storage(config): pass
|
||||
run(scenario())
|
||||
|
||||
|
||||
def test_sqlite_requires_migration_and_does_not_fall_back(tmp_path):
|
||||
async def scenario():
|
||||
config = StorageSettings(kind='sqlite',sqlite_path=str(tmp_path/'empty.db'))
|
||||
with pytest.raises(Exception):
|
||||
async with open_storage(config):
|
||||
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'
|
||||
run(scenario())
|
||||
|
||||
|
||||
def test_postgresql_locks_serialize_independent_application_instances(postgres_settings):
|
||||
settings = postgres_settings
|
||||
async def scenario():
|
||||
session = uuid4().hex
|
||||
async with open_storage(settings) as first, open_storage(settings) as second:
|
||||
entered = asyncio.Event()
|
||||
async def waiting():
|
||||
async with second.locks.hold(session): entered.set()
|
||||
async with first.locks.hold(session):
|
||||
task=asyncio.create_task(waiting())
|
||||
await asyncio.sleep(0.15)
|
||||
assert not entered.is_set()
|
||||
# 同一个锁池中，不同会话仍能执行。
|
||||
async with second.locks.hold(session+'-other'):
|
||||
pass
|
||||
await asyncio.wait_for(task,3)
|
||||
assert entered.is_set()
|
||||
run(scenario())
|
||||
|
||||
|
||||
def test_postgresql_cancelled_lock_holder_releases_server_lock(postgres_settings):
|
||||
settings = postgres_settings
|
||||
async def scenario():
|
||||
session=uuid4().hex
|
||||
async with open_storage(settings) as first, open_storage(settings) as second:
|
||||
ready=asyncio.Event()
|
||||
async def holding():
|
||||
async with first.locks.hold(session):
|
||||
ready.set()
|
||||
await asyncio.Future()
|
||||
task=asyncio.create_task(holding())
|
||||
await ready.wait()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError): await task
|
||||
async with asyncio.timeout(3):
|
||||
async with second.locks.hold(session): pass
|
||||
run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize('kind', ['unknown', 'sqlite', 'postgresql'])
|
||||
def test_invalid_storage_config_is_rejected(monkeypatch,kind):
|
||||
monkeypatch.setattr(Config,'AGENT_BACKEND','langgraph')
|
||||
monkeypatch.setattr(Config,'AGENT_STORAGE',kind)
|
||||
monkeypatch.setattr(Config,'AGENT_SQLITE_PATH',':memory:')
|
||||
monkeypatch.setattr(Config,'AGENT_POSTGRES_DSN','')
|
||||
with pytest.raises(ValueError):Config.validate()
|
||||
Reference in New Issue
Block a user