Compare commits
18
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bc68e02711 | ||
|
|
eb4ec93cea | ||
|
|
fa6f4b0ea5 | ||
|
|
267b34caaf | ||
|
|
6ab1dd06ac | ||
|
|
113d6d7e00 | ||
|
|
a66b1de477 | ||
|
|
16b698b54d | ||
|
|
4cd8b8c733 | ||
|
|
677ab01806 | ||
|
|
f4e832f35c | ||
|
|
1c5e1280cb | ||
|
|
0035d90e36 | ||
|
|
2047e0dc12 | ||
|
|
2e9c0a3c7a | ||
|
|
a246dc8b17 | ||
|
|
bb921bcc45 | ||
|
|
4f4ac27de2 |
@@ -12,6 +12,7 @@ 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. |
|
||||
@@ -23,7 +24,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: components, hooks, local types, query helpers, atoms, constants, and small utilities should live near the code that changes with them.
|
||||
- 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.
|
||||
- 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.
|
||||
@@ -32,6 +33,8 @@ 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.
|
||||
@@ -46,6 +49,7 @@ 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.
|
||||
@@ -60,8 +64,10 @@ 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.
|
||||
@@ -83,11 +89,13 @@ 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.
|
||||
@@ -96,7 +104,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. 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 or the root folder becomes noisy. 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,6 +53,8 @@ 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_dataset_permissions_to_rbac, migrate_member_roles_to_rbac
|
||||
from .rbac import migrate_member_roles_to_rbac
|
||||
from .retention import (
|
||||
archive_workflow_runs,
|
||||
archive_workflow_runs_plan,
|
||||
@@ -76,7 +76,6 @@ __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,7 +7,6 @@ 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 (
|
||||
@@ -178,4 +177,3 @@ def legacy_model_types(
|
||||
|
||||
|
||||
data_migrate.add_command(legacy_model_types)
|
||||
data_migrate.add_command(migrate_dataset_permissions_to_rbac)
|
||||
|
||||
+65
-437
@@ -1,55 +1,11 @@
|
||||
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 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
|
||||
from models import TenantAccountJoin, TenantAccountRole
|
||||
from services.enterprise.rbac_service import ListOption, RBACService
|
||||
|
||||
|
||||
def _resolve_builtin_role_id(tenant_id: str, operator_account_id: str, legacy_role: str) -> str:
|
||||
@@ -59,86 +15,26 @@ 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.
|
||||
"""
|
||||
if legacy_role not in _LEGACY_ROLE_TO_BUILTIN_TAG:
|
||||
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:
|
||||
raise ValueError(f"Unsupported legacy workspace role: {legacy_role}")
|
||||
|
||||
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(
|
||||
roles = RBACService.Roles.list(
|
||||
tenant_id=tenant_id,
|
||||
account_id=operator_account_id,
|
||||
member_account_id=member_account_id,
|
||||
role_ids=[role_id],
|
||||
)
|
||||
return member_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}")
|
||||
|
||||
|
||||
@click.command(
|
||||
@@ -146,16 +42,7 @@ def _replace_member_role(
|
||||
)
|
||||
@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.")
|
||||
@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:
|
||||
def migrate_member_roles_to_rbac(tenant_id: str | None, dry_run: bool) -> None:
|
||||
"""Backfill RBAC member-role bindings from legacy `TenantAccountJoin.role` data.
|
||||
|
||||
This is an offline migration command for workspaces that already have
|
||||
@@ -163,322 +50,63 @@ def migrate_member_roles_to_rbac(
|
||||
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")
|
||||
|
||||
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] = {}
|
||||
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)
|
||||
|
||||
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}")
|
||||
joins = list(session.scalars(stmt).all())
|
||||
|
||||
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: {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:
|
||||
if not joins:
|
||||
click.echo(click.style("No workspace members found for migration.", fg="yellow"))
|
||||
return
|
||||
|
||||
if dry_run:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Dry run completed. Scanned {scanned_count} members across {tenant_count} tenants. "
|
||||
"No RBAC bindings were written.",
|
||||
fg="yellow",
|
||||
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,
|
||||
)
|
||||
)
|
||||
else:
|
||||
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]
|
||||
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",
|
||||
)
|
||||
f"tenant={workspace_id} member={member_account_id} "
|
||||
f"legacy_role={legacy_role} -> rbac_role_id={resolved_role_id}"
|
||||
)
|
||||
|
||||
if dry_run:
|
||||
continue
|
||||
|
||||
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", exclude_unset=True
|
||||
),
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
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
|
||||
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
|
||||
|
||||
if dry_run:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Dry run completed. Scanned {scanned_count} datasets; "
|
||||
f"{partial_dataset_count} partial-member datasets would be migrated.",
|
||||
fg="yellow",
|
||||
)
|
||||
)
|
||||
click.echo(click.style("Dry run completed. No RBAC bindings were written.", fg="yellow"))
|
||||
else:
|
||||
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",
|
||||
)
|
||||
)
|
||||
click.echo(click.style(f"RBAC member-role migration completed. Migrated {migrated_count} members.", fg="green"))
|
||||
|
||||
@@ -34,12 +34,6 @@ 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):
|
||||
"""
|
||||
|
||||
@@ -43,7 +43,6 @@ from controllers.console.wraps import (
|
||||
from core.ops.ops_trace_manager import OpsTraceManager
|
||||
from core.rag.entities import PreProcessingRule, Rule, Segmentation
|
||||
from core.rag.retrieval.retrieval_methods import RetrievalMethod
|
||||
from core.rbac import RBACResourceWhitelistScope
|
||||
from core.trigger.constants import TRIGGER_NODE_TYPES
|
||||
from extensions.ext_database import db
|
||||
from fields.base import ResponseModel
|
||||
@@ -70,7 +69,6 @@ from services.entities.knowledge_entities.knowledge_entities import (
|
||||
WeightVectorSetting,
|
||||
)
|
||||
from services.feature_service import FeatureService
|
||||
from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task
|
||||
|
||||
ALLOW_CREATE_APP_MODES = ["chat", "agent-chat", "advanced-chat", "workflow", "completion"]
|
||||
|
||||
@@ -668,14 +666,6 @@ class AppListApi(Resource):
|
||||
|
||||
app_service = AppService()
|
||||
app = app_service.create_app(current_tenant_id, params, current_user)
|
||||
if dify_config.RBAC_ENABLED:
|
||||
enterprise_rbac_service.RBACService.AppAccess.replace_whitelist(
|
||||
tenant_id=str(current_tenant_id),
|
||||
account_id=current_user.id,
|
||||
app_id=str(app.id),
|
||||
payload=enterprise_rbac_service.ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_created_app_rbac_access_task.delay(current_tenant_id, current_user.id, app_id=app.id)
|
||||
permission_keys_map = enterprise_rbac_service.RBACService.AppPermissions.batch_get(
|
||||
str(current_tenant_id),
|
||||
current_user.id,
|
||||
|
||||
@@ -167,12 +167,16 @@ 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")
|
||||
|
||||
@@ -8,8 +8,11 @@ from controllers.common.schema import register_response_schema_models, register_
|
||||
from controllers.console import console_ns
|
||||
from controllers.console.app.wraps import get_app_model
|
||||
from controllers.console.wraps import (
|
||||
RBACPermission,
|
||||
RBACResourceScope,
|
||||
account_initialization_required,
|
||||
edit_permission_required,
|
||||
rbac_permission_required,
|
||||
setup_required,
|
||||
with_current_tenant_id,
|
||||
with_current_user,
|
||||
@@ -220,6 +223,7 @@ class WorkflowCommentListApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
|
||||
@get_app_model()
|
||||
def get(self, current_tenant_id: str, app_model: App):
|
||||
"""Get all comments for a workflow."""
|
||||
@@ -236,6 +240,7 @@ class WorkflowCommentListApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@get_app_model()
|
||||
@@ -268,6 +273,7 @@ class WorkflowCommentDetailApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
|
||||
@get_app_model()
|
||||
def get(self, current_tenant_id: str, app_model: App, comment_id: str):
|
||||
"""Get a specific workflow comment."""
|
||||
@@ -286,6 +292,7 @@ class WorkflowCommentDetailApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@get_app_model()
|
||||
@@ -314,6 +321,7 @@ class WorkflowCommentDetailApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@get_app_model()
|
||||
@@ -341,6 +349,7 @@ class WorkflowCommentResolveApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@get_app_model()
|
||||
@@ -369,6 +378,7 @@ class WorkflowCommentReplyApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@get_app_model()
|
||||
@@ -404,6 +414,7 @@ class WorkflowCommentReplyDetailApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@get_app_model()
|
||||
@@ -436,6 +447,7 @@ class WorkflowCommentReplyDetailApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@get_app_model()
|
||||
@@ -471,6 +483,7 @@ class WorkflowCommentMentionUsersApi(Resource):
|
||||
@setup_required
|
||||
@account_initialization_required
|
||||
@with_current_user
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
|
||||
@get_app_model()
|
||||
def get(self, current_user: Account, app_model: App):
|
||||
"""Get all users in current tenant for mentions."""
|
||||
|
||||
@@ -19,8 +19,7 @@ from controllers.console.app.error import AppNotFoundError
|
||||
from core.db.session_factory import session_factory
|
||||
from extensions.ext_database import db
|
||||
from libs.login import current_account_with_tenant
|
||||
from models import App, AppMode, TrialApp
|
||||
from services.recommended_app_service import RecommendedAppService
|
||||
from models import App, AppMode
|
||||
|
||||
|
||||
def _load_app_model(session: Session, app_id: str) -> App | None:
|
||||
@@ -42,10 +41,7 @@ def _load_app_model_from_scoped_session(app_id: str) -> App | None:
|
||||
|
||||
|
||||
def _load_app_model_with_trial(app_id: str) -> App | None:
|
||||
"""Load a normal app through its trial registration without applying current-tenant scope."""
|
||||
app_model = db.session.scalar(
|
||||
select(App).join(TrialApp, TrialApp.app_id == App.id).where(App.id == app_id, App.status == "normal").limit(1)
|
||||
)
|
||||
app_model = db.session.scalar(select(App).where(App.id == app_id, App.status == "normal").limit(1))
|
||||
return app_model
|
||||
|
||||
|
||||
@@ -199,8 +195,6 @@ def get_app_model_with_trial[**P, R](
|
||||
*,
|
||||
mode: AppMode | list[AppMode] | None = None,
|
||||
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
|
||||
"""Inject an app registered for trial or available from the recommended catalog."""
|
||||
|
||||
def decorator(view_func: Callable[P, R]) -> Callable[P, R]:
|
||||
@wraps(view_func)
|
||||
def decorated_view(*args: P.args, **kwargs: P.kwargs) -> R:
|
||||
@@ -213,8 +207,6 @@ def get_app_model_with_trial[**P, R](
|
||||
del kwargs["app_id"]
|
||||
|
||||
app_model = _load_app_model_with_trial(app_id)
|
||||
if app_model is None:
|
||||
app_model = RecommendedAppService.get_app(app_id, session=db.session())
|
||||
|
||||
if not app_model:
|
||||
raise AppNotFoundError()
|
||||
|
||||
@@ -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, SeatsLimitExceededError
|
||||
from services.errors.account import AccountRegisterError
|
||||
|
||||
from ..error import AccountInFreezeError, EmailSendIpLimitError, SeatsLimitExceeded
|
||||
from ..error import AccountInFreezeError, EmailSendIpLimitError
|
||||
from ..wraps import email_password_login_enabled, email_register_enabled, setup_required
|
||||
|
||||
|
||||
@@ -208,7 +208,5 @@ class EmailRegisterResetApi(Resource):
|
||||
timezone=timezone,
|
||||
session=db.session,
|
||||
)
|
||||
except SeatsLimitExceededError:
|
||||
raise SeatsLimitExceeded()
|
||||
except AccountRegisterError:
|
||||
raise AccountInFreezeError()
|
||||
|
||||
@@ -16,6 +16,7 @@ from controllers.console.auth.error import (
|
||||
)
|
||||
from controllers.console.error import AccountNotFound, EmailSendIpLimitError
|
||||
from controllers.console.wraps import email_password_login_enabled, setup_required
|
||||
from events.tenant_event import tenant_was_created
|
||||
from extensions.ext_database import db
|
||||
from libs.helper import EmailStr, extract_remote_ip
|
||||
from libs.password import hash_password
|
||||
@@ -200,4 +201,7 @@ class ForgotPasswordResetApi(Resource):
|
||||
not TenantService.get_join_tenants(account, session=db.session)
|
||||
and FeatureService.get_system_features().is_allow_create_workspace
|
||||
):
|
||||
TenantService.create_owner_tenant(account, session=db.session())
|
||||
tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session)
|
||||
TenantService.create_tenant_member(tenant, account, db.session, role="owner")
|
||||
account.current_tenant = tenant
|
||||
tenant_was_created.send(tenant)
|
||||
|
||||
@@ -25,7 +25,6 @@ from controllers.console.error import (
|
||||
AccountNotFound,
|
||||
EmailSendIpLimitError,
|
||||
NotAllowedCreateWorkspace,
|
||||
SeatsLimitExceeded,
|
||||
WorkspacesLimitExceeded,
|
||||
)
|
||||
from controllers.console.wraps import (
|
||||
@@ -35,6 +34,7 @@ from controllers.console.wraps import (
|
||||
setup_required,
|
||||
with_current_user,
|
||||
)
|
||||
from events.tenant_event import tenant_was_created
|
||||
from extensions.ext_database import db
|
||||
from libs.helper import EmailStr, extract_remote_ip
|
||||
from libs.helper import timezone as validate_timezone_string
|
||||
@@ -51,7 +51,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, SeatsLimitExceededError
|
||||
from services.errors.account import AccountRegisterError
|
||||
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
@@ -301,7 +301,10 @@ class EmailCodeLoginApi(Resource):
|
||||
if not FeatureService.get_system_features().is_allow_create_workspace:
|
||||
raise NotAllowedCreateWorkspace()
|
||||
else:
|
||||
TenantService.create_owner_tenant(account, session=db.session())
|
||||
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session)
|
||||
TenantService.create_tenant_member(new_tenant, account, db.session, role="owner")
|
||||
account.current_tenant = new_tenant
|
||||
tenant_was_created.send(new_tenant)
|
||||
|
||||
if account is None:
|
||||
try:
|
||||
@@ -314,8 +317,6 @@ 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()
|
||||
|
||||
@@ -11,6 +11,7 @@ from configs import dify_config
|
||||
from constants.languages import languages
|
||||
from controllers.common.fields import RedirectResponse
|
||||
from controllers.common.schema import query_params_from_model, register_response_schema_model, register_schema_models
|
||||
from events.tenant_event import tenant_was_created
|
||||
from extensions.ext_database import db
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from libs.helper import extract_remote_ip
|
||||
@@ -24,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, SeatsLimitExceededError
|
||||
from services.errors.account import AccountNotFoundError, AccountRegisterError
|
||||
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkSpaceNotFoundError
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
@@ -37,7 +38,6 @@ class OAuthLoginQuery(BaseModel):
|
||||
invite_token: str | None = Field(default=None, description="Optional invitation token")
|
||||
timezone: str | None = Field(default=None, description="Preferred timezone")
|
||||
language: str | None = Field(default=None, description="Preferred interface language")
|
||||
redirect_url: str | None = Field(default=None, description="Relative page to resume after login")
|
||||
|
||||
|
||||
class OAuthCallbackQuery(BaseModel):
|
||||
@@ -87,36 +87,6 @@ def _validated_language(value: str | None) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _url_origin(url: str) -> tuple[str, str, int] | None:
|
||||
parsed_url = urllib.parse.urlsplit(url)
|
||||
if parsed_url.scheme not in {"http", "https"} or parsed_url.hostname is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
port = parsed_url.port
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
if port is None:
|
||||
port = 443 if parsed_url.scheme == "https" else 80
|
||||
return parsed_url.scheme, parsed_url.hostname, port
|
||||
|
||||
|
||||
def _get_redirect_target(redirect_url: str | None) -> str:
|
||||
if not redirect_url:
|
||||
return dify_config.CONSOLE_WEB_URL
|
||||
|
||||
parsed_url = urllib.parse.urlsplit(redirect_url)
|
||||
normalized_path = redirect_url.lstrip().replace("\\", "/")
|
||||
if not parsed_url.scheme and not parsed_url.netloc and not normalized_path.startswith("//"):
|
||||
return redirect_url
|
||||
|
||||
redirect_origin = _url_origin(redirect_url)
|
||||
if redirect_origin is not None and redirect_origin == _url_origin(dify_config.CONSOLE_WEB_URL):
|
||||
return redirect_url
|
||||
return dify_config.CONSOLE_WEB_URL
|
||||
|
||||
|
||||
def _preferred_interface_language(language: str | None = None) -> str:
|
||||
if language:
|
||||
return language
|
||||
@@ -139,7 +109,6 @@ class OAuthLogin(Resource):
|
||||
invite_token = request.args.get("invite_token") or None
|
||||
timezone = _validated_timezone(request.args.get("timezone") or None)
|
||||
language = _validated_language(request.args.get("language") or None)
|
||||
redirect_url = request.args.get("redirect_url") or None
|
||||
OAUTH_PROVIDERS = get_oauth_providers()
|
||||
with current_app.app_context():
|
||||
oauth_provider = OAUTH_PROVIDERS.get(provider)
|
||||
@@ -150,7 +119,6 @@ class OAuthLogin(Resource):
|
||||
invite_token=invite_token,
|
||||
timezone=timezone,
|
||||
language=language,
|
||||
redirect_url=redirect_url,
|
||||
)
|
||||
return redirect(auth_url)
|
||||
|
||||
@@ -176,7 +144,6 @@ class OAuthCallback(Resource):
|
||||
invite_token = oauth_state.get("invite_token")
|
||||
timezone = _validated_timezone(oauth_state.get("timezone"))
|
||||
language = _validated_language(oauth_state.get("language"))
|
||||
redirect_url = oauth_state.get("redirect_url")
|
||||
|
||||
if not code:
|
||||
return {"error": "Authorization code is required"}, 400
|
||||
@@ -215,8 +182,6 @@ 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}")
|
||||
|
||||
@@ -245,9 +210,9 @@ class OAuthCallback(Resource):
|
||||
ip_address=extract_remote_ip(request),
|
||||
)
|
||||
|
||||
target_url = _get_redirect_target(redirect_url)
|
||||
query_char = "&" if "?" in target_url else "?"
|
||||
target_url = f"{target_url}{query_char}oauth_new_user={str(oauth_new_user).lower()}"
|
||||
base_url = dify_config.CONSOLE_WEB_URL
|
||||
query_char = "&" if "?" in base_url else "?"
|
||||
target_url = f"{base_url}{query_char}oauth_new_user={str(oauth_new_user).lower()}"
|
||||
response = redirect(target_url)
|
||||
|
||||
set_access_token_to_cookie(request, response, token_pair.access_token)
|
||||
@@ -281,7 +246,10 @@ def _generate_account(
|
||||
if not FeatureService.get_system_features().is_allow_create_workspace:
|
||||
raise WorkSpaceNotAllowedCreateError()
|
||||
else:
|
||||
TenantService.create_owner_tenant(account, session=db.session())
|
||||
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session)
|
||||
TenantService.create_tenant_member(new_tenant, account, db.session, role="owner")
|
||||
account.current_tenant = new_tenant
|
||||
tenant_was_created.send(new_tenant)
|
||||
|
||||
if not account:
|
||||
normalized_email = user_info.email.lower()
|
||||
|
||||
@@ -50,8 +50,6 @@ from models.provider_ids import ModelProviderID
|
||||
from services.api_token_service import ApiTokenCache
|
||||
from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
|
||||
from services.enterprise import rbac_service as enterprise_rbac_service
|
||||
from services.enterprise.rbac_service import RBACResourceWhitelistScope, ReplaceMemberBindings
|
||||
from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task
|
||||
|
||||
register_response_schema_models(console_ns, ApiBaseUrlResponse, SimpleResultResponse, UsageCheckResponse)
|
||||
|
||||
@@ -567,16 +565,6 @@ class DatasetListApi(Resource):
|
||||
except services.errors.dataset.DatasetNameDuplicateError:
|
||||
raise DatasetNameDuplicateError()
|
||||
|
||||
if dify_config.RBAC_ENABLED:
|
||||
if permission == DatasetPermissionEnum.ALL_TEAM:
|
||||
enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist(
|
||||
current_tenant_id,
|
||||
current_user.id,
|
||||
dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_created_app_rbac_access_task.delay(current_tenant_id, current_user.id, dataset_id=dataset.id)
|
||||
|
||||
permission_keys_map = enterprise_rbac_service.RBACService.DatasetPermissions.batch_get(
|
||||
str(current_tenant_id),
|
||||
current_user.id,
|
||||
|
||||
@@ -58,12 +58,6 @@ 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,6 +85,7 @@ 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
|
||||
@@ -123,6 +124,7 @@ 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)),
|
||||
|
||||
@@ -106,7 +106,7 @@ class RecommendedAppListApi(Resource):
|
||||
language_prefix = _resolve_language(args.language, current_user)
|
||||
|
||||
return RecommendedAppListResponse.model_validate(
|
||||
RecommendedAppService.get_recommended_apps_and_categories(language_prefix, session=db.session()),
|
||||
RecommendedAppService.get_recommended_apps_and_categories(db.session, language_prefix),
|
||||
from_attributes=True,
|
||||
).model_dump(mode="json")
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ from controllers.console.auth.error import (
|
||||
NotOwnerError,
|
||||
OwnerTransferLimitError,
|
||||
)
|
||||
from controllers.console.error import EmailSendIpLimitError, SeatsLimitExceeded, WorkspaceMembersLimitExceeded
|
||||
from controllers.console.error import EmailSendIpLimitError, WorkspaceMembersLimitExceeded
|
||||
from controllers.console.wraps import (
|
||||
account_initialization_required,
|
||||
is_allow_transfer_owner,
|
||||
@@ -128,14 +128,12 @@ def _normalize_invitee_emails(emails: list[str]) -> list[str]:
|
||||
return list(dict.fromkeys(email.lower() for email in emails))
|
||||
|
||||
|
||||
def _count_new_member_invites(tenant_id: str, emails: list[str]) -> tuple[int, int]:
|
||||
def _count_new_member_invites(tenant_id: str, emails: list[str]) -> int:
|
||||
new_member_count = 0
|
||||
new_account_count = 0
|
||||
for email in emails:
|
||||
account = AccountService.get_account_by_email_with_case_fallback(db.session, email)
|
||||
if not account:
|
||||
new_member_count += 1
|
||||
new_account_count += 1
|
||||
continue
|
||||
|
||||
exists = db.session.scalar(
|
||||
@@ -146,7 +144,7 @@ def _count_new_member_invites(tenant_id: str, emails: list[str]) -> tuple[int, i
|
||||
if not exists:
|
||||
new_member_count += 1
|
||||
|
||||
return new_member_count, new_account_count
|
||||
return new_member_count
|
||||
|
||||
|
||||
def _count_current_members(tenant_id: str) -> int:
|
||||
@@ -155,7 +153,7 @@ def _count_current_members(tenant_id: str) -> int:
|
||||
)
|
||||
|
||||
|
||||
def _check_member_invite_limits(tenant_id: str, new_member_count: int, new_account_count: int) -> None:
|
||||
def _check_member_invite_limits(tenant_id: str, new_member_count: int) -> None:
|
||||
if new_member_count <= 0:
|
||||
return
|
||||
|
||||
@@ -165,10 +163,6 @@ def _check_member_invite_limits(tenant_id: str, new_member_count: int, new_accou
|
||||
workspace_members = features.workspace_members
|
||||
if workspace_members.enabled is True and not workspace_members.is_available(new_member_count):
|
||||
raise WorkspaceMembersLimitExceeded()
|
||||
if new_account_count > 0:
|
||||
seats = FeatureService.get_system_features(is_authenticated=True).license.seats
|
||||
if not seats.is_available(new_account_count):
|
||||
raise SeatsLimitExceeded()
|
||||
return
|
||||
|
||||
if dify_config.BILLING_ENABLED and features.billing.enabled is True:
|
||||
@@ -266,8 +260,8 @@ class MemberInviteEmailApi(Resource):
|
||||
tenant_id = inviter.current_tenant.id
|
||||
with redis_client.lock(f"workspace_member_invite:{tenant_id}", timeout=60):
|
||||
if dify_config.ENTERPRISE_ENABLED is True or dify_config.BILLING_ENABLED is True:
|
||||
new_member_count, new_account_count = _count_new_member_invites(tenant_id, invitee_emails)
|
||||
_check_member_invite_limits(tenant_id, new_member_count, new_account_count)
|
||||
new_member_count = _count_new_member_invites(tenant_id, invitee_emails)
|
||||
_check_member_invite_limits(tenant_id, new_member_count)
|
||||
|
||||
for invitee_email in invitee_emails:
|
||||
try:
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from flask import request
|
||||
@@ -13,11 +14,9 @@ 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
|
||||
from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task
|
||||
|
||||
|
||||
class _RBACRoleList(svc.Paginated[svc.RBACRole]):
|
||||
@@ -512,8 +511,14 @@ class RBACAccessPolicyBindingUnlockApi(Resource):
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _AccessScope(StrEnum):
|
||||
ALL = "all"
|
||||
SPECIFIC = "specific"
|
||||
ONLY_ME = "only_me"
|
||||
|
||||
|
||||
class _ResourceAccessScopeRequest(BaseModel):
|
||||
scope: RBACResourceWhitelistScope
|
||||
scope: _AccessScope
|
||||
|
||||
|
||||
class _ReplaceBindingsRequest(BaseModel):
|
||||
@@ -579,15 +584,14 @@ class RBACAppWhitelistApi(Resource):
|
||||
def put(self, app_id):
|
||||
tenant_id, account_id = _current_ids()
|
||||
request = _payload(_ResourceAccessScopeRequest)
|
||||
result = svc.RBACService.AppAccess.replace_whitelist(
|
||||
tenant_id,
|
||||
account_id,
|
||||
str(app_id),
|
||||
svc.ReplaceMemberBindings(scope=request.scope.value),
|
||||
return _dump(
|
||||
svc.RBACService.AppAccess.replace_whitelist(
|
||||
tenant_id,
|
||||
account_id,
|
||||
str(app_id),
|
||||
svc.ReplaceMemberBindings(scope=request.scope.value),
|
||||
)
|
||||
)
|
||||
if dify_config.RBAC_ENABLED and request.scope is RBACResourceWhitelistScope.ALL:
|
||||
initialize_created_app_rbac_access_task.delay(tenant_id, account_id, str(app_id))
|
||||
return _dump(result)
|
||||
|
||||
|
||||
@console_ns.route("/workspaces/current/rbac/apps/<uuid:app_id>/user-access-policies")
|
||||
@@ -612,8 +616,8 @@ class RBACAppUserAccessPolicyAssignmentApi(Resource):
|
||||
svc.RBACService.AppAccess.replace_user_access_policies(
|
||||
tenant_id,
|
||||
account_id,
|
||||
app_id,
|
||||
target_account_id,
|
||||
str(app_id),
|
||||
str(target_account_id),
|
||||
payload,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -47,12 +47,10 @@ class EnterpriseWorkspace(Resource):
|
||||
if account is None:
|
||||
return {"message": "owner account not found."}, 404
|
||||
|
||||
tenant = TenantService.create_owner_tenant(
|
||||
account,
|
||||
name=args.name,
|
||||
is_from_dashboard=True,
|
||||
session=db.session(),
|
||||
)
|
||||
tenant = TenantService.create_tenant(args.name, is_from_dashboard=True, session=db.session)
|
||||
TenantService.create_tenant_member(tenant, account, db.session, role="owner")
|
||||
|
||||
tenant_was_created.send(tenant)
|
||||
|
||||
resp = {
|
||||
"id": tenant.id,
|
||||
|
||||
@@ -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
|
||||
from controllers.common.fields import EventStreamResponse, SimpleResultResponse
|
||||
from controllers.common.schema import register_enum_models, register_response_schema_models, register_schema_models
|
||||
from controllers.openapi._models import (
|
||||
AccountPayload,
|
||||
@@ -95,6 +95,7 @@ register_response_schema_models(
|
||||
openapi_ns,
|
||||
ErrorBody,
|
||||
EventStreamResponse,
|
||||
SimpleResultResponse,
|
||||
UsageInfo,
|
||||
MessageMetadata,
|
||||
AppListRow,
|
||||
|
||||
@@ -34,6 +34,7 @@ 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"
|
||||
@@ -223,6 +224,19 @@ 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, Iterator
|
||||
from collections.abc import Callable, Generator
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
@@ -61,7 +61,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _translate_service_errors() -> Iterator[None]:
|
||||
def _translate_service_errors() -> Generator[None, None, None]:
|
||||
try:
|
||||
yield
|
||||
except WorkflowNotFoundError as ex:
|
||||
@@ -166,6 +166,7 @@ class AppRunApi(Resource):
|
||||
surface="apps",
|
||||
)
|
||||
|
||||
# response-contract:ignore compact_generate_response
|
||||
return helper.compact_generate_response(stream_obj)
|
||||
|
||||
|
||||
|
||||
@@ -25,12 +25,13 @@ from controllers.openapi._models import (
|
||||
AppListRow,
|
||||
)
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind, 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
|
||||
@@ -166,7 +167,9 @@ 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 != "end_user" and auth_data.account_id is not None
|
||||
dify_config.RBAC_ENABLED
|
||||
and auth_data.caller_kind != CallerKind.END_USER
|
||||
and auth_data.account_id is not None
|
||||
)
|
||||
access_filter = AppAccessFilter.unrestricted()
|
||||
if apply_rbac_filter:
|
||||
@@ -203,7 +206,7 @@ class AppListApi(Resource):
|
||||
limit=query.limit,
|
||||
mode=query.mode.value if query.mode else "all", # type:ignore
|
||||
name=query.name,
|
||||
status="normal",
|
||||
status=AppStatus.NORMAL,
|
||||
# Visibility gate pushed into the query — pagination.total stays
|
||||
# consistent across pages because invisible rows never count.
|
||||
openapi_visible=True,
|
||||
|
||||
@@ -25,6 +25,7 @@ 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
|
||||
@@ -62,7 +63,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 != "normal":
|
||||
if not app or app.status != AppStatus.NORMAL:
|
||||
continue
|
||||
tenant = tenants_by_id.get(str(app.tenant_id))
|
||||
items.append(
|
||||
|
||||
@@ -2,7 +2,6 @@ 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
|
||||
@@ -21,6 +20,11 @@ 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
|
||||
@@ -78,9 +82,9 @@ class AuthData(BaseModel):
|
||||
tenant_role: TenantAccountRole | None = None
|
||||
|
||||
caller: Account | EndUser | None = None
|
||||
caller_kind: Literal["account", "end_user"] | None = None
|
||||
caller_kind: CallerKind | None = None
|
||||
|
||||
def require_app_context(self) -> tuple[App, Account | EndUser, Literal["account", "end_user"]]:
|
||||
def require_app_context(self) -> tuple[App, Account | EndUser, CallerKind]:
|
||||
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,6 +17,7 @@ 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,
|
||||
@@ -28,7 +29,9 @@ 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,
|
||||
@@ -217,7 +220,12 @@ class PipelineRouter:
|
||||
if not token:
|
||||
raise Unauthorized("bearer required")
|
||||
|
||||
identity = get_authenticator().authenticate(token)
|
||||
try:
|
||||
identity = get_authenticator().authenticate(token)
|
||||
except TokenExpiredError:
|
||||
raise SessionExpired()
|
||||
except InvalidBearerError:
|
||||
raise InvalidBearer()
|
||||
|
||||
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
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind
|
||||
from extensions.ext_database import db
|
||||
from models.account import TenantStatus
|
||||
from models.enums import EndUserType
|
||||
from models.account import AccountStatus, TenantStatus
|
||||
from models.enums import AppStatus, 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 != "normal":
|
||||
if not app or app.status != AppStatus.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 = "account"
|
||||
data.caller_kind = CallerKind.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) != "active":
|
||||
if data.caller is not None and getattr(data.caller, "status", None) != AccountStatus.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 = "end_user"
|
||||
data.caller_kind = CallerKind.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
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind
|
||||
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 != "account":
|
||||
if data.caller_kind != CallerKind.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, RBACRequirement
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind, 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 == "account":
|
||||
if caller_kind == CallerKind.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, RBACRequirement
|
||||
from controllers.openapi.auth.data import AuthData, CallerKind, 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 == "account":
|
||||
if caller_kind == CallerKind.ACCOUNT:
|
||||
if workflow_run.created_by_role != CreatorUserRole.ACCOUNT or workflow_run.created_by != caller.id:
|
||||
raise NotFound("Workflow run not found")
|
||||
else:
|
||||
|
||||
@@ -48,7 +48,6 @@ from services.errors.account import (
|
||||
MemberNotInTenantError,
|
||||
NoPermissionError,
|
||||
RoleAlreadyAssignedError,
|
||||
SeatsLimitExceededError,
|
||||
)
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
@@ -191,8 +190,6 @@ 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))
|
||||
|
||||
|
||||
@@ -43,7 +43,6 @@ from models.dataset import DatasetPermissionEnum
|
||||
from models.enums import TagType
|
||||
from models.provider_ids import ModelProviderID
|
||||
from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
|
||||
from services.enterprise.rbac_service import RBACResourceWhitelistScope, RBACService, ReplaceMemberBindings
|
||||
from services.entities.knowledge_entities.knowledge_entities import (
|
||||
ExternalRetrievalModel,
|
||||
KnowledgeProvider,
|
||||
@@ -59,7 +58,6 @@ from services.tag_service import (
|
||||
from services.tag_service import (
|
||||
UpdateTagPayload as UpdateTagServicePayload,
|
||||
)
|
||||
from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task
|
||||
|
||||
register_enum_models(service_api_ns, DatasetPermissionEnum)
|
||||
|
||||
@@ -525,15 +523,6 @@ class DatasetListApi(DatasetApiResource):
|
||||
except services.errors.dataset.DatasetNameDuplicateError:
|
||||
raise DatasetNameDuplicateError()
|
||||
|
||||
if payload.permission == DatasetPermissionEnum.ALL_TEAM and dify_config.RBAC_ENABLED:
|
||||
RBACService.DatasetAccess.replace_whitelist(
|
||||
tenant_id,
|
||||
current_user.id,
|
||||
dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_created_app_rbac_access_task.delay(tenant_id, current_user.id, dataset_id=dataset.id)
|
||||
|
||||
return _dump_service_dataset_detail(dataset), 200
|
||||
|
||||
|
||||
|
||||
@@ -12,7 +12,6 @@ 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
|
||||
@@ -270,8 +269,8 @@ def cloud_edition_billing_rate_limit_check[**P, R](
|
||||
subscription_plan=knowledge_rate_limit.subscription_plan,
|
||||
operation="knowledge",
|
||||
)
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
|
||||
session.add(rate_limit_log)
|
||||
db.session.add(rate_limit_log)
|
||||
db.session.commit()
|
||||
raise Forbidden(
|
||||
"Sorry, you have reached the knowledge base request rate limit of your subscription."
|
||||
)
|
||||
|
||||
@@ -47,7 +47,6 @@ from core.repositories import DifyCoreRepositoryFactory
|
||||
from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository
|
||||
from extensions.ext_database import db
|
||||
from factories import file_factory
|
||||
from graphon.filters import ResponseStreamFilter
|
||||
from graphon.graph_engine.layers import GraphEngineLayer
|
||||
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
|
||||
from graphon.runtime import GraphRuntimeState
|
||||
@@ -269,7 +268,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
workflow_node_execution_repository: WorkflowNodeExecutionRepository,
|
||||
graph_runtime_state: GraphRuntimeState,
|
||||
pause_state_config: PauseStateLayerConfig | None = None,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
):
|
||||
"""
|
||||
Resume a paused advanced chat execution.
|
||||
@@ -299,7 +297,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
stream=application_generate_entity.stream,
|
||||
pause_state_config=pause_state_config,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
response_stream_filter=response_stream_filter,
|
||||
)
|
||||
|
||||
def single_iteration_generate(
|
||||
@@ -494,7 +491,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
pause_state_config: PauseStateLayerConfig | None = None,
|
||||
graph_runtime_state: GraphRuntimeState | None = None,
|
||||
graph_engine_layers: Sequence[GraphEngineLayer] = (),
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
|
||||
"""
|
||||
Generate App response.
|
||||
@@ -542,14 +538,12 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
)
|
||||
|
||||
graph_layers: list[GraphEngineLayer] = list(graph_engine_layers)
|
||||
resolved_response_stream_filter = response_stream_filter or ResponseStreamFilter()
|
||||
if pause_state_config is not None:
|
||||
graph_layers.append(
|
||||
PauseStatePersistenceLayer(
|
||||
session_factory=pause_state_config.session_factory,
|
||||
generate_entity=application_generate_entity,
|
||||
state_owner_user_id=pause_state_config.state_owner_user_id,
|
||||
response_stream_filter=resolved_response_stream_filter,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -570,7 +564,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
"workflow_node_execution_repository": workflow_node_execution_repository,
|
||||
"graph_engine_layers": tuple(graph_layers),
|
||||
"graph_runtime_state": graph_runtime_state,
|
||||
"response_stream_filter": resolved_response_stream_filter,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -610,7 +603,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
workflow_node_execution_repository: WorkflowNodeExecutionRepository,
|
||||
graph_engine_layers: Sequence[GraphEngineLayer] = (),
|
||||
graph_runtime_state: GraphRuntimeState | None = None,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
):
|
||||
"""
|
||||
Generate worker in a new thread.
|
||||
@@ -670,7 +662,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
workflow_node_execution_repository=workflow_node_execution_repository,
|
||||
graph_engine_layers=graph_engine_layers,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
response_stream_filter=response_stream_filter,
|
||||
)
|
||||
|
||||
try:
|
||||
|
||||
@@ -39,7 +39,6 @@ from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from extensions.otel import WorkflowAppRunnerHandler, trace_span
|
||||
from graphon.enums import WorkflowType
|
||||
from graphon.filters import ResponseStreamFilter
|
||||
from graphon.graph_engine.command_channels import RedisChannel
|
||||
from graphon.graph_engine.layers import GraphEngineLayer
|
||||
from graphon.runtime import GraphRuntimeState, VariablePool
|
||||
@@ -74,7 +73,6 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
|
||||
workflow_node_execution_repository: WorkflowNodeExecutionRepository,
|
||||
graph_engine_layers: Sequence[GraphEngineLayer] = (),
|
||||
graph_runtime_state: GraphRuntimeState | None = None,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
):
|
||||
super().__init__(
|
||||
queue_manager=queue_manager,
|
||||
@@ -92,7 +90,6 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
|
||||
self._workflow_execution_repository = workflow_execution_repository
|
||||
self._workflow_node_execution_repository = workflow_node_execution_repository
|
||||
self._resume_graph_runtime_state = graph_runtime_state
|
||||
self._response_stream_filter = response_stream_filter
|
||||
|
||||
@trace_span(WorkflowAppRunnerHandler)
|
||||
def run(self):
|
||||
@@ -230,7 +227,6 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
|
||||
variable_pool=variable_pool,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
command_channel=command_channel,
|
||||
response_stream_filter=self._response_stream_filter,
|
||||
)
|
||||
|
||||
self._queue_manager.graph_runtime_state = graph_runtime_state
|
||||
|
||||
@@ -42,7 +42,6 @@ from core.repositories import DifyCoreRepositoryFactory
|
||||
from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository
|
||||
from extensions.ext_database import db
|
||||
from factories import file_factory
|
||||
from graphon.filters import ResponseStreamFilter
|
||||
from graphon.graph_engine.layers import GraphEngineLayer
|
||||
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
|
||||
from graphon.runtime import GraphRuntimeState
|
||||
@@ -281,7 +280,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
graph_engine_layers: Sequence[GraphEngineLayer] = (),
|
||||
pause_state_config: PauseStateLayerConfig | None = None,
|
||||
variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
|
||||
"""
|
||||
Resume a paused workflow execution using the persisted runtime state.
|
||||
@@ -312,7 +310,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
graph_engine_layers=graph_engine_layers,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
pause_state_config=pause_state_config,
|
||||
response_stream_filter=response_stream_filter,
|
||||
)
|
||||
|
||||
def _generate(
|
||||
@@ -331,7 +328,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
graph_engine_layers: Sequence[GraphEngineLayer] = (),
|
||||
graph_runtime_state: GraphRuntimeState | None = None,
|
||||
pause_state_config: PauseStateLayerConfig | None = None,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
|
||||
"""
|
||||
Generate App response.
|
||||
@@ -360,14 +356,12 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
app_mode=app_model.mode,
|
||||
)
|
||||
|
||||
resolved_response_stream_filter = response_stream_filter or ResponseStreamFilter()
|
||||
if pause_state_config is not None:
|
||||
graph_layers.append(
|
||||
PauseStatePersistenceLayer(
|
||||
session_factory=pause_state_config.session_factory,
|
||||
generate_entity=application_generate_entity,
|
||||
state_owner_user_id=pause_state_config.state_owner_user_id,
|
||||
response_stream_filter=resolved_response_stream_filter,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -390,7 +384,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
"workflow_node_execution_repository": workflow_node_execution_repository,
|
||||
"graph_engine_layers": tuple(graph_layers),
|
||||
"graph_runtime_state": graph_runtime_state,
|
||||
"response_stream_filter": resolved_response_stream_filter,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -597,7 +590,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
root_node_id: str | None = None,
|
||||
graph_engine_layers: Sequence[GraphEngineLayer] = (),
|
||||
graph_runtime_state: GraphRuntimeState | None = None,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Generate worker in a new thread.
|
||||
@@ -646,7 +638,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
root_node_id=root_node_id,
|
||||
graph_engine_layers=graph_engine_layers,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
response_stream_filter=response_stream_filter,
|
||||
)
|
||||
|
||||
try:
|
||||
|
||||
@@ -18,7 +18,6 @@ from core.workflow.workflow_entry import WorkflowEntry
|
||||
from extensions.ext_redis import redis_client
|
||||
from extensions.otel import WorkflowAppRunnerHandler, trace_span
|
||||
from graphon.enums import WorkflowType
|
||||
from graphon.filters import ResponseStreamFilter
|
||||
from graphon.graph_engine.command_channels import RedisChannel
|
||||
from graphon.graph_engine.layers import GraphEngineLayer
|
||||
from graphon.runtime import GraphRuntimeState, VariablePool
|
||||
@@ -47,7 +46,6 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
|
||||
workflow_node_execution_repository: WorkflowNodeExecutionRepository,
|
||||
graph_engine_layers: Sequence[GraphEngineLayer] = (),
|
||||
graph_runtime_state: GraphRuntimeState | None = None,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
):
|
||||
super().__init__(
|
||||
queue_manager=queue_manager,
|
||||
@@ -62,7 +60,6 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
|
||||
self._workflow_execution_repository = workflow_execution_repository
|
||||
self._workflow_node_execution_repository = workflow_node_execution_repository
|
||||
self._resume_graph_runtime_state = graph_runtime_state
|
||||
self._response_stream_filter = response_stream_filter
|
||||
|
||||
@trace_span(WorkflowAppRunnerHandler)
|
||||
def run(self):
|
||||
@@ -166,7 +163,6 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
|
||||
variable_pool=variable_pool,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
command_channel=command_channel,
|
||||
response_stream_filter=self._response_stream_filter,
|
||||
)
|
||||
|
||||
persistence_layer = WorkflowPersistenceLayer(
|
||||
|
||||
@@ -7,7 +7,6 @@ from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity
|
||||
from core.workflow.system_variables import SystemVariableKey, get_system_text
|
||||
from graphon.filters import ResponseStreamFilter
|
||||
from graphon.graph_engine.layers import GraphEngineLayer
|
||||
from graphon.graph_events import GraphEngineEvent, GraphRunPausedEvent
|
||||
from models.model import AppMode
|
||||
@@ -42,10 +41,6 @@ class WorkflowResumptionContext(BaseModel):
|
||||
# Only workflow / chatflow could be paused.
|
||||
generate_entity: _GenerateEntityUnion
|
||||
serialized_graph_runtime_state: str
|
||||
# Optional so that a workflow run paused before this field existed still
|
||||
# loads: it just degrades to fresh-filter behavior on resume for that one
|
||||
# stale run.
|
||||
serialized_response_stream_filter_state: str | None = None
|
||||
|
||||
def dumps(self) -> str:
|
||||
return self.model_dump_json()
|
||||
@@ -57,12 +52,6 @@ class WorkflowResumptionContext(BaseModel):
|
||||
def get_generate_entity(self) -> WorkflowAppGenerateEntity | AdvancedChatAppGenerateEntity:
|
||||
return self.generate_entity.entity
|
||||
|
||||
def get_response_stream_filter(self) -> ResponseStreamFilter:
|
||||
response_stream_filter = ResponseStreamFilter()
|
||||
if self.serialized_response_stream_filter_state is not None:
|
||||
response_stream_filter.loads(self.serialized_response_stream_filter_state)
|
||||
return response_stream_filter
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PauseStateLayerConfig:
|
||||
@@ -78,17 +67,11 @@ class PauseStatePersistenceLayer(GraphEngineLayer):
|
||||
session_factory: Engine | sessionmaker[Session],
|
||||
generate_entity: WorkflowAppGenerateEntity | AdvancedChatAppGenerateEntity,
|
||||
state_owner_user_id: str,
|
||||
response_stream_filter: ResponseStreamFilter,
|
||||
):
|
||||
"""Create a PauseStatePersistenceLayer.
|
||||
|
||||
The `state_owner_user_id` is used when creating state file for pause.
|
||||
It generally should id of the creator of workflow.
|
||||
|
||||
`response_stream_filter` must be the exact same instance that
|
||||
`WorkflowEntry` is using to stream this run's events — this layer
|
||||
dumps its state on pause, and a different instance would silently
|
||||
persist the wrong (empty) filter state.
|
||||
"""
|
||||
if isinstance(session_factory, Engine):
|
||||
session_factory = sessionmaker(session_factory)
|
||||
@@ -96,7 +79,6 @@ class PauseStatePersistenceLayer(GraphEngineLayer):
|
||||
self._session_maker = session_factory
|
||||
self._state_owner_user_id = state_owner_user_id
|
||||
self._generate_entity = generate_entity
|
||||
self._response_stream_filter = response_stream_filter
|
||||
|
||||
def _get_repo(self) -> APIWorkflowRunRepository:
|
||||
return DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._session_maker)
|
||||
@@ -137,7 +119,6 @@ class PauseStatePersistenceLayer(GraphEngineLayer):
|
||||
state = WorkflowResumptionContext(
|
||||
serialized_graph_runtime_state=self.graph_runtime_state.dumps(),
|
||||
generate_entity=entity_wrapper,
|
||||
serialized_response_stream_filter_state=self._response_stream_filter.dumps(),
|
||||
)
|
||||
|
||||
workflow_run_id = get_system_text(
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
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
|
||||
@@ -20,9 +19,8 @@ class DifyWorkflowCallbackHandler(DifyAgentCallbackHandler):
|
||||
trace_manager: TraceQueueManager | None = None,
|
||||
) -> Generator[ToolInvokeMessage, None, None]:
|
||||
for tool_output in tool_outputs:
|
||||
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")
|
||||
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 as CodeLanguage # noqa: PLC0414
|
||||
from graphon.nodes.code.entities import CodeLanguage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
code_execution_endpoint_url = URL(str(dify_config.CODE_EXECUTION_ENDPOINT))
|
||||
@@ -133,9 +133,7 @@ class CodeExecutor:
|
||||
return response_code.data.stdout or ""
|
||||
|
||||
@classmethod
|
||||
def execute_workflow_code_template(
|
||||
cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]):
|
||||
"""
|
||||
Execute code
|
||||
:param language: code language
|
||||
|
||||
@@ -11,7 +11,7 @@ class Jinja2TemplateTransformer(TemplateTransformer):
|
||||
|
||||
@classmethod
|
||||
@override
|
||||
def transform_response(cls, response: str) -> dict[str, Any]:
|
||||
def transform_response(cls, response: str):
|
||||
"""
|
||||
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) -> str:
|
||||
def extract_result_str_from_response(cls, response: 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) -> dict[str, Any]:
|
||||
def transform_response(cls, response: str) -> Mapping[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[str, Any]) -> dict[str, Any]:
|
||||
def _post_process_result(cls, result: dict[Any, Any]) -> dict[Any, 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 {key: convert_scientific_notation(value) for key, value in result.items()}
|
||||
return convert_scientific_notation(result)
|
||||
|
||||
@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 isinstance(claim_code, str) or not claim_code:
|
||||
if not claim_code:
|
||||
raise ValueError("Creators Platform did not return a valid claim_code")
|
||||
return claim_code
|
||||
|
||||
|
||||
@@ -10,21 +10,18 @@ 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 core.db import session_factory
|
||||
from extensions.ext_database import db
|
||||
from models.provider import ProviderCredential, ProviderModelCredential
|
||||
from models.tools import BuiltinToolProvider
|
||||
|
||||
with session_factory.create_session() as session:
|
||||
with Session(db.engine) as session:
|
||||
if credential_type == PluginCredentialType.MODEL:
|
||||
# Check both pre-defined and custom model credentials using a single UNION query
|
||||
stmt = (
|
||||
@@ -45,7 +42,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,7 +1,4 @@
|
||||
from typing import Any
|
||||
|
||||
|
||||
def download_with_size_limit(url: str, max_download_size: int, **kwargs: Any) -> bytes:
|
||||
def download_with_size_limit(url, max_download_size: int, **kwargs):
|
||||
from core.file import remote_fetcher
|
||||
|
||||
response = remote_fetcher.make_request("GET", url, follow_redirects=True, **kwargs)
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
import base64
|
||||
|
||||
from Crypto.PublicKey import RSA
|
||||
|
||||
from libs import rsa
|
||||
|
||||
|
||||
@@ -13,13 +11,13 @@ def obfuscated_token(token: str) -> str:
|
||||
return token[:6] + "*" * 12 + token[-2:]
|
||||
|
||||
|
||||
def full_mask_token(token_length: int = 20) -> str:
|
||||
def full_mask_token(token_length=20):
|
||||
return "*" * token_length
|
||||
|
||||
|
||||
def encrypt_token(tenant_id: str, token: str) -> str:
|
||||
def encrypt_token(tenant_id: str, token: str):
|
||||
from extensions.ext_database import db
|
||||
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")
|
||||
@@ -32,15 +30,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]) -> list[str]:
|
||||
def batch_decrypt_token(tenant_id: str, tokens: 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) -> tuple[RSA.RsaKey, object]:
|
||||
def get_decrypt_decoding(tenant_id: str):
|
||||
return rsa.get_decrypt_decoding(tenant_id)
|
||||
|
||||
|
||||
def decrypt_token_with_decoding(token: str, rsa_key: RSA.RsaKey, cipher_rsa: object) -> str:
|
||||
def decrypt_token_with_decoding(token: str, rsa_key, cipher_rsa):
|
||||
return rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from yarl import URL
|
||||
@@ -20,7 +19,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) -> bytes:
|
||||
def download_plugin_pkg(plugin_unique_identifier: str):
|
||||
return download_with_size_limit(get_plugin_pkg_url(plugin_unique_identifier), dify_config.PLUGIN_MAX_PACKAGE_SIZE)
|
||||
|
||||
|
||||
@@ -40,7 +39,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[str, Any]]:
|
||||
def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]:
|
||||
if not plugin_ids:
|
||||
return []
|
||||
|
||||
@@ -54,19 +53,10 @@ def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict[str, Any]]:
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
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
|
||||
return data.get("data", {}).get("plugins", [])
|
||||
|
||||
|
||||
def record_install_plugin_event(plugin_unique_identifier: str) -> None:
|
||||
def record_install_plugin_event(plugin_unique_identifier: str):
|
||||
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]) -> None:
|
||||
def set(self, credentials: dict[str, Any]):
|
||||
"""
|
||||
Cache model provider credentials.
|
||||
|
||||
@@ -43,7 +43,7 @@ class ProviderCredentialsCache:
|
||||
"""
|
||||
redis_client.setex(self.cache_key, 86400, json.dumps(credentials))
|
||||
|
||||
def delete(self) -> None:
|
||||
def delete(self):
|
||||
"""
|
||||
Delete cached model provider credentials.
|
||||
|
||||
|
||||
@@ -20,18 +20,17 @@ 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
|
||||
new_spec = importlib.util.spec_from_file_location(module_name, py_file_path)
|
||||
if not new_spec or not new_spec.loader:
|
||||
# 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:
|
||||
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
|
||||
new_spec.loader = importlib.util.LazyLoader(new_spec.loader)
|
||||
spec = new_spec
|
||||
spec.loader = importlib.util.LazyLoader(spec.loader)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
if not existed_spec:
|
||||
sys.modules[module_name] = module
|
||||
if spec.loader is not None:
|
||||
spec.loader.exec_module(module)
|
||||
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: Any) -> None:
|
||||
def __init__(self, **kwargs):
|
||||
self.cache_key = self._generate_cache_key(**kwargs)
|
||||
|
||||
@abstractmethod
|
||||
def _generate_cache_key(self, **kwargs: Any) -> str:
|
||||
def _generate_cache_key(self, **kwargs) -> 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]) -> None:
|
||||
def set(self, config: dict[str, Any]):
|
||||
"""Cache provider credentials"""
|
||||
redis_client.setex(self.cache_key, 86400, json.dumps(config))
|
||||
|
||||
def delete(self) -> None:
|
||||
def delete(self):
|
||||
"""Delete cached provider credentials"""
|
||||
redis_client.delete(self.cache_key)
|
||||
|
||||
@@ -48,7 +48,7 @@ class SingletonProviderCredentialsCache(ProviderCredentialsCache):
|
||||
)
|
||||
|
||||
@override
|
||||
def _generate_cache_key(self, **kwargs: Any) -> str:
|
||||
def _generate_cache_key(self, **kwargs) -> 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: Any) -> str:
|
||||
def _generate_cache_key(self, **kwargs) -> 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]) -> None:
|
||||
def set(self, config: dict[str, Any]):
|
||||
"""Cache provider credentials"""
|
||||
pass
|
||||
|
||||
def delete(self) -> None:
|
||||
def delete(self):
|
||||
"""Delete cached provider credentials"""
|
||||
pass
|
||||
|
||||
@@ -125,7 +125,5 @@ class ProviderConfigEncrypter:
|
||||
return data
|
||||
|
||||
|
||||
def create_provider_encrypter(
|
||||
tenant_id: str, config: list[BasicProviderConfig], cache: ProviderConfigCache
|
||||
) -> tuple[ProviderConfigEncrypter, ProviderConfigCache]:
|
||||
def create_provider_encrypter(tenant_id: str, config: list[BasicProviderConfig], cache: 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]) -> None:
|
||||
def set(self, parameters: dict[str, Any]):
|
||||
"""Cache model provider credentials."""
|
||||
redis_client.setex(self.cache_key, 86400, json.dumps(parameters))
|
||||
|
||||
def delete(self) -> None:
|
||||
def delete(self):
|
||||
"""
|
||||
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]) -> dict[str, Any]:
|
||||
def extract_external_trace_id_from_args(args: Mapping[str, Any]):
|
||||
"""
|
||||
Extract 'external_trace_id' from args.
|
||||
|
||||
|
||||
@@ -76,11 +76,7 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
if not user_id:
|
||||
user = EndUserService.get_or_create_end_user(app)
|
||||
else:
|
||||
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)
|
||||
user = cls._get_user(user_id, app)
|
||||
|
||||
conversation_id = conversation_id or ""
|
||||
|
||||
@@ -230,13 +226,6 @@ 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) -> str:
|
||||
def get_name(self):
|
||||
if self == CredentialType.API_KEY:
|
||||
return "API KEY"
|
||||
elif self == CredentialType.OAUTH2:
|
||||
|
||||
@@ -4,8 +4,6 @@ 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_*``
|
||||
@@ -16,15 +14,12 @@ metadata.
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from collections.abc import Mapping, Sequence
|
||||
from mimetypes import guess_type
|
||||
from typing import Literal, Protocol, cast
|
||||
from typing import ClassVar
|
||||
|
||||
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
|
||||
@@ -72,18 +67,14 @@ logger = logging.getLogger(__name__)
|
||||
_provider_entities_adapter: TypeAdapter[list[ProviderEntity]] = TypeAdapter(list[ProviderEntity])
|
||||
|
||||
|
||||
class _RedisLock(Protocol):
|
||||
def acquire(self, *, blocking: bool = True, blocking_timeout: float | None = None) -> bool: ...
|
||||
|
||||
def release(self) -> None: ...
|
||||
|
||||
|
||||
class PluginService:
|
||||
_plugin_model_providers_memory_cache: ClassVar[dict[str, tuple[int, float, tuple[ProviderEntity, ...]]]] = {}
|
||||
|
||||
class LatestPluginCache(BaseModel):
|
||||
plugin_id: str
|
||||
version: str
|
||||
unique_identifier: str
|
||||
status: Literal["active", "deleted"]
|
||||
status: str
|
||||
deprecated_reason: str
|
||||
alternative_plugin_id: str
|
||||
|
||||
@@ -91,12 +82,6 @@ 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.
|
||||
@@ -113,10 +98,6 @@ 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:
|
||||
"""
|
||||
@@ -148,25 +129,8 @@ class PluginService:
|
||||
return declaration
|
||||
|
||||
@classmethod
|
||||
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
|
||||
def _copy_provider_entities(cls, providers: Sequence[ProviderEntity]) -> tuple[ProviderEntity, ...]:
|
||||
return tuple(provider.model_copy(deep=True) for provider in providers)
|
||||
|
||||
@classmethod
|
||||
def _load_plugin_model_providers_generation(cls, tenant_id: str) -> int | None:
|
||||
@@ -199,35 +163,76 @@ class PluginService:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
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
|
||||
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
|
||||
|
||||
cache_keys = [cls._get_plugin_model_providers_cache_key(tenant_id, generation)]
|
||||
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
|
||||
|
||||
try:
|
||||
cached_provider_entries = redis_client.mget(cache_keys)
|
||||
except (LockError, RedisError, RuntimeError):
|
||||
except (RedisError, RuntimeError):
|
||||
logger.warning("Failed to read cached plugin model providers for tenant %s.", tenant_id, exc_info=True)
|
||||
return None, False
|
||||
return None
|
||||
|
||||
if len(cached_provider_entries) != len(cache_keys):
|
||||
logger.warning(
|
||||
"Unexpected cached plugin model providers response size for tenant %s.",
|
||||
tenant_id,
|
||||
)
|
||||
return None, False
|
||||
return None
|
||||
|
||||
for cache_key, cached_providers in zip(cache_keys, cached_provider_entries):
|
||||
if not cached_providers:
|
||||
continue
|
||||
|
||||
try:
|
||||
payload = cls._decode_plugin_model_providers_cache_payload(cached_providers)
|
||||
providers = tuple(_provider_entities_adapter.validate_json(payload))
|
||||
return providers, True
|
||||
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
|
||||
except (TypeError, ValueError, ValidationError):
|
||||
logger.warning(
|
||||
"Invalid cached plugin model providers for tenant %s; deleting cache key %s.",
|
||||
@@ -244,7 +249,7 @@ class PluginService:
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return None, True
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _store_cached_plugin_model_providers(
|
||||
@@ -252,94 +257,15 @@ class PluginService:
|
||||
) -> None:
|
||||
cache_key = cls._get_plugin_model_providers_cache_key(tenant_id, generation)
|
||||
try:
|
||||
payload = cls._encode_plugin_model_providers_cache_payload(
|
||||
_provider_entities_adapter.dump_json(list(providers))
|
||||
)
|
||||
payload = _provider_entities_adapter.dump_json(list(providers)).decode("utf-8")
|
||||
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 stored in Redis."""
|
||||
"""Invalidate tenant-scoped provider metadata across Redis and worker-local mirrors."""
|
||||
cls._plugin_model_providers_memory_cache.pop(tenant_id, None)
|
||||
cache_key = cls._get_plugin_model_providers_cache_key(tenant_id)
|
||||
generation_key = cls._get_plugin_model_providers_generation_cache_key(tenant_id)
|
||||
try:
|
||||
@@ -361,68 +287,21 @@ class PluginService:
|
||||
are intentionally owned by this service so tenant isolation and cache
|
||||
expiry are handled in one place.
|
||||
"""
|
||||
deadline = time.monotonic() + cls.PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT
|
||||
cached_providers = cls._load_cached_plugin_model_providers(tenant_id, client=client)
|
||||
if cached_providers is not None:
|
||||
return cached_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,
|
||||
)
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
def fetch_latest_plugin_version(plugin_ids: Sequence[str]) -> Mapping[str, LatestPluginCache | None]:
|
||||
@@ -461,7 +340,7 @@ class PluginService:
|
||||
plugin_id=plugin_id,
|
||||
version=manifest.latest_version,
|
||||
unique_identifier=manifest.latest_package_identifier,
|
||||
status=cast(Literal["active", "deleted"], manifest.status),
|
||||
status=manifest.status,
|
||||
deprecated_reason=manifest.deprecated_reason,
|
||||
alternative_plugin_id=manifest.alternative_plugin_id,
|
||||
)
|
||||
|
||||
@@ -1030,10 +1030,6 @@ 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
|
||||
@@ -1045,9 +1041,6 @@ 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 = []
|
||||
@@ -1062,16 +1055,13 @@ class DatasetRetrieval:
|
||||
content=json.dumps(contents),
|
||||
source=DatasetQuerySource.APP,
|
||||
source_app_id=app_id,
|
||||
created_by_role=created_by_role,
|
||||
created_by_role=CreatorUserRole(user_from),
|
||||
created_by=created_by,
|
||||
)
|
||||
dataset_queries.append(dataset_query)
|
||||
|
||||
if not dataset_queries:
|
||||
return
|
||||
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
|
||||
session.add_all(dataset_queries)
|
||||
if dataset_queries:
|
||||
db.session.add_all(dataset_queries)
|
||||
db.session.commit()
|
||||
|
||||
def _retriever(
|
||||
self,
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
from core.rbac.entities import RBACPermission, RBACResourceScope, RBACResourceWhitelistScope
|
||||
from core.rbac.entities import RBACPermission, RBACResourceScope
|
||||
|
||||
__all__ = ["RBACPermission", "RBACResourceScope", "RBACResourceWhitelistScope"]
|
||||
__all__ = ["RBACPermission", "RBACResourceScope"]
|
||||
|
||||
@@ -13,14 +13,6 @@ 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,17 +102,16 @@ 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 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}"
|
||||
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']}"
|
||||
elif api_key_header_prefix == "custom":
|
||||
pass
|
||||
|
||||
headers[api_key_header] = api_key_value
|
||||
headers[api_key_header] = credentials["api_key_value"]
|
||||
|
||||
elif credentials["auth_type"] == "api_key_query":
|
||||
# For query parameter authentication, we don't add anything to headers
|
||||
|
||||
@@ -7,7 +7,6 @@ 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
|
||||
@@ -339,49 +338,47 @@ class ToolEngine:
|
||||
user_id: str,
|
||||
) -> list[str]:
|
||||
"""
|
||||
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.
|
||||
Create message file
|
||||
|
||||
:return: message file ids
|
||||
"""
|
||||
result = []
|
||||
|
||||
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,
|
||||
)
|
||||
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
|
||||
|
||||
session.add(message_file)
|
||||
result.append(message_file.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()
|
||||
|
||||
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
|
||||
|
||||
@@ -5,10 +5,10 @@ from json import loads as json_loads
|
||||
from json.decoder import JSONDecodeError
|
||||
from typing import Any, TypedDict
|
||||
|
||||
import httpx
|
||||
from flask import has_request_context, request
|
||||
from yaml import YAMLError, safe_load
|
||||
|
||||
from core.helper import ssrf_proxy
|
||||
from core.tools.entities.common_entities import I18nObject
|
||||
from core.tools.entities.tool_bundle import ApiToolBundle
|
||||
from core.tools.entities.tool_entities import ApiProviderSchemaType, ToolParameter
|
||||
@@ -376,7 +376,7 @@ class ApiBasedToolSchemaParser:
|
||||
raise ToolNotSupportedError("Only openapi is supported now.")
|
||||
|
||||
# get openapi yaml
|
||||
response = ssrf_proxy.get(
|
||||
response = httpx.get(
|
||||
api_url, headers={"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "}, timeout=5
|
||||
)
|
||||
|
||||
|
||||
@@ -46,26 +46,18 @@ logger = logging.getLogger(__name__)
|
||||
_file_access_controller = DatabaseFileAccessController()
|
||||
|
||||
|
||||
def iter_dify_graph_engine_events(
|
||||
engine: GraphEngine,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
) -> Generator[GraphEngineEvent, None, None]:
|
||||
def iter_dify_graph_engine_events(engine: GraphEngine) -> Generator[GraphEngineEvent, None, None]:
|
||||
"""
|
||||
Apply Dify's response streaming compatibility filter to GraphEngine events.
|
||||
|
||||
Graphon v0.5.0 emits raw variable stream chunks and requires callers to opt
|
||||
into the legacy response-ordered stream behavior that Dify exposes to its
|
||||
workflow runners and tests.
|
||||
|
||||
``response_stream_filter``, when supplied, must be the same instance a
|
||||
caller intends to persist on pause (see ``PauseStatePersistenceLayer``) so
|
||||
the filter's ``paths_map`` reflects everything the engine has actually
|
||||
streamed for this run.
|
||||
"""
|
||||
yield from filter_graph_events(
|
||||
engine.run(),
|
||||
context=GraphEventFilterContext.from_engine(engine),
|
||||
filters=[response_stream_filter or ResponseStreamFilter()],
|
||||
filters=[ResponseStreamFilter()],
|
||||
)
|
||||
|
||||
|
||||
@@ -175,7 +167,6 @@ class WorkflowEntry:
|
||||
variable_pool: VariablePool,
|
||||
graph_runtime_state: GraphRuntimeState,
|
||||
command_channel: CommandChannel | None = None,
|
||||
response_stream_filter: ResponseStreamFilter | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Init workflow entry
|
||||
@@ -192,8 +183,6 @@ class WorkflowEntry:
|
||||
:param variable_pool: variable pool
|
||||
:param graph_runtime_state: pre-created graph runtime state
|
||||
:param command_channel: command channel for external control (optional, defaults to InMemoryChannel)
|
||||
:param response_stream_filter: pre-restored filter for resumed runs (optional, defaults to a fresh
|
||||
ResponseStreamFilter for runs with no prior pause)
|
||||
:param thread_pool_id: thread pool id
|
||||
"""
|
||||
# check call depth
|
||||
@@ -206,7 +195,6 @@ class WorkflowEntry:
|
||||
command_channel = InMemoryChannel()
|
||||
|
||||
self.command_channel = command_channel
|
||||
self._response_stream_filter = response_stream_filter or ResponseStreamFilter()
|
||||
execution_context = capture_current_context()
|
||||
graph_runtime_state.execution_context = execution_context
|
||||
self._child_engine_builder = _WorkflowChildEngineBuilder(tenant_id=tenant_id)
|
||||
@@ -252,7 +240,7 @@ class WorkflowEntry:
|
||||
|
||||
try:
|
||||
# Preserve Dify's response-stream semantics on top of Graphon 0.5.0.
|
||||
generator = iter_dify_graph_engine_events(graph_engine, self._response_stream_filter)
|
||||
generator = iter_dify_graph_engine_events(graph_engine)
|
||||
yield from generator
|
||||
except GenerateTaskStoppedError:
|
||||
pass
|
||||
|
||||
@@ -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, ...)``
|
||||
or ``Model.model_validate(...).model_dump()`` return.
|
||||
a documented ``@ns.response(..., Model)`` and an actual ``dump_response(Model, ...)``,
|
||||
``Model(...).model_dump()``, 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,6 +28,7 @@ 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[
|
||||
@@ -41,6 +42,7 @@ 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})
|
||||
@@ -109,18 +111,22 @@ 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) -> None:
|
||||
def add_known(self, model: str, source: ModelValueSource) -> 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) -> str | None:
|
||||
def single_known_model(self) -> tuple[str, ModelValueSource] | None:
|
||||
if self.has_unknown_assignment or len(self.known_models) != 1:
|
||||
return None
|
||||
return next(iter(self.known_models))
|
||||
model = next(iter(self.known_models))
|
||||
source: ModelValueSource = "constructor" if self.known_sources == {"constructor"} else "model_validate"
|
||||
return model, source
|
||||
|
||||
|
||||
def dotted_name(node: ast.AST) -> str | None:
|
||||
@@ -249,6 +255,12 @@ 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
|
||||
@@ -257,6 +269,12 @@ 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
|
||||
@@ -272,6 +290,10 @@ 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
|
||||
@@ -287,7 +309,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, str] | None = None
|
||||
expr: ast.AST | None, variable_models: dict[str, tuple[str, ModelValueSource]] | None = None
|
||||
) -> tuple[ActualKind, str | None]:
|
||||
if expr is None:
|
||||
return "none", None
|
||||
@@ -299,10 +321,14 @@ 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:
|
||||
# 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:
|
||||
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.
|
||||
return "model_dump_variable", model_name
|
||||
|
||||
model_dump_model = model_name_from_model_dump(expr)
|
||||
@@ -325,7 +351,9 @@ def actual_kind_from_expr(
|
||||
return "unknown", None
|
||||
|
||||
|
||||
def actual_response_from_return(return_node: ast.Return, variable_models: dict[str, str]) -> ActualResponse:
|
||||
def actual_response_from_return(
|
||||
return_node: ast.Return, variable_models: dict[str, tuple[str, ModelValueSource]]
|
||||
) -> ActualResponse:
|
||||
status: int | None = 200
|
||||
body_expr = return_node.value
|
||||
|
||||
@@ -363,18 +391,21 @@ def target_names(target: ast.AST) -> Iterable[str]:
|
||||
|
||||
|
||||
def record_assignment(
|
||||
assignments: defaultdict[str, VariableAssignmentSummary], targets: Iterable[str], model_name: str | None
|
||||
assignments: defaultdict[str, VariableAssignmentSummary],
|
||||
targets: Iterable[str],
|
||||
model_assignment: tuple[str, ModelValueSource] | None,
|
||||
) -> None:
|
||||
for target in targets:
|
||||
if model_name is None:
|
||||
if model_assignment 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:
|
||||
assignments[target].add_known(model_name)
|
||||
model_name, source = model_assignment
|
||||
assignments[target].add_known(model_name, source)
|
||||
|
||||
|
||||
def variable_model_assignments_for_method(method: MethodNode) -> dict[str, str]:
|
||||
def variable_model_assignments_for_method(method: MethodNode) -> dict[str, tuple[str, ModelValueSource]]:
|
||||
"""Infer local variables that are unambiguously assigned one response model."""
|
||||
|
||||
assignments: defaultdict[str, VariableAssignmentSummary] = defaultdict(VariableAssignmentSummary)
|
||||
@@ -385,10 +416,10 @@ def variable_model_assignments_for_method(method: MethodNode) -> dict[str, str]:
|
||||
record_assignment(
|
||||
assignments,
|
||||
(name for target in targets for name in target_names(target)),
|
||||
model_name_from_model_value(value),
|
||||
model_value_from_model_value(value),
|
||||
)
|
||||
case ast.AnnAssign(target=target, value=value) if value is not None:
|
||||
record_assignment(assignments, target_names(target), model_name_from_model_value(value))
|
||||
record_assignment(assignments, target_names(target), model_value_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)
|
||||
@@ -399,9 +430,13 @@ def variable_model_assignments_for_method(method: MethodNode) -> dict[str, str]:
|
||||
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_name_from_model_value(value))
|
||||
record_assignment(assignments, target_names(target), model_value_from_model_value(value))
|
||||
|
||||
return {name: model for name, summary in assignments.items() if (model := summary.single_known_model()) is not None}
|
||||
return {
|
||||
name: assignment
|
||||
for name, summary in assignments.items()
|
||||
if (assignment := summary.single_known_model()) is not None
|
||||
}
|
||||
|
||||
|
||||
def actual_responses_for_method(method: MethodNode) -> list[ActualResponse]:
|
||||
@@ -545,13 +580,52 @@ 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]:
|
||||
module = ast.parse(file_path.read_text(encoding="utf-8"), filename=str(file_path))
|
||||
source = file_path.read_text(encoding="utf-8")
|
||||
lines = source.splitlines()
|
||||
module = ast.parse(source, 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)
|
||||
@@ -559,6 +633,8 @@ 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:
|
||||
|
||||
@@ -35,10 +35,10 @@ if [[ "${MODE}" == "worker" ]]; then
|
||||
if [[ -z "${CELERY_QUEUES}" ]]; then
|
||||
if [[ "${EDITION}" == "CLOUD" ]]; then
|
||||
# Cloud edition: separate queues for dataset and trigger tasks
|
||||
DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution"
|
||||
DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution"
|
||||
else
|
||||
# Community edition (SELF_HOSTED): dataset, pipeline and workflow have separate queues
|
||||
DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution"
|
||||
DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution"
|
||||
fi
|
||||
else
|
||||
DEFAULT_QUEUES="${CELERY_QUEUES}"
|
||||
|
||||
@@ -153,7 +153,6 @@ def init_app(app: DifyApp) -> Celery:
|
||||
"tasks.trigger_processing_tasks", # async trigger processing
|
||||
"tasks.generate_summary_index_task", # summary index generation
|
||||
"tasks.regenerate_summary_index_task", # summary index regeneration
|
||||
"tasks.initialize_created_app_rbac_access_task", # app access initialization
|
||||
"tasks.app_generate.resume_agent_app_task", # ENG-635: Agent v2 chat ask_human resume
|
||||
]
|
||||
day = dify_config.CELERY_BEAT_SCHEDULER_TIME
|
||||
|
||||
@@ -27,7 +27,6 @@ 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,
|
||||
@@ -57,7 +56,6 @@ 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,
|
||||
|
||||
+2
-1
@@ -7,7 +7,8 @@ class ResponseModel(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
extra="ignore",
|
||||
populate_by_name=True,
|
||||
validate_by_name=True,
|
||||
validate_by_alias=True,
|
||||
serialize_by_alias=True,
|
||||
protected_namespaces=(),
|
||||
)
|
||||
|
||||
+2
-19
@@ -34,7 +34,6 @@ class OAuthState(TypedDict, total=False):
|
||||
invite_token: str
|
||||
timezone: str
|
||||
language: str
|
||||
redirect_url: str
|
||||
|
||||
|
||||
class GitHubEmailRecord(TypedDict, total=False):
|
||||
@@ -73,7 +72,6 @@ def encode_oauth_state(
|
||||
invite_token: str | None = None,
|
||||
timezone: str | None = None,
|
||||
language: str | None = None,
|
||||
redirect_url: str | None = None,
|
||||
) -> str | None:
|
||||
state: OAuthState = {}
|
||||
if invite_token:
|
||||
@@ -82,8 +80,6 @@ def encode_oauth_state(
|
||||
state["timezone"] = timezone
|
||||
if language:
|
||||
state["language"] = language
|
||||
if redirect_url:
|
||||
state["redirect_url"] = redirect_url
|
||||
if not state:
|
||||
return None
|
||||
|
||||
@@ -126,7 +122,6 @@ class OAuth:
|
||||
invite_token: str | None = None,
|
||||
timezone: str | None = None,
|
||||
language: str | None = None,
|
||||
redirect_url: str | None = None,
|
||||
) -> str:
|
||||
raise NotImplementedError()
|
||||
|
||||
@@ -156,19 +151,13 @@ class GitHubOAuth(OAuth):
|
||||
invite_token: str | None = None,
|
||||
timezone: str | None = None,
|
||||
language: str | None = None,
|
||||
redirect_url: str | None = None,
|
||||
) -> str:
|
||||
params = {
|
||||
"client_id": self.client_id,
|
||||
"redirect_uri": self.redirect_uri,
|
||||
"scope": "user:email", # Request only basic user information
|
||||
}
|
||||
state = encode_oauth_state(
|
||||
invite_token=invite_token,
|
||||
timezone=timezone,
|
||||
language=language,
|
||||
redirect_url=redirect_url,
|
||||
)
|
||||
state = encode_oauth_state(invite_token=invite_token, timezone=timezone, language=language)
|
||||
if state:
|
||||
params["state"] = state
|
||||
return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
|
||||
@@ -259,7 +248,6 @@ class GoogleOAuth(OAuth):
|
||||
invite_token: str | None = None,
|
||||
timezone: str | None = None,
|
||||
language: str | None = None,
|
||||
redirect_url: str | None = None,
|
||||
) -> str:
|
||||
params = {
|
||||
"client_id": self.client_id,
|
||||
@@ -267,12 +255,7 @@ class GoogleOAuth(OAuth):
|
||||
"redirect_uri": self.redirect_uri,
|
||||
"scope": "openid email",
|
||||
}
|
||||
state = encode_oauth_state(
|
||||
invite_token=invite_token,
|
||||
timezone=timezone,
|
||||
language=language,
|
||||
redirect_url=redirect_url,
|
||||
)
|
||||
state = encode_oauth_state(invite_token=invite_token, timezone=timezone, language=language)
|
||||
if state:
|
||||
params["state"] = state
|
||||
return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
|
||||
|
||||
+26
-10
@@ -236,6 +236,16 @@ 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
|
||||
# ============================================================================
|
||||
@@ -343,13 +353,15 @@ 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 | Literal["invalid"]:
|
||||
def cache_get(self, token_hash: str) -> ResolvedRow | None | NegativeCache:
|
||||
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
|
||||
if text == "invalid":
|
||||
return "invalid"
|
||||
try:
|
||||
return NegativeCache(text)
|
||||
except ValueError:
|
||||
pass
|
||||
try:
|
||||
return ResolvedRow.from_cache(json.loads(text))
|
||||
except (ValueError, KeyError):
|
||||
@@ -363,8 +375,8 @@ class OAuthAccessTokenResolver:
|
||||
json.dumps(row.to_cache()),
|
||||
)
|
||||
|
||||
def cache_set_negative(self, token_hash: str) -> None:
|
||||
self._redis.setex(self._cache_key(token_hash), self._negative_ttl, "invalid")
|
||||
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 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;
|
||||
@@ -385,7 +397,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)
|
||||
self.cache_set_negative(token_hash, NegativeCache.EXPIRED)
|
||||
|
||||
|
||||
class _VariantResolver:
|
||||
@@ -395,9 +407,11 @@ class _VariantResolver:
|
||||
|
||||
def resolve(self, token_hash: str) -> ResolvedRow | None:
|
||||
cached = self._parent.cache_get(token_hash)
|
||||
if cached == "invalid":
|
||||
if isinstance(cached, NegativeCache):
|
||||
if cached is NegativeCache.EXPIRED:
|
||||
raise TokenExpiredError("token_expired")
|
||||
return None
|
||||
if cached is not None and not isinstance(cached, str):
|
||||
if cached is not None:
|
||||
if not self._matches_variant(cached):
|
||||
return None
|
||||
return cached
|
||||
@@ -413,7 +427,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)
|
||||
return None
|
||||
raise TokenExpiredError("token_expired")
|
||||
|
||||
if not self._matches_variant_model(row):
|
||||
logger.error(
|
||||
@@ -472,7 +486,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 == "invalid":
|
||||
if text in (NegativeCache.INVALID, NegativeCache.EXPIRED):
|
||||
return
|
||||
try:
|
||||
data = json.loads(text)
|
||||
@@ -601,6 +615,8 @@ 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,9 +2,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
_DIAGNOSTIC_PREFIXES = ("ERROR ", "WARNING ")
|
||||
_DIAGNOSTIC_PREFIXES = ("ERROR ", "WARN ", "WARNING ")
|
||||
_LOCATION_PREFIX = "-->"
|
||||
|
||||
|
||||
@@ -13,7 +14,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 ...`` / ``WARNING ...``)
|
||||
- diagnostic headline lines (``ERROR ...`` / ``WARN ...`` / ``WARNING ...``)
|
||||
- the following location line (``--> path:line:column``), when present
|
||||
"""
|
||||
|
||||
@@ -36,11 +37,28 @@ 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(extract_diagnostics(raw_output))
|
||||
sys.stdout.write(render_diagnostics(raw_output, exit_code=args.status))
|
||||
return 0
|
||||
|
||||
|
||||
|
||||
@@ -214,18 +214,6 @@ 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"""
|
||||
|
||||
@@ -7003,7 +7003,6 @@ Initiate OAuth login process
|
||||
| provider | path | OAuth provider name (github/google) | Yes | string |
|
||||
| invite_token | query | Optional invitation token | No | string |
|
||||
| language | query | Preferred interface language | No | string |
|
||||
| redirect_url | query | Relative page to resume after login | No | string |
|
||||
| timezone | query | Preferred timezone | No | string |
|
||||
|
||||
#### Responses
|
||||
@@ -13433,7 +13432,6 @@ 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 |
|
||||
@@ -14544,8 +14542,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
|
||||
|
||||
@@ -16714,6 +16712,7 @@ Input field definition for snippet parameters.
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| description | string | | No |
|
||||
| icon | string | | No |
|
||||
| icon_background | string | | No |
|
||||
| icon_type | string | | No |
|
||||
@@ -16863,7 +16862,6 @@ 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 |
|
||||
|
||||
@@ -17084,6 +17082,7 @@ 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 |
|
||||
@@ -17097,12 +17096,11 @@ 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 |
|
||||
|
||||
@@ -17534,7 +17532,6 @@ Coarse node-level status used by Inspector to pick a banner.
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| invite_token | string | Optional invitation token | No |
|
||||
| language | string | Preferred interface language | No |
|
||||
| redirect_url | string | Relative page to resume after login | No |
|
||||
| timezone | string | Preferred timezone | No |
|
||||
|
||||
#### OAuthProviderAccountResponse
|
||||
@@ -18801,13 +18798,6 @@ Model class for provider quota configuration.
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| url | string | URL to fetch | Yes |
|
||||
|
||||
#### ReplaceUserAccessPolicies
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| access_policy_ids | [ string ] | | No |
|
||||
| account_ids | [ string ] | | No |
|
||||
|
||||
#### ReplaceUserAccessPoliciesResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
|
||||
@@ -990,6 +990,12 @@ 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,7 +1337,6 @@ 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,6 +1,5 @@
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from enum import StrEnum
|
||||
from typing import Any, override
|
||||
@@ -18,8 +17,6 @@ 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
|
||||
@@ -105,10 +102,7 @@ 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:String}}",
|
||||
parameters={"id": id},
|
||||
)
|
||||
results = self._client.query(f"SELECT id FROM {self._config.database}.{self._collection_name} WHERE id='{id}'")
|
||||
return results.row_count > 0
|
||||
|
||||
@override
|
||||
@@ -116,26 +110,20 @@ class MyScaleVector(BaseVector):
|
||||
if not ids:
|
||||
return
|
||||
self._client.command(
|
||||
f"DELETE FROM {self._config.database}.{self._collection_name} WHERE id IN {{ids:Array(String)}}",
|
||||
parameters={"ids": ids},
|
||||
f"DELETE FROM {self._config.database}.{self._collection_name} WHERE id IN {str(tuple(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} "
|
||||
f"WHERE metadata.{key}={{value:String}}",
|
||||
parameters={"value": value},
|
||||
f"SELECT DISTINCT id FROM {self._config.database}.{self._collection_name} WHERE metadata.{key}='{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:String}}",
|
||||
parameters={"value": value},
|
||||
f"DELETE FROM {self._config.database}.{self._collection_name} WHERE metadata.{key}='{value}'"
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -144,29 +132,22 @@ class MyScaleVector(BaseVector):
|
||||
|
||||
@override
|
||||
def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
|
||||
return self._search(
|
||||
"TextSearch('enable_nlq=false')(text, {query:String})",
|
||||
SortOrder.DESC,
|
||||
parameters={"query": query},
|
||||
**kwargs,
|
||||
)
|
||||
return self._search(f"TextSearch('enable_nlq=false')(text, '{query}')", SortOrder.DESC, **kwargs)
|
||||
|
||||
def _search(
|
||||
self, dist: str, order: SortOrder, parameters: dict[str, Any] | None = None, **kwargs: Any
|
||||
) -> list[Document]:
|
||||
def _search(self, dist: str, order: SortOrder, **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_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}")
|
||||
where_str = (
|
||||
f"WHERE dist < {1 - score_threshold}"
|
||||
if self._metric.upper() == "COSINE" and order == SortOrder.ASC and score_threshold > 0.0
|
||||
else ""
|
||||
)
|
||||
document_ids_filter = kwargs.get("document_ids_filter")
|
||||
if document_ids_filter:
|
||||
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 ""
|
||||
document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
|
||||
where_str = f"{where_str} AND metadata['document_id'] in ({document_ids})"
|
||||
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}
|
||||
@@ -178,17 +159,12 @@ class MyScaleVector(BaseVector):
|
||||
vector=r["vector"],
|
||||
metadata=r["metadata"],
|
||||
)
|
||||
for r in self._client.query(sql, parameters=query_parameters).named_results()
|
||||
for r in self._client.query(sql).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,41 +181,14 @@ 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' 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"}
|
||||
assert vector.text_exists("id-1") is True
|
||||
assert vector.get_ids_by_metadata_field("document_id", "doc-1") == ["id-1", "id-2"]
|
||||
|
||||
vector.delete_by_ids(["id-1", "id-2"])
|
||||
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"}
|
||||
vector.delete_by_metadata_field("document_id", "doc-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"])
|
||||
@@ -226,28 +199,6 @@ 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):
|
||||
@@ -264,8 +215,7 @@ def test_search_with_document_filter_and_exception(myscale_module):
|
||||
)
|
||||
assert len(docs) == 1
|
||||
sql = vector._client.query.call_args.args[0]
|
||||
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"]}
|
||||
assert "metadata['document_id'] in ('doc-1', 'doc-2')" in sql
|
||||
|
||||
vector._client.query.side_effect = RuntimeError("boom")
|
||||
assert vector._search("distance(vector, [0.1])", myscale_module.SortOrder.ASC, top_k=1) == []
|
||||
|
||||
@@ -33,6 +33,7 @@ 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,6 +41,7 @@ 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
|
||||
|
||||
+1
-3
@@ -42,12 +42,11 @@ 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",
|
||||
"httpx-sse==0.4.3",
|
||||
"json-repair==0.60.1",
|
||||
"json-repair==0.59.4",
|
||||
]
|
||||
# Before adding new dependency, consider place it in
|
||||
# alphabet order (a-z) and suitable group.
|
||||
@@ -108,7 +107,6 @@ override-dependencies = [
|
||||
"litellm>=1.83.10,<2.0.0",
|
||||
"pyarrow>=23.0.1,<24.0.0",
|
||||
"cryptography>=49.0.0,<50.0.0",
|
||||
"setuptools>=80.10.2,<81",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
|
||||
@@ -10,7 +10,6 @@ import json
|
||||
import logging
|
||||
import secrets
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from hashlib import sha256
|
||||
from typing import Any, NotRequired, TypedDict, cast
|
||||
@@ -65,7 +64,6 @@ from services.errors.account import (
|
||||
MemberNotInTenantError,
|
||||
NoPermissionError,
|
||||
RoleAlreadyAssignedError,
|
||||
SeatsLimitExceededError,
|
||||
TenantNotFoundError,
|
||||
)
|
||||
from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
|
||||
@@ -435,13 +433,6 @@ 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=(
|
||||
@@ -1292,7 +1283,7 @@ class TenantService:
|
||||
def create_owner_tenant_if_not_exist(
|
||||
account: Account, name: str | None = None, is_setup: bool | None = False, *, session: scoped_session | Session
|
||||
):
|
||||
"""Create an owner workspace only when the account has no membership."""
|
||||
"""Check if user have a workspace or not"""
|
||||
available_ta = session.scalar(
|
||||
select(TenantAccountJoin)
|
||||
.where(TenantAccountJoin.account_id == account.id)
|
||||
@@ -1303,28 +1294,8 @@ class TenantService:
|
||||
if available_ta:
|
||||
return
|
||||
|
||||
TenantService.create_owner_tenant(account, name=name, is_setup=is_setup, session=session)
|
||||
|
||||
@staticmethod
|
||||
def create_owner_tenant(
|
||||
account: Account,
|
||||
name: str | None = None,
|
||||
is_setup: bool | None = False,
|
||||
is_from_dashboard: bool | None = False,
|
||||
*,
|
||||
session: Session,
|
||||
) -> Tenant:
|
||||
"""Create an owner workspace and bind its owner RBAC role when enabled.
|
||||
|
||||
This is the single write path for a newly created workspace with an
|
||||
owner. It persists the legacy membership before creating the matching
|
||||
RBAC role binding, then makes the workspace current for the account.
|
||||
"""
|
||||
if (
|
||||
not FeatureService.get_system_features().is_allow_create_workspace
|
||||
and not is_setup
|
||||
and not is_from_dashboard
|
||||
):
|
||||
"""Create owner tenant if not exist"""
|
||||
if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup:
|
||||
raise WorkSpaceNotAllowedCreateError()
|
||||
|
||||
workspaces = FeatureService.get_system_features().license.workspaces
|
||||
@@ -1332,19 +1303,9 @@ class TenantService:
|
||||
raise WorkspacesLimitExceededError()
|
||||
|
||||
if name:
|
||||
tenant = TenantService.create_tenant(
|
||||
name=name,
|
||||
is_setup=is_setup,
|
||||
is_from_dashboard=is_from_dashboard,
|
||||
session=session,
|
||||
)
|
||||
tenant = TenantService.create_tenant(name=name, is_setup=is_setup, session=session)
|
||||
else:
|
||||
tenant = TenantService.create_tenant(
|
||||
name=f"{account.name}'s Workspace",
|
||||
is_setup=is_setup,
|
||||
is_from_dashboard=is_from_dashboard,
|
||||
session=session,
|
||||
)
|
||||
tenant = TenantService.create_tenant(name=f"{account.name}'s Workspace", is_setup=is_setup, session=session)
|
||||
TenantService.create_tenant_member(tenant, account, session, role="owner")
|
||||
if dify_config.RBAC_ENABLED:
|
||||
owner_role_id = AccountService._resolve_legacy_role_id(str(tenant.id), account.id, TenantAccountRole.OWNER)
|
||||
@@ -1357,7 +1318,6 @@ class TenantService:
|
||||
account.current_tenant = tenant
|
||||
session.commit()
|
||||
tenant_was_created.send(tenant)
|
||||
return tenant
|
||||
|
||||
@staticmethod
|
||||
def create_tenant_member(
|
||||
@@ -1616,25 +1576,6 @@ class TenantService:
|
||||
|
||||
return updated_accounts
|
||||
|
||||
@staticmethod
|
||||
def iter_member_account_id_batches(tenant_id: str, batch_size: int, *, session: Session) -> Iterator[list[str]]:
|
||||
"""Yield workspace member account ids in bounded, ordered batches."""
|
||||
offset = 0
|
||||
while True:
|
||||
stmt = (
|
||||
select(TenantAccountJoin.account_id)
|
||||
.where(TenantAccountJoin.tenant_id == tenant_id)
|
||||
.order_by(TenantAccountJoin.id)
|
||||
.offset(offset)
|
||||
.limit(batch_size)
|
||||
)
|
||||
account_ids = list(session.scalars(stmt).all())
|
||||
if not account_ids:
|
||||
return
|
||||
|
||||
yield account_ids
|
||||
offset += batch_size
|
||||
|
||||
@staticmethod
|
||||
def get_dataset_operator_members(tenant: Tenant, *, session: scoped_session | Session) -> list[Account]:
|
||||
"""Get dataset admin members"""
|
||||
@@ -2025,7 +1966,10 @@ class RegisterService:
|
||||
and FeatureService.get_system_features().license.workspaces.is_available()
|
||||
):
|
||||
try:
|
||||
TenantService.create_owner_tenant(account, session=session)
|
||||
tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=session)
|
||||
TenantService.create_tenant_member(tenant, account, session, role="owner")
|
||||
account.current_tenant = tenant
|
||||
tenant_was_created.send(tenant)
|
||||
except Exception:
|
||||
_try_join_enterprise_default_workspace(str(account.id))
|
||||
raise
|
||||
@@ -2037,10 +1981,6 @@ 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,26 +50,9 @@ 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):
|
||||
@@ -193,7 +176,6 @@ 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)
|
||||
@@ -233,18 +215,12 @@ 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_quota_request("GET", "/quota/info", params=params)
|
||||
cls._send_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,
|
||||
bucket: str = "",
|
||||
cls, tenant_id: str, feature_key: str, request_id: str, amount: int = 1, meta: dict | None = None
|
||||
) -> QuotaReserveResult:
|
||||
"""Reserve quota before task execution."""
|
||||
payload: dict = {
|
||||
@@ -253,21 +229,13 @@ 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_quota_request("POST", "/quota/reserve", json=payload))
|
||||
return _quota_reserve_adapter.validate_python(cls._send_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,
|
||||
bucket: str = "",
|
||||
cls, tenant_id: str, feature_key: str, reservation_id: str, actual_amount: int, meta: dict | None = None
|
||||
) -> QuotaCommitResult:
|
||||
"""Commit a reservation with actual consumption."""
|
||||
payload: dict = {
|
||||
@@ -276,57 +244,23 @@ 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_quota_request("POST", "/quota/commit", json=payload))
|
||||
return _quota_commit_adapter.validate_python(cls._send_request("POST", "/quota/commit", json=payload))
|
||||
|
||||
@classmethod
|
||||
def quota_release(
|
||||
cls, tenant_id: str, feature_key: str, reservation_id: str, bucket: str = ""
|
||||
) -> QuotaReleaseResult:
|
||||
def quota_release(cls, tenant_id: str, feature_key: str, reservation_id: str) -> QuotaReleaseResult:
|
||||
"""Release a reservation (cancel, return frozen quota)."""
|
||||
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)
|
||||
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,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -400,12 +334,6 @@ 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),
|
||||
@@ -413,17 +341,10 @@ 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,
|
||||
base_url: str | None = None,
|
||||
):
|
||||
def _send_request(cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None):
|
||||
headers = {"Content-Type": "application/json", "Billing-Api-Secret-Key": cls.secret_key}
|
||||
|
||||
url = f"{base_url or cls.base_url}{endpoint}"
|
||||
url = f"{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,8 +7,6 @@ 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
|
||||
@@ -23,37 +21,11 @@ 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"
|
||||
@@ -107,32 +79,14 @@ class CreditPoolService:
|
||||
return credit_pool
|
||||
|
||||
@classmethod
|
||||
def get_pool(
|
||||
cls, tenant_id: str, pool_type: str | ProviderQuotaType = "trial"
|
||||
) -> TenantCreditPool | CreditPoolBalance | None:
|
||||
def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> TenantCreditPool | 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 == normalized_pool_type,
|
||||
TenantCreditPool.pool_type == pool_type,
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
@@ -148,7 +102,7 @@ class CreditPoolService:
|
||||
pool = cls.get_pool(tenant_id, pool_type)
|
||||
if not pool:
|
||||
return False
|
||||
return pool.has_sufficient_credits(credits_required)
|
||||
return pool.remaining_credits >= credits_required
|
||||
|
||||
@classmethod
|
||||
def check_and_deduct_credits(
|
||||
@@ -160,54 +114,10 @@ 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=normalized_pool_type)
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
|
||||
if not pool:
|
||||
raise QuotaExceededError("Credit pool not found")
|
||||
|
||||
@@ -238,26 +148,12 @@ 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=normalized_pool_type)
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
|
||||
if not pool:
|
||||
logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, normalized_pool_type)
|
||||
logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, pool_type)
|
||||
return 0
|
||||
|
||||
deducted_credits = min(credits_required, pool.remaining_credits)
|
||||
|
||||
@@ -12,7 +12,6 @@ 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
|
||||
|
||||
@@ -255,7 +254,6 @@ class ResourceUserAccessPoliciesResponse(_RBACModel):
|
||||
|
||||
class ReplaceUserAccessPolicies(_RBACModel):
|
||||
access_policy_ids: list[str] = Field(default_factory=list)
|
||||
account_ids: list[str] = Field(default_factory=list)
|
||||
|
||||
@field_validator("access_policy_ids", mode="before")
|
||||
@classmethod
|
||||
@@ -366,6 +364,7 @@ _LEGACY_WORKSPACE_ADMIN_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_WORKSPACE_EDITOR_KEYS: list[str] = [
|
||||
"workspace.member.manage",
|
||||
"api_extension.manage",
|
||||
"plugin.install",
|
||||
"credential.use",
|
||||
@@ -436,7 +435,6 @@ _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",
|
||||
]
|
||||
|
||||
@@ -651,18 +649,17 @@ class ReplaceRoleBindings(_RBACModel):
|
||||
|
||||
|
||||
class ReplaceMemberBindings(_RBACModel):
|
||||
scope: RBACResourceWhitelistScope = RBACResourceWhitelistScope.SPECIFIC
|
||||
scope: str = "specific"
|
||||
|
||||
@field_validator("scope")
|
||||
@classmethod
|
||||
def _normalize_scope(cls, value: Any) -> RBACResourceWhitelistScope:
|
||||
def _normalize_scope(cls, value: Any) -> str:
|
||||
scope = str(value or "").strip().lower()
|
||||
if scope == "":
|
||||
return RBACResourceWhitelistScope.SPECIFIC
|
||||
try:
|
||||
return RBACResourceWhitelistScope(scope)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"invalid scope: {value}") from exc
|
||||
if scope in {"", "specific"}:
|
||||
return "specific"
|
||||
if scope in {"all", "only_me"}:
|
||||
return scope
|
||||
raise ValueError(f"invalid scope: {value}")
|
||||
|
||||
|
||||
class DeleteMemberBindings(_RBACModel):
|
||||
@@ -746,7 +743,6 @@ def _inner_call(
|
||||
account_id=account_id,
|
||||
json=json,
|
||||
params=params,
|
||||
timeout=dify_config.ENTERPRISE_RBAC_REQUEST_TIMEOUT,
|
||||
)
|
||||
|
||||
|
||||
@@ -1126,17 +1122,16 @@ class RBACService:
|
||||
tenant_id: str,
|
||||
account_id: str | None,
|
||||
app_id: str,
|
||||
target_account_id: str | None,
|
||||
target_account_id: str,
|
||||
payload: ReplaceUserAccessPolicies,
|
||||
) -> ReplaceUserAccessPoliciesResponse:
|
||||
request_data = payload.model_dump(mode="json")
|
||||
data = _inner_call(
|
||||
"PUT",
|
||||
f"{_INNER_PREFIX}/apps/user-access-policies",
|
||||
tenant_id=tenant_id,
|
||||
account_id=account_id,
|
||||
params={"app_id": app_id, "account_id": target_account_id},
|
||||
json=request_data,
|
||||
json=payload.model_dump(mode="json"),
|
||||
)
|
||||
return ReplaceUserAccessPoliciesResponse.model_validate(data or {})
|
||||
|
||||
@@ -1296,7 +1291,7 @@ class RBACService:
|
||||
tenant_id: str,
|
||||
account_id: str | None,
|
||||
dataset_id: str,
|
||||
target_account_id: str | None,
|
||||
target_account_id: str,
|
||||
payload: ReplaceUserAccessPolicies,
|
||||
) -> ReplaceUserAccessPoliciesResponse:
|
||||
data = _inner_call(
|
||||
@@ -1305,7 +1300,7 @@ class RBACService:
|
||||
tenant_id=tenant_id,
|
||||
account_id=account_id,
|
||||
params={"dataset_id": dataset_id, "account_id": target_account_id},
|
||||
json=payload.model_dump(mode="json", exclude_unset=True),
|
||||
json=payload.model_dump(mode="json"),
|
||||
)
|
||||
return ReplaceUserAccessPoliciesResponse.model_validate(data or {})
|
||||
|
||||
|
||||
@@ -37,10 +37,6 @@ class AccountAlreadyInTenantError(BaseServiceError):
|
||||
pass
|
||||
|
||||
|
||||
class SeatsLimitExceededError(BaseServiceError):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidActionError(BaseServiceError):
|
||||
pass
|
||||
|
||||
|
||||
@@ -79,7 +79,6 @@ 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):
|
||||
@@ -458,11 +457,6 @@ 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[
|
||||
|
||||
@@ -4,7 +4,6 @@ from pathlib import Path
|
||||
from typing import Any, override
|
||||
|
||||
from flask import current_app
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval
|
||||
from services.recommend_app.recommend_app_base import RecommendAppRetrievalBase
|
||||
@@ -33,8 +32,7 @@ class BuildInRecommendAppRetrieval(RecommendAppRetrievalBase):
|
||||
return result
|
||||
|
||||
@override
|
||||
def get_recommend_app_detail(self, app_id: str, *, session: Session | None = None):
|
||||
del session
|
||||
def get_recommend_app_detail(self, app_id: str):
|
||||
result = self.fetch_recommended_app_detail_from_builtin(app_id)
|
||||
return result
|
||||
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import Any, NotRequired, TypedDict, override
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from constants.languages import languages
|
||||
from extensions.ext_database import db
|
||||
@@ -56,10 +55,8 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase):
|
||||
return result
|
||||
|
||||
@override
|
||||
def get_recommend_app_detail(
|
||||
self, app_id: str, *, session: Session | None = None
|
||||
) -> RecommendedAppDetailDict | None:
|
||||
result = self.fetch_recommended_app_detail_from_db(app_id, session=session)
|
||||
def get_recommend_app_detail(self, app_id: str) -> RecommendedAppDetailDict | None:
|
||||
result = self.fetch_recommended_app_detail_from_db(app_id)
|
||||
return result
|
||||
|
||||
@override
|
||||
@@ -149,17 +146,14 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fetch_recommended_app_detail_from_db(
|
||||
cls, app_id: str, *, session: Session | None = None
|
||||
) -> RecommendedAppDetailDict | None:
|
||||
def fetch_recommended_app_detail_from_db(cls, app_id: str) -> RecommendedAppDetailDict | None:
|
||||
"""
|
||||
Fetch recommended app detail from db.
|
||||
:param app_id: App ID
|
||||
:return:
|
||||
"""
|
||||
# is in public recommended list
|
||||
query_session = session if session is not None else db.session
|
||||
recommended_app = query_session.scalar(
|
||||
recommended_app = db.session.scalar(
|
||||
select(RecommendedApp).where(RecommendedApp.is_listed == True, RecommendedApp.app_id == app_id).limit(1)
|
||||
)
|
||||
|
||||
@@ -167,7 +161,7 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase):
|
||||
return None
|
||||
|
||||
# get app detail
|
||||
app_model = query_session.get(App, app_id)
|
||||
app_model = db.session.get(App, app_id)
|
||||
if not app_model or not app_model.is_public:
|
||||
return None
|
||||
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
from typing import Any, Protocol
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
|
||||
class RecommendAppRetrievalBase(Protocol):
|
||||
"""Interface for recommend app retrieval."""
|
||||
@@ -10,6 +8,6 @@ class RecommendAppRetrievalBase(Protocol):
|
||||
|
||||
def get_learn_dify_apps(self, language: str) -> Any: ...
|
||||
|
||||
def get_recommend_app_detail(self, app_id: str, *, session: Session | None = None) -> Any: ...
|
||||
def get_recommend_app_detail(self, app_id: str) -> Any: ...
|
||||
|
||||
def get_type(self) -> str: ...
|
||||
|
||||
@@ -3,7 +3,6 @@ from typing import Any, override
|
||||
|
||||
import httpx
|
||||
from flask import has_request_context, request
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from configs import dify_config
|
||||
from services.recommend_app.buildin.buildin_retrieval import BuildInRecommendAppRetrieval
|
||||
@@ -34,8 +33,7 @@ class RemoteRecommendAppRetrieval(RecommendAppRetrievalBase):
|
||||
"""
|
||||
|
||||
@override
|
||||
def get_recommend_app_detail(self, app_id: str, *, session: Session | None = None):
|
||||
del session
|
||||
def get_recommend_app_detail(self, app_id: str):
|
||||
try:
|
||||
result = self.fetch_recommended_app_detail_from_dify_official(app_id)
|
||||
except Exception as e:
|
||||
|
||||
@@ -1,28 +1,17 @@
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session, scoped_session
|
||||
from sqlalchemy.orm import scoped_session
|
||||
|
||||
from configs import dify_config
|
||||
from models.model import AccountTrialAppRecord, App, TrialApp
|
||||
from models.model import AccountTrialAppRecord, TrialApp
|
||||
from services.feature_service import FeatureService
|
||||
from services.recommend_app.recommend_app_factory import RecommendAppRetrievalFactory
|
||||
|
||||
|
||||
class RecommendedAppService:
|
||||
@classmethod
|
||||
def get_app(cls, app_id: str, *, session: Session) -> App | None:
|
||||
"""Return a normal app only when it belongs to the recommended catalog."""
|
||||
mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE
|
||||
retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)()
|
||||
recommended_app_detail = retrieval_instance.get_recommend_app_detail(app_id, session=session)
|
||||
if recommended_app_detail is None:
|
||||
return None
|
||||
|
||||
return session.scalar(select(App).where(App.id == app_id, App.status == "normal").limit(1))
|
||||
|
||||
@classmethod
|
||||
def get_recommended_apps_and_categories(cls, language: str, *, session: Session):
|
||||
def get_recommended_apps_and_categories(cls, session: scoped_session, language: str):
|
||||
"""
|
||||
Get recommended apps and categories.
|
||||
:param language: language
|
||||
@@ -99,6 +88,6 @@ class RecommendedAppService:
|
||||
session.commit()
|
||||
|
||||
@staticmethod
|
||||
def _can_trial_app(session: Session | scoped_session, app_id: str) -> bool:
|
||||
def _can_trial_app(session: scoped_session, app_id: str) -> bool:
|
||||
trial_app_model = session.scalar(select(TrialApp).where(TrialApp.app_id == app_id).limit(1))
|
||||
return trial_app_model is not None
|
||||
|
||||
@@ -2,11 +2,11 @@ import json
|
||||
import logging
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
from httpx import get
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from core.entities.provider_entities import ProviderConfig
|
||||
from core.helper import ssrf_proxy
|
||||
from core.tools.__base.tool_runtime import ToolRuntime
|
||||
from core.tools.custom_tool.provider import ApiToolProviderController
|
||||
from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity
|
||||
@@ -212,7 +212,7 @@ class ApiToolManageService:
|
||||
}
|
||||
|
||||
try:
|
||||
response = ssrf_proxy.get(url, headers=headers, timeout=10)
|
||||
response = get(url, headers=headers, timeout=10)
|
||||
if response.status_code != 200:
|
||||
raise ValueError(f"Got status code {response.status_code}")
|
||||
schema = response.text
|
||||
|
||||
@@ -25,7 +25,6 @@ from core.repositories import DifyCoreRepositoryFactory
|
||||
from extensions.ext_database import db
|
||||
from graphon.entities import WorkflowStartReason
|
||||
from graphon.enums import WorkflowExecutionStatus
|
||||
from graphon.filters import ResponseStreamFilter
|
||||
from graphon.runtime import GraphRuntimeState
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from libs.flask_utils import set_login_user
|
||||
@@ -487,7 +486,6 @@ def _resume_app_execution(payload: dict[str, Any]) -> None:
|
||||
generate_entity = resumption_context.get_generate_entity()
|
||||
|
||||
graph_runtime_state = GraphRuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state)
|
||||
response_stream_filter = resumption_context.get_response_stream_filter()
|
||||
|
||||
conversation = None
|
||||
message = None
|
||||
@@ -564,7 +562,6 @@ def _resume_app_execution(payload: dict[str, Any]) -> None:
|
||||
message=message,
|
||||
generate_entity=generate_entity,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
response_stream_filter=response_stream_filter,
|
||||
session_factory=session_factory,
|
||||
pause_state_config=pause_config,
|
||||
workflow_run_id=workflow_run_id,
|
||||
@@ -577,7 +574,6 @@ def _resume_app_execution(payload: dict[str, Any]) -> None:
|
||||
user=user,
|
||||
generate_entity=generate_entity,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
response_stream_filter=response_stream_filter,
|
||||
session_factory=session_factory,
|
||||
pause_state_config=pause_config,
|
||||
workflow_run_id=workflow_run_id,
|
||||
@@ -596,7 +592,6 @@ def _resume_advanced_chat(
|
||||
message: Message,
|
||||
generate_entity: AdvancedChatAppGenerateEntity,
|
||||
graph_runtime_state: GraphRuntimeState,
|
||||
response_stream_filter: ResponseStreamFilter,
|
||||
session_factory: sessionmaker,
|
||||
pause_state_config: PauseStateLayerConfig,
|
||||
workflow_run_id: str,
|
||||
@@ -636,7 +631,6 @@ def _resume_advanced_chat(
|
||||
workflow_node_execution_repository=workflow_node_execution_repository,
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
pause_state_config=pause_state_config,
|
||||
response_stream_filter=response_stream_filter,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to resume chatflow execution for workflow run %s", workflow_run_id)
|
||||
@@ -660,7 +654,6 @@ def _resume_workflow(
|
||||
user: Account | EndUser,
|
||||
generate_entity: WorkflowAppGenerateEntity,
|
||||
graph_runtime_state: GraphRuntimeState,
|
||||
response_stream_filter: ResponseStreamFilter,
|
||||
session_factory: sessionmaker,
|
||||
pause_state_config: PauseStateLayerConfig,
|
||||
workflow_run_id: str,
|
||||
@@ -700,7 +693,6 @@ def _resume_workflow(
|
||||
workflow_execution_repository=workflow_execution_repository,
|
||||
workflow_node_execution_repository=workflow_node_execution_repository,
|
||||
pause_state_config=pause_state_config,
|
||||
response_stream_filter=response_stream_filter,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to resume workflow execution for workflow run %s", workflow_run_id)
|
||||
|
||||
@@ -232,7 +232,6 @@ def resume_workflow_execution(task_data_dict: dict[str, Any]) -> None:
|
||||
return
|
||||
|
||||
graph_runtime_state = GraphRuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state)
|
||||
response_stream_filter = resumption_context.get_response_stream_filter()
|
||||
|
||||
with session_factory() as session:
|
||||
workflow = session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id))
|
||||
@@ -295,7 +294,6 @@ def resume_workflow_execution(task_data_dict: dict[str, Any]) -> None:
|
||||
workflow_node_execution_repository=workflow_node_execution_repository,
|
||||
graph_engine_layers=graph_engine_layers,
|
||||
pause_state_config=pause_config,
|
||||
response_stream_filter=response_stream_filter,
|
||||
)
|
||||
workflow_run_repo.delete_workflow_pause(pause_entity)
|
||||
|
||||
|
||||
@@ -1,67 +0,0 @@
|
||||
"""Initialize default RBAC access for existing workspace members after app creation."""
|
||||
|
||||
import logging
|
||||
|
||||
from celery import shared_task
|
||||
|
||||
from configs import dify_config
|
||||
from extensions.ext_database import db
|
||||
from services.account_service import TenantService
|
||||
from services.enterprise import rbac_service as enterprise_rbac_service
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
APP_RBAC_ACCOUNT_POLICY_BATCH_SIZE = 500
|
||||
APP_RBAC_DEFAULT_ACCESS_POLICY_ID = "default"
|
||||
APP_RBAC_QUEUE = "app_rbac"
|
||||
|
||||
|
||||
@shared_task(queue=APP_RBAC_QUEUE, bind=True, max_retries=3, default_retry_delay=60)
|
||||
def initialize_created_app_rbac_access_task(
|
||||
self, tenant_id: str, account_id: str, app_id: str | None = None, dataset_id: str | None = None
|
||||
) -> None:
|
||||
"""Grant the default app policy to current workspace members.
|
||||
|
||||
App scope is persisted synchronously before this task is queued. Replacing
|
||||
member policies is idempotent, so retrying the whole synchronization is safe
|
||||
when the enterprise RBAC service is temporarily unavailable.
|
||||
"""
|
||||
if not dify_config.RBAC_ENABLED:
|
||||
return
|
||||
|
||||
try:
|
||||
for account_ids in TenantService.iter_member_account_id_batches(
|
||||
tenant_id,
|
||||
APP_RBAC_ACCOUNT_POLICY_BATCH_SIZE,
|
||||
session=db.session(),
|
||||
):
|
||||
if app_id is not None:
|
||||
enterprise_rbac_service.RBACService.AppAccess.replace_user_access_policies(
|
||||
tenant_id=tenant_id,
|
||||
account_id=account_id,
|
||||
app_id=app_id,
|
||||
target_account_id=None,
|
||||
payload=enterprise_rbac_service.ReplaceUserAccessPolicies(
|
||||
access_policy_ids=[APP_RBAC_DEFAULT_ACCESS_POLICY_ID],
|
||||
account_ids=account_ids,
|
||||
),
|
||||
)
|
||||
elif dataset_id is not None:
|
||||
enterprise_rbac_service.RBACService.DatasetAccess.replace_user_access_policies(
|
||||
tenant_id=tenant_id,
|
||||
account_id=account_id,
|
||||
dataset_id=dataset_id,
|
||||
target_account_id=None,
|
||||
payload=enterprise_rbac_service.ReplaceUserAccessPolicies(
|
||||
access_policy_ids=[APP_RBAC_DEFAULT_ACCESS_POLICY_ID],
|
||||
account_ids=account_ids,
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception(
|
||||
"Failed to initialize app RBAC access; retrying: tenant_id=%s app_id=%s attempt=%s",
|
||||
tenant_id,
|
||||
app_id,
|
||||
self.request.retries + 1,
|
||||
)
|
||||
raise self.retry(exc=exc)
|
||||
@@ -701,29 +701,16 @@ 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()
|
||||
continue
|
||||
break
|
||||
session.commit()
|
||||
|
||||
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,6 +72,7 @@ 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(
|
||||
@@ -82,7 +83,7 @@ def mint_token(flask_app: Flask):
|
||||
subject_issuer=subject_issuer,
|
||||
client_id="difyctl",
|
||||
device_label="test-device",
|
||||
expires_at=datetime.now(UTC) + timedelta(hours=1),
|
||||
expires_at=expires_at or (datetime.now(UTC) + timedelta(hours=1)),
|
||||
)
|
||||
db.session.add(row)
|
||||
db.session.commit()
|
||||
@@ -111,6 +112,21 @@ 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,6 +6,7 @@ acceptance/rejection on app-scoped routes.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Generator
|
||||
|
||||
import pytest
|
||||
@@ -16,6 +17,50 @@ 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,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user