test(agent): harden runner persistence coverage

tests/integration/persistence/test_migrations.py

@@ -10,10 +10,12 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q
 from __future__ import annotations
 
 import pytest
+from alembic.script import ScriptDirectory
 from sqlalchemy.ext.asyncio import create_async_engine
 
 from langbot.pkg.entity.persistence.base import Base
 from langbot.pkg.persistence.alembic_runner import (
+    _ALEMBIC_DIR,
     run_alembic_upgrade,
     run_alembic_stamp,
     get_alembic_current,
@@ -38,6 +40,19 @@ async def sqlite_engine(sqlite_db_url):
     await engine.dispose()
 
 
+def alembic_head_revision() -> str:
+    """Return the repository's current Alembic head revision."""
+    return ScriptDirectory.from_config(_alembic_script_config()).get_current_head()
+
+
+def _alembic_script_config():
+    from alembic.config import Config
+
+    cfg = Config()
+    cfg.set_main_option('script_location', _ALEMBIC_DIR)
+    return cfg
+
+
 class TestSQLiteMigrationBaseline:
     """Tests for baseline stamp workflow."""
 
@@ -103,8 +118,7 @@ class TestSQLiteMigrationUpgrade:
         # Verify revision
         rev = await get_alembic_current(sqlite_engine)
         assert rev is not None, "Expected a revision after upgrade"
-        # Head should be the latest migration
-        assert rev.startswith('0003'), f"Expected head to be 0003_*, got {rev}"
+        assert rev == alembic_head_revision()
 
     @pytest.mark.asyncio
     async def test_upgrade_idempotent(self, sqlite_engine):
@@ -248,4 +262,4 @@ class TestSQLiteMigrationGetCurrent:
         await run_alembic_stamp(sqlite_engine, '0001_baseline')
 
         rev = await get_alembic_current(sqlite_engine)
-        assert rev == '0001_baseline'
+        assert rev == '0001_baseline'

tests/integration/pipeline/test_full_flow.py

@@ -14,14 +14,14 @@ from __future__ import annotations
 import pytest
 import asyncio
 from unittest.mock import AsyncMock, Mock
-import sys
 
 from tests.factories import FakeApp, text_query, mock_platform_adapter
-from tests.factories.provider import FakeProvider
 from tests.factories.platform import FakePlatform
+from tests.factories.message import text_chain
 
 
 pytestmark = pytest.mark.integration
+DEFAULT_RUNNER_ID = 'plugin:langbot/local-agent/default'
 
 
 # ============== FIXTURE FOR SYS.MODULES ISOLATION ==============
@@ -47,10 +47,6 @@ def mock_circular_import_chain():
     # Mock core.app - Application class is referenced but not instantiated
     mock_core_app = Mock()
 
-    # Mock provider.runner with preregistered_runners list
-    mock_runner = Mock()
-    mock_runner.preregistered_runners = []  # Will be populated in tests
-
     # Mock utils.importutil - prevents auto-import of runners
     mock_importutil = Mock()
     mock_importutil.import_modules_in_pkg = lambda pkg: None
@@ -74,7 +70,7 @@ def mock_circular_import_chain():
         mocks={
             'langbot.pkg.core.entities': mock_core_entities,
             'langbot.pkg.core.app': mock_core_app,
-            'langbot.pkg.provider.runner': mock_runner,
+            'langbot.pkg.provider.runner': Mock(preregistered_runners=[]),
             'langbot.pkg.utils.importutil': mock_importutil,
             'langbot.pkg.pipeline.controller': Mock(),
             'langbot.pkg.pipeline.pipelinemgr': Mock(),
@@ -104,48 +100,23 @@ def mock_circular_import_chain():
 # ============== FAKE RUNNER ==============
 
 class FakeRunner:
-    """Minimal fake runner class for pipeline integration tests.
-
-    Note: preregistered_runners expects a CLASS, not an instance.
-    The handler calls runner_cls(self.ap, query.pipeline_config) to instantiate.
-    """
+    """Minimal fake runner behavior for the orchestrator-backed pipeline tests."""
 
     name = 'local-agent'
 
-    def __init__(self, app=None, config=None):
-        self.app = app
-        self.config = config or {}
-        self._provider = FakeProvider()
-        # Instance-level configuration set via class attribute
-        self._response_text = "fake response"
-        self._raise_error = None
+    def __init__(self, response_text: str = "fake response", error: Exception | None = None):
+        self._response_text = response_text
+        self._raise_error = error
 
     @classmethod
-    def returns(cls, text: str):
-        """Create a runner class configured to return specific text."""
-        # We create a subclass with configured response
-        class ConfiguredRunner(cls):
-            name = cls.name
-            _response_text = text
-            _raise_error = None
-
-            def __init__(self, app=None, config=None):
-                super().__init__(app, config)
-                self._response_text = text
-        return ConfiguredRunner
+    def returns(cls, text: str) -> "FakeRunner":
+        """Create a fake runner configured to return specific text."""
+        return cls(response_text=text)
 
     @classmethod
-    def raises(cls, error: Exception):
+    def raises(cls, error: Exception) -> "FakeRunner":
         """Create a runner class configured to raise an error."""
-        class ConfiguredRunner(cls):
-            name = cls.name
-            _response_text = None
-            _raise_error = error
-
-            def __init__(self, app=None, config=None):
-                super().__init__(app, config)
-                self._raise_error = error
-        return ConfiguredRunner
+        return cls(error=error)
 
     async def run(self, query):
         """Run the fake provider and yield messages."""
@@ -159,6 +130,22 @@ class FakeRunner:
         yield Message(role='assistant', content=self._response_text)
 
 
+class FakeAgentRunOrchestrator:
+    """Adapter that exposes FakeRunner through the current AgentRunOrchestrator surface."""
+
+    def __init__(self, runner: FakeRunner | None = None):
+        self.runner = runner or FakeRunner()
+        self.queries = []
+
+    async def run_from_query(self, query):
+        self.queries.append(query)
+        async for result in self.runner.run(query):
+            yield result
+
+    def resolve_runner_id_for_telemetry(self, query):
+        return DEFAULT_RUNNER_ID
+
+
 # ============== PIPELINE APP FIXTURE ==============
 
 @pytest.fixture
@@ -222,6 +209,7 @@ def pipeline_app():
 
     # Survey mock
     app.survey = None
+    app.agent_run_orchestrator = FakeAgentRunOrchestrator()
 
     return app
 
@@ -235,11 +223,10 @@ def fake_platform_adapter():
 
 
 @pytest.fixture
-def set_fake_runner():
-    """Factory fixture to set a fake runner CLASS in preregistered_runners."""
-    def _set_runner(runner_cls):
-        # preregistered_runners expects a list of runner classes
-        sys.modules['langbot.pkg.provider.runner'].preregistered_runners = [runner_cls]
+def set_fake_runner(pipeline_app):
+    """Factory fixture to set fake runner behavior on the orchestrator surface."""
+    def _set_runner(runner: FakeRunner):
+        pipeline_app.agent_run_orchestrator.runner = runner
     return _set_runner
 
 
@@ -249,11 +236,13 @@ def create_minimal_pipeline_config():
     """Create minimal pipeline configuration for tests."""
     return {
         'ai': {
-            'runner': {'runner': 'local-agent', 'expire-time': None},
-            'local-agent': {
-                'model': {'primary': 'test-model-uuid', 'fallbacks': []},
-                'prompt': 'default',
-                'knowledge-bases': [],
+            'runner': {'id': DEFAULT_RUNNER_ID, 'expire-time': None},
+            'runner_config': {
+                DEFAULT_RUNNER_ID: {
+                    'model': {'primary': 'test-model-uuid', 'fallbacks': []},
+                    'prompt': [{'role': 'system', 'content': 'default'}],
+                    'knowledge-bases': [],
+                },
             },
         },
         'output': {
@@ -396,7 +385,7 @@ class TestProcessorStage:
         adapter, platform = fake_platform_adapter
 
         # Set fake runner that returns pong
-        fake_runner = FakeRunner().returns("LANGBOT_FAKE_PONG")
+        fake_runner = FakeRunner.returns("LANGBOT_FAKE_PONG")
         set_fake_runner(fake_runner)
 
         # Create query
@@ -502,7 +491,7 @@ class TestRunnerExceptionFlow:
         adapter, platform = fake_platform_adapter
 
         # Set fake runner that raises exception
-        fake_runner = FakeRunner().raises(ValueError("API Error: rate limit exceeded"))
+        fake_runner = FakeRunner.raises(ValueError("API Error: rate limit exceeded"))
         set_fake_runner(fake_runner)
 
         # Create query with exception handling config
@@ -541,7 +530,7 @@ class TestRunnerExceptionFlow:
         adapter, platform = fake_platform_adapter
 
         # Set fake runner that raises specific exception
-        fake_runner = FakeRunner().raises(RuntimeError("Custom runtime error"))
+        fake_runner = FakeRunner.raises(RuntimeError("Custom runtime error"))
         set_fake_runner(fake_runner)
 
         # Create query with show-error mode
@@ -578,7 +567,7 @@ class TestRunnerExceptionFlow:
         adapter, platform = fake_platform_adapter
 
         # Set fake runner that raises exception
-        fake_runner = FakeRunner().raises(Exception("Hidden error"))
+        fake_runner = FakeRunner.raises(Exception("Hidden error"))
         set_fake_runner(fake_runner)
 
         # Create query with hide mode
@@ -666,7 +655,7 @@ class TestStageChainIntegration:
         adapter, platform = fake_platform_adapter
 
         # Set fake runner
-        fake_runner = FakeRunner().returns("LANGBOT_FAKE_PONG")
+        fake_runner = FakeRunner.returns("LANGBOT_FAKE_PONG")
         set_fake_runner(fake_runner)
 
         # Create query
@@ -710,7 +699,6 @@ class TestStageChainIntegration:
         assert len(results) >= 1
 
         # Build resp_message_chain from resp_messages
-        from tests.factories.message import text_chain
         for resp_msg in query.resp_messages:
             if resp_msg.content:
                 query.resp_message_chain.append(text_chain(resp_msg.content))
@@ -775,4 +763,4 @@ class TestStageChainIntegration:
         assert results[0].result_type == entities.ResultType.INTERRUPT
 
         # Chain stops here - no resp_messages
-        assert len(query.resp_messages) == 0
+        assert len(query.resp_messages) == 0

tests/unit_tests/agent/test_orchestrator_artifact.py

@@ -8,8 +8,7 @@ from langbot.pkg.agent.runner.orchestrator import (
     AgentRunOrchestrator,
     MAX_ARTIFACT_INLINE_BYTES,
 )
-from langbot.pkg.agent.runner.descriptor import AgentRunnerDescriptor
-from langbot.pkg.agent.runner.host_models import AgentEventEnvelope, AgentBinding
+from langbot.pkg.agent.runner.host_models import AgentEventEnvelope
 from langbot.pkg.agent.runner.errors import RunnerProtocolError
 from langbot_plugin.api.entities.builtin.agent_runner.input import AgentInput
 from langbot_plugin.api.entities.builtin.agent_runner.event import ActorContext

tests/unit_tests/agent/test_state_api_auth.py

@@ -20,14 +20,12 @@ Authorization rules:
 from __future__ import annotations
 
 import pytest
-import asyncio
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch
 from sqlalchemy.ext.asyncio import create_async_engine
 
-from langbot.pkg.agent.runner.session_registry import AgentRunSessionRegistry, get_session_registry
+from langbot.pkg.agent.runner.session_registry import AgentRunSessionRegistry
 from langbot.pkg.agent.runner.persistent_state_store import PersistentStateStore, reset_persistent_state_store
 from langbot.pkg.plugin.handler import RuntimeConnectionHandler
-from langbot_plugin.runtime.io.connection import Connection
 from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
 
 # Import shared test fixtures
@@ -72,7 +70,7 @@ async def persistent_store(db_engine):
 
     # Create the table
     from langbot.pkg.entity.persistence.agent_runner_state import AgentRunnerState
-    from sqlalchemy import text
+
     async with db_engine.begin() as conn:
         await conn.run_sync(AgentRunnerState.__table__.create, checkfirst=True)