vent_agent_2.0/tests/test_extension_components.py

"""扩展服务生命周期、中间件排序与真实图装配。"""
import asyncio
from contextlib import asynccontextmanager
from types import SimpleNamespace

import pytest
from langchain.agents.middleware.types import AgentMiddleware
from langchain_core.messages import AIMessage

import config
from agents import agent_factory, agent_memory, agent_registry
from db import db_conn, chat_store
from extensions import ExtensionRegistry, get_service
from extensions import registry as registry_module
from extensions.services import ServiceContainer, service_view
from runtime import environment
from runtime.environment import RuntimeEnvironment
from runtime.execution import execute
from tools import registry as tools


@asynccontextmanager
async def resource(dependencies):
    yield object()


def middleware(name):
    return lambda services: type(name, (AgentMiddleware,), {})()


@pytest.fixture
def isolated(monkeypatch, tmp_path):
    # 此文件验证通用扩展生命周期；临时根没有业务技能，独立关闭主通风机扩展。
    monkeypatch.setenv('MAIN_FAN_ENABLED', '0')
    (tmp_path / 'skills').mkdir()
    schema = config.ROOT / 'db' / 'schema_sqlite.sql'
    adapter = db_conn._SqliteAdapter(str(tmp_path / 'business.sqlite'))
    monkeypatch.setattr(adapter, 'schema_file', lambda: schema)
    monkeypatch.setattr(db_conn, '_adapter', adapter)
    monkeypatch.setattr(db_conn, '_schema_ready', False)
    monkeypatch.setattr(config, 'ROOT', tmp_path)
    monkeypatch.setattr(agent_memory, '_saver', None)
    monkeypatch.setattr(agent_memory, '_kind', 'memory')
    yield
    db_conn.close()


def use_catalog(monkeypatch, catalog):
    catalog.validate()
    monkeypatch.setattr(environment, 'load_extensions', lambda: catalog)
    monkeypatch.setattr(environment, 'approval_names', lambda: set())


@pytest.mark.parametrize('kind', ['service', 'middleware', 'tool'])
def test_missing_service_dependencies(kind):
    catalog = ExtensionRegistry()
    if kind == 'service':
        catalog.register_service('retriever', resource, requires=['missing'])
    elif kind == 'middleware':
        catalog.register_middleware('Recall', middleware('Recall'), requires=['missing'])
    else:
        def search(query: str) -> str:
            """检索。"""
            return query
        catalog.register_tool(search, requires=['missing'])
    with pytest.raises(ValueError, match='未注册'):
        catalog.validate()


@pytest.mark.parametrize('kind', ['service', 'middleware'])
def test_cycles_and_duplicate_names(kind):
    catalog = ExtensionRegistry()
    if kind == 'service':
        catalog.register_service('one', resource, requires=['two'])
        catalog.register_service('two', resource, requires=['one'])
        with pytest.raises(ValueError, match='冲突'):
            catalog.register_service('one', resource)
    else:
        catalog.register_middleware('One', middleware('One'), after=['Two'])
        catalog.register_middleware('Two', middleware('Two'), after=['One'])
        with pytest.raises(ValueError, match='冲突'):
            catalog.register_middleware('One', middleware('One'))
    with pytest.raises(ValueError, match='循环'):
        catalog.validate()


def test_middleware_order_scope_and_independent_instances():
    catalog = ExtensionRegistry()
    catalog.register_agent(name='researcher', description='检索', prompt='检索', tools=[])
    catalog.register_middleware('Recall', middleware('Recall'), targets=['main', 'researcher'], before=['Budget'])
    catalog.register_middleware('Budget', middleware('Budget'), targets=['main'])
    catalog.register_middleware('ChildOnly', middleware('ChildOnly'), targets=['researcher'], after=['Recall'])
    catalog.validate()
    main = catalog.build_middleware('main')
    child = catalog.build_middleware('researcher')
    assert [m.name for m in main] == ['Recall', 'Budget']
    assert [m.name for m in child] == ['Recall', 'ChildOnly']
    assert main[0] is not child[0]
    assert main[0] is not catalog.build_middleware('main')[0]
    assert catalog.build_middleware('general-purpose') == []


@pytest.mark.parametrize('kwargs,match', [
    ({'targets': ['missing']}, '未注册目标'),
    ({'after': ['missing']}, '未注册项'),
])
def test_invalid_middleware_references(kwargs, match):
    catalog = ExtensionRegistry()
    catalog.register_middleware('Recall', middleware('Recall'), **kwargs)
    with pytest.raises(ValueError, match=match):
        catalog.validate()


def test_reserved_names_and_actual_middleware_name():
    catalog = ExtensionRegistry()
    with pytest.raises(ValueError, match='冲突'):
        catalog.register_middleware('SummarizationMiddleware', middleware('SummarizationMiddleware'))
    catalog.register_middleware('Recall', middleware('WrongName'))
    with pytest.raises(ValueError, match='同名'):
        catalog.build_middleware('main')


@pytest.mark.asyncio
async def test_service_order_shared_lookup_and_reverse_cleanup(isolated, monkeypatch):
    calls = []
    catalog = ExtensionRegistry()
    @asynccontextmanager
    async def database(deps):
        assert not deps
        value = object()
        calls.append('open_db')
        try:
            yield value
        finally:
            calls.append('close_db')
    @asynccontextmanager
    async def search(deps):
        calls.append('open_search')
        assert set(deps) == {'db'}
        with pytest.raises(TypeError):
            deps['other'] = object()
        try:
            yield SimpleNamespace(database=deps['db'])
        finally:
            calls.append('close_search')
    catalog.register_service('search', search, requires=['db'])
    catalog.register_service('db', database)
    use_catalog(monkeypatch, catalog)
    old = None
    for _ in range(2):
        async with RuntimeEnvironment(prewarm=False, billing_retry=False):
            assert get_service('search').database is get_service('db')
            assert get_service('db') is not old
            old = get_service('db')
        with pytest.raises(RuntimeError, match='尚未启动'):
            get_service('db')
    assert calls == ['open_db', 'open_search', 'close_search', 'close_db'] * 2


@pytest.mark.asyncio
@pytest.mark.parametrize('fail_at', ['service', 'middleware'])
async def test_startup_failure_releases_all_services(isolated, monkeypatch, fail_at):
    calls = []
    catalog = ExtensionRegistry()
    @asynccontextmanager
    async def db(deps):
        calls.append('open')
        try:
            yield object()
        finally:
            calls.append('close')
    @asynccontextmanager
    async def broken(deps):
        raise ValueError('unavailable')
        yield
    catalog.register_service('db', db)
    if fail_at == 'service':
        catalog.register_service('broken', broken, requires=['db'])
    else:
        catalog.register_middleware('Recall', lambda deps: object(), requires=['db'])
    use_catalog(monkeypatch, catalog)
    with pytest.raises(ValueError):
        async with RuntimeEnvironment(prewarm=False, billing_retry=False):
            pytest.fail('不可进入运行体')
    assert calls == ['open', 'close']
    assert not agent_memory.is_initialized()
    assert db_conn._adapter._conn is None
    with pytest.raises(RuntimeError, match='尚未启动'):
        get_service('db')
    use_catalog(monkeypatch, ExtensionRegistry())
    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
        pass


@pytest.mark.asyncio
async def test_cancel_during_service_start_waits_and_closes(isolated, monkeypatch):
    calls = []
    started, finish = asyncio.Event(), asyncio.Event()
    catalog = ExtensionRegistry()
    @asynccontextmanager
    async def slow(deps):
        started.set()
        await finish.wait()
        calls.append('open')
        try:
            yield object()
        finally:
            calls.append('close')
    catalog.register_service('slow', slow)
    use_catalog(monkeypatch, catalog)
    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()
    finish.set()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert calls == ['open', 'close']
    with pytest.raises(RuntimeError):
        get_service('slow')


def test_disabled_extension_skips_config_validation_and_setup(monkeypatch):
    def forbidden(*args):
        pytest.fail('停用模块不能执行校验或 setup')
    fake = SimpleNamespace(enabled=lambda: False, validate_config=forbidden, setup=forbidden)
    original = registry_module.importlib.import_module
    monkeypatch.setattr(registry_module, 'discover_modules', lambda: ['extensions.core.extension', 'optional'])
    monkeypatch.setattr(registry_module.importlib, 'import_module', lambda name: fake if name == 'optional' else original(name))
    tools.reset_for_tests()
    catalog = registry_module.load_extensions()
    assert catalog.disabled_modules == ['optional']
    assert catalog.modules == ['extensions.core.extension']


def test_extension_config_failure_prevents_publication(monkeypatch):
    def invalid():
        raise ValueError('missing endpoint')
    fake = SimpleNamespace(enabled=lambda: True, validate_config=invalid,
                           setup=lambda catalog: pytest.fail('配置错误后不能注册'))
    monkeypatch.setattr(registry_module, 'discover_modules', lambda: ['broken'])
    monkeypatch.setattr(registry_module.importlib, 'import_module', lambda name: fake)
    tools.reset_for_tests()
    with pytest.raises(ValueError, match='missing endpoint'):
        registry_module.load_extensions()
    assert registry_module._loaded is None
    assert not tools.get_tool_names()


@pytest.mark.asyncio
async def test_enabled_demo_runs_real_graph_without_network(isolated, monkeypatch):
    from test_runtime_regressions import ScriptedModel
    from extensions.context_demo.extension import lookup_demo_term
    from pydantic import Field
    class Model(ScriptedModel):
        seen: list[str] = Field(default_factory=list)
        def _stream(self, messages, *args, **kwargs):
            self.seen.append(str(messages[0].content))
            yield from super()._stream(messages, *args, **kwargs)
    monkeypatch.setenv('CONTEXT_DEMO_ENABLED', '1')
    tools.reset_for_tests()
    model = Model(responses=[
        AIMessage(content='', tool_calls=[{'id': 'term', 'name': 'lookup_demo_term', 'args': {'term': '风量'}}]),
        AIMessage(content='已查询术语表。'),
    ])
    monkeypatch.setattr('agents.model_factory.get_model', lambda: model)
    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
        assert '空气体积' in lookup_demo_term('风量')
        sid = chat_store.create_session('demo-user')['id']
        result = await execute(user_name='demo-user', session_id=sid, message='什么是风量？')
        assert result.status == 'success'
        assert '已查询术语表' in result.output
        assert model.seen and all('演示术语表可查询：风速、风量' in text for text in model.seen)
        assert model.seen[0].count('演示术语表可查询') == 1
        assert model.seen[1].count('演示术语表可查询') == 1
    assert agent_registry.peek_agent() is None
    with pytest.raises(RuntimeError, match='尚未启动'):
        lookup_demo_term('风量')


def test_agent_factory_attaches_middleware_only_to_requested_targets(monkeypatch):
    catalog = registry_module.load_extensions()
    # 使用独立目录副本，保留核心工具，避免修改共享扩展缓存。
    local = ExtensionRegistry()
    local.tools.update(catalog.tools)
    local.agents.update(catalog.agents)
    local.register_middleware('ChildOnly', middleware('ChildOnly'), targets=['demo-searcher'])
    local.register_middleware('MainOnly', middleware('MainOnly'))
    local.validate()
    monkeypatch.setattr(agent_factory, 'load_extensions', lambda: local)
    monkeypatch.setattr(agent_factory.model_factory, 'get_model', lambda: object())
    monkeypatch.setattr(agent_factory, '_build_backend', lambda: None)
    monkeypatch.setattr('deepagents.create_deep_agent', lambda **kwargs: kwargs)
    agent = agent_factory.build_agent()
    assert [m.name for m in agent['middleware']] == ['DynamicTimeMiddleware', 'MainOnly']
    assert [m.name for m in agent['subagents'][0]['middleware']] == ['ChildOnly']


@pytest.mark.asyncio
async def test_cleanup_failure_does_not_skip_other_services(isolated, monkeypatch):
    calls = []
    catalog = ExtensionRegistry()
    @asynccontextmanager
    async def db(deps):
        try:
            yield object()
        finally:
            calls.append('db_closed')
    @asynccontextmanager
    async def search(deps):
        try:
            yield object()
        finally:
            calls.append('search_closed')
            raise RuntimeError('close failed')
    catalog.register_service('db', db)
    catalog.register_service('search', search, requires=['db'])
    use_catalog(monkeypatch, catalog)
    with pytest.raises(RuntimeError, match='close failed'):
        async with RuntimeEnvironment(prewarm=False, billing_retry=False):
            pass
    assert calls == ['search_closed', 'db_closed']
    assert not agent_memory.is_initialized()
    assert db_conn._adapter._conn is None
    with pytest.raises(RuntimeError):
        get_service('db')


@pytest.mark.asyncio
async def test_middleware_executes_only_on_target_in_real_delegation(isolated, monkeypatch):
    from test_runtime_regressions import ScriptedModel
    seen = []
    def factory(name):
        class Middleware(AgentMiddleware):
            async def awrap_model_call(self, request, handler):
                seen.append(name + '_before')
                result = await handler(request)
                seen.append(name + '_after')
                return result
        Middleware.__name__ = name
        return lambda deps: Middleware()
    base = registry_module.load_extensions()
    local = ExtensionRegistry()
    local.tools.update(base.tools)
    local.agents.update(base.agents)
    local.register_middleware('MainOuter', factory('MainOuter'), before=['MainInner'])
    local.register_middleware('MainInner', factory('MainInner'))
    local.register_middleware('ChildOnly', factory('ChildOnly'), targets=['demo-searcher'])
    local.validate()
    monkeypatch.setattr(registry_module, '_loaded', local)
    model = ScriptedModel(responses=[
        AIMessage(content='', tool_calls=[{'id': 'delegate', 'name': 'task',
                  'args': {'subagent_type': 'demo-searcher', 'description': '查询术语'}}]),
        AIMessage(content='检索结论。'), AIMessage(content='汇总完成。'),
    ])
    monkeypatch.setattr('agents.model_factory.get_model', lambda: model)
    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
        sid = chat_store.create_session('delegation')['id']
        result = await execute(user_name='delegation', session_id=sid, message='请检索并汇总')
        assert result.status == 'success', result.error
        assert result.output.endswith('汇总完成。')
    main_call = ['MainOuter_before', 'MainInner_before', 'MainInner_after', 'MainOuter_after']
    assert seen == main_call + ['ChildOnly_before', 'ChildOnly_after'] + main_call


@pytest.mark.parametrize('hook', ['enabled', 'validate_config', 'setup'])
def test_async_declaration_hooks_are_rejected(monkeypatch, hook):
    async def asynchronous(*args):
        return True
    fake = SimpleNamespace(enabled=lambda: True, validate_config=lambda: None, setup=lambda registry: None)
    setattr(fake, hook, asynchronous)
    monkeypatch.setattr(registry_module, 'discover_modules', lambda: ['async_extension'])
    monkeypatch.setattr(registry_module.importlib, 'import_module', lambda name: fake)
    tools.reset_for_tests()
    with pytest.raises(ValueError, match='同步'):
        registry_module.load_extensions()
    assert registry_module._loaded is None