diff --git a/src/langbot/pkg/core/stages/load_config.py b/src/langbot/pkg/core/stages/load_config.py index 49c439758..e34c740bd 100644 --- a/src/langbot/pkg/core/stages/load_config.py +++ b/src/langbot/pkg/core/stages/load_config.py @@ -41,6 +41,7 @@ _RUNTIME_POLICY_DEFAULTS = { } }, 'plugin': { + 'connect_timeout_seconds': 180.0, 'worker': { 'max_cpus': 1.0, 'max_memory_mb': 512, @@ -56,7 +57,7 @@ _RUNTIME_POLICY_DEFAULTS = { 'restart_failure_window_seconds': 30.0, 'restart_circuit_open_seconds': 60.0, 'require_hard_limits': False, - } + }, }, 'mcp': {'stdio': {'enabled': True}}, 'monitoring': { diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 9d4c1c75a..cd07565aa 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -6,6 +6,7 @@ import contextlib import contextvars import hashlib import json +import math import time import uuid from typing import Any @@ -76,7 +77,7 @@ _GITHUB_ASSET_HOSTS = frozenset( } ) _HTTP_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308}) -_CONNECT_TIMEOUT_SEC = 30.0 +_DEFAULT_CONNECT_TIMEOUT_SECONDS = 180.0 _HEARTBEAT_INTERVAL_SEC = 20.0 _HEARTBEAT_FAILURE_THRESHOLD = 3 _RECONNECT_MAX_DELAY_SEC = 60.0 @@ -206,6 +207,17 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): return f'{constants.instance_id}:plugin-runtime' + @staticmethod + def _runtime_connect_timeout(plugin_config: dict[str, Any]) -> float: + value = plugin_config.get('connect_timeout_seconds', _DEFAULT_CONNECT_TIMEOUT_SECONDS) + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value <= 0: + raise ValueError('plugin.connect_timeout_seconds must be a positive number') + return float(value) + + @staticmethod + def _runtime_connect_timeout_error(timeout_seconds: float) -> str: + return f'Plugin runtime did not become ready within {timeout_seconds:g} seconds' + def _runtime_handler(self) -> handler.RuntimeConnectionHandler: runtime_handler = getattr(self, 'handler', None) if runtime_handler is None: @@ -701,10 +713,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): """ runtime_handler = self._runtime_handler() + started_at = time.monotonic() async with self._state_lock: all_states: dict[str, PluginInstallationDesiredState] = {} workspace_installations: dict[str, set[str]] = {} + workspace_count = 0 for context in contexts: + workspace_count += 1 execution_context = await self._validate_execution_context(context) states = await self._load_workspace_desired_states(execution_context) installation_ids = {state.binding.installation_uuid for state in states} @@ -724,6 +739,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): runtime_handler.unregister_installation_binding(previous.binding) self._known_desired_states = all_states self._workspace_installations = workspace_installations + self.ap.logger.info( + 'Shared plugin runtime reconcile completed: workspaces=%d desired_installations=%d ' + 'elapsed_seconds=%.3f', + workspace_count, + len(all_states), + time.monotonic() - started_at, + ) return result async def _validate_execution_context(self, context: TenantContext) -> ExecutionContext: @@ -959,11 +981,14 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): task_coro = self.ctrl.run(new_connection_callback) self._transport_task = asyncio.create_task(task_coro) + connect_timeout_seconds = self._runtime_connect_timeout(self.ap.instance_config.data.get('plugin', {})) try: - await asyncio.wait_for(self._connected.wait(), timeout=_CONNECT_TIMEOUT_SEC) + await asyncio.wait_for(self._connected.wait(), timeout=connect_timeout_seconds) except asyncio.TimeoutError as exc: await self._stop_transport() - raise PluginRuntimeNotConnectedError('Plugin runtime did not become ready within 30 seconds') from exc + raise PluginRuntimeNotConnectedError( + self._runtime_connect_timeout_error(connect_timeout_seconds) + ) from exc if connect_errors: await self._stop_transport() raise PluginRuntimeNotConnectedError(f'Plugin runtime connection failed: {connect_errors[-1]}') diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index 1ae2ba515..a9d25086f 100644 --- a/src/langbot/templates/config.yaml +++ b/src/langbot/templates/config.yaml @@ -245,6 +245,8 @@ storage: max_concurrency: 16 plugin: enable: true + # Maximum time for the Runtime transport, handshake, and desired-state replay. + connect_timeout_seconds: 180.0 runtime_ws_url: 'ws://langbot_plugin_runtime:5400/control/ws' enable_marketplace: true display_plugin_debug_url: 'ws://localhost:5401/plugin/debug/ws' diff --git a/tests/unit_tests/core/test_load_config.py b/tests/unit_tests/core/test_load_config.py index defbcee9d..03dd7cc5d 100644 --- a/tests/unit_tests/core/test_load_config.py +++ b/tests/unit_tests/core/test_load_config.py @@ -319,6 +319,7 @@ class TestApplyEnvOverridesToConfig: load_config = get_load_config_module() cfg = { 'plugin': { + 'connect_timeout_seconds': 30.0, 'worker': { 'max_cpus': 1.0, 'max_memory_mb': 512, @@ -329,11 +330,12 @@ class TestApplyEnvOverridesToConfig: 'restart_failure_threshold': 8, 'restart_failure_window_seconds': 30.0, 'restart_circuit_open_seconds': 60.0, - } + }, }, 'mcp': {'stdio': {'enabled': True}}, } env = { + 'PLUGIN__CONNECT_TIMEOUT_SECONDS': '180', 'PLUGIN__WORKER__MAX_CPUS': '2.5', 'PLUGIN__WORKER__MAX_MEMORY_MB': '1024', 'PLUGIN__WORKER__MAX_PIDS': '64', @@ -349,6 +351,7 @@ class TestApplyEnvOverridesToConfig: with patch.dict(os.environ, env, clear=True): result = load_config._apply_env_overrides_to_config(cfg) + assert result['plugin']['connect_timeout_seconds'] == 180.0 assert result['plugin']['worker'] == { 'max_cpus': 2.5, 'max_memory_mb': 1024, @@ -393,6 +396,14 @@ class TestApplyEnvOverridesToConfig: assert isinstance(result['plugin']['worker']['max_memory_mb'], int) assert result['mcp']['stdio']['enabled'] is False + def test_runtime_policy_defaults_add_typed_plugin_connect_timeout(self): + load_config = get_load_config_module() + + completed = load_config._complete_runtime_policy_defaults({'plugin': {'enable': True}}) + + assert completed['plugin']['connect_timeout_seconds'] == 180.0 + assert isinstance(completed['plugin']['connect_timeout_seconds'], float) + def test_webhook_prefix_override(self): """Test overriding webhook_prefix via environment variable.""" load_config = get_load_config_module() diff --git a/tests/unit_tests/plugin/test_connector_reconcile.py b/tests/unit_tests/plugin/test_connector_reconcile.py index b22e8e13e..8759a537e 100644 --- a/tests/unit_tests/plugin/test_connector_reconcile.py +++ b/tests/unit_tests/plugin/test_connector_reconcile.py @@ -153,6 +153,31 @@ async def test_empty_projected_workspaces_do_not_retain_installation_sets(): connector.handler.reconcile_plugin_installations.assert_awaited_once_with(()) +@pytest.mark.asyncio +async def test_shared_reconcile_logs_workspace_installation_counts_and_elapsed_time(): + binding_a = execution_binding('workspace-a') + binding_b = execution_binding('workspace-b') + setting_a = plugin_setting('01', 'a' * 64) + setting_b = plugin_setting('02', 'b' * 64) + connector = shared_connector( + [[binding_a, binding_b]], + {'workspace-a': [setting_a], 'workspace-b': [setting_b]}, + ) + connector.handler = runtime_handler() + await connector._prepare_connected_runtime() + + matching_calls = [ + call + for call in connector.ap.logger.info.call_args_list + if call.args + and call.args[0] + == 'Shared plugin runtime reconcile completed: workspaces=%d desired_installations=%d elapsed_seconds=%.3f' + ] + assert len(matching_calls) == 1 + assert matching_calls[0].args[1:3] == (2, 2) + assert matching_calls[0].args[3] >= 0 + + @pytest.mark.asyncio async def test_fresh_shared_runtime_cache_replays_persisted_local_package(): package = b'local-lbpkg-bytes' diff --git a/tests/unit_tests/plugin/test_connector_static.py b/tests/unit_tests/plugin/test_connector_static.py index 2db33b708..eba46df6b 100644 --- a/tests/unit_tests/plugin/test_connector_static.py +++ b/tests/unit_tests/plugin/test_connector_static.py @@ -6,9 +6,10 @@ Tests cover: from __future__ import annotations -import pytest from importlib import import_module +import pytest + def get_connector_module(): """Lazy import to avoid circular import issues.""" @@ -60,3 +61,28 @@ def test_runtime_id_is_stable_across_core_restarts(monkeypatch): monkeypatch.setattr(connector.constants, 'instance_id', 'instance-a') assert connector.PluginRuntimeConnector._build_runtime_id() == 'instance-a:plugin-runtime' + + +def test_runtime_connect_timeout_defaults_to_three_minutes(): + connector = get_connector_module() + assert connector.PluginRuntimeConnector._runtime_connect_timeout({}) == 180.0 + + +def test_runtime_connect_timeout_reads_typed_plugin_config(): + connector = get_connector_module() + assert connector.PluginRuntimeConnector._runtime_connect_timeout({'connect_timeout_seconds': 45.5}) == 45.5 + + +@pytest.mark.parametrize('value', [True, False, None, 0, -1, float('nan'), float('inf'), '180', object()]) +def test_runtime_connect_timeout_rejects_invalid_values(value): + connector = get_connector_module() + with pytest.raises(ValueError, match='plugin.connect_timeout_seconds'): + connector.PluginRuntimeConnector._runtime_connect_timeout({'connect_timeout_seconds': value}) + + +def test_runtime_connect_timeout_error_displays_actual_seconds(): + connector = get_connector_module() + + assert connector.PluginRuntimeConnector._runtime_connect_timeout_error(45.5) == ( + 'Plugin runtime did not become ready within 45.5 seconds' + )