"""CORE-DIAG-6: management ownership barriers under a foreign live span.""" import asyncio from types import SimpleNamespace as NS from unittest.mock import AsyncMock, patch from uuid import uuid4 import pytest import quart from langbot.pkg.api import management_diagnostics as md from langbot.pkg.api.http.context import ( ExecutionContext, PrincipalContext, PrincipalType, RequestContext, WorkspaceContext, ) from langbot.pkg.api.http.controller.group import AuthType from langbot.pkg.api.http.controller.groups.pipelines.embed import EmbedRouterGroup from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup from langbot.pkg.api.http.service.agent import AgentService 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, Routes MODES = ['enabled', 'disabled', 'absent', 'mismatched', 'policy_off', 'beta_off', 'broken'] def app(identity): ap = NS(instance_config=NS(data={'space': {'url': 'https://example.invalid'}})) ap.diagnostics = d.DiagnosticsManager(ap, version='4.11.0b2', instance_id=identity) return ap def owners(mode): a, b = app('instance-A'), app('instance-B') b_manager = b.diagnostics if mode == 'disabled': b.instance_config.data['space']['disable_telemetry'] = True b.diagnostics = d.DiagnosticsManager(b, version='4.11.0b2', instance_id='instance-B') b_manager = b.diagnostics elif mode == 'absent': del b.diagnostics elif mode == 'mismatched': b.diagnostics = a.diagnostics elif mode in {'policy_off', 'beta_off'}: # Also test an attached structural producer whose enabled flag stays true. b.diagnostics = Recorder() b.instance_config.data['space']['disable_telemetry' if mode == 'policy_off' else 'disable_beta_diagnostics'] = ( True ) elif mode == 'broken': b.diagnostics = NS(enabled=True, emit=lambda *a, **kw: (_ for _ in ()).throw(RuntimeError('broken'))) ca = ExecutionContext('instance-A', str(uuid4()), 1) cb = ExecutionContext('instance-B', str(uuid4()), 1) d.privacy.code_value('operation', 'http.isolation.parent') parent = d.Span(a.diagnostics, 'api', 'http.isolation.parent', {'workspace_uuid': ca.workspace_uuid}) return a, b, b_manager, ca, cb, parent @d.observe('event', 'platform.target2yiri', source='platform', stage='convert') async def converter(native): return native async def nested(b, cb, seen): md.workspace(cb) seen.append(d.current_span()) with md.scope(b, 'websocket.isolation.inner', source='websocket'): md.workspace(cb) assert await converter(42) == 42 # Both an inner management boundary and an ownerless converter must be safe. assert await converter(43) == 43 return 44 def assert_isolated(a, b, b_manager, ca, cb, parent, mode): assert parent.fields['workspace_uuid'] == ca.workspace_uuid assert parent.outcome is None assert len(a.diagnostics.pending) == 1 # Only the caller's own start event. assert not b_manager.pending if mode != 'enabled' else b_manager.pending if isinstance(getattr(b, 'diagnostics', None), Recorder): assert b.diagnostics.events == [] if mode == 'enabled': assert all(e['instance_id'] == 'instance-B' for e in b_manager.pending) assert all(e['trace_id'] != parent.fields['trace_id'] for e in b_manager.pending) assert all(e['workspace_uuid'] != ca.workspace_uuid for e in b_manager.pending) assert any(e['workspace_uuid'] == cb.workspace_uuid for e in b_manager.pending) @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES) @pytest.mark.parametrize('shape', ['decorator', 'scope']) async def test_nested_management_masks_foreign_parent_and_restores_caller(mode, shape): a, b, manager, ca, cb, parent = owners(mode) seen = [] @md.observe('http.isolation.endpoint', source='http', ap=b) async def endpoint(): return await nested(b, cb, seen) with parent.activate(): if shape == 'decorator': assert await endpoint() == 44 else: with md.scope(b, 'websocket.isolation.outer', source='websocket'): assert await nested(b, cb, seen) == 44 assert d.current_span() is parent assert d.current_span() is None assert (seen[0] is not None) == (mode == 'enabled') assert_isolated(a, b, manager, ca, cb, parent, mode) @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES) async def test_actual_agent_debug_never_mutates_foreign_workspace(mode): a, b, manager, ca, cb, parent = owners(mode) service = AgentService(b) seen = [] async def missing(*args): await nested(b, cb, seen) return None service.get_agent = missing with parent.activate(): with pytest.raises(ValueError, match='^Agent not found$'): await service.debug_agent(cb, 'private-agent', {}) assert d.current_span() is parent assert_isolated(a, b, manager, ca, cb, parent, mode) @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES) @pytest.mark.parametrize('termination', ['exhaust', 'close', 'error', 'cancel']) async def test_real_http_body_keeps_barrier_during_advancement_and_cleanup(mode, termination): a, b, manager, ca, cb, parent = owners(mode) web = quart.Quart(__name__) routes = Routes(b, web) seen, closed = [], [] failure = asyncio.CancelledError('private-cancel') if termination == 'cancel' else ValueError('private-error') @routes.route('/isolated', auth_type=AuthType.NONE) async def endpoint(): await nested(b, cb, seen) async def stream(): try: await nested(b, cb, seen) yield b'one' if termination in {'cancel', 'error'}: raise failure await nested(b, cb, seen) finally: await nested(b, cb, seen) closed.append(True) return quart.Response(stream()) # Construct with no caller parent; the later body consumer has app A's span. async with web.test_request_context('/management/isolated'): response = await web.full_dispatch_request() with parent.activate(): try: async with response.response as body: iterator = body.__aiter__() assert await anext(iterator) == b'one' assert d.current_span() is parent if termination != 'close': with pytest.raises((StopAsyncIteration, type(failure))) as caught: await anext(iterator) if termination in {'cancel', 'error'}: assert caught.value is failure else: assert isinstance(caught.value, StopAsyncIteration) finally: assert d.current_span() is parent assert closed == [True] assert all((span is not None) == (mode == 'enabled') for span in seen) assert_isolated(a, b, manager, ca, cb, parent, mode) @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES) async def test_real_mcp_registered_tool_masks_foreign_parent(mode): a, b, manager, ca, cb, parent = owners(mode) seen = [] async def get_bot(*args, **kwargs): return {'value': await nested(b, cb, seen)} b.bot_service = NS(get_bot=get_bot) server = LangBotMCPServer(b) with parent.activate(), patch('langbot.pkg.api.mcp.server._authorized', return_value=cb): assert await server.mcp.call_tool('get_bot', {'bot_uuid': 'private-bot'}) assert d.current_span() is parent assert_isolated(a, b, manager, ca, cb, parent, mode) @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES) @pytest.mark.parametrize('group_class', [WebSocketChatRouterGroup, EmbedRouterGroup]) async def test_real_websocket_receive_masks_foreign_parent(mode, group_class): a, b, manager, ca, cb, parent = owners(mode) group = group_class(b, quart.Quart(__name__)) connection = NS(is_active=True, connection_id='private-id', send_queue=asyncio.Queue()) seen = [] async def handle(*args, **kwargs): await nested(b, cb, seen) adapter = NS(handle_websocket_message=handle) if group_class is WebSocketChatRouterGroup: group._revalidate_websocket_authorization = AsyncMock(return_value=cb) else: group._resolve_connected_bot = AsyncMock(return_value=NS(execution_context=cb)) async def receive(): connection.is_active = False return '{"type":"message","text":"private-prompt"}' with ( parent.activate(), patch('quart.websocket', NS(receive=receive)), patch( 'langbot.pkg.api.http.controller.groups.pipelines.websocket_chat.ws_connection_manager.update_activity', new=AsyncMock(), ), ): await group._handle_receive(connection, adapter, NS(execution_context=cb), 'private-token') assert d.current_span() is parent assert seen assert_isolated(a, b, manager, ca, cb, parent, mode) @pytest.mark.parametrize('kind', ['execution', 'request']) def test_workspace_rejects_foreign_instance_or_foreign_active_span(kind): a, b, _, ca, cb, parent = owners('enabled') context = cb if kind == 'request': context = RequestContext( cb.instance_uuid, 1, 'request', 'api_key', PrincipalContext(PrincipalType.API_KEY), WorkspaceContext(cb.workspace_uuid, None, None, frozenset()), ) with md.scope(a, 'http.isolation.annotation', source='http'): span = d.current_span() md.workspace(context) assert 'workspace_uuid' not in span.fields md.workspace(ca) assert span.fields['workspace_uuid'] == ca.workspace_uuid with parent.activate(): before = dict(parent.fields) md.workspace(ca) md.workspace(context) assert parent.fields == before def test_structural_recorder_workspace_and_scope_exception_identity(): recorder = Recorder() ap = NS(diagnostics=recorder) ctx = ExecutionContext('instance', str(uuid4()), 1) error = ValueError('private-error') with pytest.raises(ValueError) as caught: with md.scope(ap, 'http.isolation.recorder', source='http'): md.workspace(ctx) assert d.current_span().fields['workspace_uuid'] == ctx.workspace_uuid raise error assert caught.value is error assert recorder.events[-1]['outcome'] == 'failed' assert d.current_span() is None @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES[1:]) @pytest.mark.parametrize('termination', ['success', 'error', 'cancel']) async def test_nested_off_boundary_restores_outer_annotation_and_error(mode, termination): a, b, manager, ca, cb, parent = owners(mode) error = asyncio.CancelledError('private-cancel') if termination == 'cancel' else ValueError('private-error') @md.observe('http.isolation.off', source='http', ap=b) async def endpoint(): await nested(b, cb, []) md.outcome('failed') if termination != 'success': raise error return 42 with parent.activate(): with md.scope(a, 'http.isolation.outer', source='http'): outer = d.current_span() if termination == 'success': assert await endpoint() == 42 else: with pytest.raises(type(error)) as caught: await endpoint() assert caught.value is error assert d.current_span() is outer md.workspace(ca) assert outer.fields['workspace_uuid'] == ca.workspace_uuid assert outer.outcome is None assert d.current_span() is parent assert not manager.pending assert [e['operation'] for e in a.diagnostics.pending] == [ 'http.isolation.parent', 'http.isolation.outer', 'http.isolation.outer', ] @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES) @pytest.mark.parametrize('group_class', [WebSocketChatRouterGroup, EmbedRouterGroup]) async def test_websocket_send_cancel_restores_foreign_parent(mode, group_class): a, b, manager, ca, cb, parent = owners(mode) group = group_class(b, quart.Quart(__name__)) connection = NS(is_active=False, send_queue=asyncio.Queue()) await connection.send_queue.put({'text': 'private-answer'}) error = asyncio.CancelledError('private-cancel') async def send(payload): await nested(b, cb, []) raise error with parent.activate(), patch('quart.websocket', NS(send=send)): with pytest.raises(asyncio.CancelledError) as caught: await group._handle_send(connection) assert caught.value is error assert d.current_span() is parent assert_isolated(a, b, manager, ca, cb, parent, mode) @pytest.mark.asyncio @pytest.mark.parametrize('mode', MODES) async def test_sync_http_body_exhaustion_has_owned_context_without_consumer_leak(mode): a, b, manager, ca, cb, parent = owners(mode) seen, closed = [], [] def stream(): try: for value in (b'one', b'two'): md.workspace(cb) seen.append(d.current_span()) yield value finally: seen.append(d.current_span()) closed.append(True) @md.observe('http.isolation.sync_body', source='http', ap=b, http=True) async def endpoint(): return quart.Response(stream()) response = await endpoint() with parent.activate(): async with response.response as body: chunks = [] async for chunk in body: assert d.current_span() is parent chunks.append(chunk) assert d.current_span() is parent assert chunks == [b'one', b'two'] assert closed == [True] assert all((span is not None) == (mode == 'enabled') for span in seen) assert_isolated(a, b, manager, ca, cb, parent, mode) @pytest.mark.asyncio @pytest.mark.parametrize('mode', ['enabled', 'disabled', 'absent', 'mismatched']) @pytest.mark.parametrize('failure_stage', ['enter', 'iterator', 'advance', 'exit', None]) async def test_custom_http_body_protocol_and_cancellation_identity(mode, failure_stage): from quart.wrappers.response import IterableBody a, b, manager, ca, cb, parent = owners(mode) error = asyncio.CancelledError('private-body-cancel') seen = [] def visit(stage): md.workspace(cb) seen.append((stage, d.current_span())) if stage == failure_stage: raise error class Body(IterableBody): def __init__(self): pass async def __aenter__(self): visit('enter') return self def __aiter__(self): visit('iterator') return self async def __anext__(self): visit('advance') raise StopAsyncIteration async def __aexit__(self, *args): visit('exit') @md.observe('http.isolation.body_protocol', source='http', ap=b, http=True) async def endpoint(): response = quart.Response() response.response = Body() return response response = await endpoint() async def consume(): async with response.response as body: assert d.current_span() is parent async for _ in body: pytest.fail('empty body yielded') with parent.activate(): if failure_stage is None: await consume() else: with pytest.raises(asyncio.CancelledError) as caught: await consume() assert caught.value is error assert d.current_span() is parent assert all((span is not None) == (mode == 'enabled') for _, span in seen) assert [s for s, _ in seen].count('exit') == (0 if failure_stage in {'enter', 'iterator'} else 1) assert_isolated(a, b, manager, ca, cb, parent, mode)