mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-22 16:35:20 +08:00
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:
@@ -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: `<type>(<scope>): <subject>`.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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__ = (
|
||||
|
||||
@@ -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'),)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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')
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user