Compare commits
58
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c0a78f0638 | ||
|
|
d34ba75357 | ||
|
|
a4f5536106 | ||
|
|
662f706f92 | ||
|
|
a00db76ddf | ||
|
|
56b0b57ff7 | ||
|
|
99c3d7d0f0 | ||
|
|
cf1ebdadf5 | ||
|
|
a3309cd857 | ||
|
|
7fc8eed716 | ||
|
|
b3e5f29421 | ||
|
|
26639e0923 | ||
|
|
7852c273e4 | ||
|
|
ef54229d6f | ||
|
|
855bb32306 | ||
|
|
f380bbaa10 | ||
|
|
b67a04aa22 | ||
|
|
ab11083c2d | ||
|
|
0cc27dd401 | ||
|
|
7d2f25df8e | ||
|
|
e3d0320826 | ||
|
|
29b95d3ced | ||
|
|
9fd153ad99 | ||
|
|
76e587f78a | ||
|
|
99010dab3e | ||
|
|
82d08851be | ||
|
|
0d7ca17cd1 | ||
|
|
7cca8b6bb0 | ||
|
|
4065f63dce | ||
|
|
7c20ffe6c4 | ||
|
|
25b90229bc | ||
|
|
084f122814 | ||
|
|
1d74bff311 | ||
|
|
7aa20d6d94 | ||
|
|
8f6b57fe24 | ||
|
|
c83dcce1f7 | ||
|
|
47ee9f7435 | ||
|
|
762e7f7e8a | ||
|
|
c1ab6226a2 | ||
|
|
908c148667 | ||
|
|
1acf6e7eb6 | ||
|
|
ea3ef81396 | ||
|
|
7b3508e376 | ||
|
|
6a5ddc751c | ||
|
|
34c1bf1062 | ||
|
|
4f61353dc2 | ||
|
|
7e9cb50152 | ||
|
|
3e606ff0dc | ||
|
|
d5cdb2e6f1 | ||
|
|
265dc54c51 | ||
|
|
89d4fe91bc | ||
|
|
16fd55ab58 | ||
|
|
ff81e7b393 | ||
|
|
8cc690268b | ||
|
|
f06127aaa4 | ||
|
|
8c484411ea | ||
|
|
4c083e76e2 | ||
|
|
24080010c9 |
@@ -36,19 +36,30 @@ Use this as the decision guide for React/TypeScript component structure. Existin
|
||||
- Avoid prop drilling. One pass-through layer is acceptable; repeated forwarding means ownership should move down or into feature-scoped Jotai UI state. Keep server/cache state in query and API data flow.
|
||||
- Do not replace prop drilling with one top-level hook that returns a large view model and then thread that object through section props. Move each hook, query, derived value, and handler to the concrete section that consumes it, or use feature-scoped Jotai atoms for simple shared form/UI state when siblings need the same source of truth.
|
||||
- When using feature-scoped Jotai state for a form, drawer, or other secondary surface, scope the store to that surface instance when stale cross-instance state is possible. Initialize stable config at the owning boundary, then let descendants read only the atoms or purpose-named hooks they actually need.
|
||||
- For Jotai-backed surfaces, put shared query atoms, mutation atoms, derived state, and write actions in the feature state file when they coordinate multiple descendants. The lowest-owner rule still applies to independent visual surfaces that do not participate in shared state.
|
||||
- For repeated row/menu action surfaces that need reset, hydrate the stable identity at the surface entry and scope only the primitives that truly need per-instance reset, such as open flags, drafts, or selected local options.
|
||||
- Keep callbacks in a parent only for workflow coordination such as form submission, shared selection, batch behavior, or navigation. Otherwise let the child or row own its action.
|
||||
- Prefer uncontrolled DOM state and CSS variables before adding controlled props.
|
||||
|
||||
## Feature-Scoped Jotai State
|
||||
|
||||
- A module's feature-local state lives in one state file for Jotai-backed features: primitive atoms, query atoms, derived atoms, write-only action atoms, mutation atoms, submission orchestration, provider exports, and optional scope configuration.
|
||||
- Keep state local when one component owns it, even inside Jotai-backed features. Dialog open flags, menu/popover visibility, confirmation visibility, form/input drafts, row-local pending flags, and in-flight refs usually belong in component state.
|
||||
- Promote UI state to an atom only when siblings need the same source of truth, the value drives a query or mutation atom, a parent workflow coordinates the state, or the state intentionally persists across hidden or unmounted descendants within a scoped surface.
|
||||
- Reflect atom-backed surface-wide locks or invariants in every affected trigger. If only one row, menu, or dialog should be disabled, keep the pending or lock state local to that row, menu, or dialog.
|
||||
- Atom order in the state file follows the dependency graph: types/constants, editable primitives, query atoms, query-data derived atoms, readiness/business derived atoms, write actions, mutation atoms, submission orchestration, provider exports.
|
||||
- Derived atom names read as business facts. Write atom names read as user or workflow commands.
|
||||
- UI components read and write the exact atom they use with `useAtomValue` or `useSetAtom`. Repeated workflow semantics live in named derived atoms or write atoms.
|
||||
- Non-query derived atoms return a narrow value with a clear domain name. Query atoms expose the TanStack Query result object so loading, error, fetch, and pagination state stay attached to the query contract.
|
||||
- Write-only atoms own state transitions that update multiple primitives, reset dependent state, guard stale async work, or advance the workflow.
|
||||
- Non-query derived atoms return a narrow value with a clear domain name; avoid pass-through aliases or bundling unrelated UI facts. Query atoms expose the TanStack Query result object so loading, error, fetch, and pagination state stay attached to the query contract.
|
||||
- Write-only atoms own synchronous state transitions that update multiple primitives, reset dependent state, or advance the workflow. Async work with loading, error, caching, retry, or stale-result concerns should be modeled as query or mutation atoms, with write atoms only changing the inputs that drive them.
|
||||
- Avoid feature hooks that aggregate form values, query results, derived state, and commands for sibling components. Prefer named derived atoms and write atoms so UI components read the exact shared fact or command they need.
|
||||
- When a form library owns validation, keep submit orchestration in feature state when post-submit result or error state is shared by the surface. Avoid duplicating validation gates or request shaping in UI hooks.
|
||||
- `jotai-tanstack-query` atoms use the same QueryClient as the React Query provider. Query atoms belong in feature state when atoms are the feature's local state surface.
|
||||
- Jotai scope is an optional instance-isolation tool for secondary surfaces with independent local state. Query atoms keep shared cache behavior through the shared QueryClient.
|
||||
- Jotai scope is an optional instance-isolation tool for secondary surfaces with independent local state. Query and mutation atoms keep shared cache behavior through the shared QueryClient.
|
||||
- Do not put `atomWithQuery`, `atomWithInfiniteQuery`, `atomWithMutation`, or broad derived orchestration atoms in a `ScopeProvider` just to reset a surface. Scoped derived atoms implicitly scope their dependencies, which can duplicate query client access and break shared invalidation. Leave query/mutation atoms unscoped; let them read scoped primitive inputs.
|
||||
- Scope providers should list resettable primitive atoms and explicit hydration tuples. If a derived atom must be scoped, confirm that every dependency it implicitly scopes is meant to be private to that surface.
|
||||
- Keep independent dialog lifecycles separate. Avoid a single discriminated "current action dialog" atom when edit, delete, and other dialogs have their own open state, loading guard, or reset behavior.
|
||||
- Route-derived stable identities that do not need instance reset or scoped isolation can be hydrated at the route or layout boundary into a feature route atom. Use scoped atoms only when stale cross-instance state or per-surface reset semantics are needed.
|
||||
|
||||
## Components, Props, And Types
|
||||
|
||||
@@ -71,6 +82,7 @@ Use this as the decision guide for React/TypeScript component structure. Existin
|
||||
- Use generated enum objects and union types directly in props, comparisons, status logic, and i18n keys. Do not add local enum constants or parallel frontend enum/status layers unless they model real product state not represented by the API. Presentation-only tone maps should be keyed by the generated enum.
|
||||
- Normalize or coerce only at a real boundary, such as user-entered forms, search, URL/query params, file names, DOM IDs, or legacy adapters. Preserve user-entered values when whitespace or formatting can be meaningful.
|
||||
- Do not coerce nullable or optional API strings to `''` in query, derived model, or payload-building code. Keep `undefined` or `null` until the final boundary that requires a string.
|
||||
- Do not use `value || undefined` for mutation payload fields where an empty string means "clear this value". Trim or normalize at the form boundary, then preserve `''` when the API contract treats it as an intentional update.
|
||||
- Local UI models are fine for presentation, form state, select options, or guarded required-field refinements. Name them as UI concepts, not generated DTO mirrors.
|
||||
- Required-value refinements are allowed only after same-branch filtering or early return. Prefer nullable-tolerant props for render-only data.
|
||||
- When a component needs a stricter shape than a generated DTO, refine once at the API/query-to-UI boundary into a purpose-named UI type instead of hiding missing fields with generic fallback or coercion helpers.
|
||||
@@ -90,12 +102,17 @@ Use this as the decision guide for React/TypeScript component structure. Existin
|
||||
|
||||
- Keep `web/contract/*` as the single source of truth for API shape; follow existing domain/router patterns and the `{ params, query?, body? }` input shape.
|
||||
- Consume queries directly with `useQuery(consoleQuery.xxx.queryOptions(...))` or `useQuery(marketplaceQuery.xxx.queryOptions(...))`.
|
||||
- In `atomWithQuery` and `atomWithInfiniteQuery`, return generated `queryOptions()` or `infiniteOptions()` directly. Pass `enabled`, `retry`, `placeholderData`, `select`, and pagination options into that call instead of spreading generated options into a hand-built object.
|
||||
- In `atomWithMutation`, return generated `mutationOptions()` directly when using generated clients. Put request shaping and submit orchestration in write atoms; do not rebuild mutation option objects just to pass through the generated mutation function.
|
||||
- For custom query functions that do not come from generated clients, wrap the options object with TanStack `queryOptions(...)` so query atoms still return a query options contract.
|
||||
- Avoid pass-through hooks and thin `web/service/use-*` wrappers that only rename `queryOptions()` or `mutationOptions()`. Extract a small `queryOptions` helper only when repeated call-site options justify it.
|
||||
- Keep feature hooks for real orchestration, workflow state, or shared domain behavior.
|
||||
- For TanStack cache data, use generated or query-derived types; do not create local wrappers for `getQueryData` or `getQueriesData`.
|
||||
- For generated oRPC `queryOptions()` / `infiniteOptions()`, do not pass `skipToken` as `input`; keep a valid placeholder input shape and use `enabled` to gate missing required params because the OpenAPI codec encodes input eagerly.
|
||||
- For generated oRPC `queryOptions()` / `infiniteOptions()`, keep returning the generated options directly. When required input is missing, use a whole-input branch such as `input: condition ? validInput : skipToken` together with `enabled: Boolean(condition)` so no request runs and no fake payload is built.
|
||||
- Do not put `skipToken` inside a nested placeholder payload, such as `{ params: { appInstanceId: skipToken } }`. Do not create hand-written "missing queryOptions" objects or coerce required IDs to `''`.
|
||||
- Consume mutations directly with `useMutation(consoleQuery.xxx.mutationOptions(...))` or `useMutation(marketplaceQuery.xxx.mutationOptions(...))`; use oRPC clients as `mutationFn` only for custom flows.
|
||||
- Put shared cache behavior in `createTanstackQueryUtils(...experimental_defaults...)`; components may add UI feedback callbacks, but should not own shared invalidation rules.
|
||||
- Component or atom mutation callbacks can handle local UI feedback such as toasts, closing dialogs, or navigation. They should not replace shared invalidation or add local cache patches for shared server state.
|
||||
- Do not use deprecated `useInvalid` or `useReset`.
|
||||
- Prefer `mutate(...)`; use `mutateAsync(...)` only when Promise semantics are required, and wrap awaited calls in `try/catch`.
|
||||
|
||||
@@ -107,8 +124,9 @@ Use this as the decision guide for React/TypeScript component structure. Existin
|
||||
- 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 that owns the trigger, open state, and hidden content when it would obscure the parent flow.
|
||||
- Preserve composability by separating behavior ownership from layout ownership. A dropdown action may own its trigger, open state, and menu content; the caller owns placement such as slots, offsets, and alignment.
|
||||
- When a dialog, dropdown, or popover component already accepts controlled `open` state, mount the surface unconditionally unless unmounting is required for performance or reset semantics. Use keyed scope or local state reset for reset behavior instead of `{open && <Surface />}` wrappers.
|
||||
- Avoid unnecessary DOM hierarchy. Do not add wrapper elements unless they provide layout, semantics, accessibility, state ownership, or integration with a library API; prefer fragments or styling an existing element when possible.
|
||||
- Avoid shallow wrappers, hook-to-props adapter components, layout-only render-prop wrappers, and prop renaming unless the wrapper adds validation, orchestration, error handling, state ownership, or a real semantic boundary. If a component only calls a hook and forwards every returned field to one child, move the hook into that child or make the wrapper own a real surface.
|
||||
- Avoid shallow wrappers, hook-to-props adapter components, layout-only render-prop wrappers, children-as-pass-through composition, and prop renaming unless the wrapper adds validation, orchestration, error handling, state ownership, or a real semantic boundary. If a component only calls a hook, forwards props, or passes trigger/content through to one child, move the logic into that child or make the wrapper own a real surface.
|
||||
|
||||
## You Might Not Need An Effect
|
||||
|
||||
@@ -117,6 +135,7 @@ Use this as the decision guide for React/TypeScript component structure. Existin
|
||||
- Do not use Effects to handle user actions. Put action-specific logic in the event handler where the cause is known.
|
||||
- Do not use Effects to copy one state value into another state value representing the same concept. Pick one source of truth and derive the rest during render.
|
||||
- Do not reset or adjust state from props with an Effect. Prefer a `key` reset, storing a stable ID and deriving the selected object, or guarded same-component render-time adjustment when truly necessary.
|
||||
- For forms initialized from query data, prefer keyed remounts or surface-entry hydration of form/field atoms over an Effect that copies query data into form state.
|
||||
- Prefer framework data APIs or TanStack Query for data fetching instead of writing request Effects in components.
|
||||
- If an Effect still seems necessary, first name the external system it synchronizes with. If there is no external system, remove the Effect and restructure the state or event flow.
|
||||
|
||||
|
||||
@@ -8,8 +8,6 @@ on:
|
||||
- "build/**"
|
||||
- "release/e-*"
|
||||
- "hotfix/**"
|
||||
- "feat/hitl-backend"
|
||||
- "feat/rbac"
|
||||
tags:
|
||||
- "*"
|
||||
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
name: Deploy HITL
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows: ["Build and Push API & Web"]
|
||||
branches:
|
||||
- "build/feat/hitl"
|
||||
types:
|
||||
- completed
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: depot-ubuntu-24.04
|
||||
if: |
|
||||
github.event.workflow_run.conclusion == 'success' &&
|
||||
github.event.workflow_run.head_branch == 'build/feat/hitl'
|
||||
steps:
|
||||
- name: Deploy to server
|
||||
uses: appleboy/ssh-action@0ff4204d59e8e51228ff73bce53f80d53301dee2 # v1.2.5
|
||||
with:
|
||||
host: ${{ secrets.HITL_SSH_HOST }}
|
||||
username: ${{ secrets.SSH_USER }}
|
||||
key: ${{ secrets.SSH_PRIVATE_KEY }}
|
||||
script: |
|
||||
${{ vars.SSH_SCRIPT || secrets.SSH_SCRIPT }}
|
||||
@@ -768,7 +768,6 @@ EVENT_BUS_REDIS_CHANNEL_TYPE=pubsub
|
||||
# Whether to use Redis cluster mode while use redis as event bus.
|
||||
# It's highly recommended to enable this for large deployments.
|
||||
EVENT_BUS_REDIS_USE_CLUSTERS=false
|
||||
EVENT_BUS_LISTENER_JOIN_TIMEOUT_MS=2000
|
||||
|
||||
# Whether to Enable human input timeout check task
|
||||
ENABLE_HUMAN_INPUT_TIMEOUT_TASK=true
|
||||
|
||||
@@ -25,6 +25,7 @@ from .plugin import (
|
||||
from .rbac import migrate_member_roles_to_rbac
|
||||
from .retention import (
|
||||
archive_workflow_runs,
|
||||
archive_workflow_runs_plan,
|
||||
clean_expired_messages,
|
||||
clean_workflow_runs,
|
||||
cleanup_orphaned_draft_variables,
|
||||
@@ -51,6 +52,7 @@ from .vector import (
|
||||
__all__ = [
|
||||
"add_qdrant_index",
|
||||
"archive_workflow_runs",
|
||||
"archive_workflow_runs_plan",
|
||||
"backfill_plugin_auto_upgrade",
|
||||
"clean_expired_messages",
|
||||
"clean_workflow_runs",
|
||||
|
||||
+544
-95
@@ -12,10 +12,160 @@ from services.clear_free_plan_tenant_expired_logs import ClearFreePlanTenantExpi
|
||||
from services.retention.conversation.messages_clean_policy import create_message_clean_policy
|
||||
from services.retention.conversation.messages_clean_service import MessagesCleanService
|
||||
from services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup
|
||||
from services.retention.workflow_run.tenant_prefix import tenant_prefix_condition
|
||||
from tasks.remove_app_and_related_data_task import delete_draft_variables_batch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_HEX_PREFIXES = tuple("0123456789abcdef")
|
||||
|
||||
|
||||
class WorkflowRunArchivePlanRow(TypedDict):
|
||||
tenant_prefix: str
|
||||
total_tenants: int
|
||||
workflow_runs: int
|
||||
workflow_node_executions: int
|
||||
paid_tenants: int
|
||||
unpaid_tenants: int
|
||||
|
||||
|
||||
class WorkflowRunArchiveTenantPlan(TypedDict):
|
||||
archive_tenant_ids: list[str] | None
|
||||
paid_tenant_ids: list[str]
|
||||
unpaid_tenant_ids: list[str]
|
||||
|
||||
|
||||
def _parse_tenant_prefixes(prefixes: str | None) -> list[str]:
|
||||
if not prefixes:
|
||||
return []
|
||||
|
||||
parsed = []
|
||||
for raw_prefix in prefixes.split(","):
|
||||
prefix = raw_prefix.strip().lower()
|
||||
if not prefix:
|
||||
continue
|
||||
if len(prefix) != 1 or prefix not in _HEX_PREFIXES:
|
||||
raise click.UsageError("--tenant-prefixes must be a comma-separated list of hex digits, e.g. 0,1,a,f.")
|
||||
parsed.append(prefix)
|
||||
return sorted(set(parsed))
|
||||
|
||||
|
||||
def _get_archive_candidate_tenant_ids_by_prefix(
|
||||
prefix: str,
|
||||
*,
|
||||
start_from: datetime.datetime | None,
|
||||
end_before: datetime.datetime,
|
||||
) -> list[str]:
|
||||
from graphon.enums import WorkflowExecutionStatus
|
||||
from models.workflow import WorkflowRun
|
||||
from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver
|
||||
|
||||
conditions = [
|
||||
WorkflowRun.created_at < end_before,
|
||||
WorkflowRun.status.in_(WorkflowExecutionStatus.ended_values()),
|
||||
WorkflowRun.type.in_(WorkflowRunArchiver.ARCHIVED_TYPE),
|
||||
tenant_prefix_condition(WorkflowRun.tenant_id, prefix),
|
||||
]
|
||||
if start_from is not None:
|
||||
conditions.append(WorkflowRun.created_at >= start_from)
|
||||
|
||||
tenant_ids = db.session.scalars(
|
||||
sa.select(WorkflowRun.tenant_id).where(*conditions).distinct().order_by(WorkflowRun.tenant_id)
|
||||
).all()
|
||||
return list(tenant_ids)
|
||||
|
||||
|
||||
def _filter_paid_workflow_archive_tenant_ids(tenant_ids: list[str]) -> tuple[list[str], list[str]]:
|
||||
from configs import dify_config
|
||||
from enums.cloud_plan import CloudPlan
|
||||
from services.billing_service import BillingService
|
||||
|
||||
tenant_ids = sorted(set(tenant_ids))
|
||||
if not tenant_ids:
|
||||
return [], []
|
||||
if not dify_config.BILLING_ENABLED:
|
||||
return tenant_ids, []
|
||||
|
||||
plans = BillingService.get_plan_bulk_with_cache(tenant_ids)
|
||||
paid_tenant_ids = [
|
||||
tenant_id
|
||||
for tenant_id in tenant_ids
|
||||
if plans.get(tenant_id) and plans[tenant_id].get("plan") in (CloudPlan.PROFESSIONAL, CloudPlan.TEAM)
|
||||
]
|
||||
unpaid_tenant_ids = sorted(set(tenant_ids) - set(paid_tenant_ids))
|
||||
return paid_tenant_ids, unpaid_tenant_ids
|
||||
|
||||
|
||||
def _resolve_archive_tenant_ids_from_plan(
|
||||
*,
|
||||
tenant_ids: str | None,
|
||||
tenant_prefixes: list[str],
|
||||
start_from: datetime.datetime | None,
|
||||
end_before: datetime.datetime,
|
||||
) -> WorkflowRunArchiveTenantPlan:
|
||||
"""
|
||||
Resolve the archive tenant scope once before scanning workflow_runs.
|
||||
|
||||
Prefix rollout should use the tenant list collected by the same planning path, then archive by
|
||||
tenant_id IN (...). Scanning workflow_runs with a tenant prefix range in every archive run is too expensive on
|
||||
the large production table this command is meant to shrink.
|
||||
"""
|
||||
if tenant_ids:
|
||||
requested_tenant_ids = [tid.strip() for tid in tenant_ids.split(",") if tid.strip()]
|
||||
elif tenant_prefixes:
|
||||
requested_tenant_ids = []
|
||||
for prefix in tenant_prefixes:
|
||||
requested_tenant_ids.extend(
|
||||
_get_archive_candidate_tenant_ids_by_prefix(
|
||||
prefix,
|
||||
start_from=start_from,
|
||||
end_before=end_before,
|
||||
)
|
||||
)
|
||||
else:
|
||||
return WorkflowRunArchiveTenantPlan(
|
||||
archive_tenant_ids=None,
|
||||
paid_tenant_ids=[],
|
||||
unpaid_tenant_ids=[],
|
||||
)
|
||||
|
||||
paid_tenant_ids, unpaid_tenant_ids = _filter_paid_workflow_archive_tenant_ids(requested_tenant_ids)
|
||||
return WorkflowRunArchiveTenantPlan(
|
||||
archive_tenant_ids=paid_tenant_ids,
|
||||
paid_tenant_ids=paid_tenant_ids,
|
||||
unpaid_tenant_ids=unpaid_tenant_ids,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_archive_time_range(
|
||||
*,
|
||||
before_days: int,
|
||||
from_days_ago: int | None,
|
||||
to_days_ago: int | None,
|
||||
start_from: datetime.datetime | None,
|
||||
end_before: datetime.datetime | None,
|
||||
) -> tuple[int, datetime.datetime | None, datetime.datetime | None]:
|
||||
if (start_from is None) ^ (end_before is None):
|
||||
raise click.UsageError("--start-from and --end-before must be provided together.")
|
||||
|
||||
if (from_days_ago is None) ^ (to_days_ago is None):
|
||||
raise click.UsageError("--from-days-ago and --to-days-ago must be provided together.")
|
||||
|
||||
if from_days_ago is not None and to_days_ago is not None:
|
||||
if start_from or end_before:
|
||||
raise click.UsageError("Choose either day offsets or explicit dates, not both.")
|
||||
if from_days_ago <= to_days_ago:
|
||||
raise click.UsageError("--from-days-ago must be greater than --to-days-ago.")
|
||||
now = datetime.datetime.now()
|
||||
start_from = now - datetime.timedelta(days=from_days_ago)
|
||||
end_before = now - datetime.timedelta(days=to_days_ago)
|
||||
before_days = 0
|
||||
|
||||
if start_from and end_before and start_from >= end_before:
|
||||
raise click.UsageError("--start-from must be earlier than --end-before.")
|
||||
|
||||
return before_days, start_from, end_before
|
||||
|
||||
|
||||
@click.command("clear-free-plan-tenant-expired-logs", help="Clear free plan tenant expired logs.")
|
||||
@click.option("--days", prompt=True, help="The days to clear free plan tenant expired logs.", default=30)
|
||||
@@ -139,11 +289,143 @@ def clean_workflow_runs(
|
||||
)
|
||||
|
||||
|
||||
@click.command(
|
||||
"archive-workflow-runs-plan",
|
||||
help="Plan workflow run archive rollout by tenant ID first hex digit.",
|
||||
)
|
||||
@click.option("--before-days", default=90, show_default=True, help="Plan runs older than N days.")
|
||||
@click.option(
|
||||
"--from-days-ago",
|
||||
default=None,
|
||||
type=click.IntRange(min=0),
|
||||
help="Lower bound in days ago (older). Must be paired with --to-days-ago.",
|
||||
)
|
||||
@click.option(
|
||||
"--to-days-ago",
|
||||
default=None,
|
||||
type=click.IntRange(min=0),
|
||||
help="Upper bound in days ago (newer). Must be paired with --from-days-ago.",
|
||||
)
|
||||
@click.option(
|
||||
"--start-from",
|
||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||
default=None,
|
||||
help="Plan runs created at or after this timestamp (UTC if no timezone).",
|
||||
)
|
||||
@click.option(
|
||||
"--end-before",
|
||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||
default=None,
|
||||
help="Plan runs created before this timestamp (UTC if no timezone).",
|
||||
)
|
||||
@click.option(
|
||||
"--include-archived",
|
||||
is_flag=True,
|
||||
help="Compatibility no-op for V2 bundle archive; plan counts source rows in the requested window.",
|
||||
)
|
||||
def archive_workflow_runs_plan(
|
||||
before_days: int,
|
||||
from_days_ago: int | None,
|
||||
to_days_ago: int | None,
|
||||
start_from: datetime.datetime | None,
|
||||
end_before: datetime.datetime | None,
|
||||
include_archived: bool,
|
||||
):
|
||||
"""
|
||||
Print the 16 tenant-prefix rollout rows used to choose archive execution order.
|
||||
|
||||
Counts use the same workflow run eligibility as archive-workflow-runs: ended runs,
|
||||
supported workflow types, and the requested created_at window. V2 bundle archive
|
||||
does not maintain per-run archive logs, so this plan reports source-table volume.
|
||||
"""
|
||||
from graphon.enums import WorkflowExecutionStatus
|
||||
from models.workflow import WorkflowNodeExecutionModel, WorkflowRun
|
||||
from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver
|
||||
|
||||
before_days, start_from, end_before = _resolve_archive_time_range(
|
||||
before_days=before_days,
|
||||
from_days_ago=from_days_ago,
|
||||
to_days_ago=to_days_ago,
|
||||
start_from=start_from,
|
||||
end_before=end_before,
|
||||
)
|
||||
plan_end_before = end_before or datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=before_days)
|
||||
if include_archived:
|
||||
click.echo(click.style("--include-archived is a no-op for V2 bundle archive plans.", fg="yellow"))
|
||||
|
||||
rows: list[WorkflowRunArchivePlanRow] = []
|
||||
for prefix in _HEX_PREFIXES:
|
||||
tenant_ids = _get_archive_candidate_tenant_ids_by_prefix(
|
||||
prefix,
|
||||
start_from=start_from,
|
||||
end_before=plan_end_before,
|
||||
)
|
||||
total_tenants = len(tenant_ids)
|
||||
paid_tenant_ids, unpaid_tenant_ids = _filter_paid_workflow_archive_tenant_ids(tenant_ids)
|
||||
|
||||
run_conditions = [
|
||||
WorkflowRun.created_at < plan_end_before,
|
||||
WorkflowRun.status.in_(WorkflowExecutionStatus.ended_values()),
|
||||
WorkflowRun.type.in_(WorkflowRunArchiver.ARCHIVED_TYPE),
|
||||
tenant_prefix_condition(WorkflowRun.tenant_id, prefix),
|
||||
]
|
||||
if start_from is not None:
|
||||
run_conditions.append(WorkflowRun.created_at >= start_from)
|
||||
workflow_runs = (
|
||||
db.session.scalar(sa.select(sa.func.count()).select_from(WorkflowRun).where(*run_conditions)) or 0
|
||||
)
|
||||
candidate_runs = sa.select(WorkflowRun.id).where(*run_conditions).subquery()
|
||||
workflow_node_executions = (
|
||||
db.session.scalar(
|
||||
sa.select(sa.func.count())
|
||||
.select_from(WorkflowNodeExecutionModel)
|
||||
.join(candidate_runs, WorkflowNodeExecutionModel.workflow_run_id == candidate_runs.c.id)
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
rows.append(
|
||||
WorkflowRunArchivePlanRow(
|
||||
tenant_prefix=prefix,
|
||||
total_tenants=total_tenants,
|
||||
workflow_runs=workflow_runs,
|
||||
workflow_node_executions=workflow_node_executions,
|
||||
paid_tenants=len(paid_tenant_ids),
|
||||
unpaid_tenants=len(unpaid_tenant_ids),
|
||||
)
|
||||
)
|
||||
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Workflow archive plan for runs before {plan_end_before.isoformat()}"
|
||||
f"{f' and at/after {start_from.isoformat()}' if start_from else ''}.",
|
||||
fg="white",
|
||||
)
|
||||
)
|
||||
click.echo("tenant_prefix,total_tenants,workflow_runs,workflow_node_executions,paid_tenants,unpaid_tenants")
|
||||
for row in rows:
|
||||
click.echo(
|
||||
f"{row['tenant_prefix']},{row['total_tenants']},{row['workflow_runs']},"
|
||||
f"{row['workflow_node_executions']},{row['paid_tenants']},{row['unpaid_tenants']}"
|
||||
)
|
||||
|
||||
ordered_rows = sorted(
|
||||
rows,
|
||||
key=lambda row: (row["workflow_runs"] + row["workflow_node_executions"], row["tenant_prefix"]),
|
||||
)
|
||||
click.echo("suggested_execution_order=" + ",".join(row["tenant_prefix"] for row in ordered_rows))
|
||||
|
||||
|
||||
@click.command(
|
||||
"archive-workflow-runs",
|
||||
help="Archive workflow runs for paid plan tenants to S3-compatible storage.",
|
||||
)
|
||||
@click.option("--tenant-ids", default=None, help="Optional comma-separated tenant IDs for grayscale rollout.")
|
||||
@click.option(
|
||||
"--tenant-prefixes",
|
||||
default=None,
|
||||
help="Optional comma-separated tenant ID first hex digits for rollout waves, e.g. 0,1,a,f.",
|
||||
)
|
||||
@click.option("--before-days", default=90, show_default=True, help="Archive runs older than N days.")
|
||||
@click.option(
|
||||
"--from-days-ago",
|
||||
@@ -169,13 +451,36 @@ def clean_workflow_runs(
|
||||
default=None,
|
||||
help="Archive runs created before this timestamp (UTC if no timezone).",
|
||||
)
|
||||
@click.option("--batch-size", default=100, show_default=True, help="Batch size for processing.")
|
||||
@click.option("--workers", default=1, show_default=True, type=int, help="Concurrent workflow runs to archive.")
|
||||
@click.option("--batch-size", default=100, show_default=True, help="Maximum workflow runs per archive bundle.")
|
||||
@click.option(
|
||||
"--workers",
|
||||
default=1,
|
||||
show_default=True,
|
||||
type=int,
|
||||
help="Reserved; bundle archive currently runs serially.",
|
||||
)
|
||||
@click.option(
|
||||
"--run-shard-index",
|
||||
default=None,
|
||||
type=click.IntRange(min=0),
|
||||
help="Zero-based workflow run shard index for parallel cron jobs. Must be paired with --run-shard-total.",
|
||||
)
|
||||
@click.option(
|
||||
"--run-shard-total",
|
||||
default=None,
|
||||
type=click.IntRange(min=1, max=16),
|
||||
help="Total workflow run shard count for parallel cron jobs. Must be paired with --run-shard-index.",
|
||||
)
|
||||
@click.option("--limit", default=None, type=int, help="Maximum number of runs to archive.")
|
||||
@click.option("--dry-run", is_flag=True, help="Preview without archiving.")
|
||||
@click.option("--delete-after-archive", is_flag=True, help="Delete runs and related data after archiving.")
|
||||
@click.option(
|
||||
"--delete-after-archive",
|
||||
is_flag=True,
|
||||
help="Not supported by bundle archive; use a separate bundle delete workflow after validation.",
|
||||
)
|
||||
def archive_workflow_runs(
|
||||
tenant_ids: str | None,
|
||||
tenant_prefixes: str | None,
|
||||
before_days: int,
|
||||
from_days_ago: int | None,
|
||||
to_days_ago: int | None,
|
||||
@@ -183,6 +488,8 @@ def archive_workflow_runs(
|
||||
end_before: datetime.datetime | None,
|
||||
batch_size: int,
|
||||
workers: int,
|
||||
run_shard_index: int | None,
|
||||
run_shard_total: int | None,
|
||||
limit: int | None,
|
||||
dry_run: bool,
|
||||
delete_after_archive: bool,
|
||||
@@ -190,14 +497,19 @@ def archive_workflow_runs(
|
||||
"""
|
||||
Archive workflow runs for paid plan tenants older than the specified days.
|
||||
|
||||
This command archives the following tables to storage:
|
||||
This command writes V2 tenant/month/shard archive bundles. Each bundle contains Parquet snapshots from:
|
||||
- workflow_runs
|
||||
- workflow_app_logs
|
||||
- workflow_node_executions
|
||||
- workflow_node_execution_offload
|
||||
- workflow_pauses
|
||||
- workflow_pause_reasons
|
||||
- workflow_trigger_logs
|
||||
|
||||
The workflow_runs and workflow_app_logs tables are preserved for UI listing.
|
||||
Source database rows are always preserved by archive. Deletion must be handled by
|
||||
a separate bundle-level delete workflow after manifest, checksum, row-count, and
|
||||
restore-sampling validation. In --dry-run mode, no storage or database writes
|
||||
happen; the command estimates per-table Parquet bytes and object size instead.
|
||||
"""
|
||||
from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver
|
||||
|
||||
@@ -209,32 +521,58 @@ def archive_workflow_runs(
|
||||
)
|
||||
)
|
||||
|
||||
if (start_from is None) ^ (end_before is None):
|
||||
click.echo(click.style("start-from and end-before must be provided together.", fg="red"))
|
||||
return
|
||||
|
||||
if (from_days_ago is None) ^ (to_days_ago is None):
|
||||
click.echo(click.style("from-days-ago and to-days-ago must be provided together.", fg="red"))
|
||||
return
|
||||
|
||||
if from_days_ago is not None and to_days_ago is not None:
|
||||
if start_from or end_before:
|
||||
click.echo(click.style("Choose either day offsets or explicit dates, not both.", fg="red"))
|
||||
return
|
||||
if from_days_ago <= to_days_ago:
|
||||
click.echo(click.style("from-days-ago must be greater than to-days-ago.", fg="red"))
|
||||
return
|
||||
now = datetime.datetime.now()
|
||||
start_from = now - datetime.timedelta(days=from_days_ago)
|
||||
end_before = now - datetime.timedelta(days=to_days_ago)
|
||||
before_days = 0
|
||||
|
||||
if start_from and end_before and start_from >= end_before:
|
||||
click.echo(click.style("start-from must be earlier than end-before.", fg="red"))
|
||||
try:
|
||||
before_days, start_from, end_before = _resolve_archive_time_range(
|
||||
before_days=before_days,
|
||||
from_days_ago=from_days_ago,
|
||||
to_days_ago=to_days_ago,
|
||||
start_from=start_from,
|
||||
end_before=end_before,
|
||||
)
|
||||
parsed_tenant_prefixes = _parse_tenant_prefixes(tenant_prefixes)
|
||||
except click.UsageError as e:
|
||||
click.echo(click.style(e.message, fg="red"))
|
||||
return
|
||||
plan_end_before = end_before or datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=before_days)
|
||||
if workers < 1:
|
||||
click.echo(click.style("workers must be at least 1.", fg="red"))
|
||||
return
|
||||
if (run_shard_index is None) ^ (run_shard_total is None):
|
||||
click.echo(click.style("run-shard-index and run-shard-total must be provided together.", fg="red"))
|
||||
return
|
||||
if run_shard_index is not None and run_shard_total is not None and run_shard_index >= run_shard_total:
|
||||
click.echo(click.style("run-shard-index must be less than run-shard-total.", fg="red"))
|
||||
return
|
||||
if delete_after_archive:
|
||||
click.echo(click.style("delete-after-archive is not supported by bundle archive.", fg="red"))
|
||||
return
|
||||
|
||||
try:
|
||||
tenant_plan = _resolve_archive_tenant_ids_from_plan(
|
||||
tenant_ids=tenant_ids,
|
||||
tenant_prefixes=parsed_tenant_prefixes,
|
||||
start_from=start_from,
|
||||
end_before=plan_end_before,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to resolve workflow archive tenant plan")
|
||||
click.echo(click.style("Failed to resolve workflow archive tenant plan.", fg="red"))
|
||||
return
|
||||
|
||||
planned_tenant_ids = tenant_plan["archive_tenant_ids"]
|
||||
planned_paid_tenant_ids = tenant_plan["paid_tenant_ids"] if planned_tenant_ids is not None else None
|
||||
paid_tenants = len(tenant_plan["paid_tenant_ids"])
|
||||
unpaid_tenants = len(tenant_plan["unpaid_tenant_ids"])
|
||||
if planned_tenant_ids is not None:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Resolved archive tenant plan: paid_tenants={paid_tenants}, unpaid_tenants={unpaid_tenants}.",
|
||||
fg="white",
|
||||
)
|
||||
)
|
||||
if not planned_tenant_ids:
|
||||
click.echo(click.style("No paid tenants matched the archive plan; nothing to archive.", fg="yellow"))
|
||||
return
|
||||
|
||||
archiver = WorkflowRunArchiver(
|
||||
days=before_days,
|
||||
@@ -242,7 +580,11 @@ def archive_workflow_runs(
|
||||
start_from=start_from,
|
||||
end_before=end_before,
|
||||
workers=workers,
|
||||
tenant_ids=[tid.strip() for tid in tenant_ids.split(",")] if tenant_ids else None,
|
||||
tenant_ids=planned_tenant_ids,
|
||||
tenant_prefixes=parsed_tenant_prefixes,
|
||||
paid_tenant_ids=planned_paid_tenant_ids,
|
||||
run_shard_index=run_shard_index,
|
||||
run_shard_total=run_shard_total,
|
||||
limit=limit,
|
||||
dry_run=dry_run,
|
||||
delete_after_archive=delete_after_archive,
|
||||
@@ -252,7 +594,9 @@ def archive_workflow_runs(
|
||||
click.style(
|
||||
f"Summary: processed={summary.total_runs_processed}, archived={summary.runs_archived}, "
|
||||
f"skipped={summary.runs_skipped}, failed={summary.runs_failed}, "
|
||||
f"time={summary.total_elapsed_time:.2f}s",
|
||||
f"bundles_archived={summary.bundles_archived}, bundles_skipped={summary.bundles_skipped}, "
|
||||
f"bundles_failed={summary.bundles_failed}, "
|
||||
f"object_size_bytes={summary.total_object_size_bytes}, time={summary.total_elapsed_time:.2f}s",
|
||||
fg="cyan",
|
||||
)
|
||||
)
|
||||
@@ -268,6 +612,52 @@ def archive_workflow_runs(
|
||||
)
|
||||
|
||||
|
||||
def _echo_bundle_archive_operation_summary(summary) -> None:
|
||||
status = "completed successfully" if summary.bundles_failed == 0 else "completed with failures"
|
||||
fg = "green" if summary.bundles_failed == 0 else "red"
|
||||
click.echo(
|
||||
click.style(
|
||||
f"{summary.operation} {status}. "
|
||||
f"bundles_success={summary.bundles_succeeded} bundles_failed={summary.bundles_failed} "
|
||||
f"runs={summary.runs_processed} rows={summary.rows_processed} "
|
||||
f"archive_bytes={summary.archive_bytes} duration={summary.elapsed_time:.2f}s "
|
||||
f"validation_time={summary.validation_time:.2f}s "
|
||||
f"runs_per_second={summary.runs_per_second:.2f} rows_per_second={summary.rows_per_second:.2f} "
|
||||
f"bytes_per_second={summary.bytes_per_second:.2f}",
|
||||
fg=fg,
|
||||
)
|
||||
)
|
||||
click.echo(click.style("table,row_count", fg="white"))
|
||||
for table_name in [
|
||||
"workflow_runs",
|
||||
"workflow_app_logs",
|
||||
"workflow_node_executions",
|
||||
"workflow_node_execution_offload",
|
||||
"workflow_pauses",
|
||||
"workflow_pause_reasons",
|
||||
"workflow_trigger_logs",
|
||||
]:
|
||||
click.echo(f"{table_name},{summary.table_counts.get(table_name, 0)}")
|
||||
for result in summary.results:
|
||||
if result.success:
|
||||
click.echo(
|
||||
click.style(
|
||||
f" bundle={result.bundle_id} tenant={result.tenant_id} runs={result.run_count} "
|
||||
f"rows={result.row_count} archive_bytes={result.archive_bytes} "
|
||||
f"time={result.elapsed_time:.2f}s validation={result.validation_time:.2f}s",
|
||||
fg="white",
|
||||
)
|
||||
)
|
||||
else:
|
||||
click.echo(
|
||||
click.style(
|
||||
f" failed bundle={result.bundle_id} tenant={result.tenant_id} "
|
||||
f"object_prefix={result.object_prefix} error={result.error}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@click.command(
|
||||
"restore-workflow-runs",
|
||||
help="Restore archived workflow runs from S3-compatible storage.",
|
||||
@@ -290,8 +680,8 @@ def archive_workflow_runs(
|
||||
default=None,
|
||||
help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.",
|
||||
)
|
||||
@click.option("--workers", default=1, show_default=True, type=int, help="Concurrent workflow runs to restore.")
|
||||
@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of runs to restore.")
|
||||
@click.option("--workers", default=1, show_default=True, type=int, help="V1 --run-id compatibility only.")
|
||||
@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of V2 bundles to restore.")
|
||||
@click.option("--dry-run", is_flag=True, help="Preview without restoring.")
|
||||
def restore_workflow_runs(
|
||||
tenant_ids: str | None,
|
||||
@@ -303,15 +693,18 @@ def restore_workflow_runs(
|
||||
dry_run: bool,
|
||||
):
|
||||
"""
|
||||
Restore an archived workflow run from storage to the database.
|
||||
Restore archived workflow runs from storage to the database.
|
||||
|
||||
This restores the following tables:
|
||||
Batch restore uses V2 bundle metadata and validates archive objects before writing source rows. This restores:
|
||||
- workflow_runs
|
||||
- workflow_app_logs
|
||||
- workflow_node_executions
|
||||
- workflow_node_execution_offload
|
||||
- workflow_pauses
|
||||
- workflow_pause_reasons
|
||||
- workflow_trigger_logs
|
||||
"""
|
||||
from services.retention.workflow_run.bundle_archive_maintenance import WorkflowRunBundleArchiveMaintenance
|
||||
from services.retention.workflow_run.restore_archived_workflow_run import WorkflowRunRestore
|
||||
|
||||
parsed_tenant_ids = None
|
||||
@@ -335,39 +728,46 @@ def restore_workflow_runs(
|
||||
)
|
||||
)
|
||||
|
||||
restorer = WorkflowRunRestore(dry_run=dry_run, workers=workers)
|
||||
if run_id:
|
||||
restorer = WorkflowRunRestore(dry_run=dry_run, workers=workers)
|
||||
results = [restorer.restore_by_run_id(run_id)]
|
||||
else:
|
||||
assert start_from is not None
|
||||
assert end_before is not None
|
||||
results = restorer.restore_batch(
|
||||
parsed_tenant_ids,
|
||||
start_date=start_from,
|
||||
end_date=end_before,
|
||||
limit=limit,
|
||||
)
|
||||
end_time = datetime.datetime.now(datetime.UTC)
|
||||
elapsed = end_time - start_time
|
||||
|
||||
end_time = datetime.datetime.now(datetime.UTC)
|
||||
elapsed = end_time - start_time
|
||||
successes = sum(1 for result in results if result.success)
|
||||
failures = len(results) - successes
|
||||
|
||||
successes = sum(1 for result in results if result.success)
|
||||
failures = len(results) - successes
|
||||
|
||||
if failures == 0:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Restore completed successfully. success={successes} duration={elapsed}",
|
||||
fg="green",
|
||||
if failures == 0:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Restore completed successfully. success={successes} duration={elapsed}",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
)
|
||||
else:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Restore completed with failures. success={successes} failed={failures} duration={elapsed}",
|
||||
fg="red",
|
||||
else:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Restore completed with failures. success={successes} failed={failures} duration={elapsed}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
if workers != 1:
|
||||
click.echo(
|
||||
click.style("--workers is ignored for V2 bundle restore; bundles are processed serially.", fg="yellow")
|
||||
)
|
||||
assert start_from is not None
|
||||
assert end_before is not None
|
||||
bundle_restorer = WorkflowRunBundleArchiveMaintenance(dry_run=dry_run, strict_content_validation=True)
|
||||
summary = bundle_restorer.restore_batch(
|
||||
tenant_ids=parsed_tenant_ids,
|
||||
start_date=start_from,
|
||||
end_date=end_before,
|
||||
limit=limit,
|
||||
)
|
||||
_echo_bundle_archive_operation_summary(summary)
|
||||
return
|
||||
|
||||
|
||||
@click.command(
|
||||
@@ -392,8 +792,20 @@ def restore_workflow_runs(
|
||||
default=None,
|
||||
help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.",
|
||||
)
|
||||
@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of runs to delete.")
|
||||
@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of V2 bundles to delete.")
|
||||
@click.option("--dry-run", is_flag=True, help="Preview without deleting.")
|
||||
@click.option(
|
||||
"--skip-bad-archives",
|
||||
is_flag=True,
|
||||
help="Continue batch deletion when one archive object fails validation.",
|
||||
)
|
||||
@click.option(
|
||||
"--restore-sample-interval",
|
||||
type=int,
|
||||
default=0,
|
||||
show_default=True,
|
||||
help="Run restore dry-run after every N successful deletes; 0 disables restore sampling.",
|
||||
)
|
||||
def delete_archived_workflow_runs(
|
||||
tenant_ids: str | None,
|
||||
run_id: str | None,
|
||||
@@ -401,10 +813,16 @@ def delete_archived_workflow_runs(
|
||||
end_before: datetime.datetime | None,
|
||||
limit: int,
|
||||
dry_run: bool,
|
||||
skip_bad_archives: bool,
|
||||
restore_sample_interval: int,
|
||||
):
|
||||
"""
|
||||
Delete archived workflow runs from the database.
|
||||
|
||||
Batch delete uses V2 bundle metadata and validates object existence, manifest schema, object size, checksum, row
|
||||
counts, and source/archive content checksums before deleting source rows. `--run-id` keeps the V1 per-run path.
|
||||
"""
|
||||
from services.retention.workflow_run.bundle_archive_maintenance import WorkflowRunBundleArchiveMaintenance
|
||||
from services.retention.workflow_run.delete_archived_workflow_run import ArchivedWorkflowRunDeletion
|
||||
|
||||
parsed_tenant_ids = None
|
||||
@@ -417,6 +835,8 @@ def delete_archived_workflow_runs(
|
||||
raise click.UsageError("--start-from and --end-before must be provided together.")
|
||||
if run_id is None and (start_from is None or end_before is None):
|
||||
raise click.UsageError("--start-from and --end-before are required for batch delete.")
|
||||
if restore_sample_interval < 0:
|
||||
raise click.BadParameter("restore-sample-interval must be >= 0")
|
||||
|
||||
start_time = datetime.datetime.now(datetime.UTC)
|
||||
target_desc = f"workflow run {run_id}" if run_id else "workflow runs"
|
||||
@@ -427,56 +847,85 @@ def delete_archived_workflow_runs(
|
||||
)
|
||||
)
|
||||
|
||||
deleter = ArchivedWorkflowRunDeletion(dry_run=dry_run)
|
||||
if run_id:
|
||||
results = [deleter.delete_by_run_id(run_id)]
|
||||
else:
|
||||
assert start_from is not None
|
||||
assert end_before is not None
|
||||
results = deleter.delete_batch(
|
||||
parsed_tenant_ids,
|
||||
start_date=start_from,
|
||||
end_date=end_before,
|
||||
limit=limit,
|
||||
deleter = ArchivedWorkflowRunDeletion(
|
||||
dry_run=dry_run,
|
||||
skip_bad_archives=skip_bad_archives,
|
||||
restore_sample_interval=restore_sample_interval,
|
||||
)
|
||||
results = [deleter.delete_by_run_id(run_id)]
|
||||
for result in results:
|
||||
if result.success:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"{'[DRY RUN] Would delete' if dry_run else 'Deleted'} "
|
||||
f"workflow run {result.run_id} (tenant={result.tenant_id}, "
|
||||
f"archive_key={result.archive_key}, counts={result.validated_counts})",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
if result.restore_sampled:
|
||||
sample_status = "passed" if result.restore_sample_success else "failed"
|
||||
click.echo(
|
||||
click.style(
|
||||
f" restore dry-run sample {sample_status} for workflow run {result.run_id}",
|
||||
fg="green" if result.restore_sample_success else "red",
|
||||
)
|
||||
)
|
||||
else:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Failed to delete workflow run {result.run_id}: {result.error}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
click.echo(
|
||||
click.style(
|
||||
" runbook: pause this delete window, verify archive storage object and manifest/checksum, "
|
||||
"retry the same run after fixing storage or DB drift, or rerun with --skip-bad-archives "
|
||||
"to quarantine this run and continue the batch.",
|
||||
fg="yellow",
|
||||
)
|
||||
)
|
||||
|
||||
for result in results:
|
||||
if result.success:
|
||||
end_time = datetime.datetime.now(datetime.UTC)
|
||||
elapsed = end_time - start_time
|
||||
|
||||
successes = sum(1 for result in results if result.success)
|
||||
failures = len(results) - successes
|
||||
|
||||
if failures == 0:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"{'[DRY RUN] Would delete' if dry_run else 'Deleted'} "
|
||||
f"workflow run {result.run_id} (tenant={result.tenant_id})",
|
||||
f"Delete completed successfully. success={successes} duration={elapsed}",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Failed to delete workflow run {result.run_id}: {result.error}",
|
||||
f"Delete completed with failures. success={successes} failed={failures} duration={elapsed}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
end_time = datetime.datetime.now(datetime.UTC)
|
||||
elapsed = end_time - start_time
|
||||
|
||||
successes = sum(1 for result in results if result.success)
|
||||
failures = len(results) - successes
|
||||
|
||||
if failures == 0:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Delete completed successfully. success={successes} duration={elapsed}",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
else:
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Delete completed with failures. success={successes} failed={failures} duration={elapsed}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
if restore_sample_interval:
|
||||
click.echo(click.style("--restore-sample-interval is ignored for V2 bundle delete.", fg="yellow"))
|
||||
assert start_from is not None
|
||||
assert end_before is not None
|
||||
bundle_deleter = WorkflowRunBundleArchiveMaintenance(
|
||||
dry_run=dry_run,
|
||||
strict_content_validation=True,
|
||||
stop_on_error=not skip_bad_archives,
|
||||
)
|
||||
summary = bundle_deleter.delete_batch(
|
||||
tenant_ids=parsed_tenant_ids,
|
||||
start_date=start_from,
|
||||
end_date=end_before,
|
||||
limit=limit,
|
||||
)
|
||||
_echo_bundle_archive_operation_summary(summary)
|
||||
|
||||
|
||||
def _find_orphaned_draft_variables(batch_size: int = 1000) -> list[str]:
|
||||
|
||||
@@ -2,7 +2,6 @@ from typing import Literal, Protocol, cast
|
||||
from urllib.parse import quote_plus, urlunparse
|
||||
|
||||
from pydantic import AliasChoices, Field
|
||||
from pydantic.types import NonNegativeInt
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
@@ -71,24 +70,6 @@ class RedisPubSubConfig(BaseSettings):
|
||||
default=600,
|
||||
)
|
||||
|
||||
PUBSUB_LISTENER_JOIN_TIMEOUT_MS: NonNegativeInt = Field(
|
||||
validation_alias=AliasChoices("EVENT_BUS_LISTENER_JOIN_TIMEOUT_MS", "PUBSUB_LISTENER_JOIN_TIMEOUT_MS"),
|
||||
description=(
|
||||
"Maximum time (milliseconds) that ``Subscription.close()`` waits for its listener thread to "
|
||||
"finish before returning. Bounds the tail latency between a terminal event being delivered to "
|
||||
"an SSE client and the response stream actually closing.\n\n"
|
||||
"The listener thread blocks on a polling read (XREAD BLOCK for streams, get_message timeout "
|
||||
"for pubsub/sharded) with a fixed 1s window, so close() naturally has to wait up to ~1s for "
|
||||
"the thread to notice the subscription was closed. Setting this lower (e.g. 100) lets close() "
|
||||
"return promptly while the daemon listener thread cleans itself up on the next poll "
|
||||
"boundary - safe because the listener holds no critical state and exits within one poll "
|
||||
"window. Setting it higher (e.g. 5000) gives the listener more grace before close() gives up "
|
||||
"and logs a warning. Default 2000ms preserves the pre-change behaviour.\n\n"
|
||||
"Also accepts ENV: EVENT_BUS_LISTENER_JOIN_TIMEOUT_MS."
|
||||
),
|
||||
default=2000,
|
||||
)
|
||||
|
||||
def _build_default_pubsub_url(self) -> str:
|
||||
defaults = _redis_defaults(self)
|
||||
if not defaults.REDIS_HOST or not defaults.REDIS_PORT:
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from services.enterprise import rbac_service as enterprise_rbac_service
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from services.app_service import AppListBaseParams
|
||||
from services.enterprise.rbac_service import MyPermissionsResponse
|
||||
|
||||
# Permission keys (dot-notation, from MyPermissionsResponse) that grant
|
||||
# list/preview access to an app. Keep this the single source of truth for both
|
||||
# the console and OpenAPI app-list endpoints.
|
||||
APP_LIST_PERMISSION_KEYS: frozenset[str] = frozenset({"app.preview", "app.acl.preview", "app.full_access"})
|
||||
|
||||
# Workspace permission key that lets a caller see apps they maintain even when
|
||||
# those apps are not in their preview whitelist.
|
||||
_MANAGE_OWN_APPS_PERMISSION_KEY = "app.create_and_management"
|
||||
|
||||
|
||||
def has_app_list_permission(permission_keys: Sequence[str]) -> bool:
|
||||
"""Return True if any of ``permission_keys`` grants app list/preview access."""
|
||||
return any(permission_key in APP_LIST_PERMISSION_KEYS for permission_key in permission_keys)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AppAccessFilter:
|
||||
"""Resolved RBAC visibility for app list/read endpoints.
|
||||
|
||||
``accessible_app_ids`` of ``None`` means the caller can see every app in the
|
||||
workspace (unrestricted). Otherwise it is the exact set of app ids the
|
||||
caller may preview; combined with ``can_manage_own_apps`` it also covers
|
||||
apps the caller maintains.
|
||||
"""
|
||||
|
||||
accessible_app_ids: set[str] | None
|
||||
can_manage_own_apps: bool
|
||||
|
||||
@classmethod
|
||||
def unrestricted(cls) -> AppAccessFilter:
|
||||
"""Filter that imposes no restriction (RBAC disabled / not applicable)."""
|
||||
return cls(accessible_app_ids=None, can_manage_own_apps=False)
|
||||
|
||||
def is_app_accessible(self, app_id: str, maintainer: str | None, account_id: str) -> bool:
|
||||
"""Whether a single app is visible to the caller under this filter.
|
||||
|
||||
Mirrors the service-layer query gate: an app is visible when the filter
|
||||
is unrestricted, the app id is whitelisted, or the caller maintains it
|
||||
and holds ``app.create_and_management``.
|
||||
"""
|
||||
if self.accessible_app_ids is None:
|
||||
return True
|
||||
if app_id in self.accessible_app_ids:
|
||||
return True
|
||||
return self.can_manage_own_apps and maintainer is not None and maintainer == account_id
|
||||
|
||||
def apply_to_params(self, params: AppListBaseParams) -> None:
|
||||
if self.accessible_app_ids is None:
|
||||
return
|
||||
params.accessible_app_ids = sorted(self.accessible_app_ids)
|
||||
params.include_own_apps = self.can_manage_own_apps
|
||||
|
||||
|
||||
def resolve_app_access_filter(
|
||||
tenant_id: str,
|
||||
account_id: str,
|
||||
*,
|
||||
permissions: MyPermissionsResponse | None = None,
|
||||
) -> AppAccessFilter:
|
||||
"""Compute the RBAC app-access filter for ``account_id`` in ``tenant_id``.
|
||||
|
||||
Pass ``permissions`` when the caller has already fetched the snapshot (the
|
||||
console controller reuses it for per-app permission keys) to avoid a second
|
||||
inner-API round trip; otherwise it is fetched here.
|
||||
"""
|
||||
if permissions is None:
|
||||
permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id)
|
||||
whitelist_scope = enterprise_rbac_service.RBACService.AppAccess.whitelist_resources(tenant_id, account_id)
|
||||
|
||||
can_manage_own_apps = _MANAGE_OWN_APPS_PERMISSION_KEY in permissions.workspace.permission_keys
|
||||
has_default_preview = has_app_list_permission(permissions.app.default_permission_keys) or has_app_list_permission(
|
||||
permissions.workspace.permission_keys
|
||||
)
|
||||
|
||||
permission_app_ids: set[str] | None = None
|
||||
if not has_default_preview:
|
||||
# Collect apps the caller can preview via per-app permission overrides.
|
||||
permission_app_ids = {
|
||||
override.resource_id
|
||||
for override in permissions.app.overrides
|
||||
if has_app_list_permission(override.permission_keys)
|
||||
}
|
||||
|
||||
accessible_app_ids: set[str] | None
|
||||
if getattr(whitelist_scope, "unrestricted", False):
|
||||
accessible_app_ids = permission_app_ids
|
||||
else:
|
||||
accessible_app_ids = set(whitelist_scope.resource_ids)
|
||||
if permission_app_ids is not None:
|
||||
accessible_app_ids |= permission_app_ids
|
||||
elif has_default_preview:
|
||||
# Default preview overrides the whitelist restriction.
|
||||
accessible_app_ids = None
|
||||
|
||||
return AppAccessFilter(accessible_app_ids=accessible_app_ids, can_manage_own_apps=can_manage_own_apps)
|
||||
@@ -1,23 +1,3 @@
|
||||
"""Shared decorator utilities for Dify controller layers.
|
||||
|
||||
This module provides decorators that are not tied to any single API group (e.g.
|
||||
console, inner, service). Currently it exposes the RBAC permission gate, which
|
||||
can be applied to any blueprint.
|
||||
|
||||
Key exports
|
||||
-----------
|
||||
``rbac_permission_required`` – decorator that enforces enterprise RBAC access
|
||||
control. When ``RBAC_ENABLED`` is ``False`` it is a no-op.
|
||||
|
||||
``RBACPermission``, ``RBACResourceScope`` – re-exported from ``core.rbac`` so
|
||||
callers only need a single import site.
|
||||
|
||||
Private helpers
|
||||
---------------
|
||||
``_extract_resource_id``, ``_is_resource_owned_by_current_user`` – kept module-
|
||||
private but accessible via the module namespace for unit-test patching.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
|
||||
@@ -32,7 +12,57 @@ from models.dataset import Dataset
|
||||
from models.model import App
|
||||
from services.enterprise.rbac_service import RBACService
|
||||
|
||||
__all__ = ["RBACPermission", "RBACResourceScope", "rbac_permission_required"]
|
||||
__all__ = ["RBACPermission", "RBACResourceScope", "enforce_rbac_access", "rbac_permission_required"]
|
||||
|
||||
|
||||
def enforce_rbac_access(
|
||||
*,
|
||||
tenant_id: str,
|
||||
account_id: str,
|
||||
resource_type: RBACResourceScope,
|
||||
scene: RBACPermission,
|
||||
resource_required: bool = True,
|
||||
path_args: dict[str, object] | None = None,
|
||||
) -> None:
|
||||
"""Enforce enterprise RBAC for an explicit account/tenant pair.
|
||||
|
||||
This is the flask-login-independent core of the RBAC gate so it can run
|
||||
inside request-handling layers that resolve the caller themselves (e.g. the
|
||||
openapi auth pipeline, which has the account on ``AuthData`` before
|
||||
flask-login is mounted).
|
||||
|
||||
No-op when ``RBAC_ENABLED`` is ``False``. For resource-scoped checks the
|
||||
resource ID is taken from ``path_args`` merged with ``request.view_args``;
|
||||
resource ownership short-circuits the check. Raises ``Forbidden`` when
|
||||
access is denied. For workspace-level checks pass ``resource_required=False``
|
||||
so the RBAC request omits ``resource_id``.
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant the access is evaluated against.
|
||||
account_id: The account requesting access.
|
||||
resource_type: The :class:`RBACResourceScope` member (app/dataset/workspace).
|
||||
scene: The :class:`RBACPermission` permission point, e.g. ``RBACPermission.APP_DELETE``.
|
||||
resource_required: Whether a concrete resource ID is required.
|
||||
path_args: Extra path arguments to merge with ``request.view_args``.
|
||||
"""
|
||||
if not dify_config.RBAC_ENABLED:
|
||||
return
|
||||
|
||||
check_resource_type = None if resource_type == RBACResourceScope.WORKSPACE else resource_type
|
||||
resource_id = None
|
||||
if resource_required and check_resource_type:
|
||||
resource_id = _extract_resource_id(resource_type, path_args)
|
||||
if _is_resource_owned_by_current_user(tenant_id, account_id, resource_type, resource_id):
|
||||
return
|
||||
allowed = RBACService.CheckAccess.check(
|
||||
tenant_id,
|
||||
account_id,
|
||||
scene=scene,
|
||||
resource_type=check_resource_type,
|
||||
resource_id=resource_id,
|
||||
)
|
||||
if not allowed:
|
||||
raise Forbidden()
|
||||
|
||||
|
||||
def rbac_permission_required[**P, R](
|
||||
@@ -41,14 +71,12 @@ def rbac_permission_required[**P, R](
|
||||
*,
|
||||
resource_required: bool = True,
|
||||
) -> Callable[[Callable[P, R]], Callable[P, R]]:
|
||||
"""Check enterprise RBAC permissions for the current user.
|
||||
"""Check enterprise RBAC permissions for the current flask-login user.
|
||||
|
||||
When ``RBAC_ENABLED`` is ``False`` the decorator is a no-op and the
|
||||
request passes through unchanged. When enabled it extracts the resource ID
|
||||
from ``request.view_args`` for resource-scoped checks, calls the RBAC
|
||||
service ``check-access`` endpoint, and raises ``Forbidden`` if the access
|
||||
is denied. For workspace-level checks, set ``resource_required=False`` so
|
||||
the RBAC request omits ``resource_id``.
|
||||
request passes through unchanged. When enabled it resolves the current
|
||||
account/tenant and delegates to :func:`enforce_rbac_access`, raising
|
||||
``Forbidden`` if access is denied.
|
||||
|
||||
Args:
|
||||
resource_type: The :class:`RBACResourceScope` member (app/dataset/workspace).
|
||||
@@ -63,23 +91,14 @@ def rbac_permission_required[**P, R](
|
||||
return view(*args, **kwargs)
|
||||
|
||||
current_user, current_tenant_id = current_account_with_tenant()
|
||||
check_resource_type = None if resource_type == RBACResourceScope.WORKSPACE else resource_type
|
||||
resource_id = None
|
||||
if resource_required and check_resource_type:
|
||||
resource_id = _extract_resource_id(resource_type, kwargs)
|
||||
if _is_resource_owned_by_current_user(current_tenant_id, current_user.id, resource_type, resource_id):
|
||||
return view(*args, **kwargs)
|
||||
allowed = RBACService.CheckAccess.check(
|
||||
current_tenant_id,
|
||||
current_user.id,
|
||||
enforce_rbac_access(
|
||||
tenant_id=current_tenant_id,
|
||||
account_id=current_user.id,
|
||||
resource_type=resource_type,
|
||||
scene=scene,
|
||||
resource_type=check_resource_type,
|
||||
resource_id=resource_id,
|
||||
resource_required=resource_required,
|
||||
path_args=kwargs,
|
||||
)
|
||||
|
||||
if not allowed:
|
||||
raise Forbidden()
|
||||
|
||||
return view(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
@@ -3,16 +3,17 @@ from uuid import UUID
|
||||
from flask import abort, request
|
||||
from flask_restx import Resource
|
||||
from pydantic import AliasChoices, BaseModel, Field, field_validator
|
||||
from sqlalchemy import func, select
|
||||
|
||||
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
|
||||
from controllers.console import console_ns
|
||||
from controllers.console.agent.app_helpers import resolve_agent_app_model
|
||||
from controllers.console.apikey import ApiKeyItem, ApiKeyList, BaseApiKeyListResource, BaseApiKeyResource
|
||||
from controllers.console.app.app import (
|
||||
AppDetailWithSite as GenericAppDetailWithSite,
|
||||
)
|
||||
from controllers.console.app.app import (
|
||||
AppListQuery,
|
||||
CopyAppPayload,
|
||||
_normalize_app_list_query_args,
|
||||
)
|
||||
from controllers.console.app.app import (
|
||||
@@ -25,9 +26,13 @@ from controllers.console.app.app import (
|
||||
UpdateAppPayload as GenericUpdateAppPayload,
|
||||
)
|
||||
from controllers.console.wraps import (
|
||||
RBACPermission,
|
||||
RBACResourceScope,
|
||||
account_initialization_required,
|
||||
edit_permission_required,
|
||||
enterprise_license_required,
|
||||
is_admin_or_owner_required,
|
||||
rbac_permission_required,
|
||||
setup_required,
|
||||
with_current_tenant_id,
|
||||
with_current_user,
|
||||
@@ -36,6 +41,7 @@ from extensions.ext_database import db
|
||||
from fields.agent_fields import (
|
||||
AgentConfigSnapshotDetailResponse,
|
||||
AgentConfigSnapshotListResponse,
|
||||
AgentConfigSnapshotRestoreResponse,
|
||||
AgentInviteOptionsResponse,
|
||||
AgentLogListResponse,
|
||||
AgentLogMessageListResponse,
|
||||
@@ -48,7 +54,8 @@ from libs.datetime_utils import parse_time_range
|
||||
from libs.helper import dump_response
|
||||
from libs.login import login_required
|
||||
from models import Account
|
||||
from models.model import IconType
|
||||
from models.enums import ApiTokenType
|
||||
from models.model import ApiToken, App, IconType
|
||||
from services.agent.errors import AgentNotFoundError
|
||||
from services.agent.observability_service import (
|
||||
AgentLogQueryParams,
|
||||
@@ -102,6 +109,46 @@ class AgentAppUpdatePayload(GenericUpdateAppPayload):
|
||||
return role
|
||||
|
||||
|
||||
class AgentAppCopyPayload(BaseModel):
|
||||
name: str | None = Field(default=None, description="Name for the copied agent")
|
||||
description: str | None = Field(default=None, description="Description for the copied agent", max_length=400)
|
||||
role: str | None = Field(default=None, description="Role for the copied agent", max_length=255)
|
||||
icon_type: IconType | None = Field(default=None, description="Icon type")
|
||||
icon: str | None = Field(default=None, description="Icon")
|
||||
icon_background: str | None = Field(default=None, description="Icon background color")
|
||||
|
||||
@field_validator("role")
|
||||
@classmethod
|
||||
def validate_role(cls, value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
role = value.strip()
|
||||
if not role:
|
||||
raise ValueError("Agent role is required when provided.")
|
||||
return role
|
||||
|
||||
|
||||
class AgentApiStatusPayload(BaseModel):
|
||||
enable_api: bool = Field(..., description="Enable or disable Agent service API")
|
||||
|
||||
|
||||
class AgentApiAccessResponse(BaseModel):
|
||||
enabled: bool
|
||||
service_api_base_url: str
|
||||
streaming_only: bool = True
|
||||
chat_endpoint: str
|
||||
stop_endpoint: str
|
||||
conversations_endpoint: str
|
||||
messages_endpoint: str
|
||||
files_upload_endpoint: str
|
||||
parameters_endpoint: str
|
||||
info_endpoint: str
|
||||
meta_endpoint: str
|
||||
api_rpm: int
|
||||
api_rph: int
|
||||
api_key_count: int
|
||||
|
||||
|
||||
class AgentAppPublishedReferenceResponse(BaseModel):
|
||||
app_id: str
|
||||
app_name: str
|
||||
@@ -185,6 +232,7 @@ class AgentStatisticsQuery(BaseModel):
|
||||
|
||||
class AgentAppPartial(GenericAppPartial):
|
||||
app_id: str | None = None
|
||||
debug_conversation_id: str | None = None
|
||||
role: str | None = None
|
||||
active_config_is_published: bool = False
|
||||
published_reference_count: int = 0
|
||||
@@ -193,10 +241,15 @@ class AgentAppPartial(GenericAppPartial):
|
||||
|
||||
class AgentAppDetailWithSite(GenericAppDetailWithSite):
|
||||
app_id: str | None = None
|
||||
debug_conversation_id: str | None = None
|
||||
role: str | None = None
|
||||
active_config_is_published: bool = False
|
||||
|
||||
|
||||
class AgentDebugConversationRefreshResponse(BaseModel):
|
||||
debug_conversation_id: str
|
||||
|
||||
|
||||
class AgentAppPagination(GenericAppPagination):
|
||||
data: list[AgentAppPartial] = Field( # type: ignore[assignment] # pyrefly: ignore[bad-override-mutable-attribute]
|
||||
validation_alias=AliasChoices("items", "data")
|
||||
@@ -207,7 +260,8 @@ register_schema_models(
|
||||
console_ns,
|
||||
AgentAppCreatePayload,
|
||||
AgentAppUpdatePayload,
|
||||
CopyAppPayload,
|
||||
AgentAppCopyPayload,
|
||||
AgentApiStatusPayload,
|
||||
AgentInviteOptionsQuery,
|
||||
AgentLogsQuery,
|
||||
AgentStatisticsQuery,
|
||||
@@ -218,11 +272,14 @@ register_schema_models(
|
||||
register_response_schema_models(
|
||||
console_ns,
|
||||
AgentAppPagination,
|
||||
AgentApiAccessResponse,
|
||||
AgentAppPublishedReferenceResponse,
|
||||
AgentAppDetailWithSite,
|
||||
AgentAppPartial,
|
||||
AgentDebugConversationRefreshResponse,
|
||||
AgentConfigSnapshotDetailResponse,
|
||||
AgentConfigSnapshotListResponse,
|
||||
AgentConfigSnapshotRestoreResponse,
|
||||
AgentInviteOptionsResponse,
|
||||
AgentLogListResponse,
|
||||
AgentLogMessageListResponse,
|
||||
@@ -237,7 +294,7 @@ def _agent_roster_service() -> AgentRosterService:
|
||||
return AgentRosterService(db.session)
|
||||
|
||||
|
||||
def _serialize_agent_app_detail(app_model) -> dict:
|
||||
def _serialize_agent_app_detail(app_model, *, current_user: Account) -> dict:
|
||||
"""Serialize an Agent App detail using roster-only DTOs.
|
||||
|
||||
`/agent` responses are roster-shaped rather than raw app-shaped: `id`
|
||||
@@ -260,6 +317,11 @@ def _serialize_agent_app_detail(app_model) -> dict:
|
||||
payload.pop("bound_agent_id", None)
|
||||
payload["app_id"] = str(app_model.id)
|
||||
payload["id"] = agent.id
|
||||
payload["debug_conversation_id"] = roster_service.get_or_create_agent_app_debug_conversation_id(
|
||||
tenant_id=app_model.tenant_id,
|
||||
agent_id=agent.id,
|
||||
account_id=current_user.id,
|
||||
)
|
||||
payload["role"] = agent.role or ""
|
||||
payload["active_config_is_published"] = roster_service.active_config_is_published(
|
||||
tenant_id=app_model.tenant_id,
|
||||
@@ -268,7 +330,7 @@ def _serialize_agent_app_detail(app_model) -> dict:
|
||||
return payload
|
||||
|
||||
|
||||
def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str) -> dict:
|
||||
def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str, current_user: Account) -> dict:
|
||||
"""Serialize Agent App lists with roster-shaped items.
|
||||
|
||||
Each item starts from the shared App list shape, then drops
|
||||
@@ -291,6 +353,11 @@ def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str) -> dict:
|
||||
tenant_id=tenant_id,
|
||||
agent_ids=[agent.id for agent in agents_by_app_id.values()],
|
||||
)
|
||||
debug_conversation_ids_by_agent_id = roster_service.load_or_create_agent_app_debug_conversation_ids_by_agent_id(
|
||||
tenant_id=tenant_id,
|
||||
agents=list(agents_by_app_id.values()),
|
||||
account_id=current_user.id,
|
||||
)
|
||||
payload = AgentAppPagination.model_validate(app_pagination, from_attributes=True).model_dump(mode="json")
|
||||
for item in payload["data"]:
|
||||
app_id = item["id"]
|
||||
@@ -299,6 +366,7 @@ def _serialize_agent_app_pagination(app_pagination, *, tenant_id: str) -> dict:
|
||||
if agent:
|
||||
item["app_id"] = app_id
|
||||
item["id"] = agent.id
|
||||
item["debug_conversation_id"] = debug_conversation_ids_by_agent_id.get(agent.id)
|
||||
item["role"] = agent.role or ""
|
||||
item["active_config_is_published"] = active_config_is_published_by_agent_id.get(agent.id, False)
|
||||
published_references = published_references_by_agent_id.get(agent.id, [])
|
||||
@@ -323,6 +391,38 @@ def _resolve_agent_app_model(*, tenant_id: str, agent_id: UUID):
|
||||
return resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
|
||||
|
||||
def _agent_api_key_count(app_id: str) -> int:
|
||||
return (
|
||||
db.session.scalar(
|
||||
select(func.count(ApiToken.id)).where(
|
||||
ApiToken.type == ApiTokenType.APP,
|
||||
ApiToken.app_id == app_id,
|
||||
)
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
|
||||
def _serialize_agent_api_access(app_model: App) -> dict:
|
||||
base_url = app_model.api_base_url
|
||||
response = AgentApiAccessResponse(
|
||||
enabled=bool(app_model.enable_api),
|
||||
service_api_base_url=base_url,
|
||||
chat_endpoint=f"{base_url}/chat-messages",
|
||||
stop_endpoint=f"{base_url}/chat-messages/{{task_id}}/stop",
|
||||
conversations_endpoint=f"{base_url}/conversations",
|
||||
messages_endpoint=f"{base_url}/messages",
|
||||
files_upload_endpoint=f"{base_url}/files/upload",
|
||||
parameters_endpoint=f"{base_url}/parameters",
|
||||
info_endpoint=f"{base_url}/info",
|
||||
meta_endpoint=f"{base_url}/meta",
|
||||
api_rpm=app_model.api_rpm or 0,
|
||||
api_rph=app_model.api_rph or 0,
|
||||
api_key_count=_agent_api_key_count(str(app_model.id)),
|
||||
)
|
||||
return response.model_dump(mode="json")
|
||||
|
||||
|
||||
def _agent_observability_service() -> AgentObservabilityService:
|
||||
return AgentObservabilityService(db.session)
|
||||
|
||||
@@ -374,7 +474,11 @@ class AgentAppListApi(Resource):
|
||||
empty = AgentAppPagination(page=args.page, limit=args.limit, total=0, has_more=False, data=[])
|
||||
return empty.model_dump(mode="json")
|
||||
|
||||
return _serialize_agent_app_pagination(app_pagination, tenant_id=current_tenant_id)
|
||||
return _serialize_agent_app_pagination(
|
||||
app_pagination,
|
||||
tenant_id=current_tenant_id,
|
||||
current_user=current_user,
|
||||
)
|
||||
|
||||
@console_ns.expect(console_ns.models[AgentAppCreatePayload.__name__])
|
||||
@console_ns.response(201, "Agent app created successfully", console_ns.models[AgentAppDetailWithSite.__name__])
|
||||
@@ -399,7 +503,7 @@ class AgentAppListApi(Resource):
|
||||
)
|
||||
|
||||
app = AppService().create_app(current_tenant_id, params, current_user)
|
||||
return _serialize_agent_app_detail(app), 201
|
||||
return _serialize_agent_app_detail(app, current_user=current_user), 201
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>")
|
||||
@@ -409,10 +513,11 @@ class AgentAppApi(Resource):
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@enterprise_license_required
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str, agent_id: UUID):
|
||||
def get(self, tenant_id: str, current_user: Account, agent_id: UUID):
|
||||
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
return _serialize_agent_app_detail(app_model)
|
||||
return _serialize_agent_app_detail(app_model, current_user=current_user)
|
||||
|
||||
@console_ns.expect(console_ns.models[AgentAppUpdatePayload.__name__])
|
||||
@console_ns.response(200, "Agent app updated successfully", console_ns.models[AgentAppDetailWithSite.__name__])
|
||||
@@ -422,8 +527,9 @@ class AgentAppApi(Resource):
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def put(self, tenant_id: str, agent_id: UUID):
|
||||
def put(self, tenant_id: str, current_user: Account, agent_id: UUID):
|
||||
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
args = AgentAppUpdatePayload.model_validate(console_ns.payload)
|
||||
args_dict: AppService.ArgsDict = {
|
||||
@@ -437,7 +543,7 @@ class AgentAppApi(Resource):
|
||||
"role": args.role,
|
||||
}
|
||||
updated = AppService().update_app(app_model, args_dict)
|
||||
return _serialize_agent_app_detail(updated)
|
||||
return _serialize_agent_app_detail(updated, current_user=current_user)
|
||||
|
||||
@console_ns.response(204, "Agent app deleted successfully")
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@@ -452,9 +558,34 @@ class AgentAppApi(Resource):
|
||||
return "", 204
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/debug-conversation/refresh")
|
||||
class AgentDebugConversationRefreshApi(Resource):
|
||||
@console_ns.response(
|
||||
200,
|
||||
"Agent debug conversation refreshed",
|
||||
console_ns.models[AgentDebugConversationRefreshResponse.__name__],
|
||||
)
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
|
||||
debug_conversation_id = _agent_roster_service().refresh_agent_app_debug_conversation_id(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=str(agent_id),
|
||||
account_id=current_user.id,
|
||||
)
|
||||
return AgentDebugConversationRefreshResponse(debug_conversation_id=debug_conversation_id).model_dump(
|
||||
mode="json"
|
||||
)
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/copy")
|
||||
class AgentAppCopyApi(Resource):
|
||||
@console_ns.expect(console_ns.models[CopyAppPayload.__name__])
|
||||
@console_ns.expect(console_ns.models[AgentAppCopyPayload.__name__])
|
||||
@console_ns.response(201, "Agent app copied successfully", console_ns.models[AgentAppDetailWithSite.__name__])
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@console_ns.response(400, "Invalid request parameters")
|
||||
@@ -465,18 +596,88 @@ class AgentAppCopyApi(Resource):
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, current_user: Account, agent_id: UUID):
|
||||
args = CopyAppPayload.model_validate(console_ns.payload or {})
|
||||
args = AgentAppCopyPayload.model_validate(console_ns.payload or {})
|
||||
copied_app = _agent_roster_service().duplicate_agent_app(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=str(agent_id),
|
||||
account=current_user,
|
||||
name=args.name,
|
||||
description=args.description,
|
||||
role=args.role,
|
||||
icon_type=args.icon_type,
|
||||
icon=args.icon,
|
||||
icon_background=args.icon_background,
|
||||
)
|
||||
return _serialize_agent_app_detail(copied_app), 201
|
||||
return _serialize_agent_app_detail(copied_app, current_user=current_user), 201
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/api-access")
|
||||
class AgentApiAccessApi(Resource):
|
||||
@console_ns.response(200, "Agent service API access", console_ns.models[AgentApiAccessResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str, agent_id: UUID):
|
||||
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
return _serialize_agent_api_access(app_model)
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/api-enable")
|
||||
class AgentApiStatusApi(Resource):
|
||||
@console_ns.expect(console_ns.models[AgentApiStatusPayload.__name__])
|
||||
@console_ns.response(200, "Agent service API status updated", console_ns.models[AgentApiAccessResponse.__name__])
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@setup_required
|
||||
@login_required
|
||||
@is_admin_or_owner_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, agent_id: UUID):
|
||||
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
args = AgentApiStatusPayload.model_validate(console_ns.payload)
|
||||
app_model = AppService().update_app_api_status(app_model, args.enable_api)
|
||||
return _serialize_agent_api_access(app_model)
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/api-keys")
|
||||
class AgentApiKeyListApi(BaseApiKeyListResource):
|
||||
resource_type = ApiTokenType.APP
|
||||
resource_model = App
|
||||
resource_id_field = "app_id"
|
||||
token_prefix = "app-"
|
||||
|
||||
@console_ns.response(200, "Agent service API keys", console_ns.models[ApiKeyList.__name__])
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str, agent_id: UUID) -> dict[str, object]:
|
||||
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
return dump_response(ApiKeyList, self._get_api_key_list(str(app_model.id), tenant_id))
|
||||
|
||||
@console_ns.response(201, "Agent service API key created", console_ns.models[ApiKeyItem.__name__])
|
||||
@console_ns.response(400, "Maximum keys exceeded")
|
||||
@with_current_tenant_id
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
|
||||
def post(self, tenant_id: str, agent_id: UUID) -> tuple[dict[str, object], int]:
|
||||
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
return dump_response(ApiKeyItem, self._create_api_key(str(app_model.id), tenant_id)), 201
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/api-keys/<uuid:api_key_id>")
|
||||
class AgentApiKeyApi(BaseApiKeyResource):
|
||||
resource_type = ApiTokenType.APP
|
||||
resource_model = App
|
||||
resource_id_field = "app_id"
|
||||
|
||||
@console_ns.response(204, "Agent service API key deleted")
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_RELEASE_AND_VERSION)
|
||||
def delete(self, tenant_id: str, current_user: Account, agent_id: UUID, api_key_id: UUID) -> tuple[str, int]:
|
||||
app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
self._delete_api_key(str(app_model.id), str(api_key_id), tenant_id, current_user)
|
||||
return "", 204
|
||||
|
||||
|
||||
@console_ns.route("/agent/invite-options")
|
||||
@@ -649,3 +850,24 @@ class AgentRosterVersionDetailApi(Resource):
|
||||
version_id=str(version_id),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/versions/<uuid:version_id>/restore")
|
||||
class AgentRosterVersionRestoreApi(Resource):
|
||||
@console_ns.response(200, "Agent version restored", console_ns.models[AgentConfigSnapshotRestoreResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, current_user: Account, agent_id: UUID, version_id: UUID):
|
||||
return dump_response(
|
||||
AgentConfigSnapshotRestoreResponse,
|
||||
_agent_roster_service().restore_agent_version(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=str(agent_id),
|
||||
version_id=str(version_id),
|
||||
account_id=current_user.id,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -10,8 +10,12 @@ backend — drive data lives in the API's own DB/storage, served straight from
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from flask import Response
|
||||
from flask_restx import Resource
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
@@ -49,6 +53,10 @@ class AgentDriveFileByAgentQuery(BaseModel):
|
||||
key: str = Field(min_length=1, description="Drive key, e.g. tender-analyzer/SKILL.md")
|
||||
|
||||
|
||||
class AgentDriveSkillInspectQuery(BaseModel):
|
||||
node_id: str | None = Field(default=None, description="Workflow node ID (workflow composer variant)")
|
||||
|
||||
|
||||
class AgentDriveItemResponse(ResponseModel):
|
||||
key: str
|
||||
size: int | None = None
|
||||
@@ -56,12 +64,63 @@ class AgentDriveItemResponse(ResponseModel):
|
||||
hash: str | None = None
|
||||
file_kind: str
|
||||
created_at: int | None = None
|
||||
is_skill: bool | None = None
|
||||
skill_metadata: str | None = None
|
||||
|
||||
|
||||
class AgentDriveListResponse(ResponseModel):
|
||||
items: list[AgentDriveItemResponse] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentDriveSkillItemResponse(ResponseModel):
|
||||
path: str
|
||||
skill_md_key: str
|
||||
archive_key: str | None = None
|
||||
name: str
|
||||
description: str
|
||||
size: int | None = None
|
||||
mime_type: str | None = None
|
||||
hash: str | None = None
|
||||
created_at: int | None = None
|
||||
|
||||
|
||||
class AgentDriveSkillListResponse(ResponseModel):
|
||||
items: list[AgentDriveSkillItemResponse] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentDriveSkillFileResponse(ResponseModel):
|
||||
path: str
|
||||
name: str
|
||||
type: str
|
||||
drive_key: str | None = None
|
||||
available_in_drive: bool
|
||||
|
||||
|
||||
class AgentDriveSkillMarkdownResponse(ResponseModel):
|
||||
key: str
|
||||
size: int | None = None
|
||||
truncated: bool
|
||||
binary: bool
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class AgentDriveSkillInspectResponse(ResponseModel):
|
||||
path: str
|
||||
skill_md_key: str
|
||||
archive_key: str | None = None
|
||||
name: str
|
||||
description: str
|
||||
size: int | None = None
|
||||
mime_type: str | None = None
|
||||
hash: str | None = None
|
||||
created_at: int | None = None
|
||||
source: str
|
||||
files: list[AgentDriveSkillFileResponse] = Field(default_factory=list)
|
||||
file_tree: list[dict[str, Any]] = Field(default_factory=list)
|
||||
skill_md: AgentDriveSkillMarkdownResponse
|
||||
warnings: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentDrivePreviewResponse(ResponseModel):
|
||||
key: str
|
||||
size: int | None = None
|
||||
@@ -75,7 +134,12 @@ class AgentDriveDownloadResponse(ResponseModel):
|
||||
|
||||
|
||||
register_response_schema_models(
|
||||
console_ns, AgentDriveListResponse, AgentDrivePreviewResponse, AgentDriveDownloadResponse
|
||||
console_ns,
|
||||
AgentDriveDownloadResponse,
|
||||
AgentDriveListResponse,
|
||||
AgentDrivePreviewResponse,
|
||||
AgentDriveSkillInspectResponse,
|
||||
AgentDriveSkillListResponse,
|
||||
)
|
||||
|
||||
|
||||
@@ -96,6 +160,13 @@ def _handle(exc: AgentDriveError) -> tuple[dict[str, object], int]:
|
||||
return {"code": exc.code, "message": exc.message}, exc.status_code
|
||||
|
||||
|
||||
def _json_response(data: Mapping[str, Any]):
|
||||
return Response(
|
||||
response=json.dumps(data, ensure_ascii=False, separators=(",", ":")),
|
||||
content_type="application/json; charset=utf-8",
|
||||
)
|
||||
|
||||
|
||||
_WORKFLOW_APP_MODES = [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]
|
||||
|
||||
|
||||
@@ -119,6 +190,49 @@ class AgentDriveListByAgentApi(Resource):
|
||||
return {"items": [{k: v for k, v in item.items() if k != "file_id"} for item in items]}
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/drive/skills")
|
||||
class AgentDriveSkillListByAgentApi(Resource):
|
||||
@console_ns.doc("list_agent_drive_skills_by_agent")
|
||||
@console_ns.doc(description="List drive-backed skills for an Agent App")
|
||||
@console_ns.doc(params={"agent_id": "Agent ID"})
|
||||
@console_ns.response(200, "Drive skills", console_ns.models[AgentDriveSkillListResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str, agent_id: UUID):
|
||||
resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
try:
|
||||
items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=str(agent_id))
|
||||
except AgentDriveError as exc:
|
||||
return _handle(exc)
|
||||
return {"items": items}
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/drive/skills/<path:skill_path>/inspect")
|
||||
class AgentDriveSkillInspectByAgentApi(Resource):
|
||||
@console_ns.doc("inspect_agent_drive_skill_by_agent")
|
||||
@console_ns.doc(description="Inspect one drive-backed skill for slash-menu hover/detail UI")
|
||||
@console_ns.doc(params={"agent_id": "Agent ID", "skill_path": "Skill path/slug, e.g. tender-analyzer"})
|
||||
@console_ns.response(200, "Drive skill inspect view", console_ns.models[AgentDriveSkillInspectResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str, agent_id: UUID, skill_path: str):
|
||||
resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id)
|
||||
try:
|
||||
return _json_response(
|
||||
AgentDriveService().inspect_skill(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=str(agent_id),
|
||||
skill_path=skill_path,
|
||||
)
|
||||
)
|
||||
except AgentDriveError as exc:
|
||||
return _handle(exc)
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/drive/files/preview")
|
||||
class AgentDrivePreviewByAgentApi(Resource):
|
||||
@console_ns.doc("preview_agent_drive_file_by_agent")
|
||||
@@ -182,6 +296,61 @@ class AgentDriveListApi(Resource):
|
||||
return {"items": [{k: v for k, v in item.items() if k != "file_id"} for item in items]}
|
||||
|
||||
|
||||
@console_ns.route("/apps/<uuid:app_id>/agent/drive/skills")
|
||||
class AgentDriveSkillListApi(Resource):
|
||||
@console_ns.doc("list_agent_drive_skills")
|
||||
@console_ns.doc(description="List drive-backed skills for the bound agent")
|
||||
@console_ns.doc(params={"app_id": "Application ID", **query_params_from_model(AgentDriveListQuery)})
|
||||
@console_ns.response(200, "Drive skills", console_ns.models[AgentDriveSkillListResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@get_app_model(mode=_WORKFLOW_APP_MODES)
|
||||
def get(self, app_model: App):
|
||||
query = query_params_from_request(AgentDriveListQuery)
|
||||
agent_id = _resolve_agent_id(app_model, query.node_id)
|
||||
if not agent_id:
|
||||
return _agent_not_bound()
|
||||
try:
|
||||
items = AgentDriveService().list_skills(tenant_id=app_model.tenant_id, agent_id=agent_id)
|
||||
except AgentDriveError as exc:
|
||||
return _handle(exc)
|
||||
return {"items": items}
|
||||
|
||||
|
||||
@console_ns.route("/apps/<uuid:app_id>/agent/drive/skills/<path:skill_path>/inspect")
|
||||
class AgentDriveSkillInspectApi(Resource):
|
||||
@console_ns.doc("inspect_agent_drive_skill")
|
||||
@console_ns.doc(description="Inspect one drive-backed skill for slash-menu hover/detail UI")
|
||||
@console_ns.doc(
|
||||
params={
|
||||
"app_id": "Application ID",
|
||||
"skill_path": "Skill path/slug, e.g. tender-analyzer",
|
||||
**query_params_from_model(AgentDriveSkillInspectQuery),
|
||||
}
|
||||
)
|
||||
@console_ns.response(200, "Drive skill inspect view", console_ns.models[AgentDriveSkillInspectResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@get_app_model(mode=_WORKFLOW_APP_MODES)
|
||||
def get(self, app_model: App, skill_path: str):
|
||||
query = query_params_from_request(AgentDriveSkillInspectQuery)
|
||||
agent_id = _resolve_agent_id(app_model, query.node_id)
|
||||
if not agent_id:
|
||||
return _agent_not_bound()
|
||||
try:
|
||||
return _json_response(
|
||||
AgentDriveService().inspect_skill(
|
||||
tenant_id=app_model.tenant_id,
|
||||
agent_id=agent_id,
|
||||
skill_path=skill_path,
|
||||
)
|
||||
)
|
||||
except AgentDriveError as exc:
|
||||
return _handle(exc)
|
||||
|
||||
|
||||
@console_ns.route("/apps/<uuid:app_id>/agent/drive/files/preview")
|
||||
class AgentDrivePreviewApi(Resource):
|
||||
@console_ns.doc("preview_agent_drive_file")
|
||||
@@ -232,4 +401,8 @@ __all__ = [
|
||||
"AgentDriveListByAgentApi",
|
||||
"AgentDrivePreviewApi",
|
||||
"AgentDrivePreviewByAgentApi",
|
||||
"AgentDriveSkillInspectApi",
|
||||
"AgentDriveSkillInspectByAgentApi",
|
||||
"AgentDriveSkillListApi",
|
||||
"AgentDriveSkillListByAgentApi",
|
||||
]
|
||||
|
||||
@@ -14,6 +14,7 @@ from werkzeug.datastructures import MultiDict
|
||||
from werkzeug.exceptions import BadRequest, NotFound
|
||||
|
||||
from configs import dify_config
|
||||
from controllers.common.app_access import resolve_app_access_filter
|
||||
from controllers.common.fields import RedirectUrlResponse, SimpleResultResponse
|
||||
from controllers.common.helpers import FileInfo
|
||||
from controllers.common.schema import (
|
||||
@@ -78,7 +79,6 @@ _TAG_IDS_BRACKET_PATTERN = re.compile(r"^tag_ids\[(\d+)\]$")
|
||||
_CREATOR_IDS_BRACKET_PATTERN = re.compile(r"^creator_ids\[(\d+)\]$")
|
||||
AppListMode = Literal["completion", "chat", "advanced-chat", "workflow", "agent-chat", "agent", "channel", "all"]
|
||||
DEFAULT_APP_LIST_MODE: AppListMode = "all"
|
||||
APP_LIST_PERMISSION_KEYS = frozenset({"app.preview", "app.acl.preview", "app.full_access"})
|
||||
|
||||
|
||||
class AppListBaseQuery(BaseModel):
|
||||
@@ -167,10 +167,6 @@ def _normalize_app_list_query_args(query_args: MultiDict[str, str]) -> dict[str,
|
||||
return normalized
|
||||
|
||||
|
||||
def _has_app_list_permission(permission_keys: Sequence[str]) -> bool:
|
||||
return any(permission_key in APP_LIST_PERMISSION_KEYS for permission_key in permission_keys)
|
||||
|
||||
|
||||
class CreateAppPayload(BaseModel):
|
||||
name: str = Field(..., min_length=1, description="App name")
|
||||
description: str | None = Field(default=None, description="App description (max 400 chars)", max_length=400)
|
||||
@@ -612,38 +608,12 @@ class AppListApi(Resource):
|
||||
current_user_id,
|
||||
)
|
||||
if dify_config.RBAC_ENABLED:
|
||||
whitelist_scope = enterprise_rbac_service.RBACService.AppAccess.whitelist_resources(
|
||||
access_filter = resolve_app_access_filter(
|
||||
str(current_tenant_id),
|
||||
current_user_id,
|
||||
permissions=permissions,
|
||||
)
|
||||
can_manage_own_apps = "app.create_and_management" in permissions.workspace.permission_keys
|
||||
has_default_preview = _has_app_list_permission(
|
||||
permissions.app.default_permission_keys
|
||||
) or _has_app_list_permission(permissions.workspace.permission_keys)
|
||||
permission_app_ids: set[str] | None = None
|
||||
if not has_default_preview:
|
||||
permission_app_ids = {
|
||||
override.resource_id
|
||||
for override in permissions.app.overrides
|
||||
if _has_app_list_permission(override.permission_keys)
|
||||
}
|
||||
|
||||
if getattr(whitelist_scope, "unrestricted", False):
|
||||
accessible_app_ids = permission_app_ids
|
||||
else:
|
||||
accessible_app_ids = set(whitelist_scope.resource_ids)
|
||||
if permission_app_ids is not None:
|
||||
accessible_app_ids |= permission_app_ids
|
||||
elif has_default_preview:
|
||||
accessible_app_ids = None
|
||||
|
||||
if accessible_app_ids:
|
||||
params.accessible_app_ids = sorted(accessible_app_ids)
|
||||
params.include_own_apps = can_manage_own_apps
|
||||
elif accessible_app_ids is not None and can_manage_own_apps:
|
||||
params.is_created_by_me = True
|
||||
elif accessible_app_ids is not None:
|
||||
params.accessible_app_ids = []
|
||||
access_filter.apply_to_params(params)
|
||||
|
||||
# get app list
|
||||
app_service = AppService()
|
||||
|
||||
@@ -40,12 +40,15 @@ from core.errors.error import (
|
||||
QuotaExceededError,
|
||||
)
|
||||
from core.helper.trace_id_helper import get_external_trace_id
|
||||
from extensions.ext_database import db
|
||||
from graphon.model_runtime.errors.invoke import InvokeError
|
||||
from libs import helper
|
||||
from libs.helper import uuid_value
|
||||
from libs.login import login_required
|
||||
from models import Account
|
||||
from models.model import App, AppMode
|
||||
from services.agent.errors import AgentNotFoundError
|
||||
from services.agent.roster_service import AgentRosterService
|
||||
from services.app_generate_service import AppGenerateService
|
||||
from services.app_task_service import AppTaskService
|
||||
from services.errors.llm import InvokeRateLimitError
|
||||
@@ -191,10 +194,11 @@ class ChatMessageApi(Resource):
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_TEST_AND_RUN)
|
||||
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.AGENT])
|
||||
def post(self, current_user: Account, app_model: App):
|
||||
return _create_chat_message(current_user=current_user, app_model=app_model)
|
||||
def post(self, current_tenant_id: str, current_user: Account, app_model: App):
|
||||
return _create_chat_message(current_tenant_id=current_tenant_id, current_user=current_user, app_model=app_model)
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/chat-messages")
|
||||
@@ -215,7 +219,12 @@ class AgentChatMessageApi(Resource):
|
||||
@with_current_tenant_id
|
||||
def post(self, current_tenant_id: str, current_user: Account, agent_id: UUID):
|
||||
app_model = resolve_agent_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
|
||||
return _create_chat_message(current_user=current_user, app_model=app_model)
|
||||
return _create_chat_message(
|
||||
current_tenant_id=current_tenant_id,
|
||||
current_user=current_user,
|
||||
app_model=app_model,
|
||||
agent_id=str(agent_id),
|
||||
)
|
||||
|
||||
|
||||
@console_ns.route("/apps/<uuid:app_id>/chat-messages/<string:task_id>/stop")
|
||||
@@ -249,11 +258,45 @@ class AgentChatMessageStopApi(Resource):
|
||||
return _stop_chat_message(current_user_id=current_user_id, app_model=app_model, task_id=task_id)
|
||||
|
||||
|
||||
def _create_chat_message(*, current_user: Account, app_model: App):
|
||||
def _resolve_current_user_agent_debug_conversation_id(
|
||||
*, current_tenant_id: str, current_user: Account, app_model: App, agent_id: str | None
|
||||
) -> str:
|
||||
roster_service = AgentRosterService(db.session)
|
||||
if agent_id:
|
||||
return roster_service.get_or_create_agent_app_debug_conversation_id(
|
||||
tenant_id=current_tenant_id,
|
||||
agent_id=agent_id,
|
||||
account_id=current_user.id,
|
||||
)
|
||||
|
||||
agent = roster_service.get_app_backing_agent(tenant_id=current_tenant_id, app_id=str(app_model.id))
|
||||
if agent is None:
|
||||
raise AgentNotFoundError()
|
||||
return roster_service.get_or_create_agent_app_debug_conversation_id(
|
||||
tenant_id=current_tenant_id,
|
||||
agent_id=agent.id,
|
||||
account_id=current_user.id,
|
||||
)
|
||||
|
||||
|
||||
def _create_chat_message(
|
||||
*, current_user: Account, app_model: App, current_tenant_id: str | None = None, agent_id: str | None = None
|
||||
):
|
||||
raw_payload = console_ns.payload or {}
|
||||
args_model = ChatMessagePayload.model_validate(raw_payload)
|
||||
args = args_model.model_dump(exclude_none=True, by_alias=True)
|
||||
|
||||
if AppMode.value_of(app_model.mode) == AppMode.AGENT:
|
||||
debug_conversation_id = _resolve_current_user_agent_debug_conversation_id(
|
||||
current_tenant_id=current_tenant_id or app_model.tenant_id,
|
||||
current_user=current_user,
|
||||
app_model=app_model,
|
||||
agent_id=agent_id,
|
||||
)
|
||||
if args_model.conversation_id and args_model.conversation_id != debug_conversation_id:
|
||||
raise NotFound("Conversation Not Exists.")
|
||||
args["conversation_id"] = debug_conversation_id
|
||||
|
||||
streaming = _resolve_debugger_chat_streaming(
|
||||
app_mode=AppMode.value_of(app_model.mode),
|
||||
response_mode=args_model.response_mode,
|
||||
|
||||
@@ -53,6 +53,7 @@ from libs.login import login_required
|
||||
from models.account import Account
|
||||
from models.enums import FeedbackFromSource, FeedbackRating
|
||||
from models.model import App, AppMode, Conversation, Message, MessageAnnotation, MessageFeedback
|
||||
from services.conversation_service import ConversationService
|
||||
from services.errors.conversation import ConversationNotExistsError
|
||||
from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError
|
||||
from services.message_service import MessageService, attach_message_extra_contents
|
||||
@@ -186,10 +187,11 @@ class ChatMessageListApi(Resource):
|
||||
@account_initialization_required
|
||||
@setup_required
|
||||
@edit_permission_required
|
||||
@with_current_user
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
|
||||
@get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT])
|
||||
def get(self, app_model: App):
|
||||
return _list_chat_messages(app_model=app_model)
|
||||
def get(self, current_user: Account, app_model: App):
|
||||
return _list_chat_messages(app_model=app_model, current_user=current_user)
|
||||
|
||||
|
||||
@console_ns.route("/agent/<uuid:agent_id>/chat-messages")
|
||||
@@ -205,10 +207,11 @@ class AgentChatMessageListApi(Resource):
|
||||
@setup_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def get(self, current_tenant_id: str, agent_id: UUID):
|
||||
def get(self, current_tenant_id: str, current_user: Account, agent_id: UUID):
|
||||
app_model = resolve_agent_app_model(tenant_id=current_tenant_id, agent_id=agent_id)
|
||||
return _list_chat_messages(app_model=app_model)
|
||||
return _list_chat_messages(app_model=app_model, current_user=current_user)
|
||||
|
||||
|
||||
@console_ns.route("/apps/<uuid:app_id>/feedbacks")
|
||||
@@ -390,14 +393,24 @@ class AgentMessageApi(Resource):
|
||||
return _get_message_detail(app_model=app_model, message_id=message_id)
|
||||
|
||||
|
||||
def _list_chat_messages(*, app_model: App):
|
||||
def _list_chat_messages(*, app_model: App, current_user: Account | None = None):
|
||||
args = ChatMessagesQuery.model_validate(request.args.to_dict())
|
||||
|
||||
conversation = db.session.scalar(
|
||||
select(Conversation)
|
||||
.where(Conversation.id == args.conversation_id, Conversation.app_id == app_model.id)
|
||||
.limit(1)
|
||||
)
|
||||
if AppMode.value_of(app_model.mode) == AppMode.AGENT and current_user is not None:
|
||||
try:
|
||||
conversation = ConversationService.get_conversation(
|
||||
app_model=app_model,
|
||||
conversation_id=args.conversation_id,
|
||||
user=current_user,
|
||||
)
|
||||
except ConversationNotExistsError:
|
||||
raise NotFound("Conversation Not Exists.")
|
||||
else:
|
||||
conversation = db.session.scalar(
|
||||
select(Conversation)
|
||||
.where(Conversation.id == args.conversation_id, Conversation.app_id == app_model.id)
|
||||
.limit(1)
|
||||
)
|
||||
|
||||
if not conversation:
|
||||
raise NotFound("Conversation Not Exists.")
|
||||
|
||||
@@ -83,7 +83,7 @@ class ApiKeyAuthDataSourceBinding(Resource):
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_CREATE, resource_required=False)
|
||||
@console_ns.expect(console_ns.models[ApiKeyAuthBindingPayload.__name__])
|
||||
@with_current_tenant_id
|
||||
def post(self, current_tenant_id: str):
|
||||
|
||||
@@ -26,6 +26,7 @@ from controllers.console.wraps import (
|
||||
with_current_tenant_id,
|
||||
with_current_user,
|
||||
)
|
||||
from extensions.ext_database import db
|
||||
from fields.base import ResponseModel
|
||||
from fields.dataset_fields import (
|
||||
dataset_detail_fields,
|
||||
@@ -390,6 +391,7 @@ class ExternalKnowledgeHitTestingApi(Resource):
|
||||
|
||||
try:
|
||||
response = HitTestingService.external_retrieve(
|
||||
session=db.session,
|
||||
dataset=dataset,
|
||||
query=payload.query,
|
||||
account=current_user,
|
||||
|
||||
@@ -18,6 +18,7 @@ from core.errors.error import (
|
||||
ProviderTokenNotInitError,
|
||||
QuotaExceededError,
|
||||
)
|
||||
from extensions.ext_database import db
|
||||
from graphon.model_runtime.errors.invoke import InvokeError
|
||||
from libs.login import resolve_account_fallback
|
||||
from models.account import Account
|
||||
@@ -115,6 +116,7 @@ class DatasetsHitTestingBase:
|
||||
try:
|
||||
current_user, _ = resolve_account_fallback(current_user, current_tenant_id)
|
||||
response = HitTestingService.retrieve(
|
||||
session=db.session,
|
||||
dataset=dataset,
|
||||
query=cast(str, args.get("query")),
|
||||
account=current_user,
|
||||
|
||||
@@ -222,7 +222,7 @@ class DatasourceAuth(Resource):
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@edit_permission_required
|
||||
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.CREDENTIAL_CREATE, resource_required=False)
|
||||
@with_current_tenant_id
|
||||
def post(self, current_tenant_id: str, provider_id: str):
|
||||
payload = DatasourceCredentialPayload.model_validate(console_ns.payload or {})
|
||||
|
||||
@@ -5,6 +5,7 @@ from sqlalchemy import select
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
from configs import dify_config
|
||||
from extensions.ext_database import db
|
||||
from libs.login import current_account_with_tenant
|
||||
from models.account import TenantPluginPermission
|
||||
@@ -17,6 +18,9 @@ def plugin_permission_required(
|
||||
def interceptor[**P, R](view: Callable[P, R]) -> Callable[P, R]:
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
|
||||
if dify_config.RBAC_ENABLED:
|
||||
return view(*args, **kwargs)
|
||||
|
||||
current_user, current_tenant_id = current_account_with_tenant()
|
||||
user = current_user
|
||||
tenant_id = current_tenant_id
|
||||
|
||||
@@ -169,7 +169,7 @@ class ModelProviderCredentialApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_CREATE, resource_required=False)
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def post(self, current_tenant_id: str, provider: str):
|
||||
@@ -244,7 +244,7 @@ class ModelProviderCredentialSwitchApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_USE, resource_required=False)
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def post(self, current_tenant_id: str, provider: str):
|
||||
@@ -326,7 +326,7 @@ class PreferredProviderTypeUpdateApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_USE, resource_required=False)
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, provider: str):
|
||||
|
||||
@@ -395,7 +395,7 @@ class ModelProviderModelCredentialApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_CREATE, resource_required=False)
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, provider: str):
|
||||
@@ -481,7 +481,7 @@ class ModelProviderModelCredentialSwitchApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_USE, resource_required=False)
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def post(self, current_tenant_id: str, provider: str):
|
||||
|
||||
@@ -469,6 +469,7 @@ class PluginDebuggingKeyApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_DEBUG, resource_required=False)
|
||||
@plugin_permission_required(debug_required=True)
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str):
|
||||
@@ -614,6 +615,7 @@ class PluginUploadFromPkgApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -634,6 +636,7 @@ class PluginUploadFromGithubApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -653,6 +656,7 @@ class PluginUploadFromBundleApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -673,6 +677,7 @@ class PluginInstallFromPkgApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -693,6 +698,7 @@ class PluginInstallFromGithubApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -719,6 +725,7 @@ class PluginInstallFromMarketplaceApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -739,6 +746,7 @@ class PluginFetchMarketplacePkgApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str):
|
||||
@@ -764,6 +772,7 @@ class PluginFetchManifestApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str):
|
||||
@@ -784,6 +793,7 @@ class PluginFetchInstallTasksApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str):
|
||||
@@ -801,6 +811,7 @@ class PluginFetchInstallTaskApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str, task_id: str):
|
||||
@@ -816,6 +827,7 @@ class PluginDeleteInstallTaskApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, task_id: str):
|
||||
@@ -831,6 +843,7 @@ class PluginDeleteAllInstallTaskItemsApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -846,6 +859,7 @@ class PluginDeleteInstallTaskItemApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, task_id: str, identifier: str):
|
||||
@@ -862,6 +876,7 @@ class PluginUpgradeFromMarketplaceApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -884,6 +899,7 @@ class PluginUpgradeFromGithubApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -911,6 +927,7 @@ class PluginUninstallApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_INSTALL, resource_required=False)
|
||||
@plugin_permission_required(install_required=True)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
@@ -1041,10 +1058,11 @@ class PluginChangeAutoUpgradeApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_PREFERENCES, resource_required=False)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, user: Account):
|
||||
if not user.is_admin_or_owner:
|
||||
if not dify_config.RBAC_ENABLED and not user.is_admin_or_owner:
|
||||
raise Forbidden()
|
||||
|
||||
args = ParserAutoUpgradeChange.model_validate(console_ns.payload)
|
||||
@@ -1097,6 +1115,7 @@ class PluginAutoUpgradeExcludePluginApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_PREFERENCES, resource_required=False)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
# exclude one single plugin
|
||||
|
||||
@@ -211,7 +211,7 @@ def _legacy_workspace_roles(
|
||||
name=role_name,
|
||||
description="",
|
||||
is_builtin=True,
|
||||
permission_keys=list(_LEGACY_ROLE_PERMISSION_KEYS[role_name]),
|
||||
permission_keys=list(dict.fromkeys(_LEGACY_ROLE_PERMISSION_KEYS[role_name])),
|
||||
role_tag="owner" if role_name == "owner" else "",
|
||||
)
|
||||
for role_name in ("owner", "admin", "editor", "normal", "dataset_operator")
|
||||
@@ -244,11 +244,6 @@ def _legacy_workspace_roles(
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Permission catalogs.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@console_ns.route("/workspaces/current/rbac/role-permissions/catalog")
|
||||
class RBACWorkspaceCatalogApi(Resource):
|
||||
@login_required
|
||||
@@ -375,30 +370,6 @@ class RBACRoleCopyApi(Resource):
|
||||
return _dump(role), 201
|
||||
|
||||
|
||||
@console_ns.route("/workspaces/current/rbac/roles/<uuid:role_id>/members")
|
||||
class RBACRoleMembersApi(Resource):
|
||||
@login_required
|
||||
@rbac_permission_required(
|
||||
RBACResourceScope.WORKSPACE, RBACPermission.WORKSPACE_ROLE_MANAGE, resource_required=False
|
||||
)
|
||||
@console_ns.response(200, "Success", console_ns.models[_RBACRoleAccountList.__name__])
|
||||
def get(self, role_id):
|
||||
tenant_id, account_id = _current_ids()
|
||||
return _dump(
|
||||
svc.RBACService.Roles.members(
|
||||
tenant_id,
|
||||
account_id,
|
||||
str(role_id),
|
||||
options=_pagination_options(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Access policies (tenant-level permission sets).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _AccessPolicyCreateRequest(BaseModel):
|
||||
name: str
|
||||
resource_type: svc.RBACResourceType
|
||||
@@ -788,11 +759,6 @@ class RBACDatasetMemberBindingsApi(Resource):
|
||||
return {"result": "success"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Workspace-level access (Settings > Access Rules).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@console_ns.route("/workspaces/current/rbac/workspace/apps/access-policy")
|
||||
class RBACWorkspaceAppMatrixApi(Resource):
|
||||
@login_required
|
||||
|
||||
@@ -971,7 +971,7 @@ class ToolBuiltinProviderSetDefaultApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False)
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_USE, resource_required=False)
|
||||
@account_initialization_required
|
||||
@with_current_tenant_id
|
||||
def post(self, current_tenant_id: str, provider: str):
|
||||
@@ -1070,6 +1070,7 @@ class ToolProviderMCPApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.MCP_MANAGE, resource_required=False)
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, user: Account):
|
||||
@@ -1125,6 +1126,7 @@ class ToolProviderMCPApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.MCP_MANAGE, resource_required=False)
|
||||
@with_current_tenant_id
|
||||
def put(self, current_tenant_id: str):
|
||||
payload = MCPProviderUpdatePayload.model_validate(console_ns.payload or {})
|
||||
@@ -1178,6 +1180,7 @@ class ToolProviderMCPApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.MCP_MANAGE, resource_required=False)
|
||||
@with_current_tenant_id
|
||||
def delete(self, current_tenant_id: str):
|
||||
payload = MCPProviderDeletePayload.model_validate(console_ns.payload or {})
|
||||
@@ -1196,6 +1199,7 @@ class ToolMCPAuthApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.MCP_MANAGE, resource_required=False)
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str):
|
||||
payload = MCPAuthPayload.model_validate(console_ns.payload or {})
|
||||
@@ -1300,6 +1304,7 @@ class ToolMCPUpdateApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.MCP_MANAGE, resource_required=False)
|
||||
@with_current_tenant_id
|
||||
def get(self, tenant_id: str, provider_id: str):
|
||||
with sessionmaker(db.engine).begin() as session:
|
||||
|
||||
@@ -31,7 +31,7 @@ from controllers.openapi._models import (
|
||||
AppDslExportQuery,
|
||||
AppDslExportResponse,
|
||||
AppDslImportPayload,
|
||||
AppInfoResponse,
|
||||
AppInfo,
|
||||
AppListQuery,
|
||||
AppListResponse,
|
||||
AppListRow,
|
||||
@@ -62,7 +62,6 @@ from controllers.openapi._models import (
|
||||
SessionListQuery,
|
||||
SessionListResponse,
|
||||
SessionRow,
|
||||
TagItem,
|
||||
TaskStopResponse,
|
||||
UsageInfo,
|
||||
WorkflowRunData,
|
||||
@@ -96,12 +95,11 @@ register_response_schema_models(
|
||||
openapi_ns,
|
||||
ErrorBody,
|
||||
EventStreamResponse,
|
||||
TagItem,
|
||||
UsageInfo,
|
||||
MessageMetadata,
|
||||
AppListRow,
|
||||
AppListResponse,
|
||||
AppInfoResponse,
|
||||
AppInfo,
|
||||
AppDescribeInfo,
|
||||
AppDescribeResponse,
|
||||
AppDslExportResponse,
|
||||
|
||||
@@ -63,6 +63,8 @@ class OpenApiErrorCode(StrEnum):
|
||||
FILE_EXTENSION_BLOCKED = "file_extension_blocked"
|
||||
MEMBER_LIMIT_EXCEEDED = "member_limit_exceeded"
|
||||
MEMBER_LICENSE_EXCEEDED = "member_license_exceeded"
|
||||
HUMAN_INPUT_FORM_NOT_FOUND = "form_not_found"
|
||||
RECIPIENT_SURFACE_MISMATCH = "recipient_surface_mismatch"
|
||||
|
||||
|
||||
class ErrorDetail(BaseModel):
|
||||
@@ -239,3 +241,16 @@ class MemberLicenseExceeded(OpenApiError): # noqa: N818
|
||||
error_code = OpenApiErrorCode.MEMBER_LICENSE_EXCEEDED
|
||||
description = "Workspace member license capacity reached."
|
||||
hint = "Contact your workspace administrator to expand the license seat count."
|
||||
|
||||
|
||||
class HumanInputFormNotFound(OpenApiError): # noqa: N818
|
||||
code = 404
|
||||
error_code = OpenApiErrorCode.HUMAN_INPUT_FORM_NOT_FOUND
|
||||
description = "No human-input form matches this token. It may be wrong, expired, or already submitted."
|
||||
|
||||
|
||||
class RecipientSurfaceMismatch(OpenApiError): # noqa: N818
|
||||
code = 403
|
||||
error_code = OpenApiErrorCode.RECIPIENT_SURFACE_MISMATCH
|
||||
description = "This form's recipient can't be submitted via the OpenAPI surface."
|
||||
hint = "Action it through its channel (web app or console)."
|
||||
|
||||
@@ -2,7 +2,8 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
from enum import StrEnum
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
@@ -13,6 +14,30 @@ from models.model import AppMode
|
||||
MAX_PAGE_LIMIT = 200
|
||||
|
||||
|
||||
class SupportedAppType(StrEnum):
|
||||
"""App types the ``app`` usage face (``get app``) lists and filters.
|
||||
|
||||
A curated subset of :class:`AppMode`: the real, user-facing app categories.
|
||||
Excludes runtime-only mode tags that are not standalone apps
|
||||
(``rag-pipeline`` is a knowledge ``Pipeline``; ``channel`` is unused) and the
|
||||
roster-owned ``agent`` type (surfaced through the roster, not this list).
|
||||
|
||||
Members reference ``AppMode.*.value`` so the subset relationship is
|
||||
type-checked: dropping a member from ``AppMode`` breaks this at import.
|
||||
This is the single source for the listable set — params, filters, and the
|
||||
generated CLI whitelist all derive from it.
|
||||
"""
|
||||
|
||||
COMPLETION = AppMode.COMPLETION.value
|
||||
CHAT = AppMode.CHAT.value
|
||||
ADVANCED_CHAT = AppMode.ADVANCED_CHAT.value
|
||||
WORKFLOW = AppMode.WORKFLOW.value
|
||||
AGENT_CHAT = AppMode.AGENT_CHAT.value
|
||||
|
||||
|
||||
SUPPORTED_APP_TYPES: Final[tuple[AppMode, ...]] = tuple(AppMode(t.value) for t in SupportedAppType)
|
||||
|
||||
|
||||
class UsageInfo(BaseModel):
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
@@ -38,18 +63,12 @@ class PaginationEnvelope[T](BaseModel):
|
||||
return cls(page=page, limit=limit, total=total, has_more=page * limit < total, data=items)
|
||||
|
||||
|
||||
class TagItem(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class AppListRow(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
description: str | None = None
|
||||
mode: AppMode
|
||||
tags: list[TagItem] = []
|
||||
updated_at: str | None = None
|
||||
created_by_name: str | None = None
|
||||
workspace_id: str | None = None
|
||||
workspace_name: str | None = None
|
||||
|
||||
@@ -70,16 +89,14 @@ class PermittedExternalAppsListResponse(BaseModel):
|
||||
data: list[AppListRow]
|
||||
|
||||
|
||||
class AppInfoResponse(BaseModel):
|
||||
class AppInfo(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
description: str | None = None
|
||||
mode: str
|
||||
author: str | None = None
|
||||
tags: list[TagItem] = []
|
||||
|
||||
|
||||
class AppDescribeInfo(AppInfoResponse):
|
||||
class AppDescribeInfo(AppInfo):
|
||||
updated_at: str | None = None
|
||||
service_api_enabled: bool
|
||||
is_agent: bool = False
|
||||
@@ -287,14 +304,13 @@ class AppDescribeQuery(BaseModel):
|
||||
|
||||
|
||||
class AppListQuery(BaseModel):
|
||||
"""mode is a closed enum."""
|
||||
"""mode is a closed enum of listable app types."""
|
||||
|
||||
workspace_id: UUIDStr
|
||||
page: int = Field(1, ge=1)
|
||||
limit: int = Field(20, ge=1, le=MAX_PAGE_LIMIT)
|
||||
mode: AppMode | None = None
|
||||
mode: SupportedAppType | None = None
|
||||
name: str | None = Field(None, max_length=200)
|
||||
tag: str | None = Field(None, max_length=100)
|
||||
|
||||
|
||||
class AppRunRequest(BaseModel):
|
||||
@@ -344,7 +360,7 @@ class PermittedExternalAppsListQuery(BaseModel):
|
||||
|
||||
page: int = Field(1, ge=1)
|
||||
limit: int = Field(20, ge=1, le=MAX_PAGE_LIMIT)
|
||||
mode: AppMode | None = None
|
||||
mode: SupportedAppType | None = None
|
||||
name: str | None = Field(None, max_length=200)
|
||||
|
||||
|
||||
|
||||
@@ -5,11 +5,12 @@ from typing import cast
|
||||
from flask_restx import Resource
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from controllers.common.wraps import RBACPermission, RBACResourceScope
|
||||
from controllers.openapi import openapi_ns
|
||||
from controllers.openapi._contract import accepts, returns
|
||||
from controllers.openapi._models import AppDslExportQuery, AppDslExportResponse, AppDslImportPayload
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from extensions.ext_database import db
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
from models import Account, App
|
||||
@@ -37,6 +38,11 @@ class AppDslImportApi(Resource):
|
||||
scope=Scope.WORKSPACE_WRITE,
|
||||
allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}),
|
||||
allowed_roles=frozenset({TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.OWNER}),
|
||||
rbac=RBACRequirement(
|
||||
resource_type=RBACResourceScope.APP,
|
||||
scene=RBACPermission.APP_IMPORT_EXPORT_DSL,
|
||||
resource_required=False,
|
||||
),
|
||||
)
|
||||
@returns(200, Import, "Import completed")
|
||||
@returns(202, Import, "Import pending confirmation")
|
||||
@@ -89,6 +95,11 @@ class AppDslImportConfirmApi(Resource):
|
||||
scope=Scope.WORKSPACE_WRITE,
|
||||
allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}),
|
||||
allowed_roles=frozenset({TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.OWNER}),
|
||||
rbac=RBACRequirement(
|
||||
resource_type=RBACResourceScope.APP,
|
||||
scene=RBACPermission.APP_IMPORT_EXPORT_DSL,
|
||||
resource_required=False,
|
||||
),
|
||||
)
|
||||
@returns(200, Import, "Import confirmed")
|
||||
@returns(400, Import, "Import failed")
|
||||
@@ -125,6 +136,7 @@ class AppDslExportApi(Resource):
|
||||
scope=Scope.APPS_READ,
|
||||
allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}),
|
||||
allowed_roles=frozenset({TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.OWNER}),
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_IMPORT_EXPORT_DSL),
|
||||
)
|
||||
@accepts(query=AppDslExportQuery)
|
||||
@returns(200, AppDslExportResponse, "Export successful")
|
||||
@@ -155,6 +167,7 @@ class AppDslCheckDependenciesApi(Resource):
|
||||
scope=Scope.APPS_READ,
|
||||
allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}),
|
||||
allowed_roles=frozenset({TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.OWNER}),
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_IMPORT_EXPORT_DSL),
|
||||
)
|
||||
@returns(200, CheckDependenciesResult, "Dependencies checked")
|
||||
def get(self, app_id: str, *, auth_data: AuthData):
|
||||
|
||||
@@ -19,12 +19,13 @@ from werkzeug.exceptions import (
|
||||
|
||||
import services
|
||||
from controllers.common.fields import EventStreamResponse
|
||||
from controllers.common.wraps import RBACPermission, RBACResourceScope
|
||||
from controllers.openapi import openapi_ns
|
||||
from controllers.openapi._audit import emit_app_run
|
||||
from controllers.openapi._contract import accepts, returns
|
||||
from controllers.openapi._models import AppRunRequest, TaskStopResponse
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from controllers.service_api.app.error import (
|
||||
AppUnavailableError,
|
||||
CompletionRequestError,
|
||||
@@ -136,7 +137,10 @@ _DISPATCH: dict[AppMode, Callable[[App, Any, AppRunRequest], Any]] = {
|
||||
|
||||
@openapi_ns.route("/apps/<string:app_id>/run")
|
||||
class AppRunApi(Resource):
|
||||
@auth_router.guard(scope=Scope.APPS_RUN)
|
||||
@auth_router.guard(
|
||||
scope=Scope.APPS_RUN,
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_TEST_AND_RUN),
|
||||
)
|
||||
@openapi_ns.response(200, "Run result (SSE stream)", openapi_ns.models[EventStreamResponse.__name__])
|
||||
@accepts(body=AppRunRequest)
|
||||
def post(self, app_id: str, *, auth_data: AuthData, body: AppRunRequest):
|
||||
@@ -167,7 +171,10 @@ class AppRunApi(Resource):
|
||||
|
||||
@openapi_ns.route("/apps/<string:app_id>/tasks/<string:task_id>/stop")
|
||||
class AppRunTaskStopApi(Resource):
|
||||
@auth_router.guard(scope=Scope.APPS_RUN)
|
||||
@auth_router.guard(
|
||||
scope=Scope.APPS_RUN,
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_TEST_AND_RUN),
|
||||
)
|
||||
@returns(200, TaskStopResponse, description="Task stopped")
|
||||
def post(self, app_id: str, task_id: str, *, auth_data: AuthData):
|
||||
app_model, caller, caller_kind = auth_data.require_app_context()
|
||||
|
||||
@@ -8,33 +8,41 @@ from typing import Any, cast
|
||||
from flask_restx import Resource
|
||||
from werkzeug.exceptions import Conflict, NotFound, UnprocessableEntity
|
||||
|
||||
from configs import dify_config
|
||||
from controllers.common.app_access import AppAccessFilter, resolve_app_access_filter
|
||||
from controllers.common.fields import Parameters
|
||||
from controllers.common.wraps import RBACPermission, RBACResourceScope
|
||||
from controllers.openapi import openapi_ns
|
||||
from controllers.openapi._contract import accepts, returns
|
||||
from controllers.openapi._input_schema import EMPTY_INPUT_SCHEMA, build_input_schema, resolve_app_config
|
||||
from controllers.openapi._models import (
|
||||
SUPPORTED_APP_TYPES,
|
||||
AppDescribeInfo,
|
||||
AppDescribeQuery,
|
||||
AppDescribeResponse,
|
||||
AppListQuery,
|
||||
AppListResponse,
|
||||
AppListRow,
|
||||
TagItem,
|
||||
)
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from controllers.service_api.app.error import AppUnavailableError
|
||||
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
|
||||
from extensions.ext_database import db
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
from models import App
|
||||
from models.model import AppMode
|
||||
from services.account_service import TenantService
|
||||
from services.app_service import AppListParams, AppService
|
||||
from services.tag_service import TagService
|
||||
|
||||
_ALLOWED_DESCRIBE_FIELDS: frozenset[str] = frozenset({"info", "parameters", "input_schema"})
|
||||
|
||||
|
||||
def _is_listable(app: App) -> bool:
|
||||
"""Whether the openapi app face exposes this app (curated, listable types only)."""
|
||||
return app.mode in SUPPORTED_APP_TYPES
|
||||
|
||||
|
||||
_EMPTY_PARAMETERS: dict[str, Any] = {
|
||||
"opening_statement": None,
|
||||
"suggested_questions": [],
|
||||
@@ -84,54 +92,55 @@ def parameters_payload(app: App) -> dict:
|
||||
return Parameters.model_validate(parameters).model_dump(mode="json")
|
||||
|
||||
|
||||
def build_app_describe_response(app: App, fields: set[str] | None) -> AppDescribeResponse:
|
||||
"""Public projection of an app (name / params / input schema) — never internal config."""
|
||||
want_info = fields is None or "info" in fields
|
||||
want_params = fields is None or "parameters" in fields
|
||||
want_schema = fields is None or "input_schema" in fields
|
||||
|
||||
info = (
|
||||
AppDescribeInfo(
|
||||
id=str(app.id),
|
||||
name=app.name,
|
||||
mode=app.mode,
|
||||
description=app.description,
|
||||
updated_at=app.updated_at.isoformat() if app.updated_at else None,
|
||||
service_api_enabled=bool(app.enable_api),
|
||||
is_agent=app.mode in (AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT),
|
||||
)
|
||||
if want_info
|
||||
else None
|
||||
)
|
||||
|
||||
parameters: dict[str, Any] | None = None
|
||||
input_schema: dict[str, Any] | None = None
|
||||
if want_params:
|
||||
try:
|
||||
parameters = parameters_payload(app)
|
||||
except AppUnavailableError:
|
||||
parameters = dict(_EMPTY_PARAMETERS)
|
||||
if want_schema:
|
||||
try:
|
||||
input_schema = build_input_schema(app)
|
||||
except AppUnavailableError:
|
||||
input_schema = dict(EMPTY_INPUT_SCHEMA)
|
||||
|
||||
return AppDescribeResponse(info=info, parameters=parameters, input_schema=input_schema)
|
||||
|
||||
|
||||
@openapi_ns.route("/apps/<string:app_id>/describe")
|
||||
class AppDescribeApi(AppReadResource):
|
||||
@auth_router.guard(scope=Scope.APPS_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}))
|
||||
@auth_router.guard(
|
||||
scope=Scope.APPS_READ,
|
||||
allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}),
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_VIEW_LAYOUT),
|
||||
)
|
||||
@returns(200, AppDescribeResponse, description="App description")
|
||||
@accepts(query=AppDescribeQuery)
|
||||
def get(self, app_id: str, *, auth_data: AuthData, query: AppDescribeQuery):
|
||||
# describe is UUID-only (workspace_id query param dropped in #37212).
|
||||
app = self._load(app_id)
|
||||
|
||||
requested = query.fields
|
||||
want_info = requested is None or "info" in requested
|
||||
want_params = requested is None or "parameters" in requested
|
||||
want_schema = requested is None or "input_schema" in requested
|
||||
|
||||
info = (
|
||||
AppDescribeInfo(
|
||||
id=str(app.id),
|
||||
name=app.name,
|
||||
mode=app.mode,
|
||||
description=app.description,
|
||||
tags=[TagItem(name=t.name) for t in app.tags],
|
||||
author=app.author_name,
|
||||
updated_at=app.updated_at.isoformat() if app.updated_at else None,
|
||||
service_api_enabled=bool(app.enable_api),
|
||||
is_agent=app.mode in ("agent-chat", "advanced-chat"),
|
||||
)
|
||||
if want_info
|
||||
else None
|
||||
)
|
||||
|
||||
parameters: dict[str, Any] | None = None
|
||||
input_schema: dict[str, Any] | None = None
|
||||
if want_params:
|
||||
try:
|
||||
parameters = parameters_payload(app)
|
||||
except AppUnavailableError:
|
||||
parameters = dict(_EMPTY_PARAMETERS)
|
||||
if want_schema:
|
||||
try:
|
||||
input_schema = build_input_schema(app)
|
||||
except AppUnavailableError:
|
||||
input_schema = dict(EMPTY_INPUT_SCHEMA)
|
||||
|
||||
return AppDescribeResponse(
|
||||
info=info,
|
||||
parameters=parameters,
|
||||
input_schema=input_schema,
|
||||
)
|
||||
return build_app_describe_response(app, query.fields)
|
||||
|
||||
|
||||
@openapi_ns.route("/apps")
|
||||
@@ -152,45 +161,57 @@ class AppListApi(Resource):
|
||||
else:
|
||||
parsed_uuid = None
|
||||
|
||||
# Compute RBAC-accessible app IDs when RBAC is enabled and the caller is an account.
|
||||
# ``None`` means unrestricted (caller can see all apps in the workspace);
|
||||
# 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
|
||||
)
|
||||
access_filter = AppAccessFilter.unrestricted()
|
||||
if apply_rbac_filter:
|
||||
access_filter = resolve_app_access_filter(workspace_id, str(auth_data.account_id))
|
||||
|
||||
tenant_name: str | None = None
|
||||
if parsed_uuid is not None:
|
||||
app: App | None = AppService.get_visible_app_by_id(db.session, str(parsed_uuid))
|
||||
if app is None or str(app.tenant_id) != workspace_id:
|
||||
return empty
|
||||
if not _is_listable(app):
|
||||
return empty
|
||||
# Apply RBAC visibility to the UUID fast-path the same way the service
|
||||
# layer does for paginated queries (id in accessible set OR own app).
|
||||
if apply_rbac_filter and not access_filter.is_app_accessible(
|
||||
str(app.id), str(app.maintainer) if app.maintainer else None, str(auth_data.account_id)
|
||||
):
|
||||
return empty
|
||||
tenant_name = TenantService.get_tenant_name(db.session, workspace_id)
|
||||
item = AppListRow(
|
||||
id=str(app.id),
|
||||
name=app.name,
|
||||
description=app.description,
|
||||
mode=app.mode,
|
||||
tags=[TagItem(name=t.name) for t in app.tags],
|
||||
updated_at=app.updated_at.isoformat() if app.updated_at else None,
|
||||
created_by_name=getattr(app, "author_name", None),
|
||||
workspace_id=str(workspace_id),
|
||||
workspace_name=tenant_name,
|
||||
)
|
||||
env = AppListResponse(page=1, limit=1, total=1, has_more=False, data=[item])
|
||||
return env
|
||||
|
||||
tag_ids: list[str] | None = None
|
||||
if query.tag:
|
||||
tags = TagService.get_tag_by_tag_name("app", workspace_id, query.tag, db.session)
|
||||
if not tags:
|
||||
return empty
|
||||
tag_ids = [tag.id for tag in tags]
|
||||
|
||||
params = AppListParams(
|
||||
page=query.page,
|
||||
limit=query.limit,
|
||||
mode=query.mode.value if query.mode else "all", # type:ignore
|
||||
name=query.name,
|
||||
tag_ids=tag_ids,
|
||||
status="normal",
|
||||
# Visibility gate pushed into the query — pagination.total stays
|
||||
# consistent across pages because invisible rows never count.
|
||||
openapi_visible=True,
|
||||
)
|
||||
|
||||
if apply_rbac_filter:
|
||||
access_filter.apply_to_params(params)
|
||||
|
||||
pagination = AppService().get_paginate_apps(str(auth_data.account_id), workspace_id, params, db.session)
|
||||
if pagination is None:
|
||||
return empty
|
||||
@@ -205,13 +226,12 @@ class AppListApi(Resource):
|
||||
name=r.name,
|
||||
description=r.description,
|
||||
mode=r.mode,
|
||||
tags=[TagItem(name=t.name) for t in r.tags],
|
||||
updated_at=r.updated_at.isoformat() if r.updated_at else None,
|
||||
created_by_name=getattr(r, "author_name", None),
|
||||
workspace_id=str(workspace_id),
|
||||
workspace_name=tenant_name,
|
||||
)
|
||||
for r in pagination.items
|
||||
if _is_listable(r)
|
||||
]
|
||||
|
||||
env = AppListResponse(
|
||||
|
||||
@@ -8,14 +8,18 @@ EE blueprint chain so this module is unreachable there.
|
||||
from __future__ import annotations
|
||||
|
||||
from flask_restx import Resource
|
||||
from werkzeug.exceptions import NotFound
|
||||
|
||||
from controllers.openapi import openapi_ns
|
||||
from controllers.openapi._contract import accepts, returns
|
||||
from controllers.openapi._models import (
|
||||
AppDescribeQuery,
|
||||
AppDescribeResponse,
|
||||
AppListRow,
|
||||
PermittedExternalAppsListQuery,
|
||||
PermittedExternalAppsListResponse,
|
||||
)
|
||||
from controllers.openapi.apps import build_app_describe_response
|
||||
from controllers.openapi.auth.composition import auth_router
|
||||
from controllers.openapi.auth.data import AuthData, Edition
|
||||
from extensions.ext_database import db
|
||||
@@ -67,9 +71,7 @@ class PermittedExternalAppsListApi(Resource):
|
||||
name=app.name,
|
||||
description=app.description,
|
||||
mode=app.mode,
|
||||
tags=[], # tenant-scoped; not surfaced cross-tenant
|
||||
updated_at=app.updated_at.isoformat() if app.updated_at else None,
|
||||
created_by_name=None, # cross-tenant author leak prevention
|
||||
workspace_id=str(app.tenant_id),
|
||||
workspace_name=tenant.name if tenant else None,
|
||||
)
|
||||
@@ -82,3 +84,20 @@ class PermittedExternalAppsListApi(Resource):
|
||||
data=items,
|
||||
)
|
||||
return env
|
||||
|
||||
|
||||
@openapi_ns.route("/permitted-external-apps/<string:app_id>/describe")
|
||||
class PermittedExternalAppDescribeApi(Resource):
|
||||
@auth_router.guard(
|
||||
scope=Scope.APPS_READ_PERMITTED_EXTERNAL,
|
||||
allowed_token_types=frozenset({TokenType.OAUTH_EXTERNAL_SSO}),
|
||||
edition=frozenset({Edition.EE}),
|
||||
)
|
||||
@returns(200, AppDescribeResponse, description="Permitted external app description")
|
||||
@accepts(query=AppDescribeQuery)
|
||||
def get(self, app_id: str, *, auth_data: AuthData, query: AppDescribeQuery):
|
||||
# App already loaded and ACL-checked by the external_sso pipeline; project it.
|
||||
app = auth_data.app
|
||||
if app is None:
|
||||
raise NotFound("app not found")
|
||||
return build_app_describe_response(app, query.fields)
|
||||
|
||||
@@ -3,9 +3,11 @@ from __future__ import annotations
|
||||
from controllers.openapi.auth.conditions import (
|
||||
EDITION_EE,
|
||||
HAS_ALLOWED_ROLES,
|
||||
HAS_RBAC,
|
||||
LOADED_APP_IS_PRIVATE,
|
||||
PATH_HAS_APP_ID,
|
||||
WEBAPP_AUTH_ENABLED,
|
||||
WEBAPP_RUN_SCOPED,
|
||||
WORKSPACE_MEMBERSHIP_REQUIRED,
|
||||
WORKSPACE_SCOPED,
|
||||
)
|
||||
@@ -25,6 +27,7 @@ from controllers.openapi.auth.verify import (
|
||||
check_acl,
|
||||
check_app_api_enabled,
|
||||
check_private_app_permission,
|
||||
check_rbac_permission,
|
||||
check_scope,
|
||||
check_workspace_member,
|
||||
check_workspace_mismatch,
|
||||
@@ -47,8 +50,9 @@ account_pipeline = AuthPipeline(
|
||||
When(WORKSPACE_SCOPED, then=check_workspace_member),
|
||||
When(PATH_HAS_APP_ID, then=check_workspace_mismatch),
|
||||
When(HAS_ALLOWED_ROLES, then=check_workspace_role),
|
||||
When(PATH_HAS_APP_ID & EDITION_EE & WEBAPP_AUTH_ENABLED, then=check_acl),
|
||||
When(EDITION_EE & LOADED_APP_IS_PRIVATE, then=check_private_app_permission),
|
||||
When(HAS_RBAC, then=check_rbac_permission),
|
||||
When(PATH_HAS_APP_ID & EDITION_EE & WEBAPP_AUTH_ENABLED & WEBAPP_RUN_SCOPED, then=check_acl),
|
||||
When(EDITION_EE & LOADED_APP_IS_PRIVATE & WEBAPP_RUN_SCOPED, then=check_private_app_permission),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
from collections.abc import Callable
|
||||
|
||||
from controllers.openapi.auth.data import AuthData, Edition, RequestContext, current_edition
|
||||
from libs.oauth_bearer import TokenType
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
from services.enterprise.enterprise_service import WebAppAccessMode
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
@@ -50,8 +50,11 @@ EDITION_SAAS = config_cond(lambda: current_edition() == Edition.SAAS)
|
||||
|
||||
WEBAPP_AUTH_ENABLED = config_cond(lambda: FeatureService.get_system_features().webapp_auth.enabled)
|
||||
|
||||
WEBAPP_RUN_SCOPED = request_cond(lambda ctx: ctx.scope == Scope.APPS_RUN)
|
||||
|
||||
WORKSPACE_MEMBERSHIP_REQUIRED = request_cond(lambda ctx: ctx.workspace_membership)
|
||||
HAS_ALLOWED_ROLES = request_cond(lambda ctx: ctx.allowed_roles is not None)
|
||||
HAS_RBAC = request_cond(lambda ctx: ctx.rbac is not None)
|
||||
|
||||
# Caller must belong to the resolved tenant: either an app-scoped path (tenant
|
||||
# from the app) or an explicit workspace-membership path (tenant from request).
|
||||
|
||||
@@ -8,6 +8,7 @@ from pydantic import BaseModel, ConfigDict, Field
|
||||
from werkzeug.exceptions import InternalServerError
|
||||
|
||||
from configs import dify_config
|
||||
from core.rbac import RBACPermission, RBACResourceScope
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
from models.account import Account, Tenant, TenantAccountRole
|
||||
from models.model import App, EndUser
|
||||
@@ -35,6 +36,14 @@ class ExternalIdentity(BaseModel):
|
||||
issuer: str | None = None
|
||||
|
||||
|
||||
class RBACRequirement(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
resource_type: RBACResourceScope
|
||||
scene: RBACPermission
|
||||
resource_required: bool = True
|
||||
|
||||
|
||||
class RequestContext(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
@@ -43,6 +52,7 @@ class RequestContext(BaseModel):
|
||||
path_params: dict[str, str]
|
||||
workspace_membership: bool = False
|
||||
allowed_roles: frozenset[TenantAccountRole] | None = None
|
||||
rbac: RBACRequirement | None = None
|
||||
|
||||
|
||||
class AuthData(BaseModel):
|
||||
@@ -59,6 +69,7 @@ class AuthData(BaseModel):
|
||||
path_params: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
allowed_roles: frozenset[TenantAccountRole] | None = None
|
||||
rbac: RBACRequirement | None = None
|
||||
|
||||
app: App | None = None
|
||||
tenant: Tenant | None = None
|
||||
|
||||
@@ -21,6 +21,7 @@ from controllers.openapi.auth.data import (
|
||||
AuthData,
|
||||
Edition,
|
||||
ExternalIdentity,
|
||||
RBACRequirement,
|
||||
RequestContext,
|
||||
current_edition,
|
||||
)
|
||||
@@ -59,6 +60,7 @@ class AuthPipeline:
|
||||
scope: Scope | None,
|
||||
workspace_membership: bool = False,
|
||||
allowed_roles: frozenset[TenantAccountRole] | None = None,
|
||||
rbac: RBACRequirement | None = None,
|
||||
) -> Any:
|
||||
req_ctx = RequestContext(
|
||||
token_type=identity.token_type,
|
||||
@@ -66,6 +68,7 @@ class AuthPipeline:
|
||||
path_params=dict(request.view_args or {}),
|
||||
workspace_membership=workspace_membership,
|
||||
allowed_roles=allowed_roles,
|
||||
rbac=rbac,
|
||||
)
|
||||
|
||||
data = AuthData(
|
||||
@@ -77,6 +80,7 @@ class AuthPipeline:
|
||||
tenants=dict(identity.verified_tenants),
|
||||
required_scope=scope,
|
||||
allowed_roles=allowed_roles,
|
||||
rbac=rbac,
|
||||
path_params=dict(req_ctx.path_params),
|
||||
external_identity=(
|
||||
ExternalIdentity(email=identity.subject_email, issuer=identity.subject_issuer)
|
||||
@@ -129,6 +133,7 @@ class PipelineRouter:
|
||||
edition: frozenset[Edition] | None = None,
|
||||
workspace_membership: bool = False,
|
||||
allowed_roles: frozenset[TenantAccountRole] | None = None,
|
||||
rbac: RBACRequirement | None = None,
|
||||
) -> Callable:
|
||||
return self._make_decorator(
|
||||
scope=scope,
|
||||
@@ -136,6 +141,7 @@ class PipelineRouter:
|
||||
edition=edition,
|
||||
workspace_membership=workspace_membership,
|
||||
allowed_roles=allowed_roles,
|
||||
rbac=rbac,
|
||||
)
|
||||
|
||||
def guard_workspace(
|
||||
@@ -145,6 +151,7 @@ class PipelineRouter:
|
||||
allowed_token_types: frozenset[TokenType] | None = None,
|
||||
edition: frozenset[Edition] | None = None,
|
||||
allowed_roles: frozenset[TenantAccountRole] | None = None,
|
||||
rbac: RBACRequirement | None = None,
|
||||
) -> Callable:
|
||||
return self._make_decorator(
|
||||
scope=scope,
|
||||
@@ -152,6 +159,7 @@ class PipelineRouter:
|
||||
edition=edition,
|
||||
workspace_membership=True,
|
||||
allowed_roles=allowed_roles,
|
||||
rbac=rbac,
|
||||
)
|
||||
|
||||
def _make_decorator(
|
||||
@@ -162,6 +170,7 @@ class PipelineRouter:
|
||||
edition: frozenset[Edition] | None,
|
||||
workspace_membership: bool,
|
||||
allowed_roles: frozenset[TenantAccountRole] | None,
|
||||
rbac: RBACRequirement | None,
|
||||
) -> Callable:
|
||||
def decorator(view: Callable) -> Callable:
|
||||
@wraps(view)
|
||||
@@ -175,6 +184,7 @@ class PipelineRouter:
|
||||
edition=edition,
|
||||
workspace_membership=workspace_membership,
|
||||
allowed_roles=allowed_roles,
|
||||
rbac=rbac,
|
||||
)
|
||||
|
||||
return decorated
|
||||
@@ -192,6 +202,7 @@ class PipelineRouter:
|
||||
edition: frozenset[Edition] | None,
|
||||
workspace_membership: bool = False,
|
||||
allowed_roles: frozenset[TenantAccountRole] | None = None,
|
||||
rbac: RBACRequirement | None = None,
|
||||
) -> Any:
|
||||
# 404 not 403 — this edition doesn't expose the feature at all
|
||||
if edition is not None and current_edition() not in edition:
|
||||
@@ -235,6 +246,7 @@ class PipelineRouter:
|
||||
scope=scope,
|
||||
workspace_membership=workspace_membership,
|
||||
allowed_roles=allowed_roles,
|
||||
rbac=rbac,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -74,12 +74,13 @@ def accept_subjects(*accepted: SubjectType) -> Callable[[F], F]:
|
||||
|
||||
|
||||
def _coerce_subject_type(raw: object) -> SubjectType | None:
|
||||
if raw is None:
|
||||
return None
|
||||
if isinstance(raw, SubjectType):
|
||||
return raw
|
||||
if isinstance(raw, str):
|
||||
return SubjectType(raw)
|
||||
match raw:
|
||||
case None:
|
||||
return None
|
||||
case SubjectType():
|
||||
return raw
|
||||
case str():
|
||||
return SubjectType(raw)
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -3,6 +3,8 @@ from __future__ import annotations
|
||||
from flask import request
|
||||
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 extensions.ext_database import db
|
||||
from libs.oauth_bearer import Scope, TokenType
|
||||
@@ -38,6 +40,9 @@ def check_workspace_mismatch(data: AuthData) -> None:
|
||||
|
||||
|
||||
def check_workspace_role(data: AuthData) -> None:
|
||||
if dify_config.RBAC_ENABLED and data.rbac is not None:
|
||||
# fine-grained permission check is performed by RBAC
|
||||
return
|
||||
if data.allowed_roles is None:
|
||||
return
|
||||
if data.tenant_role is None:
|
||||
@@ -46,6 +51,27 @@ def check_workspace_role(data: AuthData) -> None:
|
||||
raise Forbidden("insufficient workspace role")
|
||||
|
||||
|
||||
def check_rbac_permission(data: AuthData) -> None:
|
||||
req = data.rbac
|
||||
if req is None:
|
||||
return
|
||||
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":
|
||||
return
|
||||
if data.account_id is None or data.tenant is None:
|
||||
raise Forbidden("rbac context missing")
|
||||
enforce_rbac_access(
|
||||
tenant_id=str(data.tenant.id),
|
||||
account_id=str(data.account_id),
|
||||
resource_type=req.resource_type,
|
||||
scene=req.scene,
|
||||
resource_required=req.resource_required,
|
||||
path_args=dict(data.path_params),
|
||||
)
|
||||
|
||||
|
||||
def check_app_api_enabled(data: AuthData) -> None:
|
||||
if data.app is None:
|
||||
return
|
||||
|
||||
@@ -12,16 +12,21 @@ import logging
|
||||
|
||||
from flask import Response
|
||||
from flask_restx import Resource
|
||||
from werkzeug.exceptions import BadRequest, NotFound
|
||||
from werkzeug.exceptions import BadRequest
|
||||
|
||||
from controllers.common.human_input import HumanInputFormSubmitPayload, stringify_form_default_values
|
||||
from controllers.common.schema import register_schema_models
|
||||
from controllers.common.wraps import RBACPermission, RBACResourceScope
|
||||
from controllers.openapi import openapi_ns
|
||||
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
|
||||
from core.workflow.human_input_policy import HumanInputSurface, is_recipient_type_allowed_for_surface
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from core.workflow.human_input_policy import (
|
||||
HumanInputSurface,
|
||||
is_recipient_type_allowed_for_surface,
|
||||
)
|
||||
from extensions.ext_database import db
|
||||
from libs.helper import to_timestamp
|
||||
from libs.oauth_bearer import Scope
|
||||
@@ -47,31 +52,37 @@ def _jsonify_form_definition(form) -> Response:
|
||||
|
||||
def _ensure_form_belongs_to_app(form, app_model: App) -> None:
|
||||
if form.app_id != app_model.id or form.tenant_id != app_model.tenant_id:
|
||||
raise NotFound("Form not found")
|
||||
raise HumanInputFormNotFound()
|
||||
|
||||
|
||||
def _ensure_form_is_allowed_for_openapi(form) -> None:
|
||||
if not is_recipient_type_allowed_for_surface(form.recipient_type, HumanInputSurface.OPENAPI):
|
||||
raise NotFound("Form not found")
|
||||
raise RecipientSurfaceMismatch()
|
||||
|
||||
|
||||
@openapi_ns.route("/apps/<string:app_id>/form/human_input/<string:form_token>")
|
||||
class OpenApiWorkflowHumanInputFormApi(Resource):
|
||||
@openapi_ns.response(200, "Form definition", openapi_ns.models[HumanInputFormDefinitionResponse.__name__])
|
||||
@auth_router.guard(scope=Scope.APPS_RUN)
|
||||
@auth_router.guard(
|
||||
scope=Scope.APPS_RUN,
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_TEST_AND_RUN),
|
||||
)
|
||||
def get(self, app_id: str, form_token: str, *, auth_data: AuthData):
|
||||
app_model, caller, caller_kind = auth_data.require_app_context()
|
||||
app_model, _caller, _caller_kind = auth_data.require_app_context()
|
||||
service = HumanInputService(db.engine)
|
||||
form = service.get_form_by_token(form_token)
|
||||
if form is None:
|
||||
raise NotFound("Form not found")
|
||||
raise HumanInputFormNotFound()
|
||||
|
||||
_ensure_form_belongs_to_app(form, app_model)
|
||||
_ensure_form_is_allowed_for_openapi(form)
|
||||
service.ensure_form_active(form)
|
||||
return _jsonify_form_definition(form)
|
||||
|
||||
@auth_router.guard(scope=Scope.APPS_RUN)
|
||||
@auth_router.guard(
|
||||
scope=Scope.APPS_RUN,
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_TEST_AND_RUN),
|
||||
)
|
||||
@returns(200, FormSubmitResponse, description="Form submitted")
|
||||
@accepts(body=HumanInputFormSubmitPayload)
|
||||
def post(self, app_id: str, form_token: str, *, auth_data: AuthData, body: HumanInputFormSubmitPayload):
|
||||
@@ -80,7 +91,7 @@ class OpenApiWorkflowHumanInputFormApi(Resource):
|
||||
service = HumanInputService(db.engine)
|
||||
form = service.get_form_by_token(form_token)
|
||||
if form is None:
|
||||
raise NotFound("Form not found")
|
||||
raise HumanInputFormNotFound()
|
||||
|
||||
_ensure_form_belongs_to_app(form, app_model)
|
||||
_ensure_form_is_allowed_for_openapi(form)
|
||||
@@ -106,6 +117,6 @@ class OpenApiWorkflowHumanInputFormApi(Resource):
|
||||
submission_end_user_id=submission_end_user_id,
|
||||
)
|
||||
except FormNotFoundError:
|
||||
raise NotFound("Form not found")
|
||||
raise HumanInputFormNotFound()
|
||||
|
||||
return FormSubmitResponse()
|
||||
|
||||
@@ -19,9 +19,10 @@ from werkzeug.exceptions import NotFound, UnprocessableEntity
|
||||
|
||||
from controllers.common.fields import EventStreamResponse
|
||||
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
|
||||
from controllers.openapi.auth.data import AuthData, RBACRequirement
|
||||
from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator
|
||||
from core.app.apps.base_app_generator import BaseAppGenerator
|
||||
from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter
|
||||
@@ -46,7 +47,10 @@ class WorkflowEventsQuery(BaseModel):
|
||||
class OpenApiWorkflowEventsApi(Resource):
|
||||
@openapi_ns.doc(params=query_params_from_model(WorkflowEventsQuery))
|
||||
@openapi_ns.response(200, "SSE event stream", openapi_ns.models[EventStreamResponse.__name__])
|
||||
@auth_router.guard(scope=Scope.APPS_RUN)
|
||||
@auth_router.guard(
|
||||
scope=Scope.APPS_RUN,
|
||||
rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_TEST_AND_RUN),
|
||||
)
|
||||
def get(self, app_id: str, task_id: str, *, auth_data: AuthData):
|
||||
app_model, caller, caller_kind = auth_data.require_app_context()
|
||||
app_mode = AppMode.value_of(app_model.mode)
|
||||
|
||||
@@ -2,6 +2,7 @@ from typing import Any, cast
|
||||
|
||||
from flask_restx import Resource
|
||||
from pydantic import Field
|
||||
from sqlalchemy import select
|
||||
|
||||
from controllers.common.fields import Parameters
|
||||
from controllers.common.schema import register_response_schema_models
|
||||
@@ -9,7 +10,11 @@ from controllers.service_api import service_api_ns
|
||||
from controllers.service_api.app.error import AppUnavailableError
|
||||
from controllers.service_api.wraps import validate_app_token
|
||||
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
|
||||
from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form
|
||||
from extensions.ext_database import db
|
||||
from fields.base import ResponseModel
|
||||
from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource, AgentStatus
|
||||
from models.agent_config_entities import AgentSoulConfig
|
||||
from models.model import App, AppMode
|
||||
from services.app_service import AppService
|
||||
|
||||
@@ -29,6 +34,40 @@ class AppMetaResponse(ResponseModel):
|
||||
register_response_schema_models(service_api_ns, Parameters, AppMetaResponse, AppInfoResponse)
|
||||
|
||||
|
||||
def _get_agent_app_feature_dict_and_user_input_form(app_model: App) -> tuple[dict[str, Any], list[dict[str, Any]]]:
|
||||
app_model_config = app_model.app_model_config
|
||||
features_dict = cast(dict[str, Any], app_model_config.to_dict()) if app_model_config is not None else {}
|
||||
|
||||
agent = db.session.scalar(
|
||||
select(Agent)
|
||||
.where(
|
||||
Agent.tenant_id == app_model.tenant_id,
|
||||
Agent.app_id == app_model.id,
|
||||
Agent.scope == AgentScope.ROSTER,
|
||||
Agent.source == AgentSource.AGENT_APP,
|
||||
Agent.status == AgentStatus.ACTIVE,
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
if agent is None or not agent.active_config_snapshot_id:
|
||||
raise AppUnavailableError()
|
||||
|
||||
snapshot = db.session.scalar(
|
||||
select(AgentConfigSnapshot)
|
||||
.where(
|
||||
AgentConfigSnapshot.tenant_id == app_model.tenant_id,
|
||||
AgentConfigSnapshot.agent_id == agent.id,
|
||||
AgentConfigSnapshot.id == agent.active_config_snapshot_id,
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
if snapshot is None:
|
||||
raise AppUnavailableError()
|
||||
|
||||
agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict)
|
||||
return features_dict, agent_app_variables_to_user_input_form(agent_soul.app_variables)
|
||||
|
||||
|
||||
@service_api_ns.route("/parameters")
|
||||
class AppParameterApi(Resource):
|
||||
"""Resource for app variables."""
|
||||
@@ -61,12 +100,16 @@ class AppParameterApi(Resource):
|
||||
|
||||
Returns the input form parameters and configuration for the application.
|
||||
"""
|
||||
if app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
|
||||
features_dict: dict[str, Any]
|
||||
user_input_form: list[dict[str, Any]]
|
||||
if app_model.mode == AppMode.AGENT:
|
||||
features_dict, user_input_form = _get_agent_app_feature_dict_and_user_input_form(app_model)
|
||||
elif app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
|
||||
workflow = app_model.workflow
|
||||
if workflow is None:
|
||||
raise AppUnavailableError()
|
||||
|
||||
features_dict: dict[str, Any] = workflow.features_dict
|
||||
features_dict = workflow.features_dict
|
||||
user_input_form = workflow.user_input_form(to_old_structure=True)
|
||||
else:
|
||||
app_model_config = app_model.app_model_config
|
||||
|
||||
@@ -4,7 +4,7 @@ from collections.abc import Mapping, Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfig
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager
|
||||
@@ -22,7 +22,7 @@ from core.app.entities.queue_entities import (
|
||||
from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature
|
||||
from core.app.layers.conversation_variable_persist_layer import ConversationVariablePersistenceLayer
|
||||
from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer
|
||||
from core.db.session_factory import session_factory
|
||||
from core.db.session_factory import create_session, session_factory
|
||||
from core.moderation.base import ModerationError
|
||||
from core.moderation.input_moderation import InputModeration
|
||||
from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository
|
||||
@@ -107,7 +107,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
|
||||
workflow_execution_id=self.application_generate_entity.workflow_run_id,
|
||||
)
|
||||
|
||||
with Session(db.engine, expire_on_commit=False) as session:
|
||||
with create_session() as session:
|
||||
app_record = session.scalar(select(App).where(App.id == app_config.app_id))
|
||||
|
||||
if not app_record:
|
||||
@@ -204,6 +204,8 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
|
||||
trace_session_id=self.application_generate_entity.extras.get("trace_session_id"),
|
||||
)
|
||||
|
||||
# Release the Flask scoped session before workflow execution so a checked-out DB connection
|
||||
# is not held for the lifetime of the graph run.
|
||||
db.session.close()
|
||||
|
||||
# RUN WORKFLOW
|
||||
@@ -368,7 +370,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
|
||||
|
||||
:return: List of conversation variables ready for use
|
||||
"""
|
||||
with sessionmaker(bind=db.engine).begin() as session:
|
||||
with create_session() as session, session.begin():
|
||||
existing_variables = self._load_existing_conversation_variables(session)
|
||||
|
||||
if not existing_variables:
|
||||
|
||||
@@ -21,6 +21,7 @@ from core.app.app_config.entities import (
|
||||
EasyUIBasedAppModelConfigFrom,
|
||||
PromptTemplateEntity,
|
||||
)
|
||||
from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form
|
||||
from models.agent_config_entities import AgentSoulConfig
|
||||
from models.model import App, AppMode, AppModelConfig, AppModelConfigDict, Conversation
|
||||
|
||||
@@ -98,8 +99,7 @@ class AgentAppConfigManager(BaseAppConfigManager):
|
||||
# pipeline's bookkeeping (token counting, persistence).
|
||||
base["prompt_type"] = PromptTemplateEntity.PromptType.SIMPLE.value
|
||||
base["pre_prompt"] = agent_soul.prompt.system_prompt or ""
|
||||
# Agent App takes the user message directly; no completion-style inputs form.
|
||||
base.setdefault("user_input_form", [])
|
||||
base["user_input_form"] = agent_app_variables_to_user_input_form(agent_soul.app_variables)
|
||||
return base
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from models.agent_config_entities import AppVariableConfig
|
||||
|
||||
|
||||
def agent_app_variables_to_user_input_form(app_variables: Sequence[AppVariableConfig]) -> list[dict[str, Any]]:
|
||||
"""Project Agent Soul app variables into the legacy service-API parameter form."""
|
||||
|
||||
user_input_form: list[dict[str, Any]] = []
|
||||
for variable in app_variables:
|
||||
form_type = _form_type_for_agent_variable(variable.type)
|
||||
form_item: dict[str, Any] = {
|
||||
"label": variable.name,
|
||||
"variable": variable.name,
|
||||
"required": variable.required,
|
||||
}
|
||||
if variable.default is not None:
|
||||
form_item["default"] = variable.default
|
||||
user_input_form.append({form_type: form_item})
|
||||
return user_input_form
|
||||
|
||||
|
||||
def _form_type_for_agent_variable(variable_type: str) -> str:
|
||||
normalized = variable_type.strip().lower()
|
||||
if normalized in {"number", "integer", "float"}:
|
||||
return "number"
|
||||
if normalized in {"boolean", "bool"}:
|
||||
return "checkbox"
|
||||
if normalized in {"paragraph", "long_text", "multiline"}:
|
||||
return "paragraph"
|
||||
return "text-input"
|
||||
|
||||
|
||||
__all__ = ["agent_app_variables_to_user_input_form"]
|
||||
@@ -12,10 +12,10 @@ from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
|
||||
from core.app.apps.base_app_runner import AppRunner
|
||||
from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity
|
||||
from core.app.entities.queue_entities import QueueAnnotationReplyEvent
|
||||
from core.db.session_factory import create_session
|
||||
from core.memory.token_buffer_memory import TokenBufferMemory
|
||||
from core.model_manager import ModelInstance
|
||||
from core.moderation.base import ModerationError
|
||||
from extensions.ext_database import db
|
||||
from graphon.model_runtime.entities.llm_entities import LLMMode
|
||||
from graphon.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey
|
||||
from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
|
||||
@@ -47,7 +47,10 @@ class AgentChatAppRunner(AppRunner):
|
||||
app_config = application_generate_entity.app_config
|
||||
app_config = cast(AgentChatAppConfig, app_config)
|
||||
app_stmt = select(App).where(App.id == app_config.app_id)
|
||||
app_record = db.session.scalar(app_stmt)
|
||||
with create_session() as session:
|
||||
app_record = session.scalar(app_stmt)
|
||||
if app_record:
|
||||
session.expunge(app_record)
|
||||
if not app_record:
|
||||
raise ValueError("App not found")
|
||||
|
||||
@@ -185,14 +188,18 @@ class AgentChatAppRunner(AppRunner):
|
||||
if {ModelFeature.MULTI_TOOL_CALL, ModelFeature.TOOL_CALL}.intersection(model_schema.features or []):
|
||||
agent_entity.strategy = AgentEntity.Strategy.FUNCTION_CALLING
|
||||
conversation_stmt = select(Conversation).where(Conversation.id == conversation.id)
|
||||
conversation_result = db.session.scalar(conversation_stmt)
|
||||
if conversation_result is None:
|
||||
raise ValueError("Conversation not found")
|
||||
msg_stmt = select(Message).where(Message.id == message.id)
|
||||
message_result = db.session.scalar(msg_stmt)
|
||||
with create_session() as session:
|
||||
conversation_result = session.scalar(conversation_stmt)
|
||||
if conversation_result is None:
|
||||
raise ValueError("Conversation not found")
|
||||
|
||||
message_result = session.scalar(msg_stmt)
|
||||
if message_result is not None:
|
||||
session.expunge(message_result)
|
||||
session.expunge(conversation_result)
|
||||
if message_result is None:
|
||||
raise ValueError("Message not found")
|
||||
db.session.close()
|
||||
|
||||
runner_cls: type[FunctionCallAgentRunner] | type[CotChatAgentRunner] | type[CotCompletionAgentRunner]
|
||||
# start agent runner
|
||||
|
||||
@@ -11,6 +11,7 @@ from core.app.entities.app_invoke_entities import (
|
||||
)
|
||||
from core.app.entities.queue_entities import QueueAnnotationReplyEvent
|
||||
from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
|
||||
from core.db.session_factory import create_session
|
||||
from core.memory.token_buffer_memory import TokenBufferMemory
|
||||
from core.model_manager import ModelInstance
|
||||
from core.moderation.base import ModerationError
|
||||
@@ -46,7 +47,10 @@ class ChatAppRunner(AppRunner):
|
||||
app_config = application_generate_entity.app_config
|
||||
app_config = cast(ChatAppConfig, app_config)
|
||||
stmt = select(App).where(App.id == app_config.app_id)
|
||||
app_record = db.session.scalar(stmt)
|
||||
with create_session() as session:
|
||||
app_record = session.scalar(stmt)
|
||||
if app_record:
|
||||
session.expunge(app_record)
|
||||
if not app_record:
|
||||
raise ValueError("App not found")
|
||||
|
||||
@@ -216,6 +220,8 @@ class ChatAppRunner(AppRunner):
|
||||
model=application_generate_entity.model_conf.model,
|
||||
)
|
||||
|
||||
# Release the Flask scoped session before LLM streaming so a checked-out DB connection
|
||||
# is not held for the lifetime of the provider response.
|
||||
db.session.close()
|
||||
|
||||
invoke_result = model_instance.invoke_llm(
|
||||
|
||||
@@ -51,8 +51,11 @@ from core.tools.entities.tool_entities import ToolProviderType
|
||||
from core.tools.tool_manager import ToolManager
|
||||
from core.trigger.constants import TRIGGER_PLUGIN_NODE_TYPE
|
||||
from core.trigger.trigger_manager import TriggerManager
|
||||
from core.workflow.human_input_forms import load_form_tokens_by_form_id
|
||||
from core.workflow.human_input_forms import (
|
||||
load_form_dispositions_by_form_id,
|
||||
)
|
||||
from core.workflow.human_input_policy import (
|
||||
FormDisposition,
|
||||
HumanInputSurface,
|
||||
enrich_human_input_pause_reasons,
|
||||
resolve_human_input_pause_reason_inputs,
|
||||
@@ -340,13 +343,14 @@ class WorkflowResponseConverter:
|
||||
human_input_form_ids = [reason.form_id for reason in resolved_reasons if isinstance(reason, HumanInputRequired)]
|
||||
expiration_times_by_form_id: dict[str, datetime] = {}
|
||||
display_in_ui_by_form_id: dict[str, bool] = {}
|
||||
form_token_by_form_id: dict[str, str] = {}
|
||||
dispositions_by_form_id: dict[str, FormDisposition] = {}
|
||||
if human_input_form_ids:
|
||||
stmt = select(
|
||||
HumanInputForm.id,
|
||||
HumanInputForm.expiration_time,
|
||||
HumanInputForm.form_definition,
|
||||
).where(HumanInputForm.id.in_(human_input_form_ids))
|
||||
hitl_surface = _INVOKE_FROM_TO_HITL_SURFACE.get(self._application_generate_entity.invoke_from)
|
||||
with Session(bind=db.engine) as session:
|
||||
for form_id, expiration_time, form_definition in session.execute(stmt):
|
||||
expiration_times_by_form_id[str(form_id)] = expiration_time
|
||||
@@ -355,17 +359,17 @@ class WorkflowResponseConverter:
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
definition_payload = {}
|
||||
display_in_ui_by_form_id[str(form_id)] = bool(definition_payload.get("display_in_ui"))
|
||||
form_token_by_form_id = load_form_tokens_by_form_id(
|
||||
dispositions_by_form_id = load_form_dispositions_by_form_id(
|
||||
human_input_form_ids,
|
||||
session=session,
|
||||
surface=_INVOKE_FROM_TO_HITL_SURFACE.get(self._application_generate_entity.invoke_from),
|
||||
surface=hitl_surface,
|
||||
)
|
||||
|
||||
# Reconnect paths must preserve the same pause-reason contract as live streams;
|
||||
# otherwise clients see schema drift after resume.
|
||||
pause_reasons = enrich_human_input_pause_reasons(
|
||||
pause_reasons,
|
||||
form_tokens_by_form_id=form_token_by_form_id,
|
||||
dispositions_by_form_id=dispositions_by_form_id,
|
||||
expiration_times_by_form_id={
|
||||
form_id: int(expiration_time.timestamp())
|
||||
for form_id, expiration_time in expiration_times_by_form_id.items()
|
||||
@@ -379,6 +383,7 @@ class WorkflowResponseConverter:
|
||||
expiration_time = expiration_times_by_form_id.get(reason.form_id)
|
||||
if expiration_time is None:
|
||||
raise ValueError(f"HumanInputForm not found for pause reason, form_id={reason.form_id}")
|
||||
disposition = dispositions_by_form_id.get(reason.form_id)
|
||||
responses.append(
|
||||
HumanInputRequiredResponse(
|
||||
task_id=task_id,
|
||||
@@ -391,7 +396,8 @@ class WorkflowResponseConverter:
|
||||
inputs=reason.inputs,
|
||||
actions=reason.actions,
|
||||
display_in_ui=display_in_ui_by_form_id.get(reason.form_id, False),
|
||||
form_token=form_token_by_form_id.get(reason.form_id),
|
||||
form_token=disposition.form_token if disposition else None,
|
||||
approval_channels=list(disposition.approval_channels) if disposition else [],
|
||||
resolved_default_values=reason.resolved_default_values,
|
||||
expiration_time=int(expiration_time.timestamp()),
|
||||
),
|
||||
|
||||
@@ -10,6 +10,7 @@ from core.app.entities.app_invoke_entities import (
|
||||
CompletionAppGenerateEntity,
|
||||
)
|
||||
from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
|
||||
from core.db.session_factory import create_session
|
||||
from core.model_manager import ModelInstance
|
||||
from core.moderation.base import ModerationError
|
||||
from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
|
||||
@@ -39,7 +40,10 @@ class CompletionAppRunner(AppRunner):
|
||||
app_config = application_generate_entity.app_config
|
||||
app_config = cast(CompletionAppConfig, app_config)
|
||||
stmt = select(App).where(App.id == app_config.app_id)
|
||||
app_record = db.session.scalar(stmt)
|
||||
with create_session() as session:
|
||||
app_record = session.scalar(stmt)
|
||||
if app_record:
|
||||
session.expunge(app_record)
|
||||
if not app_record:
|
||||
raise ValueError("App not found")
|
||||
|
||||
@@ -174,6 +178,8 @@ class CompletionAppRunner(AppRunner):
|
||||
model=application_generate_entity.model_conf.model,
|
||||
)
|
||||
|
||||
# Release the Flask scoped session before LLM streaming so a checked-out DB connection
|
||||
# is not held for the lifetime of the provider response.
|
||||
db.session.close()
|
||||
|
||||
invoke_result = model_instance.invoke_llm(
|
||||
|
||||
@@ -11,6 +11,7 @@ from core.app.entities.queue_entities import (
|
||||
QueueMessageEndEvent,
|
||||
QueueStopEvent,
|
||||
)
|
||||
from models.model import AppMode
|
||||
|
||||
|
||||
class MessageBasedAppQueueManager(AppQueueManager):
|
||||
@@ -47,4 +48,6 @@ class MessageBasedAppQueueManager(AppQueueManager):
|
||||
self.stop_listen()
|
||||
|
||||
if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped():
|
||||
if self._app_mode == AppMode.ADVANCED_CHAT.value:
|
||||
return
|
||||
raise GenerateTaskStoppedError()
|
||||
|
||||
@@ -3,6 +3,7 @@ import time
|
||||
from typing import cast
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager
|
||||
from core.app.apps.pipeline.pipeline_config_manager import PipelineConfig
|
||||
@@ -14,12 +15,12 @@ from core.app.entities.app_invoke_entities import (
|
||||
build_dify_run_context,
|
||||
)
|
||||
from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer
|
||||
from core.db.session_factory import create_session
|
||||
from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository
|
||||
from core.workflow.node_factory import DifyGraphInitContext, DifyNodeFactory, get_default_root_node_id
|
||||
from core.workflow.system_variables import build_bootstrap_variables, build_system_variables
|
||||
from core.workflow.variable_pool_initializer import add_node_inputs_to_pool, add_variables_to_pool
|
||||
from core.workflow.workflow_entry import WorkflowEntry
|
||||
from extensions.ext_database import db
|
||||
from graphon.enums import WorkflowType
|
||||
from graphon.graph import Graph
|
||||
from graphon.graph_events import GraphEngineEvent, GraphRunFailedEvent
|
||||
@@ -83,22 +84,24 @@ class PipelineRunner(WorkflowBasedAppRunner):
|
||||
user_from = self._resolve_user_from(invoke_from)
|
||||
|
||||
user_id = None
|
||||
if invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
|
||||
end_user = db.session.get(EndUser, self.application_generate_entity.user_id)
|
||||
if end_user:
|
||||
user_id = end_user.session_id
|
||||
else:
|
||||
user_id = self.application_generate_entity.user_id
|
||||
with create_session() as session:
|
||||
if invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
|
||||
end_user = session.get(EndUser, self.application_generate_entity.user_id)
|
||||
if end_user:
|
||||
user_id = end_user.session_id
|
||||
else:
|
||||
user_id = self.application_generate_entity.user_id
|
||||
|
||||
pipeline = db.session.get(Pipeline, app_config.app_id)
|
||||
if not pipeline:
|
||||
raise ValueError("Pipeline not found")
|
||||
pipeline = session.get(Pipeline, app_config.app_id)
|
||||
if not pipeline:
|
||||
raise ValueError("Pipeline not found")
|
||||
|
||||
workflow = self.get_workflow(pipeline=pipeline, workflow_id=app_config.workflow_id)
|
||||
if not workflow:
|
||||
raise ValueError("Workflow not initialized")
|
||||
workflow = self.get_workflow(session=session, pipeline=pipeline, workflow_id=app_config.workflow_id)
|
||||
if not workflow:
|
||||
raise ValueError("Workflow not initialized")
|
||||
|
||||
db.session.close()
|
||||
session.expunge(pipeline)
|
||||
session.expunge(workflow)
|
||||
|
||||
# if only single iteration run is requested
|
||||
if self.application_generate_entity.single_iteration_run or self.application_generate_entity.single_loop_run:
|
||||
@@ -208,12 +211,12 @@ class PipelineRunner(WorkflowBasedAppRunner):
|
||||
)
|
||||
self._handle_event(workflow_entry, event)
|
||||
|
||||
def get_workflow(self, pipeline: Pipeline, workflow_id: str) -> Workflow | None:
|
||||
def get_workflow(self, session: Session, pipeline: Pipeline, workflow_id: str) -> Workflow | None:
|
||||
"""
|
||||
Get workflow
|
||||
"""
|
||||
# fetch workflow by workflow_id
|
||||
workflow = db.session.scalar(
|
||||
workflow = session.scalar(
|
||||
select(Workflow)
|
||||
.where(Workflow.tenant_id == pipeline.tenant_id, Workflow.app_id == pipeline.id, Workflow.id == workflow_id)
|
||||
.limit(1)
|
||||
@@ -298,11 +301,11 @@ class PipelineRunner(WorkflowBasedAppRunner):
|
||||
"""
|
||||
if isinstance(event, GraphRunFailedEvent):
|
||||
if document_id and dataset_id:
|
||||
document = db.session.scalar(
|
||||
select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1)
|
||||
)
|
||||
if document:
|
||||
document.indexing_status = "error"
|
||||
document.error = event.error or "Unknown error"
|
||||
db.session.add(document)
|
||||
db.session.commit()
|
||||
with create_session() as session, session.begin():
|
||||
document = session.scalar(
|
||||
select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1)
|
||||
)
|
||||
if document:
|
||||
document.indexing_status = "error"
|
||||
document.error = event.error or "Unknown error"
|
||||
session.add(document)
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import override
|
||||
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
|
||||
from core.app.apps.exc import GenerateTaskStoppedError
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.app.entities.queue_entities import (
|
||||
AppQueueEvent,
|
||||
@@ -43,6 +42,3 @@ class WorkflowAppQueueManager(AppQueueManager):
|
||||
| QueueWorkflowPartialSuccessEvent,
|
||||
):
|
||||
self.stop_listen()
|
||||
|
||||
if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped():
|
||||
raise GenerateTaskStoppedError()
|
||||
|
||||
@@ -288,6 +288,7 @@ class HumanInputRequiredResponse(StreamResponse):
|
||||
actions: Sequence[UserActionConfig] = Field(default_factory=list)
|
||||
display_in_ui: bool = False
|
||||
form_token: str | None = None
|
||||
approval_channels: list[str] = Field(default_factory=list)
|
||||
resolved_default_values: Mapping[str, Any] = Field(default_factory=dict)
|
||||
expiration_time: int = Field(..., description="Unix timestamp in seconds")
|
||||
|
||||
@@ -311,6 +312,7 @@ class HumanInputRequiredPauseReasonPayload(BaseModel):
|
||||
actions: Sequence[UserActionConfig] = Field(default_factory=list)
|
||||
display_in_ui: bool = False
|
||||
form_token: str | None = None
|
||||
approval_channels: list[str] = Field(default_factory=list)
|
||||
resolved_default_values: Mapping[str, Any] = Field(default_factory=dict)
|
||||
expiration_time: int
|
||||
|
||||
@@ -325,6 +327,7 @@ class HumanInputRequiredPauseReasonPayload(BaseModel):
|
||||
actions=data.actions,
|
||||
display_in_ui=data.display_in_ui,
|
||||
form_token=data.form_token,
|
||||
approval_channels=data.approval_channels,
|
||||
resolved_default_values=data.resolved_default_values,
|
||||
expiration_time=data.expiration_time,
|
||||
)
|
||||
|
||||
@@ -14,6 +14,7 @@ from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from http import HTTPStatus
|
||||
from typing import Any, cast
|
||||
|
||||
import httpx
|
||||
@@ -293,28 +294,27 @@ class StreamableHTTPTransport:
|
||||
json=message.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
headers=headers,
|
||||
) as response:
|
||||
if response.status_code == 202:
|
||||
logger.debug("Received 202 Accepted")
|
||||
return
|
||||
|
||||
if response.status_code == 204:
|
||||
logger.debug("Received 204 No Content")
|
||||
return
|
||||
|
||||
if response.status_code == 404:
|
||||
if isinstance(message.root, JSONRPCRequest):
|
||||
error_msg = (
|
||||
f"MCP server URL returned 404 Not Found: {self.url} "
|
||||
"— verify the server URL is correct and the server is running"
|
||||
if is_initialization
|
||||
else "Session terminated by server"
|
||||
)
|
||||
self._send_session_terminated_error(
|
||||
ctx.server_to_client_queue,
|
||||
message.root.id,
|
||||
message=error_msg,
|
||||
)
|
||||
return
|
||||
match response.status_code:
|
||||
case HTTPStatus.ACCEPTED:
|
||||
logger.debug("Received 202 Accepted")
|
||||
return
|
||||
case HTTPStatus.NO_CONTENT:
|
||||
logger.debug("Received 204 No Content")
|
||||
return
|
||||
case HTTPStatus.NOT_FOUND:
|
||||
if isinstance(message.root, JSONRPCRequest):
|
||||
error_msg = (
|
||||
f"MCP server URL returned 404 Not Found: {self.url} "
|
||||
"— verify the server URL is correct and the server is running"
|
||||
if is_initialization
|
||||
else "Session terminated by server"
|
||||
)
|
||||
self._send_session_terminated_error(
|
||||
ctx.server_to_client_queue,
|
||||
message.root.id,
|
||||
message=error_msg,
|
||||
)
|
||||
return
|
||||
|
||||
response.raise_for_status()
|
||||
if is_initialization:
|
||||
|
||||
@@ -3,7 +3,6 @@ from collections.abc import Generator, Mapping
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
|
||||
from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator
|
||||
@@ -13,10 +12,19 @@ from core.app.apps.completion.app_generator import CompletionAppGenerator
|
||||
from core.app.apps.workflow.app_generator import WorkflowAppGenerator
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig
|
||||
from core.db.session_factory import create_session
|
||||
from core.plugin.backwards_invocation.base import BaseBackwardsInvocation
|
||||
from extensions.ext_database import db
|
||||
from models import Account
|
||||
from models.model import App, AppMode, EndUser
|
||||
from models import Account, TenantAccountJoin
|
||||
from models.model import (
|
||||
App,
|
||||
AppMode,
|
||||
AppModelConfig,
|
||||
AppModelConfigDict,
|
||||
EndUser,
|
||||
load_annotation_reply_config,
|
||||
)
|
||||
from models.workflow import Workflow
|
||||
from services.end_user_service import EndUserService
|
||||
|
||||
|
||||
@@ -30,18 +38,18 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
|
||||
"""Retrieve app parameters."""
|
||||
if app.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
|
||||
workflow = app.workflow
|
||||
workflow = cls._get_workflow(app)
|
||||
if workflow is None:
|
||||
raise ValueError("unexpected app type")
|
||||
|
||||
features_dict: dict[str, Any] = workflow.features_dict
|
||||
user_input_form = workflow.user_input_form(to_old_structure=True)
|
||||
else:
|
||||
app_model_config = app.app_model_config
|
||||
if app_model_config is None:
|
||||
app_model_config_dict = cls._get_app_model_config_dict(app)
|
||||
if app_model_config_dict is None:
|
||||
raise ValueError("unexpected app type")
|
||||
|
||||
features_dict = cast(dict[str, Any], app_model_config.to_dict())
|
||||
features_dict = cast(dict[str, Any], app_model_config_dict)
|
||||
|
||||
user_input_form = features_dict.get("user_input_form", [])
|
||||
|
||||
@@ -68,7 +76,7 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
if not user_id:
|
||||
user = EndUserService.get_or_create_end_user(app)
|
||||
else:
|
||||
user = cls._get_user(user_id)
|
||||
user = cls._get_user(user_id, app)
|
||||
|
||||
conversation_id = conversation_id or ""
|
||||
|
||||
@@ -79,7 +87,10 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
|
||||
return cls.invoke_chat_app(app, user, conversation_id, query, stream, inputs, files)
|
||||
case AppMode.WORKFLOW:
|
||||
return cls.invoke_workflow_app(app, user, stream, inputs, files)
|
||||
workflow = cls._get_workflow(app)
|
||||
if not workflow:
|
||||
raise ValueError("unexpected app type")
|
||||
return cls.invoke_workflow_app(app, workflow, user, stream, inputs, files)
|
||||
case AppMode.COMPLETION:
|
||||
return cls.invoke_completion_app(app, user, stream, inputs, files)
|
||||
case _:
|
||||
@@ -101,7 +112,7 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
"""
|
||||
match app.mode:
|
||||
case AppMode.ADVANCED_CHAT:
|
||||
workflow = app.workflow
|
||||
workflow = cls._get_workflow(app)
|
||||
if not workflow:
|
||||
raise ValueError("unexpected app type")
|
||||
|
||||
@@ -158,6 +169,7 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
def invoke_workflow_app(
|
||||
cls,
|
||||
app: App,
|
||||
workflow: Workflow,
|
||||
user: EndUser | Account,
|
||||
stream: bool,
|
||||
inputs: Mapping,
|
||||
@@ -166,10 +178,6 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
"""
|
||||
invoke workflow app
|
||||
"""
|
||||
workflow = app.workflow
|
||||
if not workflow:
|
||||
raise ValueError("unexpected app type")
|
||||
|
||||
pause_config = PauseStateLayerConfig(
|
||||
session_factory=db.engine,
|
||||
state_owner_user_id=workflow.created_by,
|
||||
@@ -207,16 +215,26 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_user(cls, user_id: str) -> EndUser | Account:
|
||||
def _get_user(cls, user_id: str, app: App) -> EndUser | Account:
|
||||
"""
|
||||
get the user by user id
|
||||
"""
|
||||
with Session(db.engine, expire_on_commit=False) as session:
|
||||
stmt = select(EndUser).where(EndUser.id == user_id)
|
||||
with create_session() as session:
|
||||
stmt = select(EndUser).where(
|
||||
EndUser.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)
|
||||
stmt = select(Account).where(
|
||||
Account.id == user_id,
|
||||
Account.id == TenantAccountJoin.account_id,
|
||||
TenantAccountJoin.tenant_id == app.tenant_id,
|
||||
)
|
||||
user = session.scalar(stmt)
|
||||
if user:
|
||||
session.expunge(user)
|
||||
|
||||
if not user:
|
||||
raise ValueError("user not found")
|
||||
@@ -229,7 +247,10 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
get app
|
||||
"""
|
||||
try:
|
||||
app = db.session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id).limit(1))
|
||||
with create_session() as session:
|
||||
app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id).limit(1))
|
||||
if app:
|
||||
session.expunge(app)
|
||||
except Exception:
|
||||
raise ValueError("app not found")
|
||||
|
||||
@@ -237,3 +258,41 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
|
||||
raise ValueError("app not found")
|
||||
|
||||
return app
|
||||
|
||||
@classmethod
|
||||
def _get_workflow(cls, app: App) -> Workflow | None:
|
||||
"""
|
||||
get workflow without relying on App.workflow's request-scoped session property
|
||||
"""
|
||||
if not app.workflow_id:
|
||||
return None
|
||||
|
||||
with create_session() as session:
|
||||
workflow = session.scalar(
|
||||
select(Workflow)
|
||||
.where(Workflow.id == app.workflow_id, Workflow.tenant_id == app.tenant_id, Workflow.app_id == app.id)
|
||||
.limit(1)
|
||||
)
|
||||
if workflow:
|
||||
session.expunge(workflow)
|
||||
return workflow
|
||||
|
||||
@classmethod
|
||||
def _get_app_model_config_dict(cls, app: App) -> AppModelConfigDict | None:
|
||||
"""
|
||||
get app model config features without relying on request-scoped session-backed model properties
|
||||
"""
|
||||
if not app.app_model_config_id:
|
||||
return None
|
||||
|
||||
with create_session() as session:
|
||||
app_model_config = session.scalar(
|
||||
select(AppModelConfig)
|
||||
.where(AppModelConfig.id == app.app_model_config_id, AppModelConfig.app_id == app.id)
|
||||
.limit(1)
|
||||
)
|
||||
if app_model_config is None:
|
||||
return None
|
||||
|
||||
annotation_reply = load_annotation_reply_config(session, app_model_config.app_id)
|
||||
return app_model_config.to_dict(annotation_reply=annotation_reply)
|
||||
|
||||
@@ -14,6 +14,12 @@ from core.rag.extractor.watercrawl.exceptions import (
|
||||
|
||||
WATERCRAWL_REQUEST_TIMEOUT: httpx.Timeout = httpx.Timeout(30.0, connect=5.0)
|
||||
|
||||
# The crawl-status stream is a long-lived SSE connection that can stay open for
|
||||
# the whole duration of a crawl, so it keeps an unbounded read while still
|
||||
# capping the initial connection. Regular requests use WATERCRAWL_REQUEST_TIMEOUT
|
||||
# so a stalled endpoint can't hang a worker forever.
|
||||
_STREAM_TIMEOUT = httpx.Timeout(None, connect=10.0)
|
||||
|
||||
|
||||
class SpiderOptions(TypedDict):
|
||||
max_depth: int
|
||||
@@ -50,6 +56,8 @@ class BaseAPIClient:
|
||||
"User-Agent": "WaterCrawl-Plugin",
|
||||
"Accept-Language": "en-US",
|
||||
}
|
||||
# Regular requests use WATERCRAWL_REQUEST_TIMEOUT; the long-lived
|
||||
# crawl-status stream overrides it with _STREAM_TIMEOUT in _request.
|
||||
return httpx.Client(headers=headers, timeout=WATERCRAWL_REQUEST_TIMEOUT)
|
||||
|
||||
def _request(
|
||||
@@ -63,7 +71,7 @@ class BaseAPIClient:
|
||||
stream = kwargs.pop("stream", False)
|
||||
url = urljoin(self.base_url, endpoint)
|
||||
if stream:
|
||||
request = self.session.build_request(method, url, params=query_params, json=data)
|
||||
request = self.session.build_request(method, url, params=query_params, json=data, timeout=_STREAM_TIMEOUT)
|
||||
return self.session.send(request, stream=True, **kwargs)
|
||||
|
||||
return self.session.request(method, url, params=query_params, json=data, **kwargs)
|
||||
|
||||
@@ -22,23 +22,35 @@ class RBACPermission(StrEnum):
|
||||
|
||||
APP_VIEW_LAYOUT = "app_view_layout"
|
||||
APP_TEST_AND_RUN = "app_test_and_run"
|
||||
APP_PREVIEW = "app_preview"
|
||||
APP_CREATE_AND_MANAGEMENT = "app_create_and_management"
|
||||
APP_RELEASE_AND_VERSION = "app_release_and_version"
|
||||
APP_IMPORT_EXPORT_DSL = "app_import_export_dsl"
|
||||
APP_EDIT = "app_edit"
|
||||
APP_MONITOR = "app_monitor"
|
||||
APP_DELETE = "app_delete"
|
||||
APP_ACCESS_CONFIG = "app_access_config"
|
||||
|
||||
DATASET_PREVIEW = "dataset_preview"
|
||||
DATASET_READONLY = "dataset_readonly"
|
||||
DATASET_EDIT = "dataset_edit"
|
||||
DATASET_CREATE_AND_MANAGEMENT = "dataset_create_and_management"
|
||||
DATASET_PIPELINE_TEST = "dataset_pipeline_test"
|
||||
DATASET_DOCUMENT_DOWNLOAD = "dataset_document_download"
|
||||
DATASET_RETRIEVAL_RECALL = "dataset_retrieval_recall"
|
||||
DATASET_USE = "dataset_use"
|
||||
DATASET_DELETE_FILE = "dataset_delete_file"
|
||||
DATASET_PIPELINE_RELEASE = "dataset_pipeline_release"
|
||||
DATASET_DELETE = "dataset_delete"
|
||||
DATASET_ACCESS_CONFIG = "dataset_access_config"
|
||||
DATASET_API_KEY_MANAGE = "dataset_api_key_manage"
|
||||
DATASET_EXTERNAL_CONNECT = "dataset_external_connect"
|
||||
DATASET_IMPORT_EXPORT_DSL = "dataset_import_export_dsl"
|
||||
|
||||
WORKSPACE_MEMBER_MANAGE = "workspace_member_manage"
|
||||
WORKSPACE_ROLE_MANAGE = "workspace_role_manage"
|
||||
API_EXTENSION_MANAGE = "api_extension_manage"
|
||||
CUSTOMIZATION_MANAGE = "customization_manage"
|
||||
|
||||
SNIPPETS_CREATE_AND_MODIFY = "snippets_create_and_modify"
|
||||
SNIPPETS_MANAGE = "snippets_management"
|
||||
@@ -49,6 +61,7 @@ class RBACPermission(StrEnum):
|
||||
PLUGIN_DEBUG = "plugin_debug"
|
||||
|
||||
CREDENTIAL_USE = "credential_use"
|
||||
CREDENTIAL_CREATE = "credential_create"
|
||||
CREDENTIAL_MANAGE = "credential_manage"
|
||||
|
||||
TOOL_MANAGE = "tool_manage"
|
||||
|
||||
@@ -359,15 +359,16 @@ class ApiTool(Tool):
|
||||
if value is None:
|
||||
return None
|
||||
elif property["type"] == "object" or property["type"] == "array":
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return json.loads(value)
|
||||
except ValueError:
|
||||
match value:
|
||||
case str():
|
||||
try:
|
||||
return json.loads(value)
|
||||
except ValueError:
|
||||
return value
|
||||
case dict():
|
||||
return value
|
||||
case _:
|
||||
return value
|
||||
elif isinstance(value, dict):
|
||||
return value
|
||||
else:
|
||||
return value
|
||||
else:
|
||||
raise ValueError(f"Invalid type {property['type']} for property {property}")
|
||||
elif "anyOf" in property and isinstance(property["anyOf"], list):
|
||||
|
||||
@@ -12,60 +12,61 @@ from collections.abc import Sequence
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.workflow.human_input_policy import HumanInputSurface, get_preferred_form_token
|
||||
from core.workflow.human_input_policy import (
|
||||
FormDisposition,
|
||||
HumanInputSurface,
|
||||
disposition_for_surface,
|
||||
)
|
||||
from extensions.ext_database import db
|
||||
from models.human_input import HumanInputFormRecipient, RecipientType
|
||||
|
||||
|
||||
def load_form_dispositions_by_form_id(
|
||||
form_ids: Sequence[str],
|
||||
*,
|
||||
session: Session | None = None,
|
||||
surface: HumanInputSurface | None = None,
|
||||
) -> dict[str, FormDisposition]:
|
||||
"""Resolve each paused form's resume token and approval channels for `surface`."""
|
||||
unique_form_ids = list(dict.fromkeys(form_ids))
|
||||
if not unique_form_ids:
|
||||
return {}
|
||||
|
||||
if session is not None:
|
||||
return _load_form_dispositions_by_form_id(session, unique_form_ids, surface=surface)
|
||||
|
||||
with Session(bind=db.engine, expire_on_commit=False) as new_session:
|
||||
return _load_form_dispositions_by_form_id(new_session, unique_form_ids, surface=surface)
|
||||
|
||||
|
||||
def _load_form_dispositions_by_form_id(
|
||||
session: Session,
|
||||
form_ids: Sequence[str],
|
||||
*,
|
||||
surface: HumanInputSurface | None,
|
||||
) -> dict[str, FormDisposition]:
|
||||
recipients_by_form_id: dict[str, list[tuple[RecipientType, str]]] = {}
|
||||
stmt = select(HumanInputFormRecipient).where(HumanInputFormRecipient.form_id.in_(form_ids))
|
||||
for recipient in session.scalars(stmt):
|
||||
recipients_by_form_id.setdefault(recipient.form_id, []).append(
|
||||
(recipient.recipient_type, recipient.access_token or "")
|
||||
)
|
||||
return {
|
||||
form_id: disposition_for_surface(recipients, surface=surface)
|
||||
for form_id, recipients in recipients_by_form_id.items()
|
||||
}
|
||||
|
||||
|
||||
def load_form_tokens_by_form_id(
|
||||
form_ids: Sequence[str],
|
||||
*,
|
||||
session: Session | None = None,
|
||||
surface: HumanInputSurface | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Load the preferred access token for each human input form."""
|
||||
unique_form_ids = list(dict.fromkeys(form_ids))
|
||||
if not unique_form_ids:
|
||||
return {}
|
||||
|
||||
if session is not None:
|
||||
return _load_form_tokens_by_form_id(session, unique_form_ids, surface=surface)
|
||||
|
||||
with Session(bind=db.engine, expire_on_commit=False) as new_session:
|
||||
return _load_form_tokens_by_form_id(new_session, unique_form_ids, surface=surface)
|
||||
|
||||
|
||||
def _load_form_tokens_by_form_id(
|
||||
session: Session,
|
||||
form_ids: Sequence[str],
|
||||
*,
|
||||
surface: HumanInputSurface | None = None,
|
||||
) -> dict[str, str]:
|
||||
recipients_by_form_id: dict[str, list[tuple[RecipientType, str]]] = {}
|
||||
stmt = select(HumanInputFormRecipient).where(HumanInputFormRecipient.form_id.in_(form_ids))
|
||||
for recipient in session.scalars(stmt):
|
||||
if not recipient.access_token:
|
||||
continue
|
||||
recipients_by_form_id.setdefault(recipient.form_id, []).append(
|
||||
(recipient.recipient_type, recipient.access_token)
|
||||
)
|
||||
|
||||
tokens_by_form_id: dict[str, str] = {}
|
||||
for form_id, recipients in recipients_by_form_id.items():
|
||||
token = _get_surface_form_token(recipients, surface=surface)
|
||||
if token is not None:
|
||||
tokens_by_form_id[form_id] = token
|
||||
return tokens_by_form_id
|
||||
|
||||
|
||||
def _get_surface_form_token(
|
||||
recipients: Sequence[tuple[RecipientType, str]],
|
||||
*,
|
||||
surface: HumanInputSurface | None,
|
||||
) -> str | None:
|
||||
if surface in {HumanInputSurface.SERVICE_API, HumanInputSurface.OPENAPI}:
|
||||
for recipient_type, token in recipients:
|
||||
if recipient_type == RecipientType.STANDALONE_WEB_APP and token:
|
||||
return token
|
||||
|
||||
return get_preferred_form_token(recipients)
|
||||
"""Resume tokens only, for callers that don't surface approval channels."""
|
||||
dispositions = load_form_dispositions_by_form_id(form_ids, session=session, surface=surface)
|
||||
return {
|
||||
form_id: disposition.form_token
|
||||
for form_id, disposition in dispositions.items()
|
||||
if disposition.form_token is not None
|
||||
}
|
||||
|
||||
@@ -2,14 +2,14 @@ from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
from graphon.entities.pause_reason import HumanInputRequired, PauseReason, PauseReasonType
|
||||
from graphon.nodes.human_input.entities import FormInputConfig, SelectInputConfig
|
||||
from graphon.nodes.human_input.enums import ValueSourceType
|
||||
from graphon.runtime.graph_runtime_state_protocol import ReadOnlyVariablePool
|
||||
from graphon.variables import ArrayStringSegment
|
||||
from models.human_input import RecipientType
|
||||
from models.human_input import ApprovalChannel, RecipientType
|
||||
|
||||
|
||||
class HumanInputSurface(StrEnum):
|
||||
@@ -20,7 +20,7 @@ class HumanInputSurface(StrEnum):
|
||||
|
||||
# SERVICE_API and OPENAPI are intentionally narrower than CONSOLE: token callers
|
||||
# should only be able to act on end-user web forms, not internal console flows.
|
||||
_ALLOWED_RECIPIENT_TYPES_BY_SURFACE: dict[HumanInputSurface, frozenset[RecipientType]] = {
|
||||
ALLOWED_RECIPIENT_TYPES_BY_SURFACE: dict[HumanInputSurface, frozenset[RecipientType]] = {
|
||||
HumanInputSurface.SERVICE_API: frozenset({RecipientType.STANDALONE_WEB_APP}),
|
||||
HumanInputSurface.CONSOLE: frozenset({RecipientType.CONSOLE, RecipientType.BACKSTAGE}),
|
||||
HumanInputSurface.OPENAPI: frozenset({RecipientType.STANDALONE_WEB_APP}),
|
||||
@@ -41,7 +41,7 @@ def is_recipient_type_allowed_for_surface(
|
||||
) -> bool:
|
||||
if recipient_type is None:
|
||||
return False
|
||||
return recipient_type in _ALLOWED_RECIPIENT_TYPES_BY_SURFACE[surface]
|
||||
return recipient_type in ALLOWED_RECIPIENT_TYPES_BY_SURFACE[surface]
|
||||
|
||||
|
||||
def get_preferred_form_token(
|
||||
@@ -59,10 +59,39 @@ def get_preferred_form_token(
|
||||
return chosen_token
|
||||
|
||||
|
||||
class FormDisposition(NamedTuple):
|
||||
"""How a paused form resolves for one API surface.
|
||||
|
||||
A form's recipients split into those the surface may act on (yielding a resume
|
||||
`form_token`) and those it may not (their channels named in `approval_channels`
|
||||
so the caller is told where approval actually happens instead).
|
||||
"""
|
||||
|
||||
form_token: str | None
|
||||
approval_channels: list[ApprovalChannel]
|
||||
|
||||
|
||||
def disposition_for_surface(
|
||||
recipients: Sequence[tuple[RecipientType, str]],
|
||||
*,
|
||||
surface: HumanInputSurface | None,
|
||||
) -> FormDisposition:
|
||||
if surface is None:
|
||||
return FormDisposition(form_token=get_preferred_form_token(recipients), approval_channels=[])
|
||||
allowed = ALLOWED_RECIPIENT_TYPES_BY_SURFACE[surface]
|
||||
actionable = [(recipient_type, token) for recipient_type, token in recipients if recipient_type in allowed]
|
||||
return FormDisposition(
|
||||
form_token=get_preferred_form_token(actionable),
|
||||
approval_channels=sorted(
|
||||
{recipient_type.approval_channel for recipient_type, _ in recipients if recipient_type not in allowed}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def enrich_human_input_pause_reasons(
|
||||
reasons: Sequence[Mapping[str, Any]],
|
||||
*,
|
||||
form_tokens_by_form_id: Mapping[str, str],
|
||||
dispositions_by_form_id: Mapping[str, FormDisposition],
|
||||
expiration_times_by_form_id: Mapping[str, int],
|
||||
) -> list[dict[str, Any]]:
|
||||
enriched: list[dict[str, Any]] = []
|
||||
@@ -71,7 +100,9 @@ def enrich_human_input_pause_reasons(
|
||||
if updated.get("TYPE") == PauseReasonType.HUMAN_INPUT_REQUIRED:
|
||||
form_id = updated.get("form_id")
|
||||
if isinstance(form_id, str):
|
||||
updated["form_token"] = form_tokens_by_form_id.get(form_id)
|
||||
disposition = dispositions_by_form_id.get(form_id)
|
||||
updated["form_token"] = disposition.form_token if disposition else None
|
||||
updated["approval_channels"] = list(disposition.approval_channels) if disposition else []
|
||||
expiration_time = expiration_times_by_form_id.get(form_id)
|
||||
if expiration_time is not None:
|
||||
updated["expiration_time"] = expiration_time
|
||||
|
||||
@@ -5,6 +5,7 @@ def init_app(app: DifyApp):
|
||||
from commands import (
|
||||
add_qdrant_index,
|
||||
archive_workflow_runs,
|
||||
archive_workflow_runs_plan,
|
||||
backfill_plugin_auto_upgrade,
|
||||
clean_expired_messages,
|
||||
clean_workflow_runs,
|
||||
@@ -72,6 +73,7 @@ def init_app(app: DifyApp):
|
||||
setup_datasource_oauth_client,
|
||||
transform_datasource_credentials,
|
||||
install_rag_pipeline_plugins,
|
||||
archive_workflow_runs_plan,
|
||||
archive_workflow_runs,
|
||||
delete_archived_workflow_runs,
|
||||
restore_workflow_runs,
|
||||
|
||||
@@ -25,7 +25,7 @@ from extensions.redis_names import (
|
||||
serialize_redis_name_args,
|
||||
)
|
||||
from libs.broadcast_channel.channel import BroadcastChannel as BroadcastChannelProtocol
|
||||
from libs.broadcast_channel.redis.channel import BroadcastChannel as RedisBroadcastChannel
|
||||
from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel
|
||||
from libs.broadcast_channel.redis.sharded_channel import ShardedRedisBroadcastChannel
|
||||
from libs.broadcast_channel.redis.streams_channel import StreamsBroadcastChannel
|
||||
|
||||
@@ -457,16 +457,14 @@ def init_app(app: DifyApp):
|
||||
|
||||
def get_pubsub_broadcast_channel() -> BroadcastChannelProtocol:
|
||||
assert _pubsub_redis_client is not None, "PubSub redis Client should be initialized here."
|
||||
join_timeout_ms = dify_config.PUBSUB_LISTENER_JOIN_TIMEOUT_MS
|
||||
if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "sharded":
|
||||
return ShardedRedisBroadcastChannel(_pubsub_redis_client, join_timeout_ms=join_timeout_ms)
|
||||
return ShardedRedisBroadcastChannel(_pubsub_redis_client)
|
||||
if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "streams":
|
||||
return StreamsBroadcastChannel(
|
||||
_pubsub_redis_client,
|
||||
retention_seconds=dify_config.PUBSUB_STREAMS_RETENTION_SECONDS,
|
||||
join_timeout_ms=join_timeout_ms,
|
||||
)
|
||||
return RedisBroadcastChannel(_pubsub_redis_client, join_timeout_ms=join_timeout_ms)
|
||||
return RedisBroadcastChannel(_pubsub_redis_client)
|
||||
|
||||
|
||||
def redis_fallback[T](default_return: T | None = None): # type: ignore
|
||||
|
||||
@@ -291,6 +291,11 @@ class AgentConfigSnapshotListResponse(ResponseModel):
|
||||
data: list[AgentConfigSnapshotSummaryResponse]
|
||||
|
||||
|
||||
class AgentConfigSnapshotRestoreResponse(ResponseModel):
|
||||
result: Literal["success"]
|
||||
active_config_snapshot_id: str
|
||||
|
||||
|
||||
class AgentComposerAgentResponse(ResponseModel):
|
||||
id: str
|
||||
name: str
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from .channel import BroadcastChannel
|
||||
from .pubsub_channel import BroadcastChannel
|
||||
from .sharded_channel import ShardedRedisBroadcastChannel
|
||||
|
||||
__all__ = ["BroadcastChannel", "ShardedRedisBroadcastChannel"]
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import Any, Self, override
|
||||
|
||||
from libs.broadcast_channel.channel import Subscription
|
||||
from libs.broadcast_channel.exc import SubscriptionClosedError
|
||||
from libs.broadcast_channel.signals import SIG_CLOSE
|
||||
from redis import Redis, RedisCluster
|
||||
from redis.client import PubSub
|
||||
|
||||
@@ -26,8 +27,6 @@ class RedisSubscriptionBase(Subscription):
|
||||
client: Redis | RedisCluster,
|
||||
pubsub: PubSub,
|
||||
topic: str,
|
||||
*,
|
||||
join_timeout_ms: int = 2000,
|
||||
):
|
||||
# The _pubsub is None only if the subscription is closed.
|
||||
self._client = client
|
||||
@@ -39,11 +38,6 @@ class RedisSubscriptionBase(Subscription):
|
||||
self._listener_thread: threading.Thread | None = None
|
||||
self._start_lock = threading.Lock()
|
||||
self._started = False
|
||||
# Max time close() will wait for the listener thread to finish before
|
||||
# returning. Bounds SSE close tail latency. The listener is a daemon
|
||||
# and exits on its own within one poll window (~1s), so a low value
|
||||
# here just means close() returns sooner without breaking anything.
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
|
||||
def _start_if_needed(self) -> None:
|
||||
"""Start the subscription if not already started."""
|
||||
@@ -90,6 +84,11 @@ class RedisSubscriptionBase(Subscription):
|
||||
if raw_message is None:
|
||||
continue
|
||||
|
||||
# If close() sent a control event to unblock us, exit immediately
|
||||
# without processing any message — the subscription is shutting down.
|
||||
if self._closed.is_set():
|
||||
break
|
||||
|
||||
if raw_message.get("type") != self._get_message_type():
|
||||
continue
|
||||
|
||||
@@ -119,6 +118,8 @@ class RedisSubscriptionBase(Subscription):
|
||||
continue
|
||||
|
||||
self._enqueue_message(payload_bytes)
|
||||
if payload_bytes == SIG_CLOSE:
|
||||
break
|
||||
|
||||
_logger.debug("%s listener thread stopped for channel %s", self._get_subscription_type().title(), self._topic)
|
||||
try:
|
||||
@@ -164,14 +165,20 @@ class RedisSubscriptionBase(Subscription):
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
if self._closed.is_set():
|
||||
return
|
||||
|
||||
yield item
|
||||
|
||||
@override
|
||||
def __iter__(self) -> Iterator[bytes]:
|
||||
"""Return an iterator over messages from the subscription."""
|
||||
if self._closed.is_set():
|
||||
raise SubscriptionClosedError(f"The Redis {self._get_subscription_type()} subscription is closed")
|
||||
self._start_if_needed()
|
||||
return iter(())
|
||||
try:
|
||||
self._start_if_needed()
|
||||
except SubscriptionClosedError:
|
||||
return iter(())
|
||||
return iter(self._message_iterator())
|
||||
|
||||
@override
|
||||
@@ -208,24 +215,55 @@ class RedisSubscriptionBase(Subscription):
|
||||
@override
|
||||
def close(self) -> None:
|
||||
"""Close the subscription and clean up resources."""
|
||||
if self._closed.is_set():
|
||||
return
|
||||
with self._start_lock:
|
||||
if self._closed.is_set():
|
||||
return
|
||||
|
||||
self._closed.set()
|
||||
listener = self._listener_thread
|
||||
self._listener_thread = None
|
||||
started = self._started
|
||||
|
||||
if started:
|
||||
self._unblock_message_iterator()
|
||||
|
||||
# Send a control event on the same Redis channel to unblock the
|
||||
self._publish_close_event()
|
||||
|
||||
self._closed.set()
|
||||
# NOTE: PubSub is not thread-safe. More specifically, the `PubSub.close` method and the
|
||||
# message retrieval method should NOT be called concurrently.
|
||||
#
|
||||
# Due to the restriction above, the PubSub cleanup logic happens inside the consumer thread.
|
||||
listener = self._listener_thread
|
||||
if listener is not None:
|
||||
listener.join(timeout=self._join_timeout_ms / 1000.0)
|
||||
self._listener_thread = None
|
||||
if listener is not None and listener.is_alive():
|
||||
listener.join(timeout=2)
|
||||
|
||||
def _unblock_message_iterator(self) -> None:
|
||||
try:
|
||||
self._queue.put_nowait(SIG_CLOSE)
|
||||
except queue.Full:
|
||||
try:
|
||||
self._queue.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
try:
|
||||
self._queue.put_nowait(SIG_CLOSE)
|
||||
except queue.Full:
|
||||
pass
|
||||
|
||||
# Abstract methods to be implemented by subclasses
|
||||
def _get_subscription_type(self) -> str:
|
||||
"""Return the subscription type (e.g., 'regular' or 'sharded')."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _publish_close_event(self) -> None:
|
||||
"""Publish a control event on the Redis channel to unblock the listener.
|
||||
|
||||
This is called by close() after setting _closed. The subclass should
|
||||
publish an empty message on the same topic so that a blocking
|
||||
get_message() call in the listener thread returns promptly.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _subscribe(self) -> None:
|
||||
"""Subscribe to the Redis topic using the appropriate command."""
|
||||
raise NotImplementedError
|
||||
|
||||
+12
-10
@@ -1,13 +1,17 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, override
|
||||
|
||||
from extensions.redis_names import serialize_redis_name
|
||||
from libs.broadcast_channel.channel import Producer, Subscriber, Subscription
|
||||
from libs.broadcast_channel.signals import SIG_CLOSE
|
||||
from redis import Redis, RedisCluster
|
||||
|
||||
from ._subscription import RedisSubscriptionBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BroadcastChannel:
|
||||
"""
|
||||
@@ -22,16 +26,11 @@ class BroadcastChannel:
|
||||
def __init__(
|
||||
self,
|
||||
redis_client: Redis | RedisCluster,
|
||||
*,
|
||||
join_timeout_ms: int = 2000,
|
||||
):
|
||||
self._client = redis_client
|
||||
# See `RedisSubscriptionBase._join_timeout_ms`: how long close()
|
||||
# waits for the listener thread before returning.
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
|
||||
def topic(self, topic: str) -> Topic:
|
||||
return Topic(self._client, topic, join_timeout_ms=self._join_timeout_ms)
|
||||
return Topic(self._client, topic)
|
||||
|
||||
|
||||
class Topic:
|
||||
@@ -39,13 +38,10 @@ class Topic:
|
||||
self,
|
||||
redis_client: Redis | RedisCluster,
|
||||
topic: str,
|
||||
*,
|
||||
join_timeout_ms: int = 2000,
|
||||
):
|
||||
self._client = redis_client
|
||||
self._topic = topic
|
||||
self._redis_topic = serialize_redis_name(topic)
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
|
||||
def as_producer(self) -> Producer:
|
||||
return self
|
||||
@@ -61,7 +57,6 @@ class Topic:
|
||||
client=self._client,
|
||||
pubsub=self._client.pubsub(),
|
||||
topic=self._redis_topic,
|
||||
join_timeout_ms=self._join_timeout_ms,
|
||||
)
|
||||
|
||||
|
||||
@@ -72,6 +67,13 @@ class _RedisSubscription(RedisSubscriptionBase):
|
||||
def _get_subscription_type(self) -> str:
|
||||
return "regular"
|
||||
|
||||
@override
|
||||
def _publish_close_event(self) -> None:
|
||||
try:
|
||||
self._client.publish(self._topic, SIG_CLOSE)
|
||||
except Exception:
|
||||
logger.exception("failed to publish close event")
|
||||
|
||||
@override
|
||||
def _subscribe(self) -> None:
|
||||
assert self._pubsub is not None
|
||||
@@ -1,13 +1,17 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, override
|
||||
|
||||
from extensions.redis_names import serialize_redis_name
|
||||
from libs.broadcast_channel.channel import Producer, Subscriber, Subscription
|
||||
from libs.broadcast_channel.signals import SIG_CLOSE
|
||||
from redis import Redis, RedisCluster
|
||||
|
||||
from ._subscription import RedisSubscriptionBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ShardedRedisBroadcastChannel:
|
||||
"""
|
||||
@@ -20,14 +24,11 @@ class ShardedRedisBroadcastChannel:
|
||||
def __init__(
|
||||
self,
|
||||
redis_client: Redis | RedisCluster,
|
||||
*,
|
||||
join_timeout_ms: int = 2000,
|
||||
):
|
||||
self._client = redis_client
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
|
||||
def topic(self, topic: str) -> ShardedTopic:
|
||||
return ShardedTopic(self._client, topic, join_timeout_ms=self._join_timeout_ms)
|
||||
return ShardedTopic(self._client, topic)
|
||||
|
||||
|
||||
class ShardedTopic:
|
||||
@@ -35,13 +36,10 @@ class ShardedTopic:
|
||||
self,
|
||||
redis_client: Redis | RedisCluster,
|
||||
topic: str,
|
||||
*,
|
||||
join_timeout_ms: int = 2000,
|
||||
):
|
||||
self._client = redis_client
|
||||
self._topic = topic
|
||||
self._redis_topic = serialize_redis_name(topic)
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
|
||||
def as_producer(self) -> Producer:
|
||||
return self
|
||||
@@ -57,7 +55,6 @@ class ShardedTopic:
|
||||
client=self._client,
|
||||
pubsub=self._client.pubsub(),
|
||||
topic=self._redis_topic,
|
||||
join_timeout_ms=self._join_timeout_ms,
|
||||
)
|
||||
|
||||
|
||||
@@ -68,6 +65,13 @@ class _RedisShardedSubscription(RedisSubscriptionBase):
|
||||
def _get_subscription_type(self) -> str:
|
||||
return "sharded"
|
||||
|
||||
@override
|
||||
def _publish_close_event(self) -> None:
|
||||
try:
|
||||
self._client.spublish(self._topic, SIG_CLOSE) # type: ignore[attr-defined,union-attr]
|
||||
except Exception:
|
||||
logger.exception("failed to publish close event")
|
||||
|
||||
@override
|
||||
def _subscribe(self) -> None:
|
||||
assert self._pubsub is not None
|
||||
|
||||
@@ -9,6 +9,7 @@ from typing import Self, override
|
||||
from extensions.redis_names import serialize_redis_name
|
||||
from libs.broadcast_channel.channel import Producer, Subscriber, Subscription
|
||||
from libs.broadcast_channel.exc import SubscriptionClosedError
|
||||
from libs.broadcast_channel.signals import SIG_CLOSE
|
||||
from redis import Redis, RedisCluster
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -29,20 +30,15 @@ class StreamsBroadcastChannel:
|
||||
redis_client: Redis | RedisCluster,
|
||||
*,
|
||||
retention_seconds: int = 600,
|
||||
join_timeout_ms: int = 2000,
|
||||
):
|
||||
self._client = redis_client
|
||||
self._retention_seconds = max(int(retention_seconds or 0), 0)
|
||||
# Max time close() will wait for the listener thread to finish.
|
||||
# See `_StreamsSubscription._join_timeout_ms` for the rationale.
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
|
||||
def topic(self, topic: str) -> StreamsTopic:
|
||||
return StreamsTopic(
|
||||
self._client,
|
||||
topic,
|
||||
retention_seconds=self._retention_seconds,
|
||||
join_timeout_ms=self._join_timeout_ms,
|
||||
)
|
||||
|
||||
|
||||
@@ -53,13 +49,11 @@ class StreamsTopic:
|
||||
topic: str,
|
||||
*,
|
||||
retention_seconds: int = 600,
|
||||
join_timeout_ms: int = 2000,
|
||||
):
|
||||
self._client = redis_client
|
||||
self._topic = topic
|
||||
self._key = serialize_redis_name(f"stream:{topic}")
|
||||
self._retention_seconds = retention_seconds
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
self.max_length = 5000
|
||||
|
||||
def as_producer(self) -> Producer:
|
||||
@@ -77,23 +71,15 @@ class StreamsTopic:
|
||||
return self
|
||||
|
||||
def subscribe(self) -> Subscription:
|
||||
return _StreamsSubscription(self._client, self._key, join_timeout_ms=self._join_timeout_ms)
|
||||
return _StreamsSubscription(self._client, self._key)
|
||||
|
||||
|
||||
class _StreamsSubscription(Subscription):
|
||||
_SENTINEL = object()
|
||||
|
||||
def __init__(self, client: Redis | RedisCluster, key: str, *, join_timeout_ms: int = 2000):
|
||||
def __init__(self, client: Redis | RedisCluster, key: str):
|
||||
self._client = client
|
||||
self._key = key
|
||||
# Max time close() will wait for the listener thread to finish before
|
||||
# returning. Bounds SSE close tail latency: the listener blocks on
|
||||
# XREAD with BLOCK=1000ms, so close() naturally waits up to ~1s for
|
||||
# the thread to notice _closed. Setting this lower lets close()
|
||||
# return promptly while the daemon listener exits on its own within
|
||||
# one BLOCK window - safe because the listener holds no critical
|
||||
# state. ``0`` means close() does not wait at all.
|
||||
self._join_timeout_ms = max(int(join_timeout_ms or 0), 0)
|
||||
|
||||
self._queue: queue.Queue[object] = queue.Queue()
|
||||
|
||||
@@ -106,7 +92,6 @@ class _StreamsSubscription(Subscription):
|
||||
# reading and writing the _listener / `_closed` attribute.
|
||||
self._lock = threading.Lock()
|
||||
self._closed: bool = False
|
||||
# self._closed = threading.Event()
|
||||
self._listener: threading.Thread | None = None
|
||||
|
||||
def _listen(self) -> None:
|
||||
@@ -144,6 +129,8 @@ class _StreamsSubscription(Subscription):
|
||||
case bytes() | bytearray():
|
||||
data_bytes = bytes(data)
|
||||
if data_bytes is not None:
|
||||
if data_bytes == SIG_CLOSE:
|
||||
break
|
||||
self._queue.put_nowait(data_bytes)
|
||||
last_id = entry_id
|
||||
finally:
|
||||
@@ -203,6 +190,13 @@ class _StreamsSubscription(Subscription):
|
||||
assert isinstance(item, (bytes, bytearray)), "Unexpected item type in stream queue"
|
||||
return bytes(item)
|
||||
|
||||
def _publish_close_event(self) -> None:
|
||||
"""Publish an empty message to the stream to unblock the listener's xread."""
|
||||
try:
|
||||
self._client.xadd(self._key, {b"data": SIG_CLOSE})
|
||||
except Exception:
|
||||
logger.exception("failed to publish close event")
|
||||
|
||||
@override
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
@@ -212,16 +206,17 @@ class _StreamsSubscription(Subscription):
|
||||
listener = self._listener
|
||||
if listener is not None:
|
||||
self._listener = None
|
||||
# We close the listener outside of the with block to avoid holding the
|
||||
# lock for a long time.
|
||||
|
||||
if listener is not None:
|
||||
self._publish_close_event()
|
||||
|
||||
if listener is not None and listener.is_alive():
|
||||
listener.join(timeout=self._join_timeout_ms / 1000.0)
|
||||
listener.join(timeout=2)
|
||||
if listener.is_alive():
|
||||
logger.debug(
|
||||
"Streams subscription listener for key %s did not stop within %dms; "
|
||||
"Streams subscription listener for key %s did not stop after join; "
|
||||
"daemon thread will exit on its own within one poll window.",
|
||||
self._key,
|
||||
self._join_timeout_ms,
|
||||
)
|
||||
|
||||
# Context manager helpers
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
SIG_CLOSE = b"__closed__"
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
"""agent drive skill metadata refactor
|
||||
|
||||
Revision ID: b2515f9d4c2a
|
||||
Revises: 4f7b2c8d9a10
|
||||
Create Date: 2026-06-18 23:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import mysql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b2515f9d4c2a"
|
||||
down_revision = "4f7b2c8d9a10"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"agent_drive_files",
|
||||
sa.Column("is_skill", sa.Boolean(), nullable=False, server_default=sa.text("false")),
|
||||
)
|
||||
op.add_column(
|
||||
"agent_drive_files",
|
||||
sa.Column("skill_metadata", sa.Text().with_variant(mysql.LONGTEXT(), "mysql"), nullable=True),
|
||||
)
|
||||
op.create_index(
|
||||
"agent_drive_files_tenant_agent_is_skill_key_idx",
|
||||
"agent_drive_files",
|
||||
["tenant_id", "agent_id", "is_skill", "key"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("agent_drive_files_tenant_agent_is_skill_key_idx", table_name="agent_drive_files")
|
||||
op.drop_column("agent_drive_files", "skill_metadata")
|
||||
op.drop_column("agent_drive_files", "is_skill")
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
"""add agent debug conversations
|
||||
|
||||
Revision ID: c8f4a6b2d3e1
|
||||
Revises: b2515f9d4c2a
|
||||
Create Date: 2026-06-22 10:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
import models
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "c8f4a6b2d3e1"
|
||||
down_revision = "b2515f9d4c2a"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _is_pg(conn) -> bool:
|
||||
return conn.dialect.name == "postgresql"
|
||||
|
||||
|
||||
def _uuid_column(name: str, *, nullable: bool = False, primary_key: bool = False) -> sa.Column:
|
||||
kwargs = {"nullable": nullable, "primary_key": primary_key}
|
||||
if primary_key and _is_pg(op.get_bind()):
|
||||
kwargs["server_default"] = sa.text("uuidv7()")
|
||||
return sa.Column(name, models.types.StringUUID(), **kwargs)
|
||||
|
||||
|
||||
def upgrade():
|
||||
op.create_table(
|
||||
"agent_debug_conversations",
|
||||
_uuid_column("id", primary_key=True),
|
||||
sa.Column("tenant_id", models.types.StringUUID(), nullable=False),
|
||||
sa.Column("agent_id", models.types.StringUUID(), nullable=False),
|
||||
sa.Column("app_id", models.types.StringUUID(), nullable=False),
|
||||
sa.Column("account_id", models.types.StringUUID(), nullable=False),
|
||||
sa.Column("conversation_id", models.types.StringUUID(), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False),
|
||||
sa.Column("updated_at", sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name=op.f("agent_debug_conversation_pkey")),
|
||||
sa.UniqueConstraint(
|
||||
"tenant_id",
|
||||
"agent_id",
|
||||
"account_id",
|
||||
name=op.f("agent_debug_conversation_agent_account_unique"),
|
||||
),
|
||||
)
|
||||
op.create_index(
|
||||
"agent_debug_conversation_conversation_idx",
|
||||
"agent_debug_conversations",
|
||||
["conversation_id"],
|
||||
)
|
||||
op.create_index(
|
||||
"agent_debug_conversation_account_idx",
|
||||
"agent_debug_conversations",
|
||||
["tenant_id", "account_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade():
|
||||
op.drop_index("agent_debug_conversation_account_idx", table_name="agent_debug_conversations")
|
||||
op.drop_index("agent_debug_conversation_conversation_idx", table_name="agent_debug_conversations")
|
||||
op.drop_table("agent_debug_conversations")
|
||||
@@ -13,6 +13,7 @@ from .agent import (
|
||||
AgentConfigRevision,
|
||||
AgentConfigRevisionOperation,
|
||||
AgentConfigSnapshot,
|
||||
AgentDebugConversation,
|
||||
AgentDriveFile,
|
||||
AgentDriveFileKind,
|
||||
AgentIconType,
|
||||
@@ -156,6 +157,7 @@ __all__ = [
|
||||
"AgentConfigRevision",
|
||||
"AgentConfigRevisionOperation",
|
||||
"AgentConfigSnapshot",
|
||||
"AgentDebugConversation",
|
||||
"AgentDriveFile",
|
||||
"AgentDriveFileKind",
|
||||
"AgentIconType",
|
||||
|
||||
+37
-2
@@ -83,6 +83,8 @@ class AgentConfigRevisionOperation(StrEnum):
|
||||
SAVE_NEW_AGENT = "save_new_agent"
|
||||
# Promotes a workflow-only Agent into the reusable Agent Roster.
|
||||
SAVE_TO_ROSTER = "save_to_roster"
|
||||
# Switches the Agent's current published config back to an existing version.
|
||||
RESTORE_VERSION = "restore_version"
|
||||
|
||||
|
||||
class WorkflowAgentBindingType(StrEnum):
|
||||
@@ -180,6 +182,34 @@ class Agent(DefaultFieldsMixin, Base):
|
||||
archived_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
|
||||
|
||||
|
||||
class AgentDebugConversation(DefaultFieldsMixin, Base):
|
||||
"""Per-account console debug conversation for an Agent App.
|
||||
|
||||
Agent App preview state must be isolated by editor account. The Agent row is
|
||||
shared by everyone in the workspace, so this table owns the user-specific
|
||||
conversation pointer used by console debug chat.
|
||||
"""
|
||||
|
||||
__tablename__ = "agent_debug_conversations"
|
||||
__table_args__ = (
|
||||
sa.PrimaryKeyConstraint("id", name="agent_debug_conversation_pkey"),
|
||||
UniqueConstraint(
|
||||
"tenant_id",
|
||||
"agent_id",
|
||||
"account_id",
|
||||
name="agent_debug_conversation_agent_account_unique",
|
||||
),
|
||||
Index("agent_debug_conversation_conversation_idx", "conversation_id"),
|
||||
Index("agent_debug_conversation_account_idx", "tenant_id", "account_id"),
|
||||
)
|
||||
|
||||
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
agent_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
account_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
conversation_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
|
||||
|
||||
class AgentConfigSnapshot(DefaultFieldsMixin, Base):
|
||||
"""Immutable Agent Soul snapshot.
|
||||
|
||||
@@ -430,14 +460,17 @@ class AgentDriveFile(DefaultFieldsMixin, Base):
|
||||
synced. ``value_owned_by_drive`` gates physical cleanup: only drive-owned values
|
||||
(created by the agent runtime or Skill standardization, not shared with other
|
||||
business records) have their storage object + record deleted when the KV entry is
|
||||
overwritten or removed; otherwise only the KV row is dropped. Lifecycle never relies
|
||||
on ``UploadFile.used/used_by`` (not a reliable refcount).
|
||||
overwritten or removed; otherwise only the KV row is dropped. Skills are represented
|
||||
by the canonical ``<path>/SKILL.md`` row with ``is_skill=True`` and a serialized
|
||||
``skill_metadata`` string. Lifecycle never relies on ``UploadFile.used/used_by``
|
||||
(not a reliable refcount).
|
||||
"""
|
||||
|
||||
__tablename__ = "agent_drive_files"
|
||||
__table_args__ = (
|
||||
sa.PrimaryKeyConstraint("id", name="agent_drive_file_pkey"),
|
||||
UniqueConstraint("tenant_id", "agent_id", "key", name="agent_drive_file_scope_key_unique"),
|
||||
Index("agent_drive_files_tenant_agent_is_skill_key_idx", "tenant_id", "agent_id", "is_skill", "key"),
|
||||
)
|
||||
|
||||
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
@@ -453,6 +486,8 @@ class AgentDriveFile(DefaultFieldsMixin, Base):
|
||||
value_owned_by_drive: Mapped[bool] = mapped_column(
|
||||
sa.Boolean, nullable=False, default=False, server_default=sa.text("false")
|
||||
)
|
||||
is_skill: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False, server_default=sa.text("false"))
|
||||
skill_metadata: Mapped[str | None] = mapped_column(LongText, nullable=True)
|
||||
size: Mapped[int | None] = mapped_column(sa.BigInteger, nullable=True)
|
||||
hash: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
||||
mime_type: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
||||
|
||||
@@ -134,20 +134,40 @@ class HumanInputDelivery(DefaultFieldsMixin, Base):
|
||||
)
|
||||
|
||||
|
||||
class ApprovalChannel(StrEnum):
|
||||
"""Where a paused human input form can be approved, surfaced to API callers."""
|
||||
|
||||
EMAIL = "email"
|
||||
WEB_APP = "web_app"
|
||||
CONSOLE = "console"
|
||||
|
||||
|
||||
class RecipientType(StrEnum):
|
||||
# EMAIL_MEMBER member means that the
|
||||
EMAIL_MEMBER = "email_member"
|
||||
EMAIL_EXTERNAL = "email_external"
|
||||
# Second value = the approval channel this recipient maps to (surfaced in `approval_channels`).
|
||||
EMAIL_MEMBER = "email_member", ApprovalChannel.EMAIL
|
||||
EMAIL_EXTERNAL = "email_external", ApprovalChannel.EMAIL
|
||||
# STANDALONE_WEB_APP is used by the standalone web app.
|
||||
#
|
||||
# It's not used while running workflows / chatflows containing HumanInput
|
||||
# node inside console.
|
||||
STANDALONE_WEB_APP = "standalone_web_app"
|
||||
STANDALONE_WEB_APP = "standalone_web_app", ApprovalChannel.WEB_APP
|
||||
# CONSOLE is used while running workflows / chatflows containing HumanInput
|
||||
# node inside console. (E.G. running installed apps or debugging workflows / chatflows)
|
||||
CONSOLE = "console"
|
||||
CONSOLE = "console", ApprovalChannel.CONSOLE
|
||||
# BACKSTAGE is used for backstage input inside console.
|
||||
BACKSTAGE = "backstage"
|
||||
BACKSTAGE = "backstage", ApprovalChannel.CONSOLE
|
||||
|
||||
_approval_channel: ApprovalChannel
|
||||
|
||||
def __new__(cls, value: str, approval_channel: ApprovalChannel) -> "RecipientType":
|
||||
member = str.__new__(cls, value)
|
||||
member._value_ = value
|
||||
member._approval_channel = approval_channel
|
||||
return member
|
||||
|
||||
@property
|
||||
def approval_channel(self) -> ApprovalChannel:
|
||||
return self._approval_channel
|
||||
|
||||
|
||||
@final
|
||||
|
||||
+27
-22
@@ -774,26 +774,7 @@ class AppModelConfig(TypeBase):
|
||||
|
||||
@property
|
||||
def annotation_reply_dict(self) -> AnnotationReplyConfig:
|
||||
annotation_setting = db.session.scalar(
|
||||
select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == self.app_id)
|
||||
)
|
||||
if annotation_setting:
|
||||
collection_binding_detail = annotation_setting.collection_binding_detail
|
||||
if not collection_binding_detail:
|
||||
raise ValueError("Collection binding detail not found")
|
||||
|
||||
return {
|
||||
"id": annotation_setting.id,
|
||||
"enabled": True,
|
||||
"score_threshold": annotation_setting.score_threshold,
|
||||
"embedding_model": {
|
||||
"embedding_provider_name": collection_binding_detail.provider_name,
|
||||
"embedding_model_name": collection_binding_detail.model_name,
|
||||
},
|
||||
}
|
||||
|
||||
else:
|
||||
return {"enabled": False}
|
||||
return load_annotation_reply_config(db.session(), self.app_id)
|
||||
|
||||
@property
|
||||
def more_like_this_dict(self) -> EnabledConfig:
|
||||
@@ -864,7 +845,7 @@ class AppModelConfig(TypeBase):
|
||||
},
|
||||
)
|
||||
|
||||
def to_dict(self) -> AppModelConfigDict:
|
||||
def to_dict(self, *, annotation_reply: AnnotationReplyConfig | None = None) -> AppModelConfigDict:
|
||||
return {
|
||||
"opening_statement": self.opening_statement,
|
||||
"suggested_questions": self.suggested_questions_list,
|
||||
@@ -872,7 +853,7 @@ class AppModelConfig(TypeBase):
|
||||
"speech_to_text": self.speech_to_text_dict,
|
||||
"text_to_speech": self.text_to_speech_dict,
|
||||
"retriever_resource": self.retriever_resource_dict,
|
||||
"annotation_reply": self.annotation_reply_dict,
|
||||
"annotation_reply": annotation_reply if annotation_reply is not None else self.annotation_reply_dict,
|
||||
"more_like_this": self.more_like_this_dict,
|
||||
"sensitive_word_avoidance": self.sensitive_word_avoidance_dict,
|
||||
"external_data_tools": self.external_data_tools_list,
|
||||
@@ -2038,6 +2019,30 @@ class AppAnnotationSetting(TypeBase):
|
||||
)
|
||||
|
||||
|
||||
def load_annotation_reply_config(session: Session, app_id: str) -> AnnotationReplyConfig:
|
||||
annotation_setting = session.scalar(select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id))
|
||||
if annotation_setting is None:
|
||||
return {"enabled": False}
|
||||
|
||||
from .dataset import DatasetCollectionBinding
|
||||
|
||||
collection_binding_detail = session.scalar(
|
||||
select(DatasetCollectionBinding).where(DatasetCollectionBinding.id == annotation_setting.collection_binding_id)
|
||||
)
|
||||
if collection_binding_detail is None:
|
||||
raise ValueError("Collection binding detail not found")
|
||||
|
||||
return {
|
||||
"id": annotation_setting.id,
|
||||
"enabled": True,
|
||||
"score_threshold": annotation_setting.score_threshold,
|
||||
"embedding_model": {
|
||||
"embedding_provider_name": collection_binding_detail.provider_name,
|
||||
"embedding_model_name": collection_binding_detail.model_name,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class OperationLog(TypeBase):
|
||||
__tablename__ = "operation_logs"
|
||||
__table_args__ = (
|
||||
|
||||
@@ -391,6 +391,80 @@ Check if activation token is valid
|
||||
| 400 | Invalid request parameters | |
|
||||
| 403 | Insufficient permissions | |
|
||||
|
||||
### [GET] /agent/{agent_id}/api-access
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Agent service API access | **application/json**: [AgentApiAccessResponse](#agentapiaccessresponse)<br> |
|
||||
|
||||
### [POST] /agent/{agent_id}/api-enable
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Request Body
|
||||
|
||||
| Required | Schema |
|
||||
| -------- | ------ |
|
||||
| Yes | **application/json**: [AgentApiStatusPayload](#agentapistatuspayload)<br> |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Agent service API status updated | **application/json**: [AgentApiAccessResponse](#agentapiaccessresponse)<br> |
|
||||
| 403 | Insufficient permissions | |
|
||||
|
||||
### [GET] /agent/{agent_id}/api-keys
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Agent service API keys | **application/json**: [ApiKeyList](#apikeylist)<br> |
|
||||
|
||||
### [POST] /agent/{agent_id}/api-keys
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 201 | Agent service API key created | **application/json**: [ApiKeyItem](#apikeyitem)<br> |
|
||||
| 400 | Maximum keys exceeded | |
|
||||
|
||||
### [DELETE] /agent/{agent_id}/api-keys/{api_key_id}
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | | Yes | string (uuid) |
|
||||
| api_key_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description |
|
||||
| ---- | ----------- |
|
||||
| 204 | Agent service API key deleted |
|
||||
|
||||
### [GET] /agent/{agent_id}/chat-messages
|
||||
Get Agent App chat messages for a conversation with pagination
|
||||
|
||||
@@ -518,7 +592,7 @@ Stop a running Agent App chat message generation
|
||||
|
||||
| Required | Schema |
|
||||
| -------- | ------ |
|
||||
| Yes | **application/json**: [CopyAppPayload](#copyapppayload)<br> |
|
||||
| Yes | **application/json**: [AgentAppCopyPayload](#agentappcopypayload)<br> |
|
||||
|
||||
#### Responses
|
||||
|
||||
@@ -528,6 +602,20 @@ Stop a running Agent App chat message generation
|
||||
| 400 | Invalid request parameters | |
|
||||
| 403 | Insufficient permissions | |
|
||||
|
||||
### [POST] /agent/{agent_id}/debug-conversation/refresh
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Agent debug conversation refreshed | **application/json**: [AgentDebugConversationRefreshResponse](#agentdebugconversationrefreshresponse)<br> |
|
||||
| 403 | Insufficient permissions | |
|
||||
|
||||
### [GET] /agent/{agent_id}/drive/files
|
||||
List agent drive entries for an Agent App
|
||||
|
||||
@@ -576,6 +664,37 @@ Truncated text preview of one Agent App drive value
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Preview | **application/json**: [AgentDrivePreviewResponse](#agentdrivepreviewresponse)<br> |
|
||||
|
||||
### [GET] /agent/{agent_id}/drive/skills
|
||||
List drive-backed skills for an Agent App
|
||||
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | Agent ID | Yes | string (uuid) |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Drive skills | **application/json**: [AgentDriveSkillListResponse](#agentdriveskilllistresponse)<br> |
|
||||
|
||||
### [GET] /agent/{agent_id}/drive/skills/{skill_path}/inspect
|
||||
Inspect one drive-backed skill for slash-menu hover/detail UI
|
||||
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | Agent ID | Yes | string (uuid) |
|
||||
| skill_path | path | Skill path/slug, e.g. tender-analyzer | Yes | string |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Drive skill inspect view | **application/json**: [AgentDriveSkillInspectResponse](#agentdriveskillinspectresponse)<br> |
|
||||
|
||||
### [POST] /agent/{agent_id}/features
|
||||
Update an Agent App's presentation features (opener, follow-up, citations, ...)
|
||||
|
||||
@@ -905,6 +1024,20 @@ Infer CLI tool + ENV suggestions from a standardized Agent App skill
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Agent version detail | **application/json**: [AgentConfigSnapshotDetailResponse](#agentconfigsnapshotdetailresponse)<br> |
|
||||
|
||||
### [POST] /agent/{agent_id}/versions/{version_id}/restore
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| agent_id | path | | Yes | string (uuid) |
|
||||
| version_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Agent version restored | **application/json**: [AgentConfigSnapshotRestoreResponse](#agentconfigsnapshotrestoreresponse)<br> |
|
||||
|
||||
### [GET] /all-workspaces
|
||||
#### Parameters
|
||||
|
||||
@@ -1454,6 +1587,40 @@ Truncated text preview of one drive value (binary-safe; SKILL.md is the main cas
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Preview | **application/json**: [AgentDrivePreviewResponse](#agentdrivepreviewresponse)<br> |
|
||||
|
||||
### [GET] /apps/{app_id}/agent/drive/skills
|
||||
List drive-backed skills for the bound agent
|
||||
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| app_id | path | Application ID | Yes | string (uuid) |
|
||||
| node_id | query | Workflow node ID (workflow composer variant) | No | string |
|
||||
| prefix | query | Key prefix filter: '<slug>/' for one skill, 'files/' for files | No | string |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Drive skills | **application/json**: [AgentDriveSkillListResponse](#agentdriveskilllistresponse)<br> |
|
||||
|
||||
### [GET] /apps/{app_id}/agent/drive/skills/{skill_path}/inspect
|
||||
Inspect one drive-backed skill for slash-menu hover/detail UI
|
||||
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| app_id | path | Application ID | Yes | string (uuid) |
|
||||
| skill_path | path | Skill path/slug, e.g. tender-analyzer | Yes | string |
|
||||
| node_id | query | Workflow node ID (workflow composer variant) | No | string |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Drive skill inspect view | **application/json**: [AgentDriveSkillInspectResponse](#agentdriveskillinspectresponse)<br> |
|
||||
|
||||
### [DELETE] /apps/{app_id}/agent/files
|
||||
Delete one drive file by key; soul ref first, then the KV row (ENG-625 D5)
|
||||
|
||||
@@ -11954,6 +12121,31 @@ Default namespace
|
||||
| chat_prompt_config | object | | No |
|
||||
| completion_prompt_config | object | | No |
|
||||
|
||||
#### AgentApiAccessResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| api_key_count | integer | | Yes |
|
||||
| api_rph | integer | | Yes |
|
||||
| api_rpm | integer | | Yes |
|
||||
| chat_endpoint | string | | Yes |
|
||||
| conversations_endpoint | string | | Yes |
|
||||
| enabled | boolean | | Yes |
|
||||
| files_upload_endpoint | string | | Yes |
|
||||
| info_endpoint | string | | Yes |
|
||||
| messages_endpoint | string | | Yes |
|
||||
| meta_endpoint | string | | Yes |
|
||||
| parameters_endpoint | string | | Yes |
|
||||
| service_api_base_url | string | | Yes |
|
||||
| stop_endpoint | string | | Yes |
|
||||
| streaming_only | boolean, <br>**Default:** true | | No |
|
||||
|
||||
#### AgentApiStatusPayload
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| enable_api | boolean | Enable or disable Agent service API | Yes |
|
||||
|
||||
#### AgentAppComposerResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@@ -11965,6 +12157,17 @@ Default namespace
|
||||
| validation | [ComposerValidationFindingsResponse](#composervalidationfindingsresponse) | | No |
|
||||
| variant | string | | Yes |
|
||||
|
||||
#### AgentAppCopyPayload
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| description | string | Description for the copied agent | No |
|
||||
| icon | string | Icon | No |
|
||||
| icon_background | string | Icon background color | No |
|
||||
| icon_type | [IconType](#icontype) | Icon type | No |
|
||||
| name | string | Name for the copied agent | No |
|
||||
| role | string | Role for the copied agent | No |
|
||||
|
||||
#### AgentAppCreatePayload
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@@ -11987,6 +12190,7 @@ Default namespace
|
||||
| bound_agent_id | string | | No |
|
||||
| created_at | integer | | No |
|
||||
| created_by | string | | No |
|
||||
| debug_conversation_id | string | | No |
|
||||
| deleted_tools | [ [DeletedTool](#deletedtool) ] | | No |
|
||||
| description | string | | No |
|
||||
| enable_api | boolean | | Yes |
|
||||
@@ -12050,6 +12254,7 @@ default (the config form sends the full desired feature state on save).
|
||||
| create_user_name | string | | No |
|
||||
| created_at | integer | | No |
|
||||
| created_by | string | | No |
|
||||
| debug_conversation_id | string | | No |
|
||||
| description | string | | No |
|
||||
| has_draft_trigger | boolean | | No |
|
||||
| icon | string | | No |
|
||||
@@ -12340,6 +12545,13 @@ Audit operation recorded for Agent Soul version/revision changes.
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| data | [ [AgentConfigSnapshotSummaryResponse](#agentconfigsnapshotsummaryresponse) ] | | Yes |
|
||||
|
||||
#### AgentConfigSnapshotRestoreResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| active_config_snapshot_id | string | | Yes |
|
||||
| result | string | | Yes |
|
||||
|
||||
#### AgentConfigSnapshotSummaryResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@@ -12375,6 +12587,12 @@ Audit operation recorded for Agent Soul version/revision changes.
|
||||
| date | string | | Yes |
|
||||
| message_count | integer | | Yes |
|
||||
|
||||
#### AgentDebugConversationRefreshResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| debug_conversation_id | string | | Yes |
|
||||
|
||||
#### AgentDriveDeleteFileByAgentQuery
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@@ -12425,9 +12643,11 @@ Audit operation recorded for Agent Soul version/revision changes.
|
||||
| created_at | integer | | No |
|
||||
| file_kind | string | | Yes |
|
||||
| hash | string | | No |
|
||||
| is_skill | boolean | | No |
|
||||
| key | string | | Yes |
|
||||
| mime_type | string | | No |
|
||||
| size | integer | | No |
|
||||
| skill_metadata | string | | No |
|
||||
|
||||
#### AgentDriveListResponse
|
||||
|
||||
@@ -12445,6 +12665,65 @@ Audit operation recorded for Agent Soul version/revision changes.
|
||||
| text | string | | No |
|
||||
| truncated | boolean | | Yes |
|
||||
|
||||
#### AgentDriveSkillFileResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| available_in_drive | boolean | | Yes |
|
||||
| drive_key | string | | No |
|
||||
| name | string | | Yes |
|
||||
| path | string | | Yes |
|
||||
| type | string | | Yes |
|
||||
|
||||
#### AgentDriveSkillInspectResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| archive_key | string | | No |
|
||||
| created_at | integer | | No |
|
||||
| description | string | | Yes |
|
||||
| file_tree | [ object ] | | No |
|
||||
| files | [ [AgentDriveSkillFileResponse](#agentdriveskillfileresponse) ] | | No |
|
||||
| hash | string | | No |
|
||||
| mime_type | string | | No |
|
||||
| name | string | | Yes |
|
||||
| path | string | | Yes |
|
||||
| size | integer | | No |
|
||||
| skill_md | [AgentDriveSkillMarkdownResponse](#agentdriveskillmarkdownresponse) | | Yes |
|
||||
| skill_md_key | string | | Yes |
|
||||
| source | string | | Yes |
|
||||
| warnings | [ string ] | | No |
|
||||
|
||||
#### AgentDriveSkillItemResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| archive_key | string | | No |
|
||||
| created_at | integer | | No |
|
||||
| description | string | | Yes |
|
||||
| hash | string | | No |
|
||||
| mime_type | string | | No |
|
||||
| name | string | | Yes |
|
||||
| path | string | | Yes |
|
||||
| size | integer | | No |
|
||||
| skill_md_key | string | | Yes |
|
||||
|
||||
#### AgentDriveSkillListResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| items | [ [AgentDriveSkillItemResponse](#agentdriveskillitemresponse) ] | | No |
|
||||
|
||||
#### AgentDriveSkillMarkdownResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| binary | boolean | | Yes |
|
||||
| key | string | | Yes |
|
||||
| size | integer | | No |
|
||||
| text | string | | No |
|
||||
| truncated | boolean | | Yes |
|
||||
|
||||
#### AgentEnvVariableConfig
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
|
||||
@@ -80,10 +80,9 @@ User-scoped operations
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| limit | query | | No | integer, <br>**Default:** 20 |
|
||||
| mode | query | | No | string, <br>**Available values:** "advanced-chat", "agent", "agent-chat", "channel", "chat", "completion", "rag-pipeline", "workflow" |
|
||||
| mode | query | App types the ``app`` usage face (``get app``) lists and filters. A curated subset of :class:`AppMode`: the real, user-facing app categories. Excludes runtime-only mode tags that are not standalone apps (``rag-pipeline`` is a knowledge ``Pipeline``; ``channel`` is unused) and the roster-owned ``agent`` type (surfaced through the roster, not this list). Members reference ``AppMode.*.value`` so the subset relationship is type-checked: dropping a member from ``AppMode`` breaks this at import. This is the single source for the listable set — params, filters, and the generated CLI whitelist all derive from it. | No | string, <br>**Available values:** "advanced-chat", "agent-chat", "chat", "completion", "workflow" |
|
||||
| name | query | | No | string |
|
||||
| page | query | | No | integer, <br>**Default:** 1 |
|
||||
| tag | query | | No | string |
|
||||
| workspace_id | query | | Yes | string |
|
||||
|
||||
#### Responses
|
||||
@@ -319,7 +318,7 @@ Upload a file to use as an input variable when running the app
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| limit | query | | No | integer, <br>**Default:** 20 |
|
||||
| mode | query | | No | string, <br>**Available values:** "advanced-chat", "agent", "agent-chat", "channel", "chat", "completion", "rag-pipeline", "workflow" |
|
||||
| mode | query | App types the ``app`` usage face (``get app``) lists and filters. A curated subset of :class:`AppMode`: the real, user-facing app categories. Excludes runtime-only mode tags that are not standalone apps (``rag-pipeline`` is a knowledge ``Pipeline``; ``channel`` is unused) and the roster-owned ``agent`` type (surfaced through the roster, not this list). Members reference ``AppMode.*.value`` so the subset relationship is type-checked: dropping a member from ``AppMode`` breaks this at import. This is the single source for the listable set — params, filters, and the generated CLI whitelist all derive from it. | No | string, <br>**Available values:** "advanced-chat", "agent-chat", "chat", "completion", "workflow" |
|
||||
| name | query | | No | string |
|
||||
| page | query | | No | integer, <br>**Default:** 1 |
|
||||
|
||||
@@ -331,6 +330,22 @@ Upload a file to use as an input variable when running the app
|
||||
| 422 | Validation error | **application/json**: [ErrorBody](#errorbody)<br> |
|
||||
| default | Error | **application/json**: [ErrorBody](#errorbody)<br> |
|
||||
|
||||
### [GET] /permitted-external-apps/{app_id}/describe
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| fields | query | | No | string |
|
||||
| app_id | path | | Yes | string |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Permitted external app description | **application/json**: [AppDescribeResponse](#appdescriberesponse)<br> |
|
||||
| 422 | Validation error | **application/json**: [ErrorBody](#errorbody)<br> |
|
||||
| default | Error | **application/json**: [ErrorBody](#errorbody)<br> |
|
||||
|
||||
### [GET] /workspaces
|
||||
#### Responses
|
||||
|
||||
@@ -507,14 +522,12 @@ Upload a file to use as an input variable when running the app
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| author | string | | No |
|
||||
| description | string | | No |
|
||||
| id | string | | Yes |
|
||||
| is_agent | boolean | | No |
|
||||
| mode | string | | Yes |
|
||||
| name | string | | Yes |
|
||||
| service_api_enabled | boolean | | Yes |
|
||||
| tags | [ [TagItem](#tagitem) ], <br>**Default:** | | No |
|
||||
| updated_at | string | | No |
|
||||
|
||||
#### AppDescribeQuery
|
||||
@@ -568,28 +581,25 @@ Request body for POST /workspaces/<workspace_id>/apps/imports.
|
||||
| yaml_content | string | Inline YAML DSL string (required when mode is yaml-content) | No |
|
||||
| yaml_url | string | Remote URL to fetch YAML from (required when mode is yaml-url) | No |
|
||||
|
||||
#### AppInfoResponse
|
||||
#### AppInfo
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| author | string | | No |
|
||||
| description | string | | No |
|
||||
| id | string | | Yes |
|
||||
| mode | string | | Yes |
|
||||
| name | string | | Yes |
|
||||
| tags | [ [TagItem](#tagitem) ], <br>**Default:** | | No |
|
||||
|
||||
#### AppListQuery
|
||||
|
||||
mode is a closed enum.
|
||||
mode is a closed enum of listable app types.
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| limit | integer, <br>**Default:** 20 | | No |
|
||||
| mode | [AppMode](#appmode) | | No |
|
||||
| mode | [SupportedAppType](#supportedapptype) | | No |
|
||||
| name | string | | No |
|
||||
| page | integer, <br>**Default:** 1 | | No |
|
||||
| tag | string | | No |
|
||||
| workspace_id | string | | Yes |
|
||||
|
||||
#### AppListResponse
|
||||
@@ -606,12 +616,10 @@ mode is a closed enum.
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| created_by_name | string | | No |
|
||||
| description | string | | No |
|
||||
| id | string | | Yes |
|
||||
| mode | [AppMode](#appmode) | | Yes |
|
||||
| name | string | | Yes |
|
||||
| tags | [ [TagItem](#tagitem) ], <br>**Default:** | | No |
|
||||
| updated_at | string | | No |
|
||||
| workspace_id | string | | No |
|
||||
| workspace_name | string | | No |
|
||||
@@ -914,7 +922,7 @@ Strict (extra='forbid').
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| limit | integer, <br>**Default:** 20 | | No |
|
||||
| mode | [AppMode](#appmode) | | No |
|
||||
| mode | [SupportedAppType](#supportedapptype) | | No |
|
||||
| name | string | | No |
|
||||
| page | integer, <br>**Default:** 1 | | No |
|
||||
|
||||
@@ -982,11 +990,23 @@ Pagination for GET /account/sessions. Strict (extra='forbid').
|
||||
| last_used_at | string | | No |
|
||||
| prefix | string | | Yes |
|
||||
|
||||
#### TagItem
|
||||
#### SupportedAppType
|
||||
|
||||
App types the ``app`` usage face (``get app``) lists and filters.
|
||||
|
||||
A curated subset of :class:`AppMode`: the real, user-facing app categories.
|
||||
Excludes runtime-only mode tags that are not standalone apps
|
||||
(``rag-pipeline`` is a knowledge ``Pipeline``; ``channel`` is unused) and the
|
||||
roster-owned ``agent`` type (surfaced through the roster, not this list).
|
||||
|
||||
Members reference ``AppMode.*.value`` so the subset relationship is
|
||||
type-checked: dropping a member from ``AppMode`` breaks this at import.
|
||||
This is the single source for the listable set — params, filters, and the
|
||||
generated CLI whitelist all derive from it.
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| name | string | | Yes |
|
||||
| SupportedAppType | string | App types the ``app`` usage face (``get app``) lists and filters. A curated subset of :class:`AppMode`: the real, user-facing app categories. Excludes runtime-only mode tags that are not standalone apps (``rag-pipeline`` is a knowledge ``Pipeline``; ``channel`` is unused) and the roster-owned ``agent`` type (surfaced through the roster, not this list). Members reference ``AppMode.*.value`` so the subset relationship is type-checked: dropping a member from ``AppMode`` breaks this at import. This is the single source for the listable set — params, filters, and the generated CLI whitelist all derive from it. | |
|
||||
|
||||
#### TaskStopResponse
|
||||
|
||||
|
||||
+11
-8
@@ -1,3 +1,4 @@
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
@@ -142,9 +143,8 @@ class TestTraceClient:
|
||||
mock_notify.assert_called_once()
|
||||
|
||||
@patch("dify_trace_aliyun.data_exporter.traceclient.OTLPSpanExporter")
|
||||
@patch("dify_trace_aliyun.data_exporter.traceclient.logger")
|
||||
def test_add_span_queue_full(
|
||||
self, mock_logger: MagicMock, mock_exporter_class: MagicMock, trace_client_factory: type[TraceClient]
|
||||
self, mock_exporter_class: MagicMock, trace_client_factory: type[TraceClient], caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
client = trace_client_factory(service_name="test-service", endpoint="http://test-endpoint", max_queue_size=1)
|
||||
|
||||
@@ -164,12 +164,15 @@ class TestTraceClient:
|
||||
client.add_span(span_data)
|
||||
assert len(client.queue) == 1
|
||||
|
||||
client.add_span(span_data)
|
||||
assert len(client.queue) == 1
|
||||
mock_logger.warning.assert_called_with("Queue is full, likely spans will be dropped.")
|
||||
with caplog.at_level(logging.WARNING):
|
||||
client.add_span(span_data)
|
||||
assert len(client.queue) == 1
|
||||
assert "Queue is full, likely spans will be dropped." in caplog.text
|
||||
|
||||
@patch("dify_trace_aliyun.data_exporter.traceclient.OTLPSpanExporter")
|
||||
def test_export_batch_error(self, mock_exporter_class: MagicMock, trace_client_factory: type[TraceClient]):
|
||||
def test_export_batch_error(
|
||||
self, mock_exporter_class: MagicMock, trace_client_factory: type[TraceClient], caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
mock_exporter = mock_exporter_class.return_value
|
||||
mock_exporter.export.side_effect = Exception("Export failed")
|
||||
|
||||
@@ -177,9 +180,9 @@ class TestTraceClient:
|
||||
mock_span = MagicMock(spec=ReadableSpan)
|
||||
client.queue.append(mock_span)
|
||||
|
||||
with patch("dify_trace_aliyun.data_exporter.traceclient.logger") as mock_logger:
|
||||
with caplog.at_level(logging.WARNING):
|
||||
client._export_batch()
|
||||
mock_logger.warning.assert_called()
|
||||
assert "Error exporting spans" in caplog.text
|
||||
|
||||
@patch("dify_trace_aliyun.data_exporter.traceclient.OTLPSpanExporter")
|
||||
def test_worker_loop(self, mock_exporter_class: MagicMock, trace_client_factory: type[TraceClient]):
|
||||
|
||||
@@ -307,13 +307,12 @@ class TestGetProjectUrl:
|
||||
monkeypatch.setattr(trace_instance, "entity", None)
|
||||
monkeypatch.setattr(trace_instance, "project_name", None)
|
||||
# Force an error by making string formatting fail
|
||||
with patch("dify_trace_weave.weave_trace.logger") as mock_logger:
|
||||
# Simulate exception via property
|
||||
original_entity = trace_instance.entity
|
||||
trace_instance.entity = None
|
||||
trace_instance.project_name = None
|
||||
url = trace_instance.get_project_url()
|
||||
assert "https://wandb.ai/" in url
|
||||
# Simulate exception via property
|
||||
original_entity = trace_instance.entity
|
||||
trace_instance.entity = None
|
||||
trace_instance.project_name = None
|
||||
url = trace_instance.get_project_url()
|
||||
assert "https://wandb.ai/" in url
|
||||
|
||||
|
||||
# ── TestTraceDispatcher ─────────────────────────────────────────────────────
|
||||
|
||||
@@ -290,7 +290,10 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol):
|
||||
batch_size: int,
|
||||
run_types: Sequence[WorkflowType] | None = None,
|
||||
tenant_ids: Sequence[str] | None = None,
|
||||
tenant_prefixes: Sequence[str] | None = None,
|
||||
workflow_ids: Sequence[str] | None = None,
|
||||
run_shard_index: int | None = None,
|
||||
run_shard_total: int | None = None,
|
||||
) -> Sequence[WorkflowRun]:
|
||||
"""
|
||||
Fetch ended workflow runs in a time window for archival and clean batching.
|
||||
@@ -298,7 +301,9 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol):
|
||||
Optional filters:
|
||||
- run_types
|
||||
- tenant_ids
|
||||
- tenant_prefixes, using the first hexadecimal digit of tenant_id for rollout waves
|
||||
- workflow_ids
|
||||
- run_shard_index/run_shard_total, using a deterministic workflow_run_id shard
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@@ -56,6 +56,7 @@ from repositories.types import (
|
||||
DailyTerminalsStats,
|
||||
DailyTokenCostStats,
|
||||
)
|
||||
from services.retention.workflow_run.tenant_prefix import tenant_prefix_condition
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -64,6 +65,40 @@ class _WorkflowRunError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
_HEX_SHARD_VALUES = {
|
||||
"0": 0,
|
||||
"1": 1,
|
||||
"2": 2,
|
||||
"3": 3,
|
||||
"4": 4,
|
||||
"5": 5,
|
||||
"6": 6,
|
||||
"7": 7,
|
||||
"8": 8,
|
||||
"9": 9,
|
||||
"a": 10,
|
||||
"b": 11,
|
||||
"c": 12,
|
||||
"d": 13,
|
||||
"e": 14,
|
||||
"f": 15,
|
||||
}
|
||||
|
||||
|
||||
def _tenant_prefix_condition(prefixes: Sequence[str]) -> sa.ColumnElement[bool]:
|
||||
conditions = [tenant_prefix_condition(WorkflowRun.tenant_id, prefix) for prefix in prefixes]
|
||||
return sa.or_(*conditions)
|
||||
|
||||
|
||||
def _workflow_run_id_shard_expr() -> sa.ColumnElement[int]:
|
||||
normalized_id = func.lower(func.replace(sa.cast(WorkflowRun.id, sa.String()), "-", ""))
|
||||
last_hex = func.substr(normalized_id, func.length(normalized_id), 1)
|
||||
return sa.case(
|
||||
*[(last_hex == hex_digit, shard_value) for hex_digit, shard_value in _HEX_SHARD_VALUES.items()],
|
||||
else_=0,
|
||||
)
|
||||
|
||||
|
||||
def _build_human_input_required_reason(
|
||||
reason_model: WorkflowPauseReason,
|
||||
form_model: HumanInputForm | None,
|
||||
@@ -378,7 +413,10 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
|
||||
batch_size: int,
|
||||
run_types: Sequence[WorkflowType] | None = None,
|
||||
tenant_ids: Sequence[str] | None = None,
|
||||
tenant_prefixes: Sequence[str] | None = None,
|
||||
workflow_ids: Sequence[str] | None = None,
|
||||
run_shard_index: int | None = None,
|
||||
run_shard_total: int | None = None,
|
||||
) -> Sequence[WorkflowRun]:
|
||||
"""
|
||||
Fetch ended workflow runs in a time window for archival and clean batching.
|
||||
@@ -387,7 +425,8 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
|
||||
- created_at in [start_from, end_before)
|
||||
- type in run_types (when provided)
|
||||
- status is an ended state
|
||||
- optional tenant_id, workflow_id filters and cursor (last_seen) for pagination
|
||||
- optional tenant_id, tenant_prefix, workflow_id filters and cursor (last_seen) for pagination
|
||||
- optional deterministic shard by the last hexadecimal digit of workflow_run_id
|
||||
"""
|
||||
with self._session_maker() as session:
|
||||
stmt = (
|
||||
@@ -410,9 +449,15 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
|
||||
if tenant_ids:
|
||||
stmt = stmt.where(WorkflowRun.tenant_id.in_(tenant_ids))
|
||||
|
||||
if tenant_prefixes:
|
||||
stmt = stmt.where(_tenant_prefix_condition(tenant_prefixes))
|
||||
|
||||
if workflow_ids:
|
||||
stmt = stmt.where(WorkflowRun.workflow_id.in_(workflow_ids))
|
||||
|
||||
if run_shard_index is not None and run_shard_total is not None:
|
||||
stmt = stmt.where((_workflow_run_id_shard_expr() % run_shard_total) == run_shard_index)
|
||||
|
||||
if last_seen:
|
||||
stmt = stmt.where(
|
||||
tuple_(WorkflowRun.created_at, WorkflowRun.id)
|
||||
|
||||
@@ -830,6 +830,16 @@ class AgentComposerService:
|
||||
) -> WorkflowAgentNodeBinding:
|
||||
node_job = payload.node_job or WorkflowNodeJobConfig()
|
||||
if binding:
|
||||
if cls._is_start_from_scratch_request(binding=binding, payload=payload):
|
||||
return cls._switch_roster_binding_to_inline_agent(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
workflow_id=workflow_id,
|
||||
node_id=node_id,
|
||||
account_id=account_id,
|
||||
binding=binding,
|
||||
payload=payload,
|
||||
)
|
||||
binding.node_job_config = node_job
|
||||
if payload.agent_soul is not None and binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT:
|
||||
current_snapshot = cls._require_version(
|
||||
@@ -880,6 +890,46 @@ class AgentComposerService:
|
||||
db.session.flush()
|
||||
return binding
|
||||
|
||||
@classmethod
|
||||
def _is_start_from_scratch_request(cls, *, binding: WorkflowAgentNodeBinding, payload: ComposerSavePayload) -> bool:
|
||||
return (
|
||||
binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT
|
||||
and payload.binding is not None
|
||||
and payload.binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT.value
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _switch_roster_binding_to_inline_agent(
|
||||
cls,
|
||||
*,
|
||||
tenant_id: str,
|
||||
app_id: str,
|
||||
workflow_id: str,
|
||||
node_id: str,
|
||||
account_id: str,
|
||||
binding: WorkflowAgentNodeBinding,
|
||||
payload: ComposerSavePayload,
|
||||
) -> WorkflowAgentNodeBinding:
|
||||
if payload.binding and (payload.binding.agent_id or payload.binding.current_snapshot_id):
|
||||
raise ValueError("Start from Scratch must not provide an existing inline agent binding.")
|
||||
|
||||
agent_soul = payload.agent_soul or AgentSoulConfig()
|
||||
agent = cls._create_workflow_only_agent(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
workflow_id=workflow_id,
|
||||
node_id=node_id,
|
||||
account_id=account_id,
|
||||
agent_soul=agent_soul,
|
||||
)
|
||||
binding.binding_type = WorkflowAgentBindingType.INLINE_AGENT
|
||||
binding.agent_id = agent.id
|
||||
binding.current_snapshot_id = agent.active_config_snapshot_id
|
||||
binding.node_job_config = payload.node_job or binding.node_job_config
|
||||
binding.updated_by = account_id
|
||||
db.session.flush()
|
||||
return binding
|
||||
|
||||
@classmethod
|
||||
def _save_to_current_version(
|
||||
cls,
|
||||
|
||||
@@ -3,6 +3,7 @@ from typing import Any, TypedDict
|
||||
from sqlalchemy import and_, func, or_, select
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from libs.helper import to_timestamp
|
||||
from models.agent import (
|
||||
@@ -10,6 +11,7 @@ from models.agent import (
|
||||
AgentConfigRevision,
|
||||
AgentConfigRevisionOperation,
|
||||
AgentConfigSnapshot,
|
||||
AgentDebugConversation,
|
||||
AgentKind,
|
||||
AgentScope,
|
||||
AgentSource,
|
||||
@@ -18,8 +20,8 @@ from models.agent import (
|
||||
WorkflowAgentNodeBinding,
|
||||
)
|
||||
from models.agent_config_entities import AgentSoulConfig
|
||||
from models.enums import AppStatus
|
||||
from models.model import App, AppMode, IconType
|
||||
from models.enums import AppStatus, ConversationFromSource, ConversationStatus
|
||||
from models.model import App, AppMode, Conversation, IconType
|
||||
from models.workflow import Workflow
|
||||
from services.agent.agent_soul_state import agent_soul_has_model
|
||||
from services.agent.composer_validator import ComposerConfigValidator
|
||||
@@ -96,6 +98,7 @@ class AgentRosterService:
|
||||
"scope": agent.scope.value,
|
||||
"source": agent.source.value,
|
||||
"app_id": agent.app_id,
|
||||
"debug_conversation_id": None,
|
||||
"workflow_id": agent.workflow_id,
|
||||
"workflow_node_id": agent.workflow_node_id,
|
||||
"active_config_snapshot_id": agent.active_config_snapshot_id,
|
||||
@@ -392,8 +395,172 @@ class AgentRosterService:
|
||||
agent.active_config_snapshot_id = version.id
|
||||
agent.active_config_has_model = agent_soul_has_model(AgentSoulConfig())
|
||||
self._session.flush()
|
||||
self._get_or_create_agent_app_debug_conversation(agent=agent, account_id=account_id)
|
||||
return agent
|
||||
|
||||
def _create_agent_app_debug_conversation(self, *, app_id: str, account_id: str) -> str:
|
||||
"""Create one console debug conversation for an Agent App editor."""
|
||||
|
||||
conversation = Conversation(
|
||||
app_id=app_id,
|
||||
app_model_config_id=None,
|
||||
model_provider=None,
|
||||
model_id="",
|
||||
override_model_configs=None,
|
||||
mode=AppMode.AGENT,
|
||||
name="Agent Debugging Conversation",
|
||||
inputs={},
|
||||
introduction="",
|
||||
system_instruction="",
|
||||
system_instruction_tokens=0,
|
||||
status=ConversationStatus.NORMAL,
|
||||
invoke_from=InvokeFrom.DEBUGGER,
|
||||
from_source=ConversationFromSource.CONSOLE,
|
||||
from_end_user_id=None,
|
||||
from_account_id=account_id,
|
||||
)
|
||||
self._session.add(conversation)
|
||||
self._session.flush()
|
||||
return conversation.id
|
||||
|
||||
def _get_or_create_agent_app_debug_conversation(self, *, agent: Agent, account_id: str) -> str:
|
||||
if not agent.app_id:
|
||||
raise AgentNotFoundError()
|
||||
|
||||
mapping = self._session.scalar(
|
||||
select(AgentDebugConversation).where(
|
||||
AgentDebugConversation.tenant_id == agent.tenant_id,
|
||||
AgentDebugConversation.agent_id == agent.id,
|
||||
AgentDebugConversation.account_id == account_id,
|
||||
)
|
||||
)
|
||||
if mapping is not None:
|
||||
conversation_id = self._session.scalar(
|
||||
select(Conversation.id).where(
|
||||
Conversation.id == mapping.conversation_id,
|
||||
Conversation.app_id == agent.app_id,
|
||||
Conversation.from_source == ConversationFromSource.CONSOLE,
|
||||
Conversation.from_account_id == account_id,
|
||||
Conversation.is_deleted.is_(False),
|
||||
)
|
||||
)
|
||||
if conversation_id:
|
||||
return conversation_id
|
||||
|
||||
mapping.conversation_id = self._create_agent_app_debug_conversation(
|
||||
app_id=agent.app_id,
|
||||
account_id=account_id,
|
||||
)
|
||||
self._session.flush()
|
||||
return mapping.conversation_id
|
||||
|
||||
conversation_id = self._create_agent_app_debug_conversation(
|
||||
app_id=agent.app_id,
|
||||
account_id=account_id,
|
||||
)
|
||||
self._session.add(
|
||||
AgentDebugConversation(
|
||||
tenant_id=agent.tenant_id,
|
||||
agent_id=agent.id,
|
||||
app_id=agent.app_id,
|
||||
account_id=account_id,
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
)
|
||||
self._session.flush()
|
||||
return conversation_id
|
||||
|
||||
def get_or_create_agent_app_debug_conversation_id(
|
||||
self, *, tenant_id: str, agent_id: str, account_id: str, commit: bool = True
|
||||
) -> str:
|
||||
"""Return the current editor's debug conversation for an Agent App."""
|
||||
|
||||
agent = self._session.scalar(
|
||||
select(Agent).where(
|
||||
Agent.tenant_id == tenant_id,
|
||||
Agent.id == agent_id,
|
||||
Agent.scope == AgentScope.ROSTER,
|
||||
Agent.source == AgentSource.AGENT_APP,
|
||||
Agent.status == AgentStatus.ACTIVE,
|
||||
)
|
||||
)
|
||||
if agent is None:
|
||||
raise AgentNotFoundError()
|
||||
|
||||
conversation_id = self._get_or_create_agent_app_debug_conversation(agent=agent, account_id=account_id)
|
||||
if commit:
|
||||
self._session.commit()
|
||||
return conversation_id
|
||||
|
||||
def refresh_agent_app_debug_conversation_id(
|
||||
self, *, tenant_id: str, agent_id: str, account_id: str, commit: bool = True
|
||||
) -> str:
|
||||
"""Start a new console debug conversation for the current Agent App editor."""
|
||||
|
||||
agent = self._session.scalar(
|
||||
select(Agent).where(
|
||||
Agent.tenant_id == tenant_id,
|
||||
Agent.id == agent_id,
|
||||
Agent.scope == AgentScope.ROSTER,
|
||||
Agent.source == AgentSource.AGENT_APP,
|
||||
Agent.status == AgentStatus.ACTIVE,
|
||||
)
|
||||
)
|
||||
if agent is None or not agent.app_id:
|
||||
raise AgentNotFoundError()
|
||||
|
||||
conversation_id = self._create_agent_app_debug_conversation(
|
||||
app_id=agent.app_id,
|
||||
account_id=account_id,
|
||||
)
|
||||
mapping = self._session.scalar(
|
||||
select(AgentDebugConversation).where(
|
||||
AgentDebugConversation.tenant_id == tenant_id,
|
||||
AgentDebugConversation.agent_id == agent_id,
|
||||
AgentDebugConversation.account_id == account_id,
|
||||
)
|
||||
)
|
||||
if mapping is None:
|
||||
self._session.add(
|
||||
AgentDebugConversation(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
app_id=agent.app_id,
|
||||
account_id=account_id,
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
mapping.app_id = agent.app_id
|
||||
mapping.conversation_id = conversation_id
|
||||
self._session.flush()
|
||||
if commit:
|
||||
self._session.commit()
|
||||
return conversation_id
|
||||
|
||||
def load_or_create_agent_app_debug_conversation_ids_by_agent_id(
|
||||
self, *, tenant_id: str, agents: list[Agent], account_id: str
|
||||
) -> dict[str, str]:
|
||||
"""Return per-account debug conversations for a page of Agent Apps."""
|
||||
|
||||
conversation_ids_by_agent_id: dict[str, str] = {}
|
||||
changed = False
|
||||
for agent in agents:
|
||||
if (
|
||||
agent.tenant_id != tenant_id
|
||||
or agent.scope != AgentScope.ROSTER
|
||||
or agent.source != AgentSource.AGENT_APP
|
||||
):
|
||||
continue
|
||||
conversation_ids_by_agent_id[agent.id] = self._get_or_create_agent_app_debug_conversation(
|
||||
agent=agent,
|
||||
account_id=account_id,
|
||||
)
|
||||
changed = True
|
||||
if changed:
|
||||
self._session.commit()
|
||||
return conversation_ids_by_agent_id
|
||||
|
||||
def load_app_backing_agents_by_app_id(self, *, tenant_id: str, app_ids: list[str]) -> dict[str, Agent]:
|
||||
"""Return active app-backed Agents keyed by Agent App id."""
|
||||
if not app_ids:
|
||||
@@ -466,6 +633,7 @@ class AgentRosterService:
|
||||
account: Any,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
role: str | None = None,
|
||||
icon_type: Any = None,
|
||||
icon: str | None = None,
|
||||
icon_background: str | None = None,
|
||||
@@ -477,6 +645,7 @@ class AgentRosterService:
|
||||
|
||||
copied_name = name or self._next_duplicate_agent_name(tenant_id=tenant_id, base_name=source_app.name)
|
||||
copied_description = description if description is not None else source_app.description
|
||||
copied_role = role if role is not None else source_agent.role or ""
|
||||
copied_icon_type = icon_type if icon_type is not None else source_app.icon_type
|
||||
copied_icon = icon if icon is not None else source_app.icon
|
||||
copied_icon_background = icon_background if icon_background is not None else source_app.icon_background
|
||||
@@ -487,7 +656,7 @@ class AgentRosterService:
|
||||
name=copied_name,
|
||||
description=copied_description,
|
||||
mode="agent",
|
||||
agent_role=source_agent.role or "",
|
||||
agent_role=copied_role,
|
||||
icon_type=self._normalize_app_icon_type(copied_icon_type),
|
||||
icon=copied_icon,
|
||||
icon_background=copied_icon_background,
|
||||
@@ -666,12 +835,16 @@ class AgentRosterService:
|
||||
@staticmethod
|
||||
def _visible_version_operations(agent: Agent) -> set[AgentConfigRevisionOperation]:
|
||||
if agent.source == AgentSource.AGENT_APP:
|
||||
return {AgentConfigRevisionOperation.SAVE_NEW_VERSION}
|
||||
return {
|
||||
AgentConfigRevisionOperation.SAVE_NEW_VERSION,
|
||||
AgentConfigRevisionOperation.RESTORE_VERSION,
|
||||
}
|
||||
return {
|
||||
AgentConfigRevisionOperation.CREATE_VERSION,
|
||||
AgentConfigRevisionOperation.SAVE_NEW_VERSION,
|
||||
AgentConfigRevisionOperation.SAVE_NEW_AGENT,
|
||||
AgentConfigRevisionOperation.SAVE_TO_ROSTER,
|
||||
AgentConfigRevisionOperation.RESTORE_VERSION,
|
||||
}
|
||||
|
||||
def active_config_is_published(self, *, tenant_id: str, agent: Agent) -> bool:
|
||||
@@ -764,6 +937,46 @@ class AgentRosterService:
|
||||
]
|
||||
return result
|
||||
|
||||
def restore_agent_version(
|
||||
self, *, tenant_id: str, agent_id: str, version_id: str, account_id: str
|
||||
) -> dict[str, Any]:
|
||||
agent = self._get_agent(tenant_id=tenant_id, agent_id=agent_id, roster_only=True)
|
||||
visible_version_ids = self._visible_version_ids_stmt(tenant_id=tenant_id, agent_id=agent_id, agent=agent)
|
||||
visible_version_id = self._session.scalar(
|
||||
select(AgentConfigSnapshot.id)
|
||||
.where(
|
||||
AgentConfigSnapshot.tenant_id == tenant_id,
|
||||
AgentConfigSnapshot.agent_id == agent_id,
|
||||
AgentConfigSnapshot.id == version_id,
|
||||
AgentConfigSnapshot.id.in_(select(visible_version_ids.c.current_snapshot_id)),
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
if not visible_version_id:
|
||||
raise AgentVersionNotFoundError()
|
||||
|
||||
version = self._get_version(tenant_id=tenant_id, agent_id=agent_id, version_id=version_id)
|
||||
if agent.active_config_snapshot_id == version.id:
|
||||
return {"result": "success", "active_config_snapshot_id": version.id}
|
||||
|
||||
previous_snapshot_id = agent.active_config_snapshot_id
|
||||
agent.active_config_snapshot_id = version.id
|
||||
agent.active_config_has_model = agent_soul_has_model(version.config_snapshot)
|
||||
agent.updated_by = account_id
|
||||
self._session.add(
|
||||
AgentConfigRevision(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
previous_snapshot_id=previous_snapshot_id,
|
||||
current_snapshot_id=version.id,
|
||||
revision=self._next_revision(tenant_id=tenant_id, agent_id=agent_id),
|
||||
operation=AgentConfigRevisionOperation.RESTORE_VERSION,
|
||||
created_by=account_id,
|
||||
)
|
||||
)
|
||||
self._session.commit()
|
||||
return {"result": "success", "active_config_snapshot_id": version.id}
|
||||
|
||||
def _get_agent(self, *, tenant_id: str, agent_id: str, roster_only: bool = False) -> Agent:
|
||||
stmt = select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id)
|
||||
if roster_only:
|
||||
@@ -789,6 +1002,17 @@ class AgentRosterService:
|
||||
raise AgentVersionNotFoundError()
|
||||
return version
|
||||
|
||||
def _next_revision(self, *, tenant_id: str, agent_id: str) -> int:
|
||||
return (
|
||||
self._session.scalar(
|
||||
select(func.max(AgentConfigRevision.revision)).where(
|
||||
AgentConfigRevision.tenant_id == tenant_id,
|
||||
AgentConfigRevision.agent_id == agent_id,
|
||||
)
|
||||
)
|
||||
or 0
|
||||
) + 1
|
||||
|
||||
def _load_published_active_snapshot_agent_ids(self, *, tenant_id: str, agents: list[Agent]) -> set[str]:
|
||||
predicates = [
|
||||
and_(
|
||||
|
||||
@@ -17,13 +17,15 @@ normalization.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mimetypes
|
||||
import posixpath
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from core.tools.tool_file_manager import ToolFileManager
|
||||
from models.agent_config_entities import AgentSkillRefConfig
|
||||
from services.agent.skill_package_service import SkillPackageService
|
||||
from services.agent_drive_service import AgentDriveService, DriveCommitItem, DriveFileRef
|
||||
from services.agent_drive_service import AgentDriveService, DriveCommitItem, DriveFileRef, DriveSkillMetadata
|
||||
|
||||
_FULL_ARCHIVE_NAME = ".DIFY-SKILL-FULL.zip"
|
||||
_SKILL_MD_NAME = "SKILL.md"
|
||||
@@ -62,7 +64,8 @@ class SkillStandardizeService:
|
||||
skill_md_bytes = self._package.read_member_bytes(content=content, member_path=manifest.entry_path)
|
||||
slug = slugify_skill_name(manifest.name)
|
||||
|
||||
# Two drive-owned ToolFiles: canonical SKILL.md + the full archive.
|
||||
# Drive-owned files: canonical SKILL.md, every inspectable archive file,
|
||||
# and the full archive for future restore/export.
|
||||
md_tool_file = self._tool_files.create_file_by_raw(
|
||||
user_id=user_id,
|
||||
tenant_id=tenant_id,
|
||||
@@ -82,6 +85,30 @@ class SkillStandardizeService:
|
||||
|
||||
skill_md_key = f"{slug}/{_SKILL_MD_NAME}"
|
||||
archive_key = f"{slug}/{_FULL_ARCHIVE_NAME}"
|
||||
member_items: list[DriveCommitItem] = []
|
||||
for member_path in sorted(set(manifest.files)):
|
||||
member_key = f"{slug}/{member_path}"
|
||||
if member_key in {skill_md_key, archive_key}:
|
||||
continue
|
||||
|
||||
member_bytes = self._package.read_member_bytes(content=content, member_path=member_path)
|
||||
mimetype = mimetypes.guess_type(member_path)[0] or "application/octet-stream"
|
||||
member_tool_file = self._tool_files.create_file_by_raw(
|
||||
user_id=user_id,
|
||||
tenant_id=tenant_id,
|
||||
conversation_id=None,
|
||||
file_binary=member_bytes,
|
||||
mimetype=mimetype,
|
||||
filename=posixpath.basename(member_path),
|
||||
)
|
||||
member_items.append(
|
||||
DriveCommitItem(
|
||||
key=member_key,
|
||||
file_ref=DriveFileRef(kind="tool_file", id=member_tool_file.id),
|
||||
value_owned_by_drive=True,
|
||||
)
|
||||
)
|
||||
|
||||
self._drive.commit(
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
@@ -91,12 +118,19 @@ class SkillStandardizeService:
|
||||
key=skill_md_key,
|
||||
file_ref=DriveFileRef(kind="tool_file", id=md_tool_file.id),
|
||||
value_owned_by_drive=True,
|
||||
is_skill=True,
|
||||
skill_metadata=DriveSkillMetadata(
|
||||
name=manifest.name,
|
||||
description=manifest.description,
|
||||
manifest_files=manifest.files,
|
||||
),
|
||||
),
|
||||
DriveCommitItem(
|
||||
key=archive_key,
|
||||
file_ref=DriveFileRef(kind="tool_file", id=archive_tool_file.id),
|
||||
value_owned_by_drive=True,
|
||||
),
|
||||
*member_items,
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@@ -17,12 +17,14 @@ ToolFile records (see ``AgentDriveFile``). This service is the control plane:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import urllib.parse
|
||||
from typing import Any, Literal
|
||||
from typing import Any, Literal, TypedDict
|
||||
from urllib.parse import unquote
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ConfigDict, field_validator
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.exc import DataError, SQLAlchemyError
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -41,6 +43,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_MAX_KEY_LENGTH = 512
|
||||
_DRIVE_REF_PREFIX = "agent-"
|
||||
_SKILL_MD_SUFFIX = "/SKILL.md"
|
||||
_SKILL_ARCHIVE_NAME = ".DIFY-SKILL-FULL.zip"
|
||||
|
||||
|
||||
class AgentDriveError(Exception):
|
||||
@@ -58,16 +62,86 @@ class AgentDriveError(Exception):
|
||||
|
||||
|
||||
class DriveFileRef(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
kind: Literal["upload_file", "tool_file"]
|
||||
id: str
|
||||
|
||||
|
||||
class DriveSkillMetadata(BaseModel):
|
||||
"""Validated skill catalog metadata stored as a JSON string on the drive row."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
name: str
|
||||
description: str = ""
|
||||
# Safe archive member paths captured during skill standardization. The drive
|
||||
# stores only canonical SKILL.md + full archive, so the UI uses this manifest
|
||||
# to show the original uploaded package contents.
|
||||
manifest_files: list[str] | None = None
|
||||
|
||||
@field_validator("name")
|
||||
@classmethod
|
||||
def _validate_name(cls, value: str) -> str:
|
||||
normalized = value.strip()
|
||||
if not normalized:
|
||||
raise ValueError("skill metadata name must not be blank")
|
||||
return normalized
|
||||
|
||||
|
||||
class DriveCommitItem(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
key: str
|
||||
file_ref: DriveFileRef
|
||||
# Drive-owned values may be physically cleaned on overwrite/removal; refs to
|
||||
# files shared with other business records should set this False.
|
||||
value_owned_by_drive: bool = True
|
||||
is_skill: bool = False
|
||||
skill_metadata: DriveSkillMetadata | None = None
|
||||
|
||||
|
||||
class AgentDriveSkillInfo(TypedDict):
|
||||
path: str
|
||||
skill_md_key: str
|
||||
archive_key: str | None
|
||||
name: str
|
||||
description: str
|
||||
size: int | None
|
||||
mime_type: str | None
|
||||
hash: str | None
|
||||
created_at: int | None
|
||||
|
||||
|
||||
class AgentDriveSkillFileInfo(TypedDict):
|
||||
path: str
|
||||
name: str
|
||||
type: str
|
||||
drive_key: str | None
|
||||
available_in_drive: bool
|
||||
|
||||
|
||||
class AgentDriveSkillInspectInfo(TypedDict):
|
||||
path: str
|
||||
skill_md_key: str
|
||||
archive_key: str | None
|
||||
name: str
|
||||
description: str
|
||||
size: int | None
|
||||
mime_type: str | None
|
||||
hash: str | None
|
||||
created_at: int | None
|
||||
source: str
|
||||
files: list[AgentDriveSkillFileInfo]
|
||||
file_tree: list[dict[str, Any]]
|
||||
skill_md: dict[str, Any]
|
||||
warnings: list[str]
|
||||
|
||||
|
||||
def decode_drive_mention_ref(ref_id: str) -> str:
|
||||
"""Decode the prompt token's URL-encoded drive-key field."""
|
||||
|
||||
return unquote(ref_id or "")
|
||||
|
||||
|
||||
def parse_agent_drive_ref(drive_ref: str) -> str:
|
||||
@@ -132,6 +206,8 @@ class AgentDriveService:
|
||||
"mime_type": row.mime_type,
|
||||
"file_kind": row.file_kind.value,
|
||||
"file_id": row.file_id,
|
||||
"is_skill": row.is_skill,
|
||||
"skill_metadata": row.skill_metadata,
|
||||
"created_at": int(row.created_at.timestamp()) if row.created_at else None,
|
||||
}
|
||||
if include_download_url:
|
||||
@@ -217,6 +293,87 @@ class AgentDriveService:
|
||||
self._delete_storage(storage_key)
|
||||
return removed_keys
|
||||
|
||||
def list_skills(self, *, tenant_id: str, agent_id: str) -> list[AgentDriveSkillInfo]:
|
||||
"""Return the drive-backed skill catalog derived from canonical ``SKILL.md`` rows."""
|
||||
|
||||
with session_factory.create_session() as session:
|
||||
self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id)
|
||||
skill_rows = list(
|
||||
session.scalars(
|
||||
select(AgentDriveFile)
|
||||
.where(
|
||||
AgentDriveFile.tenant_id == tenant_id,
|
||||
AgentDriveFile.agent_id == agent_id,
|
||||
AgentDriveFile.is_skill.is_(True),
|
||||
)
|
||||
.order_by(AgentDriveFile.key)
|
||||
)
|
||||
)
|
||||
archive_keys = set(
|
||||
session.scalars(
|
||||
select(AgentDriveFile.key).where(
|
||||
AgentDriveFile.tenant_id == tenant_id,
|
||||
AgentDriveFile.agent_id == agent_id,
|
||||
AgentDriveFile.key.in_([self._skill_archive_key(row.key) for row in skill_rows]),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
skills: list[AgentDriveSkillInfo] = []
|
||||
for row in skill_rows:
|
||||
metadata = self._parse_skill_metadata(row.key, row.skill_metadata)
|
||||
archive_key = self._skill_archive_key(row.key)
|
||||
skills.append(
|
||||
{
|
||||
"path": self._skill_path_from_key(row.key),
|
||||
"skill_md_key": row.key,
|
||||
"archive_key": archive_key if archive_key in archive_keys else None,
|
||||
"name": metadata.name,
|
||||
"description": metadata.description,
|
||||
"size": row.size,
|
||||
"mime_type": row.mime_type,
|
||||
"hash": row.hash,
|
||||
"created_at": int(row.created_at.timestamp()) if row.created_at else None,
|
||||
}
|
||||
)
|
||||
return skills
|
||||
|
||||
def inspect_skill(self, *, tenant_id: str, agent_id: str, skill_path: str) -> AgentDriveSkillInspectInfo:
|
||||
"""Return the UI-facing skill inspect view for slash-menu hover/detail."""
|
||||
|
||||
skill_path = normalize_drive_key(skill_path)
|
||||
skill_md_key = skill_path if skill_path.endswith(_SKILL_MD_SUFFIX) else f"{skill_path}{_SKILL_MD_SUFFIX}"
|
||||
skill_path = self._skill_path_from_key(skill_md_key)
|
||||
catalog = next(
|
||||
(item for item in self.list_skills(tenant_id=tenant_id, agent_id=agent_id) if item["path"] == skill_path),
|
||||
None,
|
||||
)
|
||||
if catalog is None:
|
||||
raise AgentDriveError("skill_not_found", "no drive-backed skill for this path", status_code=404)
|
||||
|
||||
manifest_files = self._manifest_files_from_skill_metadata(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
skill_md_key=skill_md_key,
|
||||
)
|
||||
drive_items = self.manifest(tenant_id=tenant_id, agent_id=agent_id, prefix=f"{skill_path}/")
|
||||
drive_keys = {item["key"] for item in drive_items}
|
||||
preview = self.preview(tenant_id=tenant_id, agent_id=agent_id, key=skill_md_key)
|
||||
files, warnings = self._skill_file_entries(
|
||||
skill_path=skill_path,
|
||||
skill_md_key=skill_md_key,
|
||||
manifest_files=manifest_files,
|
||||
drive_keys=drive_keys,
|
||||
)
|
||||
return {
|
||||
**catalog,
|
||||
"source": "skill_md",
|
||||
"files": files,
|
||||
"file_tree": self._build_file_tree(files),
|
||||
"skill_md": preview,
|
||||
"warnings": warnings,
|
||||
}
|
||||
|
||||
def _commit_one(
|
||||
self,
|
||||
session: Session,
|
||||
@@ -228,9 +385,10 @@ class AgentDriveService:
|
||||
pending_storage_deletes: list[str],
|
||||
) -> dict[str, Any]:
|
||||
key = normalize_drive_key(item.key)
|
||||
skill_metadata = self._validate_skill_commit_fields(key=key, item=item)
|
||||
file_kind = AgentDriveFileKind(item.file_ref.kind)
|
||||
file_id = item.file_ref.id
|
||||
size, mime_type = self._validate_source(
|
||||
size, mime_type, file_hash = self._validate_source(
|
||||
session, tenant_id=tenant_id, user_id=user_id, file_kind=file_kind, file_id=file_id
|
||||
)
|
||||
|
||||
@@ -245,6 +403,11 @@ class AgentDriveService:
|
||||
# Idempotent re-commit of the same value: leave it (do not clean).
|
||||
if existing.file_kind == file_kind and existing.file_id == file_id:
|
||||
existing.value_owned_by_drive = item.value_owned_by_drive
|
||||
existing.is_skill = item.is_skill
|
||||
existing.skill_metadata = skill_metadata
|
||||
existing.size = size
|
||||
existing.mime_type = mime_type
|
||||
existing.hash = file_hash
|
||||
return self._row_dict(existing)
|
||||
# Overwrite: clean the previous drive-owned value if no longer referenced.
|
||||
if existing.value_owned_by_drive:
|
||||
@@ -259,7 +422,10 @@ class AgentDriveService:
|
||||
existing.file_kind = file_kind
|
||||
existing.file_id = file_id
|
||||
existing.value_owned_by_drive = item.value_owned_by_drive
|
||||
existing.is_skill = item.is_skill
|
||||
existing.skill_metadata = skill_metadata
|
||||
existing.size = size
|
||||
existing.hash = file_hash
|
||||
existing.mime_type = mime_type
|
||||
return self._row_dict(existing)
|
||||
|
||||
@@ -271,7 +437,10 @@ class AgentDriveService:
|
||||
file_kind=file_kind,
|
||||
file_id=file_id,
|
||||
value_owned_by_drive=item.value_owned_by_drive,
|
||||
is_skill=item.is_skill,
|
||||
skill_metadata=skill_metadata,
|
||||
size=size,
|
||||
hash=file_hash,
|
||||
mime_type=mime_type,
|
||||
created_by=user_id,
|
||||
)
|
||||
@@ -287,8 +456,187 @@ class AgentDriveService:
|
||||
"size": row.size,
|
||||
"mime_type": row.mime_type,
|
||||
"value_owned_by_drive": row.value_owned_by_drive,
|
||||
"is_skill": row.is_skill,
|
||||
"skill_metadata": row.skill_metadata,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _skill_path_from_key(key: str) -> str:
|
||||
if not key.endswith(_SKILL_MD_SUFFIX):
|
||||
raise AgentDriveError(
|
||||
"invalid_skill_key",
|
||||
"skill rows must use the canonical '<path>/SKILL.md' key",
|
||||
status_code=500,
|
||||
)
|
||||
path = key[: -len(_SKILL_MD_SUFFIX)]
|
||||
if not path:
|
||||
raise AgentDriveError(
|
||||
"invalid_skill_key",
|
||||
"skill rows must use the canonical '<path>/SKILL.md' key",
|
||||
status_code=500,
|
||||
)
|
||||
return path
|
||||
|
||||
@classmethod
|
||||
def _skill_archive_key(cls, key: str) -> str:
|
||||
return f"{cls._skill_path_from_key(key)}/{_SKILL_ARCHIVE_NAME}"
|
||||
|
||||
@classmethod
|
||||
def _validate_skill_commit_fields(cls, *, key: str, item: DriveCommitItem) -> str | None:
|
||||
if not item.is_skill:
|
||||
if item.skill_metadata is not None:
|
||||
raise AgentDriveError(
|
||||
"invalid_skill_metadata",
|
||||
"skill metadata is only allowed for canonical skill rows",
|
||||
status_code=400,
|
||||
)
|
||||
return None
|
||||
cls._skill_path_from_key(key)
|
||||
if item.skill_metadata is None:
|
||||
raise AgentDriveError(
|
||||
"invalid_skill_metadata",
|
||||
"skill metadata is required for canonical skill rows",
|
||||
status_code=400,
|
||||
)
|
||||
return json.dumps(
|
||||
item.skill_metadata.model_dump(mode="json", exclude_none=True),
|
||||
separators=(",", ":"),
|
||||
sort_keys=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_skill_metadata(key: str, raw_metadata: str | None) -> DriveSkillMetadata:
|
||||
if raw_metadata is None:
|
||||
raise AgentDriveError(
|
||||
"invalid_skill_metadata",
|
||||
f"skill row '{key}' is missing required metadata",
|
||||
status_code=500,
|
||||
)
|
||||
try:
|
||||
return DriveSkillMetadata.model_validate(json.loads(raw_metadata))
|
||||
except (ValueError, TypeError) as exc:
|
||||
raise AgentDriveError(
|
||||
"invalid_skill_metadata",
|
||||
f"skill row '{key}' has invalid stored metadata",
|
||||
status_code=500,
|
||||
) from exc
|
||||
|
||||
@staticmethod
|
||||
def _manifest_files_from_skill_metadata(*, tenant_id: str, agent_id: str, skill_md_key: str) -> list[str] | None:
|
||||
with session_factory.create_session() as session:
|
||||
row = session.scalar(
|
||||
select(AgentDriveFile).where(
|
||||
AgentDriveFile.tenant_id == tenant_id,
|
||||
AgentDriveFile.agent_id == agent_id,
|
||||
AgentDriveFile.key == skill_md_key,
|
||||
AgentDriveFile.is_skill.is_(True),
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
try:
|
||||
metadata = AgentDriveService._parse_skill_metadata(row.key, row.skill_metadata)
|
||||
except Exception:
|
||||
logger.warning("drive skill inspect: malformed skill metadata for %s", skill_md_key, exc_info=True)
|
||||
return None
|
||||
return [str(item) for item in (metadata.manifest_files or []) if str(item).strip()] or None
|
||||
|
||||
@classmethod
|
||||
def _skill_file_entries(
|
||||
cls,
|
||||
*,
|
||||
skill_path: str,
|
||||
skill_md_key: str,
|
||||
manifest_files: list[str] | None,
|
||||
drive_keys: set[str],
|
||||
) -> tuple[list[AgentDriveSkillFileInfo], list[str]]:
|
||||
warnings: list[str] = []
|
||||
if manifest_files:
|
||||
paths = sorted({normalize_drive_key(path) for path in manifest_files})
|
||||
else:
|
||||
paths = sorted(
|
||||
{
|
||||
key.removeprefix(f"{skill_path}/")
|
||||
for key in drive_keys
|
||||
if not key.endswith(f"/{_SKILL_ARCHIVE_NAME}")
|
||||
}
|
||||
)
|
||||
warnings.append("manifest_files_unavailable")
|
||||
|
||||
files: list[AgentDriveSkillFileInfo] = []
|
||||
for path in paths:
|
||||
if path == _SKILL_ARCHIVE_NAME:
|
||||
continue
|
||||
drive_key = f"{skill_path}/{path}"
|
||||
files.append(
|
||||
{
|
||||
"path": path,
|
||||
"name": path.rsplit("/", 1)[-1],
|
||||
"type": "file",
|
||||
"drive_key": drive_key if drive_key in drive_keys else None,
|
||||
"available_in_drive": drive_key in drive_keys,
|
||||
}
|
||||
)
|
||||
if "SKILL.md" not in {file["path"] for file in files}:
|
||||
files.insert(
|
||||
0,
|
||||
{
|
||||
"path": "SKILL.md",
|
||||
"name": "SKILL.md",
|
||||
"type": "file",
|
||||
"drive_key": skill_md_key,
|
||||
"available_in_drive": skill_md_key in drive_keys,
|
||||
},
|
||||
)
|
||||
return files, warnings
|
||||
|
||||
@staticmethod
|
||||
def _build_file_tree(files: list[AgentDriveSkillFileInfo]) -> list[dict[str, Any]]:
|
||||
root: dict[str, Any] = {}
|
||||
for file in files:
|
||||
cursor = root
|
||||
parts = [part for part in file["path"].split("/") if part]
|
||||
path_parts: list[str] = []
|
||||
for part in parts[:-1]:
|
||||
path_parts.append(part)
|
||||
directory = cursor.setdefault(
|
||||
part,
|
||||
{
|
||||
"name": part,
|
||||
"path": "/".join(path_parts),
|
||||
"type": "directory",
|
||||
"children": {},
|
||||
},
|
||||
)
|
||||
cursor = directory["children"]
|
||||
leaf_name = parts[-1] if parts else file["name"]
|
||||
cursor[leaf_name] = {
|
||||
"name": leaf_name,
|
||||
"path": file["path"],
|
||||
"type": file["type"],
|
||||
"drive_key": file["drive_key"],
|
||||
"available_in_drive": file["available_in_drive"],
|
||||
}
|
||||
|
||||
def serialize(node: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
result: list[dict[str, Any]] = []
|
||||
for item in sorted(node.values(), key=lambda value: (value["type"] != "directory", value["name"])):
|
||||
if item["type"] == "directory":
|
||||
children = serialize(item["children"])
|
||||
result.append(
|
||||
{
|
||||
"name": item["name"],
|
||||
"path": item["path"],
|
||||
"type": "directory",
|
||||
"children": children,
|
||||
}
|
||||
)
|
||||
else:
|
||||
result.append(item)
|
||||
return result
|
||||
|
||||
return serialize(root)
|
||||
|
||||
@staticmethod
|
||||
def _assert_agent_belongs_to_tenant(session: Session, *, tenant_id: str, agent_id: str) -> None:
|
||||
try:
|
||||
@@ -309,7 +657,7 @@ class AgentDriveService:
|
||||
user_id: str,
|
||||
file_kind: AgentDriveFileKind,
|
||||
file_id: str,
|
||||
) -> tuple[int | None, str | None]:
|
||||
) -> tuple[int | None, str | None, str | None]:
|
||||
"""Verify the source file exists for the tenant (and user, for ToolFile).
|
||||
|
||||
Malformed ids (e.g. a non-UUID hitting a UUID column) are treated as a
|
||||
@@ -328,7 +676,7 @@ class AgentDriveService:
|
||||
raise AgentDriveError(
|
||||
"source_not_found", "source ToolFile not found for this tenant/user", status_code=404
|
||||
)
|
||||
return tool_file.size, tool_file.mimetype
|
||||
return tool_file.size, tool_file.mimetype, None
|
||||
upload_file = session.scalar(
|
||||
select(UploadFile).where(UploadFile.id == file_id, UploadFile.tenant_id == tenant_id)
|
||||
)
|
||||
@@ -337,7 +685,7 @@ class AgentDriveService:
|
||||
raise AgentDriveError("source_not_found", "source file ref is invalid", status_code=404) from exc
|
||||
if upload_file is None:
|
||||
raise AgentDriveError("source_not_found", "source UploadFile not found for this tenant", status_code=404)
|
||||
return upload_file.size, upload_file.mime_type
|
||||
return upload_file.size, upload_file.mime_type, upload_file.hash
|
||||
|
||||
def _cleanup_value(
|
||||
self,
|
||||
@@ -509,6 +857,8 @@ __all__ = [
|
||||
"AgentDriveService",
|
||||
"DriveCommitItem",
|
||||
"DriveFileRef",
|
||||
"DriveSkillMetadata",
|
||||
"decode_drive_mention_ref",
|
||||
"normalize_drive_key",
|
||||
"parse_agent_drive_ref",
|
||||
]
|
||||
|
||||
@@ -5,6 +5,10 @@ import httpx
|
||||
|
||||
from services.auth.api_key_auth_base import ApiKeyAuthBase, AuthCredentials
|
||||
|
||||
# Explicit bounded timeout for credential-validation requests so a slow or
|
||||
# hanging Firecrawl endpoint cannot block the worker indefinitely.
|
||||
_CREDENTIAL_TIMEOUT = httpx.Timeout(10.0)
|
||||
|
||||
|
||||
class FirecrawlAuth(ApiKeyAuthBase):
|
||||
def __init__(self, credentials: AuthCredentials):
|
||||
@@ -42,7 +46,7 @@ class FirecrawlAuth(ApiKeyAuthBase):
|
||||
return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}"
|
||||
|
||||
def _post_request(self, url, data, headers):
|
||||
return httpx.post(url, headers=headers, json=data)
|
||||
return httpx.post(url, headers=headers, json=data, timeout=_CREDENTIAL_TIMEOUT)
|
||||
|
||||
def _handle_error(self, response):
|
||||
try:
|
||||
|
||||
@@ -8,7 +8,10 @@ from services.auth.api_key_auth_base import ApiKeyAuthBase, AuthCredentials
|
||||
|
||||
_http_client: httpx.Client = get_pooled_http_client(
|
||||
"auth:jina_standalone",
|
||||
lambda: httpx.Client(limits=httpx.Limits(max_keepalive_connections=50, max_connections=100)),
|
||||
lambda: httpx.Client(
|
||||
timeout=httpx.Timeout(10.0),
|
||||
limits=httpx.Limits(max_keepalive_connections=50, max_connections=100),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -8,7 +8,10 @@ from services.auth.api_key_auth_base import ApiKeyAuthBase, AuthCredentials
|
||||
|
||||
_http_client: httpx.Client = get_pooled_http_client(
|
||||
"auth:jina",
|
||||
lambda: httpx.Client(limits=httpx.Limits(max_keepalive_connections=50, max_connections=100)),
|
||||
lambda: httpx.Client(
|
||||
timeout=httpx.Timeout(10.0),
|
||||
limits=httpx.Limits(max_keepalive_connections=50, max_connections=100),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,4 +1,12 @@
|
||||
"""Tenant credit pool accounting.
|
||||
|
||||
Credit deductions are guarded by a tenant-level Redis lock before the database
|
||||
row lock is acquired. This keeps concurrent usage accounting for one tenant
|
||||
from piling up database transactions while preserving cross-tenant concurrency.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -7,13 +15,44 @@ from configs import dify_config
|
||||
from core.db.session_factory import session_factory
|
||||
from core.errors.error import QuotaExceededError
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from models import TenantCreditPool
|
||||
from models.enums import ProviderQuotaType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS = 10
|
||||
CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS = 5
|
||||
|
||||
|
||||
class CreditPoolService:
|
||||
@staticmethod
|
||||
def _get_tenant_lock_key(tenant_id: str) -> str:
|
||||
return f"credit_pool:tenant:{tenant_id}:deduct_lock"
|
||||
|
||||
@classmethod
|
||||
def _deduct_with_tenant_lock(cls, tenant_id: str, deduct: Callable[[], int]) -> int:
|
||||
lock_key = cls._get_tenant_lock_key(tenant_id)
|
||||
lock = redis_client.lock(
|
||||
lock_key,
|
||||
timeout=CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS,
|
||||
blocking_timeout=CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS,
|
||||
)
|
||||
lock_acquired = False
|
||||
|
||||
try:
|
||||
lock_acquired = lock.acquire(blocking=True)
|
||||
if not lock_acquired:
|
||||
raise QuotaExceededError("Failed to acquire credit pool lock")
|
||||
|
||||
return deduct()
|
||||
finally:
|
||||
if lock_acquired:
|
||||
try:
|
||||
lock.release()
|
||||
except Exception:
|
||||
logger.warning("Failed to release credit pool lock, tenant_id=%s", tenant_id, exc_info=True)
|
||||
|
||||
@staticmethod
|
||||
def _get_locked_pool(session: Session, tenant_id: str, pool_type: str) -> TenantCreditPool | None:
|
||||
return session.scalar(
|
||||
@@ -76,7 +115,7 @@ class CreditPoolService:
|
||||
if credits_required <= 0:
|
||||
return 0
|
||||
|
||||
try:
|
||||
def deduct() -> int:
|
||||
with session_factory.get_session_maker().begin() as session:
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
|
||||
if not pool:
|
||||
@@ -89,14 +128,16 @@ class CreditPoolService:
|
||||
raise QuotaExceededError("Insufficient credits remaining")
|
||||
|
||||
pool.quota_used += credits_required
|
||||
return credits_required
|
||||
|
||||
try:
|
||||
return cls._deduct_with_tenant_lock(tenant_id, deduct)
|
||||
except QuotaExceededError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("Failed to deduct credits for tenant %s", tenant_id)
|
||||
raise QuotaExceededError("Failed to deduct credits")
|
||||
|
||||
return credits_required
|
||||
|
||||
@classmethod
|
||||
def deduct_credits_capped(
|
||||
cls,
|
||||
@@ -108,7 +149,7 @@ class CreditPoolService:
|
||||
if credits_required <= 0:
|
||||
return 0
|
||||
|
||||
try:
|
||||
def deduct() -> int:
|
||||
with session_factory.get_session_maker().begin() as session:
|
||||
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
|
||||
if not pool:
|
||||
@@ -121,6 +162,9 @@ class CreditPoolService:
|
||||
|
||||
pool.quota_used += deducted_credits
|
||||
return deducted_credits
|
||||
|
||||
try:
|
||||
return cls._deduct_with_tenant_lock(tenant_id, deduct)
|
||||
except QuotaExceededError:
|
||||
raise
|
||||
except Exception:
|
||||
|
||||
@@ -182,6 +182,10 @@ class EnterpriseRequest(BaseRequest):
|
||||
inner_headers: dict[str, str] = {INNER_TENANT_ID_HEADER: tenant_id}
|
||||
if account_id:
|
||||
inner_headers[INNER_ACCOUNT_ID_HEADER] = account_id
|
||||
|
||||
if not cls.base_url.startswith("http") or not cls.base_url.startswith("https") or not cls.base_url:
|
||||
raise ValueError("ENTERPRISE_RBAC_API_URL is required when RBAC_ENABLED=true")
|
||||
|
||||
url = f"{cls.rbac_base_url}{endpoint}"
|
||||
mounts = cls._build_mounts()
|
||||
|
||||
|
||||
@@ -312,15 +312,26 @@ _LEGACY_WORKSPACE_OWNER_KEYS: list[str] = [
|
||||
"plugin.manage",
|
||||
"plugin.debug",
|
||||
"credential.use",
|
||||
"credential.create",
|
||||
"credential.manage",
|
||||
"billing.view",
|
||||
"billing.subscription.manage",
|
||||
"billing.manage",
|
||||
"app.acl.preview",
|
||||
"app_library.access",
|
||||
"app.create_and_management",
|
||||
"app.tag.manage",
|
||||
"dataset.acl.preview",
|
||||
"dataset.create_and_management",
|
||||
"dataset.tag.manage",
|
||||
"dataset.external.connect",
|
||||
"dataset.api_key.manage",
|
||||
"snippets.create_and_modify",
|
||||
"snippets.management",
|
||||
"tool.manage",
|
||||
"mcp.manage",
|
||||
"snippets.create_and_modify",
|
||||
"snippets.management",
|
||||
]
|
||||
|
||||
_LEGACY_WORKSPACE_ADMIN_KEYS: list[str] = [
|
||||
@@ -334,15 +345,24 @@ _LEGACY_WORKSPACE_ADMIN_KEYS: list[str] = [
|
||||
"plugin.manage",
|
||||
"plugin.debug",
|
||||
"credential.use",
|
||||
"credential.create",
|
||||
"credential.manage",
|
||||
"billing.view",
|
||||
"billing.subscription.manage",
|
||||
"billing.manage",
|
||||
"app_library.access",
|
||||
"app.create_and_management",
|
||||
"app.tag.manage",
|
||||
"dataset.create_and_management",
|
||||
"dataset.tag.manage",
|
||||
"dataset.external.connect",
|
||||
"dataset.api_key.manage",
|
||||
"snippets.create_and_modify",
|
||||
"snippets.management",
|
||||
"tool.manage",
|
||||
"mcp.manage",
|
||||
"snippets.create_and_modify",
|
||||
"snippets.management",
|
||||
]
|
||||
|
||||
_LEGACY_WORKSPACE_EDITOR_KEYS: list[str] = [
|
||||
@@ -356,7 +376,9 @@ _LEGACY_WORKSPACE_EDITOR_KEYS: list[str] = [
|
||||
"dataset.create_and_management",
|
||||
"dataset.tag.manage",
|
||||
"dataset.external.connect",
|
||||
"snippets.create_and_modify",
|
||||
"tool.manage",
|
||||
"snippets.create_and_modify",
|
||||
]
|
||||
|
||||
_LEGACY_WORKSPACE_NORMAL_KEYS: list[str] = [
|
||||
@@ -373,6 +395,7 @@ _LEGACY_WORKSPACE_DATASET_OPERATOR_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_APP_OWNER_KEYS: list[str] = [
|
||||
"app.acl.preview",
|
||||
"app.acl.view_layout",
|
||||
"app.acl.test_and_run",
|
||||
"app.acl.edit",
|
||||
@@ -384,6 +407,7 @@ _LEGACY_APP_OWNER_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_APP_ADMIN_KEYS: list[str] = [
|
||||
"app.acl.preview",
|
||||
"app.acl.view_layout",
|
||||
"app.acl.test_and_run",
|
||||
"app.acl.edit",
|
||||
@@ -395,6 +419,7 @@ _LEGACY_APP_ADMIN_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_APP_EDITOR_KEYS: list[str] = [
|
||||
"app.acl.preview",
|
||||
"app.acl.view_layout",
|
||||
"app.acl.test_and_run",
|
||||
"app.acl.edit",
|
||||
@@ -406,12 +431,14 @@ _LEGACY_APP_EDITOR_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_APP_NORMAL_KEYS: list[str] = [
|
||||
"app.acl.preview",
|
||||
"app.acl.view_layout",
|
||||
"app.acl.test_and_run",
|
||||
"app.acl.monitor",
|
||||
]
|
||||
|
||||
_LEGACY_DATASET_OWNER_KEYS: list[str] = [
|
||||
"dataset.acl.preview",
|
||||
"dataset.acl.readonly",
|
||||
"dataset.acl.edit",
|
||||
"dataset.acl.import_export_dsl",
|
||||
@@ -427,6 +454,7 @@ _LEGACY_DATASET_OWNER_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_DATASET_ADMIN_KEYS: list[str] = [
|
||||
"dataset.acl.preview",
|
||||
"dataset.acl.readonly",
|
||||
"dataset.acl.edit",
|
||||
"dataset.acl.import_export_dsl",
|
||||
@@ -442,6 +470,7 @@ _LEGACY_DATASET_ADMIN_KEYS: list[str] = [
|
||||
]
|
||||
|
||||
_LEGACY_DATASET_EDITOR_KEYS: list[str] = [
|
||||
"dataset.acl.preview",
|
||||
"dataset.acl.readonly",
|
||||
"dataset.acl.edit",
|
||||
"dataset.acl.import_export_dsl",
|
||||
@@ -492,6 +521,19 @@ _LEGACY_MY_PERMISSIONS: dict[TenantAccountRole, dict[str, list[str]]] = {
|
||||
}
|
||||
|
||||
|
||||
def _legacy_role_permission_keys(role: TenantAccountRole) -> list[str]:
|
||||
permissions = _LEGACY_MY_PERMISSIONS.get(role, {})
|
||||
return list(
|
||||
dict.fromkeys(
|
||||
[
|
||||
*permissions.get("workspace", []),
|
||||
*permissions.get("app", []),
|
||||
*permissions.get("dataset", []),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _legacy_my_permissions(tenant_id: str, account_id: str | None) -> MyPermissionsResponse:
|
||||
if not account_id:
|
||||
return MyPermissionsResponse()
|
||||
@@ -728,6 +770,7 @@ class RBACService:
|
||||
data = _inner_call(
|
||||
"GET",
|
||||
f"{_INNER_PREFIX}/role-permissions/catalog",
|
||||
params={"billing_enabled": dify_config.BILLING_ENABLED},
|
||||
tenant_id=tenant_id,
|
||||
account_id=account_id,
|
||||
)
|
||||
@@ -1518,21 +1561,44 @@ class RBACService:
|
||||
)
|
||||
return AccessMatrixItem.model_validate(data or {})
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Member ↔ role bindings (screenshot 3: Settings > Members > Assign roles).
|
||||
# ------------------------------------------------------------------
|
||||
class MemberRoles:
|
||||
@staticmethod
|
||||
def get(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse:
|
||||
data = _inner_call(
|
||||
"GET",
|
||||
f"{_INNER_PREFIX}/members/rbac-roles",
|
||||
tenant_id=tenant_id,
|
||||
account_id=account_id,
|
||||
params={"account_id": member_account_id},
|
||||
)
|
||||
rst = MemberRolesResponse.model_validate(data or {})
|
||||
return rst
|
||||
if dify_config.RBAC_ENABLED:
|
||||
data = _inner_call(
|
||||
"GET",
|
||||
f"{_INNER_PREFIX}/members/rbac-roles",
|
||||
tenant_id=tenant_id,
|
||||
account_id=account_id,
|
||||
params={"account_id": member_account_id},
|
||||
)
|
||||
rst = MemberRolesResponse.model_validate(data or {})
|
||||
return rst
|
||||
else:
|
||||
with session_factory.create_session() as session:
|
||||
role = session.scalar(
|
||||
select(TenantAccountJoin.role).where(
|
||||
TenantAccountJoin.tenant_id == tenant_id,
|
||||
TenantAccountJoin.account_id == member_account_id,
|
||||
)
|
||||
)
|
||||
return MemberRolesResponse(
|
||||
account_id=member_account_id,
|
||||
roles=[
|
||||
RBACRole(
|
||||
id=role,
|
||||
name=role,
|
||||
description="",
|
||||
is_builtin=True,
|
||||
type="",
|
||||
permission_keys=_legacy_role_permission_keys(role),
|
||||
role_tag="owner" if role == "owner" else role,
|
||||
tenant_id=tenant_id,
|
||||
)
|
||||
]
|
||||
if role
|
||||
else [],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def batch_get(
|
||||
|
||||
@@ -4,6 +4,7 @@ import time
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session, scoped_session
|
||||
|
||||
from core.app.app_config.entities import ModelConfig
|
||||
from core.rag.datasource.retrieval_service import DefaultRetrievalModelDict, RetrievalService
|
||||
@@ -12,7 +13,6 @@ from core.rag.index_processor.constant.query_type import QueryType
|
||||
from core.rag.models.document import Document
|
||||
from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
|
||||
from core.rag.retrieval.retrieval_methods import RetrievalMethod
|
||||
from extensions.ext_database import db
|
||||
from graphon.model_runtime.entities import LLMMode
|
||||
from models import Account
|
||||
from models.dataset import Dataset, DatasetQuery
|
||||
@@ -56,7 +56,9 @@ class HitTestingService:
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _dump_retrieval_records(cls, records: list[RetrievalSegments]) -> list[dict[str, Any]]:
|
||||
def _dump_retrieval_records(
|
||||
cls, session: Session | scoped_session, records: list[RetrievalSegments]
|
||||
) -> list[dict[str, Any]]:
|
||||
document_ids = {
|
||||
document_id
|
||||
for record in records
|
||||
@@ -69,9 +71,7 @@ class HitTestingService:
|
||||
|
||||
documents = {
|
||||
document.id: cls._dump_dataset_document(document)
|
||||
for document in db.session.scalars(
|
||||
select(DatasetDocument).where(DatasetDocument.id.in_(document_ids))
|
||||
).all()
|
||||
for document in session.scalars(select(DatasetDocument).where(DatasetDocument.id.in_(document_ids))).all()
|
||||
}
|
||||
|
||||
records_with_documents: list[dict[str, Any]] = []
|
||||
@@ -105,6 +105,7 @@ class HitTestingService:
|
||||
@classmethod
|
||||
def retrieve(
|
||||
cls,
|
||||
session: Session | scoped_session,
|
||||
dataset: Dataset,
|
||||
query: str,
|
||||
account: Account,
|
||||
@@ -142,7 +143,7 @@ class HitTestingService:
|
||||
if metadata_filter_document_ids:
|
||||
document_ids_filter = metadata_filter_document_ids.get(dataset.id, [])
|
||||
if metadata_condition and not document_ids_filter:
|
||||
return cls.compact_retrieve_response(query, [])
|
||||
return cls.compact_retrieve_response(session, query, [])
|
||||
all_documents = RetrievalService.retrieve(
|
||||
retrieval_method=RetrievalMethod(
|
||||
resolved_retrieval_model.get("search_method", RetrievalMethod.SEMANTIC_SEARCH)
|
||||
@@ -181,14 +182,15 @@ class HitTestingService:
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=account.id,
|
||||
)
|
||||
db.session.add(dataset_query)
|
||||
db.session.commit()
|
||||
session.add(dataset_query)
|
||||
session.commit()
|
||||
|
||||
return cls.compact_retrieve_response(query, all_documents)
|
||||
return cls.compact_retrieve_response(session, query, all_documents)
|
||||
|
||||
@classmethod
|
||||
def external_retrieve(
|
||||
cls,
|
||||
session: Session | scoped_session,
|
||||
dataset: Dataset,
|
||||
query: str,
|
||||
account: Account,
|
||||
@@ -222,20 +224,22 @@ class HitTestingService:
|
||||
created_by=account.id,
|
||||
)
|
||||
|
||||
db.session.add(dataset_query)
|
||||
db.session.commit()
|
||||
session.add(dataset_query)
|
||||
session.commit()
|
||||
|
||||
return dict(cls.compact_external_retrieve_response(dataset, query, all_documents))
|
||||
|
||||
@classmethod
|
||||
def compact_retrieve_response(cls, query: str, documents: list[Document]) -> RetrieveResponseDict:
|
||||
def compact_retrieve_response(
|
||||
cls, session: Session | scoped_session, query: str, documents: list[Document]
|
||||
) -> RetrieveResponseDict:
|
||||
records = RetrievalService.format_retrieval_documents(documents)
|
||||
|
||||
return {
|
||||
"query": {
|
||||
"content": query,
|
||||
},
|
||||
"records": cls._dump_retrieval_records(records),
|
||||
"records": cls._dump_retrieval_records(session, records),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -3,6 +3,8 @@ from typing import TypedDict
|
||||
|
||||
import httpx
|
||||
|
||||
OPERATION_REQUEST_TIMEOUT = httpx.Timeout(10.0, connect=3.0)
|
||||
|
||||
|
||||
class UtmInfo(TypedDict, total=False):
|
||||
"""Expected shape of the utm_info dict passed to record_utm.
|
||||
@@ -26,7 +28,9 @@ class OperationService:
|
||||
headers = {"Content-Type": "application/json", "Billing-Api-Secret-Key": cls.secret_key}
|
||||
|
||||
url = f"{cls.base_url}{endpoint}"
|
||||
response = httpx.request(method, url, json=json, params=params, headers=headers)
|
||||
response = httpx.request(
|
||||
method, url, json=json, params=params, headers=headers, timeout=OPERATION_REQUEST_TIMEOUT
|
||||
)
|
||||
|
||||
return response.json()
|
||||
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
"""
|
||||
Archive Paid Plan Workflow Run Logs Service.
|
||||
|
||||
This service archives workflow run logs for paid plan users older than the configured
|
||||
retention period (default: 90 days) to S3-compatible storage.
|
||||
This service archives workflow run logs for paid plan users older than the configured retention period (default:
|
||||
90 days) to S3-compatible storage.
|
||||
|
||||
Archive V2 writes bundle-level Parquet objects. A bundle contains many workflow runs and their related table rows.
|
||||
Bundle metadata lives in the object-store manifest instead of a database table, so archive/delete/restore does not move
|
||||
the large-table retention problem into another OLTP table.
|
||||
|
||||
Archived tables:
|
||||
- workflow_runs
|
||||
@@ -16,18 +20,19 @@ Archived tables:
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import io
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import zipfile
|
||||
from collections.abc import Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, TypedDict
|
||||
|
||||
import click
|
||||
from sqlalchemy import inspect
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
from sqlalchemy import inspect, select
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from configs import dify_config
|
||||
@@ -39,12 +44,24 @@ from libs.archive_storage import (
|
||||
ArchiveStorageNotConfiguredError,
|
||||
get_archive_storage,
|
||||
)
|
||||
from models.workflow import WorkflowAppLog, WorkflowRun
|
||||
from models.trigger import WorkflowTriggerLog
|
||||
from models.workflow import (
|
||||
WorkflowAppLog,
|
||||
WorkflowNodeExecutionModel,
|
||||
WorkflowNodeExecutionOffload,
|
||||
WorkflowPause,
|
||||
WorkflowPauseReason,
|
||||
WorkflowRun,
|
||||
)
|
||||
from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository
|
||||
from repositories.api_workflow_run_repository import APIWorkflowRunRepository
|
||||
from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository
|
||||
from services.billing_service import BillingService
|
||||
from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_NAME, ARCHIVE_SCHEMA_VERSION
|
||||
from services.retention.workflow_run.constants import (
|
||||
ARCHIVE_BUNDLE_FORMAT,
|
||||
ARCHIVE_BUNDLE_MANIFEST_NAME,
|
||||
ARCHIVE_BUNDLE_SCHEMA_VERSION,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -53,17 +70,41 @@ class TableStatsManifestEntry(TypedDict):
|
||||
row_count: int
|
||||
checksum: str
|
||||
size_bytes: int
|
||||
object_key: str
|
||||
|
||||
|
||||
class ArchiveManifestDict(TypedDict):
|
||||
schema_version: str
|
||||
workflow_run_id: str
|
||||
archive_format: str
|
||||
tenant_id: str
|
||||
app_id: str
|
||||
workflow_id: str
|
||||
created_at: str
|
||||
tenant_prefix: str
|
||||
year: int
|
||||
month: int
|
||||
shard: str
|
||||
bundle_id: str
|
||||
object_prefix: str
|
||||
workflow_run_count: int
|
||||
workflow_node_execution_count: int
|
||||
min_created_at: str
|
||||
max_created_at: str
|
||||
min_run_id: str
|
||||
max_run_id: str
|
||||
archived_at: str
|
||||
tables: dict[str, TableStatsManifestEntry]
|
||||
run_ids: list[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArchiveBundleIdentity:
|
||||
"""Stable identity and object prefix for one V2 archive bundle."""
|
||||
|
||||
tenant_prefix: str
|
||||
tenant_id: str
|
||||
year: int
|
||||
month: int
|
||||
shard: str
|
||||
bundle_id: str
|
||||
object_prefix: str
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -74,16 +115,21 @@ class TableStats:
|
||||
row_count: int
|
||||
checksum: str
|
||||
size_bytes: int
|
||||
object_key: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ArchiveResult:
|
||||
"""Result of archiving a single workflow run."""
|
||||
"""Result of archiving a bundle of workflow runs."""
|
||||
|
||||
run_id: str
|
||||
bundle_id: str
|
||||
tenant_id: str
|
||||
object_prefix: str
|
||||
success: bool
|
||||
run_count: int = 0
|
||||
tables: list[TableStats] = field(default_factory=list)
|
||||
object_size_bytes: int = 0
|
||||
skipped: bool = False
|
||||
error: str | None = None
|
||||
elapsed_time: float = 0.0
|
||||
|
||||
@@ -96,6 +142,12 @@ class ArchiveSummary:
|
||||
runs_archived: int = 0
|
||||
runs_skipped: int = 0
|
||||
runs_failed: int = 0
|
||||
total_bundles_processed: int = 0
|
||||
bundles_archived: int = 0
|
||||
bundles_skipped: int = 0
|
||||
bundles_failed: int = 0
|
||||
total_object_size_bytes: int = 0
|
||||
table_stats: dict[str, TableStats] = field(default_factory=dict)
|
||||
total_elapsed_time: float = 0.0
|
||||
|
||||
|
||||
@@ -104,16 +156,20 @@ class WorkflowRunArchiver:
|
||||
Archive workflow run logs for paid plan users.
|
||||
|
||||
Storage Layout:
|
||||
{tenant_id}/app_id={app_id}/year={YYYY}/month={MM}/workflow_run_id={run_id}/
|
||||
└── archive.v1.0.zip
|
||||
workflow-runs/v2/tenant_prefix={prefix}/tenant_id={tenant_id}/year={YYYY}/month={MM}/
|
||||
shard={shard}/bundle={bundle_id}/
|
||||
├── manifest.json
|
||||
├── workflow_runs.jsonl
|
||||
├── workflow_app_logs.jsonl
|
||||
├── workflow_node_executions.jsonl
|
||||
├── workflow_node_execution_offload.jsonl
|
||||
├── workflow_pauses.jsonl
|
||||
├── workflow_pause_reasons.jsonl
|
||||
└── workflow_trigger_logs.jsonl
|
||||
├── workflow_runs.parquet
|
||||
├── workflow_app_logs.parquet
|
||||
├── workflow_node_executions.parquet
|
||||
├── workflow_node_execution_offload.parquet
|
||||
├── workflow_pauses.parquet
|
||||
├── workflow_pause_reasons.parquet
|
||||
└── workflow_trigger_logs.parquet
|
||||
|
||||
`batch_size` is the maximum workflow_runs per bundle. The current implementation groups each fetched page by
|
||||
tenant/month before writing bundles. Bundle idempotency is based on the manifest object key; the manifest is
|
||||
uploaded after all table objects, so a missing manifest means the bundle should be retried.
|
||||
"""
|
||||
|
||||
ARCHIVED_TYPE = [
|
||||
@@ -132,6 +188,10 @@ class WorkflowRunArchiver:
|
||||
|
||||
start_from: datetime.datetime | None
|
||||
end_before: datetime.datetime
|
||||
paid_tenant_ids: set[str] | None
|
||||
tenant_prefixes: list[str]
|
||||
run_shard_index: int | None
|
||||
run_shard_total: int | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -141,6 +201,10 @@ class WorkflowRunArchiver:
|
||||
end_before: datetime.datetime | None = None,
|
||||
workers: int = 1,
|
||||
tenant_ids: Sequence[str] | None = None,
|
||||
tenant_prefixes: Sequence[str] | None = None,
|
||||
paid_tenant_ids: Sequence[str] | None = None,
|
||||
run_shard_index: int | None = None,
|
||||
run_shard_total: int | None = None,
|
||||
limit: int | None = None,
|
||||
dry_run: bool = False,
|
||||
delete_after_archive: bool = False,
|
||||
@@ -156,10 +220,19 @@ class WorkflowRunArchiver:
|
||||
end_before: Optional end time (exclusive) for archiving
|
||||
workers: Number of concurrent workflow runs to archive
|
||||
tenant_ids: Optional tenant IDs for grayscale rollout
|
||||
tenant_prefixes: Optional tenant ID first-hex prefixes for rollout waves. CLI callers should resolve these
|
||||
to tenant_ids during planning so workflow_runs scan uses tenant_id IN (...) instead of a prefix range.
|
||||
paid_tenant_ids: Optional paid-tenant whitelist resolved by the archive plan. When provided, archive uses it
|
||||
for per-run paid filtering and does not call billing on every fetched page.
|
||||
run_shard_index: Optional zero-based workflow run shard index for parallel cron jobs
|
||||
run_shard_total: Optional total workflow run shard count for parallel cron jobs
|
||||
limit: Maximum number of runs to archive (None for unlimited)
|
||||
dry_run: If True, only preview without making changes
|
||||
delete_after_archive: If True, delete runs and related data after archiving
|
||||
delete_after_archive: Reserved for the V1 per-run path. Bundle archive requires a separate validated
|
||||
bundle delete workflow.
|
||||
"""
|
||||
if delete_after_archive:
|
||||
raise ValueError("delete_after_archive is not supported by bundle archive")
|
||||
self.days = days
|
||||
self.batch_size = batch_size
|
||||
if start_from or end_before:
|
||||
@@ -176,6 +249,16 @@ class WorkflowRunArchiver:
|
||||
raise ValueError("workers must be at least 1")
|
||||
self.workers = workers
|
||||
self.tenant_ids = sorted(set(tenant_ids)) if tenant_ids else []
|
||||
self.tenant_prefixes = sorted(set(tenant_prefixes)) if tenant_prefixes else []
|
||||
self.paid_tenant_ids = set(paid_tenant_ids) if paid_tenant_ids is not None else None
|
||||
if (run_shard_index is None) ^ (run_shard_total is None):
|
||||
raise ValueError("run_shard_index and run_shard_total must be provided together")
|
||||
if run_shard_total is not None and not 1 <= run_shard_total <= 16:
|
||||
raise ValueError("run_shard_total must be between 1 and 16")
|
||||
if run_shard_index is not None and run_shard_total is not None and not 0 <= run_shard_index < run_shard_total:
|
||||
raise ValueError("run_shard_index must be between 0 and run_shard_total - 1")
|
||||
self.run_shard_index = run_shard_index
|
||||
self.run_shard_total = run_shard_total
|
||||
self.limit = limit
|
||||
self.dry_run = dry_run
|
||||
self.delete_after_archive = delete_after_archive
|
||||
@@ -209,124 +292,185 @@ class WorkflowRunArchiver:
|
||||
return summary
|
||||
|
||||
session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
|
||||
repo = self._get_workflow_run_repo()
|
||||
attempted_count = 0
|
||||
|
||||
def _archive_with_session(run: WorkflowRun) -> ArchiveResult:
|
||||
with session_maker() as session:
|
||||
return self._archive_run(session, storage, run)
|
||||
|
||||
last_seen: tuple[datetime.datetime, str] | None = None
|
||||
archived_count = 0
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.workers) as executor:
|
||||
for tenant_scope in self._tenant_scan_scopes():
|
||||
last_seen: tuple[datetime.datetime, str] | None = None
|
||||
while True:
|
||||
# Check limit
|
||||
if self.limit and archived_count >= self.limit:
|
||||
if self.limit and attempted_count >= self.limit:
|
||||
click.echo(click.style(f"Reached limit of {self.limit} runs", fg="yellow"))
|
||||
break
|
||||
|
||||
# Fetch batch of runs
|
||||
runs = self._get_runs_batch(last_seen)
|
||||
|
||||
runs = self._get_runs_batch(last_seen, tenant_scope=tenant_scope)
|
||||
if not runs:
|
||||
break
|
||||
|
||||
run_ids = [run.id for run in runs]
|
||||
with session_maker() as session:
|
||||
archived_run_ids = repo.get_archived_run_ids(session, run_ids)
|
||||
|
||||
last_seen = (runs[-1].created_at, runs[-1].id)
|
||||
|
||||
# Filter to paid tenants only
|
||||
tenant_ids = {run.tenant_id for run in runs}
|
||||
paid_tenants = self._filter_paid_tenants(tenant_ids)
|
||||
|
||||
runs_to_process: list[WorkflowRun] = []
|
||||
for run in runs:
|
||||
summary.total_runs_processed += 1
|
||||
|
||||
# Skip non-paid tenants
|
||||
if run.tenant_id not in paid_tenants:
|
||||
summary.runs_skipped += 1
|
||||
continue
|
||||
|
||||
# Skip already archived runs
|
||||
if run.id in archived_run_ids:
|
||||
summary.runs_skipped += 1
|
||||
continue
|
||||
|
||||
# Check limit
|
||||
if self.limit and archived_count + len(runs_to_process) >= self.limit:
|
||||
if self.limit and attempted_count + len(runs_to_process) >= self.limit:
|
||||
break
|
||||
|
||||
runs_to_process.append(run)
|
||||
|
||||
if not runs_to_process:
|
||||
continue
|
||||
|
||||
results = list(executor.map(_archive_with_session, runs_to_process))
|
||||
for bundle_runs in self._group_runs_for_bundles(runs_to_process):
|
||||
summary.total_bundles_processed += 1
|
||||
with session_maker() as session:
|
||||
result = self._archive_bundle(session, storage, bundle_runs)
|
||||
|
||||
for run, result in zip(runs_to_process, results):
|
||||
if result.success:
|
||||
summary.runs_archived += 1
|
||||
archived_count += 1
|
||||
if result.skipped:
|
||||
attempted_count += result.run_count
|
||||
summary.bundles_skipped += 1
|
||||
summary.runs_skipped += result.run_count
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Skipped bundle {result.bundle_id} (tenant={result.tenant_id}, "
|
||||
f"runs={result.run_count}, reason={result.error or 'already handled'})",
|
||||
fg="yellow",
|
||||
)
|
||||
)
|
||||
elif result.success:
|
||||
attempted_count += result.run_count
|
||||
summary.bundles_archived += 1
|
||||
summary.runs_archived += result.run_count
|
||||
self._merge_result_stats(summary, result)
|
||||
click.echo(
|
||||
click.style(
|
||||
f"{'[DRY RUN] Would archive' if self.dry_run else 'Archived'} "
|
||||
f"run {run.id} (tenant={run.tenant_id}, "
|
||||
f"tables={len(result.tables)}, time={result.elapsed_time:.2f}s)",
|
||||
f"bundle {result.bundle_id} (tenant={result.tenant_id}, runs={result.run_count}, "
|
||||
f"tables={len(result.tables)}, object_size_bytes={result.object_size_bytes}, "
|
||||
f"time={result.elapsed_time:.2f}s)",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
if self.dry_run:
|
||||
self._echo_table_estimates(result.tables)
|
||||
else:
|
||||
summary.runs_failed += 1
|
||||
attempted_count += result.run_count
|
||||
summary.bundles_failed += 1
|
||||
summary.runs_failed += result.run_count
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Failed to archive run {run.id}: {result.error}",
|
||||
f"Failed to archive bundle {result.bundle_id}: {result.error}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
|
||||
if self.limit and attempted_count >= self.limit:
|
||||
break
|
||||
|
||||
summary.total_elapsed_time = time.time() - start_time
|
||||
click.echo(
|
||||
click.style(
|
||||
f"{'[DRY RUN] ' if self.dry_run else ''}Archive complete: "
|
||||
f"processed={summary.total_runs_processed}, archived={summary.runs_archived}, "
|
||||
f"skipped={summary.runs_skipped}, failed={summary.runs_failed}, "
|
||||
f"bundles_archived={summary.bundles_archived}, bundles_skipped={summary.bundles_skipped}, "
|
||||
f"bundles_failed={summary.bundles_failed}, "
|
||||
f"object_size_bytes={summary.total_object_size_bytes}, "
|
||||
f"time={summary.total_elapsed_time:.2f}s",
|
||||
fg="white",
|
||||
)
|
||||
)
|
||||
if self.dry_run:
|
||||
self._echo_summary_estimates(summary)
|
||||
|
||||
return summary
|
||||
|
||||
@staticmethod
|
||||
def _merge_result_stats(summary: ArchiveSummary, result: ArchiveResult) -> None:
|
||||
summary.total_object_size_bytes += result.object_size_bytes
|
||||
for table_stat in result.tables:
|
||||
summary_stat = summary.table_stats.get(table_stat.table_name)
|
||||
if summary_stat is None:
|
||||
summary.table_stats[table_stat.table_name] = TableStats(
|
||||
table_name=table_stat.table_name,
|
||||
row_count=table_stat.row_count,
|
||||
checksum="",
|
||||
size_bytes=table_stat.size_bytes,
|
||||
)
|
||||
continue
|
||||
summary_stat.row_count += table_stat.row_count
|
||||
summary_stat.size_bytes += table_stat.size_bytes
|
||||
|
||||
@staticmethod
|
||||
def _echo_table_estimates(table_stats: Sequence[TableStats]) -> None:
|
||||
for stat in table_stats:
|
||||
click.echo(
|
||||
click.style(
|
||||
f" table={stat.table_name} rows={stat.row_count} parquet_bytes={stat.size_bytes}",
|
||||
fg="white",
|
||||
)
|
||||
)
|
||||
|
||||
def _echo_summary_estimates(self, summary: ArchiveSummary) -> None:
|
||||
click.echo(click.style("[DRY RUN] Estimated archive totals by table:", fg="white"))
|
||||
for table_name in self.ARCHIVED_TABLES:
|
||||
stat = summary.table_stats.get(table_name)
|
||||
row_count = stat.row_count if stat else 0
|
||||
size_bytes = stat.size_bytes if stat else 0
|
||||
click.echo(click.style(f" table={table_name} rows={row_count} parquet_bytes={size_bytes}", fg="white"))
|
||||
|
||||
def _get_runs_batch(
|
||||
self,
|
||||
last_seen: tuple[datetime.datetime, str] | None,
|
||||
tenant_scope: Sequence[str] | None = None,
|
||||
) -> Sequence[WorkflowRun]:
|
||||
"""Fetch a batch of workflow runs to archive."""
|
||||
repo = self._get_workflow_run_repo()
|
||||
tenant_ids = list(tenant_scope) if tenant_scope is not None else self.tenant_ids or None
|
||||
return repo.get_runs_batch_by_time_range(
|
||||
start_from=self.start_from,
|
||||
end_before=self.end_before,
|
||||
last_seen=last_seen,
|
||||
batch_size=self.batch_size,
|
||||
run_types=self.ARCHIVED_TYPE,
|
||||
tenant_ids=self.tenant_ids or None,
|
||||
tenant_ids=tenant_ids,
|
||||
tenant_prefixes=None if tenant_ids else self.tenant_prefixes or None,
|
||||
run_shard_index=self.run_shard_index,
|
||||
run_shard_total=self.run_shard_total,
|
||||
)
|
||||
|
||||
def _tenant_scan_scopes(self) -> list[list[str] | None]:
|
||||
if not self.tenant_ids:
|
||||
return [None]
|
||||
return [[tenant_id] for tenant_id in self.tenant_ids]
|
||||
|
||||
def _build_start_message(self) -> str:
|
||||
range_desc = f"before {self.end_before.isoformat()}"
|
||||
if self.start_from:
|
||||
range_desc = f"between {self.start_from.isoformat()} and {self.end_before.isoformat()}"
|
||||
run_shard_desc = "all"
|
||||
if self.run_shard_index is not None and self.run_shard_total is not None:
|
||||
run_shard_desc = f"{self.run_shard_index}/{self.run_shard_total}"
|
||||
return (
|
||||
f"{'[DRY RUN] ' if self.dry_run else ''}Starting workflow run archiving "
|
||||
f"for runs {range_desc} "
|
||||
f"(batch_size={self.batch_size}, tenant_ids={','.join(self.tenant_ids) or 'all'})"
|
||||
f"(batch_size={self.batch_size}, tenant_ids={self._format_tenant_scope()}, "
|
||||
f"tenant_prefixes={','.join(self.tenant_prefixes) or 'all'}, run_shard={run_shard_desc})"
|
||||
)
|
||||
|
||||
def _format_tenant_scope(self) -> str:
|
||||
if not self.tenant_ids:
|
||||
return "all"
|
||||
if len(self.tenant_ids) <= 10:
|
||||
return ",".join(self.tenant_ids)
|
||||
return f"{len(self.tenant_ids)} planned tenants"
|
||||
|
||||
def _filter_paid_tenants(self, tenant_ids: set[str]) -> set[str]:
|
||||
"""Filter tenant IDs to only include paid tenants."""
|
||||
if self.paid_tenant_ids is not None:
|
||||
return tenant_ids & self.paid_tenant_ids
|
||||
|
||||
if not dify_config.BILLING_ENABLED:
|
||||
# If billing is not enabled, treat all tenants as paid
|
||||
return tenant_ids
|
||||
@@ -349,177 +493,293 @@ class WorkflowRunArchiver:
|
||||
|
||||
return paid
|
||||
|
||||
def _archive_run(
|
||||
def _archive_bundle(
|
||||
self,
|
||||
session: Session,
|
||||
storage: ArchiveStorage | None,
|
||||
run: WorkflowRun,
|
||||
runs: Sequence[WorkflowRun],
|
||||
) -> ArchiveResult:
|
||||
"""Archive a single workflow run."""
|
||||
"""Archive one tenant/month bundle of workflow runs."""
|
||||
if not runs:
|
||||
raise ValueError("runs must not be empty")
|
||||
start_time = time.time()
|
||||
result = ArchiveResult(run_id=run.id, tenant_id=run.tenant_id, success=False)
|
||||
identity = self._build_bundle_identity(runs)
|
||||
result = ArchiveResult(
|
||||
bundle_id=identity.bundle_id,
|
||||
tenant_id=identity.tenant_id,
|
||||
object_prefix=identity.object_prefix,
|
||||
run_count=len(runs),
|
||||
success=False,
|
||||
)
|
||||
|
||||
try:
|
||||
# Extract data from all tables
|
||||
table_data, app_logs, trigger_metadata = self._extract_data(session, run)
|
||||
if not self.dry_run:
|
||||
if storage is None:
|
||||
raise ArchiveStorageNotConfiguredError("Archive storage not configured")
|
||||
if storage.object_exists(self._get_manifest_object_key(identity)):
|
||||
result.success = True
|
||||
result.skipped = True
|
||||
result.error = "bundle already archived"
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
locked_runs = self._lock_runs_for_archive(session, [run.id for run in runs])
|
||||
if len(locked_runs) != len(runs):
|
||||
result.success = True
|
||||
result.skipped = True
|
||||
result.error = "one or more runs locked or deleted by another archiver"
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
runs = locked_runs
|
||||
|
||||
table_data = self._extract_bundle_data(session, runs)
|
||||
table_stats, table_payloads, manifest_data = self._build_archive_payload(identity, runs, table_data)
|
||||
object_size = len(manifest_data) + sum(len(payload) for payload in table_payloads.values())
|
||||
|
||||
if self.dry_run:
|
||||
# In dry run, just report what would be archived
|
||||
for table_name in self.ARCHIVED_TABLES:
|
||||
records = table_data.get(table_name, [])
|
||||
result.tables.append(
|
||||
TableStats(
|
||||
table_name=table_name,
|
||||
row_count=len(records),
|
||||
checksum="",
|
||||
size_bytes=0,
|
||||
)
|
||||
)
|
||||
result.tables = table_stats
|
||||
result.object_size_bytes = object_size
|
||||
result.success = True
|
||||
else:
|
||||
if storage is None:
|
||||
raise ArchiveStorageNotConfiguredError("Archive storage not configured")
|
||||
archive_key = self._get_archive_key(run)
|
||||
|
||||
# Serialize tables for the archive bundle
|
||||
table_stats: list[TableStats] = []
|
||||
table_payloads: dict[str, bytes] = {}
|
||||
for table_name in self.ARCHIVED_TABLES:
|
||||
records = table_data.get(table_name, [])
|
||||
data = ArchiveStorage.serialize_to_jsonl(records)
|
||||
table_payloads[table_name] = data
|
||||
checksum = ArchiveStorage.compute_checksum(data)
|
||||
|
||||
table_stats.append(
|
||||
TableStats(
|
||||
table_name=table_name,
|
||||
row_count=len(records),
|
||||
checksum=checksum,
|
||||
size_bytes=len(data),
|
||||
)
|
||||
)
|
||||
|
||||
# Generate and upload archive bundle
|
||||
manifest = self._generate_manifest(run, table_stats)
|
||||
manifest_data = json.dumps(manifest, indent=2, default=str).encode("utf-8")
|
||||
archive_data = self._build_archive_bundle(manifest_data, table_payloads)
|
||||
storage.put_object(archive_key, archive_data)
|
||||
|
||||
repo = self._get_workflow_run_repo()
|
||||
archived_log_count = repo.create_archive_logs(session, run, app_logs, trigger_metadata)
|
||||
for table_name, payload in table_payloads.items():
|
||||
storage.put_object(self._get_table_object_key(identity, table_name), payload)
|
||||
storage.put_object(self._get_manifest_object_key(identity), manifest_data)
|
||||
session.commit()
|
||||
|
||||
deleted_counts = None
|
||||
if self.delete_after_archive:
|
||||
deleted_counts = repo.delete_runs_with_related(
|
||||
[run],
|
||||
delete_node_executions=self._delete_node_executions,
|
||||
delete_trigger_logs=self._delete_trigger_logs,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Archived workflow run %s: tables=%s, archived_logs=%s, deleted=%s",
|
||||
run.id,
|
||||
"Archived workflow run bundle %s: tenant=%s runs=%s tables=%s object_prefix=%s",
|
||||
identity.bundle_id,
|
||||
identity.tenant_id,
|
||||
len(runs),
|
||||
{s.table_name: s.row_count for s in table_stats},
|
||||
archived_log_count,
|
||||
deleted_counts,
|
||||
identity.object_prefix,
|
||||
)
|
||||
|
||||
result.tables = table_stats
|
||||
result.object_size_bytes = object_size
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Failed to archive workflow run %s", run.id)
|
||||
logger.exception("Failed to archive workflow run bundle %s", identity.bundle_id)
|
||||
result.error = str(e)
|
||||
session.rollback()
|
||||
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
def _extract_data(
|
||||
def _lock_runs_for_archive(
|
||||
self,
|
||||
session: Session,
|
||||
run: WorkflowRun,
|
||||
) -> tuple[dict[str, list[dict[str, Any]]], Sequence[WorkflowAppLog], str | None]:
|
||||
run_ids: Sequence[str],
|
||||
) -> list[WorkflowRun]:
|
||||
"""
|
||||
Lock workflow runs before archiving a bundle.
|
||||
|
||||
Parallel cron jobs may select overlapping pages. Row-level SKIP LOCKED keeps duplicate archivers from uploading
|
||||
conflicting bundle objects for the same source rows.
|
||||
"""
|
||||
if not run_ids:
|
||||
return []
|
||||
stmt = (
|
||||
select(WorkflowRun)
|
||||
.where(WorkflowRun.id.in_(run_ids))
|
||||
.order_by(WorkflowRun.created_at.asc(), WorkflowRun.id.asc())
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
return list(session.scalars(stmt))
|
||||
|
||||
def _extract_bundle_data(
|
||||
self,
|
||||
session: Session,
|
||||
runs: Sequence[WorkflowRun],
|
||||
) -> dict[str, list[dict[str, Any]]]:
|
||||
"""Extract all archived table rows for a bundle."""
|
||||
run_ids = [run.id for run in runs]
|
||||
table_data: dict[str, list[dict[str, Any]]] = {}
|
||||
table_data["workflow_runs"] = [self._row_to_dict(run)]
|
||||
repo = self._get_workflow_run_repo()
|
||||
app_logs = repo.get_app_logs_by_run_id(session, run.id)
|
||||
table_data["workflow_runs"] = [self._row_to_dict(run) for run in runs]
|
||||
|
||||
app_logs = list(session.scalars(select(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids))))
|
||||
table_data["workflow_app_logs"] = [self._row_to_dict(row) for row in app_logs]
|
||||
node_exec_repo = self._get_workflow_node_execution_repo(session)
|
||||
node_exec_records = node_exec_repo.get_executions_by_workflow_run(
|
||||
tenant_id=run.tenant_id,
|
||||
app_id=run.app_id,
|
||||
workflow_run_id=run.id,
|
||||
|
||||
node_exec_records = list(
|
||||
session.scalars(
|
||||
select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.workflow_run_id.in_(run_ids))
|
||||
)
|
||||
)
|
||||
node_exec_ids = [record.id for record in node_exec_records]
|
||||
offload_records = node_exec_repo.get_offloads_by_execution_ids(session, node_exec_ids)
|
||||
offload_records = []
|
||||
if node_exec_ids:
|
||||
offload_records = list(
|
||||
session.scalars(
|
||||
select(WorkflowNodeExecutionOffload).where(
|
||||
WorkflowNodeExecutionOffload.node_execution_id.in_(node_exec_ids)
|
||||
)
|
||||
)
|
||||
)
|
||||
table_data["workflow_node_executions"] = [self._row_to_dict(row) for row in node_exec_records]
|
||||
table_data["workflow_node_execution_offload"] = [self._row_to_dict(row) for row in offload_records]
|
||||
repo = self._get_workflow_run_repo()
|
||||
pause_records = repo.get_pause_records_by_run_id(session, run.id)
|
||||
|
||||
pause_records = list(session.scalars(select(WorkflowPause).where(WorkflowPause.workflow_run_id.in_(run_ids))))
|
||||
pause_ids = [pause.id for pause in pause_records]
|
||||
pause_reason_records = repo.get_pause_reason_records_by_run_id(
|
||||
session,
|
||||
pause_ids,
|
||||
)
|
||||
pause_reason_records = []
|
||||
if pause_ids:
|
||||
pause_reason_records = list(
|
||||
session.scalars(select(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids)))
|
||||
)
|
||||
table_data["workflow_pauses"] = [self._row_to_dict(row) for row in pause_records]
|
||||
table_data["workflow_pause_reasons"] = [self._row_to_dict(row) for row in pause_reason_records]
|
||||
|
||||
trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
|
||||
trigger_records = trigger_repo.list_by_run_id(run.id)
|
||||
trigger_records: list[WorkflowTriggerLog] = []
|
||||
for run_id in run_ids:
|
||||
trigger_records.extend(trigger_repo.list_by_run_id(run_id))
|
||||
table_data["workflow_trigger_logs"] = [self._row_to_dict(row) for row in trigger_records]
|
||||
trigger_metadata = trigger_records[0].trigger_metadata if trigger_records else None
|
||||
return table_data, app_logs, trigger_metadata
|
||||
return table_data
|
||||
|
||||
@staticmethod
|
||||
def _row_to_dict(row: Any) -> dict[str, Any]:
|
||||
mapper = inspect(row).mapper
|
||||
return {str(column.name): getattr(row, mapper.get_property_by_column(column).key) for column in mapper.columns}
|
||||
|
||||
def _get_archive_key(self, run: WorkflowRun) -> str:
|
||||
"""Get the storage key for the archive bundle."""
|
||||
created_at = run.created_at
|
||||
prefix = (
|
||||
f"{run.tenant_id}/app_id={run.app_id}/year={created_at.strftime('%Y')}/"
|
||||
f"month={created_at.strftime('%m')}/workflow_run_id={run.id}"
|
||||
)
|
||||
return f"{prefix}/{ARCHIVE_BUNDLE_NAME}"
|
||||
def _build_archive_payload(
|
||||
self,
|
||||
identity: ArchiveBundleIdentity,
|
||||
runs: Sequence[WorkflowRun],
|
||||
table_data: dict[str, list[dict[str, Any]]],
|
||||
) -> tuple[list[TableStats], dict[str, bytes], bytes]:
|
||||
"""Build the archive payload and size stats without writing it to storage."""
|
||||
table_stats: list[TableStats] = []
|
||||
table_payloads: dict[str, bytes] = {}
|
||||
for table_name in self.ARCHIVED_TABLES:
|
||||
records = table_data.get(table_name, [])
|
||||
data = self._serialize_to_parquet(records)
|
||||
table_payloads[table_name] = data
|
||||
checksum = ArchiveStorage.compute_checksum(data)
|
||||
|
||||
table_stats.append(
|
||||
TableStats(
|
||||
table_name=table_name,
|
||||
row_count=len(records),
|
||||
checksum=checksum,
|
||||
size_bytes=len(data),
|
||||
object_key=self._get_table_object_key(identity, table_name),
|
||||
)
|
||||
)
|
||||
|
||||
manifest = self._generate_manifest(identity, runs, table_stats)
|
||||
manifest_data = json.dumps(manifest, indent=2, default=str).encode("utf-8")
|
||||
return table_stats, table_payloads, manifest_data
|
||||
|
||||
def _generate_manifest(
|
||||
self,
|
||||
run: WorkflowRun,
|
||||
identity: ArchiveBundleIdentity,
|
||||
runs: Sequence[WorkflowRun],
|
||||
table_stats: list[TableStats],
|
||||
) -> ArchiveManifestDict:
|
||||
"""Generate a manifest for the archived workflow run."""
|
||||
"""Generate a manifest for the archived workflow run bundle."""
|
||||
tables: dict[str, TableStatsManifestEntry] = {
|
||||
stat.table_name: {
|
||||
"row_count": stat.row_count,
|
||||
"checksum": stat.checksum,
|
||||
"size_bytes": stat.size_bytes,
|
||||
"object_key": stat.object_key,
|
||||
}
|
||||
for stat in table_stats
|
||||
}
|
||||
sorted_runs = sorted(runs, key=lambda run: (run.created_at, run.id))
|
||||
return ArchiveManifestDict(
|
||||
schema_version=ARCHIVE_SCHEMA_VERSION,
|
||||
workflow_run_id=run.id,
|
||||
tenant_id=run.tenant_id,
|
||||
app_id=run.app_id,
|
||||
workflow_id=run.workflow_id,
|
||||
created_at=run.created_at.isoformat(),
|
||||
schema_version=ARCHIVE_BUNDLE_SCHEMA_VERSION,
|
||||
archive_format=ARCHIVE_BUNDLE_FORMAT,
|
||||
tenant_id=identity.tenant_id,
|
||||
tenant_prefix=identity.tenant_prefix,
|
||||
year=identity.year,
|
||||
month=identity.month,
|
||||
shard=identity.shard,
|
||||
bundle_id=identity.bundle_id,
|
||||
object_prefix=identity.object_prefix,
|
||||
workflow_run_count=len(runs),
|
||||
workflow_node_execution_count=tables["workflow_node_executions"]["row_count"],
|
||||
min_created_at=sorted_runs[0].created_at.isoformat(),
|
||||
max_created_at=sorted_runs[-1].created_at.isoformat(),
|
||||
min_run_id=min(run.id for run in runs),
|
||||
max_run_id=max(run.id for run in runs),
|
||||
archived_at=datetime.datetime.now(datetime.UTC).isoformat(),
|
||||
tables=tables,
|
||||
run_ids=[run.id for run in sorted_runs],
|
||||
)
|
||||
|
||||
def _build_archive_bundle(self, manifest_data: bytes, table_payloads: dict[str, bytes]) -> bytes:
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
|
||||
archive.writestr("manifest.json", manifest_data)
|
||||
for table_name in self.ARCHIVED_TABLES:
|
||||
data = table_payloads.get(table_name)
|
||||
if data is None:
|
||||
raise ValueError(f"Missing archive payload for {table_name}")
|
||||
archive.writestr(f"{table_name}.jsonl", data)
|
||||
return buffer.getvalue()
|
||||
@staticmethod
|
||||
def _serialize_to_parquet(records: list[dict[str, Any]]) -> bytes:
|
||||
normalized_records = [WorkflowRunArchiver._normalize_record_for_parquet(record) for record in records]
|
||||
table = pa.Table.from_pylist(normalized_records) if normalized_records else pa.table({})
|
||||
sink = pa.BufferOutputStream()
|
||||
pq.write_table(table, sink, compression="zstd")
|
||||
return sink.getvalue().to_pybytes()
|
||||
|
||||
@staticmethod
|
||||
def _normalize_record_for_parquet(record: dict[str, Any]) -> dict[str, Any]:
|
||||
def normalize(value: Any) -> Any:
|
||||
if isinstance(value, Enum):
|
||||
return value.value
|
||||
if isinstance(value, dict | list):
|
||||
return json.dumps(value, default=str, ensure_ascii=False)
|
||||
return value
|
||||
|
||||
return {key: normalize(value) for key, value in record.items()}
|
||||
|
||||
def _group_runs_for_bundles(self, runs: Sequence[WorkflowRun]) -> list[list[WorkflowRun]]:
|
||||
"""Group a fetched page into tenant/month bundles."""
|
||||
grouped: dict[tuple[str, int, int], list[WorkflowRun]] = {}
|
||||
for run in runs:
|
||||
key = (run.tenant_id, run.created_at.year, run.created_at.month)
|
||||
grouped.setdefault(key, []).append(run)
|
||||
return [sorted(group, key=lambda run: (run.created_at, run.id)) for group in grouped.values()]
|
||||
|
||||
def _build_bundle_identity(self, runs: Sequence[WorkflowRun]) -> ArchiveBundleIdentity:
|
||||
"""Build the object-store identity for a bundle."""
|
||||
sorted_runs = sorted(runs, key=lambda run: (run.created_at, run.id))
|
||||
first_run = sorted_runs[0]
|
||||
tenant_ids = {run.tenant_id for run in sorted_runs}
|
||||
if len(tenant_ids) != 1:
|
||||
raise ValueError("archive bundle cannot span multiple tenants")
|
||||
years_months = {(run.created_at.year, run.created_at.month) for run in sorted_runs}
|
||||
if len(years_months) != 1:
|
||||
raise ValueError("archive bundle cannot span multiple months")
|
||||
|
||||
run_ids_digest = hashlib.sha256(",".join(run.id for run in sorted_runs).encode("utf-8")).hexdigest()
|
||||
tenant_prefix = first_run.tenant_id[0].lower()
|
||||
shard = self._bundle_shard_name()
|
||||
year, month = next(iter(years_months))
|
||||
bundle_id = run_ids_digest[:16]
|
||||
object_prefix = (
|
||||
f"workflow-runs/v2/tenant_prefix={tenant_prefix}/tenant_id={first_run.tenant_id}/"
|
||||
f"year={year:04d}/month={month:02d}/shard={shard}/bundle={bundle_id}"
|
||||
)
|
||||
return ArchiveBundleIdentity(
|
||||
tenant_prefix=tenant_prefix,
|
||||
tenant_id=first_run.tenant_id,
|
||||
year=year,
|
||||
month=month,
|
||||
shard=shard,
|
||||
bundle_id=bundle_id,
|
||||
object_prefix=object_prefix,
|
||||
)
|
||||
|
||||
def _bundle_shard_name(self) -> str:
|
||||
if self.run_shard_index is None or self.run_shard_total is None:
|
||||
return "00-of-01"
|
||||
return f"{self.run_shard_index:02d}-of-{self.run_shard_total:02d}"
|
||||
|
||||
@staticmethod
|
||||
def _get_table_object_key(identity: ArchiveBundleIdentity, table_name: str) -> str:
|
||||
return f"{identity.object_prefix}/{table_name}.parquet"
|
||||
|
||||
@staticmethod
|
||||
def _get_manifest_object_key(identity: ArchiveBundleIdentity) -> str:
|
||||
return f"{identity.object_prefix}/{ARCHIVE_BUNDLE_MANIFEST_NAME}"
|
||||
|
||||
def _delete_trigger_logs(self, session: Session, run_ids: Sequence[str]) -> int:
|
||||
trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
|
||||
|
||||
@@ -0,0 +1,872 @@
|
||||
"""
|
||||
Maintain V2 workflow-run archive bundles.
|
||||
|
||||
Archive V2 keeps bundle metadata in object-store manifests, not in a database table. This module discovers bundles by
|
||||
listing `manifest.json` objects, uses object-store marker files for delete/restore state, and only touches the database
|
||||
for source-table validation, deletion, and restoration.
|
||||
|
||||
Each bundle is processed in its own database transaction. A failed bundle leaves source rows unchanged unless the
|
||||
transaction has already committed; marker handling makes the next run able to reconcile the common committed-but-marker
|
||||
not-updated case.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
import pyarrow.parquet as pq
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import delete, func, inspect, select
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlalchemy.engine import CursorResult
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from extensions.ext_database import db
|
||||
from libs.archive_storage import ArchiveStorage, ArchiveStorageNotConfiguredError, get_archive_storage
|
||||
from models.trigger import WorkflowTriggerLog
|
||||
from models.workflow import (
|
||||
WorkflowAppLog,
|
||||
WorkflowNodeExecutionModel,
|
||||
WorkflowNodeExecutionOffload,
|
||||
WorkflowPause,
|
||||
WorkflowPauseReason,
|
||||
WorkflowRun,
|
||||
)
|
||||
from services.retention.workflow_run.constants import (
|
||||
ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME,
|
||||
ARCHIVE_BUNDLE_DELETED_MARKER_NAME,
|
||||
ARCHIVE_BUNDLE_FORMAT,
|
||||
ARCHIVE_BUNDLE_MANIFEST_NAME,
|
||||
ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME,
|
||||
ARCHIVE_BUNDLE_RESTORED_MARKER_NAME,
|
||||
ARCHIVE_BUNDLE_SCHEMA_VERSION,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ARCHIVE_ROOT_PREFIX = "workflow-runs/v2/"
|
||||
_CHUNK_SIZE = 5_000
|
||||
|
||||
|
||||
class TableManifestEntry(TypedDict):
|
||||
row_count: int
|
||||
checksum: str
|
||||
size_bytes: int
|
||||
object_key: str
|
||||
|
||||
|
||||
class BundleManifest(TypedDict):
|
||||
schema_version: str
|
||||
archive_format: str
|
||||
tenant_id: str
|
||||
tenant_prefix: str
|
||||
year: int
|
||||
month: int
|
||||
shard: str
|
||||
bundle_id: str
|
||||
object_prefix: str
|
||||
workflow_run_count: int
|
||||
workflow_node_execution_count: int
|
||||
min_created_at: str
|
||||
max_created_at: str
|
||||
min_run_id: str
|
||||
max_run_id: str
|
||||
archived_at: str
|
||||
tables: dict[str, TableManifestEntry]
|
||||
run_ids: list[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BundleReference:
|
||||
"""Object-store reference for one V2 archive bundle."""
|
||||
|
||||
object_prefix: str
|
||||
manifest_key: str
|
||||
manifest: BundleManifest
|
||||
|
||||
|
||||
@dataclass
|
||||
class BundleOperationResult:
|
||||
"""Result for one V2 bundle delete or restore operation."""
|
||||
|
||||
bundle_id: str
|
||||
tenant_id: str
|
||||
object_prefix: str
|
||||
success: bool = False
|
||||
table_counts: dict[str, int] = field(default_factory=dict)
|
||||
archive_bytes: int = 0
|
||||
elapsed_time: float = 0.0
|
||||
validation_time: float = 0.0
|
||||
error: str | None = None
|
||||
|
||||
@property
|
||||
def run_count(self) -> int:
|
||||
return self.table_counts.get("workflow_runs", 0)
|
||||
|
||||
@property
|
||||
def row_count(self) -> int:
|
||||
return sum(self.table_counts.values())
|
||||
|
||||
|
||||
@dataclass
|
||||
class BundleOperationSummary:
|
||||
"""Aggregate metrics for a V2 bundle maintenance command."""
|
||||
|
||||
operation: str
|
||||
bundles_processed: int = 0
|
||||
bundles_succeeded: int = 0
|
||||
bundles_failed: int = 0
|
||||
rows_processed: int = 0
|
||||
runs_processed: int = 0
|
||||
archive_bytes: int = 0
|
||||
elapsed_time: float = 0.0
|
||||
validation_time: float = 0.0
|
||||
table_counts: dict[str, int] = field(default_factory=dict)
|
||||
results: list[BundleOperationResult] = field(default_factory=list)
|
||||
|
||||
@property
|
||||
def runs_per_second(self) -> float:
|
||||
if self.elapsed_time <= 0:
|
||||
return 0.0
|
||||
return self.runs_processed / self.elapsed_time
|
||||
|
||||
@property
|
||||
def rows_per_second(self) -> float:
|
||||
if self.elapsed_time <= 0:
|
||||
return 0.0
|
||||
return self.rows_processed / self.elapsed_time
|
||||
|
||||
@property
|
||||
def bytes_per_second(self) -> float:
|
||||
if self.elapsed_time <= 0:
|
||||
return 0.0
|
||||
return self.archive_bytes / self.elapsed_time
|
||||
|
||||
|
||||
TABLE_MODELS: dict[str, Any] = {
|
||||
"workflow_runs": WorkflowRun,
|
||||
"workflow_app_logs": WorkflowAppLog,
|
||||
"workflow_node_executions": WorkflowNodeExecutionModel,
|
||||
"workflow_node_execution_offload": WorkflowNodeExecutionOffload,
|
||||
"workflow_pauses": WorkflowPause,
|
||||
"workflow_pause_reasons": WorkflowPauseReason,
|
||||
"workflow_trigger_logs": WorkflowTriggerLog,
|
||||
}
|
||||
|
||||
ARCHIVED_TABLES = [
|
||||
"workflow_runs",
|
||||
"workflow_app_logs",
|
||||
"workflow_node_executions",
|
||||
"workflow_node_execution_offload",
|
||||
"workflow_pauses",
|
||||
"workflow_pause_reasons",
|
||||
"workflow_trigger_logs",
|
||||
]
|
||||
|
||||
RESTORE_ORDER = [
|
||||
"workflow_runs",
|
||||
"workflow_app_logs",
|
||||
"workflow_node_executions",
|
||||
"workflow_node_execution_offload",
|
||||
"workflow_pauses",
|
||||
"workflow_pause_reasons",
|
||||
"workflow_trigger_logs",
|
||||
]
|
||||
|
||||
|
||||
class WorkflowRunBundleArchiveMaintenance:
|
||||
"""
|
||||
Delete and restore V2 workflow-run archive bundles.
|
||||
|
||||
Args:
|
||||
dry_run: Validate and report counts without changing source rows or object-store markers.
|
||||
strict_content_validation: Compare source-table content checksums against Parquet content before destructive
|
||||
delete and after restore. Keep enabled for real maintenance.
|
||||
stop_on_error: Stop batch processing after the first failed bundle.
|
||||
"""
|
||||
|
||||
dry_run: bool
|
||||
strict_content_validation: bool
|
||||
stop_on_error: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dry_run: bool = False,
|
||||
strict_content_validation: bool = True,
|
||||
stop_on_error: bool = True,
|
||||
) -> None:
|
||||
self.dry_run = dry_run
|
||||
self.strict_content_validation = strict_content_validation
|
||||
self.stop_on_error = stop_on_error
|
||||
|
||||
def delete_batch(
|
||||
self,
|
||||
*,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
limit: int,
|
||||
) -> BundleOperationSummary:
|
||||
"""Validate and delete source rows for archived V2 bundles in the requested created_at window."""
|
||||
return self._process_batch(
|
||||
operation="delete",
|
||||
tenant_ids=tenant_ids,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
def restore_batch(
|
||||
self,
|
||||
*,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
limit: int,
|
||||
) -> BundleOperationSummary:
|
||||
"""Restore source rows for deleted V2 bundles in the requested created_at window."""
|
||||
return self._process_batch(
|
||||
operation="restore",
|
||||
tenant_ids=tenant_ids,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
def _process_batch(
|
||||
self,
|
||||
*,
|
||||
operation: str,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
limit: int,
|
||||
) -> BundleOperationSummary:
|
||||
start_time = time.time()
|
||||
summary = BundleOperationSummary(operation=operation)
|
||||
if tenant_ids is not None and not tenant_ids:
|
||||
return summary
|
||||
|
||||
storage = self._get_archive_storage()
|
||||
bundle_refs = self._list_bundle_refs(
|
||||
storage,
|
||||
operation=operation,
|
||||
tenant_ids=tenant_ids,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
logger.info("Found %s V2 archive bundles for %s", len(bundle_refs), operation)
|
||||
session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
|
||||
for bundle_ref in bundle_refs:
|
||||
with session_maker() as session:
|
||||
if operation == "delete":
|
||||
result = self._delete_bundle(session, storage, bundle_ref)
|
||||
elif operation == "restore":
|
||||
result = self._restore_bundle(session, storage, bundle_ref)
|
||||
else:
|
||||
raise ValueError(f"Unsupported operation: {operation}")
|
||||
|
||||
self._merge_result(summary, result)
|
||||
if not result.success and self.stop_on_error:
|
||||
logger.error("Stopping V2 bundle %s after failure: %s", operation, result.error)
|
||||
break
|
||||
|
||||
summary.elapsed_time = time.time() - start_time
|
||||
return summary
|
||||
|
||||
def _list_bundle_refs(
|
||||
self,
|
||||
storage: ArchiveStorage,
|
||||
*,
|
||||
operation: str,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
limit: int,
|
||||
) -> list[BundleReference]:
|
||||
start_date = self._to_naive_utc(start_date)
|
||||
end_date = self._to_naive_utc(end_date)
|
||||
manifest_keys = self._list_manifest_keys(storage, tenant_ids)
|
||||
refs: list[BundleReference] = []
|
||||
for manifest_key in manifest_keys:
|
||||
manifest_data = self._get_checked_object(storage, manifest_key)
|
||||
object_prefix = manifest_key.removesuffix(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}")
|
||||
manifest = self._load_and_validate_manifest(manifest_data, object_prefix=object_prefix)
|
||||
min_created_at = self._parse_manifest_datetime(manifest["min_created_at"])
|
||||
max_created_at = self._parse_manifest_datetime(manifest["max_created_at"])
|
||||
if max_created_at < start_date or min_created_at >= end_date:
|
||||
continue
|
||||
if tenant_ids and manifest["tenant_id"] not in tenant_ids:
|
||||
continue
|
||||
if operation == "delete" and self._is_deleted(storage, object_prefix):
|
||||
continue
|
||||
if operation == "restore" and not self._is_deleted(storage, object_prefix):
|
||||
continue
|
||||
refs.append(BundleReference(object_prefix=object_prefix, manifest_key=manifest_key, manifest=manifest))
|
||||
|
||||
refs.sort(
|
||||
key=lambda ref: (
|
||||
self._parse_manifest_datetime(ref.manifest["min_created_at"]),
|
||||
ref.manifest["tenant_id"],
|
||||
ref.manifest["bundle_id"],
|
||||
)
|
||||
)
|
||||
return refs[:limit]
|
||||
|
||||
@staticmethod
|
||||
def _list_manifest_keys(storage: ArchiveStorage, tenant_ids: Sequence[str] | None) -> list[str]:
|
||||
keys: list[str] = []
|
||||
if tenant_ids:
|
||||
prefixes = [
|
||||
f"{_ARCHIVE_ROOT_PREFIX}tenant_prefix={tenant_id[0].lower()}/tenant_id={tenant_id}/"
|
||||
for tenant_id in tenant_ids
|
||||
]
|
||||
else:
|
||||
prefixes = [_ARCHIVE_ROOT_PREFIX]
|
||||
for prefix in prefixes:
|
||||
keys.extend(storage.list_objects(prefix))
|
||||
return sorted(key for key in keys if key.endswith(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}"))
|
||||
|
||||
def _delete_bundle(
|
||||
self,
|
||||
session: Session,
|
||||
storage: ArchiveStorage,
|
||||
bundle_ref: BundleReference,
|
||||
) -> BundleOperationResult:
|
||||
start_time = time.time()
|
||||
result = self._new_result(bundle_ref.manifest)
|
||||
try:
|
||||
validation_start = time.time()
|
||||
manifest, table_records, archive_bytes = self._validate_archive_object(storage, bundle_ref)
|
||||
result.table_counts = self._manifest_table_counts(manifest)
|
||||
result.archive_bytes = archive_bytes
|
||||
|
||||
self._lock_workflow_runs(session, manifest["run_ids"])
|
||||
if self._is_delete_started(storage, bundle_ref.object_prefix) and self._live_counts_match(
|
||||
session, manifest, expected_present=False
|
||||
):
|
||||
result.validation_time = time.time() - validation_start
|
||||
if not self.dry_run:
|
||||
self._mark_deleted(storage, bundle_ref.object_prefix)
|
||||
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
|
||||
result.success = True
|
||||
return result
|
||||
|
||||
self._validate_live_counts(session, manifest, expected_present=True)
|
||||
if self.strict_content_validation:
|
||||
self._validate_live_content(session, table_records)
|
||||
result.validation_time = time.time() - validation_start
|
||||
|
||||
if not self.dry_run:
|
||||
self._put_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
|
||||
deleted_counts = self._delete_bundle_rows(session, table_records)
|
||||
if deleted_counts != result.table_counts:
|
||||
raise ValueError(
|
||||
f"Deleted row count mismatch: expected={result.table_counts}, actual={deleted_counts}"
|
||||
)
|
||||
self._validate_live_counts(session, manifest, expected_present=False)
|
||||
session.commit()
|
||||
self._mark_deleted(storage, bundle_ref.object_prefix)
|
||||
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
|
||||
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
|
||||
result.success = True
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
result.error = str(e)
|
||||
logger.exception("Failed to delete V2 archive bundle %s", bundle_ref.object_prefix)
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
def _restore_bundle(
|
||||
self,
|
||||
session: Session,
|
||||
storage: ArchiveStorage,
|
||||
bundle_ref: BundleReference,
|
||||
) -> BundleOperationResult:
|
||||
start_time = time.time()
|
||||
result = self._new_result(bundle_ref.manifest)
|
||||
try:
|
||||
validation_start = time.time()
|
||||
manifest, table_records, archive_bytes = self._validate_archive_object(storage, bundle_ref)
|
||||
result.table_counts = self._manifest_table_counts(manifest)
|
||||
result.archive_bytes = archive_bytes
|
||||
|
||||
if self._live_counts_match(session, manifest, expected_present=True):
|
||||
if self.strict_content_validation:
|
||||
self._validate_live_content(session, table_records)
|
||||
result.validation_time = time.time() - validation_start
|
||||
if not self.dry_run:
|
||||
self._mark_restored(storage, bundle_ref.object_prefix)
|
||||
result.success = True
|
||||
return result
|
||||
|
||||
self._validate_live_counts(session, manifest, expected_present=False)
|
||||
result.validation_time = time.time() - validation_start
|
||||
|
||||
if not self.dry_run:
|
||||
self._put_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME)
|
||||
restored_counts = self._restore_bundle_rows(session, table_records)
|
||||
if restored_counts != result.table_counts:
|
||||
self._validate_live_counts(session, manifest, expected_present=True)
|
||||
self._validate_live_counts(session, manifest, expected_present=True)
|
||||
if self.strict_content_validation:
|
||||
self._validate_live_content(session, table_records)
|
||||
session.commit()
|
||||
self._mark_restored(storage, bundle_ref.object_prefix)
|
||||
result.success = True
|
||||
except Exception as e:
|
||||
session.rollback()
|
||||
result.error = str(e)
|
||||
logger.exception("Failed to restore V2 archive bundle %s", bundle_ref.object_prefix)
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _new_result(manifest: BundleManifest) -> BundleOperationResult:
|
||||
return BundleOperationResult(
|
||||
bundle_id=manifest["bundle_id"],
|
||||
tenant_id=manifest["tenant_id"],
|
||||
object_prefix=manifest["object_prefix"],
|
||||
)
|
||||
|
||||
def _validate_archive_object(
|
||||
self,
|
||||
storage: ArchiveStorage,
|
||||
bundle_ref: BundleReference,
|
||||
) -> tuple[BundleManifest, dict[str, list[dict[str, Any]]], int]:
|
||||
manifest = bundle_ref.manifest
|
||||
table_records: dict[str, list[dict[str, Any]]] = {}
|
||||
total_size = len(storage.get_object(bundle_ref.manifest_key))
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
info = manifest["tables"][table_name]
|
||||
payload = self._get_checked_object(storage, info["object_key"])
|
||||
total_size += len(payload)
|
||||
if len(payload) != info["size_bytes"]:
|
||||
raise ValueError(
|
||||
f"Archive object size mismatch for {info['object_key']}: "
|
||||
f"expected={info['size_bytes']}, actual={len(payload)}"
|
||||
)
|
||||
checksum = ArchiveStorage.compute_checksum(payload)
|
||||
if checksum != info["checksum"]:
|
||||
raise ValueError(
|
||||
f"Archive object checksum mismatch for {info['object_key']}: "
|
||||
f"expected={info['checksum']}, actual={checksum}"
|
||||
)
|
||||
records = self._deserialize_parquet(payload)
|
||||
if len(records) != info["row_count"]:
|
||||
raise ValueError(
|
||||
f"Parquet row count mismatch for {info['object_key']}: "
|
||||
f"expected={info['row_count']}, actual={len(records)}"
|
||||
)
|
||||
table_records[table_name] = records
|
||||
return manifest, table_records, total_size
|
||||
|
||||
@staticmethod
|
||||
def _get_checked_object(storage: ArchiveStorage, object_key: str) -> bytes:
|
||||
if not storage.object_exists(object_key):
|
||||
raise FileNotFoundError(f"Archive object not found: {object_key}")
|
||||
return storage.get_object(object_key)
|
||||
|
||||
@staticmethod
|
||||
def _load_and_validate_manifest(
|
||||
manifest_data: bytes,
|
||||
*,
|
||||
object_prefix: str,
|
||||
) -> BundleManifest:
|
||||
loaded = json.loads(manifest_data)
|
||||
if not isinstance(loaded, dict):
|
||||
raise ValueError("manifest.json must be an object")
|
||||
required_fields = {
|
||||
"schema_version",
|
||||
"archive_format",
|
||||
"tenant_id",
|
||||
"tenant_prefix",
|
||||
"year",
|
||||
"month",
|
||||
"shard",
|
||||
"bundle_id",
|
||||
"object_prefix",
|
||||
"workflow_run_count",
|
||||
"workflow_node_execution_count",
|
||||
"tables",
|
||||
"run_ids",
|
||||
}
|
||||
missing_fields = sorted(required_fields - set(loaded))
|
||||
if missing_fields:
|
||||
raise ValueError(f"manifest missing required fields: {', '.join(missing_fields)}")
|
||||
manifest = cast(BundleManifest, loaded)
|
||||
if manifest["schema_version"] != ARCHIVE_BUNDLE_SCHEMA_VERSION:
|
||||
raise ValueError(f"unsupported bundle schema_version: {manifest['schema_version']}")
|
||||
if manifest["archive_format"] != ARCHIVE_BUNDLE_FORMAT:
|
||||
raise ValueError(f"unsupported bundle archive_format: {manifest['archive_format']}")
|
||||
if manifest["object_prefix"] != object_prefix:
|
||||
raise ValueError("manifest object_prefix does not match object key")
|
||||
if manifest["tenant_id"][0].lower() != manifest["tenant_prefix"]:
|
||||
raise ValueError("manifest tenant_prefix does not match tenant_id")
|
||||
if len(manifest["run_ids"]) != manifest["workflow_run_count"]:
|
||||
raise ValueError("manifest run_ids count does not match workflow_run_count")
|
||||
|
||||
tables = manifest["tables"]
|
||||
if not isinstance(tables, dict):
|
||||
raise ValueError("manifest tables must be an object")
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
if table_name not in tables:
|
||||
raise ValueError(f"manifest missing table: {table_name}")
|
||||
info = tables[table_name]
|
||||
for key in ("row_count", "checksum", "size_bytes", "object_key"):
|
||||
if key not in info:
|
||||
raise ValueError(f"manifest table {table_name} missing {key}")
|
||||
expected_key = f"{object_prefix}/{table_name}.parquet"
|
||||
if info["object_key"] != expected_key:
|
||||
raise ValueError(
|
||||
f"manifest object_key mismatch for {table_name}: "
|
||||
f"expected={expected_key}, actual={info['object_key']}"
|
||||
)
|
||||
return manifest
|
||||
|
||||
@staticmethod
|
||||
def _deserialize_parquet(payload: bytes) -> list[dict[str, Any]]:
|
||||
table = pq.read_table(io.BytesIO(payload))
|
||||
return table.to_pylist()
|
||||
|
||||
def _validate_live_counts(
|
||||
self,
|
||||
session: Session,
|
||||
manifest: BundleManifest,
|
||||
*,
|
||||
expected_present: bool,
|
||||
) -> None:
|
||||
expected_counts = self._manifest_table_counts(manifest)
|
||||
actual_counts = self._count_live_rows(session, manifest["run_ids"])
|
||||
if not expected_present:
|
||||
expected_counts = dict.fromkeys(expected_counts, 0)
|
||||
if actual_counts != expected_counts:
|
||||
state = "present" if expected_present else "deleted"
|
||||
raise ValueError(
|
||||
f"Live row count mismatch for {state} bundle: expected={expected_counts}, actual={actual_counts}"
|
||||
)
|
||||
|
||||
def _live_counts_match(self, session: Session, manifest: BundleManifest, *, expected_present: bool) -> bool:
|
||||
expected_counts = self._manifest_table_counts(manifest)
|
||||
if not expected_present:
|
||||
expected_counts = dict.fromkeys(expected_counts, 0)
|
||||
return self._count_live_rows(session, manifest["run_ids"]) == expected_counts
|
||||
|
||||
@staticmethod
|
||||
def _manifest_table_counts(manifest: BundleManifest) -> dict[str, int]:
|
||||
return {table_name: manifest["tables"][table_name]["row_count"] for table_name in ARCHIVED_TABLES}
|
||||
|
||||
def _count_live_rows(self, session: Session, run_ids: Sequence[str]) -> dict[str, int]:
|
||||
node_ids = self._select_ids_by_run_ids(session, WorkflowNodeExecutionModel, run_ids)
|
||||
pause_ids = self._select_ids_by_run_ids(session, WorkflowPause, run_ids)
|
||||
return {
|
||||
"workflow_runs": self._count_by_run_ids(session, WorkflowRun, run_ids),
|
||||
"workflow_app_logs": self._count_by_run_ids(session, WorkflowAppLog, run_ids),
|
||||
"workflow_node_executions": len(node_ids),
|
||||
"workflow_node_execution_offload": self._count_by_column(
|
||||
session, WorkflowNodeExecutionOffload, WorkflowNodeExecutionOffload.node_execution_id, node_ids
|
||||
),
|
||||
"workflow_pauses": len(pause_ids),
|
||||
"workflow_pause_reasons": self._count_by_column(
|
||||
session, WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids
|
||||
),
|
||||
"workflow_trigger_logs": self._count_by_run_ids(session, WorkflowTriggerLog, run_ids),
|
||||
}
|
||||
|
||||
def _validate_live_content(
|
||||
self,
|
||||
session: Session,
|
||||
table_records: dict[str, list[dict[str, Any]]],
|
||||
) -> None:
|
||||
run_ids = [str(record["id"]) for record in table_records["workflow_runs"]]
|
||||
node_ids = [str(record["id"]) for record in table_records["workflow_node_executions"]]
|
||||
pause_ids = [str(record["id"]) for record in table_records["workflow_pauses"]]
|
||||
|
||||
live_records = {
|
||||
"workflow_runs": self._load_records_by_run_ids(session, WorkflowRun, run_ids),
|
||||
"workflow_app_logs": self._load_records_by_run_ids(session, WorkflowAppLog, run_ids),
|
||||
"workflow_node_executions": self._load_records_by_run_ids(session, WorkflowNodeExecutionModel, run_ids),
|
||||
"workflow_node_execution_offload": self._load_records_by_column(
|
||||
session, WorkflowNodeExecutionOffload, WorkflowNodeExecutionOffload.node_execution_id, node_ids
|
||||
),
|
||||
"workflow_pauses": self._load_records_by_run_ids(session, WorkflowPause, run_ids),
|
||||
"workflow_pause_reasons": self._load_records_by_column(
|
||||
session, WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids
|
||||
),
|
||||
"workflow_trigger_logs": self._load_records_by_run_ids(session, WorkflowTriggerLog, run_ids),
|
||||
}
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
live_checksum = self._records_checksum(live_records[table_name])
|
||||
archive_checksum = self._records_checksum(table_records[table_name])
|
||||
if live_checksum != archive_checksum:
|
||||
raise ValueError(
|
||||
f"Live/archive content checksum mismatch for {table_name}: "
|
||||
f"expected={archive_checksum}, actual={live_checksum}"
|
||||
)
|
||||
|
||||
def _delete_bundle_rows(
|
||||
self,
|
||||
session: Session,
|
||||
table_records: dict[str, list[dict[str, Any]]],
|
||||
) -> dict[str, int]:
|
||||
run_ids = [str(record["id"]) for record in table_records["workflow_runs"]]
|
||||
node_ids = [str(record["id"]) for record in table_records["workflow_node_executions"]]
|
||||
pause_ids = [str(record["id"]) for record in table_records["workflow_pauses"]]
|
||||
|
||||
deleted_counts = dict.fromkeys(ARCHIVED_TABLES, 0)
|
||||
deleted_counts["workflow_pause_reasons"] = self._delete_by_column(
|
||||
session, WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids
|
||||
)
|
||||
deleted_counts["workflow_node_execution_offload"] = self._delete_by_column(
|
||||
session, WorkflowNodeExecutionOffload, WorkflowNodeExecutionOffload.node_execution_id, node_ids
|
||||
)
|
||||
deleted_counts["workflow_trigger_logs"] = self._delete_by_run_ids(session, WorkflowTriggerLog, run_ids)
|
||||
deleted_counts["workflow_app_logs"] = self._delete_by_run_ids(session, WorkflowAppLog, run_ids)
|
||||
deleted_counts["workflow_node_executions"] = self._delete_by_run_ids(
|
||||
session, WorkflowNodeExecutionModel, run_ids
|
||||
)
|
||||
deleted_counts["workflow_pauses"] = self._delete_by_run_ids(session, WorkflowPause, run_ids)
|
||||
deleted_counts["workflow_runs"] = self._delete_by_run_ids(session, WorkflowRun, run_ids)
|
||||
return deleted_counts
|
||||
|
||||
def _restore_bundle_rows(
|
||||
self,
|
||||
session: Session,
|
||||
table_records: dict[str, list[dict[str, Any]]],
|
||||
) -> dict[str, int]:
|
||||
restored_counts = dict.fromkeys(ARCHIVED_TABLES, 0)
|
||||
for table_name in RESTORE_ORDER:
|
||||
restored_counts[table_name] = self._restore_table_records(session, table_name, table_records[table_name])
|
||||
return restored_counts
|
||||
|
||||
def _restore_table_records(
|
||||
self,
|
||||
session: Session,
|
||||
table_name: str,
|
||||
records: list[dict[str, Any]],
|
||||
) -> int:
|
||||
if not records:
|
||||
return 0
|
||||
model = TABLE_MODELS[table_name]
|
||||
total = 0
|
||||
for chunk in self._chunks(records, _CHUNK_SIZE):
|
||||
converted = [self._prepare_insert_record(model, record) for record in chunk]
|
||||
stmt = pg_insert(cast(Any, model.__table__)).values(converted)
|
||||
stmt = stmt.on_conflict_do_nothing(index_elements=["id"])
|
||||
result = session.execute(stmt)
|
||||
total += cast(CursorResult, result).rowcount or 0
|
||||
return total
|
||||
|
||||
def _prepare_insert_record(
|
||||
self,
|
||||
model: Any,
|
||||
record: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
table = model.__table__
|
||||
columns_by_name = {column.name: column for column in table.columns}
|
||||
prepared = {key: value for key, value in record.items() if key in columns_by_name}
|
||||
for column_name, value in list(prepared.items()):
|
||||
column = columns_by_name[column_name]
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(column.type, sa.DateTime) and isinstance(value, str):
|
||||
prepared[column_name] = datetime.datetime.fromisoformat(value)
|
||||
elif isinstance(column.type, sa.JSON) and isinstance(value, str):
|
||||
prepared[column_name] = json.loads(value)
|
||||
return prepared
|
||||
|
||||
@staticmethod
|
||||
def _row_to_dict(row: Any) -> dict[str, Any]:
|
||||
mapper = inspect(row).mapper
|
||||
return {str(column.name): getattr(row, mapper.get_property_by_column(column).key) for column in mapper.columns}
|
||||
|
||||
@staticmethod
|
||||
def _normalize_record_for_checksum(record: dict[str, Any]) -> dict[str, Any]:
|
||||
def normalize(value: Any) -> Any:
|
||||
if isinstance(value, Enum):
|
||||
return value.value
|
||||
if isinstance(value, dict | list):
|
||||
return json.dumps(value, default=str, ensure_ascii=False)
|
||||
return value
|
||||
|
||||
return {key: normalize(value) for key, value in record.items()}
|
||||
|
||||
@classmethod
|
||||
def _records_checksum(cls, records: list[dict[str, Any]]) -> str:
|
||||
normalized = [cls._normalize_record_for_checksum(record) for record in records]
|
||||
normalized.sort(key=lambda record: json.dumps(record, sort_keys=True, default=str, ensure_ascii=False))
|
||||
payload = json.dumps(normalized, sort_keys=True, default=str, ensure_ascii=False, separators=(",", ":"))
|
||||
return ArchiveStorage.compute_checksum(payload.encode("utf-8"))
|
||||
|
||||
@staticmethod
|
||||
def _lock_workflow_runs(session: Session, run_ids: Sequence[str]) -> None:
|
||||
for chunk in WorkflowRunBundleArchiveMaintenance._chunks(run_ids, _CHUNK_SIZE):
|
||||
list(session.scalars(select(WorkflowRun.id).where(WorkflowRun.id.in_(chunk)).with_for_update()))
|
||||
|
||||
@staticmethod
|
||||
def _select_ids_by_run_ids(
|
||||
session: Session,
|
||||
model: Any,
|
||||
run_ids: Sequence[str],
|
||||
) -> list[str]:
|
||||
if not run_ids:
|
||||
return []
|
||||
ids: list[str] = []
|
||||
for chunk in WorkflowRunBundleArchiveMaintenance._chunks(run_ids, _CHUNK_SIZE):
|
||||
ids.extend(
|
||||
str(row_id) for row_id in session.scalars(select(model.id).where(model.workflow_run_id.in_(chunk)))
|
||||
)
|
||||
return ids
|
||||
|
||||
@staticmethod
|
||||
def _count_by_run_ids(
|
||||
session: Session,
|
||||
model: Any,
|
||||
run_ids: Sequence[str],
|
||||
) -> int:
|
||||
return WorkflowRunBundleArchiveMaintenance._count_by_column(
|
||||
session, model, WorkflowRunBundleArchiveMaintenance._run_id_column(model), run_ids
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _count_by_column(
|
||||
session: Session,
|
||||
model: Any,
|
||||
column: Any,
|
||||
values: Sequence[str],
|
||||
) -> int:
|
||||
if not values:
|
||||
return 0
|
||||
total = 0
|
||||
for chunk in WorkflowRunBundleArchiveMaintenance._chunks(values, _CHUNK_SIZE):
|
||||
total += session.scalar(select(func.count()).select_from(model).where(column.in_(chunk))) or 0
|
||||
return total
|
||||
|
||||
def _load_records_by_run_ids(
|
||||
self,
|
||||
session: Session,
|
||||
model: Any,
|
||||
run_ids: Sequence[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
return self._load_records_by_column(session, model, self._run_id_column(model), run_ids)
|
||||
|
||||
def _load_records_by_column(
|
||||
self,
|
||||
session: Session,
|
||||
model: Any,
|
||||
column: Any,
|
||||
values: Sequence[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
if not values:
|
||||
return []
|
||||
rows: list[Any] = []
|
||||
for chunk in self._chunks(values, _CHUNK_SIZE):
|
||||
rows.extend(session.scalars(select(model).where(column.in_(chunk))))
|
||||
return [self._row_to_dict(row) for row in rows]
|
||||
|
||||
@staticmethod
|
||||
def _delete_by_run_ids(
|
||||
session: Session,
|
||||
model: Any,
|
||||
run_ids: Sequence[str],
|
||||
) -> int:
|
||||
return WorkflowRunBundleArchiveMaintenance._delete_by_column(
|
||||
session, model, WorkflowRunBundleArchiveMaintenance._run_id_column(model), run_ids
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _run_id_column(model: Any) -> Any:
|
||||
if model is WorkflowRun:
|
||||
return WorkflowRun.id
|
||||
return model.workflow_run_id
|
||||
|
||||
@staticmethod
|
||||
def _delete_by_column(
|
||||
session: Session,
|
||||
model: Any,
|
||||
column: Any,
|
||||
values: Sequence[str],
|
||||
) -> int:
|
||||
if not values:
|
||||
return 0
|
||||
total = 0
|
||||
for chunk in WorkflowRunBundleArchiveMaintenance._chunks(values, _CHUNK_SIZE):
|
||||
result = session.execute(delete(model).where(column.in_(chunk)))
|
||||
total += cast(CursorResult, result).rowcount or 0
|
||||
return total
|
||||
|
||||
@staticmethod
|
||||
def _is_deleted(storage: ArchiveStorage, object_prefix: str) -> bool:
|
||||
return storage.object_exists(f"{object_prefix}/{ARCHIVE_BUNDLE_DELETED_MARKER_NAME}")
|
||||
|
||||
@staticmethod
|
||||
def _is_delete_started(storage: ArchiveStorage, object_prefix: str) -> bool:
|
||||
return storage.object_exists(f"{object_prefix}/{ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME}")
|
||||
|
||||
@staticmethod
|
||||
def _mark_deleted(storage: ArchiveStorage, object_prefix: str) -> None:
|
||||
WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME)
|
||||
|
||||
@staticmethod
|
||||
def _mark_restored(storage: ArchiveStorage, object_prefix: str) -> None:
|
||||
WorkflowRunBundleArchiveMaintenance._delete_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME)
|
||||
WorkflowRunBundleArchiveMaintenance._delete_marker(
|
||||
storage, object_prefix, ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME
|
||||
)
|
||||
WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
|
||||
|
||||
@staticmethod
|
||||
def _put_marker(storage: ArchiveStorage, object_prefix: str, marker_name: str) -> None:
|
||||
payload = json.dumps({"created_at": datetime.datetime.now(datetime.UTC).isoformat()}).encode("utf-8")
|
||||
storage.put_object(f"{object_prefix}/{marker_name}", payload)
|
||||
|
||||
@staticmethod
|
||||
def _delete_marker(storage: ArchiveStorage, object_prefix: str, marker_name: str) -> None:
|
||||
marker_key = f"{object_prefix}/{marker_name}"
|
||||
if storage.object_exists(marker_key):
|
||||
storage.delete_object(marker_key)
|
||||
|
||||
@staticmethod
|
||||
def _parse_manifest_datetime(value: str) -> datetime.datetime:
|
||||
return WorkflowRunBundleArchiveMaintenance._to_naive_utc(datetime.datetime.fromisoformat(value))
|
||||
|
||||
@staticmethod
|
||||
def _to_naive_utc(value: datetime.datetime) -> datetime.datetime:
|
||||
if value.tzinfo is None:
|
||||
return value
|
||||
return value.astimezone(datetime.UTC).replace(tzinfo=None)
|
||||
|
||||
@staticmethod
|
||||
def _chunks(values: Sequence[Any], size: int) -> list[Sequence[Any]]:
|
||||
return [values[index : index + size] for index in range(0, len(values), size)]
|
||||
|
||||
@staticmethod
|
||||
def _get_archive_storage() -> ArchiveStorage:
|
||||
try:
|
||||
return get_archive_storage()
|
||||
except ArchiveStorageNotConfiguredError as e:
|
||||
raise RuntimeError(f"Archive storage not configured: {e}") from e
|
||||
|
||||
@staticmethod
|
||||
def _merge_result(summary: BundleOperationSummary, result: BundleOperationResult) -> None:
|
||||
summary.results.append(result)
|
||||
summary.bundles_processed += 1
|
||||
summary.validation_time += result.validation_time
|
||||
if result.success:
|
||||
summary.bundles_succeeded += 1
|
||||
summary.rows_processed += result.row_count
|
||||
summary.runs_processed += result.run_count
|
||||
summary.archive_bytes += result.archive_bytes
|
||||
for table_name, count in result.table_counts.items():
|
||||
summary.table_counts[table_name] = summary.table_counts.get(table_name, 0) + count
|
||||
else:
|
||||
summary.bundles_failed += 1
|
||||
@@ -1,2 +1,10 @@
|
||||
ARCHIVE_SCHEMA_VERSION = "1.0"
|
||||
ARCHIVE_BUNDLE_NAME = f"archive.v{ARCHIVE_SCHEMA_VERSION}.zip"
|
||||
|
||||
ARCHIVE_BUNDLE_SCHEMA_VERSION = "2.0"
|
||||
ARCHIVE_BUNDLE_FORMAT = "parquet"
|
||||
ARCHIVE_BUNDLE_MANIFEST_NAME = "manifest.json"
|
||||
ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME = "_DELETE_STARTED"
|
||||
ARCHIVE_BUNDLE_DELETED_MARKER_NAME = "_DELETED"
|
||||
ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME = "_RESTORE_STARTED"
|
||||
ARCHIVE_BUNDLE_RESTORED_MARKER_NAME = "_RESTORED"
|
||||
|
||||
@@ -2,20 +2,68 @@
|
||||
Delete Archived Workflow Run Service.
|
||||
|
||||
This service deletes archived workflow run data from the database while keeping
|
||||
archive logs intact.
|
||||
archive logs intact. Deletion is intentionally gated by archive-object validation:
|
||||
the archive bundle must exist, have a supported manifest, pass zip/member checksum
|
||||
checks, and match the live row counts for every cleanup-owned table before rows
|
||||
are removed from the primary database.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import zipfile
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import TypedDict
|
||||
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from extensions.ext_database import db
|
||||
from models.workflow import WorkflowRun
|
||||
from libs.archive_storage import ArchiveStorage, ArchiveStorageNotConfiguredError, get_archive_storage
|
||||
from models.workflow import WorkflowArchiveLog, WorkflowRun
|
||||
from repositories.api_workflow_run_repository import APIWorkflowRunRepository, RunsWithRelatedCountsDict
|
||||
from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository
|
||||
from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_NAME, ARCHIVE_SCHEMA_VERSION
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _TableManifestEntry(TypedDict):
|
||||
row_count: int
|
||||
checksum: str
|
||||
size_bytes: int
|
||||
|
||||
|
||||
class _ArchiveManifest(TypedDict):
|
||||
schema_version: str
|
||||
workflow_run_id: str
|
||||
tenant_id: str
|
||||
app_id: str
|
||||
workflow_id: str
|
||||
tables: dict[str, _TableManifestEntry]
|
||||
|
||||
|
||||
_ARCHIVED_TABLES = [
|
||||
"workflow_runs",
|
||||
"workflow_app_logs",
|
||||
"workflow_node_executions",
|
||||
"workflow_node_execution_offload",
|
||||
"workflow_pauses",
|
||||
"workflow_pause_reasons",
|
||||
"workflow_trigger_logs",
|
||||
]
|
||||
|
||||
_TABLE_TO_COUNT_KEY = {
|
||||
"workflow_runs": "runs",
|
||||
"workflow_app_logs": "app_logs",
|
||||
"workflow_node_executions": "node_executions",
|
||||
"workflow_node_execution_offload": "offloads",
|
||||
"workflow_pauses": "pauses",
|
||||
"workflow_pause_reasons": "pause_reasons",
|
||||
"workflow_trigger_logs": "trigger_logs",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -34,13 +82,49 @@ class DeleteResult:
|
||||
"pause_reasons": 0,
|
||||
}
|
||||
)
|
||||
validated_counts: RunsWithRelatedCountsDict = field(
|
||||
default_factory=lambda: { # type: ignore[assignment]
|
||||
"runs": 0,
|
||||
"node_executions": 0,
|
||||
"offloads": 0,
|
||||
"app_logs": 0,
|
||||
"trigger_logs": 0,
|
||||
"pauses": 0,
|
||||
"pause_reasons": 0,
|
||||
}
|
||||
)
|
||||
archive_key: str | None = None
|
||||
restore_sampled: bool = False
|
||||
restore_sample_success: bool | None = None
|
||||
error: str | None = None
|
||||
elapsed_time: float = 0.0
|
||||
|
||||
|
||||
class ArchivedWorkflowRunDeletion:
|
||||
def __init__(self, dry_run: bool = False):
|
||||
"""
|
||||
Delete archived workflow-run rows after validating the archive bundle.
|
||||
|
||||
Args:
|
||||
dry_run: Preview validation and row counts without deleting.
|
||||
skip_bad_archives: Continue batch deletion after a validation/delete failure.
|
||||
restore_sample_interval: Run restore dry-run for every Nth successful deletion; 0 disables sampling.
|
||||
"""
|
||||
|
||||
_delete_attempt_count: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dry_run: bool = False,
|
||||
*,
|
||||
skip_bad_archives: bool = False,
|
||||
restore_sample_interval: int = 0,
|
||||
):
|
||||
self.dry_run = dry_run
|
||||
self.skip_bad_archives = skip_bad_archives
|
||||
if restore_sample_interval < 0:
|
||||
raise ValueError("restore_sample_interval must be >= 0")
|
||||
self.restore_sample_interval = restore_sample_interval
|
||||
self._delete_attempt_count = 0
|
||||
self.workflow_run_repo: APIWorkflowRunRepository | None = None
|
||||
|
||||
def delete_by_run_id(self, run_id: str) -> DeleteResult:
|
||||
@@ -57,12 +141,13 @@ class ArchivedWorkflowRunDeletion:
|
||||
return result
|
||||
|
||||
result.tenant_id = run.tenant_id
|
||||
if not repo.get_archived_run_ids(session, [run.id]):
|
||||
archive_log = repo.get_archived_log_by_run_id(run.id)
|
||||
if archive_log is None:
|
||||
result.error = f"Workflow run {run_id} is not archived"
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
result = self._delete_run(run)
|
||||
result = self._delete_run(run, archive_log)
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
@@ -78,8 +163,8 @@ class ArchivedWorkflowRunDeletion:
|
||||
|
||||
repo = self._get_workflow_run_repo()
|
||||
with session_maker() as session:
|
||||
runs = list(
|
||||
repo.get_archived_runs_by_time_range(
|
||||
archive_logs = list(
|
||||
repo.get_archived_logs_by_time_range(
|
||||
session=session,
|
||||
tenant_ids=tenant_ids,
|
||||
start_date=start_date,
|
||||
@@ -87,14 +172,44 @@ class ArchivedWorkflowRunDeletion:
|
||||
limit=limit,
|
||||
)
|
||||
)
|
||||
for run in runs:
|
||||
results.append(self._delete_run(run))
|
||||
run_ids = [archive_log.workflow_run_id for archive_log in archive_logs]
|
||||
runs_by_id = {run.id: run for run in session.query(WorkflowRun).where(WorkflowRun.id.in_(run_ids)).all()}
|
||||
for archive_log in archive_logs:
|
||||
run = runs_by_id.get(archive_log.workflow_run_id)
|
||||
if run is None:
|
||||
result = DeleteResult(
|
||||
run_id=archive_log.workflow_run_id,
|
||||
tenant_id=archive_log.tenant_id,
|
||||
success=False,
|
||||
error=f"Workflow run {archive_log.workflow_run_id} not found",
|
||||
)
|
||||
else:
|
||||
result = self._delete_run(run, archive_log)
|
||||
results.append(result)
|
||||
if not result.success and not self.skip_bad_archives:
|
||||
logger.error("Stopping archived workflow run deletion after failure: %s", result.error)
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
def _delete_run(self, run: WorkflowRun) -> DeleteResult:
|
||||
def _delete_run(self, run: WorkflowRun, archive_log: WorkflowArchiveLog | None = None) -> DeleteResult:
|
||||
start_time = time.time()
|
||||
result = DeleteResult(run_id=run.id, tenant_id=run.tenant_id, success=False)
|
||||
if archive_log is None:
|
||||
archive_log = self._get_workflow_run_repo().get_archived_log_by_run_id(run.id)
|
||||
if archive_log is None:
|
||||
result.error = f"Workflow run {run.id} is not archived"
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
try:
|
||||
result.archive_key = self._validate_archive_before_delete(run, archive_log)
|
||||
result.validated_counts = self._count_live_related_rows(run)
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
if self.dry_run:
|
||||
result.success = True
|
||||
result.elapsed_time = time.time() - start_time
|
||||
@@ -108,17 +223,202 @@ class ArchivedWorkflowRunDeletion:
|
||||
delete_trigger_logs=self._delete_trigger_logs,
|
||||
)
|
||||
result.deleted_counts = deleted_counts
|
||||
self._verify_post_delete(run.id)
|
||||
if self._should_run_restore_sample():
|
||||
result.restore_sampled = True
|
||||
result.restore_sample_success = self._run_restore_dry_run_sample(archive_log)
|
||||
if not result.restore_sample_success:
|
||||
raise RuntimeError(f"Restore dry-run sample failed for workflow run {run.id}")
|
||||
result.success = True
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
def _validate_archive_before_delete(self, run: WorkflowRun, archive_log: WorkflowArchiveLog) -> str:
|
||||
storage = self._get_archive_storage()
|
||||
archive_key = self._get_archive_key(archive_log)
|
||||
if not storage.object_exists(archive_key):
|
||||
raise FileNotFoundError(f"Archive bundle not found: {archive_key}")
|
||||
|
||||
archive_data = storage.get_object(archive_key)
|
||||
manifest = self._validate_archive_bundle(
|
||||
archive_data,
|
||||
run_id=run.id,
|
||||
tenant_id=run.tenant_id,
|
||||
app_id=run.app_id,
|
||||
workflow_id=run.workflow_id,
|
||||
)
|
||||
expected_counts = self._counts_from_manifest(manifest)
|
||||
current_counts = self._count_live_related_rows(run)
|
||||
if current_counts != expected_counts:
|
||||
raise ValueError(
|
||||
"Archive row count mismatch before delete: "
|
||||
f"run_id={run.id}, expected={expected_counts}, current={current_counts}"
|
||||
)
|
||||
return archive_key
|
||||
|
||||
@staticmethod
|
||||
def _validate_archive_bundle(
|
||||
archive_data: bytes,
|
||||
*,
|
||||
run_id: str,
|
||||
tenant_id: str,
|
||||
app_id: str,
|
||||
workflow_id: str,
|
||||
) -> _ArchiveManifest:
|
||||
try:
|
||||
with zipfile.ZipFile(io.BytesIO(archive_data), mode="r") as archive:
|
||||
bad_member = archive.testzip()
|
||||
if bad_member:
|
||||
raise ValueError(f"zip CRC check failed for member {bad_member}")
|
||||
try:
|
||||
manifest_data = archive.read("manifest.json")
|
||||
except KeyError as e:
|
||||
raise ValueError("manifest.json missing from archive bundle") from e
|
||||
loaded = json.loads(manifest_data)
|
||||
if not isinstance(loaded, dict):
|
||||
raise ValueError("manifest.json must be an object")
|
||||
manifest = loaded
|
||||
|
||||
required_fields = {
|
||||
"schema_version",
|
||||
"workflow_run_id",
|
||||
"tenant_id",
|
||||
"app_id",
|
||||
"workflow_id",
|
||||
"tables",
|
||||
}
|
||||
missing_fields = sorted(required_fields - set(manifest))
|
||||
if missing_fields:
|
||||
raise ValueError(f"manifest missing required fields: {', '.join(missing_fields)}")
|
||||
if manifest["schema_version"] != ARCHIVE_SCHEMA_VERSION:
|
||||
raise ValueError(
|
||||
f"unsupported archive schema_version: {manifest['schema_version']} "
|
||||
f"(expected {ARCHIVE_SCHEMA_VERSION})"
|
||||
)
|
||||
if manifest["workflow_run_id"] != run_id:
|
||||
raise ValueError("manifest workflow_run_id does not match delete target")
|
||||
if manifest["tenant_id"] != tenant_id:
|
||||
raise ValueError("manifest tenant_id does not match delete target")
|
||||
if manifest["app_id"] != app_id:
|
||||
raise ValueError("manifest app_id does not match delete target")
|
||||
if manifest["workflow_id"] != workflow_id:
|
||||
raise ValueError("manifest workflow_id does not match delete target")
|
||||
|
||||
tables = manifest["tables"]
|
||||
if not isinstance(tables, dict):
|
||||
raise ValueError("manifest tables must be an object")
|
||||
missing_tables = [table_name for table_name in _ARCHIVED_TABLES if table_name not in tables]
|
||||
if missing_tables:
|
||||
raise ValueError(f"manifest missing tables: {', '.join(missing_tables)}")
|
||||
|
||||
for table_name in _ARCHIVED_TABLES:
|
||||
info = tables[table_name]
|
||||
if not isinstance(info, dict):
|
||||
raise ValueError(f"manifest table entry must be an object: {table_name}")
|
||||
for key in ("row_count", "checksum", "size_bytes"):
|
||||
if key not in info:
|
||||
raise ValueError(f"manifest table {table_name} missing {key}")
|
||||
member_path = f"{table_name}.jsonl"
|
||||
try:
|
||||
payload = archive.read(member_path)
|
||||
except KeyError as e:
|
||||
raise ValueError(f"archive member missing: {member_path}") from e
|
||||
if len(payload) != info["size_bytes"]:
|
||||
raise ValueError(
|
||||
f"archive member size mismatch for {member_path}: "
|
||||
f"expected={info['size_bytes']}, actual={len(payload)}"
|
||||
)
|
||||
checksum = ArchiveStorage.compute_checksum(payload)
|
||||
if checksum != info["checksum"]:
|
||||
raise ValueError(
|
||||
f"archive member checksum mismatch for {member_path}: "
|
||||
f"expected={info['checksum']}, actual={checksum}"
|
||||
)
|
||||
row_count = len(ArchiveStorage.deserialize_from_jsonl(payload))
|
||||
if row_count != info["row_count"]:
|
||||
raise ValueError(
|
||||
f"archive row count mismatch for {member_path}: "
|
||||
f"expected={info['row_count']}, actual={row_count}"
|
||||
)
|
||||
|
||||
return manifest # type: ignore[return-value]
|
||||
except zipfile.BadZipFile as e:
|
||||
raise ValueError("archive bundle is not a valid zip file") from e
|
||||
|
||||
@staticmethod
|
||||
def _counts_from_manifest(manifest: _ArchiveManifest) -> RunsWithRelatedCountsDict:
|
||||
counts: RunsWithRelatedCountsDict = {
|
||||
"runs": 0,
|
||||
"node_executions": 0,
|
||||
"offloads": 0,
|
||||
"app_logs": 0,
|
||||
"trigger_logs": 0,
|
||||
"pauses": 0,
|
||||
"pause_reasons": 0,
|
||||
}
|
||||
for table_name, count_key in _TABLE_TO_COUNT_KEY.items():
|
||||
counts[count_key] = manifest["tables"][table_name]["row_count"] # type: ignore[literal-required]
|
||||
return counts
|
||||
|
||||
def _count_live_related_rows(self, run: WorkflowRun) -> RunsWithRelatedCountsDict:
|
||||
repo = self._get_workflow_run_repo()
|
||||
return repo.count_runs_with_related(
|
||||
[run],
|
||||
count_node_executions=self._count_node_executions,
|
||||
count_trigger_logs=self._count_trigger_logs,
|
||||
)
|
||||
|
||||
def _verify_post_delete(self, run_id: str) -> None:
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False)() as session:
|
||||
if session.get(WorkflowRun, run_id) is not None:
|
||||
raise RuntimeError(f"Post-delete verification failed: workflow run {run_id} still exists")
|
||||
|
||||
def _should_run_restore_sample(self) -> bool:
|
||||
if self.restore_sample_interval == 0:
|
||||
return False
|
||||
self._delete_attempt_count += 1
|
||||
return self._delete_attempt_count % self.restore_sample_interval == 0
|
||||
|
||||
@staticmethod
|
||||
def _run_restore_dry_run_sample(archive_log: WorkflowArchiveLog) -> bool:
|
||||
from services.retention.workflow_run.restore_archived_workflow_run import WorkflowRunRestore
|
||||
|
||||
restorer = WorkflowRunRestore(dry_run=True, workers=1)
|
||||
# Reuse restore's dry-run path so the runbook exercises the actual restore code.
|
||||
result = restorer._restore_from_run(
|
||||
archive_log,
|
||||
session_maker=sessionmaker(bind=db.engine, expire_on_commit=False),
|
||||
)
|
||||
return result.success
|
||||
|
||||
@staticmethod
|
||||
def _get_archive_key(archive_log: WorkflowArchiveLog) -> str:
|
||||
created_at = archive_log.run_created_at
|
||||
prefix = (
|
||||
f"{archive_log.tenant_id}/app_id={archive_log.app_id}/year={created_at.strftime('%Y')}/"
|
||||
f"month={created_at.strftime('%m')}/workflow_run_id={archive_log.workflow_run_id}"
|
||||
)
|
||||
return f"{prefix}/{ARCHIVE_BUNDLE_NAME}"
|
||||
|
||||
@staticmethod
|
||||
def _get_archive_storage() -> ArchiveStorage:
|
||||
try:
|
||||
return get_archive_storage()
|
||||
except ArchiveStorageNotConfiguredError as e:
|
||||
raise RuntimeError(f"Archive storage not configured: {e}") from e
|
||||
|
||||
@staticmethod
|
||||
def _delete_trigger_logs(session: Session, run_ids: Sequence[str]) -> int:
|
||||
trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
|
||||
return trigger_repo.delete_by_run_ids(run_ids)
|
||||
|
||||
@staticmethod
|
||||
def _count_trigger_logs(session: Session, run_ids: Sequence[str]) -> int:
|
||||
trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
|
||||
return trigger_repo.count_by_run_ids(run_ids)
|
||||
|
||||
@staticmethod
|
||||
def _delete_node_executions(
|
||||
session: Session,
|
||||
@@ -132,6 +432,19 @@ class ArchivedWorkflowRunDeletion:
|
||||
)
|
||||
return repo.delete_by_runs(session, run_ids)
|
||||
|
||||
@staticmethod
|
||||
def _count_node_executions(
|
||||
session: Session,
|
||||
runs: Sequence[WorkflowRun],
|
||||
) -> tuple[int, int]:
|
||||
from repositories.factory import DifyAPIRepositoryFactory
|
||||
|
||||
run_ids = [run.id for run in runs]
|
||||
repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(
|
||||
session_maker=sessionmaker(bind=session.get_bind(), expire_on_commit=False)
|
||||
)
|
||||
return repo.count_by_runs(session, run_ids)
|
||||
|
||||
def _get_workflow_run_repo(self) -> APIWorkflowRunRepository:
|
||||
if self.workflow_run_repo is not None:
|
||||
return self.workflow_run_repo
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
def tenant_prefix_bounds(prefix: str) -> tuple[str, str | None]:
|
||||
prefix_value = int(prefix, 16)
|
||||
lower_bound = f"{prefix}0000000-0000-0000-0000-000000000000"
|
||||
if prefix_value == 15:
|
||||
return lower_bound, None
|
||||
upper_bound = f"{prefix_value + 1:x}0000000-0000-0000-0000-000000000000"
|
||||
return lower_bound, upper_bound
|
||||
|
||||
|
||||
def tenant_prefix_condition(column, prefix: str):
|
||||
lower_bound, upper_bound = tenant_prefix_bounds(prefix)
|
||||
condition = column >= lower_bound
|
||||
if upper_bound is not None:
|
||||
condition = sa.and_(condition, column < upper_bound)
|
||||
return condition
|
||||
@@ -475,7 +475,7 @@ class TriggerProviderService:
|
||||
tenant_id=tenant_id, provider_id=provider_id
|
||||
)
|
||||
# Create encrypter
|
||||
encrypter, cache = create_provider_encrypter(
|
||||
encrypter, _ = create_provider_encrypter(
|
||||
tenant_id=tenant_id,
|
||||
config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()],
|
||||
cache=NoOpProviderCredentialCache(),
|
||||
@@ -506,14 +506,20 @@ class TriggerProviderService:
|
||||
subscription.credentials = dict(encrypter.encrypt(dict(refreshed_credentials.credentials)))
|
||||
subscription.credential_expires_at = refreshed_credentials.expires_at
|
||||
|
||||
# Clear cache
|
||||
cache.delete()
|
||||
|
||||
return {
|
||||
provider_id_value = subscription.provider_id
|
||||
result = {
|
||||
"result": "success",
|
||||
"expires_at": refreshed_credentials.expires_at,
|
||||
}
|
||||
|
||||
# Clear the trigger runtime credential cache after the DB commit so dispatch uses the refreshed token.
|
||||
delete_cache_for_subscription(
|
||||
tenant_id=tenant_id,
|
||||
provider_id=provider_id_value,
|
||||
subscription_id=subscription_id,
|
||||
)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def refresh_subscription(
|
||||
cls,
|
||||
|
||||
@@ -23,8 +23,11 @@ from core.app.entities.task_entities import (
|
||||
WorkflowStartStreamResponse,
|
||||
)
|
||||
from core.app.layers.pause_state_persist_layer import WorkflowResumptionContext
|
||||
from core.workflow.human_input_forms import load_form_tokens_by_form_id
|
||||
from core.workflow.human_input_forms import (
|
||||
load_form_dispositions_by_form_id,
|
||||
)
|
||||
from core.workflow.human_input_policy import (
|
||||
FormDisposition,
|
||||
HumanInputSurface,
|
||||
enrich_human_input_pause_reasons,
|
||||
resolve_human_input_pause_reason_inputs,
|
||||
@@ -359,7 +362,7 @@ def _build_human_input_required_events(
|
||||
|
||||
expiration_times_by_form_id: dict[str, int] = {}
|
||||
display_in_ui_by_form_id: dict[str, bool] = {}
|
||||
form_tokens_by_form_id: dict[str, str] = {}
|
||||
dispositions_by_form_id: dict[str, FormDisposition] = {}
|
||||
if human_input_form_ids and session_maker is not None:
|
||||
stmt = select(HumanInputForm.id, HumanInputForm.expiration_time, HumanInputForm.form_definition).where(
|
||||
HumanInputForm.id.in_(human_input_form_ids)
|
||||
@@ -372,7 +375,7 @@ def _build_human_input_required_events(
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
definition_payload = {}
|
||||
display_in_ui_by_form_id[str(form_id)] = bool(definition_payload.get("display_in_ui"))
|
||||
form_tokens_by_form_id = load_form_tokens_by_form_id(
|
||||
dispositions_by_form_id = load_form_dispositions_by_form_id(
|
||||
human_input_form_ids,
|
||||
session=session,
|
||||
surface=human_input_surface,
|
||||
@@ -393,6 +396,7 @@ def _build_human_input_required_events(
|
||||
reason.inputs,
|
||||
variable_pool=variable_pool,
|
||||
)
|
||||
disposition = dispositions_by_form_id.get(form_id)
|
||||
|
||||
response = HumanInputRequiredResponse(
|
||||
task_id=task_id,
|
||||
@@ -405,7 +409,8 @@ def _build_human_input_required_events(
|
||||
inputs=resolved_inputs,
|
||||
actions=reason.actions,
|
||||
display_in_ui=display_in_ui_by_form_id.get(form_id, False),
|
||||
form_token=form_tokens_by_form_id.get(form_id),
|
||||
form_token=disposition.form_token if disposition else None,
|
||||
approval_channels=list(disposition.approval_channels) if disposition else [],
|
||||
resolved_default_values=reason.resolved_default_values,
|
||||
expiration_time=expiration_time,
|
||||
),
|
||||
@@ -493,11 +498,11 @@ def _build_pause_event(
|
||||
for form_id in [reason.get("form_id")]
|
||||
if isinstance(form_id, str)
|
||||
]
|
||||
form_tokens_by_form_id: dict[str, str] = {}
|
||||
dispositions_by_form_id: dict[str, FormDisposition] = {}
|
||||
expiration_times_by_form_id: dict[str, int] = {}
|
||||
if human_input_form_ids and session_maker is not None:
|
||||
with session_maker() as session:
|
||||
form_tokens_by_form_id = load_form_tokens_by_form_id(
|
||||
dispositions_by_form_id = load_form_dispositions_by_form_id(
|
||||
human_input_form_ids,
|
||||
session=session,
|
||||
surface=human_input_surface,
|
||||
@@ -512,7 +517,7 @@ def _build_pause_event(
|
||||
# otherwise clients see schema drift after resume.
|
||||
reasons = enrich_human_input_pause_reasons(
|
||||
reasons,
|
||||
form_tokens_by_form_id=form_tokens_by_form_id,
|
||||
dispositions_by_form_id=dispositions_by_form_id,
|
||||
expiration_times_by_form_id=expiration_times_by_form_id,
|
||||
)
|
||||
|
||||
|
||||
@@ -19,11 +19,16 @@ from core.app.entities.app_invoke_entities import (
|
||||
InvokeFrom,
|
||||
WorkflowAppGenerateEntity,
|
||||
)
|
||||
from core.app.entities.task_entities import WorkflowFinishStreamResponse, WorkflowStartStreamResponse
|
||||
from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig, WorkflowResumptionContext
|
||||
from core.repositories import DifyCoreRepositoryFactory
|
||||
from extensions.ext_database import db
|
||||
from graphon.entities import WorkflowStartReason
|
||||
from graphon.enums import WorkflowExecutionStatus
|
||||
from graphon.runtime import GraphRuntimeState
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from libs.flask_utils import set_login_user
|
||||
from libs.helper import to_timestamp
|
||||
from models.account import Account
|
||||
from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
|
||||
from models.model import App, AppMode, Conversation, EndUser, Message
|
||||
@@ -173,14 +178,24 @@ class _AppRunner:
|
||||
)
|
||||
except Exception as exc:
|
||||
if exec_params.streaming:
|
||||
_publish_error_event(exc, exec_params.workflow_run_id, exec_params.app_mode)
|
||||
_publish_failed_workflow_terminal_events(
|
||||
exc=exc,
|
||||
exec_params=exec_params,
|
||||
)
|
||||
raise
|
||||
|
||||
if not exec_params.streaming:
|
||||
return response
|
||||
|
||||
assert isinstance(response, Generator)
|
||||
_publish_streaming_response(response, exec_params.workflow_run_id, exec_params.app_mode)
|
||||
_publish_streaming_response(
|
||||
response,
|
||||
exec_params.workflow_run_id,
|
||||
exec_params.app_mode,
|
||||
exec_params.workflow_id,
|
||||
exec_params.args.get("inputs", {}),
|
||||
WorkflowStartReason.INITIAL,
|
||||
)
|
||||
|
||||
def _run_app(
|
||||
self,
|
||||
@@ -246,29 +261,197 @@ def _resolve_user_for_run(session: Session, workflow_run: WorkflowRun) -> Accoun
|
||||
return session.get(EndUser, workflow_run.created_by)
|
||||
|
||||
|
||||
def _publish_error_event(exc: Exception, workflow_run_id: str, app_mode: AppMode) -> None:
|
||||
topic = MessageBasedAppGenerator.get_response_topic(app_mode, workflow_run_id)
|
||||
payload = json.dumps({"event": "error", "message": str(exc), "status": 500})
|
||||
topic.publish(payload.encode())
|
||||
def _publish_failed_workflow_terminal_events(exc: Exception, exec_params: AppExecutionParams) -> None:
|
||||
"""Publish synthetic workflow lifecycle events for pre-runtime failures.
|
||||
|
||||
Early failures can happen before the app generator creates a task entity or
|
||||
emits any workflow queue events. In that window SSE consumers still need a
|
||||
normal terminal event to close their state machines, so we synthesize a
|
||||
minimal `workflow_started -> workflow_finished(failed)` sequence here.
|
||||
|
||||
`workflow_run_id` is reused as a synthetic `task_id` because no application
|
||||
task id exists yet on this failure path.
|
||||
"""
|
||||
timestamp = to_timestamp(naive_utc_now())
|
||||
assert timestamp is not None
|
||||
|
||||
topic = MessageBasedAppGenerator.get_response_topic(exec_params.app_mode, exec_params.workflow_run_id)
|
||||
started_payload = WorkflowStartStreamResponse(
|
||||
task_id=exec_params.workflow_run_id,
|
||||
workflow_run_id=exec_params.workflow_run_id,
|
||||
data=WorkflowStartStreamResponse.Data(
|
||||
id=exec_params.workflow_run_id,
|
||||
workflow_id=exec_params.workflow_id,
|
||||
inputs=exec_params.args.get("inputs", {}),
|
||||
created_at=timestamp,
|
||||
reason=WorkflowStartReason.INITIAL,
|
||||
),
|
||||
)
|
||||
topic.publish(json.dumps(started_payload.model_dump(mode="json"), ensure_ascii=False).encode())
|
||||
|
||||
finished_payload = WorkflowFinishStreamResponse(
|
||||
task_id=exec_params.workflow_run_id,
|
||||
workflow_run_id=exec_params.workflow_run_id,
|
||||
data=WorkflowFinishStreamResponse.Data(
|
||||
id=exec_params.workflow_run_id,
|
||||
workflow_id=exec_params.workflow_id,
|
||||
status=WorkflowExecutionStatus.FAILED,
|
||||
outputs=None,
|
||||
error=str(exc),
|
||||
elapsed_time=0.0,
|
||||
total_tokens=0,
|
||||
total_steps=0,
|
||||
created_by={},
|
||||
created_at=timestamp,
|
||||
finished_at=timestamp,
|
||||
exceptions_count=1,
|
||||
files=[],
|
||||
),
|
||||
)
|
||||
topic.publish(json.dumps(finished_payload.model_dump(mode="json"), ensure_ascii=False).encode())
|
||||
|
||||
|
||||
def _get_event_name(event: str | Mapping[str, Any] | BaseModel) -> str | None:
|
||||
if isinstance(event, BaseModel):
|
||||
# Temporary compatibility for legacy BaseModel stream events; remove after confirming generators always emit
|
||||
# str / Mapping responses.
|
||||
event_name = getattr(event, "event", None)
|
||||
elif isinstance(event, Mapping):
|
||||
event_name = event.get("event")
|
||||
else:
|
||||
return None
|
||||
|
||||
if event_name is None:
|
||||
return None
|
||||
return str(event_name)
|
||||
|
||||
|
||||
def _get_task_id(event: str | Mapping[str, Any] | BaseModel) -> str | None:
|
||||
if isinstance(event, BaseModel):
|
||||
# Temporary compatibility for legacy BaseModel stream events; remove after confirming generators always emit
|
||||
# str / Mapping responses.
|
||||
task_id = getattr(event, "task_id", None)
|
||||
elif isinstance(event, Mapping):
|
||||
task_id = event.get("task_id")
|
||||
else:
|
||||
return None
|
||||
|
||||
return task_id if isinstance(task_id, str) and task_id else None
|
||||
|
||||
|
||||
def _publish_streaming_response(
|
||||
response_stream: Generator[str | Mapping[str, Any] | BaseModel, None, None],
|
||||
workflow_run_id: str,
|
||||
workflow_run_id: str | uuid.UUID,
|
||||
app_mode: AppMode,
|
||||
workflow_id: str,
|
||||
inputs: Mapping[str, Any],
|
||||
started_reason: WorkflowStartReason,
|
||||
) -> None:
|
||||
topic = MessageBasedAppGenerator.get_response_topic(app_mode, workflow_run_id)
|
||||
for event in response_stream:
|
||||
try:
|
||||
if isinstance(event, BaseModel):
|
||||
payload = json.dumps(event.model_dump(mode="json"), ensure_ascii=False)
|
||||
else:
|
||||
payload = json.dumps(event, ensure_ascii=False, default=str)
|
||||
except (TypeError, ValueError):
|
||||
logger.exception("error while encoding event")
|
||||
continue
|
||||
"""Publish workflow stream events and close broken streams with a failed terminal event.
|
||||
|
||||
topic.publish(payload.encode())
|
||||
`_AppRunner.run()` only handles failures before the generator is returned.
|
||||
Once we start iterating the runtime stream, this helper becomes the last
|
||||
place that can guarantee SSE consumers eventually see a terminal workflow
|
||||
lifecycle event.
|
||||
"""
|
||||
normalized_workflow_run_id = str(workflow_run_id)
|
||||
|
||||
def _publish_failed_terminal_event(error_message: str, task_id: str, publish_started: bool) -> None:
|
||||
timestamp = to_timestamp(naive_utc_now())
|
||||
assert timestamp is not None
|
||||
|
||||
if publish_started:
|
||||
started_payload = WorkflowStartStreamResponse(
|
||||
task_id=task_id,
|
||||
workflow_run_id=normalized_workflow_run_id,
|
||||
data=WorkflowStartStreamResponse.Data(
|
||||
id=normalized_workflow_run_id,
|
||||
workflow_id=workflow_id,
|
||||
inputs=inputs,
|
||||
created_at=timestamp,
|
||||
reason=started_reason,
|
||||
),
|
||||
)
|
||||
topic.publish(
|
||||
json.dumps(
|
||||
started_payload.model_dump(mode="json", fallback=str),
|
||||
ensure_ascii=False,
|
||||
).encode()
|
||||
)
|
||||
|
||||
finished_payload = WorkflowFinishStreamResponse(
|
||||
task_id=task_id,
|
||||
workflow_run_id=normalized_workflow_run_id,
|
||||
data=WorkflowFinishStreamResponse.Data(
|
||||
id=normalized_workflow_run_id,
|
||||
workflow_id=workflow_id,
|
||||
status=WorkflowExecutionStatus.FAILED,
|
||||
outputs=None,
|
||||
error=error_message,
|
||||
elapsed_time=0.0,
|
||||
total_tokens=0,
|
||||
total_steps=0,
|
||||
created_by={},
|
||||
created_at=timestamp,
|
||||
finished_at=timestamp,
|
||||
exceptions_count=1,
|
||||
files=[],
|
||||
),
|
||||
)
|
||||
topic.publish(json.dumps(finished_payload.model_dump(mode="json"), ensure_ascii=False).encode())
|
||||
|
||||
terminal_events = {"workflow_finished", "workflow_paused"}
|
||||
unexpected_stream_end_message = "Workflow stream ended without a terminal event"
|
||||
topic = MessageBasedAppGenerator.get_response_topic(app_mode, normalized_workflow_run_id)
|
||||
started_published = False
|
||||
terminal_published = False
|
||||
last_task_id = normalized_workflow_run_id
|
||||
|
||||
try:
|
||||
for event in response_stream:
|
||||
event_name = _get_event_name(event)
|
||||
task_id = _get_task_id(event)
|
||||
if task_id is not None:
|
||||
last_task_id = task_id
|
||||
|
||||
try:
|
||||
if isinstance(event, BaseModel):
|
||||
payload = json.dumps(event.model_dump(mode="json"), ensure_ascii=False)
|
||||
else:
|
||||
payload = json.dumps(event, ensure_ascii=False, default=str)
|
||||
except (TypeError, ValueError):
|
||||
logger.exception("error while encoding event")
|
||||
continue
|
||||
|
||||
topic.publish(payload.encode())
|
||||
|
||||
if event_name == "workflow_started":
|
||||
started_published = True
|
||||
elif event_name in terminal_events:
|
||||
terminal_published = True
|
||||
except Exception as exc:
|
||||
if not terminal_published:
|
||||
logger.exception(
|
||||
"Workflow stream for run %s failed before terminal event; publishing fallback terminal event",
|
||||
normalized_workflow_run_id,
|
||||
)
|
||||
_publish_failed_terminal_event(
|
||||
error_message=str(exc) or exc.__class__.__name__,
|
||||
task_id=last_task_id,
|
||||
publish_started=not started_published,
|
||||
)
|
||||
raise
|
||||
|
||||
if not terminal_published:
|
||||
logger.warning(
|
||||
"Workflow stream for run %s ended without a terminal event; publishing fallback terminal event",
|
||||
normalized_workflow_run_id,
|
||||
)
|
||||
_publish_failed_terminal_event(
|
||||
error_message=unexpected_stream_end_message,
|
||||
task_id=last_task_id,
|
||||
publish_started=not started_published,
|
||||
)
|
||||
|
||||
|
||||
@shared_task(queue=WORKFLOW_BASED_APP_EXECUTION_QUEUE)
|
||||
@@ -454,7 +637,14 @@ def _resume_advanced_chat(
|
||||
raise
|
||||
|
||||
assert isinstance(response, Generator)
|
||||
_publish_streaming_response(response, workflow_run_id, AppMode.ADVANCED_CHAT)
|
||||
_publish_streaming_response(
|
||||
response,
|
||||
workflow_run_id,
|
||||
AppMode.ADVANCED_CHAT,
|
||||
workflow.id,
|
||||
generate_entity.inputs,
|
||||
WorkflowStartReason.RESUMPTION,
|
||||
)
|
||||
|
||||
|
||||
def _resume_workflow(
|
||||
@@ -509,7 +699,14 @@ def _resume_workflow(
|
||||
raise
|
||||
|
||||
assert isinstance(response, Generator)
|
||||
_publish_streaming_response(response, workflow_run_id, AppMode.WORKFLOW)
|
||||
_publish_streaming_response(
|
||||
response,
|
||||
workflow_run_id,
|
||||
AppMode.WORKFLOW,
|
||||
workflow.id,
|
||||
generate_entity.inputs,
|
||||
WorkflowStartReason.RESUMPTION,
|
||||
)
|
||||
|
||||
try:
|
||||
workflow_run_repo.delete_workflow_pause(pause_entity)
|
||||
|
||||
@@ -1,17 +1,17 @@
|
||||
import datetime
|
||||
import io
|
||||
import json
|
||||
import uuid
|
||||
import zipfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import pytest
|
||||
|
||||
from services.retention.workflow_run.archive_paid_plan_workflow_run import (
|
||||
ArchiveSummary,
|
||||
WorkflowRunArchiver,
|
||||
)
|
||||
from services.retention.workflow_run.constants import ARCHIVE_SCHEMA_VERSION
|
||||
from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_FORMAT, ARCHIVE_BUNDLE_SCHEMA_VERSION
|
||||
|
||||
|
||||
class TestWorkflowRunArchiverInit:
|
||||
@@ -39,6 +39,22 @@ class TestWorkflowRunArchiverInit:
|
||||
with pytest.raises(ValueError, match="workers must be at least 1"):
|
||||
WorkflowRunArchiver(workers=0)
|
||||
|
||||
def test_run_shard_index_without_total_raises(self):
|
||||
with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"):
|
||||
WorkflowRunArchiver(run_shard_index=0)
|
||||
|
||||
def test_run_shard_total_without_index_raises(self):
|
||||
with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"):
|
||||
WorkflowRunArchiver(run_shard_total=4)
|
||||
|
||||
def test_run_shard_total_above_supported_range_raises(self):
|
||||
with pytest.raises(ValueError, match="run_shard_total must be between 1 and 16"):
|
||||
WorkflowRunArchiver(run_shard_index=0, run_shard_total=17)
|
||||
|
||||
def test_run_shard_index_must_be_less_than_total(self):
|
||||
with pytest.raises(ValueError, match="run_shard_index must be between 0 and run_shard_total - 1"):
|
||||
WorkflowRunArchiver(run_shard_index=4, run_shard_total=4)
|
||||
|
||||
def test_valid_init_defaults(self):
|
||||
archiver = WorkflowRunArchiver(days=30, batch_size=50)
|
||||
assert archiver.days == 30
|
||||
@@ -55,29 +71,93 @@ class TestWorkflowRunArchiverInit:
|
||||
assert archiver.end_before is not None
|
||||
assert archiver.workers == 2
|
||||
|
||||
def test_delete_after_archive_is_not_supported_for_bundle_archive(self):
|
||||
with pytest.raises(ValueError, match="delete_after_archive is not supported by bundle archive"):
|
||||
WorkflowRunArchiver(delete_after_archive=True)
|
||||
|
||||
def test_get_runs_batch_passes_shard_options(self):
|
||||
repo = MagicMock()
|
||||
repo.get_runs_batch_by_time_range.return_value = []
|
||||
archiver = WorkflowRunArchiver(
|
||||
tenant_prefixes=["0", "a"],
|
||||
run_shard_index=1,
|
||||
run_shard_total=4,
|
||||
workflow_run_repo=repo,
|
||||
)
|
||||
|
||||
archiver._get_runs_batch(None)
|
||||
|
||||
repo.get_runs_batch_by_time_range.assert_called_once()
|
||||
assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] == ["0", "a"]
|
||||
assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_index"] == 1
|
||||
assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_total"] == 4
|
||||
|
||||
def test_get_runs_batch_prefers_planned_tenant_ids_over_prefix_filter(self):
|
||||
repo = MagicMock()
|
||||
repo.get_runs_batch_by_time_range.return_value = []
|
||||
archiver = WorkflowRunArchiver(
|
||||
tenant_ids=["0tenant"],
|
||||
tenant_prefixes=["0"],
|
||||
paid_tenant_ids=["0tenant"],
|
||||
workflow_run_repo=repo,
|
||||
)
|
||||
|
||||
archiver._get_runs_batch(None)
|
||||
|
||||
repo.get_runs_batch_by_time_range.assert_called_once()
|
||||
assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["0tenant"]
|
||||
assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] is None
|
||||
|
||||
def test_get_runs_batch_uses_current_tenant_scan_scope(self):
|
||||
repo = MagicMock()
|
||||
repo.get_runs_batch_by_time_range.return_value = []
|
||||
archiver = WorkflowRunArchiver(
|
||||
tenant_ids=["tenant-a", "tenant-b"],
|
||||
workflow_run_repo=repo,
|
||||
)
|
||||
|
||||
archiver._get_runs_batch(None, tenant_scope=["tenant-b"])
|
||||
|
||||
repo.get_runs_batch_by_time_range.assert_called_once()
|
||||
assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["tenant-b"]
|
||||
|
||||
def test_start_message_includes_shard(self):
|
||||
archiver = WorkflowRunArchiver(tenant_prefixes=["0"], run_shard_index=1, run_shard_total=4)
|
||||
|
||||
message = archiver._build_start_message()
|
||||
|
||||
assert "tenant_prefixes=0" in message
|
||||
assert "run_shard=1/4" in message
|
||||
|
||||
def test_start_message_summarizes_large_planned_tenant_list(self):
|
||||
tenant_ids = [f"tenant-{index}" for index in range(11)]
|
||||
archiver = WorkflowRunArchiver(tenant_ids=tenant_ids, tenant_prefixes=["0"])
|
||||
|
||||
message = archiver._build_start_message()
|
||||
|
||||
assert "tenant_ids=11 planned tenants" in message
|
||||
assert "tenant-10" not in message
|
||||
|
||||
|
||||
class TestBuildArchiveBundle:
|
||||
def test_bundle_contains_manifest_and_all_tables(self):
|
||||
def test_bundle_contains_manifest_and_all_table_objects(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = MagicMock()
|
||||
run.id = str(uuid.uuid4())
|
||||
run.tenant_id = str(uuid.uuid4())
|
||||
run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
|
||||
identity = archiver._build_bundle_identity([run])
|
||||
table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}
|
||||
|
||||
manifest_data = json.dumps({"schema_version": ARCHIVE_SCHEMA_VERSION}).encode("utf-8")
|
||||
table_payloads = dict.fromkeys(archiver.ARCHIVED_TABLES, b"")
|
||||
table_stats, table_payloads, manifest_data = archiver._build_archive_payload(identity, [run], table_data)
|
||||
manifest = json.loads(manifest_data)
|
||||
|
||||
bundle_bytes = archiver._build_archive_bundle(manifest_data, table_payloads)
|
||||
|
||||
with zipfile.ZipFile(io.BytesIO(bundle_bytes), "r") as zf:
|
||||
names = set(zf.namelist())
|
||||
assert "manifest.json" in names
|
||||
for table in archiver.ARCHIVED_TABLES:
|
||||
assert f"{table}.jsonl" in names, f"Missing {table}.jsonl in bundle"
|
||||
|
||||
def test_bundle_missing_table_payload_raises(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
manifest_data = b"{}"
|
||||
incomplete_payloads = {archiver.ARCHIVED_TABLES[0]: b"data"}
|
||||
|
||||
with pytest.raises(ValueError, match="Missing archive payload"):
|
||||
archiver._build_archive_bundle(manifest_data, incomplete_payloads)
|
||||
assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION
|
||||
assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT
|
||||
assert manifest["object_prefix"] == identity.object_prefix
|
||||
assert set(table_payloads) == set(archiver.ARCHIVED_TABLES)
|
||||
assert {stat.table_name for stat in table_stats} == set(archiver.ARCHIVED_TABLES)
|
||||
assert pq.read_table(pa.BufferReader(table_payloads["workflow_runs"])).num_rows == 1
|
||||
|
||||
|
||||
class TestGenerateManifest:
|
||||
@@ -88,25 +168,39 @@ class TestGenerateManifest:
|
||||
run = MagicMock()
|
||||
run.id = str(uuid.uuid4())
|
||||
run.tenant_id = str(uuid.uuid4())
|
||||
run.app_id = str(uuid.uuid4())
|
||||
run.workflow_id = str(uuid.uuid4())
|
||||
run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
|
||||
identity = archiver._build_bundle_identity([run])
|
||||
|
||||
stats = [
|
||||
TableStats(table_name="workflow_runs", row_count=1, checksum="abc123", size_bytes=512),
|
||||
TableStats(table_name="workflow_app_logs", row_count=2, checksum="def456", size_bytes=1024),
|
||||
TableStats(
|
||||
table_name="workflow_runs",
|
||||
row_count=1,
|
||||
checksum="abc123",
|
||||
size_bytes=512,
|
||||
object_key="workflow_runs.parquet",
|
||||
),
|
||||
TableStats(
|
||||
table_name="workflow_node_executions",
|
||||
row_count=2,
|
||||
checksum="def456",
|
||||
size_bytes=1024,
|
||||
object_key="workflow_node_executions.parquet",
|
||||
),
|
||||
]
|
||||
|
||||
manifest = archiver._generate_manifest(run, stats)
|
||||
manifest = archiver._generate_manifest(identity, [run], stats)
|
||||
|
||||
assert manifest["schema_version"] == ARCHIVE_SCHEMA_VERSION
|
||||
assert manifest["workflow_run_id"] == run.id
|
||||
assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION
|
||||
assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT
|
||||
assert manifest["bundle_id"] == identity.bundle_id
|
||||
assert manifest["tenant_id"] == run.tenant_id
|
||||
assert manifest["app_id"] == run.app_id
|
||||
assert manifest["workflow_run_count"] == 1
|
||||
assert manifest["workflow_node_execution_count"] == 2
|
||||
assert manifest["run_ids"] == [run.id]
|
||||
assert "tables" in manifest
|
||||
assert manifest["tables"]["workflow_runs"]["row_count"] == 1
|
||||
assert manifest["tables"]["workflow_runs"]["checksum"] == "abc123"
|
||||
assert manifest["tables"]["workflow_app_logs"]["row_count"] == 2
|
||||
assert manifest["tables"]["workflow_node_executions"]["row_count"] == 2
|
||||
|
||||
|
||||
class TestFilterPaidTenants:
|
||||
@@ -163,6 +257,19 @@ class TestFilterPaidTenants:
|
||||
|
||||
assert result == set()
|
||||
|
||||
def test_planned_paid_tenants_skip_billing_lookup(self):
|
||||
archiver = WorkflowRunArchiver(days=90, paid_tenant_ids=["t1", "t3"])
|
||||
|
||||
with (
|
||||
patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg,
|
||||
patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing,
|
||||
):
|
||||
cfg.BILLING_ENABLED = True
|
||||
result = archiver._filter_paid_tenants({"t1", "t2", "t3"})
|
||||
|
||||
billing.get_plan_bulk_with_cache.assert_not_called()
|
||||
assert result == {"t1", "t3"}
|
||||
|
||||
|
||||
class TestDryRunArchive:
|
||||
@patch("services.retention.workflow_run.archive_paid_plan_workflow_run.get_archive_storage")
|
||||
@@ -175,3 +282,81 @@ class TestDryRunArchive:
|
||||
mock_get_storage.assert_not_called()
|
||||
assert isinstance(summary, ArchiveSummary)
|
||||
assert summary.runs_failed == 0
|
||||
|
||||
def test_dry_run_estimates_table_and_object_sizes(self):
|
||||
archiver = WorkflowRunArchiver(days=90, dry_run=True)
|
||||
run = MagicMock()
|
||||
run.id = "run-1"
|
||||
run.tenant_id = "tenant-1"
|
||||
run.app_id = "app-1"
|
||||
run.workflow_id = "workflow-1"
|
||||
run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
|
||||
table_data = {
|
||||
"workflow_runs": [{"id": "run-1", "tenant_id": "tenant-1"}],
|
||||
"workflow_app_logs": [{"id": "log-1", "workflow_run_id": "run-1"}],
|
||||
}
|
||||
|
||||
with patch.object(archiver, "_extract_bundle_data", return_value=table_data):
|
||||
result = archiver._archive_bundle(MagicMock(), None, [run])
|
||||
|
||||
stats_by_table = {stat.table_name: stat for stat in result.tables}
|
||||
assert result.success is True
|
||||
assert result.object_size_bytes > 0
|
||||
assert stats_by_table["workflow_runs"].row_count == 1
|
||||
assert stats_by_table["workflow_runs"].size_bytes > 0
|
||||
assert stats_by_table["workflow_app_logs"].row_count == 1
|
||||
assert stats_by_table["workflow_app_logs"].size_bytes > 0
|
||||
assert stats_by_table["workflow_node_executions"].row_count == 0
|
||||
assert stats_by_table["workflow_node_executions"].size_bytes > 0
|
||||
|
||||
def test_summary_merges_dry_run_estimates(self):
|
||||
summary = ArchiveSummary()
|
||||
result = MagicMock()
|
||||
result.object_size_bytes = 128
|
||||
result.tables = [
|
||||
MagicMock(table_name="workflow_runs", row_count=1, size_bytes=64),
|
||||
MagicMock(table_name="workflow_app_logs", row_count=2, size_bytes=32),
|
||||
]
|
||||
|
||||
WorkflowRunArchiver._merge_result_stats(summary, result)
|
||||
|
||||
assert summary.total_object_size_bytes == 128
|
||||
assert summary.table_stats["workflow_runs"].row_count == 1
|
||||
assert summary.table_stats["workflow_runs"].size_bytes == 64
|
||||
assert summary.table_stats["workflow_app_logs"].row_count == 2
|
||||
assert summary.table_stats["workflow_app_logs"].size_bytes == 32
|
||||
|
||||
|
||||
class TestArchiveRunIdempotency:
|
||||
def test_locked_bundle_is_skipped(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = MagicMock()
|
||||
run.id = "run-1"
|
||||
run.tenant_id = "tenant-1"
|
||||
run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
|
||||
|
||||
with (
|
||||
patch.object(archiver, "_lock_runs_for_archive", return_value=[]),
|
||||
):
|
||||
storage = MagicMock()
|
||||
storage.object_exists.return_value = False
|
||||
result = archiver._archive_bundle(MagicMock(), storage, [run])
|
||||
|
||||
assert result.success is True
|
||||
assert result.skipped is True
|
||||
assert result.error == "one or more runs locked or deleted by another archiver"
|
||||
|
||||
def test_already_archived_bundle_is_skipped(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = MagicMock()
|
||||
run.id = "run-1"
|
||||
run.tenant_id = "tenant-1"
|
||||
run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
|
||||
storage = MagicMock()
|
||||
storage.object_exists.return_value = True
|
||||
|
||||
result = archiver._archive_bundle(MagicMock(), storage, [run])
|
||||
|
||||
assert result.success is True
|
||||
assert result.skipped is True
|
||||
assert result.error == "bundle already archived"
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user