LangBot/tests/unit_tests/api/test_diagnostics_management_transports.py

"""Management MCP, WebSocket and debug execution boundaries."""

import asyncio
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

import pytest
import quart

from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup
from langbot.pkg.api.http.controller.groups.pipelines.embed import EmbedRouterGroup
from langbot.pkg.api.mcp.server import LangBotMCPServer
from langbot.pkg.telemetry import diagnostics as d
from tests.unit_tests.api.test_diagnostics_management_http import Recorder


@pytest.mark.asyncio
async def test_mcp_real_registered_tool_has_source_and_never_content():
    recorder = Recorder()
    seen = []

    async def get_bot(*args, **kwargs):
        seen.append(d.current_span())
        return {'secret': 'private-result'}

    ap = SimpleNamespace(diagnostics=recorder, bot_service=SimpleNamespace(get_bot=get_bot))
    server = LangBotMCPServer(ap)
    with patch('langbot.pkg.api.mcp.server._authorized', return_value=object()):
        result = await server.mcp.call_tool('get_bot', {'bot_uuid': 'private-bot'})
    assert result
    assert [e['outcome'] for e in recorder.events] == ['started', 'succeeded']
    assert recorder.events[-1]['operation'] == 'mcp.get_bot'
    assert recorder.events[-1]['source'] == 'mcp'
    assert seen[0] is not None
    assert 'private-' not in json.dumps(recorder.events)
    assert d.current_span() is None


@pytest.mark.asyncio
async def test_mcp_tool_permission_rejection_is_observed():
    server = LangBotMCPServer(SimpleNamespace(diagnostics=Recorder()))
    with pytest.raises(Exception):
        await server.mcp.call_tool('list_bots', {})
    assert server.ap.diagnostics.events[-1]['outcome'] in {'rejected', 'failed'}
    assert server.ap.diagnostics.events[-1]['source'] == 'mcp'


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'group_class,path',
    [
        (WebSocketChatRouterGroup, '/api/v1/pipelines/private-pipeline/ws/connect'),
        (EmbedRouterGroup, '/api/v1/embed/11111111-1111-4111-8111-111111111111/ws/connect?session_id=private-session'),
    ],
)
async def test_real_websocket_session_auth_rejection_is_content_free(group_class, path):
    app = quart.Quart(__name__)
    ap = SimpleNamespace(diagnostics=Recorder())
    group = group_class(ap, app)
    group._authenticate_websocket = AsyncMock(side_effect=ValueError('private-token'))
    if group_class is EmbedRouterGroup:
        group._resolve_bot = AsyncMock(return_value=(object(), 'private-pipeline'))
    await group.initialize()
    async with app.test_client().websocket(path) as socket:
        frame = json.loads(await socket.receive())
        assert frame['type'] == 'error'
    terminal = [e for e in ap.diagnostics.events if e['outcome'] != 'started']
    assert terminal
    assert terminal[-1]['outcome'] == 'rejected'
    assert terminal[-1]['source'] == 'websocket'
    assert 'private-' not in json.dumps(ap.diagnostics.events)


@pytest.mark.asyncio
@pytest.mark.parametrize('group_class', [WebSocketChatRouterGroup, EmbedRouterGroup])
async def test_websocket_received_message_boundaries_do_not_capture_frames(group_class):
    app = quart.Quart(__name__)
    ap = SimpleNamespace(diagnostics=Recorder())
    group = group_class(ap, app)
    connection = SimpleNamespace(is_active=True, connection_id='private-id', send_queue=asyncio.Queue())
    adapter = SimpleNamespace(handle_websocket_message=AsyncMock())
    if group_class is WebSocketChatRouterGroup:
        group._revalidate_websocket_authorization = AsyncMock(return_value=object())
        args = (connection, adapter, object(), 'private-token')
    else:
        group._resolve_connected_bot = AsyncMock(return_value=object())
        args = (connection, adapter, object(), 'private-pipeline')

    async def receive():
        connection.is_active = False
        return json.dumps({'type': 'message', 'text': 'private-prompt'})

    with (
        patch('quart.websocket', SimpleNamespace(receive=receive)),
        patch(
            'langbot.pkg.api.http.controller.groups.pipelines.websocket_chat.ws_connection_manager.update_activity',
            new=AsyncMock(),
        ),
    ):
        await group._handle_receive(*args)
    adapter.handle_websocket_message.assert_awaited_once()
    assert [e['outcome'] for e in ap.diagnostics.events] == ['started', 'succeeded']
    assert ap.diagnostics.events[-1]['operation'].endswith('.message')
    assert 'private-' not in json.dumps(ap.diagnostics.events)
    assert d.current_span() is None


@pytest.mark.asyncio
async def test_debug_service_marks_synthetic_source_before_validation():
    from langbot.pkg.api.http.service.agent import AgentService

    ap = SimpleNamespace(diagnostics=Recorder())
    service = AgentService(ap)
    seen = []

    async def get_agent(*args):
        seen.append(d.current_span())
        return None

    service.get_agent = get_agent
    with pytest.raises(ValueError):
        await service.debug_agent(object(), 'private-id', {'text': 'private-prompt'})
    assert seen[0].fields['source'] == 'webui_debug'
    assert ap.diagnostics.events[-1]['operation'] == 'http.agent.debug_agent'
    assert ap.diagnostics.events[-1]['source'] == 'webui_debug'
    assert 'private-' not in json.dumps(ap.diagnostics.events)


@pytest.mark.asyncio
@pytest.mark.parametrize('group_class', [WebSocketChatRouterGroup, EmbedRouterGroup])
@pytest.mark.parametrize(
    'frame,expected', [('private-invalid-json', 'rejected'), ('{"type":"private-unknown"}', 'skipped')]
)
async def test_websocket_invalid_frames_have_finite_outcomes(group_class, frame, expected):
    app = quart.Quart(__name__)
    ap = SimpleNamespace(diagnostics=Recorder())
    group = group_class(ap, app)
    connection = SimpleNamespace(is_active=True, connection_id='private-id', send_queue=asyncio.Queue())

    async def receive():
        connection.is_active = False
        return frame

    with (
        patch('quart.websocket', SimpleNamespace(receive=receive)),
        patch(
            'langbot.pkg.api.http.controller.groups.pipelines.websocket_chat.ws_connection_manager.update_activity',
            new=AsyncMock(),
        ),
    ):
        await group._handle_receive(connection, object(), object(), 'private-token')
    assert ap.diagnostics.events[-1]['outcome'] == expected
    assert 'private-' not in json.dumps(ap.diagnostics.events)


@pytest.mark.asyncio
@pytest.mark.parametrize('group_class', [WebSocketChatRouterGroup, EmbedRouterGroup])
async def test_websocket_send_boundary_preserves_payload_and_cancellation(group_class):
    app = quart.Quart(__name__)
    ap = SimpleNamespace(diagnostics=Recorder())
    group = group_class(ap, app)
    connection = SimpleNamespace(is_active=False, send_queue=asyncio.Queue())
    await connection.send_queue.put({'text': 'private-answer'})
    send = AsyncMock(side_effect=asyncio.CancelledError('private-error'))
    with patch('quart.websocket', SimpleNamespace(send=send)):
        with pytest.raises(asyncio.CancelledError):
            await group._handle_send(connection)
    assert json.loads(send.call_args.args[0]) == {'text': 'private-answer'}
    assert ap.diagnostics.events[-1]['outcome'] == 'cancelled'
    assert ap.diagnostics.events[-1]['operation'].endswith('.send')
    assert 'private-' not in json.dumps(ap.diagnostics.events)
    assert d.current_span() is None