fix(agent): harden runner integration and QA

tests/integration/persistence/test_workspace_migration.py

@@ -149,6 +149,35 @@ async def test_workspace_upgrade_is_idempotent_and_preserves_identifiers(legacy_
     assert workspace_uuid_after == workspace_uuid_before
 
 
+async def test_workspace_upgrade_repairs_ownerless_existing_local_workspace(legacy_engine):
+    await run_alembic_upgrade(legacy_engine, '0016_agent_workspace')
+    async with legacy_engine.begin() as conn:
+        owner_account_uuid = await conn.scalar(sa.text('SELECT uuid FROM users ORDER BY id LIMIT 1'))
+        workspace_uuid = await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
+        await conn.execute(sa.text('DELETE FROM workspace_memberships'))
+        await conn.execute(
+            sa.text('UPDATE workspaces SET created_by_account_uuid = NULL WHERE uuid = :workspace_uuid'),
+            {'workspace_uuid': workspace_uuid},
+        )
+
+    await run_alembic_upgrade(legacy_engine, 'head')
+
+    async with legacy_engine.connect() as conn:
+        workspace = (
+            await conn.execute(
+                sa.text('SELECT created_by_account_uuid FROM workspaces WHERE uuid = :workspace_uuid'),
+                {'workspace_uuid': workspace_uuid},
+            )
+        ).mappings().one()
+        membership = (await conn.execute(sa.text('SELECT * FROM workspace_memberships'))).mappings().one()
+
+    assert workspace['created_by_account_uuid'] == owner_account_uuid
+    assert membership['workspace_uuid'] == workspace_uuid
+    assert membership['account_uuid'] == owner_account_uuid
+    assert membership['role'] == 'owner'
+    assert membership['status'] == 'active'
+
+
 async def test_workspace_kernel_upgrade_downgrade_upgrade_round_trip(tmp_path):
     engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-round-trip.db"}')
     try:

tests/unit_tests/agent/test_registry.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from unittest.mock import AsyncMock
+
 import pytest
 
 from langbot.pkg.agent.runner.registry import AgentRunnerRegistry
@@ -210,6 +212,27 @@ class TestRegistryGet:
 
         assert exc_info.value.runner_id == 'plugin:notexist/unknown/default'
 
+    @pytest.mark.asyncio
+    async def test_get_refreshes_partial_startup_cache_on_miss(self):
+        """A runner initialized after early discovery should become available."""
+        ap = FakeApplication()
+        ap.plugin_connector.list_agent_runners = AsyncMock(
+            side_effect=ap.plugin_connector.list_agent_runners,
+        )
+        registry = AgentRunnerRegistry(ap)
+
+        await registry.list_runners(TEST_CONTEXT)
+        cache = registry._cache[('instance-test', 'workspace-test', 1)]
+        cache.pop('plugin:alice/my-agent/custom')
+
+        descriptor = await registry.get(
+            TEST_CONTEXT,
+            'plugin:alice/my-agent/custom',
+        )
+
+        assert descriptor.id == 'plugin:alice/my-agent/custom'
+        assert ap.plugin_connector.list_agent_runners.await_count == 2
+
     @pytest.mark.asyncio
     async def test_get_runner_with_bound_plugins_filter(self):
         """Get runner with bound plugins authorization."""
@@ -257,6 +280,27 @@ class TestRegistryMetadataForPipeline:
         assert stages[0]['config'][0]['type'] == 'string'
         assert stages[0]['config'][0]['id'] == 'plugin:alice/my-agent/custom.param1'
 
+    @pytest.mark.asyncio
+    async def test_metadata_refreshes_partial_startup_cache(self):
+        """Pipeline metadata should not preserve an early partial discovery."""
+        ap = FakeApplication()
+        ap.plugin_connector.list_agent_runners = AsyncMock(
+            side_effect=ap.plugin_connector.list_agent_runners,
+        )
+        registry = AgentRunnerRegistry(ap)
+
+        await registry.list_runners(TEST_CONTEXT)
+        cache = registry._cache[('instance-test', 'workspace-test', 1)]
+        cache.pop('plugin:alice/my-agent/custom')
+
+        options, _ = await registry.get_runner_metadata_for_pipeline(TEST_CONTEXT)
+
+        assert {item['name'] for item in options} == {
+            'plugin:langbot-team/LocalAgent/default',
+            'plugin:alice/my-agent/custom',
+        }
+        assert ap.plugin_connector.list_agent_runners.await_count == 2
+
 
 class TestDescriptorValidation:
     """Tests for descriptor validation."""

tests/unit_tests/pipeline/test_controller.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import asyncio
 from types import SimpleNamespace
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
 
 from langbot.pkg.api.http.context import ExecutionContext
 from langbot.pkg.agent.runner.errors import RunnerNotFoundError
 from langbot.pkg.pipeline.controller import Controller
+from langbot.pkg.pipeline.pool import QueryPool
 
 
 def make_app():
@@ -77,3 +79,74 @@ async def test_try_claim_steering_sets_pipeline_context_before_claiming():
     assert query.pipeline_config is pipeline.pipeline_entity.config
     assert query.variables['_pipeline_bound_plugins'] == ['test/runner']
     app.agent_run_orchestrator.try_claim_steering_from_query.assert_awaited_once_with(query)
+
+
+@pytest.mark.asyncio
+async def test_consumer_transfers_query_from_queue_to_running_task():
+    app = make_app()
+    app.query_pool = QueryPool()
+    session = SimpleNamespace(_semaphore=asyncio.Semaphore(1))
+    app.sess_mgr.get_session = AsyncMock(return_value=session)
+    app.persistence_mgr = SimpleNamespace(mode=SimpleNamespace(value='oss'))
+    app.workspace_service = SimpleNamespace(
+        get_execution_binding=AsyncMock(
+            return_value=SimpleNamespace(
+                instance_uuid='instance-test',
+                workspace_uuid='workspace-test',
+                placement_generation=1,
+            )
+        )
+    )
+    runtime_pipeline = SimpleNamespace(run=AsyncMock())
+    app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=runtime_pipeline)
+
+    worker_tasks = []
+    task_created = asyncio.Event()
+
+    def create_task(coro, **_kwargs):
+        task = asyncio.create_task(coro)
+        worker_tasks.append(task)
+        task_created.set()
+        return task
+
+    app.task_mgr = SimpleNamespace(create_task=create_task)
+
+    query = Mock()
+    query.query_id = 0
+    query.bot_uuid = 'bot-test'
+    query.pipeline_uuid = 'pipeline-test'
+    context = ExecutionContext(
+        instance_uuid='instance-test',
+        workspace_uuid='workspace-test',
+        placement_generation=1,
+    )
+    with patch('langbot.pkg.pipeline.pool.pipeline_query.Query', return_value=query):
+        query = await app.query_pool.add_query(
+            bot_uuid='bot-test',
+            launcher_type=Mock(),
+            launcher_id='launcher-test',
+            sender_id='sender-test',
+            message_event=Mock(),
+            message_chain=Mock(),
+            adapter=Mock(),
+            pipeline_uuid='pipeline-test',
+            execution_context=context,
+        )
+
+    controller = Controller(app)
+    consumer_task = asyncio.create_task(controller.consumer())
+    try:
+        await asyncio.wait_for(task_created.wait(), timeout=1)
+        await asyncio.wait_for(worker_tasks[0], timeout=1)
+    finally:
+        consumer_task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await consumer_task
+
+    runtime_pipeline.run.assert_awaited_once_with(query)
+    assert app.query_pool.queries == []
+    assert app.query_pool.cached_queries == {}
+    assert app.query_pool.active_query_count_by_workspace == {}
+    assert session._semaphore._value == 1
+    assert controller.semaphore._value == 10
+    app.logger.error.assert_not_called()

tests/unit_tests/platform/test_websocket_adapter_attachments.py

@@ -2,9 +2,9 @@
 
 The web debug client uploads Image / Voice / File components carrying a storage
 key in ``path``. This helper resolves each to a base64 data URI (so multimodal
-LLM input and the Box sandbox inbox have usable bytes), then deletes the
-consumed storage object and clears ``path``. Covers mimetype selection per
-type and fail-closed error handling.
+LLM input and the Box sandbox inbox have usable bytes) while retaining the key
+for browser history. Covers mimetype selection per type and fail-closed error
+handling.
 """
 
 from __future__ import annotations
@@ -52,7 +52,7 @@ def _make_adapter(load_return=b'hello', load_side_effect=None):
 
 
 @pytest.mark.asyncio
-async def test_image_jpeg_mimetype_and_cleanup():
+async def test_image_jpeg_mimetype_and_retained_storage_key():
     adapter, storage_mgr, _ = _make_adapter(load_return=b'\xff\xd8\xff')
     path = f'{_UPLOAD_PREFIX}photo.jpg'
     chain = [{'type': 'Image', 'path': path}]
@@ -61,12 +61,24 @@ async def test_image_jpeg_mimetype_and_cleanup():
 
     expected_b64 = base64.b64encode(b'\xff\xd8\xff').decode('utf-8')
     assert chain[0]['base64'] == f'data:image/jpeg;base64,{expected_b64}'
-    assert chain[0]['path'] == ''  # consumed
-    storage_mgr.delete_scoped_object_key.assert_awaited_once_with(
-        _CONTEXT,
-        path,
-        expected_owner_type='upload_image',
-    )
+    assert chain[0]['path'] == path
+    storage_mgr.delete_scoped_object_key.assert_not_awaited()
+
+
+def test_history_retains_storage_key_without_large_base64_payload():
+    path = f'{_UPLOAD_PREFIX}photo.jpg'
+    chain = [
+        {
+            'type': 'Image',
+            'path': path,
+            'base64': 'data:image/jpeg;base64,large-payload',
+        }
+    ]
+
+    history = WebSocketAdapter._history_message_chain(chain)
+
+    assert history == [{'type': 'Image', 'path': path, 'base64': ''}]
+    assert chain[0]['base64'] == 'data:image/jpeg;base64,large-payload'
 
 
 @pytest.mark.asyncio

tests/unit_tests/platform/test_websocket_session_isolation.py

@@ -425,7 +425,7 @@ async def test_attachment_key_must_belong_to_connection_upload_scope():
     await adapter._process_image_components(connection, message_chain)
 
     assert message_chain[0]['base64'].startswith('data:image/png;base64,')
-    assert message_chain[0]['path'] == ''
+    assert message_chain[0]['path'] == 'v1/current/upload_image/key.png'
     storage_mgr.scoped_prefix.assert_called_once_with(
         connection.execution_context,
         owner_type='upload_image',
@@ -439,11 +439,7 @@ async def test_attachment_key_must_belong_to_connection_upload_scope():
         'v1/current/upload_image/key.png',
         expected_owner_type='upload_image',
     )
-    storage_mgr.delete_scoped_object_key.assert_awaited_once_with(
-        connection.execution_context,
-        'v1/current/upload_image/key.png',
-        expected_owner_type='upload_image',
-    )
+    storage_mgr.delete_scoped_object_key.assert_not_awaited()
 
     with pytest.raises(ValueError, match='does not belong'):
         await adapter._process_image_components(

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})

tests/unit_tests/provider/test_mcp_box_integration.py

@@ -1004,7 +1004,9 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo
 
     ap = _make_ap()
     ap.box_service.available = True
+    shared_workspace = tmp_path / 'shared-box-workspace' / 'tenants' / 'workspace-a'
     ap.box_service.default_workspace = str(tmp_path / 'shared-box-workspace')
+    ap.box_service.workspace_host_path = Mock(return_value=str(shared_workspace))
     ap.box_service.create_session = AsyncMock(return_value={})
     ap.box_service.build_spec = Mock(return_value='validated-spec')
     ap.box_service.client = SimpleNamespace(
@@ -1053,8 +1055,9 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo
     assert ap.box_service.build_spec.call_args.kwargs.get('skip_host_mount_validation', False) is False
     assert ap.box_service.build_spec.call_args.args[0]['host_path'] == str(host_path)
 
-    staged_file = tmp_path / 'shared-box-workspace' / '.mcp' / 'u1' / 'workspace' / 'server.py'
+    staged_file = shared_workspace / '.mcp' / 'u1' / 'workspace' / 'server.py'
     assert staged_file.read_text(encoding='utf-8') == 'print("hello")\n'
+    ap.box_service.workspace_host_path.assert_called_with(session.execution_context)
 
     assert ap.box_service.start_managed_process.await_args.args[0] == session.execution_context
     process_payload = ap.box_service.start_managed_process.await_args.args[2]