LangBot/tests/unit_tests/api/test_adapter_session_scoping.py

from __future__ import annotations

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

import lark_oapi
import pytest
import quart

from langbot.pkg.api.http.context import (
    PrincipalContext,
    PrincipalType,
    RequestContext,
    WorkspaceContext,
)
from langbot.pkg.api.http.controller.groups.platform.adapters import (
    AdaptersRouterGroup,
    _AdapterSessionScope,
    _bind_session_scope,
    _get_owned_session,
    _make_room_for_session,
    _pop_owned_session,
    _start_adapter_session_task,
)


pytestmark = pytest.mark.asyncio


SENSITIVE_ADAPTER_ROUTES = (
    ('post', '/api/v1/platform/adapters/lark/create-app'),
    ('get', '/api/v1/platform/adapters/lark/create-app/status/missing'),
    ('delete', '/api/v1/platform/adapters/lark/create-app/missing'),
    ('post', '/api/v1/platform/adapters/weixin/login'),
    ('get', '/api/v1/platform/adapters/weixin/login/status/missing'),
    ('delete', '/api/v1/platform/adapters/weixin/login/missing'),
    ('post', '/api/v1/platform/adapters/dingtalk/create-app'),
    ('get', '/api/v1/platform/adapters/dingtalk/create-app/status/missing'),
    ('delete', '/api/v1/platform/adapters/dingtalk/create-app/missing'),
    ('post', '/api/v1/platform/adapters/wecombot/create-bot'),
    ('get', '/api/v1/platform/adapters/wecombot/create-bot/status/missing'),
    ('delete', '/api/v1/platform/adapters/wecombot/create-bot/missing'),
    ('post', '/api/v1/platform/adapters/qqofficial/bind'),
    ('get', '/api/v1/platform/adapters/qqofficial/bind/status/missing'),
    ('delete', '/api/v1/platform/adapters/qqofficial/bind/missing'),
)


def _request_context(
    *,
    account_uuid: str = 'account-a',
    workspace_uuid: str = 'workspace-a',
    placement_generation: int = 1,
) -> RequestContext:
    return RequestContext(
        instance_uuid='instance-test',
        placement_generation=placement_generation,
        request_id='request-test',
        auth_type='user-token',
        principal=PrincipalContext(
            principal_type=PrincipalType.ACCOUNT,
            account_uuid=account_uuid,
        ),
        workspace=WorkspaceContext(
            workspace_uuid=workspace_uuid,
            membership_uuid='membership-test',
            role='developer',
            permissions=frozenset({'resource.manage'}),
        ),
    )


async def _create_client(*, role: str = 'developer'):
    quart_app = quart.Quart(__name__)
    accounts = {
        'owner-token': SimpleNamespace(uuid='account-a', user='owner@example.com'),
        'other-token': SimpleNamespace(uuid='account-b', user='other@example.com'),
    }

    async def get_authenticated_account(token: str):
        return accounts[token]

    async def resolve_account_workspace(account_uuid: str, requested_workspace_uuid: str | None):
        workspace_uuid = requested_workspace_uuid or 'workspace-a'
        return SimpleNamespace(
            execution=SimpleNamespace(
                instance_uuid='instance-test',
                placement_generation=1,
            ),
            workspace=SimpleNamespace(uuid=workspace_uuid),
            membership=SimpleNamespace(
                uuid=f'membership-{account_uuid}-{workspace_uuid}',
                role=role,
                projection_revision=1,
            ),
        )

    class TestTaskManager:
        def create_user_task(self, coro, **_kwargs):
            return SimpleNamespace(task=asyncio.create_task(coro))

    application = SimpleNamespace(
        user_service=SimpleNamespace(
            get_authenticated_account=AsyncMock(side_effect=get_authenticated_account),
        ),
        workspace_collaboration_service=SimpleNamespace(
            resolve_account_workspace=AsyncMock(side_effect=resolve_account_workspace),
        ),
        platform_mgr=SimpleNamespace(),
        task_mgr=TestTaskManager(),
    )
    router = AdaptersRouterGroup(application, quart_app)
    await router.initialize()
    return quart_app.test_client()


@pytest.mark.parametrize(('method', 'path'), SENSITIVE_ADAPTER_ROUTES)
async def test_sensitive_adapter_flows_require_resource_manage(method: str, path: str):
    client = await _create_client(role='viewer')

    response = await getattr(client, method)(
        path,
        headers={'Authorization': 'Bearer owner-token'},
    )

    assert response.status_code == 403
    assert (await response.get_json())['code'] == 'permission_denied'


async def test_session_scope_matches_exact_tenant_placement_and_principal():
    owner_context = _request_context()
    sessions: dict[str, dict] = {'session-test': {'status': 'waiting'}}
    _bind_session_scope(sessions['session-test'], owner_context)

    assert sessions['session-test']['scope'] == _AdapterSessionScope.from_request_context(owner_context)
    assert _get_owned_session(sessions, 'session-test', owner_context) is sessions['session-test']

    for other_context in (
        _request_context(account_uuid='account-b'),
        _request_context(workspace_uuid='workspace-b'),
        _request_context(placement_generation=2),
    ):
        assert _get_owned_session(sessions, 'session-test', other_context) is None
        assert _pop_owned_session(sessions, 'session-test', other_context) is None
        assert 'session-test' in sessions

    assert _pop_owned_session(sessions, 'session-test', owner_context) is not None
    assert sessions == {}


async def test_session_capacity_evicts_oldest_session_in_same_workspace():
    owner_context = _request_context()
    sessions: dict[str, dict] = {}
    tasks = []
    for index in range(10):
        task = SimpleNamespace(done=Mock(return_value=False), cancel=Mock())
        tasks.append(task)
        session = {'created_at': float(index), 'task': task}
        _bind_session_scope(session, owner_context)
        sessions[f'session-{index}'] = session

    _make_room_for_session(sessions, owner_context)

    assert 'session-0' not in sessions
    assert len(sessions) == 9
    tasks[0].cancel.assert_called_once_with()


async def test_adapter_session_task_uses_tenant_task_admission():
    blocker = asyncio.Event()

    async def credential_exchange():
        await blocker.wait()

    task_manager = SimpleNamespace(create_user_task=Mock())

    def create_user_task(coro, **_kwargs):
        return SimpleNamespace(task=asyncio.create_task(coro))

    task_manager.create_user_task.side_effect = create_user_task
    application = SimpleNamespace(task_mgr=task_manager)
    request_context = _request_context()

    returned = _start_adapter_session_task(
        application,
        credential_exchange(),
        adapter='lark',
        session_id='session-test',
        request_context=request_context,
    )

    assert returned is not None
    task_manager.create_user_task.assert_called_once()
    kwargs = task_manager.create_user_task.call_args.kwargs
    assert kwargs['kind'] == 'platform-adapter-credential-exchange'
    assert kwargs['instance_uuid'] == request_context.instance_uuid
    assert kwargs['workspace_uuid'] == request_context.workspace_uuid
    assert kwargs['placement_generation'] == request_context.placement_generation
    blocker.set()
    await returned


async def test_lark_session_status_and_delete_hide_cross_scope_sessions(monkeypatch):
    registration_blocker = asyncio.Event()

    async def fake_register_app(*, on_qr_code, source: str):
        assert source == 'langbot'
        on_qr_code({'url': 'https://example.test/lark-qr'})
        await registration_blocker.wait()
        raise AssertionError('registration should have been cancelled')

    monkeypatch.setattr(lark_oapi, 'aregister_app', fake_register_app)
    client = await _create_client()
    owner_headers = {
        'Authorization': 'Bearer owner-token',
        'X-Workspace-Id': 'workspace-a',
    }

    create_response = await client.post(
        '/api/v1/platform/adapters/lark/create-app',
        headers=owner_headers,
    )
    assert create_response.status_code == 200
    session_id = (await create_response.get_json())['data']['session_id']
    status_path = f'/api/v1/platform/adapters/lark/create-app/status/{session_id}'
    delete_path = f'/api/v1/platform/adapters/lark/create-app/{session_id}'

    for headers in (
        {
            'Authorization': 'Bearer other-token',
            'X-Workspace-Id': 'workspace-a',
        },
        {
            'Authorization': 'Bearer owner-token',
            'X-Workspace-Id': 'workspace-b',
        },
    ):
        status_response = await client.get(status_path, headers=headers)
        delete_response = await client.delete(delete_path, headers=headers)
        assert status_response.status_code == 404
        assert delete_response.status_code == 404
        assert (await status_response.get_json())['msg'] == 'Session not found'
        assert (await delete_response.get_json())['msg'] == 'Session not found'

    owner_status_response = await client.get(status_path, headers=owner_headers)
    assert owner_status_response.status_code == 200
    assert (await owner_status_response.get_json())['data']['status'] == 'waiting'

    owner_delete_response = await client.delete(delete_path, headers=owner_headers)
    assert owner_delete_response.status_code == 200
    missing_delete_response = await client.delete(delete_path, headers=owner_headers)
    assert missing_delete_response.status_code == 404
    await asyncio.sleep(0)