From e5e62c8fe9759fa0d46b25d6a0209e85a660b903 Mon Sep 17 00:00:00 2001 From: Hyu Date: Tue, 1 Sep 2026 12:46:16 +0800 Subject: [PATCH] fix(plugin): keep status responsive during installation --- src/langbot/pkg/plugin/connector.py | 47 ++++++----- .../plugin/test_connector_reconcile.py | 77 +++++++++++++++++-- 2 files changed, 100 insertions(+), 24 deletions(-) diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 5cfc44df0..3bc33935f 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -789,9 +789,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): if not self.is_enable_plugin or not hasattr(self, 'handler'): return runtime_handler = self._runtime_handler() - desired_states = await self._load_workspace_desired_states(execution_context) - desired_by_uuid = {state.binding.installation_uuid: state for state in desired_states} async with self._state_lock: + # Read the durable desired state while holding the same gate used by + # install/remove bookkeeping. Otherwise a request can load a stale + # pre-install snapshot, wait for the installer to publish its + # in-memory state, and then incorrectly remove that new binding. + desired_states = await self._load_workspace_desired_states(execution_context) + desired_by_uuid = {state.binding.installation_uuid: state for state in desired_states} previous_ids = set(self._workspace_installations.get(execution_context.workspace_uuid, set())) for installation_uuid in previous_ids - set(desired_by_uuid): previous = self._known_desired_states.get(installation_uuid) @@ -1751,14 +1755,27 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): artifact_digest = hashlib.sha256(file_bytes).hexdigest() await self._store_artifact_package(execution_context, artifact_digest, file_bytes) try: - binding, previous_digest, previous_was_durable = await self._persist_installation_package( - execution_context, - plugin_author=plugin_author, - plugin_name=plugin_name, - install_source=install_source, - install_info=install_info, - artifact_digest=artifact_digest, - ) + # Persist and publish the new desired generation under the same + # gate used by request-time reconciliation. This closes the small + # window where another request could observe the durable row first + # and perform the same slow Runtime apply while holding the gate. + async with self._state_lock: + binding, previous_digest, previous_was_durable = await self._persist_installation_package( + execution_context, + plugin_author=plugin_author, + plugin_name=plugin_name, + install_source=install_source, + install_info=install_info, + artifact_digest=artifact_digest, + ) + desired = PluginInstallationDesiredState(binding=binding, enabled=True) + runtime_handler.register_installation_binding( + binding, + plugin_author=plugin_author, + plugin_name=plugin_name, + ) + self._known_desired_states[binding.installation_uuid] = desired + self._workspace_installations.setdefault(binding.workspace_uuid, set()).add(binding.installation_uuid) except Exception: await self._delete_artifact_if_unreferenced(execution_context, artifact_digest) raise @@ -1770,18 +1787,10 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): pass except Exception as exc: self.ap.logger.debug(f'Legacy OSS plugin cleanup skipped: {exc}') - runtime_handler.register_installation_binding( - binding, - plugin_author=plugin_author, - plugin_name=plugin_name, - ) await self._apply_desired_state( - PluginInstallationDesiredState(binding=binding, enabled=True), + desired, artifact_package=file_bytes, ) - desired = PluginInstallationDesiredState(binding=binding, enabled=True) - self._known_desired_states[binding.installation_uuid] = desired - self._workspace_installations.setdefault(binding.workspace_uuid, set()).add(binding.installation_uuid) if previous_digest is not None and previous_digest != artifact_digest: await self._delete_artifact_if_unreferenced(execution_context, previous_digest) await self._wait_for_installed_plugin_ready(plugin_author, plugin_name, task_context) diff --git a/tests/unit_tests/plugin/test_connector_reconcile.py b/tests/unit_tests/plugin/test_connector_reconcile.py index 0a230e656..35d06a72e 100644 --- a/tests/unit_tests/plugin/test_connector_reconcile.py +++ b/tests/unit_tests/plugin/test_connector_reconcile.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import datetime import hashlib from contextlib import nullcontext @@ -13,6 +14,7 @@ from langbot_plugin.runtime.plugin.mgr import PluginInstallSource from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.plugin.connector import ( PluginInstallationFailedError, + PluginInstallationDesiredState, PluginRuntimeConnector, ) @@ -77,6 +79,7 @@ def runtime_handler( apply_plugin_installation=AsyncMock(return_value={'state': 'starting'}), installation_scope=Mock(side_effect=lambda _binding: nullcontext()), list_plugins=AsyncMock(return_value=[]), + ping=AsyncMock(return_value={'pong': 'pong'}), ) @@ -109,15 +112,15 @@ def shared_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 + 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 + assert connector.handler.reconcile_plugin_installations.await_args.kwargs['timeout'] == 900 @pytest.mark.asyncio @@ -287,6 +290,70 @@ async def test_local_install_persists_verified_package_before_runtime_apply(): ) +@pytest.mark.asyncio +async def test_workspace_reads_do_not_wait_for_an_installation_apply(): + package = b'local-lbpkg-bytes' + digest = hashlib.sha256(package).hexdigest() + execution_context = ExecutionContext( + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=1, + ) + 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=digest, + ) + app = SimpleNamespace( + instance_config=SimpleNamespace(data={'plugin': {'enable': True}}), + deployment=SimpleNamespace(mode='cloud'), + logger=Mock(), + ) + connector = PluginRuntimeConnector(app, AsyncMock()) + connector.handler = runtime_handler() + connector._current_execution_context = AsyncMock(return_value=execution_context) + connector._validate_execution_context = AsyncMock(return_value=execution_context) + connector._inspect_plugin_package = Mock(return_value=('author', 'plugin')) + connector._store_artifact_package = AsyncMock() + connector._persist_installation_package = AsyncMock(return_value=(binding, None, False)) + connector._wait_for_installed_plugin_ready = AsyncMock() + connector._load_workspace_desired_states = AsyncMock( + return_value=[PluginInstallationDesiredState(binding=binding, enabled=True)] + ) + apply_started = asyncio.Event() + release_apply = asyncio.Event() + + async def slow_apply(*_args, **_kwargs): + apply_started.set() + await release_apply.wait() + return {'state': 'starting'} + + connector.handler.apply_plugin_installation = AsyncMock(side_effect=slow_apply) + install_task = asyncio.create_task( + connector.install_plugin( + PluginInstallSource.LOCAL, + {'plugin_file': package}, + ) + ) + try: + await asyncio.wait_for(apply_started.wait(), timeout=1) + + async def check_plugin_runtime_status(): + await connector.require_workspace_context(execution_context) + return await connector.ping_plugin_runtime() + + assert await asyncio.wait_for(check_plugin_runtime_status(), timeout=0.1) == {'pong': 'pong'} + + assert connector.handler.apply_plugin_installation.await_count == 1 + assert connector._known_desired_states[binding.installation_uuid].binding == binding + finally: + release_apply.set() + await install_task + + @pytest.mark.asyncio async def test_local_install_cleans_untracked_legacy_plugin_before_runtime_apply(): package = b'local-lbpkg-bytes'