fix(monitoring): restore Cloud messages and bot-scoped sessions (#2526)

* fix(monitoring): restore Cloud message persistence and bot-scoped sessions

* fix(migrations): support partial monitoring schemas and align regression fixtures

* test(migrations): complete raw bot session fixture values

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
Hyu
2026-09-11 14:34:37 +08:00
committed by GitHub
parent ce6b647fe7
commit ff6ad6adc2
36 changed files with 2436 additions and 617 deletions
@@ -0,0 +1,19 @@
"""Identifier normalization must not rely on SQLite's permissive codecs."""
import pytest
from langbot.pkg.api.http.service import monitoring
@pytest.mark.parametrize(
('value', 'expected'),
[(None, None), ('', ''), ('00123', '00123'), (' 用户 ', ' 用户 '), (123, '123'), (-123, '-123'), (0, '0')],
)
def test_normalize_user_id_preserves_opaque_strings(value, expected):
assert monitoring._normalize_user_id(value) == expected
@pytest.mark.parametrize('value', [True, False, 1.5, b'123', ['123'], {'id': 123}])
def test_normalize_user_id_rejects_unsupported_types(value):
with pytest.raises(TypeError, match='user_id must be a string, integer, or None'):
monitoring._normalize_user_id(value)
@@ -0,0 +1,220 @@
"""Bot-scoped session regressions exercised against real SQL databases."""
import datetime as dt
import logging
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.api.http.service.monitoring import MonitoringService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence import monitoring as models
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.pipeline.monitoring_helper import MonitoringHelper
from tests.integration.persistence.test_monitoring_postgres import cloud_database # noqa: F401
pytestmark = pytest.mark.asyncio
@pytest.mark.asyncio(loop_scope='module')
async def test_postgres_upgrade_rls_and_concurrent_bot_counts(cloud_database): # noqa: F811
import asyncio
import importlib
from alembic.migration import MigrationContext
from alembic.operations import Operations
from tests.integration.persistence.test_monitoring_postgres import WORKSPACE_A, _context, _read
ap, admin = cloud_database
service = ap.monitoring_service
ctx = _context(WORKSPACE_A)
await service.record_session_start(ctx, session_id='person_42', **resource('a'))
for bot in ['a', 'b']:
await service.record_message(ctx, session_id='person_42', message_content=bot, **resource(bot))
async with admin.begin() as conn:
def migrate(connection):
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0023_bot_scoped_sessions')
with Operations.context(MigrationContext.configure(connection)):
migration.downgrade()
migration.upgrade()
rls = connection.execute(
sa.text("SELECT relrowsecurity, relforcerowsecurity FROM pg_class WHERE relname='monitoring_sessions'")
).one()
assert tuple(rls) == (True, True)
assert (
connection.execute(
sa.text("SELECT count(*) FROM pg_policies WHERE tablename='monitoring_sessions'")
).scalar_one()
== 1
)
await conn.run_sync(migrate)
rows, total = await _read(service, 'get_sessions', ctx)
assert total == 2
assert {r['bot_id']: r['message_count'] for r in rows} == {'a': 1, 'b': 1}
await asyncio.gather(*[service.record_session_start(ctx, session_id='race', **resource('a')) for _ in range(10)])
result = await _read(service, 'get_session_analysis', ctx, 'race', bot_id='a')
assert result['session']['message_count'] == 10
assert not (await _read(service, 'get_session_analysis', ctx, 'person_42'))['found']
assert (await _read(service, 'get_session_analysis', ctx, 'person_42', bot_id='b'))['message_stats']['total'] == 1
async def test_migration_reconstructs_collisions_and_preserves_indexes(service):
import importlib
from alembic.migration import MigrationContext
from alembic.operations import Operations
engine = service.ap.persistence_mgr.get_db_engine()
async with engine.begin() as conn:
def upgrade(connection):
table = models.MonitoringSession.__table__
table.drop(connection)
metadata = sa.MetaData()
legacy = table.to_metadata(metadata)
legacy.primary_key._columns.remove(legacy.c.bot_id)
legacy.c.bot_id.primary_key = False
# Resolve the unchanged Workspace FK in copied metadata.
Base.metadata.tables['workspaces'].to_metadata(metadata)
legacy.create(connection)
now = dt.datetime(2026, 1, 1)
connection.execute(
sa.insert(legacy).values(
workspace_uuid='workspace',
session_id='person_42',
**resource('a'),
message_count=99,
start_time=now,
last_activity=now,
is_active=True,
)
)
for bot in ['a', 'b']:
connection.execute(
sa.insert(models.MonitoringMessage).values(
id=bot,
workspace_uuid='workspace',
timestamp=now,
**resource(bot),
session_id='person_42',
message_content=bot,
role='user',
status='success',
level='info',
)
)
indexes = {i['name'] for i in sa.inspect(connection).get_indexes('monitoring_sessions')}
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0023_bot_scoped_sessions')
with Operations.context(MigrationContext.configure(connection)):
migration.upgrade()
migration.upgrade() # Fresh/already-upgraded schema is safe.
assert sa.inspect(connection).get_pk_constraint('monitoring_sessions')['constrained_columns'] == [
'workspace_uuid',
'bot_id',
'session_id',
]
assert indexes <= {i['name'] for i in sa.inspect(connection).get_indexes('monitoring_sessions')}
await conn.run_sync(upgrade)
rows, total = await service.get_sessions(context())
assert total == 2
assert {r['bot_id']: r['message_count'] for r in rows} == {'a': 1, 'b': 1}
assert {r['pipeline_id'] for r in rows} == {'a', 'b'}
def context(bot=None):
return ExecutionContext(instance_uuid='test', workspace_uuid='workspace', placement_generation=1, bot_uuid=bot)
def resource(bot):
return dict(bot_id=bot, bot_name=bot, pipeline_id=bot, pipeline_name=bot)
@pytest.fixture
async def service():
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
class Persistence:
serialize_model = PersistenceManager.serialize_model
def get_db_engine(self):
return engine
async def execute_async(self, stmt):
async with engine.begin() as conn:
return await conn.execute(stmt)
ap = SimpleNamespace(persistence_mgr=Persistence(), logger=logging.getLogger(__name__))
ap.monitoring_service = MonitoringService(ap)
yield ap.monitoring_service
await engine.dispose()
async def test_helper_first_message_count_and_two_bot_isolation(service):
for bot in ['a', 'b', 'a']:
query = SimpleNamespace(
_execution_context=context(bot),
launcher_type='person',
launcher_id=42,
sender_id=42,
message_chain=SimpleNamespace(model_dump=lambda: []),
)
assert await MonitoringHelper.record_query_start(service.ap, query, **resource(bot))
rows, total = await service.get_sessions(context())
assert total == 2
assert {r['bot_id']: r['message_count'] for r in rows} == {'a': 2, 'b': 1}
assert {r['pipeline_id'] for r in rows} == {'a', 'b'}
assert {r['session_id'] for r in rows} == {'person_42'}
async def test_analysis_fails_closed_and_scopes_statistics(service):
for bot in ['a', 'b']:
await service.record_session_start(context(bot), session_id='person_42', **resource(bot))
await service.record_message(context(bot), session_id='person_42', message_content=bot, **resource(bot))
assert (await service.get_session_analysis(context(), 'person_42'))['found'] is False
result = await service.get_session_analysis(context(), 'person_42', bot_id='b')
assert result['message_stats']['total'] == 1
assert result['session']['bot_id'] == 'b'
async def test_activity_requires_bot_and_upsert_counts_racing_first_queries(service):
for _ in range(2):
await service.record_session_start(context('a'), session_id='person_42', **resource('a'))
with pytest.raises(ValueError, match='bot'):
await service.update_session_activity(context(), 'person_42')
assert await service.update_session_activity(context('a'), 'person_42')
assert not await service.update_session_activity(context('b'), 'person_42')
rows, _ = await service.get_sessions(context())
assert rows[0]['message_count'] == 3
async def test_old_active_sessions_are_listed_exported_and_not_cleaned(service):
for bot in ['a', 'b']:
await service.record_session_start(context(bot), session_id='person_42', **resource(bot))
old = dt.datetime(2000, 1, 1)
await service.ap.persistence_mgr.execute_async(sa.update(models.MonitoringSession).values(start_time=old))
await service.ap.persistence_mgr.execute_async(
sa.update(models.MonitoringSession).where(models.MonitoringSession.bot_id == 'a').values(last_activity=old)
)
since = dt.datetime.now(dt.timezone.utc).replace(tzinfo=None) - dt.timedelta(days=1)
rows, total = await service.get_sessions(context(), start_time=since)
assert total == 1 and rows[0]['bot_id'] == 'b'
assert len(await service.export_sessions(context(), start_time=since)) == 1
count = await service._delete_expired_in_batches(
context(),
models.MonitoringSession,
models.MonitoringSession.last_activity,
models.MonitoringSession.session_id,
since,
1,
2,
)
assert count == 1
rows, total = await service.get_sessions(context())
assert total == 1 and rows[0]['bot_id'] == 'b'
@@ -0,0 +1,125 @@
from __future__ import annotations
import datetime
from types import SimpleNamespace
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.monitoring import MonitoringLLMCall, MonitoringMessage
from langbot.pkg.entity.persistence.workspace import Workspace
pytestmark = pytest.mark.asyncio
A = '00000000-0000-0000-0000-00000000000a'
B = '00000000-0000-0000-0000-00000000000b'
START = datetime.datetime(2026, 1, 1)
@pytest.fixture
async def traffic_app():
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
await connection.execute(
sqlalchemy.insert(Workspace),
[
{'uuid': wid, 'instance_uuid': 'instance', 'name': wid, 'slug': wid, 'source': 'cloud_projection'}
for wid in (A, B)
],
)
for wid, bot, count in [(A, 'bot-a', 60), (A, 'bot-b', 7), (B, 'bot-a', 9)]:
common = {
'workspace_uuid': wid,
'timestamp': START,
'bot_id': bot,
'bot_name': bot,
'pipeline_id': 'pipeline',
'pipeline_name': 'Pipeline',
'session_id': 'person_42',
'status': 'success',
}
await connection.execute(
sqlalchemy.insert(MonitoringMessage),
[
dict(common, id=f'{wid}-{bot}-{i}', message_content='test fixture', level='info', role='user')
for i in range(count)
],
)
await connection.execute(
sqlalchemy.insert(MonitoringLLMCall),
[
dict(
common,
id=f'{wid}-{bot}-{i}',
model_name='fixture-model',
input_tokens=1,
output_tokens=1,
total_tokens=2,
duration=1,
)
for i in range(count)
],
)
class Persistence:
def get_db_engine(self):
return engine
async def execute_async(self, statement):
async with engine.connect() as connection:
return await connection.execute(statement)
yield SimpleNamespace(persistence_mgr=Persistence())
await engine.dispose()
async def test_traffic_counts_all_rows_not_just_latest_page(traffic_app):
from langbot.pkg.api.http.service.monitoring_traffic import get_traffic_series
context = ExecutionContext(instance_uuid='instance', workspace_uuid=A, placement_generation=1)
result = await get_traffic_series(
traffic_app, context, bot_ids=['bot-a'], start_time=START, end_time=START + datetime.timedelta(hours=2)
)
assert result['bucket'] == 'hour'
assert result['truncated'] is False
assert sum(point['messages'] for point in result['points']) == 60
assert sum(point['llm_calls'] for point in result['points']) == 60
assert len(result['points']) == 3
assert result['points'][1]['messages'] == result['points'][1]['llm_calls'] == 0
assert result['points'][0]['timestamp'] == '2026-01-01T00:00:00Z'
async def test_traffic_workspace_pipeline_and_empty_filters(traffic_app):
from langbot.pkg.api.http.service.monitoring_traffic import get_traffic_series
context = ExecutionContext(instance_uuid='instance', workspace_uuid=B, placement_generation=1)
kwargs = dict(start_time=START, end_time=START + datetime.timedelta(hours=2))
result = await get_traffic_series(traffic_app, context, **kwargs)
assert sum(point['messages'] for point in result['points']) == 9
empty = await get_traffic_series(traffic_app, context, pipeline_ids=['missing'], **kwargs)
assert sum(point['messages'] for point in empty['points']) == 0
assert sum(point['llm_calls'] for point in empty['points']) == 0
async def test_traffic_bounds_large_ranges_and_marks_truncation(traffic_app):
from langbot.pkg.api.http.service.monitoring_traffic import get_traffic_series
context = ExecutionContext(instance_uuid='instance', workspace_uuid=A, placement_generation=1)
result = await get_traffic_series(
traffic_app, context, start_time=START, end_time=START + datetime.timedelta(days=5000)
)
assert result['bucket'] == 'day'
assert result['truncated'] is True
assert len(result['points']) == 1000
async def test_traffic_fails_closed_without_workspace(traffic_app):
from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.service.monitoring_traffic import get_traffic_series
with pytest.raises(WorkspaceRequiredError):
await get_traffic_series(traffic_app, None)