Files
LangBot/tests/unit_tests/api/test_execution_application_api.py
T

525 lines
21 KiB
Python

from __future__ import annotations
import asyncio
import contextlib
import contextvars
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding
from langbot_plugin.entities.io.resp import ActionResponse
from langbot.pkg.api.application.contracts import ApplicationAPIException
from langbot.pkg.api.application.execution import ExecutionAPI, ExecutionAPIContext
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
from langbot.pkg.agent.runner.session_registry import get_session_registry
@pytest.fixture
def host():
binding = InstallationBinding(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=7,
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
tenant = contextvars.ContextVar('execution_test_tenant', default=None)
class Persistence:
execute_async = AsyncMock()
@contextlib.asynccontextmanager
async def tenant_scope(self, workspace):
token = tenant.set(workspace)
try:
yield
finally:
tenant.reset(token)
ap = SimpleNamespace(
deployment=SimpleNamespace(mode='cloud'),
persistence_mgr=Persistence(),
logger=Mock(),
workspace_service=SimpleNamespace(get_execution_binding=AsyncMock(return_value=binding)),
instance_config=SimpleNamespace(data={}),
)
handler = RuntimeConnectionHandler(Mock(), AsyncMock(return_value=True), ap)
handler.register_installation_binding(binding, plugin_author='author-a', plugin_name='plugin-a')
ap.plugin_connector = SimpleNamespace(handler=handler)
return SimpleNamespace(ap=ap, handler=handler, binding=binding, api=ExecutionAPI(ap), tenant=tenant)
def run_session(host, **changes):
result = {
'run_id': 'run-a',
'plugin_identity': 'author-a/plugin-a',
'authorization': {
'workspace_id': host.binding.workspace_uuid,
'available_apis': {},
'resources': {},
'authorized_ids': {},
},
}
result.update(changes)
return result
def stub_run(monkeypatch, session):
monkeypatch.setattr(get_session_registry(), 'get', AsyncMock(return_value=session))
def install_action(host, name, action):
# The production constructor applies this wrapper after registration.
host.handler.actions[name] = action
host.handler._secure_plugin_actions()
@pytest.mark.asyncio
async def test_catalog_covers_all_core_business_actions_and_keeps_transport_internal(host):
excluded = {'initialize_plugin_settings', 'get_plugin_settings'}
expected = {name for name in host.handler.actions if not name.startswith('__') and name not in excluded}
assert {op.action for op in host.api.operations()} == expected
assert len(host.api.catalog()) == len(expected)
for operation in host.api.catalog():
assert operation['input_schema']['additionalProperties'] is False
assert operation['output_schema']
result = await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(host.binding))
assert result['version']
assert host.handler.current_action_context is None
assert host.tenant.get() is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
'name', ['__ping', '__file_chunk', 'initialize_plugin_settings', 'get_plugin_settings', 'missing']
)
async def test_runtime_control_actions_are_not_exposed(host, name):
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(f'execution.{name}', ExecutionAPIContext(host.binding))
assert exc.value.error.code == 'unknown_operation'
@pytest.mark.asyncio
@pytest.mark.parametrize(
'payload',
[
{'run_id': None},
{'run_id': 'another-run'},
{'workspace_uuid': 'workspace-b'},
{'caller_plugin_identity': 'other/plugin'},
{'action_context': {}},
{'placement_generation': 8},
],
)
async def test_payload_cannot_supply_authentication_fields(host, payload):
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(host.binding), payload)
assert exc.value.error.code == 'invalid_input'
assert host.handler.current_action_context is None
@pytest.mark.asyncio
@pytest.mark.parametrize('field,value', [('runtime_revision', 2), ('artifact_digest', 'b' * 64)])
async def test_stale_installation_is_rejected_before_business_dispatch(host, field, value):
altered = host.binding.model_copy(update={field: value})
business = AsyncMock(return_value=ActionResponse.success({'version': 'test'}))
install_action(host, 'get_langbot_version', business)
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(altered))
assert exc.value.error.code == 'permission_denied'
business.assert_not_called()
assert host.handler.current_action_context is None
assert host.tenant.get() is None
@pytest.mark.asyncio
async def test_stale_generation_and_incomplete_binding_are_rejected(host):
host.ap.workspace_service.get_execution_binding.return_value = host.binding.model_copy(
update={'placement_generation': 8}
)
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(host.binding))
assert exc.value.error.code == 'permission_denied'
incomplete = ActionContext(instance_uuid='instance-a', workspace_uuid='workspace-a', placement_generation=7)
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(incomplete))
assert exc.value.error.code == 'invalid_authentication'
@pytest.mark.asyncio
@pytest.mark.parametrize('session', [None, {'workspace': 'workspace-b'}, {'plugin': 'other/plugin'}])
async def test_run_context_cannot_cross_workspace_or_plugin(host, monkeypatch, session):
actual = run_session(host)
if session is None:
actual = None
elif 'workspace' in session:
actual['authorization']['workspace_id'] = session['workspace']
else:
actual['plugin_identity'] = session['plugin']
stub_run(monkeypatch, actual)
business = AsyncMock(return_value=ActionResponse.success({'version': 'test'}))
install_action(host, 'get_langbot_version', business)
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(host.binding, 'run-a'))
assert exc.value.error.code == 'permission_denied'
business.assert_not_called()
@pytest.mark.asyncio
async def test_run_identity_is_injected_and_cannot_fall_back_to_plugin_privileges(host, monkeypatch):
stub_run(monkeypatch, run_session(host))
business = AsyncMock(return_value=ActionResponse.success({'version': 'test'}))
install_action(host, 'get_langbot_version', business)
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(host.binding, 'run-a'))
assert business.call_args.args[0] == {'run_id': 'run-a', 'caller_plugin_identity': 'author-a/plugin-a'}
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_bots', ExecutionAPIContext(host.binding, 'run-a'))
assert exc.value.error.code == 'permission_denied'
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_prompt', ExecutionAPIContext(host.binding))
assert exc.value.error.code == 'permission_denied'
@pytest.mark.asyncio
async def test_resource_denial_uses_real_secured_host_action(host, monkeypatch):
stub_run(monkeypatch, run_session(host))
host.ap.model_mgr = SimpleNamespace(get_model_by_uuid=AsyncMock())
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.invoke_llm',
ExecutionAPIContext(host.binding, 'run-a'),
{
'llm_model_uuid': 'ungranted-model',
'messages': [],
},
)
assert 'not authorized' in exc.value.error.message
host.ap.model_mgr.get_model_by_uuid.assert_not_called()
assert host.handler.current_action_context is None
@pytest.mark.asyncio
async def test_run_capability_denial_uses_real_secured_host_action(host, monkeypatch):
stub_run(monkeypatch, run_session(host))
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.state_get', ExecutionAPIContext(host.binding, 'run-a'), {'scope': 'conversation', 'key': 'x'}
)
assert 'not authorized' in exc.value.error.message
assert host.handler.current_action_context is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
'payload',
[
{'llm_model_uuid': 1, 'messages': []},
{'llm_model_uuid': 'model', 'messages': 'invalid'},
{'llm_model_uuid': 'model', 'messages': [], 'extra_args': {'x': object()}},
],
)
async def test_input_type_validation_precedes_provider_calls(host, payload):
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.invoke_llm', ExecutionAPIContext(host.binding), payload)
assert exc.value.error.code == 'invalid_input'
@pytest.mark.asyncio
@pytest.mark.parametrize(
'response',
[
ActionResponse.success({'tokens': 'wrong'}),
ActionResponse.success({}),
ActionResponse.success({'tokens': 1, 'unknown': object()}),
{'tokens': 1},
],
)
async def test_invalid_output_is_not_reported_as_success(host, response):
install_action(host, 'count_tokens', AsyncMock(return_value=response))
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.count_tokens', ExecutionAPIContext(host.binding), {'llm_model_uuid': 'model', 'messages': []}
)
assert exc.value.error.code == 'invalid_output'
@pytest.mark.asyncio
async def test_failure_and_cancellation_restore_outer_context(host):
outer = host.binding.model_copy(update={'workspace_uuid': 'outer'})
token = host.handler._current_action_context.set(outer)
try:
install_action(host, 'get_langbot_version', AsyncMock(side_effect=RuntimeError('private provider error')))
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(host.binding))
assert exc.value.error.message == 'Execution operation failed'
assert host.handler.current_action_context is outer
install_action(host, 'get_langbot_version', AsyncMock(side_effect=asyncio.CancelledError()))
with pytest.raises(asyncio.CancelledError):
await host.api.invoke('execution.get_langbot_version', ExecutionAPIContext(host.binding))
assert host.handler.current_action_context is outer
assert host.tenant.get() is None
finally:
host.handler._current_action_context.reset(token)
@pytest.mark.asyncio
async def test_stream_isolates_context_from_consumer_and_closes_on_early_exit(host):
closed = asyncio.Event()
async def provider(data):
try:
assert host.handler.current_action_context == host.binding
assert host.tenant.get() == 'workspace-a'
yield ActionResponse.success({'chunk': {'text': 'first'}})
await asyncio.Event().wait()
finally:
assert host.handler.current_action_context == host.binding
closed.set()
install_action(host, 'invoke_llm_stream', provider)
stream = host.api.stream(
'execution.invoke_llm_stream', ExecutionAPIContext(host.binding), {'llm_model_uuid': 'model', 'messages': []}
)
async with contextlib.aclosing(stream):
assert await anext(stream) == {'chunk': {'text': 'first'}}
assert host.handler.current_action_context is None
assert host.tenant.get() is None
assert closed.is_set()
assert host.handler.current_action_context is None
assert host.tenant.get() is None
@pytest.mark.asyncio
async def test_stream_reports_late_failure_and_cancels_provider(host):
closed = asyncio.Event()
async def provider(data):
try:
yield ActionResponse.success({'chunk': {'text': 'first'}})
yield ActionResponse.error('provider refused request')
finally:
closed.set()
install_action(host, 'invoke_llm_stream', provider)
async with contextlib.aclosing(
host.api.stream(
'execution.invoke_llm_stream',
ExecutionAPIContext(host.binding),
{'llm_model_uuid': 'model', 'messages': []},
)
) as stream:
assert await anext(stream) == {'chunk': {'text': 'first'}}
with pytest.raises(ApplicationAPIException, match='provider refused request'):
await anext(stream)
assert closed.is_set()
assert host.tenant.get() is None
@pytest.mark.asyncio
async def test_stream_cancellation_closes_provider(host):
started, closed = asyncio.Event(), asyncio.Event()
async def provider(data):
try:
started.set()
await asyncio.Event().wait()
yield ActionResponse.success({'chunk': {}})
finally:
closed.set()
install_action(host, 'invoke_llm_stream', provider)
stream = host.api.stream(
'execution.invoke_llm_stream', ExecutionAPIContext(host.binding), {'llm_model_uuid': 'model', 'messages': []}
)
consumer = asyncio.create_task(anext(stream))
await asyncio.wait_for(started.wait(), timeout=2)
consumer.cancel()
with pytest.raises(asyncio.CancelledError):
await consumer
assert closed.is_set()
assert host.handler.current_action_context is None
assert host.tenant.get() is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
'operation,payload',
[
('call_tool', {'tool_name': 'ungranted-tool', 'parameters': {}}),
('get_tool_detail', {'tool_name': 'ungranted-tool'}),
('retrieve_knowledge', {'kb_id': 'ungranted-kb', 'query_text': 'test'}),
('invoke_embedding', {'embedding_model_uuid': 'ungranted-model', 'texts': ['test']}),
('invoke_rerank', {'rerank_model_uuid': 'ungranted-model', 'query': 'test', 'documents': ['test']}),
('get_binary_storage', {'owner_type': 'workspace', 'key': 'key'}),
],
)
async def test_resource_families_preserve_frozen_grants(host, monkeypatch, operation, payload):
stub_run(monkeypatch, run_session(host))
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(f'execution.{operation}', ExecutionAPIContext(host.binding, 'run-a'), payload)
assert 'not authorized' in exc.value.error.message
assert host.handler.current_action_context is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
'operation,payload',
[
('history_page', {}),
('event_get', {'event_id': 'event-a'}),
('bind_box', {'box_id': 'box-a'}),
('export_box_files', {'paths': ['report.txt']}),
],
)
async def test_run_capability_families_cannot_fall_back_to_plugin_access(host, monkeypatch, operation, payload):
stub_run(monkeypatch, run_session(host))
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(f'execution.{operation}', ExecutionAPIContext(host.binding, 'run-a'), payload)
assert 'not authorized' in exc.value.error.message
@pytest.mark.asyncio
async def test_box_export_preserves_structured_failure(host, monkeypatch):
from langbot.pkg.box.export_errors import export_error
stub_run(monkeypatch, run_session(host))
error = export_error('file_bytes', limit=100, actual=101, path='report.txt')
install_action(
host,
'export_box_files',
AsyncMock(
return_value=ActionResponse.error(
str(error),
data={'error': error.error.model_dump(mode='json')},
)
),
)
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.export_box_files', ExecutionAPIContext(host.binding, 'run-a'), {'paths': ['report.txt']}
)
assert exc.value.error.details['error']['code'] == 'box.export.limit_exceeded'
assert exc.value.error.details['error']['details']['files_retained'] is True
@pytest.mark.asyncio
async def test_frozen_tool_source_cannot_be_replaced(host, monkeypatch):
session = run_session(host)
session['authorization']['authorized_ids'] = {'tool': {'tool-a'}}
session['authorization']['resources'] = {'tools': [{'tool_name': 'tool-a'}]}
stub_run(monkeypatch, session)
host.ap.tool_mgr = SimpleNamespace(execute_func_call=AsyncMock())
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.call_tool', ExecutionAPIContext(host.binding, 'run-a'), {'tool_name': 'tool-a', 'parameters': {}}
)
assert 'source' in exc.value.error.message.lower()
host.ap.tool_mgr.execute_func_call.assert_not_called()
@pytest.mark.asyncio
async def test_runner_admin_operations_do_not_grant_admin_implicitly(host):
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke('execution.run_stats', ExecutionAPIContext(host.binding))
assert 'not authorized' in exc.value.error.message
@pytest.mark.asyncio
async def test_query_identity_cannot_be_replaced_in_run_context(host, monkeypatch):
stub_run(monkeypatch, run_session(host))
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.invoke_llm',
ExecutionAPIContext(host.binding, 'run-a'),
{
'llm_model_uuid': 'model',
'messages': [],
'query_id': 'other-query',
},
)
assert exc.value.error.code == 'invalid_input'
@pytest.mark.asyncio
async def test_provider_cancellation_cannot_leave_stream_consumer_waiting(host):
async def provider(data):
raise asyncio.CancelledError()
yield
install_action(host, 'invoke_llm_stream', provider)
stream = host.api.stream(
'execution.invoke_llm_stream', ExecutionAPIContext(host.binding), {'llm_model_uuid': 'model', 'messages': []}
)
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(anext(stream), timeout=2)
assert host.tenant.get() is None
@pytest.mark.asyncio
async def test_concurrent_streams_keep_independent_installation_contexts(host):
other = host.binding.model_copy(
update={'workspace_uuid': 'workspace-b', 'installation_uuid': '00000000-0000-4000-8000-000000000002'}
)
host.handler.register_installation_binding(other, plugin_author='author-a', plugin_name='plugin-a')
host.ap.workspace_service.get_execution_binding.side_effect = (
lambda workspace, **kwargs: host.binding if workspace == 'workspace-a' else other
)
async def provider(data):
workspace = host.handler.current_action_context.workspace_uuid
for _ in range(2):
await asyncio.sleep(0)
assert host.handler.current_action_context.workspace_uuid == workspace
assert host.tenant.get() == workspace
yield ActionResponse.success({'chunk': {'workspace': workspace}})
install_action(host, 'invoke_llm_stream', provider)
async def consume(binding):
results = []
async with contextlib.aclosing(
host.api.stream(
'execution.invoke_llm_stream', ExecutionAPIContext(binding), {'llm_model_uuid': 'model', 'messages': []}
)
) as stream:
async for chunk in stream:
assert host.handler.current_action_context is None
assert host.tenant.get() is None
results.append(chunk['chunk']['workspace'])
return results
assert await asyncio.gather(consume(host.binding), consume(other)) == [['workspace-a'] * 2, ['workspace-b'] * 2]
@pytest.mark.asyncio
async def test_tool_parameters_reach_legacy_plugin_and_run_handlers(host, monkeypatch):
business = AsyncMock(return_value=ActionResponse.success({'result': 'ok'}))
install_action(host, 'call_tool', business)
await host.api.invoke(
'execution.call_tool', ExecutionAPIContext(host.binding), {'tool_name': 'tool', 'parameters': {'x': 1}}
)
assert business.call_args.args[0]['tool_parameters'] == {'x': 1}
assert 'parameters' not in business.call_args.args[0]
stub_run(monkeypatch, run_session(host))
await host.api.invoke('execution.call_tool', ExecutionAPIContext(host.binding, 'run-a'), {'tool_name': 'tool'})
assert business.call_args.args[0]['parameters'] == {}
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.call_tool',
ExecutionAPIContext(host.binding, 'run-a'),
{'tool_name': 'tool', 'tool_parameters': {}},
)
assert exc.value.error.code == 'invalid_input'
@pytest.mark.asyncio
async def test_platform_call_requires_bot_identity_outside_run(host):
with pytest.raises(ApplicationAPIException) as exc:
await host.api.invoke(
'execution.call_platform_api', ExecutionAPIContext(host.binding), {'action': 'send_message'}
)
assert exc.value.error.code == 'invalid_input'