mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-26 18:35:38 +08:00
fix(monitoring): bound inflight snapshot queries
Require a valid owner and unexpired lease for pending Pipeline messages to appear as active tasks. Exclude historical ingress records without rewriting their stored status. Project only the fields needed by snapshots, skip linkage queries for empty results, and resolve linked runs from selected message IDs instead of a cross-table OR join that scans historical messages. Cover lease filtering, terminal transitions, Workspace isolation, payload exclusion, and indexed lookup plans with regression tests.
This commit is contained in:
BiFangKNT
@@ -739,6 +739,217 @@ async def test_pipeline_debug_is_consistent_across_rows_rollups_and_telemetry(en
|
||||
assert sum(p['llm_calls'] for p in traffic['points']) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inflight_excludes_unowned_and_expired_messages_without_rewriting_history(engine, service):
|
||||
import sqlalchemy as sa
|
||||
|
||||
now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
|
||||
Message = persistence_monitoring.MonitoringMessage
|
||||
async with AsyncSession(engine) as session:
|
||||
for message_id, owner, expiry, timestamp in [
|
||||
('legacy', None, None, _dt(1)),
|
||||
('recent-unowned', None, now + datetime.timedelta(minutes=1), now),
|
||||
('expired', 'stopped-host', now - datetime.timedelta(seconds=1), _dt(1)),
|
||||
('long-running', 'live-host', now + datetime.timedelta(minutes=1), _dt(1)),
|
||||
]:
|
||||
session.add(
|
||||
Message(
|
||||
id=message_id,
|
||||
workspace_uuid=WORKSPACE,
|
||||
timestamp=timestamp,
|
||||
bot_id='bot',
|
||||
bot_name='Bot',
|
||||
pipeline_id='pipeline',
|
||||
pipeline_name='Pipeline',
|
||||
session_id='session',
|
||||
message_content='historical message',
|
||||
status='pending',
|
||||
level='info',
|
||||
execution_owner_id=owner,
|
||||
execution_lease_expires_at=expiry,
|
||||
)
|
||||
)
|
||||
session.add(
|
||||
persistence_agent_run.AgentRun(
|
||||
run_id='remote-queue',
|
||||
workspace_id=WORKSPACE,
|
||||
runner_id='runner',
|
||||
status='queued',
|
||||
queue_name='remote',
|
||||
created_at=_dt(1),
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
# A previous connection must not reintroduce unowned pending rows either.
|
||||
snapshot = await service.get_inflight_snapshot(WORKSPACE, [{'source': 'pipeline', 'id': 'legacy'}])
|
||||
assert snapshot['active_total'] == 2
|
||||
assert {row['id'] for row in snapshot['items']} == {'long-running', 'remote-queue'}
|
||||
async with AsyncSession(engine) as session:
|
||||
statuses = (await session.execute(sa.select(Message.status))).scalars().all()
|
||||
assert statuses == ['pending'] * 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('source', ['all', 'pipeline'])
|
||||
@pytest.mark.parametrize('statuses', [['running'], ['running', 'queued']])
|
||||
async def test_active_inflight_count_and_page_use_the_indexed_lease_range(engine, service, source, statuses):
|
||||
import sqlalchemy as sa
|
||||
|
||||
await _seed(engine)
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sa.update(persistence_monitoring.MonitoringMessage)
|
||||
.where(persistence_monitoring.MonitoringMessage.id == 'msg-user')
|
||||
.values(status='pending', message_content='historical payload ' * 8192)
|
||||
)
|
||||
statements = []
|
||||
|
||||
def capture(connection, cursor, statement, parameters, context, executemany):
|
||||
if statement.startswith('SELECT '):
|
||||
statements.append((statement, parameters))
|
||||
|
||||
sa.event.listen(engine.sync_engine, 'before_cursor_execute', capture)
|
||||
try:
|
||||
result = await service.get_executions(
|
||||
WORKSPACE, statuses=statuses, source=source, include_summary=False, for_inflight=True
|
||||
)
|
||||
finally:
|
||||
sa.event.remove(engine.sync_engine, 'before_cursor_execute', capture)
|
||||
|
||||
assert result == {'items': [], 'total': 0, 'has_more': False}
|
||||
assert len(statements) == 2 # Both COUNT and the sorted page must skip historical payloads.
|
||||
async with engine.connect() as connection:
|
||||
for statement, parameters in statements:
|
||||
plan = (await connection.exec_driver_sql('EXPLAIN QUERY PLAN ' + statement, parameters)).all()
|
||||
assert any(
|
||||
'ix_monitoring_messages_workspace_status_lease' in row[3] and 'execution_lease_expires_at>?' in row[3]
|
||||
for row in plan
|
||||
), plan
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_inflight_snapshot_skips_linked_history_and_clears_tracking(engine, service, monkeypatch):
|
||||
import sqlalchemy as sa
|
||||
from langbot.pkg.utils.inflight import inflight_hub
|
||||
|
||||
await _seed(engine)
|
||||
monkeypatch.setattr(inflight_hub, 'run_workspaces', {'finished-run': WORKSPACE, 'other-run': 'other-workspace'})
|
||||
statements = []
|
||||
|
||||
def capture(connection, cursor, statement, parameters, context, executemany):
|
||||
statements.append(statement)
|
||||
|
||||
sa.event.listen(engine.sync_engine, 'before_cursor_execute', capture)
|
||||
try:
|
||||
snapshot = await service.get_inflight_snapshot(WORKSPACE)
|
||||
finally:
|
||||
sa.event.remove(engine.sync_engine, 'before_cursor_execute', capture)
|
||||
|
||||
assert snapshot == {'items': [], 'active_total': 0, 'truncated': False}
|
||||
assert inflight_hub.run_workspaces == {'other-run': 'other-workspace'}
|
||||
assert not any('JOIN monitoring_messages ON' in query or 'FROM agent_run_event' in query for query in statements)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inflight_never_loads_payloads_in_active_recent_or_terminal_rows(engine, service, monkeypatch):
|
||||
import json
|
||||
import sqlalchemy as sa
|
||||
from langbot.pkg.entity.persistence.event_log import EventLog
|
||||
|
||||
now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
|
||||
payload = json.dumps({'text': 'full input', 'attachment': 'x' * (2 * 1024 * 1024)})
|
||||
async with AsyncSession(engine) as session:
|
||||
for message_id, status, timestamp in [
|
||||
('active-message', 'pending', _dt(1)),
|
||||
('recent-message', 'success', now),
|
||||
('finished-message', 'success', _dt(1)),
|
||||
]:
|
||||
session.add(
|
||||
persistence_monitoring.MonitoringMessage(
|
||||
id=message_id,
|
||||
workspace_uuid=WORKSPACE,
|
||||
timestamp=timestamp,
|
||||
bot_id='bot',
|
||||
bot_name='Bot',
|
||||
pipeline_id='pipeline',
|
||||
pipeline_name='Pipeline',
|
||||
session_id='session',
|
||||
message_content=payload,
|
||||
variables=payload,
|
||||
status=status,
|
||||
level='info',
|
||||
event_id=message_id,
|
||||
run_id='linked-run' if status == 'pending' else None,
|
||||
execution_owner_id='host' if status == 'pending' else None,
|
||||
execution_lease_expires_at=now + datetime.timedelta(minutes=1) if status == 'pending' else None,
|
||||
)
|
||||
)
|
||||
session.add(
|
||||
EventLog(
|
||||
event_id=message_id,
|
||||
workspace_id=WORKSPACE,
|
||||
event_type='message.received',
|
||||
source='platform',
|
||||
input_summary='small summary',
|
||||
input_json=payload,
|
||||
metadata_json=json.dumps({'platform': 'test', 'unused_payload': payload}),
|
||||
)
|
||||
)
|
||||
for run_id in ['linked-run', 'standalone-run']:
|
||||
session.add(
|
||||
persistence_agent_run.AgentRun(
|
||||
run_id=run_id,
|
||||
workspace_id=WORKSPACE,
|
||||
runner_id='runner',
|
||||
status='running',
|
||||
created_at=_dt(1),
|
||||
authorization_json=payload,
|
||||
metadata_json=json.dumps(
|
||||
{'event_type': 'message.received', 'agent_name': 'Agent', 'input': {'text': payload}}
|
||||
),
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
linkage_queries = []
|
||||
|
||||
def reject_payload_reads(connection, cursor, statement, parameters, context, executemany):
|
||||
for column in [
|
||||
'monitoring_messages.message_content',
|
||||
'monitoring_messages.variables',
|
||||
'agent_run.authorization_json',
|
||||
'event_log.input_json',
|
||||
]:
|
||||
assert column not in statement
|
||||
if statement.startswith('SELECT agent_run.run_id,') and 'monitoring_messages.id AS message_id' in statement:
|
||||
linkage_queries.append((statement, parameters))
|
||||
|
||||
deserialize = service._deserialize_json
|
||||
|
||||
def bounded_metadata(raw):
|
||||
assert not raw or len(raw) < 10000
|
||||
return deserialize(raw)
|
||||
|
||||
monkeypatch.setattr(service, '_deserialize_json', bounded_metadata)
|
||||
sa.event.listen(engine.sync_engine, 'before_cursor_execute', reject_payload_reads)
|
||||
try:
|
||||
snapshot = await service.get_inflight_snapshot(WORKSPACE, [{'source': 'pipeline', 'id': 'finished-message'}])
|
||||
rows = {row['id']: row for row in snapshot['items']}
|
||||
assert rows['active-message']['input_preview'] == 'small summary'
|
||||
assert rows['recent-message']['status_group'] == 'completed'
|
||||
assert rows['finished-message']['status_group'] == 'completed'
|
||||
assert rows['standalone-run']['target_name'] == 'Agent'
|
||||
assert snapshot['active_total'] == 2
|
||||
finally:
|
||||
sa.event.remove(engine.sync_engine, 'before_cursor_execute', reject_payload_reads)
|
||||
assert linkage_queries
|
||||
async with engine.connect() as connection:
|
||||
for statement, parameters in linkage_queries:
|
||||
plan = (await connection.exec_driver_sql('EXPLAIN QUERY PLAN ' + statement, parameters)).all()
|
||||
message_lookups = [row[3] for row in plan if 'SEARCH monitoring_messages' in row[3]]
|
||||
assert message_lookups and all('(id=?)' in lookup for lookup in message_lookups), plan
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inflight_snapshot_is_scoped_and_keeps_terminal_transition(engine, service):
|
||||
import sqlalchemy as sa
|
||||
|
||||
Reference in New Issue
Block a user