langbot-plugin 0.6.4__tar.gz → 0.6.6__tar.gz

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: langbot-plugin
-Version: 0.6.4
+Version: 0.6.6
 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.6.4 → langbot_plugin-0.6.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langbot-plugin"
-version = "0.6.4"
+version = "0.6.6"
 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.6.4 → langbot_plugin-0.6.6}/src/langbot_plugin/cli/run/controller.py

@@ -315,6 +315,7 @@ class PluginRuntimeController:
 
                 # If in production mode, break the loop after first connection
                 if self.prod_mode:
+                    await self._controller_task
                     break
 
                 # In debug mode, wait for shutdown signal

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/src/langbot_plugin/runtime/io/controllers/stdio/client.py

@@ -65,12 +65,14 @@ class StdioClientController(Controller):
         working_dir: str = ".",
         *,
         capture_stderr: bool = False,
+        close_process_on_callback_return: bool = True,
     ):
         self.command = command
         self.args = args
         self.env = env
         self.working_dir = working_dir
         self.capture_stderr = capture_stderr
+        self.close_process_on_callback_return = close_process_on_callback_return
 
     async def run(
         self,
@@ -95,7 +97,8 @@ class StdioClientController(Controller):
         try:
             await new_connection_callback(self.connection)
         finally:
-            await self.close()
+            if self.close_process_on_callback_return:
+                await self.close()
 
     async def close(self) -> None:
         """Close pipes and reap the owned subprocess."""

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/src/langbot_plugin/runtime/plugin/mgr.py

@@ -26,6 +26,7 @@ from langbot_plugin.runtime.io.connection import Connection
 from langbot_plugin.runtime.io.controllers.stdio import (
     client as stdio_client_controller,
 )
+from langbot_plugin.runtime.io.controller import Controller
 from langbot_plugin.runtime.plugin.runner_service import RunnerRuntimeService
 from langbot_plugin.runtime.plugin import container as runtime_plugin_container
 from langbot_plugin.runtime.io.handlers import plugin as runtime_plugin_handler_cls
@@ -148,6 +149,7 @@ class SharedPluginWorkerRuntime:
         default_factory=dict
     )
     plugin_handler: runtime_plugin_handler_cls.PluginConnectionHandler | None = None
+    controller: Controller | None = None
     launch_task: asyncio.Task[None] | None = None
     transport_registered_event: asyncio.Event = field(default_factory=asyncio.Event)
     ready_event: asyncio.Event = field(default_factory=asyncio.Event)
@@ -1717,6 +1719,14 @@ class PluginManager:
             task.cancel()
             await asyncio.gather(task, return_exceptions=True)
         worker.launch_task = None
+        controller = worker.controller
+        if controller is not None:
+            close = getattr(controller, "close", None)
+            if close is not None:
+                with contextlib.suppress(Exception):
+                    await close()
+            if worker.controller is controller:
+                worker.controller = None
         if not worker.slots:
             transfer_root = (
                 self.artifact_store.base_path
@@ -1954,6 +1964,7 @@ class PluginManager:
                     dependency_environment=worker.dependency_environment,
                 )
             )
+            worker.controller = controller
 
             async def new_plugin_connection_callback(connection: Connection):
                 shared_transfer_root = (
@@ -1993,7 +2004,15 @@ class PluginManager:
                 finally:
                     await self.remove_plugin_handler(plugin_handler)
 
-            await controller.run(new_plugin_connection_callback)
+            try:
+                await controller.run(new_plugin_connection_callback)
+            finally:
+                close = getattr(controller, "close", None)
+                if close is not None:
+                    with contextlib.suppress(Exception):
+                        await close()
+                if worker.controller is controller:
+                    worker.controller = None
         finally:
             self._revoke_registration_capability(capability)
 

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/src/langbot_plugin/runtime/plugin/worker_launcher.py

@@ -625,6 +625,7 @@ class PluginWorkerLauncher:
             args=self.build_shared_pool_nsjail_args(launch_spec),
             env={},
             working_dir="/",
+            close_process_on_callback_return=False,
         )
 
     def _append_readonly_runtime_mounts(

langbot_plugin-0.6.6/src/langbot_plugin/version.py

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

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/tests/cli/run/test_controller.py

@@ -585,6 +585,37 @@ async def test_mount_stdio_prod_registers_plugin(monkeypatch):
     assert controller.plugin_container.status is RuntimeContainerStatus.MOUNTED
 
 
+@pytest.mark.asyncio
+async def test_mount_stdio_prod_keeps_serving_after_registration(monkeypatch):
+    capability = "registration-capability-that-is-long-enough-for-tests"
+    monkeypatch.setenv("LANGBOT_PLUGIN_REGISTRATION_CAPABILITY", capability)
+    fake_handler_cls = _fake_handler_class(run_forever=True)
+    monkeypatch.setattr(controller_module, "PluginRuntimeHandler", fake_handler_cls)
+    _install_stdio_controller(monkeypatch, "connect")
+    controller = PluginRuntimeController(
+        plugin_manifest=_manifest("Plugin", "demo"),
+        component_manifests=[],
+        stdio=True,
+        ws_debug_url="ws://runtime/plugin/ws",
+        prod_mode=True,
+    )
+
+    mount_task = asyncio.create_task(controller.mount())
+    handler = None
+    try:
+        await _wait_until(lambda: bool(getattr(fake_handler_cls, "instances")))
+        handler = getattr(fake_handler_cls, "instances")[0]
+        await handler.registered.wait()
+        await asyncio.sleep(0)
+        assert not mount_task.done()
+    finally:
+        mount_task.cancel()
+        await asyncio.gather(mount_task, return_exceptions=True)
+
+    assert handler is not None
+    assert handler.cancelled.is_set()
+
+
 @pytest.mark.asyncio
 async def test_mount_prod_connection_failure_sets_waiter_and_exits(monkeypatch):
     capability = "registration-capability-that-is-long-enough-for-tests"

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/tests/packaging/test_release_metadata.py

@@ -11,5 +11,5 @@ def test_stable_release_metadata_is_consistent() -> None:
         (root / "src/langbot_plugin/version.py").read_text(encoding="utf-8").strip()
     )
 
-    assert project["project"]["version"] == "0.6.4"
-    assert source_version == '__version__ = "0.6.4"'
+    assert project["project"]["version"] == "0.6.6"
+    assert source_version == '__version__ = "0.6.6"'

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/tests/runtime/io/test_controllers.py

@@ -20,7 +20,7 @@ class FakeProcess:
     def __init__(self, stdin=object(), stdout=object()):
         self.stdin = stdin
         self.stdout = stdout
-        self.returncode = 0
+        self.returncode: int | None = 0
 
 
 class FakeWebSocket:
@@ -114,6 +114,33 @@ async def test_stdio_client_controller_captures_stderr_only_when_requested(
     assert captured["stderr"] is stdio_client.asyncio.subprocess.PIPE
 
 
+async def test_stdio_client_controller_can_transfer_process_lifetime_ownership(
+    monkeypatch,
+):
+    process = FakeProcess()
+    process.returncode = None
+
+    async def fake_create_subprocess_exec(*args, **kwargs):
+        return process
+
+    monkeypatch.setattr(
+        stdio_client.asyncio,
+        "create_subprocess_exec",
+        fake_create_subprocess_exec,
+    )
+    controller = stdio_client.StdioClientController(
+        "python",
+        [],
+        {},
+        close_process_on_callback_return=False,
+    )
+
+    await controller.run(lambda _connection: asyncio.sleep(0))
+
+    assert controller.process is process
+    assert controller.connection is not None
+
+
 async def test_stdio_client_controller_rejects_missing_pipes(monkeypatch):
     async def fake_create_subprocess_exec(*args, **kwargs):
         return FakeProcess(stdin=None, stdout=object())

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/tests/runtime/plugin/test_installation_manager.py

@@ -2082,6 +2082,7 @@ async def test_shared_worker_ready_timeout_cancels_hung_controller_and_records_f
     )
     worker = manager.installation_runtimes[binding].shared_worker
     cancelled = asyncio.Event()
+    closed = asyncio.Event()
 
     class NeverRegisterController:
         process = None
@@ -2092,6 +2093,9 @@ async def test_shared_worker_ready_timeout_cancels_hung_controller_and_records_f
             finally:
                 cancelled.set()
 
+        async def close(self):
+            closed.set()
+
     monkeypatch.setattr(
         manager.worker_launcher,
         "create_shared_pool_controller",
@@ -2110,6 +2114,8 @@ async def test_shared_worker_ready_timeout_cancels_hung_controller_and_records_f
     await asyncio.wait_for(wait_for_failure(), timeout=1)
 
     assert cancelled.is_set()
+    assert closed.is_set()
+    assert worker.controller is None
     assert manager.restart_coordinator.snapshot()["active_launches"] == 0
     assert manager.installation_runtimes[binding].state == "failed"
     await manager._stop_shared_worker(worker)

{langbot_plugin-0.6.4 → langbot_plugin-0.6.6}/tests/runtime/plugin/test_worker_launcher.py

@@ -277,6 +277,19 @@ def test_shared_launcher_absolutizes_runtime_paths_before_fixed_root_cwd(
     assert f"{staging.tmp_path.absolute()}:/tmp" in dependency_args
 
 
+def test_shared_pool_controller_transfers_process_lifetime_ownership(tmp_path):
+    launcher = PluginWorkerLauncher(
+        nsjail_path="/usr/bin/nsjail",
+        cgroup_v2_available=True,
+        platform="linux",
+    )
+    launcher.configure(_policy(), "shared")
+
+    controller = launcher.create_shared_pool_controller(_launch_spec(tmp_path))
+
+    assert getattr(controller, "close_process_on_callback_return") is False
+
+
 def test_shared_launcher_fails_without_nsjail_or_required_cgroup():
     no_nsjail = PluginWorkerLauncher(
         nsjail_path="",

langbot_plugin-0.6.4/src/langbot_plugin/version.py

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