LangBot/tests/unit_tests/pipeline/test_controller.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, Mock, patch

import pytest

from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.agent.runner.errors import RunnerNotFoundError
from langbot.pkg.pipeline.controller import Controller
from langbot.pkg.pipeline.pool import QueryPool


def make_app():
    app = SimpleNamespace()
    app.instance_config = SimpleNamespace(data={'concurrency': {'pipeline': 10}})
    app.logger = MagicMock()
    app.pipeline_mgr = SimpleNamespace()
    app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock()
    app.sess_mgr = SimpleNamespace()
    app.sess_mgr.get_session = AsyncMock(return_value=SimpleNamespace())
    app.agent_run_orchestrator = SimpleNamespace()
    app.agent_run_orchestrator.try_claim_steering_from_query = AsyncMock()
    return app


def make_pipeline():
    return SimpleNamespace(
        pipeline_entity=SimpleNamespace(config={'ai': {'runner': {'id': 'plugin:test/runner/default'}}}),
        bound_plugins=['test/runner'],
        bound_mcp_servers=[],
    )


def make_query(query_id: int, pipeline_uuid: str):
    context = ExecutionContext(
        instance_uuid='instance-test',
        workspace_uuid='workspace-test',
        placement_generation=1,
        pipeline_uuid=pipeline_uuid,
    )
    return SimpleNamespace(
        query_id=query_id,
        pipeline_uuid=pipeline_uuid,
        variables={},
        _execution_context=context,
    )


@pytest.mark.asyncio
async def test_try_claim_steering_returns_false_when_runner_lookup_fails():
    app = make_app()
    app.pipeline_mgr.get_pipeline_by_uuid.return_value = make_pipeline()
    app.agent_run_orchestrator.try_claim_steering_from_query.side_effect = RunnerNotFoundError(
        'plugin:missing/runner/default'
    )
    controller = Controller(app)
    query = make_query(1, 'pipeline-001')

    claimed = await controller._try_claim_steering_before_session_slot(query)

    assert claimed is False
    app.logger.warning.assert_called_once()


@pytest.mark.asyncio
async def test_try_claim_steering_sets_pipeline_context_before_claiming():
    app = make_app()
    pipeline = make_pipeline()
    app.pipeline_mgr.get_pipeline_by_uuid.return_value = pipeline
    app.agent_run_orchestrator.try_claim_steering_from_query.return_value = True
    controller = Controller(app)
    query = make_query(2, 'pipeline-002')

    claimed = await controller._try_claim_steering_before_session_slot(query)

    assert claimed is True
    assert query.pipeline_config is pipeline.pipeline_entity.config
    assert query.variables['_pipeline_bound_plugins'] == ['test/runner']
    app.agent_run_orchestrator.try_claim_steering_from_query.assert_awaited_once_with(query)


@pytest.mark.asyncio
@pytest.mark.parametrize('text,expected', [('ordinary chat', False), ('qa: hello', True)])
@pytest.mark.parametrize('claim_result', [True, False])
async def test_busy_group_steering_obeys_response_rules(sample_query, text, expected, claim_result):
    from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes, Session
    from langbot_plugin.api.entities.builtin.platform.message import MessageChain, Plain

    app = make_app()
    app.sess_mgr.get_session.return_value = Session(
        instance_uuid='test-instance',
        workspace_uuid='test-workspace',
        placement_generation=1,
        launcher_type=LauncherTypes.GROUP,
        launcher_id=sample_query.launcher_id,
    )
    pipeline = make_pipeline()
    pipeline.pipeline_entity.config['trigger'] = {
        'group-respond-rules': {'at': True, 'prefix': ['qa:'], 'regexp': [], 'random': 0}
    }
    app.pipeline_mgr.get_pipeline_by_uuid.return_value = pipeline
    app.agent_run_orchestrator.try_claim_steering_from_query.return_value = claim_result
    sample_query.launcher_type = LauncherTypes.GROUP
    sample_query.message_chain = MessageChain([Plain(text=text)])
    claimed = await Controller(app)._try_claim_steering_before_session_slot(sample_query)
    assert claimed is (expected and claim_result), app.logger.warning.call_args
    assert str(sample_query.message_chain) == text
    if expected:
        forwarded = app.agent_run_orchestrator.try_claim_steering_from_query.await_args.args[0]
        assert str(forwarded.message_chain) == ' hello'
    else:
        app.agent_run_orchestrator.try_claim_steering_from_query.assert_not_awaited()


@pytest.mark.asyncio
async def test_consumer_transfers_query_from_queue_to_running_task():
    app = make_app()
    app.query_pool = QueryPool()
    session = SimpleNamespace(_semaphore=asyncio.Semaphore(1))
    app.sess_mgr.get_session = AsyncMock(return_value=session)
    app.persistence_mgr = SimpleNamespace(mode=SimpleNamespace(value='oss'))
    app.workspace_service = SimpleNamespace(
        get_execution_binding=AsyncMock(
            return_value=SimpleNamespace(
                instance_uuid='instance-test',
                workspace_uuid='workspace-test',
                placement_generation=1,
            )
        )
    )
    runtime_pipeline = SimpleNamespace(run=AsyncMock())
    app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=runtime_pipeline)

    worker_tasks = []
    task_created = asyncio.Event()

    def create_task(coro, **_kwargs):
        task = asyncio.create_task(coro)
        worker_tasks.append(task)
        task_created.set()
        return task

    app.task_mgr = SimpleNamespace(create_task=create_task)

    query = Mock()
    query.query_id = 0
    query.bot_uuid = 'bot-test'
    query.pipeline_uuid = 'pipeline-test'
    context = ExecutionContext(
        instance_uuid='instance-test',
        workspace_uuid='workspace-test',
        placement_generation=1,
    )
    with patch('langbot.pkg.pipeline.pool.pipeline_query.Query', return_value=query):
        query = await app.query_pool.add_query(
            bot_uuid='bot-test',
            launcher_type=Mock(),
            launcher_id='launcher-test',
            sender_id='sender-test',
            message_event=Mock(),
            message_chain=Mock(),
            adapter=Mock(),
            pipeline_uuid='pipeline-test',
            execution_context=context,
        )

    controller = Controller(app)
    consumer_task = asyncio.create_task(controller.consumer())
    try:
        await asyncio.wait_for(task_created.wait(), timeout=1)
        await asyncio.wait_for(worker_tasks[0], timeout=1)
    finally:
        consumer_task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await consumer_task

    runtime_pipeline.run.assert_awaited_once_with(query)
    assert app.query_pool.queries == []
    assert app.query_pool.cached_queries == {}
    assert app.query_pool.active_query_count_by_workspace == {}
    assert session._semaphore._value == 1
    assert controller.semaphore._value == 10
    app.logger.error.assert_not_called()