mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-09-17 07:17:18 +00:00
feat(telemetry): add isolated beta quality diagnostics and release identity
This commit is contained in:
dadachann
@@ -0,0 +1,194 @@
|
||||
"""Real Quart registration boundaries with content-canary payloads."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import quart
|
||||
|
||||
from langbot.pkg.api.http.controller.group import AuthType, RouterGroup
|
||||
from langbot.pkg.telemetry import diagnostics as d
|
||||
|
||||
|
||||
class Recorder:
|
||||
enabled = True
|
||||
|
||||
def __init__(self):
|
||||
self.events = []
|
||||
|
||||
def emit(self, kind, operation, outcome, **fields):
|
||||
# A real transport projects the exception class, never its message.
|
||||
error = fields.pop('error', None)
|
||||
if error is not None:
|
||||
fields['error_type'] = type(error).__name__
|
||||
self.events.append(dict(kind=kind, operation=operation, outcome=outcome, **fields))
|
||||
|
||||
|
||||
class Routes(RouterGroup):
|
||||
path = '/management'
|
||||
name = 'management'
|
||||
|
||||
async def initialize(self):
|
||||
@self.route('/ok/<identifier>', auth_type=AuthType.NONE)
|
||||
async def success(identifier):
|
||||
self.ap.seen.append(d.current_span())
|
||||
return self.success({'secret': identifier})
|
||||
|
||||
@self.route('/business', auth_type=AuthType.NONE)
|
||||
async def business():
|
||||
return self.fail('private-error-code', 'private-error-message')
|
||||
|
||||
@self.route('/auth')
|
||||
async def authenticated():
|
||||
raise AssertionError('must not run')
|
||||
|
||||
@self.route('/error', auth_type=AuthType.NONE)
|
||||
async def error():
|
||||
raise ValueError('private-exception-message')
|
||||
|
||||
@self.route('/cancel', auth_type=AuthType.NONE)
|
||||
async def cancel():
|
||||
raise asyncio.CancelledError('private-cancel-message')
|
||||
|
||||
@self.route('/stream', auth_type=AuthType.NONE)
|
||||
async def stream():
|
||||
async def body():
|
||||
self.ap.seen.append(d.current_span())
|
||||
yield b'private-stream-chunk'
|
||||
self.ap.seen.append(d.current_span())
|
||||
|
||||
return quart.Response(body())
|
||||
|
||||
|
||||
async def setup(manager=True):
|
||||
app = quart.Quart(__name__)
|
||||
ap = SimpleNamespace(seen=[])
|
||||
if manager is True:
|
||||
manager = Recorder()
|
||||
if manager is not None:
|
||||
ap.diagnostics = manager
|
||||
routes = Routes(ap, app)
|
||||
await routes.initialize()
|
||||
return app, ap, routes
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_success_uses_code_identity_not_path_or_payload():
|
||||
app, ap, _ = await setup()
|
||||
response = await app.test_client().get(
|
||||
'/management/ok/private-id?token=private-query', headers={'Authorization': 'private-token'}
|
||||
)
|
||||
assert (await response.get_json())['data']['secret'] == 'private-id'
|
||||
assert [e['outcome'] for e in ap.diagnostics.events] == ['started', 'succeeded']
|
||||
event = ap.diagnostics.events[-1]
|
||||
assert event['source'] == 'http'
|
||||
assert re.fullmatch(r'[A-Za-z_][A-Za-z0-9_.:-]{0,127}', event['operation'])
|
||||
assert ap.seen[0] is not None
|
||||
assert d.current_span() is None
|
||||
assert 'private-' not in json.dumps(ap.diagnostics.events)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'path,status,outcome', [('business', 200, 'failed'), ('auth', 401, 'rejected'), ('error', 500, 'failed')]
|
||||
)
|
||||
async def test_http_business_auth_and_exception_outcomes(path, status, outcome):
|
||||
app, ap, _ = await setup()
|
||||
response = await app.test_client().get('/management/' + path)
|
||||
assert response.status_code == status
|
||||
assert ap.diagnostics.events[-1]['outcome'] == outcome
|
||||
assert 'private-' not in json.dumps(ap.diagnostics.events)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_cancellation_propagates():
|
||||
app, ap, _ = await setup()
|
||||
async with app.test_request_context('/management/cancel'):
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await app.full_dispatch_request()
|
||||
assert ap.diagnostics.events[-1]['outcome'] == 'cancelled'
|
||||
assert d.current_span() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_auth_cancellation_not_reinterpreted_as_api_key():
|
||||
app, ap, routes = await setup()
|
||||
routes._authenticate_support_admin = AsyncMock(return_value=None)
|
||||
routes._authenticate_account = AsyncMock(side_effect=asyncio.CancelledError())
|
||||
async with app.test_request_context('/management/auth', headers={'Authorization': 'Bearer private-token'}):
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await app.full_dispatch_request()
|
||||
assert ap.diagnostics.events[-1]['outcome'] == 'cancelled'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_stream_has_parent_context_without_consumer_leak():
|
||||
app, ap, _ = await setup()
|
||||
async with app.test_request_context('/management/stream'):
|
||||
response = await app.full_dispatch_request()
|
||||
assert [e['outcome'] for e in ap.diagnostics.events] == ['started']
|
||||
assert d.current_span() is None
|
||||
async with response.response as body:
|
||||
iterator = body.__aiter__()
|
||||
assert await anext(iterator) == b'private-stream-chunk'
|
||||
assert d.current_span() is None
|
||||
assert ap.seen[-1] is not None
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await anext(iterator)
|
||||
assert ap.seen[0] is ap.seen[1]
|
||||
assert ap.diagnostics.events[-1]['outcome'] == 'succeeded'
|
||||
assert 'private-' not in json.dumps(ap.diagnostics.events)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'manager', [None, SimpleNamespace(enabled=False, emit=lambda *a, **kw: pytest.fail('disabled emission'))]
|
||||
)
|
||||
async def test_absent_disabled_manager_preserves_result_without_context(manager):
|
||||
app, ap, _ = await setup(manager)
|
||||
response = await app.test_client().get('/management/ok/private-id')
|
||||
assert response.status_code == 200
|
||||
assert ap.seen == [None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broken_diagnostic_emit_never_masks_operation():
|
||||
class Broken(Recorder):
|
||||
def emit(self, *args, **kwargs):
|
||||
raise RuntimeError('diagnostics broken')
|
||||
|
||||
app, _, _ = await setup(Broken())
|
||||
response = await app.test_client().get('/management/ok/private-id')
|
||||
assert response.status_code == 200
|
||||
async with app.test_request_context('/management/cancel'):
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await app.full_dispatch_request()
|
||||
assert d.current_span() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anonymous_handlers_have_distinct_stable_code_operations():
|
||||
app, ap, routes = await setup()
|
||||
|
||||
@routes.route('/items/<identifier>', auth_type=AuthType.NONE, methods=['GET'])
|
||||
async def _(identifier):
|
||||
return routes.success()
|
||||
|
||||
@routes.route('/items/<identifier>', auth_type=AuthType.NONE, methods=['POST'])
|
||||
async def _(identifier):
|
||||
return routes.success()
|
||||
|
||||
@routes.route('/other', auth_type=AuthType.NONE)
|
||||
async def _():
|
||||
return routes.success()
|
||||
|
||||
await app.test_client().get('/management/items/private-id')
|
||||
await app.test_client().post('/management/items/private-id')
|
||||
await app.test_client().get('/management/other')
|
||||
operations = [e['operation'] for e in ap.diagnostics.events if e['outcome'] == 'started']
|
||||
assert len(set(operations)) == 3
|
||||
assert all(re.fullmatch(r'[A-Za-z_][A-Za-z0-9_.:-]{0,127}', op) for op in operations)
|
||||
assert 'private-' not in json.dumps(ap.diagnostics.events)
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Verify identities across every source-declared management registration."""
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
|
||||
from langbot.pkg.api.management_diagnostics import operation_id
|
||||
|
||||
|
||||
def test_every_http_registration_has_a_distinct_wire_safe_operation():
|
||||
root = Path(__file__).resolve().parents[3] / 'src/langbot/pkg/api/http/controller/groups'
|
||||
identities = []
|
||||
for path in root.rglob('*.py'):
|
||||
module = 'langbot.pkg.api.http.controller.groups.' + '.'.join(path.relative_to(root).with_suffix('').parts)
|
||||
for cls in ast.walk(ast.parse(path.read_text())):
|
||||
if not isinstance(cls, ast.ClassDef):
|
||||
continue
|
||||
prefix = ''
|
||||
for decorator in cls.decorator_list:
|
||||
if (
|
||||
isinstance(decorator, ast.Call)
|
||||
and isinstance(decorator.func, ast.Attribute)
|
||||
and decorator.func.attr == 'group_class'
|
||||
):
|
||||
prefix = ast.literal_eval(decorator.args[1])
|
||||
for fn in ast.walk(cls):
|
||||
if not isinstance(fn, ast.AsyncFunctionDef):
|
||||
continue
|
||||
for decorator in fn.decorator_list:
|
||||
if not (
|
||||
isinstance(decorator, ast.Call)
|
||||
and isinstance(decorator.func, ast.Attribute)
|
||||
and decorator.func.attr == 'route'
|
||||
):
|
||||
continue
|
||||
rule = prefix + ast.literal_eval(decorator.args[0])
|
||||
methods = next(
|
||||
(ast.literal_eval(k.value) for k in decorator.keywords if k.arg == 'methods'), ['GET']
|
||||
)
|
||||
operation = operation_id(
|
||||
'http', SimpleNamespace(__module__=module, __name__=fn.name), rule=rule, methods=methods
|
||||
)
|
||||
assert re.fullmatch(r'[A-Za-z_][A-Za-z0-9_.:-]{0,127}', operation), operation
|
||||
identities.append(operation)
|
||||
assert len(identities) >= 200 # Guard against accidentally scanning an empty/subset tree.
|
||||
assert len(set(identities)) == len(identities)
|
||||
@@ -0,0 +1,431 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,163 @@
|
||||
"""HTTP streaming termination, fail-open behavior and real producer privacy."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import quart
|
||||
|
||||
from langbot.pkg.agent.runner import errors as runner_errors
|
||||
from langbot.pkg.api import management_diagnostics as md
|
||||
from langbot.pkg.api.http.controller.group import AuthType
|
||||
from langbot.pkg.api.http.controller.groups.agent_debug_stream import debug_stream_response
|
||||
from langbot.pkg.api.mcp.mount import MCPMount
|
||||
from langbot.pkg.telemetry import diagnostics as d
|
||||
from tests.unit_tests.api.test_diagnostics_management_http import Recorder, setup
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('termination', ['close', 'cancel', 'error'])
|
||||
async def test_stream_termination_preserves_cleanup_and_error(termination):
|
||||
app, ap, routes = await setup()
|
||||
closed = []
|
||||
|
||||
@routes.route('/end', auth_type=AuthType.NONE)
|
||||
async def end():
|
||||
async def body():
|
||||
try:
|
||||
yield b'one'
|
||||
if termination == 'error':
|
||||
raise ValueError('private-error')
|
||||
if termination == 'cancel':
|
||||
raise asyncio.CancelledError('private-cancel')
|
||||
yield b'two'
|
||||
finally:
|
||||
closed.append(d.current_span())
|
||||
|
||||
return quart.Response(body())
|
||||
|
||||
async with app.test_request_context('/management/end'):
|
||||
response = await app.full_dispatch_request()
|
||||
if termination == 'close':
|
||||
async with response.response as body:
|
||||
assert await anext(body.__aiter__()) == b'one'
|
||||
else:
|
||||
error = ValueError if termination == 'error' else asyncio.CancelledError
|
||||
with pytest.raises(error):
|
||||
async with response.response as body:
|
||||
iterator = body.__aiter__()
|
||||
assert await anext(iterator) == b'one'
|
||||
await anext(iterator)
|
||||
assert len(closed) == 1 and closed[0] is not None
|
||||
expected = 'failed' if termination == 'error' else 'cancelled'
|
||||
assert [e['outcome'] for e in ap.diagnostics.events] == ['started', expected]
|
||||
assert d.current_span() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_debug_ndjson_error_marks_http_business_failure():
|
||||
app, ap, routes = await setup()
|
||||
service = SimpleNamespace(debug_agent=AsyncMock(side_effect=runner_errors.RunnerNotFoundError('private-error')))
|
||||
|
||||
@routes.route('/ndjson', auth_type=AuthType.NONE)
|
||||
async def ndjson():
|
||||
return debug_stream_response(service, object(), 'private-agent', {'text': 'private-text'})
|
||||
|
||||
response = await app.test_client().get('/management/ndjson')
|
||||
assert response.status_code == 200
|
||||
assert json.loads(await response.get_data())['kind'] == 'error'
|
||||
assert ap.diagnostics.events[-1]['outcome'] == 'failed'
|
||||
assert 'private-' not in json.dumps(ap.diagnostics.events)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_diagnostic_failure_preserves_success_and_stream():
|
||||
class BrokenFinish(Recorder):
|
||||
def emit(self, kind, operation, outcome, **fields):
|
||||
if outcome != 'started':
|
||||
raise RuntimeError('private-diagnostic-error')
|
||||
super().emit(kind, operation, outcome, **fields)
|
||||
|
||||
app, _, _ = await setup(BrokenFinish())
|
||||
response = await app.test_client().get('/management/stream')
|
||||
assert await response.get_data() == b'private-stream-chunk'
|
||||
assert d.current_span() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_manager_status_privacy_and_no_network_in_request():
|
||||
app, ap, _ = await setup(None)
|
||||
ap.instance_config = SimpleNamespace(data={'space': {'url': 'https://example.invalid'}})
|
||||
ap.diagnostics = d.DiagnosticsManager(ap, version='4.11.0-beta.2', instance_id='instance-test', capacity=20)
|
||||
ap.diagnostics.credentials = AsyncMock(side_effect=AssertionError('must not await credentials'))
|
||||
for path in ['ok/private-identifier', 'business', 'auth', 'cancel']:
|
||||
async with app.test_request_context('/management/' + path):
|
||||
try:
|
||||
await app.full_dispatch_request()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
events = list(ap.diagnostics.pending)
|
||||
assert len(events) == 8
|
||||
assert {e['outcome'] for e in events} == {'started', 'succeeded', 'failed', 'rejected', 'cancelled'}
|
||||
assert all(re.fullmatch(r'[A-Za-z_][A-Za-z0-9_.:-]{0,127}', e['operation']) for e in events)
|
||||
assert 'private-' not in json.dumps(events)
|
||||
ap.diagnostics.credentials.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'config,version',
|
||||
[
|
||||
({}, '4.11.0'),
|
||||
({'disable_telemetry': True}, '4.11.0-beta.2'),
|
||||
({'disable_beta_diagnostics': True}, '4.11.0-beta.2'),
|
||||
],
|
||||
)
|
||||
async def test_real_manager_disabled_modes_do_not_install_context(config, version):
|
||||
app, ap, _ = await setup(None)
|
||||
ap.instance_config = SimpleNamespace(data={'space': config})
|
||||
ap.diagnostics = d.DiagnosticsManager(ap, version=version, instance_id='instance-test')
|
||||
response = await app.test_client().get('/management/ok/private-id')
|
||||
assert response.status_code == 200
|
||||
assert ap.seen == [None]
|
||||
assert not ap.diagnostics.pending
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_mount_observes_auth_without_headers_or_body():
|
||||
ap = SimpleNamespace(
|
||||
diagnostics=Recorder(), apikey_service=SimpleNamespace(authenticate_api_key=AsyncMock(return_value=None))
|
||||
)
|
||||
mount = MCPMount(ap)
|
||||
send = AsyncMock()
|
||||
fallback = AsyncMock()
|
||||
await mount.wrap(fallback)(
|
||||
{'type': 'http', 'path': '/mcp/private-path', 'headers': [(b'x-api-key', b'private-token')]}, AsyncMock(), send
|
||||
)
|
||||
assert send.call_args_list[0].args[0]['status'] == 401
|
||||
assert [e['outcome'] for e in ap.diagnostics.events] == ['started', 'rejected']
|
||||
assert ap.diagnostics.events[-1]['operation'] == 'mcp.request'
|
||||
assert 'private-' not in json.dumps(ap.diagnostics.events)
|
||||
fallback.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_debug_boundary_source_inherits_without_duplicate_run():
|
||||
recorder = Recorder()
|
||||
ap = SimpleNamespace(diagnostics=recorder)
|
||||
|
||||
@md.observe('http.test.debug', source='webui_debug', ap=ap)
|
||||
async def debug():
|
||||
@d.observe('run', 'runner.run', source='agent', ap=ap)
|
||||
async def run():
|
||||
return 'private-content'
|
||||
|
||||
return await run()
|
||||
|
||||
assert await debug() == 'private-content'
|
||||
assert len([e for e in recorder.events if e['kind'] == 'run' and e['outcome'] == 'started']) == 1
|
||||
assert all(e['source'] == 'webui_debug' and e['attributes']['synthetic'] for e in recorder.events)
|
||||
assert 'private-' not in json.dumps(recorder.events)
|
||||
@@ -0,0 +1,173 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user