|
|
@@ -0,0 +1,322 @@
|
|
|
+"""公共生命周期:启动失败、取消、重复进入以及真实 SQLite 资源释放。"""
|
|
|
+import asyncio
|
|
|
+from concurrent.futures import ThreadPoolExecutor
|
|
|
+import threading
|
|
|
+from types import SimpleNamespace
|
|
|
+
|
|
|
+import pytest
|
|
|
+from langchain_core.messages import AIMessageChunk
|
|
|
+
|
|
|
+import config
|
|
|
+from agents import agent_memory, agent_registry
|
|
|
+from db import db_conn, chat_store, billing_store
|
|
|
+from runtime import environment
|
|
|
+from runtime.environment import RuntimeEnvironment
|
|
|
+from runtime.execution import execute
|
|
|
+from scheduler import scheduler
|
|
|
+
|
|
|
+
|
|
|
+@pytest.fixture
|
|
|
+def lifecycle(monkeypatch):
|
|
|
+ calls = []
|
|
|
+ def mark(name):
|
|
|
+ def call(*args, **kwargs):
|
|
|
+ calls.append(name)
|
|
|
+ return call
|
|
|
+ def amark(name):
|
|
|
+ async def call(*args, **kwargs):
|
|
|
+ calls.append(name)
|
|
|
+ return call
|
|
|
+ for name in ('setup_logging', 'load_extensions', 'approval_names'):
|
|
|
+ monkeypatch.setattr(environment, name, mark(name))
|
|
|
+ monkeypatch.setattr(config, 'log_config_sources', mark('config'))
|
|
|
+ monkeypatch.setattr(db_conn, 'init_schema', mark('schema'))
|
|
|
+ monkeypatch.setattr(db_conn, 'health', lambda: {'status': 'ok'})
|
|
|
+ monkeypatch.setattr(db_conn, 'close', mark('db_close'))
|
|
|
+ monkeypatch.setattr(environment.pricing, 'seed_default_pricing', mark('pricing'))
|
|
|
+ monkeypatch.setattr(environment.run_store, 'recover_interrupted', mark('recover'))
|
|
|
+ monkeypatch.setattr(agent_registry, 'invalidate', mark('invalidate'))
|
|
|
+ monkeypatch.setattr(agent_memory, 'init_checkpointer', amark('checkpoint'))
|
|
|
+ monkeypatch.setattr(agent_memory, 'close_checkpointer', amark('checkpoint_close'))
|
|
|
+ monkeypatch.setattr(scheduler, 'start', amark('scheduler_start'))
|
|
|
+ monkeypatch.setattr(scheduler, 'shutdown', amark('scheduler_stop'))
|
|
|
+ return calls
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_default_environment_and_sequential_reuse(lifecycle):
|
|
|
+ env = RuntimeEnvironment(prewarm=False, billing_retry=False)
|
|
|
+ for _ in range(2):
|
|
|
+ lifecycle.clear()
|
|
|
+ async with env as current:
|
|
|
+ assert current is env
|
|
|
+ assert lifecycle[-4:] == ['pricing', 'recover', 'invalidate', 'checkpoint']
|
|
|
+ assert lifecycle[-3:] == ['invalidate', 'checkpoint_close', 'db_close']
|
|
|
+ assert 'scheduler_start' not in lifecycle
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_nested_environment_does_not_close_outer(lifecycle):
|
|
|
+ env = RuntimeEnvironment(prewarm=False, billing_retry=False)
|
|
|
+ async with env:
|
|
|
+ for other in (env, RuntimeEnvironment()):
|
|
|
+ with pytest.raises(RuntimeError, match='已有活动'):
|
|
|
+ async with other:
|
|
|
+ pytest.fail('不应启动嵌套环境')
|
|
|
+ assert 'checkpoint_close' not in lifecycle
|
|
|
+ assert lifecycle.count('db_close') == 1
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+@pytest.mark.parametrize('stage', ['pricing', 'checkpoint', 'scheduler'])
|
|
|
+async def test_partial_startup_failure_cleans_and_allows_restart(stage, lifecycle, monkeypatch):
|
|
|
+ def fail():
|
|
|
+ raise ValueError(stage)
|
|
|
+ async def afail():
|
|
|
+ raise ValueError(stage)
|
|
|
+ with monkeypatch.context() as patch:
|
|
|
+ if stage == 'pricing':
|
|
|
+ patch.setattr(environment.pricing, 'seed_default_pricing', fail)
|
|
|
+ elif stage == 'checkpoint':
|
|
|
+ patch.setattr(agent_memory, 'init_checkpointer', afail)
|
|
|
+ else:
|
|
|
+ patch.setattr(scheduler, 'start', afail)
|
|
|
+ with pytest.raises(ValueError, match=stage):
|
|
|
+ async with RuntimeEnvironment(scheduler=True, prewarm=False, billing_retry=False):
|
|
|
+ pytest.fail('启动失败后不应进入执行体')
|
|
|
+ assert lifecycle[-1] == 'db_close'
|
|
|
+ if stage != 'pricing':
|
|
|
+ assert 'checkpoint_close' in lifecycle
|
|
|
+ if stage == 'scheduler':
|
|
|
+ assert lifecycle.index('scheduler_stop') < lifecycle.index('checkpoint_close')
|
|
|
+ async with RuntimeEnvironment(prewarm=False, billing_retry=False):
|
|
|
+ pass
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_body_exception_and_cleanup_error_still_release_resources(lifecycle, monkeypatch):
|
|
|
+ async def fail_close():
|
|
|
+ lifecycle.append('checkpoint_close')
|
|
|
+ raise RuntimeError('close failed')
|
|
|
+ with monkeypatch.context() as patch:
|
|
|
+ patch.setattr(agent_memory, 'close_checkpointer', fail_close)
|
|
|
+ with pytest.raises(RuntimeError, match='close failed'):
|
|
|
+ async with RuntimeEnvironment(prewarm=False, billing_retry=False):
|
|
|
+ raise ValueError('body failed')
|
|
|
+ assert lifecycle[-1] == 'db_close'
|
|
|
+ async with RuntimeEnvironment(prewarm=False, billing_retry=False):
|
|
|
+ pass
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_cancel_during_startup_waits_for_initialization(lifecycle, monkeypatch):
|
|
|
+ started, finish = asyncio.Event(), asyncio.Event()
|
|
|
+ async def init():
|
|
|
+ started.set()
|
|
|
+ await finish.wait()
|
|
|
+ lifecycle.append('checkpoint')
|
|
|
+ monkeypatch.setattr(agent_memory, 'init_checkpointer', init)
|
|
|
+ async def run():
|
|
|
+ async with RuntimeEnvironment(prewarm=False, billing_retry=False):
|
|
|
+ pytest.fail('取消后不应进入执行体')
|
|
|
+ task = asyncio.create_task(run())
|
|
|
+ await started.wait()
|
|
|
+ task.cancel()
|
|
|
+ await asyncio.sleep(0)
|
|
|
+ task.cancel() # 清理期间再次取消,也要等待底层初始化收尾。
|
|
|
+ assert 'db_close' not in lifecycle
|
|
|
+ finish.set()
|
|
|
+ with pytest.raises(asyncio.CancelledError):
|
|
|
+ await task
|
|
|
+ assert lifecycle[-4:] == ['checkpoint', 'invalidate', 'checkpoint_close', 'db_close']
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_shutdown_waits_for_workers_before_closing_connections(lifecycle, monkeypatch):
|
|
|
+ from billing import meter
|
|
|
+ prewarm_started, billing_started = threading.Event(), threading.Event()
|
|
|
+ release = threading.Event()
|
|
|
+ def compile_agent():
|
|
|
+ prewarm_started.set()
|
|
|
+ assert release.wait(5)
|
|
|
+ lifecycle.append('compiled')
|
|
|
+ def retry():
|
|
|
+ billing_started.set()
|
|
|
+ assert release.wait(5)
|
|
|
+ lifecycle.append('settled')
|
|
|
+ monkeypatch.setattr(agent_registry, 'get_agent', compile_agent)
|
|
|
+ monkeypatch.setattr(meter, 'retry_pending', retry)
|
|
|
+ running = asyncio.Event()
|
|
|
+ async def run():
|
|
|
+ async with RuntimeEnvironment(scheduler=True):
|
|
|
+ running.set()
|
|
|
+ await asyncio.Event().wait()
|
|
|
+ task = asyncio.create_task(run())
|
|
|
+ try:
|
|
|
+ await running.wait()
|
|
|
+ assert await asyncio.to_thread(prewarm_started.wait, 2)
|
|
|
+ assert await asyncio.to_thread(billing_started.wait, 2)
|
|
|
+ task.cancel()
|
|
|
+ await asyncio.sleep(0)
|
|
|
+ await asyncio.sleep(0)
|
|
|
+ assert 'checkpoint_close' not in lifecycle
|
|
|
+ finally:
|
|
|
+ release.set()
|
|
|
+ with pytest.raises(asyncio.CancelledError):
|
|
|
+ await task
|
|
|
+ assert lifecycle.index('scheduler_stop') < lifecycle.index('checkpoint_close')
|
|
|
+ assert lifecycle.index('settled') < lifecycle.index('db_close')
|
|
|
+ assert lifecycle.index('compiled') < lifecycle.index('checkpoint_close')
|
|
|
+ assert not [t for t in asyncio.all_tasks() if t.get_name().startswith('runtime-')]
|
|
|
+
|
|
|
+
|
|
|
+class FakeAgent:
|
|
|
+ def __init__(self):
|
|
|
+ self.checkpointer = agent_memory.get_checkpointer()
|
|
|
+
|
|
|
+ async def aget_state(self, options):
|
|
|
+ return SimpleNamespace(tasks=[])
|
|
|
+
|
|
|
+ async def astream(self, *args, **kwargs):
|
|
|
+ yield {'type': 'messages', 'data': (AIMessageChunk(content='报告完成', usage_metadata={
|
|
|
+ 'input_tokens': 8, 'output_tokens': 2, 'total_tokens': 10}), {})}
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_sqlite_script_executes_bills_and_reopens(monkeypatch, tmp_path):
|
|
|
+ # 独立业务库与 checkpoint,不使用线上配置,也不调用模型网络。
|
|
|
+ adapter = db_conn._SqliteAdapter(str(tmp_path / 'business.sqlite'))
|
|
|
+ monkeypatch.setattr(db_conn, '_adapter', adapter)
|
|
|
+ monkeypatch.setattr(db_conn, '_schema_ready', False)
|
|
|
+ monkeypatch.setattr(agent_memory, '_saver', None)
|
|
|
+ monkeypatch.setattr(agent_memory, '_kind', 'sqlite')
|
|
|
+ monkeypatch.setattr(config, 'CHECKPOINT_DB_PATH', str(tmp_path / 'checkpoint.sqlite'))
|
|
|
+ monkeypatch.setattr(agent_registry._slot, '_compile', FakeAgent)
|
|
|
+ schema = config.ROOT / 'db' / 'schema_sqlite.sql'
|
|
|
+ monkeypatch.setattr(adapter, 'schema_file', lambda: schema)
|
|
|
+ monkeypatch.setattr(config, 'ROOT', tmp_path)
|
|
|
+ previous = None
|
|
|
+ env = RuntimeEnvironment(prewarm=False, billing_retry=False)
|
|
|
+ for _ in range(2):
|
|
|
+ async with env:
|
|
|
+ session = chat_store.create_session('runtime-test')
|
|
|
+ agent = agent_registry.get_agent()
|
|
|
+ assert agent is not previous
|
|
|
+ assert agent.checkpointer is agent_memory.get_checkpointer()
|
|
|
+ result = await execute(user_name='runtime-test', session_id=session['id'], message='生成报告')
|
|
|
+ assert result.status == 'success'
|
|
|
+ assert result.output == '报告完成'
|
|
|
+ assert result.usage['total_tokens'] == 10
|
|
|
+ assert billing_store.list_pricing()
|
|
|
+ assert db_conn.scalar('SELECT total_tokens FROM billing_usage WHERE session_id = ?',
|
|
|
+ (session['id'],)) == 10
|
|
|
+ conn = adapter._conn
|
|
|
+ saver_conn = agent.checkpointer.conn
|
|
|
+ previous = agent
|
|
|
+ assert adapter._conn is None
|
|
|
+ with pytest.raises(Exception, match='closed'):
|
|
|
+ conn.execute('SELECT 1')
|
|
|
+ with pytest.raises(ValueError, match='no active connection'):
|
|
|
+ await saver_conn.execute('SELECT 1')
|
|
|
+ assert not agent_memory.is_initialized()
|
|
|
+ assert agent_registry.peek_agent() is None
|
|
|
+
|
|
|
+
|
|
|
+def test_mysql_closes_connections_from_all_threads_and_reopens():
|
|
|
+ adapter = db_conn._MySqlAdapter()
|
|
|
+ class Connection:
|
|
|
+ def __init__(self):
|
|
|
+ self.closed = False
|
|
|
+ def close(self):
|
|
|
+ self.closed = True
|
|
|
+ def ping(self, **kwargs):
|
|
|
+ assert not self.closed
|
|
|
+ adapter._connect = lambda database: Connection()
|
|
|
+ barrier = threading.Barrier(3)
|
|
|
+ def connect():
|
|
|
+ conn = adapter._connection()
|
|
|
+ barrier.wait(timeout=3)
|
|
|
+ return conn
|
|
|
+ with ThreadPoolExecutor(max_workers=2) as pool:
|
|
|
+ tasks = [pool.submit(connect) for _ in range(2)]
|
|
|
+ main_conn = adapter._connection()
|
|
|
+ barrier.wait(timeout=3)
|
|
|
+ connections = [main_conn, *(task.result() for task in tasks)]
|
|
|
+ assert len({id(conn) for conn in connections}) == 3
|
|
|
+ adapter.close()
|
|
|
+ assert all(conn.closed for conn in connections)
|
|
|
+ assert pool.submit(adapter._connection).result() not in connections
|
|
|
+ assert adapter._connection() not in connections
|
|
|
+ adapter.close()
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_checkpoint_setup_failure_closes_partial_connection(monkeypatch, tmp_path):
|
|
|
+ import aiosqlite
|
|
|
+ from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
|
+ monkeypatch.setattr(agent_memory, '_saver', None)
|
|
|
+ monkeypatch.setattr(agent_memory, '_kind', 'sqlite')
|
|
|
+ monkeypatch.setattr(config, 'CHECKPOINT_DB_PATH', str(tmp_path / 'bad.sqlite'))
|
|
|
+ connected = []
|
|
|
+ original = aiosqlite.connect
|
|
|
+ async def connect(*args, **kwargs):
|
|
|
+ conn = await original(*args, **kwargs)
|
|
|
+ connected.append(conn)
|
|
|
+ return conn
|
|
|
+ async def fail(self):
|
|
|
+ raise ValueError('setup failed')
|
|
|
+ monkeypatch.setattr(aiosqlite, 'connect', connect)
|
|
|
+ monkeypatch.setattr(AsyncSqliteSaver, 'setup', fail)
|
|
|
+ with pytest.raises(agent_memory.CheckpointerConfigError, match='setup failed'):
|
|
|
+ await agent_memory.init_checkpointer()
|
|
|
+ assert not agent_memory.is_initialized()
|
|
|
+ with pytest.raises(ValueError, match='no active connection'):
|
|
|
+ await connected[0].execute('SELECT 1')
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_scheduler_start_cancellation_waits_for_database_read(lifecycle, monkeypatch):
|
|
|
+ started, finish = asyncio.Event(), asyncio.Event()
|
|
|
+ async def start():
|
|
|
+ started.set()
|
|
|
+ await finish.wait()
|
|
|
+ lifecycle.append('scheduler_loaded')
|
|
|
+ monkeypatch.setattr(scheduler, 'start', start)
|
|
|
+ async def run():
|
|
|
+ async with RuntimeEnvironment(scheduler=True, prewarm=False, billing_retry=False):
|
|
|
+ pytest.fail('取消后不应进入执行体')
|
|
|
+ task = asyncio.create_task(run())
|
|
|
+ await started.wait()
|
|
|
+ task.cancel()
|
|
|
+ await asyncio.sleep(0)
|
|
|
+ assert 'db_close' not in lifecycle
|
|
|
+ finish.set()
|
|
|
+ with pytest.raises(asyncio.CancelledError):
|
|
|
+ await task
|
|
|
+ assert lifecycle.index('scheduler_loaded') < lifecycle.index('scheduler_stop')
|
|
|
+ assert lifecycle.index('scheduler_stop') < lifecycle.index('db_close')
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_http_lifespan_uses_common_environment(lifecycle, monkeypatch):
|
|
|
+ import main
|
|
|
+ def env(**kwargs):
|
|
|
+ assert kwargs == {'scheduler': True}
|
|
|
+ return RuntimeEnvironment(**kwargs, prewarm=False, billing_retry=False)
|
|
|
+ monkeypatch.setattr(main, 'RuntimeEnvironment', env)
|
|
|
+ async with main.app.router.lifespan_context(main.app):
|
|
|
+ assert 'scheduler_start' in lifecycle
|
|
|
+ assert 'db_close' not in lifecycle
|
|
|
+ assert lifecycle[-4:] == ['scheduler_stop', 'invalidate', 'checkpoint_close', 'db_close']
|
|
|
+
|
|
|
+
|
|
|
+def test_environment_reopens_across_event_loops(lifecycle):
|
|
|
+ env = RuntimeEnvironment(prewarm=False, billing_retry=False, recover_interrupted=False)
|
|
|
+ async def run():
|
|
|
+ async with env:
|
|
|
+ pass
|
|
|
+ asyncio.run(run())
|
|
|
+ asyncio.run(run())
|
|
|
+ assert lifecycle.count('checkpoint') == 2
|
|
|
+ assert lifecycle.count('db_close') == 2
|
|
|
+ assert 'recover' not in lifecycle
|