LangBot/src/langbot/pkg/persistence/tenant_uow.py

from __future__ import annotations

import asyncio
import collections.abc
import contextlib
import contextvars
import dataclasses
import enum
import functools
import types
import typing

import sqlalchemy
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
import sqlalchemy.orm as sqlalchemy_orm
from pgvector.sqlalchemy import Vector
from sqlalchemy.dialects.postgresql.dml import OnConflictDoNothing as PostgreSQLOnConflictDoNothing
from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLOnConflictDoUpdate
from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing
from sqlalchemy.dialects.sqlite.dml import OnConflictDoUpdate as SQLiteOnConflictDoUpdate


TENANT_SETTING = 'langbot.workspace_uuid'
ACCOUNT_SETTING = 'langbot.account_uuid'
API_KEY_HASH_SETTING = 'langbot.api_key_hash'
INVITATION_HASH_SETTING = 'langbot.invitation_hash'
INSTANCE_SETTING = 'langbot.instance_uuid'
IDENTITY_DIGEST_SETTING = 'langbot.identity_digest'
DIRECTORY_INSTANCE_SETTING = 'langbot.directory_instance_uuid'

TENANT_POLICY_NAME = 'langbot_workspace_isolation'
LOCAL_DIRECTORY_WRITE_POLICY_NAME = 'langbot_workspace_local_directory_write'
ACCOUNT_DISCOVERY_POLICY_NAME = 'langbot_account_discovery'
API_KEY_DISCOVERY_POLICY_NAME = 'langbot_api_key_discovery'
INVITATION_DISCOVERY_POLICY_NAME = 'langbot_invitation_discovery'
INSTANCE_DISCOVERY_POLICY_NAME = 'langbot_instance_discovery'
DIRECTORY_PROJECTION_POLICY_NAME = 'langbot_directory_projection'

# Keep this contract explicit. A new tenant-owned table must be added to both
# this runtime list and the corresponding Alembic migration before release.
TENANT_TABLE_COLUMNS: dict[str, str] = {
    'workspaces': 'uuid',
    'workspace_memberships': 'workspace_uuid',
    'workspace_invitations': 'workspace_uuid',
    'workspace_execution_states': 'workspace_uuid',
    'support_admin_temporary_sessions': 'workspace_uuid',
    'workspace_metadata': 'workspace_uuid',
    'api_keys': 'workspace_uuid',
    'bots': 'workspace_uuid',
    'bot_admins': 'workspace_uuid',
    'binary_storages': 'workspace_uuid',
    'mcp_servers': 'workspace_uuid',
    'model_providers': 'workspace_uuid',
    'llm_models': 'workspace_uuid',
    'embedding_models': 'workspace_uuid',
    'rerank_models': 'workspace_uuid',
    'legacy_pipelines': 'workspace_uuid',
    'pipeline_run_records': 'workspace_uuid',
    'plugin_settings': 'workspace_uuid',
    'knowledge_bases': 'workspace_uuid',
    'knowledge_base_files': 'workspace_uuid',
    'knowledge_base_chunks': 'workspace_uuid',
    'webhooks': 'workspace_uuid',
    'monitoring_messages': 'workspace_uuid',
    'monitoring_llm_calls': 'workspace_uuid',
    'monitoring_tool_calls': 'workspace_uuid',
    'monitoring_sessions': 'workspace_uuid',
    'monitoring_errors': 'workspace_uuid',
    'monitoring_embedding_calls': 'workspace_uuid',
    'monitoring_feedback': 'workspace_uuid',
    # Created by 0013 rather than ORM metadata; it is still part of the same
    # business-database RLS contract and permits no discovery policies.
    'langbot_vectors': 'workspace_uuid',
}

DIRECTORY_PROJECTION_TABLE_COLUMNS: dict[str, str] = {
    'directory_projection_states': 'instance_uuid',
    'directory_projection_inbox': 'instance_uuid',
}

DIRECTORY_PROJECTED_TENANT_TABLES = frozenset(
    {
        'workspaces',
        'workspace_memberships',
        'workspace_execution_states',
    }
)


class PersistenceScopeKind(enum.StrEnum):
    WORKSPACE = 'workspace'
    ACCOUNT_DISCOVERY = 'account_discovery'
    API_KEY_DISCOVERY = 'api_key_discovery'
    INVITATION_DISCOVERY = 'invitation_discovery'
    INSTANCE_DISCOVERY = 'instance_discovery'
    IDENTITY_DISCOVERY = 'identity_discovery'
    DIRECTORY_PROJECTION = 'directory_projection'


@dataclasses.dataclass(frozen=True, slots=True)
class PersistenceScope:
    """A trusted, transaction-local database visibility scope."""

    kind: PersistenceScopeKind
    settings: tuple[tuple[str, str], ...]

    @classmethod
    def workspace(cls, workspace_uuid: str) -> PersistenceScope:
        return cls._one(PersistenceScopeKind.WORKSPACE, TENANT_SETTING, workspace_uuid)

    @classmethod
    def account(cls, account_uuid: str) -> PersistenceScope:
        return cls._one(PersistenceScopeKind.ACCOUNT_DISCOVERY, ACCOUNT_SETTING, account_uuid)

    @classmethod
    def api_key(cls, key_hash: str) -> PersistenceScope:
        return cls._one(PersistenceScopeKind.API_KEY_DISCOVERY, API_KEY_HASH_SETTING, key_hash)

    @classmethod
    def invitation(cls, token_hash: str) -> PersistenceScope:
        return cls._one(PersistenceScopeKind.INVITATION_DISCOVERY, INVITATION_HASH_SETTING, token_hash)

    @classmethod
    def instance(cls, instance_uuid: str) -> PersistenceScope:
        return cls._one(PersistenceScopeKind.INSTANCE_DISCOVERY, INSTANCE_SETTING, instance_uuid)

    @classmethod
    def identity(cls, identity_digest: str) -> PersistenceScope:
        return cls._one(PersistenceScopeKind.IDENTITY_DISCOVERY, IDENTITY_DIGEST_SETTING, identity_digest)

    @classmethod
    def directory(cls, instance_uuid: str) -> PersistenceScope:
        return cls._one(
            PersistenceScopeKind.DIRECTORY_PROJECTION,
            DIRECTORY_INSTANCE_SETTING,
            instance_uuid,
        )

    @classmethod
    def _one(cls, kind: PersistenceScopeKind, setting: str, value: str) -> PersistenceScope:
        return cls(kind, ((setting, cls._normalize(value, kind.value)),))

    @staticmethod
    def _normalize(value: str, label: str) -> str:
        if not isinstance(value, str):
            raise TypeError(f'{label} must be a string')
        normalized = value.strip()
        if not normalized:
            raise ValueError(f'{label} must not be empty')
        if len(normalized) > 512:
            raise ValueError(f'{label} exceeds the database scope limit')
        return normalized


class TenantScopeRequiredError(RuntimeError):
    """Raised when Cloud business data is accessed without a trusted scope."""


class CrossScopeTransactionError(RuntimeError):
    """Raised when code tries to change scope inside an active transaction."""


class TransactionRollbackOnlyError(RuntimeError):
    """Raised when a caught failure made the outer transaction unsafe to commit."""


class ScopedSessionTransactionError(RuntimeError):
    """Raised when callers try to escape a UoW-owned transaction boundary."""


@dataclasses.dataclass(slots=True)
class ActiveScopedTransaction:
    scope: PersistenceScope
    session: sqlalchemy_asyncio.AsyncSession
    engine: sqlalchemy_asyncio.AsyncEngine
    owner_task: asyncio.Task[typing.Any] | None
    depth: int = 1
    rollback_only: bool = False
    rollback_only_cause: BaseException | None = None
    after_commit_waiters: list[asyncio.Future[None]] = dataclasses.field(default_factory=list)

    def mark_rollback_only(self, cause: BaseException | None = None) -> None:
        """Remember that this transaction must never release commit-gated work."""

        self.rollback_only = True
        if self.rollback_only_cause is None:
            self.rollback_only_cause = cause


_UOW_SESSION_CONTROL_CAPABILITY = object()


@dataclasses.dataclass(slots=True)
class _ScopedSessionGuardState:
    owner_task: asyncio.Task[typing.Any] | None
    transaction_state: ActiveScopedTransaction | None = None
    active: bool = True
    session_events_blocked: bool = False


_SYNC_PROXY_CAPABILITY: contextvars.ContextVar[_ScopedSessionGuardState | None] = contextvars.ContextVar(
    'langbot_sync_session_proxy_capability',
    default=None,
)

_ALLOWED_SCOPED_BUILTIN_FUNCTION_TYPES = {
    'coalesce': sqlalchemy.sql.functions.coalesce,
    'count': sqlalchemy.sql.functions.count,
    'now': sqlalchemy.sql.functions.now,
    'sum': sqlalchemy.sql.functions.sum,
}
_ALLOWED_SCOPED_GENERIC_FUNCTIONS = frozenset({'length', 'nullif'})
_ALLOWED_SCOPED_CUSTOM_OPERATORS = frozenset({'<=>'})
_ALLOWED_SCOPED_STATEMENT_TYPES = (
    sqlalchemy.sql.dml.UpdateBase,
    sqlalchemy.sql.selectable.SelectBase,
)
_ALLOWED_SCOPED_POST_VALUES_TYPES = (
    PostgreSQLOnConflictDoNothing,
    PostgreSQLOnConflictDoUpdate,
    SQLiteOnConflictDoNothing,
    SQLiteOnConflictDoUpdate,
)


def _collect_embedded_clause_elements(value: typing.Any) -> list[sqlalchemy.sql.elements.ClauseElement]:
    """Find executable expressions in SQLAlchemy containers omitted by visitors."""

    if isinstance(value, sqlalchemy.sql.elements.ClauseElement):
        return [value]
    if isinstance(value, collections.abc.Mapping):
        elements: list[sqlalchemy.sql.elements.ClauseElement] = []
        for key, item in value.items():
            elements.extend(_collect_embedded_clause_elements(key))
            elements.extend(_collect_embedded_clause_elements(item))
        return elements
    if isinstance(value, (list, tuple)):
        elements = []
        for item in value:
            elements.extend(_collect_embedded_clause_elements(item))
        return elements
    clause_factory = getattr(value, '__clause_element__', None)
    if callable(clause_factory):
        clause_element = clause_factory()
        if isinstance(clause_element, sqlalchemy.sql.elements.ClauseElement):
            return [clause_element]
    return []


def _validate_opaque_identifier(
    value: typing.Any,
    *,
    label: str,
) -> list[sqlalchemy.sql.elements.ClauseElement]:
    """Validate identifiers stored outside SQLAlchemy's normal AST traversal."""

    if isinstance(value, sqlalchemy.sql.elements.ClauseElement):
        return [value]
    if not isinstance(value, str):
        raise ScopedSessionTransactionError(f'TenantUnitOfWork does not allow an unknown {label}')
    if isinstance(value, sqlalchemy.sql.elements.quoted_name) and value.quote is False:
        raise ScopedSessionTransactionError('TenantUnitOfWork does not allow forced-unquoted SQL identifiers')
    if not value or value[0].isdigit() or any(not (character.isalnum() or character == '_') for character in value):
        raise ScopedSessionTransactionError(f'TenantUnitOfWork does not allow an unsafe {label}')
    return []


def _validate_scoped_sql_type(
    sql_type: typing.Any,
    *,
    seen: set[int] | None = None,
) -> None:
    """Reject caller-defined SQL compilers hidden behind otherwise safe nodes."""

    if not isinstance(sql_type, sqlalchemy.types.TypeEngine):
        return
    if seen is None:
        seen = set()
    identity = id(sql_type)
    if identity in seen:
        return
    seen.add(identity)

    if type(sql_type) is Vector:
        return
    if 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'):
        identifier = getattr(sql_type, identifier_attribute, None)
        if isinstance(identifier, sqlalchemy.sql.elements.quoted_name) and identifier.quote is False:
            raise ScopedSessionTransactionError('TenantUnitOfWork does not allow forced-unquoted SQL type identifiers')

    nested_types: list[typing.Any] = [
        getattr(sql_type, 'item_type', None),
        getattr(sql_type, 'impl', None),
    ]
    nested_types.extend(getattr(sql_type, 'types', ()) or ())
    for nested_type in nested_types:
        if nested_type is not sql_type:
            _validate_scoped_sql_type(nested_type, seen=seen)


def _collect_multi_value_elements(
    multi_values: typing.Any,
) -> list[sqlalchemy.sql.elements.ClauseElement]:
    """Validate batch INSERT rows which SQLAlchemy omits from AST visitors."""

    children: list[sqlalchemy.sql.elements.ClauseElement] = []
    for value_group in multi_values or ():
        if not isinstance(value_group, (list, tuple)):
            raise ScopedSessionTransactionError('TenantUnitOfWork does not allow an unknown batch INSERT shape')
        for row in value_group:
            if isinstance(row, collections.abc.Mapping):
                for key, value in row.items():
                    children.extend(_validate_opaque_identifier(key, label='batch INSERT column'))
                    children.extend(_collect_embedded_clause_elements(value))
            elif isinstance(row, (list, tuple)):
                children.extend(_collect_embedded_clause_elements(row))
            else:
                raise ScopedSessionTransactionError('TenantUnitOfWork does not allow an unknown batch INSERT row')
    return children


def _iter_scoped_statement_elements(
    statement: sqlalchemy.sql.elements.ClauseElement,
) -> typing.Iterator[sqlalchemy.sql.elements.ClauseElement]:
    """Walk the SQL AST, including dialect clauses SQLAlchemy omits from visitors.

    PostgreSQL and SQLite ``ON CONFLICT`` clauses currently expose no children
    through ``get_children()``.  Their conflict targets and update expressions
    still form executable SQL, so this walker adds those containers explicitly
    and rejects any future, unknown post-values clause by default.
    """

    pending: list[sqlalchemy.sql.elements.ClauseElement] = [statement]
    seen: set[int] = set()
    while pending:
        element = pending.pop()
        identity = id(element)
        if identity in seen:
            continue
        seen.add(identity)
        yield element

        children = [
            child for child in element.get_children() if isinstance(child, sqlalchemy.sql.elements.ClauseElement)
        ]
        post_values_clause = getattr(element, '_post_values_clause', None)
        if post_values_clause is not None:
            if not isinstance(post_values_clause, _ALLOWED_SCOPED_POST_VALUES_TYPES):
                raise ScopedSessionTransactionError(
                    'TenantUnitOfWork does not allow an unknown dialect post-values SQL clause'
                )
            children.append(post_values_clause)

        children.extend(_collect_multi_value_elements(getattr(element, '_multi_values', ())))

        if isinstance(element, _ALLOWED_SCOPED_POST_VALUES_TYPES):
            if getattr(element, 'constraint_target', None) is not None:
                raise ScopedSessionTransactionError('TenantUnitOfWork does not allow named ON CONFLICT constraints')
            for target in getattr(element, 'inferred_target_elements', ()) or ():
                children.extend(_validate_opaque_identifier(target, label='ON CONFLICT target'))
            children.extend(_collect_embedded_clause_elements(getattr(element, 'inferred_target_whereclause', None)))
            children.extend(_collect_embedded_clause_elements(getattr(element, 'update_whereclause', None)))
            for update_pair in getattr(element, 'update_values_to_set', ()):
                if not isinstance(update_pair, (list, tuple)) or len(update_pair) != 2:
                    raise ScopedSessionTransactionError(
                        'TenantUnitOfWork does not allow an unknown ON CONFLICT update shape'
                    )
                key, value = update_pair
                children.extend(_validate_opaque_identifier(key, label='ON CONFLICT update column'))
                children.extend(_collect_embedded_clause_elements(value))

        pending.extend(children)


def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[str, typing.Any]) -> None:
    """Admit only structured SQLAlchemy statements with a small safe vocabulary.

    A textual denylist cannot be complete on PostgreSQL: ordinary built-in
    functions can execute SQL supplied in string arguments, while statement
    prefixes, suffixes, hints, and custom operators can inject syntax without
    appearing as a top-level ``TextClause``.  Tenant code therefore uses
    SQLAlchemy's structured query/DML AST exclusively.  Raw fragments and
    unknown functions/operators fail closed before compilation.
    """

    statement = args[0] if args else kwargs.get('statement')
    if statement is None:
        return
    if kwargs.get('execution_options'):
        raise ScopedSessionTransactionError('TenantUnitOfWork does not allow public SQL execution options')
    if kwargs.get('bind_arguments'):
        raise ScopedSessionTransactionError('TenantUnitOfWork does not allow foreign database bind arguments')
    if not isinstance(statement, _ALLOWED_SCOPED_STATEMENT_TYPES):
        raise ScopedSessionTransactionError(
            'TenantUnitOfWork public SQL must use a structured SQLAlchemy query, values, or DML statement'
        )

    for element in _iter_scoped_statement_elements(statement):
        if not type(element).__module__.startswith('sqlalchemy.'):
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork does not allow custom SQL AST nodes in public statements'
            )

        _validate_scoped_sql_type(getattr(element, 'type', None))

        if isinstance(element, sqlalchemy.sql.selectable.Values):
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork does not allow SQL VALUES constructs in public statements'
            )

        if isinstance(element, sqlalchemy.sql.dml.Insert) and getattr(element, 'select', None) is not None:
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork does not allow INSERT FROM SELECT in public statements'
            )

        for identifier_attribute in ('name', 'schema', 'collation'):
            identifier = getattr(element, identifier_attribute, None)
            if isinstance(identifier, sqlalchemy.sql.elements.quoted_name) and identifier.quote is False:
                raise ScopedSessionTransactionError('TenantUnitOfWork does not allow forced-unquoted SQL identifiers')

        if any(
            getattr(element, attribute_name, None)
            for attribute_name in (
                '_prefixes',
                '_suffixes',
                '_hints',
                '_statement_hints',
                '_execution_options',
                '_with_options',
            )
        ):
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork does not allow textual statement modifiers or execution options'
            )

        if (
            isinstance(element, sqlalchemy.sql.elements.TextClause)
            or getattr(
                element,
                '__visit_name__',
                None,
            )
            == 'textual_label_reference'
        ):
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork does not allow raw or textual SQL fragments in public statements'
            )

        if isinstance(element, sqlalchemy.sql.elements.ColumnClause) and element.is_literal and element.name != '*':
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork does not allow literal SQL columns in public statements'
            )

        if isinstance(element, sqlalchemy.sql.elements.Extract):
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork does not allow SQL EXTRACT fields in public statements'
            )

        if isinstance(element, sqlalchemy.sql.elements.BindParameter) and element.literal_execute:
            raise ScopedSessionTransactionError('TenantUnitOfWork does not allow literal-execute SQL parameters')

        if isinstance(element, sqlalchemy.sql.elements.Cast) and type(element.type) is not Vector:
            raise ScopedSessionTransactionError(
                'TenantUnitOfWork only allows the trusted pgvector cast used by tenant vector search'
            )

        if isinstance(element, sqlalchemy.sql.functions.FunctionElement):
            function_name = str(getattr(element, 'name', '')).casefold()
            package_names = tuple(getattr(element, 'packagenames', ()))
            generic_function_allowed = (
                type(element) is sqlalchemy.sql.functions.Function
                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):
                raise ScopedSessionTransactionError(
                    f'TenantUnitOfWork does not allow SQL function {function_name or "<unknown>"!r}'
                )

        for operator_attribute in ('operator', 'modifier'):
            operator = getattr(element, operator_attribute, None)
            if not isinstance(operator, sqlalchemy.sql.operators.custom_op):
                continue
            operator_name = str(operator.opstring)
            if not (
                isinstance(element, sqlalchemy.sql.elements.BinaryExpression)
                and operator_attribute == 'operator'
                and operator_name in _ALLOWED_SCOPED_CUSTOM_OPERATORS
            ):
                raise ScopedSessionTransactionError(
                    f'TenantUnitOfWork does not allow custom SQL operator {operator_name!r}'
                )


class _NoopSessionEventCollection:
    """Empty SQLAlchemy SessionEvents collection used after fail-closed detection."""

    def __bool__(self) -> bool:
        return False

    def __iter__(self) -> typing.Iterator[typing.Any]:
        return iter(())

    def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> None:
        del args, kwargs


class _NoopSessionEventsDispatch:
    """Prevent a rejected SessionEvents listener from firing during rollback cleanup."""

    _event_names: tuple[str, ...] = ()

    def __getattr__(self, name: str) -> _NoopSessionEventCollection:
        del name
        return _NOOP_SESSION_EVENT_COLLECTION


_NOOP_SESSION_EVENT_COLLECTION = _NoopSessionEventCollection()
_NOOP_SESSION_EVENTS_DISPATCH = _NoopSessionEventsDispatch()


class TenantScopedSyncSession(sqlalchemy_orm.Session):
    """Synchronous proxy usable only from its owning AsyncSession operation."""

    def __init__(
        self,
        bind: sqlalchemy.engine.Engine | sqlalchemy.engine.Connection | None = None,
        *,
        binds: dict[typing.Any, typing.Any] | None = None,
        langbot_guard_state: _ScopedSessionGuardState,
        **kwargs: typing.Any,
    ) -> None:
        if binds:
            raise ScopedSessionTransactionError('Tenant-scoped Sessions cannot use multiple database binds')
        object.__setattr__(self, '_langbot_guard_state', langbot_guard_state)
        object.__setattr__(self, '_langbot_expected_bind', bind)
        object.__setattr__(self, '_langbot_initializing', True)
        try:
            super().__init__(bind=bind, binds=None, **kwargs)
        finally:
            object.__setattr__(self, '_langbot_initializing', False)

    def __getattribute__(self, name: str) -> typing.Any:
        if name.startswith('_'):
            return super().__getattribute__(name)
        self._require_proxy_capability()
        guard = object.__getattribute__(self, '_langbot_guard_state')
        if name == 'dispatch' and guard.session_events_blocked:
            return _NOOP_SESSION_EVENTS_DISPATCH
        attribute = super().__getattribute__(name)
        if name == 'dispatch' and not object.__getattribute__(self, '_langbot_initializing'):
            event_name = next(
                (candidate for candidate in attribute._event_names if bool(getattr(attribute, candidate))),
                None,
            )
            if event_name is not None:
                guard.session_events_blocked = True
                self._reject_proxy_escape(f'ORM Session event listener {event_name}')
        if isinstance(attribute, types.MethodType):
            return self._guard_method_call(attribute)
        return attribute

    def __setattr__(self, name: str, value: typing.Any) -> None:
        if name.startswith('_') or object.__getattribute__(self, '_langbot_initializing'):
            super().__setattr__(name, value)
            return
        self._require_proxy_capability()
        super().__setattr__(name, value)

    def __delattr__(self, name: str) -> None:
        if name.startswith('_') or object.__getattribute__(self, '_langbot_initializing'):
            super().__delattr__(name)
            return
        self._require_proxy_capability()
        super().__delattr__(name)

    def get_bind(
        self,
        mapper: typing.Any = None,
        *,
        clause: typing.Any = None,
        bind: typing.Any = None,
        **kwargs: typing.Any,
    ) -> sqlalchemy.engine.Engine | sqlalchemy.engine.Connection:
        self._require_proxy_capability()
        expected_bind = object.__getattribute__(self, '_langbot_expected_bind')
        if bind is not None and bind is not expected_bind:
            self._reject_proxy_escape('foreign database bind')
        resolved = super().get_bind(mapper, clause=clause, bind=bind, **kwargs)
        if resolved is not expected_bind:
            self._reject_proxy_escape('foreign database bind')
        return resolved

    def flush(self, objects: typing.Sequence[typing.Any] | None = None) -> None:
        """Reject caller-supplied SQL expressions hidden in ORM state."""

        self._require_proxy_capability()
        for instance in self.new.union(self.dirty):
            state = sqlalchemy.inspect(instance)
            for column_property in state.mapper.column_attrs:
                for value in state.attrs[column_property.key].history.added:
                    if isinstance(value, sqlalchemy.sql.elements.ClauseElement):
                        self._reject_proxy_escape('ORM SQL expression attribute value')
                    clause_factory = getattr(value, '__clause_element__', None)
                    if callable(clause_factory) and isinstance(
                        clause_factory(),
                        sqlalchemy.sql.elements.ClauseElement,
                    ):
                        self._reject_proxy_escape('ORM SQL expression attribute value')
        super().flush(objects)

    def _guard_method_call(self, method: types.MethodType) -> typing.Callable[..., typing.Any]:
        @functools.wraps(method)
        def guarded(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
            self._require_proxy_capability()
            return method(*args, **kwargs)

        return guarded

    def _require_proxy_capability(self) -> None:
        if object.__getattribute__(self, '_langbot_initializing'):
            return
        guard = object.__getattribute__(self, '_langbot_guard_state')
        if not guard.active:
            raise ScopedSessionTransactionError('TenantUnitOfWork scoped Session is no longer active')
        try:
            current_task = asyncio.current_task()
        except RuntimeError:
            current_task = None
        if guard.owner_task is not current_task:
            raise CrossScopeTransactionError(
                'Scoped database sessions cannot be inherited by child tasks; open an explicit task scope'
            )
        if _SYNC_PROXY_CAPABILITY.get() is not guard:
            self._reject_proxy_escape('synchronous Session access')

    def _reject_proxy_escape(self, operation: str) -> typing.NoReturn:
        guard = object.__getattribute__(self, '_langbot_guard_state')
        if guard.transaction_state is not None:
            error = ScopedSessionTransactionError(
                f'TenantUnitOfWork owns AsyncSession lifecycle; direct {operation} is not allowed'
            )
            guard.transaction_state.mark_rollback_only(error)
            raise error
        raise ScopedSessionTransactionError(
            f'TenantUnitOfWork owns AsyncSession lifecycle; direct {operation} is not allowed'
        )


class TenantScopedAsyncSession(sqlalchemy_asyncio.AsyncSession):
    """An AsyncSession whose task and transaction lifecycle are UoW-owned.

    Normal ORM operations remain available to the task which opened the UoW.
    Transaction-control and connection-escape APIs are deliberately withheld:
    committing, rolling back, or closing the raw Session would otherwise let a
    nested caller defeat the outer unit of work.  A blanket public-attribute
    owner check also covers a Session captured in the parent task and later
    passed directly to a child task.
    """

    def __init__(
        self,
        bind: sqlalchemy_asyncio.AsyncEngine,
        *,
        owner_task: asyncio.Task[typing.Any] | None,
        **kwargs: typing.Any,
    ) -> None:
        guard = _ScopedSessionGuardState(owner_task=owner_task)
        object.__setattr__(self, '_langbot_guard_state', guard)
        object.__setattr__(self, '_langbot_bind', None)
        object.__setattr__(self, '_langbot_internal_access_depth', 0)
        object.__setattr__(self, '_langbot_sync_capability_tokens', [])
        object.__setattr__(self, '_langbot_initializing', True)
        capability_token = _SYNC_PROXY_CAPABILITY.set(guard)
        try:
            super().__init__(
                bind,
                sync_session_class=TenantScopedSyncSession,
                langbot_guard_state=guard,
                **kwargs,
            )
        finally:
            _SYNC_PROXY_CAPABILITY.reset(capability_token)
            object.__setattr__(self, '_langbot_initializing', False)

    def __getattribute__(self, name: str) -> typing.Any:
        # ContextVars are copied into asyncio child tasks, and callers can also
        # copy the Session object itself.  Guard the whole supported public
        # surface rather than trying to enumerate every current/future ORM I/O
        # method. Private attributes remain an internal implementation detail.
        if name.startswith('_'):
            return super().__getattribute__(name)

        self._require_owner_task()
        if (
            name in {'sync_session', 'binds', 'object_session'}
            and object.__getattribute__(self, '_langbot_internal_access_depth') == 0
        ):
            self._reject_transaction_escape(f'{name} access')
        self._enter_internal_access()
        try:
            attribute = super().__getattribute__(name)
        finally:
            self._exit_internal_access()
        if isinstance(attribute, types.MethodType):
            return self._guard_method_call(attribute)
        return attribute

    def __setattr__(self, name: str, value: typing.Any) -> None:
        if name.startswith('_'):
            super().__setattr__(name, value)
            return
        self._require_owner_task()
        if not object.__getattribute__(self, '_langbot_initializing') and name in {'bind', 'binds', 'sync_session'}:
            self._reject_transaction_escape(f'{name} replacement')
        super().__setattr__(name, value)

    def __delattr__(self, name: str) -> None:
        if name.startswith('_'):
            super().__delattr__(name)
            return
        self._require_owner_task()
        if name in {'bind', 'binds', 'sync_session'}:
            self._reject_transaction_escape(f'{name} deletion')
        super().__delattr__(name)

    @property
    def bind(self) -> typing.NoReturn:
        self._reject_transaction_escape('bind access')

    @bind.setter
    def bind(self, value: sqlalchemy_asyncio.AsyncEngine) -> None:
        existing = object.__getattribute__(self, '_langbot_bind')
        if existing is not None:
            self._reject_transaction_escape('bind replacement')
        object.__setattr__(self, '_langbot_bind', value)

    @property
    def no_autoflush(self) -> typing.ContextManager[None]:
        """Preserve SQLAlchemy's helper without exposing the synchronous Session."""

        @contextlib.contextmanager
        def boundary() -> typing.Iterator[None]:
            self._require_owner_task()
            sync_session = object.__getattribute__(self, '_proxied')
            previous = object.__getattribute__(sync_session, 'autoflush')
            object.__setattr__(sync_session, 'autoflush', False)
            try:
                yield
            finally:
                self._require_owner_task()
                object.__setattr__(sync_session, 'autoflush', previous)

        return boundary()

    def _bind_transaction_state(self, capability: object, state: ActiveScopedTransaction) -> None:
        self._require_control_capability(capability, 'transaction-state binding')
        self._require_owner_task()
        object.__getattribute__(self, '_langbot_guard_state').transaction_state = state

    def begin(self) -> typing.NoReturn:
        self._reject_transaction_escape('begin')

    def begin_nested(self) -> typing.NoReturn:
        self._reject_transaction_escape('begin_nested')

    async def commit(self) -> typing.NoReturn:
        self._reject_transaction_escape('commit')

    async def rollback(self) -> typing.NoReturn:
        self._reject_transaction_escape('rollback')

    async def close(self) -> typing.NoReturn:
        self._reject_transaction_escape('close')

    async def aclose(self) -> typing.NoReturn:
        self._reject_transaction_escape('aclose')

    async def reset(self) -> typing.NoReturn:
        self._reject_transaction_escape('reset')

    async def invalidate(self) -> typing.NoReturn:
        self._reject_transaction_escape('invalidate')

    async def close_all(self) -> typing.NoReturn:
        self._reject_transaction_escape('close_all')

    async def connection(self, *args: typing.Any, **kwargs: typing.Any) -> typing.NoReturn:
        del args, kwargs
        self._reject_transaction_escape('connection')

    def get_bind(self, *args: typing.Any, **kwargs: typing.Any) -> typing.NoReturn:
        del args, kwargs
        self._reject_transaction_escape('get_bind')

    def get_transaction(self) -> typing.NoReturn:
        self._reject_transaction_escape('get_transaction')

    def get_nested_transaction(self) -> typing.NoReturn:
        self._reject_transaction_escape('get_nested_transaction')

    async def run_sync(self, *args: typing.Any, **kwargs: typing.Any) -> typing.NoReturn:
        del args, kwargs
        self._reject_transaction_escape('run_sync')

    async def __aenter__(self) -> typing.NoReturn:
        self._reject_transaction_escape('context-manager entry')

    async def __aexit__(self, *args: typing.Any, **kwargs: typing.Any) -> typing.NoReturn:
        del args, kwargs
        self._reject_transaction_escape('context-manager exit')

    async def _start_owned_transaction(self, capability: object) -> sqlalchemy_asyncio.AsyncSessionTransaction:
        """Start the one root transaction; callable only by TenantUnitOfWork."""

        self._require_control_capability(capability, 'owned transaction start')
        self._require_owner_task()
        self._enter_internal_access()
        try:
            transaction = super().begin()
            await transaction
            return transaction
        finally:
            self._exit_internal_access()

    async def _close_owned_session(self, capability: object) -> None:
        """Close the Session after its root transaction has been finalized."""

        self._require_control_capability(capability, 'owned Session close')
        self._require_owner_task()
        self._enter_internal_access()
        try:
            await super().close()
        finally:
            try:
                self._exit_internal_access()
            finally:
                guard = object.__getattribute__(self, '_langbot_guard_state')
                guard.active = False
                guard.transaction_state = None
                guard.owner_task = None

    async def _commit_owned_transaction(
        self,
        capability: object,
        transaction: sqlalchemy_asyncio.AsyncSessionTransaction,
    ) -> None:
        self._require_control_capability(capability, 'owned transaction commit')
        self._require_owner_task()
        self._enter_internal_access()
        try:
            await transaction.commit()
        finally:
            self._exit_internal_access()

    async def _rollback_owned_transaction(
        self,
        capability: object,
        transaction: sqlalchemy_asyncio.AsyncSessionTransaction,
    ) -> None:
        self._require_control_capability(capability, 'owned transaction rollback')
        self._require_owner_task()
        self._enter_internal_access()
        try:
            await transaction.rollback()
        finally:
            self._exit_internal_access()

    async def _apply_owned_scope_setting(
        self,
        capability: object,
        setting_name: str,
        setting_value: str,
    ) -> None:
        self._require_control_capability(capability, 'tenant scope configuration')
        if setting_name not in {
            TENANT_SETTING,
            ACCOUNT_SETTING,
            API_KEY_HASH_SETTING,
            INVITATION_HASH_SETTING,
            INSTANCE_SETTING,
            IDENTITY_DIGEST_SETTING,
            DIRECTORY_INSTANCE_SETTING,
        }:
            self._reject_transaction_escape('unknown tenant scope configuration')
        self._require_owner_task()
        self._enter_internal_access()
        try:
            await super().execute(
                sqlalchemy.text('SELECT set_config(:setting_name, :setting_value, true)'),
                {'setting_name': setting_name, 'setting_value': setting_value},
            )
        finally:
            self._exit_internal_access()

    async def execute_on_transaction_connection(
        self,
        *args: typing.Any,
        **kwargs: typing.Any,
    ) -> sqlalchemy.engine.cursor.CursorResult[typing.Any]:
        """Execute Core SQL without exposing the transaction-bound Connection."""

        self._require_owner_task()
        self._validate_statement_or_reject(args, kwargs)
        self._enter_internal_access()
        try:
            await super().flush()
            connection = await super().connection()
            return await connection.execute(*args, **kwargs)
        finally:
            self._exit_internal_access()

    async def execute(self, *args: typing.Any, **kwargs: typing.Any) -> sqlalchemy.engine.Result[typing.Any]:
        self._validate_statement_or_reject(args, kwargs)
        return await super().execute(*args, **kwargs)

    async def get(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
        self._reject_nonempty_orm_options(
            'get',
            kwargs,
            {'options', 'with_for_update', 'identity_token', 'execution_options', 'bind_arguments'},
        )
        return await super().get(*args, **kwargs)

    async def get_one(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
        self._reject_nonempty_orm_options(
            'get_one',
            kwargs,
            {'options', 'with_for_update', 'identity_token', 'execution_options', 'bind_arguments'},
        )
        return await super().get_one(*args, **kwargs)

    async def refresh(self, *args: typing.Any, **kwargs: typing.Any) -> None:
        self._reject_nonempty_orm_options('refresh', kwargs, {'with_for_update'})
        await super().refresh(*args, **kwargs)

    async def merge(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
        self._reject_nonempty_orm_options('merge', kwargs, {'options'})
        return await super().merge(*args, **kwargs)

    async def scalar(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
        self._validate_statement_or_reject(args, kwargs)
        return await super().scalar(*args, **kwargs)

    async def scalars(self, *args: typing.Any, **kwargs: typing.Any) -> sqlalchemy.engine.ScalarResult[typing.Any]:
        self._validate_statement_or_reject(args, kwargs)
        return await super().scalars(*args, **kwargs)

    async def stream(self, *args: typing.Any, **kwargs: typing.Any) -> sqlalchemy_asyncio.AsyncResult[typing.Any]:
        del args, kwargs
        self._reject_transaction_escape('stream; live database results cannot outlive the task-owned UoW operation')

    async def stream_scalars(
        self,
        *args: typing.Any,
        **kwargs: typing.Any,
    ) -> sqlalchemy_asyncio.AsyncScalarResult[typing.Any]:
        del args, kwargs
        self._reject_transaction_escape(
            'stream_scalars; live database results cannot outlive the task-owned UoW operation'
        )

    def _validate_statement_or_reject(
        self,
        args: tuple[typing.Any, ...],
        kwargs: dict[str, typing.Any],
    ) -> None:
        try:
            _validate_scoped_statement_call(args, kwargs)
        except ScopedSessionTransactionError as exc:
            guard = object.__getattribute__(self, '_langbot_guard_state')
            if guard.transaction_state is not None:
                guard.transaction_state.mark_rollback_only(exc)
            raise

    def _reject_nonempty_orm_options(
        self,
        operation: str,
        kwargs: dict[str, typing.Any],
        option_names: set[str],
    ) -> None:
        for option_name in option_names:
            value = kwargs.get(option_name)
            if option_name == 'with_for_update':
                if value is None or value is False:
                    continue
                self._reject_transaction_escape(f'{operation} option {option_name}')
            if option_name == 'identity_token':
                if value is None:
                    continue
                self._reject_transaction_escape(f'{operation} option {option_name}')
            if value is None or value is False:
                continue
            if isinstance(value, collections.abc.Sized) and len(value) == 0:
                continue
            self._reject_transaction_escape(f'{operation} option {option_name}')

    def _guard_method_call(self, method: types.MethodType) -> typing.Callable[..., typing.Any]:
        if asyncio.iscoroutinefunction(method):

            @functools.wraps(method)
            async def guarded_async(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
                self._require_owner_task()
                self._enter_internal_access()
                try:
                    return await method(*args, **kwargs)
                finally:
                    self._exit_internal_access()

            return guarded_async

        @functools.wraps(method)
        def guarded_sync(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
            self._require_owner_task()
            self._enter_internal_access()
            try:
                return method(*args, **kwargs)
            finally:
                self._exit_internal_access()

        return guarded_sync

    def _enter_internal_access(self) -> None:
        depth = object.__getattribute__(self, '_langbot_internal_access_depth')
        guard = object.__getattribute__(self, '_langbot_guard_state')
        tokens = object.__getattribute__(self, '_langbot_sync_capability_tokens')
        tokens.append(_SYNC_PROXY_CAPABILITY.set(guard))
        object.__setattr__(self, '_langbot_internal_access_depth', depth + 1)

    def _exit_internal_access(self) -> None:
        depth = object.__getattribute__(self, '_langbot_internal_access_depth')
        if depth <= 0:
            raise RuntimeError('Scoped Session internal access stack underflow')
        tokens = object.__getattribute__(self, '_langbot_sync_capability_tokens')
        _SYNC_PROXY_CAPABILITY.reset(tokens.pop())
        object.__setattr__(self, '_langbot_internal_access_depth', depth - 1)

    def _require_control_capability(self, capability: object, operation: str) -> None:
        if capability is not _UOW_SESSION_CONTROL_CAPABILITY:
            self._reject_transaction_escape(operation)

    def _require_owner_task(self) -> None:
        guard = object.__getattribute__(self, '_langbot_guard_state')
        if not guard.active:
            raise ScopedSessionTransactionError('TenantUnitOfWork scoped Session is no longer active')
        owner_task = guard.owner_task
        try:
            current_task = asyncio.current_task()
        except RuntimeError:
            current_task = None
        if owner_task is not current_task:
            raise CrossScopeTransactionError(
                'Scoped database sessions cannot be inherited by child tasks; open an explicit task scope'
            )

    def _reject_transaction_escape(self, operation: str) -> typing.NoReturn:
        self._require_owner_task()
        error = ScopedSessionTransactionError(
            f'TenantUnitOfWork owns AsyncSession transaction lifecycle; direct {operation} is not allowed'
        )
        state = object.__getattribute__(self, '_langbot_guard_state').transaction_state
        if state is not None:
            state.mark_rollback_only(error)
        raise error


ActiveTransactionVar = contextvars.ContextVar[ActiveScopedTransaction | None]

# ``AsyncSession`` delegates database I/O to a greenlet-backed synchronous
# Engine, but Python context variables are preserved across that boundary.  A
# per-engine ``handle_error`` listener therefore covers DBAPI failures from all
# direct Session APIs (execute/scalar/get/flush) as well as callers which obtain
# the transaction-bound Connection.  The owner-task and engine checks prevent
# copied child-task contexts or unrelated database engines from poisoning this
# transaction.
_DATABASE_OPERATION_TRANSACTION: contextvars.ContextVar[ActiveScopedTransaction | None] = contextvars.ContextVar(
    'langbot_database_operation_transaction',
    default=None,
)


def _mark_current_transaction_rollback_only(exception_context: typing.Any) -> None:
    state = _DATABASE_OPERATION_TRANSACTION.get()
    if state is None:
        return
    try:
        current_task = asyncio.current_task()
    except RuntimeError:  # pragma: no cover - async engines always run in an event loop
        return
    if state.owner_task is not current_task or exception_context.engine is not state.engine.sync_engine:
        return
    cause = exception_context.sqlalchemy_exception or exception_context.original_exception
    state.mark_rollback_only(cause)


def _install_rollback_only_error_listener(engine: sqlalchemy_asyncio.AsyncEngine) -> None:
    sync_engine = engine.sync_engine
    if not sqlalchemy.event.contains(sync_engine, 'handle_error', _mark_current_transaction_rollback_only):
        sqlalchemy.event.listen(sync_engine, 'handle_error', _mark_current_transaction_rollback_only)


@dataclasses.dataclass(slots=True)
class ActivePersistenceScope:
    """A trusted visibility scope without an open database transaction."""

    scope: PersistenceScope
    owner_task: asyncio.Task[typing.Any] | None
    depth: int = 1


ActivePersistenceScopeVar = contextvars.ContextVar[ActivePersistenceScope | None]


class PersistenceScopeBoundary:
    """Carry a trusted persistence scope across long-running async work.

    Unlike :class:`TenantUnitOfWork`, this boundary never opens a database
    session. ``PersistenceManager.execute_async`` materializes the scope as a
    short transaction for each database call. This makes the boundary suitable
    for request and pipeline lifetimes which can spend substantial time waiting
    on model providers, tools, or streaming clients.

    Context variables are copied into child tasks, so implicit use from a child
    task is rejected by the manager. A child may establish its own explicit
    boundary because the scope value still comes from trusted application code.
    """

    def __init__(
        self,
        scope: PersistenceScope,
        *,
        active_scope: ActivePersistenceScopeVar,
        active_transaction: ActiveTransactionVar,
    ) -> None:
        self.scope = scope
        self._active_scope = active_scope
        self._active_transaction = active_transaction
        self._active_state: ActivePersistenceScope | None = None
        self._context_token: contextvars.Token[ActivePersistenceScope | None] | None = None
        self._used = False
        self._owns_scope = False

    async def __aenter__(self) -> PersistenceScopeBoundary:
        if self._used:
            raise RuntimeError('PersistenceScopeBoundary instances cannot be reused')
        self._used = True

        transaction = self._active_transaction.get()
        if transaction is not None:
            if transaction.owner_task is not asyncio.current_task():
                raise CrossScopeTransactionError(
                    'Scoped database transactions cannot be inherited by child tasks; open an explicit task scope'
                )
            if transaction.scope != self.scope:
                raise CrossScopeTransactionError(
                    f'Cannot enter {self.scope.kind.value} scope while {transaction.scope.kind.value} scope is active'
                )

        active = self._active_scope.get()
        if active is not None and active.owner_task is asyncio.current_task():
            if active.scope != self.scope:
                raise CrossScopeTransactionError(
                    f'Cannot enter {self.scope.kind.value} scope while {active.scope.kind.value} scope is active'
                )
            active.depth += 1
            self._active_state = active
            return self

        state = ActivePersistenceScope(
            scope=self.scope,
            owner_task=asyncio.current_task(),
        )
        self._active_state = state
        self._context_token = self._active_scope.set(state)
        self._owns_scope = True
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: types.TracebackType | None,
    ) -> bool:
        del exc_type, exc_value, traceback
        state = self._active_state
        if state is None:
            raise RuntimeError('PersistenceScopeBoundary has no active state')
        if state.owner_task is not asyncio.current_task():
            raise CrossScopeTransactionError('Scoped persistence boundaries cannot be exited by a different task')

        if not self._owns_scope:
            state.depth -= 1
            self._active_state = None
            return False

        if self._context_token is None:
            raise RuntimeError('PersistenceScopeBoundary has no context token')
        self._active_scope.reset(self._context_token)
        state.depth = 0
        self._active_state = None
        return False


class TenantUnitOfWork:
    """Bind one trusted visibility scope to one database transaction.

    PostgreSQL receives transaction-local settings consumed by RLS policies.
    SQLite keeps the same transaction boundary so OSS code exercises the same
    request/task ownership and rollback semantics. A manager-owned ContextVar
    lets all legacy ``execute_async`` helpers reuse this exact session.
    """

    def __init__(
        self,
        engine: sqlalchemy_asyncio.AsyncEngine,
        workspace_uuid: str | None = None,
        *,
        scope: PersistenceScope | None = None,
        active_transaction: ActiveTransactionVar | None = None,
        active_scope: ActivePersistenceScopeVar | None = None,
        on_pool_timeout: typing.Callable[[], None] | None = None,
    ) -> None:
        if (workspace_uuid is None) == (scope is None):
            raise ValueError('TenantUnitOfWork requires exactly one Workspace or persistence scope')

        self._engine = engine
        self.scope = scope or PersistenceScope.workspace(typing.cast(str, workspace_uuid))
        self.workspace_uuid = self.scope.settings[0][1] if self.scope.kind == PersistenceScopeKind.WORKSPACE else None
        self._active_transaction = active_transaction
        self._active_scope = active_scope
        self._on_pool_timeout = on_pool_timeout
        self._session: sqlalchemy_asyncio.AsyncSession | None = None
        self._transaction: sqlalchemy_asyncio.AsyncSessionTransaction | None = None
        self._active_state: ActiveScopedTransaction | None = None
        self._context_token: contextvars.Token[ActiveScopedTransaction | None] | None = None
        self._database_operation_token: contextvars.Token[ActiveScopedTransaction | None] | None = None
        self._used = False
        self._owns_transaction = False
        _install_rollback_only_error_listener(engine)

    @property
    def session(self) -> sqlalchemy_asyncio.AsyncSession:
        if self._session is None:
            raise RuntimeError('TenantUnitOfWork is not active')
        if self._active_state is not None:
            self._require_owner_task(self._active_state)
        return self._session

    async def __aenter__(self) -> TenantUnitOfWork:
        if self._used:
            raise RuntimeError('TenantUnitOfWork instances cannot be reused')
        self._used = True

        active_scope = self._active_scope.get() if self._active_scope is not None else None
        if active_scope is not None and active_scope.owner_task is asyncio.current_task():
            if active_scope.scope != self.scope:
                # A request/runtime Workspace boundary is transaction-free and
                # may temporarily enter a narrower trusted discovery UoW (for
                # example, resolving an account record by its verified UUID).
                # It must still never switch directly to another Workspace.
                workspace_to_discovery = (
                    active_scope.scope.kind == PersistenceScopeKind.WORKSPACE
                    and self.scope.kind
                    in {
                        PersistenceScopeKind.ACCOUNT_DISCOVERY,
                        PersistenceScopeKind.API_KEY_DISCOVERY,
                        PersistenceScopeKind.INVITATION_DISCOVERY,
                        PersistenceScopeKind.INSTANCE_DISCOVERY,
                        PersistenceScopeKind.IDENTITY_DISCOVERY,
                    }
                )
                if not workspace_to_discovery:
                    raise CrossScopeTransactionError(
                        f'Cannot enter {self.scope.kind.value} scope '
                        f'while {active_scope.scope.kind.value} scope is active'
                    )

        active = self._active_transaction.get() if self._active_transaction is not None else None
        if active is not None:
            if active.owner_task is asyncio.current_task():
                if active.scope != self.scope:
                    raise CrossScopeTransactionError(
                        f'Cannot enter {self.scope.kind.value} scope while {active.scope.kind.value} scope is active'
                    )
                active.depth += 1
                self._active_state = active
                self._session = active.session
                return self
            # ContextVars are copied into child tasks. Merely calling a helper
            # there must fail (PersistenceManager enforces that), but an
            # explicit UoW is allowed to replace the inherited pointer with an
            # independent transaction owned by the child task.

        owner_task = asyncio.current_task()
        session = TenantScopedAsyncSession(
            self._engine,
            owner_task=owner_task,
            expire_on_commit=False,
            close_resets_only=False,
        )
        self._session = session
        transaction: sqlalchemy_asyncio.AsyncSessionTransaction | None = None
        try:
            transaction = await session._start_owned_transaction(_UOW_SESSION_CONTROL_CAPABILITY)
            self._transaction = transaction
            if self._engine.dialect.name == 'postgresql':
                for setting_name, setting_value in self.scope.settings:
                    await session._apply_owned_scope_setting(
                        _UOW_SESSION_CONTROL_CAPABILITY,
                        setting_name,
                        setting_value,
                    )
            state = ActiveScopedTransaction(
                scope=self.scope,
                session=session,
                engine=self._engine,
                owner_task=owner_task,
            )
            session._bind_transaction_state(_UOW_SESSION_CONTROL_CAPABILITY, state)
            self._active_state = state
            if self._active_transaction is not None:
                self._context_token = self._active_transaction.set(state)
            self._database_operation_token = _DATABASE_OPERATION_TRANSACTION.set(state)
            self._owns_transaction = True
        except BaseException as exc:
            if isinstance(exc, sqlalchemy.exc.TimeoutError) and self._on_pool_timeout is not None:
                self._on_pool_timeout()
            if self._database_operation_token is not None:
                _DATABASE_OPERATION_TRANSACTION.reset(self._database_operation_token)
                self._database_operation_token = None
            if self._active_transaction is not None and self._context_token is not None:
                self._active_transaction.reset(self._context_token)
                self._context_token = None
            if transaction is not None and transaction.sync_transaction is not None:
                await session._rollback_owned_transaction(_UOW_SESSION_CONTROL_CAPABILITY, transaction)
            await session._close_owned_session(_UOW_SESSION_CONTROL_CAPABILITY)
            self._session = None
            self._transaction = None
            raise
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: types.TracebackType | None,
    ) -> bool:
        state = self._active_state
        if state is None:
            raise RuntimeError('TenantUnitOfWork has no active transaction state')
        self._require_owner_task(state)

        if not self._owns_transaction:
            if exc_type is not None:
                state.mark_rollback_only(exc_value)
            state.depth -= 1
            self._session = None
            self._active_state = None
            return False

        session = self.session
        transaction = self._transaction
        if transaction is None:
            raise RuntimeError('TenantUnitOfWork has no active root transaction')
        if exc_type is not None:
            state.mark_rollback_only(exc_value)
        rollback_only = state.rollback_only
        committed = False
        try:
            if exc_type is None and not rollback_only:
                await typing.cast(TenantScopedAsyncSession, session)._commit_owned_transaction(
                    _UOW_SESSION_CONTROL_CAPABILITY,
                    transaction,
                )
                committed = True
            else:
                await typing.cast(TenantScopedAsyncSession, session)._rollback_owned_transaction(
                    _UOW_SESSION_CONTROL_CAPABILITY,
                    transaction,
                )
        finally:
            try:
                if self._active_transaction is not None and self._context_token is not None:
                    self._active_transaction.reset(self._context_token)
                await typing.cast(TenantScopedAsyncSession, session)._close_owned_session(
                    _UOW_SESSION_CONTROL_CAPABILITY
                )
            finally:
                if self._database_operation_token is not None:
                    _DATABASE_OPERATION_TRANSACTION.reset(self._database_operation_token)
                self._database_operation_token = None
                self._session = None
                self._transaction = None
                self._active_state = None
                state.depth = 0
                self._complete_after_commit_waiters(state, committed=committed)

        if exc_type is None and rollback_only:
            raise TransactionRollbackOnlyError(
                'A scoped database operation failed or rolled back; '
                'the transaction was rolled back and after-commit work was cancelled'
            ) from state.rollback_only_cause
        return False

    async def execute(self, *args: typing.Any, **kwargs: typing.Any) -> sqlalchemy.engine.Result[typing.Any]:
        try:
            return await self.session.execute(*args, **kwargs)
        except BaseException as exc:
            self._mark_rollback_only(exc)
            raise

    async def flush(self) -> None:
        try:
            await self.session.flush()
        except BaseException as exc:
            self._mark_rollback_only(exc)
            raise

    def _mark_rollback_only(self, cause: BaseException) -> None:
        state = self._active_state
        if state is not None:
            state.mark_rollback_only(cause)

    @staticmethod
    def _require_owner_task(state: ActiveScopedTransaction) -> None:
        if state.owner_task is not asyncio.current_task():
            raise CrossScopeTransactionError(
                'Scoped database transactions cannot be inherited by child tasks; open an explicit task scope'
            )

    @staticmethod
    def _complete_after_commit_waiters(
        state: ActiveScopedTransaction,
        *,
        committed: bool,
    ) -> None:
        for waiter in state.after_commit_waiters:
            if waiter.done():
                continue
            if committed:
                waiter.set_result(None)
            else:
                waiter.cancel()
        state.after_commit_waiters.clear()