Compare commits
48
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7570427f51 | ||
|
|
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 |
@@ -12,7 +12,6 @@ Use this as the component decision guide for Dify web. Existing code is referenc
|
||||
| Question | Default | Promote or extract only when |
|
||||
| --- | --- | --- |
|
||||
| Where should code live? | Keep it local to the feature workflow, route, or owner. | Multiple verticals need the same stable primitive. |
|
||||
| How should route/tab folders be named? | Match the current route segment, tab name, or user-visible surface. | Keep a historical or broader parent only when it still owns multiple surfaces. |
|
||||
| Who owns state, data, and handlers? | The lowest component that uses them. | A parent coordinates shared loading, errors, empty UI, selection, submission, navigation, or one consistent snapshot. |
|
||||
| Should this become Jotai state? | Keep synchronous UI/form state in component or DOM state. | Siblings need one source of truth, the value drives atoms, or scoped workflow state must survive hidden/unmounted steps. |
|
||||
| Should URL state enter Jotai? | Let Next.js route params and `nuqs` own URL state and updates. | Query atoms or shared derived atoms need a read-only bridge hydrated at the route/surface boundary. |
|
||||
@@ -24,7 +23,7 @@ Use this as the component decision guide for Dify web. Existing code is referenc
|
||||
|
||||
- Search before adding UI, hooks, helpers, query utilities, or styling patterns. Reuse existing base components, feature components, hooks, utilities, and design styles when they fit.
|
||||
- Follow Dify's CSS-first Tailwind v4 contract from `packages/dify-ui/README.md` and `packages/dify-ui/AGENTS.md`. Prefer design-system tokens, utilities, and radius mappings over generic Tailwind choices.
|
||||
- Group feature code by workflow, route, or ownership area with route-aligned names: components, hooks, local types, query helpers, atoms, constants, tests, and small utilities should live near the code that changes with them.
|
||||
- Group feature code by workflow, route, or ownership area: components, hooks, local types, query helpers, atoms, constants, and small utilities should live near the code that changes with them.
|
||||
- Keep source/default selection, validation, dirty checks, and payload shaping close to the workflow that owns submit behavior. Do not hide flow-specific priority order, fallback behavior, or submit semantics in generic utilities.
|
||||
- Prefer direct conditionals for small branch-specific decisions, especially form source selection and request payload assembly.
|
||||
- Loading states for page sections, cards, lists, tables, forms, and drawers should be skeletons scoped to the content being loaded. Use spinners only for small inline busy indicators.
|
||||
@@ -33,8 +32,6 @@ Use this as the component decision guide for Dify web. Existing code is referenc
|
||||
|
||||
- State-heavy wizards, drawers, modals, and secondary workflows can be a small feature surface: an entry file, one feature-local state file when Jotai is actually needed, and shallow `ui/` owners that match real visual regions.
|
||||
- The entry file handles route integration, provider wiring, close behavior, and surface mounting. The composition owner handles high-level workflow branching. The closest visual owner handles section branching.
|
||||
- When a page or tab maps to a route segment, name its feature folder after that route/tab surface instead of a stale parent grouping. Remove misleading intermediate folders when only one surface remains.
|
||||
- When a tab folder grows into several independent sections or action areas, split the first level by product/visual owners. Keep the root for the entry component and cross-owner state, colocate tests with the owner folder, and put truly shared local UI under a specifically named `components/` file.
|
||||
- Repeated TanStack query calls in sibling components are acceptable when each component independently consumes the data; TanStack Query deduplicates and shares cache.
|
||||
- Pass stable domain identity across boundaries. Do not forward derived presentation state when the receiver can derive it from its own data source.
|
||||
- A component that owns a visual surface should also own data access, loading, empty, and error states for content rendered inside it unless a parent truly coordinates that state.
|
||||
@@ -49,7 +46,6 @@ Use this as the component decision guide for Dify web. Existing code is referenc
|
||||
- Use uncontrolled `@langgenius/dify-ui/form` and `@langgenius/dify-ui/field` controls for edit/create forms whose fields are read only at submit time. Initialize query-backed defaults with `defaultValue` and keyed remounts.
|
||||
- Promote form state to atoms only when another component must react to in-progress values, a draft must survive unmount/remount in the scoped workflow, or multiple steps share the same editable draft before submit.
|
||||
- Treat `useParams`, route args, and `nuqs` query state as framework-owned state. When atom logic needs those values, hydrate primitive atoms at the route or surface boundary, such as with `useHydrateAtoms(..., { dangerouslyForceHydrate: true })`; keep URL updates in the route/query-state APIs instead of write atoms.
|
||||
- Within a route-owned feature, choose one source for route identity. If route params are bridged into feature atoms, use that bridge consistently for route-derived queries and actions instead of also threading the same route id through page, tab, and section props.
|
||||
- For async work tied to atom state, use `atomWithQuery` or `atomWithMutation`; write atoms should update only the inputs that drive those atoms. This applies to pure frontend async work as well as network requests, so do not hand-roll loading/error/in-flight state with `useState` or `useRef` for atom-orchestrated async behavior. For component-owned remote work, use `useQuery` or `useMutation` directly.
|
||||
- Row-local async state belongs to the row owner unless it participates in a shared Jotai workflow or needs atom-scoped reset semantics.
|
||||
- Leave query and mutation atoms unscoped so they keep shared QueryClient cache and invalidation behavior. Scope resettable primitives and explicit hydration tuples; scope a derived atom only when every dependency should be private to that surface.
|
||||
@@ -64,10 +60,8 @@ Use this as the component decision guide for Dify web. Existing code is referenc
|
||||
- Type component signatures directly; do not use `FC` or `React.FC`.
|
||||
- Prefer `function` for top-level components and module helpers. Use arrow functions for local callbacks, handlers, and lambda-style APIs.
|
||||
- Prefer named exports. Use default exports only where the framework requires them, such as Next.js route files.
|
||||
- Avoid barrel files that only re-export secondary owners. `index.tsx` is acceptable for a route/tab entry component; import header controls, switches, sections, and row owners from their concrete owner files.
|
||||
- Type simple one-off props inline. Use a named `Props` type only when reused, exported, complex, or clearer.
|
||||
- Use API-generated or API-returned types at component boundaries. Keep small UI conversion helpers and one-off UI extensions beside the component that needs them.
|
||||
- Avoid `common.tsx` buckets for shared UI. Use a feature-local `components/` folder with concrete filenames that describe the shared role.
|
||||
- Do not create type aliases that only rename another type. Use aliases only for real UI concepts, refinements, or reusable local contracts.
|
||||
- Name values by their domain role and backend API contract, especially persistent IDs and route params. Normalize framework or route params at the boundary.
|
||||
- Put fallback and invariant checks in the lowest component that already handles that state. Do not extract helpers whose only behavior is hiding missing display data.
|
||||
@@ -89,13 +83,11 @@ Use this as the component decision guide for Dify web. Existing code is referenc
|
||||
|
||||
- Keep `web/contract/*` as the API shape source of truth and follow the `{ params, query?, body? }` input shape.
|
||||
- Consume generated queries with `useQuery(consoleQuery.xxx.queryOptions(...))` or `useQuery(marketplaceQuery.xxx.queryOptions(...))`.
|
||||
- If a generated query input comes from an atom, including a route-identity bridge atom, keep the query in `atomWithQuery`; do not unwrap the atom in a component just to call `useQuery`.
|
||||
- Consume owner-local mutations with `useMutation(consoleQuery.xxx.mutationOptions(...))` or `useMutation(marketplaceQuery.xxx.mutationOptions(...))` when pending/error state is not consumed by feature atoms.
|
||||
- In `atomWithQuery`, `atomWithInfiniteQuery`, and `atomWithMutation`, return generated `queryOptions()`, `infiniteOptions()`, or `mutationOptions()` directly. Pass `enabled`, `retry`, `placeholderData`, `select`, and pagination options into the generated call instead of spreading options into a hand-built object.
|
||||
- For generated oRPC options with missing required input, branch the whole input with `input: condition ? validInput : skipToken` and `enabled: Boolean(condition)`. Never place `skipToken` inside a nested placeholder payload or coerce required IDs to `''`.
|
||||
- When prefetch and render use the same request, extract local query options or a query-options atom so `prefetchQuery` and `useQuery`/`atomWithQuery` share the exact options.
|
||||
- For custom query or mutation functions, wrap options with TanStack `queryOptions(...)` or `mutationOptions(...)`.
|
||||
- Do not extract generated `queryOptions(...)` into a helper solely to share input construction; extract only when prefetch/render must share exact options or the helper owns real domain behavior.
|
||||
- Avoid pass-through hooks and thin `web/service/use-*` wrappers that only rename generated options. Keep feature hooks for real orchestration, workflow state, or shared domain behavior.
|
||||
- Put shared cache behavior in `createTanstackQueryUtils(...experimental_defaults...)`. Component or atom callbacks may handle local toasts, closing dialogs, and navigation, but should not replace shared invalidation or patch shared server state locally.
|
||||
- For overlays that may open heavier secondary content, prefetch from the trigger/menu open event with `queryClient.prefetchQuery(queryOptions)` when `onOpenChange` is available. Do not mount hidden subscribers just to warm cache.
|
||||
@@ -104,7 +96,7 @@ Use this as the component decision guide for Dify web. Existing code is referenc
|
||||
|
||||
## Boundaries And Overlays
|
||||
|
||||
- Use the first level below a page or tab to organize independent page sections when it adds structure or the root folder becomes noisy. This layer is layout/semantic first, not automatically the data owner.
|
||||
- Use the first level below a page or tab to organize independent page sections when it adds structure. This layer is layout/semantic first, not automatically the data owner.
|
||||
- Treat component names, semantic roles, and user- or design-marked visual regions as boundary constraints. Keep adjacent UI as a sibling owner or introduce a correctly named broader owner.
|
||||
- Keep cohesive forms, menu bodies, and one-off helpers local unless they need their own state, reuse, or semantic boundary.
|
||||
- Separate hidden secondary surfaces from the trigger's main flow. For dialogs, dropdowns, popovers, and similar branches, extract a small local component when hidden content would obscure the parent.
|
||||
|
||||
@@ -53,8 +53,6 @@ jobs:
|
||||
|
||||
- name: Run Type Checks
|
||||
if: steps.changed-files.outputs.any_changed == 'true'
|
||||
env:
|
||||
PYREFLY_OUTPUT_FORMAT: github
|
||||
run: make type-check-core
|
||||
|
||||
- name: Dotenv check
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -167,16 +167,12 @@ register_schema_models(
|
||||
ChatMessagesQuery,
|
||||
MessageFeedbackPayload,
|
||||
FeedbackExportQuery,
|
||||
)
|
||||
register_response_schema_models(
|
||||
console_ns,
|
||||
AnnotationCountResponse,
|
||||
SuggestedQuestionsResponse,
|
||||
MessageDetailResponse,
|
||||
MessageInfiniteScrollPaginationResponse,
|
||||
SimpleResultResponse,
|
||||
TextFileResponse,
|
||||
)
|
||||
register_response_schema_models(console_ns, SimpleResultResponse, TextFileResponse)
|
||||
|
||||
|
||||
@console_ns.route("/apps/<uuid:app_id>/chat-messages")
|
||||
|
||||
@@ -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."
|
||||
|
||||
@@ -85,7 +85,6 @@ def _published_app_filter():
|
||||
class InstalledAppInfoResponse(ResponseModel):
|
||||
id: str
|
||||
name: str | None = None
|
||||
description: str | None = None
|
||||
mode: str | None = None
|
||||
icon_type: str | None = None
|
||||
icon: str | None = None
|
||||
@@ -124,7 +123,6 @@ class InstalledAppResponse(ResponseModel):
|
||||
return {
|
||||
"id": _safe_primitive(getattr(value, "id", "")) or "",
|
||||
"name": _safe_primitive(getattr(value, "name", None)),
|
||||
"description": _safe_primitive(getattr(value, "description", None)),
|
||||
"mode": _safe_primitive(getattr(value, "mode", None)),
|
||||
"icon_type": _safe_primitive(getattr(value, "icon_type", None)),
|
||||
"icon": _safe_primitive(getattr(value, "icon", None)),
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -20,7 +20,7 @@ openapi_ns = Namespace("openapi", description="User-scoped operations", path="/"
|
||||
|
||||
# Register response/query models BEFORE importing controller modules so that
|
||||
# @openapi_ns.response / @openapi_ns.expect decorators can resolve model names.
|
||||
from controllers.common.fields import EventStreamResponse, SimpleResultResponse
|
||||
from controllers.common.fields import EventStreamResponse
|
||||
from controllers.common.schema import register_enum_models, register_response_schema_models, register_schema_models
|
||||
from controllers.openapi._models import (
|
||||
AccountPayload,
|
||||
@@ -95,7 +95,6 @@ register_response_schema_models(
|
||||
openapi_ns,
|
||||
ErrorBody,
|
||||
EventStreamResponse,
|
||||
SimpleResultResponse,
|
||||
UsageInfo,
|
||||
MessageMetadata,
|
||||
AppListRow,
|
||||
|
||||
@@ -34,7 +34,6 @@ class OpenApiErrorCode(StrEnum):
|
||||
# transport-generic (resolved from HTTP status for plain werkzeug raises)
|
||||
BAD_REQUEST = "bad_request"
|
||||
UNAUTHORIZED = "unauthorized"
|
||||
TOKEN_EXPIRED = "token_expired"
|
||||
FORBIDDEN = "forbidden"
|
||||
NOT_FOUND = "not_found"
|
||||
METHOD_NOT_ALLOWED = "method_not_allowed"
|
||||
@@ -224,19 +223,6 @@ class OpenApiErrorFormatter:
|
||||
return isinstance(part, (str, int)) and not isinstance(part, bool)
|
||||
|
||||
|
||||
class InvalidBearer(OpenApiError): # noqa: N818
|
||||
code = 401
|
||||
error_code = OpenApiErrorCode.UNAUTHORIZED
|
||||
description = "Invalid or unknown bearer token."
|
||||
|
||||
|
||||
class SessionExpired(OpenApiError): # noqa: N818
|
||||
code = 401
|
||||
error_code = OpenApiErrorCode.TOKEN_EXPIRED
|
||||
description = "Your session has expired."
|
||||
hint = "Re-authenticate to continue (e.g. re-run your login command)."
|
||||
|
||||
|
||||
class FilenameNotExists(OpenApiError): # noqa: N818
|
||||
code = 400
|
||||
error_code = OpenApiErrorCode.FILENAME_NOT_EXISTS
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable, Generator
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
@@ -61,7 +61,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _translate_service_errors() -> Generator[None, None, None]:
|
||||
def _translate_service_errors() -> Iterator[None]:
|
||||
try:
|
||||
yield
|
||||
except WorkflowNotFoundError as ex:
|
||||
@@ -166,7 +166,6 @@ class AppRunApi(Resource):
|
||||
surface="apps",
|
||||
)
|
||||
|
||||
# response-contract:ignore compact_generate_response
|
||||
return helper.compact_generate_response(stream_obj)
|
||||
|
||||
|
||||
|
||||
@@ -25,13 +25,12 @@ from controllers.openapi._models import (
|
||||
AppListRow,
|
||||
)
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind, RBACRequirement
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from controllers.service_api.app.error import AppUnavailableError
|
||||
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
|
||||
from extensions.ext_database import db
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
from models import App
|
||||
from models.enums import AppStatus
|
||||
from models.model import AppMode
|
||||
from services.account_service import TenantService
|
||||
from services.app_service import AppListParams, AppService
|
||||
@@ -167,9 +166,7 @@ class AppListApi(Resource):
|
||||
# an empty set or list means the caller has no accessible apps.
|
||||
# End-users bypass RBAC here — their access is controlled by scope upstream.
|
||||
apply_rbac_filter = (
|
||||
dify_config.RBAC_ENABLED
|
||||
and auth_data.caller_kind != CallerKind.END_USER
|
||||
and auth_data.account_id is not None
|
||||
dify_config.RBAC_ENABLED and auth_data.caller_kind != "end_user" and auth_data.account_id is not None
|
||||
)
|
||||
access_filter = AppAccessFilter.unrestricted()
|
||||
if apply_rbac_filter:
|
||||
@@ -206,7 +203,7 @@ class AppListApi(Resource):
|
||||
limit=query.limit,
|
||||
mode=query.mode.value if query.mode else "all", # type:ignore
|
||||
name=query.name,
|
||||
status=AppStatus.NORMAL,
|
||||
status="normal",
|
||||
# Visibility gate pushed into the query — pagination.total stays
|
||||
# consistent across pages because invisible rows never count.
|
||||
openapi_visible=True,
|
||||
|
||||
@@ -25,7 +25,6 @@ from controllers.openapi.auth.data import AuthData, Edition
|
||||
from extensions.ext_database import db
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
from models import App
|
||||
from models.enums import AppStatus
|
||||
from services.account_service import TenantService
|
||||
from services.app_service import AppService
|
||||
from services.enterprise.app_permitted_service import list_permitted_apps
|
||||
@@ -63,7 +62,7 @@ class PermittedExternalAppsListApi(Resource):
|
||||
items: list[AppListRow] = []
|
||||
for app_id in page_result.app_ids:
|
||||
app = apps_by_id.get(app_id)
|
||||
if not app or app.status != AppStatus.NORMAL:
|
||||
if not app or app.status != "normal":
|
||||
continue
|
||||
tenant = tenants_by_id.get(str(app.tenant_id))
|
||||
items.append(
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from enum import StrEnum
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from werkzeug.exceptions import InternalServerError
|
||||
@@ -20,11 +21,6 @@ class Edition(StrEnum):
|
||||
SAAS = "saas"
|
||||
|
||||
|
||||
class CallerKind(StrEnum):
|
||||
ACCOUNT = "account"
|
||||
END_USER = "end_user"
|
||||
|
||||
|
||||
def current_edition() -> Edition:
|
||||
if dify_config.EDITION == "CLOUD":
|
||||
return Edition.SAAS
|
||||
@@ -82,9 +78,9 @@ class AuthData(BaseModel):
|
||||
tenant_role: TenantAccountRole | None = None
|
||||
|
||||
caller: Account | EndUser | None = None
|
||||
caller_kind: CallerKind | None = None
|
||||
caller_kind: Literal["account", "end_user"] | None = None
|
||||
|
||||
def require_app_context(self) -> tuple[App, Account | EndUser, CallerKind]:
|
||||
def require_app_context(self) -> tuple[App, Account | EndUser, Literal["account", "end_user"]]:
|
||||
if self.app is None or self.caller is None or self.caller_kind is None:
|
||||
raise InternalServerError("pipeline_invariant_violated: app context missing")
|
||||
return self.app, self.caller, self.caller_kind
|
||||
|
||||
@@ -17,7 +17,6 @@ from flask_login import user_logged_in
|
||||
from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
|
||||
|
||||
from controllers.openapi._audit import emit_wrong_surface
|
||||
from controllers.openapi._errors import InvalidBearer, SessionExpired
|
||||
from controllers.openapi.auth.data import (
|
||||
AuthData,
|
||||
Edition,
|
||||
@@ -29,9 +28,7 @@ from controllers.openapi.auth.data import (
|
||||
from controllers.openapi.auth.flow import When
|
||||
from libs.oauth_bearer import (
|
||||
AuthContext,
|
||||
InvalidBearerError,
|
||||
Scope,
|
||||
TokenExpiredError,
|
||||
TokenType,
|
||||
extract_bearer,
|
||||
get_authenticator,
|
||||
@@ -220,12 +217,7 @@ class PipelineRouter:
|
||||
if not token:
|
||||
raise Unauthorized("bearer required")
|
||||
|
||||
try:
|
||||
identity = get_authenticator().authenticate(token)
|
||||
except TokenExpiredError:
|
||||
raise SessionExpired()
|
||||
except InvalidBearerError:
|
||||
raise InvalidBearer()
|
||||
identity = get_authenticator().authenticate(token)
|
||||
|
||||
if allowed_token_types is not None and identity.token_type not in allowed_token_types:
|
||||
emit_wrong_surface(
|
||||
|
||||
@@ -5,10 +5,10 @@ import uuid
|
||||
from flask import request
|
||||
from werkzeug.exceptions import Forbidden, InternalServerError, NotFound, Unauthorized
|
||||
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind
|
||||
from controllers.openapi.auth.data import AuthData
|
||||
from extensions.ext_database import db
|
||||
from models.account import AccountStatus, TenantStatus
|
||||
from models.enums import AppStatus, EndUserType
|
||||
from models.account import TenantStatus
|
||||
from models.enums import EndUserType
|
||||
from services.account_service import AccountService, TenantService
|
||||
from services.app_service import AppService
|
||||
from services.end_user_service import EndUserService
|
||||
@@ -24,7 +24,7 @@ def load_app(data: AuthData) -> None:
|
||||
except ValueError:
|
||||
raise NotFound("app not found")
|
||||
app = AppService.get_app_by_id(db.session, app_id)
|
||||
if not app or app.status != AppStatus.NORMAL:
|
||||
if not app or app.status != "normal":
|
||||
raise NotFound("app not found")
|
||||
data.app = app
|
||||
|
||||
@@ -65,7 +65,7 @@ def load_account(data: AuthData) -> None:
|
||||
if data.tenant:
|
||||
account.current_tenant = data.tenant
|
||||
data.caller = account
|
||||
data.caller_kind = CallerKind.ACCOUNT
|
||||
data.caller_kind = "account"
|
||||
|
||||
|
||||
def load_workspace_role(data: AuthData) -> None:
|
||||
@@ -73,7 +73,7 @@ def load_workspace_role(data: AuthData) -> None:
|
||||
return
|
||||
if data.tenant is None or data.account_id is None:
|
||||
return
|
||||
if data.caller is not None and getattr(data.caller, "status", None) != AccountStatus.ACTIVE:
|
||||
if data.caller is not None and getattr(data.caller, "status", None) != "active":
|
||||
return
|
||||
role = TenantService.get_account_role_in_tenant(db.session, str(data.account_id), str(data.tenant.id))
|
||||
if role is None:
|
||||
@@ -91,7 +91,7 @@ def resolve_external_user(data: AuthData) -> None:
|
||||
user_id=data.external_identity.email,
|
||||
)
|
||||
data.caller = end_user
|
||||
data.caller_kind = CallerKind.END_USER
|
||||
data.caller_kind = "end_user"
|
||||
|
||||
|
||||
def load_app_access_mode(data: AuthData) -> None:
|
||||
|
||||
@@ -5,7 +5,7 @@ from werkzeug.exceptions import Forbidden, NotFound, UnprocessableEntity
|
||||
|
||||
from configs import dify_config
|
||||
from controllers.common.wraps import enforce_rbac_access
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind
|
||||
from controllers.openapi.auth.data import AuthData
|
||||
from extensions.ext_database import db
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
from services.account_service import AccountService, TenantService
|
||||
@@ -58,7 +58,7 @@ def check_rbac_permission(data: AuthData) -> None:
|
||||
if not dify_config.RBAC_ENABLED:
|
||||
return
|
||||
# Only account callers are subject to RBAC; end_user access is scope-controlled.
|
||||
if data.caller_kind != CallerKind.ACCOUNT:
|
||||
if data.caller_kind != "account":
|
||||
return
|
||||
if data.account_id is None or data.tenant is None:
|
||||
raise Forbidden("rbac context missing")
|
||||
|
||||
@@ -22,7 +22,7 @@ from controllers.openapi._contract import accepts, returns
|
||||
from controllers.openapi._errors import HumanInputFormNotFound, RecipientSurfaceMismatch
|
||||
from controllers.openapi._models import FormSubmitResponse, HumanInputFormDefinitionResponse
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind, RBACRequirement
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from core.workflow.human_input_policy import (
|
||||
HumanInputSurface,
|
||||
is_recipient_type_allowed_for_surface,
|
||||
@@ -98,7 +98,7 @@ class OpenApiWorkflowHumanInputFormApi(Resource):
|
||||
|
||||
submission_user_id: str | None = None
|
||||
submission_end_user_id: str | None = None
|
||||
if caller_kind == CallerKind.ACCOUNT:
|
||||
if caller_kind == "account":
|
||||
submission_user_id = caller.id
|
||||
else:
|
||||
submission_end_user_id = caller.id
|
||||
|
||||
@@ -22,7 +22,7 @@ from controllers.common.schema import query_params_from_model
|
||||
from controllers.common.wraps import RBACPermission, RBACResourceScope
|
||||
from controllers.openapi import openapi_ns
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind, RBACRequirement
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator
|
||||
from core.app.apps.base_app_generator import BaseAppGenerator
|
||||
from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter
|
||||
@@ -70,7 +70,7 @@ class OpenApiWorkflowEventsApi(Resource):
|
||||
if workflow_run.app_id != app_model.id:
|
||||
raise NotFound("Workflow run not found")
|
||||
|
||||
if caller_kind == CallerKind.ACCOUNT:
|
||||
if caller_kind == "account":
|
||||
if workflow_run.created_by_role != CreatorUserRole.ACCOUNT or workflow_run.created_by != caller.id:
|
||||
raise NotFound("Workflow run not found")
|
||||
else:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -2,8 +2,8 @@
|
||||
|
||||
This checker intentionally stays conservative. It only reports a hard schema
|
||||
mismatch when both sides are statically known for the same 2xx status code:
|
||||
a documented ``@ns.response(..., Model)`` and an actual ``dump_response(Model, ...)``,
|
||||
``Model(...).model_dump()``, or ``Model.model_validate(...).model_dump()`` return.
|
||||
a documented ``@ns.response(..., Model)`` and an actual ``dump_response(Model, ...)``
|
||||
or ``Model.model_validate(...).model_dump()`` return.
|
||||
|
||||
Raw dictionaries, raw lists, ``None`` responses, streaming helpers, missing
|
||||
response schemas, and returns with non-literal status codes are classified as
|
||||
@@ -28,7 +28,6 @@ from typing import Any, Literal
|
||||
HTTP_METHODS = {"delete", "get", "head", "options", "patch", "post", "put"}
|
||||
NO_BODY_STATUSES = {HTTPStatus.NO_CONTENT.value, HTTPStatus.RESET_CONTENT.value, HTTPStatus.NOT_MODIFIED.value}
|
||||
DEFAULT_CONTROLLER_DIRS = ("controllers/console", "controllers/service_api", "controllers/web")
|
||||
IGNORE_COMMENT_MARKERS = ("response-contract:ignore",)
|
||||
|
||||
type Classification = Literal["valid", "mismatch", "unknown", "refactorable"]
|
||||
type ActualKind = Literal[
|
||||
@@ -42,7 +41,6 @@ type ActualKind = Literal[
|
||||
"unknown",
|
||||
]
|
||||
type MethodNode = ast.FunctionDef | ast.AsyncFunctionDef
|
||||
type ModelValueSource = Literal["constructor", "model_validate"]
|
||||
|
||||
HTTP_STATUS_NAMES = {status.name: status.value for status in HTTPStatus}
|
||||
HTTP_STATUS_NAMES.update({f"HTTP_{status.value}_{status.name}": status.value for status in HTTPStatus})
|
||||
@@ -111,22 +109,18 @@ class VariableAssignmentSummary:
|
||||
"""Track whether a local name is safe to treat as one specific response model."""
|
||||
|
||||
known_models: set[str] = field(default_factory=set)
|
||||
known_sources: set[ModelValueSource] = field(default_factory=set)
|
||||
has_unknown_assignment: bool = False
|
||||
|
||||
def add_known(self, model: str, source: ModelValueSource) -> None:
|
||||
def add_known(self, model: str) -> None:
|
||||
self.known_models.add(model)
|
||||
self.known_sources.add(source)
|
||||
|
||||
def add_unknown(self) -> None:
|
||||
self.has_unknown_assignment = True
|
||||
|
||||
def single_known_model(self) -> tuple[str, ModelValueSource] | None:
|
||||
def single_known_model(self) -> str | None:
|
||||
if self.has_unknown_assignment or len(self.known_models) != 1:
|
||||
return None
|
||||
model = next(iter(self.known_models))
|
||||
source: ModelValueSource = "constructor" if self.known_sources == {"constructor"} else "model_validate"
|
||||
return model, source
|
||||
return next(iter(self.known_models))
|
||||
|
||||
|
||||
def dotted_name(node: ast.AST) -> str | None:
|
||||
@@ -255,12 +249,6 @@ def model_name_from_model_validate_call(node: ast.AST) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def model_value_from_model_validate_call(node: ast.AST) -> tuple[str, ModelValueSource] | None:
|
||||
if model_name := model_name_from_model_validate_call(node):
|
||||
return model_name, "model_validate"
|
||||
return None
|
||||
|
||||
|
||||
def model_name_from_constructor_call(node: ast.AST) -> str | None:
|
||||
if not isinstance(node, ast.Call):
|
||||
return None
|
||||
@@ -269,12 +257,6 @@ def model_name_from_constructor_call(node: ast.AST) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def model_value_from_constructor_call(node: ast.AST) -> tuple[str, ModelValueSource] | None:
|
||||
if model_name := model_name_from_constructor_call(node):
|
||||
return model_name, "constructor"
|
||||
return None
|
||||
|
||||
|
||||
def model_name_from_model_dump(node: ast.AST) -> str | None:
|
||||
if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute) or node.func.attr != "model_dump":
|
||||
return None
|
||||
@@ -290,10 +272,6 @@ def model_name_from_model_value(node: ast.AST) -> str | None:
|
||||
return model_name_from_model_validate_call(node) or model_name_from_constructor_call(node)
|
||||
|
||||
|
||||
def model_value_from_model_value(node: ast.AST) -> tuple[str, ModelValueSource] | None:
|
||||
return model_value_from_model_validate_call(node) or model_value_from_constructor_call(node)
|
||||
|
||||
|
||||
def model_name_from_dump_response(node: ast.AST) -> str | None:
|
||||
if not isinstance(node, ast.Call):
|
||||
return None
|
||||
@@ -309,7 +287,7 @@ def model_name_from_dump_response(node: ast.AST) -> str | None:
|
||||
|
||||
|
||||
def actual_kind_from_expr(
|
||||
expr: ast.AST | None, variable_models: dict[str, tuple[str, ModelValueSource]] | None = None
|
||||
expr: ast.AST | None, variable_models: dict[str, str] | None = None
|
||||
) -> tuple[ActualKind, str | None]:
|
||||
if expr is None:
|
||||
return "none", None
|
||||
@@ -321,14 +299,10 @@ def actual_kind_from_expr(
|
||||
if isinstance(expr, ast.Call) and isinstance(expr.func, ast.Attribute) and expr.func.attr == "model_dump":
|
||||
dumped_value = expr.func.value
|
||||
if isinstance(dumped_value, ast.Name) and variable_models:
|
||||
model_assignment = variable_models.get(dumped_value.id)
|
||||
if model_assignment:
|
||||
model_name, source = model_assignment
|
||||
if source == "constructor":
|
||||
return "model", model_name
|
||||
# A variable dump from model_validate can match today, but it
|
||||
# bypasses dump_response and is easier to drift; keep it visible
|
||||
# as refactorable.
|
||||
# A variable dump can match today, but it bypasses dump_response and
|
||||
# is easier to drift; keep it visible as refactorable.
|
||||
model_name = variable_models.get(dumped_value.id)
|
||||
if model_name:
|
||||
return "model_dump_variable", model_name
|
||||
|
||||
model_dump_model = model_name_from_model_dump(expr)
|
||||
@@ -351,9 +325,7 @@ def actual_kind_from_expr(
|
||||
return "unknown", None
|
||||
|
||||
|
||||
def actual_response_from_return(
|
||||
return_node: ast.Return, variable_models: dict[str, tuple[str, ModelValueSource]]
|
||||
) -> ActualResponse:
|
||||
def actual_response_from_return(return_node: ast.Return, variable_models: dict[str, str]) -> ActualResponse:
|
||||
status: int | None = 200
|
||||
body_expr = return_node.value
|
||||
|
||||
@@ -391,21 +363,18 @@ def target_names(target: ast.AST) -> Iterable[str]:
|
||||
|
||||
|
||||
def record_assignment(
|
||||
assignments: defaultdict[str, VariableAssignmentSummary],
|
||||
targets: Iterable[str],
|
||||
model_assignment: tuple[str, ModelValueSource] | None,
|
||||
assignments: defaultdict[str, VariableAssignmentSummary], targets: Iterable[str], model_name: str | None
|
||||
) -> None:
|
||||
for target in targets:
|
||||
if model_assignment is None:
|
||||
if model_name is None:
|
||||
# Once a name receives an unknown value, later model_dump() calls on it
|
||||
# are no longer a reliable signal for the returned schema.
|
||||
assignments[target].add_unknown()
|
||||
else:
|
||||
model_name, source = model_assignment
|
||||
assignments[target].add_known(model_name, source)
|
||||
assignments[target].add_known(model_name)
|
||||
|
||||
|
||||
def variable_model_assignments_for_method(method: MethodNode) -> dict[str, tuple[str, ModelValueSource]]:
|
||||
def variable_model_assignments_for_method(method: MethodNode) -> dict[str, str]:
|
||||
"""Infer local variables that are unambiguously assigned one response model."""
|
||||
|
||||
assignments: defaultdict[str, VariableAssignmentSummary] = defaultdict(VariableAssignmentSummary)
|
||||
@@ -416,10 +385,10 @@ def variable_model_assignments_for_method(method: MethodNode) -> dict[str, tuple
|
||||
record_assignment(
|
||||
assignments,
|
||||
(name for target in targets for name in target_names(target)),
|
||||
model_value_from_model_value(value),
|
||||
model_name_from_model_value(value),
|
||||
)
|
||||
case ast.AnnAssign(target=target, value=value) if value is not None:
|
||||
record_assignment(assignments, target_names(target), model_value_from_model_value(value))
|
||||
record_assignment(assignments, target_names(target), model_name_from_model_value(value))
|
||||
case ast.AugAssign(target=target) | ast.For(target=target) | ast.AsyncFor(target=target):
|
||||
# Mutation and loop targets overwrite prior values with runtime-dependent data.
|
||||
record_assignment(assignments, target_names(target), None)
|
||||
@@ -430,13 +399,9 @@ def variable_model_assignments_for_method(method: MethodNode) -> dict[str, tuple
|
||||
case ast.ExceptHandler(name=name) if name:
|
||||
assignments[name].add_unknown()
|
||||
case ast.NamedExpr(target=target, value=value):
|
||||
record_assignment(assignments, target_names(target), model_value_from_model_value(value))
|
||||
record_assignment(assignments, target_names(target), model_name_from_model_value(value))
|
||||
|
||||
return {
|
||||
name: assignment
|
||||
for name, summary in assignments.items()
|
||||
if (assignment := summary.single_known_model()) is not None
|
||||
}
|
||||
return {name: model for name, summary in assignments.items() if (model := summary.single_known_model()) is not None}
|
||||
|
||||
|
||||
def actual_responses_for_method(method: MethodNode) -> list[ActualResponse]:
|
||||
@@ -580,52 +545,13 @@ def iter_controller_files(paths: Iterable[Path]) -> Iterable[Path]:
|
||||
yield from sorted(child for child in path.rglob("*.py") if child.is_file())
|
||||
|
||||
|
||||
def node_start_lineno(node: ast.ClassDef | MethodNode) -> int:
|
||||
decorator_lines = [decorator.lineno for decorator in node.decorator_list]
|
||||
if decorator_lines:
|
||||
return min(decorator_lines)
|
||||
return node.lineno
|
||||
|
||||
|
||||
def line_has_ignore_marker(line: str) -> bool:
|
||||
_, marker, comment = line.partition("#")
|
||||
if not marker:
|
||||
return False
|
||||
normalized = comment.lower()
|
||||
return any(ignore_marker in normalized for ignore_marker in IGNORE_COMMENT_MARKERS)
|
||||
|
||||
|
||||
def node_has_ignore_comment(lines: Sequence[str], node: ast.ClassDef | MethodNode) -> bool:
|
||||
start = node_start_lineno(node)
|
||||
end = node.end_lineno or node.lineno
|
||||
if any(line_has_ignore_marker(line) for line in lines[start - 1 : end]):
|
||||
return True
|
||||
|
||||
line_index = start - 2
|
||||
while line_index >= 0:
|
||||
stripped = lines[line_index].strip()
|
||||
if not stripped:
|
||||
line_index -= 1
|
||||
continue
|
||||
if not stripped.startswith("#"):
|
||||
break
|
||||
if line_has_ignore_marker(lines[line_index]):
|
||||
return True
|
||||
line_index -= 1
|
||||
return False
|
||||
|
||||
|
||||
def checks_for_file(file_path: Path, repo_root: Path) -> list[ContractCheck]:
|
||||
source = file_path.read_text(encoding="utf-8")
|
||||
lines = source.splitlines()
|
||||
module = ast.parse(source, filename=str(file_path))
|
||||
module = ast.parse(file_path.read_text(encoding="utf-8"), filename=str(file_path))
|
||||
checks: list[ContractCheck] = []
|
||||
|
||||
for node in module.body:
|
||||
if not isinstance(node, ast.ClassDef):
|
||||
continue
|
||||
if node_has_ignore_comment(lines, node):
|
||||
continue
|
||||
|
||||
class_routes = routes_from_decorators(node.decorator_list)
|
||||
class_documented = response_docs_from_decorators(node.decorator_list)
|
||||
@@ -633,8 +559,6 @@ def checks_for_file(file_path: Path, repo_root: Path) -> list[ContractCheck]:
|
||||
for item in node.body:
|
||||
if not isinstance(item, ast.FunctionDef | ast.AsyncFunctionDef) or item.name not in HTTP_METHODS:
|
||||
continue
|
||||
if node_has_ignore_comment(lines, item):
|
||||
continue
|
||||
|
||||
routes = routes_from_decorators(item.decorator_list) or class_routes
|
||||
if not routes:
|
||||
|
||||
@@ -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,
|
||||
|
||||
+1
-2
@@ -7,8 +7,7 @@ class ResponseModel(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
extra="ignore",
|
||||
validate_by_name=True,
|
||||
validate_by_alias=True,
|
||||
populate_by_name=True,
|
||||
serialize_by_alias=True,
|
||||
protected_namespaces=(),
|
||||
)
|
||||
|
||||
+10
-26
@@ -236,16 +236,6 @@ class TokenExpiredError(Exception):
|
||||
"""Hard-expire bookkeeping is the resolver's job before raising."""
|
||||
|
||||
|
||||
class NegativeCache(StrEnum):
|
||||
"""Negative cache markers. ``EXPIRED`` is distinct from ``INVALID`` so a
|
||||
retry inside ``NEGATIVE_TTL`` still reports expiry instead of collapsing
|
||||
into a generic unknown-token miss.
|
||||
"""
|
||||
|
||||
INVALID = "invalid"
|
||||
EXPIRED = "expired"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Registry
|
||||
# ============================================================================
|
||||
@@ -353,15 +343,13 @@ class OAuthAccessTokenResolver:
|
||||
def _cache_key(self, token_hash: str) -> str:
|
||||
return TOKEN_CACHE_KEY_FMT.format(hash=token_hash)
|
||||
|
||||
def cache_get(self, token_hash: str) -> ResolvedRow | None | NegativeCache:
|
||||
def cache_get(self, token_hash: str) -> ResolvedRow | None | Literal["invalid"]:
|
||||
raw = self._redis.get(self._cache_key(token_hash))
|
||||
if raw is None:
|
||||
return None
|
||||
text = raw.decode() if isinstance(raw, (bytes, bytearray)) else raw
|
||||
try:
|
||||
return NegativeCache(text)
|
||||
except ValueError:
|
||||
pass
|
||||
if text == "invalid":
|
||||
return "invalid"
|
||||
try:
|
||||
return ResolvedRow.from_cache(json.loads(text))
|
||||
except (ValueError, KeyError):
|
||||
@@ -375,8 +363,8 @@ class OAuthAccessTokenResolver:
|
||||
json.dumps(row.to_cache()),
|
||||
)
|
||||
|
||||
def cache_set_negative(self, token_hash: str, marker: NegativeCache = NegativeCache.INVALID) -> None:
|
||||
self._redis.setex(self._cache_key(token_hash), self._negative_ttl, str(marker))
|
||||
def cache_set_negative(self, token_hash: str) -> None:
|
||||
self._redis.setex(self._cache_key(token_hash), self._negative_ttl, "invalid")
|
||||
|
||||
def hard_expire(self, session: Session, row_id: uuid.UUID | str, token_hash: str) -> None:
|
||||
"""Atomic CAS — only the worker that flips revoked_at emits audit;
|
||||
@@ -397,7 +385,7 @@ class OAuthAccessTokenResolver:
|
||||
extra={"audit": True, "token_id": str(row_id)},
|
||||
)
|
||||
self._redis.delete(self._cache_key(token_hash))
|
||||
self.cache_set_negative(token_hash, NegativeCache.EXPIRED)
|
||||
self.cache_set_negative(token_hash)
|
||||
|
||||
|
||||
class _VariantResolver:
|
||||
@@ -407,11 +395,9 @@ class _VariantResolver:
|
||||
|
||||
def resolve(self, token_hash: str) -> ResolvedRow | None:
|
||||
cached = self._parent.cache_get(token_hash)
|
||||
if isinstance(cached, NegativeCache):
|
||||
if cached is NegativeCache.EXPIRED:
|
||||
raise TokenExpiredError("token_expired")
|
||||
if cached == "invalid":
|
||||
return None
|
||||
if cached is not None:
|
||||
if cached is not None and not isinstance(cached, str):
|
||||
if not self._matches_variant(cached):
|
||||
return None
|
||||
return cached
|
||||
@@ -427,7 +413,7 @@ class _VariantResolver:
|
||||
now = datetime.now(UTC)
|
||||
if row.expires_at is not None and row.expires_at <= now:
|
||||
self._parent.hard_expire(session, row.id, token_hash)
|
||||
raise TokenExpiredError("token_expired")
|
||||
return None
|
||||
|
||||
if not self._matches_variant_model(row):
|
||||
logger.error(
|
||||
@@ -486,7 +472,7 @@ def record_layer0_verdict(token_hash: str, tenant_id: str, verdict: bool) -> Non
|
||||
if raw is None:
|
||||
return
|
||||
text = raw.decode() if isinstance(raw, (bytes, bytearray)) else raw
|
||||
if text in (NegativeCache.INVALID, NegativeCache.EXPIRED):
|
||||
if text == "invalid":
|
||||
return
|
||||
try:
|
||||
data = json.loads(text)
|
||||
@@ -615,8 +601,6 @@ def validate_bearer(*, accept: frozenset[Accepts]) -> Callable[[Callable[_DP, _D
|
||||
|
||||
try:
|
||||
ctx = get_authenticator().authenticate(token)
|
||||
except TokenExpiredError:
|
||||
raise Unauthorized("token_expired")
|
||||
except InvalidBearerError as e:
|
||||
raise Unauthorized(str(e))
|
||||
|
||||
|
||||
@@ -2,10 +2,9 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
_DIAGNOSTIC_PREFIXES = ("ERROR ", "WARN ", "WARNING ")
|
||||
_DIAGNOSTIC_PREFIXES = ("ERROR ", "WARNING ")
|
||||
_LOCATION_PREFIX = "-->"
|
||||
|
||||
|
||||
@@ -14,7 +13,7 @@ def extract_diagnostics(raw_output: str) -> str:
|
||||
|
||||
The full pyrefly output includes code excerpts and carets, which create noisy
|
||||
diffs. This helper keeps only:
|
||||
- diagnostic headline lines (``ERROR ...`` / ``WARN ...`` / ``WARNING ...``)
|
||||
- diagnostic headline lines (``ERROR ...`` / ``WARNING ...``)
|
||||
- the following location line (``--> path:line:column``), when present
|
||||
"""
|
||||
|
||||
@@ -37,28 +36,11 @@ def extract_diagnostics(raw_output: str) -> str:
|
||||
return "\n".join(diagnostics) + "\n"
|
||||
|
||||
|
||||
def render_diagnostics(raw_output: str, exit_code: int) -> str:
|
||||
"""Render concise diagnostics and fall back to raw output on unmatched failures."""
|
||||
|
||||
diagnostics = extract_diagnostics(raw_output)
|
||||
if diagnostics:
|
||||
return diagnostics
|
||||
|
||||
if exit_code != 0:
|
||||
return raw_output
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Read pyrefly output from stdin and print normalized diagnostics."""
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--status", type=int, default=0)
|
||||
args = parser.parse_args()
|
||||
|
||||
raw_output = sys.stdin.read()
|
||||
sys.stdout.write(render_diagnostics(raw_output, exit_code=args.status))
|
||||
sys.stdout.write(extract_diagnostics(raw_output))
|
||||
return 0
|
||||
|
||||
|
||||
|
||||
@@ -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"""
|
||||
|
||||
@@ -13432,6 +13432,7 @@ Soft lifecycle state for Agent records.
|
||||
| created_at | integer | | No |
|
||||
| files | [ string ] | | Yes |
|
||||
| id | string | | Yes |
|
||||
| message_chain_id | string | | No |
|
||||
| message_id | string | | Yes |
|
||||
| observation | string | | No |
|
||||
| position | integer | | Yes |
|
||||
@@ -14542,8 +14543,8 @@ Enum class for configurate method of provider model.
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| annotation_create_account | [SimpleAccount](#simpleaccount) | | No |
|
||||
| annotation_id | string | | Yes |
|
||||
| created_at | integer | | No |
|
||||
| id | string | | Yes |
|
||||
|
||||
#### ConversationDetail
|
||||
|
||||
@@ -16712,7 +16713,6 @@ Input field definition for snippet parameters.
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| description | string | | No |
|
||||
| icon | string | | No |
|
||||
| icon_background | string | | No |
|
||||
| icon_type | string | | No |
|
||||
@@ -16862,6 +16862,7 @@ Enum class for large language model mode.
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| expired_at | string | | Yes |
|
||||
| seats | [LicenseLimitationModel](#licenselimitationmodel) | | Yes |
|
||||
| status | [LicenseStatus](#licensestatus) | | Yes |
|
||||
| workspaces | [LicenseLimitationModel](#licenselimitationmodel) | | Yes |
|
||||
|
||||
@@ -17082,7 +17083,6 @@ Enum class for large language model mode.
|
||||
| agent_thoughts | [ [AgentThought](#agentthought) ] | | No |
|
||||
| annotation | [ConversationAnnotation](#conversationannotation) | | No |
|
||||
| annotation_hit_history | [ConversationAnnotationHitHistory](#conversationannotationhithistory) | | No |
|
||||
| answer | string | | Yes |
|
||||
| answer_tokens | integer | | No |
|
||||
| conversation_id | string | | Yes |
|
||||
| created_at | integer | | No |
|
||||
@@ -17096,11 +17096,12 @@ Enum class for large language model mode.
|
||||
| inputs | object | | Yes |
|
||||
| message | [JSONValue](#jsonvalue) | | No |
|
||||
| message_files | [ [MessageFile](#messagefile) ] | | No |
|
||||
| message_metadata_dict | [JSONValue](#jsonvalue) | | No |
|
||||
| message_tokens | integer | | No |
|
||||
| metadata | [JSONValue](#jsonvalue) | | No |
|
||||
| parent_message_id | string | | No |
|
||||
| provider_response_latency | number | | No |
|
||||
| query | string | | Yes |
|
||||
| re_sign_file_url_answer | string | | Yes |
|
||||
| status | string | | Yes |
|
||||
| workflow_run_id | string | | No |
|
||||
|
||||
|
||||
@@ -990,12 +990,6 @@ Pagination for GET /account/sessions. Strict (extra='forbid').
|
||||
| last_used_at | string | | No |
|
||||
| prefix | string | | Yes |
|
||||
|
||||
#### SimpleResultResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| result | string | | Yes |
|
||||
|
||||
#### SupportedAppType
|
||||
|
||||
App types the ``app`` usage face (``get app``) lists and filters.
|
||||
|
||||
@@ -1337,6 +1337,7 @@ Parsed multipart form fields for HITL uploads.
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| expired_at | string | | Yes |
|
||||
| seats | [LicenseLimitationModel](#licenselimitationmodel) | | Yes |
|
||||
| status | [LicenseStatus](#licensestatus) | | Yes |
|
||||
| workspaces | [LicenseLimitationModel](#licenselimitationmodel) | | Yes |
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -72,7 +72,6 @@ def mint_token(flask_app: Flask):
|
||||
prefix: str,
|
||||
subject_email: str,
|
||||
subject_issuer: str | None,
|
||||
expires_at: datetime | None = None,
|
||||
) -> OAuthAccessToken:
|
||||
with flask_app.app_context():
|
||||
row = OAuthAccessToken(
|
||||
@@ -83,7 +82,7 @@ def mint_token(flask_app: Flask):
|
||||
subject_issuer=subject_issuer,
|
||||
client_id="difyctl",
|
||||
device_label="test-device",
|
||||
expires_at=expires_at or (datetime.now(UTC) + timedelta(hours=1)),
|
||||
expires_at=datetime.now(UTC) + timedelta(hours=1),
|
||||
)
|
||||
db.session.add(row)
|
||||
db.session.commit()
|
||||
@@ -112,21 +111,6 @@ def account_token(workspace_account, mint_token) -> str:
|
||||
return token
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def expired_account_token(workspace_account, mint_token) -> str:
|
||||
account, _, _ = workspace_account
|
||||
token = "dfoa_" + uuid.uuid4().hex
|
||||
mint_token(
|
||||
token,
|
||||
account_id=account.id,
|
||||
prefix="dfoa_",
|
||||
subject_email=account.email,
|
||||
subject_issuer="dify:account",
|
||||
expires_at=datetime.now(UTC) - timedelta(minutes=1),
|
||||
)
|
||||
return token
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flush_auth_redis(flask_app: Flask) -> Generator[None, None, None]:
|
||||
def _flush():
|
||||
|
||||
@@ -6,7 +6,6 @@ acceptance/rejection on app-scoped routes.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Generator
|
||||
|
||||
import pytest
|
||||
@@ -17,50 +16,6 @@ from extensions.ext_database import db
|
||||
from models import App, Tenant
|
||||
|
||||
|
||||
def test_expired_token_returns_401_token_expired(
|
||||
test_client: FlaskClient,
|
||||
expired_account_token: str,
|
||||
) -> None:
|
||||
"""An expired bearer is distinguishable from an unknown one: 401 with the
|
||||
domain code ``token_expired`` (+ actionable hint), not a generic 401 or 500."""
|
||||
res = test_client.get(
|
||||
"/openapi/v1/account",
|
||||
headers={"Authorization": f"Bearer {expired_account_token}"},
|
||||
)
|
||||
assert res.status_code == 401
|
||||
assert res.json["code"] == "token_expired"
|
||||
assert res.json["hint"]
|
||||
|
||||
|
||||
def test_expired_token_replay_stays_token_expired(
|
||||
test_client: FlaskClient,
|
||||
expired_account_token: str,
|
||||
) -> None:
|
||||
"""The distinct ``expired`` negative-cache marker keeps the second hit (served
|
||||
from cache, inside NEGATIVE_TTL) reporting ``token_expired`` rather than
|
||||
collapsing into a generic unknown-token 401."""
|
||||
headers = {"Authorization": f"Bearer {expired_account_token}"}
|
||||
first = test_client.get("/openapi/v1/account", headers=headers)
|
||||
second = test_client.get("/openapi/v1/account", headers=headers)
|
||||
assert first.json["code"] == "token_expired"
|
||||
assert second.status_code == 401
|
||||
assert second.json["code"] == "token_expired"
|
||||
|
||||
|
||||
def test_unknown_token_returns_401_unauthorized_not_500(
|
||||
test_client: FlaskClient,
|
||||
workspace_account,
|
||||
) -> None:
|
||||
"""An unknown bearer is a clean 401 ``unauthorized`` — not the latent 500 the
|
||||
pipeline used to leak for unmapped InvalidBearerError."""
|
||||
res = test_client.get(
|
||||
"/openapi/v1/account",
|
||||
headers={"Authorization": "Bearer dfoa_" + uuid.uuid4().hex},
|
||||
)
|
||||
assert res.status_code == 401
|
||||
assert res.json["code"] == "unauthorized"
|
||||
|
||||
|
||||
def test_info_accepts_account_bearer_with_apps_read_scope(
|
||||
test_client: FlaskClient,
|
||||
app_in_workspace: App,
|
||||
|
||||
@@ -27,6 +27,7 @@ extend-select = ["ANN401", "ARG", "TID251"]
|
||||
"controllers/web/test_wraps.py" = ["ARG"]
|
||||
"core/app/layers/test_pause_state_persist_layer.py" = ["ARG"]
|
||||
"core/rag/retrieval/test_dataset_retrieval_integration.py" = ["ARG"]
|
||||
"models/test_account.py" = ["ARG"]
|
||||
"models/test_conversation_message_inputs.py" = ["ARG"]
|
||||
"models/test_conversation_status_count.py" = ["ARG"]
|
||||
"repositories/test_sqlalchemy_api_workflow_run_repository.py" = ["ARG"]
|
||||
|
||||
@@ -3,7 +3,6 @@ Integration tests for Account and Tenant model methods that interact with the da
|
||||
|
||||
Migrated from unit_tests/models/test_account_models.py, replacing
|
||||
@patch("models.account.db") mock patches with real PostgreSQL operations.
|
||||
Also absorbs unit_tests/models/test_account.py role helper coverage.
|
||||
|
||||
Covers:
|
||||
- Account.current_tenant setter (sets _current_tenant and role from TenantAccountJoin)
|
||||
@@ -13,7 +12,6 @@ Covers:
|
||||
"""
|
||||
|
||||
from collections.abc import Generator
|
||||
from typing import cast
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
@@ -22,10 +20,8 @@ from sqlalchemy.orm import Session
|
||||
|
||||
from models.account import Account, AccountIntegrate, Tenant, TenantAccountJoin, TenantAccountRole
|
||||
|
||||
TrackedRow = Account | AccountIntegrate | Tenant | TenantAccountJoin
|
||||
|
||||
|
||||
def _cleanup_tracked_rows(db_session: Session, tracked: list[TrackedRow]) -> None:
|
||||
def _cleanup_tracked_rows(db_session: Session, tracked: list) -> None:
|
||||
"""Delete rows tracked during the test so committed state does not leak into the DB.
|
||||
|
||||
Rolls back any pending (uncommitted) session state first, then issues DELETE
|
||||
@@ -56,7 +52,7 @@ def _build_account(email_prefix: str = "account") -> Account:
|
||||
class _DBTrackingTestBase:
|
||||
"""Base class providing a tracker list and shared row factories for account/tenant tests."""
|
||||
|
||||
_tracked: list[TrackedRow]
|
||||
_tracked: list
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _setup_cleanup(self, db_session_with_containers: Session) -> Generator[None, None, None]:
|
||||
@@ -88,22 +84,6 @@ class _DBTrackingTestBase:
|
||||
return join
|
||||
|
||||
|
||||
class TestTenantAccountRole:
|
||||
"""Tests for TenantAccountRole helper methods."""
|
||||
|
||||
def test_account_is_privileged_role(self) -> None:
|
||||
assert TenantAccountRole.ADMIN == "admin"
|
||||
assert TenantAccountRole.OWNER == "owner"
|
||||
assert TenantAccountRole.EDITOR == "editor"
|
||||
assert TenantAccountRole.NORMAL == "normal"
|
||||
|
||||
assert TenantAccountRole.is_privileged_role(TenantAccountRole.ADMIN)
|
||||
assert TenantAccountRole.is_privileged_role(TenantAccountRole.OWNER)
|
||||
assert not TenantAccountRole.is_privileged_role(TenantAccountRole.NORMAL)
|
||||
assert not TenantAccountRole.is_privileged_role(TenantAccountRole.EDITOR)
|
||||
assert not TenantAccountRole.is_privileged_role(cast(TenantAccountRole, ""))
|
||||
|
||||
|
||||
class TestAccountCurrentTenantSetter(_DBTrackingTestBase):
|
||||
"""Integration tests for Account.current_tenant property setter."""
|
||||
|
||||
@@ -196,7 +176,7 @@ class TestAccountGetByOpenId(_DBTrackingTestBase):
|
||||
assert result is not None
|
||||
assert result.id == account.id
|
||||
|
||||
def test_get_by_openid_returns_none_when_no_integrate_exists(self) -> None:
|
||||
def test_get_by_openid_returns_none_when_no_integrate_exists(self, db_session_with_containers: Session) -> None:
|
||||
"""get_by_openid returns None when no AccountIntegrate row matches."""
|
||||
result = Account.get_by_openid("github", f"github_{uuid4()}")
|
||||
|
||||
|
||||
@@ -50,6 +50,7 @@ project-excludes = [
|
||||
"libs/broadcast_channel/redis/test_streams_channel.py",
|
||||
"libs/test_auto_renew_redis_lock_integration.py",
|
||||
"libs/test_rate_limiter_integration.py",
|
||||
"models/test_account.py",
|
||||
"models/test_conversation_message_inputs.py",
|
||||
"models/test_types_enum_text.py",
|
||||
"repositories/test_sqlalchemy_api_workflow_node_execution_repository.py",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -77,25 +77,6 @@ class AnnotationApi(Resource):
|
||||
assert "prefer dump_response" in checks[0].reason
|
||||
|
||||
|
||||
def test_constructor_variable_model_dump_is_valid(tmp_path: Path):
|
||||
checks = _checks_for_source(
|
||||
tmp_path,
|
||||
"""
|
||||
@ns.route("/annotations")
|
||||
class AnnotationApi(Resource):
|
||||
@ns.response(201, "Created", ns.models[AnnotationResponse.__name__])
|
||||
def post(self):
|
||||
response = AnnotationResponse(id="new", name=name)
|
||||
return response.model_dump(mode="json"), 201
|
||||
""",
|
||||
)
|
||||
|
||||
assert len(checks) == 1
|
||||
assert checks[0].classification == "valid"
|
||||
assert checks[0].actual[0].kind == "model"
|
||||
assert checks[0].actual[0].model == "AnnotationResponse"
|
||||
|
||||
|
||||
def test_variable_model_dump_with_wrong_documented_schema_is_mismatch(tmp_path: Path):
|
||||
checks = _checks_for_source(
|
||||
tmp_path,
|
||||
@@ -136,38 +117,6 @@ class StreamApi(Resource):
|
||||
assert {actual.model for actual in checks[0].actual} == {"StreamResponse"}
|
||||
|
||||
|
||||
def test_response_contract_ignore_comment_skips_route_method(tmp_path: Path):
|
||||
checks = _checks_for_source(
|
||||
tmp_path,
|
||||
"""
|
||||
@ns.route("/binary")
|
||||
class BinaryApi(Resource):
|
||||
# response-contract:ignore binary response
|
||||
@ns.response(200, "Binary file")
|
||||
def get(self):
|
||||
return send_file(path)
|
||||
|
||||
|
||||
# response-contract:ignore compact Flask response
|
||||
@ns.route("/compact")
|
||||
class CompactApi(Resource):
|
||||
def get(self):
|
||||
return make_response({"url": "https://example.com"})
|
||||
|
||||
|
||||
@ns.route("/regular")
|
||||
class RegularApi(Resource):
|
||||
@ns.response(200, "OK", ns.models[RegularResponse.__name__])
|
||||
def get(self):
|
||||
return dump_response(RegularResponse, {})
|
||||
""",
|
||||
)
|
||||
|
||||
assert len(checks) == 1
|
||||
assert checks[0].class_name == "RegularApi"
|
||||
assert checks[0].classification == "valid"
|
||||
|
||||
|
||||
def test_main_is_report_only_by_default_for_mismatches(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
module = _load_lint_response_contracts_module()
|
||||
controller_path = tmp_path / "controllers" / "sample.py"
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -22,17 +22,23 @@ class TestSpecSchemaDefinitionsApi:
|
||||
assert status == 200
|
||||
assert resp == schema_definitions
|
||||
|
||||
def test_get_exception_returns_empty_list(self, caplog):
|
||||
def test_get_exception_returns_empty_list(self):
|
||||
api = spec_module.SpecSchemaDefinitionsApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
with patch.object(
|
||||
spec_module,
|
||||
"SchemaManager",
|
||||
side_effect=Exception("boom"),
|
||||
with (
|
||||
patch.object(
|
||||
spec_module,
|
||||
"SchemaManager",
|
||||
side_effect=Exception("boom"),
|
||||
),
|
||||
patch.object(
|
||||
spec_module.logger,
|
||||
"exception",
|
||||
) as log_exception,
|
||||
):
|
||||
resp, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert resp == []
|
||||
assert "boom" in caplog.text
|
||||
log_exception.assert_called_once()
|
||||
|
||||
@@ -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"})
|
||||
|
||||
@@ -321,56 +321,3 @@ def test_guard_no_external_identity_when_subject_email_absent(app):
|
||||
view()
|
||||
|
||||
assert received["data"].external_identity is None
|
||||
|
||||
|
||||
# --- auth-failure mapping (no raw 500 leak) ---
|
||||
|
||||
|
||||
def test_guard_expired_token_raises_session_expired_401(app):
|
||||
from controllers.openapi._errors import OpenApiErrorCode, SessionExpired
|
||||
from libs.oauth_bearer import TokenExpiredError
|
||||
|
||||
router = _make_router()
|
||||
|
||||
with app.test_request_context("/test", headers={"Authorization": "Bearer tok"}):
|
||||
with (
|
||||
patch("controllers.openapi.auth.pipeline.extract_bearer", return_value="tok"),
|
||||
patch("controllers.openapi.auth.pipeline.get_authenticator") as mock_auth,
|
||||
patch("controllers.openapi.auth.pipeline.current_edition", return_value=Edition.CE),
|
||||
):
|
||||
mock_auth.return_value.authenticate.side_effect = TokenExpiredError("token_expired")
|
||||
|
||||
@router.guard(scope=Scope.FULL)
|
||||
def view(*, auth_data):
|
||||
pass
|
||||
|
||||
with pytest.raises(SessionExpired) as exc:
|
||||
view()
|
||||
|
||||
assert exc.value.code == 401
|
||||
assert exc.value.error_code == OpenApiErrorCode.TOKEN_EXPIRED
|
||||
|
||||
|
||||
def test_guard_invalid_token_raises_unified_401_not_500(app):
|
||||
from controllers.openapi._errors import InvalidBearer, OpenApiErrorCode
|
||||
from libs.oauth_bearer import InvalidBearerError
|
||||
|
||||
router = _make_router()
|
||||
|
||||
with app.test_request_context("/test", headers={"Authorization": "Bearer tok"}):
|
||||
with (
|
||||
patch("controllers.openapi.auth.pipeline.extract_bearer", return_value="tok"),
|
||||
patch("controllers.openapi.auth.pipeline.get_authenticator") as mock_auth,
|
||||
patch("controllers.openapi.auth.pipeline.current_edition", return_value=Edition.CE),
|
||||
):
|
||||
mock_auth.return_value.authenticate.side_effect = InvalidBearerError("invalid_bearer")
|
||||
|
||||
@router.guard(scope=Scope.FULL)
|
||||
def view(*, auth_data):
|
||||
pass
|
||||
|
||||
with pytest.raises(InvalidBearer) as exc:
|
||||
view()
|
||||
|
||||
assert exc.value.code == 401
|
||||
assert exc.value.error_code == OpenApiErrorCode.UNAUTHORIZED
|
||||
|
||||
@@ -3,36 +3,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from controllers.openapi._models import AppRunRequest
|
||||
from models import Account
|
||||
from models.model import App, AppMode
|
||||
|
||||
_TEST_APP_ID = str(uuid.uuid4())
|
||||
_TEST_TENANT_ID = str(uuid.uuid4())
|
||||
_TEST_ACCOUNT_ID = str(uuid.uuid4())
|
||||
|
||||
|
||||
def _make_app() -> App:
|
||||
app = App()
|
||||
app.id = _TEST_APP_ID
|
||||
app.tenant_id = _TEST_TENANT_ID
|
||||
app.name = "Streaming app"
|
||||
app.mode = AppMode.CHAT
|
||||
app.enable_site = False
|
||||
app.enable_api = True
|
||||
return app
|
||||
|
||||
|
||||
def _make_account() -> Account:
|
||||
account = Account(name="OpenAPI caller", email="caller@example.com")
|
||||
account.id = _TEST_ACCOUNT_ID
|
||||
return account
|
||||
|
||||
|
||||
def test_app_run_request_has_no_response_mode_field():
|
||||
@@ -63,19 +40,15 @@ def test_run_chat_always_calls_generate_with_streaming_true(
|
||||
from controllers.openapi.app_run import _run_chat
|
||||
|
||||
generate_mock = Mock(return_value=iter([]))
|
||||
|
||||
class GenerateService:
|
||||
generate = generate_mock
|
||||
|
||||
monkeypatch.setattr(
|
||||
sys.modules["controllers.openapi.app_run"],
|
||||
"AppGenerateService",
|
||||
GenerateService,
|
||||
SimpleNamespace(generate=generate_mock),
|
||||
)
|
||||
with app.test_request_context(f"/openapi/v1/apps/{_TEST_APP_ID}/run", method="POST"):
|
||||
with app.test_request_context("/openapi/v1/apps/app-1/run", method="POST"):
|
||||
_run_chat(
|
||||
_make_app(),
|
||||
_make_account(),
|
||||
SimpleNamespace(id="app-1", tenant_id="t-1"),
|
||||
SimpleNamespace(id="acct-1"),
|
||||
AppRunRequest(inputs={}, query="hello"),
|
||||
)
|
||||
_, kwargs = generate_mock.call_args
|
||||
@@ -107,11 +80,11 @@ def test_stop_task_calls_queue_manager_and_graph_engine(app: Flask, bypass_pipel
|
||||
|
||||
auth_data = AuthData.model_construct(
|
||||
token_type=TokenType.OAUTH_ACCOUNT,
|
||||
account_id=uuid.UUID(_TEST_ACCOUNT_ID),
|
||||
account_id=uuid.uuid4(),
|
||||
token_hash="test",
|
||||
scopes=frozenset({Scope.FULL}),
|
||||
app=_make_app(),
|
||||
caller=_make_account(),
|
||||
app=SimpleNamespace(id="app-1", tenant_id="t-1"),
|
||||
caller=SimpleNamespace(id="acct-1"),
|
||||
caller_kind="account",
|
||||
)
|
||||
|
||||
|
||||
@@ -5,7 +5,6 @@ view function decorated with @accepts/@returns, driven inside a request context.
|
||||
"""
|
||||
|
||||
from functools import wraps
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
@@ -101,7 +100,7 @@ def test_accepts_validation_error_is_sanitized_and_structured(app):
|
||||
with pytest.raises(UnprocessableEntity) as exc_info:
|
||||
view()
|
||||
|
||||
data = cast(dict[str, Any], cast(Any, exc_info.value).data)
|
||||
data = exc_info.value.data
|
||||
assert data["message"] == "Request validation failed"
|
||||
assert isinstance(data["errors"], list)
|
||||
assert data["errors"]
|
||||
|
||||
@@ -33,7 +33,6 @@ from controllers.openapi._errors import (
|
||||
OpenApiErrorCode,
|
||||
OpenApiErrorFormatter,
|
||||
RecipientSurfaceMismatch,
|
||||
SessionExpired,
|
||||
)
|
||||
from controllers.service_api.app.error import (
|
||||
AppUnavailableError,
|
||||
@@ -354,20 +353,3 @@ class TestErrorCodeEnumRegistration:
|
||||
schema = model.__schema__
|
||||
assert schema["type"] == "string"
|
||||
assert set(schema["enum"]) == {member.value for member in OpenApiErrorCode}
|
||||
|
||||
|
||||
class TestSessionExpired:
|
||||
def test_session_expired_emits_token_expired_401_with_hint(self):
|
||||
fmt = OpenApiErrorFormatter()
|
||||
e = SessionExpired()
|
||||
data = {"code": "unauthorized", "message": e.description, "status": 401}
|
||||
|
||||
wire = fmt.finalize(e, data, 401)
|
||||
|
||||
assert wire["code"] == OpenApiErrorCode.TOKEN_EXPIRED
|
||||
assert wire["status"] == 401
|
||||
assert wire["hint"]
|
||||
|
||||
def test_session_expired_code_is_401(self):
|
||||
assert SessionExpired.code == 401
|
||||
assert SessionExpired.error_code == OpenApiErrorCode.TOKEN_EXPIRED
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -1,87 +0,0 @@
|
||||
"""Resolver-level expiry signalling.
|
||||
|
||||
An expired token must be distinguishable from an unknown/revoked one: the
|
||||
resolver raises ``TokenExpiredError`` for expiry and returns ``None`` for
|
||||
everything else. The signal survives the negative-cache window via a distinct
|
||||
``expired`` marker so a retry inside ``NEGATIVE_TTL`` still reports expiry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.oauth_bearer import (
|
||||
OAuthAccessTokenResolver,
|
||||
TokenExpiredError,
|
||||
)
|
||||
|
||||
|
||||
def _row(expires_at: datetime):
|
||||
row = MagicMock()
|
||||
row.id = "11111111-1111-1111-1111-111111111111"
|
||||
row.account_id = "22222222-2222-2222-2222-222222222222"
|
||||
row.prefix = "dfoa_"
|
||||
row.subject_email = None
|
||||
row.subject_issuer = None
|
||||
row.client_id = None
|
||||
row.expires_at = expires_at
|
||||
return row
|
||||
|
||||
|
||||
def _resolver(redis: MagicMock, db_row=None) -> OAuthAccessTokenResolver:
|
||||
session = MagicMock()
|
||||
session.query.return_value.filter.return_value.one_or_none.return_value = db_row
|
||||
session.execute.return_value.rowcount = 1
|
||||
return OAuthAccessTokenResolver(session_factory=lambda: session, redis_client=redis)
|
||||
|
||||
|
||||
def test_resolve_raises_token_expired_for_expired_db_row():
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = None # cache miss -> DB path
|
||||
past = datetime.now(UTC) - timedelta(minutes=1)
|
||||
resolver = _resolver(redis, db_row=_row(past))
|
||||
|
||||
with pytest.raises(TokenExpiredError):
|
||||
resolver.for_account().resolve("expiredhash")
|
||||
|
||||
|
||||
def test_resolve_raises_token_expired_for_expired_cache_marker():
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = b"expired" # negative-cache replay
|
||||
resolver = _resolver(redis, db_row=None)
|
||||
|
||||
with pytest.raises(TokenExpiredError):
|
||||
resolver.for_account().resolve("expiredhash")
|
||||
|
||||
|
||||
def test_resolve_returns_none_for_invalid_cache_marker():
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = b"invalid"
|
||||
resolver = _resolver(redis, db_row=None)
|
||||
|
||||
assert resolver.for_account().resolve("revokedhash") is None
|
||||
|
||||
|
||||
def test_resolve_returns_none_for_unknown_token():
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = None # cache miss
|
||||
resolver = _resolver(redis, db_row=None) # no DB row
|
||||
|
||||
assert resolver.for_account().resolve("unknownhash") is None
|
||||
|
||||
|
||||
def test_hard_expire_caches_expired_marker_not_invalid():
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = None
|
||||
past = datetime.now(UTC) - timedelta(minutes=1)
|
||||
resolver = _resolver(redis, db_row=_row(past))
|
||||
|
||||
with pytest.raises(TokenExpiredError):
|
||||
resolver.for_account().resolve("expiredhash")
|
||||
|
||||
setex_values = [call.args[2] for call in redis.setex.call_args_list]
|
||||
assert "expired" in setex_values
|
||||
assert "invalid" not in setex_values
|
||||
@@ -1,4 +1,4 @@
|
||||
from libs.pyrefly_diagnostics import extract_diagnostics, render_diagnostics
|
||||
from libs.pyrefly_diagnostics import extract_diagnostics
|
||||
|
||||
|
||||
def test_extract_diagnostics_keeps_only_summary_and_location_lines() -> None:
|
||||
@@ -40,37 +40,6 @@ def test_extract_diagnostics_handles_error_without_location_line() -> None:
|
||||
assert diagnostics == "ERROR unexpected pyrefly output format [bad-format]\n"
|
||||
|
||||
|
||||
def test_extract_diagnostics_keeps_warn_headlines_and_location_lines() -> None:
|
||||
# Arrange
|
||||
raw_output = """INFO Checking project configured at `/tmp/project/pyrefly.toml`
|
||||
WARN Skipping include pattern `/tmp/project/tests` because it is matched by `project-excludes`.
|
||||
--> tests/test_containers_integration_tests/pyrefly.toml:3:1
|
||||
"""
|
||||
|
||||
# Act
|
||||
diagnostics = extract_diagnostics(raw_output)
|
||||
|
||||
# Assert
|
||||
assert diagnostics == (
|
||||
"WARN Skipping include pattern `/tmp/project/tests` because it is matched by `project-excludes`.\n"
|
||||
" --> tests/test_containers_integration_tests/pyrefly.toml:3:1\n"
|
||||
)
|
||||
|
||||
|
||||
def test_render_diagnostics_falls_back_to_raw_output_for_nonzero_exit_without_matches() -> None:
|
||||
# Arrange
|
||||
raw_output = (
|
||||
"INFO Checking project configured at `/tmp/project/pyrefly.toml`\n"
|
||||
"No Python files matched pattern `/tmp/project/tests/test_containers_integration_tests`\n"
|
||||
)
|
||||
|
||||
# Act
|
||||
diagnostics = render_diagnostics(raw_output, exit_code=1)
|
||||
|
||||
# Assert
|
||||
assert diagnostics == raw_output
|
||||
|
||||
|
||||
def test_extract_diagnostics_returns_empty_for_non_error_output() -> None:
|
||||
# Arrange
|
||||
raw_output = "INFO Checking project configured at `/tmp/project/pyrefly.toml`\n"
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
from models.account import TenantAccountRole
|
||||
|
||||
|
||||
def test_account_is_privileged_role():
|
||||
assert TenantAccountRole.ADMIN == "admin"
|
||||
assert TenantAccountRole.OWNER == "owner"
|
||||
assert TenantAccountRole.EDITOR == "editor"
|
||||
assert TenantAccountRole.NORMAL == "normal"
|
||||
|
||||
assert TenantAccountRole.is_privileged_role(TenantAccountRole.ADMIN)
|
||||
assert TenantAccountRole.is_privileged_role(TenantAccountRole.OWNER)
|
||||
assert not TenantAccountRole.is_privileged_role(TenantAccountRole.NORMAL)
|
||||
assert not TenantAccountRole.is_privileged_role(TenantAccountRole.EDITOR)
|
||||
assert not TenantAccountRole.is_privileged_role("")
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user