From 5d774ff555b0c83c42fd1819315261c63ec66a04 Mon Sep 17 00:00:00 2001
From: BiFangKNT <1320414964@qq.com>
Date: Fri, 23 Oct 2026 14:39:38 +0800
Subject: [PATCH] fix(persistence): resolve resource URL migration and scoped
 SQL failures

- Reflect migration tables without resolving foreign-key targets
- Allow validated resource URL expansion through the tenant SQL AST guard
- Add regression coverage for trusted and malformed expansion expressions
---
 .../alembic/versions/0041_resource_urls.py    |  3 ++-
 src/langbot/pkg/persistence/tenant_uow.py     | 18 ++++++++--------
 .../unit_tests/persistence/test_tenant_uow.py | 21 +++++++++++++++++++
 3 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/langbot/pkg/persistence/alembic/versions/0041_resource_urls.py b/src/langbot/pkg/persistence/alembic/versions/0041_resource_urls.py
index 0030a73b8..aaa90bd0a 100644
--- a/src/langbot/pkg/persistence/alembic/versions/0041_resource_urls.py
+++ b/src/langbot/pkg/persistence/alembic/versions/0041_resource_urls.py
@@ -35,7 +35,8 @@ def upgrade():
         if not retained:
             continue
         # Reflection intentionally reads literal stored values without expansion.
-        source = sa.Table(table.name, sa.MetaData(), autoload_with=connection)
+        # Only local columns are needed; historical schemas may omit FK targets.
+        source = sa.Table(table.name, sa.MetaData(), autoload_with=connection, resolve_fks=False)
         retained = [name for name in retained if name in source.c]
         keys = list(source.primary_key.columns)
         if len(keys) != 1:
diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py
index 2b6a1093d..e8ca31765 100644
--- a/src/langbot/pkg/persistence/tenant_uow.py
+++ b/src/langbot/pkg/persistence/tenant_uow.py
@@ -412,8 +412,16 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s
             'TenantUnitOfWork public SQL must use a structured SQLAlchemy query, values, or DML statement'
         )
 
+    from .resource_urls import ExpandedResourceURLs
+
     for element in _iter_scoped_statement_elements(statement):
-        if not type(element).__module__.startswith('sqlalchemy.'):
+        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
+        )
+        if not type(element).__module__.startswith('sqlalchemy.') and not resource_expansion_allowed:
             raise ScopedSessionTransactionError(
                 'TenantUnitOfWork does not allow custom SQL AST nodes in public statements'
             )
@@ -503,14 +511,6 @@ 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 = (
diff --git a/tests/unit_tests/persistence/test_tenant_uow.py b/tests/unit_tests/persistence/test_tenant_uow.py
index ee4c0b570..9a8b1c827 100644
--- a/tests/unit_tests/persistence/test_tenant_uow.py
+++ b/tests/unit_tests/persistence/test_tenant_uow.py
@@ -980,6 +980,27 @@ async def test_scoped_sql_structure_allows_only_the_production_vocabulary(statem
     _validate_scoped_statement_call((statement,), {})
 
 
+async def test_resource_expansion_allows_only_the_fixed_function_and_media_column() -> None:
+    from langbot.pkg.entity.persistence.event_log import EventLog
+    from langbot.pkg.persistence.resource_urls import ExpandedResourceURLs
+
+    statement = sa.update(EventLog).where(EventLog.input_json == '{}').values(input_json='[]')
+    _validate_scoped_statement_call((statement,), {})
+
+    class SpoofedExpansion(ExpandedResourceURLs):
+        inherit_cache = True
+
+    for expression in (
+        ExpandedResourceURLs(),
+        ExpandedResourceURLs(sa.literal('{}')),
+        ExpandedResourceURLs(User.uuid),
+        ExpandedResourceURLs(EventLog.input_json, sa.literal('{}')),
+        SpoofedExpansion(EventLog.input_json),
+    ):
+        with pytest.raises(ScopedSessionTransactionError, match='custom SQL AST'):
+            _validate_scoped_statement_call((sa.select(expression),), {})
+
+
 async def test_monitoring_json_casts_are_limited_to_journal_metadata() -> None:
     from langbot.pkg.entity.persistence.agent_run import AgentRun
     from langbot.pkg.entity.persistence.event_log import EventLog