mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-17 01:46:07 +00:00
feat(agent-runner): enforce 4.x host-owned execution
This commit is contained in:
huanghuoguoguo
@@ -15,6 +15,20 @@ from mcp import types as mcp_types
|
||||
from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_loopback_transport_from_proxy_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Keep real loopback transport tests independent of workstation proxies."""
|
||||
for variable in (
|
||||
'ALL_PROXY',
|
||||
'all_proxy',
|
||||
'HTTP_PROXY',
|
||||
'http_proxy',
|
||||
'HTTPS_PROXY',
|
||||
'https_proxy',
|
||||
):
|
||||
monkeypatch.delenv(variable, raising=False)
|
||||
|
||||
|
||||
class _TransportProbe:
|
||||
def __init__(self, streamable_status: int | None) -> None:
|
||||
self.streamable_status = streamable_status
|
||||
|
||||
@@ -8,9 +8,12 @@ import httpx
|
||||
import pytest
|
||||
from mcp import types as mcp_types
|
||||
|
||||
from langbot.pkg.agent.runner.execution_context import project_mcp_resource_config
|
||||
from langbot.pkg.provider.tools.errors import ToolExecutionDeniedError
|
||||
from langbot.pkg.provider.tools.loaders.mcp import (
|
||||
MCP_RESOURCE_CONTEXT_QUERY_KEY,
|
||||
MCP_RESOURCE_TRACE_QUERY_KEY,
|
||||
MCP_READ_RESOURCE_SCHEMA,
|
||||
MCP_TOOL_LIST_RESOURCES,
|
||||
MCP_TOOL_READ_RESOURCE,
|
||||
MCPLoader,
|
||||
@@ -257,24 +260,43 @@ async def test_mcp_loader_can_hide_synthetic_resource_tools():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_loader_refuses_resource_tool_calls_when_agent_read_disabled():
|
||||
async def test_mcp_loader_get_tool_returns_synthetic_resource_schema():
|
||||
loader = MCPLoader(_app())
|
||||
loader.sessions = {'docs': _connected_session()}
|
||||
|
||||
tool = await loader.get_tool(MCP_TOOL_READ_RESOURCE)
|
||||
|
||||
assert tool is not None
|
||||
assert tool.name == MCP_TOOL_READ_RESOURCE
|
||||
assert tool.parameters == MCP_READ_RESOURCE_SCHEMA
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('read_enabled', [False, 0, None, 'false'])
|
||||
@pytest.mark.parametrize(
|
||||
('tool_name', 'parameters'),
|
||||
[
|
||||
(MCP_TOOL_LIST_RESOURCES, {'server_name': 'docs'}),
|
||||
(MCP_TOOL_READ_RESOURCE, {'server_name': 'docs', 'uri': 'file:///README.md'}),
|
||||
],
|
||||
)
|
||||
async def test_mcp_loader_refuses_resource_tool_calls_when_agent_read_disabled(
|
||||
read_enabled,
|
||||
tool_name,
|
||||
parameters,
|
||||
):
|
||||
loader = MCPLoader(_app())
|
||||
session = _connected_session()
|
||||
loader.sessions = {'docs': session}
|
||||
query = SimpleNamespace(
|
||||
variables={
|
||||
'_pipeline_bound_mcp_servers': ['srv-1'],
|
||||
'_pipeline_mcp_resource_agent_read_enabled': False,
|
||||
}
|
||||
)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
MCP_TOOL_READ_RESOURCE,
|
||||
{'server_name': 'docs', 'uri': 'file:///README.md'},
|
||||
query = SimpleNamespace(variables={'_pipeline_bound_mcp_servers': ['srv-1']})
|
||||
project_mcp_resource_config(
|
||||
query,
|
||||
{'mcp-resource-agent-read-enabled': read_enabled},
|
||||
)
|
||||
|
||||
assert result[0].text == 'Error: MCP resource agent reads are disabled.'
|
||||
with pytest.raises(ToolExecutionDeniedError, match='MCP resource agent reads are disabled'):
|
||||
await loader.invoke_tool(tool_name, parameters, query)
|
||||
|
||||
session.session.read_resource.assert_not_called()
|
||||
|
||||
|
||||
@@ -321,3 +343,93 @@ async def test_build_resource_context_for_query_uses_only_bound_attached_text_re
|
||||
assert query.variables[MCP_RESOURCE_CONTEXT_QUERY_KEY]['resource_count'] == 1
|
||||
docs.session.read_resource.assert_awaited_once()
|
||||
other.session.read_resource.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
('enabled_config', 'expected_enabled'),
|
||||
[
|
||||
pytest.param({}, True, id='missing'),
|
||||
pytest.param({'enabled': True}, True, id='true'),
|
||||
pytest.param({'enabled': False}, False, id='false'),
|
||||
pytest.param({'enabled': 0}, False, id='zero'),
|
||||
pytest.param({'enabled': None}, False, id='none'),
|
||||
pytest.param({'enabled': 'false'}, False, id='string-false'),
|
||||
],
|
||||
)
|
||||
async def test_build_resource_context_attachment_enabled_fails_closed(enabled_config, expected_enabled):
|
||||
loader = MCPLoader(_app())
|
||||
session = _connected_session(name='docs', uuid='srv-1')
|
||||
session.session.read_resource.return_value = mcp_types.ReadResourceResult(
|
||||
contents=[
|
||||
mcp_types.TextResourceContents(
|
||||
uri='file:///README.md',
|
||||
mimeType='text/markdown',
|
||||
text='enabled attachment',
|
||||
)
|
||||
]
|
||||
)
|
||||
loader.sessions = {'docs': session}
|
||||
attachment = {
|
||||
'server_uuid': 'srv-1',
|
||||
'server_name': 'docs',
|
||||
'uri': 'file:///README.md',
|
||||
'mode': 'pinned',
|
||||
**enabled_config,
|
||||
}
|
||||
query = SimpleNamespace(
|
||||
variables={
|
||||
'_pipeline_bound_mcp_servers': ['srv-1'],
|
||||
'_pipeline_mcp_resource_attachments': [attachment],
|
||||
}
|
||||
)
|
||||
|
||||
context = await loader.build_resource_context_for_query(query)
|
||||
|
||||
if expected_enabled:
|
||||
assert 'enabled attachment' in context
|
||||
session.session.read_resource.assert_awaited_once()
|
||||
else:
|
||||
assert context == ''
|
||||
session.session.read_resource.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
('field_name', 'invalid_value'),
|
||||
[
|
||||
pytest.param('max_tokens', '100', id='string-tokens'),
|
||||
pytest.param('max_tokens', 0, id='zero-tokens'),
|
||||
pytest.param('max_tokens', -1, id='negative-tokens'),
|
||||
pytest.param('max_tokens', True, id='boolean-tokens'),
|
||||
pytest.param('max_bytes', '4096', id='string-bytes'),
|
||||
pytest.param('max_bytes', 0, id='zero-bytes'),
|
||||
pytest.param('max_bytes', -1, id='negative-bytes'),
|
||||
pytest.param('max_bytes', True, id='boolean-bytes'),
|
||||
],
|
||||
)
|
||||
async def test_build_resource_context_invalid_attachment_limits_fail_closed(field_name, invalid_value):
|
||||
ap = _app()
|
||||
loader = MCPLoader(ap)
|
||||
session = _connected_session(name='docs', uuid='srv-1')
|
||||
loader.sessions = {'docs': session}
|
||||
query = SimpleNamespace(
|
||||
variables={
|
||||
'_pipeline_bound_mcp_servers': ['srv-1'],
|
||||
'_pipeline_mcp_resource_attachments': [
|
||||
{
|
||||
'server_uuid': 'srv-1',
|
||||
'server_name': 'docs',
|
||||
'uri': 'file:///README.md',
|
||||
'mode': 'pinned',
|
||||
field_name: invalid_value,
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
context = await loader.build_resource_context_for_query(query)
|
||||
|
||||
assert context == ''
|
||||
session.session.read_resource.assert_not_awaited()
|
||||
ap.logger.warning.assert_called_once()
|
||||
|
||||
@@ -193,11 +193,13 @@ class TestPersistActivatedSkill:
|
||||
|
||||
query = SimpleNamespace(variables={ACTIVATED_SKILLS_KEY: {'pdf': {'name': 'pdf'}}})
|
||||
query._agent_run_session = {
|
||||
'runner_id': 'plugin:langbot-team/LocalAgent/default',
|
||||
'state_context': {
|
||||
'scope_keys': {'conversation': 'conv-scope-key'},
|
||||
'binding_identity': 'binding-1',
|
||||
'conversation_id': 'c1',
|
||||
'runner_id': 'plugin:test/runner/default',
|
||||
'authorization': {
|
||||
'state_context': {
|
||||
'scope_keys': {'conversation': 'conv-scope-key'},
|
||||
'binding_identity': 'binding-1',
|
||||
'conversation_id': 'c1',
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -214,7 +216,7 @@ class TestPersistActivatedSkill:
|
||||
assert kwargs['state_key'] == ACTIVATED_SKILL_NAMES_STATE_KEY
|
||||
assert kwargs['value'] == ['pdf']
|
||||
assert kwargs['scope'] == 'conversation'
|
||||
assert kwargs['runner_id'] == 'plugin:langbot-team/LocalAgent/default'
|
||||
assert kwargs['runner_id'] == 'plugin:test/runner/default'
|
||||
assert kwargs['binding_identity'] == 'binding-1'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -7,12 +7,15 @@ Tests cover:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from importlib import import_module
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.provider.tools.errors import ToolExecutionDeniedError
|
||||
|
||||
|
||||
def get_toolmgr_module():
|
||||
@@ -259,6 +262,119 @@ class TestToolManagerExecuteFuncCall:
|
||||
mock_plugin_loader.invoke_tool.assert_called_once()
|
||||
mock_mcp_loader.invoke_tool.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bound_mcp_source_wins_over_unbound_plugin_name_collision(self, mock_app_with_loaders):
|
||||
toolmgr = get_toolmgr_module()
|
||||
mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
|
||||
mock_plugin_loader.has_tool = AsyncMock(return_value=True)
|
||||
mock_mcp_loader.has_tool = AsyncMock(return_value=True)
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
|
||||
query = Mock(variables={})
|
||||
|
||||
result = await manager.execute_func_call(
|
||||
'shared_tool',
|
||||
{'value': 1},
|
||||
query,
|
||||
source_ref={'source': 'mcp', 'source_id': 'bound-mcp'},
|
||||
)
|
||||
|
||||
assert result == 'mcp_result'
|
||||
mock_plugin_loader.has_tool.assert_not_awaited()
|
||||
mock_plugin_loader.invoke_tool.assert_not_awaited()
|
||||
mock_mcp_loader.has_tool.assert_awaited_once_with('shared_tool', source_id='bound-mcp')
|
||||
mock_mcp_loader.invoke_tool.assert_awaited_once_with(
|
||||
'shared_tool',
|
||||
{'value': 1},
|
||||
query,
|
||||
source_id='bound-mcp',
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bound_plugin_source_selects_one_of_two_same_name_plugins(self, mock_app_with_loaders):
|
||||
toolmgr = get_toolmgr_module()
|
||||
mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
|
||||
mock_plugin_loader.has_tool = AsyncMock(return_value=True)
|
||||
mock_mcp_loader.has_tool = AsyncMock(return_value=True)
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
|
||||
query = Mock(variables={})
|
||||
|
||||
result = await manager.execute_func_call(
|
||||
'shared_tool',
|
||||
{},
|
||||
query,
|
||||
source_ref={'source': 'plugin', 'source_id': 'authorized/plugin'},
|
||||
)
|
||||
|
||||
assert result == 'plugin_result'
|
||||
mock_plugin_loader.has_tool.assert_awaited_once_with('shared_tool', source_id='authorized/plugin')
|
||||
mock_plugin_loader.invoke_tool.assert_awaited_once_with(
|
||||
'shared_tool',
|
||||
{},
|
||||
query,
|
||||
source_id='authorized/plugin',
|
||||
)
|
||||
mock_mcp_loader.has_tool.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_synthetic_mcp_tool_is_monitored_as_error(
|
||||
self,
|
||||
mock_app_with_loaders,
|
||||
sample_query,
|
||||
):
|
||||
toolmgr = get_toolmgr_module()
|
||||
mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
|
||||
mock_app.monitoring_service = SimpleNamespace(record_tool_call=AsyncMock())
|
||||
sample_query.variables = {}
|
||||
sample_query.launcher_type = None
|
||||
sample_query.launcher_id = None
|
||||
sample_query.bot_uuid = 'bot-1'
|
||||
mock_mcp_loader.has_tool = AsyncMock(return_value=True)
|
||||
mock_mcp_loader.invoke_tool = AsyncMock(
|
||||
side_effect=ToolExecutionDeniedError(
|
||||
'langbot_mcp_read_resource',
|
||||
'MCP resource agent reads are disabled',
|
||||
)
|
||||
)
|
||||
manager = toolmgr.ToolManager(mock_app)
|
||||
self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
|
||||
|
||||
with pytest.raises(ToolExecutionDeniedError):
|
||||
await manager.execute_func_call(
|
||||
'langbot_mcp_read_resource',
|
||||
{'server_name': 'docs', 'uri': 'file:///README.md'},
|
||||
sample_query,
|
||||
source_ref={'source': 'mcp', 'source_id': None},
|
||||
)
|
||||
|
||||
record = mock_app.monitoring_service.record_tool_call
|
||||
record.assert_awaited_once()
|
||||
assert record.await_args.kwargs['status'] == 'error'
|
||||
assert record.await_args.kwargs['tool_source'] == 'mcp'
|
||||
assert 'MCP resource agent reads are disabled' in record.await_args.kwargs['error_message']
|
||||
|
||||
|
||||
class TestToolManagerSourceResolution:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('source', ['mcp', 'plugin'])
|
||||
async def test_same_name_implementations_in_one_scope_are_hidden(self, source):
|
||||
toolmgr = get_toolmgr_module()
|
||||
app = Mock(logger=Mock())
|
||||
manager = toolmgr.ToolManager(app)
|
||||
manager.get_tool_catalog = AsyncMock(
|
||||
return_value=[
|
||||
{'name': 'shared_tool', 'source': source, 'source_id': 'source-a'},
|
||||
{'name': 'shared_tool', 'source': source, 'source_id': 'source-b'},
|
||||
{'name': 'unique_tool', 'source': 'builtin'},
|
||||
]
|
||||
)
|
||||
|
||||
catalog = await manager.get_resolved_tool_catalog()
|
||||
|
||||
assert [item['name'] for item in catalog] == ['unique_tool']
|
||||
app.logger.warning.assert_called_once()
|
||||
|
||||
|
||||
class TestToolManagerShutdown:
|
||||
"""Tests for shutdown method."""
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.provider.tools.loaders.mcp import (
|
||||
MCP_TOOL_LIST_RESOURCES,
|
||||
MCP_TOOL_READ_RESOURCE,
|
||||
MCPLoader,
|
||||
MCPSessionStatus,
|
||||
)
|
||||
from langbot.pkg.provider.tools.loaders.plugin import PluginToolLoader
|
||||
from langbot_plugin.api.entities.builtin.resource.tool import LLMTool
|
||||
|
||||
|
||||
def make_tool(name: str) -> LLMTool:
|
||||
return LLMTool(
|
||||
name=name,
|
||||
human_desc=name,
|
||||
description=name,
|
||||
parameters={'type': 'object', 'properties': {}},
|
||||
func=lambda parameters: parameters,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_two_mcp_servers_with_same_tool_route_to_authorized_server():
|
||||
tool = make_tool('shared_tool')
|
||||
first = SimpleNamespace(
|
||||
server_uuid='mcp-a',
|
||||
get_tools=Mock(return_value=[tool]),
|
||||
invoke_mcp_tool=AsyncMock(return_value='from-a'),
|
||||
)
|
||||
second = SimpleNamespace(
|
||||
server_uuid='mcp-b',
|
||||
get_tools=Mock(return_value=[tool]),
|
||||
invoke_mcp_tool=AsyncMock(return_value='from-b'),
|
||||
)
|
||||
loader = MCPLoader(SimpleNamespace(logger=Mock()))
|
||||
loader.sessions = {'first': first, 'second': second}
|
||||
query = SimpleNamespace(variables={})
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'shared_tool',
|
||||
{'value': 1},
|
||||
query,
|
||||
source_id='mcp-b',
|
||||
)
|
||||
|
||||
assert result == 'from-b'
|
||||
first.invoke_mcp_tool.assert_not_awaited()
|
||||
second.invoke_mcp_tool.assert_awaited_once_with(
|
||||
'shared_tool',
|
||||
{'value': 1},
|
||||
query=query,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('reserved_name', [MCP_TOOL_LIST_RESOURCES, MCP_TOOL_READ_RESOURCE])
|
||||
async def test_reserved_name_routes_to_real_mcp_tool_when_source_id_is_present(reserved_name):
|
||||
real_tool = make_tool(reserved_name)
|
||||
session = SimpleNamespace(
|
||||
server_uuid='srv-real',
|
||||
get_tools=Mock(return_value=[real_tool]),
|
||||
invoke_mcp_tool=AsyncMock(return_value='real-result'),
|
||||
)
|
||||
loader = MCPLoader(SimpleNamespace(logger=Mock()))
|
||||
loader.sessions = {'real': session}
|
||||
query = SimpleNamespace(variables={})
|
||||
|
||||
assert await loader.has_tool(reserved_name, source_id='srv-real') is True
|
||||
assert await loader.get_tool(reserved_name, source_id='srv-real') is real_tool
|
||||
result = await loader.invoke_tool(reserved_name, {}, query, source_id='srv-real')
|
||||
|
||||
assert result == 'real-result'
|
||||
session.invoke_mcp_tool.assert_awaited_once_with(reserved_name, {}, query=query)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
('reserved_name', 'invoke_method'),
|
||||
[
|
||||
(MCP_TOOL_LIST_RESOURCES, '_invoke_mcp_list_resources'),
|
||||
(MCP_TOOL_READ_RESOURCE, '_invoke_mcp_read_resource'),
|
||||
],
|
||||
)
|
||||
async def test_reserved_name_uses_host_synthetic_tool_when_source_id_is_none(
|
||||
reserved_name,
|
||||
invoke_method,
|
||||
):
|
||||
real_tool = make_tool(reserved_name)
|
||||
session = SimpleNamespace(
|
||||
server_uuid='srv-real',
|
||||
server_name='real',
|
||||
enable=True,
|
||||
status=MCPSessionStatus.CONNECTED,
|
||||
session=object(),
|
||||
get_tools=Mock(return_value=[real_tool]),
|
||||
has_resource_support=Mock(return_value=True),
|
||||
invoke_mcp_tool=AsyncMock(return_value='real-result'),
|
||||
)
|
||||
loader = MCPLoader(SimpleNamespace(logger=Mock()))
|
||||
loader.sessions = {'real': session}
|
||||
synthetic_invoke = AsyncMock(return_value='synthetic-result')
|
||||
setattr(loader, invoke_method, synthetic_invoke)
|
||||
query = SimpleNamespace(variables={})
|
||||
|
||||
assert await loader.has_tool(reserved_name, source_id=None) is True
|
||||
tool = await loader.get_tool(reserved_name, source_id=None)
|
||||
result = await loader.invoke_tool(reserved_name, {}, query, source_id=None)
|
||||
|
||||
assert tool is not None
|
||||
assert tool is not real_tool
|
||||
assert result == 'synthetic-result'
|
||||
synthetic_invoke.assert_awaited_once_with({}, query)
|
||||
session.invoke_mcp_tool.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_two_plugins_with_same_tool_forward_only_authorized_plugin():
|
||||
manifest = SimpleNamespace(
|
||||
owner='authorized/plugin',
|
||||
metadata=SimpleNamespace(
|
||||
name='shared_tool',
|
||||
description=SimpleNamespace(en_US='Shared tool'),
|
||||
),
|
||||
spec={
|
||||
'llm_prompt': 'Shared tool',
|
||||
'parameters': {'type': 'object', 'properties': {}},
|
||||
},
|
||||
)
|
||||
connector = SimpleNamespace(
|
||||
list_tools=AsyncMock(return_value=[manifest]),
|
||||
call_tool=AsyncMock(return_value='authorized-result'),
|
||||
)
|
||||
loader = PluginToolLoader(SimpleNamespace(plugin_connector=connector, logger=Mock()))
|
||||
query = SimpleNamespace(session=SimpleNamespace(), query_id=7)
|
||||
query.session.model_dump = Mock(return_value={})
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'shared_tool',
|
||||
{},
|
||||
query,
|
||||
source_id='authorized/plugin',
|
||||
)
|
||||
|
||||
assert result == 'authorized-result'
|
||||
connector.call_tool.assert_awaited_once_with(
|
||||
'shared_tool',
|
||||
{},
|
||||
session=query.session,
|
||||
query_id=7,
|
||||
bound_plugins=['authorized/plugin'],
|
||||
)
|
||||
Reference in New Issue
Block a user