mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-20 20:44:21 +00:00
huanghuoguoguo
431 lines
14 KiB
Python
431 lines
14 KiB
Python
"""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 'claim_token' not in renewed
|
|
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 'claim_token' not in released
|
|
assert released['claim_lease_expires_at'] is None
|
|
assert released['dispatch_attempts'] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_claim_next_run_applies_scope_filters(store):
|
|
await store.create_run(
|
|
run_id='run-other-runner',
|
|
event_id='evt-other-runner',
|
|
binding_id='binding-1',
|
|
runner_id='runner-b',
|
|
conversation_id='conv-a',
|
|
bot_id='bot-a',
|
|
workspace_id='workspace-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
priority=30,
|
|
)
|
|
await store.create_run(
|
|
run_id='run-other-conversation',
|
|
event_id='evt-other-conversation',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
conversation_id='conv-b',
|
|
bot_id='bot-a',
|
|
workspace_id='workspace-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
priority=20,
|
|
)
|
|
await store.create_run(
|
|
run_id='run-allowed',
|
|
event_id='evt-allowed',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
conversation_id='conv-a',
|
|
bot_id='bot-a',
|
|
workspace_id='workspace-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
priority=10,
|
|
)
|
|
|
|
claimed = await store.claim_next_run(
|
|
runtime_id='runtime-a',
|
|
queue_name='default',
|
|
runner_ids=['runner-a'],
|
|
conversation_id='conv-a',
|
|
bot_id='bot-a',
|
|
workspace_id='workspace-a',
|
|
thread_id=None,
|
|
strict_thread=True,
|
|
)
|
|
|
|
assert claimed is not None
|
|
assert claimed['run_id'] == 'run-allowed'
|
|
|
|
|
|
@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_release_expired_claims_requeues_runs(store, db_engine):
|
|
await store.create_run(
|
|
run_id='run-expired-release',
|
|
event_id='evt-3',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
await store.create_run(
|
|
run_id='run-active-claim',
|
|
event_id='evt-4',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
expired_claim = await store.claim_next_run(runtime_id='runtime-a', queue_name='default', lease_seconds=60)
|
|
active_claim = await store.claim_next_run(runtime_id='runtime-b', queue_name='default', lease_seconds=60)
|
|
assert expired_claim is not None
|
|
assert active_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-release')
|
|
.values(claim_lease_expires_at=datetime.datetime.now(UTC) - datetime.timedelta(seconds=1))
|
|
)
|
|
await session.commit()
|
|
|
|
released = await store.release_expired_claims()
|
|
|
|
assert [run['run_id'] for run in released] == ['run-expired-release']
|
|
assert released[0]['status'] == 'queued'
|
|
assert released[0]['status_reason'] == 'claim lease expired'
|
|
assert released[0]['claimed_by_runtime_id'] is None
|
|
assert 'claim_token' not in released[0]
|
|
assert released[0]['claim_lease_expires_at'] is None
|
|
|
|
active = await store.get_run('run-active-claim')
|
|
assert active is not None
|
|
assert active['status'] == 'claimed'
|
|
assert active['claimed_by_runtime_id'] == active_claim['claimed_by_runtime_id']
|
|
assert 'claim_token' not in active
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expired_claim_cannot_renew_or_release(store, db_engine):
|
|
await store.create_run(
|
|
run_id='run-stale-claim',
|
|
event_id='evt-stale',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
claimed = await store.claim_next_run(runtime_id='runtime-a', queue_name='default', lease_seconds=60)
|
|
assert claimed is not None
|
|
token = claimed['claim_token']
|
|
|
|
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-stale-claim')
|
|
.values(claim_lease_expires_at=datetime.datetime.now(UTC) - datetime.timedelta(seconds=1))
|
|
)
|
|
await session.commit()
|
|
|
|
assert await store.renew_claim(run_id='run-stale-claim', claim_token=token, runtime_id='runtime-a') is None
|
|
assert await store.release_claim(run_id='run-stale-claim', claim_token=token, runtime_id='runtime-a') is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_status_validation_and_terminal_transition_rules(store):
|
|
with pytest.raises(ValueError, match='Unknown run status'):
|
|
await store.create_run(
|
|
run_id='run-invalid-create',
|
|
event_id='evt-invalid',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='bogus',
|
|
)
|
|
|
|
await store.create_run(
|
|
run_id='run-invalid-release',
|
|
event_id='evt-release',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
claim = await store.claim_next_run(runtime_id='runtime-a', queue_name='default')
|
|
assert claim is not None
|
|
with pytest.raises(ValueError, match='Unknown run status'):
|
|
await store.release_claim(
|
|
run_id='run-invalid-release',
|
|
claim_token=claim['claim_token'],
|
|
runtime_id='runtime-a',
|
|
status='bogus',
|
|
)
|
|
|
|
await store.create_run(
|
|
run_id='run-terminal',
|
|
event_id='evt-terminal',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
)
|
|
with pytest.raises(ValueError, match='Unknown run status'):
|
|
await store.finalize_run(run_id='run-terminal', status='bogus')
|
|
|
|
completed = await store.finalize_run(
|
|
run_id='run-terminal',
|
|
status='completed',
|
|
metadata={'attempt': 1},
|
|
)
|
|
assert completed is not None
|
|
assert completed['status'] == 'completed'
|
|
|
|
merged = await store.finalize_run(
|
|
run_id='run-terminal',
|
|
status='completed',
|
|
metadata={'retry_observed': True},
|
|
)
|
|
assert merged is not None
|
|
assert merged['metadata'] == {'attempt': 1, 'retry_observed': True}
|
|
|
|
with pytest.raises(ValueError, match='Cannot transition terminal run'):
|
|
await store.finalize_run(run_id='run-terminal', status='failed')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_audit_event_uses_next_sequence(store):
|
|
await store.create_run(
|
|
run_id='run-audit',
|
|
event_id='evt-5',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
)
|
|
await store.append_event(
|
|
run_id='run-audit',
|
|
sequence=1,
|
|
event_type='message.completed',
|
|
data={'ok': True},
|
|
)
|
|
|
|
event = await store.append_audit_event(
|
|
run_id='run-audit',
|
|
event_type='admin.run_cancel',
|
|
data={'action': 'run_cancel'},
|
|
metadata={'permission': 'agent_run:admin'},
|
|
)
|
|
|
|
assert event is not None
|
|
assert event['sequence'] == 2
|
|
assert event['type'] == 'admin.run_cancel'
|
|
assert event['source'] == 'host'
|
|
assert event['data'] == {'action': 'run_cancel'}
|
|
assert event['metadata'] == {'permission': 'agent_run:admin'}
|
|
assert await store.append_audit_event(run_id='missing', event_type='admin.missing') is None
|
|
|
|
|
|
@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, total_count = await store.list_runtimes(statuses=['online'])
|
|
assert [runtime['runtime_id'] for runtime in runtimes] == ['runtime-a']
|
|
assert total_count == 1
|
|
|
|
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'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_stats_splits_active_and_claimed_runs(store):
|
|
await store.register_runtime(runtime_id='runtime-a')
|
|
await store.create_run(
|
|
run_id='run-running',
|
|
event_id='evt-running',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='running',
|
|
)
|
|
await store.create_run(
|
|
run_id='run-claimed',
|
|
event_id='evt-claimed',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
assert await store.claim_next_run(runtime_id='runtime-a', queue_name='default') is not None
|
|
|
|
stats = await store.get_runtime_stats()
|
|
|
|
assert stats['active_runs'] == 2
|
|
assert stats['claimed_runs'] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runner_stats_reports_zero_success_rate_for_failed_only_runner(store):
|
|
now = int(datetime.datetime.now(UTC).timestamp())
|
|
await store.create_run(
|
|
run_id='run-failed',
|
|
event_id='evt-failed',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='failed',
|
|
)
|
|
|
|
stats = await store.get_runner_stats(start_time=now - 10, end_time=now + 10)
|
|
|
|
assert stats[0]['runner_id'] == 'runner-a'
|
|
assert stats[0]['failed_runs'] == 1
|
|
assert stats[0]['success_rate'] == 0.0
|