Files
LangBot/tests/unit_tests/agent/test_run_ledger_store.py
T
BiFangKNT ce5b7ffea1 fix(agent): preserve terminal events when cancelling runner execution
Emit and persist a Host cancellation event before closing the runner stream.
Preserve existing terminal events and record cancelled runs consistently.

Serialize Host event sequence allocation to prevent conflicts with admin
cancellation audits, and notify observers with the persisted sequence.

Add deterministic cancellation and concurrent event allocation regressions.
2026-10-23 17:23:44 +08:00

686 lines
24 KiB
Python

"""Tests for RunLedgerStore host primitives."""
from __future__ import annotations
import datetime
import asyncio
import json
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.fixture
async def legacy_run(store, db_engine):
await store.create_run(
run_id='legacy',
event_id='event',
binding_id='binding',
runner_id='runner',
workspace_id='ws',
)
metadata = {
'input': [
{'type': 'Image', 'base64': 'eA=='},
{'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA==', 'detail': 'low'}},
{'type': 'image_url', 'image_url': {'url': 'https://example.com/image.png'}},
{'type': 'text', 'text': 'keep this'},
],
'attempt': 0,
'activated_skills': {'existing': 'v1'},
}
async with db_engine.begin() as conn:
# Explicit TEXT binding simulates data persisted before the media contract.
await conn.execute(
sqlalchemy.update(AgentRun)
.where(AgentRun.run_id == 'legacy')
.values(
metadata_json=sqlalchemy.bindparam('legacy_metadata', json.dumps(metadata), type_=sqlalchemy.Text())
)
)
return metadata
@pytest.mark.asyncio
@pytest.mark.parametrize('operation', ['pin', 'finalize'])
async def test_metadata_update_redacts_legacy_images(store, legacy_run, operation):
if operation == 'pin':
assert await store.pin_skill_revision('legacy', 'ws', 'demo', 'v2') == 'v2'
else:
await store.finalize_run(run_id='legacy', status='completed', metadata={'note': 'done'})
run = await store.get_run('legacy')
expected = {
**legacy_run,
'input': [
{'type': 'Image', 'base64': None},
{'type': 'image_url', 'image_url': {'url': None, 'detail': 'low'}},
*legacy_run['input'][2:],
],
}
if operation == 'pin':
expected['activated_skills'] = {'existing': 'v1', 'demo': 'v2'}
assert run['status'] == 'running'
else:
expected['note'] = 'done'
assert run['status'] == 'completed'
assert run['metadata'] == expected
@pytest.mark.asyncio
async def test_finalize_without_metadata_preserves_legacy_snapshot(store, legacy_run):
run = await store.finalize_run(run_id='legacy', status='completed')
assert run['status'] == 'completed'
assert run['metadata'] == legacy_run
@pytest.mark.asyncio
async def test_metadata_update_still_rejects_new_inline_images(store, legacy_run):
with pytest.raises(sqlalchemy.exc.StatementError, match='Inline images cannot be persisted'):
await store.finalize_run(
run_id='legacy', status='completed', metadata={'new_image': {'type': 'Image', 'base64': 'eQ=='}}
)
run = await store.get_run('legacy')
assert run['status'] == 'running'
assert run['metadata'] == legacy_run
@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_host_events_allocate_sequences_across_stores(store, db_engine):
await store.create_run(run_id='concurrent-host', event_id=None, binding_id=None, runner_id='runner')
other = RunLedgerStore(db_engine)
events = await asyncio.gather(
store.append_audit_event(run_id='concurrent-host', event_type='admin.run_cancel'),
other.append_event(
run_id='concurrent-host',
sequence=None,
event_type='run.failed',
data={'code': 'cancelled'},
source='host',
),
)
assert sorted(event['sequence'] for event in events) == [1, 2]
persisted, *_ = await store.page_run_events(run_id='concurrent-host', limit=10)
assert {event['type'] for event in persisted} == {'admin.run_cancel', 'run.failed'}
@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
@pytest.mark.asyncio
async def test_processor_instance_history_filters_count_and_pages(store):
for index, (workspace, binding) in enumerate(
[
('one', 'processor-a'),
('one', 'processor-b'),
('two', 'processor-a'),
('one', 'processor-a'),
]
):
await store.create_run(
run_id=f'run-{index}',
event_id=f'event-{index}',
binding_id=binding,
runner_id='same-component',
workspace_id=workspace,
)
first, cursor, more, total = await store.list_runs(workspace_id='one', binding_id='processor-a', limit=1)
assert (total, more) == (2, True)
assert first[0]['run_id'] == 'run-3'
second, _, more, total = await store.list_runs(
workspace_id='one', binding_id='processor-a', limit=1, before_id=cursor
)
assert (total, more) == (2, False)
assert second[0]['run_id'] == 'run-0'
@pytest.mark.asyncio
async def test_run_lifecycle_retains_milliseconds(store, monkeypatch):
started = datetime.datetime(2026, 9, 8, 0, 0, 0, 123000, tzinfo=UTC)
monkeypatch.setattr('langbot.pkg.agent.runner.run_ledger_store._utc_now', lambda: started)
run = await store.create_run(
run_id='run-ms', event_id='evt-ms', binding_id='binding-ms', runner_id='runner-ms', status='running'
)
assert run['created_at_ms'] == round(started.timestamp() * 1000)
assert run['started_at_ms'] == run['created_at_ms']
assert run['finished_at_ms'] is None
finished = started + datetime.timedelta(milliseconds=275)
monkeypatch.setattr('langbot.pkg.agent.runner.run_ledger_store._utc_now', lambda: finished)
await store.finalize_run(run_id='run-ms', status='completed')
saved = await store.get_run('run-ms')
assert saved['finished_at_ms'] - saved['started_at_ms'] == 275
@pytest.mark.asyncio
async def test_agent_run_pages_include_old_bindings_without_leaking_other_agents(store):
for run_id, agent_id, binding, workspace in [
('old', None, 'agent_one_plugin:a/old/default', 'w'),
('debug', None, 'debug:one:plugin:a/new/default', 'w'),
('new', 'one', 'new-binding', 'w'),
('other', 'two', 'agent_two_plugin:a/old/default', 'w'),
('conflicting-id', 'two', 'agent_one_plugin:a/old/default', 'w'),
('other-workspace', 'one', 'agent_one_plugin:a/old/default', 'other'),
('plugin-processor', None, 'event_processor:one', 'w'),
]:
await store.create_run(
run_id=run_id,
event_id=None,
agent_id=agent_id,
binding_id=binding,
workspace_id=workspace,
runner_id='runner',
)
page, cursor, more, total = await store.list_runs(agent_id='one', workspace_id='w', limit=2)
assert [run['run_id'] for run in page] == ['new', 'debug']
assert more and total == 3
page, cursor, more, total = await store.list_runs(agent_id='one', workspace_id='w', before_id=cursor, limit=2)
assert [run['run_id'] for run in page] == ['old']
assert not more and total == 3
@pytest.mark.asyncio
async def test_skill_pins_belong_to_run_and_survive_store_restart(store, db_engine):
for run_id in ('old', 'new'):
await store.create_run(
run_id=run_id,
event_id=run_id,
binding_id='binding',
runner_id='runner',
workspace_id='ws',
conversation_id='same',
status='running',
)
assert await store.pin_skill_revision('old', 'ws', 'demo', 'v1') == 'v1'
restarted = RunLedgerStore(db_engine)
assert await restarted.pin_skill_revision('old', 'ws', 'demo', 'v2') == 'v1'
assert await restarted.pin_skill_revision('new', 'ws', 'demo', 'v2') == 'v2'
assert (await restarted.get_run('old'))['metadata']['activated_skills'] == {'demo': 'v1'}
with pytest.raises(ValueError, match='current Workspace'):
await restarted.pin_skill_revision('old', 'other', 'demo', 'v2')
await store.finalize_run(run_id='old', status='completed')
with pytest.raises(ValueError, match='active run'):
await restarted.pin_skill_revision('old', 'ws', 'demo', 'v2')
@pytest.mark.asyncio
async def test_activation_uses_new_revision_only_for_new_run(store, db_engine, tmp_path):
from types import SimpleNamespace
from unittest.mock import Mock
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.skill.repository import SkillRepository
from langbot.pkg.skill.manager import SkillManager
from langbot.pkg.provider.tools.loaders.skill_authoring import SkillToolLoader
context = ExecutionContext(instance_uuid='instance', workspace_uuid='ws', placement_generation=1)
async def binding(workspace_uuid, **kwargs):
return context
app = SimpleNamespace(
logger=Mock(),
instance_config=SimpleNamespace(data={'skills': {'root': str(tmp_path / 'skills')}}),
workspace_service=SimpleNamespace(get_execution_binding=binding),
persistence_mgr=SimpleNamespace(get_db_engine=lambda: db_engine),
)
app.skill_repository = SkillRepository(app)
app.skill_mgr = SkillManager(app)
loader = SkillToolLoader(app)
first = await app.skill_repository.create_skill(context, {'name': 'demo', 'instructions': 'version one'})
await app.skill_mgr.reload_skills(context)
def query(run_id):
return SimpleNamespace(
instance_uuid='instance',
workspace_uuid='ws',
placement_generation=1,
variables={'_pipeline_bound_skills': ['demo']},
_execution_context=context,
_agent_run_session={'run_id': run_id},
)
for name in ('old', 'new'):
await store.create_run(
run_id=name,
event_id=name,
binding_id='binding',
runner_id='runner',
workspace_id='ws',
conversation_id='same',
status='running',
)
old = query('old')
assert 'version one' in (await loader._invoke_activate_skill({'skill_name': 'demo'}, old))['content']
second = await app.skill_repository.update_skill(
context, 'demo', {'instructions': 'version two'}, base_revision=first['revision']
)
await app.skill_mgr.reload_skills(context)
assert 'version two' in (await loader._invoke_activate_skill({'skill_name': 'demo'}, query('new')))['content']
assert 'version one' in (await loader._invoke_activate_skill({'skill_name': 'demo'}, old))['content']
await app.skill_repository.delete_skill(context, 'demo')
await app.skill_mgr.reload_skills(context)
assert 'version one' in (await loader._invoke_activate_skill({'skill_name': 'demo'}, query('old')))['content']
assert (await store.get_run('new'))['metadata']['activated_skills']['demo'] == second['revision']
revoked = query('old')
revoked.variables['_pipeline_bound_skills'] = []
with pytest.raises(ValueError, match='not authorized'):
await loader._invoke_activate_skill({'skill_name': 'demo'}, revoked)