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
2026-10-17 19:45:19 +08:00
parent 1318efc248
commit d663331889
3 changed files with 347 additions and 28 deletions
+74 -18
View File
@@ -1291,15 +1291,16 @@ class MonitoringService(ExecutionDetailsMixin):
'input_preview': _message_preview(json.dumps(metadata.get('input') or {}, ensure_ascii=False), 200),
}
def _serialize_pipeline_execution(self, row) -> dict:
def _serialize_pipeline_execution(self, row, *, summary_only=False) -> dict:
created_at_ms = _epoch_ms(row.timestamp)
preview = '' if summary_only else _message_preview(row.message_content)
return {
'source': 'pipeline',
'id': row.id,
'event_id': row.event_id,
'status': row.status,
'status_group': _status_group(row.status, 'pipeline'),
'title': _message_preview(row.message_content),
'title': row.pipeline_name if summary_only else preview,
'target_kind': 'pipeline',
'target_id': row.pipeline_id,
'target_name': row.pipeline_name,
@@ -1319,7 +1320,7 @@ class MonitoringService(ExecutionDetailsMixin):
'queue_name': None,
'debug': row.bot_id == 'websocket-proxy-bot',
'has_error': row.status == 'error',
'input_preview': _message_preview(row.message_content),
'input_preview': preview,
'platform': row.platform,
'user_id': row.user_id,
'user_name': row.user_name,
@@ -1533,6 +1534,23 @@ class MonitoringService(ExecutionDetailsMixin):
'real': total - debug_count,
}
@staticmethod
def _inflight_pipeline_condition(*, pending_only: bool = False):
"""Historical pending records without a live Host lease are not tasks."""
Message = persistence_monitoring.MonitoringMessage
live_lease = sqlalchemy.and_(
Message.execution_owner_id.is_not(None),
Message.execution_lease_expires_at > datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),
)
# Active-only queries must expose the lease range to the composite index.
# The status OR otherwise reads historical payload pages to check role.
if pending_only:
return live_lease
return sqlalchemy.or_(
Message.status != 'pending',
live_lease,
)
async def get_executions(
self,
context: TenantContext,
@@ -1548,6 +1566,7 @@ class MonitoringService(ExecutionDetailsMixin):
limit: int = 50,
offset: int = 0,
include_summary: bool = True,
for_inflight: bool = False,
) -> dict:
"""List executions from both object types (agent runs and pipeline queries)."""
workspace_uuid = require_workspace_uuid(context)
@@ -1579,15 +1598,18 @@ class MonitoringService(ExecutionDetailsMixin):
).where(*conditions)
)
if source in {'all', 'pipeline'} and (not agent_ids or pipeline_ids):
pipeline_statuses = self._raw_pipeline_statuses(statuses)
conditions = self._pipeline_execution_conditions(
workspace_uuid,
bot_ids=bot_ids,
pipeline_ids=pipeline_ids,
statuses=self._raw_pipeline_statuses(statuses),
statuses=pipeline_statuses,
start_time=start_time,
end_time=end_time,
debug=debug,
)
if for_inflight:
conditions.append(self._inflight_pipeline_condition(pending_only=pipeline_statuses == ['pending']))
branches.append(
sqlalchemy.select(
sqlalchemy.literal('pipeline').label('source'),
@@ -1595,7 +1617,7 @@ class MonitoringService(ExecutionDetailsMixin):
Message.timestamp.label('time'),
).where(*conditions)
)
if source in {'all', 'event'} and mode != 'debug' and not pipeline_ids and not agent_ids:
if source in {'all', 'event'} and mode != 'debug' and not pipeline_ids and not agent_ids and not for_inflight:
conditions = self._unhandled_event_conditions(workspace_uuid, bot_ids, start_time, end_time)
if statuses:
conditions.append(self._execution_json_text(EventLog.metadata_json, 'status').in_(statuses))
@@ -1634,20 +1656,23 @@ class MonitoringService(ExecutionDetailsMixin):
workspace_column = model.workspace_uuid if kind == 'pipeline' else model.workspace_id
records = await self._execution_rows(
workspace_uuid,
sqlalchemy.select(model).where(workspace_column == workspace_uuid, identity_column.in_(ids)),
self._execution_statement(model, summary_only=for_inflight).where(
workspace_column == workspace_uuid, identity_column.in_(ids)
),
summary_only=for_inflight,
)
names = await self._agent_names(workspace_uuid, records) if kind == 'agent' else {}
for record in records:
item = (
self._serialize_agent_execution(record, names)
if kind == 'agent'
else self._serialize_pipeline_execution(record)
else self._serialize_pipeline_execution(record, summary_only=for_inflight)
if kind == 'pipeline'
else self._serialize_event_execution(record)
)
by_key[(kind, item['id'])] = item
page = [by_key[(ref.source, ref.id)] for ref in refs if (ref.source, ref.id) in by_key]
await self.enrich_execution_rows(workspace_uuid, page)
await self.enrich_execution_rows(workspace_uuid, page, summary_only=for_inflight)
if not include_summary:
return {'items': page, 'total': total, 'has_more': total > offset + len(page)}
return {
@@ -1692,7 +1717,9 @@ class MonitoringService(ExecutionDetailsMixin):
from ....utils.inflight import inflight_hub
workspace = require_workspace_uuid(context)
active = await self.get_executions(context, statuses=['running', 'queued'], limit=100, include_summary=False)
active = await self.get_executions(
context, statuses=['running', 'queued'], limit=100, include_summary=False, for_inflight=True
)
# Execution timestamps are stored as UTC without a timezone. asyncpg
# rejects aware datetimes for these columns, even though SQLite accepts them.
recent_start = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) - datetime.timedelta(
@@ -1703,6 +1730,7 @@ class MonitoringService(ExecutionDetailsMixin):
start_time=recent_start,
limit=100,
include_summary=False,
for_inflight=True,
)
items = {(r['source'], r['id']): r for r in recent['items'] + active['items'] if r['source'] != 'event'}
# Keep the actual terminal state of older runs that just left the active set.
@@ -1722,28 +1750,56 @@ class MonitoringService(ExecutionDetailsMixin):
]:
missing = [r['id'] for r in (previous or [])[:100] if r['source'] == kind and (kind, r['id']) not in items]
if missing:
records = await self._execution_rows(
workspace, sqlalchemy.select(model).where(scope == workspace, column.in_(missing))
statement = self._execution_statement(model, summary_only=True).where(
scope == workspace, column.in_(missing)
)
if kind == 'pipeline':
statement = statement.where(self._inflight_pipeline_condition())
records = await self._execution_rows(workspace, statement, summary_only=True)
names = await self._agent_names(workspace, records) if kind == 'agent' else {}
for record in records:
row = (
self._serialize_agent_execution(record, names)
if kind == 'agent'
else self._serialize_pipeline_execution(record)
else self._serialize_pipeline_execution(record, summary_only=True)
)
items[(kind, row['id'])] = row
rows = list(items.values())
snapshot = {
'items': rows,
'active_total': active['total'],
'truncated': active['has_more'] or recent['has_more'],
}
if not rows:
inflight_hub.track(workspace, [])
return snapshot
Run = persistence_agent_run.AgentRun
Message = persistence_monitoring.MonitoringMessage
agent_ids = [r['id'] for r in rows if r['source'] == 'agent']
pipeline_ids = [r['id'] for r in rows if r['source'] == 'pipeline']
linked = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Run.run_id, Message.id.label('message_id'))
.outerjoin(Message, sqlalchemy.and_(Message.workspace_uuid == workspace, Message.run_id == Run.run_id))
.where(
Run.workspace_id == workspace, sqlalchemy.or_(Run.run_id.in_(agent_ids), Message.id.in_(pipeline_ids))
links = []
if agent_ids:
links.append(
sqlalchemy.select(Run.run_id, sqlalchemy.literal(None).label('message_id')).where(
Run.workspace_id == workspace, Run.run_id.in_(agent_ids)
)
)
if pipeline_ids:
# Start from the selected message IDs. A cross-table OR after a
# LEFT JOIN makes SQLite scan the attachment-heavy message ledger
# whenever the snapshot contains even one Pipeline execution.
selected_messages = sqlalchemy.union_all(
*[sqlalchemy.select(sqlalchemy.literal(message_id).label('message_id')) for message_id in pipeline_ids]
).subquery('selected_messages')
links.append(
sqlalchemy.select(Run.run_id, Message.id.label('message_id'))
.select_from(selected_messages)
.join(Message, Message.id == selected_messages.c.message_id)
.join(Run, sqlalchemy.and_(Run.workspace_id == workspace, Run.run_id == Message.run_id))
.where(Message.workspace_uuid == workspace)
)
linked = await self.ap.persistence_mgr.execute_async(
links[0] if len(links) == 1 else sqlalchemy.union_all(*links)
)
identities = {
r.run_id: ('pipeline', r.message_id) if r.message_id in pipeline_ids else ('agent', r.run_id)
@@ -1766,7 +1822,7 @@ class MonitoringService(ExecutionDetailsMixin):
row = items.get(identities[run_id])
if row is not None:
row['progress_event'] = event_type
return {'items': rows, 'active_total': active['total'], 'truncated': active['has_more'] or recent['has_more']}
return snapshot
async def get_token_statistics(
self,
@@ -4,6 +4,7 @@ from __future__ import annotations
import datetime
import json
from types import SimpleNamespace
from langbot_plugin.api.entities.builtin.platform.events import event_summary
import sqlalchemy as sa
@@ -354,8 +355,44 @@ class ExecutionDetailsMixin:
async with AsyncSession(self.ap.persistence_mgr.get_db_engine()) as session:
return await session.execute(statement)
async def _execution_rows(self, workspace, statement):
return list((await self._execution_query(workspace, statement)).scalars().all())
async def _execution_rows(self, workspace, statement, *, summary_only=False):
result = await self._execution_query(workspace, statement)
if not summary_only:
return list(result.scalars().all())
records = []
for row in result.mappings():
fields = dict(row)
metadata = {
key.removeprefix('display_metadata_'): fields.pop(key)
for key in row
if key.startswith('display_metadata_')
}
if metadata:
fields['metadata_json'] = json.dumps(metadata, ensure_ascii=False)
records.append(SimpleNamespace(**fields))
return records
def _execution_statement(self, model, *, summary_only=False):
"""Keep periodic snapshots independent of message and attachment sizes."""
if not summary_only:
return sa.select(model)
columns = [
column
for column in model.__table__.columns
if column.name not in {'message_content', 'variables', 'input_json', 'authorization_json', 'metadata_json'}
]
if model is models.MonitoringMessage:
return sa.select(*columns)
keys = (
('processor_id', 'processor_type', 'event_type', 'agent_name', 'target_name', 'bot_name')
if model is AgentRun
else ('platform', 'bot_name')
)
# Metadata also contains complete inputs. Project only the small display
# fields in SQL, rather than deserializing an attachment-bearing object.
for key in keys:
columns.append(self._execution_json_text(model.metadata_json, key).label(f'display_metadata_{key}'))
return sa.select(*columns)
async def _execution_page(self, workspace, statement, convert, offset, limit):
rows = await self._execution_rows(workspace, statement.offset(offset).limit(limit + 1))
@@ -445,14 +482,21 @@ class ExecutionDetailsMixin:
or None,
}
async def enrich_execution_rows(self, workspace, items):
async def enrich_execution_rows(self, workspace, items, *, summary_only=False):
pipeline_items = {item['id']: item for item in items if item.get('source') == 'pipeline'}
if pipeline_items:
# Pipeline rows own the UI identity; their linked Runner ledger owns
# timing and usage. Read them in one batch for both list and detail.
linked = await self._execution_query(
workspace,
sa.select(models.MonitoringMessage.id, AgentRun)
sa.select(
models.MonitoringMessage.id,
AgentRun.started_at,
AgentRun.finished_at,
AgentRun.usage_json,
AgentRun.cost_json,
)
.select_from(models.MonitoringMessage)
.join(
AgentRun,
sa.and_(
@@ -465,10 +509,15 @@ class ExecutionDetailsMixin:
models.MonitoringMessage.id.in_(pipeline_items),
),
)
for message_id, run in linked.all():
metrics = self._serialize_agent_execution(run)
pipeline_items[message_id].update(
{key: metrics[key] for key in ('started_at_ms', 'finished_at_ms', 'duration_ms', 'usage', 'cost')}
for run in linked.all():
pipeline_items[run.id].update(
started_at_ms=epoch(run.started_at),
finished_at_ms=epoch(run.finished_at),
duration_ms=int(round((run.finished_at - run.started_at).total_seconds() * 1000))
if run.started_at and run.finished_at
else None,
usage=self._deserialize_json(run.usage_json) or None,
cost=self._deserialize_json(run.cost_json) or None,
)
calls = models.MonitoringLLMCall
@@ -498,10 +547,11 @@ class ExecutionDetailsMixin:
return
events = await self._execution_rows(
workspace,
sa.select(EventLog).where(
self._execution_statement(EventLog, summary_only=summary_only).where(
EventLog.workspace_id == workspace,
EventLog.event_id.in_(ids),
),
summary_only=summary_only,
)
by_id = {row.event_id: row for row in events}
for item in items:
@@ -510,7 +560,9 @@ class ExecutionDetailsMixin:
meta = parsed(event.metadata_json, {})
item.update(
event_type=event.event_type,
input_preview=event_preview(parsed(event.input_json), event.input_summary, event.event_type),
input_preview=event_preview(
None if summary_only else parsed(event.input_json), event.input_summary, event.event_type
),
user_id=event.actor_id,
user_name=event.actor_name,
platform=meta.get('platform') or event.source,
@@ -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