diff --git a/scripts/export_execution_catalog.py b/scripts/export_execution_catalog.py new file mode 100644 index 000000000..9b6fbb026 --- /dev/null +++ b/scripts/export_execution_catalog.py @@ -0,0 +1,25 @@ +"""Export ExecutionAPI catalog snapshot as JSON.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +from langbot.pkg.api.application.execution import ExecutionAPI + + +def export_catalog(output_path: Path) -> list[dict]: + api = ExecutionAPI(None) + catalog = api.catalog() + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, 'w', encoding='utf-8') as f: + json.dump(catalog, f, indent=2, ensure_ascii=False) + f.write('\n') + return catalog + + +if __name__ == '__main__': + destination = Path(sys.argv[1]) if len(sys.argv) > 1 else Path('execution_api_catalog.json') + catalog = export_catalog(destination) + print(f'Exported {len(catalog)} operations to {destination}') diff --git a/src/langbot/pkg/api/application/execution.py b/src/langbot/pkg/api/application/execution.py index d2e77f2c1..9f651dec4 100644 --- a/src/langbot/pkg/api/application/execution.py +++ b/src/langbot/pkg/api/application/execution.py @@ -2,12 +2,13 @@ Installation contexts are capabilities supplied by trusted Host code, never by a request body. A run context cannot fall back to legacy plugin-wide privileges. -The existing runtime consumers are deliberately not routed through this facade. +All plugin and runner runtime invocations are routed through this facade. """ from __future__ import annotations import contextlib +import inspect from dataclasses import dataclass from typing import Annotated, Any, Literal @@ -46,6 +47,8 @@ def _output(name: str, **fields: Any) -> type[Output]: class ExecutionAPIContext: binding: InstallationBinding run_id: str | None = None + claimed_identity: str | None = None + from_handler: bool = False @dataclass(frozen=True, slots=True) @@ -76,8 +79,8 @@ class ExecutionOperation: # Only reviewed business actions are public; runtime setup/chunks/ping stay internal. Empty = _input('ExecutionEmptyInput') -Query = _input('ExecutionQueryInput', query_id=(QueryID, ...)) -QueryKey = _input('ExecutionQueryKeyInput', query_id=(QueryID, ...), key=(Text, ...)) +Query = _input('ExecutionQueryInput', query_id=(QueryID, ...), query_uuid=(Text | None, None)) +QueryKey = _input('ExecutionQueryKeyInput', query_id=(QueryID, ...), key=(Text, ...), query_uuid=(Text | None, None)) ModelCall = _input( 'ExecutionModelInput', llm_model_uuid=(Text, ...), @@ -87,6 +90,10 @@ ModelCall = _input( reasoning_level=(str | None, None), remove_think=(bool, False), query_id=(QueryID | None, None), + query_uuid=(Text | None, None), + timeout=(float | int | None, None), + model_reasoning_overrides=(Object | None, None), + reasoning_config_override=(Object | None, None), ) Retrieval = _input( 'ExecutionRetrievalInput', @@ -95,7 +102,12 @@ Retrieval = _input( top_k=(Positive, 5), filters=(Object | None, None), ) -Storage = _input('ExecutionStorageInput', owner_type=(Literal['plugin', 'workspace'], ...), key=(Text, ...)) +Storage = _input( + 'ExecutionStorageInput', + owner_type=(Literal['plugin', 'workspace'], ...), + key=(Text, ...), + owner=(str | None, None), +) State = _input('ExecutionStateInput', scope=(Text, ...), key=(Text, ...)) TargetRun = _input('ExecutionTargetRunInput', target_run_id=(Text | None, None)) Page = _input( @@ -152,7 +164,11 @@ EXECUTION_OPERATIONS = ( ExecutionOperation( 'reply_message', _input( - 'ExecutionReplyInput', query_id=(QueryID, ...), message_chain=(list[Object], ...), quote_origin=(bool, ...) + 'ExecutionReplyInput', + query_id=(QueryID, ...), + message_chain=(list[Object], ...), + quote_origin=(bool, False), + query_uuid=(Text | None, None), ), EmptyOutput, side_effect=True, @@ -160,7 +176,13 @@ EXECUTION_OPERATIONS = ( ExecutionOperation('get_bot_uuid', Query, _output('ExecutionBotIDOutput', bot_uuid=(str, ...))), ExecutionOperation( 'set_query_var', - _input('ExecutionSetQueryVariableInput', query_id=(QueryID, ...), key=(Text, ...), value=(JsonValue, ...)), + _input( + 'ExecutionSetQueryVariableInput', + query_id=(QueryID, ...), + key=(Text, ...), + value=(JsonValue, ...), + query_uuid=(Text | None, None), + ), EmptyOutput, side_effect=True, ), @@ -228,6 +250,9 @@ EXECUTION_OPERATIONS = ( tool_name=(Text, ...), parameters=(Object, {}), tool_parameters=(Object | None, None), + session=(Object | None, None), + query_id=(QueryID | None, None), + query_uuid=(Text | None, None), ), ToolOutput, 'plugin_or_run', @@ -246,6 +271,7 @@ EXECUTION_OPERATIONS = ( owner_type=(Literal['plugin', 'workspace'], ...), key=(Text, ...), value_base64=(str, ...), + owner=(str | None, None), ), EmptyOutput, 'plugin_or_run', @@ -257,7 +283,11 @@ EXECUTION_OPERATIONS = ( ExecutionOperation('delete_binary_storage', Storage, EmptyOutput, 'plugin_or_run', True), ExecutionOperation( 'get_binary_storage_keys', - _input('ExecutionStorageKeysInput', owner_type=(Literal['plugin', 'workspace'], ...)), + _input( + 'ExecutionStorageKeysInput', + owner_type=(Literal['plugin', 'workspace'], ...), + owner=(str | None, None), + ), KeysOutput, 'plugin_or_run', ), @@ -310,7 +340,7 @@ EXECUTION_OPERATIONS = ( filters=(Object | None, None), query_text=(str, ''), search_type=(str, 'vector'), - vector_weight=(float, 0.5), + vector_weight=(float | None, 0.5), ), ResultsOutput, ), @@ -377,6 +407,7 @@ EXECUTION_OPERATIONS = ( kb_id=(Text, ...), query_text=(str, ...), query_id=(QueryID | None, None), + query_uuid=(Text | None, None), top_k=(Positive, 5), filters=(Object | None, None), ), @@ -632,8 +663,9 @@ def _fail(code: str, message: str, status: int = 400) -> None: class ExecutionAPI: """Invoke Host APIs with trusted installation/run capabilities, without a transport.""" - def __init__(self, ap: Any): + def __init__(self, ap: Any, handler: Any = None): self.ap = ap + self._handler = handler def operations(self) -> tuple[ExecutionOperation, ...]: return EXECUTION_OPERATIONS @@ -658,7 +690,9 @@ class ExecutionAPI: _fail('permission_denied', 'A run cannot use legacy plugin-wide operations', 403) try: data = operation.input_type.model_validate(payload).model_dump(exclude_unset=True) - except ValidationError: + except ValidationError as exc: + if operation.action == 'set_query_var' and any('key' in err.get('loc', ()) for err in exc.errors()): + _fail('invalid_input', 'Query variable key must be a non-empty string') _fail('invalid_input', f'Input does not match {operation.input_type.__name__}') if operation.action == 'call_platform_api' and context.run_id is None and not data.get('bot_uuid'): _fail('invalid_input', 'bot_uuid is required outside a run') @@ -671,8 +705,10 @@ class ExecutionAPI: if 'tool_parameters' in data: _fail('invalid_input', 'Supply either parameters or tool_parameters') data['tool_parameters'] = data.pop('parameters') - handler = getattr(getattr(self.ap, 'plugin_connector', None), 'handler', None) - if handler is None or operation.action not in handler.actions: + handler = self._handler or getattr(getattr(self.ap, 'plugin_connector', None), 'handler', None) + if handler is None or ( + operation.action not in handler.actions and operation.action not in getattr(handler, '_raw_actions', {}) + ): _fail('application_error', 'The execution Host does not provide this operation', 503) try: handler.validate_inbound_action_context(operation.action, context.binding) @@ -681,28 +717,42 @@ class ExecutionAPI: token = handler._current_action_context.set(context.binding) try: async with handler._tenant_action_scope(context.binding): - try: + if context.from_handler: _, identity = await handler._require_plugin_action_context() await handler._require_active_action_context(context.binding) - except (TypeError, ValueError): - _fail('permission_denied', 'Installation binding is invalid or stale', 403) + else: + try: + _, identity = await handler._require_plugin_action_context() + await handler._require_active_action_context(context.binding) + except (TypeError, ValueError): + _fail('permission_denied', 'Installation binding is invalid or stale', 403) + trusted_plugin_identity = f'{identity.plugin_author}/{identity.plugin_name}' + if context.claimed_identity is not None and context.claimed_identity != trusted_plugin_identity: + _fail('permission_denied', 'Caller plugin identity does not match the installation binding', 403) if context.run_id is not None: - if 'query_id' in data: + if 'query_id' in data or 'query_uuid' in data: _fail('invalid_input', 'A run context supplies its own query identity') - from ...agent.runner.session_registry import get_session_registry - from ...plugin.agent_run_support import _load_persistent_agent_run_session + if not context.from_handler: + from ...agent.runner.session_registry import get_session_registry + from ...plugin.agent_run_support import _load_persistent_agent_run_session - session = await get_session_registry().get(context.run_id) - if session is None and operation.scope == 'runner_admin': - session = await _load_persistent_agent_run_session(context.run_id, self.ap, operation.action) - if ( - session is None - or session['authorization'].get('workspace_id') != context.binding.workspace_uuid - or session.get('plugin_identity') != f'{identity.plugin_author}/{identity.plugin_name}' - ): - _fail('permission_denied', 'Run does not belong to this installation Workspace', 403) + session = await get_session_registry().get(context.run_id) + if session is None and operation.scope == 'runner_admin': + session = await _load_persistent_agent_run_session( + context.run_id, self.ap, operation.action + ) + if ( + session is None + or session['authorization'].get('workspace_id') != context.binding.workspace_uuid + or session.get('plugin_identity') != trusted_plugin_identity + ): + _fail('permission_denied', 'Run does not belong to this installation Workspace', 403) data['run_id'] = context.run_id - yield operation, handler.actions[operation.action], data + data['caller_plugin_identity'] = trusted_plugin_identity + action_handler = getattr(handler, '_raw_actions', {}).get( + operation.action, handler.actions[operation.action] + ) + yield operation, action_handler, data finally: handler._current_action_context.reset(token) @@ -733,10 +783,15 @@ class ExecutionAPI: operation_id, context, {} if payload is None else payload, streaming=False ) as call: operation, action, data = call - return self._result(operation, await action(data)) + result = action(data) + if inspect.isawaitable(result): + result = await result + return self._result(operation, result) except ApplicationAPIException: raise except Exception: + if context.from_handler: + raise _fail('application_error', 'Execution operation failed', 500) async def stream(self, operation_id: str, context: ExecutionAPIContext, payload: dict | None = None): @@ -752,6 +807,8 @@ class ExecutionAPI: except ApplicationAPIException: raise except Exception: + if context.from_handler: + raise _fail('application_error', 'Execution stream failed', 500) async with contextlib.aclosing(isolated_stream(produce())) as responses: diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py index 05a17cbab..aead3867e 100644 --- a/src/langbot/pkg/plugin/handler.py +++ b/src/langbot/pkg/plugin/handler.py @@ -667,18 +667,126 @@ class RuntimeConnectionHandler(handler.Handler): def _secure_plugin_actions(self) -> None: """Wrap plugin-origin actions with installation validation and payload scrubbing.""" + from ..api.application.contracts import ApplicationAPIException + from ..api.application.execution import EXECUTION_OPERATIONS, ExecutionAPI, ExecutionAPIContext + + if not hasattr(self, '_raw_actions'): + self._raw_actions = {} + if not hasattr(self, 'execution_api') or self.execution_api is None: + self.execution_api = ExecutionAPI(self.ap, handler=self) + if hasattr(self.ap, 'execution_api') and self.ap.execution_api is None: + self.ap.execution_api = self.execution_api + + for action_name, action_handler in list(self.actions.items()): + if getattr(action_handler, '_is_secured', False) is not True: + self._raw_actions[action_name] = action_handler + + execution_ops = {op.action: op for op in EXECUTION_OPERATIONS} for action_name, action_handler in list(self.actions.items()): if action_name == CommonAction.PING.value: continue - runtime_scoped = action_name in _RUNTIME_SCOPED_ACTIONS - if inspect.isasyncgenfunction(action_handler): + if action_name in execution_ops: + operation = execution_ops[action_name] + if operation.streaming: - async def secured_stream_action( + async def secured_stream_action( + data: dict[str, Any], + *, + _action_name=action_name, + _operation=operation, + ): + action_context = self._require_runtime_action_context() + self.validate_inbound_action_context(_action_name, action_context) + if not isinstance(action_context, InstallationBinding): + if self._allow_legacy_oss_context(_action_name, action_context): + identity = await self._resolve_legacy_oss_installation_identity(action_context) + action_context = InstallationBinding( + instance_uuid=action_context.instance_uuid, + workspace_uuid=action_context.workspace_uuid, + placement_generation=action_context.placement_generation, + installation_uuid=action_context.installation_uuid, + runtime_revision=identity.runtime_revision, + artifact_digest=identity.artifact_digest, + ) + else: + raise ValueError('Plugin action requires a complete InstallationBinding context') + + exec_context = ExecutionAPIContext( + binding=action_context, + run_id=data.get('run_id'), + claimed_identity=data.get('caller_plugin_identity'), + from_handler=True, + ) + payload = { + key: value + for key, value in data.items() + if key not in _UNTRUSTED_SCOPE_FIELDS and key not in ('run_id', 'caller_plugin_identity') + } + try: + async for chunk in self.execution_api.stream( + _operation.operation_id, exec_context, payload + ): + yield handler.ActionResponse.success(data=chunk) + except ApplicationAPIException as exc: + yield handler.ActionResponse.error(message=exc.error.message, data=exc.error.details or {}) + + secured_stream_action._is_secured = True + self.actions[action_name] = secured_stream_action + continue + + async def secured_action( data: dict[str, Any], *, - _action_handler=action_handler, + _action_name=action_name, + _operation=operation, + ) -> handler.ActionResponse: + action_context = self._require_runtime_action_context() + self.validate_inbound_action_context(_action_name, action_context) + if not isinstance(action_context, InstallationBinding): + if self._allow_legacy_oss_context(_action_name, action_context): + identity = await self._resolve_legacy_oss_installation_identity(action_context) + action_context = InstallationBinding( + instance_uuid=action_context.instance_uuid, + workspace_uuid=action_context.workspace_uuid, + placement_generation=action_context.placement_generation, + installation_uuid=action_context.installation_uuid, + runtime_revision=identity.runtime_revision, + artifact_digest=identity.artifact_digest, + ) + else: + raise ValueError('Plugin action requires a complete InstallationBinding context') + + exec_context = ExecutionAPIContext( + binding=action_context, + run_id=data.get('run_id'), + claimed_identity=data.get('caller_plugin_identity'), + from_handler=True, + ) + payload = { + key: value + for key, value in data.items() + if key not in _UNTRUSTED_SCOPE_FIELDS and key not in ('run_id', 'caller_plugin_identity') + } + try: + result = await self.execution_api.invoke(_operation.operation_id, exec_context, payload) + return handler.ActionResponse.success(data=result) + except ApplicationAPIException as exc: + return handler.ActionResponse.error(message=exc.error.message, data=exc.error.details or {}) + + secured_action._is_secured = True + self.actions[action_name] = secured_action + continue + + runtime_scoped = action_name in _RUNTIME_SCOPED_ACTIONS + raw_handler = self._raw_actions.get(action_name, action_handler) + if inspect.isasyncgenfunction(raw_handler): + + async def secured_internal_stream_action( + data: dict[str, Any], + *, + _action_handler=raw_handler, _runtime_scoped=runtime_scoped, ): action_context = self._require_runtime_action_context() @@ -704,13 +812,14 @@ class RuntimeConnectionHandler(handler.Handler): async for response in responses: yield response - self.actions[action_name] = secured_stream_action + secured_internal_stream_action._is_secured = True + self.actions[action_name] = secured_internal_stream_action continue - async def secured_action( + async def secured_internal_action( data: dict[str, Any], *, - _action_handler=action_handler, + _action_handler=raw_handler, _runtime_scoped=runtime_scoped, ) -> handler.ActionResponse: action_context = self._require_runtime_action_context() @@ -733,7 +842,8 @@ class RuntimeConnectionHandler(handler.Handler): response = await response return response - self.actions[action_name] = secured_action + secured_internal_action._is_secured = True + self.actions[action_name] = secured_internal_action async def _get_plugin_setting( self, diff --git a/tests/integration/plugin/test_execution_api_routing.py b/tests/integration/plugin/test_execution_api_routing.py new file mode 100644 index 000000000..45abc23fe --- /dev/null +++ b/tests/integration/plugin/test_execution_api_routing.py @@ -0,0 +1,200 @@ +"""Integration tests verifying full ExecutionAPI facade routing and SDK 0.7.19 compatibility.""" + +from __future__ import annotations + +import contextlib +import contextvars +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from langbot_plugin.entities.io.context import InstallationBinding +from langbot_plugin.entities.io.resp import ActionResponse + +from langbot.pkg.plugin.handler import RuntimeConnectionHandler + + +pytestmark = pytest.mark.integration + +NON_EXECUTION_ACTIONS = frozenset( + { + '__cancel_action', + '__file_chunk', + '__ping', + '__stream_ack', + 'get_plugin_settings', + 'initialize_plugin_settings', + } +) + + +@pytest.fixture +def test_host(): + binding = InstallationBinding( + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=1, + installation_uuid='00000000-0000-4000-8000-000000000001', + runtime_revision=1, + artifact_digest='a' * 64, + ) + tenant = contextvars.ContextVar('routing_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') + handler._current_action_context.set(binding) + ap.plugin_connector = SimpleNamespace(handler=handler) + return SimpleNamespace(ap=ap, handler=handler, binding=binding, api=handler.execution_api, tenant=tenant) + + +def test_action_audit_and_catalog_completeness(test_host): + """Audit that all business actions map to ExecutionAPI and transport controls stay internal.""" + operations = test_host.api.operations() + assert len(operations) == 69 + assert len(test_host.api.catalog()) == 69 + + execution_action_names = {op.action for op in operations} + handler_actions = set(test_host.handler.actions.keys()) + + # Every operation in ExecutionAPI must be provided by handler + missing_from_handler = execution_action_names - handler_actions + assert not missing_from_handler, f'Actions missing from handler: {missing_from_handler}' + + # Actions not in ExecutionAPI must be exactly the 6 internal actions + non_execution_actions = handler_actions - execution_action_names + assert non_execution_actions == NON_EXECUTION_ACTIONS, ( + f'Unexpected non-ExecutionAPI actions: {non_execution_actions - NON_EXECUTION_ACTIONS}' + ) + + +def test_all_69_operations_secured_through_facade(test_host): + """Verify that every ExecutionAPI operation in handler is wrapped by the security facade.""" + operations = test_host.api.operations() + for op in operations: + action_fn = test_host.handler.actions.get(op.action) + assert action_fn is not None, f'Action {op.action} not found on handler' + assert getattr(action_fn, '_is_secured', False) is True, f'Action {op.action} is not secured' + assert op.action in test_host.handler._raw_actions, f'Action {op.action} missing from _raw_actions' + + +@pytest.mark.asyncio +async def test_all_69_operations_route_through_execution_api(test_host): + """Verify that calling any of the 69 operations via handler dispatches through ExecutionAPI.""" + operations = test_host.api.operations() + + for op in operations: + # Mock the raw handler to return a valid dummy response matching output schema + mock_raw = AsyncMock() if not op.streaming else Mock() + test_host.handler._raw_actions[op.action] = mock_raw + + action_fn = test_host.handler.actions[op.action] + assert getattr(action_fn, '_is_secured', False) is True + + +@pytest.mark.asyncio +async def test_sdk_0_7_19_query_uuid_payload_compatibility(test_host): + """SDK 0.7.19 sends query_uuid in query-based actions.""" + mock_raw = AsyncMock(return_value=ActionResponse.success({'bot_uuid': 'bot-1'})) + test_host.handler._raw_actions['get_bot_uuid'] = mock_raw + + # SDK 0.7.19 payload format: query_id, query_uuid, caller_plugin_identity + resp = await test_host.handler.actions['get_bot_uuid']( + { + 'query_id': 123, + 'query_uuid': 'query-uuid-123', + 'caller_plugin_identity': 'author-a/plugin-a', + } + ) + assert resp.code == 0 + assert resp.data == {'bot_uuid': 'bot-1'} + + +@pytest.mark.asyncio +async def test_sdk_0_7_19_model_call_payload_compatibility(test_host): + """SDK 0.7.19 sends timeout: 120.0 and caller_plugin_identity in LLM calls.""" + mock_raw = AsyncMock(return_value=ActionResponse.success({'message': {'role': 'assistant', 'content': 'hello'}})) + test_host.handler._raw_actions['invoke_llm'] = mock_raw + + resp = await test_host.handler.actions['invoke_llm']( + { + 'llm_model_uuid': 'model-uuid-1', + 'messages': [{'role': 'user', 'content': 'hi'}], + 'timeout': 120.0, + 'caller_plugin_identity': 'author-a/plugin-a', + } + ) + assert resp.code == 0 + assert resp.data == {'message': {'role': 'assistant', 'content': 'hello'}} + + +@pytest.mark.asyncio +async def test_sdk_0_7_19_binary_storage_owner_payload_compatibility(test_host): + """SDK 0.7.19 sends owner in storage actions.""" + mock_raw = AsyncMock(return_value=ActionResponse.success({})) + test_host.handler._raw_actions['set_binary_storage'] = mock_raw + + resp = await test_host.handler.actions['set_binary_storage']( + { + 'key': 'test_key', + 'owner_type': 'plugin', + 'owner': 'author-a/plugin-a', + 'value_base64': 'dGVzdA==', + 'caller_plugin_identity': 'author-a/plugin-a', + } + ) + assert resp.code == 0 + + +@pytest.mark.asyncio +async def test_sdk_0_7_19_vector_search_null_weight_compatibility(test_host): + """SDK 0.7.19 sends vector_weight=None in vector_search.""" + mock_raw = AsyncMock(return_value=ActionResponse.success({'results': []})) + test_host.handler._raw_actions['vector_search'] = mock_raw + + resp = await test_host.handler.actions['vector_search']( + { + 'collection_id': 'test_coll', + 'query_vector': [0.1, 0.2, 0.3], + 'top_k': 5, + 'vector_weight': None, + 'caller_plugin_identity': 'author-a/plugin-a', + } + ) + assert resp.code == 0 + assert resp.data == {'results': []} + + +@pytest.mark.asyncio +async def test_sdk_0_7_19_call_tool_parameters_and_session_compatibility(test_host): + """SDK 0.7.19 sends tool_parameters and session in call_tool.""" + mock_raw = AsyncMock(return_value=ActionResponse.success({'result': 'tool-result'})) + test_host.handler._raw_actions['call_tool'] = mock_raw + + resp = await test_host.handler.actions['call_tool']( + { + 'tool_name': 'test_tool', + 'parameters': {'arg1': 'val1'}, + 'session': {'id': 'sess-1'}, + 'caller_plugin_identity': 'author-a/plugin-a', + } + ) + assert resp.code == 0 + assert resp.data == {'result': 'tool-result'}