mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-17 03:04:20 +00:00
feat(agent-runner): add host run ledger primitives
This commit is contained in:
huanghuoguoguo
@@ -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
|
||||
Reference in New Issue
Block a user