From b73f66a2343dee6ecdb5d27db644661bbf63b4c1 Mon Sep 17 00:00:00 2001 From: BiFangKNT <1320414964@qq.com> Date: Sun, 18 Oct 2026 12:26:09 +0800 Subject: [PATCH] feat(persistence): enforce image reference storage contract - Reject structured inline images on ORM and Core writes - Route plugin run results through the media-aware journal - Redact legacy inline images before merging run metadata - Cover storage contracts and regressions on SQLite and PostgreSQL --- AGENTS.md | 1 + ARCHITECTURE.md | 1 + .../pkg/agent/runner/run_ledger_store.py | 6 +- .../pkg/api/http/service/monitoring.py | 4 +- .../pkg/entity/persistence/agent_run.py | 7 +- .../pkg/entity/persistence/event_log.py | 5 +- .../pkg/entity/persistence/monitoring.py | 5 +- .../pkg/entity/persistence/transcript.py | 7 +- src/langbot/pkg/persistence/media.py | 46 +++++ src/langbot/pkg/persistence/tenant_uow.py | 6 +- src/langbot/pkg/plugin/runner_actions.py | 12 +- .../test_media_reference_postgres.py | 190 ++++++++++++++++++ .../agent/test_run_ledger_api_auth.py | 48 ++++- .../unit_tests/agent/test_run_ledger_store.py | 76 +++++++ .../persistence/test_media_reference_text.py | 139 +++++++++++++ .../storage/test_inline_media_cleanup.py | 2 + .../storage/test_media_persistence.py | 43 ++++ 17 files changed, 577 insertions(+), 21 deletions(-) create mode 100644 src/langbot/pkg/persistence/media.py create mode 100644 tests/integration/persistence/test_media_reference_postgres.py create mode 100644 tests/unit_tests/persistence/test_media_reference_text.py diff --git a/AGENTS.md b/AGENTS.md index bc016a26e..96f8c6903 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -88,6 +88,7 @@ Config keys to verify in `data/config.yaml` / `src/langbot/templates/config.yaml - New schema changes use Alembic under `src/langbot/pkg/persistence/alembic/versions/`. LangBot 4.x does not support upgrading 3.x databases. - New platform behavior belongs in platform adapters only for platform translation; pipeline/business logic belongs in `pkg/pipeline/` or services. - User-facing strings must support i18n; follow the backend language rules below for server-generated content. +- Conversation payload columns, including new ones, must use `MediaReferenceText` and persist image references only. See the [media persistence contract](ARCHITECTURE.md#persistence-and-configuration). - Code comments and docstrings must be English. - Keep compatibility and security in mind; LangBot is used in both self-hosted/community and toB deployments. - Commit message format: `(): `. diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index e4d12fc51..f99f926a8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -267,6 +267,7 @@ Persistence is centered on `pkg/persistence/mgr.py`. - SQLite is the default database; PostgreSQL is supported. - Models live under `pkg/entity/persistence/`. +- `pkg/persistence/media.py::MediaReferenceText` validates bound conversation payloads on ORM/Core writes while preserving the existing TEXT schema and string API. It rejects recognized structured inline images: retained images go through `MediaCache.externalize_chain_dump()` before serialization, while metadata-only snapshots and transcripts use `redact_inline_images()`. Reads and comparisons preserve historical values for incremental cleanup. Runtime/adapter messages, arbitrary diagnostic text, and opaque strings are outside this contract. - Timezone-less SQL `DateTime` columns store UTC-naive values on both backends. Normalize aware values with `pkg/persistence/datetime_utils.py::as_naive_utc` at the bind boundary, including deadlines, leases, event times, and query/retention cutoffs; do not merely strip an offset. Existing naive rows represent UTC. Restore UTC awareness for application comparisons and epoch serialization. Runner ledger/event/transcript stores follow this existing schema contract without a data migration or changes to journal/tenant authorization. - Fresh schemas are created from current metadata, then Alembic migrations run to head. LangBot 4.x does not upgrade 3.x databases. - New schema changes use Alembic under `pkg/persistence/alembic/versions/`; there is no legacy migration chain in 4.x. diff --git a/src/langbot/pkg/agent/runner/run_ledger_store.py b/src/langbot/pkg/agent/runner/run_ledger_store.py index 2a8b9c75b..06ec1d969 100644 --- a/src/langbot/pkg/agent/runner/run_ledger_store.py +++ b/src/langbot/pkg/agent/runner/run_ledger_store.py @@ -12,6 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession from sqlalchemy.orm import sessionmaker from ...persistence.datetime_utils import as_naive_utc +from ...storage.media import MediaCache from ...utils.inflight import inflight_hub from ...utils.execution_lease import host_execution_lease @@ -384,7 +385,7 @@ class RunLedgerStore: if run is None or run.workspace_id != workspace_id or run.status in TERMINAL_STATUSES: raise ValueError('Skill activation requires an active run in the current Workspace') original = run.metadata_json - metadata = _json_loads(original, {}) + metadata = MediaCache.redact_inline_images(_json_loads(original, {})) bindings = metadata.setdefault('activated_skills', {}) if name in bindings: return str(bindings[name]) @@ -439,7 +440,8 @@ class RunLedgerStore: if cost is not None: run.cost_json = _json_dumps(cost) if metadata is not None: - existing_metadata = _json_loads(run.metadata_json, {}) + # Discard legacy image bytes before merging; new metadata still uses the write contract. + existing_metadata = MediaCache.redact_inline_images(_json_loads(run.metadata_json, {})) if isinstance(existing_metadata, dict): existing_metadata.update(metadata) run.metadata_json = _json_dumps(existing_metadata) diff --git a/src/langbot/pkg/api/http/service/monitoring.py b/src/langbot/pkg/api/http/service/monitoring.py index 4cd5628ec..0480a7d1b 100644 --- a/src/langbot/pkg/api/http/service/monitoring.py +++ b/src/langbot/pkg/api/http/service/monitoring.py @@ -580,7 +580,7 @@ class MonitoringService(ExecutionDetailsMixin): 'user_id': _normalize_user_id(user_id), 'user_name': user_name, 'runner_name': runner_name, - 'variables': variables, + 'variables': self._sanitize_message_content(variables) if variables is not None else None, 'role': role, 'event_id': event_id, 'run_id': run_id, @@ -888,7 +888,7 @@ class MonitoringService(ExecutionDetailsMixin): if level is not None: update_values['level'] = level if variables is not None: - update_values['variables'] = variables + update_values['variables'] = self._sanitize_message_content(variables) await self.ap.persistence_mgr.execute_async( sqlalchemy.update(persistence_monitoring.MonitoringMessage) diff --git a/src/langbot/pkg/entity/persistence/agent_run.py b/src/langbot/pkg/entity/persistence/agent_run.py index c68ee01f5..05d5c2b00 100644 --- a/src/langbot/pkg/entity/persistence/agent_run.py +++ b/src/langbot/pkg/entity/persistence/agent_run.py @@ -6,6 +6,7 @@ import datetime import sqlalchemy +from ...persistence.media import MediaReferenceText from .base import Base @@ -107,7 +108,7 @@ class AgentRun(Base): authorization_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) """Run-scoped authorization snapshot JSON.""" - metadata_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + metadata_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Additional metadata JSON.""" __table_args__ = ( @@ -185,7 +186,7 @@ class AgentRunEvent(Base): type = sqlalchemy.Column(sqlalchemy.String(100), nullable=False, index=True) """Result event type.""" - data_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + data_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Result event payload JSON.""" usage_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) @@ -197,7 +198,7 @@ class AgentRunEvent(Base): source = sqlalchemy.Column(sqlalchemy.String(50), nullable=True) """Source that appended the event.""" - metadata_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + metadata_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Additional metadata JSON.""" __table_args__ = ( diff --git a/src/langbot/pkg/entity/persistence/event_log.py b/src/langbot/pkg/entity/persistence/event_log.py index da16ee1b5..8132d7ce9 100644 --- a/src/langbot/pkg/entity/persistence/event_log.py +++ b/src/langbot/pkg/entity/persistence/event_log.py @@ -6,6 +6,7 @@ import sqlalchemy import datetime from .base import Base +from ...persistence.media import MediaReferenceText class EventLog(Base): @@ -66,7 +67,7 @@ class EventLog(Base): input_summary = sqlalchemy.Column(sqlalchemy.Text, nullable=True) """Brief summary of input (truncated text, max 1000 chars).""" - input_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + input_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Full input JSON if reasonably sized (AgentInput as JSON string).""" # Raw event reference @@ -82,7 +83,7 @@ class EventLog(Base): created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, default=datetime.datetime.utcnow) """When this record was created.""" - metadata_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + metadata_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Additional metadata as JSON string.""" __table_args__ = (sqlalchemy.Index('ix_event_log_media_ttl', 'workspace_id', 'created_at', 'id'),) diff --git a/src/langbot/pkg/entity/persistence/monitoring.py b/src/langbot/pkg/entity/persistence/monitoring.py index 5f9fdd257..358c1791e 100644 --- a/src/langbot/pkg/entity/persistence/monitoring.py +++ b/src/langbot/pkg/entity/persistence/monitoring.py @@ -1,5 +1,6 @@ import sqlalchemy +from ...persistence.media import MediaReferenceText from .base import Base @@ -19,7 +20,7 @@ class MonitoringMessage(Base): bot_name = sqlalchemy.Column(sqlalchemy.String(255), nullable=False) pipeline_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=False, index=True) pipeline_name = sqlalchemy.Column(sqlalchemy.String(255), nullable=False) - message_content = sqlalchemy.Column(sqlalchemy.Text, nullable=False) + message_content = sqlalchemy.Column(MediaReferenceText, nullable=False) session_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=False, index=True) execution_owner_id = sqlalchemy.Column(sqlalchemy.String(36), nullable=True, index=True) execution_lease_expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True, index=True) @@ -31,7 +32,7 @@ class MonitoringMessage(Base): user_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True) user_name = sqlalchemy.Column(sqlalchemy.String(255), nullable=True) # User display name runner_name = sqlalchemy.Column(sqlalchemy.String(255), nullable=True) # Runner name for this query - variables = sqlalchemy.Column(sqlalchemy.Text, nullable=True) # Query variables as JSON string + variables = sqlalchemy.Column(MediaReferenceText, nullable=True) # Query variables as JSON string role = sqlalchemy.Column(sqlalchemy.String(50), nullable=True, default='user') # user, assistant event_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True) run_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True) diff --git a/src/langbot/pkg/entity/persistence/transcript.py b/src/langbot/pkg/entity/persistence/transcript.py index f19c87d8a..d61be2e3b 100644 --- a/src/langbot/pkg/entity/persistence/transcript.py +++ b/src/langbot/pkg/entity/persistence/transcript.py @@ -6,6 +6,7 @@ import sqlalchemy import datetime from .base import Base +from ...persistence.media import MediaReferenceText class Transcript(Base): @@ -48,11 +49,11 @@ class Transcript(Base): content = sqlalchemy.Column(sqlalchemy.Text, nullable=True) """Text content summary (may be truncated for large messages, max 4000 chars).""" - content_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + content_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Full structured content as JSON string (Message model dump).""" # Attachment references - attachment_refs_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + attachment_refs_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Attachment references as JSON string.""" # Sequence for cursor-based pagination @@ -69,7 +70,7 @@ class Transcript(Base): created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, default=datetime.datetime.utcnow) """When this item was created.""" - metadata_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True) + metadata_json = sqlalchemy.Column(MediaReferenceText, nullable=True) """Additional metadata as JSON string (sender_id, platform, etc.).""" # Indexes diff --git a/src/langbot/pkg/persistence/media.py b/src/langbot/pkg/persistence/media.py new file mode 100644 index 000000000..337c63834 --- /dev/null +++ b/src/langbot/pkg/persistence/media.py @@ -0,0 +1,46 @@ +"""The storage contract for conversation payloads containing image references.""" + +from __future__ import annotations + +import json +from typing import Any + +from sqlalchemy import Text +from sqlalchemy.engine import Dialect +from sqlalchemy.sql.operators import OperatorType +from sqlalchemy.types import TypeDecorator + + +class MediaReferenceText(TypeDecorator[str]): + """Keep the existing TEXT format while rejecting structured inline images. + + Writers must externalize retained images through MediaCache or redact + metadata-only images before binding a value. Both ORM and Core writes use + this contract; reads preserve historical values for incremental cleanup. + Plain text remains valid for fields such as monitoring message content. + """ + + impl = Text + cache_ok = True + + def process_bind_param(self, value: str | None, dialect: Dialect) -> str | None: + if value is None: + return None + try: + payload = json.loads(value) + except json.JSONDecodeError: + return value + + # Model imports must not initialize the media cache's runtime dependencies. + from ..storage.media import MediaCache + + if MediaCache.has_inline_images(payload): + raise ValueError( + 'Inline images cannot be persisted; use MediaCache.externalize_chain_dump() ' + 'for retained content or MediaCache.redact_inline_images() for metadata' + ) + return value + + def coerce_compared_value(self, op: OperatorType | None, value: Any) -> Text: + # Legacy payloads must remain usable in searches and compare-and-swap updates. + return self.impl diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 396b20103..c55b7a252 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -19,6 +19,8 @@ from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLO from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing from sqlalchemy.dialects.sqlite.dml import OnConflictDoUpdate as SQLiteOnConflictDoUpdate +from .media import MediaReferenceText + TENANT_SETTING = 'langbot.workspace_uuid' ACCOUNT_SETTING = 'langbot.account_uuid' @@ -292,7 +294,7 @@ def _validate_scoped_sql_type( if type(sql_type) in {Vector, HALFVEC}: return - if not type(sql_type).__module__.startswith('sqlalchemy.'): + if type(sql_type) is not MediaReferenceText and not type(sql_type).__module__.startswith('sqlalchemy.'): raise ScopedSessionTransactionError('TenantUnitOfWork does not allow custom SQL types in public statements') for identifier_attribute in ('name', 'schema', 'collation'): @@ -486,7 +488,7 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s monitoring_metadata = ( type(element.type) is sqlalchemy.JSON and isinstance(source, sqlalchemy.Column) - and type(source.type) is sqlalchemy.Text + and type(source.type) in {sqlalchemy.Text, MediaReferenceText} and getattr(getattr(source, 'table', None), 'name', None) in {'event_log', 'agent_run'} and ( source.name == 'metadata_json' or (source.table.name == 'event_log' and source.name == 'input_json') diff --git a/src/langbot/pkg/plugin/runner_actions.py b/src/langbot/pkg/plugin/runner_actions.py index bca1a6c07..659ac2f4c 100644 --- a/src/langbot/pkg/plugin/runner_actions.py +++ b/src/langbot/pkg/plugin/runner_actions.py @@ -469,12 +469,16 @@ def register(h): if payload_error: return handler.ActionResponse.error(message=payload_error) - event = await store.append_event( + from ..agent.runner.run_journal import AgentRunJournal + + event = await AgentRunJournal(h.ap).append_run_result( run_id=str(target_run_id), sequence=sequence, - event_type=str(event_type), - data=event_payload, - usage=usage if isinstance(usage, dict) else None, + result_dict={ + 'type': str(event_type), + 'data': event_payload, + 'usage': usage if isinstance(usage, dict) else None, + }, source=str(data.get('source') or result.get('source') or 'runner'), metadata=metadata, ) diff --git a/tests/integration/persistence/test_media_reference_postgres.py b/tests/integration/persistence/test_media_reference_postgres.py new file mode 100644 index 000000000..a9bea397c --- /dev/null +++ b/tests/integration/persistence/test_media_reference_postgres.py @@ -0,0 +1,190 @@ +"""Media writes against real asyncpg and existing TEXT columns. + +Run with TEST_POSTGRES_URL pointing to a disposable PostgreSQL database. +Each case creates and removes its own schema. +""" + +import importlib +import json +import os +import uuid +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine + +from langbot.pkg.agent.runner.run_journal import AgentRunJournal +from langbot.pkg.agent.runner.run_ledger_store import RunLedgerStore +from langbot.pkg.entity.persistence.agent_run import AgentRun, AgentRunEvent +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.persistence.tenant_uow import _validate_scoped_statement_call +from langbot.pkg.storage.media import MediaCache + +pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio] + + +def _published_schema(connection): + with Operations.context(MigrationContext.configure(connection)): + for revision in ( + '58846a8d7a81_add_event_log_and_transcript_tables', + '8d3a1f2c4b6e_add_agent_run_ledger', + '0036_execution_host_leases', + ): + importlib.import_module(f'langbot.pkg.persistence.alembic.versions.{revision}').upgrade() + + +@pytest.fixture(params=['metadata', 'published_migrations']) +async def pg_engine(request): + url = os.environ.get('TEST_POSTGRES_URL') + if not url: + pytest.skip('TEST_POSTGRES_URL is required (disposable PostgreSQL)') + namespace = f'media_contract_{uuid.uuid4().hex}' + admin = create_async_engine(url) + engine = None + try: + async with admin.begin() as conn: + await conn.execute(sa.text(f'CREATE SCHEMA {namespace}')) + engine = create_async_engine(url, connect_args={'server_settings': {'search_path': namespace}}) + async with engine.begin() as conn: + if request.param == 'metadata': + await conn.run_sync( + lambda sync: Base.metadata.create_all(sync, tables=[AgentRun.__table__, AgentRunEvent.__table__]) + ) + else: + await conn.run_sync(_published_schema) + yield engine + finally: + if engine is not None: + await engine.dispose() + async with admin.begin() as conn: + await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS {namespace} CASCADE')) + await admin.dispose() + + +@pytest.fixture +async def store(pg_engine): + store = RunLedgerStore(pg_engine) + await store.create_run( + run_id='run', event_id='event', binding_id=None, runner_id='runner', workspace_id='ws', metadata={'attempt': 0} + ) + return store + + +@pytest.mark.parametrize('writer', ['orm', 'insert', 'update']) +@pytest.mark.parametrize( + 'image', + [ + {'type': 'Image', 'base64': 'eA=='}, + {'type': 'image_base64', 'image_base64': 'eA=='}, + {'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA=='}}, + ], +) +async def test_unprocessed_images_roll_back(pg_engine, store, writer, image): + content = json.dumps({'nested': [image]}) + with pytest.raises(sa.exc.StatementError, match='Inline images cannot be persisted'): + if writer == 'orm': + async with AsyncSession(pg_engine) as session: + run = (await session.execute(sa.select(AgentRun).where(AgentRun.run_id == 'run'))).scalar_one() + run.status = 'completed' + run.metadata_json = content + await session.commit() + else: + statement = ( + sa.insert(AgentRun).values(run_id='new', runner_id='runner', status='completed', metadata_json=content) + if writer == 'insert' + else sa.update(AgentRun) + .where(AgentRun.run_id == 'run') + .values(status='completed', metadata_json=content) + ) + async with pg_engine.begin() as conn: + await conn.execute(statement) + run = await store.get_run('run') + assert run['status'] == 'running' + assert run['metadata'] == {'attempt': 0} + assert await store.get_run('new') is None + + +@pytest.mark.parametrize('operation', ['pin', 'finalize', 'status_only']) +async def test_legacy_metadata_updates_preserve_the_write_contract(pg_engine, store, operation): + legacy = {'image': {'type': 'Image', 'base64': 'eA=='}, 'attempt': 0} + original = json.dumps(legacy) + async with pg_engine.begin() as conn: + # Seed the pre-contract format while keeping current model mappings. + await conn.execute( + sa.update(AgentRun) + .where(AgentRun.run_id == 'run') + .values(metadata_json=sa.bindparam('legacy', original, type_=sa.Text())) + ) + assert ( + await conn.execute(sa.select(AgentRun.run_id).where(AgentRun.metadata_json == original)) + ).scalar_one() == 'run' + + with pytest.raises(sa.exc.StatementError, match='Inline images cannot be persisted'): + await store.finalize_run( + run_id='run', status='completed', metadata={'new': {'type': 'Image', 'base64': 'eQ=='}} + ) + assert (await store.get_run('run'))['status'] == 'running' + assert (await store.get_run('run'))['metadata'] == legacy + + if operation == 'pin': + assert await store.pin_skill_revision('run', 'ws', 'demo', 'v1') == 'v1' + assert (await store.get_run('run'))['metadata'] == { + 'image': {'type': 'Image', 'base64': None}, + 'attempt': 0, + 'activated_skills': {'demo': 'v1'}, + } + else: + metadata = {'note': 'done'} if operation == 'finalize' else None + run = await store.finalize_run(run_id='run', status='completed', metadata=metadata) + assert run['status'] == 'completed' + assert run['metadata'] == ( + {'image': {'type': 'Image', 'base64': None}, 'attempt': 0, 'note': 'done'} + if metadata is not None + else legacy + ) + + +async def test_journal_retains_references_and_redacts_metadata(pg_engine, store): + provider = SimpleNamespace(exists=AsyncMock(return_value=False), save=AsyncMock()) + storage = SimpleNamespace(storage_provider=provider) + app = SimpleNamespace(storage_mgr=storage, persistence_mgr=SimpleNamespace(get_db_engine=lambda: pg_engine)) + storage.media_cache = MediaCache(app, storage) + image = {'type': 'Image', 'base64': 'eA=='} + event = await AgentRunJournal(app).append_run_result( + run_id='run', + sequence=1, + result_dict={'type': 'tool.call.completed', 'data': {'image': image}}, + metadata={'preview': image}, + ) + assert event['data']['image']['base64'] is None + assert event['data']['image']['url'].startswith('/api/v1/files/media/') + assert event['metadata'] == {'preview': {'type': 'Image', 'base64': None}} + provider.save.assert_awaited_once() + async with pg_engine.connect() as conn: + row = (await conn.execute(sa.select(AgentRunEvent.data_json, AgentRunEvent.metadata_json))).one() + assert json.loads(row.data_json) == event['data'] + assert json.loads(row.metadata_json) == event['metadata'] + + +async def test_text_schema_and_monitoring_json_casts(pg_engine, store): + statement = sa.select(sa.cast(AgentRun.metadata_json, sa.JSON)['attempt'].as_integer()) + _validate_scoped_statement_call((statement,), {}) + async with pg_engine.connect() as conn: + assert (await conn.execute(statement)).scalar_one() == 0 + columns = ( + ( + await conn.execute( + sa.text( + 'SELECT data_type FROM information_schema.columns WHERE table_schema = current_schema() ' + "AND table_name IN ('agent_run', 'agent_run_event') AND column_name IN ('data_json', 'metadata_json')" + ) + ) + ) + .scalars() + .all() + ) + assert columns == ['text'] * 3 diff --git a/tests/unit_tests/agent/test_run_ledger_api_auth.py b/tests/unit_tests/agent/test_run_ledger_api_auth.py index 992bbcdfd..d3e6b8f5e 100644 --- a/tests/unit_tests/agent/test_run_ledger_api_auth.py +++ b/tests/unit_tests/agent/test_run_ledger_api_auth.py @@ -3,8 +3,9 @@ from __future__ import annotations import datetime +import json from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest import sqlalchemy @@ -17,6 +18,7 @@ from langbot.pkg.agent.runner.session_registry import AgentRunSessionRegistry from langbot.pkg.entity.persistence import agent_run as agent_run_model from langbot.pkg.entity.persistence.base import Base from langbot.pkg.plugin.handler import RuntimeConnectionHandler +from langbot.pkg.storage.media import MediaCache from langbot_plugin.api.entities.builtin.runner.run_ledger import ( AgentRun, AgentRunEvent, @@ -1045,6 +1047,50 @@ async def test_run_append_result_basic_path(session_registry, db_engine): assert event.usage == {'output_tokens': 1} +@pytest.mark.asyncio +async def test_run_append_result_uses_cas_and_redacts_metadata(session_registry, db_engine): + await _register_session(session_registry, available_apis={'run_append_result': True}) + await _create_run(db_engine) + handler = _handler(db_engine) + objects = {} + + async def save(key, data): + objects[key] = data + + storage = SimpleNamespace( + storage_provider=SimpleNamespace(exists=AsyncMock(return_value=False), save=AsyncMock(side_effect=save)) + ) + storage.media_cache = MediaCache(handler.ap, storage) + handler.ap.storage_mgr = storage + image = {'type': 'Image', 'base64': 'eA=='} + result = await handler.actions[PluginToRuntimeAction.RUN_APPEND_RESULT.value]( + { + 'run_id': 'run_1', + 'caller_plugin_identity': 'test/runner', + 'sequence': 2, + 'event_type': 'tool.call.completed', + 'data': {'tool_call_id': 'call_1', 'tool_name': 'image', 'result': {'images': [image]}}, + 'metadata': {'preview': image}, + } + ) + assert result.code == 0 + event = AgentRunEvent.model_validate(result.data) + saved = event.data['result']['images'][0] + assert saved['base64'] is None + assert saved['url'].startswith('/api/v1/files/media/') + assert objects == {saved['storage_key']: b'x'} + async with db_engine.connect() as conn: + metadata = ( + await conn.execute( + sqlalchemy.select(agent_run_model.AgentRunEvent.metadata_json).where( + agent_run_model.AgentRunEvent.sequence == 2 + ) + ) + ).scalar_one() + assert json.loads(metadata)['preview']['base64'] is None + assert image['base64'] == 'eA==' + + @pytest.mark.asyncio async def test_run_append_result_rejects_side_effecting_result_type(session_registry, db_engine): await _register_session(session_registry, available_apis={'run_append_result': True}) diff --git a/tests/unit_tests/agent/test_run_ledger_store.py b/tests/unit_tests/agent/test_run_ledger_store.py index 22c97baea..2195f86f7 100644 --- a/tests/unit_tests/agent/test_run_ledger_store.py +++ b/tests/unit_tests/agent/test_run_ledger_store.py @@ -3,6 +3,7 @@ from __future__ import annotations import datetime +import json import pytest import sqlalchemy @@ -36,6 +37,81 @@ def store(db_engine): return RunLedgerStore(db_engine) +@pytest.fixture +async def legacy_run(store, db_engine): + await store.create_run( + run_id='legacy', + event_id='event', + binding_id='binding', + runner_id='runner', + workspace_id='ws', + ) + metadata = { + 'input': [ + {'type': 'Image', 'base64': 'eA=='}, + {'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA==', 'detail': 'low'}}, + {'type': 'image_url', 'image_url': {'url': 'https://example.com/image.png'}}, + {'type': 'text', 'text': 'keep this'}, + ], + 'attempt': 0, + 'activated_skills': {'existing': 'v1'}, + } + async with db_engine.begin() as conn: + # Explicit TEXT binding simulates data persisted before the media contract. + await conn.execute( + sqlalchemy.update(AgentRun) + .where(AgentRun.run_id == 'legacy') + .values( + metadata_json=sqlalchemy.bindparam('legacy_metadata', json.dumps(metadata), type_=sqlalchemy.Text()) + ) + ) + return metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize('operation', ['pin', 'finalize']) +async def test_metadata_update_redacts_legacy_images(store, legacy_run, operation): + if operation == 'pin': + assert await store.pin_skill_revision('legacy', 'ws', 'demo', 'v2') == 'v2' + else: + await store.finalize_run(run_id='legacy', status='completed', metadata={'note': 'done'}) + + run = await store.get_run('legacy') + expected = { + **legacy_run, + 'input': [ + {'type': 'Image', 'base64': None}, + {'type': 'image_url', 'image_url': {'url': None, 'detail': 'low'}}, + *legacy_run['input'][2:], + ], + } + if operation == 'pin': + expected['activated_skills'] = {'existing': 'v1', 'demo': 'v2'} + assert run['status'] == 'running' + else: + expected['note'] = 'done' + assert run['status'] == 'completed' + assert run['metadata'] == expected + + +@pytest.mark.asyncio +async def test_finalize_without_metadata_preserves_legacy_snapshot(store, legacy_run): + run = await store.finalize_run(run_id='legacy', status='completed') + assert run['status'] == 'completed' + assert run['metadata'] == legacy_run + + +@pytest.mark.asyncio +async def test_metadata_update_still_rejects_new_inline_images(store, legacy_run): + with pytest.raises(sqlalchemy.exc.StatementError, match='Inline images cannot be persisted'): + await store.finalize_run( + run_id='legacy', status='completed', metadata={'new_image': {'type': 'Image', 'base64': 'eQ=='}} + ) + run = await store.get_run('legacy') + assert run['status'] == 'running' + assert run['metadata'] == legacy_run + + @pytest.mark.asyncio async def test_create_queued_run_claim_renew_release(store): run = await store.create_run( diff --git a/tests/unit_tests/persistence/test_media_reference_text.py b/tests/unit_tests/persistence/test_media_reference_text.py new file mode 100644 index 000000000..9d2d7c899 --- /dev/null +++ b/tests/unit_tests/persistence/test_media_reference_text.py @@ -0,0 +1,139 @@ +"""Image reference contracts apply to ORM and Core writes, including new callers.""" + +import json + +import pytest +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql, sqlite +from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine +from sqlalchemy.orm import DeclarativeBase + +from langbot.pkg.entity.persistence.agent_run import AgentRun, AgentRunEvent +from langbot.pkg.entity.persistence.event_log import EventLog +from langbot.pkg.entity.persistence.monitoring import MonitoringMessage +from langbot.pkg.entity.persistence.transcript import Transcript +from langbot.pkg.persistence.media import MediaReferenceText +from langbot.pkg.persistence.tenant_uow import ScopedSessionTransactionError, _validate_scoped_statement_call + + +class PayloadBase(DeclarativeBase): + pass + + +class Payload(PayloadBase): + __tablename__ = 'payload' + id = sa.Column(sa.Integer, primary_key=True) + content = sa.Column(MediaReferenceText) + + +@pytest.fixture +async def engine(): + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + async with engine.begin() as conn: + await conn.run_sync(PayloadBase.metadata.create_all) + yield engine + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('writer', ['orm', 'insert', 'update']) +@pytest.mark.parametrize( + 'image', + [ + {'type': 'Image', 'base64': 'eA=='}, + {'type': 'Image', 'url': 'data:image/png;base64,eA=='}, + {'type': 'image_base64', 'image_base64': 'eA=='}, + {'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA=='}}, + {'type': 'image', 'image_url': 'data:image/png;base64,eA=='}, + ], +) +async def test_unprocessed_images_are_rejected_at_the_bind_boundary(engine, writer, image): + content = json.dumps({'nested': [image]}) + async with engine.begin() as conn: + await conn.execute(sa.insert(Payload).values(id=1, content='original')) + with pytest.raises(sa.exc.StatementError, match='use MediaCache.externalize_chain_dump'): + if writer == 'orm': + async with AsyncSession(engine) as session: + session.add(Payload(id=2, content=content)) + await session.commit() + else: + statement = ( + sa.insert(Payload).values(id=2, content=content) + if writer == 'insert' + else sa.update(Payload).where(Payload.id == 1).values(content=content) + ) + async with engine.begin() as conn: + await conn.execute(statement) + async with engine.connect() as conn: + assert (await conn.execute(sa.select(Payload.content))).scalars().all() == ['original'] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'content', + [ + None, + 'plain text', + '{historical invalid JSON', + json.dumps({'type': 'Plain', 'text': 'base64 is a transport encoding'}), + json.dumps({'type': 'Image', 'url': 'https://example.com/image.png'}), + json.dumps({'type': 'Image', 'url': '/api/v1/files/media/hash.png', 'base64': None}), + ], +) +async def test_plain_text_and_references_round_trip_unchanged(engine, content): + async with AsyncSession(engine) as session: + session.add(Payload(id=1, content=content)) + await session.commit() + async with engine.connect() as conn: + assert (await conn.execute(sa.select(Payload.content))).scalar_one() == content + + +@pytest.mark.asyncio +async def test_legacy_images_can_be_read_compared_and_redacted(engine): + content = json.dumps({'type': 'Image', 'base64': 'eA=='}) + async with engine.begin() as conn: + # Simulate a row created before the contract was introduced. + await conn.execute(sa.text('INSERT INTO payload (id, content) VALUES (1, :content)'), {'content': content}) + assert ( + await conn.execute(sa.select(Payload.content).where(Payload.content == content)) + ).scalar_one() == content + assert (await conn.execute(sa.select(Payload.id).where(Payload.content.like('%base64%')))).scalar_one() == 1 + updated = await conn.execute( + sa.update(Payload).where(Payload.content == content).values(content='{"type": "Image", "base64": null}') + ) + assert updated.rowcount == 1 + + +@pytest.mark.parametrize('dialect', [sqlite.dialect(), postgresql.dialect()], ids=['sqlite', 'postgresql']) +def test_contract_keeps_text_schema_and_handles_escaped_json_keys(dialect): + column_type = MediaReferenceText() + assert column_type.compile(dialect=dialect) == 'TEXT' + process = column_type.dialect_impl(dialect).bind_processor(dialect) + with pytest.raises(ValueError, match='Inline images cannot be persisted'): + process(r'{"type": "Image", "base\u00364": "eA=="}') + + +@pytest.mark.parametrize( + 'model, fields', + [ + (MonitoringMessage, ['message_content', 'variables']), + (EventLog, ['input_json', 'metadata_json']), + (AgentRun, ['metadata_json']), + (AgentRunEvent, ['data_json', 'metadata_json']), + (Transcript, ['content_json', 'attachment_refs_json', 'metadata_json']), + ], +) +def test_conversation_payload_columns_declare_the_contract(model, fields): + for field in fields: + assert isinstance(model.__table__.c[field].type, MediaReferenceText), f'{model.__tablename__}.{field}' + + +def test_tenant_queries_admit_only_the_exact_trusted_media_type(): + _validate_scoped_statement_call((sa.select(Payload.content),), {}) + + class UntrustedMediaReferenceText(MediaReferenceText): + pass + + for column_type in (UntrustedMediaReferenceText(), MediaReferenceText(collation=sa.sql.quoted_name('x', False))): + with pytest.raises(ScopedSessionTransactionError): + _validate_scoped_statement_call((sa.select(sa.column('content', column_type)),), {}) diff --git a/tests/unit_tests/storage/test_inline_media_cleanup.py b/tests/unit_tests/storage/test_inline_media_cleanup.py index 3c798217a..e3aa1bafe 100644 --- a/tests/unit_tests/storage/test_inline_media_cleanup.py +++ b/tests/unit_tests/storage/test_inline_media_cleanup.py @@ -126,6 +126,8 @@ async def insert_record(ap, model, number, workspace=WORKSPACE_A, timestamp=OLD, content_json=payload, ) field = 'content_json' + # Seed the pre-contract database format without applying today's write policy. + values[field] = sa.bindparam('historical_payload', payload, type_=sa.Text()) await ap.persistence_mgr.execute_async(sa.insert(model).values(**values)) return field diff --git a/tests/unit_tests/storage/test_media_persistence.py b/tests/unit_tests/storage/test_media_persistence.py index 6b77e89b2..0b1fb9318 100644 --- a/tests/unit_tests/storage/test_media_persistence.py +++ b/tests/unit_tests/storage/test_media_persistence.py @@ -11,6 +11,7 @@ import pytest import sqlalchemy as sa from langbot.pkg.agent.runner.run_journal import AgentRunJournal +from langbot.pkg.agent.runner.run_ledger_store import RunLedgerStore from langbot.pkg.entity.persistence.agent_run import AgentRunEvent from langbot.pkg.entity.persistence.event_log import EventLog from langbot.pkg.entity.persistence.monitoring import MonitoringMessage @@ -181,6 +182,48 @@ async def test_cas_failure_does_not_write_inline_payload_to_database(media_servi assert not objects +@pytest.mark.asyncio +@pytest.mark.parametrize('field', ['data', 'metadata']) +async def test_direct_ledger_writes_cannot_bypass_image_persistence(media_service, field): + service, objects, _ = media_service + engine = service.ap.persistence_mgr.get_db_engine() + with pytest.raises(sa.exc.StatementError, match='Inline images cannot be persisted'): + await RunLedgerStore(engine).append_event( + run_id='run', sequence=1, event_type='tool.call.completed', **{field: {'image': image()}} + ) + async with engine.connect() as conn: + assert (await conn.execute(sa.select(sa.func.count()).select_from(AgentRunEvent))).scalar_one() == 0 + assert not objects + + +@pytest.mark.asyncio +@pytest.mark.parametrize('write_kind', ['create', 'update']) +async def test_monitoring_variables_follow_metadata_only_policy(media_service, write_kind): + service, objects, _ = media_service + variables = json.dumps({'preview': image(), 'count': 0}) + context = _context(WORKSPACE_A) + if write_kind == 'create': + message_id = await service.record_message( + context, + bot_id='same-bot', + bot_name='Bot', + pipeline_id='pipeline', + pipeline_name='Pipeline', + message_content='message', + session_id='session', + variables=variables, + ) + else: + message_id = await _record_message(service, context, 'message') + await service.update_message_status(context, message_id, 'success', variables=variables) + async with service.ap.persistence_mgr.get_db_engine().connect() as conn: + stored = ( + await conn.execute(sa.select(MonitoringMessage.variables).where(MonitoringMessage.id == message_id)) + ).scalar_one() + assert json.loads(stored) == {'preview': {'type': 'Image', 'base64': None}, 'count': 0} + assert not objects + + @pytest.mark.asyncio @pytest.mark.parametrize('role', ['user', 'assistant']) async def test_transcript_remains_metadata_only_without_creating_cas_objects(media_service, role):