diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index 4b82b3efc..f1bff1c5b 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -301,11 +301,36 @@ class Application: async def initialize(self): pass + async def _initialize_plugin_runtime(self) -> None: + try: + await self.plugin_connector.initialize() + except asyncio.CancelledError: + raise + except Exception as exc: + self.logger.warning(f'Plugin runtime unavailable during startup; reconnecting in background: {exc}') + self.plugin_connector.schedule_reconnect() + + def _start_plugin_runtime_initialization(self) -> asyncio.Task | None: + task = getattr(self, '_plugin_runtime_initialization_task', None) + if task is not None and not task.done(): + return task + # This is application lifecycle work, not a request side effect. It must + # not wait on PersistenceManager's after-commit gate at boot. + task = asyncio.create_task( + self._initialize_plugin_runtime(), + name='plugin-runtime-initialization', + ) + self._plugin_runtime_initialization_task = task + return task + async def run(self): self.event_loop_monitor.start() try: - if self.directory_projection_service is not None: - self.task_mgr.create_task( + if ( + self.directory_projection_service is not None + and getattr(self, 'directory_projection_task', None) is None + ): + self.directory_projection_task = self.task_mgr.create_task( self.directory_projection_service.run(), name='cloud-directory-projection', scopes=[core_entities.LifecycleControlScope.APPLICATION], @@ -322,7 +347,6 @@ class Application: name='cloud-manifest-refresh', scopes=[core_entities.LifecycleControlScope.APPLICATION], ) - await self.plugin_connector.initialize_plugins() # 后续可能会允许动态重启其他任务 # 故为了防止程序在非 Ctrl-C 情况下退出，这里创建一个不会结束的协程 @@ -348,6 +372,7 @@ class Application: name='http-api-controller', scopes=[core_entities.LifecycleControlScope.APPLICATION], ) + self._start_plugin_runtime_initialization() # Telemetry instance heartbeat (startup + daily); respects # space.disable_telemetry via TelemetryManager.send(). @@ -529,6 +554,11 @@ class Application: if self.task_mgr is not None: self.task_mgr.cancel_by_scope(core_entities.LifecycleControlScope.APPLICATION) + plugin_runtime_task = getattr(self, '_plugin_runtime_initialization_task', None) + if plugin_runtime_task is not None and not plugin_runtime_task.done(): + plugin_runtime_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await plugin_runtime_task with contextlib.suppress(Exception): await self.event_loop_monitor.stop() mcp_mount = getattr(self.http_ctrl, 'mcp_mount', None) diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index 14d533dcc..7941a5553 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -1,6 +1,6 @@ from __future__ import annotations -from .. import stage, app +from .. import stage, app, entities as core_entities from ...utils import version, proxy, constants from ...pipeline import pool, controller, pipelinemgr from ...pipeline import aggregator as message_aggregator @@ -292,14 +292,17 @@ class BuildAppStage(stage.BootingStage): async def runtime_disconnect_callback(connector: plugin_connector.PluginRuntimeConnector) -> None: connector.schedule_reconnect() + if ap.directory_projection_service is not None: + # Keep the projection fresh while shared Runtime cold restore runs. + # BuildApp initializes the connector before Application.run() starts + # its long-lived tasks, so start the single refresh task here. + ap.directory_projection_task = ap.task_mgr.create_task( + ap.directory_projection_service.run(), + name='cloud-directory-projection', + scopes=[core_entities.LifecycleControlScope.APPLICATION], + ) + plugin_connector_inst = plugin_connector.PluginRuntimeConnector(ap, runtime_disconnect_callback) - try: - await plugin_connector_inst.initialize() - except Exception as exc: - # Keep the API/UI available while an external or managed runtime is - # starting, then recover in the background with bounded backoff. - ap.logger.warning(f'Plugin runtime unavailable during startup; reconnecting in background: {exc}') - plugin_connector_inst.schedule_reconnect() ap.plugin_connector = plugin_connector_inst workspace_service_inst.release_startup_execution_bindings() diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 0279ab285..03893cf4c 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -701,7 +701,13 @@ 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 +742,13 @@ 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/core/test_app_shutdown.py b/tests/unit_tests/core/test_app_shutdown.py index ad6dcfad6..2c953aaac 100644 --- a/tests/unit_tests/core/test_app_shutdown.py +++ b/tests/unit_tests/core/test_app_shutdown.py @@ -144,3 +144,39 @@ async def test_runtime_resource_stats_are_aggregate_and_constant_time() -> None: assert stats['models']['providers'] == 1 assert stats['runtimes']['plugin_installations'] == 1 assert stats['runtimes']['plugin_runtime_connected'] is True + + +@pytest.mark.asyncio +async def test_start_plugin_runtime_initialization_bypasses_after_commit_gate() -> None: + app = Application() + app.plugin_connector = SimpleNamespace(initialize=AsyncMock()) + app.task_mgr = SimpleNamespace(create_task=AsyncMock()) + + task = app._start_plugin_runtime_initialization() + await task + + app.plugin_connector.initialize.assert_awaited_once_with() + app.task_mgr.create_task.assert_not_called() + + +@pytest.mark.asyncio +async def test_shutdown_cancels_plugin_runtime_initialization_task() -> None: + app = Application() + app._plugin_runtime_initialization_task = asyncio.create_task(asyncio.sleep(60)) + app.task_mgr = SimpleNamespace(cancel_by_scope=lambda *_: None, tasks=[]) + app.event_loop_monitor = SimpleNamespace(stop=AsyncMock()) + app.http_ctrl = SimpleNamespace(mcp_mount=None) + app.platform_mgr = None + app.tool_mgr = None + app.model_mgr = None + app.box_service = None + app.plugin_connector = None + app.telemetry = None + app.vector_db_mgr = None + app.storage_mgr = None + app.persistence_mgr = SimpleNamespace(db=SimpleNamespace(engine=SimpleNamespace(dispose=AsyncMock()))) + app.deployment = None + + await app.shutdown() + + assert app._plugin_runtime_initialization_task.cancelled() 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."""