fix(agent): harden runner integration and QA

tests/unit_tests/plugin/test_connector_methods.py

@@ -20,6 +20,7 @@ from tests.factories import text_query
 from langbot_plugin.entities.io.context import InstallationBinding
 
 from langbot.pkg.api.http.context import ExecutionContext
+from langbot.pkg.workspace.errors import WorkspaceNotFoundError
 
 
 TEST_EXECUTION_CONTEXT = ExecutionContext(
@@ -76,6 +77,63 @@ def configure_handler(connector, runtime_handler):
     return runtime_handler
 
 
+async def _collect_agent_results(connector, context):
+    return [
+        result
+        async for result in connector.run_agent(
+            'qa',
+            'agent-runner',
+            'default',
+            context,
+        )
+    ]
+
+
+class TestRunAgent:
+    @pytest.mark.asyncio
+    async def test_revalidates_trusted_execution_context(self):
+        connector = create_mock_connector()
+        connector._current_execution_context = AsyncMock(
+            return_value=TEST_EXECUTION_CONTEXT
+        )
+
+        class RuntimeHandler:
+            installation_scope = Mock(
+                side_effect=lambda _binding: nullcontext()
+            )
+
+            async def run_agent(self, *_args):
+                yield {'type': 'run.completed'}
+
+        configure_handler(connector, RuntimeHandler())
+
+        results = await _collect_agent_results(
+            connector,
+            {'conversation': {'workspace_id': TEST_EXECUTION_CONTEXT.workspace_uuid}},
+        )
+
+        assert results == [{'type': 'run.completed'}]
+        connector.require_workspace_context.assert_awaited_once_with(
+            TEST_EXECUTION_CONTEXT
+        )
+
+    @pytest.mark.asyncio
+    async def test_rejects_payload_workspace_mismatch(self):
+        connector = create_mock_connector()
+        connector._current_execution_context = AsyncMock(
+            return_value=TEST_EXECUTION_CONTEXT
+        )
+        configure_handler(connector, AsyncMock())
+
+        with pytest.raises(WorkspaceNotFoundError, match='Plugin resource not found'):
+            await _collect_agent_results(
+                connector,
+                {'conversation': {'workspace_id': 'workspace-other'}},
+            )
+
+        connector.require_workspace_context.assert_not_awaited()
+
+
 class TestListPlugins:
     """Tests for list_plugins method."""
 

tests/unit_tests/plugin/test_connector_ping.py

@@ -72,6 +72,16 @@ async def test_stop_transport_tolerates_handler_callback_removing_attribute():
     assert not hasattr(connector, 'handler')
 
 
+@pytest.mark.asyncio
+async def test_stop_transport_tolerates_cancelled_controller_close():
+    connector = make_connector()
+    connector.ctrl = SimpleNamespace(close=AsyncMock(side_effect=asyncio.CancelledError))
+
+    await connector._stop_transport()
+
+    connector.ctrl.close.assert_awaited_once_with()
+
+
 @pytest.mark.asyncio
 async def test_stdio_runtime_connection_does_not_capture_unconsumed_stderr(
     monkeypatch: pytest.MonkeyPatch,

tests/unit_tests/plugin/test_connector_reconcile.py

@@ -249,6 +249,57 @@ async def test_local_install_persists_verified_package_before_runtime_apply():
     )
 
 
+@pytest.mark.asyncio
+async def test_local_install_cleans_untracked_legacy_plugin_before_runtime_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='oss'),
+        logger=Mock(),
+    )
+    connector = PluginRuntimeConnector(app, AsyncMock())
+    connector.handler = runtime_handler()
+    connector._current_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()
+    events: list[str] = []
+
+    async def delete_legacy_plugin(plugin_author: str, plugin_name: str):
+        assert (plugin_author, plugin_name) == ('author', 'plugin')
+        events.append('cleanup')
+        yield {'current_action': 'plugin deleted'}
+
+    async def apply_installation(*_args, **_kwargs):
+        events.append('apply')
+        return {'state': 'starting'}
+
+    connector.handler.delete_plugin = delete_legacy_plugin
+    connector.handler.apply_plugin_installation = AsyncMock(side_effect=apply_installation)
+
+    await connector.install_plugin(
+        PluginInstallSource.LOCAL,
+        {'plugin_file': package},
+    )
+
+    assert events == ['cleanup', 'apply']
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(('remaining_references', 'statement_count'), [(1, 1), (0, 2)])
 async def test_artifact_cleanup_is_reference_counted_within_workspace(

tests/unit_tests/plugin/test_handler_tenancy.py

@@ -313,6 +313,21 @@ def test_runtime_connection_is_instance_scoped_and_unbound():
     assert runtime_handler.bound_action_context is None
 
 
+def test_explicit_installation_scope_overrides_nested_inbound_context():
+    runtime_handler, _app, target_binding = make_handler()
+    caller_context = workspace_context().for_installation('legacy-caller')
+    token = runtime_handler._current_action_context.set(caller_context)
+    try:
+        assert runtime_handler.resolve_outbound_action_context(None) == caller_context
+        with runtime_handler.installation_scope(target_binding):
+            assert runtime_handler.resolve_outbound_action_context(None) == target_binding
+        with runtime_handler.installation_scope(None):
+            assert runtime_handler.resolve_outbound_action_context(None) is None
+        assert runtime_handler.resolve_outbound_action_context(None) == caller_context
+    finally:
+        runtime_handler._current_action_context.reset(token)
+
+
 def test_inbound_tenant_action_requires_complete_installation_envelope():
     runtime_handler, _app, installation_context = make_handler()
 
@@ -363,6 +378,44 @@ async def test_legacy_oss_worker_capability_remains_usable_after_identity_migrat
     assert response.code == 0
 
 
+@pytest.mark.asyncio
+async def test_legacy_oss_knowledge_file_reply_uses_complete_installation_binding():
+    runtime_handler, app, installation_context = make_handler()
+    app.deployment = SimpleNamespace(mode='oss')
+    setting = SimpleNamespace(
+        plugin_author='author-a',
+        plugin_name='plugin-a',
+        installation_uuid=installation_context.installation_uuid,
+        runtime_revision=installation_context.runtime_revision,
+        artifact_digest=installation_context.artifact_digest,
+    )
+    result = Mock()
+    result.first.return_value = setting
+    app.persistence_mgr.execute_async.return_value = result
+    app.rag_runtime_service = SimpleNamespace(
+        get_file_stream=AsyncMock(return_value=b'knowledge-file'),
+    )
+    runtime_handler.send_file = AsyncMock(return_value='knowledge-file-key')
+    legacy_context = workspace_context().for_installation(
+        installation_context.installation_uuid
+    )
+
+    response = await invoke_with_context(
+        runtime_handler,
+        legacy_context,
+        PluginToRuntimeAction.GET_KNOWLEDEGE_FILE_STREAM,
+        {'storage_path': 'knowledge/file.txt'},
+    )
+
+    assert response.code == 0
+    assert response.data == {'file_key': 'knowledge-file-key'}
+    runtime_handler.send_file.assert_awaited_once_with(
+        b'knowledge-file',
+        '',
+        action_context=installation_context,
+    )
+
+
 def test_installation_uuid_cannot_move_between_workspaces():
     runtime_handler, _app, binding = make_handler()
     moved = binding.model_copy(update={'workspace_uuid': 'workspace-b', 'runtime_revision': 2})