langbot-plugin 0.5.2__tar.gz → 0.5.4__tar.gz

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: langbot-plugin
-Version: 0.5.2
+Version: 0.5.4
 Summary: This package contains the SDK, CLI for building plugins for LangBot, plus the runtime for hosting LangBot plugins
 Project-URL: Homepage, https://langbot.app
 Project-URL: Repository, https://github.com/langbot-app/langbot-plugin-sdk

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langbot-plugin"
-version = "0.5.2"
+version = "0.5.4"
 description = "This package contains the SDK, CLI for building plugins for LangBot, plus the runtime for hosting LangBot plugins"
 readme = "README.md"
 authors = [

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/src/langbot_plugin/runtime/io/handler.py

@@ -34,7 +34,10 @@ from langbot_plugin.entities.io.errors import (
     ActionCallError,
 )
 from langbot_plugin.entities.io.actions.enums import ActionType, CommonAction
-from langbot_plugin.runtime.security import PLUGIN_RUNTIME_PROFILE_ENV
+from langbot_plugin.runtime.security import (
+    PLUGIN_FILE_STORAGE_DIR_ENV,
+    PLUGIN_RUNTIME_PROFILE_ENV,
+)
 from langbot_plugin.runtime.bounded_executor import blocking_work_scope
 
 logger = logging.getLogger(__name__)
@@ -124,7 +127,7 @@ class Handler(abc.ABC):
 
         if file_storage_dir is None:
             runtime_profile = os.environ.get(PLUGIN_RUNTIME_PROFILE_ENV, "oss_dev")
-            file_storage_dir = (
+            file_storage_dir = os.environ.get(PLUGIN_FILE_STORAGE_DIR_ENV) or (
                 SHARED_WORKER_FILE_STORAGE_DIR
                 if runtime_profile == "shared"
                 else FILE_STORAGE_DIR

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/src/langbot_plugin/runtime/plugin/mgr.py

@@ -219,6 +219,23 @@ class PluginManager:
                 initial_policy.max_concurrent_restarts
             )
 
+    async def _complete_installation_transition(
+        self,
+        operation: typing.Coroutine[typing.Any, typing.Any, None],
+    ) -> None:
+        """Join one lifecycle mutation through repeated caller cancellation."""
+
+        task = asyncio.create_task(operation)
+        cancelled = False
+        while not task.done():
+            try:
+                await asyncio.shield(task)
+            except asyncio.CancelledError:
+                cancelled = True
+        task.result()
+        if cancelled:
+            raise asyncio.CancelledError
+
     @property
     def installation_runtimes(
         self,
@@ -1058,11 +1075,20 @@ class PluginManager:
         self,
         binding: InstallationBinding,
     ) -> dict[str, typing.Any]:
-        binding = self.context.deactivate_installation_binding(binding)
-        async with self._installation_lifecycle_semaphore():
-            await self._revoke_installation_runtime(binding)
-        if self._active_binding_by_uuid.get(binding.installation_uuid) == binding:
-            self._active_binding_by_uuid.pop(binding.installation_uuid, None)
+        async def remove() -> None:
+            removed_binding = self.context.deactivate_installation_binding(binding)
+            async with self._installation_lifecycle_semaphore():
+                await self._revoke_installation_runtime(removed_binding)
+            if (
+                self._active_binding_by_uuid.get(removed_binding.installation_uuid)
+                == removed_binding
+            ):
+                self._active_binding_by_uuid.pop(
+                    removed_binding.installation_uuid,
+                    None,
+                )
+
+        await self._complete_installation_transition(remove())
         return {
             "installation_uuid": binding.installation_uuid,
             "state": "removed",
@@ -1108,22 +1134,26 @@ class PluginManager:
                 lock = self._retain_installation_operation_lock(installation_uuid)
                 try:
                     async with lock:
-                        if self.context.is_current_installation_binding(
-                            current_binding
-                        ):
-                            self.context.deactivate_installation_binding(
+
+                        async def remove_current() -> None:
+                            if self.context.is_current_installation_binding(
                                 current_binding
-                            )
-                        async with self._installation_lifecycle_semaphore():
-                            await self._revoke_installation_runtime(current_binding)
-                        if (
-                            self._active_binding_by_uuid.get(installation_uuid)
-                            == current_binding
-                        ):
-                            self._active_binding_by_uuid.pop(
-                                installation_uuid,
-                                None,
-                            )
+                            ):
+                                self.context.deactivate_installation_binding(
+                                    current_binding
+                                )
+                            async with self._installation_lifecycle_semaphore():
+                                await self._revoke_installation_runtime(current_binding)
+                            if (
+                                self._active_binding_by_uuid.get(installation_uuid)
+                                == current_binding
+                            ):
+                                self._active_binding_by_uuid.pop(
+                                    installation_uuid,
+                                    None,
+                                )
+
+                        await self._complete_installation_transition(remove_current())
                         removed.append(installation_uuid)
                 finally:
                     self._forget_installation_operation_lock(
@@ -1469,10 +1499,12 @@ class PluginManager:
         binding: InstallationBinding,
     ) -> None:
         self._revoke_registration_capabilities_for_binding(binding)
-        runtime = self._installations.pop(binding, None)
+        runtime = self._installations.get(binding)
         if runtime is None:
             return
         await self._stop_installation_worker(runtime)
+        if self._installations.get(binding) is runtime:
+            self._installations.pop(binding, None)
 
     async def _stop_installation_worker(
         self,
@@ -1486,20 +1518,29 @@ class PluginManager:
             except Exception as exc:
                 logger.warning("Failed to notify revoked plugin worker: %s", exc)
             close = getattr(handler, "close", None)
-            if close is not None:
-                await close()
-            else:
-                await handler.conn.close()
-            if handler in self.plugin_handlers:
-                self.plugin_handlers.remove(handler)
+            try:
+                if close is not None:
+                    await close()
+                else:
+                    await handler.conn.close()
+            except Exception as exc:
+                logger.warning("Failed to close revoked plugin worker: %s", exc)
+            finally:
+                if handler in self.plugin_handlers:
+                    self.plugin_handlers.remove(handler)
+                runtime.plugin_handler = None
             process = handler.stdio_process
             if process is not None and process.returncode is None:
-                stopped = await stdio_client_controller.stop_process(process)
-                if not stopped:
-                    logger.error(
-                        "Plugin worker process did not exit after SIGKILL: %s",
-                        runtime.binding.installation_uuid,
-                    )
+                try:
+                    stopped = await stdio_client_controller.stop_process(process)
+                except Exception as exc:
+                    logger.error("Failed to stop revoked plugin worker: %s", exc)
+                else:
+                    if not stopped:
+                        logger.error(
+                            "Plugin worker process did not exit after SIGKILL: %s",
+                            runtime.binding.installation_uuid,
+                        )
         runtime.plugin_handler = None
         if runtime.plugin_container is not None:
             self._binding_by_container_id.pop(id(runtime.plugin_container), None)

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/src/langbot_plugin/runtime/plugin/restart_coordinator.py

@@ -3,14 +3,16 @@ from __future__ import annotations
 import asyncio
 import collections
 import time
-from collections.abc import Awaitable
+from collections.abc import Coroutine
 from dataclasses import dataclass
 from typing import Callable
 
 from langbot_plugin.entities.io.context import PluginWorkerPolicy
 
 
-async def _complete_state_transition(operation: Awaitable[None]) -> None:
+async def _complete_state_transition(
+    operation: Coroutine[object, object, None],
+) -> None:
     """Finish one tiny coordinator mutation even if its caller is cancelled.
 
     A supervisor cancellation must not strand the circuit in half-open state.
@@ -19,14 +21,19 @@ async def _complete_state_transition(operation: Awaitable[None]) -> None:
     """
 
     task = asyncio.create_task(operation)
-    try:
-        await asyncio.shield(task)
-    except asyncio.CancelledError:
-        # The cancellation has already been delivered to this caller; the
-        # shielded mutation remains live and can now be joined directly before
-        # propagating the original cancellation.
-        await task
-        raise
+    cancelled = False
+    while not task.done():
+        try:
+            await asyncio.shield(task)
+        except asyncio.CancelledError:
+            # A caller can be cancelled repeatedly while the mutation is
+            # waiting for the state lock. Keep joining the shielded task until
+            # the transition is complete; otherwise a half-open probe can be
+            # stranded permanently.
+            cancelled = True
+    task.result()
+    if cancelled:
+        raise asyncio.CancelledError
 
 
 @dataclass(slots=True)

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/src/langbot_plugin/runtime/plugin/worker_launcher.py

@@ -28,6 +28,7 @@ from langbot_plugin.runtime.plugin.dependency_environment import (
 from langbot_plugin.runtime.helper import pkgmgr as pkgmgr_helper
 from langbot_plugin.runtime import bounded_executor
 from langbot_plugin.runtime.security import (
+    PLUGIN_FILE_STORAGE_DIR_ENV,
     PLUGIN_REGISTRATION_CAPABILITY_ENV,
     PLUGIN_RUNTIME_PROFILE_ENV,
 )
@@ -140,6 +141,9 @@ class PluginWorkerLauncher:
                 "--prod",
             ],
             env={
+                PLUGIN_FILE_STORAGE_DIR_ENV: str(
+                    launch_spec.paths.root_path / "rpc-transfer"
+                ),
                 PLUGIN_REGISTRATION_CAPABILITY_ENV: (
                     launch_spec.registration_capability
                 ),

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/src/langbot_plugin/runtime/security.py

@@ -15,6 +15,7 @@ PLUGIN_DEBUG_KEY_HEADER = "X-LangBot-Plugin-Debug-Key"
 PLUGIN_REGISTRATION_CAPABILITY_ENV = "LANGBOT_PLUGIN_REGISTRATION_CAPABILITY"
 PLUGIN_REGISTRATION_CAPABILITY_HEADER = "X-LangBot-Plugin-Registration-Capability"
 PLUGIN_RUNTIME_PROFILE_ENV = "LANGBOT_PLUGIN_RUNTIME_PROFILE"
+PLUGIN_FILE_STORAGE_DIR_ENV = "LANGBOT_PLUGIN_FILE_STORAGE_DIR"
 RUNTIME_SECRET_MIN_LENGTH = 32
 WORKSPACE_DEBUG_TOKEN_TTL_SECONDS = 2 * 60 * 60
 

langbot_plugin-0.5.4/src/langbot_plugin/version.py

@@ -0,0 +1 @@
+__version__ = "0.5.4"

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/tests/runtime/io/test_handler.py

@@ -511,6 +511,21 @@ def test_handler_file_storage_dir_is_created_for_instances(tmp_path, monkeypatch
     assert (tmp_path / "data" / "temp" / "lbp").is_dir()
 
 
+def test_managed_worker_file_storage_uses_launcher_directory(tmp_path, monkeypatch):
+    artifact = tmp_path / "artifact" / "code"
+    artifact.mkdir(parents=True)
+    artifact.chmod(0o555)
+    rpc_transfer = tmp_path / "installations" / "installation-a" / "rpc-transfer"
+    monkeypatch.chdir(artifact)
+    monkeypatch.setenv("LANGBOT_PLUGIN_FILE_STORAGE_DIR", str(rpc_transfer))
+
+    handler = Handler(QueueConnection())
+
+    assert handler.file_storage_dir == str(rpc_transfer)
+    assert rpc_transfer.is_dir()
+    assert not (artifact / "data").exists()
+
+
 def test_shared_worker_file_storage_uses_private_writable_tmp(tmp_path, monkeypatch):
     worker_tmp = tmp_path / "lbp-rpc"
     monkeypatch.setenv(PLUGIN_RUNTIME_PROFILE_ENV, "shared")

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/tests/runtime/plugin/test_installation_manager.py

@@ -131,6 +131,110 @@ async def test_manager_indexes_same_artifact_installations_by_complete_binding(
     assert context.is_current_installation_binding(binding_b)
 
 
+async def test_cancelled_remove_finishes_worker_revoke_before_propagating(
+    tmp_path,
+):
+    context, manager = _manager(tmp_path)
+    package = _package()
+    digest = hashlib.sha256(package).hexdigest()
+    binding = _binding(
+        "installation-a",
+        digest,
+        workspace_uuid="workspace-a",
+    )
+    await manager.apply_plugin_installation(
+        binding,
+        artifact_package=package,
+        enabled=False,
+    )
+    runtime = manager.installation_runtimes[binding]
+    shutdown_started = asyncio.Event()
+    allow_shutdown = asyncio.Event()
+
+    class Handler:
+        stdio_process = None
+        conn = None
+
+        def cancel_inflight_messages(self):
+            return None
+
+        async def shutdown_plugin(self):
+            shutdown_started.set()
+            await allow_shutdown.wait()
+
+        async def close(self):
+            return None
+
+    handler = Handler()
+    runtime.plugin_handler = handler
+    manager.plugin_handlers.append(handler)
+
+    remove_task = asyncio.create_task(manager.remove_plugin_installation(binding))
+    await asyncio.wait_for(shutdown_started.wait(), timeout=1)
+    remove_task.cancel()
+    await asyncio.sleep(0)
+    remove_task.cancel()
+    allow_shutdown.set()
+
+    with pytest.raises(asyncio.CancelledError):
+        await remove_task
+    assert not context.is_current_installation_binding(binding)
+    assert binding not in manager.installation_runtimes
+    assert binding.installation_uuid not in manager._active_binding_by_uuid
+    assert handler not in manager.plugin_handlers
+    assert runtime.plugin_handler is None
+
+    with pytest.raises(ValueError, match="stale"):
+        await manager.reconcile_plugin_installations(
+            (PluginInstallationDesiredState(binding=binding, enabled=False),)
+        )
+
+
+async def test_remove_finishes_state_cleanup_when_handler_close_fails(
+    tmp_path,
+):
+    context, manager = _manager(tmp_path)
+    package = _package()
+    digest = hashlib.sha256(package).hexdigest()
+    binding = _binding(
+        "installation-a",
+        digest,
+        workspace_uuid="workspace-a",
+    )
+    await manager.apply_plugin_installation(
+        binding,
+        artifact_package=package,
+        enabled=False,
+    )
+    runtime = manager.installation_runtimes[binding]
+
+    class Handler:
+        stdio_process = None
+        conn = None
+
+        def cancel_inflight_messages(self):
+            return None
+
+        async def shutdown_plugin(self):
+            return None
+
+        async def close(self):
+            raise RuntimeError("close failed")
+
+    handler = Handler()
+    runtime.plugin_handler = handler
+    manager.plugin_handlers.append(handler)
+
+    result = await manager.remove_plugin_installation(binding)
+
+    assert result["state"] == "removed"
+    assert not context.is_current_installation_binding(binding)
+    assert binding not in manager.installation_runtimes
+    assert binding.installation_uuid not in manager._active_binding_by_uuid
+    assert handler not in manager.plugin_handlers
+    assert runtime.plugin_handler is None
+
+
 async def test_desired_state_does_not_reject_on_aggregate_worker_capacity(
     tmp_path,
     monkeypatch,

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/tests/runtime/plugin/test_restart_coordinator.py

@@ -198,6 +198,42 @@ async def test_abandoned_probe_releases_slot_without_counting_failure():
     await replacement_probe.abandon()
 
 
+async def test_repeated_cancellation_cannot_strand_half_open_probe():
+    coordinator = PluginRestartCoordinator()
+    coordinator.configure(
+        _policy(
+            max_concurrent_restarts=1,
+            restart_failure_threshold=1,
+            restart_circuit_open_seconds=0.01,
+        )
+    )
+    failed = await coordinator.acquire("installation-a")
+    await failed.record_failure()
+    probe = await asyncio.wait_for(
+        coordinator.acquire("installation-a"),
+        timeout=1,
+    )
+    assert probe.is_half_open_probe is True
+
+    await coordinator._state_lock.acquire()
+    abandon_task = asyncio.create_task(probe.abandon())
+    await asyncio.sleep(0)
+    abandon_task.cancel()
+    await asyncio.sleep(0)
+    abandon_task.cancel()
+    coordinator._state_lock.release()
+
+    with pytest.raises(asyncio.CancelledError):
+        await abandon_task
+    assert coordinator.snapshot()["half_open_probe_inflight"] is False
+
+    replacement = await asyncio.wait_for(
+        coordinator.acquire("installation-a"),
+        timeout=1,
+    )
+    await replacement.abandon()
+
+
 async def test_cancelled_acquire_does_not_leak_launch_slot():
     coordinator = PluginRestartCoordinator()
     coordinator.configure(

{langbot_plugin-0.5.2 → langbot_plugin-0.5.4}/tests/runtime/plugin/test_worker_launcher.py

@@ -22,6 +22,7 @@ from langbot_plugin.runtime.plugin.worker_launcher import (
     PluginWorkerLauncher,
 )
 from langbot_plugin.runtime.security import (
+    PLUGIN_FILE_STORAGE_DIR_ENV,
     PLUGIN_REGISTRATION_CAPABILITY_ENV,
     PLUGIN_RUNTIME_PROFILE_ENV,
 )
@@ -327,6 +328,7 @@ def test_launcher_builds_profile_specific_controllers(tmp_path):
     assert oss_controller.command == sys.executable
     assert oss_controller.working_dir == str(launch_spec.artifact.code_path)
     assert oss_controller.env == {
+        PLUGIN_FILE_STORAGE_DIR_ENV: str(launch_spec.paths.root_path / "rpc-transfer"),
         PLUGIN_REGISTRATION_CAPABILITY_ENV: "capability-value",
         PLUGIN_RUNTIME_PROFILE_ENV: "oss_dev",
     }

langbot_plugin-0.5.2/src/langbot_plugin/version.py

@@ -1 +0,0 @@
-__version__ = "0.5.2"