feat(agent-runner): add host run ledger primitives

tests/unit_tests/agent/test_run_ledger_api_auth.py

@@ -0,0 +1,490 @@
+"""Tests for AgentRunner run ledger pull API authorization."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from langbot.pkg.agent.runner.run_ledger_store import RunLedgerStore
+from langbot.pkg.agent.runner.session_registry import AgentRunSessionRegistry
+from langbot.pkg.entity.persistence import agent_run as agent_run_model
+from langbot.pkg.entity.persistence.base import Base
+from langbot.pkg.plugin.handler import RuntimeConnectionHandler
+from langbot_plugin.api.entities.builtin.agent_runner.run_ledger import (
+    AgentRun,
+    AgentRunEvent,
+    RunEventPage,
+    RunPage,
+)
+from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
+
+from .conftest import make_resources
+
+
+class FakeConnection:
+    pass
+
+
+class FakeApplication:
+    def __init__(self, db_engine):
+        self.logger = MagicMock()
+        self.persistence_mgr = MagicMock()
+        self.persistence_mgr.get_db_engine = MagicMock(return_value=db_engine)
+
+
+@pytest.fixture
+def session_registry(monkeypatch):
+    registry = AgentRunSessionRegistry()
+    monkeypatch.setattr(
+        'langbot.pkg.agent.runner.session_registry._global_registry',
+        registry,
+    )
+    return registry
+
+
+@pytest.fixture
+async def db_engine():
+    engine = create_async_engine('sqlite+aiosqlite:///:memory:')
+    assert agent_run_model.AgentRun.__tablename__ == 'agent_run'
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+    yield engine
+    await engine.dispose()
+
+
+def _handler(db_engine):
+    async def fake_disconnect():
+        return True
+
+    fake_app = FakeApplication(db_engine)
+    return RuntimeConnectionHandler(FakeConnection(), fake_disconnect, fake_app)
+
+
+async def _register_session(
+    session_registry,
+    *,
+    run_id='run_1',
+    conversation_id='conv_1',
+    available_apis=None,
+):
+    await session_registry.register(
+        run_id=run_id,
+        runner_id='plugin:test/runner/default',
+        query_id=None,
+        plugin_identity='test/runner',
+        resources=make_resources(),
+        conversation_id=conversation_id,
+        bot_id='bot_1',
+        workspace_id='workspace_1',
+        thread_id=None,
+        available_apis=available_apis or {},
+    )
+
+
+async def _create_run(
+    db_engine,
+    *,
+    run_id='run_1',
+    conversation_id='conv_1',
+    bot_id='bot_1',
+    workspace_id='workspace_1',
+    thread_id=None,
+    plugin_identity='test/runner',
+    available_apis=None,
+):
+    store = RunLedgerStore(db_engine)
+    await store.create_run(
+        run_id=run_id,
+        event_id='evt_1',
+        binding_id='binding_1',
+        runner_id='plugin:test/runner/default',
+        conversation_id=conversation_id,
+        bot_id=bot_id,
+        workspace_id=workspace_id,
+        thread_id=thread_id,
+        authorization={
+            'runner_id': 'plugin:test/runner/default',
+            'binding_id': 'binding_1',
+            'plugin_identity': plugin_identity,
+            'resources': make_resources(),
+            'available_apis': available_apis or {},
+            'conversation_id': conversation_id,
+            'bot_id': bot_id,
+            'workspace_id': workspace_id,
+            'thread_id': thread_id,
+            'state_policy': {'enable_state': True, 'state_scopes': ['conversation', 'actor']},
+            'state_context': {},
+        },
+    )
+    await store.append_event(
+        run_id=run_id,
+        sequence=1,
+        event_type='message.completed',
+        data={'message': {'role': 'assistant', 'content': 'ok'}},
+    )
+
+
+@pytest.mark.asyncio
+async def test_run_get_returns_current_run(session_registry, db_engine):
+    await _register_session(session_registry, available_apis={'run_get': True})
+    await _create_run(db_engine)
+    handler = _handler(db_engine)
+    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]
+
+    result = await run_get(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code == 0
+    run = AgentRun.model_validate(result.data)
+    assert run.run_id == 'run_1'
+    assert run.status == 'running'
+
+
+@pytest.mark.asyncio
+async def test_run_list_rejects_cross_conversation(session_registry, db_engine):
+    await _register_session(session_registry, available_apis={'run_list': True})
+    handler = _handler(db_engine)
+    run_list = handler.actions[PluginToRuntimeAction.RUN_LIST.value]
+
+    result = await run_list(
+        {
+            'run_id': 'run_1',
+            'conversation_id': 'conv_other',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code != 0
+    assert 'not accessible' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_run_list_returns_scoped_runs(session_registry, db_engine):
+    await _register_session(session_registry, available_apis={'run_list': True})
+    await _create_run(db_engine)
+    await _create_run(db_engine, run_id='run_other', conversation_id='conv_other')
+    handler = _handler(db_engine)
+    run_list = handler.actions[PluginToRuntimeAction.RUN_LIST.value]
+
+    result = await run_list(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code == 0
+    page = RunPage.model_validate(result.data)
+    assert [run.run_id for run in page.items] == ['run_1']
+
+
+@pytest.mark.asyncio
+async def test_run_events_page_returns_events(session_registry, db_engine):
+    await _register_session(session_registry, available_apis={'run_events_page': True})
+    await _create_run(db_engine)
+    handler = _handler(db_engine)
+    run_events_page = handler.actions[PluginToRuntimeAction.RUN_EVENTS_PAGE.value]
+
+    result = await run_events_page(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code == 0
+    page = RunEventPage.model_validate(result.data)
+    assert [item.sequence for item in page.items] == [1]
+    assert page.items[0].type == 'message.completed'
+
+
+@pytest.mark.asyncio
+async def test_run_get_uses_persistent_authorization_after_session_expired(db_engine):
+    await _create_run(db_engine, available_apis={'run_get': True})
+    handler = _handler(db_engine)
+    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]
+
+    result = await run_get(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code == 0
+    run = AgentRun.model_validate(result.data)
+    assert run.run_id == 'run_1'
+
+
+@pytest.mark.asyncio
+async def test_persistent_run_get_rejects_cross_scope(db_engine):
+    await _create_run(db_engine, available_apis={'run_get': True})
+    await _create_run(
+        db_engine,
+        run_id='run_other',
+        conversation_id='conv_other',
+        available_apis={'run_get': True},
+    )
+    handler = _handler(db_engine)
+    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]
+
+    result = await run_get(
+        {
+            'run_id': 'run_1',
+            'target_run_id': 'run_other',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code != 0
+    assert 'not accessible' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_persistent_run_get_requires_capability(db_engine):
+    await _create_run(db_engine, available_apis={'run_get': False})
+    handler = _handler(db_engine)
+    run_get = handler.actions[PluginToRuntimeAction.RUN_GET.value]
+
+    result = await run_get(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code != 0
+    assert 'not authorized' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_persistent_authorization_does_not_reopen_artifact_api(db_engine):
+    await _create_run(db_engine, available_apis={'artifact_metadata': True})
+    handler = _handler(db_engine)
+    artifact_metadata = handler.actions[PluginToRuntimeAction.ARTIFACT_METADATA.value]
+
+    result = await artifact_metadata(
+        {
+            'run_id': 'run_1',
+            'artifact_id': 'artifact_1',
+            'caller_plugin_identity': 'test/runner',
+        }
+    )
+
+    assert result.code != 0
+    assert 'not found or expired' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_run_cancel_basic_path(session_registry, db_engine):
+    await _register_session(session_registry, available_apis={'run_cancel': True})
+    await _create_run(db_engine)
+    handler = _handler(db_engine)
+    run_cancel = handler.actions[PluginToRuntimeAction.RUN_CANCEL.value]
+
+    result = await run_cancel(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+            'reason': 'user requested',
+        }
+    )
+
+    assert result.code == 0
+    run = AgentRun.model_validate(result.data)
+    assert run.run_id == 'run_1'
+    assert run.cancel_requested_at is not None
+    assert run.status_reason == 'user requested'
+
+
+@pytest.mark.asyncio
+async def test_run_append_result_basic_path(session_registry, db_engine):
+    await _register_session(session_registry, available_apis={'run_append_result': True})
+    await _create_run(db_engine)
+    handler = _handler(db_engine)
+    run_append_result = handler.actions[PluginToRuntimeAction.RUN_APPEND_RESULT.value]
+
+    result = await run_append_result(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+            'sequence': 2,
+            'result': {
+                'type': 'message.delta',
+                'data': {'delta': 'hello'},
+                'usage': {'output_tokens': 1},
+            },
+        }
+    )
+
+    assert result.code == 0
+    event = AgentRunEvent.model_validate(result.data)
+    assert event.run_id == 'run_1'
+    assert event.sequence == 2
+    assert event.type == 'message.delta'
+    assert event.data == {'delta': 'hello'}
+    assert event.usage == {'output_tokens': 1}
+
+
+@pytest.mark.asyncio
+async def test_run_finalize_basic_path(session_registry, db_engine):
+    await _register_session(session_registry, available_apis={'run_finalize': True})
+    await _create_run(db_engine)
+    handler = _handler(db_engine)
+    run_finalize = handler.actions[PluginToRuntimeAction.RUN_FINALIZE.value]
+
+    result = await run_finalize(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+            'status': 'completed',
+            'status_reason': 'done',
+            'usage': {'total_tokens': 3},
+        }
+    )
+
+    assert result.code == 0
+    run = AgentRun.model_validate(result.data)
+    assert run.status == 'completed'
+    assert run.status_reason == 'done'
+    assert run.finished_at is not None
+    assert run.usage == {'total_tokens': 3}
+
+
+@pytest.mark.asyncio
+async def test_runtime_register_heartbeat_and_list_actions(session_registry, db_engine):
+    await _register_session(
+        session_registry,
+        available_apis={
+            'runtime_register': True,
+            'runtime_heartbeat': True,
+            'runtime_list': True,
+        },
+    )
+    handler = _handler(db_engine)
+    runtime_register = handler.actions[PluginToRuntimeAction.RUNTIME_REGISTER.value]
+    runtime_heartbeat = handler.actions[PluginToRuntimeAction.RUNTIME_HEARTBEAT.value]
+    runtime_list = handler.actions[PluginToRuntimeAction.RUNTIME_LIST.value]
+
+    registered = await runtime_register(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+            'runtime_id': 'runtime_1',
+            'display_name': 'Runtime 1',
+            'capabilities': {'runner': True},
+            'labels': {'region': 'test'},
+            'metadata': {'slots': 2},
+        }
+    )
+
+    assert registered.code == 0
+    assert registered.data['runtime_id'] == 'runtime_1'
+    assert registered.data['capabilities'] == {'runner': True}
+
+    heartbeat = await runtime_heartbeat(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+            'runtime_id': 'runtime_1',
+            'capabilities': {'runner': True, 'stream': True},
+            'labels': {'region': 'test'},
+            'metadata': {'active_runs': 1},
+        }
+    )
+
+    assert heartbeat.code == 0
+    assert heartbeat.data['capabilities'] == {'runner': True, 'stream': True}
+    assert heartbeat.data['metadata'] == {'slots': 2, 'active_runs': 1}
+
+    page = await runtime_list(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+            'statuses': ['online'],
+            'labels': {'region': 'test'},
+        }
+    )
+
+    assert page.code == 0
+    assert [item['runtime_id'] for item in page.data['items']] == ['runtime_1']
+
+
+@pytest.mark.asyncio
+async def test_run_claim_renew_and_release_actions(session_registry, db_engine):
+    await _register_session(
+        session_registry,
+        available_apis={
+            'run_claim': True,
+            'run_renew_claim': True,
+            'run_release_claim': True,
+        },
+    )
+    await RunLedgerStore(db_engine).create_run(
+        run_id='queued_run',
+        event_id='evt_queued',
+        binding_id='binding_1',
+        runner_id='plugin:test/runner/default',
+        conversation_id='conv_1',
+        bot_id='bot_1',
+        workspace_id='workspace_1',
+        status='queued',
+        queue_name='default',
+        priority=5,
+    )
+    handler = _handler(db_engine)
+    run_claim = handler.actions[PluginToRuntimeAction.RUN_CLAIM.value]
+    run_renew_claim = handler.actions[PluginToRuntimeAction.RUN_RENEW_CLAIM.value]
+    run_release_claim = handler.actions[PluginToRuntimeAction.RUN_RELEASE_CLAIM.value]
+
+    claimed = await run_claim(
+        {
+            'run_id': 'run_1',
+            'caller_plugin_identity': 'test/runner',
+            'runtime_id': 'runtime_1',
+            'queue_name': 'default',
+            'lease_seconds': 30,
+        }
+    )
+
+    assert claimed.code == 0
+    assert claimed.data['run_id'] == 'queued_run'
+    assert claimed.data['status'] == 'claimed'
+    assert claimed.data['claimed_by_runtime_id'] == 'runtime_1'
+    claim_token = claimed.data['claim_token']
+
+    renewed = await run_renew_claim(
+        {
+            'run_id': 'run_1',
+            'target_run_id': 'queued_run',
+            'caller_plugin_identity': 'test/runner',
+            'runtime_id': 'runtime_1',
+            'claim_token': claim_token,
+            'lease_seconds': 60,
+        }
+    )
+
+    assert renewed.code == 0
+    assert renewed.data['claim_token'] == claim_token
+
+    released = await run_release_claim(
+        {
+            'run_id': 'run_1',
+            'target_run_id': 'queued_run',
+            'caller_plugin_identity': 'test/runner',
+            'runtime_id': 'runtime_1',
+            'claim_token': claim_token,
+            'reason': 'done with lease',
+        }
+    )
+
+    assert released.code == 0
+    assert released.data['status'] == 'queued'
+    assert released.data['claimed_by_runtime_id'] is None
+    assert released.data['claim_token'] is None

tests/unit_tests/agent/test_run_ledger_store.py

@@ -0,0 +1,167 @@
+"""Tests for RunLedgerStore host primitives."""
+
+from __future__ import annotations
+
+import datetime
+
+import pytest
+import sqlalchemy
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import sessionmaker
+
+from langbot.pkg.agent.runner.run_ledger_store import RunLedgerStore
+from langbot.pkg.entity.persistence.agent_run import AgentRun
+from langbot.pkg.entity.persistence.base import Base
+
+
+UTC = datetime.timezone.utc
+
+
+@pytest.fixture
+async def db_engine(tmp_path):
+    db_path = tmp_path / 'run_ledger_store.db'
+    engine = create_async_engine(f'sqlite+aiosqlite:///{db_path}', echo=False)
+
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+    yield engine
+
+    await engine.dispose()
+
+
+@pytest.fixture
+def store(db_engine):
+    return RunLedgerStore(db_engine)
+
+
+@pytest.mark.asyncio
+async def test_create_queued_run_claim_renew_release(store):
+    run = await store.create_run(
+        run_id='run-queued',
+        event_id='evt-1',
+        binding_id='binding-1',
+        runner_id='runner-a',
+        status='queued',
+        queue_name='default',
+        priority=10,
+        requested_runtime_id='runtime-a',
+    )
+
+    assert run['status'] == 'queued'
+    assert run['started_at'] is None
+    assert run['queue_name'] == 'default'
+    assert run['priority'] == 10
+    assert run['requested_runtime_id'] == 'runtime-a'
+
+    assert await store.claim_next_run(runtime_id='runtime-b', queue_name='default') is None
+
+    claimed = await store.claim_next_run(runtime_id='runtime-a', queue_name='default', lease_seconds=30)
+
+    assert claimed is not None
+    assert claimed['run_id'] == 'run-queued'
+    assert claimed['status'] == 'claimed'
+    assert claimed['claimed_by_runtime_id'] == 'runtime-a'
+    assert claimed['claim_token']
+    assert claimed['dispatch_attempts'] == 1
+    assert claimed['claim_lease_expires_at'] is not None
+    assert claimed['last_claimed_at'] is not None
+
+    token = claimed['claim_token']
+    assert await store.renew_claim(run_id='run-queued', claim_token='wrong-token') is None
+
+    renewed = await store.renew_claim(run_id='run-queued', claim_token=token, lease_seconds=90)
+
+    assert renewed is not None
+    assert renewed['claim_token'] == token
+    assert renewed['claim_lease_expires_at'] >= claimed['claim_lease_expires_at']
+
+    released = await store.release_claim(
+        run_id='run-queued',
+        claim_token=token,
+        status='queued',
+        status_reason='runtime released capacity',
+    )
+
+    assert released is not None
+    assert released['status'] == 'queued'
+    assert released['status_reason'] == 'runtime released capacity'
+    assert released['claimed_by_runtime_id'] is None
+    assert released['claim_token'] is None
+    assert released['claim_lease_expires_at'] is None
+    assert released['dispatch_attempts'] == 1
+
+
+@pytest.mark.asyncio
+async def test_expired_claim_can_be_reclaimed(store, db_engine):
+    await store.create_run(
+        run_id='run-expired',
+        event_id='evt-2',
+        binding_id='binding-1',
+        runner_id='runner-a',
+        status='queued',
+        queue_name='default',
+    )
+    first_claim = await store.claim_next_run(runtime_id='runtime-a', queue_name='default', lease_seconds=60)
+    assert first_claim is not None
+
+    session_factory = sessionmaker(db_engine, class_=AsyncSession, expire_on_commit=False)
+    async with session_factory() as session:
+        await session.execute(
+            sqlalchemy.update(AgentRun)
+            .where(AgentRun.run_id == 'run-expired')
+            .values(claim_lease_expires_at=datetime.datetime.now(UTC) - datetime.timedelta(seconds=1))
+        )
+        await session.commit()
+
+    reclaimed = await store.claim_next_run(runtime_id='runtime-b', queue_name='default', lease_seconds=60)
+
+    assert reclaimed is not None
+    assert reclaimed['run_id'] == 'run-expired'
+    assert reclaimed['claimed_by_runtime_id'] == 'runtime-b'
+    assert reclaimed['claim_token'] != first_claim['claim_token']
+    assert reclaimed['dispatch_attempts'] == 2
+
+
+@pytest.mark.asyncio
+async def test_runtime_register_heartbeat_list_and_mark_stale(store):
+    registered = await store.register_runtime(
+        runtime_id='runtime-a',
+        display_name='Runtime A',
+        endpoint='http://runtime-a',
+        version='1.0.0',
+        capabilities={'stream': True},
+        labels={'region': 'test'},
+        metadata={'slot_count': 2},
+        heartbeat_deadline_seconds=30,
+    )
+
+    assert registered['runtime_id'] == 'runtime-a'
+    assert registered['status'] == 'online'
+    assert registered['display_name'] == 'Runtime A'
+    assert registered['capabilities'] == {'stream': True}
+    assert registered['labels'] == {'region': 'test'}
+    assert registered['metadata'] == {'slot_count': 2}
+    assert registered['last_heartbeat_at'] is not None
+    assert registered['heartbeat_deadline_at'] is not None
+
+    heartbeat = await store.heartbeat_runtime(
+        runtime_id='runtime-a',
+        metadata={'active_runs': 1},
+        heartbeat_deadline_seconds=30,
+    )
+
+    assert heartbeat is not None
+    assert heartbeat['metadata'] == {'slot_count': 2, 'active_runs': 1}
+
+    runtimes = await store.list_runtimes(statuses=['online'])
+    assert [runtime['runtime_id'] for runtime in runtimes] == ['runtime-a']
+
+    stale = await store.mark_stale_runtimes(
+        now=datetime.datetime.now(UTC) + datetime.timedelta(seconds=31),
+    )
+
+    assert [runtime['runtime_id'] for runtime in stale] == ['runtime-a']
+    assert stale[0]['status'] == 'stale'
+    assert (await store.get_runtime('runtime-a'))['status'] == 'stale'