fix(cloud): tolerate slow plugin runtime reconciliation (#2405)

Co-authored-by: Chan <dadachann@users.noreply.github.com>

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': {

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]}')

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'

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()

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'

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'
+    )