feat(media): centralize resource URLs and persist refreshed image URLs

- Deduplicate message and history resource URLs using stable catalog references
- Persist renewed URLs with concurrent refresh protection
- Refresh OneBot images via get_msg intervals or NapCat rkey TTL
- Migrate historical payloads and cover SQLite and PostgreSQL behavior
This commit is contained in:
BiFangKNT
2026-10-23 12:49:56 +08:00
parent 1be47ae925
commit 3456668497
20 changed files with 1116 additions and 47 deletions
@@ -7,6 +7,7 @@ import datetime
import sqlalchemy
from .base import Base
from ...persistence.media import MediaReferenceText
class AgentInteraction(Base):
@@ -29,15 +30,15 @@ class AgentInteraction(Base):
actor_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True)
status = sqlalchemy.Column(sqlalchemy.String(50), nullable=False, index=True)
request_json = sqlalchemy.Column(sqlalchemy.Text, nullable=False)
request_json = sqlalchemy.Column(MediaReferenceText, nullable=False)
delivery_target_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
delivery_result_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
delivery_result_json = sqlalchemy.Column(MediaReferenceText, nullable=True)
replaces_interaction_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
callback_token_hash = sqlalchemy.Column(sqlalchemy.String(64), nullable=False, unique=True, index=True)
expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True, index=True)
submitted_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
submission_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
submission_json = sqlalchemy.Column(MediaReferenceText, nullable=True)
status_reason = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, default=datetime.datetime.utcnow)
@@ -1,6 +1,7 @@
import sqlalchemy as sa
from .base import Base
from ...persistence.media import MediaReferenceJSON
class AssistantConversation(Base):
@@ -13,7 +14,7 @@ class AssistantConversation(Base):
account_uuid = sa.Column(sa.String(36), sa.ForeignKey('users.uuid', ondelete='CASCADE'), nullable=False)
revision = sa.Column(sa.Integer, nullable=False, default=0)
status = sa.Column(sa.String(20), nullable=False, default='ready')
messages = sa.Column(sa.JSON, nullable=False, default=list)
messages = sa.Column(MediaReferenceJSON, nullable=False, default=list)
error = sa.Column(sa.Text, nullable=True)
model_name = sa.Column(sa.String(255), nullable=True)
model_uuid = sa.Column(sa.String(36), nullable=True)
@@ -1,6 +1,6 @@
import sqlalchemy
from ...persistence.media import MediaReferenceText
from ...persistence.media import MediaReferenceText, MediaReferenceJSON
from .base import Base
@@ -107,7 +107,7 @@ class MonitoringLLMCall(Base):
error_message = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
message_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True) # Associated message ID
details = sqlalchemy.Column(sqlalchemy.JSON, nullable=True)
details = sqlalchemy.Column(MediaReferenceJSON, nullable=True)
__table_args__ = (
sqlalchemy.Index('ix_monitoring_llm_calls_workspace_timestamp', 'workspace_uuid', 'timestamp'),
@@ -137,8 +137,8 @@ class MonitoringToolCall(Base):
pipeline_name = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
session_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True)
message_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True)
arguments = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
result = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
arguments = sqlalchemy.Column(MediaReferenceText, nullable=True)
result = sqlalchemy.Column(MediaReferenceText, nullable=True)
error_message = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
__table_args__ = (
@@ -1,6 +1,7 @@
import sqlalchemy
from .base import Base
from ...persistence.media import MediaReferenceJSON
class LegacyPipeline(Base):
@@ -74,7 +75,7 @@ class PipelineRunRecord(Base):
)
started_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
finished_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
result = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
result = sqlalchemy.Column(MediaReferenceJSON, nullable=False)
knowledge_base_uuid = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
__table_args__ = (
@@ -0,0 +1,26 @@
"""Shared, stable identities for resource URLs embedded in conversation payloads."""
import sqlalchemy as sa
from .base import Base
class ResourceURL(Base):
__tablename__ = 'resource_urls'
id = sa.Column(sa.String(36), primary_key=True)
url = sa.Column(sa.Text, nullable=True)
url_hash = sa.Column(sa.String(64), nullable=True, unique=True)
canonical_id = sa.Column(sa.String(36), sa.ForeignKey('resource_urls.id'), nullable=True, index=True)
created_at = sa.Column(sa.DateTime, nullable=False)
last_checked_at = sa.Column(sa.DateTime, nullable=True)
updated_at = sa.Column(sa.DateTime, nullable=False)
__table_args__ = (
sa.CheckConstraint(
'(canonical_id IS NULL AND url IS NOT NULL AND url_hash IS NOT NULL) '
'OR (canonical_id IS NOT NULL AND url IS NULL AND url_hash IS NULL)',
name='ck_resource_url_canonical',
),
sa.Index('ix_resource_urls_last_checked', 'last_checked_at'),
)
@@ -0,0 +1,92 @@
"""Normalize retained resource URLs behind stable shared references."""
import datetime
import json
from alembic import op
import sqlalchemy as sa
revision = '0041_resource_urls'
down_revision = '0040_session_summaries'
branch_labels = None
depends_on = None
def upgrade():
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.resource_url import ResourceURL
from langbot.pkg.persistence.media import MediaReferenceText, MediaReferenceJSON
from langbot.pkg.persistence.resource_urls import create_postgres_expansion, intern, install
connection = op.get_bind()
ResourceURL.__table__.create(connection, checkfirst=True)
install(connection.engine)
if connection.dialect.name == 'postgresql':
create_postgres_expansion(connection)
inspector = sa.inspect(connection)
existing = set(inspector.get_table_names())
for table in sorted(Base.metadata.tables.values(), key=lambda table: table.name):
if table.name not in existing:
continue
retained = [
column.name for column in table.columns if isinstance(column.type, (MediaReferenceText, MediaReferenceJSON))
]
json_columns = {column.name for column in table.columns if isinstance(column.type, MediaReferenceJSON)}
if not retained:
continue
# Reflection intentionally reads literal stored values without expansion.
source = sa.Table(table.name, sa.MetaData(), autoload_with=connection)
retained = [name for name in retained if name in source.c]
keys = list(source.primary_key.columns)
if len(keys) != 1:
raise RuntimeError(f'Resource URL migration requires one primary key: {table.name}')
key = keys[0]
cursor = None
stamp = source.c.get('created_at')
if stamp is None:
stamp = source.c.get('timestamp')
while True:
query = sa.select(source).order_by(key).limit(100)
if cursor is not None:
query = query.where(key > cursor)
rows = connection.execute(query).mappings().all()
if not rows:
break
for row in rows:
created = row[stamp.name] if stamp is not None else None
if created is not None and created.tzinfo is not None:
created = created.astimezone(datetime.timezone.utc).replace(tzinfo=None)
updates = {}
for name in retained:
raw = json.dumps(row[name]) if name in json_columns else row[name]
converted = intern(connection, raw, created)
updates[name] = json.loads(converted) if name in json_columns else converted
updates = {name: value for name, value in updates.items() if value != row[name]}
if updates:
connection.execute(source.update().where(key == row[key.name]).values(**updates))
cursor = rows[-1][key.name]
def downgrade():
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.media import MediaReferenceText, MediaReferenceJSON
from langbot.pkg.persistence.resource_urls import uninstall
connection = op.get_bind()
existing = set(sa.inspect(connection).get_table_names())
for table in sorted(Base.metadata.tables.values(), key=lambda table: table.name):
if table.name not in existing:
continue
columns = {column['name'] for column in sa.inspect(connection).get_columns(table.name)}
for column in table.columns:
if column.name in columns and isinstance(column.type, (MediaReferenceText, MediaReferenceJSON)):
name = connection.dialect.identifier_preparer.quote(table.name)
field = connection.dialect.identifier_preparer.quote(column.name)
expanded = f'langbot_expand_resource_urls(CAST({field} AS TEXT))'
if isinstance(column.type, MediaReferenceJSON) and connection.dialect.name == 'postgresql':
expanded = f'CAST({expanded} AS JSON)'
connection.exec_driver_sql(f'UPDATE {name} SET {field} = {expanded}')
op.drop_table('resource_urls')
if connection.dialect.name == 'postgresql':
connection.exec_driver_sql('DROP FUNCTION langbot_expand_resource_urls(text)')
uninstall(connection.engine)
+47 -1
View File
@@ -5,7 +5,8 @@ from __future__ import annotations
import json
from typing import Any
from sqlalchemy import Text
from sqlalchemy import Text, JSON, cast
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.engine import Dialect
from sqlalchemy.sql.operators import OperatorType
from sqlalchemy.types import TypeDecorator
@@ -23,9 +24,20 @@ class MediaReferenceText(TypeDecorator[str]):
impl = Text
cache_ok = True
class comparator_factory(Text.Comparator):
def operate(self, op, *other, **kw):
from .resource_urls import ExpandedResourceURLs
# ORM identity maps compare annotated copies of mapped columns.
if other and isinstance(other[0], ColumnElement) and self.expr.shares_lineage(other[0]):
return super().operate(op, *other, **kw)
# CAS and textual filters compare the same expanded value readers see.
return op(ExpandedResourceURLs(self.expr), *other, **kw)
def process_bind_param(self, value: str | None, dialect: Dialect) -> str | None:
if value is None:
return None
reject_internal_references(value)
try:
payload = json.loads(value)
except json.JSONDecodeError:
@@ -44,3 +56,37 @@ class MediaReferenceText(TypeDecorator[str]):
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
def column_expression(self, column):
from .resource_urls import ExpandedResourceURLs
return ExpandedResourceURLs(column)
class MediaReferenceJSON(TypeDecorator):
"""Preserve native JSON API values while interning structured resource fields."""
impl = JSON
cache_ok = True
def process_bind_param(self, value, dialect):
if value is not None:
reject_internal_references(json.dumps(value))
return value
def column_expression(self, column):
from .resource_urls import ExpandedResourceJSON
return ExpandedResourceJSON(cast(column, Text))
def reject_internal_references(value):
"""Only the transactional normalizer may introduce catalog references."""
from .resource_urls import REFERENCE
if REFERENCE.search(value):
raise ValueError('Internal resource references cannot be supplied as conversation content')
# Register the catalog and its database hooks before engines are opened.
from . import resource_urls as _resource_urls # noqa: E402,F401
+9
View File
@@ -183,16 +183,25 @@ class PersistenceManager:
if self.mode == PersistenceMode.CLOUD_RUNTIME:
await self._validate_cloud_runtime()
from .resource_urls import install
install(engine.sync_engine)
return
if self.mode == PersistenceMode.RELEASE_MIGRATION:
async with self._release_migration_lock():
await self._initialize_managed_schema()
await self._validate_release_schema()
from .resource_urls import install
install(engine.sync_engine)
return
await self._initialize_managed_schema()
await self._enable_sqlite_foreign_keys_after_migration()
from .resource_urls import install
install(engine.sync_engine)
if self.mode == PersistenceMode.OSS_COMPAT:
await self.write_space_model_providers()
@@ -0,0 +1,391 @@
"""Normalize resource URLs at the typed SQL boundary, within the source transaction.
Only typed conversation payload writes are interned. Reads expand references in SQL so
ORM and Core callers retain their existing string/JSON contracts. The catalog
is shared; authorization continues to belong to the enclosing business record.
"""
import datetime
import hashlib
import json
import re
import uuid
import sqlite3
from collections import OrderedDict
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql, sqlite
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.functions import FunctionElement
from ..entity.persistence.resource_url import ResourceURL
PREFIX = 'langbot-resource://'
REFERENCE = re.compile(r'langbot-resource://([a-f0-9-]{36})')
FIELDS = frozenset({'url', 'image_url', 'audio_url', 'file_url', 'video_url', 'original_url'})
def now():
return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
def digest(url):
return hashlib.sha256(url.encode()).hexdigest()
def fields(value):
if isinstance(value, list):
for child in value:
yield from fields(child)
elif isinstance(value, dict):
for key, child in value.items():
if key in FIELDS and isinstance(child, str) and child.startswith(('http://', 'https://', '/api/v1/files/')):
yield value, key, child
else:
yield from fields(child)
def intern(connection, raw, created_at=None, identities_out=None):
"""Intern resource fields synchronously in the current SQLAlchemy transaction."""
if not isinstance(raw, str):
return raw
try:
payload = json.loads(raw)
except ValueError:
return raw
nodes = list(fields(payload))
if not nodes:
return raw
table = ResourceURL.__table__
stamp = now()
urls = {url for _, _, url in nodes}
aliases = connection.engine._langbot_resource_aliases
known = {url: aliases[url] for url in urls if url in aliases}
valid = (
set(connection.execute(sa.select(table.c.id).where(table.c.id.in_(set(known.values())))).scalars())
if known
else set()
)
identities = {url: identity for url, identity in known.items() if identity in valid}
urls -= identities.keys()
hashes = dict(sorted((digest(url), url) for url in urls))
insert = {'sqlite': sqlite.insert, 'postgresql': postgresql.insert}[connection.dialect.name]
if hashes:
connection.execute(
insert(table)
.values(
[
dict(id=str(uuid.uuid4()), url=url, url_hash=key, created_at=created_at or stamp, updated_at=stamp)
for key, url in hashes.items()
]
)
.on_conflict_do_nothing(index_elements=['url_hash'])
)
for identity, url, key in connection.execute(
sa.select(table.c.id, table.c.url, table.c.url_hash).where(table.c.url_hash.in_(hashes))
):
if hashes[key] != url:
raise ValueError('Resource URL hash collision')
identities[url] = identity
if created_at is not None and identities:
# Migration order must not make an old shared URL appear newly created.
connection.execute(
table.update()
.where(table.c.id.in_(set(identities.values())), table.c.created_at > created_at)
.values(created_at=created_at)
)
for node, key, url in nodes:
node[key] = PREFIX + identities[url]
aliases[url] = identities[url]
aliases.move_to_end(url)
while len(aliases) > 8192:
aliases.popitem(last=False)
if identities_out is not None:
identities_out.update(identities)
return json.dumps(payload, ensure_ascii=False)
class ExpandedResourceURLs(FunctionElement):
type = sa.Text()
inherit_cache = True
class ExpandedResourceJSON(FunctionElement):
type = sa.JSON()
inherit_cache = True
@compiles(ExpandedResourceJSON)
def compile_json_expansion(element, compiler, **kw):
expression = compile_expansion(element, compiler, **kw)
if compiler.dialect.name == 'postgresql':
return f'CAST({expression} AS JSON)'
return expression
@compiles(ExpandedResourceURLs)
def compile_expansion(element, compiler, **kw):
source = compiler.process(list(element.clauses)[0], **kw)
if not getattr(compiler.dialect, '_langbot_resource_urls', False):
return source
return f'langbot_expand_resource_urls({source})'
def install(engine):
"""Enable normalization after schema initialization, before producers start."""
engine.dialect._langbot_resource_urls = True
if not hasattr(engine, '_langbot_resource_aliases'):
engine._langbot_resource_aliases = OrderedDict()
engine.clear_compiled_cache()
if not sa.event.contains(engine, 'before_cursor_execute', normalize_parameters):
sa.event.listen(engine, 'before_cursor_execute', normalize_parameters, retval=True)
def uninstall(engine):
"""Restore literal payload reads and writes after catalog downgrade."""
if sa.event.contains(engine, 'before_cursor_execute', normalize_parameters):
sa.event.remove(engine, 'before_cursor_execute', normalize_parameters)
engine.dialect._langbot_resource_urls = False
engine.clear_compiled_cache()
engine._langbot_resource_aliases.clear()
def normalize_parameters(connection, cursor, statement, parameters, context, executemany):
"""Replace only typed payload binds; never inspect arbitrary SQL or configuration."""
from .media import MediaReferenceText, MediaReferenceJSON
compiled = context.compiled
if compiled is None or not (context.isinsert or context.isupdate):
return statement, parameters
names = {
name
for bind, name in compiled.bind_names.items()
if isinstance(bind.type, (MediaReferenceText, MediaReferenceJSON))
}
if not names:
return statement, parameters
def normalize(row):
if isinstance(row, dict):
return {key: intern(connection, value) if key in names else value for key, value in row.items()}
positions = compiled.positiontup or []
if not positions:
return row
return tuple(
intern(connection, value) if positions[index % len(positions)] in names else value
for index, value in enumerate(row)
)
batched = executemany and getattr(compiled, '_insertmanyvalues', None) is None
return statement, [normalize(row) for row in parameters] if batched else normalize(parameters)
def register_sqlite(dbapi, record):
"""The UDF runs on aiosqlite's worker, including its catalog SELECTs."""
driver = getattr(dbapi, 'driver_connection', None)
raw_connection = driver._conn if hasattr(driver, '_conn') else dbapi
if not isinstance(raw_connection, sqlite3.Connection):
return
def expand(raw):
if not isinstance(raw, str) or PREFIX not in raw:
return raw
resolved = {}
def lookup(match):
identity = match.group(1)
if identity in resolved:
return resolved[identity]
visited = set()
while identity not in visited:
visited.add(identity)
row = raw_connection.execute(
'SELECT url, canonical_id FROM resource_urls WHERE id=?', (identity,)
).fetchone()
if row is None:
raise ValueError('Missing resource URL')
if row[1] is None:
value = json.dumps(row[0], ensure_ascii=False)[1:-1]
for reference in visited:
resolved[reference] = value
return value
identity = row[1]
raise ValueError('Cyclic resource URL reference')
return REFERENCE.sub(lookup, raw)
dbapi.create_function('langbot_expand_resource_urls', 1, expand)
POSTGRES_EXPANSION = r"""
CREATE OR REPLACE FUNCTION public.langbot_expand_resource_urls(payload text)
RETURNS text LANGUAGE plpgsql STABLE SECURITY INVOKER SET search_path = public, pg_temp AS $$
DECLARE token text; identity text; target text; parent text; result text := payload; visited text[];
BEGIN
IF payload IS NULL OR position('langbot-resource://' in payload) = 0 THEN RETURN payload; END IF;
FOR token IN SELECT DISTINCT m[1] FROM regexp_matches(payload, 'langbot-resource://([a-f0-9-]{36})', 'g') m LOOP
identity := token; visited := ARRAY[]::text[];
LOOP
IF identity = ANY(visited) THEN RAISE EXCEPTION 'Cyclic resource URL reference'; END IF;
visited := array_append(visited, identity);
SELECT url, canonical_id INTO target, parent FROM resource_urls WHERE id = identity;
IF NOT FOUND THEN RAISE EXCEPTION 'Missing resource URL'; END IF;
EXIT WHEN parent IS NULL;
identity := parent;
END LOOP;
result := replace(result, 'langbot-resource://' || token, substr(to_json(target)::text, 2, length(to_json(target)::text)-2));
END LOOP;
RETURN result;
END $$;
"""
def catalog_created(table, connection, **kw):
if connection.dialect.name == 'postgresql':
create_postgres_expansion(connection)
install(connection.engine)
def create_postgres_expansion(connection):
schema = connection.scalar(sa.select(sa.func.current_schema()))
quoted = connection.dialect.identifier_preparer.quote(schema)
connection.exec_driver_sql(POSTGRES_EXPANSION.replace('public', quoted))
async def refresh_images(engine, peer, payload, *, is_current=None):
"""Serialize refreshes on catalog identities and overwrite current URLs atomically."""
from ..platform.adapters.aiocqhttp.image_urls import image_nodes
nodes = list(image_nodes(payload))
if not nodes:
return
def connected():
return not peer.closed and (is_current is None or is_current())
table = ResourceURL.__table__
# Establish catalog identities for runtime-only inputs too. Existing history
# already has identities created in its original write or migration.
async with engine.begin() as db:
references = {}
await db.run_sync(lambda sync: intern(sync, json.dumps(payload), identities_out=references))
identities = set(references.values())
async with engine.connect() as db:
if engine.dialect.name == 'sqlite':
await db.exec_driver_sql('BEGIN IMMEDIATE')
else:
await db.begin()
try:
rows = (
(
await db.execute(
sa.select(table).where(table.c.id.in_(identities)).order_by(table.c.id).with_for_update()
)
)
.mappings()
.all()
)
by_id = {row['id']: row for row in rows}
for identity in list(by_id):
row = by_id[identity]
while row['canonical_id'] is not None:
row = (
(await db.execute(sa.select(table).where(table.c.id == row['canonical_id']).with_for_update()))
.mappings()
.one()
)
by_id[identity] = row
by_url = {node['url']: by_id[references[node['url']]] for node, _, _ in nodes}
# Capture origins even when renewal is not due or identification fails.
for node, message_id, image_id in nodes:
peer.image_urls.remember(by_url[node['url']]['url'], message_id, image_id)
# Identification and upstream credential handling remain connection-local.
if not connected() or not await peer.image_urls.identify():
return
napcat = peer.image_urls.app_name == 'NapCat.Onebot'
due = []
stamp = now()
for node, message_id, image_id in nodes:
row = by_url.get(node['url'])
if row is None:
continue
node['url'] = row['url']
checked = row['last_checked_at'] or row['created_at']
if napcat or (stamp - checked).total_seconds() >= 1800:
due.append((node, message_id, image_id, row))
if due:
# Keep Source/Image identities, including those recovered from events.
batch = [
dict(message_id=message_id, image={'type': 'Image', 'image_id': image_id, 'url': node['url']})
for node, message_id, image_id, _ in due
]
old = [entry['image']['url'] for entry in batch]
completed = await peer.image_urls.refresh(batch)
if not completed:
return
if not connected():
return
stamp = now()
updated = set()
for (node, _, _, row), entry, previous in zip(due, batch, old):
if previous not in completed:
continue
fresh = entry['image']['url']
node['url'] = fresh
if napcat and fresh == previous:
continue
if row['id'] in updated:
continue
updated.add(row['id'])
values = dict(last_checked_at=stamp)
if fresh != previous:
fresh_hash = digest(fresh)
# Lock the destination before merging identical current URLs.
destination = (
await db.execute(
sa.select(table.c.id, table.c.url)
.where(table.c.url_hash == fresh_hash)
.with_for_update()
)
).first()
if destination is None:
try:
# Another source may discover the same renewed URL
# concurrently. Let the unique index arbitrate it.
async with db.begin_nested():
await db.execute(
table.update()
.where(table.c.id == row['id'])
.values(url=fresh, url_hash=fresh_hash, updated_at=stamp, **values)
)
continue
except sa.exc.IntegrityError as exc:
if getattr(exc.orig, 'sqlstate', None) != '23505':
raise
destination = (
await db.execute(
sa.select(table.c.id, table.c.url)
.where(table.c.url_hash == fresh_hash)
.with_for_update()
)
).one()
if destination is not None and destination.id != row['id']:
if destination.url != fresh:
raise ValueError('Resource URL hash collision')
values.update(url=None, url_hash=None, canonical_id=destination.id, updated_at=stamp)
await db.execute(
table.update().where(table.c.id == destination.id).values(last_checked_at=stamp)
)
else:
values.update(url=fresh, url_hash=fresh_hash, updated_at=stamp)
await db.execute(table.update().where(table.c.id == row['id']).values(**values))
await db.commit()
except BaseException:
await db.rollback()
raise
sa.event.listen(sa.engine.Engine, 'connect', register_sqlite)
sa.event.listen(ResourceURL.__table__, 'after_create', catalog_created)
+15 -3
View File
@@ -19,7 +19,7 @@ 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
from .media import MediaReferenceText, MediaReferenceJSON
TENANT_SETTING = 'langbot.workspace_uuid'
@@ -295,7 +295,9 @@ def _validate_scoped_sql_type(
if type(sql_type) in {Vector, HALFVEC}:
return
if type(sql_type) is not MediaReferenceText and not type(sql_type).__module__.startswith('sqlalchemy.'):
if type(sql_type) not in {MediaReferenceText, MediaReferenceJSON} 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'):
@@ -501,6 +503,14 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s
)
if isinstance(element, sqlalchemy.sql.functions.FunctionElement):
from .resource_urls import ExpandedResourceURLs
resource_expansion_allowed = (
type(element) is ExpandedResourceURLs
and len(list(element.clauses)) == 1
and isinstance(list(element.clauses)[0], sqlalchemy.Column)
and type(list(element.clauses)[0].type) is MediaReferenceText
)
function_name = str(getattr(element, 'name', '')).casefold()
package_names = tuple(getattr(element, 'packagenames', ()))
generic_function_allowed = (
@@ -508,7 +518,9 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s
and function_name in _ALLOWED_SCOPED_GENERIC_FUNCTIONS
)
builtin_function_allowed = _ALLOWED_SCOPED_BUILTIN_FUNCTION_TYPES.get(function_name) is type(element)
if package_names or not (generic_function_allowed or builtin_function_allowed):
if package_names or not (
generic_function_allowed or builtin_function_allowed or resource_expansion_allowed
):
raise ScopedSessionTransactionError(
f'TenantUnitOfWork does not allow SQL function {function_name or "<unknown>"!r}'
)
@@ -77,6 +77,7 @@ class OneBotImageURLs:
self.keys = {}
self.references = OrderedDict()
self.messages = OrderedDict()
self.urls = OrderedDict()
self.retry_after = 0
def remember(self, url, message_id, image_id):
@@ -87,23 +88,37 @@ class OneBotImageURLs:
while len(self.references) > 2048:
self.references.popitem(last=False)
async def identify(self):
"""Identify lazily with the same backoff as resource renewal."""
async with self.lock:
if self.connection.closed or time.monotonic() < self.retry_after:
return False
try:
await self._identify()
return True
except Exception:
self.retry_after = time.monotonic() + 30
raise
async def _identify(self):
if self.app_name is None:
info = await self.connection.call_action('get_version_info')
self.app_name = info.get('app_name', '')
async def refresh(self, payload):
nodes = list(image_nodes(payload))
if not nodes or self.connection.closed:
return
return False
for node, message_id, image_id in nodes:
self.remember(node['url'], message_id, image_id)
async with self.lock:
if self.connection.closed or time.monotonic() < self.retry_after:
return
return False
try:
if self.app_name is None:
info = await self.connection.call_action('get_version_info')
self.app_name = info.get('app_name', '')
await self._identify()
if self.app_name == 'NapCat.Onebot':
await self._napcat(nodes)
else:
await self._standard(nodes)
return await self._napcat(nodes)
return await self._standard(nodes)
except Exception:
# Failed/unsupported APIs must not break history reads or hammer the peer.
self.retry_after = time.monotonic() + 30
@@ -112,7 +127,7 @@ class OneBotImageURLs:
async def _napcat(self, nodes):
wanted = {kind for node, _, _ in nodes if (kind := nt_image_type(node['url'])) is not None}
if not wanted:
return
return set()
if any(self.keys.get(kind, ('', 0))[1] <= time.time() for kind in wanted):
result = await self.connection.call_action('nc_get_rkey')
keys = {}
@@ -129,10 +144,12 @@ class OneBotImageURLs:
if not wanted.issubset(keys):
raise ValueError('NapCat returned no valid image credentials')
self.keys.update(keys)
completed = set()
for node, _, _ in nodes:
kind = nt_image_type(node['url'])
if kind not in wanted:
continue
completed.add(node['url'])
parsed = urlsplit(node['url'])
parts = parsed.query.split('&')
# Preserve every other parameter byte-for-byte, including signed resource IDs.
@@ -143,10 +160,17 @@ class OneBotImageURLs:
for part in parts
]
node['url'] = urlunsplit(parsed._replace(query='&'.join(parts)))
return completed
async def _standard(self, nodes):
completed = set()
for node, _, _ in nodes:
old_url = node['url']
refreshed = self.urls.get(old_url)
if refreshed is not None and refreshed[1] > time.monotonic():
node['url'] = refreshed[0]
completed.add(old_url)
continue
ref = self.references.get(old_url)
if ref is None:
continue
@@ -160,12 +184,19 @@ class OneBotImageURLs:
data = segment.data
if data.get('file') and remote_image_url(data.get('url')):
images[str(data['file'])] = data['url']
cached = (images, time.monotonic() + 60)
cached = (images, time.monotonic() + 1800)
self.messages[str(message_id)] = cached
self.messages.move_to_end(str(message_id))
while len(self.messages) > 256:
self.messages.popitem(last=False)
new_url = cached[0].get(image_id)
if new_url:
completed.add(old_url)
self.remember(new_url, message_id, image_id)
deadline = cached[1]
self.urls[old_url] = (new_url, deadline)
self.urls[new_url] = (new_url, deadline)
while len(self.urls) > 2048:
self.urls.popitem(last=False)
node['url'] = new_url
return completed
@@ -102,6 +102,7 @@ class OneBotConnection:
self.image_urls.keys.clear()
self.image_urls.references.clear()
self.image_urls.messages.clear()
self.image_urls.urls.clear()
for future in self.pending.values():
if not future.done():
future.set_exception(ApiNotAvailable('OneBot connection closed'))
+9 -2
View File
@@ -1,4 +1,4 @@
"""Refresh image references at authorized read/use boundaries, without persistence."""
"""Refresh image references at authorized read/use boundaries."""
import copy
import json
@@ -21,7 +21,14 @@ async def refresh_image_payload(ap, workspace_uuid, bot_id, payload):
return payload
result = copy.deepcopy(payload)
try:
await connection.image_urls.refresh(result)
from ..persistence.resource_urls import refresh_images
await refresh_images(
ap.persistence_mgr.get_db_engine(),
connection,
result,
is_current=lambda: bot.adapter.bot.connection is connection,
)
if bot.adapter.bot.connection is not connection or connection.closed:
return payload
return result