From c4bad508d278bc76590cfb6218a7d192d8593151 Mon Sep 17 00:00:00 2001
From: dadachann <185672915+dadachann@users.noreply.github.com>
Date: Wed, 19 Aug 2026 15:44:03 +0000
Subject: [PATCH] fix(runtime): honor configured cold reconcile timeout

---
 src/langbot/pkg/plugin/connector.py               | 12 ++++++++++--
 src/langbot/pkg/plugin/handler.py                 |  4 +++-
 .../unit_tests/plugin/test_connector_reconcile.py | 15 ++++++++++++++-
 tests/unit_tests/plugin/test_handler.py           | 12 ++++++++++++
 4 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py
index 0279ab285..bcdf4fe8e 100644
--- a/src/langbot/pkg/plugin/connector.py
+++ b/src/langbot/pkg/plugin/connector.py
@@ -701,7 +701,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
             }
             self._known_desired_states.update({state.binding.installation_uuid: state for state in desired_states})
 
-        result = await runtime_handler.reconcile_plugin_installations(tuple(self._known_desired_states.values()))
+        reconcile_timeout_seconds = max(300.0, self._runtime_connect_timeout(self.ap.instance_config.data.get("plugin", {})))
+        result = await runtime_handler.reconcile_plugin_installations(
+            tuple(self._known_desired_states.values()),
+            timeout=reconcile_timeout_seconds,
+        )
         await self._repair_reconcile_missing_artifacts(self._known_desired_states, result)
         self._record_reconcile_failures(self._known_desired_states, result)
 
@@ -736,7 +740,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
                     if state.binding.installation_uuid in all_states:
                         raise ValueError('Duplicate plugin installation UUID across projected Workspaces')
                     all_states[state.binding.installation_uuid] = state
-            result = await runtime_handler.reconcile_plugin_installations(tuple(all_states.values()))
+            reconcile_timeout_seconds = max(300.0, self._runtime_connect_timeout(self.ap.instance_config.data.get("plugin", {})))
+            result = await runtime_handler.reconcile_plugin_installations(
+                tuple(all_states.values()),
+                timeout=reconcile_timeout_seconds,
+            )
             await self._repair_reconcile_missing_artifacts(all_states, result)
             self._record_reconcile_failures(all_states, result)
             for installation_uuid, previous in tuple(self._known_desired_states.items()):
diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py
index 8dae74f6f..8463750f0 100644
--- a/src/langbot/pkg/plugin/handler.py
+++ b/src/langbot/pkg/plugin/handler.py
@@ -1677,13 +1677,15 @@ class RuntimeConnectionHandler(handler.Handler):
     async def reconcile_plugin_installations(
         self,
         installations: tuple[PluginInstallationDesiredState, ...],
+        *,
+        timeout: float = 300,
     ) -> dict[str, Any]:
         request = ReconcilePluginInstallationsRequest(installations=installations)
         with self.installation_scope(None):
             return await self.call_action(
                 LangBotToRuntimeAction.RECONCILE_PLUGIN_INSTALLATIONS,
                 request.model_dump(),
-                timeout=300,
+                timeout=timeout,
             )
 
     async def apply_plugin_installation(
diff --git a/tests/unit_tests/plugin/test_connector_reconcile.py b/tests/unit_tests/plugin/test_connector_reconcile.py
index 8759a537e..d4d9b6b5b 100644
--- a/tests/unit_tests/plugin/test_connector_reconcile.py
+++ b/tests/unit_tests/plugin/test_connector_reconcile.py
@@ -107,6 +107,19 @@ def shared_connector(
     return connector
 
 
+@pytest.mark.asyncio
+async def test_shared_reconcile_uses_configured_cold_start_timeout():
+    binding = execution_binding("workspace-a")
+    setting = plugin_setting("01", "a" * 64)
+    connector = shared_connector([[binding]], {"workspace-a": [setting]})
+    connector.ap.instance_config.data["plugin"]["connect_timeout_seconds"] = 900
+    connector.handler = runtime_handler()
+
+    await connector._prepare_connected_runtime()
+
+    assert connector.handler.reconcile_plugin_installations.await_args.kwargs["timeout"] == 900
+
+
 @pytest.mark.asyncio
 async def test_shared_reconnect_replays_two_workspaces_and_removes_missing_projection():
     binding_a = execution_binding('workspace-a')
@@ -150,7 +163,7 @@ async def test_empty_projected_workspaces_do_not_retain_installation_sets():
 
     assert connector._workspace_installations == {}
     assert connector._known_desired_states == {}
-    connector.handler.reconcile_plugin_installations.assert_awaited_once_with(())
+    connector.handler.reconcile_plugin_installations.assert_awaited_once_with((), timeout=300.0)
 
 
 @pytest.mark.asyncio
diff --git a/tests/unit_tests/plugin/test_handler.py b/tests/unit_tests/plugin/test_handler.py
index 952d59134..844dcdc12 100644
--- a/tests/unit_tests/plugin/test_handler.py
+++ b/tests/unit_tests/plugin/test_handler.py
@@ -81,6 +81,18 @@ async def test_reconcile_plugin_installations_allows_cloud_cold_start_to_finish(
     assert runtime_handler.call_action.await_args.kwargs['timeout'] == 300
 
 
+@pytest.mark.asyncio
+async def test_reconcile_plugin_installations_accepts_configured_cold_start_timeout():
+    runtime_handler = make_handler(SimpleNamespace())
+    runtime_handler.call_action = AsyncMock(return_value={})
+    binding = next(iter(runtime_handler._installation_bindings.values()))[0]
+    desired = PluginInstallationDesiredState(binding=binding, enabled=True)
+
+    await runtime_handler.reconcile_plugin_installations((desired,), timeout=900)
+
+    assert runtime_handler.call_action.await_args.kwargs["timeout"] == 900
+
+
 class TestHandlerQueryVariables:
     """Tests for handler query variable logic."""