mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-24 01:15:24 +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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user