[autofix.ci] apply automated fixes

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Byron.wang <byron@dify.ai>

api/.env.example

@@ -666,6 +666,7 @@ PLUGIN_REMOTE_INSTALL_PORT=5003
 PLUGIN_REMOTE_INSTALL_HOST=localhost
 PLUGIN_MAX_PACKAGE_SIZE=15728640
 PLUGIN_MODEL_SCHEMA_CACHE_TTL=3600
+PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED=true
 PLUGIN_MODEL_PROVIDERS_CACHE_TTL=86400
 # Comma-separated marketplace plugin IDs whose latest versions are installed for newly registered users.
 # Example: langgenius/openai,langgenius/gemini
@@ -677,6 +678,8 @@ INNER_API_KEY_FOR_PLUGIN=QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y
 
 # Dify Agent backend
 AGENT_BACKEND_BASE_URL=http://localhost:5050
+# Bearer token sent to the Agent backend /runs API. Must match DIFY_AGENT_API_TOKEN on the server side.
+AGENT_BACKEND_API_TOKEN=dify-agent-run-token-for-dev-only
 AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS=30
 AGENT_BACKEND_STREAM_MAX_RECONNECTS=3
 AGENT_BACKEND_RUN_TIMEOUT_SECONDS=1200

api/app_factory.py

@@ -1,10 +1,13 @@
 import logging
 import time
+from collections.abc import Callable
+from typing import NamedTuple
 
 import socketio
 from flask import request
 from opentelemetry.trace import get_current_span
 from opentelemetry.trace.span import INVALID_SPAN_ID, INVALID_TRACE_ID
+from werkzeug.exceptions import Forbidden, HTTPException, ServiceUnavailable
 
 from configs import dify_config
 from contexts.wrapper import RecyclableContextVar
@@ -42,6 +45,53 @@ _CONSOLE_EXEMPT_PREFIXES = (
     "/console/api/activate/check",
 )
 
+_WEBAPP_EXEMPT_PREFIXES = ("/api/system-features",)
+
+_INVALID_LICENSE_STATUSES = (LicenseStatus.INACTIVE, LicenseStatus.EXPIRED, LicenseStatus.LOST)
+
+
+def _session_surface_error(license_status: LicenseStatus | None) -> HTTPException:
+    if license_status is None:
+        return UnauthorizedAndForceLogout("Unable to verify enterprise license. Please contact your administrator.")
+    return UnauthorizedAndForceLogout(f"Enterprise license is {license_status}. Please contact your administrator.")
+
+
+def _bearer_surface_error(license_status: LicenseStatus | None) -> HTTPException:
+    """Token-authed: forcing a logout is meaningless and license state must not leak."""
+    return Forbidden(description="license_required")
+
+
+def _retryable_surface_error(license_status: LicenseStatus | None) -> HTTPException:
+    """Webhook senders retry on 5xx but treat 4xx as permanent, disabling the subscription."""
+    return ServiceUnavailable(description="license_required")
+
+
+class _LicenseGatedSurface(NamedTuple):
+    prefix: str
+    exempt_prefixes: tuple[str, ...]
+    build_error: Callable[[LicenseStatus | None], HTTPException]
+
+
+# /files (plugin-daemon data plane), /inner/api (enterprise control plane) and /health
+# stay ungated: blocking them breaks workflow execution or license recovery itself.
+_LICENSE_GATED_SURFACES = (
+    _LicenseGatedSurface("/console/api/", _CONSOLE_EXEMPT_PREFIXES, _session_surface_error),
+    _LicenseGatedSurface("/api/", _WEBAPP_EXEMPT_PREFIXES, _session_surface_error),
+    _LicenseGatedSurface("/v1", (), _bearer_surface_error),
+    _LicenseGatedSurface("/mcp", (), _bearer_surface_error),
+    _LicenseGatedSurface("/triggers", (), _retryable_surface_error),
+)
+
+
+def _match_license_gated_surface(path: str) -> _LicenseGatedSurface | None:
+    for surface in _LICENSE_GATED_SURFACES:
+        if not path.startswith(surface.prefix):
+            continue
+        if any(path.startswith(exempt) for exempt in surface.exempt_prefixes):
+            return None
+        return surface
+    return None
+
 
 # ----------------------------
 # Application Factory Function
@@ -62,38 +112,17 @@ def create_flask_app_with_configs() -> DifyApp:
         init_request_context()
         RecyclableContextVar.increment_thread_recycles()
 
-        # Enterprise license validation for API endpoints (both console and webapp)
-        # When license expires, block all API access except bootstrap endpoints needed
-        # for the frontend to load the license expiration page without infinite reloads.
         if dify_config.ENTERPRISE_ENABLED:
-            is_console_api = request.path.startswith("/console/api/")
-            is_webapp_api = request.path.startswith("/api/")
+            surface = _match_license_gated_surface(request.path)
+            if surface is not None:
+                try:
+                    license_status = EnterpriseService.get_cached_license_status()
+                except Exception:
+                    logger.exception("Failed to check enterprise license status")
+                    license_status = None
 
-            if is_console_api or is_webapp_api:
-                if is_console_api:
-                    is_exempt = any(request.path.startswith(p) for p in _CONSOLE_EXEMPT_PREFIXES)
-                else:  # webapp API
-                    is_exempt = request.path.startswith("/api/system-features")
-
-                if not is_exempt:
-                    try:
-                        # Check license status (cached — see EnterpriseService for TTL details)
-                        license_status = EnterpriseService.get_cached_license_status()
-                        if license_status in (LicenseStatus.INACTIVE, LicenseStatus.EXPIRED, LicenseStatus.LOST):
-                            raise UnauthorizedAndForceLogout(
-                                f"Enterprise license is {license_status}. Please contact your administrator."
-                            )
-                        if license_status is None:
-                            raise UnauthorizedAndForceLogout(
-                                "Unable to verify enterprise license. Please contact your administrator."
-                            )
-                    except UnauthorizedAndForceLogout:
-                        raise
-                    except Exception:
-                        logger.exception("Failed to check enterprise license status")
-                        raise UnauthorizedAndForceLogout(
-                            "Unable to verify enterprise license. Please contact your administrator."
-                        )
+                if license_status is None or license_status in _INVALID_LICENSE_STATUSES:
+                    raise surface.build_error(license_status)
 
     # add after request hook for injecting trace headers from OpenTelemetry span context
     # Only adds headers when OTEL is enabled and has valid context

api/clients/agent_backend/factory.py

@@ -11,6 +11,7 @@ from clients.agent_backend.fake_client import FakeAgentBackendRunClient, FakeAge
 def create_agent_backend_run_client(
     *,
     base_url: str | None = None,
+    api_token: str | None = None,
     use_fake: bool = False,
     fake_scenario: str | FakeAgentBackendScenario = FakeAgentBackendScenario.SUCCESS,
     stream_read_timeout_seconds: float = 30,
@@ -22,8 +23,11 @@ def create_agent_backend_run_client(
         return FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario(fake_scenario))
     if base_url is None:
         raise ValueError("base_url is required when creating a real Agent backend client")
+    headers: dict[str, str] = {}
+    if api_token:
+        headers["Authorization"] = f"Bearer {api_token}"
     return DifyAgentBackendRunClient(
-        Client(base_url=base_url, stream_timeout=stream_read_timeout_seconds),
+        Client(base_url=base_url, stream_timeout=stream_read_timeout_seconds, headers=headers),
         stream_max_reconnects=stream_max_reconnects,
         stream_timeout_seconds=stream_run_timeout_seconds,
     )

api/commands/__init__.py

@@ -3,6 +3,7 @@ CLI command modules extracted from `commands.py`.
 """
 
 from .account import create_tenant, reset_email, reset_password
+from .app_maintenance import convert_to_agent_apps, fix_app_site_missing
 from .data_migrate import data_migrate, legacy_model_types
 from .data_migration import (
     export_migration_data,
@@ -10,6 +11,7 @@ from .data_migration import (
     import_migration_data,
     migration_data_wizard,
 )
+from .database import upgrade_db
 from .plugin import (
     backfill_plugin_auto_upgrade,
     extract_plugins,
@@ -36,12 +38,6 @@ from .retention import (
     restore_workflow_runs,
 )
 from .storage import clear_orphaned_file_records, file_usage, migrate_oss, remove_orphaned_files_on_storage
-from .system import (
-    convert_to_agent_apps,
-    fix_app_site_missing,
-    reset_encrypt_key_pair,
-    upgrade_db,
-)
 from .vector import (
     add_qdrant_index,
     migrate_annotation_vector_database,
@@ -49,6 +45,8 @@ from .vector import (
     old_metadata_migration,
     vdb_migrate,
 )
+from .workflow_migration import migrate_legacy_sys_files_workflows
+from .workspace import reset_encrypt_key_pair
 
 __all__ = [
     "add_qdrant_index",
@@ -80,6 +78,7 @@ __all__ = [
     "migrate_data_for_plugin",
     "migrate_dataset_permissions_to_rbac",
     "migrate_knowledge_vector_database",
+    "migrate_legacy_sys_files_workflows",
     "migrate_member_roles_to_rbac",
     "migrate_oss",
     "migration_data_wizard",

api/commands/app_maintenance.py

@@ -1,86 +1,27 @@
+"""App data maintenance CLI commands."""
+
 import logging
 
 import click
 import sqlalchemy as sa
-from sqlalchemy import delete, select, update
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy import select, update
 
-from configs import dify_config
-from enums.deployment_edition import DeploymentEdition
 from events.app_event import app_was_created
 from extensions.ext_database import db
-from extensions.ext_redis import redis_client
-from libs.db_migration_lock import DbMigrationAutoRenewLock
-from libs.rsa import generate_key_pair
 from models import Tenant
 from models.model import App, AppMode, Conversation
-from models.provider import Provider, ProviderModel
-from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider
 
 logger = logging.getLogger(__name__)
 
-DB_UPGRADE_LOCK_TTL_SECONDS = 60
-
-
-@click.command(
-    "reset-encrypt-key-pair",
-    help="Reset the asymmetric key pair of workspace for encrypt LLM credentials. "
-    "After the reset, all LLM credentials and tool provider credentials "
-    "(builtin / API / MCP) will be purged, requiring re-entry. "
-    "Only support SELF_HOSTED mode.",
-)
-@click.confirmation_option(
-    prompt=click.style(
-        "Are you sure you want to reset encrypt key pair? "
-        "This will also purge builtin / API / MCP tool provider records for every tenant. "
-        "This operation cannot be rolled back!",
-        fg="red",
-    )
-)
-def reset_encrypt_key_pair():
-    """
-    Reset the encrypted key pair of workspace for encrypt LLM credentials.
-    After the reset, all LLM credentials will become invalid, requiring re-entry.
-    Only support SELF_HOSTED mode.
-    """
-    if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD:
-        click.echo(click.style("This command is only for SELF_HOSTED installations.", fg="red"))
-        return
-    with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
-        tenants = session.scalars(select(Tenant)).all()
-        for tenant in tenants:
-            if not tenant:
-                click.echo(click.style("No workspaces found. Run /install first.", fg="red"))
-                return
-
-            tenant.encrypt_public_key = generate_key_pair(tenant.id)
-
-            session.execute(delete(Provider).where(Provider.provider_type == "custom", Provider.tenant_id == tenant.id))
-            session.execute(delete(ProviderModel).where(ProviderModel.tenant_id == tenant.id))
-
-            # Purge tool provider records that hold credentials encrypted under the
-            # tenant key. Leaving them in place causes /console/api/workspaces/current/
-            # tool-providers to 500 because decryption fails on stale ciphertext (#35396).
-            session.execute(delete(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant.id))
-            session.execute(delete(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant.id))
-            session.execute(delete(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant.id))
-
-            click.echo(
-                click.style(
-                    f"Congratulations! The asymmetric key pair of workspace {tenant.id} has been reset.",
-                    fg="green",
-                )
-            )
-
 
 @click.command("convert-to-agent-apps", help="Convert Agent Assistant to Agent App.")
-def convert_to_agent_apps():
+def convert_to_agent_apps() -> None:
     """
     Convert Agent Assistant to Agent App.
     """
     click.echo(click.style("Starting convert to agent apps.", fg="green"))
 
-    proceeded_app_ids = []
+    proceeded_app_ids: list[str] = []
 
     while True:
         # fetch first 1000 apps
@@ -133,48 +74,14 @@ def convert_to_agent_apps():
     click.echo(click.style(f"Conversion complete. Converted {len(proceeded_app_ids)} agent apps.", fg="green"))
 
 
-@click.command("upgrade-db", help="Upgrade the database")
-def upgrade_db():
-    click.echo("Preparing database migration...")
-    lock = DbMigrationAutoRenewLock(
-        redis_client=redis_client,
-        name="db_upgrade_lock",
-        ttl_seconds=DB_UPGRADE_LOCK_TTL_SECONDS,
-        logger=logger,
-        log_context="db_migration",
-    )
-    if lock.acquire(blocking=False):
-        migration_succeeded = False
-        try:
-            click.echo(click.style("Starting database migration.", fg="green"))
-
-            # run db migration
-            import flask_migrate
-
-            flask_migrate.upgrade()
-
-            migration_succeeded = True
-            click.echo(click.style("Database migration successful!", fg="green"))
-
-        except Exception as e:
-            logger.exception("Failed to execute database migration")
-            click.echo(click.style(f"Database migration failed: {e}", fg="red"))
-            raise SystemExit(1)
-        finally:
-            status = "successful" if migration_succeeded else "failed"
-            lock.release_safely(status=status)
-    else:
-        click.echo("Database migration skipped")
-
-
 @click.command("fix-app-site-missing", help="Fix app related site missing issue.")
-def fix_app_site_missing():
+def fix_app_site_missing() -> None:
     """
     Fix app related site missing issue.
     """
     click.echo(click.style("Starting fix for missing app-related sites.", fg="green"))
 
-    failed_app_ids = []
+    failed_app_ids: list[str] = []
     while True:
         sql = """select apps.id as id from apps left join sites on sites.app_id=apps.id
 where sites.id is null limit 1000"""

api/commands/database.py

@@ -0,0 +1,45 @@
+"""Database schema migration CLI commands."""
+
+import logging
+
+import click
+
+from extensions.ext_redis import redis_client
+from libs.db_migration_lock import DbMigrationAutoRenewLock
+
+logger = logging.getLogger(__name__)
+
+DB_UPGRADE_LOCK_TTL_SECONDS = 60
+
+
+@click.command("upgrade-db", help="Upgrade the database")
+def upgrade_db() -> None:
+    click.echo("Preparing database migration...")
+    lock = DbMigrationAutoRenewLock(
+        redis_client=redis_client,
+        name="db_upgrade_lock",
+        ttl_seconds=DB_UPGRADE_LOCK_TTL_SECONDS,
+        logger=logger,
+        log_context="db_migration",
+    )
+    if lock.acquire(blocking=False):
+        migration_succeeded = False
+        try:
+            click.echo(click.style("Starting database migration.", fg="green"))
+
+            import flask_migrate
+
+            flask_migrate.upgrade()
+
+            migration_succeeded = True
+            click.echo(click.style("Database migration successful!", fg="green"))
+
+        except Exception as e:
+            logger.exception("Failed to execute database migration")
+            click.echo(click.style(f"Database migration failed: {e}", fg="red"))
+            raise SystemExit(1)
+        finally:
+            status = "successful" if migration_succeeded else "failed"
+            lock.release_safely(status=status)
+    else:
+        click.echo("Database migration skipped")

api/commands/workflow_migration.py

@@ -0,0 +1,173 @@
+"""Workflow data migration CLI commands.
+
+TODO: Remove the legacy system file workflow migration command after the production migration is complete.
+"""
+
+import logging
+from dataclasses import dataclass
+
+import click
+from sqlalchemy import select
+from sqlalchemy.orm import Session, load_only, sessionmaker
+
+from extensions.ext_database import db
+from models.workflow import Workflow, WorkflowType
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LegacySysFilesWorkflowMigrationStats:
+    scanned: int = 0
+    migrated: int = 0
+    failed: int = 0
+    batches: int = 0
+    last_id: str | None = None
+
+
+def _build_legacy_sys_files_workflow_query(
+    *,
+    start_after_id: str | None,
+    batch_size: int,
+    tenant_id: str | None,
+    app_id: str | None,
+):
+    # Workflow IDs are UUID4, so this is not chronological pagination. The migration only needs a stable total
+    # order that matches the resume cursor; ordering by the same primary-key column used in the `id > cursor`
+    # predicate lets each batch continue deterministically without offset scans.
+    stmt = (
+        select(Workflow)
+        .options(load_only(Workflow.id, Workflow.type, Workflow.graph))
+        .where(Workflow.type.in_((WorkflowType.WORKFLOW, WorkflowType.CHAT)))
+        .order_by(Workflow.id)
+        .limit(batch_size)
+    )
+    if start_after_id:
+        stmt = stmt.where(Workflow.id > start_after_id)
+    if tenant_id:
+        stmt = stmt.where(Workflow.tenant_id == tenant_id)
+    if app_id:
+        stmt = stmt.where(Workflow.app_id == app_id)
+    return stmt
+
+
+def _migrate_legacy_sys_files_workflow_batch(
+    *,
+    session: Session,
+    start_after_id: str | None,
+    batch_size: int,
+    tenant_id: str | None,
+    app_id: str | None,
+    dry_run: bool,
+) -> LegacySysFilesWorkflowMigrationStats:
+    stats = LegacySysFilesWorkflowMigrationStats()
+    workflows = session.scalars(
+        _build_legacy_sys_files_workflow_query(
+            start_after_id=start_after_id,
+            batch_size=batch_size,
+            tenant_id=tenant_id,
+            app_id=app_id,
+        )
+    ).all()
+
+    for workflow in workflows:
+        stats.scanned += 1
+        stats.last_id = workflow.id
+        try:
+            if workflow.migrate_legacy_sys_files_graph_in_place():
+                stats.migrated += 1
+        except Exception:
+            stats.failed += 1
+            logger.exception("Failed to migrate legacy sys.files workflow, workflow_id=%s", workflow.id)
+
+    if dry_run:
+        session.rollback()
+    else:
+        session.commit()
+    return stats
+
+
+def run_legacy_sys_files_workflow_migration(
+    *,
+    batch_size: int,
+    limit: int | None,
+    start_after_id: str | None,
+    tenant_id: str | None,
+    app_id: str | None,
+    dry_run: bool,
+) -> LegacySysFilesWorkflowMigrationStats:
+    """Scan Workflow and Advanced Chat graphs in keyset-paginated batches."""
+    if batch_size <= 0:
+        raise click.UsageError("--batch-size must be greater than 0")
+    if limit is not None and limit <= 0:
+        raise click.UsageError("--limit must be greater than 0 when provided")
+
+    session_maker = sessionmaker(db.engine, expire_on_commit=False)
+    total = LegacySysFilesWorkflowMigrationStats(last_id=start_after_id)
+    next_start_after_id = start_after_id
+
+    while limit is None or total.scanned < limit:
+        remaining = None if limit is None else limit - total.scanned
+        current_batch_size = batch_size if remaining is None else min(batch_size, remaining)
+        if current_batch_size <= 0:
+            break
+
+        with session_maker() as session:
+            batch_stats = _migrate_legacy_sys_files_workflow_batch(
+                session=session,
+                start_after_id=next_start_after_id,
+                batch_size=current_batch_size,
+                tenant_id=tenant_id,
+                app_id=app_id,
+                dry_run=dry_run,
+            )
+
+        if batch_stats.scanned == 0:
+            break
+
+        total.scanned += batch_stats.scanned
+        total.migrated += batch_stats.migrated
+        total.failed += batch_stats.failed
+        total.batches += 1
+        total.last_id = batch_stats.last_id
+        next_start_after_id = batch_stats.last_id
+
+        if batch_stats.scanned < current_batch_size:
+            break
+
+    return total
+
+
+@click.command(
+    "migrate-legacy-sys-files-workflows",
+    help="Migrate Workflow and Advanced Chat graphs that still reference deprecated sys.files.",
+)
+@click.option("--batch-size", default=1000, show_default=True, type=int, help="Number of workflows to scan per batch.")
+@click.option("--limit", default=None, type=int, help="Maximum number of workflows to scan in this run.")
+@click.option("--start-after-id", default=None, help="Resume scanning after this workflow ID.")
+@click.option("--tenant-id", default=None, help="Limit migration to one tenant.")
+@click.option("--app-id", default=None, help="Limit migration to one app.")
+@click.option("--dry-run", is_flag=True, default=False, help="Scan and report without saving changes.")
+def migrate_legacy_sys_files_workflows(
+    batch_size: int,
+    limit: int | None,
+    start_after_id: str | None,
+    tenant_id: str | None,
+    app_id: str | None,
+    dry_run: bool,
+) -> None:
+    stats = run_legacy_sys_files_workflow_migration(
+        batch_size=batch_size,
+        limit=limit,
+        start_after_id=start_after_id,
+        tenant_id=tenant_id,
+        app_id=app_id,
+        dry_run=dry_run,
+    )
+    click.echo(
+        "Legacy sys.files workflow migration finished: "
+        f"scanned={stats.scanned} migrated={stats.migrated} failed={stats.failed} "
+        f"batches={stats.batches} last_id={stats.last_id or ''}"
+    )
+    if dry_run:
+        click.echo("Dry run only: no workflow graph changes were saved.")

api/commands/workspace.py

@@ -0,0 +1,64 @@
+"""Workspace maintenance CLI commands."""
+
+import click
+from sqlalchemy import delete, select
+from sqlalchemy.orm import sessionmaker
+
+from configs import dify_config
+from enums.deployment_edition import DeploymentEdition
+from extensions.ext_database import db
+from libs.rsa import generate_key_pair
+from models import Tenant
+from models.provider import Provider, ProviderModel
+from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider
+
+
+@click.command(
+    "reset-encrypt-key-pair",
+    help="Reset the asymmetric key pair of workspace for encrypt LLM credentials. "
+    "After the reset, all LLM credentials and tool provider credentials "
+    "(builtin / API / MCP) will be purged, requiring re-entry. "
+    "Only support SELF_HOSTED mode.",
+)
+@click.confirmation_option(
+    prompt=click.style(
+        "Are you sure you want to reset encrypt key pair? "
+        "This will also purge builtin / API / MCP tool provider records for every tenant. "
+        "This operation cannot be rolled back!",
+        fg="red",
+    )
+)
+def reset_encrypt_key_pair() -> None:
+    """
+    Reset the encrypted key pair of workspace for encrypt LLM credentials.
+    After the reset, all LLM credentials will become invalid, requiring re-entry.
+    Only support SELF_HOSTED mode.
+    """
+    if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD:
+        click.echo(click.style("This command is only for SELF_HOSTED installations.", fg="red"))
+        return
+    with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
+        tenants = session.scalars(select(Tenant)).all()
+        for tenant in tenants:
+            if not tenant:
+                click.echo(click.style("No workspaces found. Run /install first.", fg="red"))
+                return
+
+            tenant.encrypt_public_key = generate_key_pair(tenant.id)
+
+            session.execute(delete(Provider).where(Provider.provider_type == "custom", Provider.tenant_id == tenant.id))
+            session.execute(delete(ProviderModel).where(ProviderModel.tenant_id == tenant.id))
+
+            # Purge tool provider records that hold credentials encrypted under the
+            # tenant key. Leaving them in place causes /console/api/workspaces/current/
+            # tool-providers to 500 because decryption fails on stale ciphertext (#35396).
+            session.execute(delete(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant.id))
+            session.execute(delete(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant.id))
+            session.execute(delete(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant.id))
+
+            click.echo(
+                click.style(
+                    f"Congratulations! The asymmetric key pair of workspace {tenant.id} has been reset.",
+                    fg="green",
+                )
+            )

api/configs/extra/agent_backend_config.py

@@ -12,6 +12,11 @@ class AgentBackendConfig(BaseSettings):
         default=None,
     )
 
+    AGENT_BACKEND_API_TOKEN: str | None = Field(
+        description="Bearer token for authenticating with the Agent backend /runs API.",
+        default=None,
+    )
+
     AGENT_BACKEND_USE_FAKE: bool = Field(
         description="Use the deterministic in-process fake Agent backend client.",
         default=False,

api/configs/feature/__init__.py

@@ -266,6 +266,12 @@ class PluginConfig(BaseSettings):
         default=60 * 60,
     )
 
+    PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED: bool = Field(
+        description="Whether tenant plugin model providers are cached in Redis. Disable when plugins are installed "
+        "by a system other than this one, which cannot invalidate the cache when a tenant's plugins change.",
+        default=True,
+    )
+
     PLUGIN_MODEL_PROVIDERS_CACHE_TTL: PositiveInt = Field(
         description="TTL in seconds for caching tenant plugin model providers in Redis",
         default=60 * 60 * 24,

api/controllers/console/app/app.py

@@ -58,6 +58,7 @@ from services.app_service import (
     AppResponseView,
     AppService,
     CreateAppParams,
+    RecentAppMode,
     StarredAppListParams,
 )
 from services.enterprise import rbac_service as enterprise_rbac_service
@@ -139,6 +140,10 @@ class AppListBaseQuery(BaseModel):
             raise ValueError("Invalid UUID format in creator_ids.") from exc
 
 
+class RecentAppListQuery(BaseModel):
+    limit: int = Field(default=8, ge=1, le=8, description="Number of recently modified apps to return (1-8)")
+
+
 class AppListQuery(AppListBaseQuery):
     pass
 
@@ -411,6 +416,33 @@ class AppPartial(AppResponseModel):
         return to_timestamp(value)
 
 
+class RecentAppResponse(ResponseModel):
+    id: str
+    name: str
+    icon_type: IconType | None = None
+    icon: str | None = None
+    icon_background: str | None = None
+    mode: RecentAppMode
+    author_name: str | None = None
+    updated_at: int
+    permission_keys: list[str] = Field(default_factory=list)
+    maintainer: str | None = None
+
+    @computed_field(return_type=str | None)  # type: ignore[prop-decorator]
+    @property
+    def icon_url(self) -> str | None:
+        return build_icon_url(self.icon_type, self.icon)
+
+    @field_validator("updated_at", mode="before")
+    @classmethod
+    def _normalize_timestamp(cls, value: datetime | int) -> int:
+        return to_timestamp(value)
+
+
+class RecentAppListResponse(ResponseModel):
+    data: list[RecentAppResponse]
+
+
 class AppDetail(AppResponseModel):
     id: str
     name: str
@@ -575,6 +607,8 @@ register_schema_models(
 register_response_schema_models(
     console_ns,
     AppPartial,
+    RecentAppResponse,
+    RecentAppListResponse,
     AppDetailWithSite,
     AppPagination,
 )
@@ -699,6 +733,49 @@ class AppListApi(Resource):
         return app_detail.model_dump(mode="json"), 201
 
 
+@console_ns.route("/apps/recent")
+class RecentAppListApi(Resource):
+    @console_ns.doc("list_recent_apps")
+    @console_ns.doc(description="Get recently modified apps for the home Continue Work section")
+    @console_ns.doc(params=query_params_from_model(RecentAppListQuery))
+    @console_ns.response(200, "Success", console_ns.models[RecentAppListResponse.__name__])
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    @with_session(write=False)
+    @with_current_user_id
+    @with_current_tenant_id
+    def get(self, current_tenant_id: str, current_user_id: str, session: Session):
+        """Return the lightweight app cards needed by the Explore home page."""
+        args = query_params_from_request(RecentAppListQuery)
+        params = AppListParams(limit=args.limit)
+
+        permissions = enterprise_rbac_service.RBACService.MyPermissions.get(
+            current_tenant_id,
+            current_user_id,
+            session=session,
+        )
+        if dify_config.RBAC_ENABLED:
+            access_filter = resolve_app_access_filter(
+                current_tenant_id,
+                current_user_id,
+                session=session,
+                permissions=permissions,
+            )
+            access_filter.apply_to_params(params)
+
+        recent_apps = AppService().get_recent_apps(current_user_id, current_tenant_id, params, session)
+        permission_keys_map = permissions.app.permission_keys_by_resource_ids([app.id for app in recent_apps])
+        response_items = [
+            RecentAppResponse.model_validate(app, from_attributes=True).model_copy(
+                update={"permission_keys": permission_keys_map.get(app.id, [])}
+            )
+            for app in recent_apps
+        ]
+        return dump_response(RecentAppListResponse, {"data": response_items}), 200
+
+
 @console_ns.route("/apps/starred")
 class StarredAppListApi(Resource):
     @console_ns.doc("list_starred_apps")

api/controllers/inner_api/__init__.py

@@ -23,7 +23,6 @@ from .knowledge import retrieval as _knowledge_retrieval
 from .plugin import agent_config as _agent_config
 from .plugin import agent_drive as _agent_drive
 from .plugin import plugin as _plugin
-from .workspace import plugin_model_providers as _plugin_model_providers
 from .workspace import workspace as _workspace
 
 api.add_namespace(inner_api_ns)
@@ -36,7 +35,6 @@ __all__ = [
     "_knowledge_retrieval",
     "_mail",
     "_plugin",
-    "_plugin_model_providers",
     "_runtime_credentials",
     "_workspace",
     "api",

api/controllers/inner_api/workspace/plugin_model_providers.py

@@ -1,39 +0,0 @@
-from flask_restx import Resource
-from pydantic import BaseModel, ConfigDict, Field
-
-from controllers.common.schema import register_schema_model
-from controllers.console.wraps import setup_required
-from controllers.inner_api import inner_api_ns
-from controllers.inner_api.wraps import enterprise_inner_api_only
-from core.plugin.plugin_service import PluginService
-
-
-class InvalidatePluginModelProvidersCachePayload(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    tenant_ids: list[str] = Field(default_factory=list, description="Workspace ids whose cache should be invalidated")
-
-
-register_schema_model(inner_api_ns, InvalidatePluginModelProvidersCachePayload)
-
-
-@inner_api_ns.route("/enterprise/workspace/plugin-model-providers/invalidate")
-class EnterprisePluginModelProvidersCacheInvalidate(Resource):
-    @setup_required
-    @enterprise_inner_api_only
-    @inner_api_ns.doc(
-        "enterprise_invalidate_plugin_model_providers_cache",
-        responses={
-            200: "Cache invalidated",
-            400: "Invalid request",
-            401: "Unauthorized - invalid API key",
-        },
-    )
-    @inner_api_ns.expect(inner_api_ns.models[InvalidatePluginModelProvidersCachePayload.__name__])
-    def post(self):
-        args = InvalidatePluginModelProvidersCachePayload.model_validate(inner_api_ns.payload or {})
-
-        for tenant_id in args.tenant_ids:
-            PluginService.invalidate_plugin_model_providers_cache(tenant_id)
-
-        return {"result": "success"}, 200

api/controllers/service_api/app/completion.py

@@ -26,6 +26,10 @@ from controllers.service_api.app.error import (
     ProviderQuotaExceededError,
     WorkflowVersionExecutionNotAllowedError,
 )
+from controllers.service_api.app.legacy_system_files import (
+    attach_legacy_system_file_warning_for_service_api,
+    normalize_legacy_system_file_args_for_service_api,
+)
 from controllers.service_api.schema import (
     InputFileList,
     expect_user_json,
@@ -390,6 +394,7 @@ class ChatApi(Resource):
             args["external_trace_id"] = external_trace_id
 
         streaming = _resolve_agent_app_streaming(app_mode=app_mode, response_mode=payload.response_mode)
+        legacy_system_file_compat = None
 
         try:
             # Eagerly validate conversation to avoid hanging on invalid conversation_id
@@ -401,6 +406,14 @@ class ChatApi(Resource):
                     session=session,
                 )
 
+            if app_mode == AppMode.ADVANCED_CHAT:
+                args, legacy_system_file_compat = normalize_legacy_system_file_args_for_service_api(
+                    session=session,
+                    app_model=app_model,
+                    args=args,
+                    raw_payload=service_api_ns.payload,
+                    workflow_id=args.get("workflow_id"),
+                )
             response = AppGenerateService.generate(
                 session=session,
                 app_model=app_model,
@@ -409,6 +422,7 @@ class ChatApi(Resource):
                 invoke_from=InvokeFrom.SERVICE_API,
                 streaming=streaming,
             )
+            response = attach_legacy_system_file_warning_for_service_api(response, legacy_system_file_compat)
 
             # response-contract:ignore compact_generate_response
             return helper.compact_generate_response(response)

api/controllers/service_api/app/legacy_system_files.py

@@ -0,0 +1,72 @@
+"""Temporary Service API adapter for the deprecated workflow file input."""
+
+from collections.abc import Generator, Mapping
+from typing import Any
+
+from sqlalchemy.orm import Session
+
+from core.app.entities.app_invoke_entities import InvokeFrom
+from core.app.features.rate_limiting.rate_limit import RateLimitGenerator
+from core.workflow.legacy_system_files import (
+    LegacySysFilesCompatVariable,
+    attach_legacy_sys_files_warning,
+    normalize_legacy_sys_files_args,
+)
+from models.model import App
+from services.app_generate_service import AppGenerateService
+
+type ServiceAPIGenerateResponse = Mapping[str, Any] | Generator[str, None, None] | RateLimitGenerator
+
+
+def normalize_legacy_system_file_args_for_service_api(
+    *,
+    session: Session,
+    app_model: App,
+    args: dict[str, Any],
+    raw_payload: Mapping[str, Any] | None,
+    workflow_id: str | None = None,
+) -> tuple[dict[str, Any], LegacySysFilesCompatVariable | None]:
+    # TODO: Remove this hidden Service API compatibility path after all persisted workflows are migrated.
+    args_with_hidden_system = _copy_hidden_system_files_arg(args=args, raw_payload=raw_payload)
+    if not _has_legacy_file_arg(args_with_hidden_system):
+        return args, None
+
+    workflow = AppGenerateService.get_workflow(
+        app_model,
+        InvokeFrom.SERVICE_API,
+        workflow_id,
+        session=session,
+    )
+    return normalize_legacy_sys_files_args(graph=workflow.graph_dict, args=args_with_hidden_system)
+
+
+def attach_legacy_system_file_warning_for_service_api(
+    response: ServiceAPIGenerateResponse,
+    compat_variable: LegacySysFilesCompatVariable | None,
+) -> ServiceAPIGenerateResponse:
+    # TODO: Remove this warning once Service API clients no longer need the legacy migration notice.
+    if compat_variable is None:
+        return response
+    return attach_legacy_sys_files_warning(response, compat_variable)
+
+
+def _copy_hidden_system_files_arg(
+    *,
+    args: dict[str, Any],
+    raw_payload: Mapping[str, Any] | None,
+) -> dict[str, Any]:
+    system = raw_payload.get("system") if isinstance(raw_payload, Mapping) else None
+    if not isinstance(system, Mapping) or "files" not in system or system["files"] is None:
+        return args
+
+    copied_args = dict(args)
+    copied_args["system"] = {"files": system["files"]}
+    return copied_args
+
+
+def _has_legacy_file_arg(args: Mapping[str, Any]) -> bool:
+    if args.get("files") is not None:
+        return True
+
+    system = args.get("system")
+    return isinstance(system, Mapping) and system.get("files") is not None

api/controllers/service_api/app/workflow.py

@@ -30,6 +30,10 @@ from controllers.service_api.app.error import (
     ProviderQuotaExceededError,
     WorkflowVersionExecutionNotAllowedError,
 )
+from controllers.service_api.app.legacy_system_files import (
+    attach_legacy_system_file_warning_for_service_api,
+    normalize_legacy_system_file_args_for_service_api,
+)
 from controllers.service_api.schema import (
     expect_user_json,
     expect_with_user,
@@ -344,6 +348,12 @@ class WorkflowRunApi(Resource):
         streaming = payload.response_mode == "streaming"
 
         try:
+            args, legacy_system_file_compat = normalize_legacy_system_file_args_for_service_api(
+                session=session,
+                app_model=app_model,
+                args=args,
+                raw_payload=service_api_ns.payload,
+            )
             response = AppGenerateService.generate(
                 session=session,
                 app_model=app_model,
@@ -352,6 +362,7 @@ class WorkflowRunApi(Resource):
                 invoke_from=InvokeFrom.SERVICE_API,
                 streaming=streaming,
             )
+            response = attach_legacy_system_file_warning_for_service_api(response, legacy_system_file_compat)
 
             # response-contract:ignore compact_generate_response
             return helper.compact_generate_response(response)
@@ -471,6 +482,13 @@ class WorkflowRunByIdApi(Resource):
         streaming = payload.response_mode == "streaming"
 
         try:
+            args, legacy_system_file_compat = normalize_legacy_system_file_args_for_service_api(
+                session=session,
+                app_model=app_model,
+                args=args,
+                raw_payload=service_api_ns.payload,
+                workflow_id=workflow_id,
+            )
             response = AppGenerateService.generate(
                 session=session,
                 app_model=app_model,
@@ -479,6 +497,7 @@ class WorkflowRunByIdApi(Resource):
                 invoke_from=InvokeFrom.SERVICE_API,
                 streaming=streaming,
             )
+            response = attach_legacy_system_file_warning_for_service_api(response, legacy_system_file_compat)
 
             # response-contract:ignore compact_generate_response
             return helper.compact_generate_response(response)

api/core/app/apps/advanced_chat/app_generator.py

@@ -46,6 +46,7 @@ from core.ops.ops_trace_manager import TraceQueueManager
 from core.prompt.utils.get_thread_messages_length import get_thread_messages_length
 from core.repositories import DifyCoreRepositoryFactory
 from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository
+from core.workflow.legacy_system_files import normalize_legacy_sys_files_args
 from extensions.ext_database import db
 from factories import file_factory
 from graphon.filters import ResponseStreamFilter
@@ -147,6 +148,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         if not args.get("query"):
             raise ValueError("query is required")
 
+        # TODO: Remove this compatibility normalization after all persisted workflows are migrated.
+        args, _ = normalize_legacy_sys_files_args(graph=workflow.graph_dict, args=args)
         query = args["query"]
         if not isinstance(query, str):
             raise ValueError("query must be a string")
@@ -616,23 +619,34 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
             message_snapshot = MessageSnapshot.from_message(message)
             session.close()
 
-            # return response or stream generator
-            response = self._handle_advanced_chat_response(
-                application_generate_entity=application_generate_entity,
-                workflow=workflow_snapshot,
-                queue_manager=queue_manager,
-                conversation=conversation_snapshot,
-                message=message_snapshot,
-                user=user,
-                stream=stream,
-                draft_var_saver_factory=self._get_draft_var_saver_factory(
-                    invoke_from,
-                    account=user,
-                    tenant_id=application_generate_entity.app_config.tenant_id,
-                ),
-            )
+            try:
+                response = self._handle_advanced_chat_response(
+                    application_generate_entity=application_generate_entity,
+                    workflow=workflow_snapshot,
+                    queue_manager=queue_manager,
+                    conversation=conversation_snapshot,
+                    message=message_snapshot,
+                    user=user,
+                    stream=stream,
+                    draft_var_saver_factory=self._get_draft_var_saver_factory(
+                        invoke_from,
+                        account=user,
+                        tenant_id=application_generate_entity.app_config.tenant_id,
+                    ),
+                )
+                converted_response = AdvancedChatAppGenerateResponseConverter.convert(
+                    response=response,
+                    invoke_from=invoke_from,
+                )
+            except BaseException:
+                self._join_worker_thread(worker_thread)
+                raise
 
-            return AdvancedChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
+            if isinstance(converted_response, Generator):
+                return self._wrap_stream_with_worker_thread_join(converted_response, worker_thread)
+
+            self._join_worker_thread(worker_thread)
+            return converted_response
 
     def _generate_worker(
         self,

api/core/app/apps/agent_app/app_generator.py

@@ -538,6 +538,7 @@ class AgentAppGenerator(MessageBasedAppGenerator):
             request_builder=AgentAppRuntimeRequestBuilder(credentials_provider=credentials_provider),
             agent_backend_client=create_agent_backend_run_client(
                 base_url=dify_config.AGENT_BACKEND_BASE_URL,
+                api_token=dify_config.AGENT_BACKEND_API_TOKEN,
                 use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
                 fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
                 stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS,

api/core/app/apps/base_app_generator.py

@@ -1,3 +1,5 @@
+import logging
+import threading
 from collections.abc import Generator, Mapping, Sequence
 from contextlib import AbstractContextManager, nullcontext
 from typing import TYPE_CHECKING, Any, Union, final
@@ -23,6 +25,10 @@ from services.workflow_draft_variable_service import DraftVariableSaver as Draft
 if TYPE_CHECKING:
     from graphon.variables.input_entities import VariableEntity
 
+logger = logging.getLogger(__name__)
+
+_WORKER_THREAD_JOIN_TIMEOUT_SECONDS = 300
+
 
 @final
 class _DebuggerDraftVariableSaver:
@@ -64,6 +70,29 @@ class _DebuggerDraftVariableSaver:
 class BaseAppGenerator:
     _file_access_controller: DatabaseFileAccessController = DatabaseFileAccessController()
 
+    @staticmethod
+    def _join_worker_thread(worker_thread: threading.Thread) -> None:
+        # Bound the wait so a leaked app worker cannot occupy an execution slot indefinitely.
+        worker_thread.join(timeout=_WORKER_THREAD_JOIN_TIMEOUT_SECONDS)
+        if worker_thread.is_alive():
+            logger.warning(
+                "Possible app worker thread leak: thread_name=%s timeout_seconds=%s; "
+                "continuing without waiting further to avoid occupying an execution slot indefinitely",
+                worker_thread.name,
+                _WORKER_THREAD_JOIN_TIMEOUT_SECONDS,
+            )
+
+    @staticmethod
+    def _wrap_stream_with_worker_thread_join[ResponseT](
+        response_stream: Generator[ResponseT, None, None],
+        worker_thread: threading.Thread,
+    ) -> Generator[ResponseT, None, None]:
+        """Keep the producer owned by the response stream until both finish."""
+        try:
+            yield from response_stream
+        finally:
+            BaseAppGenerator._join_worker_thread(worker_thread)
+
     @staticmethod
     def _bind_file_access_scope(
         *,

api/core/app/apps/common/workflow_response_converter.py

@@ -155,8 +155,10 @@ class WorkflowResponseConverter:
             # TODO(@future-refactor): store system variables separately from user inputs so we don't
             # need to flatten `sys.*` entries into the input payload just for rerun/export tooling.
             if field_name == SystemVariableKey.CONVERSATION_ID:
-                # Conversation IDs are session-scoped; omitting them keeps workflow inputs
-                # reusable without pinning new runs to a prior conversation.
+                # Conversation IDs are session-scoped; omitting them keeps workflow inputs reusable.
+                continue
+            if field_name == SystemVariableKey.FILES:
+                # When files are exposed as an input, application inputs use the canonical `userinput.files` key.
                 continue
             inputs[f"sys.{field_name}"] = value
         handled = WorkflowEntry.handle_special_values(inputs)

api/core/app/apps/pipeline/pipeline_generator.py

@@ -351,17 +351,28 @@ class PipelineGenerator(BaseAppGenerator):
                 user,
                 tenant_id=pipeline.tenant_id,
             )
-            # return response or stream generator
-            response = self._handle_response(
-                application_generate_entity=application_generate_entity,
-                workflow=workflow,
-                queue_manager=queue_manager,
-                user=user,
-                stream=streaming,
-                draft_var_saver_factory=draft_var_saver_factory,
-            )
+            try:
+                response = self._handle_response(
+                    application_generate_entity=application_generate_entity,
+                    workflow=workflow,
+                    queue_manager=queue_manager,
+                    user=user,
+                    stream=streaming,
+                    draft_var_saver_factory=draft_var_saver_factory,
+                )
+                converted_response = WorkflowAppGenerateResponseConverter.convert(
+                    response=response,
+                    invoke_from=invoke_from,
+                )
+            except BaseException:
+                self._join_worker_thread(worker_thread)
+                raise
 
-            return WorkflowAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
+            if isinstance(converted_response, Generator):
+                return self._wrap_stream_with_worker_thread_join(converted_response, worker_thread)
+
+            self._join_worker_thread(worker_thread)
+            return converted_response
 
     def single_iteration_generate(
         self,

api/core/app/apps/workflow/app_generator.py

@@ -41,6 +41,7 @@ from core.helper.trace_id_helper import (
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.repositories import DifyCoreRepositoryFactory
 from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository
+from core.workflow.legacy_system_files import normalize_legacy_sys_files_args
 from extensions.ext_database import db
 from factories import file_factory
 from graphon.filters import ResponseStreamFilter
@@ -164,6 +165,8 @@ class WorkflowAppGenerator(BaseAppGenerator):
         pause_state_config: PauseStateLayerConfig | None = None,
     ) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str, None, None]:
         with self._bind_file_access_scope(tenant_id=app_model.tenant_id, user=user, invoke_from=invoke_from):
+            # TODO: Remove this compatibility normalization after all persisted workflows are migrated.
+            args, _ = normalize_legacy_sys_files_args(graph=workflow.graph_dict, args=args)
             files: Sequence[Mapping[str, Any]] = args.get("files") or []
 
             # parse files
@@ -405,17 +408,28 @@ class WorkflowAppGenerator(BaseAppGenerator):
                 tenant_id=app_model.tenant_id,
             )
 
-            # return response or stream generator
-            response = self._handle_response(
-                application_generate_entity=application_generate_entity,
-                workflow=workflow,
-                queue_manager=queue_manager,
-                user=user,
-                draft_var_saver_factory=draft_var_saver_factory,
-                stream=streaming,
-            )
+            try:
+                response = self._handle_response(
+                    application_generate_entity=application_generate_entity,
+                    workflow=workflow,
+                    queue_manager=queue_manager,
+                    user=user,
+                    draft_var_saver_factory=draft_var_saver_factory,
+                    stream=streaming,
+                )
+                converted_response = WorkflowAppGenerateResponseConverter.convert(
+                    response=response,
+                    invoke_from=invoke_from,
+                )
+            except BaseException:
+                self._join_worker_thread(worker_thread)
+                raise
 
-            return WorkflowAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
+            if isinstance(converted_response, Generator):
+                return self._wrap_stream_with_worker_thread_join(converted_response, worker_thread)
+
+            self._join_worker_thread(worker_thread)
+            return converted_response
 
     def single_iteration_generate(
         self,

api/core/plugin/plugin_service.py

@@ -66,7 +66,7 @@ from services.enterprise.plugin_manager_service import (
     PreUninstallPluginRequest,
 )
 from services.errors.plugin import PluginInstallationForbiddenError
-from services.feature_service import FeatureService, PluginInstallationScope
+from services.feature_service import FeatureService, PluginInstallationPermissionModel, PluginInstallationScope
 
 logger = logging.getLogger(__name__)
 _provider_entities_adapter: TypeAdapter[list[ProviderEntity]] = TypeAdapter(list[ProviderEntity])
@@ -434,14 +434,18 @@ class PluginService:
                     exc_info=True,
                 )
 
+    @classmethod
+    def _fetch_plugin_model_providers_uncached(
+        cls, tenant_id: str, client: PluginModelClient | None
+    ) -> tuple[ProviderEntity, ...]:
+        model_client = client or PluginModelClient()
+        return tuple(cls._to_provider_entity(provider) for provider in model_client.fetch_model_providers(tenant_id))
+
     @classmethod
     def _fetch_and_cache_plugin_model_providers(
         cls, tenant_id: str, client: PluginModelClient | None, *, refresh_generation: int | None
     ) -> tuple[ProviderEntity, ...]:
-        model_client = client or PluginModelClient()
-        providers = tuple(
-            cls._to_provider_entity(provider) for provider in model_client.fetch_model_providers(tenant_id)
-        )
+        providers = cls._fetch_plugin_model_providers_uncached(tenant_id, client)
         generation = cls._load_plugin_model_providers_generation(tenant_id)
         if generation is not None and generation == refresh_generation:
             cls._store_cached_plugin_model_providers(tenant_id, generation, providers)
@@ -471,6 +475,9 @@ class PluginService:
         are intentionally owned by this service so tenant isolation and cache
         expiry are handled in one place.
         """
+        if not dify_config.PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED:
+            return cls._fetch_plugin_model_providers_uncached(tenant_id, client)
+
         deadline = time.monotonic() + cls.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT
 
         while True:
@@ -597,22 +604,30 @@ class PluginService:
             return result
 
     @staticmethod
-    def _check_marketplace_only_permission():
+    def _check_marketplace_only_permission() -> None:
         """
         Check if the marketplace only permission is enabled
         """
-        features = FeatureService.get_system_features()
-        if features.plugin_installation_permission.restrict_to_marketplace_only:
+        permission = PluginService._get_plugin_installation_permission()
+        if permission.restrict_to_marketplace_only:
             raise PluginInstallationForbiddenError("Plugin installation is restricted to marketplace only")
 
     @staticmethod
-    def _check_plugin_installation_scope(plugin_verification: PluginVerification | None):
+    def _get_plugin_installation_permission() -> PluginInstallationPermissionModel:
+        """Resolve the validated policy and reject deny-all before any installation side effect."""
+        permission = FeatureService.get_plugin_installation_permission()
+        if permission.plugin_installation_scope == PluginInstallationScope.NONE:
+            raise PluginInstallationForbiddenError("Installing plugins is not allowed")
+        return permission
+
+    @staticmethod
+    def _check_plugin_installation_scope(plugin_verification: PluginVerification | None) -> None:
         """
         Check the plugin installation scope
         """
-        features = FeatureService.get_system_features()
+        permission = PluginService._get_plugin_installation_permission()
 
-        match features.plugin_installation_permission.plugin_installation_scope:
+        match permission.plugin_installation_scope:
             case PluginInstallationScope.OFFICIAL_ONLY:
                 if (
                     plugin_verification is None
@@ -627,10 +642,10 @@ class PluginService:
                     raise PluginInstallationForbiddenError(
                         "Plugin installation is restricted to official and specific partners"
                     )
-            case PluginInstallationScope.NONE:
-                raise PluginInstallationForbiddenError("Installing plugins is not allowed")
             case PluginInstallationScope.ALL:
                 pass
+            case _:
+                raise PluginInstallationForbiddenError("Plugin installation policy is invalid")
 
     @staticmethod
     def get_debugging_key(tenant_id: str) -> str:
@@ -900,7 +915,7 @@ class PluginService:
         # check if plugin pkg is already downloaded
         manager = PluginInstaller()
 
-        features = FeatureService.get_system_features()
+        permission = PluginService._get_plugin_installation_permission()
 
         try:
             manager.fetch_plugin_manifest(tenant_id, new_plugin_unique_identifier)
@@ -912,7 +927,7 @@ class PluginService:
             response = manager.upload_pkg(
                 tenant_id,
                 pkg,
-                verify_signature=features.plugin_installation_permission.restrict_to_marketplace_only,
+                verify_signature=permission.restrict_to_marketplace_only,
             )
 
             # check if the plugin is available to install
@@ -967,11 +982,11 @@ class PluginService:
         """
         PluginService._check_marketplace_only_permission()
         manager = PluginInstaller()
-        features = FeatureService.get_system_features()
+        permission = PluginService._get_plugin_installation_permission()
         response = manager.upload_pkg(
             tenant_id,
             pkg,
-            verify_signature=features.plugin_installation_permission.restrict_to_marketplace_only,
+            verify_signature=permission.restrict_to_marketplace_only,
         )
         PluginService._check_plugin_installation_scope(response.verification)
 
@@ -989,13 +1004,13 @@ class PluginService:
         pkg = download_with_size_limit(
             f"https://github.com/{repo}/releases/download/{version}/{package}", dify_config.PLUGIN_MAX_PACKAGE_SIZE
         )
-        features = FeatureService.get_system_features()
+        permission = PluginService._get_plugin_installation_permission()
 
         manager = PluginInstaller()
         response = manager.upload_pkg(
             tenant_id,
             pkg,
-            verify_signature=features.plugin_installation_permission.restrict_to_marketplace_only,
+            verify_signature=permission.restrict_to_marketplace_only,
         )
         PluginService._check_plugin_installation_scope(response.verification)
 
@@ -1069,7 +1084,7 @@ class PluginService:
         if not dify_config.MARKETPLACE_ENABLED:
             raise ValueError("marketplace is not enabled")
 
-        features = FeatureService.get_system_features()
+        permission = PluginService._get_plugin_installation_permission()
 
         manager = PluginInstaller()
         try:
@@ -1079,7 +1094,7 @@ class PluginService:
             response = manager.upload_pkg(
                 tenant_id,
                 pkg,
-                verify_signature=features.plugin_installation_permission.restrict_to_marketplace_only,
+                verify_signature=permission.restrict_to_marketplace_only,
             )
             # check if the plugin is available to install
             PluginService._check_plugin_installation_scope(response.verification)
@@ -1101,7 +1116,7 @@ class PluginService:
         # collect actual plugin_unique_identifiers
         actual_plugin_unique_identifiers = []
         metas = []
-        features = FeatureService.get_system_features()
+        permission = PluginService._get_plugin_installation_permission()
 
         # check if already downloaded
         for plugin_unique_identifier in plugin_unique_identifiers:
@@ -1119,7 +1134,7 @@ class PluginService:
                 response = manager.upload_pkg(
                     tenant_id,
                     pkg,
-                    verify_signature=features.plugin_installation_permission.restrict_to_marketplace_only,
+                    verify_signature=permission.restrict_to_marketplace_only,
                 )
                 # check if the plugin is available to install
                 PluginService._check_plugin_installation_scope(response.verification)

api/core/rag/retrieval/dataset_retrieval.py

@@ -2023,6 +2023,8 @@ class DatasetRetrieval:
             redis_client.zremrangebyscore(key, 0, current_time - 60000)
             request_count = redis_client.zcard(key)
             if request_count > knowledge_rate_limit.limit:
+                # The rate-limit exception is raised after this block, so commit the audit row
+                # explicitly instead of relying on the Session context, which only closes it.
                 with session_factory.create_session() as session:
                     rate_limit_log = RateLimitLog(
                         tenant_id=tenant_id,
@@ -2030,6 +2032,7 @@ class DatasetRetrieval:
                         operation="knowledge",
                     )
                     session.add(rate_limit_log)
+                    session.commit()
                 raise exc.RateLimitExceededError(
                     "you have reached the knowledge base request rate limit of your subscription."
                 )

api/core/tools/mcp_tool/tool.py

@@ -107,6 +107,8 @@ class MCPTool(Tool):
         if self.entity.output_schema and result.structuredContent:
             for k, v in result.structuredContent.items():
                 yield self.create_variable_message(k, v)
+        elif result.structuredContent:
+            yield self.create_json_message(result.structuredContent)
 
     def _process_text_content(self, content: TextContent) -> Generator[ToolInvokeMessage, None, None]:
         """Process text content and yield appropriate messages."""

api/core/workflow/generator/prompts/builder_prompts.py

@@ -42,9 +42,9 @@ _NODE_SNIPPETS: dict[str, str] = {
     ["local_file", "remote_url"]. Only when you include "custom" must you
     also set ``allowed_file_extensions`` to a non-empty list like
     [".epub", ".rtf"]; otherwise leave it [].
-    In Advanced-Chat mode ``sys.query`` and ``sys.files`` are automatic
-    system variables — downstream nodes may reference them; do NOT add
-    them to ``variables``.""",
+    In Advanced-Chat mode ``sys.query`` is automatic. ``userinput.files`` is
+    the automatic file-upload variable in both app modes. Downstream nodes
+    may reference these variables; do NOT add them to ``variables``.""",
     "end": """\
 - end (Workflow mode only):
     {"outputs": [
@@ -168,8 +168,8 @@ _NODE_SNIPPETS: dict[str, str] = {
     Single output variable ``text``: a string when ``is_array_file`` is false,
     an array of strings (one per file) when it is true. ``variable_selector``
     MUST point at a ``start`` variable declared with type "file" / "file-list"
-    (or ``sys.files`` in Advanced-Chat mode). That start variable MUST set a
-    non-empty ``allowed_file_types`` (use ["document"] for document text).""",
+    (or the automatic ``userinput.files`` variable). A declared start variable
+    MUST set a non-empty ``allowed_file_types`` (use ["document"] for document text).""",
     "variable-aggregator": """\
 - variable-aggregator  (merge mutually-exclusive branches into one output):
     {"output_type": "string",        # VarType of the merged value — one of

api/core/workflow/generator/prompts/node_builder_prompts.py

@@ -91,21 +91,20 @@ def format_parallel_plan(
 def format_mode_section(mode: str) -> str:
     """Tell each builder which app mode it is configuring for.
 
-    Matters most in advanced-chat, where ``sys.query`` / ``sys.files`` are the
-    sanctioned way to reference the user's message — without this the model
-    invents start-node variables that postprocess then materializes as
-    spurious form inputs.
+    ``sys.query`` is available in advanced-chat, while ``userinput.files`` is
+    available in both app modes. Without this guidance the model invents
+    start-node variables that postprocess then materializes as spurious inputs.
     """
     if mode == "advanced-chat":
         return (
             "# App mode\n\n"
             "advanced-chat: the user's chat message is available as sys.query and uploaded files "
-            'as sys.files — placeholder {{#sys.query#}}, selector ["sys", "query"]. Reference them '
-            "directly; do NOT invent start-node variables for the chat message.\n\n"
+            "as userinput.files. Use placeholder {{#sys.query#}} or selector "
+            '["userinput", "files"] directly; do NOT invent start-node variables for them.\n\n'
         )
     return (
         "# App mode\n\n"
-        "workflow: there are NO automatic system variables; reference user input only through "
+        "workflow: uploaded files are available as userinput.files; all other user input must use "
         "the start node's declared variables.\n\n"
     )
 

api/core/workflow/generator/prompts/planner_prompts.py

@@ -96,10 +96,10 @@ minimum set of Dify workflow nodes needed to fulfil it, in execution order.
       - "text-input" for short single-line values (URLs, names),
       - "paragraph" for free-form multi-line text (descriptions, queries),
       - "number" / "select" / "file" / "file-list" for the obvious cases.
-    In Advanced-Chat mode the ``sys.query`` / ``sys.files`` system
-    variables are automatic — downstream nodes may reference them without
-    a ``start_inputs`` entry. In Workflow mode there is NO automatic
-    variable; everything the user supplies must be in ``start_inputs``.
+    In Advanced-Chat mode ``sys.query`` is automatic. ``userinput.files`` is
+    automatic in both app modes. Downstream nodes may reference these values
+    without a ``start_inputs`` entry; every other user-supplied Workflow value
+    must be declared in ``start_inputs``.
 11. Give every node a unique runtime-safe ``id`` using only letters, digits,
     and underscores. In create mode use ``node1``, ``node2``, ... in node-list
     order. In refine mode preserve the existing id for every retained node.

api/core/workflow/generator/runner.py

@@ -1382,8 +1382,8 @@ class WorkflowGenerator:
         multiple outputs remain untouched so validation fails closed instead
         of guessing which value the workflow should consume.
 
-        For Advanced-Chat mode, ``sys.query`` and ``sys.files`` are always
-        treated as resolved without any declaration. Tool nodes' parameter
+        ``sys.query`` in Advanced-Chat mode and ``userinput.files`` in either mode
+        are treated as resolved without declarations. Tool nodes' parameter
         references aren't validated here because we don't know each tool's
         schema — the run time validates those.
         """
@@ -1398,9 +1398,12 @@ class WorkflowGenerator:
         for node in nodes:
             cls._collect_refs_in_data(node.get("data") or {}, refs)
 
+        automatic_refs = {("userinput", "files")}
+        if mode == "advanced-chat":
+            automatic_refs.add(("sys", "query"))
+
         for node_id, var in refs:
-            # Advanced-Chat system variables are always resolved.
-            if mode == "advanced-chat" and node_id == "sys":
+            if (node_id, var) in automatic_refs:
                 continue
             target = nodes_by_id.get(node_id)
             if target is None:

api/core/workflow/legacy_system_files.py

@@ -0,0 +1,218 @@
+"""Compatibility helpers for workflows that still reference deprecated `sys.files`.
+
+TODO: Remove this module after all persisted Workflow and Advanced Chat graphs
+have been migrated from the deprecated system file variable to `userinput.files`.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Generator, Iterable, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+_LEGACY_SYSTEM_NODE_ID = "sys"
+_USER_INPUT_NODE_ID = "userinput"
+_LEGACY_FILES_VARIABLE = "files"
+_USER_INPUT_FILE_SELECTOR = [_USER_INPUT_NODE_ID, _LEGACY_FILES_VARIABLE]
+_USER_INPUT_FILE_INPUT_KEY = ".".join(_USER_INPUT_FILE_SELECTOR)
+_LEGACY_FILES_TEMPLATE = "{{#sys.files#}}"
+_USER_INPUT_FILES_TEMPLATE = "{{#userinput.files#}}"
+
+
+@dataclass(frozen=True)
+class LegacySysFilesCompatVariable:
+    node_id: str
+    variable_name: str
+
+
+@dataclass(frozen=True)
+class LegacySysFilesGraphMigrationResult:
+    graph: dict[str, Any]
+    changed: bool
+
+
+def migrate_legacy_sys_files_graph_with_result(
+    graph: Mapping[str, Any],
+) -> LegacySysFilesGraphMigrationResult:
+    """Return the migrated graph and whether any legacy reference was rewritten."""
+
+    graph_copy = dict(graph)
+    nodes = graph_copy.get("nodes")
+    if not isinstance(nodes, list):
+        return LegacySysFilesGraphMigrationResult(graph=graph_copy, changed=False)
+
+    # Legacy references are stored in node data. Restricting both search and replacement to `nodes`
+    # avoids recursively scanning graph-level metadata and edges for every workflow load.
+    if not _contains_legacy_sys_files_reference(nodes):
+        return LegacySysFilesGraphMigrationResult(graph=graph_copy, changed=False)
+
+    graph_copy["nodes"] = _replace_legacy_sys_files_references(nodes)
+    return LegacySysFilesGraphMigrationResult(graph=graph_copy, changed=True)
+
+
+def resolve_legacy_sys_files_compat_variable(graph: Mapping[str, Any]) -> LegacySysFilesCompatVariable | None:
+    """Resolve the target variable used by the `sys.files` compatibility layer."""
+
+    nodes = graph.get("nodes")
+    if not isinstance(nodes, list):
+        return None
+    if not _contains_file_input_reference(nodes):
+        return None
+    return LegacySysFilesCompatVariable(node_id=_USER_INPUT_NODE_ID, variable_name=_LEGACY_FILES_VARIABLE)
+
+
+def normalize_legacy_sys_files_args(
+    *,
+    graph: Mapping[str, Any],
+    args: Mapping[str, Any],
+) -> tuple[dict[str, Any], LegacySysFilesCompatVariable | None]:
+    """Map Service/Web API file arguments onto the `userinput.files` system alias.
+
+    The top-level `files` argument and hidden `system.files` payload both feed
+    the same runtime file collection. After graph references are migrated, the
+    file collection is exposed in the variable pool as `userinput.files`.
+    """
+
+    normalized_args = dict(args)
+    files_from_input, input_files_used = _extract_userinput_files(args)
+    if input_files_used:
+        normalized_args["files"] = files_from_input
+        return normalized_args, None
+
+    compat_variable = resolve_legacy_sys_files_compat_variable(graph)
+    if compat_variable is None:
+        return normalized_args, None
+
+    files, legacy_files_used = _extract_legacy_files(args)
+    if not legacy_files_used:
+        return normalized_args, None
+
+    if normalized_args.get("files") is None:
+        normalized_args["files"] = files
+
+    raw_inputs = normalized_args.get("inputs")
+    inputs = dict(raw_inputs) if isinstance(raw_inputs, Mapping) else {}
+    inputs.setdefault(_USER_INPUT_FILE_INPUT_KEY, files)
+    normalized_args["inputs"] = inputs
+    return normalized_args, compat_variable
+
+
+def attach_legacy_sys_files_warning(
+    response: Mapping[str, Any] | Iterable[Any],
+    compat_variable: LegacySysFilesCompatVariable,
+) -> Mapping[str, Any] | Generator[str, None, None]:
+    warning = build_legacy_sys_files_warning(compat_variable)
+    if isinstance(response, Mapping):
+        response_with_warning = dict(response)
+        existing_warnings = response_with_warning.get("warnings")
+        warnings = list(existing_warnings) if isinstance(existing_warnings, list) else []
+        warnings.append(warning)
+        response_with_warning["warnings"] = warnings
+        return response_with_warning
+
+    def _with_warning() -> Generator[str, None, None]:
+        try:
+            yield f"data: {json.dumps({'event': 'warning', 'warning': warning})}\n\n"
+            yield from response
+        finally:
+            close = getattr(response, "close", None)
+            if callable(close):
+                close()
+
+    return _with_warning()
+
+
+def build_legacy_sys_files_warning(compat_variable: LegacySysFilesCompatVariable) -> str:
+    variable_selector = ".".join((compat_variable.node_id, compat_variable.variable_name))
+    return (
+        "sys.files is deprecated. This workflow now reads files from "
+        f"`{variable_selector}`; update Service API calls to pass files in "
+        f"`inputs.{variable_selector}` instead of `system.files` or top-level `files`."
+    )
+
+
+def _contains_legacy_sys_files_reference(value: Any) -> bool:
+    if _is_legacy_sys_files_selector(value):
+        return True
+
+    if isinstance(value, str):
+        return _LEGACY_FILES_TEMPLATE in value
+
+    if isinstance(value, Mapping):
+        return any(_contains_legacy_sys_files_reference(item) for item in value.values())
+
+    if isinstance(value, list):
+        return any(_contains_legacy_sys_files_reference(item) for item in value)
+
+    return False
+
+
+def _contains_file_input_reference(value: Any) -> bool:
+    if _is_legacy_sys_files_selector(value) or _is_userinput_files_selector(value):
+        return True
+
+    if isinstance(value, str):
+        return _LEGACY_FILES_TEMPLATE in value or _USER_INPUT_FILES_TEMPLATE in value
+
+    if isinstance(value, Mapping):
+        return any(_contains_file_input_reference(item) for item in value.values())
+
+    if isinstance(value, list):
+        return any(_contains_file_input_reference(item) for item in value)
+
+    return False
+
+
+def _replace_legacy_sys_files_references(value: Any) -> Any:
+    if _is_legacy_sys_files_selector(value):
+        return list(_USER_INPUT_FILE_SELECTOR)
+
+    if isinstance(value, str):
+        return value.replace(_LEGACY_FILES_TEMPLATE, _USER_INPUT_FILES_TEMPLATE)
+
+    if isinstance(value, Mapping):
+        return {key: _replace_legacy_sys_files_references(item) for key, item in value.items()}
+
+    if isinstance(value, list):
+        return [_replace_legacy_sys_files_references(item) for item in value]
+
+    return value
+
+
+def _is_legacy_sys_files_selector(value: Any) -> bool:
+    return (
+        isinstance(value, list)
+        and len(value) == 2
+        and value[0] == _LEGACY_SYSTEM_NODE_ID
+        and value[1] == _LEGACY_FILES_VARIABLE
+    )
+
+
+def _is_userinput_files_selector(value: Any) -> bool:
+    return isinstance(value, list) and value == _USER_INPUT_FILE_SELECTOR
+
+
+def serialized_graph_may_contain_legacy_sys_files(serialized_graph: str) -> bool:
+    """Cheaply reject stored graphs that cannot contain a legacy file reference."""
+
+    return _LEGACY_FILES_TEMPLATE in serialized_graph or ('"sys"' in serialized_graph and '"files"' in serialized_graph)
+
+
+def _extract_legacy_files(args: Mapping[str, Any]) -> tuple[Any, bool]:
+    if "files" in args and args["files"] is not None:
+        return args["files"], True
+
+    system = args.get("system")
+    if isinstance(system, Mapping) and "files" in system and system["files"] is not None:
+        return system["files"], True
+
+    return None, False
+
+
+def _extract_userinput_files(args: Mapping[str, Any]) -> tuple[Any, bool]:
+    inputs = args.get("inputs")
+    if isinstance(inputs, Mapping) and inputs.get(_USER_INPUT_FILE_INPUT_KEY) is not None:
+        return inputs[_USER_INPUT_FILE_INPUT_KEY], True
+
+    return None, False

api/core/workflow/node_factory.py

@@ -497,6 +497,7 @@ class DifyNodeFactory(NodeFactory):
                 ),
                 "agent_backend_client": create_agent_backend_run_client(
                     base_url=dify_config.AGENT_BACKEND_BASE_URL,
+                    api_token=dify_config.AGENT_BACKEND_API_TOKEN,
                     use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
                     fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
                     stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS,

api/core/workflow/nodes/agent_v2/runtime_request_builder.py

@@ -368,7 +368,7 @@ class WorkflowAgentRuntimeRequestBuilder:
 
         if uploaded_files is not None:
             lines.append("- Uploaded workflow files:")
-            lines.append(f"  - sys.files: {uploaded_files}")
+            lines.append(f"  - userinput.files: {uploaded_files}")
 
         if resolved_outputs:
             lines.append("- Previous node outputs:")

api/core/workflow/system_variables.py

@@ -16,6 +16,7 @@ from .variable_prefixes import (
     ENVIRONMENT_VARIABLE_NODE_ID,
     RAG_PIPELINE_VARIABLE_NODE_ID,
     SYSTEM_VARIABLE_NODE_ID,
+    USER_INPUT_VARIABLE_NODE_ID,
 )
 
 
@@ -118,6 +119,12 @@ def build_bootstrap_variables(
         *(_with_selector(variable, ENVIRONMENT_VARIABLE_NODE_ID) for variable in environment_variables),
         *(_with_selector(variable, CONVERSATION_VARIABLE_NODE_ID) for variable in conversation_variables),
     ]
+    # TODO: Stop emitting the legacy `sys.files` selector after stored graphs and Service API callers are migrated.
+    # `userinput.files` remains the canonical file-upload variable.
+    for variable in system_variables:
+        if variable.name == SystemVariableKey.FILES.value:
+            variables.append(_with_selector(variable, USER_INPUT_VARIABLE_NODE_ID))
+            break
 
     rag_pipeline_variables_map: defaultdict[str, dict[str, Any]] = defaultdict(dict)
     for rag_var in rag_pipeline_variables:

api/core/workflow/variable_prefixes.py

@@ -1,4 +1,5 @@
 SYSTEM_VARIABLE_NODE_ID = "sys"
+USER_INPUT_VARIABLE_NODE_ID = "userinput"
 ENVIRONMENT_VARIABLE_NODE_ID = "env"
 CONVERSATION_VARIABLE_NODE_ID = "conversation"
 RAG_PIPELINE_VARIABLE_NODE_ID = "rag"

api/extensions/ext_commands.py

@@ -29,6 +29,7 @@ def init_app(app: DifyApp):
         install_rag_pipeline_plugins,
         migrate_data_for_plugin,
         migrate_dataset_permissions_to_rbac,
+        migrate_legacy_sys_files_workflows,
         migrate_member_roles_to_rbac,
         migrate_oss,
         migration_data_wizard,
@@ -57,6 +58,7 @@ def init_app(app: DifyApp):
         data_migrate,
         upgrade_db,
         fix_app_site_missing,
+        migrate_legacy_sys_files_workflows,
         migrate_data_for_plugin,
         migrate_dataset_permissions_to_rbac,
         migrate_member_roles_to_rbac,

api/models/model.py

@@ -1117,14 +1117,14 @@ class ExporleBanner(TypeBase):
     status: Mapped[BannerStatus] = mapped_column(
         EnumText(BannerStatus, length=255),
         nullable=False,
-        server_default=sa.text("'enabled'::character varying"),
+        server_default=sa.text("'enabled'"),
         default=BannerStatus.ENABLED,
     )
     created_at: Mapped[datetime] = mapped_column(
         sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
     )
     language: Mapped[str] = mapped_column(
-        String(255), nullable=False, server_default=sa.text("'en-US'::character varying"), default="en-US"
+        String(255), nullable=False, server_default=sa.text("'en-US'"), default="en-US"
     )
 
 

api/models/workflow.py

@@ -25,6 +25,10 @@ from typing_extensions import deprecated
 
 from core.trigger.constants import TRIGGER_PLUGIN_NODE_TYPE
 from core.workflow.human_input_adapter import adapt_node_config_for_graph
+from core.workflow.legacy_system_files import (
+    migrate_legacy_sys_files_graph_with_result,
+    serialized_graph_may_contain_legacy_sys_files,
+)
 from core.workflow.nodes.human_input.pause_reason import (
     HumanInputRequired,
 )
@@ -325,7 +329,39 @@ class Workflow(Base):  # bug
         # Currently, the following functions / methods would mutate the returned dict:
         #
         # - `_get_graph_and_variable_pool_for_single_node_run`.
-        return json.loads(self.graph) if self.graph else {}
+        if not self.graph:
+            return {}
+
+        graph = json.loads(self.graph)
+        if not self._supports_legacy_sys_files_compatibility() or not serialized_graph_may_contain_legacy_sys_files(
+            self.graph
+        ):
+            return graph
+
+        # TODO: Remove this load-time compatibility rewrite after all persisted workflows are migrated.
+        return migrate_legacy_sys_files_graph_with_result(graph).graph
+
+    def migrate_legacy_sys_files_graph_in_place(self) -> bool:
+        if (
+            not self.graph
+            or not self._supports_legacy_sys_files_compatibility()
+            or not serialized_graph_may_contain_legacy_sys_files(self.graph)
+        ):
+            return False
+
+        # TODO: Remove this in-place compatibility rewrite after all persisted workflows are migrated.
+        migration_result = migrate_legacy_sys_files_graph_with_result(json.loads(self.graph))
+        if migration_result.changed:
+            self.graph = json.dumps(migration_result.graph)
+        return migration_result.changed
+
+    def _supports_legacy_sys_files_compatibility(self) -> bool:
+        return self.type in {
+            WorkflowType.WORKFLOW,
+            WorkflowType.CHAT,
+            WorkflowType.WORKFLOW.value,
+            WorkflowType.CHAT.value,
+        }
 
     def get_node_config_by_id(self, node_id: str) -> NodeConfigDict:
         """Extract a node configuration from the workflow graph by node ID.
@@ -487,7 +523,7 @@ class Workflow(Base):  # bug
             "memory":
               {
                 "window": { "enabled": false, "size": 10 },
-                "query_prompt_template": "{{#sys.query#}}\n\n{{#sys.files#}}",
+                "query_prompt_template": "{{#sys.query#}}\n\n{{#userinput.files#}}",
                 "role_prefix": { "user": "", "assistant": "" },
               },
             "selected": false,
@@ -1520,7 +1556,7 @@ class ConversationVariable(TypeBase):
         return variable_factory.build_conversation_variable_from_mapping(mapping)
 
 
-# Only `sys.query` and `sys.files` could be modified.
+# TODO: Remove file-system-variable editability after all persisted workflows are migrated.
 _EDITABLE_SYSTEM_VARIABLE = frozenset(("query", "files"))
 
 

api/openapi/markdown/console-openapi.md

@@ -1672,6 +1672,23 @@ Create a new application
 | 200 | Import confirmed | **application/json**: [Import](#import)<br> |
 | 400 | Import failed | **application/json**: [Import](#import)<br> |
 
+### [GET] /apps/recent
+**Return the lightweight app cards needed by the Explore home page**
+
+Get recently modified apps for the home Continue Work section
+
+#### Parameters
+
+| Name | Located in | Description | Required | Schema |
+| ---- | ---------- | ----------- | -------- | ------ |
+| limit | query | Number of recently modified apps to return (1-8) | No | integer, <br>**Default:** 8 |
+
+#### Responses
+
+| Code | Description | Schema |
+| ---- | ----------- | ------ |
+| 200 | Success | **application/json**: [RecentAppListResponse](#recentapplistresponse)<br> |
+
 ### [GET] /apps/starred
 Get applications starred by the current account
 
@@ -21018,6 +21035,28 @@ Whitelist scopes accepted by RBAC app and dataset access config APIs.
 | result | string |  | Yes |
 | updated_at | integer |  | Yes |
 
+#### RecentAppListResponse
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| data | [ [RecentAppResponse](#recentappresponse) ] |  | Yes |
+
+#### RecentAppResponse
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| author_name | string |  | No |
+| icon | string |  | No |
+| icon_background | string |  | No |
+| icon_type | [IconType](#icontype) |  | No |
+| icon_url | string |  | Yes |
+| id | string |  | Yes |
+| maintainer | string |  | No |
+| mode | string, <br>**Available values:** "advanced-chat", "agent-chat", "chat", "completion", "workflow" | *Enum:* `"advanced-chat"`, `"agent-chat"`, `"chat"`, `"completion"`, `"workflow"` | Yes |
+| name | string |  | Yes |
+| permission_keys | [ string ] |  | No |
+| updated_at | integer |  | Yes |
+
 #### RecommendedAppDetailNullableResponse
 
 | Name | Type | Description | Required |

api/pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "dify-api"
-version = "1.16.0"
+version = "1.16.1"
 requires-python = "~=3.12.0"
 
 dependencies = [
     # Legacy: mature and widely deployed
     "bleach>=6.4.0,<7.0.0",
-    "boto3>=1.43.46,<2.0.0",
+    "boto3>=1.43.56,<2.0.0",
     "celery>=5.6.3,<6.0.0",
     "croniter>=6.2.2,<7.0.0",
     "dify-agent",
@@ -193,10 +193,10 @@ dev = [
 ############################################################
 storage = [
     "azure-storage-blob>=12.30.0,<13.0.0",
-    "bce-python-sdk==0.9.72",
+    "bce-python-sdk==0.9.76",
     "cos-python-sdk-v5>=1.9.44,<2.0.0",
     "esdk-obs-python>=3.26.6,<4.0.0",
-    "google-cloud-storage>=3.12.1,<4.0.0",
+    "google-cloud-storage>=3.13.0,<4.0.0",
     "opendal==0.46.0",
     "oss2>=2.19.1,<3.0.0",
     "supabase>=2.31.0,<3.0.0",

api/services/app_generate_service.py

@@ -488,7 +488,7 @@ class AppGenerateService:
         )
 
     @classmethod
-    def _get_workflow(
+    def get_workflow(
         cls,
         app_model: App,
         invoke_from: InvokeFrom,
@@ -533,6 +533,17 @@ class AppGenerateService:
 
         return workflow
 
+    @classmethod
+    def _get_workflow(
+        cls,
+        app_model: App,
+        invoke_from: InvokeFrom,
+        workflow_id: str | None = None,
+        *,
+        session: Session,
+    ) -> Workflow:
+        return cls.get_workflow(app_model, invoke_from, workflow_id, session=session)
+
     @classmethod
     def get_response_generator(
         cls,

api/services/app_service.py

@@ -1,6 +1,7 @@
 import json
 import logging
 from collections.abc import Sequence
+from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Literal, NotRequired, TypedDict, cast, override
 
@@ -41,6 +42,20 @@ from tasks.remove_app_and_related_data_task import remove_app_and_related_data_t
 logger = logging.getLogger(__name__)
 
 AppListSortBy = Literal["last_modified", "recently_created", "earliest_created"]
+RecentAppMode = Literal[
+    AppMode.COMPLETION,
+    AppMode.WORKFLOW,
+    AppMode.CHAT,
+    AppMode.ADVANCED_CHAT,
+    AppMode.AGENT_CHAT,
+]
+RECENT_APP_MODES: tuple[RecentAppMode, ...] = (
+    AppMode.COMPLETION,
+    AppMode.WORKFLOW,
+    AppMode.CHAT,
+    AppMode.ADVANCED_CHAT,
+    AppMode.AGENT_CHAT,
+)
 
 
 class AppListBaseParams(BaseModel):
@@ -65,6 +80,19 @@ class StarredAppListParams(AppListBaseParams):
     pass
 
 
+@dataclass(frozen=True)
+class RecentAppListItem:
+    id: str
+    name: str
+    icon_type: IconType | None
+    icon: str | None
+    icon_background: str | None
+    mode: RecentAppMode
+    author_name: str | None
+    updated_at: datetime
+    maintainer: str | None
+
+
 class CreateAppParams(BaseModel):
     name: str = Field(min_length=1)
     description: str | None = None
@@ -323,6 +351,62 @@ class AppService:
 
         return app_models
 
+    def get_recent_apps(
+        self,
+        user_id: str,
+        tenant_id: str,
+        params: AppListParams,
+        session: Session,
+    ) -> list[RecentAppListItem]:
+        """Return recently modified apps as one lightweight, non-paginated projection."""
+        filters = self._build_app_list_filters(user_id, tenant_id, params, session)
+        if not filters:
+            return []
+
+        stmt = (
+            sa.select(
+                App.id,
+                App.name,
+                App.icon_type,
+                App.icon,
+                App.icon_background,
+                App.mode,
+                Account.name.label("author_name"),
+                App.updated_at,
+                App.maintainer,
+            )
+            .outerjoin(Account, Account.id == App.created_by)
+            .where(*filters, App.mode.in_(RECENT_APP_MODES))
+            .order_by(App.updated_at.desc())
+            .limit(params.limit)
+        )
+        rows = session.execute(stmt).all()
+
+        return [
+            RecentAppListItem(
+                id=str(app_id),
+                name=name,
+                icon_type=icon_type,
+                icon=icon,
+                icon_background=icon_background,
+                mode=cast(RecentAppMode, mode),
+                author_name=author_name,
+                updated_at=updated_at,
+                maintainer=maintainer,
+            )
+            for (
+                app_id,
+                name,
+                icon_type,
+                icon,
+                icon_background,
+                mode,
+                author_name,
+                updated_at,
+                maintainer,
+            ) in rows
+        ]
+
     def get_paginate_starred_apps(
         self,
         user_id: str,

api/services/feature_service.py

@@ -1,6 +1,8 @@
+import logging
+from collections.abc import Mapping
 from enum import StrEnum
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from configs import dify_config
 from constants.dsl_version import CURRENT_APP_DSL_VERSION
@@ -10,6 +12,8 @@ from enums.hosted_provider import HostedTrialProvider
 from services.billing_service import BillingInfo, BillingService
 from services.enterprise.enterprise_service import EnterpriseService
 
+logger = logging.getLogger(__name__)
+
 
 class FeatureResponseModel(BaseModel):
     model_config = ConfigDict(json_schema_serialization_defaults_required=True, protected_namespaces=())
@@ -131,6 +135,13 @@ class PluginInstallationPermissionModel(FeatureResponseModel):
     restrict_to_marketplace_only: bool = False
 
 
+class _EnterprisePluginInstallationPermission(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    plugin_installation_scope: PluginInstallationScope = Field(alias="pluginInstallationScope")
+    restrict_to_marketplace_only: bool = Field(alias="restrictToMarketplaceOnly", strict=True)
+
+
 class FeatureModel(FeatureResponseModel):
     billing: BillingModel = BillingModel()
     education: EducationModel = EducationModel()
@@ -285,6 +296,14 @@ class FeatureService:
         """Return whether Enterprise plugin credential policies must be enforced."""
         return dify_config.ENTERPRISE_ENABLED
 
+    @classmethod
+    def get_plugin_installation_permission(cls) -> PluginInstallationPermissionModel:
+        """Resolve the validated deployment-wide plugin installation policy."""
+        if not dify_config.ENTERPRISE_ENABLED:
+            return PluginInstallationPermissionModel()
+
+        return cls._resolve_plugin_installation_permission(EnterpriseService.get_info())
+
     @classmethod
     def get_license(cls) -> LicenseModel:
         """Return full license detail. Enterprise-only; requires an authenticated caller.
@@ -452,6 +471,33 @@ class FeatureService:
                 )
         return license_model
 
+    @classmethod
+    def _resolve_plugin_installation_permission(
+        cls, enterprise_info: Mapping[str, object]
+    ) -> PluginInstallationPermissionModel:
+        if "PluginInstallationPermission" not in enterprise_info:
+            return PluginInstallationPermissionModel()
+
+        try:
+            permission = _EnterprisePluginInstallationPermission.model_validate(
+                enterprise_info["PluginInstallationPermission"]
+            )
+        except ValidationError as exc:
+            # Do not attach the exception because it may contain raw Enterprise configuration values.
+            logger.error(  # noqa: TRY400
+                "Invalid Enterprise plugin installation permission; denying all plugin installations: %s",
+                exc.errors(include_input=False),
+            )
+            return PluginInstallationPermissionModel(
+                plugin_installation_scope=PluginInstallationScope.NONE,
+                restrict_to_marketplace_only=True,
+            )
+
+        return PluginInstallationPermissionModel(
+            plugin_installation_scope=permission.plugin_installation_scope,
+            restrict_to_marketplace_only=permission.restrict_to_marketplace_only,
+        )
+
     @classmethod
     def _fulfill_params_from_enterprise(cls, features: SystemFeatureModel):
         enterprise_info = EnterpriseService.get_info()
@@ -499,11 +545,4 @@ class FeatureService:
                 status=LicenseStatus(license_info.get("status", LicenseStatus.INACTIVE))
             )
 
-        if "PluginInstallationPermission" in enterprise_info:
-            plugin_installation_info = enterprise_info["PluginInstallationPermission"]
-            features.plugin_installation_permission.plugin_installation_scope = plugin_installation_info[
-                "pluginInstallationScope"
-            ]
-            features.plugin_installation_permission.restrict_to_marketplace_only = plugin_installation_info[
-                "restrictToMarketplaceOnly"
-            ]
+        features.plugin_installation_permission = cls._resolve_plugin_installation_permission(enterprise_info)

api/services/workflow/workflow_converter.py

@@ -18,6 +18,7 @@ from core.app.apps.completion.app_config_manager import CompletionAppConfigManag
 from core.helper import encrypter
 from core.prompt.simple_prompt_transform import SimplePromptTransform
 from core.prompt.utils.prompt_template_parser import PromptTemplateParser
+from core.workflow.variable_prefixes import USER_INPUT_VARIABLE_NODE_ID
 from events.app_event import app_was_created
 from graphon.file import FileUploadConfig
 from graphon.model_runtime.entities.llm_entities import LLMMode
@@ -597,7 +598,7 @@ class WorkflowConverter:
                 },
                 "vision": {
                     "enabled": file_upload is not None,
-                    "variable_selector": ["sys", "files"] if file_upload is not None else None,
+                    "variable_selector": [USER_INPUT_VARIABLE_NODE_ID, "files"] if file_upload is not None else None,
                     "configs": {"detail": file_upload.image_config.detail}
                     if file_upload is not None and file_upload.image_config is not None
                     else None,

api/services/workflow_service.py

@@ -6,8 +6,10 @@ from collections.abc import Callable, Generator, Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any, cast
 
-from sqlalchemy import exists, select
+from sqlalchemy import exists, inspect, select, update
+from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm.attributes import set_committed_value
 
 from configs import dify_config
 from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfigManager
@@ -143,6 +145,8 @@ from .human_input_delivery_test_service import (
 from .workflow_draft_variable_service import DraftVariableSaver, DraftVarLoader, WorkflowDraftVariableService
 from .workflow_restore import apply_published_workflow_snapshot_to_draft
 
+logger = logging.getLogger(__name__)
+
 _file_access_controller = DatabaseFileAccessController()
 
 
@@ -155,6 +159,7 @@ class WorkflowService:
         """Initialize WorkflowService with repository dependencies."""
         if session_maker is None:
             session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
+        self._session_maker = session_maker
         self._node_execution_service_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(
             session_maker
         )
@@ -211,7 +216,7 @@ class WorkflowService:
         )
 
         # return draft workflow
-        return workflow
+        return self._persist_legacy_sys_files_migration_on_load(workflow)
 
     def get_published_workflow_by_id(self, app_model: App, workflow_id: str, *, session: Session) -> Workflow | None:
         """
@@ -236,6 +241,7 @@ class WorkflowService:
                 f"Cannot use draft workflow version. Workflow ID: {workflow_id}. "
                 f"Please use a published workflow version or leave workflow_id empty."
             )
+        self._persist_legacy_sys_files_migration_on_load(workflow)
         return workflow
 
     def get_published_workflow(self, app_model: App, *, session: Session) -> Workflow | None:
@@ -259,6 +265,53 @@ class WorkflowService:
             .limit(1)
         )
 
+        return self._persist_legacy_sys_files_migration_on_load(workflow)
+
+    def _persist_legacy_sys_files_migration_on_load(self, workflow: Workflow | None) -> Workflow | None:
+        """Persist a load-time graph rewrite without joining or dirtying the caller's transaction."""
+
+        if workflow is None:
+            return None
+        if inspect(workflow, raiseerr=False) is None:
+            return workflow
+
+        # TODO: Remove this load-time persistence path after the historical workflow migration is complete.
+        original_graph = workflow.graph
+        if not workflow.migrate_legacy_sys_files_graph_in_place():
+            return workflow
+
+        migrated_graph = workflow.graph
+        try:
+            with self._session_maker.begin() as session:
+                result = session.execute(
+                    update(Workflow)
+                    .where(
+                        Workflow.id == workflow.id,
+                        Workflow.tenant_id == workflow.tenant_id,
+                        Workflow.graph == original_graph,
+                    )
+                    .values(graph=migrated_graph)
+                )
+                if getattr(result, "rowcount", None) == 0:
+                    logger.warning(
+                        "Skipped persisting legacy sys.files workflow migration because the workflow changed "
+                        "concurrently, "
+                        "workflow_id=%s tenant_id=%s",
+                        workflow.id,
+                        workflow.tenant_id,
+                    )
+        except SQLAlchemyError:
+            logger.warning(
+                "Failed to persist legacy sys.files workflow migration, workflow_id=%s tenant_id=%s",
+                workflow.id,
+                workflow.tenant_id,
+                exc_info=True,
+            )
+        finally:
+            # The conditional update owns persistence. Mark the caller's instance clean so its later flush cannot
+            # overwrite a concurrent workflow edit with the compatibility rewrite.
+            set_committed_value(workflow, "graph", migrated_graph)
+
         return workflow
 
     def get_accessible_app_ids(self, app_ids: Sequence[str], tenant_id: str, *, session: Session) -> set[str]:

api/tasks/agent_backend_session_cleanup_task.py

@@ -22,6 +22,7 @@ def _create_agent_backend_client():
         return None
     return create_agent_backend_run_client(
         base_url=dify_config.AGENT_BACKEND_BASE_URL,
+        api_token=dify_config.AGENT_BACKEND_API_TOKEN,
         use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
         fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
         stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS,

api/tasks/app_generate/workflow_execute_task.py

@@ -457,7 +457,7 @@ def _publish_streaming_response(
 @shared_task(queue=WORKFLOW_BASED_APP_EXECUTION_QUEUE)
 def workflow_based_app_execution_task(
     payload: str,
-) -> Generator[Mapping[str, Any] | str, None, None] | Mapping[str, Any] | None:
+) -> Mapping[str, Any] | None:
     exec_params = AppExecutionParams.model_validate_json(payload)
 
     logger.info("workflow_based_app_execution_task run with params: %s", exec_params)

api/tests/fixtures/workflow/basic_chatflow.yml

@@ -61,7 +61,7 @@ workflow:
           query_prompt_template: '{{#sys.query#}}
 
 
-            {{#sys.files#}}'
+            {{#userinput.files#}}'
           window:
             enabled: false
             size: 10

api/tests/fixtures/workflow/multilingual_parallel_llm_streaming_workflow.yml

@@ -162,7 +162,7 @@ workflow:
           query_prompt_template: '{{#sys.query#}}
 
 
-            {{#sys.files#}}'
+            {{#userinput.files#}}'
           role_prefix:
             assistant: ''
             user: ''
@@ -207,7 +207,7 @@ workflow:
           query_prompt_template: '{{#sys.query#}}
 
 
-            {{#sys.files#}}'
+            {{#userinput.files#}}'
           role_prefix:
             assistant: ''
             user: ''

api/tests/fixtures/workflow/test_complex_branch.yml

@@ -178,7 +178,7 @@ workflow:
           query_prompt_template: '{{#sys.query#}}
 
 
-            {{#sys.files#}}'
+            {{#userinput.files#}}'
           role_prefix:
             assistant: ''
             user: ''

api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py

@@ -1,494 +1,85 @@
-"""Testcontainers integration tests for controllers.console.datasets.data_source endpoints."""
+"""Integration coverage for Notion page bindings backed by persisted documents."""
 
-from __future__ import annotations
-
-import inspect
-from collections.abc import Iterator
-from datetime import UTC, datetime
-from unittest.mock import MagicMock, PropertyMock, patch
+from inspect import unwrap
+from unittest.mock import MagicMock, patch
 from uuid import uuid4
 
-import pytest
 from flask import Flask
 from sqlalchemy.orm import Session
-from werkzeug.exceptions import NotFound
 
-from controllers.console.datasets import data_source
-from controllers.console.datasets.data_source import (
-    DataSourceApi,
-    DataSourceNotionDatasetSyncApi,
-    DataSourceNotionDocumentSyncApi,
-    DataSourceNotionIndexingEstimateApi,
-    DataSourceNotionListApi,
-    DataSourceNotionPreviewApi,
-)
-from core.rag.index_processor.constant.index_type import IndexStructureType
-from models import Account, DataSourceOauthBinding
+from controllers.console.datasets.data_source import DataSourceNotionListApi
+from models import Account
 from models.dataset import Document
 from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
 
 
-@pytest.fixture
-def current_user() -> Account:
-    account = Account(name="Test User", email="u1@example.com")
-    account.id = "u1"
-    return account
+def test_notion_page_is_marked_bound_from_persisted_document(
+    flask_app_with_containers: Flask,
+    db_session_with_containers: Session,
+) -> None:
+    tenant_id = str(uuid4())
+    dataset_id = str(uuid4())
+    account = Account(name="Test User", email="user@example.com")
+    account.id = str(uuid4())
+    document = Document(
+        tenant_id=tenant_id,
+        dataset_id=dataset_id,
+        position=1,
+        data_source_type=DataSourceType.NOTION_IMPORT,
+        data_source_info='{"notion_page_id": "page-1"}',
+        batch=f"batch-{uuid4()}",
+        name="Notion Page",
+        created_from=DocumentCreatedFrom.WEB,
+        created_by=str(uuid4()),
+        indexing_status=IndexingStatus.COMPLETED,
+        enabled=True,
+    )
+    db_session_with_containers.add(document)
+    db_session_with_containers.commit()
+    runtime = MagicMock(
+        get_online_document_pages=lambda **_kwargs: iter(
+            [
+                MagicMock(
+                    result=[
+                        MagicMock(
+                            workspace_id="workspace-1",
+                            workspace_name="Workspace",
+                            workspace_icon=None,
+                            pages=[
+                                MagicMock(
+                                    page_id="page-1",
+                                    page_name="Page",
+                                    type="page",
+                                    parent_id="parent",
+                                    page_icon=None,
+                                )
+                            ],
+                        )
+                    ]
+                )
+            ]
+        ),
+        datasource_provider_type=lambda: None,
+    )
 
-
-@pytest.fixture
-def mock_engine() -> Iterator[None]:
-    with patch.object(
-        type(data_source.db),
-        "engine",
-        new_callable=PropertyMock,
-        return_value=MagicMock(),
+    with (
+        flask_app_with_containers.test_request_context(f"/?credential_id=c1&dataset_id={dataset_id}"),
+        patch(
+            "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
+            return_value={"token": "token"},
+        ),
+        patch(
+            "controllers.console.datasets.data_source.DatasetService.get_dataset",
+            return_value=MagicMock(data_source_type="notion_import"),
+        ),
+        patch(
+            "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime",
+            return_value=runtime,
+        ),
     ):
-        yield
-
-
-class TestDataSourceApi:
-    @pytest.fixture
-    def app(self, flask_app_with_containers: Flask) -> Flask:
-        return flask_app_with_containers
-
-    def test_get_success(self, app: Flask) -> None:
-        api = DataSourceApi()
-        method = inspect.unwrap(api.get)
-
-        binding = DataSourceOauthBinding(
-            tenant_id="tenant-1",
-            access_token="token",
-            provider="notion",
-            source_info={
-                "workspace_name": "Workspace",
-                "workspace_id": "workspace-1",
-                "workspace_icon": None,
-                "total": 1,
-                "pages": [
-                    {
-                        "page_id": "page-1",
-                        "page_name": "Page",
-                        "page_icon": {"type": "emoji", "emoji": "P", "url": None},
-                        "parent_id": "parent-1",
-                        "type": "page",
-                    }
-                ],
-            },
-        )
-        binding.id = "b1"
-        binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
-        binding.disabled = False
-
-        with (
-            app.test_request_context("/"),
-            patch(
-                "controllers.console.datasets.data_source.db.session.scalars",
-                return_value=MagicMock(all=lambda: [binding]),
-            ),
-        ):
-            response, status = method(api, "tenant-1")
-
-        assert status == 200
-        assert response["data"][0] == {
-            "id": "b1",
-            "provider": "notion",
-            "created_at": 1779670923,
-            "is_bound": True,
-            "disabled": False,
-            "source_info": {
-                "workspace_name": "Workspace",
-                "workspace_id": "workspace-1",
-                "workspace_icon": None,
-                "pages": [
-                    {
-                        "page_name": "Page",
-                        "page_id": "page-1",
-                        "page_icon": {"type": "emoji", "url": None, "emoji": "P"},
-                        "parent_id": "parent-1",
-                        "type": "page",
-                    }
-                ],
-                "total": 1,
-            },
-            "link": "http://localhost/console/api/oauth/data-source/notion",
-        }
-
-    def test_get_no_bindings(self, app: Flask) -> None:
-        api = DataSourceApi()
-        method = inspect.unwrap(api.get)
-
-        with (
-            app.test_request_context("/"),
-            patch(
-                "controllers.console.datasets.data_source.db.session.scalars",
-                return_value=MagicMock(all=lambda: []),
-            ),
-        ):
-            response, status = method(api, "tenant-1")
-
-        assert status == 200
-        assert response["data"] == []
-
-    def test_patch_enable_binding(self, app: Flask) -> None:
-        api = DataSourceApi()
-        method = inspect.unwrap(api.patch)
-
-        binding = MagicMock(id="b1", disabled=True)
-        session = MagicMock()
-        session.scalar.return_value = binding
-
-        with app.test_request_context("/"):
-            response, status = method(api, session, "tenant-1", "b1", "enable")
-
-        assert status == 200
-        assert binding.disabled is False
-
-    def test_patch_disable_binding(self, app: Flask) -> None:
-        api = DataSourceApi()
-        method = inspect.unwrap(api.patch)
-
-        binding = MagicMock(id="b1", disabled=False)
-        session = MagicMock()
-        session.scalar.return_value = binding
-
-        with app.test_request_context("/"):
-            response, status = method(api, session, "tenant-1", "b1", "disable")
-
-        assert status == 200
-        assert binding.disabled is True
-
-    def test_patch_binding_not_found(self, app: Flask) -> None:
-        api = DataSourceApi()
-        method = inspect.unwrap(api.patch)
-        session = MagicMock()
-        session.scalar.return_value = None
-
-        with app.test_request_context("/"):
-            with pytest.raises(NotFound):
-                method(api, session, "tenant-1", "b1", "enable")
-
-    def test_patch_enable_already_enabled(self, app: Flask) -> None:
-        api = DataSourceApi()
-        method = inspect.unwrap(api.patch)
-
-        binding = MagicMock(id="b1", disabled=False)
-        session = MagicMock()
-        session.scalar.return_value = binding
-
-        with app.test_request_context("/"):
-            with pytest.raises(ValueError):
-                method(api, session, "tenant-1", "b1", "enable")
-
-    def test_patch_disable_already_disabled(self, app: Flask) -> None:
-        api = DataSourceApi()
-        method = inspect.unwrap(api.patch)
-
-        binding = MagicMock(id="b1", disabled=True)
-        session = MagicMock()
-        session.scalar.return_value = binding
-
-        with app.test_request_context("/"):
-            with pytest.raises(ValueError):
-                method(api, session, "tenant-1", "b1", "disable")
-
-
-class TestDataSourceNotionListApi:
-    @pytest.fixture
-    def app(self, flask_app_with_containers: Flask) -> Flask:
-        return flask_app_with_containers
-
-    def test_get_credential_not_found(self, app: Flask, current_user: Account) -> None:
-        api = DataSourceNotionListApi()
-        method = inspect.unwrap(api.get)
-
-        with (
-            app.test_request_context("/?credential_id=c1"),
-            patch(
-                "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
-                return_value=None,
-            ),
-        ):
-            with pytest.raises(NotFound):
-                method(api, MagicMock(), "tenant-1", current_user)
-
-    def test_get_success_no_dataset_id(self, app: Flask, current_user: Account, mock_engine: None) -> None:
-        api = DataSourceNotionListApi()
-        method = inspect.unwrap(api.get)
-
-        page = MagicMock(
-            page_id="p1",
-            page_name="Page 1",
-            type="page",
-            parent_id="parent",
-            page_icon=None,
+        response, status = unwrap(DataSourceNotionListApi().get)(
+            DataSourceNotionListApi(), db_session_with_containers, tenant_id, account
         )
 
-        online_document_message = MagicMock(
-            result=[
-                MagicMock(
-                    workspace_id="w1",
-                    workspace_name="My Workspace",
-                    workspace_icon="icon",
-                    pages=[page],
-                )
-            ]
-        )
-
-        with (
-            app.test_request_context("/?credential_id=c1"),
-            patch(
-                "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
-                return_value={"token": "t"},
-            ),
-            patch(
-                "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime",
-                return_value=MagicMock(
-                    get_online_document_pages=lambda **kw: iter([online_document_message]),
-                    datasource_provider_type=lambda: None,
-                ),
-            ),
-        ):
-            response, status = method(api, MagicMock(), "tenant-1", current_user)
-
-        assert status == 200
-
-    def test_get_success_with_dataset_id(
-        self, app: Flask, current_user: Account, mock_engine: None, db_session_with_containers: Session
-    ) -> None:
-        api = DataSourceNotionListApi()
-        method = inspect.unwrap(api.get)
-        tenant_id = str(uuid4())
-        dataset_id = str(uuid4())
-
-        page = MagicMock(
-            page_id="p1",
-            page_name="Page 1",
-            type="page",
-            parent_id="parent",
-            page_icon=None,
-        )
-
-        online_document_message = MagicMock(
-            result=[
-                MagicMock(
-                    workspace_id="w1",
-                    workspace_name="My Workspace",
-                    workspace_icon="icon",
-                    pages=[page],
-                )
-            ]
-        )
-
-        dataset = MagicMock(data_source_type="notion_import")
-        document = Document(
-            tenant_id=tenant_id,
-            dataset_id=dataset_id,
-            position=1,
-            data_source_type=DataSourceType.NOTION_IMPORT,
-            data_source_info='{"notion_page_id": "p1"}',
-            batch=f"batch-{uuid4()}",
-            name="Notion Page",
-            created_from=DocumentCreatedFrom.WEB,
-            created_by=str(uuid4()),
-            indexing_status=IndexingStatus.COMPLETED,
-            enabled=True,
-        )
-        db_session_with_containers.add(document)
-        db_session_with_containers.commit()
-
-        with (
-            app.test_request_context(f"/?credential_id=c1&dataset_id={dataset_id}"),
-            patch(
-                "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
-                return_value={"token": "t"},
-            ),
-            patch(
-                "controllers.console.datasets.data_source.DatasetService.get_dataset",
-                return_value=dataset,
-            ),
-            patch(
-                "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime",
-                return_value=MagicMock(
-                    get_online_document_pages=lambda **kw: iter([online_document_message]),
-                    datasource_provider_type=lambda: None,
-                ),
-            ),
-        ):
-            response, status = method(api, db_session_with_containers, tenant_id, current_user)
-
-        assert status == 200
-
-    def test_get_invalid_dataset_type(self, app: Flask, current_user: Account, mock_engine: None) -> None:
-        api = DataSourceNotionListApi()
-        method = inspect.unwrap(api.get)
-
-        dataset = MagicMock(data_source_type="other_type")
-
-        with (
-            app.test_request_context("/?credential_id=c1&dataset_id=ds1"),
-            patch(
-                "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
-                return_value={"token": "t"},
-            ),
-            patch(
-                "controllers.console.datasets.data_source.DatasetService.get_dataset",
-                return_value=dataset,
-            ),
-        ):
-            with pytest.raises(ValueError):
-                method(api, MagicMock(), "tenant-1", current_user)
-
-
-class TestDataSourceNotionPreviewApi:
-    @pytest.fixture
-    def app(self, flask_app_with_containers: Flask) -> Flask:
-        return flask_app_with_containers
-
-    def test_get_preview_success(self, app: Flask) -> None:
-        api = DataSourceNotionPreviewApi()
-        method = inspect.unwrap(api.get)
-
-        extractor = MagicMock(extract=lambda: [MagicMock(page_content="hello")])
-
-        with (
-            app.test_request_context("/?credential_id=c1"),
-            patch(
-                "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
-                return_value={"integration_secret": "t"},
-            ),
-            patch(
-                "controllers.console.datasets.data_source.NotionExtractor",
-                return_value=extractor,
-            ),
-        ):
-            response, status = method(api, "tenant-1", "p1", "page")
-
-        assert status == 200
-
-
-class TestDataSourceNotionIndexingEstimateApi:
-    @pytest.fixture
-    def app(self, flask_app_with_containers: Flask) -> Flask:
-        return flask_app_with_containers
-
-    def test_post_indexing_estimate_success(self, app: Flask) -> None:
-        api = DataSourceNotionIndexingEstimateApi()
-        method = inspect.unwrap(api.post)
-
-        empty_rules: dict[str, object] = {}
-        payload: dict[str, object] = {
-            "notion_info_list": [
-                {
-                    "workspace_id": "w1",
-                    "credential_id": "c1",
-                    "pages": [{"page_id": "p1", "type": "page"}],
-                }
-            ],
-            "process_rule": {"rules": empty_rules},
-            "doc_form": IndexStructureType.PARAGRAPH_INDEX,
-            "doc_language": "English",
-        }
-
-        with (
-            app.test_request_context("/", method="POST", json=payload, headers={"Content-Type": "application/json"}),
-            patch(
-                "controllers.console.datasets.data_source.DocumentService.estimate_args_validate",
-            ),
-            patch(
-                "controllers.console.datasets.data_source.IndexingRunner.indexing_estimate",
-                return_value=MagicMock(model_dump=lambda: {"total_pages": 1}),
-            ),
-        ):
-            response, status = method(api, MagicMock(), "tenant-1")
-
-        assert status == 200
-
-
-class TestDataSourceNotionDatasetSyncApi:
-    @pytest.fixture
-    def app(self, flask_app_with_containers: Flask) -> Flask:
-        return flask_app_with_containers
-
-    def test_get_success(self, app: Flask) -> None:
-        api = DataSourceNotionDatasetSyncApi()
-        method = inspect.unwrap(api.get)
-
-        with (
-            app.test_request_context("/"),
-            patch(
-                "controllers.console.datasets.data_source.DatasetService.get_dataset",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "controllers.console.datasets.data_source.DocumentService.get_document_by_dataset_id",
-                return_value=[MagicMock(id="d1")],
-            ),
-            patch(
-                "controllers.console.datasets.data_source.document_indexing_sync_task.delay",
-                return_value=None,
-            ),
-        ):
-            response, status = method(api, MagicMock(), "ds-1")
-
-        assert status == 200
-
-    def test_get_dataset_not_found(self, app: Flask) -> None:
-        api = DataSourceNotionDatasetSyncApi()
-        method = inspect.unwrap(api.get)
-
-        with (
-            app.test_request_context("/"),
-            patch(
-                "controllers.console.datasets.data_source.DatasetService.get_dataset",
-                return_value=None,
-            ),
-        ):
-            with pytest.raises(NotFound):
-                method(api, MagicMock(), "ds-1")
-
-
-class TestDataSourceNotionDocumentSyncApi:
-    @pytest.fixture
-    def app(self, flask_app_with_containers: Flask) -> Flask:
-        return flask_app_with_containers
-
-    def test_get_success(self, app: Flask) -> None:
-        api = DataSourceNotionDocumentSyncApi()
-        method = inspect.unwrap(api.get)
-
-        with (
-            app.test_request_context("/"),
-            patch(
-                "controllers.console.datasets.data_source.DatasetService.get_dataset",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "controllers.console.datasets.data_source.DocumentService.get_document",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "controllers.console.datasets.data_source.document_indexing_sync_task.delay",
-                return_value=None,
-            ),
-        ):
-            response, status = method(api, MagicMock(), "ds-1", "doc-1")
-
-        assert status == 200
-
-    def test_get_document_not_found(self, app: Flask) -> None:
-        api = DataSourceNotionDocumentSyncApi()
-        method = inspect.unwrap(api.get)
-
-        with (
-            app.test_request_context("/"),
-            patch(
-                "controllers.console.datasets.data_source.DatasetService.get_dataset",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "controllers.console.datasets.data_source.DocumentService.get_document",
-                return_value=None,
-            ),
-        ):
-            with pytest.raises(NotFound):
-                method(api, MagicMock(), "ds-1", "doc-1")
+    assert status == 200
+    assert response["notion_info"][0]["pages"][0]["is_bound"] is True

api/tests/test_containers_integration_tests/services/test_messages_clean_service.py

@@ -4,7 +4,7 @@ import datetime
 import json
 import uuid
 from decimal import Decimal
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import pytest
 from faker import Faker
@@ -1172,65 +1172,8 @@ class TestMessagesCleanServiceIntegration:
         # Verify all messages were deleted
         assert db_session_with_containers.query(Message).where(Message.id.in_(msg_ids)).count() == 0
 
-    def test_from_time_range_validation(self):
-        """Test that from_time_range raises ValueError for invalid inputs."""
-        policy = MagicMock(spec=BillingDisabledPolicy)
-        now = datetime.datetime.now()
-
-        with pytest.raises(ValueError, match="start_from .* must be less than end_before"):
-            MessagesCleanService.from_time_range(policy, now, now)
-
-        with pytest.raises(ValueError, match="batch_size .* must be greater than 0"):
-            MessagesCleanService.from_time_range(policy, now - datetime.timedelta(days=1), now, batch_size=0)
-
-    def test_from_time_range_success(self):
-        """Test that from_time_range creates a service with correct parameters."""
-        policy = MagicMock(spec=BillingDisabledPolicy)
-        start = datetime.datetime(2024, 1, 1)
-        end = datetime.datetime(2024, 2, 1)
-
-        service = MessagesCleanService.from_time_range(policy, start, end)
-        assert service._start_from == start
-        assert service._end_before == end
-
-    def test_from_days_validation(self):
-        """Test that from_days raises ValueError for invalid inputs."""
-        policy = MagicMock(spec=BillingDisabledPolicy)
-
-        with pytest.raises(ValueError, match="days .* must be greater than or equal to 0"):
-            MessagesCleanService.from_days(policy, days=-1)
-
-        with pytest.raises(ValueError, match="batch_size .* must be greater than 0"):
-            MessagesCleanService.from_days(policy, days=30, batch_size=0)
-
-    def test_from_days_success(self):
-        """Test that from_days creates a service with correct parameters."""
-        policy = MagicMock(spec=BillingDisabledPolicy)
-
-        with patch("services.retention.conversation.messages_clean_service.naive_utc_now") as mock_now:
-            fixed_now = datetime.datetime(2024, 6, 1)
-            mock_now.return_value = fixed_now
-
-            service = MessagesCleanService.from_days(policy, days=10)
-            assert service._start_from is None
-            assert service._end_before == fixed_now - datetime.timedelta(days=10)
-
     def test_batch_delete_message_relations_empty(self, db_session_with_containers: Session):
         """Test that batch_delete_message_relations with empty list does nothing."""
         # Get execute call count before
         MessagesCleanService._batch_delete_message_relations(db_session_with_containers, [])
         # No exception means success — empty list is a no-op
-
-    def test_run_calls_clean_messages(self):
-        """Test that run() delegates to _clean_messages_by_time_range."""
-        policy = MagicMock(spec=BillingDisabledPolicy)
-        service = MessagesCleanService(
-            policy=policy,
-            end_before=datetime.datetime.now(),
-            batch_size=10,
-        )
-        with patch.object(service, "_clean_messages_by_time_range") as mock_clean:
-            mock_clean.return_value = {"total_deleted": 5}
-            result = service.run()
-            assert result == {"total_deleted": 5}
-            mock_clean.assert_called_once()

api/tests/unit_tests/commands/test_fix_app_site_missing.py

@@ -4,7 +4,7 @@ from unittest.mock import MagicMock
 import pytest
 from sqlalchemy.orm import Session
 
-from commands import system as system_commands
+from commands import app_maintenance as app_maintenance_commands
 
 
 def test_fix_app_site_missing_passes_loaded_session_to_signal(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -30,15 +30,15 @@ def test_fix_app_site_missing_passes_loaded_session_to_signal(monkeypatch: pytes
     engine = MagicMock()
     engine.begin.return_value.__enter__.return_value = connection
 
-    monkeypatch.setattr(system_commands, "db", SimpleNamespace(engine=engine, session=scoped_session))
+    monkeypatch.setattr(app_maintenance_commands, "db", SimpleNamespace(engine=engine, session=scoped_session))
     send = MagicMock(side_effect=lambda *_args, **_kwargs: phase_events.append("signal"))
-    monkeypatch.setattr(system_commands.app_was_created, "send", send)
+    monkeypatch.setattr(app_maintenance_commands.app_was_created, "send", send)
 
-    system_commands.fix_app_site_missing.callback()
+    app_maintenance_commands.fix_app_site_missing.callback()
 
     scoped_session.assert_called_once_with()
     scalar.assert_called_once()
-    get.assert_called_once_with(system_commands.Tenant, app.tenant_id)
+    get.assert_called_once_with(app_maintenance_commands.Tenant, app.tenant_id)
     tenant.get_accounts.assert_called_once_with(session=session)
     send.assert_called_once_with(app, account=account, session=session)
     commit.assert_called_once_with()
@@ -62,15 +62,19 @@ def test_fix_app_site_missing_rolls_back_when_signal_fails(monkeypatch: pytest.M
     engine = MagicMock()
     engine.begin.return_value.__enter__.return_value = connection
 
-    monkeypatch.setattr(system_commands, "db", SimpleNamespace(engine=engine, session=MagicMock(return_value=session)))
+    monkeypatch.setattr(
+        app_maintenance_commands,
+        "db",
+        SimpleNamespace(engine=engine, session=MagicMock(return_value=session)),
+    )
 
     def fail_signal(*_args, **_kwargs) -> None:
         phase_events.append("signal")
         raise RuntimeError("failed")
 
-    monkeypatch.setattr(system_commands.app_was_created, "send", MagicMock(side_effect=fail_signal))
+    monkeypatch.setattr(app_maintenance_commands.app_was_created, "send", MagicMock(side_effect=fail_signal))
 
-    system_commands.fix_app_site_missing.callback()
+    app_maintenance_commands.fix_app_site_missing.callback()
 
     session.rollback.assert_called_once_with()
     session.commit.assert_not_called()

api/tests/unit_tests/commands/test_legacy_model_type_migration.py

@@ -6,6 +6,7 @@ import json
 import os
 import threading
 import time
+from collections.abc import Iterator
 from datetime import datetime, timedelta
 from pathlib import Path
 from types import SimpleNamespace
@@ -15,9 +16,12 @@ import pytest
 import sqlalchemy as sa
 from click.testing import CliRunner
 from sqlalchemy.exc import OperationalError
+from sqlalchemy.orm import Session, SessionTransaction, sessionmaker
 
 from graphon.model_runtime.entities.model_entities import ModelType
+from models import Dataset, DatasetPermission, DatasetPermissionEnum
 from models.account import Tenant
+from models.base import TypeBase
 from models.enums import CredentialSourceType
 from models.provider import ProviderModel
 from tests.helpers.legacy_model_type_migration import (
@@ -59,6 +63,40 @@ def command_module():
         )
 
 
+@pytest.fixture
+def rbac_session(sqlite_engine: sa.Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]:
+    """Bind RBAC command reads to persisted SQLite dataset rows."""
+
+    TypeBase.metadata.create_all(
+        sqlite_engine,
+        tables=[Dataset.__table__, DatasetPermission.__table__],
+    )
+    factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
+    monkeypatch.setattr("commands.rbac.session_factory.create_session", factory)
+    with factory() as session:
+        yield session
+
+
+def _persist_dataset(
+    session: Session,
+    *,
+    dataset_id: str = "dataset-1",
+    tenant_id: str = "tenant-1",
+    permission: DatasetPermissionEnum = DatasetPermissionEnum.ONLY_ME,
+    created_by: str = "creator-account-1",
+) -> Dataset:
+    dataset = Dataset(
+        id=dataset_id,
+        tenant_id=tenant_id,
+        name=f"Dataset {dataset_id}",
+        permission=permission,
+        created_by=created_by,
+    )
+    session.add(dataset)
+    session.commit()
+    return dataset
+
+
 def _parse_json_lines(output: io.StringIO) -> list[dict[str, object]]:
     return [json.loads(line) for line in output.getvalue().splitlines() if line.strip()]
 
@@ -363,56 +401,35 @@ def test_dataset_permission_rbac_migration_maps_legacy_permissions_to_enum_scope
 
 def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator(
     command_module,
+    rbac_session: Session,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     rbac_module = importlib.import_module("commands.rbac")
-    dataset_row = SimpleNamespace(
-        id="dataset-1",
-        tenant_id="tenant-1",
-        permission="only_me",
-        created_by="creator-account-1",
-    )
-    execute_results = [[dataset_row], [], []]
+    _persist_dataset(rbac_session)
     calls: list[dict[str, object]] = []
-    session_closed = False
-
-    class FakeExecuteResult:
-        def __init__(self, rows: list[object]) -> None:
-            self._rows = rows
-
-        def all(self) -> list[object]:
-            return self._rows
-
-    class FakeSession:
-        def __enter__(self):
-            return self
-
-        def __exit__(self, exc_type, exc, traceback) -> None:
-            nonlocal session_closed
-            session_closed = True
-            pass
-
-        def execute(self, stmt):
-            return FakeExecuteResult(execute_results.pop(0))
-
-    class FakeSessionFactory:
-        @staticmethod
-        def create_session() -> FakeSession:
-            return FakeSession()
+    read_transaction_ended = False
 
     def fake_replace_whitelist(**kwargs):
-        assert session_closed is True
+        assert read_transaction_ended is True
         calls.append(kwargs)
 
-    monkeypatch.setattr(rbac_module, "session_factory", FakeSessionFactory)
-    monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist)
+    def _record_transaction_end(session: Session, transaction: object) -> None:
+        nonlocal read_transaction_ended
+        del transaction
+        if session.get_bind() is rbac_session.get_bind():
+            read_transaction_ended = True
 
-    command_module.migrate_dataset_permissions_to_rbac.callback(
-        tenant_id=None,
-        dataset_id=None,
-        batch_size=500,
-        dry_run=False,
-    )
+    sa.event.listen(Session, "after_transaction_end", _record_transaction_end)
+    monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist)
+    try:
+        command_module.migrate_dataset_permissions_to_rbac.callback(
+            tenant_id=None,
+            dataset_id=None,
+            batch_size=500,
+            dry_run=False,
+        )
+    finally:
+        sa.event.remove(Session, "after_transaction_end", _record_transaction_end)
 
     assert calls[0]["tenant_id"] == "tenant-1"
     assert calls[0]["account_id"] == "creator-account-1"
@@ -422,41 +439,19 @@ def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator(
 
 def test_dataset_permission_rbac_migration_dry_run_outputs_structured_proposed_changes(
     command_module,
+    rbac_session: Session,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     rbac_module = importlib.import_module("commands.rbac")
-    dataset_row = SimpleNamespace(
-        id="dataset-1",
-        tenant_id="tenant-1",
-        permission="partial_members",
-        created_by="creator-account-1",
+    dataset = _persist_dataset(rbac_session, permission=DatasetPermissionEnum.PARTIAL_TEAM)
+    rbac_session.add(
+        DatasetPermission(
+            dataset_id=dataset.id,
+            account_id="member-account-1",
+            tenant_id=dataset.tenant_id,
+        )
     )
-    permission_row = SimpleNamespace(dataset_id="dataset-1", account_id="member-account-1")
-    execute_results = [[dataset_row], [permission_row], []]
-
-    class FakeExecuteResult:
-        def __init__(self, rows: list[object]) -> None:
-            self._rows = rows
-
-        def all(self) -> list[object]:
-            return self._rows
-
-    class FakeSession:
-        def __enter__(self):
-            return self
-
-        def __exit__(self, exc_type, exc, traceback) -> None:
-            pass
-
-        def execute(self, stmt):
-            return FakeExecuteResult(execute_results.pop(0))
-
-    class FakeSessionFactory:
-        @staticmethod
-        def create_session() -> FakeSession:
-            return FakeSession()
-
-    monkeypatch.setattr(rbac_module, "session_factory", FakeSessionFactory)
+    rbac_session.commit()
     monkeypatch.setattr(
         rbac_module.RBACService.DatasetAccess,
         "replace_whitelist",
@@ -1306,50 +1301,36 @@ def test_provider_models_processing_uses_same_plan_locking_and_transaction_entry
     begin_calls: list[str] = []
     configure_calls: list[str] = []
 
-    class _FakeBeginContext:
-        def __init__(self, phase: str) -> None:
-            self._phase = phase
+    def _record_begin(session: Session, transaction: SessionTransaction) -> None:
+        if session.get_bind() is sqlite_engine and transaction.parent is None:
+            begin_calls.append(current_phase["name"])
 
-        def __enter__(self) -> None:
-            begin_calls.append(self._phase)
-
-        def __exit__(self, exc_type, exc, tb) -> bool:
-            return False
-
-    class _FakeSession:
-        def __init__(self, phase: str) -> None:
-            self._phase = phase
-
-        def __enter__(self):
-            return self
-
-        def __exit__(self, exc_type, exc, tb) -> bool:
-            return False
-
-        def begin(self) -> _FakeBeginContext:
-            return _FakeBeginContext(self._phase)
-
-    def _fake_session_factory(engine: sa.Engine) -> _FakeSession:
-        return _FakeSession(current_phase["name"])
-
-    def _fake_build_plan(self, session, candidate, *, lock_rows: bool):
+    def _fake_build_plan(self, session: Session, candidate, *, lock_rows: bool):
+        assert session.get_bind() is sqlite_engine
         lock_rows_seen.append((current_phase["name"], lock_rows))
-        return SimpleNamespace(group_row_ids=[str(candidate.row.id)], winner=None, loser_rows=[])
+        return migration_module._ProviderModelGroupPlan(
+            group_row_ids=[str(candidate.row.id)],
+            winner=None,
+            loser_rows=[],
+        )
 
     def _fake_emit_plan(self, plan, *, session, tx_id: str, business_key: dict[str, object]) -> None:
         return None
 
-    def _fake_configure(self, session) -> None:
+    def _fake_configure(self, session: Session) -> None:
+        assert session.get_bind() is sqlite_engine
         configure_calls.append(current_phase["name"])
 
-    monkeypatch.setattr(migration_module, "_session_factory", _fake_session_factory)
     monkeypatch.setattr(migration_module.Migration, "_build_provider_model_group_plan", _fake_build_plan)
     monkeypatch.setattr(migration_module.Migration, "_emit_provider_model_group_plan", _fake_emit_plan)
     monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure)
-
-    dry_migration._process_provider_model_group(candidate, business_key)
-    current_phase["name"] = "apply"
-    apply_migration._process_provider_model_group(candidate, business_key)
+    sa.event.listen(Session, "after_transaction_create", _record_begin)
+    try:
+        dry_migration._process_provider_model_group(candidate, business_key)
+        current_phase["name"] = "apply"
+        apply_migration._process_provider_model_group(candidate, business_key)
+    finally:
+        sa.event.remove(Session, "after_transaction_create", _record_begin)
 
     assert [phase for phase, _ in lock_rows_seen] == ["dry", "apply"]
     assert lock_rows_seen[0][1] == lock_rows_seen[1][1]
@@ -1392,6 +1373,22 @@ def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeou
     sqlite_engine: sa.Engine,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    create_minimal_legacy_model_type_schema(sqlite_engine)
+    created_at = datetime(2025, 1, 1, 12, 0, 0)
+    _insert_load_balancing_model_config(
+        sqlite_engine,
+        row_id="40000000-0000-0000-0000-000000000001",
+        tenant_id="tenant-1",
+        provider_name="openai",
+        model_name="gpt-4o-mini",
+        model_type="text-generation",
+        name="credential",
+        encrypted_config="{}",
+        credential_id="50000000-0000-0000-0000-000000000001",
+        enabled=True,
+        created_at=created_at,
+        updated_at=created_at,
+    )
     output = io.StringIO()
     migration = migration_module.Migration(
         tenant_id="tenant-1",
@@ -1401,37 +1398,18 @@ def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeou
         model_types=(ModelType.LLM,),
         orm_models=(migration_module.LoadBalancingModelConfig,),
     )
-    candidate = migration_module._RowWithRawModelType(
-        row=SimpleNamespace(id="lb-row-1"),
-        raw_model_type="text-generation",
-        canonical_model_type=ModelType.LLM,
-    )
+    candidate = migration._load_load_balancing_model_config_candidates(None)[0]
     lock_timeout_exc = OperationalError("SELECT 1", {}, SimpleNamespace(pgcode="55P03"))
+    transaction_begins = 0
 
-    class _FakeBeginContext:
-        def __enter__(self) -> None:
-            return None
-
-        def __exit__(self, exc_type, exc, tb) -> bool:
-            return False
-
-    class _FakeSession:
-        def __enter__(self):
-            return self
-
-        def __exit__(self, exc_type, exc, tb) -> bool:
-            return False
-
-        def begin(self) -> _FakeBeginContext:
-            return _FakeBeginContext()
-
-    def _fake_session_factory(engine: sa.Engine) -> _FakeSession:
-        return _FakeSession()
+    def _record_begin(session: Session, transaction: SessionTransaction) -> None:
+        nonlocal transaction_begins
+        if session.get_bind() is sqlite_engine and transaction.parent is None:
+            transaction_begins += 1
 
     def _fake_reload(self, session, original_candidate, *, lock_rows: bool):
         raise lock_timeout_exc
 
-    monkeypatch.setattr(migration_module, "_session_factory", _fake_session_factory)
     monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", lambda self, session: None)
     monkeypatch.setattr(
         migration_module.Migration,
@@ -1439,17 +1417,22 @@ def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeou
         _fake_reload,
     )
 
-    migration._process_load_balancing_model_config_row(candidate)
+    sa.event.listen(Session, "after_transaction_create", _record_begin)
+    try:
+        migration._process_load_balancing_model_config_row(candidate)
+    finally:
+        sa.event.remove(Session, "after_transaction_create", _record_begin)
 
     lines = _parse_json_lines(output)
     assert len(lines) == 1
     assert lines[0]["event"] == "lock_timeout_skipped"
     attrs = cast(dict[str, object], lines[0]["attrs"])
     assert attrs["table_name"] == "load_balancing_model_configs"
-    assert attrs["id"] == "lb-row-1"
+    assert attrs["id"] == str(candidate.row.id)
     assert attrs["error"] == str(lock_timeout_exc)
     assert isinstance(attrs["stacktrace"], str)
     assert "OperationalError" in attrs["stacktrace"]
+    assert transaction_begins == 1
 
 
 def test_process_load_balancing_model_config_row_logs_update_after_sql_execution(
@@ -1457,6 +1440,23 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution
     sqlite_engine: sa.Engine,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    create_minimal_legacy_model_type_schema(sqlite_engine)
+    created_at = datetime(2025, 1, 1, 12, 0, 0)
+    row_id = "40000000-0000-0000-0000-000000000002"
+    _insert_load_balancing_model_config(
+        sqlite_engine,
+        row_id=row_id,
+        tenant_id="tenant-1",
+        provider_name="openai",
+        model_name="gpt-4o-mini",
+        model_type="text-generation",
+        name="credential",
+        encrypted_config="{}",
+        credential_id="50000000-0000-0000-0000-000000000002",
+        enabled=True,
+        created_at=created_at,
+        updated_at=created_at,
+    )
     migration = migration_module.Migration(
         tenant_id="tenant-1",
         engine=sqlite_engine,
@@ -1465,42 +1465,33 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution
         model_types=(ModelType.LLM,),
         orm_models=(migration_module.LoadBalancingModelConfig,),
     )
-    candidate = migration_module._RowWithRawModelType(
-        row=SimpleNamespace(id="lb-row-1"),
-        raw_model_type="text-generation",
-        canonical_model_type=ModelType.LLM,
-    )
+    candidate = migration._load_load_balancing_model_config_candidates(None)[0]
     action_log: list[str] = []
 
-    class _FakeBeginContext:
-        def __enter__(self) -> None:
+    def _record_begin(session: Session, transaction: SessionTransaction) -> None:
+        if session.get_bind() is sqlite_engine and transaction.parent is None:
             action_log.append("begin")
 
-        def __exit__(self, exc_type, exc, tb) -> bool:
-            return False
-
-    class _FakeSession:
-        def __enter__(self):
-            return self
-
-        def __exit__(self, exc_type, exc, tb) -> bool:
-            return False
-
-        def begin(self) -> _FakeBeginContext:
-            return _FakeBeginContext()
-
-        def execute(self, stmt) -> None:
+    def _record_sql(
+        connection: sa.Connection,
+        cursor: object,
+        statement: str,
+        parameters: object,
+        context: object,
+        executemany: bool,
+    ) -> None:
+        del connection, cursor, parameters, context, executemany
+        if statement.lstrip().upper().startswith("UPDATE"):
             action_log.append("sql_execute")
 
-    def _fake_session_factory(engine: sa.Engine) -> _FakeSession:
-        return _FakeSession()
-
     def _fake_configure(self, session) -> None:
         action_log.append("configure_lock_timeout")
 
-    def _fake_reload(self, session, original_candidate, *, lock_rows: bool):
+    original_reload = migration_module.Migration._reload_load_balancing_model_config_candidate
+
+    def _record_reload(self, session: Session, original_candidate, *, lock_rows: bool):
         action_log.append(f"reload_candidate:{lock_rows}")
-        return candidate
+        return original_reload(self, session, original_candidate, lock_rows=lock_rows)
 
     def _fake_log_row_updated(self, *args, **kwargs) -> None:
         action_log.append("log_row_updated")
@@ -1508,12 +1499,11 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution
     def _fake_cache_cleanup(self, *, row_id: str, tx_id: str) -> None:
         action_log.append("cache_cleanup")
 
-    monkeypatch.setattr(migration_module, "_session_factory", _fake_session_factory)
     monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure)
     monkeypatch.setattr(
         migration_module.Migration,
         "_reload_load_balancing_model_config_candidate",
-        _fake_reload,
+        _record_reload,
     )
     monkeypatch.setattr(migration_module.Migration, "_log_row_updated", _fake_log_row_updated)
     monkeypatch.setattr(
@@ -1522,7 +1512,13 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution
         _fake_cache_cleanup,
     )
 
-    migration._process_load_balancing_model_config_row(candidate)
+    sa.event.listen(Session, "after_transaction_create", _record_begin)
+    sa.event.listen(sqlite_engine, "before_cursor_execute", _record_sql)
+    try:
+        migration._process_load_balancing_model_config_row(candidate)
+    finally:
+        sa.event.remove(sqlite_engine, "before_cursor_execute", _record_sql)
+        sa.event.remove(Session, "after_transaction_create", _record_begin)
 
     assert action_log == [
         "begin",
@@ -1532,6 +1528,10 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution
         "log_row_updated",
         "cache_cleanup",
     ]
+    with Session(sqlite_engine) as session:
+        persisted = session.get(migration_module.LoadBalancingModelConfig, row_id)
+        assert persisted is not None
+        assert persisted.model_type == ModelType.LLM
 
 
 def test_load_balancing_model_config_cache_delete_failure_logs_stacktrace(

api/tests/unit_tests/commands/test_migrate_legacy_sys_files_workflows.py

@@ -0,0 +1,236 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import click
+import pytest
+
+from commands import migrate_legacy_sys_files_workflows
+from commands import workflow_migration as workflow_migration_commands
+
+
+def test_migrate_legacy_sys_files_workflows_command_passes_batch_options(mocker, capsys):
+    runner = mocker.patch.object(
+        workflow_migration_commands,
+        "run_legacy_sys_files_workflow_migration",
+        return_value=workflow_migration_commands.LegacySysFilesWorkflowMigrationStats(
+            scanned=10,
+            migrated=2,
+            failed=0,
+            batches=1,
+            last_id="workflow-10",
+        ),
+    )
+
+    migrate_legacy_sys_files_workflows.callback(
+        batch_size=200,
+        limit=500,
+        start_after_id="workflow-1",
+        tenant_id="tenant-1",
+        app_id="app-1",
+        dry_run=True,
+    )
+
+    runner.assert_called_once_with(
+        batch_size=200,
+        limit=500,
+        start_after_id="workflow-1",
+        tenant_id="tenant-1",
+        app_id="app-1",
+        dry_run=True,
+    )
+    captured = capsys.readouterr()
+    assert "scanned=10" in captured.out
+    assert "migrated=2" in captured.out
+    assert "last_id=workflow-10" in captured.out
+
+
+def test_migrate_legacy_sys_files_workflows_rejects_non_positive_batch_size():
+    with pytest.raises(click.UsageError, match="batch-size"):
+        migrate_legacy_sys_files_workflows.callback(
+            batch_size=0,
+            limit=None,
+            start_after_id=None,
+            tenant_id=None,
+            app_id=None,
+            dry_run=False,
+        )
+
+
+def test_migrate_legacy_sys_files_workflows_rejects_non_positive_limit():
+    with pytest.raises(click.UsageError, match="limit"):
+        migrate_legacy_sys_files_workflows.callback(
+            batch_size=100,
+            limit=0,
+            start_after_id=None,
+            tenant_id=None,
+            app_id=None,
+            dry_run=False,
+        )
+
+
+def test_build_legacy_sys_files_workflow_query_uses_keyset_pagination():
+    stmt = workflow_migration_commands._build_legacy_sys_files_workflow_query(
+        start_after_id="workflow-1",
+        batch_size=200,
+        tenant_id="tenant-1",
+        app_id="app-1",
+    )
+    compiled = str(stmt.compile(compile_kwargs={"literal_binds": True}))
+
+    assert "workflows.id > 'workflow-1'" in compiled
+    assert "workflows.tenant_id = 'tenant-1'" in compiled
+    assert "workflows.app_id = 'app-1'" in compiled
+    assert "ORDER BY workflows.id" in compiled
+    assert "LIMIT 200" in compiled
+    assert "workflows.environment_variables" not in compiled
+
+
+def test_migrate_legacy_sys_files_workflow_batch_dry_run_rolls_back():
+    migrated_workflow = MagicMock()
+    migrated_workflow.id = "workflow-1"
+    migrated_workflow.migrate_legacy_sys_files_graph_in_place.return_value = True
+    untouched_workflow = MagicMock()
+    untouched_workflow.id = "workflow-2"
+    untouched_workflow.migrate_legacy_sys_files_graph_in_place.return_value = False
+    session = MagicMock()
+    session.scalars.return_value.all.return_value = [migrated_workflow, untouched_workflow]
+
+    stats = workflow_migration_commands._migrate_legacy_sys_files_workflow_batch(
+        session=session,
+        start_after_id=None,
+        batch_size=200,
+        tenant_id=None,
+        app_id=None,
+        dry_run=True,
+    )
+
+    assert stats.scanned == 2
+    assert stats.migrated == 1
+    assert stats.failed == 0
+    assert stats.last_id == "workflow-2"
+    session.rollback.assert_called_once()
+    session.commit.assert_not_called()
+
+
+def test_migrate_legacy_sys_files_workflow_batch_commits_and_counts_failures(caplog):
+    migrated_workflow = MagicMock()
+    migrated_workflow.id = "workflow-1"
+    migrated_workflow.migrate_legacy_sys_files_graph_in_place.return_value = True
+    failing_workflow = MagicMock()
+    failing_workflow.id = "workflow-2"
+    failing_workflow.migrate_legacy_sys_files_graph_in_place.side_effect = RuntimeError("boom")
+    session = MagicMock()
+    session.scalars.return_value.all.return_value = [migrated_workflow, failing_workflow]
+
+    stats = workflow_migration_commands._migrate_legacy_sys_files_workflow_batch(
+        session=session,
+        start_after_id=None,
+        batch_size=200,
+        tenant_id=None,
+        app_id=None,
+        dry_run=False,
+    )
+
+    assert stats.scanned == 2
+    assert stats.migrated == 1
+    assert stats.failed == 1
+    assert stats.last_id == "workflow-2"
+    assert "Failed to migrate legacy" in caplog.text
+    session.commit.assert_called_once()
+    session.rollback.assert_not_called()
+
+
+def test_run_legacy_sys_files_workflow_migration_uses_keyset_batches(mocker):
+    session_maker = MagicMock()
+    sessions = [MagicMock(), MagicMock()]
+    session_maker.side_effect = sessions
+    mocker.patch.object(workflow_migration_commands, "sessionmaker", return_value=session_maker)
+    mocker.patch.object(workflow_migration_commands, "db", SimpleNamespace(engine=object()))
+    migrate_batch = mocker.patch.object(
+        workflow_migration_commands,
+        "_migrate_legacy_sys_files_workflow_batch",
+        side_effect=[
+            workflow_migration_commands.LegacySysFilesWorkflowMigrationStats(
+                scanned=2,
+                migrated=1,
+                failed=0,
+                last_id="workflow-2",
+            ),
+            workflow_migration_commands.LegacySysFilesWorkflowMigrationStats(
+                scanned=1,
+                migrated=1,
+                failed=0,
+                last_id="workflow-3",
+            ),
+        ],
+    )
+
+    stats = workflow_migration_commands.run_legacy_sys_files_workflow_migration(
+        batch_size=2,
+        limit=3,
+        start_after_id="workflow-0",
+        tenant_id="tenant-1",
+        app_id="app-1",
+        dry_run=True,
+    )
+
+    assert stats.scanned == 3
+    assert stats.migrated == 2
+    assert stats.batches == 2
+    assert stats.last_id == "workflow-3"
+    assert migrate_batch.call_args_list[0].kwargs["start_after_id"] == "workflow-0"
+    assert migrate_batch.call_args_list[0].kwargs["batch_size"] == 2
+    assert migrate_batch.call_args_list[1].kwargs["start_after_id"] == "workflow-2"
+    assert migrate_batch.call_args_list[1].kwargs["batch_size"] == 1
+
+
+def test_run_legacy_sys_files_workflow_migration_stops_on_empty_batch(mocker):
+    session_maker = MagicMock(return_value=MagicMock())
+    mocker.patch.object(workflow_migration_commands, "sessionmaker", return_value=session_maker)
+    mocker.patch.object(workflow_migration_commands, "db", SimpleNamespace(engine=object()))
+    mocker.patch.object(
+        workflow_migration_commands,
+        "_migrate_legacy_sys_files_workflow_batch",
+        return_value=workflow_migration_commands.LegacySysFilesWorkflowMigrationStats(scanned=0),
+    )
+
+    stats = workflow_migration_commands.run_legacy_sys_files_workflow_migration(
+        batch_size=2,
+        limit=None,
+        start_after_id=None,
+        tenant_id=None,
+        app_id=None,
+        dry_run=False,
+    )
+
+    assert stats.scanned == 0
+    assert stats.batches == 0
+
+
+def test_run_legacy_sys_files_workflow_migration_stops_on_short_batch(mocker):
+    session_maker = MagicMock(return_value=MagicMock())
+    mocker.patch.object(workflow_migration_commands, "sessionmaker", return_value=session_maker)
+    mocker.patch.object(workflow_migration_commands, "db", SimpleNamespace(engine=object()))
+    migrate_batch = mocker.patch.object(
+        workflow_migration_commands,
+        "_migrate_legacy_sys_files_workflow_batch",
+        return_value=workflow_migration_commands.LegacySysFilesWorkflowMigrationStats(
+            scanned=1,
+            migrated=1,
+            failed=0,
+            last_id="workflow-1",
+        ),
+    )
+
+    stats = workflow_migration_commands.run_legacy_sys_files_workflow_migration(
+        batch_size=2,
+        limit=None,
+        start_after_id=None,
+        tenant_id=None,
+        app_id=None,
+        dry_run=False,
+    )
+
+    assert stats.scanned == 1
+    assert stats.batches == 1
+    migrate_batch.assert_called_once()

api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py

@@ -14,7 +14,7 @@ from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 import commands
-from commands import system as system_commands
+from commands import workspace as workspace_commands
 from core.tools.entities.tool_entities import ApiProviderSchemaType
 from graphon.model_runtime.entities.model_entities import ModelType
 from models import Tenant
@@ -83,11 +83,11 @@ def _encrypted_rows(tenant_id: str, *, suffix: str = "1") -> tuple[object, ...]:
 
 
 def _bind_command_to_sqlite(monkeypatch: pytest.MonkeyPatch, session: Session) -> None:
-    monkeypatch.setattr(system_commands, "db", SimpleNamespace(engine=session.get_bind()))
+    monkeypatch.setattr(workspace_commands, "db", SimpleNamespace(engine=session.get_bind()))
 
 
 def test_reset_aborts_when_not_self_hosted(monkeypatch, capsys):
-    monkeypatch.setattr(system_commands.dify_config, "EDITION", "CLOUD")
+    monkeypatch.setattr(workspace_commands.dify_config, "EDITION", "CLOUD")
 
     exit_code = _invoke_reset()
     captured = capsys.readouterr()
@@ -106,8 +106,8 @@ def test_reset_purges_provider_and_tool_tables_for_each_tenant(
 ) -> None:
     """The command must purge LLM provider rows AND every tool provider table
     that stores ciphertext encrypted under the tenant key (#35396)."""
-    monkeypatch.setattr(system_commands.dify_config, "EDITION", "SELF_HOSTED")
-    monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}")
+    monkeypatch.setattr(workspace_commands.dify_config, "EDITION", "SELF_HOSTED")
+    monkeypatch.setattr(workspace_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}")
     _bind_command_to_sqlite(monkeypatch, sqlite_session)
 
     tenant = _tenant(TENANT_ID)
@@ -146,8 +146,8 @@ def test_reset_purges_provider_and_tool_tables_for_each_tenant(
 )
 def test_reset_iterates_all_tenants(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None:
     """Multi-tenant deployments must purge every tenant, not just the first."""
-    monkeypatch.setattr(system_commands.dify_config, "EDITION", "SELF_HOSTED")
-    monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}")
+    monkeypatch.setattr(workspace_commands.dify_config, "EDITION", "SELF_HOSTED")
+    monkeypatch.setattr(workspace_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}")
 
     _bind_command_to_sqlite(monkeypatch, sqlite_session)
     tenant_ids = [f"11111111-1111-1111-1111-{index:012d}" for index in range(3)]

api/tests/unit_tests/commands/test_upgrade_db.py

@@ -4,7 +4,7 @@ import types
 from unittest.mock import MagicMock
 
 import commands
-from commands import system as system_commands
+from commands import database as database_commands
 from libs.db_migration_lock import LockNotOwnedError, RedisError
 
 HEARTBEAT_WAIT_TIMEOUT_SECONDS = 5.0
@@ -25,11 +25,11 @@ def _invoke_upgrade_db() -> int:
 
 
 def test_upgrade_db_skips_when_lock_not_acquired(monkeypatch, capsys):
-    monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 1234)
+    monkeypatch.setattr(database_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 1234)
 
     lock = MagicMock()
     lock.acquire.return_value = False
-    system_commands.redis_client.lock.return_value = lock
+    database_commands.redis_client.lock.return_value = lock
 
     exit_code = _invoke_upgrade_db()
     captured = capsys.readouterr()
@@ -37,18 +37,20 @@ def test_upgrade_db_skips_when_lock_not_acquired(monkeypatch, capsys):
     assert exit_code == 0
     assert "Database migration skipped" in captured.out
 
-    system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=1234, thread_local=False)
+    database_commands.redis_client.lock.assert_called_once_with(
+        name="db_upgrade_lock", timeout=1234, thread_local=False
+    )
     lock.acquire.assert_called_once_with(blocking=False)
     lock.release.assert_not_called()
 
 
 def test_upgrade_db_failure_not_masked_by_lock_release(monkeypatch, capsys):
-    monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 321)
+    monkeypatch.setattr(database_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 321)
 
     lock = MagicMock()
     lock.acquire.return_value = True
     lock.release.side_effect = LockNotOwnedError("simulated")
-    system_commands.redis_client.lock.return_value = lock
+    database_commands.redis_client.lock.return_value = lock
 
     def _upgrade():
         raise RuntimeError("boom")
@@ -61,18 +63,18 @@ def test_upgrade_db_failure_not_masked_by_lock_release(monkeypatch, capsys):
     assert exit_code == 1
     assert "Database migration failed: boom" in captured.out
 
-    system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=321, thread_local=False)
+    database_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=321, thread_local=False)
     lock.acquire.assert_called_once_with(blocking=False)
     lock.release.assert_called_once()
 
 
 def test_upgrade_db_success_ignores_lock_not_owned_on_release(monkeypatch, capsys):
-    monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 999)
+    monkeypatch.setattr(database_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 999)
 
     lock = MagicMock()
     lock.acquire.return_value = True
     lock.release.side_effect = LockNotOwnedError("simulated")
-    system_commands.redis_client.lock.return_value = lock
+    database_commands.redis_client.lock.return_value = lock
 
     _install_fake_flask_migrate(monkeypatch, lambda: None)
 
@@ -82,7 +84,7 @@ def test_upgrade_db_success_ignores_lock_not_owned_on_release(monkeypatch, capsy
     assert exit_code == 0
     assert "Database migration successful!" in captured.out
 
-    system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=999, thread_local=False)
+    database_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=999, thread_local=False)
     lock.acquire.assert_called_once_with(blocking=False)
     lock.release.assert_called_once()
 
@@ -93,11 +95,11 @@ def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys):
     """
 
     # Use a small TTL so the heartbeat interval triggers quickly.
-    monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3)
+    monkeypatch.setattr(database_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3)
 
     lock = MagicMock()
     lock.acquire.return_value = True
-    system_commands.redis_client.lock.return_value = lock
+    database_commands.redis_client.lock.return_value = lock
 
     renewed = threading.Event()
 
@@ -121,11 +123,11 @@ def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys):
 
 def test_upgrade_db_ignores_reacquire_errors(monkeypatch, capsys):
     # Use a small TTL so heartbeat runs during the upgrade call.
-    monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3)
+    monkeypatch.setattr(database_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3)
 
     lock = MagicMock()
     lock.acquire.return_value = True
-    system_commands.redis_client.lock.return_value = lock
+    database_commands.redis_client.lock.return_value = lock
 
     attempted = threading.Event()
 

api/tests/unit_tests/commands/test_workspace.py

@@ -0,0 +1,49 @@
+from unittest.mock import MagicMock
+
+from commands import reset_encrypt_key_pair
+from commands import workspace as workspace_commands
+
+
+def test_reset_encrypt_key_pair_skips_non_self_hosted(monkeypatch, capsys):
+    monkeypatch.setattr(workspace_commands.dify_config, "EDITION", "CLOUD")
+
+    reset_encrypt_key_pair.callback()
+
+    captured = capsys.readouterr()
+    assert "only for SELF_HOSTED" in captured.out
+
+
+def test_reset_encrypt_key_pair_rotates_keys_and_removes_custom_provider_data(monkeypatch, capsys):
+    monkeypatch.setattr(workspace_commands.dify_config, "EDITION", "SELF_HOSTED")
+    monkeypatch.setattr(workspace_commands, "generate_key_pair", lambda tenant_id: f"public-key-{tenant_id}")
+    tenant = MagicMock()
+    tenant.id = "tenant-1"
+    session = MagicMock()
+    session.scalars.return_value.all.return_value = [tenant]
+    session_manager = MagicMock()
+    session_manager.begin.return_value.__enter__.return_value = session
+    monkeypatch.setattr(workspace_commands, "sessionmaker", lambda *args, **kwargs: session_manager)
+    monkeypatch.setattr(workspace_commands, "db", MagicMock(engine=object()))
+
+    reset_encrypt_key_pair.callback()
+
+    assert tenant.encrypt_public_key == "public-key-tenant-1"
+    assert session.execute.call_count == 5
+    captured = capsys.readouterr()
+    assert "tenant-1 has been reset" in captured.out
+
+
+def test_reset_encrypt_key_pair_stops_when_workspace_record_is_missing(monkeypatch, capsys):
+    monkeypatch.setattr(workspace_commands.dify_config, "EDITION", "SELF_HOSTED")
+    session = MagicMock()
+    session.scalars.return_value.all.return_value = [None]
+    session_manager = MagicMock()
+    session_manager.begin.return_value.__enter__.return_value = session
+    monkeypatch.setattr(workspace_commands, "sessionmaker", lambda *args, **kwargs: session_manager)
+    monkeypatch.setattr(workspace_commands, "db", MagicMock(engine=object()))
+
+    reset_encrypt_key_pair.callback()
+
+    session.execute.assert_not_called()
+    captured = capsys.readouterr()
+    assert "No workspaces found" in captured.out

api/tests/unit_tests/configs/test_env_consistency.py

@@ -4,6 +4,7 @@ from dotenv import dotenv_values
 
 BASE_API_AND_DOCKER_CONFIG_SET_DIFF: frozenset[str] = frozenset(
     (
+        "AGENT_BACKEND_API_TOKEN",
         "APP_MAX_EXECUTION_TIME",
         "BATCH_UPLOAD_LIMIT",
         "CELERY_BEAT_SCHEDULER_TIME",
@@ -43,6 +44,7 @@ BASE_API_AND_DOCKER_CONFIG_SET_DIFF: frozenset[str] = frozenset(
 
 BASE_API_AND_DOCKER_COMPOSE_CONFIG_SET_DIFF: frozenset[str] = frozenset(
     (
+        "AGENT_BACKEND_API_TOKEN",
         "BATCH_UPLOAD_LIMIT",
         "CELERY_BEAT_SCHEDULER_TIME",
         "HTTP_REQUEST_MAX_CONNECT_TIMEOUT",

api/tests/unit_tests/conftest.py

@@ -1,11 +1,13 @@
 import os
+import shutil
 from collections.abc import Iterator
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import pytest
 from flask import Flask
 from sqlalchemy import create_engine
-from sqlalchemy.engine import Engine
+from sqlalchemy.engine import URL, Engine
 from sqlalchemy.orm import Session, sessionmaker
 
 # Getting the absolute path of the current file's directory
@@ -35,7 +37,7 @@ os.environ.setdefault("OPENDAL_SCHEME", "fs")
 os.environ.setdefault("OPENDAL_FS_ROOT", "/tmp/dify-storage")
 os.environ.setdefault("STORAGE_TYPE", "opendal")
 
-from core.db.session_factory import configure_session_factory, session_factory
+import core.db.session_factory as session_factory_module
 from extensions import ext_redis
 from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
 from models.base import TypeBase
@@ -111,42 +113,70 @@ def reset_secret_key():
         dify_config.SECRET_KEY = original
 
 
-@pytest.fixture(scope="session")
-def _unit_test_engine():
-    engine = create_engine("sqlite:///:memory:")
-    yield engine
-    engine.dispose()
-
-
 @pytest.fixture
-def sqlite_engine() -> Iterator[Engine]:
-    """Create an isolated in-memory SQLite engine for tests that need a disposable database."""
+def _sqlite_engine(_sqlite_database_template: Path, tmp_path: Path) -> Iterator[Engine]:
+    """Create an engine over a pristine per-test copy of the SQLite schema."""
+
+    database_path = tmp_path / "unit-tests.sqlite3"
+    shutil.copyfile(_sqlite_database_template, database_path)
+    engine = create_engine(URL.create("sqlite", database=str(database_path)))
 
-    engine = create_engine("sqlite:///:memory:")
     try:
         yield engine
     finally:
         engine.dispose()
+        database_path.unlink(missing_ok=True)
 
 
-@pytest.fixture
-def sqlite_session(request: pytest.FixtureRequest, sqlite_engine: Engine) -> Iterator[Session]:
-    """Yield a SQLite session after creating the model tables passed through ``request.param``."""
+@pytest.fixture(scope="session")
+def _sqlite_database_template(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    """Create one empty full-schema SQLite database per pytest worker."""
 
-    models: tuple[type[TypeBase], ...] = request.param
-    tables = [model.metadata.tables[model.__tablename__] for model in models]
-    TypeBase.metadata.create_all(sqlite_engine, tables=tables)
-    session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
-    with session_factory() as session:
-        yield session
+    database_path = tmp_path_factory.mktemp("sqlite-template") / "unit-tests.sqlite3"
+    engine = create_engine(URL.create("sqlite", database=str(database_path)))
+    try:
+        TypeBase.metadata.create_all(engine)
+    finally:
+        engine.dispose()
+    return database_path
 
 
 @pytest.fixture(autouse=True)
-def _configure_session_factory(_unit_test_engine):
-    try:
-        session_factory.get_session_maker()
-    except RuntimeError:
-        configure_session_factory(_unit_test_engine, expire_on_commit=False)
+def _sqlite_session_factory(
+    _sqlite_engine: Engine,
+    monkeypatch: pytest.MonkeyPatch,
+) -> sessionmaker[Session]:
+    """Bind all unit-test Sessions to the pristine full-schema SQLite database."""
+
+    factory = sessionmaker(bind=_sqlite_engine, expire_on_commit=False)
+    monkeypatch.setattr(session_factory_module, "_session_maker", factory)
+    return factory
+
+
+@pytest.fixture
+def sqlite_engine(_sqlite_engine: Engine) -> Engine:
+    """Expose the pristine full-schema SQLite engine to tests."""
+
+    return _sqlite_engine
+
+
+@pytest.fixture
+def sqlite_session_factory(_sqlite_session_factory: sessionmaker[Session]) -> sessionmaker[Session]:
+    """Expose the shared SQLite session factory to tests."""
+
+    return _sqlite_session_factory
+
+
+@pytest.fixture
+def sqlite_session(_sqlite_session_factory: sessionmaker[Session]) -> Iterator[Session]:
+    """Yield a session over the pristine full-schema SQLite database.
+
+    Legacy indirect model parameters remain accepted by pytest but are ignored.
+    Remove those decorators as their test files receive individual review.
+    """
+
+    with _sqlite_session_factory() as session:
+        yield session
 
 
 def persist_service_api_tenant_owner(session: Session, tenant: Tenant, owner: Account) -> TenantAccountJoin:

api/tests/unit_tests/controllers/console/app/test_app_response_models.py

@@ -708,6 +708,121 @@ def test_app_list_api_attaches_permission_keys(app, app_module):
     assert resp["data"][0]["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"]
 
 
+def test_recent_app_list_api_returns_only_home_card_fields(app, app_module):
+    method = app_module.RecentAppListApi.get
+    while hasattr(method, "__wrapped__"):
+        method = method.__wrapped__
+
+    recent_app = SimpleNamespace(
+        id="app-1",
+        name="Recent App",
+        icon_type="emoji",
+        icon="🚀",
+        icon_background="#FFFFFF",
+        mode="chat",
+        author_name="Recent Author",
+        updated_at=_ts(15),
+        maintainer="acct-1",
+    )
+    get_recent_apps = MagicMock(return_value=[recent_app])
+
+    with app.test_request_context("/apps/recent?limit=8"):
+        with pytest.MonkeyPatch.context() as monkeypatch:
+            monkeypatch.setattr(dify_config, "RBAC_ENABLED", False)
+            monkeypatch.setattr(app_module.AppService, "get_recent_apps", get_recent_apps)
+            monkeypatch.setattr(
+                app_module.enterprise_rbac_service.RBACService.MyPermissions,
+                "get",
+                lambda tenant_id, account_id, session: app_module.enterprise_rbac_service.MyPermissionsResponse(
+                    app=app_module.enterprise_rbac_service.ResourcePermissionSnapshot(
+                        overrides=[
+                            app_module.enterprise_rbac_service.ResourcePermissionKeys(
+                                resource_id="app-1",
+                                permission_keys=["app.acl.monitor"],
+                            )
+                        ]
+                    )
+                ),
+            )
+
+            resp, status = method(app_module.RecentAppListApi(), "tenant-1", "acct-1", MagicMock())
+
+    assert status == 200
+    assert resp == {
+        "data": [
+            {
+                "id": "app-1",
+                "name": "Recent App",
+                "icon_type": "emoji",
+                "icon": "🚀",
+                "icon_background": "#FFFFFF",
+                "mode": "chat",
+                "author_name": "Recent Author",
+                "updated_at": int(_ts(15).timestamp()),
+                "permission_keys": ["app.acl.monitor"],
+                "maintainer": "acct-1",
+                "icon_url": None,
+            }
+        ]
+    }
+    params = get_recent_apps.call_args.args[2]
+    assert params.limit == 8
+    assert "total" not in resp
+    assert "description" not in resp["data"][0]
+    assert "tags" not in resp["data"][0]
+    assert "workflow" not in resp["data"][0]
+
+
+@pytest.mark.parametrize("mode", ["channel", "rag-pipeline", "agent"])
+def test_recent_app_response_rejects_non_home_app_modes(app_module, mode: str) -> None:
+    with pytest.raises(ValidationError):
+        app_module.RecentAppResponse.model_validate(
+            {
+                "id": "app-1",
+                "name": "Recent App",
+                "mode": mode,
+                "updated_at": _ts(),
+            }
+        )
+
+
+def test_recent_app_list_api_applies_rbac_visibility_filter(app, app_module):
+    method = app_module.RecentAppListApi.get
+    while hasattr(method, "__wrapped__"):
+        method = method.__wrapped__
+
+    get_recent_apps = MagicMock(return_value=[])
+    with app.test_request_context("/apps/recent"):
+        with pytest.MonkeyPatch.context() as monkeypatch:
+            monkeypatch.setattr(dify_config, "RBAC_ENABLED", True)
+            monkeypatch.setattr(app_module.AppService, "get_recent_apps", get_recent_apps)
+            monkeypatch.setattr(
+                app_module.enterprise_rbac_service.RBACService.MyPermissions,
+                "get",
+                lambda tenant_id, account_id, session: app_module.enterprise_rbac_service.MyPermissionsResponse(
+                    workspace=app_module.enterprise_rbac_service.WorkspacePermissionSnapshot(
+                        permission_keys=["app.create_and_management"]
+                    )
+                ),
+            )
+            monkeypatch.setattr(
+                app_module.enterprise_rbac_service.RBACService.AppAccess,
+                "whitelist_resources",
+                lambda tenant_id, account_id: SimpleNamespace(
+                    unrestricted=False,
+                    resource_ids=["app-shared"],
+                ),
+            )
+
+            resp, status = method(app_module.RecentAppListApi(), "tenant-1", "acct-1", MagicMock())
+
+    assert status == 200
+    assert resp == {"data": []}
+    params = get_recent_apps.call_args.args[2]
+    assert params.accessible_app_ids == ["app-shared"]
+    assert params.include_own_apps is True
+
+
 def test_app_list_api_limits_to_apps_created_by_current_user_without_view_permission(app, app_module):
     method = app_module.AppListApi.get
     while hasattr(method, "__wrapped__"):

api/tests/unit_tests/controllers/console/datasets/test_data_source.py

@@ -1,18 +1,22 @@
 from __future__ import annotations
 
 import inspect
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from datetime import UTC, datetime
-from typing import cast
+from typing import Literal, cast
 from unittest.mock import MagicMock, PropertyMock, patch
-from uuid import uuid4
+from uuid import UUID
 
 import pytest
 from flask import Flask
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+from werkzeug.exceptions import NotFound
 
 from controllers.console.datasets import data_source as module
 from controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi
 from models import Account, DataSourceOauthBinding
+from models.engine import db
 
 ControllerMethod = Callable[..., tuple[dict[str, object], int]]
 
@@ -22,10 +26,15 @@ def unwrap(func: object) -> ControllerMethod:
 
 
 @pytest.fixture
-def flask_app() -> Flask:
+def flask_app() -> Iterator[Flask]:
     app = Flask(__name__)
     app.config["TESTING"] = True
-    return app
+    app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
+    db.init_app(app)
+
+    with app.app_context():
+        DataSourceOauthBinding.__table__.create(db.engine)
+        yield app
 
 
 @pytest.fixture
@@ -35,9 +44,13 @@ def current_user() -> Account:
     return account
 
 
-def test_get_data_source_integrates_serializes_orm_binding(flask_app: Flask) -> None:
+TENANT_ID = "11111111-1111-1111-1111-111111111111"
+BINDING_ID = "22222222-2222-2222-2222-222222222222"
+
+
+def _add_binding(session: Session, *, disabled: bool) -> DataSourceOauthBinding:
     binding = DataSourceOauthBinding(
-        tenant_id="tenant-1",
+        tenant_id=TENANT_ID,
         access_token="token",
         provider="notion",
         source_info={
@@ -55,24 +68,31 @@ def test_get_data_source_integrates_serializes_orm_binding(flask_app: Flask) ->
                 }
             ],
         },
+        disabled=disabled,
     )
-    binding.id = "binding-1"
+    binding.id = BINDING_ID
     binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
-    binding.disabled = False
+    session.add(binding)
+    session.commit()
+    return binding
 
-    with (
-        flask_app.test_request_context("/"),
-        patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [binding])),
-    ):
-        response, status = unwrap(DataSourceApi().get)(DataSourceApi(), "tenant-1")
+
+def test_get_data_source_integrates_serializes_orm_binding(
+    flask_app: Flask,
+) -> None:
+    binding = _add_binding(db.session, disabled=False)
+    expected_created_at = int(binding.created_at.timestamp())
+
+    with flask_app.test_request_context("/"):
+        response, status = unwrap(DataSourceApi().get)(DataSourceApi(), TENANT_ID)
 
     assert status == 200
     assert response == {
         "data": [
             {
-                "id": "binding-1",
+                "id": BINDING_ID,
                 "provider": "notion",
-                "created_at": 1779670923,
+                "created_at": expected_created_at,
                 "is_bound": True,
                 "disabled": False,
                 "source_info": {
@@ -96,34 +116,75 @@ def test_get_data_source_integrates_serializes_orm_binding(flask_app: Flask) ->
     }
 
 
-def test_get_data_source_integrates_preserves_empty_list_when_no_binding(flask_app: Flask) -> None:
-    with (
-        flask_app.test_request_context("/"),
-        patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [])),
-    ):
-        response, status = unwrap(DataSourceApi().get)(DataSourceApi(), "tenant-1")
+def test_get_data_source_integrates_preserves_empty_list_when_no_binding(
+    flask_app: Flask,
+) -> None:
+    with flask_app.test_request_context("/"):
+        response, status = unwrap(DataSourceApi().get)(DataSourceApi(), TENANT_ID)
 
     assert status == 200
     assert response == {"data": []}
 
 
-def test_patch_data_source_binding_uses_injected_session(flask_app: Flask) -> None:
-    binding = MagicMock(disabled=True)
-    session = MagicMock()
-    session.scalar.return_value = binding
+@pytest.mark.parametrize(
+    ("disabled", "action", "expected_disabled"),
+    [(True, "enable", False), (False, "disable", True)],
+)
+@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True)
+def test_patch_data_source_binding_updates_state(
+    flask_app: Flask,
+    sqlite_session: Session,
+    disabled: bool,
+    action: Literal["enable", "disable"],
+    expected_disabled: bool,
+) -> None:
+    _add_binding(sqlite_session, disabled=disabled)
+    sqlite_session.expunge_all()
 
     with flask_app.test_request_context("/"):
-        response, status = unwrap(DataSourceApi().patch)(DataSourceApi(), session, "tenant-1", uuid4(), "enable")
+        response, status = unwrap(DataSourceApi().patch)(
+            DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), action
+        )
 
+    sqlite_session.flush()
+    sqlite_session.expire_all()
+    binding = sqlite_session.scalar(select(DataSourceOauthBinding).where(DataSourceOauthBinding.id == BINDING_ID))
     assert status == 200
     assert response == {"result": "success"}
-    assert binding.disabled is False
-    session.scalar.assert_called_once()
-    session.add.assert_not_called()
-    session.commit.assert_not_called()
+    assert binding is not None
+    assert binding.disabled is expected_disabled
 
 
-def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask, current_user: Account) -> None:
+@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True)
+def test_patch_data_source_binding_rejects_unknown_binding(
+    flask_app: Flask,
+    sqlite_session: Session,
+) -> None:
+    with flask_app.test_request_context("/"), pytest.raises(NotFound, match="Data source binding not found"):
+        unwrap(DataSourceApi().patch)(DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), "enable")
+
+
+@pytest.mark.parametrize(("disabled", "action"), [(False, "enable"), (True, "disable")])
+@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True)
+def test_patch_data_source_binding_rejects_current_state(
+    flask_app: Flask,
+    sqlite_session: Session,
+    disabled: bool,
+    action: Literal["enable", "disable"],
+) -> None:
+    _add_binding(sqlite_session, disabled=disabled)
+    sqlite_session.expunge_all()
+
+    with flask_app.test_request_context("/"), pytest.raises(ValueError):
+        unwrap(DataSourceApi().patch)(DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), action)
+
+
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_notion_pre_import_pages_serializes_frontend_list_shape(
+    flask_app: Flask,
+    current_user: Account,
+    sqlite_session: Session,
+) -> None:
     page = MagicMock(
         page_id="page-1",
         page_name="Page",
@@ -145,8 +206,6 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask
         get_online_document_pages=MagicMock(return_value=iter([online_document_message])),
         datasource_provider_type=MagicMock(return_value="online_document"),
     )
-    session = MagicMock()
-
     with (
         flask_app.test_request_context("/?credential_id=credential-1"),
         patch.object(
@@ -158,7 +217,7 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask
         patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime),
     ):
         response, status = unwrap(DataSourceNotionListApi().get)(
-            DataSourceNotionListApi(), session, "tenant-1", current_user
+            DataSourceNotionListApi(), sqlite_session, "tenant-1", current_user
         )
 
     assert status == 200
@@ -183,3 +242,38 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask
     }
     runtime.get_online_document_pages.assert_called_once()
     assert runtime.get_online_document_pages.call_args.kwargs["datasource_parameters"] == {}
+
+
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_notion_pre_import_pages_rejects_missing_credential(
+    flask_app: Flask,
+    current_user: Account,
+    sqlite_session: Session,
+) -> None:
+    with (
+        flask_app.test_request_context("/?credential_id=credential-1"),
+        patch.object(module.DatasourceProviderService, "get_datasource_credentials", return_value=None),
+        pytest.raises(NotFound, match="Credential not found"),
+    ):
+        unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi(), sqlite_session, TENANT_ID, current_user)
+
+
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_notion_pre_import_pages_rejects_non_notion_dataset(
+    flask_app: Flask,
+    current_user: Account,
+    sqlite_session: Session,
+) -> None:
+    dataset = MagicMock(data_source_type="other_type")
+
+    with (
+        flask_app.test_request_context("/?credential_id=credential-1&dataset_id=dataset-1"),
+        patch.object(
+            module.DatasourceProviderService,
+            "get_datasource_credentials",
+            return_value={"token": "token"},
+        ),
+        patch.object(module.DatasetService, "get_dataset", return_value=dataset),
+        pytest.raises(ValueError, match="Dataset is not notion type"),
+    ):
+        unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi(), sqlite_session, TENANT_ID, current_user)

api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py

@@ -0,0 +1,171 @@
+"""Unit tests for controllers.console.datasets.data_source Notion endpoints."""
+
+from __future__ import annotations
+
+import inspect
+from unittest.mock import MagicMock, patch
+
+import pytest
+from flask import Flask
+from sqlalchemy.orm import Session
+from werkzeug.exceptions import NotFound
+
+from controllers.console.datasets.data_source import (
+    DataSourceNotionDatasetSyncApi,
+    DataSourceNotionDocumentSyncApi,
+    DataSourceNotionIndexingEstimateApi,
+    DataSourceNotionPreviewApi,
+)
+from core.rag.index_processor.constant.index_type import IndexStructureType
+from models import Account
+
+
+@pytest.fixture
+def current_user() -> Account:
+    account = Account(name="Test User", email="u1@example.com")
+    account.id = "u1"
+    return account
+
+
+class TestDataSourceNotionPreviewApi:
+    def test_get_preview_success(self, app: Flask) -> None:
+        api = DataSourceNotionPreviewApi()
+        method = inspect.unwrap(api.get)
+
+        extractor = MagicMock(extract=lambda: [MagicMock(page_content="hello")])
+
+        with (
+            app.test_request_context("/?credential_id=c1"),
+            patch(
+                "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
+                return_value={"integration_secret": "t"},
+            ),
+            patch(
+                "controllers.console.datasets.data_source.NotionExtractor",
+                return_value=extractor,
+            ),
+        ):
+            response, status = method(api, "tenant-1", "p1", "page")
+
+        assert status == 200
+
+
+class TestDataSourceNotionIndexingEstimateApi:
+    @pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+    def test_post_indexing_estimate_success(self, app: Flask, sqlite_session: Session) -> None:
+        api = DataSourceNotionIndexingEstimateApi()
+        method = inspect.unwrap(api.post)
+
+        empty_rules: dict[str, object] = {}
+        payload: dict[str, object] = {
+            "notion_info_list": [
+                {
+                    "workspace_id": "w1",
+                    "credential_id": "c1",
+                    "pages": [{"page_id": "p1", "type": "page"}],
+                }
+            ],
+            "process_rule": {"rules": empty_rules},
+            "doc_form": IndexStructureType.PARAGRAPH_INDEX,
+            "doc_language": "English",
+        }
+
+        with (
+            app.test_request_context("/", method="POST", json=payload, headers={"Content-Type": "application/json"}),
+            patch(
+                "controllers.console.datasets.data_source.DocumentService.estimate_args_validate",
+            ),
+            patch(
+                "controllers.console.datasets.data_source.IndexingRunner.indexing_estimate",
+                return_value=MagicMock(model_dump=lambda: {"total_pages": 1}),
+            ),
+        ):
+            response, status = method(api, sqlite_session, "tenant-1")
+
+        assert status == 200
+
+
+class TestDataSourceNotionDatasetSyncApi:
+    @pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+    def test_get_success(self, app: Flask, sqlite_session: Session) -> None:
+        api = DataSourceNotionDatasetSyncApi()
+        method = inspect.unwrap(api.get)
+
+        with (
+            app.test_request_context("/"),
+            patch(
+                "controllers.console.datasets.data_source.DatasetService.get_dataset",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "controllers.console.datasets.data_source.DocumentService.get_document_by_dataset_id",
+                return_value=[MagicMock(id="d1")],
+            ),
+            patch(
+                "controllers.console.datasets.data_source.document_indexing_sync_task.delay",
+                return_value=None,
+            ),
+        ):
+            response, status = method(api, sqlite_session, "ds-1")
+
+        assert status == 200
+
+    @pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+    def test_get_dataset_not_found(self, app: Flask, sqlite_session: Session) -> None:
+        api = DataSourceNotionDatasetSyncApi()
+        method = inspect.unwrap(api.get)
+
+        with (
+            app.test_request_context("/"),
+            patch(
+                "controllers.console.datasets.data_source.DatasetService.get_dataset",
+                return_value=None,
+            ),
+        ):
+            with pytest.raises(NotFound):
+                method(api, sqlite_session, "ds-1")
+
+
+class TestDataSourceNotionDocumentSyncApi:
+    @pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+    def test_get_success(self, app: Flask, sqlite_session: Session) -> None:
+        api = DataSourceNotionDocumentSyncApi()
+        method = inspect.unwrap(api.get)
+
+        with (
+            app.test_request_context("/"),
+            patch(
+                "controllers.console.datasets.data_source.DatasetService.get_dataset",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "controllers.console.datasets.data_source.DocumentService.get_document",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "controllers.console.datasets.data_source.document_indexing_sync_task.delay",
+                return_value=None,
+            ),
+        ):
+            response, status = method(api, sqlite_session, "ds-1", "doc-1")
+
+        assert status == 200
+
+    @pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+    def test_get_document_not_found(self, app: Flask, sqlite_session: Session) -> None:
+        api = DataSourceNotionDocumentSyncApi()
+        method = inspect.unwrap(api.get)
+
+        with (
+            app.test_request_context("/"),
+            patch(
+                "controllers.console.datasets.data_source.DatasetService.get_dataset",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "controllers.console.datasets.data_source.DocumentService.get_document",
+                return_value=None,
+            ),
+        ):
+            with pytest.raises(NotFound):
+                method(api, sqlite_session, "ds-1", "doc-1")

api/tests/unit_tests/controllers/inner_api/workspace/test_plugin_model_providers.py

@@ -1,64 +0,0 @@
-import inspect
-from unittest.mock import call, patch
-
-import pytest
-from flask import Flask
-from pydantic import ValidationError
-
-from controllers.inner_api.workspace.plugin_model_providers import (
-    EnterprisePluginModelProvidersCacheInvalidate,
-    InvalidatePluginModelProvidersCachePayload,
-)
-
-
-class TestInvalidatePluginModelProvidersCachePayload:
-    def test_valid_payload(self):
-        payload = InvalidatePluginModelProvidersCachePayload.model_validate(
-            {"tenant_ids": ["tenant-alpha", "tenant-beta"]}
-        )
-        assert payload.tenant_ids == ["tenant-alpha", "tenant-beta"]
-
-    def test_missing_tenant_ids_defaults_to_empty(self):
-        payload = InvalidatePluginModelProvidersCachePayload.model_validate({})
-        assert payload.tenant_ids == []
-
-    def test_unknown_field_rejected(self):
-        with pytest.raises(ValidationError):
-            InvalidatePluginModelProvidersCachePayload.model_validate({"tenant_ids": ["tenant-alpha"], "generation": 7})
-
-
-class TestEnterprisePluginModelProvidersCacheInvalidate:
-    @pytest.fixture
-    def api_instance(self):
-        return EnterprisePluginModelProvidersCacheInvalidate()
-
-    def _post(self, api_instance, app: Flask, payload):
-        unwrapped_post = inspect.unwrap(api_instance.post)
-        with app.test_request_context():
-            with patch("controllers.inner_api.workspace.plugin_model_providers.inner_api_ns") as mock_ns:
-                mock_ns.payload = payload
-                return unwrapped_post(api_instance)
-
-    @patch("controllers.inner_api.workspace.plugin_model_providers.PluginService")
-    def test_post_invalidates_once_per_tenant(self, mock_plugin_service, api_instance, app: Flask):
-        result = self._post(api_instance, app, {"tenant_ids": ["tenant-alpha", "tenant-beta"]})
-
-        assert result == ({"result": "success"}, 200)
-        assert mock_plugin_service.invalidate_plugin_model_providers_cache.call_args_list == [
-            call("tenant-alpha"),
-            call("tenant-beta"),
-        ]
-
-    @patch("controllers.inner_api.workspace.plugin_model_providers.PluginService")
-    def test_post_with_empty_list_is_a_no_op(self, mock_plugin_service, api_instance, app: Flask):
-        result = self._post(api_instance, app, {"tenant_ids": []})
-
-        assert result == ({"result": "success"}, 200)
-        mock_plugin_service.invalidate_plugin_model_providers_cache.assert_not_called()
-
-    @patch("controllers.inner_api.workspace.plugin_model_providers.PluginService")
-    def test_post_with_missing_payload_is_a_no_op(self, mock_plugin_service, api_instance, app: Flask):
-        result = self._post(api_instance, app, None)
-
-        assert result == ({"result": "success"}, 200)
-        mock_plugin_service.invalidate_plugin_model_providers_cache.assert_not_called()

api/tests/unit_tests/controllers/service_api/app/test_file_preview.py

@@ -1,17 +1,31 @@
-"""
-Unit tests for Service API File Preview endpoint
+"""Unit tests for the Service API file-preview endpoint.
+
+Ownership checks run against persisted message, file, app, and upload rows so the
+tests exercise the same SQLAlchemy statements and tenant boundary as production.
+Storage remains mocked because it is the external I/O boundary of the endpoint.
 """
 
 import logging
-import uuid
+from collections.abc import Iterator
+from dataclasses import dataclass
+from datetime import datetime
+from decimal import Decimal
 from typing import Protocol, cast
 from unittest.mock import Mock, patch
+from uuid import uuid4
 
 import pytest
+from sqlalchemy import event
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
 
 from controllers.service_api.app.error import FileAccessDeniedError, FileNotFoundError
 from controllers.service_api.app.file_preview import FilePreviewApi
-from models.model import App, EndUser, Message, MessageFile, UploadFile
+from extensions.storage.storage_type import StorageType
+from graphon.file import FileTransferMethod, FileType
+from models.base import TypeBase
+from models.enums import ConversationFromSource, CreatorUserRole
+from models.model import App, AppMode, Message, MessageFile, UploadFile
 
 
 class _FilePreviewLogRecord(Protocol):
@@ -20,367 +34,252 @@ class _FilePreviewLogRecord(Protocol):
     error: str
 
 
+@dataclass(frozen=True)
+class _Database:
+    """Expose the real test session through the interface used by the controller."""
+
+    session: Session
+
+
+@dataclass(frozen=True)
+class _PreviewRecords:
+    app: App
+    message: Message
+    message_file: MessageFile
+    upload_file: UploadFile
+
+
+@pytest.fixture
+def database(sqlite_engine: Engine) -> Iterator[_Database]:
+    """Create only the tables required by file ownership validation."""
+
+    models = (App, Message, MessageFile, UploadFile)
+    tables = [TypeBase.metadata.tables[model.__tablename__] for model in models]
+    TypeBase.metadata.create_all(sqlite_engine, tables=tables)
+    with Session(sqlite_engine, expire_on_commit=False) as session:
+        yield _Database(session)
+
+
+@pytest.fixture
+def file_preview_api() -> FilePreviewApi:
+    """Create the resource instance under test."""
+
+    return FilePreviewApi()
+
+
+def _upload_file(*, tenant_id: str, file_id: str | None = None) -> UploadFile:
+    upload_file = UploadFile(
+        tenant_id=tenant_id,
+        storage_type=StorageType.LOCAL,
+        key="storage/key/test_file.jpg",
+        name="test_file.jpg",
+        size=1024,
+        extension="jpg",
+        mime_type="image/jpeg",
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by=str(uuid4()),
+        created_at=datetime(2026, 1, 1),
+        used=True,
+    )
+    if file_id is not None:
+        upload_file.id = file_id
+    return upload_file
+
+
+def _persist_preview_records(
+    session: Session,
+    *,
+    app_id: str | None = None,
+    app_tenant_id: str | None = None,
+    upload_tenant_id: str | None = None,
+) -> _PreviewRecords:
+    app_id = app_id or str(uuid4())
+    app_tenant_id = app_tenant_id or str(uuid4())
+    upload_file = _upload_file(tenant_id=upload_tenant_id or app_tenant_id)
+    app = App(
+        id=app_id,
+        tenant_id=app_tenant_id,
+        name="Preview app",
+        description="",
+        mode=AppMode.CHAT,
+        icon_type=None,
+        icon="",
+        icon_background=None,
+        enable_site=True,
+        enable_api=True,
+    )
+    message = Message(
+        id=str(uuid4()),
+        app_id=app_id,
+        conversation_id=str(uuid4()),
+        _inputs={},
+        query="preview",
+        message={},
+        message_unit_price=Decimal(0),
+        answer="answer",
+        answer_unit_price=Decimal(0),
+        currency="USD",
+        from_source=ConversationFromSource.API,
+    )
+    message_file = MessageFile(
+        message_id=message.id,
+        type=FileType.IMAGE,
+        transfer_method=FileTransferMethod.LOCAL_FILE,
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by=str(uuid4()),
+        upload_file_id=upload_file.id,
+    )
+    session.add_all([app, message, message_file, upload_file])
+    session.commit()
+    return _PreviewRecords(app=app, message=message, message_file=message_file, upload_file=upload_file)
+
+
 class TestFilePreviewApi:
-    """Test suite for FilePreviewApi"""
+    """Exercise ownership validation and response construction."""
 
-    @pytest.fixture
-    def file_preview_api(self):
-        """Create FilePreviewApi instance for testing"""
-        return FilePreviewApi()
+    def test_validate_file_ownership_success(self, file_preview_api: FilePreviewApi, database: _Database):
+        records = _persist_preview_records(database.session)
 
-    @pytest.fixture
-    def mock_app(self):
-        """Mock App model"""
-        app = Mock(spec=App)
-        app.id = str(uuid.uuid4())
-        app.tenant_id = str(uuid.uuid4())
-        return app
+        with patch("controllers.service_api.app.file_preview.db", database):
+            message_file, upload_file = file_preview_api._validate_file_ownership(
+                records.upload_file.id, records.app.id
+            )
 
-    @pytest.fixture
-    def mock_end_user(self):
-        """Mock EndUser model"""
-        end_user = Mock(spec=EndUser)
-        end_user.id = str(uuid.uuid4())
-        return end_user
+        assert message_file.id == records.message_file.id
+        assert upload_file.id == records.upload_file.id
+        assert upload_file.tenant_id == records.app.tenant_id
 
-    @pytest.fixture
-    def mock_upload_file(self):
-        """Mock UploadFile model"""
-        upload_file = Mock(spec=UploadFile)
-        upload_file.id = str(uuid.uuid4())
-        upload_file.name = "test_file.jpg"
-        upload_file.extension = "jpg"
-        upload_file.mime_type = "image/jpeg"
-        upload_file.size = 1024
-        upload_file.key = "storage/key/test_file.jpg"
-        upload_file.tenant_id = str(uuid.uuid4())
-        return upload_file
+    def test_validate_file_ownership_file_not_found(self, file_preview_api: FilePreviewApi, database: _Database):
+        with patch("controllers.service_api.app.file_preview.db", database):
+            with pytest.raises(FileNotFoundError, match="File not found in message context"):
+                file_preview_api._validate_file_ownership(str(uuid4()), str(uuid4()))
 
-    @pytest.fixture
-    def mock_message_file(self):
-        """Mock MessageFile model"""
-        message_file = Mock(spec=MessageFile)
-        message_file.id = str(uuid.uuid4())
-        message_file.upload_file_id = str(uuid.uuid4())
-        message_file.message_id = str(uuid.uuid4())
-        return message_file
+    def test_validate_file_ownership_access_denied(self, file_preview_api: FilePreviewApi, database: _Database):
+        records = _persist_preview_records(database.session)
 
-    @pytest.fixture
-    def mock_message(self):
-        """Mock Message model"""
-        message = Mock(spec=Message)
-        message.id = str(uuid.uuid4())
-        message.app_id = str(uuid.uuid4())
-        return message
+        with patch("controllers.service_api.app.file_preview.db", database):
+            with pytest.raises(FileAccessDeniedError, match="not owned by requesting app"):
+                file_preview_api._validate_file_ownership(records.upload_file.id, str(uuid4()))
 
-    def test_validate_file_ownership_success(
-        self, file_preview_api: FilePreviewApi, mock_app, mock_upload_file, mock_message_file, mock_message
-    ):
-        """Test successful file ownership validation"""
-        file_id = str(uuid.uuid4())
-        app_id = mock_app.id
+    def test_validate_file_ownership_upload_file_not_found(self, file_preview_api: FilePreviewApi, database: _Database):
+        records = _persist_preview_records(database.session)
+        database.session.delete(records.upload_file)
+        database.session.commit()
 
-        # Set up the mocks
-        mock_upload_file.tenant_id = mock_app.tenant_id
-        mock_message.app_id = app_id
-        mock_message_file.upload_file_id = file_id
-        mock_message_file.message_id = mock_message.id
+        with patch("controllers.service_api.app.file_preview.db", database):
+            with pytest.raises(FileNotFoundError, match="Upload file record not found"):
+                file_preview_api._validate_file_ownership(records.upload_file.id, records.app.id)
 
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock scalar() for MessageFile and Message queries
-            mock_db.session.scalar.side_effect = [
-                mock_message_file,  # MessageFile query
-                mock_message,  # Message query
-            ]
-            # Mock get() for UploadFile and App PK lookups
-            mock_db.session.get.side_effect = [
-                mock_upload_file,  # UploadFile query
-                mock_app,  # App query for tenant validation
-            ]
+    def test_validate_file_ownership_tenant_mismatch(self, file_preview_api: FilePreviewApi, database: _Database):
+        records = _persist_preview_records(database.session, upload_tenant_id=str(uuid4()))
 
-            # Execute the method
-            result_message_file, result_upload_file = file_preview_api._validate_file_ownership(file_id, app_id)
-
-            # Assertions
-            assert result_message_file == mock_message_file
-            assert result_upload_file == mock_upload_file
-
-    def test_validate_file_ownership_file_not_found(self, file_preview_api: FilePreviewApi):
-        """Test file ownership validation when MessageFile not found"""
-        file_id = str(uuid.uuid4())
-        app_id = str(uuid.uuid4())
-
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock MessageFile not found via scalar()
-            mock_db.session.scalar.return_value = None
-
-            # Execute and assert exception
-            with pytest.raises(FileNotFoundError) as exc_info:
-                file_preview_api._validate_file_ownership(file_id, app_id)
-
-            assert "File not found in message context" in str(exc_info.value)
-
-    def test_validate_file_ownership_access_denied(self, file_preview_api: FilePreviewApi, mock_message_file):
-        """Test file ownership validation when Message not owned by app"""
-        file_id = str(uuid.uuid4())
-        app_id = str(uuid.uuid4())
-
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock MessageFile found but Message not owned by app via scalar()
-            mock_db.session.scalar.side_effect = [
-                mock_message_file,  # MessageFile query - found
-                None,  # Message query - not found (access denied)
-            ]
-
-            # Execute and assert exception
-            with pytest.raises(FileAccessDeniedError) as exc_info:
-                file_preview_api._validate_file_ownership(file_id, app_id)
-
-            assert "not owned by requesting app" in str(exc_info.value)
-
-    def test_validate_file_ownership_upload_file_not_found(
-        self, file_preview_api: FilePreviewApi, mock_message_file, mock_message
-    ):
-        """Test file ownership validation when UploadFile not found"""
-        file_id = str(uuid.uuid4())
-        app_id = str(uuid.uuid4())
-
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock scalar() for MessageFile and Message
-            mock_db.session.scalar.side_effect = [
-                mock_message_file,  # MessageFile query - found
-                mock_message,  # Message query - found
-            ]
-            # Mock get() for UploadFile - not found
-            mock_db.session.get.return_value = None
-
-            # Execute and assert exception
-            with pytest.raises(FileNotFoundError) as exc_info:
-                file_preview_api._validate_file_ownership(file_id, app_id)
-
-            assert "Upload file record not found" in str(exc_info.value)
-
-    def test_validate_file_ownership_tenant_mismatch(
-        self, file_preview_api: FilePreviewApi, mock_app, mock_upload_file, mock_message_file, mock_message
-    ):
-        """Test file ownership validation with tenant mismatch"""
-        file_id = str(uuid.uuid4())
-        app_id = mock_app.id
-
-        # Set up tenant mismatch
-        mock_upload_file.tenant_id = "different_tenant_id"
-        mock_app.tenant_id = "app_tenant_id"
-        mock_message.app_id = app_id
-        mock_message_file.upload_file_id = file_id
-        mock_message_file.message_id = mock_message.id
-
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock scalar() for MessageFile and Message queries
-            mock_db.session.scalar.side_effect = [
-                mock_message_file,  # MessageFile query
-                mock_message,  # Message query
-            ]
-            # Mock get() for UploadFile and App PK lookups
-            mock_db.session.get.side_effect = [
-                mock_upload_file,  # UploadFile query
-                mock_app,  # App query for tenant validation
-            ]
-
-            # Execute and assert exception
-            with pytest.raises(FileAccessDeniedError) as exc_info:
-                file_preview_api._validate_file_ownership(file_id, app_id)
-
-            assert "tenant mismatch" in str(exc_info.value)
+        with patch("controllers.service_api.app.file_preview.db", database):
+            with pytest.raises(FileAccessDeniedError, match="tenant mismatch"):
+                file_preview_api._validate_file_ownership(records.upload_file.id, records.app.id)
 
     def test_validate_file_ownership_invalid_input(self, file_preview_api: FilePreviewApi):
-        """Test file ownership validation with invalid input"""
-
-        # Test with empty file_id
-        with pytest.raises(FileAccessDeniedError) as exc_info:
+        with pytest.raises(FileAccessDeniedError, match="Invalid file or app identifier"):
             file_preview_api._validate_file_ownership("", "app_id")
-        assert "Invalid file or app identifier" in str(exc_info.value)
 
-        # Test with empty app_id
-        with pytest.raises(FileAccessDeniedError) as exc_info:
+        with pytest.raises(FileAccessDeniedError, match="Invalid file or app identifier"):
             file_preview_api._validate_file_ownership("file_id", "")
-        assert "Invalid file or app identifier" in str(exc_info.value)
 
-    def test_build_file_response_basic(self, file_preview_api: FilePreviewApi, mock_upload_file):
-        """Test basic file response building"""
-        mock_generator = Mock()
+    @pytest.mark.parametrize(
+        ("as_attachment", "mime_type", "name", "extension", "size"),
+        [
+            (False, "image/jpeg", "test_file.jpg", "jpg", 1024),
+            (True, "image/jpeg", "test_file.jpg", "jpg", 1024),
+            (False, "text/html", "unsafe.html", "html", 1024),
+            (False, "video/mp4", "test_file.mp4", "mp4", 1024),
+            (False, "image/jpeg", "test_file.jpg", "jpg", 0),
+        ],
+    )
+    def test_build_file_response(
+        self,
+        file_preview_api: FilePreviewApi,
+        as_attachment: bool,
+        mime_type: str,
+        name: str,
+        extension: str,
+        size: int,
+    ):
+        upload_file = _upload_file(tenant_id=str(uuid4()))
+        upload_file.mime_type = mime_type
+        upload_file.name = name
+        upload_file.extension = extension
+        upload_file.size = size
 
-        response = file_preview_api._build_file_response(mock_generator, mock_upload_file, False)
+        response = file_preview_api._build_file_response(Mock(), upload_file, as_attachment)
 
-        # Check response properties
-        assert response.mimetype == mock_upload_file.mime_type
         assert response.direct_passthrough is True
-        assert response.headers["Content-Length"] == str(mock_upload_file.size)
         assert "Cache-Control" in response.headers
-
-    def test_build_file_response_as_attachment(self, file_preview_api: FilePreviewApi, mock_upload_file):
-        """Test file response building with attachment flag"""
-        mock_generator = Mock()
-
-        response = file_preview_api._build_file_response(mock_generator, mock_upload_file, True)
-
-        # Check attachment-specific headers
-        assert "attachment" in response.headers["Content-Disposition"]
-        assert mock_upload_file.name in response.headers["Content-Disposition"]
-        assert response.headers["Content-Type"] == "application/octet-stream"
-
-    def test_build_file_response_html_forces_attachment(self, file_preview_api: FilePreviewApi, mock_upload_file):
-        """Test HTML files are forced to download"""
-        mock_generator = Mock()
-        mock_upload_file.mime_type = "text/html"
-        mock_upload_file.name = "unsafe.html"
-        mock_upload_file.extension = "html"
-
-        response = file_preview_api._build_file_response(mock_generator, mock_upload_file, False)
-
-        assert "attachment" in response.headers["Content-Disposition"]
-        assert response.headers["Content-Type"] == "application/octet-stream"
-        assert response.headers["X-Content-Type-Options"] == "nosniff"
-
-    def test_build_file_response_audio_video(self, file_preview_api: FilePreviewApi, mock_upload_file):
-        """Test file response building for audio/video files"""
-        mock_generator = Mock()
-        mock_upload_file.mime_type = "video/mp4"
-
-        response = file_preview_api._build_file_response(mock_generator, mock_upload_file, False)
-
-        # Check Range support for media files
-        assert response.headers["Accept-Ranges"] == "bytes"
-
-    def test_build_file_response_no_size(self, file_preview_api: FilePreviewApi, mock_upload_file):
-        """Test file response building when size is unknown"""
-        mock_generator = Mock()
-        mock_upload_file.size = 0  # Unknown size
-
-        response = file_preview_api._build_file_response(mock_generator, mock_upload_file, False)
-
-        # Content-Length should not be set when size is unknown
-        assert "Content-Length" not in response.headers
+        assert ("Content-Length" in response.headers) is bool(size)
+        if as_attachment or mime_type == "text/html":
+            assert "attachment" in response.headers["Content-Disposition"]
+            assert response.headers["Content-Type"] == "application/octet-stream"
+        else:
+            assert response.mimetype == mime_type
+        if mime_type == "text/html":
+            assert response.headers["X-Content-Type-Options"] == "nosniff"
+        if mime_type.startswith("video/"):
+            assert response.headers["Accept-Ranges"] == "bytes"
 
     @patch("controllers.service_api.app.file_preview.storage")
-    def test_get_method_integration(
-        self,
-        mock_storage,
-        file_preview_api: FilePreviewApi,
-        mock_app,
-        mock_end_user,
-        mock_upload_file,
-        mock_message_file,
-        mock_message,
+    def test_components_use_validated_file(
+        self, mock_storage: Mock, file_preview_api: FilePreviewApi, database: _Database
     ):
-        """Test the full GET method integration (without decorator)"""
-        file_id = str(uuid.uuid4())
-        app_id = mock_app.id
+        records = _persist_preview_records(database.session)
+        generator = Mock()
 
-        # Set up mocks
-        mock_upload_file.tenant_id = mock_app.tenant_id
-        mock_message.app_id = app_id
-        mock_message_file.upload_file_id = file_id
-        mock_message_file.message_id = mock_message.id
+        with patch("controllers.service_api.app.file_preview.db", database):
+            message_file, upload_file = file_preview_api._validate_file_ownership(
+                records.upload_file.id, records.app.id
+            )
+            response = file_preview_api._build_file_response(generator, upload_file, False)
 
-        mock_generator = Mock()
-        mock_storage.load.return_value = mock_generator
-
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock scalar() for MessageFile and Message queries
-            mock_db.session.scalar.side_effect = [
-                mock_message_file,  # MessageFile query
-                mock_message,  # Message query
-            ]
-            # Mock get() for UploadFile and App PK lookups
-            mock_db.session.get.side_effect = [
-                mock_upload_file,  # UploadFile query
-                mock_app,  # App query for tenant validation
-            ]
-
-            # Test the core logic directly without Flask decorators
-            # Validate file ownership
-            result_message_file, result_upload_file = file_preview_api._validate_file_ownership(file_id, app_id)
-            assert result_message_file == mock_message_file
-            assert result_upload_file == mock_upload_file
-
-            # Test file response building
-            response = file_preview_api._build_file_response(mock_generator, mock_upload_file, False)
-            assert response is not None
-
-            # Verify storage was called correctly
-            mock_storage.load.assert_not_called()  # Since we're testing components separately
+        assert message_file.id == records.message_file.id
+        assert response.mimetype == "image/jpeg"
+        mock_storage.load.assert_not_called()
 
     @patch("controllers.service_api.app.file_preview.storage")
-    def test_storage_error_handling(
-        self,
-        mock_storage,
-        file_preview_api: FilePreviewApi,
-        mock_app,
-        mock_upload_file,
-        mock_message_file,
-        mock_message,
+    def test_storage_error_remains_external(
+        self, mock_storage: Mock, file_preview_api: FilePreviewApi, database: _Database
     ):
-        """Test storage error handling in the core logic"""
-        file_id = str(uuid.uuid4())
-        app_id = mock_app.id
+        records = _persist_preview_records(database.session)
+        mock_storage.load.side_effect = OSError("Storage error")
 
-        # Set up mocks
-        mock_upload_file.tenant_id = mock_app.tenant_id
-        mock_message.app_id = app_id
-        mock_message_file.upload_file_id = file_id
-        mock_message_file.message_id = mock_message.id
+        with patch("controllers.service_api.app.file_preview.db", database):
+            _, upload_file = file_preview_api._validate_file_ownership(records.upload_file.id, records.app.id)
 
-        # Mock storage error
-        mock_storage.load.side_effect = Exception("Storage error")
-
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock scalar() for MessageFile and Message queries
-            mock_db.session.scalar.side_effect = [
-                mock_message_file,  # MessageFile query
-                mock_message,  # Message query
-            ]
-            # Mock get() for UploadFile and App PK lookups
-            mock_db.session.get.side_effect = [
-                mock_upload_file,  # UploadFile query
-                mock_app,  # App query for tenant validation
-            ]
-
-            # First validate file ownership works
-            result_message_file, result_upload_file = file_preview_api._validate_file_ownership(file_id, app_id)
-            assert result_message_file == mock_message_file
-            assert result_upload_file == mock_upload_file
-
-            # Test storage error handling
-            with pytest.raises(Exception) as exc_info:
-                mock_storage.load(mock_upload_file.key, stream=True)
-
-            assert "Storage error" in str(exc_info.value)
+        with pytest.raises(OSError, match="Storage error"):
+            mock_storage.load(upload_file.key, stream=True)
 
     def test_validate_file_ownership_unexpected_error_logging(
-        self, file_preview_api: FilePreviewApi, caplog: pytest.LogCaptureFixture
+        self,
+        file_preview_api: FilePreviewApi,
+        database: _Database,
+        sqlite_engine: Engine,
+        caplog: pytest.LogCaptureFixture,
     ):
-        """Test that unexpected errors are logged properly"""
-        file_id = str(uuid.uuid4())
-        app_id = str(uuid.uuid4())
+        file_id = str(uuid4())
+        app_id = str(uuid4())
 
-        with patch("controllers.service_api.app.file_preview.db") as mock_db:
-            # Mock database scalar to raise unexpected exception
-            mock_db.session.scalar.side_effect = Exception("Unexpected database error")
+        def fail_statement(*_args: object) -> None:
+            raise RuntimeError("Unexpected database error")
 
-            # Execute and assert exception
-            with caplog.at_level(logging.ERROR, logger="controllers.service_api.app.file_preview"):
-                with pytest.raises(FileAccessDeniedError) as exc_info:
-                    file_preview_api._validate_file_ownership(file_id, app_id)
+        event.listen(sqlite_engine, "before_cursor_execute", fail_statement)
+        try:
+            with patch("controllers.service_api.app.file_preview.db", database):
+                with caplog.at_level(logging.ERROR, logger="controllers.service_api.app.file_preview"):
+                    with pytest.raises(FileAccessDeniedError, match="File access validation failed"):
+                        file_preview_api._validate_file_ownership(file_id, app_id)
+        finally:
+            event.remove(sqlite_engine, "before_cursor_execute", fail_statement)
 
-            # Verify error message
-            assert "File access validation failed" in str(exc_info.value)
-
-            # Verify logging was called with the structured context fields. The ``extra`` keys
-            # are attached to the LogRecord as attributes, so they are not in ``caplog.text``.
-            assert len(caplog.records) == 1
-            log_record = caplog.records[0]
-            assert log_record.getMessage() == "Unexpected error during file ownership validation"
-            record = cast(_FilePreviewLogRecord, log_record)
-            assert record.file_id == file_id
-            assert record.app_id == app_id
-            assert record.error == "Unexpected database error"
+        assert len(caplog.records) == 1
+        log_record = caplog.records[0]
+        assert log_record.getMessage() == "Unexpected error during file ownership validation"
+        record = cast(_FilePreviewLogRecord, log_record)
+        assert record.file_id == file_id
+        assert record.app_id == app_id
+        assert record.error == "Unexpected database error"

api/tests/unit_tests/controllers/service_api/app/test_legacy_system_files.py

@@ -0,0 +1,89 @@
+from unittest.mock import MagicMock
+
+from controllers.service_api.app.legacy_system_files import (
+    attach_legacy_system_file_warning_for_service_api,
+    normalize_legacy_system_file_args_for_service_api,
+)
+from core.app.entities.app_invoke_entities import InvokeFrom
+from services.app_generate_service import AppGenerateService
+
+_LEGACY_FILE_TEMPLATE = "{{#" + ".".join(("sys", "files")) + "#}}"
+_USER_INPUT_FILE_INPUT_KEY = ".".join(("userinput", "files"))
+
+
+def _legacy_file_graph() -> dict:
+    return {
+        "nodes": [
+            {"id": "start", "data": {"type": "start", "variables": []}},
+            {"id": "answer", "data": {"type": "answer", "answer": _LEGACY_FILE_TEMPLATE}},
+        ],
+        "edges": [],
+    }
+
+
+def test_hidden_service_api_file_payload_maps_to_userinput_files(mocker):
+    workflow = MagicMock()
+    workflow.graph_dict = _legacy_file_graph()
+    get_workflow = mocker.patch.object(AppGenerateService, "get_workflow", return_value=workflow)
+    app_model = MagicMock()
+    session = MagicMock()
+    files = [{"transfer_method": "remote_url", "url": "https://example.com/a.png"}]
+
+    args, compat_variable = normalize_legacy_system_file_args_for_service_api(
+        session=session,
+        app_model=app_model,
+        args={"inputs": {}, "files": None},
+        raw_payload={"system": {"files": files}},
+    )
+
+    get_workflow.assert_called_once_with(app_model, InvokeFrom.SERVICE_API, None, session=session)
+    assert compat_variable is not None
+    assert args["files"] == files
+    assert args["inputs"][_USER_INPUT_FILE_INPUT_KEY] == files
+
+
+def test_service_api_file_payload_is_ignored_when_absent(mocker):
+    get_workflow = mocker.patch.object(AppGenerateService, "get_workflow")
+    app_model = MagicMock()
+    original_args = {"inputs": {}}
+
+    args, compat_variable = normalize_legacy_system_file_args_for_service_api(
+        session=MagicMock(),
+        app_model=app_model,
+        args=original_args,
+        raw_payload={},
+    )
+
+    assert args is original_args
+    assert compat_variable is None
+    get_workflow.assert_not_called()
+
+
+def test_top_level_service_api_file_payload_still_checks_workflow_graph(mocker):
+    workflow = MagicMock()
+    workflow.graph_dict = {"nodes": []}
+    get_workflow = mocker.patch.object(AppGenerateService, "get_workflow", return_value=workflow)
+    app_model = MagicMock()
+    session = MagicMock()
+    files = [{"id": "file-1"}]
+
+    args, compat_variable = normalize_legacy_system_file_args_for_service_api(
+        session=session,
+        app_model=app_model,
+        args={"inputs": {}, "files": files},
+        raw_payload={},
+    )
+
+    get_workflow.assert_called_once_with(app_model, InvokeFrom.SERVICE_API, None, session=session)
+    assert args["files"] == files
+    assert compat_variable is None
+
+
+def test_service_api_warning_is_attached_only_when_compatibility_was_used():
+    compat_variable = MagicMock(node_id="userinput", variable_name="files")
+
+    response = attach_legacy_system_file_warning_for_service_api({"answer": "ok"}, compat_variable)
+    response_without_warning = attach_legacy_system_file_warning_for_service_api({"answer": "ok"}, None)
+
+    assert response["warnings"]
+    assert response_without_warning == {"answer": "ok"}

api/tests/unit_tests/controllers/service_api/app/test_message.py

@@ -15,12 +15,15 @@ Focus on:
 """
 
 import uuid
+from collections.abc import Iterator
 from inspect import unwrap
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import pytest
 from flask import Flask
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import BadRequest, InternalServerError, NotFound
 
 from controllers.service_api.app.error import NotChatAppError
@@ -44,6 +47,14 @@ from services.errors.message import (
 from services.message_service import MessageService
 
 
+@pytest.fixture
+def orm_session(sqlite_engine: Engine) -> Iterator[Session]:
+    """Provide a real caller-owned session for MessageService interface tests."""
+
+    with Session(sqlite_engine, expire_on_commit=False) as session:
+        yield session
+
+
 class TestMessageListQuery:
     """Test suite for MessageListQuery Pydantic model."""
 
@@ -253,7 +264,7 @@ class TestMessageService:
         assert callable(MessageService.get_suggested_questions_after_answer)
 
     @patch.object(MessageService, "pagination_by_first_id")
-    def test_pagination_by_first_id_returns_pagination_result(self, mock_pagination):
+    def test_pagination_by_first_id_returns_pagination_result(self, mock_pagination, orm_session: Session):
         """Test pagination_by_first_id returns expected format."""
         mock_result = Mock()
         mock_result.data = []
@@ -267,7 +278,7 @@ class TestMessageService:
             conversation_id=str(uuid.uuid4()),
             first_id=None,
             limit=20,
-            session=Mock(),
+            session=orm_session,
         )
 
         assert hasattr(result, "data")
@@ -275,7 +286,7 @@ class TestMessageService:
         assert hasattr(result, "has_more")
 
     @patch.object(MessageService, "pagination_by_first_id")
-    def test_pagination_raises_conversation_not_exists_error(self, mock_pagination):
+    def test_pagination_raises_conversation_not_exists_error(self, mock_pagination, orm_session: Session):
         """Test pagination raises ConversationNotExistsError."""
         import services.errors.conversation
 
@@ -288,11 +299,11 @@ class TestMessageService:
                 conversation_id="invalid_id",
                 first_id=None,
                 limit=20,
-                session=Mock(),
+                session=orm_session,
             )
 
     @patch.object(MessageService, "pagination_by_first_id")
-    def test_pagination_raises_first_message_not_exists_error(self, mock_pagination):
+    def test_pagination_raises_first_message_not_exists_error(self, mock_pagination, orm_session: Session):
         """Test pagination raises FirstMessageNotExistsError."""
         mock_pagination.side_effect = FirstMessageNotExistsError()
 
@@ -303,11 +314,11 @@ class TestMessageService:
                 conversation_id=str(uuid.uuid4()),
                 first_id="invalid_first_id",
                 limit=20,
-                session=Mock(),
+                session=orm_session,
             )
 
     @patch.object(MessageService, "create_feedback")
-    def test_create_feedback_with_rating_and_content(self, mock_create_feedback):
+    def test_create_feedback_with_rating_and_content(self, mock_create_feedback, orm_session: Session):
         """Test create_feedback with rating and content."""
         mock_create_feedback.return_value = None
 
@@ -317,13 +328,13 @@ class TestMessageService:
             user=Mock(spec=EndUser),
             rating=FeedbackRating.LIKE,
             content="Great response!",
-            session=Mock(),
+            session=orm_session,
         )
 
         mock_create_feedback.assert_called_once()
 
     @patch.object(MessageService, "create_feedback")
-    def test_create_feedback_raises_message_not_exists_error(self, mock_create_feedback):
+    def test_create_feedback_raises_message_not_exists_error(self, mock_create_feedback, orm_session: Session):
         """Test create_feedback raises MessageNotExistsError."""
         mock_create_feedback.side_effect = MessageNotExistsError()
 
@@ -334,11 +345,11 @@ class TestMessageService:
                 user=Mock(spec=EndUser),
                 rating=FeedbackRating.LIKE,
                 content=None,
-                session=Mock(),
+                session=orm_session,
             )
 
     @patch.object(MessageService, "get_all_messages_feedbacks")
-    def test_get_all_messages_feedbacks_returns_list(self, mock_get_feedbacks):
+    def test_get_all_messages_feedbacks_returns_list(self, mock_get_feedbacks, orm_session: Session):
         """Test get_all_messages_feedbacks returns list of feedbacks."""
         mock_feedbacks = [
             {"message_id": str(uuid.uuid4()), "rating": "like"},
@@ -346,13 +357,15 @@ class TestMessageService:
         ]
         mock_get_feedbacks.return_value = mock_feedbacks
 
-        result = MessageService.get_all_messages_feedbacks(app_model=Mock(spec=App), page=1, limit=20, session=Mock())
+        result = MessageService.get_all_messages_feedbacks(
+            app_model=Mock(spec=App), page=1, limit=20, session=orm_session
+        )
 
         assert len(result) == 2
         assert result[0]["rating"] == "like"
 
     @patch.object(MessageService, "get_suggested_questions_after_answer")
-    def test_get_suggested_questions_returns_questions_list(self, mock_get_questions):
+    def test_get_suggested_questions_returns_questions_list(self, mock_get_questions, orm_session: Session):
         """Test get_suggested_questions_after_answer returns list of questions."""
         mock_questions = ["What about this aspect?", "Can you elaborate on that?", "How does this relate to...?"]
         mock_get_questions.return_value = mock_questions
@@ -362,14 +375,14 @@ class TestMessageService:
             user=Mock(spec=EndUser),
             message_id=str(uuid.uuid4()),
             invoke_from=Mock(),
-            session=Mock(),
+            session=orm_session,
         )
 
         assert len(result) == 3
         assert isinstance(result[0], str)
 
     @patch.object(MessageService, "get_suggested_questions_after_answer")
-    def test_get_suggested_questions_raises_disabled_error(self, mock_get_questions):
+    def test_get_suggested_questions_raises_disabled_error(self, mock_get_questions, orm_session: Session):
         """Test get_suggested_questions_after_answer raises SuggestedQuestionsAfterAnswerDisabledError."""
         mock_get_questions.side_effect = SuggestedQuestionsAfterAnswerDisabledError()
 
@@ -379,11 +392,11 @@ class TestMessageService:
                 user=Mock(spec=EndUser),
                 message_id=str(uuid.uuid4()),
                 invoke_from=Mock(),
-                session=Mock(),
+                session=orm_session,
             )
 
     @patch.object(MessageService, "get_suggested_questions_after_answer")
-    def test_get_suggested_questions_raises_message_not_exists_error(self, mock_get_questions):
+    def test_get_suggested_questions_raises_message_not_exists_error(self, mock_get_questions, orm_session: Session):
         """Test get_suggested_questions_after_answer raises MessageNotExistsError."""
         mock_get_questions.side_effect = MessageNotExistsError()
 
@@ -393,7 +406,7 @@ class TestMessageService:
                 user=Mock(spec=EndUser),
                 message_id="invalid_message_id",
                 invoke_from=Mock(),
-                session=Mock(),
+                session=orm_session,
             )
 
 

api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py

@@ -24,6 +24,7 @@ from unittest.mock import Mock, patch
 
 import pytest
 from flask import Flask
+from sqlalchemy.orm import Session
 from werkzeug.datastructures import FileStorage
 from werkzeug.exceptions import Forbidden, NotFound
 
@@ -38,6 +39,7 @@ from controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow import (
 )
 from core.app.entities.app_invoke_entities import InvokeFrom
 from models.account import Account
+from models.dataset import Dataset
 from services.errors.file import FileTooLargeError, UnsupportedFileTypeError
 from services.rag_pipeline.entity.pipeline_service_api_entities import (
     DatasourceNodeRunApiEntity,
@@ -46,6 +48,20 @@ from services.rag_pipeline.entity.pipeline_service_api_entities import (
 from services.rag_pipeline.rag_pipeline import RagPipelineService
 
 
+def _persist_dataset(session: Session, *, tenant_id: str, dataset_id: str) -> Dataset:
+    dataset = Dataset(
+        id=dataset_id,
+        tenant_id=tenant_id,
+        name="Pipeline dataset",
+        created_by="account-1",
+        data_source_type=None,
+        indexing_technique=None,
+    )
+    session.add(dataset)
+    session.commit()
+    return dataset
+
+
 class TestDatasourceNodeRunPayload:
     """Test suite for DatasourceNodeRunPayload Pydantic model."""
 
@@ -550,13 +566,15 @@ class TestPipelineRunApiPost:
     )
     @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService")
     @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns")
-    def test_post_success_streaming(self, mock_ns, mock_svc_cls, mock_current_user, mock_gen_svc, mock_helper, app):
+    @pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
+    def test_post_success_streaming(
+        self, mock_ns, mock_svc_cls, mock_current_user, mock_gen_svc, mock_helper, app, sqlite_session: Session
+    ):
         """Test successful pipeline run with streaming response."""
         tenant_id = str(uuid.uuid4())
         dataset_id = str(uuid.uuid4())
 
-        session = Mock()
-        session.scalar.return_value = Mock()
+        _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id)
 
         mock_ns.payload = {
             "inputs": {"key": "val"},
@@ -577,33 +595,33 @@ class TestPipelineRunApiPost:
 
         with app.test_request_context("/datasets/test/pipeline/run", method="POST"):
             api = PipelineRunApi()
-            response = api.post.__wrapped__(api, session, tenant_id=tenant_id, dataset_id=dataset_id)
+            response = api.post.__wrapped__(api, sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id)
 
         assert response == {"result": "ok"}
-        mock_svc_cls.assert_called_once_with(session)
+        mock_svc_cls.assert_called_once_with(sqlite_session)
         mock_gen_svc.generate.assert_called_once()
 
-    def test_post_not_found(self, app: Flask):
+    @pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
+    def test_post_not_found(self, app: Flask, sqlite_session: Session):
         """Test NotFound when dataset check fails."""
-        session = Mock()
-        session.scalar.return_value = None
-
         with app.test_request_context("/datasets/test/pipeline/run", method="POST"):
             api = PipelineRunApi()
             with pytest.raises(NotFound):
                 api.post.__wrapped__(
                     api,
-                    session,
+                    sqlite_session,
                     tenant_id=str(uuid.uuid4()),
                     dataset_id=str(uuid.uuid4()),
                 )
 
     @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", new="not_account")
     @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns")
-    def test_post_forbidden_non_account_user(self, mock_ns, app: Flask):
+    @pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
+    def test_post_forbidden_non_account_user(self, mock_ns, app: Flask, sqlite_session: Session):
         """Test Forbidden when current_user is not an Account."""
-        session = Mock()
-        session.scalar.return_value = Mock()
+        tenant_id = str(uuid.uuid4())
+        dataset_id = str(uuid.uuid4())
+        _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id)
         mock_ns.payload = {
             "inputs": {},
             "datasource_type": "online_document",
@@ -618,9 +636,9 @@ class TestPipelineRunApiPost:
             with pytest.raises(Forbidden):
                 api.post.__wrapped__(
                     api,
-                    session,
-                    tenant_id=str(uuid.uuid4()),
-                    dataset_id=str(uuid.uuid4()),
+                    sqlite_session,
+                    tenant_id=tenant_id,
+                    dataset_id=dataset_id,
                 )
 
 

api/tests/unit_tests/controllers/service_api/test_wraps.py

@@ -3,10 +3,13 @@ Unit tests for Service API wraps (authentication decorators)
 """
 
 import uuid
+from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 from flask import Flask
+from sqlalchemy import select
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
 
 from controllers.service_api.wraps import (
@@ -21,12 +24,11 @@ from controllers.service_api.wraps import (
     validate_dataset_token,
 )
 from enums.cloud_plan import CloudPlan
-from models.account import TenantStatus
-from models.model import ApiToken
-from tests.unit_tests.conftest import (
-    setup_mock_dataset_owner_execute_result,
-    setup_mock_tenant_owner_execute_result,
-)
+from models import Account, Tenant, TenantAccountJoin
+from models.account import TenantAccountRole
+from models.dataset import Dataset, RateLimitLog
+from models.enums import ApiTokenType
+from models.model import ApiToken, App, AppMode, IconType
 
 
 def _configure_current_app_mock(mock_current_app):
@@ -34,6 +36,51 @@ def _configure_current_app_mock(mock_current_app):
     mock_current_app._get_current_object = Mock(return_value=Mock())
 
 
+def _session_proxy(session: Session) -> MagicMock:
+    """Emulate Flask-SQLAlchemy's callable scoped-session proxy around a test session."""
+    proxy = MagicMock(wraps=session)
+    proxy.return_value = session
+    return proxy
+
+
+def _api_token(*, tenant_id: str, app_id: str | None = None, token_type: ApiTokenType) -> ApiToken:
+    return ApiToken(
+        id=str(uuid.uuid4()),
+        tenant_id=tenant_id,
+        app_id=app_id,
+        type=token_type,
+        token="test_token",
+    )
+
+
+def _persist_workspace(session: Session) -> tuple[Tenant, Account, TenantAccountJoin]:
+    tenant = Tenant(name="Workspace")
+    account = Account(name="Owner", email=f"owner-{uuid.uuid4()}@example.com")
+    membership = TenantAccountJoin(
+        tenant_id=tenant.id,
+        account_id=account.id,
+        current=True,
+        role=TenantAccountRole.OWNER,
+    )
+    session.add_all([tenant, account, membership])
+    session.commit()
+    return tenant, account, membership
+
+
+def _app_model(*, tenant_id: str, enable_api: bool = True) -> App:
+    return App(
+        id=str(uuid.uuid4()),
+        tenant_id=tenant_id,
+        name="Service API App",
+        mode=AppMode.CHAT,
+        icon_type=IconType.EMOJI,
+        icon="chat",
+        icon_background="#FFFFFF",
+        enable_site=False,
+        enable_api=enable_api,
+    )
+
+
 class TestValidateAndGetApiToken:
     """Test suite for validate_and_get_api_token function"""
 
@@ -70,21 +117,24 @@ class TestValidateAndGetApiToken:
     def test_valid_token_returns_api_token(self, mock_fetch_token, mock_cache_cls, mock_record_usage, app: Flask):
         """Test that valid token returns the ApiToken object."""
         # Arrange
-        mock_api_token = Mock(spec=ApiToken)
-        mock_api_token.token = "valid_token_123"
-        mock_api_token.type = "app"
+        api_token = _api_token(
+            tenant_id=str(uuid.uuid4()),
+            app_id=str(uuid.uuid4()),
+            token_type=ApiTokenType.APP,
+        )
+        api_token.token = "valid_token_123"
 
         mock_cache_instance = Mock()
         mock_cache_instance.get.return_value = None  # Cache miss
         mock_cache_cls.get = mock_cache_instance.get
-        mock_fetch_token.return_value = mock_api_token
+        mock_fetch_token.return_value = api_token
 
         # Act
         with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer valid_token_123"}):
             result = validate_and_get_api_token("app")
 
         # Assert
-        assert result == mock_api_token
+        assert result == api_token
 
     @patch("controllers.service_api.wraps.record_token_usage")
     @patch("controllers.service_api.wraps.ApiTokenCache")
@@ -117,116 +167,124 @@ class TestValidateAppToken:
         return app
 
     @patch("controllers.service_api.wraps.user_logged_in")
-    @patch("controllers.service_api.wraps.db")
     @patch("controllers.service_api.wraps.validate_and_get_api_token")
     @patch("controllers.service_api.wraps.current_app")
+    @pytest.mark.parametrize(
+        "sqlite_session",
+        [(App, ApiToken, Tenant, Account, TenantAccountJoin)],
+        indirect=True,
+    )
     def test_valid_app_token_allows_access(
-        self, mock_current_app, mock_validate_token, mock_db, mock_user_logged_in, app
+        self,
+        mock_current_app,
+        mock_validate_token,
+        mock_user_logged_in,
+        app: Flask,
+        sqlite_session: Session,
     ):
         """Test that valid app token allows access to decorated view."""
         # Arrange
         _configure_current_app_mock(mock_current_app)
 
-        mock_api_token = Mock()
-        mock_api_token.app_id = str(uuid.uuid4())
-        mock_api_token.tenant_id = str(uuid.uuid4())
-        mock_validate_token.return_value = mock_api_token
-
-        mock_app = Mock()
-        mock_app.id = mock_api_token.app_id
-        mock_app.status = "normal"
-        mock_app.enable_api = True
-        mock_app.tenant_id = mock_api_token.tenant_id
-
-        mock_tenant = Mock()
-        mock_tenant.status = TenantStatus.NORMAL
-        mock_tenant.id = mock_api_token.tenant_id
-
-        mock_account = Mock()
-        mock_account.id = str(uuid.uuid4())
-
-        # Use side_effect to return app first, then tenant via session.get()
-        mock_db.session.get.side_effect = [mock_app, mock_tenant]
-
-        # Mock the tenant owner execute result (execute(select(...)).one_or_none())
-        setup_mock_tenant_owner_execute_result(mock_db, mock_tenant, mock_account)
+        tenant, account, _ = _persist_workspace(sqlite_session)
+        app_model = _app_model(tenant_id=tenant.id)
+        api_token = _api_token(tenant_id=tenant.id, app_id=app_model.id, token_type=ApiTokenType.APP)
+        sqlite_session.add_all([app_model, api_token])
+        sqlite_session.commit()
+        mock_validate_token.return_value = api_token
 
         @validate_app_token
         def protected_view(app_model):
             return {"success": True, "app_id": app_model.id}
 
         # Act
-        with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}):
+        with (
+            app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}),
+            patch("controllers.service_api.wraps.db.session", _session_proxy(sqlite_session)),
+        ):
             result = protected_view()
 
         # Assert
         assert result["success"] is True
-        assert result["app_id"] == mock_app.id
+        assert result["app_id"] == app_model.id
+        assert account.current_tenant_id == tenant.id
 
-    @patch("controllers.service_api.wraps.db")
     @patch("controllers.service_api.wraps.validate_and_get_api_token")
-    def test_app_not_found_raises_forbidden(self, mock_validate_token, mock_db, app: Flask):
+    @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
+    def test_app_not_found_raises_forbidden(self, mock_validate_token, app: Flask, sqlite_session: Session):
         """Test that Forbidden is raised when app no longer exists."""
         # Arrange
-        mock_api_token = Mock()
-        mock_api_token.app_id = str(uuid.uuid4())
-        mock_validate_token.return_value = mock_api_token
-
-        mock_db.session.get.return_value = None
+        api_token = _api_token(
+            tenant_id=str(uuid.uuid4()),
+            app_id=str(uuid.uuid4()),
+            token_type=ApiTokenType.APP,
+        )
+        mock_validate_token.return_value = api_token
 
         @validate_app_token
         def protected_view(**kwargs):
             return {"success": True}
 
         # Act & Assert
-        with app.test_request_context("/", method="GET"):
+        with (
+            app.test_request_context("/", method="GET"),
+            patch("controllers.service_api.wraps.db.session", sqlite_session),
+        ):
             with pytest.raises(Forbidden) as exc_info:
                 protected_view()
             assert "no longer exists" in str(exc_info.value)
 
-    @patch("controllers.service_api.wraps.db")
     @patch("controllers.service_api.wraps.validate_and_get_api_token")
-    def test_app_status_abnormal_raises_forbidden(self, mock_validate_token, mock_db, app: Flask):
+    @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
+    def test_app_status_abnormal_raises_forbidden(self, mock_validate_token, app: Flask, sqlite_session: Session):
         """Test that Forbidden is raised when app status is abnormal."""
         # Arrange
-        mock_api_token = Mock()
-        mock_api_token.app_id = str(uuid.uuid4())
-        mock_validate_token.return_value = mock_api_token
-
-        mock_app = Mock()
-        mock_app.status = "abnormal"
-        mock_db.session.get.return_value = mock_app
+        app_model = _app_model(tenant_id=str(uuid.uuid4()))
+        sqlite_session.add(app_model)
+        sqlite_session.commit()
+        app_model.status = "abnormal"
+        mock_validate_token.return_value = _api_token(
+            tenant_id=app_model.tenant_id,
+            app_id=app_model.id,
+            token_type=ApiTokenType.APP,
+        )
 
         @validate_app_token
         def protected_view(**kwargs):
             return {"success": True}
 
         # Act & Assert
-        with app.test_request_context("/", method="GET"):
+        with (
+            app.test_request_context("/", method="GET"),
+            patch("controllers.service_api.wraps.db.session", sqlite_session),
+        ):
             with pytest.raises(Forbidden) as exc_info:
                 protected_view()
             assert "status is abnormal" in str(exc_info.value)
 
-    @patch("controllers.service_api.wraps.db")
     @patch("controllers.service_api.wraps.validate_and_get_api_token")
-    def test_app_api_disabled_raises_forbidden(self, mock_validate_token, mock_db, app: Flask):
+    @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
+    def test_app_api_disabled_raises_forbidden(self, mock_validate_token, app: Flask, sqlite_session: Session):
         """Test that Forbidden is raised when app API is disabled."""
         # Arrange
-        mock_api_token = Mock()
-        mock_api_token.app_id = str(uuid.uuid4())
-        mock_validate_token.return_value = mock_api_token
-
-        mock_app = Mock()
-        mock_app.status = "normal"
-        mock_app.enable_api = False
-        mock_db.session.get.return_value = mock_app
+        app_model = _app_model(tenant_id=str(uuid.uuid4()), enable_api=False)
+        sqlite_session.add(app_model)
+        sqlite_session.commit()
+        mock_validate_token.return_value = _api_token(
+            tenant_id=app_model.tenant_id,
+            app_id=app_model.id,
+            token_type=ApiTokenType.APP,
+        )
 
         @validate_app_token
         def protected_view(**kwargs):
             return {"success": True}
 
         # Act & Assert
-        with app.test_request_context("/", method="GET"):
+        with (
+            app.test_request_context("/", method="GET"),
+            patch("controllers.service_api.wraps.db.session", sqlite_session),
+        ):
             with pytest.raises(Forbidden) as exc_info:
                 protected_view()
             assert "API service has been disabled" in str(exc_info.value)
@@ -468,26 +526,35 @@ class TestCloudEditionBillingRateLimitCheck:
 
     @patch("controllers.service_api.wraps.validate_and_get_api_token")
     @patch("controllers.service_api.wraps.FeatureService.get_knowledge_rate_limit")
-    @patch("controllers.service_api.wraps.db")
-    @patch("controllers.service_api.wraps.sessionmaker")
+    @pytest.mark.parametrize("sqlite_session", [(RateLimitLog,)], indirect=True)
     def test_rejects_over_rate_limit(
-        self, mock_sessionmaker, mock_db, mock_get_rate_limit, mock_validate_token, app: Flask
+        self,
+        mock_get_rate_limit,
+        mock_validate_token,
+        app: Flask,
+        sqlite_session: Session,
     ):
         """Test that Forbidden is raised when over rate limit."""
         # Arrange
-        mock_validate_token.return_value = Mock(tenant_id="tenant123")
+        tenant_id = str(uuid.uuid4())
+        mock_validate_token.return_value = _api_token(
+            tenant_id=tenant_id,
+            token_type=ApiTokenType.DATASET,
+        )
 
         mock_rate_limit = Mock()
         mock_rate_limit.enabled = True
         mock_rate_limit.limit = 10
         mock_rate_limit.subscription_plan = "pro"
         mock_get_rate_limit.return_value = mock_rate_limit
-        rate_limit_log_session = MagicMock()
-        session_factory = MagicMock()
-        session_factory.begin.return_value.__enter__.return_value = rate_limit_log_session
-        mock_sessionmaker.return_value = session_factory
 
-        with patch("controllers.service_api.wraps.redis_client") as mock_redis:
+        with (
+            patch("controllers.service_api.wraps.redis_client") as mock_redis,
+            patch(
+                "controllers.service_api.wraps.db",
+                SimpleNamespace(engine=sqlite_session.get_bind()),
+            ),
+        ):
             mock_redis.zcard.return_value = 15  # Over limit
 
             @cloud_edition_billing_rate_limit_check("knowledge", "dataset")
@@ -499,9 +566,12 @@ class TestCloudEditionBillingRateLimitCheck:
                 with pytest.raises(Forbidden) as exc_info:
                     knowledge_request()
                 assert "rate limit" in str(exc_info.value)
-            mock_sessionmaker.assert_called_once_with(bind=mock_db.engine, expire_on_commit=False)
-            rate_limit_log_session.add.assert_called_once()
-            mock_db.session.commit.assert_not_called()
+
+        persisted_logs = sqlite_session.scalars(select(RateLimitLog)).all()
+        assert len(persisted_logs) == 1
+        assert persisted_logs[0].tenant_id == tenant_id
+        assert persisted_logs[0].subscription_plan == "pro"
+        assert persisted_logs[0].operation == "knowledge"
 
 
 class TestValidateDatasetToken:
@@ -515,65 +585,62 @@ class TestValidateDatasetToken:
         return app
 
     @patch("controllers.service_api.wraps.user_logged_in")
-    @patch("controllers.service_api.wraps.db")
     @patch("controllers.service_api.wraps.validate_and_get_api_token")
     @patch("controllers.service_api.wraps.current_app")
-    def test_valid_dataset_token(self, mock_current_app, mock_validate_token, mock_db, mock_user_logged_in, app: Flask):
+    @pytest.mark.parametrize(
+        "sqlite_session",
+        [(Tenant, Account, TenantAccountJoin)],
+        indirect=True,
+    )
+    def test_valid_dataset_token(
+        self,
+        mock_current_app,
+        mock_validate_token,
+        mock_user_logged_in,
+        app: Flask,
+        sqlite_session: Session,
+    ):
         """Test that valid dataset token allows access."""
         # Arrange
         _configure_current_app_mock(mock_current_app)
 
-        tenant_id = str(uuid.uuid4())
-        mock_api_token = Mock()
-        mock_api_token.tenant_id = tenant_id
-        mock_validate_token.return_value = mock_api_token
-
-        mock_tenant = Mock()
-        mock_tenant.id = tenant_id
-        mock_tenant.status = TenantStatus.NORMAL
-
-        mock_ta = Mock()
-        mock_ta.account_id = str(uuid.uuid4())
-
-        mock_account = Mock()
-        mock_account.id = mock_ta.account_id
-        mock_account.current_tenant = mock_tenant
-
-        # Mock the tenant account join query (execute(select(...)).one_or_none())
-        setup_mock_dataset_owner_execute_result(mock_db, mock_tenant, mock_ta)
-
-        # Mock the account lookup via session.get()
-        mock_db.session.get.return_value = mock_account
+        tenant, account, _ = _persist_workspace(sqlite_session)
+        api_token = _api_token(tenant_id=tenant.id, token_type=ApiTokenType.DATASET)
+        mock_validate_token.return_value = api_token
 
         @validate_dataset_token
         def protected_view(tenant_id):
             return {"success": True, "tenant_id": tenant_id}
 
         # Act
-        with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}):
+        with (
+            app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}),
+            patch("controllers.service_api.wraps.db.session", _session_proxy(sqlite_session)),
+        ):
             result = protected_view()
 
         # Assert
         assert result["success"] is True
-        assert result["tenant_id"] == tenant_id
+        assert result["tenant_id"] == tenant.id
+        assert account.current_tenant_id == tenant.id
 
-    @patch("controllers.service_api.wraps.db")
     @patch("controllers.service_api.wraps.validate_and_get_api_token")
-    def test_dataset_not_found_raises_not_found(self, mock_validate_token, mock_db, app: Flask):
+    @pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
+    def test_dataset_not_found_raises_not_found(self, mock_validate_token, app: Flask, sqlite_session: Session):
         """Test that NotFound is raised when dataset doesn't exist."""
         # Arrange
-        mock_api_token = Mock()
-        mock_api_token.tenant_id = str(uuid.uuid4())
-        mock_validate_token.return_value = mock_api_token
-
-        mock_db.session.scalar.return_value = None
+        api_token = _api_token(tenant_id=str(uuid.uuid4()), token_type=ApiTokenType.DATASET)
+        mock_validate_token.return_value = api_token
 
         @validate_dataset_token
         def protected_view(dataset_id=None, **kwargs):
             return {"success": True}
 
         # Act & Assert
-        with app.test_request_context("/", method="GET"):
+        with (
+            app.test_request_context("/", method="GET"),
+            patch("controllers.service_api.wraps.db.session", sqlite_session),
+        ):
             with pytest.raises(NotFound) as exc_info:
                 protected_view(dataset_id=str(uuid.uuid4()))
             assert "Dataset not found" in str(exc_info.value)

api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py

@@ -46,7 +46,7 @@ class TestAdvancedChatAppGeneratorValidation:
         with pytest.raises(ValueError, match="query must be a string"):
             generator.generate(
                 app_model=SimpleNamespace(),
-                workflow=SimpleNamespace(),
+                workflow=SimpleNamespace(graph_dict={"nodes": []}),
                 user=SimpleNamespace(),
                 args={"inputs": {}, "query": 123},
                 invoke_from=InvokeFrom.WEB_APP,
@@ -186,7 +186,7 @@ class TestAdvancedChatAppGeneratorInternals:
 
         result = generator.generate(
             app_model=SimpleNamespace(id="app", tenant_id="tenant"),
-            workflow=SimpleNamespace(features_dict={}),
+            workflow=SimpleNamespace(features_dict={}, graph_dict={"nodes": []}),
             user=user,
             args={
                 "query": "hello",
@@ -442,6 +442,13 @@ class TestAdvancedChatAppGeneratorInternals:
             def start(self):
                 thread_data["started"] = True
 
+            def join(self, timeout):
+                thread_data["joined"] = True
+                thread_data["join_timeout"] = timeout
+
+            def is_alive(self):
+                return False
+
         monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.threading.Thread", _Thread)
         monkeypatch.setattr(
             "core.app.apps.advanced_chat.app_generator.db", SimpleNamespace(engine=object(), session=db_session)
@@ -475,6 +482,8 @@ class TestAdvancedChatAppGeneratorInternals:
 
         assert response["response"] == {"raw": True}
         assert thread_data["started"] is True
+        assert thread_data["joined"] is True
+        assert thread_data["join_timeout"] == 300
         assert "pause-layer" in thread_data["kwargs"]["graph_engine_layers"]
         assert generator._dialogue_count == 3
         assert init_records.call_args.kwargs["session"] is db_session
@@ -542,6 +551,13 @@ class TestAdvancedChatAppGeneratorInternals:
             def start(self):
                 thread_data["started"] = True
 
+            def join(self, timeout):
+                thread_data["joined"] = True
+                thread_data["join_timeout"] = timeout
+
+            def is_alive(self):
+                return False
+
         monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.threading.Thread", _Thread)
         monkeypatch.setattr(
             "core.app.apps.advanced_chat.app_generator.db", SimpleNamespace(engine=object(), session=db_session)
@@ -574,6 +590,8 @@ class TestAdvancedChatAppGeneratorInternals:
         init_records.assert_not_called()
         get_thread_messages_length.assert_called_once_with(conversation.id, session=db_session)
         assert thread_data["started"] is True
+        assert thread_data["joined"] is True
+        assert thread_data["join_timeout"] == 300
         db_session.commit.assert_not_called()
         db_session.refresh.assert_not_called()
         db_session.close.assert_called_once()
@@ -1191,7 +1209,7 @@ class TestAdvancedChatAppGeneratorInternals:
         monkeypatch.setattr(generator, "_generate", _fake_generate)
 
         app_model = SimpleNamespace(id="app", tenant_id="tenant")
-        workflow = SimpleNamespace(features_dict={})
+        workflow = SimpleNamespace(features_dict={}, graph_dict={"nodes": []})
         from models import Account
 
         user = Account(name="Tester", email="tester@example.com")
@@ -1271,7 +1289,7 @@ class TestAdvancedChatAppGeneratorInternals:
         monkeypatch.setattr(generator, "_generate", _fake_generate)
 
         app_model = SimpleNamespace(id="app", tenant_id="tenant")
-        workflow = SimpleNamespace(features_dict={})
+        workflow = SimpleNamespace(features_dict={}, graph_dict={"nodes": []})
         from models.model import EndUser
 
         user = EndUser(tenant_id="tenant", type="session", name="tester", session_id="session")

api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_resumption.py

@@ -22,7 +22,7 @@ def _build_converter() -> WorkflowResponseConverter:
         app_config=SimpleNamespace(app_id="app-1", tenant_id="tenant-1"),
         invoke_from=InvokeFrom.EXPLORE,
         files=[],
-        inputs={},
+        inputs={"userinput.files": []},
         workflow_execution_id="run-1",
         call_depth=0,
     )
@@ -54,3 +54,17 @@ def test_workflow_start_stream_response_carries_initial_reason():
         reason=WorkflowStartReason.INITIAL,
     )
     assert resp.data.reason is WorkflowStartReason.INITIAL
+
+
+def test_workflow_start_stream_response_exposes_only_canonical_file_input():
+    converter = _build_converter()
+
+    resp = converter.workflow_start_to_stream_response(
+        task_id="task-1",
+        workflow_run_id="run-1",
+        workflow_id="wf-1",
+        reason=WorkflowStartReason.INITIAL,
+    )
+
+    assert resp.data.inputs["userinput.files"] == []
+    assert "sys.files" not in resp.data.inputs

api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py

@@ -435,6 +435,7 @@ def test_generate_success_returns_converted(generator, mocker: MockerFixture):
     mocker.patch.object(module, "PipelineQueueManager", return_value=queue_manager)
 
     worker_thread = MagicMock()
+    worker_thread.is_alive.return_value = False
     mocker.patch.object(module.threading, "Thread", return_value=worker_thread)
 
     mocker.patch.object(generator, "_get_draft_var_saver_factory", return_value=MagicMock())
@@ -461,6 +462,7 @@ def test_generate_success_returns_converted(generator, mocker: MockerFixture):
     )
 
     assert result == "converted"
+    worker_thread.join.assert_called_once_with(timeout=300)
 
 
 def test_single_iteration_generate_validates_inputs(generator, mocker: MockerFixture):

api/tests/unit_tests/core/app/apps/test_base_app_generator.py

@@ -1,3 +1,6 @@
+import logging
+from unittest.mock import Mock
+
 import pytest
 
 from core.app.apps.base_app_generator import BaseAppGenerator
@@ -369,6 +372,58 @@ def test_validate_inputs_optional_file_with_empty_string_ignores_default():
 
 
 class TestBaseAppGeneratorExtras:
+    def test_wrap_stream_joins_worker_after_stream_exhaustion(self):
+        base_app_generator = BaseAppGenerator()
+        worker_thread = Mock()
+        worker_thread.is_alive.return_value = False
+
+        def response_stream():
+            yield {"event": "workflow_finished"}
+
+        managed_stream = base_app_generator._wrap_stream_with_worker_thread_join(
+            response_stream(),
+            worker_thread,
+        )
+
+        assert next(managed_stream) == {"event": "workflow_finished"}
+        worker_thread.join.assert_not_called()
+
+        with pytest.raises(StopIteration):
+            next(managed_stream)
+
+        worker_thread.join.assert_called_once_with(timeout=300)
+
+    def test_wrap_stream_joins_worker_when_stream_closes(self):
+        base_app_generator = BaseAppGenerator()
+        worker_thread = Mock()
+        worker_thread.is_alive.return_value = False
+
+        def response_stream():
+            yield {"event": "workflow_started"}
+            yield {"event": "workflow_finished"}
+
+        managed_stream = base_app_generator._wrap_stream_with_worker_thread_join(
+            response_stream(),
+            worker_thread,
+        )
+
+        assert next(managed_stream) == {"event": "workflow_started"}
+        managed_stream.close()
+
+        worker_thread.join.assert_called_once_with(timeout=300)
+
+    def test_join_worker_thread_warns_when_thread_remains_alive(self, caplog: pytest.LogCaptureFixture):
+        worker_thread = Mock()
+        worker_thread.name = "leaked-app-worker"
+        worker_thread.is_alive.return_value = True
+
+        with caplog.at_level(logging.WARNING, logger="core.app.apps.base_app_generator"):
+            BaseAppGenerator._join_worker_thread(worker_thread)
+
+        worker_thread.join.assert_called_once_with(timeout=300)
+        assert "Possible app worker thread leak" in caplog.text
+        assert "leaked-app-worker" in caplog.text
+
     def test_prepare_user_inputs_converts_files_and_lists(self, monkeypatch: pytest.MonkeyPatch):
         base_app_generator = BaseAppGenerator()
 

api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py

@@ -111,7 +111,7 @@ def test_generate_includes_parent_trace_context_in_extras(monkeypatch):
 
     result = generator.generate(
         app_model=SimpleNamespace(tenant_id="tenant-1", id="app-1"),
-        workflow=SimpleNamespace(features_dict={}),
+        workflow=SimpleNamespace(features_dict={}, graph_dict={"nodes": []}),
         user=SimpleNamespace(id="user-1", session_id="session-1"),
         args={
             "inputs": {"query": "hello"},
@@ -211,6 +211,13 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture):
         def start(self):
             return None
 
+        def join(self, timeout):
+            worker_kwargs["joined"] = True
+            worker_kwargs["join_timeout"] = timeout
+
+        def is_alive(self):
+            return False
+
     mocker.patch("core.app.apps.workflow.app_generator.threading.Thread", DummyThread)
 
     app_model = SimpleNamespace(mode="workflow", tenant_id="tenant")
@@ -244,6 +251,8 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture):
     assert result == "converted"
     assert worker_kwargs["kwargs"]["graph_engine_layers"] == ("base-layer", pause_layer)
     assert worker_kwargs["kwargs"]["graph_runtime_state"] is graph_runtime_state
+    assert worker_kwargs["joined"] is True
+    assert worker_kwargs["join_timeout"] == 300
     assert draft_saver_factory.call_args.kwargs["tenant_id"] == app_model.tenant_id
 
 
@@ -286,6 +295,8 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
 
     mocker.patch("core.app.apps.workflow.app_generator.WorkflowAppRunner", side_effect=runner_ctor)
 
+    worker_lifecycle: dict[str, bool] = {}
+
     class ImmediateThread:
         def __init__(self, target, kwargs):
             target(**kwargs)
@@ -293,6 +304,13 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
         def start(self):
             return None
 
+        def join(self, timeout):
+            worker_lifecycle["joined"] = True
+            worker_lifecycle["join_timeout"] = timeout
+
+        def is_alive(self):
+            return False
+
     mocker.patch("core.app.apps.workflow.app_generator.threading.Thread", ImmediateThread)
 
     mocker.patch(
@@ -331,5 +349,7 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
     )
 
     assert result == "raw-response"
+    assert worker_lifecycle["joined"] is True
+    assert worker_lifecycle["join_timeout"] == 300
     runner_instance.run.assert_called_once()
     queue_manager.graph_runtime_state = runtime_state

api/tests/unit_tests/core/app/apps/workflow/test_active_workflow_tasks.py

@@ -1,5 +1,9 @@
+import threading
+from collections.abc import Generator
+
 import pytest
 
+from core.app.apps.base_app_generator import BaseAppGenerator
 from core.app.apps.workflow.active_workflow_tasks import (
     active_workflow_task,
     get_active_workflow_task_count,
@@ -28,3 +32,51 @@ def test_active_workflow_task_rejects_duplicate_task_id() -> None:
         with pytest.raises(ValueError, match="already active"):
             with active_workflow_task("task-a"):
                 pass
+
+
+def test_managed_stream_waits_for_active_worker_cleanup() -> None:
+    worker_started = threading.Event()
+    release_worker = threading.Event()
+    stream_exhausted = threading.Event()
+    consumer_finished = threading.Event()
+    consumer_errors: list[BaseException] = []
+
+    def run_worker() -> None:
+        with active_workflow_task("task-a"):
+            worker_started.set()
+            release_worker.wait()
+
+    def response_stream() -> Generator[dict[str, str], None, None]:
+        yield {"event": "workflow_finished"}
+        stream_exhausted.set()
+
+    worker_thread = threading.Thread(target=run_worker)
+    worker_thread.start()
+    assert worker_started.wait(timeout=2)
+
+    managed_stream = BaseAppGenerator._wrap_stream_with_worker_thread_join(response_stream(), worker_thread)
+    assert next(managed_stream) == {"event": "workflow_finished"}
+
+    def finish_stream() -> None:
+        try:
+            list(managed_stream)
+        except BaseException as exc:
+            consumer_errors.append(exc)
+        finally:
+            consumer_finished.set()
+
+    consumer_thread = threading.Thread(target=finish_stream)
+    consumer_thread.start()
+    try:
+        assert stream_exhausted.wait(timeout=2)
+        assert not consumer_finished.is_set()
+        assert get_active_workflow_task_count() == 1
+    finally:
+        release_worker.set()
+        consumer_thread.join(timeout=2)
+        worker_thread.join(timeout=2)
+
+    assert not consumer_thread.is_alive()
+    assert not worker_thread.is_alive()
+    assert consumer_errors == []
+    assert get_active_workflow_task_count() == 0

api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py

@@ -15,6 +15,70 @@ from models.model import AppMode
 
 
 class TestWorkflowAppGeneratorValidation:
+    def test_generate_stream_joins_worker_after_response_exhaustion(self, monkeypatch: pytest.MonkeyPatch):
+        generator = WorkflowAppGenerator()
+        worker_thread = Mock()
+        worker_thread.is_alive.return_value = False
+        app_config = WorkflowUIBasedAppConfig(
+            tenant_id="tenant",
+            app_id="app",
+            app_mode=AppMode.WORKFLOW,
+            additional_features=AppAdditionalFeatures(),
+            variables=[],
+            workflow_id="workflow-id",
+        )
+        application_generate_entity = WorkflowAppGenerateEntity.model_construct(
+            task_id="task",
+            app_config=app_config,
+            inputs={},
+            files=[],
+            user_id="user",
+            stream=True,
+            invoke_from=InvokeFrom.WEB_APP,
+            extras={},
+        )
+
+        def response_stream():
+            yield {"event": "workflow_finished"}
+
+        monkeypatch.setattr(generator, "_bind_file_access_scope", lambda **kwargs: contextlib.nullcontext())
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.WorkflowAppQueueManager",
+            lambda **kwargs: SimpleNamespace(**kwargs),
+        )
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.current_app",
+            SimpleNamespace(_get_current_object=lambda: SimpleNamespace(name="flask")),
+        )
+        monkeypatch.setattr("core.app.apps.workflow.app_generator.contextvars.copy_context", lambda: "ctx")
+        monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", lambda **kwargs: worker_thread)
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.db",
+            SimpleNamespace(session=SimpleNamespace(close=Mock())),
+        )
+        monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory")
+        monkeypatch.setattr(generator, "_handle_response", lambda **kwargs: response_stream())
+        monkeypatch.setattr(
+            "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert",
+            lambda response, invoke_from: response,
+        )
+
+        managed_stream = generator._generate(
+            app_model=SimpleNamespace(mode=AppMode.WORKFLOW, tenant_id="tenant"),
+            workflow=SimpleNamespace(id="workflow-id"),
+            user=SimpleNamespace(id="user"),
+            application_generate_entity=application_generate_entity,
+            invoke_from=InvokeFrom.WEB_APP,
+            workflow_execution_repository=SimpleNamespace(),
+            workflow_node_execution_repository=SimpleNamespace(),
+            streaming=True,
+        )
+
+        worker_thread.start.assert_called_once_with()
+        worker_thread.join.assert_not_called()
+        assert list(managed_stream) == [{"event": "workflow_finished"}]
+        worker_thread.join.assert_called_once_with(timeout=300)
+
     def test_ensure_snippet_start_node_returns_original_for_non_snippet_workflow(self):
         workflow = SimpleNamespace(kind_or_standard="workflow")
         session = SimpleNamespace(scalar=Mock())
@@ -313,7 +377,7 @@ class TestWorkflowAppGeneratorGenerate:
 
         result = generator.generate(
             app_model=SimpleNamespace(id="app", tenant_id="tenant"),
-            workflow=SimpleNamespace(features_dict={}),
+            workflow=SimpleNamespace(features_dict={}, graph_dict={"nodes": []}),
             user=SimpleNamespace(id="user", session_id="session"),
             args={"inputs": {}, SKIP_PREPARE_USER_INPUTS_KEY: True},
             invoke_from=InvokeFrom.WEB_APP,

api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py

@@ -1,24 +1,95 @@
 """Unit tests for the message cycle manager optimization."""
 
 import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import pytest
 from flask import Flask, current_app
+from sqlalchemy import Engine, event, select
+from sqlalchemy.orm import Session, sessionmaker
 
 from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueRetrieverResourcesEvent
 from core.app.entities.task_entities import MessageStreamResponse, StreamEvent, TaskStateMetadata
+from core.app.task_pipeline import message_cycle_manager as message_cycle_manager_module
 from core.app.task_pipeline.message_cycle_manager import MessageCycleManager
 from core.rag.entities import RetrievalSourceMetadata
-from models.model import App, AppMode
+from graphon.file import FileTransferMethod, FileType
+from models import model as model_module
+from models.base import TypeBase
+from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo
+from models.model import App, AppMode, Conversation, MessageFile
 
 
-def _patch_create_session(mock_session):
-    session_cm = Mock()
-    session_cm.__enter__ = Mock(return_value=mock_session)
-    session_cm.__exit__ = Mock(return_value=False)
-    return patch("core.app.task_pipeline.message_cycle_manager.session_factory.create_session", return_value=session_cm)
+@dataclass(frozen=True)
+class _SQLiteDb:
+    engine: Engine
+    session: Session
+
+
+@pytest.fixture
+def cycle_db(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]:
+    """Bind request-owned and cycle-manager-owned sessions to isolated SQLite."""
+    TypeBase.metadata.create_all(
+        sqlite_engine,
+        tables=[App.__table__, Conversation.__table__, MessageFile.__table__],
+    )
+    owned_session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
+    with owned_session_factory() as request_session:
+        sqlite_db = _SQLiteDb(engine=sqlite_engine, session=request_session)
+        monkeypatch.setattr(message_cycle_manager_module, "db", sqlite_db)
+        monkeypatch.setattr(model_module, "db", sqlite_db)
+        monkeypatch.setattr(message_cycle_manager_module.session_factory, "create_session", owned_session_factory)
+        yield request_session
+
+
+def _app(*, app_id: str = "app-id", tenant_id: str = "tenant-1") -> App:
+    return App(
+        id=app_id,
+        tenant_id=tenant_id,
+        name="Test App",
+        description="",
+        mode=AppMode.CHAT,
+        enable_site=True,
+        enable_api=True,
+        max_active_requests=0,
+    )
+
+
+def _conversation(*, conversation_id: str = "conv-1", app_id: str = "app-id") -> Conversation:
+    conversation = Conversation(
+        app_id=app_id,
+        mode=AppMode.CHAT,
+        name="",
+        status="normal",
+        from_source=ConversationFromSource.API,
+        inputs={},
+    )
+    conversation.id = conversation_id
+    return conversation
+
+
+def _message_file(
+    *,
+    file_id: str = "file-1",
+    message_id: str = "test-message-id",
+    belongs_to: MessageFileBelongsTo | None = MessageFileBelongsTo.ASSISTANT,
+    url: str | None = "http://example.com/image.png",
+    file_type: FileType = FileType.IMAGE,
+) -> MessageFile:
+    message_file = MessageFile(
+        message_id=message_id,
+        type=file_type,
+        transfer_method=FileTransferMethod.TOOL_FILE,
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by="account-id",
+        belongs_to=belongs_to,
+        url=url,
+    )
+    message_file.id = file_id
+    return message_file
 
 
 class TestMessageCycleManagerOptimization:
@@ -37,30 +108,22 @@ class TestMessageCycleManagerOptimization:
         task_state = Mock()
         return MessageCycleManager(application_generate_entity=mock_application_generate_entity, task_state=task_state)
 
-    def test_get_message_event_type_with_assistant_file(self, message_cycle_manager):
+    def test_get_message_event_type_with_assistant_file(self, message_cycle_manager, cycle_db: Session):
         """Test get_message_event_type returns MESSAGE_FILE when message has assistant-generated files.
 
         This ensures that AI-generated images (belongs_to='assistant') trigger the MESSAGE_FILE event,
         allowing the frontend to properly display generated image files with url field.
         """
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
-            # Setup mock session and message file
-            mock_session = Mock()
-            mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
+        cycle_db.add(_message_file())
+        cycle_db.commit()
 
-            mock_message_file = Mock()
-            mock_message_file.belongs_to = "assistant"
-            mock_session.scalar.return_value = mock_message_file
+        with current_app.app_context():
+            result = message_cycle_manager.get_message_event_type("test-message-id")
 
-            # Execute
-            with current_app.app_context():
-                result = message_cycle_manager.get_message_event_type("test-message-id")
+        assert result == StreamEvent.MESSAGE_FILE
+        assert "test-message-id" in message_cycle_manager._message_has_file
 
-            # Assert
-            assert result == StreamEvent.MESSAGE_FILE
-            mock_session.scalar.assert_called_once()
-
-    def test_get_message_event_type_with_user_file(self, message_cycle_manager):
+    def test_get_message_event_type_with_user_file(self, message_cycle_manager, cycle_db: Session):
         """Test get_message_event_type returns MESSAGE when message only has user-uploaded files.
 
         This is a regression test for the issue where user-uploaded images (belongs_to='user')
@@ -68,90 +131,81 @@ class TestMessageCycleManagerOptimization:
         resulting in broken images in the chat UI. The query filters for belongs_to='assistant',
         so when only user files exist, the database query returns None, resulting in MESSAGE event type.
         """
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
-            # Setup mock session and message file
-            mock_session = Mock()
-            mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
+        cycle_db.add(_message_file(belongs_to=MessageFileBelongsTo.USER))
+        cycle_db.commit()
 
-            # When querying for assistant files with only user files present, return None
-            # (simulates database query with belongs_to='assistant' filter returning no results)
-            mock_session.scalar.return_value = None
+        with current_app.app_context():
+            result = message_cycle_manager.get_message_event_type("test-message-id")
 
-            # Execute
-            with current_app.app_context():
-                result = message_cycle_manager.get_message_event_type("test-message-id")
+        assert result == StreamEvent.MESSAGE
+        assert "test-message-id" not in message_cycle_manager._message_has_file
 
-            # Assert
-            assert result == StreamEvent.MESSAGE
-            mock_session.scalar.assert_called_once()
-
-    def test_get_message_event_type_without_message_file(self, message_cycle_manager):
+    def test_get_message_event_type_without_message_file(self, message_cycle_manager, cycle_db: Session):
         """Test get_message_event_type returns MESSAGE when message has no files."""
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
-            # Setup mock session and no message file
-            mock_session = Mock()
-            mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
-            # Current implementation uses session.scalar(select(...))
-            mock_session.scalar.return_value = None
+        assert list(cycle_db.scalars(select(MessageFile)).all()) == []
 
-            # Execute
-            with current_app.app_context():
-                result = message_cycle_manager.get_message_event_type("test-message-id")
+        with current_app.app_context():
+            result = message_cycle_manager.get_message_event_type("test-message-id")
 
-            # Assert
-            assert result == StreamEvent.MESSAGE
-            mock_session.scalar.assert_called_once()
+        assert result == StreamEvent.MESSAGE
 
-    def test_get_message_event_type_uses_cache_without_query(self, message_cycle_manager):
+    def test_get_message_event_type_uses_cache_without_query(
+        self, message_cycle_manager, cycle_db: Session, sqlite_engine: Engine
+    ):
         """Return MESSAGE_FILE directly from in-memory cache without opening a DB session."""
         message_cycle_manager._message_has_file.add("cached-message")
+        statements: list[str] = []
 
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
+        def record_statement(_conn, _cursor, statement, _parameters, _context, _executemany) -> None:
+            statements.append(statement)
+
+        event.listen(sqlite_engine, "before_cursor_execute", record_statement)
+        try:
             result = message_cycle_manager.get_message_event_type("cached-message")
+        finally:
+            event.remove(sqlite_engine, "before_cursor_execute", record_statement)
 
         assert result == StreamEvent.MESSAGE_FILE
-        mock_session_factory.create_session.assert_not_called()
+        assert statements == []
 
-    def test_message_to_stream_response_with_precomputed_event_type(self, message_cycle_manager):
+    def test_message_to_stream_response_with_precomputed_event_type(self, message_cycle_manager, cycle_db: Session):
         """MessageCycleManager.message_to_stream_response expects a valid event_type; callers should precompute it."""
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
-            # Setup mock session and message file
-            mock_session = Mock()
-            mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
+        cycle_db.add(_message_file())
+        cycle_db.commit()
 
-            mock_message_file = Mock()
-            mock_message_file.belongs_to = "assistant"
-            mock_session.scalar.return_value = mock_message_file
+        with current_app.app_context():
+            event_type = message_cycle_manager.get_message_event_type("test-message-id")
+            result = message_cycle_manager.message_to_stream_response(
+                answer="Hello world", message_id="test-message-id", event_type=event_type
+            )
 
-            # Execute: compute event type once, then pass to message_to_stream_response
-            with current_app.app_context():
-                event_type = message_cycle_manager.get_message_event_type("test-message-id")
-                result = message_cycle_manager.message_to_stream_response(
-                    answer="Hello world", message_id="test-message-id", event_type=event_type
-                )
+        assert isinstance(result, MessageStreamResponse)
+        assert result.answer == "Hello world"
+        assert result.id == "test-message-id"
+        assert result.event == StreamEvent.MESSAGE_FILE
 
-            # Assert
-            assert isinstance(result, MessageStreamResponse)
-            assert result.answer == "Hello world"
-            assert result.id == "test-message-id"
-            assert result.event == StreamEvent.MESSAGE_FILE
-            mock_session.scalar.assert_called_once()
-
-    def test_message_to_stream_response_with_event_type_skips_query(self, message_cycle_manager):
+    def test_message_to_stream_response_with_event_type_skips_query(
+        self, message_cycle_manager, cycle_db: Session, sqlite_engine: Engine
+    ):
         """Test that message_to_stream_response skips database query when event_type is provided."""
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
-            # Execute with event_type provided
+        statements: list[str] = []
+
+        def record_statement(_conn, _cursor, statement, _parameters, _context, _executemany) -> None:
+            statements.append(statement)
+
+        event.listen(sqlite_engine, "before_cursor_execute", record_statement)
+        try:
             result = message_cycle_manager.message_to_stream_response(
                 answer="Hello world", message_id="test-message-id", event_type=StreamEvent.MESSAGE
             )
+        finally:
+            event.remove(sqlite_engine, "before_cursor_execute", record_statement)
 
-            # Assert
-            assert isinstance(result, MessageStreamResponse)
-            assert result.answer == "Hello world"
-            assert result.id == "test-message-id"
-            assert result.event == StreamEvent.MESSAGE
-            # Should not open a session when event_type is provided
-            mock_session_factory.create_session.assert_not_called()
+        assert isinstance(result, MessageStreamResponse)
+        assert result.answer == "Hello world"
+        assert result.id == "test-message-id"
+        assert result.event == StreamEvent.MESSAGE
+        assert statements == []
 
     def test_message_to_stream_response_with_from_variable_selector(self, message_cycle_manager):
         """Test message_to_stream_response with from_variable_selector parameter."""
@@ -168,40 +222,32 @@ class TestMessageCycleManagerOptimization:
         assert result.from_variable_selector == ["var1", "var2"]
         assert result.event == StreamEvent.MESSAGE
 
-    def test_optimization_usage_example(self, message_cycle_manager):
+    def test_optimization_usage_example(self, message_cycle_manager, cycle_db: Session, sqlite_engine: Engine):
         """Test the optimization pattern that should be used by callers."""
-        # Step 1: Get event type once (this queries database)
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
-            mock_session = Mock()
-            mock_session_factory.create_session.return_value.__enter__.return_value = mock_session
-            # Current implementation uses session.scalar(select(...))
-            mock_session.scalar.return_value = None  # No files
+        statements: list[str] = []
+
+        def record_statement(_conn, _cursor, statement, _parameters, _context, _executemany) -> None:
+            statements.append(statement)
+
+        event.listen(sqlite_engine, "before_cursor_execute", record_statement)
+        try:
             with current_app.app_context():
                 event_type = message_cycle_manager.get_message_event_type("test-message-id")
-
-        # Should open session once
-        mock_session_factory.create_session.assert_called_once()
-        assert event_type == StreamEvent.MESSAGE
-
-        # Step 2: Use event_type for multiple calls (no additional queries)
-        with patch("core.app.task_pipeline.message_cycle_manager.session_factory") as mock_session_factory:
-            mock_session_factory.create_session.return_value.__enter__.return_value = Mock()
-
             chunk1_response = message_cycle_manager.message_to_stream_response(
                 answer="Chunk 1", message_id="test-message-id", event_type=event_type
             )
-
             chunk2_response = message_cycle_manager.message_to_stream_response(
                 answer="Chunk 2", message_id="test-message-id", event_type=event_type
             )
+        finally:
+            event.remove(sqlite_engine, "before_cursor_execute", record_statement)
 
-            # Should not open session again when event_type provided
-            mock_session_factory.create_session.assert_not_called()
-
-            assert chunk1_response.event == StreamEvent.MESSAGE
-            assert chunk2_response.event == StreamEvent.MESSAGE
-            assert chunk1_response.answer == "Chunk 1"
-            assert chunk2_response.answer == "Chunk 2"
+        assert event_type == StreamEvent.MESSAGE
+        assert len([statement for statement in statements if statement.lstrip().upper().startswith("SELECT")]) == 1
+        assert chunk1_response.event == StreamEvent.MESSAGE
+        assert chunk2_response.event == StreamEvent.MESSAGE
+        assert chunk1_response.answer == "Chunk 1"
+        assert chunk2_response.answer == "Chunk 2"
 
     def test_generate_conversation_name_returns_none_for_completion(self, message_cycle_manager):
         """Return None when completion entities are used for conversation naming.
@@ -269,51 +315,38 @@ class TestMessageCycleManagerOptimization:
         assert message_cycle_manager._application_generate_entity.is_new_conversation is False
         mock_timer.assert_not_called()
 
-    def test_generate_conversation_name_worker_returns_when_conversation_missing(self, message_cycle_manager):
+    def test_generate_conversation_name_worker_returns_when_conversation_missing(
+        self, message_cycle_manager, cycle_db: Session
+    ):
         """Return early when the conversation cannot be found."""
         flask_app = Flask(__name__)
-        db_session = Mock()
-        db_session.scalar.return_value = None
+        assert list(cycle_db.scalars(select(Conversation)).all()) == []
 
-        with _patch_create_session(db_session):
-            message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-missing", "hello")
+        message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-missing", "hello")
 
-        db_session.commit.assert_not_called()
+        assert list(cycle_db.scalars(select(Conversation)).all()) == []
 
-    def test_generate_conversation_name_worker_returns_when_app_missing(self, message_cycle_manager):
+    def test_generate_conversation_name_worker_returns_when_app_missing(self, message_cycle_manager, cycle_db: Session):
         """Return early when non-completion conversation has no app relation."""
         flask_app = Flask(__name__)
-        conversation = SimpleNamespace(mode=AppMode.CHAT, app=None, app_id="app-id")
-        db_session = Mock()
-        db_session.scalar.return_value = conversation
-        db_session.get.return_value = None
+        conversation = _conversation()
+        cycle_db.add(conversation)
+        cycle_db.commit()
 
-        with _patch_create_session(db_session):
-            message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", "hello")
+        message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", "hello")
 
-        db_session.commit.assert_not_called()
+        assert cycle_db.get(Conversation, "conv-1").name == ""
+        assert cycle_db.get(App, "app-id") is None
 
-    def test_generate_conversation_name_worker_uses_cached_name(self, message_cycle_manager):
+    def test_generate_conversation_name_worker_uses_cached_name(
+        self, message_cycle_manager, cycle_db: Session, sqlite_engine: Engine
+    ):
         """Use cached conversation name when present and avoid LLM call."""
         flask_app = Flask(__name__)
-
-        class ConversationWithPoisonedApp:
-            mode = AppMode.CHAT
-            app_id = "app-id"
-            name = ""
-
-            @property
-            def app(self):
-                raise AssertionError("conversation.app must not open an implicit session")
-
-        conversation = ConversationWithPoisonedApp()
-        app_model = SimpleNamespace(tenant_id="tenant-1")
-        db_session = Mock()
-        db_session.scalar.return_value = conversation
-        db_session.get.return_value = app_model
+        cycle_db.add_all([_app(), _conversation()])
+        cycle_db.commit()
 
         with (
-            _patch_create_session(db_session) as create_session,
             patch("core.app.task_pipeline.message_cycle_manager.redis_client") as mock_redis,
             patch("core.app.task_pipeline.message_cycle_manager.LLMGenerator") as mock_llm_generator,
         ):
@@ -321,27 +354,23 @@ class TestMessageCycleManagerOptimization:
 
             message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", "hello")
 
+        assert cycle_db.in_transaction() is False
+        with Session(sqlite_engine) as verification_session:
+            conversation = verification_session.get(Conversation, "conv-1")
+        assert conversation is not None
         assert conversation.name == "cached-title"
-        create_session.assert_called_once_with()
-        db_session.get.assert_called_once_with(App, "app-id")
-        db_session.commit.assert_called_once()
         mock_llm_generator.generate_conversation_name.assert_not_called()
         mock_redis.setex.assert_not_called()
 
-    def test_generate_conversation_name_worker_generates_and_caches_name(self, message_cycle_manager):
+    def test_generate_conversation_name_worker_generates_and_caches_name(
+        self, message_cycle_manager, cycle_db: Session, sqlite_engine: Engine
+    ):
         """Generate conversation name and write it to redis cache on cache miss."""
         flask_app = Flask(__name__)
-        conversation = SimpleNamespace(
-            mode=AppMode.CHAT,
-            app=SimpleNamespace(tenant_id="tenant-1"),
-            app_id="app-id",
-            name="",
-        )
-        db_session = Mock()
-        db_session.scalar.return_value = conversation
+        cycle_db.add_all([_app(), _conversation()])
+        cycle_db.commit()
 
         with (
-            _patch_create_session(db_session),
             patch("core.app.task_pipeline.message_cycle_manager.redis_client") as mock_redis,
             patch("core.app.task_pipeline.message_cycle_manager.LLMGenerator") as mock_llm_generator,
         ):
@@ -350,27 +379,27 @@ class TestMessageCycleManagerOptimization:
 
             message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", "hello")
 
+        assert cycle_db.in_transaction() is False
+        with Session(sqlite_engine) as verification_session:
+            conversation = verification_session.get(Conversation, "conv-1")
+        assert conversation is not None
         assert conversation.name == "generated-title"
-        db_session.commit.assert_called_once()
         mock_redis.setex.assert_called_once()
 
     def test_generate_conversation_name_worker_falls_back_when_generation_fails(
-        self, message_cycle_manager, caplog: pytest.LogCaptureFixture
+        self,
+        message_cycle_manager,
+        cycle_db: Session,
+        sqlite_engine: Engine,
+        caplog: pytest.LogCaptureFixture,
     ):
         """Fallback to truncated query when LLM generation fails."""
         flask_app = Flask(__name__)
-        conversation = SimpleNamespace(
-            mode=AppMode.CHAT,
-            app=SimpleNamespace(tenant_id="tenant-1"),
-            app_id="app-id",
-            name="",
-        )
-        db_session = Mock()
-        db_session.scalar.return_value = conversation
+        cycle_db.add_all([_app(), _conversation()])
+        cycle_db.commit()
         long_query = "q" * 60
 
         with (
-            _patch_create_session(db_session),
             patch("core.app.task_pipeline.message_cycle_manager.redis_client") as mock_redis,
             patch("core.app.task_pipeline.message_cycle_manager.LLMGenerator") as mock_llm_generator,
             patch("core.app.task_pipeline.message_cycle_manager.dify_config") as mock_dify_config,
@@ -382,8 +411,11 @@ class TestMessageCycleManagerOptimization:
             with caplog.at_level(logging.ERROR, logger="core.app.task_pipeline.message_cycle_manager"):
                 message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", long_query)
 
+        assert cycle_db.in_transaction() is False
+        with Session(sqlite_engine) as verification_session:
+            conversation = verification_session.get(Conversation, "conv-1")
+        assert conversation is not None
         assert conversation.name == (long_query[:47] + "...")
-        db_session.commit.assert_called_once()
         assert any(record.levelno == logging.ERROR for record in caplog.records)
 
     def test_handle_annotation_reply_sets_metadata(self, message_cycle_manager):
@@ -454,33 +486,25 @@ class TestMessageCycleManagerOptimization:
         assert message_cycle_manager._task_state.metadata.retriever_resources[0].position == 1
         assert message_cycle_manager._task_state.metadata.retriever_resources[1].position == 2
 
-    def test_message_file_to_stream_response_builds_signed_url(self, message_cycle_manager):
+    def test_message_file_to_stream_response_builds_signed_url(self, message_cycle_manager, cycle_db: Session):
         """Build a stream response with a signed tool file URL.
 
-        Args: message_cycle_manager with mocked Session/db and sign_tool_file.
+        Args: message_cycle_manager with a persisted MessageFile and mocked sign_tool_file.
         Returns: MessageStreamResponse with signed url and belongs_to normalized to user.
         Side effects: Calls sign_tool_file for tool file ids.
         """
         message_cycle_manager._application_generate_entity.task_id = "task-1"
-
-        message_file = SimpleNamespace(
-            id="file-1",
-            type="image",
-            belongs_to=None,
-            url="tool://file.verylongextension",
-            message_id="msg-1",
+        cycle_db.add(
+            _message_file(
+                file_id="file-1",
+                message_id="msg-1",
+                belongs_to=None,
+                url="tool://file.verylongextension",
+            )
         )
+        cycle_db.commit()
 
-        session = Mock()
-        session.scalar.return_value = message_file
-
-        with (
-            patch("core.app.task_pipeline.message_cycle_manager.Session") as mock_session_cls,
-            patch("core.app.task_pipeline.message_cycle_manager.sign_tool_file") as mock_sign,
-            patch("core.app.task_pipeline.message_cycle_manager.db") as mock_db,
-        ):
-            mock_db.engine = Mock()
-            mock_session_cls.return_value.__enter__.return_value = session
+        with patch("core.app.task_pipeline.message_cycle_manager.sign_tool_file") as mock_sign:
             mock_sign.return_value = "signed-url"
 
             response = message_cycle_manager.message_file_to_stream_response(SimpleNamespace(message_file_id="file-1"))
@@ -514,56 +538,42 @@ class TestMessageCycleManagerOptimization:
         assert len(message_cycle_manager._task_state.metadata.retriever_resources) == 1
         assert message_cycle_manager._task_state.metadata.retriever_resources[0].position == 1
 
-    def test_message_file_to_stream_response_uses_http_url_directly(self, message_cycle_manager):
+    def test_message_file_to_stream_response_uses_http_url_directly(self, message_cycle_manager, cycle_db: Session):
         """Use original URL when message file URL is already HTTP."""
         message_cycle_manager._application_generate_entity.task_id = "task-http"
-        message_file = SimpleNamespace(
-            id="file-http",
-            type="image",
-            belongs_to="assistant",
-            url="http://example.com/pic.png",
-            message_id="msg-http",
-        )
-
-        session = Mock()
-        session.scalar.return_value = message_file
-
-        with (
-            patch("core.app.task_pipeline.message_cycle_manager.Session") as mock_session_cls,
-            patch("core.app.task_pipeline.message_cycle_manager.db") as mock_db,
-        ):
-            mock_db.engine = Mock()
-            mock_session_cls.return_value.__enter__.return_value = session
-
-            response = message_cycle_manager.message_file_to_stream_response(
-                SimpleNamespace(message_file_id="file-http")
+        cycle_db.add(
+            _message_file(
+                file_id="file-http",
+                message_id="msg-http",
+                belongs_to=MessageFileBelongsTo.ASSISTANT,
+                url="http://example.com/pic.png",
             )
+        )
+        cycle_db.commit()
+
+        response = message_cycle_manager.message_file_to_stream_response(SimpleNamespace(message_file_id="file-http"))
 
         assert response is not None
         assert response.url == "http://example.com/pic.png"
         assert "msg-http" in message_cycle_manager._message_has_file
 
-    def test_message_file_to_stream_response_defaults_extension_to_bin_without_dot(self, message_cycle_manager):
+    def test_message_file_to_stream_response_defaults_extension_to_bin_without_dot(
+        self, message_cycle_manager, cycle_db: Session
+    ):
         """Default tool file extension to .bin when URL has no extension part."""
         message_cycle_manager._application_generate_entity.task_id = "task-bin"
-        message_file = SimpleNamespace(
-            id="file-bin",
-            type="file",
-            belongs_to="assistant",
-            url="tool-file-id",
-            message_id="msg-bin",
+        cycle_db.add(
+            _message_file(
+                file_id="file-bin",
+                message_id="msg-bin",
+                belongs_to=MessageFileBelongsTo.ASSISTANT,
+                url="tool-file-id",
+                file_type=FileType.CUSTOM,
+            )
         )
+        cycle_db.commit()
 
-        session = Mock()
-        session.scalar.return_value = message_file
-
-        with (
-            patch("core.app.task_pipeline.message_cycle_manager.Session") as mock_session_cls,
-            patch("core.app.task_pipeline.message_cycle_manager.sign_tool_file") as mock_sign,
-            patch("core.app.task_pipeline.message_cycle_manager.db") as mock_db,
-        ):
-            mock_db.engine = Mock()
-            mock_session_cls.return_value.__enter__.return_value = session
+        with patch("core.app.task_pipeline.message_cycle_manager.sign_tool_file") as mock_sign:
             mock_sign.return_value = "signed-bin-url"
 
             response = message_cycle_manager.message_file_to_stream_response(
@@ -574,19 +584,13 @@ class TestMessageCycleManagerOptimization:
         assert response.url == "signed-bin-url"
         mock_sign.assert_called_once_with(tool_file_id="tool-file-id", extension=".bin")
 
-    def test_message_file_to_stream_response_returns_none_when_file_missing(self, message_cycle_manager):
+    def test_message_file_to_stream_response_returns_none_when_file_missing(
+        self, message_cycle_manager, cycle_db: Session
+    ):
         """Return None when message file lookup does not find a record."""
-        session = Mock()
-        session.scalar.return_value = None
+        assert list(cycle_db.scalars(select(MessageFile)).all()) == []
 
-        with (
-            patch("core.app.task_pipeline.message_cycle_manager.Session") as mock_session_cls,
-            patch("core.app.task_pipeline.message_cycle_manager.db") as mock_db,
-        ):
-            mock_db.engine = Mock()
-            mock_session_cls.return_value.__enter__.return_value = session
-
-            response = message_cycle_manager.message_file_to_stream_response(SimpleNamespace(message_file_id="missing"))
+        response = message_cycle_manager.message_file_to_stream_response(SimpleNamespace(message_file_id="missing"))
 
         assert response is None
 

api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py

@@ -1,7 +1,57 @@
-import sys
+from datetime import datetime, timedelta
 from unittest.mock import MagicMock, patch
 
+import pytest
+from sqlalchemy import event
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import Session, sessionmaker
+
+import core.llm_generator.llm_generator as generator_module
 from core.llm_generator.llm_generator import LLMGenerator, _parse_string_list
+from core.model_manager import ModelInstance, ModelManager
+from core.workflow.generator import tool_catalogue as tool_catalogue_module
+from core.workflow.generator.tool_catalogue import ToolCatalogueEntry
+from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
+from graphon.model_runtime.entities.message_entities import AssistantPromptMessage
+from models.dataset import Dataset
+from services.workflow_service import WorkflowService
+
+
+@pytest.fixture
+def dataset_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session:
+    """Bind the real SQLite session to the production database extension."""
+
+    monkeypatch.setattr(generator_module.db, "session", sqlite_session)
+    return sqlite_session
+
+
+def _llm_result(content: str) -> LLMResult:
+    """Build a real non-streaming LLM response around deterministic test content."""
+
+    return LLMResult(
+        model="test-model",
+        message=AssistantPromptMessage(content=content),
+        usage=LLMUsage.empty_usage(),
+    )
+
+
+def _model_manager() -> tuple[MagicMock, MagicMock]:
+    """Build spec-constrained mocks for the model-manager boundary and its default model."""
+
+    model_manager = MagicMock(spec=ModelManager)
+    model_instance = MagicMock(spec=ModelInstance)
+    model_manager.get_default_model_instance.return_value = model_instance
+    return model_manager, model_instance
+
+
+def _dataset(*, dataset_id: str, tenant_id: str, name: str, created_at: datetime) -> Dataset:
+    return Dataset(
+        id=dataset_id,
+        tenant_id=tenant_id,
+        name=name,
+        created_by="account-id",
+        created_at=created_at,
+    )
 
 
 class TestParseStringList:
@@ -34,95 +84,115 @@ class TestParseStringList:
 class TestGenerateWorkflowInstructionSuggestions:
     @patch("core.llm_generator.llm_generator.ModelManager.for_tenant")
     def test_no_default_model(self, mock_for_tenant):
-        mock_for_tenant.return_value.get_default_model_instance.side_effect = Exception("No model")
+        model_manager, _ = _model_manager()
+        model_manager.get_default_model_instance.side_effect = RuntimeError("no default model")
+        mock_for_tenant.return_value = model_manager
+
         assert LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") == []
 
     @patch("core.llm_generator.llm_generator.ModelManager.for_tenant")
     @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context")
     def test_llm_success(self, mock_build_context, mock_for_tenant):
         mock_build_context.return_value = "context"
-
-        mock_model = MagicMock()
-        mock_model.invoke_llm.return_value = MagicMock()
-        mock_model.invoke_llm.return_value.message.get_text_content.return_value = '["idea 1", "idea 2"]'
-
-        mock_for_tenant.return_value.get_default_model_instance.return_value = mock_model
+        model_manager, model_instance = _model_manager()
+        model_instance.invoke_llm.return_value = _llm_result('["idea 1", "idea 2"]')
+        mock_for_tenant.return_value = model_manager
 
         result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow")
         assert result == ["idea 1", "idea 2"]
+        model_instance.invoke_llm.assert_called_once()
 
     @patch("core.llm_generator.llm_generator.ModelManager.for_tenant")
     @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context")
     def test_llm_error(self, mock_build_context, mock_for_tenant):
         mock_build_context.return_value = "context"
+        model_manager, model_instance = _model_manager()
+        model_instance.invoke_llm.side_effect = RuntimeError("API error")
+        mock_for_tenant.return_value = model_manager
 
-        mock_model = MagicMock()
-        mock_model.invoke_llm.side_effect = Exception("API error")
-
-        mock_for_tenant.return_value.get_default_model_instance.return_value = mock_model
-
-        assert LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") == []
+        result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow")
+        assert result == []
+        model_instance.invoke_llm.assert_called_once()
 
     @patch("core.llm_generator.llm_generator.ModelManager.for_tenant")
     @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context")
     def test_llm_bad_output(self, mock_build_context, mock_for_tenant):
         mock_build_context.return_value = "context"
+        model_manager, model_instance = _model_manager()
+        model_instance.invoke_llm.return_value = _llm_result("Not a list")
+        mock_for_tenant.return_value = model_manager
 
-        mock_model = MagicMock()
-        mock_model.invoke_llm.return_value = MagicMock()
-        mock_model.invoke_llm.return_value.message.get_text_content.return_value = "Not a list"
-
-        mock_for_tenant.return_value.get_default_model_instance.return_value = mock_model
-
-        assert LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") == []
+        result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow")
+        assert result == []
+        model_instance.invoke_llm.assert_called_once()
 
 
+@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True)
 class TestBuildSuggestionContext:
-    @patch("core.llm_generator.llm_generator.db.session.scalars")
-    def test_both_success(self, mock_scalars, monkeypatch):
-        mock_scalars.return_value.all.return_value = ["kb1", "kb2"]
+    def test_both_success(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch):
+        now = datetime.now()
+        dataset_session.add_all(
+            (
+                _dataset(dataset_id="kb-1", tenant_id="tenant", name="kb1", created_at=now),
+                _dataset(
+                    dataset_id="kb-2",
+                    tenant_id="tenant",
+                    name="kb2",
+                    created_at=now - timedelta(seconds=1),
+                ),
+                _dataset(dataset_id="other-kb", tenant_id="other", name="private", created_at=now),
+            )
+        )
+        dataset_session.commit()
 
-        # ``_build_suggestion_context`` imports the tool catalogue lazily, so we
-        # stub the module in ``sys.modules``. Use ``monkeypatch.setitem`` so the
-        # ORIGINAL module is RESTORED on teardown — a bare ``del`` would evict it
-        # from sys.modules entirely, after which a sibling test that imported
-        # ``build_tool_catalogue`` at collection time (e.g. test_tool_catalogue)
-        # diverges from a freshly re-imported module and its @patch targets stop
-        # applying, silently breaking it under xdist.
-        mock_tool_catalogue = MagicMock()
-        mock_tool_catalogue.build_tool_catalogue.return_value = "catalog"
-        mock_tool_catalogue.format_tool_catalogue.return_value = "tool1\ntool2"
-        monkeypatch.setitem(sys.modules, "core.workflow.generator.tool_catalogue", mock_tool_catalogue)
+        def build_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]:
+            return [
+                ToolCatalogueEntry(
+                    provider_name="provider",
+                    provider_type="builtin",
+                    plugin_id="",
+                    tool_name="tool1",
+                    tool_label="tool1",
+                    description="First tool",
+                ),
+                ToolCatalogueEntry(
+                    provider_name="provider",
+                    provider_type="builtin",
+                    plugin_id="",
+                    tool_name="tool2",
+                    tool_label="tool2",
+                    description="Second tool",
+                ),
+            ]
+
+        # Keep the real module and formatter; only isolate provider/plugin discovery.
+        monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", build_tool_catalogue)
 
         result = LLMGenerator._build_suggestion_context("tenant")
         assert "Knowledge bases:\n- kb1\n- kb2" in result
-        assert "Installed tools:\ntool1\ntool2" in result
+        assert "Installed tools:\n- provider/tool1 — First tool\n- provider/tool2 — Second tool" in result
 
-    @patch("core.llm_generator.llm_generator.db.session.scalars")
-    def test_both_fail(self, mock_scalars, monkeypatch):
-        mock_scalars.side_effect = Exception("DB error")
+    def test_both_fail(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch):
+        def fail_query(_orm_execute_state: object) -> None:
+            raise SQLAlchemyError("DB error")
 
-        # See ``test_both_success``: restore the original module via monkeypatch
-        # rather than ``del``-ing it, so we don't evict it for sibling tests.
-        mock_tool_catalogue = MagicMock()
-        mock_tool_catalogue.build_tool_catalogue.side_effect = Exception("Tool error")
-        monkeypatch.setitem(sys.modules, "core.workflow.generator.tool_catalogue", mock_tool_catalogue)
+        def fail_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]:
+            raise RuntimeError("Tool error")
 
-        assert LLMGenerator._build_suggestion_context("tenant") == ""
+        event.listen(dataset_session, "do_orm_execute", fail_query)
+        monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", fail_tool_catalogue)
+
+        try:
+            assert LLMGenerator._build_suggestion_context("tenant") == ""
+        finally:
+            event.remove(dataset_session, "do_orm_execute", fail_query)
 
 
 class TestWorkflowServiceInterface:
-    def test_protocol_methods(self):
-        # Just to cover the 'pass' statements in the Protocol definition
+    def test_real_workflow_service_exposes_protocol_methods(self):
         from core.llm_generator.llm_generator import WorkflowServiceInterface
 
-        class MockService(WorkflowServiceInterface):
-            def get_draft_workflow(self, app_model, workflow_id=None, *, session):
-                return super().get_draft_workflow(app_model, workflow_id, session=session)
+        service: WorkflowServiceInterface = WorkflowService(sessionmaker())
 
-            def get_node_last_run(self, app_model, workflow, node_id):
-                return super().get_node_last_run(app_model, workflow, node_id)
-
-        service = MockService()
-        service.get_draft_workflow(None, session=None)
-        service.get_node_last_run(None, None, "node")
+        assert callable(service.get_draft_workflow)
+        assert callable(service.get_node_last_run)

api/tests/unit_tests/core/memory/test_token_buffer_memory.py

@@ -1,11 +1,19 @@
-"""Comprehensive unit tests for core/memory/token_buffer_memory.py"""
+"""Comprehensive SQLite-backed tests for token-buffer memory."""
 
+from collections.abc import Iterator
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+from decimal import Decimal
 from unittest.mock import MagicMock, patch
 from uuid import uuid4
 
 import pytest
+from sqlalchemy import Engine, event
+from sqlalchemy.orm import Session
 
+from core.memory import token_buffer_memory as memory_module
 from core.memory.token_buffer_memory import TokenBufferMemory
+from graphon.file import FileTransferMethod, FileType
 from graphon.model_runtime.entities import (
     AssistantPromptMessage,
     ImagePromptMessageContent,
@@ -13,13 +21,44 @@ from graphon.model_runtime.entities import (
     TextPromptMessageContent,
     UserPromptMessage,
 )
-from models.model import AppMode
+from models.base import TypeBase
+from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo
+from models.model import AppMode, Message, MessageFile
+from models.workflow import Workflow, WorkflowType
 
 # ---------------------------------------------------------------------------
 # Helpers / shared fixtures
 # ---------------------------------------------------------------------------
 
 
+@dataclass(frozen=True)
+class Database:
+    """Typed SQLite binding plus executed SQL for query-count assertions."""
+
+    engine: Engine
+    session: Session
+    statements: list[tuple[str, object]]
+
+
+@pytest.fixture
+def database(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Database]:
+    TypeBase.metadata.create_all(
+        sqlite_engine,
+        tables=[Message.__table__, MessageFile.__table__, Workflow.__table__],
+    )
+    statements: list[tuple[str, object]] = []
+
+    def record_statement(_connection, _cursor, statement, parameters, _context, _executemany) -> None:
+        statements.append((statement, parameters))
+
+    event.listen(sqlite_engine, "before_cursor_execute", record_statement)
+    with Session(sqlite_engine, expire_on_commit=False) as session:
+        database = Database(engine=sqlite_engine, session=session, statements=statements)
+        monkeypatch.setattr(memory_module, "db", database)
+        yield database
+    event.remove(sqlite_engine, "before_cursor_execute", record_statement)
+
+
 def _make_conversation(mode: AppMode = AppMode.CHAT) -> MagicMock:
     """Return a minimal Conversation mock."""
     conv = MagicMock()
@@ -47,6 +86,73 @@ def _make_message(answer: str = "hello", answer_tokens: int = 5) -> MagicMock:
     return msg
 
 
+def _persist_message(
+    database: Database,
+    conversation_id: str,
+    *,
+    query: str = "user query",
+    answer: str = "hello",
+    answer_tokens: int = 5,
+    created_at: datetime | None = None,
+    workflow_run_id: str | None = None,
+) -> Message:
+    message = Message(
+        id=str(uuid4()),
+        app_id="app-1",
+        conversation_id=conversation_id,
+        _inputs={},
+        query=query,
+        message={},
+        message_unit_price=Decimal(0),
+        answer=answer,
+        answer_tokens=answer_tokens,
+        answer_unit_price=Decimal(0),
+        currency="USD",
+        from_source=ConversationFromSource.API,
+        workflow_run_id=workflow_run_id,
+        created_at=created_at or datetime.now(UTC).replace(tzinfo=None),
+    )
+    database.session.add(message)
+    database.session.commit()
+    return message
+
+
+def _persist_message_file(
+    database: Database,
+    message: Message,
+    *,
+    belongs_to: MessageFileBelongsTo | None,
+) -> MessageFile:
+    message_file = MessageFile(
+        message_id=message.id,
+        type=FileType.IMAGE,
+        transfer_method=FileTransferMethod.REMOTE_URL,
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by="account-1",
+        belongs_to=belongs_to,
+        url="https://example.com/image.png",
+    )
+    database.session.add(message_file)
+    database.session.commit()
+    return message_file
+
+
+def _persist_workflow(database: Database, *, workflow_id: str) -> Workflow:
+    workflow = Workflow(
+        id=workflow_id,
+        tenant_id="tenant-1",
+        app_id="app-1",
+        type=WorkflowType.CHAT,
+        version="1",
+        graph="{}",
+        features="{}",
+        created_by="account-1",
+    )
+    database.session.add(workflow)
+    database.session.commit()
+    return workflow
+
+
 # ===========================================================================
 # Tests for __init__ and workflow_run_repo property
 # ===========================================================================
@@ -61,25 +167,25 @@ class TestInit:
         assert mem.model_instance is mi
         assert mem._workflow_run_repo is None
 
-    def test_workflow_run_repo_is_created_lazily(self):
+    def test_workflow_run_repo_is_created_lazily(self, database: Database):
         conv = _make_conversation()
         mi = _make_model_instance()
         mem = TokenBufferMemory(conversation=conv, model_instance=mi)
 
         mock_repo = MagicMock()
-        with (
-            patch("core.memory.token_buffer_memory.sessionmaker") as mock_sm,
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.DifyAPIRepositoryFactory.create_api_workflow_run_repository",
-                return_value=mock_repo,
-            ),
-        ):
-            mock_db.engine = MagicMock()
+        with patch(
+            "core.memory.token_buffer_memory.DifyAPIRepositoryFactory.create_api_workflow_run_repository",
+            return_value=mock_repo,
+        ) as repository_factory:
             repo = mem.workflow_run_repo
             assert repo is mock_repo
             assert mem._workflow_run_repo is mock_repo
 
+        session_factory = repository_factory.call_args.args[0]
+        with session_factory() as session:
+            assert isinstance(session, Session)
+            assert session.get_bind() is database.engine
+
     def test_workflow_run_repo_cached_after_first_access(self):
         conv = _make_conversation()
         mi = _make_model_instance()
@@ -410,7 +516,7 @@ class TestBuildPromptMessageWithFiles:
             )
 
     @pytest.mark.parametrize("mode", [AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
-    def test_workflow_mode_workflow_not_found_raises(self, mode):
+    def test_workflow_mode_workflow_not_found_raises(self, mode, database: Database):
         """Raises ValueError when Workflow lookup returns None."""
         conv = _make_conversation(mode)
         conv.app = MagicMock()
@@ -422,22 +528,17 @@ class TestBuildPromptMessageWithFiles:
         mem._workflow_run_repo = MagicMock()
         mem._workflow_run_repo.get_workflow_run_by_id.return_value = mock_workflow_run
 
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-        ):
-            mock_db.session.scalar.return_value = None  # workflow not found
-
-            with pytest.raises(ValueError, match="Workflow not found"):
-                mem._build_prompt_message_with_files(
-                    message_files=[],
-                    text_content="text",
-                    message=_make_message(),
-                    app_record=MagicMock(),
-                    is_user_message=True,
-                )
+        with pytest.raises(ValueError, match="Workflow not found"):
+            mem._build_prompt_message_with_files(
+                message_files=[],
+                text_content="text",
+                message=_make_message(),
+                app_record=MagicMock(),
+                is_user_message=True,
+            )
 
     @pytest.mark.parametrize("mode", [AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
-    def test_workflow_mode_success_no_files_user(self, mode):
+    def test_workflow_mode_success_no_files_user(self, mode, database: Database):
         """Happy path: workflow mode, no message files → plain UserPromptMessage."""
         conv = _make_conversation(mode)
         conv.app = MagicMock()
@@ -445,22 +546,16 @@ class TestBuildPromptMessageWithFiles:
         mock_workflow_run = MagicMock()
         mock_workflow_run.workflow_id = str(uuid4())
 
-        mock_workflow = MagicMock()
-        mock_workflow.features_dict = {}
+        workflow = _persist_workflow(database, workflow_id=mock_workflow_run.workflow_id)
 
         mem = TokenBufferMemory(conversation=conv, model_instance=_make_model_instance())
         mem._workflow_run_repo = MagicMock()
         mem._workflow_run_repo.get_workflow_run_by_id.return_value = mock_workflow_run
 
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
-                return_value=None,
-            ),
+        with patch(
+            "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
+            return_value=None,
         ):
-            mock_db.session.scalar.return_value = mock_workflow
-
             result = mem._build_prompt_message_with_files(
                 message_files=[],
                 text_content="wf text",
@@ -471,6 +566,7 @@ class TestBuildPromptMessageWithFiles:
 
         assert isinstance(result, UserPromptMessage)
         assert result.content == "wf text"
+        assert database.session.get(Workflow, workflow.id) is workflow
 
     # ------------------------------------------------------------------
     # Invalid mode
@@ -498,417 +594,140 @@ class TestBuildPromptMessageWithFiles:
 
 
 class TestGetHistoryPromptMessages:
-    """Tests for get_history_prompt_messages."""
+    """Tests for persisted history retrieval, file batching, and pruning."""
 
     def _make_memory(self, mode: AppMode = AppMode.CHAT) -> TokenBufferMemory:
         conv = _make_conversation(mode)
         conv.app = MagicMock()
         return TokenBufferMemory(conversation=conv, model_instance=_make_model_instance())
 
-    def test_returns_empty_when_no_messages(self):
+    def test_returns_empty_when_no_messages(self, database: Database) -> None:
+        assert self._make_memory().get_history_prompt_messages() == []
+
+    def test_skips_newest_message_without_answer(self, database: Database) -> None:
         mem = self._make_memory()
-        with patch("core.memory.token_buffer_memory.db") as mock_db:
-            mock_db.session.scalars.return_value.all.return_value = []
-            result = mem.get_history_prompt_messages()
-        assert result == []
+        message = _persist_message(database, mem.conversation.id, answer="", answer_tokens=0)
 
-    def test_skips_first_message_without_answer(self):
-        """The newest message (index 0 after extraction) without answer and tokens==0 is skipped."""
+        assert mem.get_history_prompt_messages() == []
+        assert database.session.get(Message, message.id) is message
+
+    def test_message_with_answer_returns_user_and_assistant_prompts(self, database: Database) -> None:
         mem = self._make_memory()
+        _persist_message(database, mem.conversation.id, query="My query", answer="My answer", answer_tokens=10)
 
-        msg_no_answer = _make_message(answer="", answer_tokens=0)
-        msg_no_answer.parent_message_id = None  # ensures extract_thread_messages returns it
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg_no_answer],
-            ),
-        ):
-            mock_db.session.scalars.return_value.all.side_effect = [
-                [msg_no_answer],  # first call: messages query
-                [],  # second call: user files query (never hit, but safe)
-            ]
-            result = mem.get_history_prompt_messages()
-
-        assert result == []
-
-    def test_message_with_answer_not_skipped(self):
-        """A message with a non-empty answer is NOT popped."""
-        mem = self._make_memory()
-
-        msg = _make_message(answer="some answer", answer_tokens=10)
-        msg.parent_message_id = None
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg],
-            ),
-            patch(
-                "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
-                return_value=None,
-            ),
-        ):
-            # user files query → empty; assistant files query → empty
-            mock_db.session.scalars.return_value.all.return_value = []
-            result = mem.get_history_prompt_messages()
-
-        assert len(result) == 2  # one user + one assistant
-
-    def test_message_limit_default_is_500(self):
-        """When message_limit is None the stmt is limited to 500."""
-        mem = self._make_memory()
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch("core.memory.token_buffer_memory.select") as mock_select,
-            patch("core.memory.token_buffer_memory.extract_thread_messages", return_value=[]),
-        ):
-            mock_stmt = MagicMock()
-            mock_select.return_value.where.return_value.order_by.return_value = mock_stmt
-            mock_stmt.limit.return_value = mock_stmt
-            mock_db.session.scalars.return_value.all.return_value = []
-
-            mem.get_history_prompt_messages(message_limit=None)
-            mock_stmt.limit.assert_called_with(500)
-
-    def test_message_limit_clipped_to_500(self):
-        """A message_limit > 500 is clamped to 500."""
-        mem = self._make_memory()
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch("core.memory.token_buffer_memory.select") as mock_select,
-            patch("core.memory.token_buffer_memory.extract_thread_messages", return_value=[]),
-        ):
-            mock_stmt = MagicMock()
-            mock_select.return_value.where.return_value.order_by.return_value = mock_stmt
-            mock_stmt.limit.return_value = mock_stmt
-            mock_db.session.scalars.return_value.all.return_value = []
-
-            mem.get_history_prompt_messages(message_limit=9999)
-            mock_stmt.limit.assert_called_with(500)
-
-    def test_message_limit_positive_used(self):
-        """A positive message_limit < 500 is used as-is."""
-        mem = self._make_memory()
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch("core.memory.token_buffer_memory.select") as mock_select,
-            patch("core.memory.token_buffer_memory.extract_thread_messages", return_value=[]),
-        ):
-            mock_stmt = MagicMock()
-            mock_select.return_value.where.return_value.order_by.return_value = mock_stmt
-            mock_stmt.limit.return_value = mock_stmt
-            mock_db.session.scalars.return_value.all.return_value = []
-
-            mem.get_history_prompt_messages(message_limit=10)
-            mock_stmt.limit.assert_called_with(10)
-
-    def test_message_limit_zero_uses_default(self):
-        """message_limit=0 triggers the else branch → default 500."""
-        mem = self._make_memory()
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch("core.memory.token_buffer_memory.select") as mock_select,
-            patch("core.memory.token_buffer_memory.extract_thread_messages", return_value=[]),
-        ):
-            mock_stmt = MagicMock()
-            mock_select.return_value.where.return_value.order_by.return_value = mock_stmt
-            mock_stmt.limit.return_value = mock_stmt
-            mock_db.session.scalars.return_value.all.return_value = []
-
-            mem.get_history_prompt_messages(message_limit=0)
-            mock_stmt.limit.assert_called_with(500)
-
-    def test_user_files_cause_build_with_files_call(self):
-        """When user_files is non-empty _build_prompt_message_with_files is invoked."""
-        mem = self._make_memory()
-        msg = _make_message()
-        msg.parent_message_id = None
-
-        mock_user_file = MagicMock()
-        mock_user_file.message_id = msg.id  # must match so batched grouping keys it to this message
-        mock_user_prompt = UserPromptMessage(content="from build")
-        mock_assistant_prompt = AssistantPromptMessage(content="answer")
-
-        call_count = {"n": 0}
-
-        def scalars_side_effect(stmt):
-            r = MagicMock()
-            if call_count["n"] == 0:
-                # messages query
-                r.all.return_value = [msg]
-            elif call_count["n"] == 1:
-                # user files
-                r.all.return_value = [mock_user_file]
-            else:
-                # assistant files
-                r.all.return_value = []
-            call_count["n"] += 1
-            return r
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg],
-            ),
-            patch.object(
-                mem,
-                "_build_prompt_message_with_files",
-                side_effect=[mock_user_prompt, mock_assistant_prompt],
-            ) as mock_build,
-            patch(
-                "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
-                return_value=None,
-            ),
-        ):
-            mock_db.session.scalars.side_effect = scalars_side_effect
-            result = mem.get_history_prompt_messages()
-
-        assert mock_build.call_count >= 1
-        # First call should be user message
-        first_call_kwargs = mock_build.call_args_list[0][1]
-        assert first_call_kwargs["is_user_message"] is True
-
-    def test_assistant_files_cause_build_with_files_call(self):
-        """When assistant_files is non-empty, build is called with is_user_message=False."""
-        mem = self._make_memory()
-        msg = _make_message()
-        msg.parent_message_id = None
-
-        mock_assistant_file = MagicMock()
-        mock_assistant_file.message_id = msg.id  # must match so batched grouping keys it to this message
-        mock_user_prompt = UserPromptMessage(content="query")
-        mock_assistant_prompt = AssistantPromptMessage(content="built")
-
-        call_count = {"n": 0}
-
-        def scalars_side_effect(stmt):
-            r = MagicMock()
-            if call_count["n"] == 0:
-                r.all.return_value = [msg]
-            elif call_count["n"] == 1:
-                r.all.return_value = []  # no user files
-            else:
-                r.all.return_value = [mock_assistant_file]
-            call_count["n"] += 1
-            return r
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg],
-            ),
-            patch.object(
-                mem,
-                "_build_prompt_message_with_files",
-                return_value=mock_assistant_prompt,
-            ) as mock_build,
-        ):
-            mock_db.session.scalars.side_effect = scalars_side_effect
-            result = mem.get_history_prompt_messages()
-
-        mock_build.assert_called_once()
-        call_kwargs = mock_build.call_args[1]
-        assert call_kwargs["is_user_message"] is False
-
-    def test_message_files_loaded_with_constant_query_count(self):
-        """Regression guard against N+1: message files must be batch-loaded.
-
-        Regardless of the number of messages in the thread, file loading must use a
-        constant number of queries (1 messages query + 2 batched file queries),
-        never 2 queries per message.
-        """
-        mem = self._make_memory()
-
-        messages = [_make_message() for _ in range(5)]
-        for m in messages:
-            m.parent_message_id = None
-
-        scalars_calls = {"n": 0}
-
-        def scalars_side_effect(stmt):
-            r = MagicMock()
-            # First call returns the thread messages; the batched file queries return none.
-            r.all.return_value = messages if scalars_calls["n"] == 0 else []
-            scalars_calls["n"] += 1
-            return r
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch("core.memory.token_buffer_memory.extract_thread_messages", return_value=messages),
-            patch("core.memory.token_buffer_memory.FileUploadConfigManager.convert", return_value=None),
-        ):
-            mock_db.session.scalars.side_effect = scalars_side_effect
-            mem.get_history_prompt_messages()
-
-        # 1 (messages) + 2 (batched user/assistant files) = 3, independent of message count.
-        # Before this fix it would have been 1 + 2 * 5 = 11 (an N+1 pattern).
-        assert scalars_calls["n"] == 3
-
-    def test_token_pruning_removes_oldest_messages(self):
-        """If tokens exceed limit, oldest messages are removed until within limit."""
-        conv = _make_conversation()
-        conv.app = MagicMock()
-
-        # Model returns tokens that decrease only after removing pairs
-        token_values = [3000, 1500]  # first call over limit, second within
-        mi = MagicMock()
-        mi.get_llm_num_tokens.side_effect = token_values
-
-        mem = TokenBufferMemory(conversation=conv, model_instance=mi)
-
-        msg = _make_message()
-        msg.parent_message_id = None
-
-        call_count = {"n": 0}
-
-        def scalars_side_effect(stmt):
-            r = MagicMock()
-            if call_count["n"] == 0:
-                r.all.return_value = [msg]
-            else:
-                r.all.return_value = []
-            call_count["n"] += 1
-            return r
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg],
-            ),
-            patch(
-                "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
-                return_value=None,
-            ),
-        ):
-            mock_db.session.scalars.side_effect = scalars_side_effect
-            result = mem.get_history_prompt_messages(max_token_limit=2000)
-
-        # After pruning, we should have fewer than the 2 initial messages
-        assert len(result) <= 1
-
-    def test_token_pruning_stops_at_single_message(self):
-        """Pruning stops when only 1 message remains (to prevent empty list)."""
-        conv = _make_conversation()
-        conv.app = MagicMock()
-
-        # Always over limit
-        mi = MagicMock()
-        mi.get_llm_num_tokens.return_value = 99999
-
-        mem = TokenBufferMemory(conversation=conv, model_instance=mi)
-
-        msg = _make_message()
-        msg.parent_message_id = None
-
-        call_count = {"n": 0}
-
-        def scalars_side_effect(stmt):
-            r = MagicMock()
-            if call_count["n"] == 0:
-                r.all.return_value = [msg]
-            else:
-                r.all.return_value = []
-            call_count["n"] += 1
-            return r
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg],
-            ),
-            patch(
-                "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
-                return_value=None,
-            ),
-        ):
-            mock_db.session.scalars.side_effect = scalars_side_effect
-            result = mem.get_history_prompt_messages(max_token_limit=1)
-
-        # At least 1 message should remain
-        assert len(result) >= 1
-
-    def test_no_pruning_when_within_limit(self):
-        """When tokens ≤ limit, no pruning occurs."""
-        mem = self._make_memory()
-        mem.model_instance.get_llm_num_tokens.return_value = 50  # well under default 2000
-
-        msg = _make_message()
-        msg.parent_message_id = None
-
-        call_count = {"n": 0}
-
-        def scalars_side_effect(stmt):
-            r = MagicMock()
-            if call_count["n"] == 0:
-                r.all.return_value = [msg]
-            else:
-                r.all.return_value = []
-            call_count["n"] += 1
-            return r
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg],
-            ),
-            patch(
-                "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
-                return_value=None,
-            ),
-        ):
-            mock_db.session.scalars.side_effect = scalars_side_effect
-            result = mem.get_history_prompt_messages(max_token_limit=2000)
-
-        assert len(result) == 2  # user + assistant
-
-    def test_plain_user_and_assistant_messages_returned(self):
-        """Without files, plain UserPromptMessage and AssistantPromptMessage appear."""
-        mem = self._make_memory()
-
-        msg = _make_message(answer="My answer")
-        msg.query = "My query"
-        msg.parent_message_id = None
-
-        call_count = {"n": 0}
-
-        def scalars_side_effect(stmt):
-            r = MagicMock()
-            if call_count["n"] == 0:
-                r.all.return_value = [msg]
-            else:
-                r.all.return_value = []
-            call_count["n"] += 1
-            return r
-
-        with (
-            patch("core.memory.token_buffer_memory.db") as mock_db,
-            patch(
-                "core.memory.token_buffer_memory.extract_thread_messages",
-                return_value=[msg],
-            ),
-            patch(
-                "core.memory.token_buffer_memory.FileUploadConfigManager.convert",
-                return_value=None,
-            ),
-        ):
-            mock_db.session.scalars.side_effect = scalars_side_effect
-            result = mem.get_history_prompt_messages()
+        result = mem.get_history_prompt_messages()
 
         assert len(result) == 2
-        user_msg, ai_msg = result
-        assert isinstance(user_msg, UserPromptMessage)
-        assert user_msg.content == "My query"
-        assert isinstance(ai_msg, AssistantPromptMessage)
-        assert ai_msg.content == "My answer"
+        assert isinstance(result[0], UserPromptMessage)
+        assert result[0].content == "My query"
+        assert isinstance(result[1], AssistantPromptMessage)
+        assert result[1].content == "My answer"
+
+    def test_history_is_conversation_scoped(self, database: Database) -> None:
+        mem = self._make_memory()
+        _persist_message(database, mem.conversation.id, answer="visible")
+        _persist_message(database, "other-conversation", answer="hidden")
+
+        result = mem.get_history_prompt_messages()
+
+        assert [prompt.content for prompt in result] == ["user query", "visible"]
+
+    @pytest.mark.parametrize(
+        ("message_limit", "expected_limit"),
+        [(None, 500), (9999, 500), (10, 10), (0, 500)],
+    )
+    def test_message_limit_is_applied_to_executable_query(
+        self,
+        database: Database,
+        message_limit: int | None,
+        expected_limit: int,
+    ) -> None:
+        mem = self._make_memory()
+        before = len(database.statements)
+
+        mem.get_history_prompt_messages(message_limit=message_limit)
+
+        statements = database.statements[before:]
+        assert len(statements) == 1
+        sql, parameters = statements[0]
+        assert "LIMIT" in sql
+        assert expected_limit in parameters
+
+    @pytest.mark.parametrize(
+        ("belongs_to", "is_user_message"),
+        [
+            (MessageFileBelongsTo.USER, True),
+            (None, True),
+            (MessageFileBelongsTo.ASSISTANT, False),
+        ],
+    )
+    def test_message_files_use_persisted_ownership(
+        self,
+        database: Database,
+        belongs_to: MessageFileBelongsTo | None,
+        is_user_message: bool,
+    ) -> None:
+        mem = self._make_memory()
+        message = _persist_message(database, mem.conversation.id)
+        message_file = _persist_message_file(database, message, belongs_to=belongs_to)
+        built_prompt = (
+            UserPromptMessage(content="built user")
+            if is_user_message
+            else AssistantPromptMessage(content="built assistant")
+        )
+
+        with patch.object(mem, "_build_prompt_message_with_files", return_value=built_prompt) as build_prompt:
+            result = mem.get_history_prompt_messages()
+
+        build_prompt.assert_called_once()
+        assert build_prompt.call_args.kwargs["message_files"] == [message_file]
+        assert build_prompt.call_args.kwargs["is_user_message"] is is_user_message
+        assert built_prompt in result
+
+    def test_message_files_are_batch_loaded_with_constant_query_count(self, database: Database) -> None:
+        mem = self._make_memory()
+        base_time = datetime.now(UTC).replace(tzinfo=None)
+        messages = [
+            _persist_message(
+                database,
+                mem.conversation.id,
+                query=f"query-{index}",
+                answer=f"answer-{index}",
+                created_at=base_time + timedelta(seconds=index),
+            )
+            for index in range(5)
+        ]
+        before = len(database.statements)
+
+        with patch("core.memory.token_buffer_memory.extract_thread_messages", return_value=messages):
+            result = mem.get_history_prompt_messages()
+
+        selects = [sql for sql, _ in database.statements[before:] if sql.lstrip().upper().startswith("SELECT")]
+        assert len(selects) == 3
+        assert len(result) == 10
+
+    @pytest.mark.parametrize(
+        ("token_values", "max_token_limit", "expected_length"),
+        [
+            ([3000, 1500], 2000, 1),
+            ([99999, 99999], 1, 1),
+            ([50], 2000, 2),
+        ],
+    )
+    def test_token_pruning_uses_persisted_history(
+        self,
+        database: Database,
+        token_values: list[int],
+        max_token_limit: int,
+        expected_length: int,
+    ) -> None:
+        mem = self._make_memory()
+        mem.model_instance.get_llm_num_tokens.side_effect = token_values
+        _persist_message(database, mem.conversation.id)
+
+        result = mem.get_history_prompt_messages(max_token_limit=max_token_limit)
+
+        assert len(result) == expected_length
 
 
 # ===========================================================================

api/tests/unit_tests/core/moderation/api/test_api.py

@@ -2,13 +2,24 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 from pydantic import ValidationError
+from sqlalchemy.orm import Session
 
+import core.moderation.api.api as moderation_module
 from core.extension.api_based_extension_requestor import APIBasedExtensionPoint
 from core.moderation.api.api import ApiModeration, ModerationInputParams, ModerationOutputParams
 from core.moderation.base import ModerationAction, ModerationInputsResult, ModerationOutputsResult
 from models.api_based_extension import APIBasedExtension
 
 
+class _DatabaseBinding:
+    """Expose the real SQLite session used by extension lookup."""
+
+    session: Session
+
+    def __init__(self, session: Session) -> None:
+        self.session = session
+
+
 class TestApiModeration:
     @pytest.fixture
     def api_config(self):
@@ -165,17 +176,27 @@ class TestApiModeration:
         with pytest.raises(ValueError, match="API-based Extension not found"):
             api_moderation._get_config_by_requestor(APIBasedExtensionPoint.APP_MODERATION_INPUT, {})
 
-    @patch("core.moderation.api.api.db.session.scalar")
-    def test_get_api_based_extension(self, mock_scalar):
-        mock_ext = MagicMock(spec=APIBasedExtension)
-        mock_scalar.return_value = mock_ext
+    @pytest.mark.parametrize("sqlite_session", [(APIBasedExtension,)], indirect=True)
+    def test_get_api_based_extension(self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> None:
+        target = APIBasedExtension(
+            tenant_id="tenant-1",
+            name="Target extension",
+            api_endpoint="https://example.com/moderate",
+            api_key="encrypted-key",
+        )
+        target.id = "ext-1"
+        other_tenant = APIBasedExtension(
+            tenant_id="tenant-2",
+            name="Other extension",
+            api_endpoint="https://example.com/other",
+            api_key="other-key",
+        )
+        other_tenant.id = "ext-2"
+        sqlite_session.add_all((target, other_tenant))
+        sqlite_session.commit()
+        monkeypatch.setattr(moderation_module, "db", _DatabaseBinding(sqlite_session))
 
         result = ApiModeration._get_api_based_extension("tenant-1", "ext-1")
 
-        assert result == mock_ext
-        mock_scalar.assert_called_once()
-        # Verify the call has the correct filters
-        args, kwargs = mock_scalar.call_args
-        stmt = args[0]
-        # We can't easily inspect the statement without complex sqlalchemy tricks,
-        # but calling it is usually enough for unit tests if we mock the result.
+        assert result is target
+        assert ApiModeration._get_api_based_extension("tenant-1", "ext-2") is None

api/tests/unit_tests/core/ops/test_lookup_helpers.py

@@ -7,36 +7,212 @@ Covers:
 - TraceTask._get_user_id_from_metadata
 """
 
-from unittest.mock import MagicMock, patch
+import uuid
+from collections.abc import Iterator
+from contextlib import contextmanager
+from unittest.mock import PropertyMock, patch
 
 import pytest
+from sqlalchemy import Engine, event
+from sqlalchemy.orm import Session
+
+from core.tools.entities.tool_entities import ApiProviderSchemaType
+from extensions.ext_database import db
+from graphon.model_runtime.entities.model_entities import ModelType
+from models.account import Tenant
+from models.base import TypeBase
+from models.model import App, AppMode, IconType
+from models.provider import Provider, ProviderCredential, ProviderModel, ProviderModelCredential, ProviderType
+from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider
 
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
 
 
-def _make_db_and_session_patches(scalar_side_effect=None, scalar_return_value=None):
-    """Return (mock_db, cm, session) ready to patch 'core.ops.ops_trace_manager.db'
-    and 'core.ops.ops_trace_manager.Session'.
+@pytest.fixture
+def orm_session(sqlite_engine: Engine) -> Iterator[Session]:
+    models = (
+        App,
+        Tenant,
+        Provider,
+        ProviderCredential,
+        ProviderModel,
+        ProviderModelCredential,
+        BuiltinToolProvider,
+        ApiToolProvider,
+        WorkflowToolProvider,
+        MCPToolProvider,
+    )
+    tables = [model.metadata.tables[model.__tablename__] for model in models]
+    TypeBase.metadata.create_all(sqlite_engine, tables=tables)
 
-    Provide either scalar_side_effect (list, for multiple calls) or
-    scalar_return_value (single value).
-    """
-    mock_db = MagicMock()
-    mock_db.engine = MagicMock()
+    with patch.object(type(db), "engine", new_callable=PropertyMock, return_value=sqlite_engine):
+        with Session(sqlite_engine, expire_on_commit=False) as session:
+            yield session
 
-    session = MagicMock()
-    if scalar_side_effect is not None:
-        session.scalar.side_effect = scalar_side_effect
+
+def _persist_app(session: Session, *, tenant_id: str, name: str = "MyApp") -> App:
+    app = App(
+        id=str(uuid.uuid4()),
+        tenant_id=tenant_id,
+        name=name,
+        mode=AppMode.WORKFLOW,
+        icon_type=IconType.EMOJI,
+        icon="workflow",
+        icon_background="#FFFFFF",
+        enable_site=True,
+        enable_api=False,
+    )
+    session.add(app)
+    session.commit()
+    return app
+
+
+def _persist_tenant(session: Session, *, name: str = "MyWorkspace") -> Tenant:
+    tenant = Tenant(name=name)
+    session.add(tenant)
+    session.commit()
+    return tenant
+
+
+def _persist_tool_provider(
+    session: Session, provider_type: str
+) -> BuiltinToolProvider | ApiToolProvider | WorkflowToolProvider | MCPToolProvider:
+    tenant_id = str(uuid.uuid4())
+    user_id = str(uuid.uuid4())
+    if provider_type in {"builtin", "plugin"}:
+        provider = BuiltinToolProvider(
+            name="CredentialA",
+            tenant_id=tenant_id,
+            user_id=user_id,
+            provider="test/provider",
+        )
+    elif provider_type == "api":
+        provider = ApiToolProvider(
+            name="CredentialA",
+            icon="icon.svg",
+            schema="{}",
+            schema_type_str=ApiProviderSchemaType.OPENAPI,
+            user_id=user_id,
+            tenant_id=tenant_id,
+            description="API provider",
+            tools_str="[]",
+            credentials_str="{}",
+        )
+    elif provider_type == "workflow":
+        provider = WorkflowToolProvider(
+            name="CredentialA",
+            label="CredentialA",
+            icon="icon.svg",
+            app_id=str(uuid.uuid4()),
+            version="1",
+            user_id=user_id,
+            tenant_id=tenant_id,
+            description="Workflow provider",
+        )
+    elif provider_type == "mcp":
+        provider = MCPToolProvider(
+            name="CredentialA",
+            server_identifier="credential-a",
+            server_url="https://example.com/mcp",
+            server_url_hash="credential-a-hash",
+            icon="icon.svg",
+            tenant_id=tenant_id,
+            user_id=user_id,
+        )
     else:
-        session.scalar.return_value = scalar_return_value
+        raise ValueError(f"unsupported provider type: {provider_type}")
 
-    cm = MagicMock()
-    cm.__enter__ = MagicMock(return_value=session)
-    cm.__exit__ = MagicMock(return_value=False)
+    session.add(provider)
+    session.commit()
+    return provider
 
-    return mock_db, cm, session
+
+def _persist_provider_credential(
+    session: Session,
+    *,
+    tenant_id: str,
+    credential_name: str = "ProvCredName",
+) -> ProviderCredential:
+    credential = ProviderCredential(
+        tenant_id=tenant_id,
+        provider_name="openai",
+        credential_name=credential_name,
+        encrypted_config="{}",
+    )
+    session.add(credential)
+    session.commit()
+    return credential
+
+
+def _persist_model_credential(
+    session: Session,
+    *,
+    tenant_id: str,
+    credential_name: str = "ModelCredName",
+) -> ProviderModelCredential:
+    credential = ProviderModelCredential(
+        tenant_id=tenant_id,
+        provider_name="openai",
+        model_name="gpt-4",
+        model_type=ModelType.LLM,
+        credential_name=credential_name,
+        encrypted_config="{}",
+    )
+    session.add(credential)
+    session.commit()
+    return credential
+
+
+def _persist_provider(
+    session: Session,
+    *,
+    tenant_id: str,
+    credential_id: str | None,
+) -> Provider:
+    provider = Provider(
+        tenant_id=tenant_id,
+        provider_name="openai",
+        provider_type=ProviderType.CUSTOM,
+        credential_id=credential_id,
+    )
+    session.add(provider)
+    session.commit()
+    return provider
+
+
+def _persist_provider_model(
+    session: Session,
+    *,
+    tenant_id: str,
+    credential_id: str | None,
+) -> ProviderModel:
+    model = ProviderModel(
+        tenant_id=tenant_id,
+        provider_name="openai",
+        model_name="gpt-4",
+        model_type=ModelType.LLM,
+        credential_id=credential_id,
+    )
+    session.add(model)
+    session.commit()
+    return model
+
+
+@contextmanager
+def _raise_on_table(engine: Engine, table_name: str) -> Iterator[None]:
+    """Raise only when SQL targets the named table, leaving other real lookups intact."""
+
+    def fail_target_query(_conn, _cursor, statement, _parameters, _context, _executemany):
+        if f"FROM {table_name}" in statement:
+            raise RuntimeError(f"forced failure for {table_name}")
+
+    event.listen(engine, "before_cursor_execute", fail_target_query)
+    try:
+        yield
+    finally:
+        event.remove(engine, "before_cursor_execute", fail_target_query)
 
 
 # ---------------------------------------------------------------------------
@@ -47,62 +223,42 @@ def _make_db_and_session_patches(scalar_side_effect=None, scalar_return_value=No
 class TestLookupAppAndWorkspaceNames:
     """Tests for _lookup_app_and_workspace_names(app_id, tenant_id)."""
 
-    def test_both_found(self):
+    def test_both_found(self, orm_session: Session):
         """Returns (app_name, workspace_name) when both records exist."""
         from core.ops.ops_trace_manager import _lookup_app_and_workspace_names
 
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_side_effect=["MyApp", "MyWorkspace"])
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            app_name, workspace_name = _lookup_app_and_workspace_names("app-123", "tenant-456")
+        tenant = _persist_tenant(orm_session)
+        app = _persist_app(orm_session, tenant_id=tenant.id)
+        app_name, workspace_name = _lookup_app_and_workspace_names(app.id, tenant.id)
 
         assert app_name == "MyApp"
         assert workspace_name == "MyWorkspace"
 
-    def test_app_only_found(self):
+    def test_app_only_found(self, orm_session: Session):
         """Returns (app_name, '') when tenant record is absent."""
         from core.ops.ops_trace_manager import _lookup_app_and_workspace_names
 
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_side_effect=["MyApp", None])
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            app_name, workspace_name = _lookup_app_and_workspace_names("app-123", "tenant-456")
+        app = _persist_app(orm_session, tenant_id=str(uuid.uuid4()))
+        app_name, workspace_name = _lookup_app_and_workspace_names(app.id, str(uuid.uuid4()))
 
         assert app_name == "MyApp"
         assert workspace_name == ""
 
-    def test_tenant_only_found(self):
+    def test_tenant_only_found(self, orm_session: Session):
         """Returns ('', workspace_name) when app record is absent."""
         from core.ops.ops_trace_manager import _lookup_app_and_workspace_names
 
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_side_effect=[None, "MyWorkspace"])
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            app_name, workspace_name = _lookup_app_and_workspace_names("app-123", "tenant-456")
+        tenant = _persist_tenant(orm_session)
+        app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), tenant.id)
 
         assert app_name == ""
         assert workspace_name == "MyWorkspace"
 
-    def test_neither_found(self):
+    def test_neither_found(self, orm_session: Session):
         """Returns ('', '') when both DB lookups return None."""
         from core.ops.ops_trace_manager import _lookup_app_and_workspace_names
 
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_side_effect=[None, None])
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            app_name, workspace_name = _lookup_app_and_workspace_names("app-123", "tenant-456")
+        app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), str(uuid.uuid4()))
 
         assert app_name == ""
         assert workspace_name == ""
@@ -111,50 +267,30 @@ class TestLookupAppAndWorkspaceNames:
         """Returns ('', '') immediately when both IDs are None — no DB access."""
         from core.ops.ops_trace_manager import _lookup_app_and_workspace_names
 
-        mock_db = MagicMock()
-        mock_session_cls = MagicMock()
+        app_name, workspace_name = _lookup_app_and_workspace_names(None, None)
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", mock_session_cls),
-        ):
-            app_name, workspace_name = _lookup_app_and_workspace_names(None, None)
-
-        mock_session_cls.assert_not_called()
         assert app_name == ""
         assert workspace_name == ""
 
-    def test_app_id_none_only_queries_tenant(self):
+    def test_app_id_none_only_queries_tenant(self, orm_session: Session):
         """When app_id is None, only the tenant query is issued."""
         from core.ops.ops_trace_manager import _lookup_app_and_workspace_names
 
-        mock_db, cm, session = _make_db_and_session_patches(scalar_return_value="OnlyWorkspace")
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            app_name, workspace_name = _lookup_app_and_workspace_names(None, "tenant-456")
+        tenant = _persist_tenant(orm_session, name="OnlyWorkspace")
+        app_name, workspace_name = _lookup_app_and_workspace_names(None, tenant.id)
 
         assert app_name == ""
         assert workspace_name == "OnlyWorkspace"
-        assert session.scalar.call_count == 1
 
-    def test_tenant_id_none_only_queries_app(self):
+    def test_tenant_id_none_only_queries_app(self, orm_session: Session):
         """When tenant_id is None, only the app query is issued."""
         from core.ops.ops_trace_manager import _lookup_app_and_workspace_names
 
-        mock_db, cm, session = _make_db_and_session_patches(scalar_return_value="OnlyApp")
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            app_name, workspace_name = _lookup_app_and_workspace_names("app-123", None)
+        app = _persist_app(orm_session, tenant_id=str(uuid.uuid4()), name="OnlyApp")
+        app_name, workspace_name = _lookup_app_and_workspace_names(app.id, None)
 
         assert app_name == "OnlyApp"
         assert workspace_name == ""
-        assert session.scalar.call_count == 1
 
 
 # ---------------------------------------------------------------------------
@@ -166,32 +302,20 @@ class TestLookupCredentialName:
     """Tests for _lookup_credential_name(credential_id, provider_type)."""
 
     @pytest.mark.parametrize("provider_type", ["builtin", "plugin", "api", "workflow", "mcp"])
-    def test_known_provider_types_return_name(self, provider_type):
+    def test_known_provider_types_return_name(self, provider_type: str, orm_session: Session):
         """Each valid provider_type results in a DB query and returns the credential name."""
         from core.ops.ops_trace_manager import _lookup_credential_name
 
-        mock_db, cm, session = _make_db_and_session_patches(scalar_return_value="CredentialA")
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            result = _lookup_credential_name("cred-123", provider_type)
+        provider = _persist_tool_provider(orm_session, provider_type)
+        result = _lookup_credential_name(provider.id, provider_type)
 
         assert result == "CredentialA"
-        session.scalar.assert_called_once()
 
-    def test_credential_not_found_returns_empty_string(self):
+    def test_credential_not_found_returns_empty_string(self, orm_session: Session):
         """Returns '' when DB yields None for the given credential_id."""
         from core.ops.ops_trace_manager import _lookup_credential_name
 
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_return_value=None)
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            result = _lookup_credential_name("cred-999", "api")
+        result = _lookup_credential_name(str(uuid.uuid4()), "api")
 
         assert result == ""
 
@@ -199,48 +323,24 @@ class TestLookupCredentialName:
         """Returns '' immediately for an unrecognised provider_type — no DB access."""
         from core.ops.ops_trace_manager import _lookup_credential_name
 
-        mock_db = MagicMock()
-        mock_session_cls = MagicMock()
+        result = _lookup_credential_name(str(uuid.uuid4()), "unknown_type")
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", mock_session_cls),
-        ):
-            result = _lookup_credential_name("cred-123", "unknown_type")
-
-        mock_session_cls.assert_not_called()
         assert result == ""
 
     def test_none_credential_id_returns_empty_string_without_db(self):
         """Returns '' immediately when credential_id is None — no DB access."""
         from core.ops.ops_trace_manager import _lookup_credential_name
 
-        mock_db = MagicMock()
-        mock_session_cls = MagicMock()
+        result = _lookup_credential_name(None, "api")
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", mock_session_cls),
-        ):
-            result = _lookup_credential_name(None, "api")
-
-        mock_session_cls.assert_not_called()
         assert result == ""
 
     def test_none_provider_type_returns_empty_string_without_db(self):
         """Returns '' immediately when provider_type is None — no DB access."""
         from core.ops.ops_trace_manager import _lookup_credential_name
 
-        mock_db = MagicMock()
-        mock_session_cls = MagicMock()
+        result = _lookup_credential_name(str(uuid.uuid4()), None)
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", mock_session_cls),
-        ):
-            result = _lookup_credential_name("cred-123", None)
-
-        mock_session_cls.assert_not_called()
         assert result == ""
 
     def test_builtin_and_plugin_map_to_same_model(self):
@@ -281,106 +381,78 @@ class TestLookupCredentialName:
 class TestLookupLlmCredentialInfo:
     """Tests for _lookup_llm_credential_info(tenant_id, provider, model, model_type)."""
 
-    def _provider_record(self, credential_id: str | None = None) -> MagicMock:
-        record = MagicMock()
-        record.credential_id = credential_id
-        return record
-
-    def _model_record(self, credential_id: str | None = None) -> MagicMock:
-        record = MagicMock()
-        record.credential_id = credential_id
-        return record
-
-    def test_model_level_credential_found(self):
+    def test_model_level_credential_found(self, orm_session: Session):
         """Returns model-level credential_id and name when ProviderModel has a credential."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        provider_record = self._provider_record(credential_id=None)
-        model_record = self._model_record(credential_id="model-cred-id")
+        tenant_id = str(uuid.uuid4())
+        model_credential = _persist_model_credential(orm_session, tenant_id=tenant_id)
+        _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None)
+        _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=model_credential.id)
 
-        # scalar calls: (1) Provider, (2) ProviderModel, (3) ProviderModelCredential.credential_name
-        mock_db, cm, _session = _make_db_and_session_patches(
-            scalar_side_effect=[provider_record, model_record, "ModelCredName"]
+        decoy_tenant_id = str(uuid.uuid4())
+        decoy_credential = _persist_model_credential(
+            orm_session,
+            tenant_id=decoy_tenant_id,
+            credential_name="WrongTenantCredential",
         )
+        _persist_provider(orm_session, tenant_id=decoy_tenant_id, credential_id=None)
+        _persist_provider_model(orm_session, tenant_id=decoy_tenant_id, credential_id=decoy_credential.id)
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", "gpt-4")
+        cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4")
 
-        assert cred_id == "model-cred-id"
+        assert cred_id == model_credential.id
         assert cred_name == "ModelCredName"
 
-    def test_provider_level_fallback_when_no_model_credential(self):
+    def test_provider_level_fallback_when_no_model_credential(self, orm_session: Session):
         """Falls back to provider-level credential when ProviderModel has no credential_id."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        provider_record = self._provider_record(credential_id="prov-cred-id")
-        model_record = self._model_record(credential_id=None)
+        tenant_id = str(uuid.uuid4())
+        provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id)
+        _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id)
+        _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None)
 
-        # scalar calls: (1) Provider, (2) ProviderModel (no cred), (3) ProviderCredential.credential_name
-        mock_db, cm, _session = _make_db_and_session_patches(
-            scalar_side_effect=[provider_record, model_record, "ProvCredName"]
-        )
+        cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4")
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", "gpt-4")
-
-        assert cred_id == "prov-cred-id"
+        assert cred_id == provider_credential.id
         assert cred_name == "ProvCredName"
 
-    def test_provider_level_fallback_when_no_model_record(self):
+    def test_provider_level_fallback_when_no_model_record(self, orm_session: Session):
         """Falls back to provider-level credential when no ProviderModel row exists."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        provider_record = self._provider_record(credential_id="prov-cred-id")
+        tenant_id = str(uuid.uuid4())
+        provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id)
+        _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id)
 
-        # scalar calls: (1) Provider, (2) ProviderModel → None, (3) ProviderCredential.credential_name
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_side_effect=[provider_record, None, "ProvCredName"])
+        cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4")
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", "gpt-4")
-
-        assert cred_id == "prov-cred-id"
+        assert cred_id == provider_credential.id
         assert cred_name == "ProvCredName"
 
-    def test_no_model_arg_uses_provider_level_only(self):
+    def test_no_model_arg_uses_provider_level_only(self, orm_session: Session):
         """When model is None, skips ProviderModel query and uses provider credential."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        provider_record = self._provider_record(credential_id="prov-cred-id")
+        tenant_id = str(uuid.uuid4())
+        provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id)
+        _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id)
 
-        # scalar calls: (1) Provider, (2) ProviderCredential.credential_name — no ProviderModel
-        mock_db, cm, session = _make_db_and_session_patches(scalar_side_effect=[provider_record, "ProvCredName"])
+        cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", None)
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", None)
-
-        assert cred_id == "prov-cred-id"
+        assert cred_id == provider_credential.id
         assert cred_name == "ProvCredName"
-        assert session.scalar.call_count == 2
 
-    def test_provider_not_found_returns_none_and_empty(self):
+    def test_provider_not_found_returns_none_and_empty(self, orm_session: Session):
         """Returns (None, '') when Provider record does not exist."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_return_value=None)
+        other_tenant_id = str(uuid.uuid4())
+        _persist_provider(orm_session, tenant_id=other_tenant_id, credential_id=None)
+        tenant_id = str(uuid.uuid4())
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", "gpt-4")
+        cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4")
 
         assert cred_id is None
         assert cred_name == ""
@@ -389,16 +461,8 @@ class TestLookupLlmCredentialInfo:
         """Returns (None, '') immediately when tenant_id is None — no DB access."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        mock_db = MagicMock()
-        mock_session_cls = MagicMock()
+        cred_id, cred_name = _lookup_llm_credential_info(None, "openai", "gpt-4")
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", mock_session_cls),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info(None, "openai", "gpt-4")
-
-        mock_session_cls.assert_not_called()
         assert cred_id is None
         assert cred_name == ""
 
@@ -406,69 +470,46 @@ class TestLookupLlmCredentialInfo:
         """Returns (None, '') immediately when provider is None — no DB access."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        mock_db = MagicMock()
-        mock_session_cls = MagicMock()
+        cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), None, "gpt-4")
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", mock_session_cls),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", None, "gpt-4")
-
-        mock_session_cls.assert_not_called()
         assert cred_id is None
         assert cred_name == ""
 
-    def test_db_error_on_outer_query_returns_none_and_empty(self):
+    def test_db_error_on_outer_query_returns_none_and_empty(self, orm_session: Session, sqlite_engine: Engine):
         """Returns (None, '') and logs a warning when the outer DB query raises."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        mock_db, cm, session = _make_db_and_session_patches()
-        session.scalar.side_effect = Exception("DB connection failed")
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", "gpt-4")
+        with _raise_on_table(sqlite_engine, "providers"):
+            cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), "openai", "gpt-4")
 
         assert cred_id is None
         assert cred_name == ""
 
-    def test_credential_name_lookup_failure_returns_id_with_empty_name(self):
+    def test_credential_name_lookup_failure_returns_id_with_empty_name(
+        self, orm_session: Session, sqlite_engine: Engine
+    ):
         """When credential name sub-query fails, returns cred_id but '' for name."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        provider_record = self._provider_record(credential_id="prov-cred-id")
+        tenant_id = str(uuid.uuid4())
+        provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id)
+        _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id)
 
-        # Provider found, no model record, then name lookup raises
-        mock_db, cm, _session = _make_db_and_session_patches(
-            scalar_side_effect=[provider_record, None, Exception("deleted")]
-        )
+        with _raise_on_table(sqlite_engine, "provider_credentials"):
+            cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4")
 
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", "gpt-4")
-
-        assert cred_id == "prov-cred-id"
+        assert cred_id == provider_credential.id
         assert cred_name == ""
 
-    def test_no_credential_on_provider_or_model_returns_none_id(self):
+    def test_no_credential_on_provider_or_model_returns_none_id(self, orm_session: Session):
         """Returns (None, '') when neither provider nor model has a credential_id."""
         from core.ops.ops_trace_manager import _lookup_llm_credential_info
 
-        provider_record = self._provider_record(credential_id=None)
-        model_record = self._model_record(credential_id=None)
+        tenant_id = str(uuid.uuid4())
+        _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None)
+        _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None)
 
-        mock_db, cm, _session = _make_db_and_session_patches(scalar_side_effect=[provider_record, model_record])
-
-        with (
-            patch("core.ops.ops_trace_manager.db", mock_db),
-            patch("core.ops.ops_trace_manager.Session", return_value=cm),
-        ):
-            cred_id, cred_name = _lookup_llm_credential_info("tenant-1", "openai", "gpt-4")
+        cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4")
 
         assert cred_id is None
         assert cred_name == ""

api/tests/unit_tests/core/ops/test_utils.py

@@ -1,9 +1,11 @@
 import re
 from datetime import datetime
-from unittest.mock import MagicMock, patch
+from decimal import Decimal
 
 import pytest
+from sqlalchemy.orm import Session
 
+import core.ops.utils as utils_module
 from core.ops.utils import (
     filter_none_values,
     generate_dotted_order,
@@ -15,6 +17,42 @@ from core.ops.utils import (
     validate_url,
     validate_url_with_path,
 )
+from models.enums import ConversationFromSource
+from models.model import Message
+
+
+class _DatabaseBinding:
+    """Expose the real SQLite session used by the message lookup helper."""
+
+    session: Session
+
+    def __init__(self, session: Session) -> None:
+        self.session = session
+
+
+@pytest.fixture
+def message_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session:
+    """Bind the message lookup helper to the shared SQLite test session."""
+
+    monkeypatch.setattr(utils_module, "db", _DatabaseBinding(sqlite_session))
+    return sqlite_session
+
+
+def _message(message_id: str) -> Message:
+    message = Message(
+        id=message_id,
+        app_id="app-id",
+        conversation_id="conversation-id",
+        query="question",
+        message={"role": "user", "content": "question"},
+        answer="answer",
+        message_unit_price=Decimal("0.0001"),
+        answer_unit_price=Decimal("0.0001"),
+        currency="USD",
+        from_source=ConversationFromSource.API,
+    )
+    message._inputs = {}
+    return message
 
 
 class TestValidateUrl:
@@ -220,22 +258,20 @@ class TestFilterNoneValues:
         assert filter_none_values({}) == {}
 
 
+@pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True)
 class TestGetMessageData:
     """Test cases for get_message_data function"""
 
-    @patch("core.ops.utils.db")
-    @patch("core.ops.utils.Message")
-    @patch("core.ops.utils.select")
-    def test_get_message_data(self, mock_select, mock_message, mock_db):
-        mock_scalar = mock_db.session.scalar
-        mock_msg_instance = MagicMock()
-        mock_scalar.return_value = mock_msg_instance
+    def test_get_message_data(self, message_session: Session):
+        target = _message("message-id")
+        unrelated = _message("other-message-id")
+        message_session.add_all((target, unrelated))
+        message_session.commit()
 
         result = get_message_data("message-id")
 
-        assert result == mock_msg_instance
-        mock_select.assert_called_once()
-        mock_scalar.assert_called_once()
+        assert result is target
+        assert result.id == "message-id"
 
 
 class TestMeasureTime:

api/tests/unit_tests/core/prompt/test_extract_thread_messages.py

@@ -1,9 +1,42 @@
-from unittest.mock import MagicMock
+from datetime import datetime, timedelta
+from decimal import Decimal
 from uuid import uuid4
 
+import pytest
+from sqlalchemy.orm import Session
+
 from constants import UUID_NIL
 from core.prompt.utils.extract_thread_messages import extract_thread_messages
 from core.prompt.utils.get_thread_messages_length import get_thread_messages_length
+from models.enums import ConversationFromSource
+from models.model import Message
+
+
+def _persisted_message(
+    *,
+    message_id: str,
+    conversation_id: str,
+    parent_message_id: str,
+    answer: str,
+    created_at: datetime,
+) -> Message:
+    message = Message(
+        id=message_id,
+        app_id="app-id",
+        conversation_id=conversation_id,
+        query="question",
+        message={"role": "user", "content": "question"},
+        answer=answer,
+        message_unit_price=Decimal("0.0001"),
+        answer_unit_price=Decimal("0.0001"),
+        currency="USD",
+        from_source=ConversationFromSource.API,
+        parent_message_id=parent_message_id,
+        created_at=created_at,
+        updated_at=created_at,
+    )
+    message._inputs = {}
+    return message
 
 
 class MockMessage:
@@ -104,33 +137,64 @@ def test_extract_thread_messages_breaks_when_parent_is_none():
     assert result[0].id == id2
 
 
-def test_get_thread_messages_length_excludes_newly_created_empty_answer():
+@pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True)
+def test_get_thread_messages_length_excludes_newly_created_empty_answer(sqlite_session: Session):
     id1, id2 = str(uuid4()), str(uuid4())
+    now = datetime.now()
     messages = [
-        MockMessage(id2, id1, answer=""),  # newest generated message should be excluded
-        MockMessage(id1, UUID_NIL, answer="ok"),
+        _persisted_message(
+            message_id=id2,
+            conversation_id="conversation-1",
+            parent_message_id=id1,
+            answer="",
+            created_at=now,
+        ),
+        _persisted_message(
+            message_id=id1,
+            conversation_id="conversation-1",
+            parent_message_id=UUID_NIL,
+            answer="ok",
+            created_at=now - timedelta(seconds=1),
+        ),
+        _persisted_message(
+            message_id=str(uuid4()),
+            conversation_id="other-conversation",
+            parent_message_id=UUID_NIL,
+            answer="unrelated",
+            created_at=now + timedelta(seconds=1),
+        ),
     ]
+    sqlite_session.add_all(messages)
+    sqlite_session.commit()
 
-    session = MagicMock()
-    session.scalars.return_value.all.return_value = messages
-
-    length = get_thread_messages_length("conversation-1", session=session)
+    length = get_thread_messages_length("conversation-1", session=sqlite_session)
 
     assert length == 1
-    session.scalars.assert_called_once()
 
 
-def test_get_thread_messages_length_keeps_non_empty_latest_answer():
+@pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True)
+def test_get_thread_messages_length_keeps_non_empty_latest_answer(sqlite_session: Session):
     id1, id2 = str(uuid4()), str(uuid4())
+    now = datetime.now()
     messages = [
-        MockMessage(id2, id1, answer="latest-answer"),
-        MockMessage(id1, UUID_NIL, answer="older-answer"),
+        _persisted_message(
+            message_id=id2,
+            conversation_id="conversation-2",
+            parent_message_id=id1,
+            answer="latest-answer",
+            created_at=now,
+        ),
+        _persisted_message(
+            message_id=id1,
+            conversation_id="conversation-2",
+            parent_message_id=UUID_NIL,
+            answer="older-answer",
+            created_at=now - timedelta(seconds=1),
+        ),
     ]
+    sqlite_session.add_all(messages)
+    sqlite_session.commit()
 
-    session = MagicMock()
-    session.scalars.return_value.all.return_value = messages
-
-    length = get_thread_messages_length("conversation-2", session=session)
+    length = get_thread_messages_length("conversation-2", session=sqlite_session)
 
     assert length == 2
-    session.scalars.assert_called_once()

api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py

@@ -1,16 +1,16 @@
-"""Unit tests for workflow-as-tool behavior.
-
-StubSession/StubScalars emulate SQLAlchemy session/scalars with minimal methods
-(`scalar`, `scalars`, `expunge`, `commit`, `refresh`, context manager) to keep
-database access mocked and predictable in tests.
-"""
+"""Unit tests for workflow-as-tool behavior with real SQLite ORM boundaries."""
 
 import json
+import uuid
+from collections.abc import Iterator
+from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
+from sqlalchemy import Engine, inspect
+from sqlalchemy.orm import Session, sessionmaker
 
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.tools.__base.tool_runtime import ToolRuntime
@@ -23,74 +23,142 @@ from core.tools.entities.tool_entities import (
     ToolProviderType,
 )
 from core.tools.errors import ToolInvokeError
+from core.tools.workflow_as_tool import tool as workflow_tool_module
 from core.tools.workflow_as_tool.tool import WorkflowTool
 from graphon.file import FILE_MODEL_IDENTITY, FileTransferMethod, FileType
+from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
+from models.base import TypeBase
+from models.enums import EndUserType
+from models.model import App, AppMode, EndUser
+from models.workflow import Workflow, WorkflowType
+
+TENANT_ID = "00000000-0000-0000-0000-000000000001"
+OTHER_TENANT_ID = "00000000-0000-0000-0000-000000000002"
+APP_ID = "00000000-0000-0000-0000-000000000003"
+ACCOUNT_ID = "00000000-0000-0000-0000-000000000004"
+END_USER_ID = "00000000-0000-0000-0000-000000000005"
+CREATOR_ID = "00000000-0000-0000-0000-000000000006"
 
 
-class StubScalars:
-    """Minimal stub for SQLAlchemy scalar results."""
-
-    _value: Any
-
-    def __init__(self, value: Any) -> None:
-        self._value = value
-
-    def first(self) -> Any:
-        return self._value
+@dataclass(frozen=True)
+class SqliteToolDb:
+    engine: Engine
+    session_maker: sessionmaker[Session]
+    caller_session: Session
 
 
-class StubSession:
-    """Minimal stub for session_factory-created sessions."""
+@pytest.fixture
+def sqlite_tool_db(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_engine: Engine,
+) -> Iterator[SqliteToolDb]:
+    """Bind service-owned sessions and Account tenant reloads to SQLite."""
+    models = (App, Workflow, EndUser, Account, Tenant, TenantAccountJoin)
+    TypeBase.metadata.create_all(sqlite_engine, tables=[model.__table__ for model in models])
+    session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
+    monkeypatch.setattr(workflow_tool_module.session_factory, "create_session", session_maker)
 
-    scalar_results: list[Any]
-    scalars_results: list[Any]
-    expunge_calls: list[object]
+    from models import account as account_module
 
-    def __init__(self, *, scalar_results: list[Any] | None = None, scalars_results: list[Any] | None = None) -> None:
-        self.scalar_results = list(scalar_results or [])
-        self.scalars_results = list(scalars_results or [])
-        self.expunge_calls: list[object] = []
-
-    def scalar(self, _stmt: Any) -> Any:
-        return self.scalar_results.pop(0)
-
-    def scalars(self, _stmt: Any) -> StubScalars:
-        return StubScalars(self.scalars_results.pop(0))
-
-    def expunge(self, value: Any) -> None:
-        self.expunge_calls.append(value)
-
-    def begin(self) -> "StubSession":
-        return self
-
-    def commit(self) -> None:
-        pass
-
-    def refresh(self, _value: Any) -> None:
-        pass
-
-    def close(self) -> None:
-        pass
-
-    def __enter__(self) -> "StubSession":
-        return self
-
-    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
-        return False
+    monkeypatch.setattr(account_module, "db", SimpleNamespace(engine=sqlite_engine))
+    with session_maker() as caller_session:
+        yield SqliteToolDb(engine=sqlite_engine, session_maker=session_maker, caller_session=caller_session)
 
 
-def _build_tool() -> WorkflowTool:
+def _persist_tenant(db: SqliteToolDb, *, tenant_id: str = TENANT_ID) -> Tenant:
+    tenant = Tenant(name="Tenant")
+    tenant.id = tenant_id
+    db.caller_session.add(tenant)
+    db.caller_session.commit()
+    return tenant
+
+
+def _persist_account(db: SqliteToolDb, *, tenant_id: str = TENANT_ID) -> Account:
+    account = Account(name="Account", email="account@example.com")
+    account.id = ACCOUNT_ID
+    join = TenantAccountJoin(
+        tenant_id=tenant_id,
+        account_id=account.id,
+        current=True,
+        role=TenantAccountRole.NORMAL,
+    )
+    db.caller_session.add_all([account, join])
+    db.caller_session.commit()
+    return account
+
+
+def _persist_end_user(
+    db: SqliteToolDb,
+    *,
+    end_user_id: str = END_USER_ID,
+    tenant_id: str = TENANT_ID,
+) -> EndUser:
+    end_user = EndUser(
+        id=end_user_id,
+        tenant_id=tenant_id,
+        app_id=APP_ID,
+        type=EndUserType.SERVICE_API,
+        name="End user",
+        session_id="end-user-session",
+    )
+    db.caller_session.add(end_user)
+    db.caller_session.commit()
+    return end_user
+
+
+def _persist_app(db: SqliteToolDb) -> App:
+    app = App(
+        id=APP_ID,
+        tenant_id=TENANT_ID,
+        name="Workflow app",
+        description="",
+        mode=AppMode.WORKFLOW,
+        icon_type=None,
+        icon="",
+        icon_background=None,
+        app_model_config_id=None,
+        workflow_id=None,
+        enable_site=False,
+        enable_api=True,
+        max_active_requests=None,
+        created_by=CREATOR_ID,
+    )
+    db.caller_session.add(app)
+    db.caller_session.commit()
+    return app
+
+
+def _persist_workflow(db: SqliteToolDb, *, version: str, workflow_id: str | None = None) -> Workflow:
+    workflow = Workflow.new(
+        tenant_id=TENANT_ID,
+        app_id=APP_ID,
+        type=WorkflowType.WORKFLOW.value,
+        version=version,
+        graph=json.dumps({"nodes": [], "edges": []}),
+        features="{}",
+        created_by=CREATOR_ID,
+        environment_variables=[],
+        conversation_variables=[],
+        rag_pipeline_variables=[],
+    )
+    workflow.id = workflow_id or str(uuid.uuid4())
+    db.caller_session.add(workflow)
+    db.caller_session.commit()
+    return workflow
+
+
+def _build_tool(*, tenant_id: str = "test_tool", workflow_app_id: str = "app-1", version: str = "1") -> WorkflowTool:
     entity = ToolEntity(
         identity=ToolIdentity(author="test", name="test tool", label=I18nObject(en_US="test tool"), provider="test"),
         parameters=[],
         description=None,
         has_runtime_parameters=False,
     )
-    runtime = ToolRuntime(tenant_id="test_tool", invoke_from=InvokeFrom.EXPLORE)
+    runtime = ToolRuntime(tenant_id=tenant_id, invoke_from=InvokeFrom.EXPLORE)
     return WorkflowTool(
-        workflow_app_id="app-1",
+        workflow_app_id=workflow_app_id,
         workflow_as_tool_id="wf-tool-1",
-        version="1",
+        version=version,
         workflow_entities={},
         workflow_call_depth=1,
         entity=entity,
@@ -98,7 +166,10 @@ def _build_tool() -> WorkflowTool:
     )
 
 
-def test_workflow_tool_should_raise_tool_invoke_error_when_result_has_error_field(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_should_raise_tool_invoke_error_when_result_has_error_field(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure that WorkflowTool will throw a `ToolInvokeError` exception when
     `WorkflowAppGenerator.generate` returns a result with `error` key inside
     the `data` element.
@@ -122,11 +193,14 @@ def test_workflow_tool_should_raise_tool_invoke_error_when_result_has_error_fiel
     with pytest.raises(ToolInvokeError) as exc_info:
         # WorkflowTool always returns a generator, so we need to iterate to
         # actually `run` the tool.
-        list(tool.invoke(MagicMock(), "test_user", {}))
+        list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
     assert exc_info.value.args == ("oops",)
 
 
-def test_workflow_tool_does_not_use_pause_state_config(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_does_not_use_pause_state_config(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure pause_state_config is passed as None."""
     tool = _build_tool()
 
@@ -140,14 +214,17 @@ def test_workflow_tool_does_not_use_pause_state_config(monkeypatch: pytest.Monke
     monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate", generate_mock)
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
-    list(tool.invoke(MagicMock(), "test_user", {}))
+    list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
 
     call_kwargs = generate_mock.call_args.kwargs
     assert "pause_state_config" in call_kwargs
     assert call_kwargs["pause_state_config"] is None
 
 
-def test_workflow_tool_passes_parent_trace_context_from_runtime(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_passes_parent_trace_context_from_runtime(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure nested workflow runtime metadata is forwarded as parent trace context."""
     tool = _build_tool()
     tool.set_parent_trace_context(
@@ -165,7 +242,7 @@ def test_workflow_tool_passes_parent_trace_context_from_runtime(monkeypatch: pyt
     monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate", generate_mock)
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
-    list(tool.invoke(MagicMock(), "test_user", {}))
+    list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
 
     call_kwargs = generate_mock.call_args.kwargs
     assert call_kwargs["args"]["parent_trace_context"].model_dump() == {
@@ -174,7 +251,10 @@ def test_workflow_tool_passes_parent_trace_context_from_runtime(monkeypatch: pyt
     }
 
 
-def test_workflow_tool_passes_parent_trace_session_id(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_passes_parent_trace_session_id(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure nested workflows inherit the parent observability session ID."""
     tool = _build_tool()
     tool.entity.parameters = [
@@ -197,14 +277,17 @@ def test_workflow_tool_passes_parent_trace_session_id(monkeypatch: pytest.Monkey
     monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate", generate_mock)
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
-    list(tool.invoke(MagicMock(), "test_user", {"trace_session_id": "user-input-session"}))
+    list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {"trace_session_id": "user-input-session"}))
 
     call_kwargs = generate_mock.call_args.kwargs
     assert call_kwargs["args"]["inputs"]["trace_session_id"] == "user-input-session"
     assert call_kwargs["args"]["trace_session_id"] == "session-1"
 
 
-def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure private trace context does not overwrite same-named workflow inputs."""
     tool = _build_tool()
     tool.entity.parameters = [
@@ -238,7 +321,7 @@ def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys(monkeypat
 
     list(
         tool.invoke(
-            MagicMock(),
+            sqlite_tool_db.caller_session,
             "test_user",
             {
                 "outer_workflow_run_id": "user-workflow-input",
@@ -256,7 +339,10 @@ def test_workflow_tool_keeps_user_inputs_named_like_trace_runtime_keys(monkeypat
     }
 
 
-def test_workflow_tool_can_clear_parent_trace_context(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_can_clear_parent_trace_context(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure reused WorkflowTool instances do not keep stale parent trace context."""
     tool = _build_tool()
     tool.set_parent_trace_context(
@@ -275,13 +361,16 @@ def test_workflow_tool_can_clear_parent_trace_context(monkeypatch: pytest.Monkey
     monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate", generate_mock)
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
-    list(tool.invoke(MagicMock(), "test_user", {}))
+    list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
 
     call_kwargs = generate_mock.call_args.kwargs
     assert "parent_trace_context" not in call_kwargs["args"]
 
 
-def test_workflow_tool_can_clear_trace_session_id(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_can_clear_trace_session_id(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure reused WorkflowTool instances do not keep stale trace session IDs."""
     tool = _build_tool()
     tool.set_trace_session_id("session-1")
@@ -297,7 +386,7 @@ def test_workflow_tool_can_clear_trace_session_id(monkeypatch: pytest.MonkeyPatc
     monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate", generate_mock)
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
-    list(tool.invoke(MagicMock(), "test_user", {}))
+    list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
 
     call_kwargs = generate_mock.call_args.kwargs
     assert "trace_session_id" not in call_kwargs["args"]
@@ -315,6 +404,7 @@ def test_workflow_tool_can_clear_trace_session_id(monkeypatch: pytest.MonkeyPatc
 def test_workflow_tool_omits_parent_trace_context_when_runtime_is_incomplete(
     monkeypatch: pytest.MonkeyPatch,
     runtime_parameters: dict[str, Any],
+    sqlite_tool_db: SqliteToolDb,
 ):
     """Ensure incomplete runtime metadata does not leak parent trace context into generator args."""
     tool = _build_tool()
@@ -330,13 +420,16 @@ def test_workflow_tool_omits_parent_trace_context_when_runtime_is_incomplete(
     monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate", generate_mock)
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
-    list(tool.invoke(MagicMock(), "test_user", {}))
+    list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
 
     call_kwargs = generate_mock.call_args.kwargs
     assert "parent_trace_context" not in call_kwargs["args"]
 
 
-def test_workflow_tool_should_generate_variable_messages_for_outputs(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_should_generate_variable_messages_for_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Test that WorkflowTool should generate variable messages when there are outputs"""
     tool = _build_tool()
 
@@ -359,7 +452,7 @@ def test_workflow_tool_should_generate_variable_messages_for_outputs(monkeypatch
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
     # Execute tool invocation
-    messages = list(tool.invoke(MagicMock(), "test_user", {}))
+    messages = list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
 
     # Verify variable messages
     variable_messages = [msg for msg in messages if msg.type == ToolInvokeMessage.MessageType.VARIABLE]
@@ -382,7 +475,10 @@ def test_workflow_tool_should_generate_variable_messages_for_outputs(monkeypatch
     assert json_messages[0].message.json_object == mock_outputs
 
 
-def test_workflow_tool_should_handle_empty_outputs(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_should_handle_empty_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Test that WorkflowTool should handle empty outputs correctly"""
     tool = _build_tool()
 
@@ -402,7 +498,7 @@ def test_workflow_tool_should_handle_empty_outputs(monkeypatch: pytest.MonkeyPat
     monkeypatch.setattr("libs.login.current_user", lambda *args, **kwargs: None)
 
     # Execute tool invocation
-    messages = list(tool.invoke(MagicMock(), "test_user", {}))
+    messages = list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {}))
 
     # Verify generated messages
     # Should contain: 0 variable messages + 1 text message + 1 JSON message = 2 messages
@@ -458,41 +554,32 @@ def test_create_file_message_should_include_file_marker():
     assert message.meta == {"file": file_obj}
 
 
-def test_resolve_user_from_database_falls_back_to_end_user(monkeypatch: pytest.MonkeyPatch):
+def test_resolve_user_from_database_falls_back_to_end_user(sqlite_tool_db: SqliteToolDb):
     """Ensure worker context can resolve EndUser when Account is missing."""
-
-    tenant = SimpleNamespace(id="tenant_id")
-    end_user = SimpleNamespace(id="end_user_id", tenant_id="tenant_id")
-
-    # Monkeypatch session factory to return our stub session
-    stub_session = StubSession(scalar_results=[tenant, None, end_user])
-    monkeypatch.setattr(
-        "core.tools.workflow_as_tool.tool.session_factory.create_session",
-        lambda: stub_session,
+    _persist_tenant(sqlite_tool_db)
+    end_user = _persist_end_user(sqlite_tool_db)
+    other_tenant_end_user = _persist_end_user(
+        sqlite_tool_db,
+        end_user_id="00000000-0000-0000-0000-000000000007",
+        tenant_id=OTHER_TENANT_ID,
     )
 
-    tool = _build_tool()
+    tool = _build_tool(tenant_id=TENANT_ID)
     tool.runtime.invoke_from = InvokeFrom.SERVICE_API
-    tool.runtime.tenant_id = "tenant_id"
 
     resolved_user = tool._resolve_user_from_database(user_id=end_user.id)
 
-    assert resolved_user is end_user
-    assert stub_session.expunge_calls == [end_user]
+    assert isinstance(resolved_user, EndUser)
+    assert resolved_user.id == end_user.id
+    assert resolved_user.tenant_id == TENANT_ID
+    assert inspect(resolved_user).detached is True
+    assert tool._resolve_user_from_database(user_id=other_tenant_end_user.id) is None
 
 
-def test_resolve_user_from_database_returns_none_when_no_tenant(monkeypatch: pytest.MonkeyPatch):
+def test_resolve_user_from_database_returns_none_when_no_tenant(sqlite_tool_db: SqliteToolDb):
     """Return None if tenant cannot be found in worker context."""
-
-    # Monkeypatch session factory to return our stub session with no tenant
-    monkeypatch.setattr(
-        "core.tools.workflow_as_tool.tool.session_factory.create_session",
-        lambda: StubSession(scalar_results=[None]),
-    )
-
-    tool = _build_tool()
+    tool = _build_tool(tenant_id=OTHER_TENANT_ID)
     tool.runtime.invoke_from = InvokeFrom.SERVICE_API
-    tool.runtime.tenant_id = "missing_tenant"
 
     resolved_user = tool._resolve_user_from_database(user_id="any")
 
@@ -544,7 +631,10 @@ def test_extract_usage_from_nested():
     assert nested == {"total_tokens": 3}
 
 
-def test_invoke_raises_when_user_not_found(monkeypatch: pytest.MonkeyPatch):
+def test_invoke_raises_when_user_not_found(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Raise ToolInvokeError when user resolution fails."""
     tool = _build_tool()
     monkeypatch.setattr(tool, "_get_app", lambda *args, **kwargs: None)
@@ -552,58 +642,45 @@ def test_invoke_raises_when_user_not_found(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(tool, "_resolve_user", lambda *args, **kwargs: None)
 
     with pytest.raises(ToolInvokeError, match="User not found"):
-        list(tool.invoke(MagicMock(), "missing", {}))
+        list(tool.invoke(sqlite_tool_db.caller_session, "missing", {}))
 
 
-def test_resolve_user_from_database_returns_account(monkeypatch: pytest.MonkeyPatch):
+def test_resolve_user_from_database_returns_account(sqlite_tool_db: SqliteToolDb):
     """Resolve Account and set tenant in worker context."""
-    tenant = SimpleNamespace(id="tenant_id")
-    account = SimpleNamespace(id="account_id", current_tenant=None)
-    set_current_tenant = Mock(side_effect=lambda tenant, *, session: setattr(account, "current_tenant", tenant))
-    account.set_current_tenant_with_session = set_current_tenant
-    session = StubSession(scalar_results=[tenant, account])
+    tenant = _persist_tenant(sqlite_tool_db)
+    account = _persist_account(sqlite_tool_db)
+    tool = _build_tool(tenant_id=TENANT_ID)
 
-    monkeypatch.setattr("core.tools.workflow_as_tool.tool.session_factory.create_session", lambda: session)
-    tool = _build_tool()
-    tool.runtime.tenant_id = "tenant_id"
-
-    resolved = tool._resolve_user_from_database(user_id="account_id")
-    assert resolved is account
-    assert account.current_tenant is tenant
-    set_current_tenant.assert_called_once_with(tenant, session=session)
-    assert session.expunge_calls == [account]
+    resolved = tool._resolve_user_from_database(user_id=account.id)
+    assert isinstance(resolved, Account)
+    assert resolved.id == account.id
+    assert resolved.current_tenant_id == tenant.id
+    assert inspect(resolved).detached is True
 
 
-def test_get_workflow_and_get_app_db_branches(monkeypatch: pytest.MonkeyPatch):
+def test_get_workflow_and_get_app_db_branches(sqlite_tool_db: SqliteToolDb):
     """Cover workflow/app retrieval branches and error cases."""
-    tool = _build_tool()
-    latest_workflow = SimpleNamespace(id="wf-latest")
-    specific_workflow = SimpleNamespace(id="wf-v1")
-    app = SimpleNamespace(id="app-1")
-    sessions = iter(
-        [
-            StubSession(scalar_results=[], scalars_results=[latest_workflow]),
-            StubSession(scalar_results=[specific_workflow], scalars_results=[]),
-            StubSession(scalar_results=[app], scalars_results=[]),
-        ]
-    )
-    monkeypatch.setattr(
-        "core.tools.workflow_as_tool.tool.session_factory.create_session",
-        lambda: next(sessions),
-    )
+    app = _persist_app(sqlite_tool_db)
+    specific_workflow = _persist_workflow(sqlite_tool_db, version="1")
+    latest_workflow = _persist_workflow(sqlite_tool_db, version="2")
+    _persist_workflow(sqlite_tool_db, version=Workflow.VERSION_DRAFT)
+    tool = _build_tool(tenant_id=TENANT_ID, workflow_app_id=APP_ID)
 
-    assert tool._get_workflow("app-1", "") is latest_workflow
-    assert tool._get_workflow("app-1", "1") is specific_workflow
-    assert tool._get_app("app-1") is app
+    latest = tool._get_workflow(APP_ID, "")
+    specific = tool._get_workflow(APP_ID, "1")
+    resolved_app = tool._get_app(APP_ID)
+
+    assert latest.id == latest_workflow.id
+    assert specific.id == specific_workflow.id
+    assert resolved_app.id == app.id
+    assert inspect(latest).detached is True
+    assert inspect(specific).detached is True
+    assert inspect(resolved_app).detached is True
 
-    monkeypatch.setattr(
-        "core.tools.workflow_as_tool.tool.session_factory.create_session",
-        lambda: StubSession(scalar_results=[None, None], scalars_results=[None]),
-    )
     with pytest.raises(ValueError, match="workflow not found"):
-        tool._get_workflow("app-1", "1")
+        tool._get_workflow(APP_ID, "missing")
     with pytest.raises(ValueError, match="app not found"):
-        tool._get_app("app-1")
+        tool._get_app("00000000-0000-0000-0000-000000000099")
 
 
 def _setup_transform_args_tool(monkeypatch: pytest.MonkeyPatch) -> WorkflowTool:
@@ -722,7 +799,10 @@ def test_transform_args_normalizes_optional_files_parameter(
     assert files == []
 
 
-def test_workflow_tool_invocation_normalizes_optional_files_parameter(monkeypatch: pytest.MonkeyPatch):
+def test_workflow_tool_invocation_normalizes_optional_files_parameter(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_tool_db: SqliteToolDb,
+):
     """Ensure casted empty FILES values do not reach workflow input validation as [None]."""
     tool = _build_tool()
     images_param = ToolParameter.get_simple_instance(
@@ -741,7 +821,7 @@ def test_workflow_tool_invocation_normalizes_optional_files_parameter(monkeypatc
     generate_mock = MagicMock(return_value={"data": {}})
     monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppGenerator.generate", generate_mock)
 
-    list(tool.invoke(MagicMock(), "test_user", {"images": None}))
+    list(tool.invoke(sqlite_tool_db.caller_session, "test_user", {"images": None}))
 
     call_kwargs = generate_mock.call_args.kwargs
     assert call_kwargs["args"]["inputs"]["images"] == []

api/tests/unit_tests/core/workflow/generator/test_prompts.py

@@ -26,6 +26,7 @@ class TestPlannerSystemPrompt:
         """Auto-mode resolution rides on the planner echoing its mode choice."""
         assert '"mode": "workflow | advanced-chat"' in PLANNER_SYSTEM_PROMPT
         assert "When the ``# Mode`` section says auto, YOU decide" in PLANNER_SYSTEM_PROMPT
+        assert "userinput.files" in PLANNER_SYSTEM_PROMPT
 
 
 class TestFormatIdealOutputSection:
@@ -75,6 +76,9 @@ class TestNodeBuilderPrompt:
         assert '"viewport":' not in prompt
         assert '"positionAbsolute":' not in prompt
 
+    def test_start_uses_user_input_files(self):
+        assert "userinput.files" in get_node_builder_system_prompt("start")
+
     def test_supports_main_human_input_and_assigner_contracts(self):
         human_input = get_node_builder_system_prompt("human-input")
         assigner = get_node_builder_system_prompt("assigner")
@@ -130,17 +134,19 @@ class TestNodeBuilderUserSections:
 
 
 class TestModeSection:
-    def test_advanced_chat_documents_system_variables(self):
+    def test_advanced_chat_documents_built_in_variables(self):
         out = format_mode_section("advanced-chat")
 
         assert "sys.query" in out
-        assert '["sys", "query"]' in out
+        assert "userinput.files" in out
+        assert '["userinput", "files"]' in out
         assert "do NOT invent start-node variables" in out
 
-    def test_workflow_mode_forbids_system_variables(self):
+    def test_workflow_mode_documents_file_input(self):
         out = format_mode_section("workflow")
 
-        assert "NO automatic system variables" in out
+        assert "userinput.files" in out
+        assert "start node's declared variables" in out
 
 
 class TestExistingGraphSection:

api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py

@@ -228,12 +228,12 @@ def _previous_node_prompt_payload(result, selector: str) -> object:
 
 
 def _uploaded_workflow_files_prompt_payload(result) -> object:
-    prefix = "  - sys.files: "
+    prefix = "  - userinput.files: "
     user_prompt = _workflow_user_prompt(result)
     for line in user_prompt.splitlines():
         if line.startswith(prefix):
             return json.loads(line.removeprefix(prefix))
-    raise AssertionError("missing prompt payload for sys.files")
+    raise AssertionError("missing prompt payload for userinput.files")
 
 
 def test_builds_create_run_request_from_agent_soul_and_node_job():

api/tests/unit_tests/core/workflow/test_legacy_system_files.py

@@ -0,0 +1,161 @@
+from core.workflow.legacy_system_files import (
+    LegacySysFilesCompatVariable,
+    attach_legacy_sys_files_warning,
+    migrate_legacy_sys_files_graph_with_result,
+    normalize_legacy_sys_files_args,
+    resolve_legacy_sys_files_compat_variable,
+)
+
+_LEGACY_NODE_ID = "sys"
+_LEGACY_ALIAS_NODE_ID = "userinput"
+_LEGACY_VARIABLE_NAME = "files"
+_LEGACY_SELECTOR = [_LEGACY_NODE_ID, _LEGACY_VARIABLE_NAME]
+_LEGACY_TEMPLATE = "{{#" + ".".join((_LEGACY_NODE_ID, _LEGACY_VARIABLE_NAME)) + "#}}"
+_LEGACY_ALIAS_SELECTOR = [_LEGACY_ALIAS_NODE_ID, _LEGACY_VARIABLE_NAME]
+_LEGACY_ALIAS_TEMPLATE = "{{#" + ".".join((_LEGACY_ALIAS_NODE_ID, _LEGACY_VARIABLE_NAME)) + "#}}"
+_LEGACY_ALIAS_INPUT_KEY = ".".join((_LEGACY_ALIAS_NODE_ID, _LEGACY_VARIABLE_NAME))
+
+
+def test_migrate_legacy_sys_files_graph_ignores_invalid_or_unrelated_graphs():
+    assert not migrate_legacy_sys_files_graph_with_result({}).changed
+    assert not migrate_legacy_sys_files_graph_with_result({"nodes": [], "edges": [_LEGACY_SELECTOR]}).changed
+    assert not migrate_legacy_sys_files_graph_with_result({"nodes": [{"data": {"value": ["sys", "query"]}}]}).changed
+
+
+def test_migrate_legacy_sys_files_graph_rewrites_sys_files_to_userinput_files_without_start_variable():
+    graph = {
+        "nodes": [
+            {"id": "start", "data": {"type": "start", "variables": [{"variable": "sys_files"}]}},
+            {
+                "id": "answer",
+                "data": {
+                    "type": "answer",
+                    "answer": _LEGACY_SELECTOR,
+                    "template": _LEGACY_TEMPLATE,
+                },
+            },
+        ],
+    }
+
+    result = migrate_legacy_sys_files_graph_with_result(graph)
+
+    assert result.changed
+    start_data = result.graph["nodes"][0]["data"]
+    assert start_data["variables"] == [{"variable": "sys_files"}]
+    assert result.graph["nodes"][1]["data"]["answer"] == _LEGACY_ALIAS_SELECTOR
+    assert result.graph["nodes"][1]["data"]["template"] == _LEGACY_ALIAS_TEMPLATE
+    assert graph["nodes"][1]["data"]["answer"] == _LEGACY_SELECTOR
+    assert graph["nodes"][1]["data"]["template"] == _LEGACY_TEMPLATE
+
+
+def test_migrate_legacy_sys_files_graph_leaves_userinput_files_target_unchanged():
+    graph = {
+        "nodes": [
+            {"id": "start", "data": {"type": "start", "variables": []}},
+            {
+                "id": "answer",
+                "data": {
+                    "type": "answer",
+                    "answer": _LEGACY_ALIAS_SELECTOR,
+                    "template": _LEGACY_ALIAS_TEMPLATE,
+                },
+            },
+        ],
+    }
+
+    result = migrate_legacy_sys_files_graph_with_result(graph)
+
+    assert not result.changed
+    assert result.graph == graph
+
+
+def test_resolve_legacy_sys_files_compat_variable_returns_userinput_files_target():
+    assert resolve_legacy_sys_files_compat_variable({}) is None
+    assert resolve_legacy_sys_files_compat_variable({"nodes": [{"data": {"value": ["sys", "query"]}}]}) is None
+
+    compat_variable = resolve_legacy_sys_files_compat_variable({"nodes": [{"data": {"value": _LEGACY_SELECTOR}}]})
+
+    assert compat_variable == LegacySysFilesCompatVariable(
+        node_id=_LEGACY_ALIAS_NODE_ID,
+        variable_name=_LEGACY_VARIABLE_NAME,
+    )
+    assert (
+        resolve_legacy_sys_files_compat_variable({"nodes": [{"data": {"value": _LEGACY_ALIAS_SELECTOR}}]})
+        == compat_variable
+    )
+
+
+def test_normalize_legacy_sys_files_args_handles_no_compat_and_top_level_files():
+    args_without_legacy, compat_without_legacy = normalize_legacy_sys_files_args(
+        graph={"nodes": []},
+        args={"inputs": {}},
+    )
+    assert args_without_legacy == {"inputs": {}}
+    assert compat_without_legacy is None
+
+    files = [{"id": "file-1"}]
+    graph = {
+        "nodes": [
+            {"id": "start", "data": {"type": "start", "variables": []}},
+            {"id": "answer", "data": {"type": "answer", "answer": _LEGACY_TEMPLATE}},
+        ],
+    }
+    normalized_args, compat_variable = normalize_legacy_sys_files_args(
+        graph=graph,
+        args={"inputs": {}, "files": files},
+    )
+
+    assert compat_variable is not None
+    assert normalized_args["files"] == files
+    assert normalized_args["inputs"][".".join((compat_variable.node_id, compat_variable.variable_name))] == files
+
+
+def test_normalize_legacy_sys_files_args_maps_userinput_files_to_top_level_files_without_warning():
+    files = [{"id": "file-1"}]
+    normalized_args, compat_variable = normalize_legacy_sys_files_args(
+        graph={"nodes": []},
+        args={"inputs": {_LEGACY_ALIAS_INPUT_KEY: files}},
+    )
+
+    assert compat_variable is None
+    assert normalized_args["files"] == files
+    assert normalized_args["inputs"] == {_LEGACY_ALIAS_INPUT_KEY: files}
+
+
+def test_normalize_legacy_sys_files_args_prefers_userinput_files_over_legacy_files():
+    legacy_files = [{"id": "legacy-file"}]
+    userinput_files = [{"id": "userinput-file"}]
+
+    normalized_args, compat_variable = normalize_legacy_sys_files_args(
+        graph={"nodes": [{"data": {"type": "answer", "answer": _LEGACY_ALIAS_TEMPLATE}}]},
+        args={
+            "inputs": {_LEGACY_ALIAS_INPUT_KEY: userinput_files},
+            "files": legacy_files,
+        },
+    )
+
+    assert compat_variable is None
+    assert normalized_args["files"] == userinput_files
+
+
+def test_attach_legacy_sys_files_warning_wraps_stream_and_closes_source():
+    class CloseableStream:
+        closed = False
+
+        def __iter__(self):
+            yield "data: payload\n\n"
+
+        def close(self):
+            self.closed = True
+
+    stream = CloseableStream()
+    wrapped = attach_legacy_sys_files_warning(
+        stream,
+        LegacySysFilesCompatVariable(node_id=_LEGACY_ALIAS_NODE_ID, variable_name=_LEGACY_VARIABLE_NAME),
+    )
+
+    chunks = list(wrapped)
+
+    assert "warning" in chunks[0]
+    assert chunks[1] == "data: payload\n\n"
+    assert stream.closed

api/tests/unit_tests/core/workflow/test_system_variable.py

@@ -1,6 +1,7 @@
 from types import SimpleNamespace
 
 from core.workflow.system_variables import (
+    build_bootstrap_variables,
     build_system_variables,
     default_system_variables,
     get_node_creation_preload_selectors,
@@ -56,6 +57,25 @@ def test_build_system_variables_preserves_file_values():
     assert system_values["files"] == [file]
 
 
+def test_build_bootstrap_variables_adds_userinput_files_alias():
+    file = File(
+        file_type=FileType.DOCUMENT,
+        transfer_method=FileTransferMethod.LOCAL_FILE,
+        related_id="file-id",
+        filename="test.txt",
+        extension=".txt",
+        mime_type="text/plain",
+        size=1,
+        storage_key="storage-key",
+    )
+
+    bootstrap_variables = build_bootstrap_variables(system_variables=build_system_variables(files=[file]))
+    file_variables_by_selector = {tuple(variable.selector): variable for variable in bootstrap_variables}
+
+    assert file_variables_by_selector[("sys", "files")].value == [file]
+    assert file_variables_by_selector[("userinput", "files")].value == [file]
+
+
 def test_default_system_variables_generates_workflow_run_id():
     system_variables = default_system_variables()
     system_values = system_variables_to_mapping(system_variables)

api/tests/unit_tests/events/test_update_provider_when_message_created.py

@@ -18,7 +18,6 @@ from models.provider import ProviderType
 @pytest.fixture
 def credit_pool_session_factory(sqlite_engine: Engine) -> Iterator[sessionmaker[Session]]:
     """Bind message-created accounting to fixture-owned SQLite sessions."""
-    TenantCreditPool.__table__.create(sqlite_engine)
     session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
     with patch("events.event_handlers.update_provider_when_message_created.db.session", session_factory):
         yield session_factory

api/tests/unit_tests/models/test_workflow.py

@@ -18,6 +18,11 @@ from models.workflow import (
     is_system_variable_editable,
 )
 
+_LEGACY_FILE_TEMPLATE = "{{#" + ".".join(("sys", "files")) + "#}}"
+_LEGACY_FILE_SELECTOR = ["sys", "files"]
+_USER_INPUT_FILE_TEMPLATE = "{{#" + ".".join(("userinput", "files")) + "#}}"
+_USER_INPUT_FILE_SELECTOR = ["userinput", "files"]
+
 
 def test_environment_variables():
     # tenant_id context variable removed - using current_user.current_tenant_id directly
@@ -245,6 +250,144 @@ class TestIsSystemVariableEditable:
         assert is_system_variable_editable("invalid_or_new_system_variable") == False
 
 
+class TestWorkflowLegacySysFilesCompatibility:
+    def _make_workflow(self, graph: dict, *, features: dict | None = None) -> Workflow:
+        return Workflow(
+            tenant_id="tenant_id",
+            app_id="app_id",
+            type="workflow",
+            version="draft",
+            graph=json.dumps(graph),
+            features=json.dumps(features or {}),
+            created_by="account_id",
+            environment_variables=[],
+            conversation_variables=[],
+        )
+
+    def test_graph_dict_rewrites_legacy_sys_files_references_to_userinput_files(self):
+        workflow = self._make_workflow(
+            {
+                "nodes": [
+                    {
+                        "id": "start",
+                        "data": {
+                            "type": "start",
+                            "title": "Start",
+                            "variables": [],
+                        },
+                    },
+                    {
+                        "id": "llm",
+                        "data": {
+                            "type": "llm",
+                            "prompt_template": [{"role": "user", "text": f"files: {_LEGACY_FILE_TEMPLATE}"}],
+                            "context": {"variable_selector": _LEGACY_FILE_SELECTOR},
+                        },
+                    },
+                ],
+                "edges": [],
+            }
+        )
+
+        stored_graph_before_read = workflow.graph
+
+        graph = workflow.graph_dict
+        start_node = next(node for node in graph["nodes"] if node["id"] == "start")
+        llm_node = next(node for node in graph["nodes"] if node["id"] == "llm")
+
+        assert start_node["data"]["variables"] == []
+        assert llm_node["data"]["prompt_template"][0]["text"] == f"files: {_USER_INPUT_FILE_TEMPLATE}"
+        assert llm_node["data"]["context"]["variable_selector"] == _USER_INPUT_FILE_SELECTOR
+
+        assert workflow.graph == stored_graph_before_read
+
+    def test_migrate_legacy_sys_files_graph_in_place_updates_stored_graph(self):
+        workflow = self._make_workflow(
+            {
+                "nodes": [
+                    {"id": "answer", "data": {"type": "answer", "answer": _LEGACY_FILE_TEMPLATE}},
+                ],
+                "edges": [],
+            }
+        )
+
+        assert workflow.migrate_legacy_sys_files_graph_in_place()
+        assert _LEGACY_FILE_TEMPLATE not in workflow.graph
+        assert _USER_INPUT_FILE_TEMPLATE in workflow.graph
+
+    def test_graph_dict_preserves_existing_start_variables_when_migrating_legacy_sys_files(self):
+        workflow = self._make_workflow(
+            {
+                "nodes": [
+                    {
+                        "id": "start",
+                        "data": {
+                            "type": "start",
+                            "title": "Start",
+                            "variables": [
+                                {"variable": "sys_files", "label": "Existing", "type": "text-input"},
+                            ],
+                        },
+                    },
+                    {
+                        "id": "answer",
+                        "data": {
+                            "type": "answer",
+                            "answer": _LEGACY_FILE_TEMPLATE,
+                        },
+                    },
+                ],
+                "edges": [],
+            }
+        )
+
+        graph = workflow.graph_dict
+        start_node = next(node for node in graph["nodes"] if node["id"] == "start")
+        answer_node = next(node for node in graph["nodes"] if node["id"] == "answer")
+
+        assert [variable["variable"] for variable in start_node["data"]["variables"]] == ["sys_files"]
+        assert answer_node["data"]["answer"] == _USER_INPUT_FILE_TEMPLATE
+
+    def test_graph_dict_leaves_userinput_files_references_unchanged(self):
+        workflow = self._make_workflow(
+            {
+                "nodes": [
+                    {
+                        "id": "start",
+                        "data": {
+                            "type": "start",
+                            "title": "Start",
+                            "variables": [],
+                        },
+                    },
+                    {
+                        "id": "answer",
+                        "data": {
+                            "type": "answer",
+                            "answer": _USER_INPUT_FILE_TEMPLATE,
+                        },
+                    },
+                ],
+                "edges": [],
+            },
+            features={
+                "file_upload": {
+                    "enabled": True,
+                    "allowed_file_upload_methods": ["remote_url"],
+                    "allowed_file_types": ["document", "custom"],
+                    "allowed_file_extensions": [".pdf"],
+                    "number_limits": 8,
+                }
+            },
+        )
+
+        graph = workflow.graph_dict
+        start_node = next(node for node in graph["nodes"] if node["id"] == "start")
+
+        assert start_node["data"]["variables"] == []
+        assert json.loads(workflow.graph) == graph
+
+
 class TestWorkflowDraftVariableGetValue:
     def test_get_value_by_case(self):
         @dataclasses.dataclass

api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py

@@ -1,89 +1,105 @@
 import logging
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
+from uuid import uuid4
 
 import pytest
+from sqlalchemy import select
+from sqlalchemy.orm import Session
 
 from models.account import (
     TenantPluginAutoUpgradeCategory,
     TenantPluginAutoUpgradeMode,
+    TenantPluginAutoUpgradeStrategy,
     TenantPluginAutoUpgradeStrategySetting,
 )
+from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
 
 MODULE = "services.plugin.plugin_auto_upgrade_service"
 PLUGIN_CATEGORY = TenantPluginAutoUpgradeCategory.TOOL
+STRATEGY_MODELS = (TenantPluginAutoUpgradeStrategy,)
 
 
-def _patched_session():
-    """Return a mock SQLAlchemy session for service calls."""
-    session = MagicMock()
-    return session
+def _strategy(
+    tenant_id: str,
+    *,
+    category: TenantPluginAutoUpgradeCategory = PLUGIN_CATEGORY,
+    setting: TenantPluginAutoUpgradeStrategySetting = TenantPluginAutoUpgradeStrategySetting.FIX_ONLY,
+    mode: TenantPluginAutoUpgradeMode = TenantPluginAutoUpgradeMode.EXCLUDE,
+    exclude: list[str] | None = None,
+    include: list[str] | None = None,
+    upgrade_time: int = 0,
+) -> TenantPluginAutoUpgradeStrategy:
+    return TenantPluginAutoUpgradeStrategy(
+        tenant_id=tenant_id,
+        category=category,
+        strategy_setting=setting,
+        upgrade_time_of_day=upgrade_time,
+        upgrade_mode=mode,
+        exclude_plugins=exclude or [],
+        include_plugins=include or [],
+    )
 
 
 class TestGetStrategy:
-    def test_returns_strategy_when_found(self):
-        session = _patched_session()
-        strategy = MagicMock()
-        session.scalar.return_value = strategy
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_returns_strategy_when_found(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
+        strategy = _strategy(tenant_id)
+        sqlite_session.add(strategy)
+        sqlite_session.commit()
 
-        from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-        result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY, session=session)
+        result = PluginAutoUpgradeService.get_strategy(tenant_id, PLUGIN_CATEGORY, session=sqlite_session)
 
         assert result is strategy
 
-    def test_returns_none_when_not_found(self):
-        session = _patched_session()
-        session.scalar.return_value = None
-
-        from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-        result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY, session=session)
-
-        assert result is None
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_returns_none_when_not_found(self, sqlite_session: Session) -> None:
+        assert PluginAutoUpgradeService.get_strategy(str(uuid4()), PLUGIN_CATEGORY, session=sqlite_session) is None
 
 
 class TestChangeStrategy:
-    def test_creates_new_strategy(self):
-        session = _patched_session()
-        session.scalar.return_value = None
-
-        with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls:
-            strat_cls.return_value = MagicMock()
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-            result = PluginAutoUpgradeService.change_strategy(
-                "t1",
-                TenantPluginAutoUpgradeStrategySetting.FIX_ONLY,
-                3,
-                TenantPluginAutoUpgradeMode.ALL,
-                [],
-                [],
-                category=PLUGIN_CATEGORY,
-                session=session,
-            )
-
-        assert result is True
-        session.add.assert_called_once()
-
-    def test_updates_existing_strategy(self):
-        session = _patched_session()
-        existing = MagicMock()
-        session.scalar.return_value = existing
-
-        from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_creates_new_strategy(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
 
         result = PluginAutoUpgradeService.change_strategy(
-            "t1",
+            tenant_id,
+            TenantPluginAutoUpgradeStrategySetting.FIX_ONLY,
+            3,
+            TenantPluginAutoUpgradeMode.ALL,
+            [],
+            [],
+            category=PLUGIN_CATEGORY,
+            session=sqlite_session,
+        )
+
+        strategy = sqlite_session.scalar(select(TenantPluginAutoUpgradeStrategy))
+        assert result is True
+        assert strategy is not None
+        assert strategy.tenant_id == tenant_id
+        assert strategy.upgrade_time_of_day == 3
+        assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.ALL
+
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_updates_existing_strategy(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
+        existing = _strategy(tenant_id)
+        sqlite_session.add(existing)
+        sqlite_session.commit()
+
+        result = PluginAutoUpgradeService.change_strategy(
+            tenant_id,
             TenantPluginAutoUpgradeStrategySetting.LATEST,
             5,
             TenantPluginAutoUpgradeMode.PARTIAL,
             ["p1"],
             ["p2"],
             category=PLUGIN_CATEGORY,
-            session=session,
+            session=sqlite_session,
         )
 
+        sqlite_session.refresh(existing)
         assert result is True
         assert existing.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST
         assert existing.upgrade_time_of_day == 5
@@ -93,157 +109,115 @@ class TestChangeStrategy:
 
 
 class TestExcludePlugin:
-    def test_creates_default_strategy_when_none_exists(self):
-        session = _patched_session()
-        session.scalar.return_value = None
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_creates_default_strategy_when_none_exists(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
 
-        with (
-            patch(f"{MODULE}.select"),
-            patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy"),
-        ):
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-            result = PluginAutoUpgradeService.exclude_plugin(
-                "t1",
-                "plugin-1",
-                PLUGIN_CATEGORY,
-                session=session,
-            )
+        result = PluginAutoUpgradeService.exclude_plugin(tenant_id, "plugin-1", PLUGIN_CATEGORY, session=sqlite_session)
 
+        strategy = sqlite_session.scalar(select(TenantPluginAutoUpgradeStrategy))
         assert result is True
-        session.add.assert_called_once()
+        assert strategy is not None
+        assert strategy.exclude_plugins == ["plugin-1"]
 
-    def test_appends_to_exclude_list_in_exclude_mode(self):
-        session = _patched_session()
-        existing = MagicMock()
-        existing.upgrade_mode = TenantPluginAutoUpgradeMode.EXCLUDE
-        existing.exclude_plugins = ["p-existing"]
-        session.scalar.return_value = existing
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_appends_to_exclude_list_in_exclude_mode(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
+        existing = _strategy(tenant_id, exclude=["p-existing"])
+        sqlite_session.add(existing)
+        sqlite_session.commit()
 
-        with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls:
-            strat_cls.UpgradeMode.EXCLUDE = "exclude"
-            strat_cls.UpgradeMode.PARTIAL = "partial"
-            strat_cls.UpgradeMode.ALL = "all"
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
+        PluginAutoUpgradeService.exclude_plugin(tenant_id, "p-new", PLUGIN_CATEGORY, session=sqlite_session)
 
-            result = PluginAutoUpgradeService.exclude_plugin("t1", "p-new", PLUGIN_CATEGORY, session=session)
-
-        assert result is True
+        sqlite_session.refresh(existing)
         assert existing.exclude_plugins == ["p-existing", "p-new"]
 
-    def test_removes_from_include_list_in_partial_mode(self):
-        session = _patched_session()
-        existing = MagicMock()
-        existing.upgrade_mode = TenantPluginAutoUpgradeMode.PARTIAL
-        existing.include_plugins = ["p1", "p2"]
-        session.scalar.return_value = existing
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_removes_from_include_list_in_partial_mode(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
+        existing = _strategy(tenant_id, mode=TenantPluginAutoUpgradeMode.PARTIAL, include=["p1", "p2"])
+        sqlite_session.add(existing)
+        sqlite_session.commit()
 
-        with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls:
-            strat_cls.UpgradeMode.EXCLUDE = "exclude"
-            strat_cls.UpgradeMode.PARTIAL = "partial"
-            strat_cls.UpgradeMode.ALL = "all"
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
+        PluginAutoUpgradeService.exclude_plugin(tenant_id, "p1", PLUGIN_CATEGORY, session=sqlite_session)
 
-            result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY, session=session)
-
-        assert result is True
+        sqlite_session.refresh(existing)
         assert existing.include_plugins == ["p2"]
 
-    def test_switches_to_exclude_mode_from_all(self):
-        session = _patched_session()
-        existing = MagicMock()
-        existing.upgrade_mode = TenantPluginAutoUpgradeMode.ALL
-        session.scalar.return_value = existing
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_switches_to_exclude_mode_from_all(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
+        existing = _strategy(tenant_id, mode=TenantPluginAutoUpgradeMode.ALL)
+        sqlite_session.add(existing)
+        sqlite_session.commit()
 
-        with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls:
-            strat_cls.UpgradeMode.EXCLUDE = "exclude"
-            strat_cls.UpgradeMode.PARTIAL = "partial"
-            strat_cls.UpgradeMode.ALL = "all"
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
+        PluginAutoUpgradeService.exclude_plugin(tenant_id, "p1", PLUGIN_CATEGORY, session=sqlite_session)
 
-            result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY, session=session)
-
-        assert result is True
+        sqlite_session.refresh(existing)
         assert existing.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE
         assert existing.exclude_plugins == ["p1"]
 
-    def test_no_duplicate_in_exclude_list(self):
-        session = _patched_session()
-        existing = MagicMock()
-        existing.upgrade_mode = TenantPluginAutoUpgradeMode.EXCLUDE
-        existing.exclude_plugins = ["p1"]
-        session.scalar.return_value = existing
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_no_duplicate_in_exclude_list(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
+        existing = _strategy(tenant_id, exclude=["p1"])
+        sqlite_session.add(existing)
+        sqlite_session.commit()
 
-        with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls:
-            strat_cls.UpgradeMode.EXCLUDE = "exclude"
-            strat_cls.UpgradeMode.PARTIAL = "partial"
-            strat_cls.UpgradeMode.ALL = "all"
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-            PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY, session=session)
+        PluginAutoUpgradeService.exclude_plugin(tenant_id, "p1", PLUGIN_CATEGORY, session=sqlite_session)
 
+        sqlite_session.refresh(existing)
         assert existing.exclude_plugins == ["p1"]
 
 
 class TestBackfillStrategyCategories:
-    def test_creates_default_missing_categories_without_fetching_daemon(self):
-        session = _patched_session()
-        tool_strategy = SimpleNamespace(
-            category=TenantPluginAutoUpgradeCategory.TOOL,
-            strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY,
-            upgrade_time_of_day=0,
-            upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE,
-            exclude_plugins=[],
-            include_plugins=[],
-        )
-        session.scalars.return_value.all.return_value = [tool_strategy]
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_creates_default_missing_categories_without_fetching_daemon(self, sqlite_session: Session) -> None:
+        tenant_id = str(uuid4())
+        tool_strategy = _strategy(tenant_id)
+        sqlite_session.add(tool_strategy)
+        sqlite_session.commit()
         installer = MagicMock()
 
         with patch(f"{MODULE}.PluginInstaller", return_value=installer):
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-            result = PluginAutoUpgradeService.backfill_strategy_categories("t1", session=session)
-            expected_time = PluginAutoUpgradeService.default_upgrade_time_of_day("t1")
+            result = PluginAutoUpgradeService.backfill_strategy_categories(tenant_id, session=sqlite_session)
+            expected_time = PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id)
 
+        strategies = list(sqlite_session.scalars(select(TenantPluginAutoUpgradeStrategy)).all())
         assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 1
         assert result.normalized is False
         installer.list_plugins.assert_not_called()
+        assert len(strategies) == len(TenantPluginAutoUpgradeCategory)
         assert tool_strategy.upgrade_time_of_day == expected_time
-        created_strategies = [call.args[0] for call in session.add.call_args_list]
         model_strategy = next(
-            strategy for strategy in created_strategies if strategy.category == TenantPluginAutoUpgradeCategory.MODEL
+            strategy for strategy in strategies if strategy.category == TenantPluginAutoUpgradeCategory.MODEL
         )
         assert model_strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST
         assert model_strategy.upgrade_time_of_day == expected_time
 
-    def test_default_upgrade_time_is_aligned_to_fifteen_minutes(self):
-        from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-        default_time = PluginAutoUpgradeService.default_upgrade_time_of_day("t1")
-
+    def test_default_upgrade_time_is_aligned_to_fifteen_minutes(self) -> None:
+        default_time = PluginAutoUpgradeService.default_upgrade_time_of_day(str(uuid4()))
         assert default_time % (15 * 60) == 0
         assert 0 <= default_time < 24 * 60 * 60
 
-    def test_creates_missing_categories_and_splits_known_plugins(self, caplog: pytest.LogCaptureFixture):
-        session = _patched_session()
-        tool_strategy = SimpleNamespace(
-            category=TenantPluginAutoUpgradeCategory.TOOL,
-            strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY,
-            upgrade_time_of_day=0,
-            upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE,
-            exclude_plugins=["tool-plugin", "model-plugin", "unknown-plugin"],
-            include_plugins=["model-plugin", "tool-plugin"],
+    @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True)
+    def test_creates_missing_categories_and_splits_known_plugins(
+        self, sqlite_session: Session, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        tenant_id = str(uuid4())
+        tool_strategy = _strategy(
+            tenant_id,
+            exclude=["tool-plugin", "model-plugin", "unknown-plugin"],
+            include=["model-plugin", "tool-plugin"],
         )
-        model_strategy = SimpleNamespace(
+        model_strategy = _strategy(
+            tenant_id,
             category=TenantPluginAutoUpgradeCategory.MODEL,
-            strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY,
-            upgrade_time_of_day=0,
-            upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE,
-            exclude_plugins=["tool-plugin", "model-plugin", "unknown-plugin"],
-            include_plugins=["model-plugin", "tool-plugin"],
+            exclude=["tool-plugin", "model-plugin", "unknown-plugin"],
+            include=["model-plugin", "tool-plugin"],
         )
-        session.scalars.return_value.all.return_value = [tool_strategy, model_strategy]
-
+        sqlite_session.add_all([tool_strategy, model_strategy])
+        sqlite_session.commit()
         installed_plugins = [
             SimpleNamespace(
                 plugin_id="tool-plugin",
@@ -261,18 +235,17 @@ class TestBackfillStrategyCategories:
             patch(f"{MODULE}.PluginInstaller", return_value=installer),
             caplog.at_level(logging.WARNING, logger=MODULE),
         ):
-            from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
-
-            result = PluginAutoUpgradeService.backfill_strategy_categories("t1", session=session)
+            result = PluginAutoUpgradeService.backfill_strategy_categories(tenant_id, session=sqlite_session)
 
+        strategies = list(sqlite_session.scalars(select(TenantPluginAutoUpgradeStrategy)).all())
         assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 2
         assert result.normalized is True
-        assert session.add.call_count == len(TenantPluginAutoUpgradeCategory) - 2
+        assert len(strategies) == len(TenantPluginAutoUpgradeCategory)
         assert tool_strategy.exclude_plugins == ["tool-plugin"]
         assert tool_strategy.include_plugins == ["tool-plugin"]
         assert model_strategy.exclude_plugins == ["model-plugin"]
         assert model_strategy.include_plugins == ["model-plugin"]
         assert (
             "Skipped unknown plugin IDs while backfilling plugin auto-upgrade strategies: "
-            "tenant_id=t1, field=exclude_plugins, plugin_ids=['unknown-plugin']" in caplog.messages
+            f"tenant_id={tenant_id}, field=exclude_plugins, plugin_ids=['unknown-plugin']" in caplog.messages
         )

api/tests/unit_tests/services/plugin/test_plugin_service.py

@@ -7,28 +7,19 @@ import pytest
 import zstandard
 from pydantic import TypeAdapter
 from redis import RedisError
+from sqlalchemy.orm import Session
 
+from core.helper.model_provider_cache import ProviderCredentialsCacheType
 from core.plugin.entities.plugin import PluginCategory, PluginInstallationSource
 from core.plugin.entities.plugin_daemon import PluginInstallTask, PluginInstallTaskStatus, PluginModelProviderEntity
 from graphon.model_runtime.entities.common_entities import I18nObject
 from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity
+from models.provider import Provider, ProviderCredential, ProviderType, TenantPreferredModelProvider
 
 MODULE = "core.plugin.plugin_service"
-
-
-class _FakeSession:
-    def __init__(self) -> None:
-        self.execute = Mock()
-        self.scalars = Mock(return_value=SimpleNamespace(all=Mock(return_value=[])))
-
-    def __enter__(self) -> "_FakeSession":
-        return self
-
-    def __exit__(self, exc_type, exc, traceback) -> None:
-        return None
-
-    def begin(self) -> "_FakeSession":
-        return self
+TENANT_ID = "11111111-1111-1111-1111-111111111111"
+OTHER_TENANT_ID = "22222222-2222-2222-2222-222222222222"
+USER_ID = "33333333-3333-3333-3333-333333333333"
 
 
 def _build_provider_entity(provider: str = "openai") -> ProviderEntity:
@@ -246,6 +237,26 @@ class TestPluginModelProviderCache:
             call([cache_key]),
         ]
 
+    def test_fetch_plugin_model_providers_bypasses_redis_when_cache_disabled(self) -> None:
+        """With the cache disabled the daemon is the only source, and Redis is never touched."""
+        with patch(f"{MODULE}.redis_client") as redis_client, patch(f"{MODULE}.dify_config") as config:
+            config.PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED = False
+            client = Mock()
+            client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
+
+            from core.plugin.plugin_service import PluginService
+
+            first = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
+            second = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
+
+        assert [provider.provider for provider in first] == ["langgenius/openai/openai"]
+        assert [provider.provider for provider in second] == ["langgenius/openai/openai"]
+        assert client.fetch_model_providers.call_count == 2
+        redis_client.get.assert_not_called()
+        redis_client.mget.assert_not_called()
+        redis_client.setex.assert_not_called()
+        redis_client.lock.assert_not_called()
+
     def test_fetch_plugin_model_providers_refetches_when_cache_read_fails(self) -> None:
         """Redis read failures do not block provider discovery for the tenant."""
         with patch(f"{MODULE}.redis_client") as redis_client:
@@ -1146,19 +1157,72 @@ class TestPluginModelProviderCacheInvalidation:
         assert result is True
         invalidate_cache.assert_called_once_with("tenant-1")
 
-    def test_uninstall_existing_plugin_invalidates_cache_after_credential_cleanup(self) -> None:
+    @pytest.mark.parametrize(
+        "sqlite_session", [(Provider, ProviderCredential, TenantPreferredModelProvider)], indirect=True
+    )
+    def test_uninstall_existing_plugin_invalidates_cache_after_credential_cleanup(
+        self, sqlite_session: Session
+    ) -> None:
         """Successful uninstall with plugin metadata also invalidates the mutated tenant provider cache."""
+        plugin_id = "langgenius/openai"
+        provider_name = f"{plugin_id}/openai"
         plugin = SimpleNamespace(
             installation_id="installation-1",
-            plugin_id="langgenius/openai",
+            plugin_id=plugin_id,
             plugin_unique_identifier="langgenius/openai:1.0.0",
         )
-        session = _FakeSession()
+        credential = ProviderCredential(
+            tenant_id=TENANT_ID,
+            provider_name=provider_name,
+            credential_name="Target credential",
+            encrypted_config="{}",
+            user_id=USER_ID,
+        )
+        other_credential = ProviderCredential(
+            tenant_id=OTHER_TENANT_ID,
+            provider_name=provider_name,
+            credential_name="Other credential",
+            encrypted_config="{}",
+            user_id=USER_ID,
+        )
+        sqlite_session.add_all([credential, other_credential])
+        sqlite_session.flush()
+        provider = Provider(
+            tenant_id=TENANT_ID,
+            provider_name=provider_name,
+            provider_type=ProviderType.CUSTOM,
+            credential_id=credential.id,
+        )
+        other_provider = Provider(
+            tenant_id=OTHER_TENANT_ID,
+            provider_name=provider_name,
+            provider_type=ProviderType.CUSTOM,
+            credential_id=other_credential.id,
+        )
+        preferred_provider = TenantPreferredModelProvider(
+            tenant_id=TENANT_ID,
+            provider_name=provider_name,
+            preferred_provider_type=ProviderType.CUSTOM,
+        )
+        other_preferred_provider = TenantPreferredModelProvider(
+            tenant_id=OTHER_TENANT_ID,
+            provider_name=provider_name,
+            preferred_provider_type=ProviderType.CUSTOM,
+        )
+        sqlite_session.add_all([provider, other_provider, preferred_provider, other_preferred_provider])
+        sqlite_session.commit()
+        credential_id = credential.id
+        other_credential_id = other_credential.id
+        provider_id = provider.id
+        other_provider_id = other_provider.id
+        preferred_provider_id = preferred_provider.id
+        other_preferred_provider_id = other_preferred_provider.id
+
         with (
-            patch(f"{MODULE}.db", SimpleNamespace(engine=object())),
+            patch(f"{MODULE}.db", SimpleNamespace(engine=sqlite_session.get_bind())),
             patch(f"{MODULE}.dify_config") as mock_config,
             patch(f"{MODULE}.PluginInstaller") as installer_cls,
-            patch(f"{MODULE}.Session", return_value=session),
+            patch(f"{MODULE}.ProviderCredentialsCache") as credentials_cache,
             patch(f"{MODULE}.PluginService.invalidate_plugin_model_providers_cache") as invalidate_cache,
         ):
             mock_config.ENTERPRISE_ENABLED = False
@@ -1168,8 +1232,26 @@ class TestPluginModelProviderCacheInvalidation:
 
             from core.plugin.plugin_service import PluginService
 
-            result = PluginService.uninstall("tenant-1", "installation-1")
+            result = PluginService.uninstall(TENANT_ID, "installation-1")
 
         assert result is True
-        installer.uninstall.assert_called_once_with("tenant-1", "installation-1")
-        invalidate_cache.assert_called_once_with("tenant-1")
+        installer.uninstall.assert_called_once_with(TENANT_ID, "installation-1")
+        invalidate_cache.assert_called_once_with(TENANT_ID)
+        credentials_cache.assert_called_once_with(
+            tenant_id=TENANT_ID,
+            identity_id=provider_id,
+            cache_type=ProviderCredentialsCacheType.PROVIDER,
+        )
+        credentials_cache.return_value.delete.assert_called_once_with()
+
+        sqlite_session.expunge_all()
+        assert sqlite_session.get(ProviderCredential, credential_id) is None
+        persisted_provider = sqlite_session.get(Provider, provider_id)
+        assert persisted_provider is not None
+        assert persisted_provider.credential_id is None
+        assert sqlite_session.get(TenantPreferredModelProvider, preferred_provider_id) is None
+        assert sqlite_session.get(ProviderCredential, other_credential_id) is not None
+        persisted_other_provider = sqlite_session.get(Provider, other_provider_id)
+        assert persisted_other_provider is not None
+        assert persisted_other_provider.credential_id == other_credential_id
+        assert sqlite_session.get(TenantPreferredModelProvider, other_preferred_provider_id) is not None

api/tests/unit_tests/services/plugin/test_plugin_service_installation.py

@@ -8,6 +8,7 @@ verification, marketplace upgrade flows, and uninstall with credential cleanup.
 from __future__ import annotations
 
 from collections.abc import Iterator
+from typing import cast
 from unittest.mock import MagicMock, patch
 from uuid import uuid4
 
@@ -19,7 +20,6 @@ from sqlalchemy.orm import Session
 from core.plugin.entities.plugin import PluginInstallationSource
 from core.plugin.entities.plugin_daemon import PluginVerification
 from core.plugin.plugin_service import PluginService
-from enums.deployment_edition import DeploymentEdition
 from models import ProviderType
 from models.engine import db
 from models.provider import Provider, ProviderCredential, TenantPreferredModelProvider
@@ -27,20 +27,16 @@ from services.errors.plugin import PluginInstallationForbiddenError
 from services.feature_service import (
     PluginInstallationPermissionModel,
     PluginInstallationScope,
-    SystemFeatureModel,
 )
 
 
-def _make_features(
+def _make_permission(
     restrict_to_marketplace: bool = False,
     scope: PluginInstallationScope = PluginInstallationScope.ALL,
-) -> SystemFeatureModel:
-    return SystemFeatureModel(
-        deployment_edition=DeploymentEdition.COMMUNITY,
-        plugin_installation_permission=PluginInstallationPermissionModel(
-            restrict_to_marketplace_only=restrict_to_marketplace,
-            plugin_installation_scope=scope,
-        ),
+) -> PluginInstallationPermissionModel:
+    return PluginInstallationPermissionModel(
+        restrict_to_marketplace_only=restrict_to_marketplace,
+        plugin_installation_scope=scope,
     )
 
 
@@ -119,22 +115,31 @@ class TestFetchLatestPluginVersion:
 class TestCheckMarketplaceOnlyPermission:
     @patch("core.plugin.plugin_service.FeatureService")
     def test_raises_when_restricted(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(restrict_to_marketplace=True)
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(restrict_to_marketplace=True)
 
         with pytest.raises(PluginInstallationForbiddenError):
             PluginService._check_marketplace_only_permission()
 
     @patch("core.plugin.plugin_service.FeatureService")
     def test_passes_when_not_restricted(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(restrict_to_marketplace=False)
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(restrict_to_marketplace=False)
 
         PluginService._check_marketplace_only_permission()  # should not raise
 
+    @patch("core.plugin.plugin_service.FeatureService")
+    def test_raises_when_scope_denies_all(self, mock_fs):
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(scope=PluginInstallationScope.NONE)
+
+        with pytest.raises(PluginInstallationForbiddenError, match="not allowed"):
+            PluginService._check_marketplace_only_permission()
+
 
 class TestCheckPluginInstallationScope:
     @patch("core.plugin.plugin_service.FeatureService")
     def test_official_only_allows_langgenius(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(scope=PluginInstallationScope.OFFICIAL_ONLY)
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(
+            scope=PluginInstallationScope.OFFICIAL_ONLY
+        )
         verification = MagicMock()
         verification.authorized_category = PluginVerification.AuthorizedCategory.Langgenius
 
@@ -142,14 +147,16 @@ class TestCheckPluginInstallationScope:
 
     @patch("core.plugin.plugin_service.FeatureService")
     def test_official_only_rejects_third_party(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(scope=PluginInstallationScope.OFFICIAL_ONLY)
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(
+            scope=PluginInstallationScope.OFFICIAL_ONLY
+        )
 
         with pytest.raises(PluginInstallationForbiddenError):
             PluginService._check_plugin_installation_scope(None)
 
     @patch("core.plugin.plugin_service.FeatureService")
     def test_official_and_partners_allows_partner(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(
             scope=PluginInstallationScope.OFFICIAL_AND_SPECIFIC_PARTNERS
         )
         verification = MagicMock()
@@ -159,7 +166,7 @@ class TestCheckPluginInstallationScope:
 
     @patch("core.plugin.plugin_service.FeatureService")
     def test_official_and_partners_rejects_none(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(
             scope=PluginInstallationScope.OFFICIAL_AND_SPECIFIC_PARTNERS
         )
 
@@ -168,7 +175,7 @@ class TestCheckPluginInstallationScope:
 
     @patch("core.plugin.plugin_service.FeatureService")
     def test_none_scope_always_raises(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(scope=PluginInstallationScope.NONE)
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(scope=PluginInstallationScope.NONE)
         verification = MagicMock()
         verification.authorized_category = PluginVerification.AuthorizedCategory.Langgenius
 
@@ -177,10 +184,19 @@ class TestCheckPluginInstallationScope:
 
     @patch("core.plugin.plugin_service.FeatureService")
     def test_all_scope_passes_any(self, mock_fs):
-        mock_fs.get_system_features.return_value = _make_features(scope=PluginInstallationScope.ALL)
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission(scope=PluginInstallationScope.ALL)
 
         PluginService._check_plugin_installation_scope(None)  # should not raise
 
+    @patch("core.plugin.plugin_service.FeatureService")
+    def test_unknown_scope_always_raises(self, mock_fs):
+        permission = _make_permission()
+        permission.plugin_installation_scope = cast(PluginInstallationScope, "unknown-scope")
+        mock_fs.get_plugin_installation_permission.return_value = permission
+
+        with pytest.raises(PluginInstallationForbiddenError, match="policy is invalid"):
+            PluginService._check_plugin_installation_scope(None)
+
 
 class TestGetPluginIconUrl:
     @patch("core.plugin.plugin_service.dify_config")
@@ -248,7 +264,7 @@ class TestUpgradePluginWithMarketplace:
     @patch("core.plugin.plugin_service.dify_config")
     def test_skips_download_when_already_installed(self, mock_config, mock_installer_cls, mock_fs, mock_marketplace):
         mock_config.MARKETPLACE_ENABLED = True
-        mock_fs.get_system_features.return_value = _make_features()
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission()
         installer = mock_installer_cls.return_value
         installer.fetch_plugin_manifest.return_value = MagicMock()
         installer.upgrade_plugin.return_value = MagicMock()
@@ -264,7 +280,7 @@ class TestUpgradePluginWithMarketplace:
     @patch("core.plugin.plugin_service.dify_config")
     def test_downloads_when_not_installed(self, mock_config, mock_installer_cls, mock_fs, mock_download):
         mock_config.MARKETPLACE_ENABLED = True
-        mock_fs.get_system_features.return_value = _make_features()
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission()
         installer = mock_installer_cls.return_value
         installer.fetch_plugin_manifest.side_effect = RuntimeError("not found")
         mock_download.return_value = b"pkg-bytes"
@@ -283,7 +299,7 @@ class TestUpgradePluginWithGithub:
     @patch("core.plugin.plugin_service.FeatureService")
     @patch("core.plugin.plugin_service.PluginInstaller")
     def test_checks_marketplace_permission_and_delegates(self, mock_installer_cls: MagicMock, mock_fs: MagicMock):
-        mock_fs.get_system_features.return_value = _make_features()
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission()
         installer = mock_installer_cls.return_value
         installer.upgrade_plugin.return_value = MagicMock()
 
@@ -298,7 +314,7 @@ class TestUploadPkg:
     @patch("core.plugin.plugin_service.FeatureService")
     @patch("core.plugin.plugin_service.PluginInstaller")
     def test_runs_permission_and_scope_checks(self, mock_installer_cls: MagicMock, mock_fs: MagicMock):
-        mock_fs.get_system_features.return_value = _make_features()
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission()
         upload_resp = MagicMock()
         upload_resp.verification = None
         mock_installer_cls.return_value.upload_pkg.return_value = upload_resp
@@ -322,7 +338,7 @@ class TestInstallFromMarketplacePkg:
     @patch("core.plugin.plugin_service.dify_config")
     def test_downloads_when_not_cached(self, mock_config, mock_installer_cls, mock_fs, mock_download):
         mock_config.MARKETPLACE_ENABLED = True
-        mock_fs.get_system_features.return_value = _make_features()
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission()
         installer = mock_installer_cls.return_value
         installer.fetch_plugin_manifest.side_effect = RuntimeError("not found")
         mock_download.return_value = b"pkg"
@@ -344,7 +360,7 @@ class TestInstallFromMarketplacePkg:
     @patch("core.plugin.plugin_service.dify_config")
     def test_uses_cached_when_already_downloaded(self, mock_config, mock_installer_cls: MagicMock, mock_fs: MagicMock):
         mock_config.MARKETPLACE_ENABLED = True
-        mock_fs.get_system_features.return_value = _make_features()
+        mock_fs.get_plugin_installation_permission.return_value = _make_permission()
         installer = mock_installer_cls.return_value
         installer.fetch_plugin_manifest.return_value = MagicMock()
         decode_resp = MagicMock()

api/tests/unit_tests/services/test_account_service.py

@@ -1,5 +1,4 @@
 import json
-from collections.abc import Iterator
 from datetime import datetime, timedelta
 from unittest.mock import MagicMock, patch
 from uuid import UUID
@@ -11,7 +10,6 @@ from sqlalchemy.orm import Session
 from configs import dify_config
 from models.account import (
     Account,
-    AccountIntegrate,
     AccountStatus,
     Tenant,
     TenantAccountJoin,
@@ -114,22 +112,6 @@ class TestAccountService:
     - Error conditions and edge cases
     """
 
-    @pytest.fixture
-    def sqlite_session(self, sqlite_engine) -> Iterator[Session]:
-        """SQLite session with the account/workspace tables these service tests touch."""
-        tables = [
-            model.metadata.tables[model.__tablename__]
-            for model in (
-                Account,
-                Tenant,
-                TenantAccountJoin,
-                TenantPluginAutoUpgradeStrategy,
-            )
-        ]
-        Account.metadata.create_all(sqlite_engine, tables=tables)
-        with Session(sqlite_engine, expire_on_commit=False) as session:
-            yield session
-
     @pytest.fixture
     def mock_password_dependencies(self):
         """Mock setup for password-related functions."""
@@ -1264,24 +1246,6 @@ class TestRegisterService:
     - Error conditions and edge cases
     """
 
-    @pytest.fixture
-    def sqlite_session(self, sqlite_engine) -> Iterator[Session]:
-        """SQLite session with the account/workspace tables registration flows touch."""
-        tables = [
-            model.metadata.tables[model.__tablename__]
-            for model in (
-                Account,
-                AccountIntegrate,
-                Tenant,
-                TenantAccountJoin,
-                TenantPluginAutoUpgradeStrategy,
-                DifySetup,
-            )
-        ]
-        Account.metadata.create_all(sqlite_engine, tables=tables)
-        with Session(sqlite_engine, expire_on_commit=False) as session:
-            yield session
-
     @pytest.fixture
     def mock_redis_dependencies(self):
         """Mock setup for Redis-related functions."""

api/tests/unit_tests/services/test_app_service.py

@@ -1,19 +1,23 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from datetime import datetime
 from types import SimpleNamespace
 from typing import cast
 from unittest.mock import MagicMock, patch
+from uuid import uuid4
 
 import pytest
+from sqlalchemy import event
 from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import Session
 
 from graphon.model_runtime.entities.model_entities import ModelType
 from models import Account
-from models.model import App, AppMode, AppModelConfig
+from models.model import App, AppMode, AppModelConfig, IconType
 from models.workflow import Workflow
 from services.agent.errors import AgentNameConflictError
-from services.app_service import AppService, CreateAppParams
+from services.app_service import AppListParams, AppService, CreateAppParams
 
 
 class TestCreateAppTransactionBoundary:
@@ -236,6 +240,92 @@ class TestOpenapiVisibilityHelpers:
         mock_session.execute.assert_called_once()
 
 
+@pytest.mark.parametrize("sqlite_session", [(Account, App, AppModelConfig)], indirect=True)
+def test_get_recent_apps_uses_one_tenant_scoped_projection_query(sqlite_session: Session) -> None:
+    tenant_id = str(uuid4())
+    other_tenant_id = str(uuid4())
+    account = Account(name="Recent Apps Author", email="recent-apps@example.com")
+    sqlite_session.add(account)
+    sqlite_session.flush()
+
+    def create_app(*, name: str, tenant_id: str, updated_at: datetime, mode: AppMode = AppMode.CHAT) -> App:
+        app = App()
+        app.id = str(uuid4())
+        app.tenant_id = tenant_id
+        app.name = name
+        app.description = ""
+        app.mode = mode
+        app.icon_type = IconType.EMOJI
+        app.icon = "🚀"
+        app.icon_background = "#FFFFFF"
+        app.enable_site = False
+        app.enable_api = False
+        app.created_by = account.id
+        app.maintainer = account.id
+        app.created_at = updated_at
+        app.updated_at = updated_at
+        app.use_icon_as_answer_icon = False
+        return app
+
+    newest = create_app(name="Newest", tenant_id=tenant_id, updated_at=datetime(2026, 7, 3))
+    legacy_agent = AppModelConfig(app_id=newest.id)
+    legacy_agent.agent_mode = '{"enabled": true, "strategy": "react"}'
+    newest.app_model_config_id = legacy_agent.id
+    second = create_app(
+        name="Second",
+        tenant_id=tenant_id,
+        updated_at=datetime(2026, 7, 2),
+        mode=AppMode.WORKFLOW,
+    )
+    second.icon_type = None
+    second.icon = None
+    second.icon_background = None
+    second.created_by = None
+    second.maintainer = None
+    channel = create_app(
+        name="Channel",
+        tenant_id=tenant_id,
+        updated_at=datetime(2026, 7, 5),
+        mode=AppMode.CHANNEL,
+    )
+    rag_pipeline = create_app(
+        name="RAG Pipeline",
+        tenant_id=tenant_id,
+        updated_at=datetime(2026, 7, 4),
+        mode=AppMode.RAG_PIPELINE,
+    )
+    oldest = create_app(name="Oldest", tenant_id=tenant_id, updated_at=datetime(2026, 7, 1))
+    foreign = create_app(name="Foreign", tenant_id=other_tenant_id, updated_at=datetime(2026, 7, 4))
+    sqlite_session.add_all([newest, legacy_agent, second, channel, rag_pipeline, oldest, foreign])
+    sqlite_session.commit()
+
+    statements: list[str] = []
+    bind = sqlite_session.get_bind()
+
+    def record_sql(_conn, _cursor, statement, _parameters, _context, _executemany) -> None:
+        statements.append(statement)
+
+    event.listen(bind, "before_cursor_execute", record_sql)
+    try:
+        recent_apps = AppService().get_recent_apps(
+            account.id,
+            tenant_id,
+            AppListParams(limit=2),
+            sqlite_session,
+        )
+    finally:
+        event.remove(bind, "before_cursor_execute", record_sql)
+
+    assert [(app.name, app.mode, app.icon_type, app.author_name, app.maintainer) for app in recent_apps] == [
+        ("Newest", AppMode.CHAT, IconType.EMOJI, "Recent Apps Author", account.id),
+        ("Second", AppMode.WORKFLOW, None, None, None),
+    ]
+    select_statements = [statement for statement in statements if statement.lstrip().upper().startswith("SELECT")]
+    assert len(select_statements) == 1
+    assert "count(" not in select_statements[0].lower()
+    assert "app_model_configs" not in select_statements[0].lower()
+
+
 class TestAppMeta:
     def test_loads_workflow_with_caller_session(self):
         session = MagicMock()

api/tests/unit_tests/services/test_feature_service_plugin_installation_permission.py

@@ -0,0 +1,92 @@
+import logging
+
+import pytest
+
+from enums.deployment_edition import DeploymentEdition
+from services import feature_service as feature_service_module
+from services.feature_service import FeatureService, PluginInstallationScope, SystemFeatureModel
+
+
+def test_get_plugin_installation_permission_defaults_to_all_for_non_enterprise(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(feature_service_module.dify_config, "ENTERPRISE_ENABLED", False)
+
+    permission = FeatureService.get_plugin_installation_permission()
+
+    assert permission.plugin_installation_scope is PluginInstallationScope.ALL
+    assert permission.restrict_to_marketplace_only is False
+
+
+def test_get_plugin_installation_permission_parses_enterprise_policy(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(feature_service_module.dify_config, "ENTERPRISE_ENABLED", True)
+    monkeypatch.setattr(
+        feature_service_module.EnterpriseService,
+        "get_info",
+        staticmethod(
+            lambda: {
+                "PluginInstallationPermission": {
+                    "pluginInstallationScope": "official_only",
+                    "restrictToMarketplaceOnly": True,
+                }
+            }
+        ),
+    )
+
+    permission = FeatureService.get_plugin_installation_permission()
+
+    assert permission.plugin_installation_scope is PluginInstallationScope.OFFICIAL_ONLY
+    assert permission.restrict_to_marketplace_only is True
+
+
+@pytest.mark.parametrize(
+    "invalid_permission",
+    [
+        {
+            "pluginInstallationScope": "unknown-scope",
+            "restrictToMarketplaceOnly": False,
+        },
+        {
+            "pluginInstallationScope": "all",
+            "restrictToMarketplaceOnly": "false",
+        },
+    ],
+    ids=["unknown_scope", "non_boolean_marketplace_restriction"],
+)
+def test_invalid_enterprise_policy_denies_all_plugin_installations(
+    caplog: pytest.LogCaptureFixture,
+    invalid_permission: dict[str, object],
+) -> None:
+    with caplog.at_level(logging.ERROR, logger="services.feature_service"):
+        permission = FeatureService._resolve_plugin_installation_permission(
+            {"PluginInstallationPermission": invalid_permission}
+        )
+
+    assert permission.plugin_installation_scope is PluginInstallationScope.NONE
+    assert permission.restrict_to_marketplace_only is True
+    assert "denying all plugin installations" in caplog.text
+
+
+def test_system_features_exposes_only_validated_plugin_installation_policy(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        feature_service_module.EnterpriseService,
+        "get_info",
+        staticmethod(
+            lambda: {
+                "PluginInstallationPermission": {
+                    "pluginInstallationScope": "unknown-scope",
+                    "restrictToMarketplaceOnly": False,
+                }
+            }
+        ),
+    )
+    features = SystemFeatureModel(deployment_edition=DeploymentEdition.ENTERPRISE)
+
+    FeatureService._fulfill_params_from_enterprise(features)
+
+    assert features.plugin_installation_permission.plugin_installation_scope is PluginInstallationScope.NONE
+    assert features.plugin_installation_permission.restrict_to_marketplace_only is True

api/tests/unit_tests/services/test_file_service.py

@@ -1,6 +1,8 @@
 import base64
 import hashlib
 import os
+from collections.abc import Iterator
+from datetime import UTC, datetime
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -9,6 +11,8 @@ from sqlalchemy.orm import Session, sessionmaker
 from werkzeug.exceptions import NotFound
 
 from configs import dify_config
+from extensions.storage.storage_type import StorageType
+from models.base import TypeBase
 from models.enums import CreatorUserRole
 from models.model import Account, EndUser, UploadFile
 from services.errors.file import BlockedFileExtensionError, FileTooLargeError, UnsupportedFileTypeError
@@ -17,31 +21,54 @@ from services.file_service import FileService
 
 class TestFileService:
     @pytest.fixture
-    def mock_db_session(self):
-        session = MagicMock(spec=Session)
-        # Mock context manager behavior
-        session.__enter__.return_value = session
-        return session
+    def sqlite_session_maker(self, sqlite_engine: Engine) -> sessionmaker[Session]:
+        TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[UploadFile.__tablename__]])
+        return sessionmaker(bind=sqlite_engine, expire_on_commit=False)
 
     @pytest.fixture
-    def mock_session_maker(self, mock_db_session):
-        maker = MagicMock(spec=sessionmaker)
-        maker.return_value = mock_db_session
-        return maker
+    def db_session(self, sqlite_session_maker: sessionmaker[Session]) -> Iterator[Session]:
+        with sqlite_session_maker() as session:
+            yield session
 
     @pytest.fixture
-    def file_service(self, mock_session_maker):
-        return FileService(session_factory=mock_session_maker)
+    def file_service(self, sqlite_session_maker: sessionmaker[Session]) -> FileService:
+        return FileService(session_factory=sqlite_session_maker)
 
-    def test_init_with_engine(self):
-        engine = MagicMock(spec=Engine)
-        service = FileService(session_factory=engine)
+    @staticmethod
+    def _persist_upload_file(
+        session: Session,
+        *,
+        file_id: str = "file_id",
+        tenant_id: str = "tenant_id",
+        extension: str = "txt",
+        mime_type: str = "text/plain",
+        key: str = "key",
+    ) -> UploadFile:
+        upload_file = UploadFile(
+            tenant_id=tenant_id,
+            storage_type=StorageType.LOCAL,
+            key=key,
+            name=f"test.{extension}",
+            size=10,
+            extension=extension,
+            mime_type=mime_type,
+            created_by_role=CreatorUserRole.ACCOUNT,
+            created_by="user_id",
+            created_at=datetime(2024, 1, 1, tzinfo=UTC),
+            used=False,
+        )
+        upload_file.id = file_id
+        session.add(upload_file)
+        session.commit()
+        return upload_file
+
+    def test_init_with_engine(self, sqlite_engine: Engine):
+        service = FileService(session_factory=sqlite_engine)
         assert isinstance(service._session_maker, sessionmaker)
 
-    def test_init_with_sessionmaker(self):
-        maker = MagicMock(spec=sessionmaker)
-        service = FileService(session_factory=maker)
-        assert service._session_maker == maker
+    def test_init_with_sessionmaker(self, sqlite_session_maker: sessionmaker[Session]):
+        service = FileService(session_factory=sqlite_session_maker)
+        assert service._session_maker == sqlite_session_maker
 
     def test_init_invalid_factory(self):
         with pytest.raises(AssertionError, match="must be a sessionmaker or an Engine."):
@@ -52,11 +79,11 @@ class TestFileService:
     @patch("services.file_service.extract_tenant_id")
     @patch("services.file_service.file_helpers.get_signed_file_url")
     def test_upload_file_success(
-        self, mock_get_url, mock_tenant_id, mock_now, mock_storage, file_service: FileService, mock_db_session
+        self, mock_get_url, mock_tenant_id, mock_now, mock_storage, file_service: FileService, db_session: Session
     ):
         # Setup
         mock_tenant_id.return_value = "tenant_id"
-        mock_now.return_value = "2024-01-01"
+        mock_now.return_value = datetime(2024, 1, 1, tzinfo=UTC)
         mock_get_url.return_value = "http://signed-url"
 
         user = MagicMock(spec=Account)
@@ -81,8 +108,9 @@ class TestFileService:
         assert result.source_url == "http://signed-url"
 
         mock_storage.save.assert_called_once()
-        mock_db_session.add.assert_called_once_with(result)
-        mock_db_session.commit.assert_called_once()
+        persisted = db_session.get(UploadFile, result.id)
+        assert persisted is not None
+        assert persisted.hash == result.hash
 
     def test_upload_file_uses_explicit_resource_tenant(self, file_service: FileService):
         user = MagicMock(spec=Account)
@@ -109,7 +137,7 @@ class TestFileService:
         with pytest.raises(ValueError, match="Filename contains invalid characters"):
             file_service.upload_file(filename="invalid/file.txt", content=b"", mimetype="text/plain", user=MagicMock())
 
-    def test_upload_file_long_filename(self, file_service: FileService, mock_db_session):
+    def test_upload_file_long_filename(self, file_service: FileService, db_session: Session):
         # Setup
         long_name = "a" * 210 + ".txt"
         user = MagicMock(spec=Account)
@@ -124,6 +152,7 @@ class TestFileService:
             result = file_service.upload_file(filename=long_name, content=b"test", mimetype="text/plain", user=user)
             assert len(result.name) <= 205  # 200 + . + extension
             assert result.name.endswith(".txt")
+            assert db_session.get(UploadFile, result.id) is not None
 
     def test_upload_file_blocked_extension(self, file_service):
         with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", "exe"):
@@ -145,7 +174,7 @@ class TestFileService:
             with pytest.raises(FileTooLargeError):
                 file_service.upload_file(filename="test.jpg", content=content, mimetype="image/jpeg", user=MagicMock())
 
-    def test_upload_file_end_user(self, file_service: FileService, mock_db_session):
+    def test_upload_file_end_user(self, file_service: FileService, db_session: Session):
         user = MagicMock(spec=EndUser)
         user.id = "end_user_id"
 
@@ -157,6 +186,7 @@ class TestFileService:
             mock_tenant.return_value = "tenant"
             result = file_service.upload_file(filename="test.txt", content=b"test", mimetype="text/plain", user=user)
             assert result.created_by_role == CreatorUserRole.END_USER
+            assert db_session.get(UploadFile, result.id) is not None
 
     def test_is_file_size_within_limit(self):
         with (
@@ -181,12 +211,8 @@ class TestFileService:
             assert FileService.is_file_size_within_limit(extension="txt", file_size=5 * 1024 * 1024) is True
             assert FileService.is_file_size_within_limit(extension="pdf", file_size=6 * 1024 * 1024) is False
 
-    def test_get_file_base64_success(self, file_service: FileService, mock_db_session):
-        # Setup
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.key = "test_key"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_file_base64_success(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session, key="test_key")
 
         with patch("services.file_service.storage") as mock_storage:
             mock_storage.load_once.return_value = b"test content"
@@ -198,16 +224,17 @@ class TestFileService:
             assert result == base64.b64encode(b"test content").decode()
             mock_storage.load_once.assert_called_once_with("test_key")
 
-    def test_get_file_base64_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
+    def test_get_file_base64_not_found(self, file_service: FileService):
         with pytest.raises(NotFound, match="File not found"):
             file_service.get_file_base64("non_existent")
 
-    def test_get_file_presigned_url_success(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.key = "upload_files/tenant_id/icon.png"
-        upload_file.mime_type = "image/png"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_file_presigned_url_success(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(
+            db_session,
+            extension="png",
+            mime_type="image/png",
+            key="upload_files/tenant_id/icon.png",
+        )
 
         with (
             patch.object(dify_config, "FILES_ACCESS_TIMEOUT", 300),
@@ -224,13 +251,11 @@ class TestFileService:
             content_type="image/png",
         )
 
-    def test_get_file_presigned_url_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
-
+    def test_get_file_presigned_url_not_found(self, file_service: FileService):
         with pytest.raises(NotFound, match="File not found"):
             file_service.get_file_presigned_url(file_id="file_id", tenant_id="tenant_id")
 
-    def test_upload_text_success(self, file_service: FileService, mock_db_session):
+    def test_upload_text_success(self, file_service: FileService, db_session: Session):
         # Setup
         text = "sample text"
         text_name = "test.txt"
@@ -249,21 +274,17 @@ class TestFileService:
             assert result.used is True
             assert result.extension == "txt"
             mock_storage.save.assert_called_once()
-            mock_db_session.add.assert_called_once()
-            mock_db_session.commit.assert_called_once()
+            assert db_session.get(UploadFile, result.id) is not None
 
-    def test_upload_text_long_name(self, file_service: FileService, mock_db_session):
+    def test_upload_text_long_name(self, file_service: FileService, db_session: Session):
         long_name = "a" * 210
         with patch("services.file_service.storage"):
             result = file_service.upload_text("text", long_name, "user", "tenant")
             assert len(result.name) == 200
+            assert db_session.get(UploadFile, result.id) is not None
 
-    def test_get_file_preview_success(self, file_service: FileService, mock_db_session):
-        # Setup
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.extension = "pdf"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_file_preview_success(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session, extension="pdf", mime_type="application/pdf")
 
         with patch("services.file_service.ExtractProcessor.load_from_upload_file") as mock_extract:
             mock_extract.return_value = "Extracted text content"
@@ -274,27 +295,17 @@ class TestFileService:
             # Assert
             assert result == "Extracted text content"
 
-    def test_get_file_preview_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
+    def test_get_file_preview_not_found(self, file_service: FileService):
         with pytest.raises(NotFound, match="File not found"):
             file_service.get_file_preview("non_existent", "tenant_id")
 
-    def test_get_file_preview_unsupported_type(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.extension = "exe"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_file_preview_unsupported_type(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session, extension="exe", mime_type="application/octet-stream")
         with pytest.raises(UnsupportedFileTypeError):
             file_service.get_file_preview("file_id", "tenant_id")
 
-    def test_get_image_preview_success(self, file_service: FileService, mock_db_session):
-        # Setup
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.extension = "jpg"
-        upload_file.mime_type = "image/jpeg"
-        upload_file.key = "key"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_image_preview_success(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session, extension="jpg", mime_type="image/jpeg")
 
         with (
             patch("services.file_service.file_helpers.verify_image_signature") as mock_verify,
@@ -316,28 +327,21 @@ class TestFileService:
             with pytest.raises(NotFound, match="File not found or signature is invalid"):
                 file_service.get_image_preview("file_id", "ts", "nonce", "sign")
 
-    def test_get_image_preview_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
+    def test_get_image_preview_not_found(self, file_service: FileService):
         with patch("services.file_service.file_helpers.verify_image_signature") as mock_verify:
             mock_verify.return_value = True
             with pytest.raises(NotFound, match="File not found or signature is invalid"):
                 file_service.get_image_preview("file_id", "ts", "nonce", "sign")
 
-    def test_get_image_preview_unsupported_type(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.extension = "txt"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_image_preview_unsupported_type(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session)
         with patch("services.file_service.file_helpers.verify_image_signature") as mock_verify:
             mock_verify.return_value = True
             with pytest.raises(UnsupportedFileTypeError):
                 file_service.get_image_preview("file_id", "ts", "nonce", "sign")
 
-    def test_get_file_generator_by_file_id_success(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.key = "key"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_file_generator_by_file_id_success(self, file_service: FileService, db_session: Session):
+        upload_file = self._persist_upload_file(db_session)
 
         with (
             patch("services.file_service.file_helpers.verify_file_signature") as mock_verify,
@@ -348,7 +352,8 @@ class TestFileService:
 
             gen, file = file_service.get_file_generator_by_file_id("file_id", "ts", "nonce", "sign")
             assert list(gen) == [b"chunk"]
-            assert file == upload_file
+            assert file.id == upload_file.id
+            assert file.key == upload_file.key
 
     def test_get_file_generator_by_file_id_invalid_sig(self, file_service):
         with patch("services.file_service.file_helpers.verify_file_signature") as mock_verify:
@@ -356,20 +361,14 @@ class TestFileService:
             with pytest.raises(NotFound, match="File not found or signature is invalid"):
                 file_service.get_file_generator_by_file_id("file_id", "ts", "nonce", "sign")
 
-    def test_get_file_generator_by_file_id_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
+    def test_get_file_generator_by_file_id_not_found(self, file_service: FileService):
         with patch("services.file_service.file_helpers.verify_file_signature") as mock_verify:
             mock_verify.return_value = True
             with pytest.raises(NotFound, match="File not found or signature is invalid"):
                 file_service.get_file_generator_by_file_id("file_id", "ts", "nonce", "sign")
 
-    def test_get_public_image_preview_success(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.extension = "png"
-        upload_file.mime_type = "image/png"
-        upload_file.key = "key"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_public_image_preview_success(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session, extension="png", mime_type="image/png")
 
         with patch("services.file_service.storage") as mock_storage:
             mock_storage.load.return_value = b"image content"
@@ -377,66 +376,56 @@ class TestFileService:
             assert gen == b"image content"
             assert mime == "image/png"
 
-    def test_get_public_image_preview_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
+    def test_get_public_image_preview_not_found(self, file_service: FileService):
         with pytest.raises(NotFound, match="File not found or signature is invalid"):
             file_service.get_public_image_preview("file_id")
 
-    def test_get_public_image_preview_unsupported_type(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.extension = "txt"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_public_image_preview_unsupported_type(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session)
         with pytest.raises(UnsupportedFileTypeError):
             file_service.get_public_image_preview("file_id")
 
-    def test_get_file_content_success(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.key = "key"
-        mock_db_session.scalar.return_value = upload_file
+    def test_get_file_content_success(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session)
 
         with patch("services.file_service.storage") as mock_storage:
             mock_storage.load.return_value = b"hello world"
             result = file_service.get_file_content("file_id")
             assert result == "hello world"
 
-    def test_get_file_content_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
+    def test_get_file_content_not_found(self, file_service: FileService):
         with pytest.raises(NotFound, match="File not found"):
             file_service.get_file_content("file_id")
 
-    def test_delete_file_success(self, file_service: FileService, mock_db_session):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "file_id"
-        upload_file.key = "key"
-        # For session.scalar(select(...))
-        mock_db_session.scalar.return_value = upload_file
+    def test_delete_file_success(self, file_service: FileService, db_session: Session):
+        self._persist_upload_file(db_session)
 
         with patch("services.file_service.storage") as mock_storage:
             file_service.delete_file("file_id")
             mock_storage.delete.assert_called_once_with("key")
-            mock_db_session.delete.assert_called_once_with(upload_file)
+            db_session.expire_all()
+            assert db_session.get(UploadFile, "file_id") is None
 
-    def test_delete_file_not_found(self, file_service: FileService, mock_db_session):
-        mock_db_session.scalar.return_value = None
+    def test_delete_file_not_found(self, file_service: FileService):
         file_service.delete_file("file_id")
         # Should return without doing anything
 
-    def test_get_upload_files_by_ids_empty(self):
-        session = MagicMock()
-        result = FileService.get_upload_files_by_ids("tenant_id", [], session=session)
+    def test_get_upload_files_by_ids_empty(self, db_session: Session):
+        result = FileService.get_upload_files_by_ids("tenant_id", [], session=db_session)
         assert result == {}
 
-    def test_get_upload_files_by_ids(self):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.id = "550e8400-e29b-41d4-a716-446655440000"
-        upload_file.tenant_id = "tenant_id"
-        session = MagicMock()
-        session.scalars().all.return_value = [upload_file]
+    def test_get_upload_files_by_ids(self, db_session: Session):
+        upload_file = self._persist_upload_file(db_session, file_id="550e8400-e29b-41d4-a716-446655440000")
+        self._persist_upload_file(
+            db_session,
+            file_id="550e8400-e29b-41d4-a716-446655440001",
+            tenant_id="other-tenant",
+        )
 
         result = FileService.get_upload_files_by_ids(
-            "tenant_id", ["550e8400-e29b-41d4-a716-446655440000"], session=session
+            "tenant_id",
+            ["550e8400-e29b-41d4-a716-446655440000", "550e8400-e29b-41d4-a716-446655440001"],
+            session=db_session,
         )
         assert result["550e8400-e29b-41d4-a716-446655440000"] == upload_file
 
@@ -453,10 +442,8 @@ class TestFileService:
         used.add("a (1).txt")
         assert FileService._dedupe_zip_entry_name("a.txt", used) == "a (2).txt"
 
-    def test_build_upload_files_zip_tempfile(self):
-        upload_file = MagicMock(spec=UploadFile)
-        upload_file.name = "test.txt"
-        upload_file.key = "key"
+    def test_build_upload_files_zip_tempfile(self, db_session: Session):
+        upload_file = self._persist_upload_file(db_session)
 
         with (
             patch("services.file_service.storage") as mock_storage,

api/tests/unit_tests/services/test_human_input_service.py

@@ -1,6 +1,5 @@
 import dataclasses
 import logging
-from collections.abc import Iterator
 from datetime import datetime, timedelta
 from unittest.mock import MagicMock
 
@@ -42,14 +41,13 @@ from services.human_input_service import (
 
 
 @pytest.fixture
-def sqlite_session_factory(sqlite_engine: Engine) -> Iterator[tuple[sessionmaker[Session], Session]]:
-    factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
-    with factory() as session:
-        yield factory, session
+def unbound_session_factory() -> sessionmaker[Session]:
+    """Supply the required constructor dependency without enabling database access."""
+    return sessionmaker()
 
 
-def _persist_app(sqlite_session: Session, mode: AppMode) -> App:
-    app = App(
+def _make_app(mode: AppMode) -> App:
+    return App(
         id="app-id",
         tenant_id="tenant-id",
         name="Test App",
@@ -60,9 +58,6 @@ def _persist_app(sqlite_session: Session, mode: AppMode) -> App:
         enable_api=True,
         max_active_requests=0,
     )
-    sqlite_session.add(app)
-    sqlite_session.commit()
-    return app
 
 
 @pytest.fixture
@@ -97,14 +92,11 @@ def sample_form_record():
     )
 
 
-@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
 def test_enqueue_resume_dispatches_task_for_workflow(
     mocker: MockerFixture,
-    sqlite_session_factory,
-    sqlite_session: Session,
+    sqlite_session_factory: sessionmaker[Session],
 ):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+    service = HumanInputService(sqlite_session_factory)
 
     workflow_run = MagicMock()
     workflow_run.app_id = "app-id"
@@ -116,7 +108,8 @@ def test_enqueue_resume_dispatches_task_for_workflow(
         return_value=workflow_run_repo,
     )
 
-    _persist_app(sqlite_session, AppMode.WORKFLOW)
+    with sqlite_session_factory.begin() as arrange_session:
+        arrange_session.add(_make_app(AppMode.WORKFLOW))
 
     resume_task = mocker.patch("services.human_input_service.resume_app_execution")
 
@@ -128,10 +121,9 @@ def test_enqueue_resume_dispatches_task_for_workflow(
 
 
 def test_ensure_form_active_respects_global_timeout(
-    monkeypatch, sample_form_record: HumanInputFormRecord, sqlite_session_factory
+    monkeypatch, sample_form_record: HumanInputFormRecord, unbound_session_factory
 ):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+    service = HumanInputService(unbound_session_factory)
     expired_record = dataclasses.replace(
         sample_form_record,
         created_at=naive_utc_now() - timedelta(hours=2),
@@ -143,14 +135,11 @@ def test_ensure_form_active_respects_global_timeout(
         service.ensure_form_active(Form(expired_record))
 
 
-@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
 def test_enqueue_resume_dispatches_task_for_advanced_chat(
     mocker: MockerFixture,
-    sqlite_session_factory,
-    sqlite_session: Session,
+    sqlite_session_factory: sessionmaker[Session],
 ):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+    service = HumanInputService(sqlite_session_factory)
 
     workflow_run = MagicMock()
     workflow_run.app_id = "app-id"
@@ -162,7 +151,8 @@ def test_enqueue_resume_dispatches_task_for_advanced_chat(
         return_value=workflow_run_repo,
     )
 
-    _persist_app(sqlite_session, AppMode.ADVANCED_CHAT)
+    with sqlite_session_factory.begin() as arrange_session:
+        arrange_session.add(_make_app(AppMode.ADVANCED_CHAT))
 
     resume_task = mocker.patch("services.human_input_service.resume_app_execution")
 
@@ -173,14 +163,11 @@ def test_enqueue_resume_dispatches_task_for_advanced_chat(
     assert call_kwargs["kwargs"]["payload"]["workflow_run_id"] == "workflow-run-id"
 
 
-@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
 def test_enqueue_resume_skips_unsupported_app_mode(
     mocker: MockerFixture,
-    sqlite_session_factory,
-    sqlite_session: Session,
+    sqlite_session_factory: sessionmaker[Session],
 ):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+    service = HumanInputService(sqlite_session_factory)
 
     workflow_run = MagicMock()
     workflow_run.app_id = "app-id"
@@ -192,7 +179,8 @@ def test_enqueue_resume_skips_unsupported_app_mode(
         return_value=workflow_run_repo,
     )
 
-    _persist_app(sqlite_session, AppMode.COMPLETION)
+    with sqlite_session_factory.begin() as arrange_session:
+        arrange_session.add(_make_app(AppMode.COMPLETION))
 
     resume_task = mocker.patch("services.human_input_service.resume_app_execution")
 
@@ -202,14 +190,13 @@ def test_enqueue_resume_skips_unsupported_app_mode(
 
 
 def test_get_form_definition_by_token_for_console_uses_repository(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory
+    sample_form_record: HumanInputFormRecord, unbound_session_factory
 ):
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     console_record = dataclasses.replace(sample_form_record, recipient_type=RecipientType.CONSOLE)
     repo.get_by_token.return_value = console_record
 
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     form = service.get_form_definition_by_token_for_console("token")
 
     repo.get_by_token.assert_called_once_with("token")
@@ -245,9 +232,8 @@ def _build_resumption_context_state(*, options: list[str], workflow_run_id: str)
 
 
 def test_resolve_form_inputs_uses_runtime_select_options(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ):
-    session_factory, _ = sqlite_session_factory
     configured_input = SelectInputConfig(
         output_variable_name="decision",
         option_source=StringListSource(
@@ -272,7 +258,7 @@ def test_resolve_form_inputs_uses_runtime_select_options(
         "services.human_input_service.DifyAPIRepositoryFactory.create_api_workflow_run_repository",
         return_value=workflow_run_repo,
     )
-    service = HumanInputService(session_factory)
+    service = HumanInputService(unbound_session_factory)
 
     resolved_inputs = service.resolve_form_inputs(Form(record))
 
@@ -284,13 +270,12 @@ def test_resolve_form_inputs_uses_runtime_select_options(
 
 
 def test_submit_form_by_token_calls_repository_and_enqueue(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ):
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = sample_form_record
     repo.mark_submitted.return_value = sample_form_record
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
 
     service.submit_form_by_token(
@@ -313,11 +298,10 @@ def test_submit_form_by_token_calls_repository_and_enqueue(
 
 
 def test_submit_form_by_token_enqueues_agent_app_resume_for_conversation_form(
-    sample_form_record, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record, unbound_session_factory, mocker: MockerFixture
 ):
     # ENG-635: a conversation-owned (Agent v2 chat) form routes to the chat
     # resume, not the workflow resume.
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     conversation_record = dataclasses.replace(
         sample_form_record,
@@ -326,7 +310,7 @@ def test_submit_form_by_token_enqueues_agent_app_resume_for_conversation_form(
     )
     repo.get_by_token.return_value = conversation_record
     repo.mark_submitted.return_value = conversation_record
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     workflow_enqueue_spy = mocker.patch.object(service, "enqueue_resume")
     chat_enqueue_spy = mocker.patch.object(service, "enqueue_agent_app_resume")
 
@@ -343,9 +327,8 @@ def test_submit_form_by_token_enqueues_agent_app_resume_for_conversation_form(
 
 
 def test_submit_form_by_token_skips_enqueue_for_delivery_test(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ):
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     test_record = dataclasses.replace(
         sample_form_record,
@@ -354,7 +337,7 @@ def test_submit_form_by_token_skips_enqueue_for_delivery_test(
     )
     repo.get_by_token.return_value = test_record
     repo.mark_submitted.return_value = test_record
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
 
     service.submit_form_by_token(
@@ -368,13 +351,12 @@ def test_submit_form_by_token_skips_enqueue_for_delivery_test(
 
 
 def test_submit_form_by_token_passes_submission_user_id(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ):
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = sample_form_record
     repo.mark_submitted.return_value = sample_form_record
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
 
     service.submit_form_by_token(
@@ -391,11 +373,10 @@ def test_submit_form_by_token_passes_submission_user_id(
     enqueue_spy.assert_called_once_with(sample_form_record.workflow_run_id)
 
 
-def test_submit_form_by_token_invalid_action(sample_form_record: HumanInputFormRecord, sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
+def test_submit_form_by_token_invalid_action(sample_form_record: HumanInputFormRecord, unbound_session_factory):
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = dataclasses.replace(sample_form_record)
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError) as exc_info:
         service.submit_form_by_token(
@@ -409,8 +390,7 @@ def test_submit_form_by_token_invalid_action(sample_form_record: HumanInputFormR
     repo.mark_submitted.assert_not_called()
 
 
-def test_submit_form_by_token_missing_inputs(sample_form_record: HumanInputFormRecord, sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
+def test_submit_form_by_token_missing_inputs(sample_form_record: HumanInputFormRecord, unbound_session_factory):
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
 
     definition_with_input = FormDefinition(
@@ -422,7 +402,7 @@ def test_submit_form_by_token_missing_inputs(sample_form_record: HumanInputFormR
     )
     form_with_input = dataclasses.replace(sample_form_record, definition=definition_with_input)
     repo.get_by_token.return_value = form_with_input
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError) as exc_info:
         service.submit_form_by_token(
@@ -436,42 +416,6 @@ def test_submit_form_by_token_missing_inputs(sample_form_record: HumanInputFormR
     repo.mark_submitted.assert_not_called()
 
 
-def test_validate_human_input_submission_accepts_select_file_and_file_list(sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
-    definition = FormDefinition.model_validate(
-        {
-            "form_content": "Pick one and upload files",
-            "inputs": [
-                {
-                    "type": "select",
-                    "output_variable_name": "decision",
-                    "option_source": {
-                        "type": "constant",
-                        "value": ["approve", "reject"],
-                    },
-                },
-                {
-                    "type": "file",
-                    "output_variable_name": "attachment",
-                    "allowed_file_types": ["document"],
-                    "allowed_file_upload_methods": ["remote_url"],
-                },
-                {
-                    "type": "file-list",
-                    "output_variable_name": "attachments",
-                    "allowed_file_types": ["document"],
-                    "allowed_file_upload_methods": ["remote_url"],
-                    "number_limits": 3,
-                },
-            ],
-            "user_actions": [{"id": "submit", "title": "Submit"}],
-            "rendered_content": "<p>Pick one and upload files</p>",
-            "expiration_time": naive_utc_now() + timedelta(hours=1),
-        }
-    )
-
-
 @pytest.mark.parametrize(
     ("input_definition", "submitted_value", "expected_message"),
     [
@@ -522,12 +466,11 @@ def test_validate_human_input_submission_accepts_select_file_and_file_list(sqlit
 )
 def test_validate_human_input_submission_rejects_invalid_select_and_file_payloads(
     sample_form_record,
-    sqlite_session_factory,
+    unbound_session_factory,
     input_definition,
     submitted_value,
     expected_message,
 ):
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition.model_validate(
         {
@@ -539,7 +482,7 @@ def test_validate_human_input_submission_rejects_invalid_select_and_file_payload
         }
     )
     repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError) as exc_info:
         service.submit_form_by_token(
@@ -569,7 +512,7 @@ def test_form_properties(sample_form_record: HumanInputFormRecord):
 
 def test_form_submitted_error_init():
     error = FormSubmittedError(form_id="test-form")
-    assert "form_id=test-form" in error.description
+    assert error.description == "This form has already been submitted by another user, form_id=test-form"
     assert error.code == 412
 
 
@@ -580,61 +523,55 @@ def test_human_input_service_init_with_engine(sqlite_engine: Engine):
     assert service._session_factory.kw["bind"] is sqlite_engine
 
 
-def test_get_form_by_token_none(sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
+def test_get_form_by_token_none(unbound_session_factory):
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = None
 
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     assert service.get_form_by_token("invalid") is None
 
 
-def test_get_form_definition_by_token_mismatch(sample_form_record: HumanInputFormRecord, sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
+def test_get_form_definition_by_token_mismatch(sample_form_record: HumanInputFormRecord, unbound_session_factory):
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = sample_form_record
 
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     # RecipientType mismatch
     assert service.get_form_definition_by_token(RecipientType.CONSOLE, "token") is None
 
 
-def test_get_form_definition_by_token_success(sample_form_record: HumanInputFormRecord, sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
+def test_get_form_definition_by_token_success(sample_form_record: HumanInputFormRecord, unbound_session_factory):
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = sample_form_record
 
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     form = service.get_form_definition_by_token(RecipientType.STANDALONE_WEB_APP, "token")
     assert form is not None
     assert form.id == sample_form_record.form_id
 
 
 def test_get_form_definition_by_token_for_console_mismatch(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory
+    sample_form_record: HumanInputFormRecord, unbound_session_factory
 ):
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = sample_form_record  # is STANDALONE_WEB_APP
 
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     assert service.get_form_definition_by_token_for_console("token") is None
 
 
-def test_submit_form_by_token_delivery_not_enabled(sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
+def test_submit_form_by_token_delivery_not_enabled(unbound_session_factory):
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = None
 
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     with pytest.raises(human_input_service_module.WebAppDeliveryNotEnabledError):
         service.submit_form_by_token(RecipientType.STANDALONE_WEB_APP, "token", "action", {})
 
 
 def test_submit_form_by_token_no_workflow_run_id(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ):
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     repo.get_by_token.return_value = sample_form_record
 
@@ -642,16 +579,15 @@ def test_submit_form_by_token_no_workflow_run_id(
     result_record = dataclasses.replace(sample_form_record, workflow_run_id=None)
     repo.mark_submitted.return_value = result_record
 
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
 
     service.submit_form_by_token(RecipientType.STANDALONE_WEB_APP, "token", "submit", {})
     enqueue_spy.assert_not_called()
 
 
-def test_ensure_form_active_errors(sample_form_record: HumanInputFormRecord, sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+def test_ensure_form_active_errors(sample_form_record: HumanInputFormRecord, unbound_session_factory):
+    service = HumanInputService(unbound_session_factory)
 
     # Submitted
     submitted_record = dataclasses.replace(sample_form_record, submitted_at=naive_utc_now())
@@ -671,18 +607,16 @@ def test_ensure_form_active_errors(sample_form_record: HumanInputFormRecord, sql
         service.ensure_form_active(Form(expired_time_record))
 
 
-def test_ensure_not_submitted_raises(sample_form_record: HumanInputFormRecord, sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+def test_ensure_not_submitted_raises(sample_form_record: HumanInputFormRecord, unbound_session_factory):
+    service = HumanInputService(unbound_session_factory)
     submitted_record = dataclasses.replace(sample_form_record, submitted_at=naive_utc_now())
 
     with pytest.raises(human_input_service_module.FormSubmittedError):
         service._ensure_not_submitted(Form(submitted_record))
 
 
-def test_enqueue_resume_workflow_not_found(mocker: MockerFixture, sqlite_session_factory):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+def test_enqueue_resume_workflow_not_found(mocker: MockerFixture, unbound_session_factory):
+    service = HumanInputService(unbound_session_factory)
 
     workflow_run_repo = MagicMock()
     workflow_run_repo.get_workflow_run_by_id_without_tenant.return_value = None
@@ -696,15 +630,12 @@ def test_enqueue_resume_workflow_not_found(mocker: MockerFixture, sqlite_session
     assert "WorkflowRun not found" in str(excinfo.value)
 
 
-@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
 def test_enqueue_resume_app_not_found(
     mocker,
-    sqlite_session_factory,
-    sqlite_session: Session,
+    sqlite_session_factory: sessionmaker[Session],
     caplog: pytest.LogCaptureFixture,
 ):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+    service = HumanInputService(sqlite_session_factory)
 
     workflow_run = MagicMock()
     workflow_run.app_id = "app-id"
@@ -715,26 +646,31 @@ def test_enqueue_resume_app_not_found(
         "services.human_input_service.DifyAPIRepositoryFactory.create_api_workflow_run_repository",
         return_value=workflow_run_repo,
     )
+    resume_task = mocker.patch("services.human_input_service.resume_app_execution")
 
     with caplog.at_level(logging.ERROR, logger="services.human_input_service"):
         service.enqueue_resume("workflow-run-id")
-        assert any(r.levelno >= logging.ERROR for r in caplog.records)
+
+    assert (
+        "services.human_input_service",
+        logging.ERROR,
+        "App not found for WorkflowRun, workflow_run_id=workflow-run-id, app_id=app-id",
+    ) in caplog.record_tuples
+    resume_task.apply_async.assert_not_called()
 
 
 def test_is_globally_expired_zero_timeout(
-    monkeypatch: pytest.MonkeyPatch, sample_form_record: HumanInputFormRecord, sqlite_session_factory
+    monkeypatch: pytest.MonkeyPatch, sample_form_record: HumanInputFormRecord, unbound_session_factory
 ):
-    session_factory, _ = sqlite_session_factory
-    service = HumanInputService(session_factory)
+    service = HumanInputService(unbound_session_factory)
 
     monkeypatch.setattr(human_input_service_module.dify_config, "HUMAN_INPUT_GLOBAL_TIMEOUT_SECONDS", 0)
     assert service._is_globally_expired(Form(sample_form_record)) is False
 
 
 def test_submit_form_by_token_normalizes_select_and_files(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ) -> None:
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
@@ -753,7 +689,7 @@ def test_submit_form_by_token_normalizes_select_and_files(
     form_with_inputs = dataclasses.replace(sample_form_record, definition=definition)
     repo.get_by_token.return_value = form_with_inputs
     repo.mark_submitted.return_value = form_with_inputs
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     single_file = File(
         file_id="file-1",
@@ -815,9 +751,8 @@ def test_submit_form_by_token_normalizes_select_and_files(
 
 
 def test_submit_form_by_token_invalid_select_value(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory
+    sample_form_record: HumanInputFormRecord, unbound_session_factory
 ) -> None:
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
@@ -832,7 +767,7 @@ def test_submit_form_by_token_invalid_select_value(
         expiration_time=sample_form_record.expiration_time,
     )
     repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError, match="Invalid value for select input 'decision'"):
         service.submit_form_by_token(
@@ -844,9 +779,8 @@ def test_submit_form_by_token_invalid_select_value(
 
 
 def test_submit_form_by_token_invalid_file_list_item(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory
+    sample_form_record: HumanInputFormRecord, unbound_session_factory
 ) -> None:
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
@@ -856,7 +790,7 @@ def test_submit_form_by_token_invalid_file_list_item(
         expiration_time=sample_form_record.expiration_time,
     )
     repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(
         InvalidFormDataError,
@@ -871,9 +805,8 @@ def test_submit_form_by_token_invalid_file_list_item(
 
 
 def test_submit_form_by_token_rejects_cross_tenant_file(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ) -> None:
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
@@ -883,7 +816,7 @@ def test_submit_form_by_token_rejects_cross_tenant_file(
         expiration_time=sample_form_record.expiration_time,
     )
     repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     mocker.patch("services.human_input_service.build_from_mapping", side_effect=ValueError("Invalid upload file"))
 
     with pytest.raises(InvalidFormDataError, match="Invalid value for file input 'attachment'"):
@@ -904,9 +837,8 @@ def test_submit_form_by_token_rejects_cross_tenant_file(
 
 
 def test_submit_form_by_token_rejects_cross_tenant_file_list(
-    sample_form_record: HumanInputFormRecord, sqlite_session_factory, mocker: MockerFixture
+    sample_form_record: HumanInputFormRecord, unbound_session_factory, mocker: MockerFixture
 ) -> None:
-    session_factory, _ = sqlite_session_factory
     repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
@@ -916,7 +848,7 @@ def test_submit_form_by_token_rejects_cross_tenant_file_list(
         expiration_time=sample_form_record.expiration_time,
     )
     repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
-    service = HumanInputService(session_factory, form_repository=repo)
+    service = HumanInputService(unbound_session_factory, form_repository=repo)
     mocker.patch("services.human_input_service.build_from_mappings", side_effect=ValueError("Invalid upload file"))
 
     with pytest.raises(

api/tests/unit_tests/services/test_oauth_server_service.py

@@ -1,16 +1,20 @@
-"""Testcontainers integration tests for OAuthServerService."""
+"""Unit tests for OAuthServerService with SQLite-backed database access."""
 
 from __future__ import annotations
 
 import uuid
+from collections.abc import Iterator
 from typing import cast
 from unittest.mock import MagicMock, patch
 from uuid import uuid4
 
 import pytest
+from flask import Flask
+from sqlalchemy import Engine
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import BadRequest
 
+from models.engine import db
 from models.model import OAuthProviderApp
 from services.oauth_server import (
     OAUTH_ACCESS_TOKEN_EXPIRES_IN,
@@ -23,10 +27,24 @@ from services.oauth_server import (
 )
 
 
-class TestOAuthServerServiceGetProviderApp:
-    """DB-backed tests for get_oauth_provider_app."""
+@pytest.fixture
+def oauth_db() -> Iterator[Session]:
+    """Provide the production database extension with an isolated SQLite provider table."""
+    app = Flask(__name__)
+    app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
+    db.init_app(app)
 
-    def _create_oauth_provider_app(self, db_session_with_containers: Session, *, client_id: str) -> OAuthProviderApp:
+    with app.app_context():
+        OAuthProviderApp.__table__.create(db.engine)
+        with Session(db.engine, expire_on_commit=False) as session:
+            yield session
+
+
+class TestOAuthServerServiceGetProviderApp:
+    """Verify provider lookup against a real SQLAlchemy database."""
+
+    def test_get_oauth_provider_app_returns_app_when_exists(self, oauth_db: Session) -> None:
+        client_id = f"client-{uuid4()}"
         app = OAuthProviderApp(
             app_icon="icon.png",
             client_id=client_id,
@@ -35,35 +53,30 @@ class TestOAuthServerServiceGetProviderApp:
             redirect_uris=["https://example.com/callback"],
             scope="read",
         )
-        db_session_with_containers.add(app)
-        db_session_with_containers.commit()
-        return app
-
-    def test_get_oauth_provider_app_returns_app_when_exists(self, db_session_with_containers: Session):
-        client_id = f"client-{uuid4()}"
-        created = self._create_oauth_provider_app(db_session_with_containers, client_id=client_id)
+        oauth_db.add(app)
+        oauth_db.commit()
 
         result = OAuthServerService.get_oauth_provider_app(client_id)
 
         assert result is not None
         assert result.client_id == client_id
-        assert result.id == created.id
+        assert result.id == app.id
 
-    def test_get_oauth_provider_app_returns_none_when_not_exists(self, db_session_with_containers: Session):
+    def test_get_oauth_provider_app_returns_none_when_not_exists(self, oauth_db: Session) -> None:
         result = OAuthServerService.get_oauth_provider_app(f"nonexistent-{uuid4()}")
 
         assert result is None
 
 
 class TestOAuthServerServiceTokenOperations:
-    """Redis-backed tests for token sign/validate operations."""
+    """Verify Redis-backed token signing and validation branches."""
 
     @pytest.fixture
     def mock_redis(self):
         with patch("services.oauth_server.redis_client") as mock:
             yield mock
 
-    def test_sign_authorization_code_stores_and_returns_code(self, mock_redis):
+    def test_sign_authorization_code_stores_and_returns_code(self, mock_redis) -> None:
         deterministic_uuid = uuid.UUID("00000000-0000-0000-0000-000000000111")
         with patch("services.oauth_server.uuid.uuid4", return_value=deterministic_uuid):
             code = OAuthServerService.sign_oauth_authorization_code("client-1", "user-1")
@@ -75,7 +88,7 @@ class TestOAuthServerServiceTokenOperations:
             ex=600,
         )
 
-    def test_sign_access_token_raises_bad_request_for_invalid_code(self, mock_redis):
+    def test_sign_access_token_raises_bad_request_for_invalid_code(self, mock_redis) -> None:
         mock_redis.get.return_value = None
 
         with pytest.raises(BadRequest, match="invalid code"):
@@ -85,14 +98,13 @@ class TestOAuthServerServiceTokenOperations:
                 client_id="client-1",
             )
 
-    def test_sign_access_token_issues_tokens_for_valid_code(self, mock_redis):
+    def test_sign_access_token_issues_tokens_for_valid_code(self, mock_redis) -> None:
         token_uuids = [
             uuid.UUID("00000000-0000-0000-0000-000000000201"),
             uuid.UUID("00000000-0000-0000-0000-000000000202"),
         ]
         with patch("services.oauth_server.uuid.uuid4", side_effect=token_uuids):
             mock_redis.get.return_value = b"user-1"
-
             access_token, refresh_token = OAuthServerService.sign_oauth_access_token(
                 grant_type=OAuthGrantType.AUTHORIZATION_CODE,
                 code="code-1",
@@ -114,7 +126,7 @@ class TestOAuthServerServiceTokenOperations:
             ex=OAUTH_REFRESH_TOKEN_EXPIRES_IN,
         )
 
-    def test_sign_access_token_raises_bad_request_for_invalid_refresh_token(self, mock_redis):
+    def test_sign_access_token_raises_bad_request_for_invalid_refresh_token(self, mock_redis) -> None:
         mock_redis.get.return_value = None
 
         with pytest.raises(BadRequest, match="invalid refresh token"):
@@ -124,11 +136,10 @@ class TestOAuthServerServiceTokenOperations:
                 client_id="client-1",
             )
 
-    def test_sign_access_token_issues_new_token_for_valid_refresh(self, mock_redis):
+    def test_sign_access_token_issues_new_token_for_valid_refresh(self, mock_redis) -> None:
         deterministic_uuid = uuid.UUID("00000000-0000-0000-0000-000000000301")
         with patch("services.oauth_server.uuid.uuid4", return_value=deterministic_uuid):
             mock_redis.get.return_value = b"user-1"
-
             access_token, returned_refresh = OAuthServerService.sign_oauth_access_token(
                 grant_type=OAuthGrantType.REFRESH_TOKEN,
                 refresh_token="refresh-1",
@@ -138,14 +149,14 @@ class TestOAuthServerServiceTokenOperations:
         assert access_token == str(deterministic_uuid)
         assert returned_refresh == "refresh-1"
 
-    def test_sign_access_token_returns_none_for_unknown_grant_type(self, mock_redis):
+    def test_sign_access_token_returns_none_for_unknown_grant_type(self, mock_redis) -> None:
         grant_type = cast(OAuthGrantType, "invalid-grant-type")
 
         result = OAuthServerService.sign_oauth_access_token(grant_type=grant_type, client_id="client-1")
 
         assert result is None
 
-    def test_sign_refresh_token_stores_with_expected_expiry(self, mock_redis):
+    def test_sign_refresh_token_stores_with_expected_expiry(self, mock_redis) -> None:
         deterministic_uuid = uuid.UUID("00000000-0000-0000-0000-000000000401")
         with patch("services.oauth_server.uuid.uuid4", return_value=deterministic_uuid):
             refresh_token = OAuthServerService._sign_oauth_refresh_token("client-2", "user-2")
@@ -157,22 +168,21 @@ class TestOAuthServerServiceTokenOperations:
             ex=OAUTH_REFRESH_TOKEN_EXPIRES_IN,
         )
 
-    def test_validate_access_token_returns_none_when_not_found(self, mock_redis, db_session_with_containers: Session):
+    def test_validate_access_token_returns_none_when_not_found(self, mock_redis, sqlite_engine: Engine) -> None:
         mock_redis.get.return_value = None
-        session = MagicMock()
 
-        result = OAuthServerService.validate_oauth_access_token("client-1", "missing-token", db_session_with_containers)
+        with Session(sqlite_engine) as session:
+            result = OAuthServerService.validate_oauth_access_token("client-1", "missing-token", session)
 
         assert result is None
 
-    def test_validate_access_token_loads_user_when_exists(self, mock_redis, db_session_with_containers: Session):
+    def test_validate_access_token_loads_user_when_exists(self, mock_redis, sqlite_engine: Engine) -> None:
         mock_redis.get.return_value = b"user-88"
         expected_user = MagicMock()
 
-        with patch("services.oauth_server.AccountService.load_user", return_value=expected_user) as mock_load:
-            result = OAuthServerService.validate_oauth_access_token(
-                "client-1", "access-token", db_session_with_containers
-            )
+        with Session(sqlite_engine) as session:
+            with patch("services.oauth_server.AccountService.load_user", return_value=expected_user) as mock_load:
+                result = OAuthServerService.validate_oauth_access_token("client-1", "access-token", session)
+                mock_load.assert_called_once_with("user-88", session)
 
         assert result is expected_user
-        mock_load.assert_called_once_with("user-88", db_session_with_containers)