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
This commit is contained in:
BiFangKNT
2026-10-18 12:26:09 +08:00
parent 1e5aeb6032
commit b73f66a234
17 changed files with 577 additions and 21 deletions
@@ -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
@@ -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})
@@ -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(
@@ -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)),), {})
@@ -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
@@ -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):