From 3b1ea56f4fdfac9887f98aebcc469030457437b9 Mon Sep 17 00:00:00 2001 From: BiFangKNT <1320414964@qq.com> Date: Tue, 20 Oct 2026 10:45:44 +0800 Subject: [PATCH] fix(monitoring): avoid historical payload reads in execution filters - Add covering indexes through migration 0040 - Preserve PostgreSQL FORCE RLS state when marking historical summaries - Add query-plan and migration regression tests to CI --- .../pkg/entity/persistence/monitoring.py | 27 ++- .../versions/0040_session_summaries.py | 80 +++++++- .../persistence/test_migrations.py | 57 ++++++ .../persistence/test_migrations_postgres.py | 190 ++++++++++++++++++ .../api/service/test_monitoring_executions.py | 84 ++++++++ 5 files changed, 426 insertions(+), 12 deletions(-) diff --git a/src/langbot/pkg/entity/persistence/monitoring.py b/src/langbot/pkg/entity/persistence/monitoring.py index 72047c04b..6392363da 100644 --- a/src/langbot/pkg/entity/persistence/monitoring.py +++ b/src/langbot/pkg/entity/persistence/monitoring.py @@ -40,9 +40,34 @@ class MonitoringMessage(Base): __table_args__ = ( sqlalchemy.Index('ix_monitoring_messages_summary_event', 'workspace_uuid', 'bot_id', 'event_id', 'role'), - sqlalchemy.Index('ix_monitoring_messages_workspace_timestamp', 'workspace_uuid', 'timestamp'), + # Time-windowed aggregates must not fetch execution metadata past historical payloads. + sqlalchemy.Index( + 'ix_monitoring_messages_workspace_timestamp', + 'workspace_uuid', + 'timestamp', + 'id', + 'status', + 'role', + 'bot_id', + 'session_id', + 'parent_message_id', + 'run_id', + 'pipeline_id', + ), sqlalchemy.Index('ix_monitoring_messages_media_ttl', 'workspace_uuid', 'timestamp', 'id'), sqlalchemy.Index('ix_monitoring_messages_workspace_session', 'workspace_uuid', 'session_id'), + # Execution membership must not traverse historical payload overflow pages. + sqlalchemy.Index( + 'ix_monitoring_messages_execution_filter', + 'workspace_uuid', + 'status', + 'role', + 'bot_id', + 'id', + 'session_id', + 'parent_message_id', + 'run_id', + ), sqlalchemy.Index( 'ix_monitoring_messages_conversation_page', 'workspace_uuid', 'bot_id', 'session_id', 'timestamp', 'id' ), diff --git a/src/langbot/pkg/persistence/alembic/versions/0040_session_summaries.py b/src/langbot/pkg/persistence/alembic/versions/0040_session_summaries.py index 038779bfd..17b9b25d2 100644 --- a/src/langbot/pkg/persistence/alembic/versions/0040_session_summaries.py +++ b/src/langbot/pkg/persistence/alembic/versions/0040_session_summaries.py @@ -1,5 +1,7 @@ """Persist idempotent conversation summaries independently of message history.""" +import contextlib + from alembic import op import sqlalchemy as sa @@ -9,19 +11,72 @@ branch_labels = None depends_on = None TABLE = 'monitoring_session_receipts' +TIME_INDEX = 'ix_monitoring_messages_workspace_timestamp' INDEXES = ( ('monitoring_messages', 'ix_monitoring_messages_summary_event', ['workspace_uuid', 'bot_id', 'event_id', 'role']), ('agent_run', 'ix_agent_run_summary_event', ['workspace_id', 'bot_id', 'event_id', 'id']), + ( + 'monitoring_messages', + 'ix_monitoring_messages_execution_filter', + ['workspace_uuid', 'status', 'role', 'bot_id', 'id', 'session_id', 'parent_message_id', 'run_id'], + ), + ( + 'monitoring_messages', + TIME_INDEX, + [ + 'workspace_uuid', + 'timestamp', + 'id', + 'status', + 'role', + 'bot_id', + 'session_id', + 'parent_message_id', + 'run_id', + 'pipeline_id', + ], + ), ) +@contextlib.contextmanager +def _suspend_summary_forced_rls(connection): + """Let the table-owning migrator mark history without a Workspace scope.""" + if connection.dialect.name != 'postgresql': + yield + return + tables = ( + connection.execute( + sa.text( + 'SELECT relname FROM pg_class WHERE relnamespace = current_schema()::regnamespace ' + 'AND relname IN :tables AND relforcerowsecurity' + ).bindparams(sa.bindparam('tables', expanding=True)), + {'tables': ('workspaces', 'workspace_metadata', 'monitoring_sessions', 'monitoring_messages', 'event_log')}, + ) + .scalars() + .all() + ) + quote = connection.dialect.identifier_preparer.quote + try: + for table in tables: + connection.execute(sa.text(f'ALTER TABLE {quote(table)} NO FORCE ROW LEVEL SECURITY')) + yield + finally: + for table in tables: + connection.execute(sa.text(f'ALTER TABLE {quote(table)} FORCE ROW LEVEL SECURITY')) + + def upgrade(): connection = op.get_bind() for table, index, fields in INDEXES: if not sa.inspect(connection).has_table(table): continue - if index not in {item['name'] for item in sa.inspect(connection).get_indexes(table)}: - op.create_index(index, table, fields) + existing = {item['name']: item['column_names'] for item in sa.inspect(connection).get_indexes(table)} + if existing.get(index) == fields: + continue + if index in existing: + op.drop_index(index, table_name=table) + op.create_index(index, table, fields) # Bare Alembic and partial-schema upgrades rely on ORM bootstrap for core tables. if not sa.inspect(connection).has_table('monitoring_sessions'): return @@ -70,15 +125,16 @@ def upgrade(): sources.append('EXISTS (SELECT 1 FROM monitoring_messages m WHERE m.workspace_uuid=w.uuid)') if sa.inspect(connection).has_table('event_log'): sources.append('EXISTS (SELECT 1 FROM event_log e WHERE e.workspace_id=w.uuid)') - connection.execute( - sa.text( - 'INSERT INTO workspace_metadata (workspace_uuid, key, value) ' - "SELECT w.uuid, 'session_summary_backfill_v1', :cursor FROM workspaces w WHERE (" - + ' OR '.join(sources) - + ') ON CONFLICT (workspace_uuid, key) DO UPDATE SET value=excluded.value' - ), - {'cursor': '{"source": "messages"}'}, - ) + with _suspend_summary_forced_rls(connection): + connection.execute( + sa.text( + 'INSERT INTO workspace_metadata (workspace_uuid, key, value) ' + "SELECT w.uuid, 'session_summary_backfill_v1', :cursor FROM workspaces w WHERE (" + + ' OR '.join(sources) + + ') ON CONFLICT (workspace_uuid, key) DO UPDATE SET value=excluded.value' + ), + {'cursor': '{"source": "messages"}'}, + ) collation = '"C"' if connection.dialect.name == 'postgresql' else 'BINARY' op.execute( sa.text( @@ -106,6 +162,8 @@ def downgrade(): if sa.inspect(connection).has_table(table): if index in {item['name'] for item in sa.inspect(connection).get_indexes(table)}: op.drop_index(index, table_name=table) + if index == TIME_INDEX: + op.create_index(index, table, ['workspace_uuid', 'timestamp']) if sa.inspect(connection).has_table(TABLE): op.drop_table(TABLE) if sa.inspect(connection).has_table('monitoring_sessions'): diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py index 602899195..cbe7a8cba 100644 --- a/tests/integration/persistence/test_migrations.py +++ b/tests/integration/persistence/test_migrations.py @@ -122,6 +122,63 @@ class TestSQLiteMigrationBaseline: class TestSQLiteMigrationUpgrade: """Tests for upgrade to head workflow.""" + @pytest.mark.asyncio + async def test_execution_filter_index_round_trip_preserves_history(self, sqlite_engine): + from langbot.pkg.entity.persistence.monitoring import MonitoringMessage + + index_name = 'ix_monitoring_messages_execution_filter' + index = next(item for item in MonitoringMessage.__table__.indexes if item.name == index_name) + + def indexes(connection): + return { + item['name']: item['column_names'] + for item in sqlalchemy.inspect(connection).get_indexes('monitoring_messages') + } + + async with sqlite_engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + await connection.execute( + sqlalchemy.insert(MonitoringMessage).values( + id='historical-root', + timestamp=sqlalchemy.func.current_timestamp(), + workspace_uuid='workspace', + bot_id='bot', + bot_name='Bot', + pipeline_id='pipeline', + pipeline_name='Pipeline', + session_id='session', + message_content='preserve historical payload', + status='success', + role='user', + level='info', + ) + ) + original = await connection.run_sync(indexes) + + await run_alembic_stamp(sqlite_engine, '0039_execution_lease_indexes') + # Fresh installations already have the ORM index before Alembic runs. + await run_alembic_upgrade(sqlite_engine, 'head') + await run_alembic_downgrade(sqlite_engine, '0039_execution_lease_indexes') + async with sqlite_engine.connect() as connection: + remaining = await connection.run_sync(indexes) + expected = { + name: columns + for name, columns in original.items() + if name not in {index_name, 'ix_monitoring_messages_summary_event'} + } + expected['ix_monitoring_messages_workspace_timestamp'] = ['workspace_uuid', 'timestamp'] + assert remaining == expected + + # Existing installations receive the same covering index through migration. + await run_alembic_upgrade(sqlite_engine, 'head') + async with sqlite_engine.connect() as connection: + upgraded = await connection.run_sync(indexes) + assert upgraded == original + assert upgraded[index_name] == [column.name for column in index.columns] + assert await connection.scalar(sqlalchemy.select(MonitoringMessage.message_content)) == ( + 'preserve historical payload' + ) + @pytest.mark.asyncio async def test_upgrade_removes_agent_enabled_column(self, sqlite_engine): async with sqlite_engine.begin() as conn: diff --git a/tests/integration/persistence/test_migrations_postgres.py b/tests/integration/persistence/test_migrations_postgres.py index cb1bd807a..6821b3bea 100644 --- a/tests/integration/persistence/test_migrations_postgres.py +++ b/tests/integration/persistence/test_migrations_postgres.py @@ -50,6 +50,7 @@ from langbot.pkg.persistence.tenant_uow import ( ) from langbot.pkg.persistence.alembic_runner import ( run_alembic_upgrade, + run_alembic_downgrade, run_alembic_stamp, get_alembic_current, _ALEMBIC_DIR, @@ -424,6 +425,195 @@ class TestPostgreSQLMigrationBaseline: class TestPostgreSQLMigrationUpgrade: """Tests for upgrade to head workflow on PostgreSQL.""" + @pytest.mark.asyncio + @pytest.mark.postgres + @pytest.mark.session_summary + @pytest.mark.contract('summary_schema') + @pytest.mark.contract('migration_history') + @pytest.mark.parametrize('mixed_rls', [False, True], ids=['forced-rls', 'mixed-rls']) + async def test_summary_migration_marks_all_history_as_table_owner( + self, postgres_url, postgres_engine, clean_tables, clean_alembic_version, mixed_rls + ): + role_name = f'lb_summary_{uuid.uuid4().hex[:12]}' + password = uuid.uuid4().hex + quote = postgres_engine.dialect.identifier_preparer.quote + owner_engine = create_async_engine( + sa.engine.make_url(postgres_url).set(username=role_name, password=password), + connect_args={'server_settings': {'search_path': f'{role_name},public'}}, + ) + tables = { + 'workspaces': 'uuid', + 'workspace_metadata': 'workspace_uuid', + 'monitoring_sessions': 'workspace_uuid', + 'monitoring_messages': 'workspace_uuid', + 'event_log': 'workspace_id', + } + states = {name: (True, True) for name in tables} + if mixed_rls: + states.update( + workspaces=(False, True), monitoring_sessions=(True, False), monitoring_messages=(False, False) + ) + try: + async with postgres_engine.begin() as connection: + await connection.execute( + text(f"CREATE ROLE {quote(role_name)} LOGIN NOSUPERUSER NOBYPASSRLS PASSWORD '{password}'") + ) + await connection.execute(text(f'CREATE SCHEMA {quote(role_name)} AUTHORIZATION {quote(role_name)}')) + async with owner_engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + await run_alembic_stamp(owner_engine, '0040_session_summaries') + await run_alembic_downgrade(owner_engine, '0039_execution_lease_indexes') + async with owner_engine.begin() as connection: + assert ( + await connection.execute( + text('SELECT rolsuper, rolbypassrls FROM pg_roles WHERE rolname = current_user') + ) + ).one() == (False, False) + for name in ('sessions', 'messages', 'events', 'empty'): + await connection.execute( + sa.insert(Workspace).values( + uuid=name, instance_uuid='instance', name=name, slug=name, source='cloud_projection' + ) + ) + now = datetime.datetime(2026, 1, 1) + # Insert only legacy columns after downgrading the ORM-created schema. + await connection.execute( + text( + 'INSERT INTO monitoring_sessions (workspace_uuid, bot_id, session_id, bot_name, ' + 'pipeline_id, pipeline_name, message_count, start_time, last_activity, is_active) ' + "VALUES ('sessions', 'bot', 'session', 'Bot', 'pipeline', 'Pipeline', 7, :now, :now, true)" + ), + {'now': now}, + ) + await connection.execute( + Base.metadata.tables['monitoring_messages'] + .insert() + .values( + id='message', + workspace_uuid='messages', + timestamp=now, + bot_id='bot', + bot_name='Bot', + pipeline_id='pipeline', + pipeline_name='Pipeline', + session_id='session', + message_content='history', + role='user', + status='success', + level='info', + ) + ) + await connection.execute( + Base.metadata.tables['event_log'] + .insert() + .values( + event_id='event', + event_type='message.received', + source='platform', + workspace_id='events', + bot_id='bot', + conversation_id='session', + ) + ) + for name, column in tables.items(): + enabled, forced = states[name] + await connection.execute( + text(f'ALTER TABLE {quote(name)} {"ENABLE" if enabled else "DISABLE"} ROW LEVEL SECURITY') + ) + await connection.execute( + text(f'ALTER TABLE {quote(name)} {"FORCE" if forced else "NO FORCE"} ROW LEVEL SECURITY') + ) + await connection.execute( + text( + f'CREATE POLICY {TENANT_POLICY_NAME} ON {quote(name)} ' + f"USING ({column}::text = NULLIF(current_setting('langbot.workspace_uuid', true), '')) " + f"WITH CHECK ({column}::text = NULLIF(current_setting('langbot.workspace_uuid', true), ''))" + ) + ) + await run_alembic_upgrade(owner_engine, 'head') + async with owner_engine.connect() as connection: + rows = await connection.execute( + text( + 'SELECT relname, relrowsecurity, relforcerowsecurity FROM pg_class ' + 'WHERE relnamespace = current_schema()::regnamespace AND relname IN :tables' + ).bindparams(sa.bindparam('tables', expanding=True)), + {'tables': tuple(tables)}, + ) + assert {name: (enabled, forced) for name, enabled, forced in rows} == states + assert await connection.scalar(sa.select(sa.func.count()).select_from(WorkspaceMetadata)) == 0 + for name in ('sessions', 'messages', 'events', 'empty'): + async with TenantUnitOfWork(owner_engine, name) as uow: + pending = ( + await uow.execute( + sa.select(WorkspaceMetadata.value).where( + WorkspaceMetadata.key == 'session_summary_backfill_v1' + ) + ) + ).scalar_one_or_none() + assert pending == (None if name == 'empty' else '{"source": "messages"}') + finally: + await owner_engine.dispose() + async with postgres_engine.begin() as connection: + await connection.execute(text(f'DROP SCHEMA IF EXISTS {quote(role_name)} CASCADE')) + await connection.execute(text(f'DROP ROLE IF EXISTS {quote(role_name)}')) + + @pytest.mark.asyncio + @pytest.mark.postgres + @pytest.mark.contract('migration_history') + async def test_execution_filter_index_upgrades_existing_history( + self, postgres_engine, clean_tables, clean_alembic_version + ): + from langbot.pkg.entity.persistence.monitoring import MonitoringMessage + + index_name = 'ix_monitoring_messages_execution_filter' + index = next(item for item in MonitoringMessage.__table__.indexes if item.name == index_name) + + async with postgres_engine.begin() as connection: + await connection.execute( + text( + 'CREATE TABLE monitoring_messages (' + 'id VARCHAR(255) PRIMARY KEY, workspace_uuid VARCHAR(36) NOT NULL, ' + 'status VARCHAR(50), role VARCHAR(50), bot_id VARCHAR(255), ' + 'session_id VARCHAR(255), parent_message_id VARCHAR(255), run_id VARCHAR(255), ' + 'timestamp TIMESTAMP, pipeline_id VARCHAR(255), event_id VARCHAR(255), message_content TEXT)' + ) + ) + await connection.execute( + text( + "INSERT INTO monitoring_messages (id, workspace_uuid, message_content) VALUES ('old', 'w', 'keep')" + ) + ) + await connection.execute( + text( + 'CREATE INDEX ix_monitoring_messages_workspace_timestamp ' + 'ON monitoring_messages (workspace_uuid, timestamp)' + ) + ) + await run_alembic_stamp(postgres_engine, '0039_execution_lease_indexes') + await run_alembic_upgrade(postgres_engine, 'head') + async with postgres_engine.connect() as connection: + indexes = await connection.run_sync(lambda sync: sa.inspect(sync).get_indexes('monitoring_messages')) + assert {item['name']: item['column_names'] for item in indexes}[index_name] == [ + column.name for column in index.columns + ] + time_index = next( + item + for item in MonitoringMessage.__table__.indexes + if item.name == 'ix_monitoring_messages_workspace_timestamp' + ) + assert {item['name']: item['column_names'] for item in indexes}[time_index.name] == [ + column.name for column in time_index.columns + ] + await run_alembic_downgrade(postgres_engine, '0039_execution_lease_indexes') + async with postgres_engine.connect() as connection: + indexes = await connection.run_sync(lambda sync: sa.inspect(sync).get_indexes('monitoring_messages')) + assert index_name not in {item['name'] for item in indexes} + assert {item['name']: item['column_names'] for item in indexes}[time_index.name] == [ + 'workspace_uuid', + 'timestamp', + ] + assert await connection.scalar(text('SELECT message_content FROM monitoring_messages')) == 'keep' + @pytest.mark.asyncio @pytest.mark.parametrize('settings_type', ['TEXT', 'JSON']) async def test_legacy_knowledge_restore_statement_accepts_historical_and_fresh_settings_types( diff --git a/tests/unit_tests/api/service/test_monitoring_executions.py b/tests/unit_tests/api/service/test_monitoring_executions.py index 39587f245..ca376baa8 100644 --- a/tests/unit_tests/api/service/test_monitoring_executions.py +++ b/tests/unit_tests/api/service/test_monitoring_executions.py @@ -403,6 +403,90 @@ async def test_overview_metrics_keep_message_rollup(engine, service): assert overview['active_sessions'] == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize('filtered', [False, True]) +@pytest.mark.parametrize( + ('mode', 'statuses', 'expected'), + [ + ('real', None, (3, 2, 1)), + ('debug', None, (0, 0, 0)), + ('all', ['completed'], (2, 1, 1)), + ('real', ['completed'], (2, 1, 1)), + ('real', ['failed'], (0, 0, 0)), + ('debug', ['completed'], (0, 0, 0)), + ], +) +async def test_overview_execution_filters_do_not_read_historical_payload_pages( + engine, service, mode, statuses, expected, filtered +): + """Execution membership must read root and link metadata without accessing wide history rows.""" + import re + import sqlalchemy as sa + + await _seed(engine) + message = persistence_monitoring.MonitoringMessage + payload = 'historical payload ' * 65536 + async with engine.begin() as connection: + await connection.execute(sa.update(message).values(message_content=payload, variables=payload)) + await connection.execute( + sa.update(message).where(message.id == 'msg-assistant').values(parent_message_id='msg-user') + ) + # Vary status selectivity so the planner cannot hide table reads behind a tiny fixture. + await connection.execute( + sa.insert(message), + [ + dict( + id=f'unlinked-{index}', + workspace_uuid=WORKSPACE, + timestamp=_dt(12), + bot_id='bot-1', + bot_name='Bot', + pipeline_id='pipeline-1', + pipeline_name='Pipeline', + session_id='session-1', + message_content='unlinked assistant output', + role='assistant', + status=('success', 'error', 'discarded', 'skipped', 'rejected', 'running', 'pending')[index % 7], + level='info', + ) + for index in range(64) + ], + ) + await connection.exec_driver_sql('ANALYZE') + + 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: + filters = ( + dict(bot_ids=['bot-1'], pipeline_ids=['pipeline-1'], start_time=_dt(9), end_time=_dt(11)) + if filtered + else {} + ) + overview = await service.get_overview_metrics(WORKSPACE, execution_statuses=statuses, mode=mode, **filters) + finally: + sa.event.remove(engine.sync_engine, 'before_cursor_execute', capture) + + assert (overview['total_messages'], overview['llm_calls'], overview['active_sessions']) == expected + checked_accesses = 0 + async with engine.connect() as connection: + for sql, parameters in statements: + # Discover aliases from the executed SQL rather than pinning an alias or index name. + sources = {message.__tablename__, *re.findall(r'FROM monitoring_messages AS (\w+)', sql)} + plan = (await connection.exec_driver_sql('EXPLAIN QUERY PLAN ' + sql, parameters)).all() + accesses = [ + row[3] for row in plan if row[3].split()[0] in {'SEARCH', 'SCAN'} and row[3].split()[1] in sources + ] + assert accesses, plan + assert all('COVERING INDEX' in access for access in accesses), plan + checked_accesses += len(accesses) + assert checked_accesses > 0 + + @pytest.mark.asyncio async def test_execution_detail_for_agent_and_pipeline(engine, service): await _seed(engine)