Compare commits
48
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0a7061a952 | ||
|
|
48717e5981 | ||
|
|
5805b94e23 | ||
|
|
c64c772b5b | ||
|
|
b4537455cf | ||
|
|
5fe010b653 | ||
|
|
66cd69b07e | ||
|
|
09eca00370 | ||
|
|
12cf69e5c8 | ||
|
|
4332023806 | ||
|
|
dd9f1e041e | ||
|
|
2271257c2f | ||
|
|
8adbda988b | ||
|
|
876ac24496 | ||
|
|
8b0892fd79 | ||
|
|
50d2ad8479 | ||
|
|
1ef8d3a2a9 | ||
|
|
b8adb0cac2 | ||
|
|
df88c45e3c | ||
|
|
d3df1b2847 | ||
|
|
709aac9337 | ||
|
|
865affdd9b | ||
|
|
71840c3f71 | ||
|
|
dc09aea7b3 | ||
|
|
99507a5ae5 | ||
|
|
55ae0aadf8 | ||
|
|
6c4ba71a23 | ||
|
|
e5f53053e3 | ||
|
|
3e392527e9 | ||
|
|
36d3a9c7fb | ||
|
|
914598d7e3 | ||
|
|
602c24cf8c | ||
|
|
e0b5ecc50f | ||
|
|
a83c7e89e0 | ||
|
|
7a4252b3de | ||
|
|
5c3dd25b32 | ||
|
|
30c8e9b08d | ||
|
|
3e3a16feaf | ||
|
|
0e1fad88c0 | ||
|
|
f442559543 | ||
|
|
2a63f26796 | ||
|
|
7e325b2ff2 | ||
|
|
26487d19b4 | ||
|
|
71a936e89a | ||
|
|
f2c19b456b | ||
|
|
2d3da2a49f | ||
|
|
11b393db9f | ||
|
|
9a01b16d6c |
@@ -22,7 +22,7 @@ from .plugin import (
|
||||
setup_system_trigger_oauth_client,
|
||||
transform_datasource_credentials,
|
||||
)
|
||||
from .rbac import migrate_member_roles_to_rbac
|
||||
from .rbac import migrate_dataset_permissions_to_rbac, migrate_member_roles_to_rbac
|
||||
from .retention import (
|
||||
archive_workflow_runs,
|
||||
archive_workflow_runs_plan,
|
||||
@@ -76,6 +76,7 @@ __all__ = [
|
||||
"legacy_model_types",
|
||||
"migrate_annotation_vector_database",
|
||||
"migrate_data_for_plugin",
|
||||
"migrate_dataset_permissions_to_rbac",
|
||||
"migrate_knowledge_vector_database",
|
||||
"migrate_member_roles_to_rbac",
|
||||
"migrate_oss",
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import cast
|
||||
|
||||
import click
|
||||
|
||||
from commands.rbac import migrate_dataset_permissions_to_rbac
|
||||
from extensions.ext_database import db
|
||||
from graphon.model_runtime.entities.model_entities import ModelType
|
||||
from services.legacy_model_type_migration import (
|
||||
@@ -177,3 +178,4 @@ def legacy_model_types(
|
||||
|
||||
|
||||
data_migrate.add_command(legacy_model_types)
|
||||
data_migrate.add_command(migrate_dataset_permissions_to_rbac)
|
||||
|
||||
+438
-66
@@ -1,11 +1,55 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
import click
|
||||
from sqlalchemy import select
|
||||
|
||||
from configs import dify_config
|
||||
from core.db.session_factory import session_factory
|
||||
from models import TenantAccountJoin, TenantAccountRole
|
||||
from services.enterprise.rbac_service import ListOption, RBACService
|
||||
from core.rbac import RBACResourceWhitelistScope
|
||||
from models import Dataset, DatasetPermission, DatasetPermissionEnum, TenantAccountJoin, TenantAccountRole
|
||||
from services.enterprise.rbac_service import ListOption, RBACService, ReplaceMemberBindings, ReplaceUserAccessPolicies
|
||||
|
||||
_RBAC_DEFAULT_ACCESS_POLICY_ID = "default"
|
||||
|
||||
_LEGACY_ROLE_TO_BUILTIN_TAG = {
|
||||
TenantAccountRole.OWNER.value: "owner",
|
||||
TenantAccountRole.ADMIN.value: "admin",
|
||||
TenantAccountRole.EDITOR.value: "editor",
|
||||
TenantAccountRole.NORMAL.value: "normal",
|
||||
TenantAccountRole.DATASET_OPERATOR.value: "dataset_operator",
|
||||
}
|
||||
|
||||
|
||||
def _resolve_builtin_role_ids(tenant_id: str, operator_account_id: str) -> dict[str, str]:
|
||||
"""Resolve every legacy workspace role to the current tenant's builtin RBAC role id.
|
||||
|
||||
The migration replays the old `TenantAccountJoin.role` values onto the
|
||||
RBAC member-role binding API. Builtin RBAC roles are tenant-scoped and
|
||||
identified by runtime ids, so the command must look them up per tenant.
|
||||
"""
|
||||
roles = RBACService.Roles.list(
|
||||
tenant_id=tenant_id,
|
||||
account_id=operator_account_id,
|
||||
options=ListOption(page_number=1, results_per_page=100),
|
||||
).data
|
||||
role_id_by_tag = {
|
||||
role.role_tag: role.id
|
||||
for role in roles
|
||||
if role.is_builtin and role.category == "global_system_default" and role.role_tag
|
||||
}
|
||||
resolved: dict[str, str] = {}
|
||||
for legacy_role, expected_builtin_tag in _LEGACY_ROLE_TO_BUILTIN_TAG.items():
|
||||
role_id = role_id_by_tag.get(expected_builtin_tag)
|
||||
if expected_builtin_tag == "dataset_operator" and not dify_config.DATASET_OPERATOR_ENABLED:
|
||||
continue
|
||||
if not role_id:
|
||||
raise ValueError(f"Builtin RBAC role not found for tenant={tenant_id}, legacy_role={legacy_role}")
|
||||
resolved[legacy_role] = role_id
|
||||
return resolved
|
||||
|
||||
|
||||
def _resolve_builtin_role_id(tenant_id: str, operator_account_id: str, legacy_role: str) -> str:
|
||||
@@ -15,26 +59,86 @@ def _resolve_builtin_role_id(tenant_id: str, operator_account_id: str, legacy_ro
|
||||
RBAC member-role binding API. Builtin RBAC roles are tenant-scoped and
|
||||
identified by runtime ids, so the command must look them up per tenant.
|
||||
"""
|
||||
expected_builtin_tag = {
|
||||
TenantAccountRole.OWNER.value: "owner",
|
||||
TenantAccountRole.ADMIN.value: "admin",
|
||||
TenantAccountRole.EDITOR.value: "editor",
|
||||
TenantAccountRole.NORMAL.value: "normal",
|
||||
TenantAccountRole.DATASET_OPERATOR.value: "dataset_operator",
|
||||
}.get(legacy_role)
|
||||
if not expected_builtin_tag:
|
||||
if legacy_role not in _LEGACY_ROLE_TO_BUILTIN_TAG:
|
||||
raise ValueError(f"Unsupported legacy workspace role: {legacy_role}")
|
||||
|
||||
roles = RBACService.Roles.list(
|
||||
return _resolve_builtin_role_ids(tenant_id, operator_account_id)[legacy_role]
|
||||
|
||||
|
||||
def _iter_tenant_member_batches(
|
||||
tenant_id: str | None,
|
||||
*,
|
||||
db_batch_size: int,
|
||||
api_batch_size: int,
|
||||
) -> Iterator[tuple[str, str, list[tuple[str, str]]]]:
|
||||
"""Yield legacy member roles in tenant-scoped API-sized batches.
|
||||
|
||||
Rows are projected to primitive values and streamed from the database, so
|
||||
the command never materializes every TenantAccountJoin ORM object. The
|
||||
iterator only keeps one tenant's API-sized batches in memory while it
|
||||
finds that tenant's owner account.
|
||||
"""
|
||||
with session_factory.create_session() as session:
|
||||
stmt = (
|
||||
select(TenantAccountJoin.tenant_id, TenantAccountJoin.account_id, TenantAccountJoin.role)
|
||||
.order_by(TenantAccountJoin.tenant_id.asc(), TenantAccountJoin.id.asc())
|
||||
.execution_options(yield_per=db_batch_size)
|
||||
)
|
||||
if tenant_id:
|
||||
stmt = stmt.where(TenantAccountJoin.tenant_id == tenant_id)
|
||||
|
||||
current_tenant_id: str | None = None
|
||||
owner_account_id: str | None = None
|
||||
batches: list[list[tuple[str, str]]] = []
|
||||
batch: list[tuple[str, str]] = []
|
||||
|
||||
def flush_current_tenant() -> Iterator[tuple[str, str, list[tuple[str, str]]]]:
|
||||
if current_tenant_id is None:
|
||||
return
|
||||
if batch:
|
||||
batches.append(batch.copy())
|
||||
if not owner_account_id:
|
||||
raise ValueError(f"Workspace owner not found for tenant={current_tenant_id}")
|
||||
for item in batches:
|
||||
yield current_tenant_id, owner_account_id, item
|
||||
|
||||
for row in session.execute(stmt):
|
||||
workspace_id = str(row.tenant_id)
|
||||
if current_tenant_id is not None and workspace_id != current_tenant_id:
|
||||
yield from flush_current_tenant()
|
||||
owner_account_id = None
|
||||
batches = []
|
||||
batch = []
|
||||
current_tenant_id = workspace_id
|
||||
account_id = str(row.account_id)
|
||||
role = str(row.role)
|
||||
if role == TenantAccountRole.OWNER.value:
|
||||
owner_account_id = account_id
|
||||
batch.append((account_id, role))
|
||||
if len(batch) >= api_batch_size:
|
||||
batches.append(batch)
|
||||
batch = []
|
||||
|
||||
yield from flush_current_tenant()
|
||||
|
||||
|
||||
def _member_already_has_role(current_roles_by_account_id: dict[str, set[str]], account_id: str, role_id: str) -> bool:
|
||||
return current_roles_by_account_id.get(account_id) == {role_id}
|
||||
|
||||
|
||||
def _replace_member_role(
|
||||
tenant_id: str,
|
||||
operator_account_id: str,
|
||||
member_account_id: str,
|
||||
role_id: str,
|
||||
) -> str:
|
||||
RBACService.MemberRoles.replace(
|
||||
tenant_id=tenant_id,
|
||||
account_id=operator_account_id,
|
||||
options=ListOption(page_number=1, results_per_page=100),
|
||||
).data
|
||||
for role in roles:
|
||||
if role.is_builtin and role.category == "global_system_default" and role.role_tag == expected_builtin_tag:
|
||||
return str(role.id)
|
||||
|
||||
raise ValueError(f"Builtin RBAC role not found for tenant={tenant_id}, legacy_role={legacy_role}")
|
||||
member_account_id=member_account_id,
|
||||
role_ids=[role_id],
|
||||
)
|
||||
return member_account_id
|
||||
|
||||
|
||||
@click.command(
|
||||
@@ -42,7 +146,16 @@ def _resolve_builtin_role_id(tenant_id: str, operator_account_id: str, legacy_ro
|
||||
)
|
||||
@click.option("--tenant-id", help="Only migrate a single workspace.")
|
||||
@click.option("--dry-run", is_flag=True, default=False, help="Preview the migration without writing RBAC bindings.")
|
||||
def migrate_member_roles_to_rbac(tenant_id: str | None, dry_run: bool) -> None:
|
||||
@click.option("--db-batch-size", default=5000, show_default=True, help="Rows fetched per database batch.")
|
||||
@click.option("--api-batch-size", default=200, show_default=True, help="Members checked per RBAC batch_get call.")
|
||||
@click.option("--workers", default=1, show_default=True, help="Concurrent member role replace calls per tenant batch.")
|
||||
def migrate_member_roles_to_rbac(
|
||||
tenant_id: str | None,
|
||||
dry_run: bool,
|
||||
db_batch_size: int,
|
||||
api_batch_size: int,
|
||||
workers: int,
|
||||
) -> None:
|
||||
"""Backfill RBAC member-role bindings from legacy `TenantAccountJoin.role` data.
|
||||
|
||||
This is an offline migration command for workspaces that already have
|
||||
@@ -50,63 +163,322 @@ def migrate_member_roles_to_rbac(tenant_id: str | None, dry_run: bool) -> None:
|
||||
member-role binding store.
|
||||
"""
|
||||
click.echo(click.style("Starting RBAC member-role migration.", fg="green"))
|
||||
if workers < 1:
|
||||
raise click.BadParameter("workers must be >= 1", param_hint="--workers")
|
||||
|
||||
with session_factory.create_session() as session:
|
||||
stmt = select(TenantAccountJoin).order_by(TenantAccountJoin.tenant_id.asc(), TenantAccountJoin.id.asc())
|
||||
if tenant_id:
|
||||
stmt = stmt.where(TenantAccountJoin.tenant_id == tenant_id)
|
||||
tenant_count = 0
|
||||
scanned_count = 0
|
||||
skipped_count = 0
|
||||
migrated_count = 0
|
||||
current_tenant_id: str | None = None
|
||||
role_ids_by_legacy_role: dict[str, str] = {}
|
||||
|
||||
joins = list(session.scalars(stmt).all())
|
||||
for workspace_id, owner_account_id, batch in _iter_tenant_member_batches(
|
||||
tenant_id,
|
||||
db_batch_size=db_batch_size,
|
||||
api_batch_size=api_batch_size,
|
||||
):
|
||||
scanned_count += len(batch)
|
||||
if workspace_id != current_tenant_id:
|
||||
tenant_count += 1
|
||||
current_tenant_id = workspace_id
|
||||
role_ids_by_legacy_role = _resolve_builtin_role_ids(workspace_id, owner_account_id)
|
||||
click.echo(f"tenant={workspace_id}")
|
||||
|
||||
if not joins:
|
||||
current_roles_by_account_id: dict[str, set[str]] = {}
|
||||
if not dry_run:
|
||||
current_roles = RBACService.MemberRoles.batch_get(
|
||||
tenant_id=workspace_id,
|
||||
account_id=owner_account_id,
|
||||
member_account_ids=[account_id for account_id, _ in batch],
|
||||
)
|
||||
current_roles_by_account_id = {
|
||||
item.account_id: {str(role.id) for role in item.roles} for item in current_roles
|
||||
}
|
||||
|
||||
replace_jobs: list[tuple[str, str]] = []
|
||||
for member_account_id, legacy_role in batch:
|
||||
resolved_role_id = role_ids_by_legacy_role.get(legacy_role)
|
||||
if not resolved_role_id:
|
||||
raise ValueError(f"Unsupported legacy workspace role: {legacy_role}")
|
||||
|
||||
if dry_run:
|
||||
click.echo(
|
||||
f"tenant={workspace_id} member={member_account_id} "
|
||||
f"legacy_role={legacy_role} -> rbac_role_id={resolved_role_id}"
|
||||
)
|
||||
continue
|
||||
|
||||
if _member_already_has_role(current_roles_by_account_id, member_account_id, resolved_role_id):
|
||||
skipped_count += 1
|
||||
continue
|
||||
|
||||
replace_jobs.append((member_account_id, resolved_role_id))
|
||||
|
||||
if replace_jobs:
|
||||
if workers == 1:
|
||||
for member_account_id, resolved_role_id in replace_jobs:
|
||||
_replace_member_role(workspace_id, owner_account_id, member_account_id, resolved_role_id)
|
||||
migrated_count += 1
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [
|
||||
executor.submit(
|
||||
_replace_member_role,
|
||||
workspace_id,
|
||||
owner_account_id,
|
||||
member_account_id,
|
||||
resolved_role_id,
|
||||
)
|
||||
for member_account_id, resolved_role_id in replace_jobs
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
future.result()
|
||||
migrated_count += 1
|
||||
|
||||
if scanned_count % 10000 == 0:
|
||||
click.echo(
|
||||
f"progress scanned={scanned_count} migrated={migrated_count} skipped={skipped_count}",
|
||||
err=True,
|
||||
)
|
||||
|
||||
if scanned_count == 0:
|
||||
click.echo(click.style("No workspace members found for migration.", fg="yellow"))
|
||||
return
|
||||
|
||||
owner_account_by_tenant: dict[str, str] = {}
|
||||
resolved_role_ids: dict[tuple[str, str], str] = {}
|
||||
migrated_count = 0
|
||||
|
||||
for join in joins:
|
||||
workspace_id = str(join.tenant_id)
|
||||
member_account_id = str(join.account_id)
|
||||
legacy_role = str(join.role)
|
||||
|
||||
if workspace_id not in owner_account_by_tenant:
|
||||
owner_join = next(
|
||||
(
|
||||
item
|
||||
for item in joins
|
||||
if str(item.tenant_id) == workspace_id and str(item.role) == TenantAccountRole.OWNER.value
|
||||
),
|
||||
None,
|
||||
)
|
||||
if not owner_join:
|
||||
raise ValueError(f"Workspace owner not found for tenant={workspace_id}")
|
||||
owner_account_by_tenant[workspace_id] = str(owner_join.account_id)
|
||||
|
||||
operator_account_id = owner_account_by_tenant[workspace_id]
|
||||
cache_key = (workspace_id, legacy_role)
|
||||
if cache_key not in resolved_role_ids:
|
||||
resolved_role_ids[cache_key] = _resolve_builtin_role_id(workspace_id, operator_account_id, legacy_role)
|
||||
|
||||
resolved_role_id = resolved_role_ids[cache_key]
|
||||
if dry_run:
|
||||
click.echo(
|
||||
f"tenant={workspace_id} member={member_account_id} "
|
||||
f"legacy_role={legacy_role} -> rbac_role_id={resolved_role_id}"
|
||||
click.style(
|
||||
f"Dry run completed. Scanned {scanned_count} members across {tenant_count} tenants. "
|
||||
"No RBAC bindings were written.",
|
||||
fg="yellow",
|
||||
)
|
||||
)
|
||||
else:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"RBAC member-role migration completed. Scanned {scanned_count} members across {tenant_count} tenants, "
|
||||
f"migrated {migrated_count}, skipped {skipped_count} already up-to-date.",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
|
||||
if dry_run:
|
||||
continue
|
||||
|
||||
RBACService.MemberRoles.replace(
|
||||
tenant_id=workspace_id,
|
||||
account_id=operator_account_id,
|
||||
member_account_id=member_account_id,
|
||||
role_ids=[resolved_role_id],
|
||||
)
|
||||
migrated_count += 1
|
||||
def _dataset_permission_enum(permission: DatasetPermissionEnum | str | None) -> DatasetPermissionEnum:
|
||||
if permission is None:
|
||||
return DatasetPermissionEnum.ONLY_ME
|
||||
try:
|
||||
return DatasetPermissionEnum(permission)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"Unsupported legacy dataset permission: {permission}") from exc
|
||||
|
||||
|
||||
def _rbac_dataset_scope_for_legacy_permission(permission: DatasetPermissionEnum) -> RBACResourceWhitelistScope:
|
||||
if permission is DatasetPermissionEnum.ALL_TEAM:
|
||||
return RBACResourceWhitelistScope.ALL
|
||||
if permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.PARTIAL_TEAM}:
|
||||
return RBACResourceWhitelistScope.SPECIFIC
|
||||
raise ValueError(f"Unsupported legacy dataset permission: {permission}")
|
||||
|
||||
|
||||
def _emit_dataset_permission_migration_event(payload: dict[str, object]) -> None:
|
||||
click.echo(json.dumps(payload, sort_keys=True))
|
||||
|
||||
|
||||
@click.command(
|
||||
"rbac-migrate-dataset-permissions",
|
||||
help=(
|
||||
"Migrate legacy dataset permission scopes and partial members into RBAC dataset access bindings. "
|
||||
"Side effect: replacing each dataset whitelist clears existing per-user policy bindings; "
|
||||
"the command then recreates legacy partial-member default bindings."
|
||||
),
|
||||
)
|
||||
@click.option("--tenant-id", help="Only migrate datasets in a single workspace.")
|
||||
@click.option("--dataset-id", help="Only migrate a single dataset.")
|
||||
@click.option("--batch-size", default=500, show_default=True, type=click.IntRange(min=1))
|
||||
@click.option(
|
||||
"--dry-run/--apply",
|
||||
default=True,
|
||||
show_default=True,
|
||||
help="Preview the migration without writing RBAC bindings. Use --apply to write changes.",
|
||||
)
|
||||
def migrate_dataset_permissions_to_rbac(
|
||||
tenant_id: str | None,
|
||||
dataset_id: str | None,
|
||||
batch_size: int,
|
||||
dry_run: bool,
|
||||
) -> None:
|
||||
"""Backfill RBAC dataset access config from legacy `Dataset.permission`.
|
||||
|
||||
Legacy mapping:
|
||||
- all_team_members -> RBAC dataset whitelist scope "all"
|
||||
- partial_members -> RBAC dataset whitelist scope "specific" plus each partial member gets the
|
||||
virtual default policy
|
||||
- only_me -> RBAC dataset whitelist scope "specific" with no member policy bindings
|
||||
|
||||
The command replaces each dataset's RBAC whitelist scope first. RBAC clears
|
||||
existing per-user policy bindings during that replace, then this command
|
||||
recreates the legacy partial-member default bindings. Re-running it is
|
||||
therefore idempotent for a dataset's current legacy configuration.
|
||||
"""
|
||||
click.echo(click.style("Starting RBAC dataset permission migration.", fg="green"))
|
||||
|
||||
scanned_count = 0
|
||||
scope_migrated_count = 0
|
||||
user_policy_migrated_count = 0
|
||||
partial_dataset_count = 0
|
||||
|
||||
last_dataset_id: str | None = None
|
||||
while True:
|
||||
with session_factory.create_session() as session:
|
||||
stmt = (
|
||||
select(Dataset.id, Dataset.tenant_id, Dataset.permission, Dataset.created_by)
|
||||
.order_by(Dataset.id.asc())
|
||||
.limit(batch_size)
|
||||
)
|
||||
if tenant_id:
|
||||
stmt = stmt.where(Dataset.tenant_id == tenant_id)
|
||||
if dataset_id:
|
||||
stmt = stmt.where(Dataset.id == dataset_id)
|
||||
if last_dataset_id:
|
||||
stmt = stmt.where(Dataset.id > last_dataset_id)
|
||||
|
||||
dataset_rows = list(session.execute(stmt).all())
|
||||
if not dataset_rows:
|
||||
break
|
||||
|
||||
dataset_ids = [str(row.id) for row in dataset_rows]
|
||||
partial_members_by_dataset_id: dict[str, list[str]] = {item: [] for item in dataset_ids}
|
||||
permission_rows = session.execute(
|
||||
select(DatasetPermission.dataset_id, DatasetPermission.account_id).where(
|
||||
DatasetPermission.dataset_id.in_(dataset_ids)
|
||||
)
|
||||
).all()
|
||||
for row in permission_rows:
|
||||
partial_members_by_dataset_id[str(row.dataset_id)].append(str(row.account_id))
|
||||
|
||||
for dataset in dataset_rows:
|
||||
workspace_id = str(dataset.tenant_id)
|
||||
current_dataset_id = str(dataset.id)
|
||||
operator_account_id = str(dataset.created_by)
|
||||
permission_value = _dataset_permission_enum(dataset.permission)
|
||||
scope = _rbac_dataset_scope_for_legacy_permission(permission_value)
|
||||
partial_member_ids = sorted(set(partial_members_by_dataset_id[current_dataset_id]))
|
||||
should_bind_partial_members = permission_value is DatasetPermissionEnum.PARTIAL_TEAM
|
||||
|
||||
click.echo(
|
||||
f"tenant={workspace_id} dataset={current_dataset_id} "
|
||||
f"operator={operator_account_id} "
|
||||
f"legacy_permission={permission_value} -> rbac_scope={scope} "
|
||||
f"partial_members={len(partial_member_ids) if should_bind_partial_members else 0}"
|
||||
)
|
||||
|
||||
scanned_count += 1
|
||||
replace_whitelist_payload = ReplaceMemberBindings(scope=scope)
|
||||
if dry_run:
|
||||
_emit_dataset_permission_migration_event(
|
||||
{
|
||||
"event": "dataset_permission_migration_proposed_change",
|
||||
"action": "replace_whitelist",
|
||||
"dry_run": True,
|
||||
"tenant_id": workspace_id,
|
||||
"dataset_id": current_dataset_id,
|
||||
"operator_account_id": operator_account_id,
|
||||
"before": {
|
||||
"legacy_dataset_permission": permission_value.value,
|
||||
"legacy_partial_member_ids": partial_member_ids if should_bind_partial_members else [],
|
||||
},
|
||||
"after": {
|
||||
"rbac_whitelist_scope": scope.value,
|
||||
},
|
||||
"call": {
|
||||
"method": "RBACService.DatasetAccess.replace_whitelist",
|
||||
"kwargs": {
|
||||
"tenant_id": workspace_id,
|
||||
"account_id": operator_account_id,
|
||||
"dataset_id": current_dataset_id,
|
||||
"payload": replace_whitelist_payload.model_dump(mode="json"),
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
if not dry_run:
|
||||
RBACService.DatasetAccess.replace_whitelist(
|
||||
tenant_id=workspace_id,
|
||||
account_id=operator_account_id,
|
||||
dataset_id=current_dataset_id,
|
||||
payload=replace_whitelist_payload,
|
||||
)
|
||||
scope_migrated_count += 1
|
||||
|
||||
if should_bind_partial_members:
|
||||
partial_dataset_count += 1
|
||||
for member_account_id in partial_member_ids:
|
||||
replace_user_access_policies_payload = ReplaceUserAccessPolicies(
|
||||
access_policy_ids=[_RBAC_DEFAULT_ACCESS_POLICY_ID],
|
||||
)
|
||||
if dry_run:
|
||||
_emit_dataset_permission_migration_event(
|
||||
{
|
||||
"event": "dataset_permission_migration_proposed_change",
|
||||
"action": "replace_user_access_policies",
|
||||
"dry_run": True,
|
||||
"tenant_id": workspace_id,
|
||||
"dataset_id": current_dataset_id,
|
||||
"operator_account_id": operator_account_id,
|
||||
"target_account_id": member_account_id,
|
||||
"before": {
|
||||
"legacy_dataset_permission": permission_value.value,
|
||||
"legacy_partial_member_id": member_account_id,
|
||||
},
|
||||
"after": {
|
||||
"rbac_user_access_policy_ids": [_RBAC_DEFAULT_ACCESS_POLICY_ID],
|
||||
},
|
||||
"call": {
|
||||
"method": "RBACService.DatasetAccess.replace_user_access_policies",
|
||||
"kwargs": {
|
||||
"tenant_id": workspace_id,
|
||||
"account_id": operator_account_id,
|
||||
"dataset_id": current_dataset_id,
|
||||
"target_account_id": member_account_id,
|
||||
"payload": replace_user_access_policies_payload.model_dump(mode="json"),
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
continue
|
||||
RBACService.DatasetAccess.replace_user_access_policies(
|
||||
tenant_id=workspace_id,
|
||||
account_id=operator_account_id,
|
||||
dataset_id=current_dataset_id,
|
||||
target_account_id=member_account_id,
|
||||
payload=replace_user_access_policies_payload,
|
||||
)
|
||||
user_policy_migrated_count += 1
|
||||
|
||||
last_dataset_id = dataset_ids[-1]
|
||||
|
||||
if dataset_id:
|
||||
break
|
||||
|
||||
if scanned_count == 0:
|
||||
click.echo(click.style("No datasets found for migration.", fg="yellow"))
|
||||
return
|
||||
|
||||
if dry_run:
|
||||
click.echo(click.style("Dry run completed. No RBAC bindings were written.", fg="yellow"))
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Dry run completed. Scanned {scanned_count} datasets; "
|
||||
f"{partial_dataset_count} partial-member datasets would be migrated.",
|
||||
fg="yellow",
|
||||
)
|
||||
)
|
||||
else:
|
||||
click.echo(click.style(f"RBAC member-role migration completed. Migrated {migrated_count} members.", fg="green"))
|
||||
click.echo(
|
||||
click.style(
|
||||
"RBAC dataset permission migration completed. "
|
||||
f"Scanned {scanned_count} datasets, migrated {scope_migrated_count} scopes, "
|
||||
f"wrote {user_policy_migrated_count} user default-policy bindings.",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
|
||||
@@ -34,6 +34,12 @@ class EnterpriseFeatureConfig(BaseSettings):
|
||||
default=False,
|
||||
)
|
||||
|
||||
ENTERPRISE_RBAC_REQUEST_TIMEOUT: int = Field(
|
||||
ge=1,
|
||||
description="Maximum timeout in seconds for inner RBAC requests.",
|
||||
default=30,
|
||||
)
|
||||
|
||||
|
||||
class EnterpriseTelemetryConfig(BaseSettings):
|
||||
"""
|
||||
|
||||
@@ -13,6 +13,7 @@ from controllers.console.wraps import (
|
||||
RBACPermission,
|
||||
RBACResourceScope,
|
||||
account_initialization_required,
|
||||
edit_permission_required,
|
||||
rbac_permission_required,
|
||||
setup_required,
|
||||
)
|
||||
@@ -95,9 +96,12 @@ class TraceAppConfigApi(Resource):
|
||||
console_ns.models[TraceAppConfigResponse.__name__],
|
||||
)
|
||||
@console_ns.response(400, "Invalid request parameters or configuration already exists")
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TRACING_CONFIG)
|
||||
@get_app_model
|
||||
def post(self, app_model: App):
|
||||
"""Create a new trace app configuration"""
|
||||
@@ -125,9 +129,12 @@ class TraceAppConfigApi(Resource):
|
||||
console_ns.models[TraceAppConfigResponse.__name__],
|
||||
)
|
||||
@console_ns.response(400, "Invalid request parameters or configuration not found")
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TRACING_CONFIG)
|
||||
@get_app_model
|
||||
def patch(self, app_model: App):
|
||||
"""Update an existing trace app configuration"""
|
||||
@@ -149,9 +156,12 @@ class TraceAppConfigApi(Resource):
|
||||
@console_ns.doc(params=query_params_from_model(TraceProviderQuery))
|
||||
@console_ns.response(204, "Tracing configuration deleted successfully")
|
||||
@console_ns.response(400, "Invalid request parameters or configuration not found")
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TRACING_CONFIG)
|
||||
@get_app_model
|
||||
def delete(self, app_model: App):
|
||||
"""Delete an existing trace app configuration"""
|
||||
|
||||
@@ -23,9 +23,9 @@ from libs.password import valid_password
|
||||
from models import Account
|
||||
from services.account_service import AccountService
|
||||
from services.billing_service import BillingService
|
||||
from services.errors.account import AccountRegisterError
|
||||
from services.errors.account import AccountRegisterError, SeatsLimitExceededError
|
||||
|
||||
from ..error import AccountInFreezeError, EmailSendIpLimitError
|
||||
from ..error import AccountInFreezeError, EmailSendIpLimitError, SeatsLimitExceeded
|
||||
from ..wraps import email_password_login_enabled, email_register_enabled, setup_required
|
||||
|
||||
|
||||
@@ -208,5 +208,7 @@ class EmailRegisterResetApi(Resource):
|
||||
timezone=timezone,
|
||||
session=db.session,
|
||||
)
|
||||
except SeatsLimitExceededError:
|
||||
raise SeatsLimitExceeded()
|
||||
except AccountRegisterError:
|
||||
raise AccountInFreezeError()
|
||||
|
||||
@@ -25,6 +25,7 @@ from controllers.console.error import (
|
||||
AccountNotFound,
|
||||
EmailSendIpLimitError,
|
||||
NotAllowedCreateWorkspace,
|
||||
SeatsLimitExceeded,
|
||||
WorkspacesLimitExceeded,
|
||||
)
|
||||
from controllers.console.wraps import (
|
||||
@@ -51,7 +52,7 @@ from models.account import Account
|
||||
from services.account_service import AccountService, InvitationDetailDict, RegisterService, TenantService
|
||||
from services.billing_service import BillingService
|
||||
from services.entities.auth_entities import LoginFailureReason, LoginPayloadBase
|
||||
from services.errors.account import AccountRegisterError
|
||||
from services.errors.account import AccountRegisterError, SeatsLimitExceededError
|
||||
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
@@ -317,6 +318,8 @@ class EmailCodeLoginApi(Resource):
|
||||
)
|
||||
except WorkSpaceNotAllowedCreateError:
|
||||
raise NotAllowedCreateWorkspace()
|
||||
except SeatsLimitExceededError:
|
||||
raise SeatsLimitExceeded()
|
||||
except AccountRegisterError:
|
||||
_log_console_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE)
|
||||
raise AccountInFreezeError()
|
||||
|
||||
@@ -25,7 +25,7 @@ from libs.token import (
|
||||
from models import Account, AccountStatus
|
||||
from services.account_service import AccountService, RegisterService, TenantService
|
||||
from services.billing_service import BillingService
|
||||
from services.errors.account import AccountNotFoundError, AccountRegisterError
|
||||
from services.errors.account import AccountNotFoundError, AccountRegisterError, SeatsLimitExceededError
|
||||
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkSpaceNotFoundError
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
@@ -182,6 +182,8 @@ class OAuthCallback(Resource):
|
||||
f"{dify_config.CONSOLE_WEB_URL}/signin"
|
||||
"?message=Workspace not found, please contact system admin to invite you to join in a workspace."
|
||||
)
|
||||
except SeatsLimitExceededError:
|
||||
return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Licensed seats limit exceeded.")
|
||||
except AccountRegisterError as e:
|
||||
return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message={e.description}")
|
||||
|
||||
|
||||
@@ -58,6 +58,12 @@ class WorkspacesLimitExceeded(BaseHTTPException):
|
||||
code = 400
|
||||
|
||||
|
||||
class SeatsLimitExceeded(BaseHTTPException):
|
||||
error_code = "limit_exceeded"
|
||||
description = "Unable to create account because the licensed seats limit was exceeded"
|
||||
code = 400
|
||||
|
||||
|
||||
class AccountBannedError(BaseHTTPException):
|
||||
error_code = "account_banned"
|
||||
description = "Account is banned."
|
||||
|
||||
@@ -36,7 +36,7 @@ from libs.login import current_account_with_tenant, login_required
|
||||
from models.account import Account, TenantAccountJoin, TenantAccountRole
|
||||
from services.account_service import AccountService, RegisterService, TenantService
|
||||
from services.enterprise import rbac_service as enterprise_rbac_service
|
||||
from services.errors.account import AccountAlreadyInTenantError
|
||||
from services.errors.account import AccountAlreadyInTenantError, SeatsLimitExceededError
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
|
||||
@@ -291,6 +291,14 @@ class MemberInviteEmailApi(Resource):
|
||||
"message": "Account already in workspace.",
|
||||
}
|
||||
)
|
||||
except SeatsLimitExceededError:
|
||||
invitation_results.append(
|
||||
MemberInviteFailedResponse(
|
||||
status="failed",
|
||||
email=invitee_email,
|
||||
message="Licensed seats limit exceeded.",
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
invitation_results.append({"status": "failed", "email": invitee_email, "message": str(e)})
|
||||
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from flask import request
|
||||
@@ -14,6 +13,7 @@ from controllers.common.schema import register_response_schema_models
|
||||
from controllers.console import console_ns
|
||||
from controllers.console.wraps import RBACPermission, RBACResourceScope, rbac_permission_required
|
||||
from core.db.session_factory import session_factory
|
||||
from core.rbac import RBACResourceWhitelistScope
|
||||
from libs.login import current_account_with_tenant, login_required
|
||||
from models import Account
|
||||
from services.enterprise import rbac_service as svc
|
||||
@@ -511,14 +511,8 @@ class RBACAccessPolicyBindingUnlockApi(Resource):
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _AccessScope(StrEnum):
|
||||
ALL = "all"
|
||||
SPECIFIC = "specific"
|
||||
ONLY_ME = "only_me"
|
||||
|
||||
|
||||
class _ResourceAccessScopeRequest(BaseModel):
|
||||
scope: _AccessScope
|
||||
scope: RBACResourceWhitelistScope
|
||||
|
||||
|
||||
class _ReplaceBindingsRequest(BaseModel):
|
||||
|
||||
@@ -48,6 +48,7 @@ from services.errors.account import (
|
||||
MemberNotInTenantError,
|
||||
NoPermissionError,
|
||||
RoleAlreadyAssignedError,
|
||||
SeatsLimitExceededError,
|
||||
)
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
@@ -190,6 +191,8 @@ class WorkspaceMembersApi(Resource):
|
||||
raise BadRequest(str(exc))
|
||||
except NoPermissionError as exc:
|
||||
raise BadRequest(str(exc))
|
||||
except SeatsLimitExceededError:
|
||||
raise BadRequest("licensed seats limit exceeded")
|
||||
except AccountRegisterError as exc:
|
||||
raise BadRequest(str(exc))
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from flask_restx import Resource
|
||||
from flask_restx.utils import merge
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
|
||||
|
||||
from configs import dify_config
|
||||
@@ -269,8 +270,8 @@ def cloud_edition_billing_rate_limit_check[**P, R](
|
||||
subscription_plan=knowledge_rate_limit.subscription_plan,
|
||||
operation="knowledge",
|
||||
)
|
||||
db.session.add(rate_limit_log)
|
||||
db.session.commit()
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
|
||||
session.add(rate_limit_log)
|
||||
raise Forbidden(
|
||||
"Sorry, you have reached the knowledge base request rate limit of your subscription."
|
||||
)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from collections.abc import Generator, Iterable, Mapping
|
||||
from typing import Any
|
||||
|
||||
from configs import dify_config
|
||||
from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler, print_text
|
||||
from core.ops.ops_trace_manager import TraceQueueManager
|
||||
from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||
@@ -19,8 +20,9 @@ class DifyWorkflowCallbackHandler(DifyAgentCallbackHandler):
|
||||
trace_manager: TraceQueueManager | None = None,
|
||||
) -> Generator[ToolInvokeMessage, None, None]:
|
||||
for tool_output in tool_outputs:
|
||||
print_text("\n[on_tool_execution]\n", color=self.color)
|
||||
print_text("Tool: " + tool_name + "\n", color=self.color)
|
||||
print_text("Outputs: " + tool_output.model_dump_json()[:1000] + "\n", color=self.color)
|
||||
print_text("\n")
|
||||
if dify_config.DEBUG:
|
||||
print_text("\n[on_tool_execution]\n", color=self.color)
|
||||
print_text("Tool: " + tool_name + "\n", color=self.color)
|
||||
print_text("Outputs: " + tool_output.model_dump_json()[:1000] + "\n", color=self.color)
|
||||
print_text("\n")
|
||||
yield tool_output
|
||||
|
||||
@@ -13,7 +13,7 @@ from core.helper.code_executor.jinja2.jinja2_transformer import Jinja2TemplateTr
|
||||
from core.helper.code_executor.python3.python3_transformer import Python3TemplateTransformer
|
||||
from core.helper.code_executor.template_transformer import TemplateTransformer
|
||||
from core.helper.http_client_pooling import get_pooled_http_client
|
||||
from graphon.nodes.code.entities import CodeLanguage
|
||||
from graphon.nodes.code.entities import CodeLanguage as CodeLanguage # noqa: PLC0414
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
code_execution_endpoint_url = URL(str(dify_config.CODE_EXECUTION_ENDPOINT))
|
||||
@@ -133,7 +133,9 @@ class CodeExecutor:
|
||||
return response_code.data.stdout or ""
|
||||
|
||||
@classmethod
|
||||
def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]):
|
||||
def execute_workflow_code_template(
|
||||
cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Execute code
|
||||
:param language: code language
|
||||
|
||||
@@ -11,7 +11,7 @@ class Jinja2TemplateTransformer(TemplateTransformer):
|
||||
|
||||
@classmethod
|
||||
@override
|
||||
def transform_response(cls, response: str):
|
||||
def transform_response(cls, response: str) -> dict[str, Any]:
|
||||
"""
|
||||
Transform response to dict
|
||||
:param response: response
|
||||
|
||||
@@ -36,14 +36,14 @@ class TemplateTransformer(ABC):
|
||||
return runner_script, preload_script
|
||||
|
||||
@classmethod
|
||||
def extract_result_str_from_response(cls, response: str):
|
||||
def extract_result_str_from_response(cls, response: str) -> str:
|
||||
result = re.search(rf"{cls._result_tag}(.*){cls._result_tag}", response, re.DOTALL)
|
||||
if not result:
|
||||
raise ValueError(f"Failed to parse result: no result tag found in response. Response: {response[:200]}...")
|
||||
return result.group(1)
|
||||
|
||||
@classmethod
|
||||
def transform_response(cls, response: str) -> Mapping[str, Any]:
|
||||
def transform_response(cls, response: str) -> dict[str, Any]:
|
||||
"""
|
||||
Transform response to dict
|
||||
:param response: response
|
||||
@@ -71,7 +71,7 @@ class TemplateTransformer(ABC):
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def _post_process_result(cls, result: dict[Any, Any]) -> dict[Any, Any]:
|
||||
def _post_process_result(cls, result: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Post-process the result to convert scientific notation strings back to numbers
|
||||
"""
|
||||
@@ -89,7 +89,7 @@ class TemplateTransformer(ABC):
|
||||
return [convert_scientific_notation(v) for v in value]
|
||||
return value
|
||||
|
||||
return convert_scientific_notation(result)
|
||||
return {key: convert_scientific_notation(value) for key, value in result.items()}
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
|
||||
@@ -24,7 +24,7 @@ def upload_dsl(dsl_file_bytes: bytes, filename: str = "template.yaml") -> str:
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
claim_code = data.get("data", {}).get("claim_code")
|
||||
if not claim_code:
|
||||
if not isinstance(claim_code, str) or not claim_code:
|
||||
raise ValueError("Creators Platform did not return a valid claim_code")
|
||||
return claim_code
|
||||
|
||||
|
||||
@@ -10,18 +10,21 @@ def is_credential_exists(credential_id: str, credential_type: "PluginCredentialT
|
||||
"""
|
||||
Check if the credential still exists in the database.
|
||||
|
||||
Uses the configured SQLAlchemy session factory instead of Flask-SQLAlchemy's
|
||||
``db.engine`` because workflow graph node construction may run without an
|
||||
active Flask application context.
|
||||
|
||||
:param credential_id: The credential ID to check
|
||||
:param credential_type: The type of credential (MODEL or TOOL)
|
||||
:return: True if credential exists, False otherwise
|
||||
"""
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from extensions.ext_database import db
|
||||
from core.db import session_factory
|
||||
from models.provider import ProviderCredential, ProviderModelCredential
|
||||
from models.tools import BuiltinToolProvider
|
||||
|
||||
with Session(db.engine) as session:
|
||||
with session_factory.create_session() as session:
|
||||
if credential_type == PluginCredentialType.MODEL:
|
||||
# Check both pre-defined and custom model credentials using a single UNION query
|
||||
stmt = (
|
||||
@@ -42,7 +45,7 @@ def is_credential_exists(credential_id: str, credential_type: "PluginCredentialT
|
||||
|
||||
def runtime_check_credential_policy_compliance(
|
||||
credential_id: str, provider: str, credential_type: "PluginCredentialType", check_existence: bool = True
|
||||
):
|
||||
) -> None:
|
||||
if dify_config.ENTERPRISE_DISABLE_RUNTIME_CREDENTIAL_CHECK:
|
||||
return
|
||||
check_credential_policy_compliance(
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
def download_with_size_limit(url, max_download_size: int, **kwargs):
|
||||
from typing import Any
|
||||
|
||||
|
||||
def download_with_size_limit(url: str, max_download_size: int, **kwargs: Any) -> bytes:
|
||||
from core.file import remote_fetcher
|
||||
|
||||
response = remote_fetcher.make_request("GET", url, follow_redirects=True, **kwargs)
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
import base64
|
||||
|
||||
from Crypto.PublicKey import RSA
|
||||
|
||||
from libs import rsa
|
||||
|
||||
|
||||
@@ -11,13 +13,13 @@ def obfuscated_token(token: str) -> str:
|
||||
return token[:6] + "*" * 12 + token[-2:]
|
||||
|
||||
|
||||
def full_mask_token(token_length=20):
|
||||
def full_mask_token(token_length: int = 20) -> str:
|
||||
return "*" * token_length
|
||||
|
||||
|
||||
def encrypt_token(tenant_id: str, token: str):
|
||||
from extensions.ext_database import db
|
||||
def encrypt_token(tenant_id: str, token: str) -> str:
|
||||
from models.account import Tenant
|
||||
from models.engine import db
|
||||
|
||||
if not (tenant := db.session.get(Tenant, tenant_id)):
|
||||
raise ValueError(f"Tenant with id {tenant_id} not found")
|
||||
@@ -30,15 +32,15 @@ def decrypt_token(tenant_id: str, token: str) -> str:
|
||||
return rsa.decrypt(base64.b64decode(token), tenant_id)
|
||||
|
||||
|
||||
def batch_decrypt_token(tenant_id: str, tokens: list[str]):
|
||||
def batch_decrypt_token(tenant_id: str, tokens: list[str]) -> list[str]:
|
||||
rsa_key, cipher_rsa = rsa.get_decrypt_decoding(tenant_id)
|
||||
|
||||
return [rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa) for token in tokens]
|
||||
|
||||
|
||||
def get_decrypt_decoding(tenant_id: str):
|
||||
def get_decrypt_decoding(tenant_id: str) -> tuple[RSA.RsaKey, object]:
|
||||
return rsa.get_decrypt_decoding(tenant_id)
|
||||
|
||||
|
||||
def decrypt_token_with_decoding(token: str, rsa_key, cipher_rsa):
|
||||
def decrypt_token_with_decoding(token: str, rsa_key: RSA.RsaKey, cipher_rsa: object) -> str:
|
||||
return rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from yarl import URL
|
||||
@@ -19,7 +20,7 @@ def get_plugin_pkg_url(plugin_unique_identifier: str) -> str:
|
||||
return str((marketplace_api_url / "api/v1/plugins/download").with_query(unique_identifier=plugin_unique_identifier))
|
||||
|
||||
|
||||
def download_plugin_pkg(plugin_unique_identifier: str):
|
||||
def download_plugin_pkg(plugin_unique_identifier: str) -> bytes:
|
||||
return download_with_size_limit(get_plugin_pkg_url(plugin_unique_identifier), dify_config.PLUGIN_MAX_PACKAGE_SIZE)
|
||||
|
||||
|
||||
@@ -39,7 +40,7 @@ def batch_fetch_plugin_manifests(plugin_ids: list[str]) -> Sequence[MarketplaceP
|
||||
return [MarketplacePluginDeclaration.model_validate(plugin) for plugin in response.json()["data"]["plugins"]]
|
||||
|
||||
|
||||
def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]:
|
||||
def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict[str, Any]]:
|
||||
if not plugin_ids:
|
||||
return []
|
||||
|
||||
@@ -53,10 +54,19 @@ def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]:
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
return data.get("data", {}).get("plugins", [])
|
||||
plugins = data.get("data", {}).get("plugins", [])
|
||||
if not isinstance(plugins, list):
|
||||
raise ValueError("Marketplace did not return a valid plugins list")
|
||||
|
||||
result: list[dict[str, Any]] = []
|
||||
for plugin in plugins:
|
||||
if not isinstance(plugin, dict) or not all(isinstance(key, str) for key in plugin):
|
||||
raise ValueError("Marketplace did not return a valid plugins list")
|
||||
result.append(plugin)
|
||||
return result
|
||||
|
||||
|
||||
def record_install_plugin_event(plugin_unique_identifier: str):
|
||||
def record_install_plugin_event(plugin_unique_identifier: str) -> None:
|
||||
url = str(marketplace_api_url / "api/v1/stats/plugins/install_count")
|
||||
response = httpx.post(url, json={"unique_identifier": plugin_unique_identifier}, timeout=MARKETPLACE_TIMEOUT)
|
||||
response.raise_for_status()
|
||||
|
||||
@@ -34,7 +34,7 @@ class ProviderCredentialsCache:
|
||||
else:
|
||||
return None
|
||||
|
||||
def set(self, credentials: dict[str, Any]):
|
||||
def set(self, credentials: dict[str, Any]) -> None:
|
||||
"""
|
||||
Cache model provider credentials.
|
||||
|
||||
@@ -43,7 +43,7 @@ class ProviderCredentialsCache:
|
||||
"""
|
||||
redis_client.setex(self.cache_key, 86400, json.dumps(credentials))
|
||||
|
||||
def delete(self):
|
||||
def delete(self) -> None:
|
||||
"""
|
||||
Delete cached model provider credentials.
|
||||
|
||||
|
||||
@@ -20,17 +20,18 @@ def import_module_from_source[T: (str, bytes)](
|
||||
raise Exception(f"Failed to load module {module_name} from {py_file_path!r}")
|
||||
else:
|
||||
# Refer to: https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
|
||||
# FIXME: mypy does not support the type of spec.loader
|
||||
spec = importlib.util.spec_from_file_location(module_name, py_file_path) # type: ignore[assignment]
|
||||
if not spec or not spec.loader:
|
||||
new_spec = importlib.util.spec_from_file_location(module_name, py_file_path)
|
||||
if not new_spec or not new_spec.loader:
|
||||
raise Exception(f"Failed to load module {module_name} from {py_file_path!r}")
|
||||
if use_lazy_loader:
|
||||
# Refer to: https://docs.python.org/3/library/importlib.html#implementing-lazy-imports
|
||||
spec.loader = importlib.util.LazyLoader(spec.loader)
|
||||
new_spec.loader = importlib.util.LazyLoader(new_spec.loader)
|
||||
spec = new_spec
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
if not existed_spec:
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
if spec.loader is not None:
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
except Exception as e:
|
||||
logger.exception("Failed to load module %s from script file '%s'", module_name, repr(py_file_path))
|
||||
|
||||
@@ -9,11 +9,11 @@ from extensions.ext_redis import redis_client
|
||||
class ProviderCredentialsCache(ABC):
|
||||
"""Base class for provider credentials cache"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
self.cache_key = self._generate_cache_key(**kwargs)
|
||||
|
||||
@abstractmethod
|
||||
def _generate_cache_key(self, **kwargs) -> str:
|
||||
def _generate_cache_key(self, **kwargs: Any) -> str:
|
||||
"""Generate cache key based on subclass implementation"""
|
||||
pass
|
||||
|
||||
@@ -28,11 +28,11 @@ class ProviderCredentialsCache(ABC):
|
||||
return None
|
||||
return None
|
||||
|
||||
def set(self, config: dict[str, Any]):
|
||||
def set(self, config: dict[str, Any]) -> None:
|
||||
"""Cache provider credentials"""
|
||||
redis_client.setex(self.cache_key, 86400, json.dumps(config))
|
||||
|
||||
def delete(self):
|
||||
def delete(self) -> None:
|
||||
"""Delete cached provider credentials"""
|
||||
redis_client.delete(self.cache_key)
|
||||
|
||||
@@ -48,7 +48,7 @@ class SingletonProviderCredentialsCache(ProviderCredentialsCache):
|
||||
)
|
||||
|
||||
@override
|
||||
def _generate_cache_key(self, **kwargs) -> str:
|
||||
def _generate_cache_key(self, **kwargs: Any) -> str:
|
||||
tenant_id = kwargs["tenant_id"]
|
||||
provider_type = kwargs["provider_type"]
|
||||
identity_name = kwargs["provider_identity"]
|
||||
@@ -63,7 +63,7 @@ class ToolProviderCredentialsCache(ProviderCredentialsCache):
|
||||
super().__init__(tenant_id=tenant_id, provider=provider, credential_id=credential_id)
|
||||
|
||||
@override
|
||||
def _generate_cache_key(self, **kwargs) -> str:
|
||||
def _generate_cache_key(self, **kwargs: Any) -> str:
|
||||
tenant_id = kwargs["tenant_id"]
|
||||
provider = kwargs["provider"]
|
||||
credential_id = kwargs["credential_id"]
|
||||
@@ -77,10 +77,10 @@ class NoOpProviderCredentialCache:
|
||||
"""Get cached provider credentials"""
|
||||
return None
|
||||
|
||||
def set(self, config: dict[str, Any]):
|
||||
def set(self, config: dict[str, Any]) -> None:
|
||||
"""Cache provider credentials"""
|
||||
pass
|
||||
|
||||
def delete(self):
|
||||
def delete(self) -> None:
|
||||
"""Delete cached provider credentials"""
|
||||
pass
|
||||
|
||||
@@ -125,5 +125,7 @@ class ProviderConfigEncrypter:
|
||||
return data
|
||||
|
||||
|
||||
def create_provider_encrypter(tenant_id: str, config: list[BasicProviderConfig], cache: ProviderConfigCache):
|
||||
def create_provider_encrypter(
|
||||
tenant_id: str, config: list[BasicProviderConfig], cache: ProviderConfigCache
|
||||
) -> tuple[ProviderConfigEncrypter, ProviderConfigCache]:
|
||||
return ProviderConfigEncrypter(tenant_id=tenant_id, config=config, provider_config_cache=cache), cache
|
||||
|
||||
@@ -37,11 +37,11 @@ class ToolParameterCache:
|
||||
else:
|
||||
return None
|
||||
|
||||
def set(self, parameters: dict[str, Any]):
|
||||
def set(self, parameters: dict[str, Any]) -> None:
|
||||
"""Cache model provider credentials."""
|
||||
redis_client.setex(self.cache_key, 86400, json.dumps(parameters))
|
||||
|
||||
def delete(self):
|
||||
def delete(self) -> None:
|
||||
"""
|
||||
Delete cached model provider credentials.
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ def get_external_trace_id(request: Any) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def extract_external_trace_id_from_args(args: Mapping[str, Any]):
|
||||
def extract_external_trace_id_from_args(args: Mapping[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Extract 'external_trace_id' from args.
|
||||
|
||||
|
||||
@@ -76,7 +76,11 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
if not user_id:
|
||||
user = EndUserService.get_or_create_end_user(app)
|
||||
else:
|
||||
user = cls._get_user(user_id, app)
|
||||
try:
|
||||
user = cls._get_user(user_id, app)
|
||||
except ValueError:
|
||||
# Plugins such as WeCom Bot pass external sender IDs rather than EndUser UUIDs.
|
||||
user = EndUserService.get_or_create_end_user(app, user_id=user_id)
|
||||
|
||||
conversation_id = conversation_id or ""
|
||||
|
||||
@@ -226,6 +230,13 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
EndUser.app_id == app.id,
|
||||
)
|
||||
user = session.scalar(stmt)
|
||||
if not user:
|
||||
stmt = select(EndUser).where(
|
||||
EndUser.session_id == user_id,
|
||||
EndUser.tenant_id == app.tenant_id,
|
||||
EndUser.app_id == app.id,
|
||||
)
|
||||
user = session.scalar(stmt)
|
||||
if not user:
|
||||
stmt = select(Account).where(
|
||||
Account.id == user_id,
|
||||
|
||||
@@ -228,7 +228,7 @@ class CredentialType(enum.StrEnum):
|
||||
OAUTH2 = "oauth2"
|
||||
UNAUTHORIZED = "unauthorized"
|
||||
|
||||
def get_name(self):
|
||||
def get_name(self) -> str:
|
||||
if self == CredentialType.API_KEY:
|
||||
return "API KEY"
|
||||
elif self == CredentialType.OAUTH2:
|
||||
|
||||
@@ -4,6 +4,8 @@ This module owns plugin daemon management calls that are shared by API services
|
||||
and core runtimes. Plugin model provider discovery is cached here, alongside
|
||||
plugin install, uninstall, and upgrade invalidation, so all cache mutations for
|
||||
plugin-owned provider metadata stay tenant-scoped and in one place.
|
||||
Provider cache payloads may be stored as prefixed zstd bytes; readers also
|
||||
accept legacy plain JSON payloads for rolling upgrades and existing Redis keys.
|
||||
|
||||
The console plugin list also normalizes endpoint setup counters against live
|
||||
endpoint records. Some plugin daemon builds return stale ``endpoints_*``
|
||||
@@ -14,12 +16,15 @@ metadata.
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from mimetypes import guess_type
|
||||
from typing import ClassVar
|
||||
from typing import Literal, Protocol
|
||||
|
||||
import zstandard
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from redis import RedisError
|
||||
from redis.exceptions import LockError
|
||||
from sqlalchemy import delete, select, update
|
||||
from sqlalchemy.orm import Session
|
||||
from yarl import URL
|
||||
@@ -67,14 +72,18 @@ logger = logging.getLogger(__name__)
|
||||
_provider_entities_adapter: TypeAdapter[list[ProviderEntity]] = TypeAdapter(list[ProviderEntity])
|
||||
|
||||
|
||||
class PluginService:
|
||||
_plugin_model_providers_memory_cache: ClassVar[dict[str, tuple[int, float, tuple[ProviderEntity, ...]]]] = {}
|
||||
class _RedisLock(Protocol):
|
||||
def acquire(self, *, blocking: bool = True, blocking_timeout: float | None = None) -> bool: ...
|
||||
|
||||
def release(self) -> None: ...
|
||||
|
||||
|
||||
class PluginService:
|
||||
class LatestPluginCache(BaseModel):
|
||||
plugin_id: str
|
||||
version: str
|
||||
unique_identifier: str
|
||||
status: str
|
||||
status: Literal["active", "deleted"]
|
||||
deprecated_reason: str
|
||||
alternative_plugin_id: str
|
||||
|
||||
@@ -82,6 +91,12 @@ class PluginService:
|
||||
REDIS_TTL = 60 * 5 # 5 minutes
|
||||
PLUGIN_MODEL_PROVIDERS_REDIS_KEY_PREFIX = "plugin_model_providers:tenant_id:"
|
||||
PLUGIN_MODEL_PROVIDERS_GENERATION_REDIS_KEY_PREFIX = "plugin_model_providers_generation:tenant_id:"
|
||||
PLUGIN_MODEL_PROVIDERS_LOCK_REDIS_KEY_PREFIX = "plugin_model_providers_refresh_lock:tenant_id:"
|
||||
PLUGIN_MODEL_PROVIDERS_LOCK_TTL = 30
|
||||
PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT = 2.0
|
||||
PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL = 0.05
|
||||
PLUGIN_MODEL_PROVIDERS_CACHE_COMPRESSION_PREFIX = b"\x00dify-plugin-model-providers-zstd-v1:"
|
||||
PLUGIN_MODEL_PROVIDERS_CACHE_COMPRESSION_MIN_BYTES = 64 * 1024
|
||||
PLUGIN_INSTALL_TASK_TERMINAL_STATUSES = (PluginInstallTaskStatus.Success, PluginInstallTaskStatus.Failed)
|
||||
# Mirror the detail-panel endpoint query size so list reconciliation and
|
||||
# the visible endpoint drawer exercise the same daemon pagination path.
|
||||
@@ -98,6 +113,10 @@ class PluginService:
|
||||
def _get_plugin_model_providers_generation_cache_key(cls, tenant_id: str) -> str:
|
||||
return f"{cls.PLUGIN_MODEL_PROVIDERS_GENERATION_REDIS_KEY_PREFIX}{tenant_id}"
|
||||
|
||||
@classmethod
|
||||
def _get_plugin_model_providers_lock_key(cls, tenant_id: str, generation: int) -> str:
|
||||
return f"{cls.PLUGIN_MODEL_PROVIDERS_LOCK_REDIS_KEY_PREFIX}{tenant_id}:generation:{generation}"
|
||||
|
||||
@staticmethod
|
||||
def _get_provider_short_name_alias(provider: PluginModelProviderEntity) -> str:
|
||||
"""
|
||||
@@ -129,8 +148,25 @@ class PluginService:
|
||||
return declaration
|
||||
|
||||
@classmethod
|
||||
def _copy_provider_entities(cls, providers: Sequence[ProviderEntity]) -> tuple[ProviderEntity, ...]:
|
||||
return tuple(provider.model_copy(deep=True) for provider in providers)
|
||||
def _encode_plugin_model_providers_cache_payload(cls, payload: bytes) -> bytes:
|
||||
if len(payload) < cls.PLUGIN_MODEL_PROVIDERS_CACHE_COMPRESSION_MIN_BYTES:
|
||||
return payload
|
||||
|
||||
return cls.PLUGIN_MODEL_PROVIDERS_CACHE_COMPRESSION_PREFIX + zstandard.compress(payload, level=1)
|
||||
|
||||
@classmethod
|
||||
def _decode_plugin_model_providers_cache_payload(cls, payload: bytes | bytearray | str) -> bytes | bytearray | str:
|
||||
if isinstance(payload, str):
|
||||
return payload
|
||||
|
||||
prefix = cls.PLUGIN_MODEL_PROVIDERS_CACHE_COMPRESSION_PREFIX
|
||||
if not payload.startswith(prefix):
|
||||
return payload
|
||||
|
||||
try:
|
||||
return zstandard.decompress(payload[len(prefix) :])
|
||||
except zstandard.ZstdError as exc:
|
||||
raise ValueError("Invalid compressed plugin model providers cache payload.") from exc
|
||||
|
||||
@classmethod
|
||||
def _load_plugin_model_providers_generation(cls, tenant_id: str) -> int | None:
|
||||
@@ -163,76 +199,35 @@ class PluginService:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _load_in_memory_plugin_model_providers(
|
||||
cls, memory_cache_key: str, generation: int
|
||||
) -> tuple[ProviderEntity, ...] | None:
|
||||
cached_entry = cls._plugin_model_providers_memory_cache.get(memory_cache_key)
|
||||
if cached_entry is None:
|
||||
return None
|
||||
def _load_cached_plugin_model_providers_for_generation(
|
||||
cls, tenant_id: str, generation: int | None
|
||||
) -> tuple[tuple[ProviderEntity, ...] | None, bool]:
|
||||
if generation is None:
|
||||
return None, False
|
||||
|
||||
cached_generation, expires_at, providers = cached_entry
|
||||
if cached_generation != generation or time.monotonic() >= expires_at:
|
||||
cls._plugin_model_providers_memory_cache.pop(memory_cache_key, None)
|
||||
return None
|
||||
|
||||
return cls._copy_provider_entities(providers)
|
||||
|
||||
@classmethod
|
||||
def _store_in_memory_plugin_model_providers(
|
||||
cls, memory_cache_key: str, generation: int, providers: Sequence[ProviderEntity]
|
||||
) -> None:
|
||||
ttl = dify_config.PLUGIN_MODEL_PROVIDERS_CACHE_TTL
|
||||
if ttl <= 0:
|
||||
cls._plugin_model_providers_memory_cache.pop(memory_cache_key, None)
|
||||
return
|
||||
|
||||
cls._plugin_model_providers_memory_cache[memory_cache_key] = (
|
||||
generation,
|
||||
time.monotonic() + ttl,
|
||||
cls._copy_provider_entities(providers),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _load_cached_plugin_model_providers(
|
||||
cls, tenant_id: str, *, client: PluginModelClient | None = None
|
||||
) -> tuple[ProviderEntity, ...] | None:
|
||||
generation = cls._load_plugin_model_providers_generation(tenant_id)
|
||||
if generation is not None:
|
||||
in_memory_cached_providers = cls._load_in_memory_plugin_model_providers(tenant_id, generation)
|
||||
if in_memory_cached_providers is not None:
|
||||
return in_memory_cached_providers
|
||||
|
||||
cache_keys = []
|
||||
if generation is not None:
|
||||
cache_keys.append(cls._get_plugin_model_providers_cache_key(tenant_id, generation))
|
||||
if generation == 0:
|
||||
cache_keys.append(cls._get_plugin_model_providers_cache_key(tenant_id))
|
||||
|
||||
if not cache_keys:
|
||||
return None
|
||||
cache_keys = [cls._get_plugin_model_providers_cache_key(tenant_id, generation)]
|
||||
|
||||
try:
|
||||
cached_provider_entries = redis_client.mget(cache_keys)
|
||||
except (RedisError, RuntimeError):
|
||||
except (LockError, RedisError, RuntimeError):
|
||||
logger.warning("Failed to read cached plugin model providers for tenant %s.", tenant_id, exc_info=True)
|
||||
return None
|
||||
return None, False
|
||||
|
||||
if len(cached_provider_entries) != len(cache_keys):
|
||||
logger.warning(
|
||||
"Unexpected cached plugin model providers response size for tenant %s.",
|
||||
tenant_id,
|
||||
)
|
||||
return None
|
||||
return None, False
|
||||
|
||||
for cache_key, cached_providers in zip(cache_keys, cached_provider_entries):
|
||||
if not cached_providers:
|
||||
continue
|
||||
|
||||
try:
|
||||
providers = tuple(_provider_entities_adapter.validate_json(cached_providers))
|
||||
if generation is not None:
|
||||
cls._store_in_memory_plugin_model_providers(tenant_id, generation, providers)
|
||||
return providers
|
||||
payload = cls._decode_plugin_model_providers_cache_payload(cached_providers)
|
||||
providers = tuple(_provider_entities_adapter.validate_json(payload))
|
||||
return providers, True
|
||||
except (TypeError, ValueError, ValidationError):
|
||||
logger.warning(
|
||||
"Invalid cached plugin model providers for tenant %s; deleting cache key %s.",
|
||||
@@ -249,7 +244,7 @@ class PluginService:
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return None
|
||||
return None, True
|
||||
|
||||
@classmethod
|
||||
def _store_cached_plugin_model_providers(
|
||||
@@ -257,15 +252,94 @@ class PluginService:
|
||||
) -> None:
|
||||
cache_key = cls._get_plugin_model_providers_cache_key(tenant_id, generation)
|
||||
try:
|
||||
payload = _provider_entities_adapter.dump_json(list(providers)).decode("utf-8")
|
||||
payload = cls._encode_plugin_model_providers_cache_payload(
|
||||
_provider_entities_adapter.dump_json(list(providers))
|
||||
)
|
||||
redis_client.setex(cache_key, dify_config.PLUGIN_MODEL_PROVIDERS_CACHE_TTL, payload)
|
||||
except (RedisError, RuntimeError):
|
||||
logger.warning("Failed to cache plugin model providers for tenant %s.", tenant_id, exc_info=True)
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def _plugin_model_providers_refresh_lock(
|
||||
cls, tenant_id: str, generation: int, *, wait_timeout: float
|
||||
) -> Iterator[bool]:
|
||||
lock_key = cls._get_plugin_model_providers_lock_key(tenant_id, generation)
|
||||
try:
|
||||
refresh_lock: _RedisLock = redis_client.lock(
|
||||
lock_key,
|
||||
timeout=cls.PLUGIN_MODEL_PROVIDERS_LOCK_TTL,
|
||||
sleep=cls.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL,
|
||||
)
|
||||
except (RedisError, RuntimeError):
|
||||
logger.warning(
|
||||
"Failed to create plugin model providers refresh lock for tenant %s.",
|
||||
tenant_id,
|
||||
exc_info=True,
|
||||
)
|
||||
yield False
|
||||
return
|
||||
|
||||
try:
|
||||
lock_acquired = refresh_lock.acquire(blocking=True, blocking_timeout=wait_timeout)
|
||||
except LockError:
|
||||
logger.warning(
|
||||
"Provider refresh lock timed out; direct daemon fallback. tenant_id=%s generation=%s",
|
||||
tenant_id,
|
||||
generation,
|
||||
exc_info=True,
|
||||
)
|
||||
yield False
|
||||
return
|
||||
except (RedisError, RuntimeError):
|
||||
# Redis failures should not block provider discovery; callers fetch directly from the daemon.
|
||||
logger.warning(
|
||||
"Failed to acquire plugin model providers refresh lock for tenant %s.",
|
||||
tenant_id,
|
||||
exc_info=True,
|
||||
)
|
||||
yield False
|
||||
return
|
||||
|
||||
if not lock_acquired:
|
||||
logger.warning(
|
||||
"Provider refresh lock timed out; direct daemon fallback. tenant_id=%s generation=%s",
|
||||
tenant_id,
|
||||
generation,
|
||||
)
|
||||
yield False
|
||||
return
|
||||
|
||||
try:
|
||||
yield True
|
||||
finally:
|
||||
try:
|
||||
refresh_lock.release()
|
||||
except (LockError, RedisError, RuntimeError):
|
||||
# Release failures must not hide the daemon result or the original exception.
|
||||
logger.warning(
|
||||
"Failed to release plugin model providers refresh lock for tenant %s generation %s.",
|
||||
tenant_id,
|
||||
generation,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
@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)
|
||||
)
|
||||
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)
|
||||
return providers
|
||||
|
||||
@classmethod
|
||||
def invalidate_plugin_model_providers_cache(cls, tenant_id: str) -> None:
|
||||
"""Invalidate tenant-scoped provider metadata across Redis and worker-local mirrors."""
|
||||
cls._plugin_model_providers_memory_cache.pop(tenant_id, None)
|
||||
"""Invalidate tenant-scoped provider metadata stored in Redis."""
|
||||
cache_key = cls._get_plugin_model_providers_cache_key(tenant_id)
|
||||
generation_key = cls._get_plugin_model_providers_generation_cache_key(tenant_id)
|
||||
try:
|
||||
@@ -287,21 +361,68 @@ class PluginService:
|
||||
are intentionally owned by this service so tenant isolation and cache
|
||||
expiry are handled in one place.
|
||||
"""
|
||||
cached_providers = cls._load_cached_plugin_model_providers(tenant_id, client=client)
|
||||
if cached_providers is not None:
|
||||
return cached_providers
|
||||
deadline = time.monotonic() + cls.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT
|
||||
|
||||
model_client = client or PluginModelClient()
|
||||
providers = tuple(
|
||||
cls._to_provider_entity(provider) for provider in model_client.fetch_model_providers(tenant_id)
|
||||
)
|
||||
if not providers:
|
||||
return providers
|
||||
generation = cls._load_plugin_model_providers_generation(tenant_id)
|
||||
if generation is not None:
|
||||
cls._store_in_memory_plugin_model_providers(tenant_id, generation, providers)
|
||||
cls._store_cached_plugin_model_providers(tenant_id, generation, providers)
|
||||
return providers
|
||||
while True:
|
||||
generation = cls._load_plugin_model_providers_generation(tenant_id)
|
||||
cached_providers, cache_available = cls._load_cached_plugin_model_providers_for_generation(
|
||||
tenant_id, generation
|
||||
)
|
||||
if cached_providers is not None:
|
||||
return cached_providers
|
||||
|
||||
if generation is None or not cache_available:
|
||||
return cls._fetch_and_cache_plugin_model_providers(
|
||||
tenant_id,
|
||||
client,
|
||||
refresh_generation=generation,
|
||||
)
|
||||
|
||||
wait_timeout = deadline - time.monotonic()
|
||||
if wait_timeout < 0:
|
||||
logger.warning(
|
||||
"Provider refresh lock timed out; direct daemon fallback. tenant_id=%s generation=%s",
|
||||
tenant_id,
|
||||
generation,
|
||||
)
|
||||
return cls._fetch_and_cache_plugin_model_providers(
|
||||
tenant_id,
|
||||
client,
|
||||
refresh_generation=generation,
|
||||
)
|
||||
|
||||
with cls._plugin_model_providers_refresh_lock(
|
||||
tenant_id,
|
||||
generation,
|
||||
wait_timeout=wait_timeout,
|
||||
) as lock_acquired:
|
||||
if not lock_acquired:
|
||||
return cls._fetch_and_cache_plugin_model_providers(
|
||||
tenant_id,
|
||||
client,
|
||||
refresh_generation=generation,
|
||||
)
|
||||
|
||||
latest_generation = cls._load_plugin_model_providers_generation(tenant_id)
|
||||
cached_providers, cache_available = cls._load_cached_plugin_model_providers_for_generation(
|
||||
tenant_id, latest_generation
|
||||
)
|
||||
if cached_providers is not None:
|
||||
return cached_providers
|
||||
if latest_generation is None or not cache_available:
|
||||
return cls._fetch_and_cache_plugin_model_providers(
|
||||
tenant_id,
|
||||
client,
|
||||
refresh_generation=latest_generation,
|
||||
)
|
||||
if latest_generation != generation:
|
||||
continue
|
||||
|
||||
return cls._fetch_and_cache_plugin_model_providers(
|
||||
tenant_id,
|
||||
client,
|
||||
refresh_generation=generation,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def fetch_latest_plugin_version(plugin_ids: Sequence[str]) -> Mapping[str, LatestPluginCache | None]:
|
||||
|
||||
@@ -1030,6 +1030,10 @@ class DatasetRetrieval:
|
||||
):
|
||||
"""
|
||||
Persist dataset query audit rows for retrieval requests.
|
||||
|
||||
Query audit logging is a side effect of retrieval. Keep it in an
|
||||
independent transaction so failures or commits here do not affect the
|
||||
request/workflow transaction that called the retriever.
|
||||
"""
|
||||
if not query and not attachment_ids:
|
||||
return
|
||||
@@ -1041,6 +1045,9 @@ class DatasetRetrieval:
|
||||
app_id,
|
||||
)
|
||||
return
|
||||
created_by_role = self._resolve_creator_user_role(user_from)
|
||||
if created_by_role is None:
|
||||
return
|
||||
dataset_queries = []
|
||||
for dataset_id in dataset_ids:
|
||||
contents = []
|
||||
@@ -1055,13 +1062,16 @@ class DatasetRetrieval:
|
||||
content=json.dumps(contents),
|
||||
source=DatasetQuerySource.APP,
|
||||
source_app_id=app_id,
|
||||
created_by_role=CreatorUserRole(user_from),
|
||||
created_by_role=created_by_role,
|
||||
created_by=created_by,
|
||||
)
|
||||
dataset_queries.append(dataset_query)
|
||||
if dataset_queries:
|
||||
db.session.add_all(dataset_queries)
|
||||
db.session.commit()
|
||||
|
||||
if not dataset_queries:
|
||||
return
|
||||
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
|
||||
session.add_all(dataset_queries)
|
||||
|
||||
def _retriever(
|
||||
self,
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
from core.rbac.entities import RBACPermission, RBACResourceScope
|
||||
from core.rbac.entities import RBACPermission, RBACResourceScope, RBACResourceWhitelistScope
|
||||
|
||||
__all__ = ["RBACPermission", "RBACResourceScope"]
|
||||
__all__ = ["RBACPermission", "RBACResourceScope", "RBACResourceWhitelistScope"]
|
||||
|
||||
@@ -13,6 +13,14 @@ class RBACResourceScope(StrEnum):
|
||||
WORKSPACE = "workspace"
|
||||
|
||||
|
||||
class RBACResourceWhitelistScope(StrEnum):
|
||||
"""Whitelist scopes accepted by RBAC app and dataset access config APIs."""
|
||||
|
||||
ALL = "all"
|
||||
SPECIFIC = "specific"
|
||||
ONLY_ME = "only_me"
|
||||
|
||||
|
||||
class RBACPermission(StrEnum):
|
||||
"""Permission points (RBAC scenes) checked by ``rbac_permission_required``.
|
||||
|
||||
|
||||
@@ -102,16 +102,17 @@ class ApiTool(Tool):
|
||||
elif not isinstance(credentials["api_key_value"], str):
|
||||
raise ToolProviderCredentialValidationError("api_key_value must be a string")
|
||||
|
||||
api_key_value = credentials["api_key_value"]
|
||||
if "api_key_header_prefix" in credentials:
|
||||
api_key_header_prefix = credentials["api_key_header_prefix"]
|
||||
if api_key_header_prefix == "basic" and credentials["api_key_value"]:
|
||||
credentials["api_key_value"] = f"Basic {credentials['api_key_value']}"
|
||||
elif api_key_header_prefix == "bearer" and credentials["api_key_value"]:
|
||||
credentials["api_key_value"] = f"Bearer {credentials['api_key_value']}"
|
||||
if api_key_header_prefix == "basic" and api_key_value:
|
||||
api_key_value = f"Basic {api_key_value}"
|
||||
elif api_key_header_prefix == "bearer" and api_key_value:
|
||||
api_key_value = f"Bearer {api_key_value}"
|
||||
elif api_key_header_prefix == "custom":
|
||||
pass
|
||||
|
||||
headers[api_key_header] = credentials["api_key_value"]
|
||||
headers[api_key_header] = api_key_value
|
||||
|
||||
elif credentials["auth_type"] == "api_key_query":
|
||||
# For query parameter authentication, we don't add anything to headers
|
||||
|
||||
@@ -7,6 +7,7 @@ from datetime import UTC, datetime
|
||||
from mimetypes import guess_type
|
||||
from typing import Any, Union, cast
|
||||
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from yarl import URL
|
||||
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
@@ -338,47 +339,49 @@ class ToolEngine:
|
||||
user_id: str,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Create message file
|
||||
Create message files produced by a tool call.
|
||||
|
||||
Tool file persistence is a side effect of agent execution. Use an
|
||||
independent transaction so this helper never commits or closes the
|
||||
caller's request-scoped session.
|
||||
|
||||
:return: message file ids
|
||||
"""
|
||||
result = []
|
||||
|
||||
for message in tool_messages:
|
||||
if "image" in message.mimetype:
|
||||
file_type = FileType.IMAGE
|
||||
elif "video" in message.mimetype:
|
||||
file_type = FileType.VIDEO
|
||||
elif "audio" in message.mimetype:
|
||||
file_type = FileType.AUDIO
|
||||
elif "text" in message.mimetype or "pdf" in message.mimetype:
|
||||
file_type = FileType.DOCUMENT
|
||||
else:
|
||||
file_type = FileType.CUSTOM
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
|
||||
for message in tool_messages:
|
||||
# extract tool file id from url
|
||||
tool_file_id = message.url.split("/")[-1].split(".")[0]
|
||||
message_file = MessageFile(
|
||||
message_id=agent_message.id,
|
||||
type=ToolEngine._resolve_tool_file_type(message),
|
||||
transfer_method=FileTransferMethod.TOOL_FILE,
|
||||
belongs_to=MessageFileBelongsTo.ASSISTANT,
|
||||
url=message.url,
|
||||
upload_file_id=tool_file_id,
|
||||
created_by_role=(
|
||||
CreatorUserRole.ACCOUNT
|
||||
if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}
|
||||
else CreatorUserRole.END_USER
|
||||
),
|
||||
created_by=user_id,
|
||||
)
|
||||
|
||||
# extract tool file id from url
|
||||
tool_file_id = message.url.split("/")[-1].split(".")[0]
|
||||
message_file = MessageFile(
|
||||
message_id=agent_message.id,
|
||||
type=file_type,
|
||||
transfer_method=FileTransferMethod.TOOL_FILE,
|
||||
belongs_to=MessageFileBelongsTo.ASSISTANT,
|
||||
url=message.url,
|
||||
upload_file_id=tool_file_id,
|
||||
created_by_role=(
|
||||
CreatorUserRole.ACCOUNT
|
||||
if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}
|
||||
else CreatorUserRole.END_USER
|
||||
),
|
||||
created_by=user_id,
|
||||
)
|
||||
|
||||
db.session.add(message_file)
|
||||
db.session.commit()
|
||||
db.session.refresh(message_file)
|
||||
|
||||
result.append(message_file.id)
|
||||
|
||||
db.session.close()
|
||||
session.add(message_file)
|
||||
result.append(message_file.id)
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _resolve_tool_file_type(message: ToolInvokeMessageBinary) -> FileType:
|
||||
if "image" in message.mimetype:
|
||||
return FileType.IMAGE
|
||||
elif "video" in message.mimetype:
|
||||
return FileType.VIDEO
|
||||
elif "audio" in message.mimetype:
|
||||
return FileType.AUDIO
|
||||
elif "text" in message.mimetype or "pdf" in message.mimetype:
|
||||
return FileType.DOCUMENT
|
||||
else:
|
||||
return FileType.CUSTOM
|
||||
|
||||
@@ -27,6 +27,7 @@ def init_app(app: DifyApp):
|
||||
install_plugins,
|
||||
install_rag_pipeline_plugins,
|
||||
migrate_data_for_plugin,
|
||||
migrate_dataset_permissions_to_rbac,
|
||||
migrate_member_roles_to_rbac,
|
||||
migrate_oss,
|
||||
migration_data_wizard,
|
||||
@@ -56,6 +57,7 @@ def init_app(app: DifyApp):
|
||||
upgrade_db,
|
||||
fix_app_site_missing,
|
||||
migrate_data_for_plugin,
|
||||
migrate_dataset_permissions_to_rbac,
|
||||
migrate_member_roles_to_rbac,
|
||||
backfill_plugin_auto_upgrade,
|
||||
extract_plugins,
|
||||
|
||||
@@ -214,6 +214,18 @@ class EndUserType(StrEnum):
|
||||
SERVICE_API = "service-api"
|
||||
TRIGGER = "trigger"
|
||||
|
||||
@classmethod
|
||||
@override
|
||||
def _missing_(cls, value):
|
||||
# Legacy rows persisted the service-api type with an underscore before it
|
||||
# was normalized to the hyphenated value. The
|
||||
# `4f7b2c8d9a10_normalize_legacy_end_user_type` migration rewrites those
|
||||
# rows, but tolerate the old value here as well so an unmigrated end user
|
||||
# keeps loading instead of failing enum validation on every request.
|
||||
if value == "service_api":
|
||||
return cls.SERVICE_API
|
||||
return super()._missing_(value)
|
||||
|
||||
|
||||
class DocumentDocType(StrEnum):
|
||||
"""Document doc_type classification"""
|
||||
|
||||
@@ -2868,6 +2868,7 @@ Delete an existing tracing configuration for an application
|
||||
| ---- | ----------- |
|
||||
| 204 | Tracing configuration deleted successfully |
|
||||
| 400 | Invalid request parameters or configuration not found |
|
||||
| 403 | Insufficient permissions |
|
||||
|
||||
### [GET] /apps/{app_id}/trace-config
|
||||
Get tracing configuration for an application
|
||||
@@ -2909,6 +2910,7 @@ Update an existing tracing configuration for an application
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Tracing configuration updated successfully | **application/json**: [TraceAppConfigResponse](#traceappconfigresponse)<br> |
|
||||
| 400 | Invalid request parameters or configuration not found | |
|
||||
| 403 | Insufficient permissions | |
|
||||
|
||||
### [POST] /apps/{app_id}/trace-config
|
||||
**Create a new trace app configuration**
|
||||
@@ -2933,6 +2935,7 @@ Create a new tracing configuration for an application
|
||||
| ---- | ----------- | ------ |
|
||||
| 201 | Tracing configuration created successfully | **application/json**: [TraceAppConfigResponse](#traceappconfigresponse)<br> |
|
||||
| 400 | Invalid request parameters or configuration already exists | |
|
||||
| 403 | Insufficient permissions | |
|
||||
|
||||
### [POST] /apps/{app_id}/trigger-enable
|
||||
**Update app trigger (enable/disable)**
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from enum import StrEnum
|
||||
from typing import Any, override
|
||||
@@ -17,6 +18,8 @@ from models.dataset import Dataset
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
METADATA_KEY_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
|
||||
|
||||
class MyScaleConfig(BaseModel):
|
||||
host: str
|
||||
@@ -102,7 +105,10 @@ class MyScaleVector(BaseVector):
|
||||
|
||||
@override
|
||||
def text_exists(self, id: str) -> bool:
|
||||
results = self._client.query(f"SELECT id FROM {self._config.database}.{self._collection_name} WHERE id='{id}'")
|
||||
results = self._client.query(
|
||||
f"SELECT id FROM {self._config.database}.{self._collection_name} WHERE id={{id:String}}",
|
||||
parameters={"id": id},
|
||||
)
|
||||
return results.row_count > 0
|
||||
|
||||
@override
|
||||
@@ -110,20 +116,26 @@ class MyScaleVector(BaseVector):
|
||||
if not ids:
|
||||
return
|
||||
self._client.command(
|
||||
f"DELETE FROM {self._config.database}.{self._collection_name} WHERE id IN {str(tuple(ids))}"
|
||||
f"DELETE FROM {self._config.database}.{self._collection_name} WHERE id IN {{ids:Array(String)}}",
|
||||
parameters={"ids": ids},
|
||||
)
|
||||
|
||||
@override
|
||||
def get_ids_by_metadata_field(self, key: str, value: str):
|
||||
self._validate_metadata_key(key)
|
||||
rows = self._client.query(
|
||||
f"SELECT DISTINCT id FROM {self._config.database}.{self._collection_name} WHERE metadata.{key}='{value}'"
|
||||
f"SELECT DISTINCT id FROM {self._config.database}.{self._collection_name} "
|
||||
f"WHERE metadata.{key}={{value:String}}",
|
||||
parameters={"value": value},
|
||||
).result_rows
|
||||
return [row[0] for row in rows]
|
||||
|
||||
@override
|
||||
def delete_by_metadata_field(self, key: str, value: str):
|
||||
self._validate_metadata_key(key)
|
||||
self._client.command(
|
||||
f"DELETE FROM {self._config.database}.{self._collection_name} WHERE metadata.{key}='{value}'"
|
||||
f"DELETE FROM {self._config.database}.{self._collection_name} WHERE metadata.{key}={{value:String}}",
|
||||
parameters={"value": value},
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -132,22 +144,29 @@ class MyScaleVector(BaseVector):
|
||||
|
||||
@override
|
||||
def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
|
||||
return self._search(f"TextSearch('enable_nlq=false')(text, '{query}')", SortOrder.DESC, **kwargs)
|
||||
return self._search(
|
||||
"TextSearch('enable_nlq=false')(text, {query:String})",
|
||||
SortOrder.DESC,
|
||||
parameters={"query": query},
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def _search(self, dist: str, order: SortOrder, **kwargs: Any) -> list[Document]:
|
||||
def _search(
|
||||
self, dist: str, order: SortOrder, parameters: dict[str, Any] | None = None, **kwargs: Any
|
||||
) -> list[Document]:
|
||||
top_k = kwargs.get("top_k", 4)
|
||||
if not isinstance(top_k, int) or top_k <= 0:
|
||||
raise ValueError("top_k must be a positive integer")
|
||||
score_threshold = float(kwargs.get("score_threshold") or 0.0)
|
||||
where_str = (
|
||||
f"WHERE dist < {1 - score_threshold}"
|
||||
if self._metric.upper() == "COSINE" and order == SortOrder.ASC and score_threshold > 0.0
|
||||
else ""
|
||||
)
|
||||
where_conditions = []
|
||||
query_parameters = dict(parameters or {})
|
||||
if self._metric.upper() == "COSINE" and order == SortOrder.ASC and score_threshold > 0.0:
|
||||
where_conditions.append(f"dist < {1 - score_threshold}")
|
||||
document_ids_filter = kwargs.get("document_ids_filter")
|
||||
if document_ids_filter:
|
||||
document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
|
||||
where_str = f"{where_str} AND metadata['document_id'] in ({document_ids})"
|
||||
where_conditions.append("metadata['document_id'] IN {document_ids_filter:Array(String)}")
|
||||
query_parameters["document_ids_filter"] = document_ids_filter
|
||||
where_str = f"WHERE {' AND '.join(where_conditions)}" if where_conditions else ""
|
||||
sql = f"""
|
||||
SELECT text, vector, metadata, {dist} as dist FROM {self._config.database}.{self._collection_name}
|
||||
{where_str} ORDER BY dist {order.value} LIMIT {top_k}
|
||||
@@ -159,12 +178,17 @@ class MyScaleVector(BaseVector):
|
||||
vector=r["vector"],
|
||||
metadata=r["metadata"],
|
||||
)
|
||||
for r in self._client.query(sql).named_results()
|
||||
for r in self._client.query(sql, parameters=query_parameters).named_results()
|
||||
]
|
||||
except Exception:
|
||||
logger.exception("Vector search operation failed")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _validate_metadata_key(key: str) -> None:
|
||||
if not METADATA_KEY_PATTERN.match(key):
|
||||
raise ValueError("metadata key must be a valid identifier")
|
||||
|
||||
@override
|
||||
def delete(self):
|
||||
self._client.command(f"DROP TABLE IF EXISTS {self._config.database}.{self._collection_name}")
|
||||
|
||||
@@ -181,14 +181,41 @@ def test_text_exists_and_metadata_operations(myscale_module):
|
||||
vector = myscale_module.MyScaleVector("collection_1", _config(myscale_module))
|
||||
vector._client.query.return_value = SimpleNamespace(row_count=1, result_rows=[("id-1",), ("id-2",)])
|
||||
|
||||
assert vector.text_exists("id-1") is True
|
||||
assert vector.get_ids_by_metadata_field("document_id", "doc-1") == ["id-1", "id-2"]
|
||||
assert vector.text_exists("id-1' OR '1'='1") is True
|
||||
text_exists_call = vector._client.query.call_args
|
||||
assert "id={id:String}" in text_exists_call.args[0]
|
||||
assert "id-1' OR '1'='1" not in text_exists_call.args[0]
|
||||
assert text_exists_call.kwargs["parameters"] == {"id": "id-1' OR '1'='1"}
|
||||
|
||||
assert vector.get_ids_by_metadata_field("document_id", "doc-1' OR '1'='1") == ["id-1", "id-2"]
|
||||
metadata_query_call = vector._client.query.call_args
|
||||
assert "metadata.document_id={value:String}" in metadata_query_call.args[0]
|
||||
assert "doc-1' OR '1'='1" not in metadata_query_call.args[0]
|
||||
assert metadata_query_call.kwargs["parameters"] == {"value": "doc-1' OR '1'='1"}
|
||||
|
||||
vector.delete_by_ids(["id-1", "id-2"])
|
||||
vector.delete_by_metadata_field("document_id", "doc-1")
|
||||
delete_ids_call = vector._client.command.call_args
|
||||
assert "id IN {ids:Array(String)}" in delete_ids_call.args[0]
|
||||
assert delete_ids_call.kwargs["parameters"] == {"ids": ["id-1", "id-2"]}
|
||||
|
||||
vector.delete_by_metadata_field("document_id", "doc-1' OR '1'='1")
|
||||
delete_metadata_call = vector._client.command.call_args
|
||||
assert "metadata.document_id={value:String}" in delete_metadata_call.args[0]
|
||||
assert "doc-1' OR '1'='1" not in delete_metadata_call.args[0]
|
||||
assert delete_metadata_call.kwargs["parameters"] == {"value": "doc-1' OR '1'='1"}
|
||||
assert vector._client.command.call_count >= 2
|
||||
|
||||
|
||||
def test_metadata_operations_reject_invalid_key(myscale_module):
|
||||
vector = myscale_module.MyScaleVector("collection_1", _config(myscale_module))
|
||||
|
||||
with pytest.raises(ValueError, match="metadata key must be a valid identifier"):
|
||||
vector.get_ids_by_metadata_field("document_id) OR 1=1 --", "doc-1")
|
||||
|
||||
with pytest.raises(ValueError, match="metadata key must be a valid identifier"):
|
||||
vector.delete_by_metadata_field("document_id) OR 1=1 --", "doc-1")
|
||||
|
||||
|
||||
def test_search_delegation_methods(myscale_module):
|
||||
vector = myscale_module.MyScaleVector("collection_1", _config(myscale_module))
|
||||
vector._search = MagicMock(return_value=["result"])
|
||||
@@ -199,6 +226,28 @@ def test_search_delegation_methods(myscale_module):
|
||||
assert result_vector == ["result"]
|
||||
assert result_text == ["result"]
|
||||
assert vector._search.call_count == 2
|
||||
vector._search.assert_any_call(
|
||||
"TextSearch('enable_nlq=false')(text, {query:String})",
|
||||
myscale_module.SortOrder.DESC,
|
||||
parameters={"query": "hello"},
|
||||
top_k=2,
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_full_text_uses_query_parameters(myscale_module):
|
||||
vector = myscale_module.MyScaleVector("collection_1", _config(myscale_module))
|
||||
vector._client.query.return_value = SimpleNamespace(
|
||||
named_results=lambda: [{"text": "doc", "vector": [0.1], "metadata": {"doc_id": "1"}}]
|
||||
)
|
||||
payload = "x') AS dist FROM dify.collection_1 UNION ALL SELECT secret FROM users --"
|
||||
|
||||
docs = vector.search_by_full_text(payload, top_k=2)
|
||||
|
||||
assert len(docs) == 1
|
||||
sql = vector._client.query.call_args.args[0]
|
||||
assert payload not in sql
|
||||
assert "TextSearch('enable_nlq=false')(text, {query:String})" in sql
|
||||
assert vector._client.query.call_args.kwargs["parameters"] == {"query": payload}
|
||||
|
||||
|
||||
def test_search_with_document_filter_and_exception(myscale_module):
|
||||
@@ -215,7 +264,8 @@ def test_search_with_document_filter_and_exception(myscale_module):
|
||||
)
|
||||
assert len(docs) == 1
|
||||
sql = vector._client.query.call_args.args[0]
|
||||
assert "metadata['document_id'] in ('doc-1', 'doc-2')" in sql
|
||||
assert "WHERE metadata['document_id'] IN {document_ids_filter:Array(String)}" in sql
|
||||
assert vector._client.query.call_args.kwargs["parameters"] == {"document_ids_filter": ["doc-1", "doc-2"]}
|
||||
|
||||
vector._client.query.side_effect = RuntimeError("boom")
|
||||
assert vector._search("distance(vector, [0.1])", myscale_module.SortOrder.ASC, top_k=1) == []
|
||||
|
||||
@@ -33,7 +33,6 @@ from models.dataset import Dataset, DatasetCollectionBinding
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from qdrant_client.conversions import common_types
|
||||
from qdrant_client.http import models as rest
|
||||
|
||||
type DictFilter = dict[str, str | int | bool | dict | list]
|
||||
type MetadataFilter = DictFilter | common_types.Filter
|
||||
|
||||
-1
@@ -41,7 +41,6 @@ from models.enums import TidbAuthBindingStatus
|
||||
if TYPE_CHECKING:
|
||||
from qdrant_client import grpc # noqa
|
||||
from qdrant_client.conversions import common_types
|
||||
from qdrant_client.http import models as rest
|
||||
|
||||
type DictFilter = dict[str, str | int | bool | dict | list]
|
||||
type MetadataFilter = DictFilter | common_types.Filter
|
||||
|
||||
@@ -42,6 +42,7 @@ dependencies = [
|
||||
"opentelemetry-propagator-b3>=1.41.1,<2.0.0",
|
||||
"readabilipy==0.3.0",
|
||||
"resend>=2.27.0,<3.0.0",
|
||||
"zstandard==0.25.0",
|
||||
# Emerging: newer and fast-moving, use compatible pins
|
||||
"fastopenapi[flask]==0.7.0",
|
||||
"graphon==0.5.3",
|
||||
|
||||
@@ -64,6 +64,7 @@ from services.errors.account import (
|
||||
MemberNotInTenantError,
|
||||
NoPermissionError,
|
||||
RoleAlreadyAssignedError,
|
||||
SeatsLimitExceededError,
|
||||
TenantNotFoundError,
|
||||
)
|
||||
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
|
||||
@@ -433,6 +434,13 @@ class AccountService:
|
||||
|
||||
raise AccountNotFound()
|
||||
|
||||
# A licensed seat is one Account row, deployment-wide; joining an existing
|
||||
# account into another workspace does not pass through here and costs no seat.
|
||||
# is_authenticated=True: server-side enforcement needs the full license payload,
|
||||
# which the enterprise fill withholds from unauthenticated (browser-facing) calls.
|
||||
if not FeatureService.get_system_features(is_authenticated=True).license.seats.is_available():
|
||||
raise SeatsLimitExceededError("licensed seats limit exceeded")
|
||||
|
||||
if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email):
|
||||
raise AccountRegisterError(
|
||||
description=(
|
||||
@@ -1981,6 +1989,10 @@ class RegisterService:
|
||||
session.rollback()
|
||||
logger.exception("Register failed")
|
||||
raise AccountRegisterError("Workspace is not allowed to create.")
|
||||
except SeatsLimitExceededError:
|
||||
session.rollback()
|
||||
logger.exception("Register failed")
|
||||
raise
|
||||
except AccountRegisterError as are:
|
||||
session.rollback()
|
||||
logger.exception("Register failed")
|
||||
|
||||
@@ -50,9 +50,26 @@ class QuotaReleaseResult(TypedDict):
|
||||
released: int
|
||||
|
||||
|
||||
class QuotaBalanceResult(TypedDict):
|
||||
available: int
|
||||
reserved: int
|
||||
quota: int
|
||||
usage: int
|
||||
|
||||
|
||||
class QuotaConsumeCappedResult(TypedDict):
|
||||
deducted: int
|
||||
available: int
|
||||
reserved: int
|
||||
quota: int
|
||||
usage: int
|
||||
|
||||
|
||||
_quota_reserve_adapter = TypeAdapter(QuotaReserveResult)
|
||||
_quota_commit_adapter = TypeAdapter(QuotaCommitResult)
|
||||
_quota_release_adapter = TypeAdapter(QuotaReleaseResult)
|
||||
_quota_balance_adapter = TypeAdapter(QuotaBalanceResult)
|
||||
_quota_consume_capped_adapter = TypeAdapter(QuotaConsumeCappedResult)
|
||||
|
||||
|
||||
class _TenantFeatureQuota(TypedDict):
|
||||
@@ -176,6 +193,7 @@ class DismissNotificationDict(TypedDict):
|
||||
|
||||
class BillingService:
|
||||
base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
|
||||
quota_base_url = os.environ.get("BILLING_QUOTA_API_URL") or base_url
|
||||
secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
|
||||
|
||||
compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 4, 60)
|
||||
@@ -215,12 +233,18 @@ class BillingService:
|
||||
def get_quota_info(cls, tenant_id: str) -> TenantFeatureQuotaInfo:
|
||||
params = {"tenant_id": tenant_id}
|
||||
return _tenant_feature_quota_info_adapter.validate_python(
|
||||
cls._send_request("GET", "/quota/info", params=params)
|
||||
cls._send_quota_request("GET", "/quota/info", params=params)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def quota_reserve(
|
||||
cls, tenant_id: str, feature_key: str, request_id: str, amount: int = 1, meta: dict | None = None
|
||||
cls,
|
||||
tenant_id: str,
|
||||
feature_key: str,
|
||||
request_id: str,
|
||||
amount: int = 1,
|
||||
meta: dict | None = None,
|
||||
bucket: str = "",
|
||||
) -> QuotaReserveResult:
|
||||
"""Reserve quota before task execution."""
|
||||
payload: dict = {
|
||||
@@ -229,13 +253,21 @@ class BillingService:
|
||||
"request_id": request_id,
|
||||
"amount": amount,
|
||||
}
|
||||
if bucket:
|
||||
payload["bucket"] = bucket
|
||||
if meta:
|
||||
payload["meta"] = meta
|
||||
return _quota_reserve_adapter.validate_python(cls._send_request("POST", "/quota/reserve", json=payload))
|
||||
return _quota_reserve_adapter.validate_python(cls._send_quota_request("POST", "/quota/reserve", json=payload))
|
||||
|
||||
@classmethod
|
||||
def quota_commit(
|
||||
cls, tenant_id: str, feature_key: str, reservation_id: str, actual_amount: int, meta: dict | None = None
|
||||
cls,
|
||||
tenant_id: str,
|
||||
feature_key: str,
|
||||
reservation_id: str,
|
||||
actual_amount: int,
|
||||
meta: dict | None = None,
|
||||
bucket: str = "",
|
||||
) -> QuotaCommitResult:
|
||||
"""Commit a reservation with actual consumption."""
|
||||
payload: dict = {
|
||||
@@ -244,23 +276,57 @@ class BillingService:
|
||||
"reservation_id": reservation_id,
|
||||
"actual_amount": actual_amount,
|
||||
}
|
||||
if bucket:
|
||||
payload["bucket"] = bucket
|
||||
if meta:
|
||||
payload["meta"] = meta
|
||||
return _quota_commit_adapter.validate_python(cls._send_request("POST", "/quota/commit", json=payload))
|
||||
return _quota_commit_adapter.validate_python(cls._send_quota_request("POST", "/quota/commit", json=payload))
|
||||
|
||||
@classmethod
|
||||
def quota_release(cls, tenant_id: str, feature_key: str, reservation_id: str) -> QuotaReleaseResult:
|
||||
def quota_release(
|
||||
cls, tenant_id: str, feature_key: str, reservation_id: str, bucket: str = ""
|
||||
) -> QuotaReleaseResult:
|
||||
"""Release a reservation (cancel, return frozen quota)."""
|
||||
return _quota_release_adapter.validate_python(
|
||||
cls._send_request(
|
||||
"POST",
|
||||
"/quota/release",
|
||||
json={
|
||||
"tenant_id": tenant_id,
|
||||
"feature_key": feature_key,
|
||||
"reservation_id": reservation_id,
|
||||
},
|
||||
)
|
||||
payload = {
|
||||
"tenant_id": tenant_id,
|
||||
"feature_key": feature_key,
|
||||
"reservation_id": reservation_id,
|
||||
}
|
||||
if bucket:
|
||||
payload["bucket"] = bucket
|
||||
return _quota_release_adapter.validate_python(cls._send_quota_request("POST", "/quota/release", json=payload))
|
||||
|
||||
@classmethod
|
||||
def quota_get_balance(cls, tenant_id: str, feature_key: str, bucket: str = "") -> QuotaBalanceResult:
|
||||
"""Get quota balance for a feature bucket."""
|
||||
params = {"tenant_id": tenant_id, "feature_key": feature_key}
|
||||
if bucket:
|
||||
params["bucket"] = bucket
|
||||
return _quota_balance_adapter.validate_python(cls._send_quota_request("GET", "/quota/balance", params=params))
|
||||
|
||||
@classmethod
|
||||
def quota_consume_capped(
|
||||
cls,
|
||||
tenant_id: str,
|
||||
feature_key: str,
|
||||
request_id: str,
|
||||
amount: int,
|
||||
meta: dict | None = None,
|
||||
bucket: str = "",
|
||||
) -> QuotaConsumeCappedResult:
|
||||
"""Consume up to the available quota and return the actual deducted amount."""
|
||||
payload: dict = {
|
||||
"tenant_id": tenant_id,
|
||||
"feature_key": feature_key,
|
||||
"request_id": request_id,
|
||||
"amount": amount,
|
||||
}
|
||||
if bucket:
|
||||
payload["bucket"] = bucket
|
||||
if meta:
|
||||
payload["meta"] = meta
|
||||
return _quota_consume_capped_adapter.validate_python(
|
||||
cls._send_quota_request("POST", "/quota/consume-capped", json=payload)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -334,6 +400,12 @@ class BillingService:
|
||||
params = {"tenant_id": tenant_id, "feature_key": feature_key}
|
||||
return cls._send_request("GET", "/billing/tenant_feature_plan/usage", params=params)
|
||||
|
||||
@classmethod
|
||||
def _send_quota_request(
|
||||
cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None
|
||||
):
|
||||
return cls._send_request(method, endpoint, json=json, params=params, base_url=cls.quota_base_url)
|
||||
|
||||
@classmethod
|
||||
@retry(
|
||||
wait=wait_fixed(2),
|
||||
@@ -341,10 +413,17 @@ class BillingService:
|
||||
retry=retry_if_exception_type(httpx.RequestError),
|
||||
reraise=True,
|
||||
)
|
||||
def _send_request(cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None):
|
||||
def _send_request(
|
||||
cls,
|
||||
method: Literal["GET", "POST", "DELETE", "PUT"],
|
||||
endpoint: str,
|
||||
json=None,
|
||||
params=None,
|
||||
base_url: str | None = None,
|
||||
):
|
||||
headers = {"Content-Type": "application/json", "Billing-Api-Secret-Key": cls.secret_key}
|
||||
|
||||
url = f"{cls.base_url}{endpoint}"
|
||||
url = f"{base_url or cls.base_url}{endpoint}"
|
||||
response = _http_client.request(method, url, json=json, params=params, headers=headers, follow_redirects=True)
|
||||
if method == "GET" and response.status_code != httpx.codes.OK:
|
||||
raise ValueError("Unable to retrieve billing information. Please try again later or contact support.")
|
||||
|
||||
@@ -7,6 +7,8 @@ from piling up database transactions while preserving cross-tenant concurrency.
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -21,11 +23,37 @@ from models.enums import ProviderQuotaType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
FEATURE_KEY_CREDIT_POOL = "credit_pool"
|
||||
CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS = 10
|
||||
CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CreditPoolBalance:
|
||||
tenant_id: str
|
||||
pool_type: str
|
||||
quota_limit: int
|
||||
quota_used: int
|
||||
|
||||
@property
|
||||
def remaining_credits(self) -> int:
|
||||
if self.quota_limit == -1:
|
||||
return -1
|
||||
return max(0, self.quota_limit - self.quota_used)
|
||||
|
||||
def has_sufficient_credits(self, required_credits: int) -> bool:
|
||||
return self.quota_limit == -1 or self.remaining_credits >= required_credits
|
||||
|
||||
|
||||
class CreditPoolService:
|
||||
@staticmethod
|
||||
def _normalize_pool_type(pool_type: str | ProviderQuotaType) -> str:
|
||||
return pool_type.value if isinstance(pool_type, ProviderQuotaType) else str(pool_type)
|
||||
|
||||
@staticmethod
|
||||
def _use_billing_quota() -> bool:
|
||||
return bool(dify_config.BILLING_ENABLED)
|
||||
|
||||
@staticmethod
|
||||
def _get_tenant_lock_key(tenant_id: str) -> str:
|
||||
return f"credit_pool:tenant:{tenant_id}:deduct_lock"
|
||||
@@ -79,14 +107,32 @@ class CreditPoolService:
|
||||
return credit_pool
|
||||
|
||||
@classmethod
|
||||
def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> TenantCreditPool | None:
|
||||
def get_pool(
|
||||
cls, tenant_id: str, pool_type: str | ProviderQuotaType = "trial"
|
||||
) -> TenantCreditPool | CreditPoolBalance | None:
|
||||
"""get tenant credit pool"""
|
||||
normalized_pool_type = cls._normalize_pool_type(pool_type)
|
||||
if cls._use_billing_quota():
|
||||
from services.billing_service import BillingService
|
||||
|
||||
balance = BillingService.quota_get_balance(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket=normalized_pool_type,
|
||||
)
|
||||
return CreditPoolBalance(
|
||||
tenant_id=tenant_id,
|
||||
pool_type=normalized_pool_type,
|
||||
quota_limit=balance["quota"],
|
||||
quota_used=balance["usage"],
|
||||
)
|
||||
|
||||
with session_factory.get_session_maker().begin() as session:
|
||||
return session.scalar(
|
||||
select(TenantCreditPool)
|
||||
.where(
|
||||
TenantCreditPool.tenant_id == tenant_id,
|
||||
TenantCreditPool.pool_type == pool_type,
|
||||
TenantCreditPool.pool_type == normalized_pool_type,
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
@@ -102,7 +148,7 @@ class CreditPoolService:
|
||||
pool = cls.get_pool(tenant_id, pool_type)
|
||||
if not pool:
|
||||
return False
|
||||
return pool.remaining_credits >= credits_required
|
||||
return pool.has_sufficient_credits(credits_required)
|
||||
|
||||
@classmethod
|
||||
def check_and_deduct_credits(
|
||||
@@ -114,10 +160,54 @@ class CreditPoolService:
|
||||
"""Deduct exactly the requested credits or raise without mutating the pool."""
|
||||
if credits_required <= 0:
|
||||
return 0
|
||||
normalized_pool_type = cls._normalize_pool_type(pool_type)
|
||||
|
||||
if cls._use_billing_quota():
|
||||
from services.billing_service import BillingService
|
||||
|
||||
request_id = str(uuid4())
|
||||
result = BillingService.quota_reserve(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket=normalized_pool_type,
|
||||
request_id=request_id,
|
||||
amount=credits_required,
|
||||
meta={"source": "credit_pool.check_and_deduct"},
|
||||
)
|
||||
reservation_id = result.get("reservation_id", "")
|
||||
if not reservation_id:
|
||||
raise QuotaExceededError("Insufficient credits remaining")
|
||||
try:
|
||||
BillingService.quota_commit(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket=normalized_pool_type,
|
||||
reservation_id=reservation_id,
|
||||
actual_amount=credits_required,
|
||||
meta={"source": "credit_pool.check_and_deduct"},
|
||||
)
|
||||
except Exception:
|
||||
try:
|
||||
BillingService.quota_release(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket=normalized_pool_type,
|
||||
reservation_id=reservation_id,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to release reserved credit pool quota, tenant_id=%s, pool_type=%s, reservation_id=%s",
|
||||
tenant_id,
|
||||
normalized_pool_type,
|
||||
reservation_id,
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
return credits_required
|
||||
|
||||
def deduct() -> int:
|
||||
with session_factory.get_session_maker().begin() as session:
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type)
|
||||
if not pool:
|
||||
raise QuotaExceededError("Credit pool not found")
|
||||
|
||||
@@ -148,12 +238,26 @@ class CreditPoolService:
|
||||
"""Deduct up to the available balance and return the actual deducted credits."""
|
||||
if credits_required <= 0:
|
||||
return 0
|
||||
normalized_pool_type = cls._normalize_pool_type(pool_type)
|
||||
|
||||
if cls._use_billing_quota():
|
||||
from services.billing_service import BillingService
|
||||
|
||||
result = BillingService.quota_consume_capped(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket=normalized_pool_type,
|
||||
request_id=str(uuid4()),
|
||||
amount=credits_required,
|
||||
meta={"source": "credit_pool.deduct_capped"},
|
||||
)
|
||||
return result["deducted"]
|
||||
|
||||
def deduct() -> int:
|
||||
with session_factory.get_session_maker().begin() as session:
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type)
|
||||
if not pool:
|
||||
logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, pool_type)
|
||||
logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, normalized_pool_type)
|
||||
return 0
|
||||
|
||||
deducted_credits = min(credits_required, pool.remaining_credits)
|
||||
|
||||
@@ -12,6 +12,7 @@ from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from configs import dify_config
|
||||
from core.db.session_factory import session_factory
|
||||
from core.rbac import RBACResourceWhitelistScope
|
||||
from models import TenantAccountJoin, TenantAccountRole
|
||||
from services.enterprise.base import EnterpriseRequest
|
||||
|
||||
@@ -364,7 +365,6 @@ _LEGACY_WORKSPACE_ADMIN_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_WORKSPACE_EDITOR_KEYS: list[str] = [
|
||||
"workspace.member.manage",
|
||||
"api_extension.manage",
|
||||
"plugin.install",
|
||||
"credential.use",
|
||||
@@ -435,6 +435,7 @@ _LEGACY_APP_EDITOR_KEYS: list[str] = [
|
||||
"app.acl.delete",
|
||||
"app.acl.release_and_version",
|
||||
"app.acl.monitor",
|
||||
"app.acl.log_and_annotation",
|
||||
"app.acl.access_config",
|
||||
]
|
||||
|
||||
@@ -649,17 +650,18 @@ class ReplaceRoleBindings(_RBACModel):
|
||||
|
||||
|
||||
class ReplaceMemberBindings(_RBACModel):
|
||||
scope: str = "specific"
|
||||
scope: RBACResourceWhitelistScope = RBACResourceWhitelistScope.SPECIFIC
|
||||
|
||||
@field_validator("scope")
|
||||
@classmethod
|
||||
def _normalize_scope(cls, value: Any) -> str:
|
||||
def _normalize_scope(cls, value: Any) -> RBACResourceWhitelistScope:
|
||||
scope = str(value or "").strip().lower()
|
||||
if scope in {"", "specific"}:
|
||||
return "specific"
|
||||
if scope in {"all", "only_me"}:
|
||||
return scope
|
||||
raise ValueError(f"invalid scope: {value}")
|
||||
if scope == "":
|
||||
return RBACResourceWhitelistScope.SPECIFIC
|
||||
try:
|
||||
return RBACResourceWhitelistScope(scope)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"invalid scope: {value}") from exc
|
||||
|
||||
|
||||
class DeleteMemberBindings(_RBACModel):
|
||||
@@ -743,6 +745,7 @@ def _inner_call(
|
||||
account_id=account_id,
|
||||
json=json,
|
||||
params=params,
|
||||
timeout=dify_config.ENTERPRISE_RBAC_REQUEST_TIMEOUT,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -37,6 +37,10 @@ class AccountAlreadyInTenantError(BaseServiceError):
|
||||
pass
|
||||
|
||||
|
||||
class SeatsLimitExceededError(BaseServiceError):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidActionError(BaseServiceError):
|
||||
pass
|
||||
|
||||
|
||||
@@ -79,6 +79,7 @@ class LicenseModel(FeatureResponseModel):
|
||||
status: LicenseStatus = LicenseStatus.NONE
|
||||
expired_at: str = ""
|
||||
workspaces: LicenseLimitationModel = LicenseLimitationModel(enabled=False, size=0, limit=0)
|
||||
seats: LicenseLimitationModel = LicenseLimitationModel(enabled=False, size=0, limit=0)
|
||||
|
||||
|
||||
class BrandingModel(FeatureResponseModel):
|
||||
@@ -457,6 +458,11 @@ class FeatureService:
|
||||
features.license.workspaces.limit = workspaces_info.get("limit", 0)
|
||||
features.license.workspaces.size = workspaces_info.get("used", 0)
|
||||
|
||||
if seats_info := license_info.get("licensedSeats"):
|
||||
features.license.seats.enabled = seats_info.get("enabled", False)
|
||||
features.license.seats.limit = seats_info.get("limit", 0)
|
||||
features.license.seats.size = seats_info.get("used", 0)
|
||||
|
||||
if "PluginInstallationPermission" in enterprise_info:
|
||||
plugin_installation_info = enterprise_info["PluginInstallationPermission"]
|
||||
features.plugin_installation_permission.plugin_installation_scope = plugin_installation_info[
|
||||
|
||||
@@ -701,16 +701,29 @@ def _delete_records(query_sql: str, params: dict[str, Any], delete_func: Callabl
|
||||
if not rows:
|
||||
break
|
||||
|
||||
success_count = 0
|
||||
for i in rows:
|
||||
record_id = str(i.id)
|
||||
try:
|
||||
delete_func(session, record_id)
|
||||
logger.info(click.style(f"Deleted {name} {record_id}", fg="green"))
|
||||
session.commit()
|
||||
success_count += 1
|
||||
except Exception:
|
||||
logger.exception("Error occurred while deleting %s %s", name, record_id)
|
||||
# continue with next record even if one deletion fails
|
||||
session.rollback()
|
||||
break
|
||||
session.commit()
|
||||
continue
|
||||
|
||||
rs.close()
|
||||
|
||||
# If we couldn't delete ANY records in this batch, we must break out of the while loop
|
||||
# to prevent an infinite loop where we keep fetching the same failing records.
|
||||
if success_count == 0:
|
||||
logger.warning(
|
||||
click.style(
|
||||
f"Failed to delete any {name} in the current batch. Stopping to prevent infinite loop.",
|
||||
fg="yellow",
|
||||
)
|
||||
)
|
||||
break
|
||||
|
||||
@@ -17,6 +17,7 @@ from services.errors.account import (
|
||||
AccountPasswordError,
|
||||
AccountRegisterError,
|
||||
CurrentPasswordIncorrectError,
|
||||
SeatsLimitExceededError,
|
||||
TenantNotFoundError,
|
||||
)
|
||||
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
|
||||
@@ -477,6 +478,32 @@ class TestAccountService:
|
||||
session=db_session_with_containers,
|
||||
)
|
||||
|
||||
def test_create_account_seats_limit_exceeded(
|
||||
self, db_session_with_containers: Session, mock_external_service_dependencies
|
||||
):
|
||||
"""
|
||||
Test account creation when the licensed seats limit is exceeded.
|
||||
"""
|
||||
fake = Faker()
|
||||
email = fake.email()
|
||||
name = fake.name()
|
||||
password = generate_valid_password(fake)
|
||||
# Setup mocks
|
||||
mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
|
||||
mock_external_service_dependencies[
|
||||
"feature_service"
|
||||
].get_system_features.return_value.license.seats.is_available.return_value = False
|
||||
mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
|
||||
|
||||
with pytest.raises(SeatsLimitExceededError):
|
||||
AccountService.create_account(
|
||||
email=email,
|
||||
name=name,
|
||||
interface_language="en-US",
|
||||
password=password,
|
||||
session=db_session_with_containers,
|
||||
)
|
||||
|
||||
def test_link_account_integrate_new_provider(
|
||||
self, db_session_with_containers: Session, mock_external_service_dependencies
|
||||
):
|
||||
|
||||
@@ -336,6 +336,174 @@ def _insert_load_balancing_model_config(
|
||||
)
|
||||
|
||||
|
||||
def test_data_migrate_group_registers_dataset_permission_rbac_migration(command_module) -> None:
|
||||
command = command_module.data_migrate.commands["rbac-migrate-dataset-permissions"]
|
||||
|
||||
assert command is command_module.migrate_dataset_permissions_to_rbac
|
||||
assert "operator_account_id" not in {param.name for param in command.params}
|
||||
|
||||
|
||||
def test_dataset_permission_rbac_migration_help_mentions_binding_clear_side_effect(command_module) -> None:
|
||||
result = CliRunner().invoke(
|
||||
command_module.data_migrate,
|
||||
["rbac-migrate-dataset-permissions", "--help"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
normalized_output = " ".join(result.output.split())
|
||||
assert "clears existing per-user policy bindings" in normalized_output
|
||||
assert "recreates legacy partial-member default bindings" in normalized_output
|
||||
|
||||
|
||||
def test_dataset_permission_rbac_migration_maps_legacy_permissions_to_enum_scopes() -> None:
|
||||
rbac_module = importlib.import_module("commands.rbac")
|
||||
|
||||
assert (
|
||||
rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.ALL_TEAM)
|
||||
is rbac_module.RBACResourceWhitelistScope.ALL
|
||||
)
|
||||
assert (
|
||||
rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.PARTIAL_TEAM)
|
||||
is rbac_module.RBACResourceWhitelistScope.SPECIFIC
|
||||
)
|
||||
assert rbac_module._dataset_permission_enum("partial_members") is rbac_module.DatasetPermissionEnum.PARTIAL_TEAM
|
||||
assert rbac_module._dataset_permission_enum(None) is rbac_module.DatasetPermissionEnum.ONLY_ME
|
||||
|
||||
|
||||
def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator(
|
||||
command_module,
|
||||
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], [], []]
|
||||
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()
|
||||
|
||||
def fake_replace_whitelist(**kwargs):
|
||||
assert session_closed is True
|
||||
calls.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(rbac_module, "session_factory", FakeSessionFactory)
|
||||
monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist)
|
||||
|
||||
command_module.migrate_dataset_permissions_to_rbac.callback(
|
||||
tenant_id=None,
|
||||
dataset_id=None,
|
||||
batch_size=500,
|
||||
dry_run=False,
|
||||
)
|
||||
|
||||
assert calls[0]["tenant_id"] == "tenant-1"
|
||||
assert calls[0]["account_id"] == "creator-account-1"
|
||||
assert calls[0]["dataset_id"] == "dataset-1"
|
||||
assert calls[0]["payload"].scope is rbac_module.RBACResourceWhitelistScope.SPECIFIC
|
||||
|
||||
|
||||
def test_dataset_permission_rbac_migration_dry_run_outputs_structured_proposed_changes(
|
||||
command_module,
|
||||
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",
|
||||
)
|
||||
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)
|
||||
monkeypatch.setattr(
|
||||
rbac_module.RBACService.DatasetAccess,
|
||||
"replace_whitelist",
|
||||
lambda **kwargs: pytest.fail("dry-run must not replace whitelist"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rbac_module.RBACService.DatasetAccess,
|
||||
"replace_user_access_policies",
|
||||
lambda **kwargs: pytest.fail("dry-run must not replace user access policies"),
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
command_module.data_migrate,
|
||||
["rbac-migrate-dataset-permissions", "--dry-run"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
events = [json.loads(line) for line in result.output.splitlines() if line.startswith("{")]
|
||||
assert [event["action"] for event in events] == ["replace_whitelist", "replace_user_access_policies"]
|
||||
assert events[0]["before"] == {
|
||||
"legacy_dataset_permission": "partial_members",
|
||||
"legacy_partial_member_ids": ["member-account-1"],
|
||||
}
|
||||
assert events[0]["after"] == {"rbac_whitelist_scope": "specific"}
|
||||
assert events[0]["call"] == {
|
||||
"method": "RBACService.DatasetAccess.replace_whitelist",
|
||||
"kwargs": {
|
||||
"tenant_id": "tenant-1",
|
||||
"account_id": "creator-account-1",
|
||||
"dataset_id": "dataset-1",
|
||||
"payload": {"scope": "specific"},
|
||||
},
|
||||
}
|
||||
assert events[1]["target_account_id"] == "member-account-1"
|
||||
assert events[1]["after"] == {"rbac_user_access_policy_ids": ["default"]}
|
||||
assert events[1]["call"]["kwargs"]["payload"] == {"access_policy_ids": ["default"]}
|
||||
|
||||
|
||||
def test_data_migrate_command_defaults_output_to_stdout_stream(
|
||||
command_module,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, PropertyMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
from controllers.common import wraps as common_wraps
|
||||
from controllers.console import console_ns
|
||||
from controllers.console import wraps as console_wraps
|
||||
from controllers.console.app import ops_trace as ops_trace_module
|
||||
from controllers.console.app import wraps as app_wraps
|
||||
from libs import login as login_lib
|
||||
from models.account import Account, AccountStatus, TenantAccountRole
|
||||
|
||||
|
||||
def _make_account(role: TenantAccountRole) -> Account:
|
||||
account = Account(name="tester", email="tester@example.com")
|
||||
account.id = "account-123" # type: ignore[assignment]
|
||||
account.status = AccountStatus.ACTIVE
|
||||
account.role = role
|
||||
account._current_tenant = SimpleNamespace(id="tenant-123") # type: ignore[assignment]
|
||||
account._get_current_object = lambda: account # type: ignore[attr-defined]
|
||||
return account
|
||||
|
||||
|
||||
def _make_app() -> SimpleNamespace:
|
||||
return SimpleNamespace(id="app-123", tenant_id="tenant-123", status="normal", mode="chat")
|
||||
|
||||
|
||||
def _patch_console_guards(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
account: Account,
|
||||
app_model: SimpleNamespace,
|
||||
*,
|
||||
rbac_enabled: bool = False,
|
||||
) -> None:
|
||||
monkeypatch.setattr(login_lib.dify_config, "LOGIN_DISABLED", True)
|
||||
monkeypatch.setattr(login_lib.dify_config, "RBAC_ENABLED", rbac_enabled)
|
||||
monkeypatch.setattr(console_wraps.dify_config, "EDITION", "CLOUD")
|
||||
monkeypatch.setattr(login_lib, "current_user", account)
|
||||
monkeypatch.setattr(login_lib, "current_account_with_tenant", lambda: (account, account.current_tenant_id))
|
||||
monkeypatch.setattr(console_wraps, "current_account_with_tenant", lambda: (account, account.current_tenant_id))
|
||||
monkeypatch.setattr(common_wraps, "current_account_with_tenant", lambda: (account, account.current_tenant_id))
|
||||
monkeypatch.setattr(app_wraps, "_load_app_model_from_scoped_session", lambda _app_id: app_model)
|
||||
|
||||
|
||||
def _patch_payload(payload: dict[str, object] | None):
|
||||
if payload is None:
|
||||
return nullcontext()
|
||||
return patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method_name", "path", "payload", "service_method_name", "service_result"),
|
||||
[
|
||||
(
|
||||
"post",
|
||||
"/console/api/apps/app-123/trace-config",
|
||||
{"tracing_provider": "mlflow", "tracing_config": {"endpoint": "https://trace.example.com"}},
|
||||
"create_tracing_app_config",
|
||||
{"id": "trace-config-1"},
|
||||
),
|
||||
(
|
||||
"patch",
|
||||
"/console/api/apps/app-123/trace-config",
|
||||
{"tracing_provider": "mlflow", "tracing_config": {"endpoint": "https://trace.example.com"}},
|
||||
"update_tracing_app_config",
|
||||
True,
|
||||
),
|
||||
(
|
||||
"delete",
|
||||
"/console/api/apps/app-123/trace-config?tracing_provider=mlflow",
|
||||
None,
|
||||
"delete_tracing_app_config",
|
||||
True,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_trace_config_mutations_require_edit_permission(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method_name: str,
|
||||
path: str,
|
||||
payload: dict[str, object] | None,
|
||||
service_method_name: str,
|
||||
service_result: object,
|
||||
) -> None:
|
||||
app.config.setdefault("RESTX_MASK_HEADER", "X-Fields")
|
||||
account = _make_account(TenantAccountRole.NORMAL)
|
||||
_patch_console_guards(monkeypatch, account, _make_app())
|
||||
service_mock = MagicMock(return_value=service_result)
|
||||
monkeypatch.setattr(ops_trace_module.OpsService, service_method_name, service_mock)
|
||||
|
||||
with app.test_request_context(path, method=method_name.upper(), json=payload):
|
||||
with _patch_payload(payload):
|
||||
with pytest.raises(Forbidden):
|
||||
getattr(ops_trace_module.TraceAppConfigApi(), method_name)(app_id="app-123")
|
||||
|
||||
service_mock.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method_name", "path", "payload", "service_method_name", "service_result"),
|
||||
[
|
||||
(
|
||||
"post",
|
||||
"/console/api/apps/app-123/trace-config",
|
||||
{"tracing_provider": "mlflow", "tracing_config": {"endpoint": "https://trace.example.com"}},
|
||||
"create_tracing_app_config",
|
||||
{"id": "trace-config-1"},
|
||||
),
|
||||
(
|
||||
"patch",
|
||||
"/console/api/apps/app-123/trace-config",
|
||||
{"tracing_provider": "mlflow", "tracing_config": {"endpoint": "https://trace.example.com"}},
|
||||
"update_tracing_app_config",
|
||||
True,
|
||||
),
|
||||
(
|
||||
"delete",
|
||||
"/console/api/apps/app-123/trace-config?tracing_provider=mlflow",
|
||||
None,
|
||||
"delete_tracing_app_config",
|
||||
True,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_trace_config_mutations_require_rbac_permission(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method_name: str,
|
||||
path: str,
|
||||
payload: dict[str, object] | None,
|
||||
service_method_name: str,
|
||||
service_result: object,
|
||||
) -> None:
|
||||
app.config.setdefault("RESTX_MASK_HEADER", "X-Fields")
|
||||
account = _make_account(TenantAccountRole.NORMAL)
|
||||
_patch_console_guards(monkeypatch, account, _make_app(), rbac_enabled=True)
|
||||
monkeypatch.setattr(common_wraps.db, "session", SimpleNamespace(scalar=lambda _stmt: "other-account"))
|
||||
monkeypatch.setattr(common_wraps.RBACService.CheckAccess, "check", MagicMock(return_value=False))
|
||||
service_mock = MagicMock(return_value=service_result)
|
||||
monkeypatch.setattr(ops_trace_module.OpsService, service_method_name, service_mock)
|
||||
|
||||
with app.test_request_context(path, method=method_name.upper(), json=payload):
|
||||
with _patch_payload(payload):
|
||||
with pytest.raises(Forbidden):
|
||||
getattr(ops_trace_module.TraceAppConfigApi(), method_name)(app_id="app-123")
|
||||
|
||||
service_mock.assert_not_called()
|
||||
@@ -26,10 +26,11 @@ from controllers.console.auth.login import EmailCodeLoginApi, LoginApi, LogoutAp
|
||||
from controllers.console.error import (
|
||||
AccountBannedError,
|
||||
AccountInFreezeError,
|
||||
SeatsLimitExceeded,
|
||||
WorkspacesLimitExceeded,
|
||||
)
|
||||
from services.entities.auth_entities import LoginFailureReason
|
||||
from services.errors.account import AccountLoginError, AccountPasswordError
|
||||
from services.errors.account import AccountLoginError, AccountPasswordError, SeatsLimitExceededError
|
||||
|
||||
|
||||
def encode_password(password: str) -> str:
|
||||
@@ -487,6 +488,45 @@ class TestLoginApi:
|
||||
assert warn_records[0].args[0] == "user@example.com"
|
||||
assert warn_records[0].args[1] == LoginFailureReason.ACCOUNT_BANNED
|
||||
|
||||
@patch("controllers.console.wraps.db")
|
||||
@patch("controllers.console.auth.login.db")
|
||||
@patch("controllers.console.auth.login.AccountService.create_account_and_tenant")
|
||||
@patch("controllers.console.auth.login.AccountService.get_email_code_login_data")
|
||||
@patch("controllers.console.auth.login.AccountService.revoke_email_code_login_token")
|
||||
@patch("controllers.console.auth.login._get_account_with_case_fallback")
|
||||
def test_email_code_login_fails_when_seats_limit_exceeded(
|
||||
self,
|
||||
mock_get_account: MagicMock,
|
||||
mock_revoke_token: MagicMock,
|
||||
mock_get_token_data: MagicMock,
|
||||
mock_create_account: MagicMock,
|
||||
mock_login_db: MagicMock,
|
||||
mock_db: MagicMock,
|
||||
app: Flask,
|
||||
):
|
||||
"""
|
||||
Test email-code login failure when creating the account would exceed the licensed seats.
|
||||
|
||||
Verifies that:
|
||||
- the new-account path is taken when no account exists for the email
|
||||
- the service-layer SeatsLimitExceededError is translated to the SeatsLimitExceeded HTTP error
|
||||
"""
|
||||
# Arrange: valid token, no existing account -> account-creation path
|
||||
mock_get_token_data.return_value = {"email": "User@Example.com", "code": "123456"}
|
||||
mock_get_account.return_value = None
|
||||
mock_create_account.side_effect = SeatsLimitExceededError("licensed seats limit exceeded")
|
||||
|
||||
# Act & Assert
|
||||
with app.test_request_context(
|
||||
"/email-code-login/validity",
|
||||
method="POST",
|
||||
json={"email": "User@Example.com", "code": encode_code("123456"), "token": "token-123"},
|
||||
):
|
||||
with pytest.raises(SeatsLimitExceeded):
|
||||
EmailCodeLoginApi().post()
|
||||
|
||||
mock_create_account.assert_called_once()
|
||||
|
||||
|
||||
class TestLogoutApi:
|
||||
"""Test cases for the LogoutApi endpoint."""
|
||||
|
||||
@@ -136,7 +136,7 @@ class TestPydanticModels:
|
||||
|
||||
def test_resource_access_scope_defaults_empty_account_ids(self):
|
||||
parsed = rbac_mod._ResourceAccessScopeRequest.model_validate({"scope": "specific"})
|
||||
assert parsed.scope is rbac_mod._AccessScope.SPECIFIC
|
||||
assert parsed.scope is rbac_mod.RBACResourceWhitelistScope.SPECIFIC
|
||||
|
||||
def test_resource_access_scope_coerce_null_account_ids(self):
|
||||
rbac_mod._ResourceAccessScopeRequest.model_validate({"scope": "all"})
|
||||
|
||||
@@ -3,7 +3,7 @@ Unit tests for Service API wraps (authentication decorators)
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
@@ -469,7 +469,10 @@ 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")
|
||||
def test_rejects_over_rate_limit(self, mock_db, mock_get_rate_limit, mock_validate_token, app: Flask):
|
||||
@patch("controllers.service_api.wraps.sessionmaker")
|
||||
def test_rejects_over_rate_limit(
|
||||
self, mock_sessionmaker, mock_db, mock_get_rate_limit, mock_validate_token, app: Flask
|
||||
):
|
||||
"""Test that Forbidden is raised when over rate limit."""
|
||||
# Arrange
|
||||
mock_validate_token.return_value = Mock(tenant_id="tenant123")
|
||||
@@ -479,6 +482,10 @@ class TestCloudEditionBillingRateLimitCheck:
|
||||
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:
|
||||
mock_redis.zcard.return_value = 15 # Over limit
|
||||
@@ -492,6 +499,9 @@ 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()
|
||||
|
||||
|
||||
class TestValidateDatasetToken:
|
||||
|
||||
@@ -32,8 +32,16 @@ def mock_print_text(mocker: MockerFixture):
|
||||
return mocker.patch("core.callback_handler.workflow_tool_callback_handler.print_text")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def enable_debug(mocker: MockerFixture):
|
||||
"""Force DEBUG on so the handler emits its verbose stdout traces."""
|
||||
mocker.patch("core.callback_handler.workflow_tool_callback_handler.dify_config.DEBUG", True)
|
||||
|
||||
|
||||
class TestDifyWorkflowCallbackHandler:
|
||||
def test_on_tool_execution_single_output_success(self, handler: DifyWorkflowCallbackHandler, mock_print_text):
|
||||
def test_on_tool_execution_single_output_success(
|
||||
self, handler: DifyWorkflowCallbackHandler, mock_print_text, enable_debug
|
||||
):
|
||||
# Arrange
|
||||
tool_name = "test_tool"
|
||||
tool_inputs = {"a": 1}
|
||||
@@ -63,7 +71,9 @@ class TestDifyWorkflowCallbackHandler:
|
||||
]
|
||||
)
|
||||
|
||||
def test_on_tool_execution_multiple_outputs(self, handler: DifyWorkflowCallbackHandler, mock_print_text):
|
||||
def test_on_tool_execution_multiple_outputs(
|
||||
self, handler: DifyWorkflowCallbackHandler, mock_print_text, enable_debug
|
||||
):
|
||||
# Arrange
|
||||
tool_name = "multi_tool"
|
||||
outputs = [
|
||||
@@ -101,6 +111,29 @@ class TestDifyWorkflowCallbackHandler:
|
||||
assert results == []
|
||||
mock_print_text.assert_not_called()
|
||||
|
||||
def test_on_tool_execution_skips_print_when_debug_disabled(
|
||||
self, handler: DifyWorkflowCallbackHandler, mock_print_text, mocker: MockerFixture
|
||||
):
|
||||
"""When DEBUG is off, outputs are still yielded but nothing is printed
|
||||
and model_dump_json() is never invoked."""
|
||||
# Arrange
|
||||
mocker.patch("core.callback_handler.workflow_tool_callback_handler.dify_config.DEBUG", False)
|
||||
message = MagicMock()
|
||||
|
||||
# Act
|
||||
results = list(
|
||||
handler.on_tool_execution(
|
||||
tool_name="quiet_tool",
|
||||
tool_inputs={},
|
||||
tool_outputs=[message],
|
||||
)
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert results == [message]
|
||||
mock_print_text.assert_not_called()
|
||||
message.model_dump_json.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("invalid_outputs", "expected_exception"),
|
||||
[
|
||||
@@ -110,7 +143,7 @@ class TestDifyWorkflowCallbackHandler:
|
||||
],
|
||||
)
|
||||
def test_on_tool_execution_invalid_outputs_type(
|
||||
self, handler: DifyWorkflowCallbackHandler, invalid_outputs, expected_exception
|
||||
self, handler: DifyWorkflowCallbackHandler, invalid_outputs, expected_exception, enable_debug
|
||||
):
|
||||
# Arrange
|
||||
tool_name = "invalid_tool"
|
||||
@@ -125,7 +158,9 @@ class TestDifyWorkflowCallbackHandler:
|
||||
)
|
||||
)
|
||||
|
||||
def test_on_tool_execution_long_json_truncation(self, handler: DifyWorkflowCallbackHandler, mock_print_text):
|
||||
def test_on_tool_execution_long_json_truncation(
|
||||
self, handler: DifyWorkflowCallbackHandler, mock_print_text, enable_debug
|
||||
):
|
||||
# Arrange
|
||||
tool_name = "long_json_tool"
|
||||
long_json = "x" * 1500
|
||||
@@ -147,7 +182,9 @@ class TestDifyWorkflowCallbackHandler:
|
||||
color="blue",
|
||||
)
|
||||
|
||||
def test_on_tool_execution_model_dump_json_exception(self, handler: DifyWorkflowCallbackHandler, mock_print_text):
|
||||
def test_on_tool_execution_model_dump_json_exception(
|
||||
self, handler: DifyWorkflowCallbackHandler, mock_print_text, enable_debug
|
||||
):
|
||||
# Arrange
|
||||
tool_name = "exception_tool"
|
||||
bad_message = MagicMock()
|
||||
@@ -167,7 +204,7 @@ class TestDifyWorkflowCallbackHandler:
|
||||
assert mock_print_text.call_count >= 2
|
||||
|
||||
def test_on_tool_execution_none_message_id_and_trace_manager(
|
||||
self, handler: DifyWorkflowCallbackHandler, mock_print_text
|
||||
self, handler: DifyWorkflowCallbackHandler, mock_print_text, enable_debug
|
||||
):
|
||||
# Arrange
|
||||
tool_name = "optional_params_tool"
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
import json
|
||||
from base64 import b64decode
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
@@ -44,7 +43,7 @@ def test_serialize_inputs_encodes_payload() -> None:
|
||||
def test_transform_response_parses_json_result_and_converts_scientific_notation() -> None:
|
||||
response = '<<RESULT>>{"value": "1e+3", "nested": {"x": "2E-2"}, "arr": ["3e+1"]}<<RESULT>>'
|
||||
|
||||
result: Mapping[str, Any] = _DummyTransformer.transform_response(response)
|
||||
result: dict[str, Any] = _DummyTransformer.transform_response(response)
|
||||
|
||||
assert result == {"value": 1000.0, "nested": {"x": 0.02}, "arr": [30.0]}
|
||||
|
||||
|
||||
@@ -46,6 +46,18 @@ class TestUploadDSL:
|
||||
with pytest.raises(ValueError, match="claim_code"):
|
||||
upload_dsl(b"app: demo")
|
||||
|
||||
@patch("core.helper.creators.httpx.post")
|
||||
def test_raises_on_non_string_claim_code(self, mock_post):
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.json.return_value = {"data": {"claim_code": 123}}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
from core.helper.creators import upload_dsl
|
||||
|
||||
with pytest.raises(ValueError, match="claim_code"):
|
||||
upload_dsl(b"app: demo")
|
||||
|
||||
@patch("core.helper.creators.httpx.post")
|
||||
def test_raises_on_http_error(self, mock_post):
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
|
||||
@@ -114,10 +114,11 @@ def test_is_credential_exists_by_type(
|
||||
scalar_result: str | None,
|
||||
expected: bool,
|
||||
) -> None:
|
||||
mocker.patch("extensions.ext_database.db", new=SimpleNamespace(engine=object()))
|
||||
session_cls = mocker.patch("sqlalchemy.orm.Session")
|
||||
session = session_cls.return_value.__enter__.return_value
|
||||
session = mocker.MagicMock()
|
||||
session.scalar.return_value = scalar_result
|
||||
session_context = mocker.MagicMock()
|
||||
session_context.__enter__.return_value = session
|
||||
mocker.patch("core.db.session_factory.create_session", return_value=session_context)
|
||||
|
||||
result = is_credential_exists("cred-1", credential_type)
|
||||
|
||||
@@ -128,11 +129,33 @@ def test_is_credential_exists_by_type(
|
||||
def test_is_credential_exists_returns_false_for_unknown_type(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
mocker.patch("extensions.ext_database.db", new=SimpleNamespace(engine=object()))
|
||||
session_cls = mocker.patch("sqlalchemy.orm.Session")
|
||||
session = session_cls.return_value.__enter__.return_value
|
||||
session = mocker.MagicMock()
|
||||
session_context = mocker.MagicMock()
|
||||
session_context.__enter__.return_value = session
|
||||
mocker.patch("core.db.session_factory.create_session", return_value=session_context)
|
||||
|
||||
result = is_credential_exists("cred-1", cast(PluginCredentialType, "unknown"))
|
||||
|
||||
assert result is False
|
||||
session.scalar.assert_not_called()
|
||||
|
||||
|
||||
def test_is_credential_exists_uses_configured_session_factory_without_flask_app_context(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
class RaisingDB:
|
||||
@property
|
||||
def engine(self):
|
||||
raise RuntimeError("Working outside of application context.")
|
||||
|
||||
session = mocker.MagicMock()
|
||||
session.scalar.return_value = "model-credential"
|
||||
session_context = mocker.MagicMock()
|
||||
session_context.__enter__.return_value = session
|
||||
create_session = mocker.patch("core.db.session_factory.create_session", return_value=session_context)
|
||||
mocker.patch("extensions.ext_database.db", new=RaisingDB())
|
||||
|
||||
result = is_credential_exists("cred-1", PluginCredentialType.MODEL)
|
||||
|
||||
assert result is True
|
||||
create_session.assert_called_once_with()
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from core.helper.marketplace import (
|
||||
@@ -53,6 +54,16 @@ def test_batch_fetch_plugin_by_ids_returns_plugins_from_response(mocker: MockerF
|
||||
response.raise_for_status.assert_called_once()
|
||||
|
||||
|
||||
def test_batch_fetch_plugin_by_ids_rejects_invalid_plugins_response(mocker: MockerFixture) -> None:
|
||||
response = MagicMock()
|
||||
response.json.return_value = {"data": {"plugins": ["p1"]}}
|
||||
response.raise_for_status.return_value = None
|
||||
mocker.patch("core.helper.marketplace.httpx.post", return_value=response)
|
||||
|
||||
with pytest.raises(ValueError, match="plugins list"):
|
||||
batch_fetch_plugin_by_ids(["p1"])
|
||||
|
||||
|
||||
def test_batch_fetch_plugin_manifests_returns_empty_for_empty_input(mocker: MockerFixture) -> None:
|
||||
post_mock = mocker.patch("core.helper.marketplace.httpx.post")
|
||||
|
||||
|
||||
@@ -346,14 +346,32 @@ class TestPluginAppBackwardsInvocation:
|
||||
assert "end_users.app_id" in compiled
|
||||
assert stmt.compile().params == {"id_1": "uid", "tenant_id_1": "tenant-1", "app_id_1": "app-1"}
|
||||
|
||||
def test_get_user_returns_end_user_by_session_id(self, mocker: MockerFixture):
|
||||
session = self.patch_create_session(mocker, side_effect=[None, MagicMock(id="session-user")])
|
||||
app = SimpleNamespace(id="app-1", tenant_id="tenant-1")
|
||||
|
||||
user = PluginAppBackwardsInvocation._get_user("wecom-sender-1", app)
|
||||
|
||||
assert user.id == "session-user"
|
||||
stmt = session.scalar.call_args_list[1].args[0]
|
||||
compiled = str(stmt.compile(dialect=postgresql.dialect()))
|
||||
assert "end_users.session_id" in compiled
|
||||
assert "end_users.tenant_id" in compiled
|
||||
assert "end_users.app_id" in compiled
|
||||
assert stmt.compile().params == {
|
||||
"session_id_1": "wecom-sender-1",
|
||||
"tenant_id_1": "tenant-1",
|
||||
"app_id_1": "app-1",
|
||||
}
|
||||
|
||||
def test_get_user_falls_back_to_account_user(self, mocker: MockerFixture):
|
||||
session = self.patch_create_session(mocker, side_effect=[None, MagicMock(id="account-user")])
|
||||
session = self.patch_create_session(mocker, side_effect=[None, None, MagicMock(id="account-user")])
|
||||
app = SimpleNamespace(id="app-1", tenant_id="tenant-1")
|
||||
|
||||
user = PluginAppBackwardsInvocation._get_user("uid", app)
|
||||
|
||||
assert user.id == "account-user"
|
||||
stmt = session.scalar.call_args_list[1].args[0]
|
||||
stmt = session.scalar.call_args_list[2].args[0]
|
||||
compiled = str(stmt.compile(dialect=postgresql.dialect()))
|
||||
assert "accounts.id" in compiled
|
||||
assert "tenant_account_joins.account_id" in compiled
|
||||
@@ -361,12 +379,41 @@ class TestPluginAppBackwardsInvocation:
|
||||
assert stmt.compile().params == {"id_1": "uid", "tenant_id_1": "tenant-1"}
|
||||
|
||||
def test_get_user_raises_when_user_not_found(self, mocker: MockerFixture):
|
||||
self.patch_create_session(mocker, side_effect=[None, None])
|
||||
self.patch_create_session(mocker, side_effect=[None, None, None])
|
||||
app = SimpleNamespace(id="app-1", tenant_id="tenant-1")
|
||||
|
||||
with pytest.raises(ValueError, match="user not found"):
|
||||
PluginAppBackwardsInvocation._get_user("uid", app)
|
||||
|
||||
def test_invoke_app_creates_end_user_for_unknown_external_user_id(self, mocker: MockerFixture):
|
||||
app = MagicMock(mode=AppMode.WORKFLOW)
|
||||
end_user = MagicMock()
|
||||
workflow = MagicMock()
|
||||
mocker.patch.object(PluginAppBackwardsInvocation, "_get_app", return_value=app)
|
||||
mocker.patch.object(PluginAppBackwardsInvocation, "_get_workflow", return_value=workflow)
|
||||
mocker.patch.object(PluginAppBackwardsInvocation, "_get_user", side_effect=ValueError("user not found"))
|
||||
get_or_create = mocker.patch(
|
||||
"core.plugin.backwards_invocation.app.EndUserService.get_or_create_end_user",
|
||||
return_value=end_user,
|
||||
)
|
||||
route = mocker.patch.object(PluginAppBackwardsInvocation, "invoke_workflow_app", return_value={"ok": True})
|
||||
|
||||
result = PluginAppBackwardsInvocation.invoke_app(
|
||||
MagicMock(),
|
||||
app_id="app",
|
||||
user_id="wecom-sender-1",
|
||||
tenant_id="tenant",
|
||||
conversation_id="",
|
||||
query=None,
|
||||
stream=True,
|
||||
inputs={},
|
||||
files=[],
|
||||
)
|
||||
|
||||
assert result == {"ok": True}
|
||||
get_or_create.assert_called_once_with(app, user_id="wecom-sender-1")
|
||||
assert route.call_args.args[2] is end_user
|
||||
|
||||
def test_get_app_returns_app(self, mocker: MockerFixture):
|
||||
app_obj = MagicMock(id="app")
|
||||
self.patch_create_session(mocker, return_value=app_obj)
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
import datetime
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch, sentinel
|
||||
from unittest.mock import MagicMock, Mock, patch, sentinel
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -44,12 +44,34 @@ class _FakeRedis:
|
||||
def delete(self, key: str) -> None:
|
||||
self._values.pop(key, None)
|
||||
|
||||
def lock(
|
||||
self,
|
||||
key: str,
|
||||
*,
|
||||
timeout: int,
|
||||
sleep: float,
|
||||
) -> "_FakeRedisLock":
|
||||
return _FakeRedisLock(self, key)
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_plugin_model_provider_memory_cache() -> None:
|
||||
PluginService._plugin_model_providers_memory_cache.clear()
|
||||
yield
|
||||
PluginService._plugin_model_providers_memory_cache.clear()
|
||||
|
||||
class _FakeRedisLock:
|
||||
def __init__(self, redis: _FakeRedis, key: str) -> None:
|
||||
self._redis = redis
|
||||
self._key = key
|
||||
self._acquired = False
|
||||
|
||||
def acquire(self, *, blocking: bool = True, blocking_timeout: float | None = None) -> bool:
|
||||
if self._key in self._redis._values:
|
||||
return False
|
||||
|
||||
self._redis._values[self._key] = "locked"
|
||||
self._acquired = True
|
||||
return True
|
||||
|
||||
def release(self) -> None:
|
||||
if self._acquired:
|
||||
self._redis.delete(self._key)
|
||||
self._acquired = False
|
||||
|
||||
|
||||
def _build_model_schema() -> AIModelEntity:
|
||||
@@ -413,12 +435,13 @@ class TestPluginModelRuntime:
|
||||
"redis_client",
|
||||
SimpleNamespace(
|
||||
get=Mock(return_value=None),
|
||||
mget=Mock(return_value=[None, None]),
|
||||
mget=Mock(return_value=[None]),
|
||||
delete=Mock(),
|
||||
setex=Mock(),
|
||||
lock=Mock(return_value=MagicMock()),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(plugin_service_module.dify_config, "PLUGIN_MODEL_PROVIDERS_CACHE_TTL", 300)
|
||||
monkeypatch.setattr(plugin_service_module.dify_config, "PLUGIN_MODEL_PROVIDERS_CACHE_TTL", 0)
|
||||
runtime = PluginModelRuntime(tenant_id="tenant", user_id="user", client=client, plugin_service=PluginService)
|
||||
|
||||
runtime.fetch_model_providers()
|
||||
|
||||
@@ -3888,7 +3888,17 @@ class TestDatasetRetrievalAdditionalHelpers:
|
||||
trace_manager.add_trace_task.assert_not_called()
|
||||
|
||||
def test_on_query(self, retrieval: DatasetRetrieval) -> None:
|
||||
with patch("core.rag.retrieval.dataset_retrieval.db.session") as mock_session:
|
||||
db_mock = Mock()
|
||||
audit_session = MagicMock()
|
||||
session_factory = MagicMock()
|
||||
session_factory.begin.return_value.__enter__.return_value = audit_session
|
||||
|
||||
with (
|
||||
patch("core.rag.retrieval.dataset_retrieval.db", db_mock),
|
||||
patch(
|
||||
"core.rag.retrieval.dataset_retrieval.sessionmaker", return_value=session_factory
|
||||
) as sessionmaker_mock,
|
||||
):
|
||||
retrieval._on_query(
|
||||
query=None,
|
||||
attachment_ids=None,
|
||||
@@ -3897,7 +3907,7 @@ class TestDatasetRetrievalAdditionalHelpers:
|
||||
user_from="account",
|
||||
user_id="u1",
|
||||
)
|
||||
mock_session.add_all.assert_not_called()
|
||||
audit_session.add_all.assert_not_called()
|
||||
|
||||
retrieval._on_query(
|
||||
query="python",
|
||||
@@ -3907,11 +3917,22 @@ class TestDatasetRetrievalAdditionalHelpers:
|
||||
user_from="account",
|
||||
user_id="u1",
|
||||
)
|
||||
mock_session.add_all.assert_called()
|
||||
mock_session.commit.assert_called()
|
||||
sessionmaker_mock.assert_called_once_with(bind=db_mock.engine, expire_on_commit=False)
|
||||
audit_session.add_all.assert_called_once()
|
||||
added_queries = audit_session.add_all.call_args.args[0]
|
||||
assert len(added_queries) == 2
|
||||
db_mock.session.commit.assert_not_called()
|
||||
|
||||
def test_on_query_normalizes_workflow_end_user_role(self, retrieval: DatasetRetrieval) -> None:
|
||||
with patch("core.rag.retrieval.dataset_retrieval.db.session") as mock_session:
|
||||
db_mock = Mock()
|
||||
audit_session = MagicMock()
|
||||
session_factory = MagicMock()
|
||||
session_factory.begin.return_value.__enter__.return_value = audit_session
|
||||
|
||||
with (
|
||||
patch("core.rag.retrieval.dataset_retrieval.db", db_mock),
|
||||
patch("core.rag.retrieval.dataset_retrieval.sessionmaker", return_value=session_factory),
|
||||
):
|
||||
retrieval._on_query(
|
||||
query="python",
|
||||
attachment_ids=None,
|
||||
@@ -3921,12 +3942,11 @@ class TestDatasetRetrievalAdditionalHelpers:
|
||||
user_id="u1",
|
||||
)
|
||||
|
||||
mock_session.add_all.assert_called_once()
|
||||
added_queries = mock_session.add_all.call_args.args[0]
|
||||
audit_session.add_all.assert_called_once()
|
||||
added_queries = audit_session.add_all.call_args.args[0]
|
||||
|
||||
assert len(added_queries) == 1
|
||||
assert added_queries[0].created_by_role == CreatorUserRole.END_USER
|
||||
mock_session.commit.assert_called_once()
|
||||
|
||||
def test_handle_invoke_result(self, retrieval: DatasetRetrieval) -> None:
|
||||
usage = LLMUsage.empty_usage()
|
||||
|
||||
@@ -84,6 +84,10 @@ def test_assembling_request_auth_header_assembly():
|
||||
assert headers["Authorization"] == "Bearer abc"
|
||||
|
||||
tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_header_prefix": "basic", "api_key_value": "abc"}
|
||||
headers = tool.assembling_request(parameters={})
|
||||
assert headers["Authorization"] == "Basic abc"
|
||||
assert tool.runtime.credentials["api_key_value"] == "abc"
|
||||
|
||||
headers = tool.assembling_request(parameters={})
|
||||
assert headers["Authorization"] == "Basic abc"
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
from collections.abc import Generator
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -131,18 +131,25 @@ def test_create_message_files_and_invoke_generator():
|
||||
created.append(obj)
|
||||
return obj
|
||||
|
||||
with patch("core.tools.tool_engine.MessageFile", side_effect=_message_file_factory):
|
||||
with patch("core.tools.tool_engine.db") as mock_db:
|
||||
ids = ToolEngine._create_message_files(
|
||||
tool_messages=binaries,
|
||||
agent_message=SimpleNamespace(id="msg-1"),
|
||||
invoke_from=InvokeFrom.DEBUGGER,
|
||||
user_id="user-1",
|
||||
)
|
||||
file_session = MagicMock()
|
||||
session_factory = MagicMock()
|
||||
session_factory.begin.return_value.__enter__.return_value = file_session
|
||||
with (
|
||||
patch("core.tools.tool_engine.MessageFile", side_effect=_message_file_factory),
|
||||
patch("core.tools.tool_engine.db") as mock_db,
|
||||
patch("core.tools.tool_engine.sessionmaker", return_value=session_factory) as mock_sessionmaker,
|
||||
):
|
||||
ids = ToolEngine._create_message_files(
|
||||
tool_messages=binaries,
|
||||
agent_message=SimpleNamespace(id="msg-1"),
|
||||
invoke_from=InvokeFrom.DEBUGGER,
|
||||
user_id="user-1",
|
||||
)
|
||||
|
||||
assert ids == ["mf-1", "mf-2"]
|
||||
assert mock_db.session.add.call_count == 2
|
||||
mock_db.session.close.assert_called_once()
|
||||
mock_sessionmaker.assert_called_once_with(bind=mock_db.engine, expire_on_commit=False)
|
||||
assert file_session.add.call_count == 2
|
||||
mock_db.session.close.assert_not_called()
|
||||
|
||||
tool = _build_tool()
|
||||
invoked = list(ToolEngine._invoke(tool, {"a": 1}, user_id="u"))
|
||||
|
||||
@@ -2,6 +2,9 @@ import ast
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.engine.default import DefaultDialect
|
||||
|
||||
from models.enums import EndUserType
|
||||
from models.model import EndUser
|
||||
from models.types import EnumText
|
||||
@@ -24,6 +27,24 @@ def test_end_user_type_is_plain_persisted_value_enum():
|
||||
assert not hasattr(EndUserType, "from_invoke_from")
|
||||
|
||||
|
||||
def test_end_user_type_accepts_legacy_service_api_value():
|
||||
# Legacy rows stored the service-api type with an underscore before it was
|
||||
# normalized to the hyphenated value; loading them must not raise. See #38201.
|
||||
assert EndUserType("service_api") is EndUserType.SERVICE_API
|
||||
assert EndUserType("service-api") is EndUserType.SERVICE_API
|
||||
|
||||
|
||||
def test_end_user_type_still_rejects_unknown_values():
|
||||
with pytest.raises(ValueError):
|
||||
EndUserType("not-a-real-end-user-type")
|
||||
|
||||
|
||||
def test_enum_text_deserializes_legacy_service_api_value():
|
||||
# Exercises the actual column load path that raised for unmigrated rows.
|
||||
column_type = EnumText(EndUserType)
|
||||
assert column_type.process_result_value("service_api", DefaultDialect()) is EndUserType.SERVICE_API
|
||||
|
||||
|
||||
def test_end_user_service_creation_methods_accept_end_user_type():
|
||||
assert inspect.signature(EndUserService.get_or_create_end_user_by_type).parameters["type"].annotation is EndUserType
|
||||
assert inspect.signature(EndUserService.create_end_user_batch).parameters["type"].annotation is EndUserType
|
||||
|
||||
@@ -633,6 +633,8 @@ class TestMyPermissions:
|
||||
assert "dataset.acl.preview" in out.workspace.permission_keys
|
||||
assert "app.acl.preview" in out.app.default_permission_keys
|
||||
assert "dataset.acl.preview" in out.dataset.default_permission_keys
|
||||
if role == "editor":
|
||||
assert "app.acl.log_and_annotation" in out.app.default_permission_keys
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("role", "expected_snippet_keys"),
|
||||
|
||||
@@ -4,6 +4,7 @@ from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, call, patch
|
||||
|
||||
import pytest
|
||||
import zstandard
|
||||
from pydantic import TypeAdapter
|
||||
from redis import RedisError
|
||||
|
||||
@@ -14,15 +15,6 @@ from graphon.model_runtime.entities.provider_entities import ConfigurateMethod,
|
||||
MODULE = "core.plugin.plugin_service"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_plugin_model_provider_memory_cache() -> None:
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
PluginService._plugin_model_providers_memory_cache.clear()
|
||||
yield
|
||||
PluginService._plugin_model_providers_memory_cache.clear()
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self) -> None:
|
||||
self.execute = Mock()
|
||||
@@ -137,17 +129,71 @@ class TestFetchLatestPluginVersion:
|
||||
|
||||
|
||||
class TestPluginModelProviderCache:
|
||||
def test_fetch_plugin_model_providers_returns_cached_provider_without_calling_daemon(self) -> None:
|
||||
"""A valid tenant cache entry is reused across runtime calls without plugin daemon access."""
|
||||
def test_store_cached_plugin_model_providers_compresses_large_payload(self) -> None:
|
||||
"""Large provider metadata payloads are compressed before being stored in Redis."""
|
||||
large_provider = _build_provider_entity()
|
||||
large_provider.label = I18nObject(en_US="OpenAI " * 10000)
|
||||
raw_payload = TypeAdapter(list[ProviderEntity]).dump_json([large_provider])
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.dify_config") as mock_config,
|
||||
):
|
||||
mock_config.PLUGIN_MODEL_PROVIDERS_CACHE_TTL = 86400
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
PluginService._store_cached_plugin_model_providers("tenant-1", 0, [large_provider])
|
||||
|
||||
redis_client.setex.assert_called_once()
|
||||
stored_key, ttl, stored_payload = redis_client.setex.call_args.args
|
||||
assert stored_key == cache_key
|
||||
assert ttl == 86400
|
||||
assert isinstance(stored_payload, bytes)
|
||||
prefix = PluginService.PLUGIN_MODEL_PROVIDERS_CACHE_COMPRESSION_PREFIX
|
||||
assert stored_payload.startswith(prefix)
|
||||
assert len(stored_payload) < len(raw_payload)
|
||||
assert zstandard.decompress(stored_payload[len(prefix) :]) == raw_payload
|
||||
|
||||
def test_fetch_plugin_model_providers_reads_compressed_cached_provider_without_calling_daemon(self) -> None:
|
||||
"""Compressed tenant cache entries are decoded before provider schema validation."""
|
||||
cached_provider = _build_provider_entity()
|
||||
cached_payload = TypeAdapter(list[ProviderEntity]).dump_json([cached_provider]).decode("utf-8")
|
||||
cached_provider.label = I18nObject(en_US="OpenAI " * 10000)
|
||||
cached_payload = TypeAdapter(list[ProviderEntity]).dump_json([cached_provider])
|
||||
generation_key = _provider_generation_key("tenant-1")
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
legacy_cache_key = _provider_cache_key("tenant-1")
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
compressed_payload = PluginService.PLUGIN_MODEL_PROVIDERS_CACHE_COMPRESSION_PREFIX + zstandard.compress(
|
||||
cached_payload, level=1
|
||||
)
|
||||
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [cached_payload, None]
|
||||
redis_client.mget.return_value = [compressed_payload]
|
||||
client = Mock()
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
assert result[0].label.en_us == "OpenAI " * 10000
|
||||
client.fetch_model_providers.assert_not_called()
|
||||
redis_client.setex.assert_not_called()
|
||||
redis_client.get.assert_called_once_with(generation_key)
|
||||
redis_client.mget.assert_called_once_with([cache_key])
|
||||
|
||||
def test_fetch_plugin_model_providers_returns_cached_provider_without_calling_daemon(self) -> None:
|
||||
"""A valid tenant cache entry is reused across runtime calls without plugin daemon access."""
|
||||
cached_provider = _build_provider_entity()
|
||||
cached_payload = TypeAdapter(list[ProviderEntity]).dump_json([cached_provider])
|
||||
generation_key = _provider_generation_key("tenant-1")
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [cached_payload]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
@@ -158,19 +204,18 @@ class TestPluginModelProviderCache:
|
||||
client.fetch_model_providers.assert_not_called()
|
||||
redis_client.setex.assert_not_called()
|
||||
redis_client.get.assert_called_once_with(generation_key)
|
||||
redis_client.mget.assert_called_once_with([cache_key, legacy_cache_key])
|
||||
redis_client.mget.assert_called_once_with([cache_key])
|
||||
|
||||
def test_fetch_plugin_model_providers_deletes_invalid_cache_and_refetches(self) -> None:
|
||||
"""Invalid generation-scoped cache payloads are removed before falling back to the daemon."""
|
||||
generation_key = _provider_generation_key("tenant-1")
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
legacy_cache_key = _provider_cache_key("tenant-1")
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.dify_config") as mock_config,
|
||||
):
|
||||
redis_client.get.side_effect = [None, None]
|
||||
redis_client.mget.return_value = ["not-json", None]
|
||||
redis_client.get.side_effect = [None, None, None]
|
||||
redis_client.mget.side_effect = [["not-json"], [None]]
|
||||
mock_config.PLUGIN_MODEL_PROVIDERS_CACHE_TTL = 86400
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
@@ -184,8 +229,11 @@ class TestPluginModelProviderCache:
|
||||
assert redis_client.setex.call_args.args[0] == cache_key
|
||||
assert redis_client.setex.call_args.args[1] == 86400
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
redis_client.get.assert_has_calls([call(generation_key), call(generation_key)])
|
||||
redis_client.mget.assert_called_once_with([cache_key, legacy_cache_key])
|
||||
redis_client.get.assert_has_calls([call(generation_key), call(generation_key), call(generation_key)])
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([cache_key]),
|
||||
call([cache_key]),
|
||||
]
|
||||
|
||||
def test_fetch_plugin_model_providers_refetches_when_cache_read_fails(self) -> None:
|
||||
"""Redis read failures do not block provider discovery for the tenant."""
|
||||
@@ -204,7 +252,6 @@ class TestPluginModelProviderCache:
|
||||
def test_fetch_plugin_model_providers_refetches_when_cached_payload_batch_read_fails(self) -> None:
|
||||
"""Redis mget failures do not block provider discovery for the tenant."""
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
legacy_cache_key = _provider_cache_key("tenant-1")
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.side_effect = RedisError("redis unavailable")
|
||||
@@ -216,14 +263,14 @@ class TestPluginModelProviderCache:
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.mget.assert_called_once_with([cache_key, legacy_cache_key])
|
||||
redis_client.mget.assert_called_once_with([cache_key])
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_returns_fresh_result_when_cache_write_fails(self) -> None:
|
||||
"""Redis write failures are non-fatal after fresh provider data has been fetched."""
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None, None]
|
||||
redis_client.mget.return_value = [None]
|
||||
redis_client.setex.side_effect = RedisError("redis unavailable")
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
@@ -235,6 +282,352 @@ class TestPluginModelProviderCache:
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_waits_for_concurrent_refresh_cache_fill(self) -> None:
|
||||
"""A cache miss waits for the active tenant refresh instead of stampeding the daemon."""
|
||||
cached_provider = _build_provider_entity()
|
||||
cached_payload = TypeAdapter(list[ProviderEntity]).dump_json([cached_provider])
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.monotonic", return_value=100.0),
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.side_effect = [[None], [cached_payload]]
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider(provider="anthropic")]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
redis_client.lock.assert_called_once_with(
|
||||
PluginService._get_plugin_model_providers_lock_key("tenant-1", 0),
|
||||
timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_TTL,
|
||||
sleep=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL,
|
||||
)
|
||||
redis_client.lock.return_value.acquire.assert_called_once_with(
|
||||
blocking=True,
|
||||
blocking_timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT,
|
||||
)
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([cache_key]),
|
||||
call([cache_key]),
|
||||
]
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
client.fetch_model_providers.assert_not_called()
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_falls_back_when_refresh_lock_wait_times_out(self) -> None:
|
||||
"""A request should stop waiting and fetch directly instead of surfacing lock contention."""
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT", 0),
|
||||
patch(f"{MODULE}.time.monotonic", return_value=100.0),
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None]
|
||||
redis_client.lock.return_value.acquire.return_value = False
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
redis_client.lock.assert_called_once_with(
|
||||
PluginService._get_plugin_model_providers_lock_key("tenant-1", 0),
|
||||
timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_TTL,
|
||||
sleep=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL,
|
||||
)
|
||||
redis_client.lock.return_value.acquire.assert_called_once_with(blocking=True, blocking_timeout=0)
|
||||
redis_client.lock.return_value.release.assert_not_called()
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.setex.assert_called_once()
|
||||
assert redis_client.setex.call_args.args[0] == cache_key
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_restarts_lock_path_after_generation_changes(self) -> None:
|
||||
"""Waiters re-read provider generation before trying to become the next refresh owner."""
|
||||
generation_key = _provider_generation_key("tenant-1")
|
||||
stale_cache_key = _provider_cache_key("tenant-1", 0)
|
||||
new_cache_key = _provider_cache_key("tenant-1", 1)
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.monotonic", side_effect=[100.0, 100.0, 100.5, 101.0]),
|
||||
):
|
||||
redis_client.get.side_effect = [None, b"1", b"1", b"1", b"1"]
|
||||
redis_client.mget.side_effect = [[None], [None], [None], [None]]
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider(provider="anthropic")]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert redis_client.get.call_args_list == [
|
||||
call(generation_key),
|
||||
call(generation_key),
|
||||
call(generation_key),
|
||||
call(generation_key),
|
||||
call(generation_key),
|
||||
]
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([stale_cache_key]),
|
||||
call([new_cache_key]),
|
||||
call([new_cache_key]),
|
||||
call([new_cache_key]),
|
||||
]
|
||||
assert redis_client.lock.call_args_list == [
|
||||
call(
|
||||
PluginService._get_plugin_model_providers_lock_key("tenant-1", 0),
|
||||
timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_TTL,
|
||||
sleep=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL,
|
||||
),
|
||||
call(
|
||||
PluginService._get_plugin_model_providers_lock_key("tenant-1", 1),
|
||||
timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_TTL,
|
||||
sleep=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL,
|
||||
),
|
||||
]
|
||||
assert redis_client.lock.return_value.acquire.call_args_list == [
|
||||
call(blocking=True, blocking_timeout=2.0),
|
||||
call(blocking=True, blocking_timeout=1.5),
|
||||
]
|
||||
assert redis_client.lock.return_value.release.call_count == 2
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.setex.assert_called_once()
|
||||
assert redis_client.setex.call_args.args[0] == new_cache_key
|
||||
assert [provider.provider for provider in result] == ["langgenius/anthropic/anthropic"]
|
||||
|
||||
def test_fetch_plugin_model_providers_falls_back_when_generation_retries_exhaust_wait_budget(self) -> None:
|
||||
"""Generation retry loops share one request-local lock wait deadline before direct fetch fallback."""
|
||||
generation_key = _provider_generation_key("tenant-1")
|
||||
stale_cache_key = _provider_cache_key("tenant-1", 0)
|
||||
new_cache_key = _provider_cache_key("tenant-1", 1)
|
||||
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.monotonic", side_effect=[100.0, 100.0, 102.1]),
|
||||
):
|
||||
redis_client.get.side_effect = [None, b"1", b"1", b"1"]
|
||||
redis_client.mget.side_effect = [[None], [None], [None]]
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider(provider="anthropic")]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert redis_client.get.call_args_list == [
|
||||
call(generation_key),
|
||||
call(generation_key),
|
||||
call(generation_key),
|
||||
call(generation_key),
|
||||
]
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([stale_cache_key]),
|
||||
call([new_cache_key]),
|
||||
call([new_cache_key]),
|
||||
]
|
||||
redis_client.lock.assert_called_once_with(
|
||||
PluginService._get_plugin_model_providers_lock_key("tenant-1", 0),
|
||||
timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_TTL,
|
||||
sleep=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL,
|
||||
)
|
||||
redis_client.lock.return_value.acquire.assert_called_once_with(blocking=True, blocking_timeout=2.0)
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.setex.assert_called_once()
|
||||
assert redis_client.setex.call_args.args[0] == new_cache_key
|
||||
assert [provider.provider for provider in result] == ["langgenius/anthropic/anthropic"]
|
||||
|
||||
def test_fetch_plugin_model_providers_releases_owned_refresh_lock_after_store(self) -> None:
|
||||
"""The refresh owner releases only its token after storing provider metadata."""
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.monotonic", return_value=100.0),
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None]
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
redis_client.lock.assert_called_once_with(
|
||||
PluginService._get_plugin_model_providers_lock_key("tenant-1", 0),
|
||||
timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_TTL,
|
||||
sleep=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL,
|
||||
)
|
||||
redis_client.lock.return_value.acquire.assert_called_once_with(
|
||||
blocking=True,
|
||||
blocking_timeout=PluginService.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT,
|
||||
)
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([cache_key]),
|
||||
call([cache_key]),
|
||||
]
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
redis_client.eval.assert_not_called()
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_returns_fresh_result_when_refresh_lock_release_fails(self) -> None:
|
||||
"""Release failures are logged, not allowed to hide a successful daemon refresh."""
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.monotonic", return_value=100.0),
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None]
|
||||
redis_client.lock.return_value.release.side_effect = RedisError("release failed")
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([cache_key]),
|
||||
call([cache_key]),
|
||||
]
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.setex.assert_called_once()
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_releases_owned_refresh_lock_when_fetch_fails(self) -> None:
|
||||
"""Release failures must not hide the daemon failure that happened while owning the lock."""
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.monotonic", return_value=100.0),
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None]
|
||||
redis_client.lock.return_value.release.side_effect = RedisError("release failed")
|
||||
client = Mock()
|
||||
client.fetch_model_providers.side_effect = RuntimeError("daemon failed")
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
with pytest.raises(RuntimeError, match="daemon failed"):
|
||||
PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([cache_key]),
|
||||
call([cache_key]),
|
||||
]
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.setex.assert_not_called()
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
|
||||
def test_fetch_plugin_model_providers_falls_back_when_refresh_lock_acquire_fails(self) -> None:
|
||||
"""Redis acquire failures degrade to a direct daemon fetch instead of hiding provider data."""
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.monotonic", return_value=100.0),
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None]
|
||||
redis_client.lock.return_value.acquire.side_effect = RedisError("redis unavailable")
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
redis_client.lock.return_value.acquire.assert_called_once()
|
||||
redis_client.lock.return_value.release.assert_not_called()
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_skips_wait_when_refresh_lock_fails(self) -> None:
|
||||
"""Lock API failures should fall back directly instead of adding timeout latency."""
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.time.sleep") as sleep,
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None]
|
||||
redis_client.lock.side_effect = RedisError("redis unavailable")
|
||||
redis_client.set.side_effect = AssertionError("raw redis set must not be used for refresh locks")
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
sleep.assert_not_called()
|
||||
redis_client.lock.assert_called_once()
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_caches_empty_provider_list(self) -> None:
|
||||
"""An empty provider list is still a valid refresh result for single-flight waiters."""
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None]
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = []
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert result == ()
|
||||
redis_client.setex.assert_called_once()
|
||||
assert redis_client.setex.call_args.args[0] == cache_key
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
|
||||
def test_fetch_plugin_model_providers_skips_cache_write_when_generation_changes_during_refresh(self) -> None:
|
||||
"""A refresh that started before invalidation must not populate the newer generation cache."""
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.side_effect = [None, None, "1"]
|
||||
redis_client.mget.return_value = [None]
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = []
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert result == ()
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.setex.assert_not_called()
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
|
||||
def test_fetch_plugin_model_providers_reuses_cached_empty_provider_list(self) -> None:
|
||||
"""A cached empty list should prevent repeated daemon fetches for tenants without plugin models."""
|
||||
empty_payload = TypeAdapter(list[ProviderEntity]).dump_json([])
|
||||
cache_key = _provider_cache_key("tenant-1", 0)
|
||||
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [empty_payload]
|
||||
client = Mock()
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
assert result == ()
|
||||
redis_client.mget.assert_called_once_with([cache_key])
|
||||
client.fetch_model_providers.assert_not_called()
|
||||
|
||||
def test_fetch_plugin_model_providers_creates_default_client_on_cache_miss(self) -> None:
|
||||
"""The service owns plugin daemon access when no runtime-provided client is injected."""
|
||||
with (
|
||||
@@ -242,7 +635,7 @@ class TestPluginModelProviderCache:
|
||||
patch(f"{MODULE}.PluginModelClient") as client_cls,
|
||||
):
|
||||
redis_client.get.return_value = None
|
||||
redis_client.mget.return_value = [None, None]
|
||||
redis_client.mget.return_value = [None]
|
||||
client = client_cls.return_value
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
|
||||
@@ -254,35 +647,6 @@ class TestPluginModelProviderCache:
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
assert [provider.provider for provider in result] == ["langgenius/openai/openai"]
|
||||
|
||||
def test_fetch_plugin_model_providers_reuses_process_local_cache(self) -> None:
|
||||
generation_key = _provider_generation_key("tenant-1")
|
||||
with (
|
||||
patch(f"{MODULE}.redis_client") as redis_client,
|
||||
patch(f"{MODULE}.PluginModelClient") as client_cls,
|
||||
):
|
||||
redis_client.get.side_effect = [None, None, None]
|
||||
redis_client.mget.return_value = [None, None]
|
||||
client = client_cls.return_value
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider()]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
first_result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1")
|
||||
redis_client.get.reset_mock()
|
||||
redis_client.mget.reset_mock()
|
||||
redis_client.setex.reset_mock()
|
||||
client.fetch_model_providers.reset_mock()
|
||||
|
||||
second_result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1")
|
||||
|
||||
redis_client.get.assert_called_once_with(generation_key)
|
||||
redis_client.mget.assert_not_called()
|
||||
redis_client.setex.assert_not_called()
|
||||
client.fetch_model_providers.assert_not_called()
|
||||
assert [provider.provider for provider in second_result] == ["langgenius/openai/openai"]
|
||||
assert second_result[0] == first_result[0]
|
||||
assert second_result[0] is not first_result[0]
|
||||
|
||||
def test_invalidate_plugin_model_providers_cache_uses_redis_pipeline(self) -> None:
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
pipe = redis_client.pipeline.return_value
|
||||
@@ -310,41 +674,29 @@ class TestPluginModelProviderCache:
|
||||
pipe.incr.assert_called_once_with(_provider_generation_key("tenant-1"))
|
||||
pipe.execute.assert_called_once_with()
|
||||
|
||||
def test_invalidate_plugin_model_providers_cache_clears_process_local_cache(self) -> None:
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
pipe = redis_client.pipeline.return_value
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
PluginService._store_in_memory_plugin_model_providers("tenant-1", 0, [_build_provider_entity()])
|
||||
PluginService.invalidate_plugin_model_providers_cache("tenant-1")
|
||||
|
||||
assert PluginService._plugin_model_providers_memory_cache == {}
|
||||
redis_client.pipeline.assert_called_once_with(transaction=False)
|
||||
pipe.delete.assert_called_once_with(_provider_cache_key("tenant-1"))
|
||||
pipe.incr.assert_called_once_with(_provider_generation_key("tenant-1"))
|
||||
pipe.execute.assert_called_once_with()
|
||||
|
||||
def test_fetch_plugin_model_providers_ignores_stale_process_local_cache_after_generation_bump(self) -> None:
|
||||
def test_fetch_plugin_model_providers_uses_new_generation_cache_after_generation_bump(self) -> None:
|
||||
generation_key = _provider_generation_key("tenant-1")
|
||||
new_cache_key = _provider_cache_key("tenant-1", 1)
|
||||
with patch(f"{MODULE}.redis_client") as redis_client:
|
||||
redis_client.get.side_effect = [b"1", b"1"]
|
||||
redis_client.get.side_effect = [b"1", b"1", b"1"]
|
||||
redis_client.mget.return_value = [None]
|
||||
client = Mock()
|
||||
client.fetch_model_providers.return_value = [_build_plugin_model_provider(provider="anthropic")]
|
||||
|
||||
from core.plugin.plugin_service import PluginService
|
||||
|
||||
PluginService._store_in_memory_plugin_model_providers("tenant-1", 0, [_build_provider_entity()])
|
||||
result = PluginService.fetch_plugin_model_providers(tenant_id="tenant-1", client=client)
|
||||
|
||||
client.fetch_model_providers.assert_called_once_with("tenant-1")
|
||||
redis_client.get.assert_has_calls([call(generation_key), call(generation_key)])
|
||||
redis_client.mget.assert_called_once_with([new_cache_key])
|
||||
redis_client.get.assert_has_calls([call(generation_key), call(generation_key), call(generation_key)])
|
||||
assert redis_client.mget.call_args_list == [
|
||||
call([new_cache_key]),
|
||||
call([new_cache_key]),
|
||||
]
|
||||
redis_client.setex.assert_called_once()
|
||||
assert redis_client.setex.call_args.args[0] == new_cache_key
|
||||
assert PluginService._plugin_model_providers_memory_cache["tenant-1"][0] == 1
|
||||
redis_client.lock.return_value.acquire.assert_called_once()
|
||||
redis_client.lock.return_value.release.assert_called_once()
|
||||
assert [provider.provider for provider in result] == ["langgenius/anthropic/anthropic"]
|
||||
|
||||
|
||||
|
||||
@@ -73,6 +73,23 @@ class TestBillingServiceSendRequest:
|
||||
assert call_args[1]["headers"]["Billing-Api-Secret-Key"] == "test-secret-key"
|
||||
assert call_args[1]["headers"]["Content-Type"] == "application/json"
|
||||
|
||||
def test_send_request_with_base_url_override(self, mock_httpx_request, mock_billing_config):
|
||||
"""Quota APIs can use the new billing service without changing legacy billing calls."""
|
||||
# Arrange
|
||||
expected_response = {"result": "success"}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = httpx.codes.OK
|
||||
mock_response.json.return_value = expected_response
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act
|
||||
result = BillingService._send_request("GET", "/quota/balance", base_url="https://quota.example.com")
|
||||
|
||||
# Assert
|
||||
assert result == expected_response
|
||||
call_args = mock_httpx_request.call_args
|
||||
assert call_args[0][1] == "https://quota.example.com/quota/balance"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code", [httpx.codes.NOT_FOUND, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.BAD_REQUEST]
|
||||
)
|
||||
@@ -393,6 +410,20 @@ class TestBillingServiceSubscriptionInfo:
|
||||
params={"tenant_id": tenant_id},
|
||||
)
|
||||
|
||||
def test_quota_get_balance_uses_quota_request(self):
|
||||
tenant_id = "tenant-123"
|
||||
with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request:
|
||||
mock_send_quota_request.return_value = {"quota": "200", "usage": "6", "available": "194", "reserved": "0"}
|
||||
|
||||
result = BillingService.quota_get_balance(tenant_id, "credit_pool", bucket="trial")
|
||||
|
||||
assert result == {"quota": 200, "usage": 6, "available": 194, "reserved": 0}
|
||||
mock_send_quota_request.assert_called_once_with(
|
||||
"GET",
|
||||
"/quota/balance",
|
||||
params={"tenant_id": tenant_id, "feature_key": "credit_pool", "bucket": "trial"},
|
||||
)
|
||||
|
||||
def test_get_knowledge_rate_limit_with_defaults(self, mock_send_request):
|
||||
"""Test knowledge rate limit retrieval with default values."""
|
||||
# Arrange
|
||||
@@ -518,19 +549,20 @@ class TestBillingServiceUsageCalculation:
|
||||
assert result == expected_response
|
||||
mock_send_request.assert_called_once_with("GET", "/tenant-feature-usage/info", params={"tenant_id": tenant_id})
|
||||
|
||||
def test_get_quota_info(self, mock_send_request):
|
||||
def test_get_quota_info(self):
|
||||
"""Test retrieval of quota info from new endpoint."""
|
||||
# Arrange
|
||||
tenant_id = "tenant-123"
|
||||
expected_response = {"trigger_event": {"limit": 100, "usage": 30}, "api_rate_limit": {"limit": -1, "usage": 0}}
|
||||
mock_send_request.return_value = expected_response
|
||||
with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request:
|
||||
mock_send_quota_request.return_value = expected_response
|
||||
|
||||
# Act
|
||||
result = BillingService.get_quota_info(tenant_id)
|
||||
# Act
|
||||
result = BillingService.get_quota_info(tenant_id)
|
||||
|
||||
# Assert
|
||||
assert result == expected_response
|
||||
mock_send_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id})
|
||||
mock_send_quota_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id})
|
||||
|
||||
def test_update_tenant_feature_plan_usage_positive_delta(self, mock_send_request):
|
||||
"""Test updating tenant feature usage with positive delta (adding credits)."""
|
||||
@@ -614,7 +646,7 @@ class TestBillingServiceQuotaOperations:
|
||||
|
||||
@pytest.fixture
|
||||
def mock_send_request(self):
|
||||
with patch.object(BillingService, "_send_request") as mock:
|
||||
with patch.object(BillingService, "_send_quota_request") as mock:
|
||||
yield mock
|
||||
|
||||
def test_quota_reserve_success(self, mock_send_request):
|
||||
@@ -652,6 +684,16 @@ class TestBillingServiceQuotaOperations:
|
||||
call_json = mock_send_request.call_args[1]["json"]
|
||||
assert call_json["meta"] == {"source": "webhook"}
|
||||
|
||||
def test_quota_reserve_with_bucket(self, mock_send_request):
|
||||
mock_send_request.return_value = {"reservation_id": "rid-2", "available": 98, "reserved": 1}
|
||||
|
||||
BillingService.quota_reserve(
|
||||
tenant_id="t1", feature_key="credit_pool", request_id="req-2", amount=1, bucket="trial"
|
||||
)
|
||||
|
||||
call_json = mock_send_request.call_args[1]["json"]
|
||||
assert call_json["bucket"] == "trial"
|
||||
|
||||
def test_quota_commit_success(self, mock_send_request):
|
||||
expected = {"available": 98, "reserved": 0, "refunded": 0}
|
||||
mock_send_request.return_value = expected
|
||||
@@ -696,6 +738,20 @@ class TestBillingServiceQuotaOperations:
|
||||
call_json = mock_send_request.call_args[1]["json"]
|
||||
assert call_json["meta"] == {"reason": "partial"}
|
||||
|
||||
def test_quota_commit_with_bucket(self, mock_send_request):
|
||||
mock_send_request.return_value = {"available": 97, "reserved": 0, "refunded": 0}
|
||||
|
||||
BillingService.quota_commit(
|
||||
tenant_id="t1",
|
||||
feature_key="credit_pool",
|
||||
reservation_id="rid-1",
|
||||
actual_amount=1,
|
||||
bucket="paid",
|
||||
)
|
||||
|
||||
call_json = mock_send_request.call_args[1]["json"]
|
||||
assert call_json["bucket"] == "paid"
|
||||
|
||||
def test_quota_release_success(self, mock_send_request):
|
||||
expected = {"available": 100, "reserved": 0, "released": 1}
|
||||
mock_send_request.return_value = expected
|
||||
@@ -720,6 +776,64 @@ class TestBillingServiceQuotaOperations:
|
||||
assert result["released"] == 1
|
||||
assert isinstance(result["released"], int)
|
||||
|
||||
def test_quota_release_with_bucket(self, mock_send_request):
|
||||
mock_send_request.return_value = {"available": 100, "reserved": 0, "released": 1}
|
||||
|
||||
BillingService.quota_release(tenant_id="t1", feature_key="credit_pool", reservation_id="rid-1", bucket="trial")
|
||||
|
||||
call_json = mock_send_request.call_args[1]["json"]
|
||||
assert call_json["bucket"] == "trial"
|
||||
|
||||
def test_quota_consume_capped_success(self, mock_send_request):
|
||||
mock_send_request.return_value = {
|
||||
"deducted": "2",
|
||||
"available": "8",
|
||||
"reserved": "0",
|
||||
"quota": "10",
|
||||
"usage": "2",
|
||||
}
|
||||
|
||||
result = BillingService.quota_consume_capped(
|
||||
tenant_id="t1",
|
||||
feature_key="credit_pool",
|
||||
request_id="req-1",
|
||||
amount=5,
|
||||
bucket="paid",
|
||||
meta={"source": "test"},
|
||||
)
|
||||
|
||||
assert result == {"deducted": 2, "available": 8, "reserved": 0, "quota": 10, "usage": 2}
|
||||
mock_send_request.assert_called_once_with(
|
||||
"POST",
|
||||
"/quota/consume-capped",
|
||||
json={
|
||||
"tenant_id": "t1",
|
||||
"feature_key": "credit_pool",
|
||||
"request_id": "req-1",
|
||||
"amount": 5,
|
||||
"bucket": "paid",
|
||||
"meta": {"source": "test"},
|
||||
},
|
||||
)
|
||||
|
||||
def test_send_quota_request_uses_quota_base_url(self):
|
||||
with (
|
||||
patch.object(BillingService, "quota_base_url", "https://quota.example.com/v1"),
|
||||
patch.object(BillingService, "_send_request") as mock_send_request,
|
||||
):
|
||||
mock_send_request.return_value = {"ok": True}
|
||||
|
||||
result = BillingService._send_quota_request("GET", "/quota/info", params={"tenant_id": "t1"})
|
||||
|
||||
assert result == {"ok": True}
|
||||
mock_send_request.assert_called_once_with(
|
||||
"GET",
|
||||
"/quota/info",
|
||||
json=None,
|
||||
params={"tenant_id": "t1"},
|
||||
base_url="https://quota.example.com/v1",
|
||||
)
|
||||
|
||||
def test_get_quota_info_coerces_string_to_int(self, mock_send_request):
|
||||
"""Test that TypeAdapter coerces string values to int for get_quota_info."""
|
||||
mock_send_request.return_value = {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
@@ -15,6 +15,8 @@ from models.enums import ProviderQuotaType
|
||||
from services.credit_pool_service import (
|
||||
CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS,
|
||||
CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS,
|
||||
FEATURE_KEY_CREDIT_POOL,
|
||||
CreditPoolBalance,
|
||||
CreditPoolService,
|
||||
)
|
||||
|
||||
@@ -64,6 +66,12 @@ def _make_redis_lock() -> MagicMock:
|
||||
return lock
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _disable_billing_quota_by_default() -> Generator[None, None, None]:
|
||||
with patch("services.credit_pool_service.dify_config.BILLING_ENABLED", False):
|
||||
yield
|
||||
|
||||
|
||||
def test_get_pool_uses_configured_session_factory_without_flask_app_context() -> None:
|
||||
engine, tenant_id, _ = _create_engine_with_pool(quota_limit=10, quota_used=2)
|
||||
|
||||
@@ -75,6 +83,13 @@ def test_get_pool_uses_configured_session_factory_without_flask_app_context() ->
|
||||
assert pool.quota_used == 2
|
||||
|
||||
|
||||
def test_credit_pool_balance_unlimited_remaining_and_sufficiency() -> None:
|
||||
pool = CreditPoolBalance(tenant_id="tenant-1", pool_type="paid", quota_limit=-1, quota_used=999)
|
||||
|
||||
assert pool.remaining_credits == -1
|
||||
assert pool.has_sufficient_credits(10_000)
|
||||
|
||||
|
||||
def test_check_and_deduct_credits_deducts_exact_amount_when_sufficient() -> None:
|
||||
engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2)
|
||||
|
||||
@@ -224,7 +239,7 @@ def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() -
|
||||
)
|
||||
redis_lock.acquire.assert_called_once_with(blocking=True)
|
||||
redis_lock.release.assert_called_once_with()
|
||||
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL)
|
||||
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="trial")
|
||||
|
||||
|
||||
def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> None:
|
||||
@@ -254,7 +269,152 @@ def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> N
|
||||
)
|
||||
redis_lock.acquire.assert_called_once_with(blocking=True)
|
||||
redis_lock.release.assert_called_once_with()
|
||||
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID)
|
||||
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="paid")
|
||||
|
||||
|
||||
def test_get_pool_uses_billing_quota_balance_when_enabled() -> None:
|
||||
tenant_id = "tenant-1"
|
||||
with (
|
||||
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
|
||||
patch("services.billing_service.BillingService.quota_get_balance") as quota_get_balance,
|
||||
):
|
||||
quota_get_balance.return_value = {"quota": 1000, "usage": 250, "available": 750, "reserved": 0}
|
||||
|
||||
pool = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID)
|
||||
|
||||
assert isinstance(pool, CreditPoolBalance)
|
||||
assert pool.quota_limit == 1000
|
||||
assert pool.quota_used == 250
|
||||
assert pool.remaining_credits == 750
|
||||
quota_get_balance.assert_called_once_with(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket="paid",
|
||||
)
|
||||
|
||||
|
||||
def test_check_and_deduct_credits_uses_billing_reserve_and_commit_when_enabled() -> None:
|
||||
tenant_id = "tenant-1"
|
||||
with (
|
||||
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
|
||||
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
|
||||
patch("services.billing_service.BillingService.quota_commit") as quota_commit,
|
||||
patch("services.billing_service.BillingService.quota_release") as quota_release,
|
||||
):
|
||||
quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3}
|
||||
|
||||
result = CreditPoolService.check_and_deduct_credits(
|
||||
tenant_id=tenant_id,
|
||||
credits_required=3,
|
||||
pool_type=ProviderQuotaType.TRIAL,
|
||||
)
|
||||
|
||||
assert result == 3
|
||||
quota_reserve.assert_called_once_with(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket="trial",
|
||||
request_id=ANY,
|
||||
amount=3,
|
||||
meta={"source": "credit_pool.check_and_deduct"},
|
||||
)
|
||||
quota_commit.assert_called_once_with(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket="trial",
|
||||
reservation_id="reservation-1",
|
||||
actual_amount=3,
|
||||
meta={"source": "credit_pool.check_and_deduct"},
|
||||
)
|
||||
quota_release.assert_not_called()
|
||||
|
||||
|
||||
def test_check_and_deduct_credits_raises_when_billing_reserve_is_insufficient() -> None:
|
||||
with (
|
||||
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
|
||||
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
|
||||
):
|
||||
quota_reserve.return_value = {"reservation_id": "", "available": 1, "reserved": 0}
|
||||
|
||||
with pytest.raises(QuotaExceededError, match="Insufficient credits remaining"):
|
||||
CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3)
|
||||
|
||||
|
||||
def test_check_and_deduct_credits_releases_billing_reservation_when_commit_fails() -> None:
|
||||
with (
|
||||
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
|
||||
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
|
||||
patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")),
|
||||
patch("services.billing_service.BillingService.quota_release") as quota_release,
|
||||
):
|
||||
quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3}
|
||||
|
||||
with pytest.raises(RuntimeError, match="commit failed"):
|
||||
CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3)
|
||||
|
||||
quota_release.assert_called_once_with(
|
||||
tenant_id="tenant-1",
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket="trial",
|
||||
reservation_id="reservation-1",
|
||||
)
|
||||
|
||||
|
||||
def test_check_and_deduct_credits_logs_when_billing_release_fails() -> None:
|
||||
with (
|
||||
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
|
||||
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
|
||||
patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")),
|
||||
patch(
|
||||
"services.billing_service.BillingService.quota_release", side_effect=RuntimeError("release failed")
|
||||
) as quota_release,
|
||||
patch("services.credit_pool_service.logger.warning") as logger_warning,
|
||||
):
|
||||
quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3}
|
||||
|
||||
with pytest.raises(RuntimeError, match="commit failed"):
|
||||
CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3)
|
||||
|
||||
quota_release.assert_called_once_with(
|
||||
tenant_id="tenant-1",
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket="trial",
|
||||
reservation_id="reservation-1",
|
||||
)
|
||||
logger_warning.assert_called_once()
|
||||
assert logger_warning.call_args.args[3] == "reservation-1"
|
||||
assert logger_warning.call_args.kwargs["exc_info"] is True
|
||||
|
||||
|
||||
def test_deduct_credits_capped_uses_billing_consume_capped_when_enabled() -> None:
|
||||
tenant_id = "tenant-1"
|
||||
with (
|
||||
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
|
||||
patch("services.billing_service.BillingService.quota_consume_capped") as quota_consume_capped,
|
||||
):
|
||||
quota_consume_capped.return_value = {
|
||||
"deducted": 2,
|
||||
"available": 0,
|
||||
"reserved": 0,
|
||||
"quota": 10,
|
||||
"usage": 10,
|
||||
}
|
||||
|
||||
result = CreditPoolService.deduct_credits_capped(
|
||||
tenant_id=tenant_id,
|
||||
credits_required=5,
|
||||
pool_type=ProviderQuotaType.PAID,
|
||||
)
|
||||
|
||||
assert result == 2
|
||||
quota_consume_capped.assert_called_once_with(
|
||||
tenant_id=tenant_id,
|
||||
feature_key=FEATURE_KEY_CREDIT_POOL,
|
||||
bucket="paid",
|
||||
request_id=ANY,
|
||||
amount=5,
|
||||
meta={"source": "credit_pool.deduct_capped"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
import pytest
|
||||
|
||||
from services import feature_service as feature_service_module
|
||||
from services.feature_service import FeatureService, SystemFeatureModel
|
||||
|
||||
_ENTERPRISE_INFO = {"License": {"licensedSeats": {"enabled": True, "limit": 3, "used": 1}}}
|
||||
|
||||
|
||||
def test_fulfill_params_from_enterprise_parses_licensed_seats(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Authenticated fill copies the licensed-seat quota out of the enterprise payload."""
|
||||
monkeypatch.setattr(
|
||||
feature_service_module.EnterpriseService,
|
||||
"get_info",
|
||||
staticmethod(lambda: _ENTERPRISE_INFO),
|
||||
)
|
||||
|
||||
features = SystemFeatureModel()
|
||||
FeatureService._fulfill_params_from_enterprise(features, is_authenticated=True)
|
||||
|
||||
assert features.license.seats.enabled is True
|
||||
assert features.license.seats.limit == 3
|
||||
assert features.license.seats.size == 1
|
||||
|
||||
|
||||
def test_fulfill_params_from_enterprise_withholds_seats_when_unauthenticated(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Seat counts are auth-gated: unauthenticated callers keep the zeroed default."""
|
||||
monkeypatch.setattr(
|
||||
feature_service_module.EnterpriseService,
|
||||
"get_info",
|
||||
staticmethod(lambda: _ENTERPRISE_INFO),
|
||||
)
|
||||
|
||||
features = SystemFeatureModel()
|
||||
FeatureService._fulfill_params_from_enterprise(features, is_authenticated=False)
|
||||
|
||||
assert features.license.seats.enabled is False
|
||||
assert features.license.seats.limit == 0
|
||||
assert features.license.seats.size == 0
|
||||
Generated
+11
-9
@@ -1374,6 +1374,7 @@ dependencies = [
|
||||
{ name = "resend" },
|
||||
{ name = "sendgrid" },
|
||||
{ name = "sseclient-py" },
|
||||
{ name = "zstandard" },
|
||||
]
|
||||
|
||||
[package.dev-dependencies]
|
||||
@@ -1657,6 +1658,7 @@ requires-dist = [
|
||||
{ name = "resend", specifier = ">=2.27.0,<3.0.0" },
|
||||
{ name = "sendgrid", specifier = ">=6.12.5,<7.0.0" },
|
||||
{ name = "sseclient-py", specifier = ">=1.8.0,<2.0.0" },
|
||||
{ name = "zstandard", specifier = "==0.25.0" },
|
||||
]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
@@ -5582,14 +5584,14 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "python-engineio"
|
||||
version = "4.13.1"
|
||||
version = "4.13.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "simple-websocket" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/34/12/bdef9dbeedbe2cdeba2a2056ad27b1fb081557d34b69a97f574843462cae/python_engineio-4.13.1.tar.gz", hash = "sha256:0a853fcef52f5b345425d8c2b921ac85023a04dfcf75d7b74696c61e940fd066", size = 92348, upload-time = "2026-02-06T23:38:06.12Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/fb/a0/f75491f942184d9960b15e763270f765fe9f239745ca5f9e16289011aed4/python_engineio-4.13.3.tar.gz", hash = "sha256:572b7783e341fed21edbc7cea297ccd378dad79265fdde96aa4664420a7c06c9", size = 79734, upload-time = "2026-06-20T22:53:52.197Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/aa/54/0cce26da03a981f949bb8449c9778537f75f5917c172e1d2992ff25cb57d/python_engineio-4.13.1-py3-none-any.whl", hash = "sha256:f32ad10589859c11053ad7d9bb3c9695cdf862113bfb0d20bc4d890198287399", size = 59847, upload-time = "2026-02-06T23:38:04.861Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5b/96/82f6328e410515fab21d5602ba35b9377a47b5a141a0c1f9efa00ce21eb4/python_engineio-4.13.3-py3-none-any.whl", hash = "sha256:1f60ecaf1358190f0e26c48c578a60428dc02a8f1295bc3dbf53d1b31116821f", size = 59993, upload-time = "2026-06-20T22:53:50.775Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5650,15 +5652,15 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "python-socketio"
|
||||
version = "5.16.1"
|
||||
version = "5.16.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "bidict" },
|
||||
{ name = "python-engineio" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/59/81/cf8284f45e32efa18d3848ed82cdd4dcc1b657b082458fbe01ad3e1f2f8d/python_socketio-5.16.1.tar.gz", hash = "sha256:f863f98eacce81ceea2e742f6388e10ca3cdd0764be21d30d5196470edf5ea89", size = 128508, upload-time = "2026-02-06T23:42:07Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/32/2d/ffce71017c106b75099fea569df6518c63fee5d6202ce0cfe7b01e6f22c3/python_socketio-5.16.3.tar.gz", hash = "sha256:89b136f677ae65607a84cecda9b4d6c5377b40a97582c504c25df89af16d520e", size = 128095, upload-time = "2026-06-15T22:07:04.003Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/07/c7/deb8c5e604404dbf10a3808a858946ca3547692ff6316b698945bb72177e/python_socketio-5.16.1-py3-none-any.whl", hash = "sha256:a3eb1702e92aa2f2b5d3ba00261b61f062cce51f1cfb6900bf3ab4d1934d2d35", size = 82054, upload-time = "2026-02-06T23:42:05.772Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0a/38/8c5e72d53ff8eb27497c4f268a7f6d9121e727a50b65248288ad79a93053/python_socketio-5.16.3-py3-none-any.whl", hash = "sha256:e7ad14202a5e6448824c7c2f86161d04e13dec05992257df5c709e6a2798c041", size = 82087, upload-time = "2026-06-15T22:07:02.498Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6112,11 +6114,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "soupsieve"
|
||||
version = "2.8"
|
||||
version = "2.8.4"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/6d/e6/21ccce3262dd4889aa3332e5a119a3491a95e8f60939870a3a035aabac0d/soupsieve-2.8.tar.gz", hash = "sha256:e2dd4a40a628cb5f28f6d4b0db8800b8f581b65bb380b97de22ba5ca8d72572f", size = 103472, upload-time = "2025-08-27T15:39:51.78Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/47/2c/0a5f6f8ee0d5589e48c7640213ed5175d52cf540a06725b628cc1a45d6ce/soupsieve-2.8.4.tar.gz", hash = "sha256:e121fd02e975c695e4e9e8774a5ee35d74714b59307868dcc5319ad2d9e3328e", size = 121110, upload-time = "2026-05-24T13:55:57.154Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/14/a0/bb38d3b76b8cae341dad93a2dd83ab7462e6dbcdd84d43f54ee60a8dc167/soupsieve-2.8-py3-none-any.whl", hash = "sha256:0cc76456a30e20f5d7f2e14a98a4ae2ee4e5abdc7c5ea0aafe795f344bc7984c", size = 36679, upload-time = "2025-08-27T15:39:50.179Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5e/f5/0c41cb68dcae6b7de4fac4188a3a9589e21fb31df21ea3a2e888db95e6c9/soupsieve-2.8.4-py3-none-any.whl", hash = "sha256:e7e6b0769c8f51ed59acab6e994b00621096cfb1c640a7509295987388fbaf65", size = 37304, upload-time = "2026-05-24T13:55:55.406Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -310,11 +310,6 @@
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"web/app/account/oauth/authorize/layout.tsx": {
|
||||
"ts/no-explicit-any": {
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"web/app/account/oauth/authorize/page.tsx": {
|
||||
"ts/no-explicit-any": {
|
||||
"count": 1
|
||||
@@ -3335,14 +3330,6 @@
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"web/app/components/explore/sidebar/app-nav-item/index.tsx": {
|
||||
"jsx-a11y/click-events-have-key-events": {
|
||||
"count": 1
|
||||
},
|
||||
"jsx-a11y/no-static-element-interactions": {
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"web/app/components/explore/try-app/app/text-generation.tsx": {
|
||||
"jsx-a11y/click-events-have-key-events": {
|
||||
"count": 1
|
||||
@@ -3619,11 +3606,6 @@
|
||||
"count": 2
|
||||
}
|
||||
},
|
||||
"web/app/components/header/account-setting/model-provider-page/model-provider-page-body.tsx": {
|
||||
"jsx-a11y/anchor-has-content": {
|
||||
"count": 1
|
||||
}
|
||||
},
|
||||
"web/app/components/header/account-setting/model-provider-page/provider-added-card/cooldown-timer.tsx": {
|
||||
"react/set-state-in-effect": {
|
||||
"count": 1
|
||||
|
||||
@@ -4496,6 +4496,7 @@ export type DeleteAppsByAppIdTraceConfigData = {
|
||||
|
||||
export type DeleteAppsByAppIdTraceConfigErrors = {
|
||||
400: unknown
|
||||
403: unknown
|
||||
}
|
||||
|
||||
export type DeleteAppsByAppIdTraceConfigResponses = {
|
||||
@@ -4538,6 +4539,7 @@ export type PatchAppsByAppIdTraceConfigData = {
|
||||
|
||||
export type PatchAppsByAppIdTraceConfigErrors = {
|
||||
400: unknown
|
||||
403: unknown
|
||||
}
|
||||
|
||||
export type PatchAppsByAppIdTraceConfigResponses = {
|
||||
@@ -4558,6 +4560,7 @@ export type PostAppsByAppIdTraceConfigData = {
|
||||
|
||||
export type PostAppsByAppIdTraceConfigErrors = {
|
||||
400: unknown
|
||||
403: unknown
|
||||
}
|
||||
|
||||
export type PostAppsByAppIdTraceConfigResponses = {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -3,19 +3,16 @@ import { toast, ToastHost } from '../index'
|
||||
|
||||
const asHTMLElement = (element: HTMLElement | SVGElement) => element as HTMLElement
|
||||
|
||||
type BaseUIAnimationGlobal = typeof globalThis & {
|
||||
BASE_UI_ANIMATIONS_DISABLED: boolean
|
||||
}
|
||||
|
||||
const dispatchToastMouseOver = (element: HTMLElement | SVGElement) => {
|
||||
element.dispatchEvent(new MouseEvent('mouseover', {
|
||||
bubbles: true,
|
||||
}))
|
||||
}
|
||||
|
||||
const dispatchToastMouseOut = (element: HTMLElement | SVGElement) => {
|
||||
element.dispatchEvent(new MouseEvent('mouseout', {
|
||||
bubbles: true,
|
||||
relatedTarget: document.body,
|
||||
}))
|
||||
}
|
||||
|
||||
describe('@langgenius/dify-ui/toast', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
@@ -38,13 +35,10 @@ describe('@langgenius/dify-ui/toast', () => {
|
||||
|
||||
await expect.element(screen.getByText('Saved')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('Your changes are available now.')).toBeInTheDocument()
|
||||
await expect.element(screen.getByRole('region', { name: 'Notifications' })).toHaveAttribute('aria-live', 'polite')
|
||||
await expect.element(screen.getByRole('region', { name: 'Notifications' })).toHaveClass('z-60')
|
||||
expect(screen.getByRole('region', { name: 'Notifications' }).element().firstElementChild).toHaveClass('top-4')
|
||||
expect(screen.getByText('Saved').element().closest('[class*="transition-opacity"]')).toHaveClass('motion-reduce:transition-none')
|
||||
expect(screen.getByRole('dialog').element()).not.toHaveClass('outline-hidden')
|
||||
const viewport = screen.getByRole('region', { name: 'Notifications' })
|
||||
await expect.element(viewport).toHaveAttribute('aria-live', 'polite')
|
||||
await expect.element(viewport).toHaveClass('z-60')
|
||||
expect(document.body.querySelector('[aria-hidden="true"].i-ri-checkbox-circle-fill')).toBeInTheDocument()
|
||||
expect(document.body.querySelector('button[aria-label="Close notification"][aria-hidden="true"]')).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('should keep multiple toast roots mounted in a collapsed stack', async () => {
|
||||
@@ -58,37 +52,6 @@ describe('@langgenius/dify-ui/toast', () => {
|
||||
|
||||
await expect.element(screen.getByText('Third toast')).toBeInTheDocument()
|
||||
expect(document.body.querySelectorAll('[role="dialog"]')).toHaveLength(3)
|
||||
expect(document.body.querySelectorAll('button[aria-label="Close notification"][aria-hidden="true"]')).toHaveLength(3)
|
||||
|
||||
const viewport = screen.getByRole('region', { name: 'Notifications' }).element()
|
||||
dispatchToastMouseOver(viewport)
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(document.body.querySelector('button[aria-label="Close notification"][aria-hidden="true"]')).not.toBeInTheDocument()
|
||||
})
|
||||
dispatchToastMouseOut(viewport)
|
||||
})
|
||||
|
||||
it('should clamp varying-height toasts to the frontmost stack height when collapsed', async () => {
|
||||
const screen = await render(<ToastHost />)
|
||||
|
||||
toast.info('Long background toast', {
|
||||
description: 'This longer toast intentionally spans multiple lines so it would overflow the collapsed stack without matching the frontmost toast height.',
|
||||
})
|
||||
toast.success('Short front toast', {
|
||||
description: 'Short message.',
|
||||
})
|
||||
|
||||
await expect.element(screen.getByText('Short front toast')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('Long background toast')).toBeInTheDocument()
|
||||
await expect.element(screen.getByRole('region', { name: 'Notifications' })).toHaveAttribute('aria-live', 'polite')
|
||||
await expect.element(screen.getByRole('dialog', { name: 'Short front toast' })).toBeInTheDocument()
|
||||
await expect.element(screen.getByRole('dialog', { name: 'Long background toast' })).toBeInTheDocument()
|
||||
|
||||
const longToastContent = screen.getByText('Long background toast').element().closest('[class*="transition-opacity"]')
|
||||
expect(longToastContent).toHaveAttribute('data-behind')
|
||||
expect(longToastContent).toHaveClass('h-full')
|
||||
expect(longToastContent?.parentElement).toHaveClass('h-full')
|
||||
})
|
||||
|
||||
it('should render a neutral toast when called directly', async () => {
|
||||
@@ -157,7 +120,6 @@ describe('@langgenius/dify-ui/toast', () => {
|
||||
dispatchToastMouseOver(viewport)
|
||||
|
||||
await expect.element(screen.getByRole('button', { name: 'Close notification' })).toBeInTheDocument()
|
||||
dispatchToastMouseOut(viewport)
|
||||
asHTMLElement(screen.getByRole('button', { name: 'Close notification' }).element()).click()
|
||||
|
||||
await vi.waitFor(() => {
|
||||
@@ -166,15 +128,112 @@ describe('@langgenius/dify-ui/toast', () => {
|
||||
expect(onClose).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('should keep zero-timeout toasts persistent', async () => {
|
||||
const screen = await render(<ToastHost />)
|
||||
it('should let pointer events pass through a toast while it is exiting', async () => {
|
||||
const onClick = vi.fn()
|
||||
const baseUIAnimationGlobal = globalThis as BaseUIAnimationGlobal
|
||||
const animationState = baseUIAnimationGlobal.BASE_UI_ANIMATIONS_DISABLED
|
||||
baseUIAnimationGlobal.BASE_UI_ANIMATIONS_DISABLED = false
|
||||
|
||||
try {
|
||||
const screen = await render(
|
||||
<>
|
||||
<style>
|
||||
{`
|
||||
[role="dialog"] {
|
||||
transition: opacity 10000s, transform 10000s !important;
|
||||
}
|
||||
[role="dialog"][data-ending-style] {
|
||||
opacity: 0 !important;
|
||||
transform: translateY(-150%) !important;
|
||||
}
|
||||
.data-ending-style\\:pointer-events-none[data-ending-style] {
|
||||
pointer-events: none;
|
||||
}
|
||||
.data-ending-style\\:after\\:pointer-events-none[data-ending-style]::after {
|
||||
pointer-events: none;
|
||||
}
|
||||
`}
|
||||
</style>
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClick}
|
||||
style={{
|
||||
position: 'fixed',
|
||||
top: '16px',
|
||||
right: '32px',
|
||||
width: '360px',
|
||||
height: '96px',
|
||||
}}
|
||||
>
|
||||
Underlying action
|
||||
</button>
|
||||
<ToastHost />
|
||||
</>,
|
||||
)
|
||||
|
||||
toast('Dismiss me', {
|
||||
timeout: 0,
|
||||
})
|
||||
|
||||
await expect.element(screen.getByRole('dialog', { name: 'Dismiss me' })).toBeInTheDocument()
|
||||
asHTMLElement(screen.getByRole('dialog', { name: 'Dismiss me' }).element()).click()
|
||||
|
||||
const viewport = screen.getByRole('region', { name: 'Notifications' }).element()
|
||||
dispatchToastMouseOver(viewport)
|
||||
|
||||
await expect.element(screen.getByRole('button', { name: 'Close notification' })).toBeInTheDocument()
|
||||
asHTMLElement(screen.getByRole('button', { name: 'Close notification' }).element()).click()
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(screen.getByRole('dialog', { name: 'Dismiss me' }).element()).toHaveAttribute('data-ending-style')
|
||||
})
|
||||
|
||||
asHTMLElement(screen.getByRole('dialog', { name: 'Dismiss me' }).element()).click()
|
||||
|
||||
const underlyingAction = asHTMLElement(screen.getByRole('button', { name: 'Underlying action' }).element())
|
||||
const rect = underlyingAction.getBoundingClientRect()
|
||||
const x = rect.left + rect.width / 2
|
||||
const y = rect.top + rect.height / 2
|
||||
|
||||
document.elementFromPoint(x, y)?.dispatchEvent(new MouseEvent('click', {
|
||||
bubbles: true,
|
||||
cancelable: true,
|
||||
clientX: x,
|
||||
clientY: y,
|
||||
}))
|
||||
|
||||
expect(onClick).toHaveBeenCalledTimes(1)
|
||||
}
|
||||
finally {
|
||||
baseUIAnimationGlobal.BASE_UI_ANIMATIONS_DISABLED = animationState
|
||||
}
|
||||
})
|
||||
|
||||
it('should pass the host timeout to added toasts', async () => {
|
||||
const screen = await render(<ToastHost timeout={1000} />)
|
||||
|
||||
toast('Auto dismiss')
|
||||
await expect.element(screen.getByText('Auto dismiss')).toBeInTheDocument()
|
||||
|
||||
await vi.advanceTimersByTimeAsync(999)
|
||||
expect(document.body).toHaveTextContent('Auto dismiss')
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1)
|
||||
await vi.waitFor(() => {
|
||||
expect(document.body).not.toHaveTextContent('Auto dismiss')
|
||||
})
|
||||
})
|
||||
|
||||
it('should keep a toast persistent when its timeout is zero', async () => {
|
||||
const screen = await render(<ToastHost timeout={1000} />)
|
||||
|
||||
toast('Persistent', {
|
||||
timeout: 0,
|
||||
})
|
||||
|
||||
await expect.element(screen.getByText('Persistent')).toBeInTheDocument()
|
||||
|
||||
await vi.advanceTimersByTimeAsync(10000)
|
||||
await vi.advanceTimersByTimeAsync(5000)
|
||||
expect(document.body).toHaveTextContent('Persistent')
|
||||
})
|
||||
|
||||
@@ -185,6 +244,7 @@ describe('@langgenius/dify-ui/toast', () => {
|
||||
description: 'Preparing your data…',
|
||||
})
|
||||
await expect.element(screen.getByText('Loading')).toBeInTheDocument()
|
||||
expect(document.body.querySelector('[aria-hidden="true"].i-ri-information-2-fill')).toBeInTheDocument()
|
||||
|
||||
toast.update(toastId, {
|
||||
title: 'Done',
|
||||
@@ -195,27 +255,28 @@ describe('@langgenius/dify-ui/toast', () => {
|
||||
await expect.element(screen.getByText('Done')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('Your data is ready.')).toBeInTheDocument()
|
||||
expect(document.body).not.toHaveTextContent('Loading')
|
||||
expect(document.body.querySelector('[aria-hidden="true"].i-ri-checkbox-circle-fill')).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('should upsert an existing toast when add is called with the same id', async () => {
|
||||
const screen = await render(<ToastHost />)
|
||||
|
||||
toast('Syncing', {
|
||||
id: 'sync-job',
|
||||
description: 'Uploading changes…',
|
||||
toast('Draft saving', {
|
||||
id: 'draft-save-status',
|
||||
description: 'Saving changes…',
|
||||
})
|
||||
await expect.element(screen.getByText('Syncing')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('Draft saving')).toBeInTheDocument()
|
||||
|
||||
toast.success('Synced', {
|
||||
id: 'sync-job',
|
||||
description: 'All changes are uploaded.',
|
||||
toast.success('Draft saved', {
|
||||
id: 'draft-save-status',
|
||||
description: 'All changes are saved.',
|
||||
})
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(document.body).not.toHaveTextContent('Syncing')
|
||||
expect(document.body).not.toHaveTextContent('Draft saving')
|
||||
})
|
||||
await expect.element(screen.getByText('Synced')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('All changes are uploaded.')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('Draft saved')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('All changes are saved.')).toBeInTheDocument()
|
||||
expect(document.body.querySelectorAll('[role="dialog"]')).toHaveLength(1)
|
||||
})
|
||||
|
||||
@@ -261,5 +322,6 @@ describe('@langgenius/dify-ui/toast', () => {
|
||||
|
||||
await expect.element(screen.getByText('Saved')).toBeInTheDocument()
|
||||
await expect.element(screen.getByText('Your changes are available now.')).toBeInTheDocument()
|
||||
expect(document.body.querySelector('[aria-hidden="true"].i-ri-checkbox-circle-fill')).toBeInTheDocument()
|
||||
})
|
||||
})
|
||||
|
||||
@@ -153,13 +153,13 @@ function ToastCard({
|
||||
<BaseToast.Root
|
||||
toast={toastItem}
|
||||
className={cn(
|
||||
'pointer-events-auto absolute top-0 right-0 w-90 max-w-[calc(100vw-2rem)] origin-top cursor-default rounded-xl select-none focus-visible:ring-2 focus-visible:ring-state-accent-solid focus-visible:outline-hidden',
|
||||
'pointer-events-auto absolute top-0 right-0 w-full origin-top cursor-default rounded-xl select-none focus-visible:ring-2 focus-visible:ring-state-accent-solid focus-visible:outline-hidden',
|
||||
'[--toast-current-height:var(--toast-frontmost-height,var(--toast-height))] [--toast-gap:8px] [--toast-peek:5px] [--toast-scale:calc(1-(var(--toast-index)*0.0225))] [--toast-shrink:calc(1-var(--toast-scale))]',
|
||||
'z-[calc(100-var(--toast-index))] h-(--toast-current-height)',
|
||||
'[transition:transform_500ms_cubic-bezier(0.22,1,0.36,1),opacity_500ms,height_150ms] motion-reduce:transition-none',
|
||||
'transform-[translateX(var(--toast-swipe-movement-x))_translateY(calc(var(--toast-swipe-movement-y)+(var(--toast-index)*var(--toast-peek))+(var(--toast-shrink)*var(--toast-current-height))))_scale(var(--toast-scale))]',
|
||||
'data-expanded:h-(--toast-height) data-expanded:transform-[translateX(var(--toast-swipe-movement-x))_translateY(calc(var(--toast-offset-y)+var(--toast-swipe-movement-y)+(var(--toast-index)*8px)))_scale(1)]',
|
||||
'data-ending-style:transform-[translateY(-150%)] data-ending-style:opacity-0',
|
||||
'data-ending-style:pointer-events-none data-ending-style:after:pointer-events-none data-ending-style:transform-[translateY(-150%)] data-ending-style:opacity-0',
|
||||
'data-ending-style:data-[swipe-direction=down]:transform-[translateY(calc(var(--toast-swipe-movement-y)+150%))]',
|
||||
'data-ending-style:data-[swipe-direction=right]:transform-[translateX(calc(var(--toast-swipe-movement-x)+150%))]',
|
||||
'data-limited:pointer-events-none data-limited:opacity-0 data-starting-style:transform-[translateY(-150%)] data-starting-style:opacity-0',
|
||||
@@ -178,13 +178,13 @@ function ToastCard({
|
||||
<div className="min-w-0 flex-1 p-1">
|
||||
<div className="flex w-full min-w-0 items-center gap-1">
|
||||
{toastItem.title != null && (
|
||||
<BaseToast.Title className="min-w-0 flex-1 system-sm-semibold wrap-break-word [overflow-wrap:anywhere] text-text-primary">
|
||||
<BaseToast.Title className="min-w-0 flex-1 system-sm-semibold wrap-break-word text-text-primary">
|
||||
{toastItem.title}
|
||||
</BaseToast.Title>
|
||||
)}
|
||||
</div>
|
||||
{toastItem.description != null && (
|
||||
<BaseToast.Description className="mt-1 min-w-0 system-xs-regular wrap-break-word [overflow-wrap:anywhere] text-text-secondary">
|
||||
<BaseToast.Description className="mt-1 min-w-0 system-xs-regular wrap-break-word text-text-secondary">
|
||||
{toastItem.description}
|
||||
</BaseToast.Description>
|
||||
)}
|
||||
@@ -222,21 +222,15 @@ function ToastViewport() {
|
||||
<BaseToast.Viewport
|
||||
aria-label={toastViewportLabel}
|
||||
className={cn(
|
||||
'group/toast-viewport pointer-events-none fixed inset-0 z-60 overflow-visible',
|
||||
'group/toast-viewport pointer-events-none fixed top-4 right-4 z-60 w-90 max-w-[calc(100vw-2rem)] overflow-visible sm:right-8',
|
||||
)}
|
||||
>
|
||||
<div
|
||||
className={cn(
|
||||
'pointer-events-none absolute top-4 right-4 w-90 max-w-[calc(100vw-2rem)] sm:right-8',
|
||||
)}
|
||||
>
|
||||
{toasts.map(toastItem => (
|
||||
<ToastCard
|
||||
key={toastItem.id}
|
||||
toast={toastItem}
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
{toasts.map(toastItem => (
|
||||
<ToastCard
|
||||
key={toastItem.id}
|
||||
toast={toastItem}
|
||||
/>
|
||||
))}
|
||||
</BaseToast.Viewport>
|
||||
)
|
||||
}
|
||||
|
||||
@@ -232,7 +232,20 @@ describe('Billing Page + Plan Integration', () => {
|
||||
|
||||
// Verify billing URL button visibility and behavior
|
||||
describe('Billing URL button', () => {
|
||||
it('should show billing button when subscription management permission is granted', () => {
|
||||
it('should show billing button when manager has subscription management permission', () => {
|
||||
setupProviderContext({ type: Plan.sandbox })
|
||||
setupAppContext({
|
||||
isCurrentWorkspaceManager: true,
|
||||
workspacePermissionKeys: ['billing.subscription.manage'],
|
||||
})
|
||||
|
||||
render(<Billing />)
|
||||
|
||||
expect(screen.getByText(/viewBillingTitle/i)).toBeInTheDocument()
|
||||
expect(screen.getByText(/viewBillingAction/i)).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('should hide billing button when subscription management permission is granted without manager role', () => {
|
||||
setupProviderContext({ type: Plan.sandbox })
|
||||
setupAppContext({
|
||||
isCurrentWorkspaceManager: false,
|
||||
@@ -241,8 +254,7 @@ describe('Billing Page + Plan Integration', () => {
|
||||
|
||||
render(<Billing />)
|
||||
|
||||
expect(screen.getByText(/viewBillingTitle/i)).toBeInTheDocument()
|
||||
expect(screen.getByText(/viewBillingAction/i)).toBeInTheDocument()
|
||||
expect(screen.queryByText(/viewBillingTitle/i)).not.toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('should hide billing button when subscription management permission is missing', () => {
|
||||
|
||||
@@ -204,12 +204,5 @@ describe('Sidebar Lifecycle Flow', () => {
|
||||
|
||||
expect(screen.getByText('explore.sidebar.noApps.title')).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('should hide NoApps on mobile', () => {
|
||||
mockMediaType = MediaType.mobile
|
||||
renderSidebar()
|
||||
|
||||
expect(screen.queryByText('explore.sidebar.noApps.title')).not.toBeInTheDocument()
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -6,6 +6,9 @@ import PluginPage from '@/app/components/plugins/plugin-page'
|
||||
import { createNuqsTestWrapper } from '@/test/nuqs-testing'
|
||||
|
||||
const mockFetchManifestFromMarketPlace = vi.fn()
|
||||
const { mockRouterReplace } = vi.hoisted(() => ({
|
||||
mockRouterReplace: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('react-i18next', () => ({
|
||||
useTranslation: () => ({
|
||||
@@ -26,6 +29,12 @@ vi.mock('@/hooks/use-document-title', () => ({
|
||||
default: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/next/navigation', () => ({
|
||||
useRouter: () => ({
|
||||
replace: mockRouterReplace,
|
||||
}),
|
||||
}))
|
||||
|
||||
vi.mock('@/context/i18n', () => ({
|
||||
useDocLink: () => (path: string) => `https://docs.example.com${path}`,
|
||||
}))
|
||||
|
||||
@@ -6,6 +6,7 @@ import { GoogleAnalyticsScripts } from '@/app/components/base/ga'
|
||||
import Zendesk from '@/app/components/base/zendesk'
|
||||
import { EducationVerifyActionRecorder } from '@/app/components/education-verify-action-recorder'
|
||||
import { GotoAnything } from '@/app/components/goto-anything'
|
||||
import MaintenanceNotice from '@/app/components/header/maintenance-notice'
|
||||
import MainNavLayout from '@/app/components/main-nav/layout'
|
||||
import { NextRouteStateBridge } from '@/app/components/next-route-state'
|
||||
import { OAuthRegistrationAnalytics } from '@/app/components/oauth-registration-analytics'
|
||||
@@ -21,34 +22,37 @@ import { RoleRouteGuard } from './role-route-guard'
|
||||
|
||||
export default async function Layout({ children }: { children: ReactNode }) {
|
||||
return (
|
||||
<>
|
||||
<React.Fragment>
|
||||
<GoogleAnalyticsScripts />
|
||||
<AmplitudeProvider />
|
||||
<OAuthRegistrationAnalytics />
|
||||
<EducationVerifyActionRecorder />
|
||||
<CommonLayoutHydrationBoundary>
|
||||
<NextRouteStateBridge>
|
||||
<AppContextProvider>
|
||||
<EventEmitterContextProvider>
|
||||
<ProviderContextProvider>
|
||||
<ModalContextProvider>
|
||||
<MainNavLayout>
|
||||
<RoleRouteGuard>
|
||||
{children}
|
||||
</RoleRouteGuard>
|
||||
</MainNavLayout>
|
||||
<InSiteMessageNotification />
|
||||
<PartnerStack />
|
||||
<ReadmePanel />
|
||||
<GotoAnything />
|
||||
<WorkflowGeneratorMount />
|
||||
</ModalContextProvider>
|
||||
</ProviderContextProvider>
|
||||
</EventEmitterContextProvider>
|
||||
</AppContextProvider>
|
||||
<div className="flex h-full flex-col overflow-hidden">
|
||||
<MaintenanceNotice />
|
||||
<AppContextProvider>
|
||||
<EventEmitterContextProvider>
|
||||
<ProviderContextProvider>
|
||||
<ModalContextProvider>
|
||||
<MainNavLayout>
|
||||
<RoleRouteGuard>
|
||||
{children}
|
||||
</RoleRouteGuard>
|
||||
</MainNavLayout>
|
||||
<InSiteMessageNotification />
|
||||
<PartnerStack />
|
||||
<ReadmePanel />
|
||||
<GotoAnything />
|
||||
<WorkflowGeneratorMount />
|
||||
</ModalContextProvider>
|
||||
</ProviderContextProvider>
|
||||
</EventEmitterContextProvider>
|
||||
</AppContextProvider>
|
||||
</div>
|
||||
</NextRouteStateBridge>
|
||||
</CommonLayoutHydrationBoundary>
|
||||
<Zendesk />
|
||||
</>
|
||||
</React.Fragment>
|
||||
)
|
||||
}
|
||||
|
||||
@@ -2,17 +2,65 @@ import type { LegacyPluginsSearchParams } from '@/app/components/plugins/plugin-
|
||||
import Marketplace from '@/app/components/plugins/marketplace'
|
||||
import PluginPage from '@/app/components/plugins/plugin-page'
|
||||
import PluginsPanel from '@/app/components/plugins/plugin-page/plugins-panel'
|
||||
import { getLegacyPluginRedirectPath } from '@/app/components/plugins/plugin-routes'
|
||||
import {
|
||||
getFirstPackageIdFromSearchParams,
|
||||
getInstallRedirectPathByPluginCategory,
|
||||
getInstallRedirectPathFromSearchParams,
|
||||
getLegacyPluginRedirectPath,
|
||||
shouldResolveInstallCategoryRedirect,
|
||||
} from '@/app/components/plugins/plugin-routes'
|
||||
import { MARKETPLACE_API_PREFIX } from '@/config'
|
||||
import { redirect } from '@/next/navigation'
|
||||
|
||||
type PluginListProps = {
|
||||
searchParams?: Promise<LegacyPluginsSearchParams>
|
||||
}
|
||||
|
||||
type MarketplaceManifestCategoryResponse = {
|
||||
data?: {
|
||||
plugin?: {
|
||||
category?: string
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const fetchPluginCategoryFromMarketplace = async (packageId: string) => {
|
||||
try {
|
||||
const response = await fetch(
|
||||
`${MARKETPLACE_API_PREFIX}/plugins/identifier?unique_identifier=${encodeURIComponent(packageId)}`,
|
||||
{ cache: 'no-store' },
|
||||
)
|
||||
|
||||
if (!response.ok)
|
||||
return undefined
|
||||
|
||||
const payload = await response.json() as MarketplaceManifestCategoryResponse
|
||||
return payload.data?.plugin?.category
|
||||
}
|
||||
catch {
|
||||
return undefined
|
||||
}
|
||||
}
|
||||
|
||||
const PluginList = async ({
|
||||
searchParams,
|
||||
}: PluginListProps) => {
|
||||
const redirectPath = getLegacyPluginRedirectPath(await searchParams)
|
||||
const resolvedSearchParams = await searchParams ?? {}
|
||||
const installRedirectPathFromSearchParams = getInstallRedirectPathFromSearchParams(resolvedSearchParams)
|
||||
|
||||
if (installRedirectPathFromSearchParams)
|
||||
redirect(installRedirectPathFromSearchParams)
|
||||
|
||||
if (shouldResolveInstallCategoryRedirect(resolvedSearchParams)) {
|
||||
const packageId = getFirstPackageIdFromSearchParams(resolvedSearchParams)
|
||||
const category = packageId ? await fetchPluginCategoryFromMarketplace(packageId) : undefined
|
||||
const installRedirectPath = getInstallRedirectPathByPluginCategory(category, resolvedSearchParams)
|
||||
|
||||
if (installRedirectPath)
|
||||
redirect(installRedirectPath)
|
||||
}
|
||||
|
||||
const redirectPath = getLegacyPluginRedirectPath(resolvedSearchParams)
|
||||
|
||||
if (redirectPath)
|
||||
redirect(redirectPath)
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
import type { ReactNode } from 'react'
|
||||
import * as React from 'react'
|
||||
import { CommonLayoutHydrationBoundary } from '@/app/(commonLayout)/hydration-boundary'
|
||||
import AmplitudeProvider from '@/app/components/base/amplitude'
|
||||
import { GoogleAnalyticsScripts } from '@/app/components/base/ga'
|
||||
import { EducationVerifyActionRecorder } from '@/app/components/education-verify-action-recorder'
|
||||
import HeaderWrapper from '@/app/components/header/header-wrapper'
|
||||
import MaintenanceNotice from '@/app/components/header/maintenance-notice'
|
||||
import { OAuthRegistrationAnalytics } from '@/app/components/oauth-registration-analytics'
|
||||
import { AppContextProvider } from '@/context/app-context-provider'
|
||||
import { EventEmitterContextProvider } from '@/context/event-emitter-provider'
|
||||
@@ -12,30 +12,32 @@ import { ModalContextProvider } from '@/context/modal-context-provider'
|
||||
import { ProviderContextProvider } from '@/context/provider-context-provider'
|
||||
import Header from './header'
|
||||
|
||||
const Layout = async ({ children }: { children: ReactNode }) => {
|
||||
export default async function Layout({ children }: { children: React.ReactNode }) {
|
||||
return (
|
||||
<>
|
||||
<React.Fragment>
|
||||
<GoogleAnalyticsScripts />
|
||||
<AmplitudeProvider />
|
||||
<OAuthRegistrationAnalytics />
|
||||
<EducationVerifyActionRecorder />
|
||||
<CommonLayoutHydrationBoundary>
|
||||
<AppContextProvider>
|
||||
<EventEmitterContextProvider>
|
||||
<ProviderContextProvider>
|
||||
<ModalContextProvider>
|
||||
<HeaderWrapper>
|
||||
<Header />
|
||||
</HeaderWrapper>
|
||||
<div className="relative flex h-0 shrink-0 grow flex-col overflow-y-auto bg-components-panel-bg">
|
||||
{children}
|
||||
</div>
|
||||
</ModalContextProvider>
|
||||
</ProviderContextProvider>
|
||||
</EventEmitterContextProvider>
|
||||
</AppContextProvider>
|
||||
<div className="flex h-full flex-col overflow-hidden bg-background-body">
|
||||
<MaintenanceNotice />
|
||||
<AppContextProvider>
|
||||
<EventEmitterContextProvider>
|
||||
<ProviderContextProvider>
|
||||
<ModalContextProvider>
|
||||
<HeaderWrapper>
|
||||
<Header />
|
||||
</HeaderWrapper>
|
||||
<div className="relative flex h-0 min-h-0 shrink-0 grow flex-col overflow-y-auto bg-components-panel-bg">
|
||||
{children}
|
||||
</div>
|
||||
</ModalContextProvider>
|
||||
</ProviderContextProvider>
|
||||
</EventEmitterContextProvider>
|
||||
</AppContextProvider>
|
||||
</div>
|
||||
</CommonLayoutHydrationBoundary>
|
||||
</>
|
||||
</React.Fragment>
|
||||
)
|
||||
}
|
||||
export default Layout
|
||||
|
||||
@@ -1,60 +1,40 @@
|
||||
'use client'
|
||||
import { cn } from '@langgenius/dify-ui/cn'
|
||||
import { useQuery, useSuspenseQuery } from '@tanstack/react-query'
|
||||
import type { ReactNode } from 'react'
|
||||
import { useSuspenseQuery } from '@tanstack/react-query'
|
||||
|
||||
import Loading from '@/app/components/base/loading'
|
||||
import Header from '@/app/signin/_header'
|
||||
import { AppContextProvider } from '@/context/app-context-provider'
|
||||
import { isLegacyBase401, userProfileQueryOptions } from '@/features/account-profile/client'
|
||||
import { systemFeaturesQueryOptions } from '@/features/system-features/client'
|
||||
import useDocumentTitle from '@/hooks/use-document-title'
|
||||
|
||||
export default function SignInLayout({ children }: any) {
|
||||
type Props = {
|
||||
children: ReactNode
|
||||
}
|
||||
|
||||
const copyrightYear = new Date().getFullYear()
|
||||
|
||||
export default function OAuthAuthorizeLayout({ children }: Props) {
|
||||
const { data: systemFeatures } = useSuspenseQuery(systemFeaturesQueryOptions())
|
||||
useDocumentTitle('')
|
||||
// Probe login state. 401 stays as `error` (not thrown) so this layout can render
|
||||
// the signin/oauth UI for unauthenticated users; other errors bubble to error.tsx.
|
||||
// (When unauthenticated, service/base.ts's auto-redirect to /signin still fires.)
|
||||
const { isPending, data: userResp, error } = useQuery({
|
||||
...userProfileQueryOptions(),
|
||||
throwOnError: err => !isLegacyBase401(err),
|
||||
})
|
||||
const isLoggedIn = !!userResp && !error
|
||||
|
||||
if (isPending) {
|
||||
return (
|
||||
<div className="flex min-h-screen w-full justify-center bg-background-default-burn">
|
||||
<Loading />
|
||||
</div>
|
||||
)
|
||||
}
|
||||
return (
|
||||
<>
|
||||
<div className={cn('flex min-h-screen w-full justify-center bg-background-default-burn p-6')}>
|
||||
<div className={cn('flex w-full shrink-0 flex-col items-center rounded-2xl border border-effects-highlight bg-background-default-subtle')}>
|
||||
<Header />
|
||||
<div className={cn('flex w-full grow flex-col items-center justify-center px-6 md:px-[108px]')}>
|
||||
<div className="flex flex-col md:w-[400px]">
|
||||
{isLoggedIn
|
||||
? (
|
||||
<AppContextProvider>
|
||||
{children}
|
||||
</AppContextProvider>
|
||||
)
|
||||
: children}
|
||||
</div>
|
||||
<div className="flex min-h-screen w-full justify-center bg-background-default-burn p-6">
|
||||
<div className="flex w-full shrink-0 flex-col items-center rounded-2xl border border-effects-highlight bg-background-default-subtle">
|
||||
<Header />
|
||||
<div className="flex w-full grow flex-col items-center justify-center px-6 md:px-[108px]">
|
||||
<div className="flex flex-col md:w-[400px]">
|
||||
{children}
|
||||
</div>
|
||||
{systemFeatures.branding.enabled === false && (
|
||||
<div className="px-8 py-6 system-xs-regular text-text-tertiary">
|
||||
©
|
||||
{' '}
|
||||
{new Date().getFullYear()}
|
||||
{' '}
|
||||
LangGenius, Inc. All rights reserved.
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
{systemFeatures.branding.enabled === false && (
|
||||
<div className="px-8 py-6 system-xs-regular text-text-tertiary">
|
||||
©
|
||||
{' '}
|
||||
{copyrightYear}
|
||||
{' '}
|
||||
LangGenius, Inc. All rights reserved.
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ let mockChatConversationDetail: Record<string, unknown> | undefined
|
||||
let mockCompletionConversationDetail: Record<string, unknown> | undefined
|
||||
let mockShowMessageLogModal = false
|
||||
let mockShowPromptLogModal = false
|
||||
let mockShowAgentLogModal = false
|
||||
let mockCurrentLogItem: Record<string, unknown> | undefined
|
||||
let mockCurrentLogModalActiveTab = 'messages'
|
||||
|
||||
@@ -81,6 +82,7 @@ vi.mock('@/app/components/app/store', () => ({
|
||||
setShowAgentLogModal: mockSetShowAgentLogModal,
|
||||
setShowMessageLogModal: mockSetShowMessageLogModal,
|
||||
showPromptLogModal: mockShowPromptLogModal,
|
||||
showAgentLogModal: mockShowAgentLogModal,
|
||||
currentLogModalActiveTab: mockCurrentLogModalActiveTab,
|
||||
}),
|
||||
}))
|
||||
@@ -126,6 +128,7 @@ vi.mock('@/app/components/base/chat/chat', () => ({
|
||||
onAnnotationEdited,
|
||||
onAnnotationRemoved,
|
||||
switchSibling,
|
||||
hideLogModal,
|
||||
}: {
|
||||
chatList: Array<{ id: string }>
|
||||
onFeedback: (mid: string, value: { rating: string, content?: string }) => Promise<boolean>
|
||||
@@ -133,8 +136,9 @@ vi.mock('@/app/components/base/chat/chat', () => ({
|
||||
onAnnotationEdited: (query: string, answer: string, index: number) => void
|
||||
onAnnotationRemoved: (index: number) => Promise<boolean>
|
||||
switchSibling: (siblingMessageId: string) => void
|
||||
hideLogModal?: boolean
|
||||
}) => (
|
||||
<div data-testid="chat-panel">
|
||||
<div data-testid="chat-panel" data-hide-log-modal={String(hideLogModal)}>
|
||||
<div>{chatList.length}</div>
|
||||
<button onClick={() => void onFeedback('message-1', { rating: 'like', content: 'nice' })}>chat-feedback</button>
|
||||
<button onClick={() => onAnnotationAdded('annotation-2', 'Admin', 'Edited question', 'Edited answer', 1)}>chat-add-annotation</button>
|
||||
@@ -145,6 +149,14 @@ vi.mock('@/app/components/base/chat/chat', () => ({
|
||||
),
|
||||
}))
|
||||
|
||||
vi.mock('@/app/components/base/agent-log-modal', () => ({
|
||||
default: ({ floating, onCancel }: { floating?: boolean, onCancel: () => void }) => (
|
||||
<div data-testid="agent-log-modal" data-floating={String(floating)}>
|
||||
<button onClick={onCancel}>close-agent-log-modal</button>
|
||||
</div>
|
||||
),
|
||||
}))
|
||||
|
||||
vi.mock('@/app/components/base/message-log-modal', () => ({
|
||||
default: ({ onCancel }: { onCancel: () => void }) => (
|
||||
<div data-testid="message-log-modal">
|
||||
@@ -255,6 +267,7 @@ describe('ConversationList', () => {
|
||||
mockCompletionConversationDetail = undefined
|
||||
mockShowMessageLogModal = false
|
||||
mockShowPromptLogModal = false
|
||||
mockShowAgentLogModal = false
|
||||
mockCurrentLogItem = undefined
|
||||
mockCurrentLogModalActiveTab = 'messages'
|
||||
mockDelAnnotation.mockResolvedValue(undefined)
|
||||
@@ -383,6 +396,7 @@ describe('ConversationList', () => {
|
||||
|
||||
expect(screen.getByTestId('var-panel')).toHaveTextContent('query:Latest question')
|
||||
expect(screen.getByTestId('model-info')).toHaveTextContent('gpt-4o')
|
||||
expect(screen.getByTestId('chat-panel')).toHaveAttribute('data-hide-log-modal', 'true')
|
||||
expect(screen.getByTestId('message-log-modal')).toBeInTheDocument()
|
||||
|
||||
fireEvent.click(screen.getByText('chat-feedback'))
|
||||
@@ -399,6 +413,61 @@ describe('ConversationList', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('should mount agent log modals from the detail panel instead of the nested chat layout', async () => {
|
||||
mockChatConversationDetail = {
|
||||
id: 'conversation-1',
|
||||
created_at: 1710000000,
|
||||
model_config: {
|
||||
model: 'gpt-4o',
|
||||
configs: {
|
||||
introduction: 'Hello there',
|
||||
},
|
||||
user_input_form: [],
|
||||
},
|
||||
message: {
|
||||
inputs: {},
|
||||
},
|
||||
}
|
||||
mockShowAgentLogModal = true
|
||||
mockCurrentLogItem = {
|
||||
id: 'message-1',
|
||||
conversationId: 'conversation-1',
|
||||
}
|
||||
mockFetchChatMessages.mockResolvedValue({
|
||||
data: [
|
||||
{
|
||||
id: 'message-1',
|
||||
answer: 'Assistant reply',
|
||||
query: 'Latest question',
|
||||
created_at: 1710000000,
|
||||
inputs: {},
|
||||
feedbacks: [],
|
||||
message: [],
|
||||
message_files: [],
|
||||
agent_thoughts: [{ id: 'thought-1' }],
|
||||
},
|
||||
],
|
||||
has_more: false,
|
||||
})
|
||||
|
||||
renderConversationList({
|
||||
searchParams: '?page=2&conversation_id=conversation-1',
|
||||
})
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId('chat-panel')).toBeInTheDocument()
|
||||
})
|
||||
|
||||
expect(screen.getByTestId('chat-panel')).toHaveAttribute('data-hide-log-modal', 'true')
|
||||
expect(screen.getByTestId('agent-log-modal')).toBeInTheDocument()
|
||||
expect(screen.getByTestId('agent-log-modal')).toHaveAttribute('data-floating', 'true')
|
||||
|
||||
fireEvent.click(screen.getByText('close-agent-log-modal'))
|
||||
|
||||
expect(mockSetCurrentLogItem).toHaveBeenCalled()
|
||||
expect(mockSetShowAgentLogModal).toHaveBeenCalledWith(false)
|
||||
})
|
||||
|
||||
it('should render completion details and refetch after feedback updates', async () => {
|
||||
mockCompletionConversationDetail = {
|
||||
id: 'conversation-1',
|
||||
@@ -424,7 +493,7 @@ describe('ConversationList', () => {
|
||||
},
|
||||
}
|
||||
mockShowPromptLogModal = true
|
||||
mockCurrentLogItem = { id: 'log-2' }
|
||||
mockCurrentLogItem = { id: 'log-2', log: [{ role: 'user', text: 'Prompt body' }] }
|
||||
|
||||
renderConversationList({
|
||||
appDetail: { id: 'app-1', mode: AppModeEnum.COMPLETION } as any,
|
||||
@@ -626,7 +695,7 @@ describe('ConversationList', () => {
|
||||
},
|
||||
}
|
||||
mockShowPromptLogModal = true
|
||||
mockCurrentLogItem = { id: 'log-2' }
|
||||
mockCurrentLogItem = { id: 'log-2', log: [{ role: 'user', text: 'Prompt body' }] }
|
||||
|
||||
renderConversationList({
|
||||
appDetail: { id: 'app-1', mode: AppModeEnum.COMPLETION } as any,
|
||||
|
||||
@@ -36,6 +36,7 @@ import ModelInfo from '@/app/components/app/log/model-info'
|
||||
import { useStore as useAppStore } from '@/app/components/app/store'
|
||||
import TextGeneration from '@/app/components/app/text-generate/item'
|
||||
import ActionButton from '@/app/components/base/action-button'
|
||||
import AgentLogModal from '@/app/components/base/agent-log-modal'
|
||||
import Chat from '@/app/components/base/chat/chat'
|
||||
import CopyIcon from '@/app/components/base/copy-icon'
|
||||
import Loading from '@/app/components/base/loading'
|
||||
@@ -165,13 +166,25 @@ function DetailPanel({ detail, onFeedback }: IDetailPanel) {
|
||||
})
|
||||
const { formatTime } = useTimestamp()
|
||||
const { onClose, appDetail } = useContext(DrawerContext)
|
||||
const { currentLogItem, setCurrentLogItem, showMessageLogModal, setShowMessageLogModal, showPromptLogModal, setShowPromptLogModal, currentLogModalActiveTab } = useAppStore(useShallow((state: AppStoreState) => ({
|
||||
const {
|
||||
currentLogItem,
|
||||
setCurrentLogItem,
|
||||
showMessageLogModal,
|
||||
setShowMessageLogModal,
|
||||
showPromptLogModal,
|
||||
setShowPromptLogModal,
|
||||
showAgentLogModal,
|
||||
setShowAgentLogModal,
|
||||
currentLogModalActiveTab,
|
||||
} = useAppStore(useShallow((state: AppStoreState) => ({
|
||||
currentLogItem: state.currentLogItem,
|
||||
setCurrentLogItem: state.setCurrentLogItem,
|
||||
showMessageLogModal: state.showMessageLogModal,
|
||||
setShowMessageLogModal: state.setShowMessageLogModal,
|
||||
showPromptLogModal: state.showPromptLogModal,
|
||||
setShowPromptLogModal: state.setShowPromptLogModal,
|
||||
showAgentLogModal: state.showAgentLogModal,
|
||||
setShowAgentLogModal: state.setShowAgentLogModal,
|
||||
currentLogModalActiveTab: state.currentLogModalActiveTab,
|
||||
})))
|
||||
const { t } = useTranslation()
|
||||
@@ -395,6 +408,7 @@ function DetailPanel({ detail, onFeedback }: IDetailPanel) {
|
||||
|
||||
const isChatMode = appDetail?.mode !== AppModeEnum.COMPLETION
|
||||
const isAdvanced = appDetail?.mode === AppModeEnum.ADVANCED_CHAT
|
||||
const shouldShowPromptLogModal = showPromptLogModal && !!currentLogItem?.log
|
||||
|
||||
const varList = getDetailVarList(detail, varValues)
|
||||
const message_files = getCompletionMessageFiles(detail, isChatMode)
|
||||
@@ -507,6 +521,7 @@ function DetailPanel({ detail, onFeedback }: IDetailPanel) {
|
||||
noChatInput
|
||||
showPromptLog
|
||||
hideProcessDetail
|
||||
hideLogModal
|
||||
chatContainerInnerClassName="px-3"
|
||||
switchSibling={switchSibling}
|
||||
/>
|
||||
@@ -546,6 +561,7 @@ function DetailPanel({ detail, onFeedback }: IDetailPanel) {
|
||||
noChatInput
|
||||
showPromptLog
|
||||
hideProcessDetail
|
||||
hideLogModal
|
||||
chatContainerInnerClassName="px-3"
|
||||
switchSibling={switchSibling}
|
||||
/>
|
||||
@@ -574,7 +590,18 @@ function DetailPanel({ detail, onFeedback }: IDetailPanel) {
|
||||
/>
|
||||
</WorkflowContextProvider>
|
||||
)}
|
||||
{!isChatMode && showPromptLogModal && (
|
||||
{showAgentLogModal && (
|
||||
<AgentLogModal
|
||||
floating
|
||||
width={width}
|
||||
currentLogItem={currentLogItem}
|
||||
onCancel={() => {
|
||||
setCurrentLogItem()
|
||||
setShowAgentLogModal(false)
|
||||
}}
|
||||
/>
|
||||
)}
|
||||
{shouldShowPromptLogModal && (
|
||||
<PromptLogModal
|
||||
width={width}
|
||||
currentLogItem={currentLogItem}
|
||||
|
||||
@@ -1058,11 +1058,20 @@ export function AppCard({ app, onlineUsers = [], onRefresh, onOpenTagManagement
|
||||
</div>
|
||||
</div>
|
||||
<div className="flex h-[26px] shrink-0 items-start px-3" />
|
||||
<div className="flex min-w-0 shrink-0 items-center pt-2 pr-4 pb-3 pl-4 system-xs-regular text-text-tertiary">
|
||||
<div
|
||||
className={cn(
|
||||
'flex min-w-0 shrink-0 items-center overflow-hidden pt-2 pb-3 pl-4 system-xs-regular text-text-tertiary',
|
||||
app.access_mode ? 'pr-9' : 'pr-4',
|
||||
)}
|
||||
>
|
||||
<div className="flex min-w-0 flex-1 items-center gap-1 whitespace-nowrap">
|
||||
<div className="truncate">{app.author_name}</div>
|
||||
<div className="shrink-0">·</div>
|
||||
<div className="truncate">{editTimeText}</div>
|
||||
{app.author_name && (
|
||||
<>
|
||||
<div className="min-w-0 truncate">{app.author_name}</div>
|
||||
<div className="shrink-0">·</div>
|
||||
</>
|
||||
)}
|
||||
<div className="min-w-0 truncate">{editTimeText}</div>
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
|
||||
@@ -119,6 +119,17 @@ describe('AgentLogModal', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('should render the floating modal through a dialog portal', () => {
|
||||
vi.mocked(fetchAgentLogDetail).mockReturnValue(new Promise(() => {}))
|
||||
|
||||
const { container } = render(<AgentLogModal {...mockProps} floating />)
|
||||
|
||||
const modal = screen.getByRole('dialog')
|
||||
expect(container).not.toContainElement(modal)
|
||||
expect(document.body).toContainElement(modal)
|
||||
expect(modal).toHaveClass('fixed', 'z-50', 'w-[480px]!', 'left-[max(8px,calc(100vw-1136px))]!')
|
||||
})
|
||||
|
||||
it('should call onCancel when close button is clicked', () => {
|
||||
vi.mocked(fetchAgentLogDetail).mockReturnValue(new Promise(() => {}))
|
||||
|
||||
@@ -158,4 +169,18 @@ describe('AgentLogModal', () => {
|
||||
|
||||
expect(mockProps.onCancel).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should not use click-away to close the floating dialog', () => {
|
||||
vi.mocked(fetchAgentLogDetail).mockReturnValue(new Promise(() => {}))
|
||||
|
||||
let clickAwayHandler!: (event: Event) => void
|
||||
vi.mocked(useClickAway).mockImplementation((callback) => {
|
||||
clickAwayHandler = callback
|
||||
})
|
||||
|
||||
render(<AgentLogModal {...mockProps} floating />)
|
||||
clickAwayHandler(new Event('click'))
|
||||
|
||||
expect(mockProps.onCancel).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import type { FC } from 'react'
|
||||
import type { IChatItem } from '@/app/components/base/chat/chat/type'
|
||||
import { cn } from '@langgenius/dify-ui/cn'
|
||||
import { Dialog, DialogContent, DialogTitle } from '@langgenius/dify-ui/dialog'
|
||||
import { RiCloseLine } from '@remixicon/react'
|
||||
import { useClickAway } from 'ahooks'
|
||||
import { useEffect, useRef, useState } from 'react'
|
||||
@@ -10,11 +11,13 @@ import AgentLogDetail from './detail'
|
||||
type AgentLogModalProps = Readonly<{
|
||||
currentLogItem?: IChatItem
|
||||
width: number
|
||||
floating?: boolean
|
||||
onCancel: () => void
|
||||
}>
|
||||
const AgentLogModal: FC<AgentLogModalProps> = ({
|
||||
currentLogItem,
|
||||
width,
|
||||
floating,
|
||||
onCancel,
|
||||
}) => {
|
||||
const { t } = useTranslation()
|
||||
@@ -22,7 +25,7 @@ const AgentLogModal: FC<AgentLogModalProps> = ({
|
||||
const [mounted, setMounted] = useState(false)
|
||||
|
||||
useClickAway(() => {
|
||||
if (mounted)
|
||||
if (mounted && !floating)
|
||||
onCancel()
|
||||
}, ref)
|
||||
|
||||
@@ -33,6 +36,44 @@ const AgentLogModal: FC<AgentLogModalProps> = ({
|
||||
if (!currentLogItem || !currentLogItem.conversationId)
|
||||
return null
|
||||
|
||||
const detailContent = (
|
||||
<>
|
||||
<AgentLogDetail
|
||||
conversationID={currentLogItem.conversationId}
|
||||
messageID={currentLogItem.id}
|
||||
log={currentLogItem}
|
||||
/>
|
||||
</>
|
||||
)
|
||||
|
||||
if (floating) {
|
||||
return (
|
||||
<Dialog
|
||||
open
|
||||
onOpenChange={(open) => {
|
||||
if (!open)
|
||||
onCancel()
|
||||
}}
|
||||
>
|
||||
<DialogContent
|
||||
backdropClassName="bg-transparent!"
|
||||
className="top-16! bottom-4! left-[max(8px,calc(100vw-1136px))]! flex max-h-none! w-[480px]! max-w-[calc(100vw-16px)]! translate-x-0! translate-y-0! flex-col overflow-hidden! rounded-xl! border-[0.5px]! border-components-panel-border! bg-components-panel-bg! p-0! pt-3! pb-3! shadow-xl!"
|
||||
>
|
||||
<DialogTitle className="text-md shrink-0 px-4 py-1 font-semibold text-text-primary">{t('runDetail.workflowTitle', { ns: 'appLog' })}</DialogTitle>
|
||||
<button
|
||||
type="button"
|
||||
aria-label={t('operation.close', { ns: 'common' })}
|
||||
className="absolute top-4 right-3 z-20 cursor-pointer border-none bg-transparent p-1 focus-visible:ring-1 focus-visible:ring-components-input-border-active focus-visible:outline-hidden"
|
||||
onClick={onCancel}
|
||||
>
|
||||
<RiCloseLine className="size-4 text-text-tertiary" aria-hidden="true" />
|
||||
</button>
|
||||
{detailContent}
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
)
|
||||
}
|
||||
|
||||
return (
|
||||
<div
|
||||
className={cn('relative z-10 flex flex-col rounded-xl border-[0.5px] border-components-panel-border bg-components-panel-bg py-3 shadow-xl')}
|
||||
@@ -54,11 +95,7 @@ const AgentLogModal: FC<AgentLogModalProps> = ({
|
||||
>
|
||||
<RiCloseLine className="size-4 text-text-tertiary" aria-hidden="true" />
|
||||
</button>
|
||||
<AgentLogDetail
|
||||
conversationID={currentLogItem.conversationId}
|
||||
messageID={currentLogItem.id}
|
||||
log={currentLogItem}
|
||||
/>
|
||||
{detailContent}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
@@ -57,6 +57,9 @@ const ChatWrapper = () => {
|
||||
} = useChatWithHistoryContext()
|
||||
|
||||
const appSourceType = isInstalledApp ? AppSourceType.installedApp : AppSourceType.webApp
|
||||
const timezone = appSourceType === AppSourceType.webApp
|
||||
? Intl.DateTimeFormat().resolvedOptions().timeZone
|
||||
: undefined
|
||||
|
||||
// Semantic variable for better code readability
|
||||
const isHistoryConversation = !!currentConversationId
|
||||
@@ -91,6 +94,8 @@ const ChatWrapper = () => {
|
||||
taskId => stopChatMessageResponding('', taskId, appSourceType, appId),
|
||||
clearChatList,
|
||||
setClearChatList,
|
||||
undefined,
|
||||
{ timezone },
|
||||
)
|
||||
const inputsFormValue = currentConversationId ? currentConversationInputs : newConversationInputsRef?.current
|
||||
const inputDisabled = useMemo(() => {
|
||||
|
||||
@@ -6,6 +6,10 @@ import { useParams, usePathname } from '@/next/navigation'
|
||||
import { sseGet, ssePost } from '@/service/base'
|
||||
import { useChat } from '../hooks'
|
||||
|
||||
const useTimestampMock = vi.hoisted(() =>
|
||||
vi.fn(() => ({ formatTime: vi.fn().mockReturnValue('10:00 AM') })),
|
||||
)
|
||||
|
||||
vi.mock('@/service/base', () => ({
|
||||
sseGet: vi.fn(),
|
||||
ssePost: vi.fn(),
|
||||
@@ -31,7 +35,7 @@ vi.mock('@langgenius/dify-ui/toast', () => ({
|
||||
}))
|
||||
|
||||
vi.mock('@/hooks/use-timestamp', () => ({
|
||||
default: () => ({ formatTime: vi.fn().mockReturnValue('10:00 AM') }),
|
||||
default: useTimestampMock,
|
||||
}))
|
||||
|
||||
vi.mock('@/next/navigation', () => ({
|
||||
@@ -91,6 +95,21 @@ describe('useChat', () => {
|
||||
expect(result.current.suggestedQuestions).toEqual([])
|
||||
})
|
||||
|
||||
it('should pass timestamp options to timestamp formatter', () => {
|
||||
renderHook(() => useChat(
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ timezone: 'UTC' },
|
||||
))
|
||||
|
||||
expect(useTimestampMock).toHaveBeenCalledWith({ timezone: 'UTC' })
|
||||
})
|
||||
|
||||
it('should initialize with opening statement and suggested questions', () => {
|
||||
const config = {
|
||||
opening_statement: 'Hello {{name}}',
|
||||
|
||||
@@ -55,6 +55,10 @@ type SendCallback = {
|
||||
isPublicAPI?: boolean
|
||||
}
|
||||
|
||||
type UseChatOptions = {
|
||||
timezone?: string
|
||||
}
|
||||
|
||||
export const useChat = (
|
||||
config?: ChatConfig,
|
||||
formSettings?: {
|
||||
@@ -66,9 +70,10 @@ export const useChat = (
|
||||
clearChatList?: boolean,
|
||||
clearChatListCallback?: (state: boolean) => void,
|
||||
initialConversationId?: string,
|
||||
options: UseChatOptions = {},
|
||||
) => {
|
||||
const { t } = useTranslation()
|
||||
const { formatTime } = useTimestamp()
|
||||
const { formatTime } = useTimestamp({ timezone: options.timezone })
|
||||
const conversationIdRef = useRef(initialConversationId ?? '')
|
||||
const initialConversationIdRef = useRef(initialConversationId ?? '')
|
||||
const hasStopRespondedRef = useRef(false)
|
||||
|
||||
@@ -81,6 +81,9 @@ const ChatWrapper = () => {
|
||||
opening_statement: currentConversationItem?.introduction || (config as any).opening_statement,
|
||||
} as ChatConfig
|
||||
}, [appParams, currentConversationItem?.introduction])
|
||||
const timezone = appSourceType === AppSourceType.webApp
|
||||
? Intl.DateTimeFormat().resolvedOptions().timeZone
|
||||
: undefined
|
||||
const {
|
||||
chatList,
|
||||
handleSend,
|
||||
@@ -98,6 +101,8 @@ const ChatWrapper = () => {
|
||||
taskId => stopChatMessageResponding('', taskId, appSourceType, appId),
|
||||
clearChatList,
|
||||
setClearChatList,
|
||||
undefined,
|
||||
{ timezone },
|
||||
)
|
||||
const inputsFormValue = currentConversationId ? currentConversationInputs : newConversationInputsRef?.current
|
||||
const inputDisabled = useMemo(() => {
|
||||
|
||||
+5
-1
@@ -109,7 +109,11 @@ describe('InputField', () => {
|
||||
const scrollBody = panel?.children[1]
|
||||
const footer = panel?.lastElementChild
|
||||
|
||||
expect(panel).toHaveClass('max-h-(--shortcut-popup-max-height)', 'overflow-hidden')
|
||||
// The max-height falls back to a viewport unit so the panel stays bounded
|
||||
// (and the footer/actions reachable via the internal scroll) even when it is
|
||||
// rendered outside the shortcuts popup that defines --shortcut-popup-max-height,
|
||||
// e.g. inside the edit dialog. See issue #37979.
|
||||
expect(panel).toHaveClass('max-h-[var(--shortcut-popup-max-height,80dvh)]', 'overflow-hidden')
|
||||
expect(header).toHaveClass('shrink-0', 'pb-2')
|
||||
expect(scrollBody).toHaveClass('min-h-0', 'flex-1', 'overflow-y-auto')
|
||||
expect(footer).toHaveClass('shrink-0', 'bg-components-panel-bg')
|
||||
|
||||
@@ -216,7 +216,7 @@ const InputField: React.FC<InputFieldProps> = ({
|
||||
}, [handleSave])
|
||||
|
||||
return (
|
||||
<div className="flex max-h-(--shortcut-popup-max-height) w-[372px] flex-col overflow-hidden rounded-xl border-[0.5px] border-components-panel-border bg-components-panel-bg-blur shadow-lg backdrop-blur-[5px]">
|
||||
<div className="flex max-h-[var(--shortcut-popup-max-height,80dvh)] w-[372px] flex-col overflow-hidden rounded-xl border-[0.5px] border-components-panel-border bg-components-panel-bg-blur shadow-lg backdrop-blur-[5px]">
|
||||
<div className="shrink-0 p-3 pb-2">
|
||||
<div className="system-md-semibold text-text-primary">{t(`${i18nPrefix}.title`, { ns: 'workflow' })}</div>
|
||||
</div>
|
||||
|
||||
@@ -6,6 +6,7 @@ let fetching = false
|
||||
let isManager = true
|
||||
let enableBilling = true
|
||||
let workspacePermissionKeys: string[] = ['billing.subscription.manage']
|
||||
let billingUrlEnabled = false
|
||||
|
||||
const refetchMock = vi.fn()
|
||||
const openAsyncWindowMock = vi.fn()
|
||||
@@ -19,11 +20,14 @@ type BillingWindowOptions = {
|
||||
type OpenAsyncWindowCall = [BillingUrlCallback, BillingWindowOptions]
|
||||
|
||||
vi.mock('@/service/use-billing', () => ({
|
||||
useBillingUrl: () => ({
|
||||
data: currentBillingUrl,
|
||||
isFetching: fetching,
|
||||
refetch: refetchMock,
|
||||
}),
|
||||
useBillingUrl: (enabled: boolean) => {
|
||||
billingUrlEnabled = enabled
|
||||
return {
|
||||
data: currentBillingUrl,
|
||||
isFetching: fetching,
|
||||
refetch: refetchMock,
|
||||
}
|
||||
},
|
||||
}))
|
||||
|
||||
vi.mock('@/hooks/use-async-window-open', () => ({
|
||||
@@ -54,28 +58,32 @@ describe('Billing', () => {
|
||||
fetching = false
|
||||
isManager = true
|
||||
enableBilling = true
|
||||
billingUrlEnabled = false
|
||||
workspacePermissionKeys = ['billing.subscription.manage']
|
||||
refetchMock.mockResolvedValue({ data: 'https://billing' })
|
||||
})
|
||||
|
||||
it('shows the billing action when subscription management permission is granted without manager role', () => {
|
||||
it('hides the billing action when subscription management permission is granted without manager role', () => {
|
||||
isManager = false
|
||||
|
||||
render(<Billing />)
|
||||
|
||||
expect(screen.getByRole('button', { name: /billing\.viewBillingTitle/ })).toBeInTheDocument()
|
||||
expect(screen.queryByRole('button', { name: /billing\.viewBillingTitle/ })).not.toBeInTheDocument()
|
||||
expect(billingUrlEnabled).toBe(false)
|
||||
})
|
||||
|
||||
it('hides the billing action when subscription management permission is missing or billing is disabled', () => {
|
||||
workspacePermissionKeys = []
|
||||
render(<Billing />)
|
||||
expect(screen.queryByRole('button', { name: /billing\.viewBillingTitle/ })).not.toBeInTheDocument()
|
||||
expect(billingUrlEnabled).toBe(false)
|
||||
|
||||
vi.clearAllMocks()
|
||||
workspacePermissionKeys = ['billing.subscription.manage']
|
||||
enableBilling = false
|
||||
render(<Billing />)
|
||||
expect(screen.queryByRole('button', { name: /billing\.viewBillingTitle/ })).not.toBeInTheDocument()
|
||||
expect(billingUrlEnabled).toBe(false)
|
||||
})
|
||||
|
||||
it('opens the billing window with the immediate url when the button is clicked', async () => {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user