mirror of
https://github.com/langgenius/dify.git
synced 2026-06-24 13:01:16 +08:00
merge main
This commit is contained in:
commit
8f78ce3f52
@ -37,12 +37,16 @@ Use this as the decision guide for React/TypeScript component structure. Existin
|
||||
- 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.
|
||||
@ -51,7 +55,11 @@ Use this as the decision guide for React/TypeScript component structure. Existin
|
||||
- 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
|
||||
|
||||
@ -74,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.
|
||||
@ -93,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`.
|
||||
|
||||
@ -110,6 +124,7 @@ 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, 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.
|
||||
|
||||
@ -120,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.
|
||||
|
||||
|
||||
21
.github/workflows/build-push.yml
vendored
21
.github/workflows/build-push.yml
vendored
@ -21,6 +21,7 @@ env:
|
||||
DIFY_WEB_IMAGE_NAME: ${{ vars.DIFY_WEB_IMAGE_NAME || 'langgenius/dify-web' }}
|
||||
DIFY_API_IMAGE_NAME: ${{ vars.DIFY_API_IMAGE_NAME || 'langgenius/dify-api' }}
|
||||
DIFY_AGENT_IMAGE_NAME: ${{ vars.DIFY_AGENT_IMAGE_NAME || 'langgenius/dify-agent-backend' }}
|
||||
DIFY_AGENT_LOCAL_SANDBOX_IMAGE_NAME: ${{ vars.DIFY_AGENT_LOCAL_SANDBOX_IMAGE_NAME || 'langgenius/dify-agent-local-sandbox' }}
|
||||
|
||||
jobs:
|
||||
build:
|
||||
@ -74,6 +75,20 @@ jobs:
|
||||
file: "dify-agent/Dockerfile"
|
||||
platform: linux/arm64
|
||||
runs_on: depot-ubuntu-24.04-4
|
||||
- service_name: "build-agent-local-sandbox-amd64"
|
||||
image_name_env: "DIFY_AGENT_LOCAL_SANDBOX_IMAGE_NAME"
|
||||
artifact_context: "local-sandbox"
|
||||
build_context: "{{defaultContext}}:dify-agent"
|
||||
file: "docker/local-sandbox/Dockerfile"
|
||||
platform: linux/amd64
|
||||
runs_on: depot-ubuntu-24.04-4
|
||||
- service_name: "build-agent-local-sandbox-arm64"
|
||||
image_name_env: "DIFY_AGENT_LOCAL_SANDBOX_IMAGE_NAME"
|
||||
artifact_context: "local-sandbox"
|
||||
build_context: "{{defaultContext}}:dify-agent"
|
||||
file: "docker/local-sandbox/Dockerfile"
|
||||
platform: linux/arm64
|
||||
runs_on: depot-ubuntu-24.04-4
|
||||
|
||||
steps:
|
||||
- name: Prepare
|
||||
@ -139,6 +154,9 @@ jobs:
|
||||
- service_name: "validate-agent-amd64"
|
||||
build_context: "{{defaultContext}}"
|
||||
file: "dify-agent/Dockerfile"
|
||||
- service_name: "validate-agent-local-sandbox-amd64"
|
||||
build_context: "{{defaultContext}}:dify-agent"
|
||||
file: "docker/local-sandbox/Dockerfile"
|
||||
steps:
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@d7f5e7f509e45cec5c76c4d5afdd7de93d0b3df5 # v4.1.0
|
||||
@ -167,6 +185,9 @@ jobs:
|
||||
- service_name: "merge-agent-images"
|
||||
image_name_env: "DIFY_AGENT_IMAGE_NAME"
|
||||
context: "agent"
|
||||
- service_name: "merge-agent-local-sandbox-images"
|
||||
image_name_env: "DIFY_AGENT_LOCAL_SANDBOX_IMAGE_NAME"
|
||||
context: "local-sandbox"
|
||||
steps:
|
||||
- name: Download digests
|
||||
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
|
||||
|
||||
@ -78,6 +78,13 @@ def _filter_snapshot_to_specs(
|
||||
return CompositorSessionSnapshot(schema_version=snapshot.schema_version, layers=filtered_layers)
|
||||
|
||||
|
||||
def _shell_layer_deps(*, include_drive: bool) -> dict[str, str]:
|
||||
deps = {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID}
|
||||
if include_drive:
|
||||
deps["drive"] = DIFY_DRIVE_LAYER_ID
|
||||
return deps
|
||||
|
||||
|
||||
class AgentBackendModelConfig(BaseModel):
|
||||
"""API-side model/plugin selection before it is converted to Dify Agent layers."""
|
||||
|
||||
@ -263,6 +270,7 @@ class AgentBackendRunRequestBuilder:
|
||||
RunLayerSpec(
|
||||
name=DIFY_DRIVE_LAYER_ID,
|
||||
type=DIFY_DRIVE_LAYER_TYPE_ID,
|
||||
deps={"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID},
|
||||
metadata=run_input.metadata,
|
||||
config=run_input.drive_config,
|
||||
)
|
||||
@ -329,14 +337,15 @@ class AgentBackendRunRequestBuilder:
|
||||
)
|
||||
|
||||
if run_input.include_shell:
|
||||
# Sandboxed bash workspace (dify.shell). Depends on execution_context so
|
||||
# the agent server can mint per-command Agent Stub env (back proxy);
|
||||
# Sandboxed bash workspace (dify.shell). Depends on execution_context
|
||||
# so the agent server can mint per-command Agent Stub env, and on
|
||||
# drive when present so that env points at /mnt/drive/<drive_ref>.
|
||||
# shellctl connection itself is server-injected.
|
||||
layers.append(
|
||||
RunLayerSpec(
|
||||
name=DIFY_SHELL_LAYER_ID,
|
||||
type=DIFY_SHELL_LAYER_TYPE_ID,
|
||||
deps={"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID},
|
||||
deps=_shell_layer_deps(include_drive=run_input.drive_config is not None),
|
||||
metadata=run_input.metadata,
|
||||
config=run_input.shell_config or DifyShellLayerConfig(),
|
||||
)
|
||||
@ -460,6 +469,7 @@ class AgentBackendRunRequestBuilder:
|
||||
RunLayerSpec(
|
||||
name=DIFY_DRIVE_LAYER_ID,
|
||||
type=DIFY_DRIVE_LAYER_TYPE_ID,
|
||||
deps={"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID},
|
||||
metadata=run_input.metadata,
|
||||
config=run_input.drive_config,
|
||||
)
|
||||
@ -528,14 +538,15 @@ class AgentBackendRunRequestBuilder:
|
||||
)
|
||||
|
||||
if run_input.include_shell:
|
||||
# Sandboxed bash workspace (dify.shell). Depends on execution_context so
|
||||
# the agent server can mint per-command Agent Stub env (back proxy);
|
||||
# Sandboxed bash workspace (dify.shell). Depends on execution_context
|
||||
# so the agent server can mint per-command Agent Stub env, and on
|
||||
# drive when present so that env points at /mnt/drive/<drive_ref>.
|
||||
# shellctl connection itself is server-injected.
|
||||
layers.append(
|
||||
RunLayerSpec(
|
||||
name=DIFY_SHELL_LAYER_ID,
|
||||
type=DIFY_SHELL_LAYER_TYPE_ID,
|
||||
deps={"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID},
|
||||
deps=_shell_layer_deps(include_drive=run_input.drive_config is not None),
|
||||
metadata=run_input.metadata,
|
||||
config=run_input.shell_config or DifyShellLayerConfig(),
|
||||
)
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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]:
|
||||
|
||||
@ -104,7 +104,7 @@ class WorkflowAgentComposerValidateApi(Resource):
|
||||
@with_current_tenant_id
|
||||
def post(self, tenant_id: str, app_model: App, node_id: str):
|
||||
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
|
||||
ComposerConfigValidator.validate_save_payload(payload)
|
||||
ComposerConfigValidator.validate_publish_payload(payload)
|
||||
findings = AgentComposerService.collect_validation_findings(
|
||||
tenant_id=tenant_id,
|
||||
payload=payload,
|
||||
@ -238,7 +238,7 @@ class AgentComposerValidateApi(Resource):
|
||||
def post(self, tenant_id: str, agent_id: UUID):
|
||||
_resolve_agent_app_id(tenant_id=tenant_id, agent_id=agent_id)
|
||||
payload = ComposerSavePayload.model_validate(console_ns.payload or {})
|
||||
ComposerConfigValidator.validate_save_payload(payload)
|
||||
ComposerConfigValidator.validate_publish_payload(payload)
|
||||
findings = AgentComposerService.collect_validation_findings(
|
||||
tenant_id=tenant_id,
|
||||
payload=payload,
|
||||
|
||||
@ -14,7 +14,6 @@ from controllers.console.app.app import (
|
||||
)
|
||||
from controllers.console.app.app import (
|
||||
AppListQuery,
|
||||
CopyAppPayload,
|
||||
_normalize_app_list_query_args,
|
||||
)
|
||||
from controllers.console.app.app import (
|
||||
@ -110,6 +109,25 @@ 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")
|
||||
|
||||
@ -228,6 +246,10 @@ class AgentAppDetailWithSite(GenericAppDetailWithSite):
|
||||
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")
|
||||
@ -238,8 +260,8 @@ register_schema_models(
|
||||
console_ns,
|
||||
AgentAppCreatePayload,
|
||||
AgentAppUpdatePayload,
|
||||
AgentAppCopyPayload,
|
||||
AgentApiStatusPayload,
|
||||
CopyAppPayload,
|
||||
AgentInviteOptionsQuery,
|
||||
AgentLogsQuery,
|
||||
AgentStatisticsQuery,
|
||||
@ -254,6 +276,7 @@ register_response_schema_models(
|
||||
AgentAppPublishedReferenceResponse,
|
||||
AgentAppDetailWithSite,
|
||||
AgentAppPartial,
|
||||
AgentDebugConversationRefreshResponse,
|
||||
AgentConfigSnapshotDetailResponse,
|
||||
AgentConfigSnapshotListResponse,
|
||||
AgentConfigSnapshotRestoreResponse,
|
||||
@ -535,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")
|
||||
@ -548,13 +596,14 @@ 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,
|
||||
|
||||
@ -1,4 +1,3 @@
|
||||
import logging
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
@ -30,7 +29,6 @@ from fields.base import ResponseModel
|
||||
from libs.helper import uuid_value
|
||||
from libs.login import login_required
|
||||
from models import Account
|
||||
from models.agent_config_entities import AgentFileRefConfig, AgentSkillRefConfig
|
||||
from models.model import App, AppMode, UploadFile
|
||||
from services.agent.composer_service import AgentComposerService
|
||||
from services.agent.skill_package_service import SkillManifest, SkillPackageError
|
||||
@ -49,8 +47,6 @@ from services.agent_drive_service import (
|
||||
)
|
||||
from services.agent_service import AgentService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_WORKFLOW_AGENT_DRIVE_APP_MODES = [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]
|
||||
_AGENT_SKILL_UPLOAD_PARAMS = {
|
||||
"file": {
|
||||
@ -130,8 +126,16 @@ class AgentLogResponse(ResponseModel):
|
||||
files: list[Any] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentUploadedSkillResponse(ResponseModel):
|
||||
name: str
|
||||
description: str
|
||||
path: str
|
||||
skill_md_key: str
|
||||
archive_key: str | None = None
|
||||
|
||||
|
||||
class AgentSkillUploadResponse(ResponseModel):
|
||||
skill: AgentSkillRefConfig
|
||||
skill: AgentUploadedSkillResponse
|
||||
manifest: SkillManifest
|
||||
|
||||
|
||||
@ -145,13 +149,11 @@ class AgentDriveFileResponse(ResponseModel):
|
||||
|
||||
class AgentDriveFileCommitResponse(ResponseModel):
|
||||
file: AgentDriveFileResponse
|
||||
config_version_id: str | None = None
|
||||
|
||||
|
||||
class AgentDriveDeleteResponse(ResponseModel):
|
||||
result: str
|
||||
removed_keys: list[str] = Field(default_factory=list)
|
||||
config_version_id: str | None = None
|
||||
|
||||
|
||||
register_schema_models(console_ns, AgentLogQuery, AgentDriveFilePayload, AgentDriveDeleteFileByAgentQuery)
|
||||
@ -161,6 +163,7 @@ register_response_schema_models(
|
||||
AgentDriveFileCommitResponse,
|
||||
AgentDriveFileResponse,
|
||||
AgentLogResponse,
|
||||
AgentUploadedSkillResponse,
|
||||
AgentSkillUploadResponse,
|
||||
SkillToolInferenceResult,
|
||||
)
|
||||
@ -242,24 +245,6 @@ def _commit_drive_file_for_app(*, current_user: Account, app_model: App, allow_n
|
||||
return {"code": exc.code, "message": exc.message}, exc.status_code
|
||||
|
||||
row = committed[0]
|
||||
file_ref = AgentFileRefConfig.model_validate(
|
||||
{
|
||||
"id": row["key"],
|
||||
"name": upload_file.name,
|
||||
"file_id": upload_file.id,
|
||||
"drive_key": row["key"],
|
||||
"type": row.get("mime_type"),
|
||||
"size": row.get("size"),
|
||||
}
|
||||
)
|
||||
config_version_id = AgentComposerService.add_drive_file_ref(
|
||||
tenant_id=app_model.tenant_id,
|
||||
agent_id=agent_id,
|
||||
account_id=current_user.id,
|
||||
file_ref=file_ref,
|
||||
app_id=app_model.id,
|
||||
node_id=node_id,
|
||||
)
|
||||
return {
|
||||
"file": {
|
||||
"name": upload_file.name,
|
||||
@ -268,7 +253,6 @@ def _commit_drive_file_for_app(*, current_user: Account, app_model: App, allow_n
|
||||
"size": row.get("size"),
|
||||
"mime_type": row.get("mime_type"),
|
||||
},
|
||||
"config_version_id": config_version_id,
|
||||
}, 201
|
||||
|
||||
|
||||
@ -283,24 +267,17 @@ def _delete_drive_file_for_app(*, current_user: Account, app_model: App, allow_n
|
||||
except AgentDriveError as exc:
|
||||
return {"code": exc.code, "message": exc.message}, exc.status_code
|
||||
|
||||
config_version_id = AgentComposerService.remove_drive_refs(
|
||||
tenant_id=app_model.tenant_id,
|
||||
agent_id=agent_id,
|
||||
account_id=current_user.id,
|
||||
file_key=key,
|
||||
app_id=app_model.id,
|
||||
node_id=node_id,
|
||||
)
|
||||
removed_keys: list[str] = []
|
||||
try:
|
||||
removed_keys = AgentDriveService().delete(tenant_id=app_model.tenant_id, agent_id=agent_id, key=key)
|
||||
result = AgentDriveService().commit(
|
||||
tenant_id=app_model.tenant_id,
|
||||
user_id=current_user.id,
|
||||
agent_id=agent_id,
|
||||
items=[DriveCommitItem(key=key, file_ref=None)],
|
||||
)
|
||||
except AgentDriveError as exc:
|
||||
return {"code": exc.code, "message": exc.message}, exc.status_code
|
||||
except Exception:
|
||||
# Soul-first ordering: the ref is already gone; orphan KV rows are
|
||||
# harmless and an idempotent DELETE retry cleans them.
|
||||
logger.exception("agent drive delete failed for key %s (soul already updated)", key)
|
||||
return {"result": "success", "removed_keys": removed_keys, "config_version_id": config_version_id}
|
||||
removed_keys = [item["key"] for item in result if item.get("removed")]
|
||||
return {"result": "success", "removed_keys": removed_keys}
|
||||
|
||||
|
||||
def _delete_skill_for_app(*, current_user: Account, app_model: App, slug: str, allow_node_id: bool = True):
|
||||
@ -312,22 +289,20 @@ def _delete_skill_for_app(*, current_user: Account, app_model: App, slug: str, a
|
||||
if "/" in slug or not slug.strip():
|
||||
return {"code": "drive_key_invalid", "message": "skill slug must be a single path segment"}, 400
|
||||
|
||||
config_version_id = AgentComposerService.remove_drive_refs(
|
||||
tenant_id=app_model.tenant_id,
|
||||
agent_id=agent_id,
|
||||
account_id=current_user.id,
|
||||
skill_slug=slug,
|
||||
app_id=app_model.id,
|
||||
node_id=node_id,
|
||||
)
|
||||
removed_keys: list[str] = []
|
||||
try:
|
||||
removed_keys = AgentDriveService().delete(tenant_id=app_model.tenant_id, agent_id=agent_id, prefix=f"{slug}/")
|
||||
result = AgentDriveService().commit(
|
||||
tenant_id=app_model.tenant_id,
|
||||
user_id=current_user.id,
|
||||
agent_id=agent_id,
|
||||
items=[
|
||||
DriveCommitItem(key=f"{slug}/SKILL.md", file_ref=None),
|
||||
DriveCommitItem(key=f"{slug}/.DIFY-SKILL-FULL.zip", file_ref=None),
|
||||
],
|
||||
)
|
||||
except AgentDriveError as exc:
|
||||
return {"code": exc.code, "message": exc.message}, exc.status_code
|
||||
except Exception:
|
||||
logger.exception("agent drive delete failed for skill %s (soul already updated)", slug)
|
||||
return {"result": "success", "removed_keys": removed_keys, "config_version_id": config_version_id}
|
||||
removed_keys = [item["key"] for item in result if item.get("removed")]
|
||||
return {"result": "success", "removed_keys": removed_keys}
|
||||
|
||||
|
||||
def _infer_skill_tools_for_app(*, app_model: App, slug: str):
|
||||
@ -455,7 +430,7 @@ class AgentDriveFilesApi(Resource):
|
||||
return _commit_drive_file_for_app(current_user=current_user, app_model=app_model)
|
||||
|
||||
@console_ns.doc("delete_agent_drive_file")
|
||||
@console_ns.doc(description="Delete one drive file by key; soul ref first, then the KV row (ENG-625 D5)")
|
||||
@console_ns.doc(description="Delete one drive file by key via drive commit-null semantics")
|
||||
@console_ns.doc(params={"app_id": "Application ID", **query_params_from_model(AgentDriveDeleteFileQuery)})
|
||||
@console_ns.response(200, "File removed", console_ns.models[AgentDriveDeleteResponse.__name__])
|
||||
@setup_required
|
||||
@ -486,9 +461,7 @@ class AgentSkillByAgentApi(Resource):
|
||||
@console_ns.route("/apps/<uuid:app_id>/agent/skills/<string:slug>")
|
||||
class AgentSkillApi(Resource):
|
||||
@console_ns.doc("delete_agent_skill")
|
||||
@console_ns.doc(
|
||||
description="Delete a standardized skill: soul ref first, then the <slug>/ drive prefix (ENG-625 D5)"
|
||||
)
|
||||
@console_ns.doc(description="Delete a standardized skill by removing its known drive keys via commit-null")
|
||||
@console_ns.doc(
|
||||
params={
|
||||
"app_id": "Application ID",
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -80,8 +80,11 @@ class SnippetDraftConfigResponse(BaseModel):
|
||||
parallel_depth_limit: int
|
||||
|
||||
|
||||
class SnippetWorkflowPaginationResponse(WorkflowPaginationResponse):
|
||||
class SnippetWorkflowPaginationResponse(BaseModel):
|
||||
items: list[SnippetWorkflowResponse]
|
||||
page: int
|
||||
limit: int
|
||||
has_more: bool
|
||||
|
||||
|
||||
register_schema_models(
|
||||
|
||||
@ -1,9 +1,10 @@
|
||||
"""Inner API endpoint for tenant-scoped knowledge retrieval.
|
||||
"""Plugin inner API endpoint for tenant-scoped knowledge retrieval.
|
||||
|
||||
This controller is a thin HTTP wrapper around
|
||||
``services.knowledge_retrieval_inner_service.InnerKnowledgeRetrievalService``.
|
||||
It intentionally keeps authorization simple: shared inner API key plus
|
||||
tenant-scoped app/dataset validation in the service layer.
|
||||
It uses the plugin inner API key because dify-agent calls this endpoint through
|
||||
the same trusted Dify API bridge as other agent/plugin inner calls; tenant-scoped
|
||||
app/dataset validation remains in the service layer.
|
||||
"""
|
||||
|
||||
from flask_restx import Resource
|
||||
@ -11,7 +12,7 @@ from pydantic import ValidationError
|
||||
|
||||
from controllers.common.schema import register_response_schema_models, register_schema_models
|
||||
from controllers.inner_api import inner_api_ns
|
||||
from controllers.inner_api.wraps import inner_api_only
|
||||
from controllers.inner_api.wraps import plugin_inner_api_only
|
||||
from core.workflow.nodes.knowledge_retrieval import exc as retrieval_exc
|
||||
from libs.exception import BaseHTTPException
|
||||
from services.entities.knowledge_retrieval_inner import InnerKnowledgeRetrieveRequest, InnerKnowledgeRetrieveResponse
|
||||
@ -48,7 +49,7 @@ register_response_schema_models(inner_api_ns, InnerKnowledgeRetrieveResponse)
|
||||
class InnerKnowledgeRetrieveApi(Resource):
|
||||
"""Retrieve knowledge from one or more datasets within the caller tenant."""
|
||||
|
||||
@inner_api_only
|
||||
@plugin_inner_api_only
|
||||
@inner_api_ns.doc("inner_knowledge_retrieve")
|
||||
@inner_api_ns.doc(description="Retrieve knowledge for trusted internal callers")
|
||||
@inner_api_ns.expect(inner_api_ns.models[InnerKnowledgeRetrieveRequest.__name__])
|
||||
@ -60,9 +61,8 @@ class InnerKnowledgeRetrieveApi(Resource):
|
||||
@inner_api_ns.doc(
|
||||
responses={
|
||||
400: "Invalid request body",
|
||||
401: "Unauthorized - invalid inner API key",
|
||||
403: "Caller tenant does not own the requested resource",
|
||||
404: "App or dataset not found",
|
||||
404: "Invalid plugin inner API key, app not found, or dataset not found",
|
||||
422: "Invalid retrieval configuration",
|
||||
429: "Knowledge retrieval rate limited",
|
||||
502: "External knowledge retrieval failed",
|
||||
|
||||
@ -1,10 +1,12 @@
|
||||
"""Inner API for the agent drive (agent 网盘) control plane — ENG-591.
|
||||
"""Inner API for the agent drive (agent 网盘) control plane.
|
||||
|
||||
Two endpoints, called by the dify-agent server (not the sandbox) with the inner
|
||||
API key. The drive ref is the URL segment ``agent-<agent_id>``; the path-like
|
||||
file key travels in the query/body, never as a URL path segment (so its ``/``
|
||||
characters do not collide with routing). Drive-owned semantics: tenant scoped,
|
||||
no user-level FileAccessScope.
|
||||
These endpoints are called by the dify-agent server (not the sandbox) with the
|
||||
inner API key. The drive ref is the URL segment ``agent-<agent_id>``; the
|
||||
path-like file key travels in the query/body, never as a URL path segment (so
|
||||
its ``/`` characters do not collide with routing). Drive-owned semantics:
|
||||
tenant scoped, no user-level FileAccessScope. Commit still canonicalizes the
|
||||
trusted execution-context user through the same EndUser lookup as plugin file
|
||||
upload before validating ToolFile ownership.
|
||||
"""
|
||||
|
||||
from flask import request
|
||||
@ -13,6 +15,7 @@ from pydantic import BaseModel, ValidationError
|
||||
|
||||
from controllers.console.wraps import setup_required
|
||||
from controllers.inner_api import inner_api_ns
|
||||
from controllers.inner_api.plugin.wraps import get_user
|
||||
from controllers.inner_api.wraps import plugin_inner_api_only
|
||||
from services.agent_drive_service import (
|
||||
AgentDriveError,
|
||||
@ -56,6 +59,24 @@ class AgentDriveManifestApi(Resource):
|
||||
return {"items": items}
|
||||
|
||||
|
||||
@inner_api_ns.route("/drive/<string:drive_ref>/skills")
|
||||
class AgentDriveSkillsApi(Resource):
|
||||
@setup_required
|
||||
@plugin_inner_api_only
|
||||
@inner_api_ns.doc("agent_drive_skills")
|
||||
@inner_api_ns.doc(description="List the skill catalog of an agent drive")
|
||||
def get(self, drive_ref: str):
|
||||
try:
|
||||
agent_id = parse_agent_drive_ref(drive_ref)
|
||||
tenant_id = (request.args.get("tenant_id") or "").strip()
|
||||
if not tenant_id:
|
||||
raise AgentDriveError("missing_tenant_id", "tenant_id is required", status_code=400)
|
||||
items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=agent_id)
|
||||
except AgentDriveError as exc:
|
||||
return _error_response(exc)
|
||||
return {"items": items}
|
||||
|
||||
|
||||
@inner_api_ns.route("/drive/<string:drive_ref>/commit")
|
||||
class AgentDriveCommitApi(Resource):
|
||||
@setup_required
|
||||
@ -69,9 +90,10 @@ class AgentDriveCommitApi(Resource):
|
||||
body = _CommitRequest.model_validate(request.get_json(silent=True) or {})
|
||||
except ValidationError as exc:
|
||||
raise AgentDriveError("invalid_request", str(exc), status_code=400) from exc
|
||||
user = get_user(body.tenant_id, body.user_id)
|
||||
items = AgentDriveService().commit(
|
||||
tenant_id=body.tenant_id,
|
||||
user_id=body.user_id,
|
||||
user_id=user.id,
|
||||
agent_id=agent_id,
|
||||
items=body.items,
|
||||
)
|
||||
|
||||
@ -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
|
||||
@ -279,12 +304,12 @@ 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)
|
||||
|
||||
|
||||
@ -335,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)
|
||||
|
||||
|
||||
|
||||
@ -16,6 +16,7 @@ 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,
|
||||
@ -37,6 +38,11 @@ from services.app_service import AppListParams, AppService
|
||||
_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": [],
|
||||
@ -171,6 +177,8 @@ class AppListApi(Resource):
|
||||
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(
|
||||
@ -223,6 +231,7 @@ class AppListApi(Resource):
|
||||
workspace_name=tenant_name,
|
||||
)
|
||||
for r in pagination.items
|
||||
if _is_listable(r)
|
||||
]
|
||||
|
||||
env = AppListResponse(
|
||||
|
||||
@ -37,6 +37,7 @@ from core.workflow.nodes.agent_v2.plugin_tools_builder import (
|
||||
from core.workflow.nodes.agent_v2.runtime_request_builder import (
|
||||
append_runtime_warnings,
|
||||
build_ask_human_layer_config,
|
||||
build_drive_aware_soul_mention_resolver,
|
||||
build_drive_layer_config,
|
||||
build_knowledge_layer_config,
|
||||
build_shell_layer_config,
|
||||
@ -123,9 +124,19 @@ class AgentAppRuntimeRequestBuilder:
|
||||
}
|
||||
|
||||
drive_config = None
|
||||
soul_prompt_resolver = build_soul_mention_resolver(agent_soul)
|
||||
if dify_config.AGENT_DRIVE_MANIFEST_ENABLED:
|
||||
drive_config, drive_warnings = build_drive_layer_config(agent_soul, agent_id=context.agent_id)
|
||||
drive_config, drive_warnings = build_drive_layer_config(
|
||||
agent_soul,
|
||||
tenant_id=context.dify_context.tenant_id,
|
||||
agent_id=context.agent_id,
|
||||
)
|
||||
append_runtime_warnings(metadata, drive_warnings)
|
||||
soul_prompt_resolver = build_drive_aware_soul_mention_resolver(
|
||||
agent_soul,
|
||||
tenant_id=context.dify_context.tenant_id,
|
||||
agent_id=context.agent_id,
|
||||
)
|
||||
knowledge_config = build_knowledge_layer_config(agent_soul)
|
||||
|
||||
request = self._request_builder.build_for_agent_app(
|
||||
@ -154,9 +165,7 @@ class AgentAppRuntimeRequestBuilder:
|
||||
),
|
||||
# ENG-616: expand slash-menu mention tokens to canonical names so
|
||||
# no frontend-internal {{#…#}} marker ever reaches the model.
|
||||
agent_soul_prompt=expand_prompt_mentions(
|
||||
agent_soul.prompt.system_prompt, build_soul_mention_resolver(agent_soul)
|
||||
).strip()
|
||||
agent_soul_prompt=expand_prompt_mentions(agent_soul.prompt.system_prompt, soul_prompt_resolver).strip()
|
||||
or None,
|
||||
user_prompt=context.user_query,
|
||||
tools=tools_layer,
|
||||
|
||||
@ -16,9 +16,6 @@ SUPPORTED_AGENT_BACKEND_FEATURES = frozenset(
|
||||
"knowledge",
|
||||
"env",
|
||||
"sandbox",
|
||||
# ENG-623: exposed at runtime as the dify.drive declaration layer
|
||||
# (an index the agent pulls through the back proxy).
|
||||
"skills_files",
|
||||
# ENG-635: human involvement is exposed at runtime as the dify.ask_human
|
||||
# deferred tool; a call pauses via the existing HITL form mechanism.
|
||||
"human",
|
||||
@ -32,11 +29,7 @@ RESERVED_AGENT_BACKEND_FEATURES = frozenset(
|
||||
)
|
||||
|
||||
|
||||
def build_runtime_feature_manifest(
|
||||
agent_soul: AgentSoulConfig,
|
||||
*,
|
||||
drive_manifest_enabled: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
def build_runtime_feature_manifest(agent_soul: AgentSoulConfig) -> dict[str, Any]:
|
||||
"""Describe PRD capabilities supported by or still reserved from Agent backend runtime."""
|
||||
warnings: list[dict[str, str]] = []
|
||||
soul_dump = agent_soul.model_dump(mode="json", exclude_none=True, exclude_defaults=True)
|
||||
@ -54,38 +47,10 @@ def build_runtime_feature_manifest(
|
||||
}
|
||||
)
|
||||
|
||||
has_skills_files = bool(agent_soul.skills_files.skills or agent_soul.skills_files.files)
|
||||
if has_skills_files and not drive_manifest_enabled:
|
||||
warnings.append(
|
||||
{
|
||||
"section": "agent_soul.skills_files",
|
||||
"code": "drive_manifest_disabled",
|
||||
"message": (
|
||||
"skills_files is configured but AGENT_DRIVE_MANIFEST_ENABLED is off; "
|
||||
"the drive declaration layer is not injected into this run."
|
||||
),
|
||||
}
|
||||
)
|
||||
for skill in agent_soul.skills_files.skills:
|
||||
if not skill.skill_md_key:
|
||||
warnings.append(
|
||||
{
|
||||
"section": "agent_soul.skills_files",
|
||||
"code": "skill_ref_dangling",
|
||||
"message": (
|
||||
f"skill_ref_dangling: skill '{skill.name or skill.id or 'unknown'}' has no drive key; "
|
||||
"re-standardize it to expose it at runtime."
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
reserved_status = dict.fromkeys(sorted(RESERVED_AGENT_BACKEND_FEATURES), "reserved_not_executed")
|
||||
reserved_status["knowledge"] = (
|
||||
"supported_by_knowledge_layer" if list_configured_knowledge_dataset_ids(agent_soul) else "not_configured"
|
||||
)
|
||||
reserved_status["skills_files"] = (
|
||||
"supported_by_drive_manifest" if drive_manifest_enabled else "drive_manifest_disabled"
|
||||
)
|
||||
reserved_status["tools.dify_tools"] = "supported_when_config_valid"
|
||||
reserved_status["tools.cli_tools"] = "supported_by_shell_bootstrap"
|
||||
reserved_status["env"] = "supported_by_shell_bootstrap"
|
||||
|
||||
@ -7,7 +7,6 @@ from typing import Any, Literal, Protocol, assert_never, cast
|
||||
from agenton.compositor import CompositorSessionSnapshot
|
||||
from dify_agent.layers.ask_human import DifyAskHumanLayerConfig
|
||||
from dify_agent.layers.drive import (
|
||||
DifyDriveFileConfig,
|
||||
DifyDriveLayerConfig,
|
||||
DifyDriveSkillConfig,
|
||||
)
|
||||
@ -55,10 +54,13 @@ from models.agent_config_entities import (
|
||||
)
|
||||
from models.provider_ids import ModelProviderID
|
||||
from services.agent.prompt_mentions import (
|
||||
MentionKind,
|
||||
build_node_job_mention_resolver,
|
||||
build_soul_mention_resolver,
|
||||
expand_prompt_mentions,
|
||||
parse_prompt_mentions,
|
||||
)
|
||||
from services.agent_drive_service import AgentDriveService, decode_drive_mention_ref
|
||||
|
||||
from .output_failure_orchestrator import retry_idempotency_key
|
||||
from .plugin_tools_builder import WorkflowAgentPluginToolsBuilder, WorkflowAgentPluginToolsBuildError
|
||||
@ -153,9 +155,6 @@ class WorkflowAgentRuntimeRequestBuilder:
|
||||
expand_prompt_mentions(node_job.workflow_prompt, build_node_job_mention_resolver(node_job)).strip()
|
||||
or "Run this workflow Agent Node for the current run."
|
||||
)
|
||||
soul_prompt = expand_prompt_mentions(
|
||||
agent_soul.prompt.system_prompt, build_soul_mention_resolver(agent_soul)
|
||||
).strip()
|
||||
user_prompt = workflow_context_prompt.strip() or "Use the current workflow context."
|
||||
credentials = self._credentials_provider.fetch(agent_soul.model.model_provider, agent_soul.model.model)
|
||||
try:
|
||||
@ -182,9 +181,20 @@ class WorkflowAgentRuntimeRequestBuilder:
|
||||
}
|
||||
|
||||
drive_config: DifyDriveLayerConfig | None = None
|
||||
soul_prompt_resolver = build_soul_mention_resolver(agent_soul)
|
||||
if dify_config.AGENT_DRIVE_MANIFEST_ENABLED:
|
||||
drive_config, drive_warnings = build_drive_layer_config(agent_soul, agent_id=context.agent.id)
|
||||
drive_config, drive_warnings = build_drive_layer_config(
|
||||
agent_soul,
|
||||
tenant_id=context.dify_context.tenant_id,
|
||||
agent_id=context.agent.id,
|
||||
)
|
||||
append_runtime_warnings(metadata, drive_warnings)
|
||||
soul_prompt_resolver = build_drive_aware_soul_mention_resolver(
|
||||
agent_soul,
|
||||
tenant_id=context.dify_context.tenant_id,
|
||||
agent_id=context.agent.id,
|
||||
)
|
||||
soul_prompt = expand_prompt_mentions(agent_soul.prompt.system_prompt, soul_prompt_resolver).strip()
|
||||
knowledge_config = build_knowledge_layer_config(agent_soul)
|
||||
|
||||
request = self._request_builder.build_for_workflow_node(
|
||||
@ -292,10 +302,7 @@ class WorkflowAgentRuntimeRequestBuilder:
|
||||
"agent_config_snapshot_id": context.snapshot.id,
|
||||
"binding_id": context.binding.id,
|
||||
"workflow_node_job_mode": node_job.mode.value,
|
||||
"runtime_support": build_runtime_feature_manifest(
|
||||
agent_soul,
|
||||
drive_manifest_enabled=dify_config.AGENT_DRIVE_MANIFEST_ENABLED,
|
||||
),
|
||||
"runtime_support": build_runtime_feature_manifest(agent_soul),
|
||||
}
|
||||
|
||||
def _build_workflow_context_prompt(
|
||||
@ -603,76 +610,107 @@ def append_runtime_warnings(metadata: dict[str, Any], warnings: list[dict[str, s
|
||||
existing.extend(warnings)
|
||||
|
||||
|
||||
def build_drive_aware_soul_mention_resolver(
|
||||
agent_soul: AgentSoulConfig,
|
||||
*,
|
||||
tenant_id: str,
|
||||
agent_id: str,
|
||||
):
|
||||
"""Resolve skill/file mentions against the agent drive and everything else via Agent Soul."""
|
||||
|
||||
base_resolver = build_soul_mention_resolver(agent_soul)
|
||||
drive_service = AgentDriveService()
|
||||
skill_catalog = drive_service.list_skills(tenant_id=tenant_id, agent_id=agent_id)
|
||||
skill_names_by_key = {skill["skill_md_key"]: skill["name"] for skill in skill_catalog}
|
||||
drive_keys = {item["key"] for item in drive_service.manifest(tenant_id=tenant_id, agent_id=agent_id)}
|
||||
|
||||
def _resolve(mention: object) -> str | None:
|
||||
if not hasattr(mention, "kind") or not hasattr(mention, "ref_id"):
|
||||
return None
|
||||
kind = cast(MentionKind, mention.kind)
|
||||
ref_id = cast(str, mention.ref_id)
|
||||
label = cast(str | None, getattr(mention, "label", None))
|
||||
if kind == MentionKind.SKILL:
|
||||
decoded_key = decode_drive_mention_ref(ref_id)
|
||||
return skill_names_by_key.get(decoded_key) or label or decoded_key
|
||||
if kind == MentionKind.FILE:
|
||||
decoded_key = decode_drive_mention_ref(ref_id)
|
||||
if decoded_key in drive_keys:
|
||||
return decoded_key.rsplit("/", 1)[-1]
|
||||
return label or decoded_key
|
||||
return base_resolver(cast(Any, mention))
|
||||
|
||||
return _resolve
|
||||
|
||||
|
||||
def build_drive_layer_config(
|
||||
agent_soul: AgentSoulConfig,
|
||||
*,
|
||||
tenant_id: str,
|
||||
agent_id: str | None,
|
||||
) -> tuple[DifyDriveLayerConfig | None, list[dict[str, str]]]:
|
||||
"""Catalog the soul's drive-backed Skills & Files into the dify.drive declaration.
|
||||
"""Derive drive runtime catalog + prompt-mentioned eager-pull keys from the drive."""
|
||||
|
||||
Returns ``(config, warnings)`` — ``config is None`` means nothing to inject
|
||||
(no skills/files configured, or no agent identity to address the drive by).
|
||||
Refs that predate standardization (no drive key) are skipped with a warning
|
||||
instead of failing the run, so historic souls keep running.
|
||||
"""
|
||||
skill_refs = agent_soul.skills_files.skills
|
||||
file_refs = agent_soul.skills_files.files
|
||||
if not skill_refs and not file_refs:
|
||||
return None, []
|
||||
|
||||
warnings: list[dict[str, str]] = []
|
||||
mentioned_drive_refs = [
|
||||
decode_drive_mention_ref(mention.ref_id)
|
||||
for mention in parse_prompt_mentions(agent_soul.prompt.system_prompt)
|
||||
if mention.kind in {MentionKind.SKILL, MentionKind.FILE}
|
||||
]
|
||||
ordered_mentions = list(dict.fromkeys(ref for ref in mentioned_drive_refs if ref))
|
||||
if not agent_id:
|
||||
if not ordered_mentions:
|
||||
return None, []
|
||||
return None, [
|
||||
{
|
||||
"section": "agent_soul.prompt.system_prompt",
|
||||
"code": "drive_ref_dangling",
|
||||
"message": "drive mentions are configured but the run has no bound agent to address a drive by.",
|
||||
}
|
||||
]
|
||||
|
||||
drive_service = AgentDriveService()
|
||||
skills_catalog = drive_service.list_skills(tenant_id=tenant_id, agent_id=agent_id)
|
||||
manifest_items = drive_service.manifest(tenant_id=tenant_id, agent_id=agent_id)
|
||||
manifest_by_key = {item["key"]: item for item in manifest_items}
|
||||
skill_keys = {skill["skill_md_key"] for skill in skills_catalog}
|
||||
warnings: list[dict[str, str]] = []
|
||||
mentioned_skill_keys: list[str] = []
|
||||
mentioned_file_keys: list[str] = []
|
||||
for drive_key in ordered_mentions:
|
||||
if drive_key in skill_keys:
|
||||
mentioned_skill_keys.append(drive_key)
|
||||
continue
|
||||
if drive_key in manifest_by_key:
|
||||
mentioned_file_keys.append(drive_key)
|
||||
continue
|
||||
warnings.append(
|
||||
{
|
||||
"section": "agent_soul.skills_files",
|
||||
"code": "skill_ref_dangling",
|
||||
"message": "skills_files is configured but the run has no bound agent to address a drive by.",
|
||||
"section": "agent_soul.prompt.system_prompt",
|
||||
"code": "mention_target_missing",
|
||||
"message": f"drive mention '{drive_key}' has no matching drive entry.",
|
||||
}
|
||||
)
|
||||
return None, warnings
|
||||
|
||||
skills: list[DifyDriveSkillConfig] = []
|
||||
for skill in skill_refs:
|
||||
if not skill.skill_md_key:
|
||||
warnings.append(
|
||||
{
|
||||
"section": "agent_soul.skills_files",
|
||||
"code": "skill_ref_dangling",
|
||||
"message": (
|
||||
f"skill_ref_dangling: skill '{skill.name or skill.id or 'unknown'}' has no drive key; "
|
||||
"re-standardize it to expose it at runtime."
|
||||
),
|
||||
}
|
||||
)
|
||||
continue
|
||||
skills.append(
|
||||
DifyDriveSkillConfig(
|
||||
name=skill.name or skill.skill_md_key.split("/", 1)[0],
|
||||
description=skill.description or "",
|
||||
skill_md_key=skill.skill_md_key,
|
||||
archive_key=skill.full_archive_key,
|
||||
)
|
||||
skills = [
|
||||
DifyDriveSkillConfig(
|
||||
path=skill["path"],
|
||||
name=skill["name"],
|
||||
description=skill["description"],
|
||||
skill_md_key=skill["skill_md_key"],
|
||||
archive_key=skill["archive_key"],
|
||||
)
|
||||
for skill in skills_catalog
|
||||
]
|
||||
|
||||
files: list[DifyDriveFileConfig] = []
|
||||
for file in file_refs:
|
||||
if not file.drive_key:
|
||||
# Plain upload references (pre-ENG-625) are not drive-backed; they are
|
||||
# simply invisible to the manifest rather than a defect worth warning on.
|
||||
continue
|
||||
size = file.get("size")
|
||||
files.append(
|
||||
DifyDriveFileConfig(
|
||||
name=file.name or file.drive_key.rsplit("/", 1)[-1],
|
||||
key=file.drive_key,
|
||||
size=size if isinstance(size, int) else None,
|
||||
mime_type=file.type,
|
||||
)
|
||||
)
|
||||
|
||||
if not skills and not files:
|
||||
return None, warnings
|
||||
return DifyDriveLayerConfig(drive_ref=f"agent-{agent_id}", skills=skills, files=files), warnings
|
||||
return (
|
||||
DifyDriveLayerConfig(
|
||||
drive_ref=f"agent-{agent_id}",
|
||||
skills=skills,
|
||||
mentioned_skill_keys=mentioned_skill_keys,
|
||||
mentioned_file_keys=mentioned_file_keys,
|
||||
),
|
||||
warnings,
|
||||
)
|
||||
|
||||
|
||||
def _cli_tool_enabled(item: object) -> bool:
|
||||
|
||||
@ -35,7 +35,6 @@ class WorkflowAgentNodeValidator:
|
||||
"soul",
|
||||
"prompt",
|
||||
"system_prompt",
|
||||
"skills_files",
|
||||
"skills",
|
||||
"files",
|
||||
"tools",
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -1,5 +1,5 @@
|
||||
from datetime import datetime
|
||||
from typing import Annotated, Literal
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field, field_validator
|
||||
|
||||
@ -16,10 +16,8 @@ from models.agent import (
|
||||
)
|
||||
from models.agent_config_entities import (
|
||||
AgentCliToolConfig,
|
||||
AgentFileRefConfig,
|
||||
AgentHumanContactConfig,
|
||||
AgentKnowledgeDatasetConfig,
|
||||
AgentSkillRefConfig,
|
||||
AgentSoulConfig,
|
||||
DeclaredOutputConfig,
|
||||
DeclaredOutputType,
|
||||
@ -396,20 +394,6 @@ class AgentComposerDifyToolCandidateResponse(ResponseModel):
|
||||
tools_count: int | None = None
|
||||
|
||||
|
||||
class AgentComposerSkillCandidateResponse(AgentSkillRefConfig):
|
||||
kind: Literal["skill"] = "skill"
|
||||
|
||||
|
||||
class AgentComposerFileCandidateResponse(AgentFileRefConfig):
|
||||
kind: Literal["file"] = "file"
|
||||
|
||||
|
||||
AgentComposerSkillFileCandidateResponse = Annotated[
|
||||
AgentComposerSkillCandidateResponse | AgentComposerFileCandidateResponse,
|
||||
Field(discriminator="kind"),
|
||||
]
|
||||
|
||||
|
||||
class AgentComposerNodeJobCandidatesResponse(ResponseModel):
|
||||
previous_node_outputs: list[WorkflowPreviousNodeOutputRef] = Field(default_factory=list)
|
||||
declare_output_types: list[DeclaredOutputType] = Field(default_factory=list)
|
||||
@ -417,7 +401,6 @@ class AgentComposerNodeJobCandidatesResponse(ResponseModel):
|
||||
|
||||
|
||||
class AgentComposerSoulCandidatesResponse(ResponseModel):
|
||||
skills_files: list[AgentComposerSkillFileCandidateResponse] = Field(default_factory=list)
|
||||
dify_tools: list[AgentComposerDifyToolCandidateResponse] = Field(default_factory=list)
|
||||
cli_tools: list[AgentCliToolConfig] = Field(default_factory=list)
|
||||
knowledge_datasets: list[AgentKnowledgeDatasetConfig] = Field(default_factory=list)
|
||||
|
||||
@ -165,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
|
||||
@ -209,10 +215,18 @@ 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()
|
||||
|
||||
self._closed.set()
|
||||
# Send a control event on the same Redis channel to unblock the
|
||||
self._publish_close_event()
|
||||
|
||||
@ -220,10 +234,21 @@ class RedisSubscriptionBase(Subscription):
|
||||
# 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:
|
||||
if listener is not None and listener.is_alive():
|
||||
listener.join(timeout=2)
|
||||
self._listener_thread = None
|
||||
|
||||
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:
|
||||
|
||||
@ -7,7 +7,7 @@ Create Date: 2026-06-05 11:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from alembic import context, op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b7c2d9e8a1f4"
|
||||
@ -17,10 +17,23 @@ depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
if _has_last_opened_at_column():
|
||||
return
|
||||
with op.batch_alter_table("tenant_account_joins", schema=None) as batch_op:
|
||||
batch_op.add_column(sa.Column("last_opened_at", sa.DateTime(), nullable=True))
|
||||
|
||||
|
||||
def downgrade():
|
||||
if not _has_last_opened_at_column():
|
||||
return
|
||||
with op.batch_alter_table("tenant_account_joins", schema=None) as batch_op:
|
||||
batch_op.drop_column("last_opened_at")
|
||||
|
||||
|
||||
def _has_last_opened_at_column() -> bool:
|
||||
if context.is_offline_mode():
|
||||
# Offline SQL generation cannot inspect the target schema. Assume the
|
||||
# linear migration path so generated SQL stays explicit.
|
||||
return False
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
return "last_opened_at" in {column["name"] for column in inspector.get_columns("tenant_account_joins")}
|
||||
|
||||
@ -6,9 +6,15 @@ Create Date: 2026-06-18 23:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import mysql
|
||||
from sqlalchemy.engine.mock import MockConnection
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b2515f9d4c2a"
|
||||
@ -31,9 +37,46 @@ def upgrade() -> None:
|
||||
"agent_drive_files",
|
||||
["tenant_id", "agent_id", "is_skill", "key"],
|
||||
)
|
||||
_remove_skills_files_from_snapshots()
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _remove_skills_files_from_snapshots() -> None:
|
||||
connection = op.get_bind()
|
||||
if connection is None or isinstance(connection, MockConnection):
|
||||
return
|
||||
snapshots = sa.table(
|
||||
"agent_config_snapshots",
|
||||
sa.column("id", sa.String()),
|
||||
sa.column("config_snapshot", sa.Text()),
|
||||
)
|
||||
rows = connection.execute(sa.select(snapshots.c.id, snapshots.c.config_snapshot)).fetchall()
|
||||
for row in rows:
|
||||
cleaned = _strip_skills_files(row.config_snapshot)
|
||||
if cleaned is None:
|
||||
continue
|
||||
connection.execute(
|
||||
snapshots.update()
|
||||
.where(snapshots.c.id == row.id)
|
||||
.values(config_snapshot=json.dumps(cleaned, separators=(",", ":"), sort_keys=True))
|
||||
)
|
||||
|
||||
|
||||
def _strip_skills_files(raw_snapshot: Any) -> dict[str, Any] | None:
|
||||
if raw_snapshot is None:
|
||||
return None
|
||||
if isinstance(raw_snapshot, str):
|
||||
snapshot = json.loads(raw_snapshot)
|
||||
elif isinstance(raw_snapshot, dict):
|
||||
snapshot = dict(raw_snapshot)
|
||||
else:
|
||||
snapshot = dict(raw_snapshot)
|
||||
if not isinstance(snapshot, dict) or "skills_files" not in snapshot:
|
||||
return None
|
||||
snapshot.pop("skills_files", None)
|
||||
return snapshot
|
||||
|
||||
@ -361,11 +361,6 @@ class AgentSoulPromptConfig(BaseModel):
|
||||
system_prompt: str = ""
|
||||
|
||||
|
||||
class AgentSoulSkillsFilesConfig(BaseModel):
|
||||
files: list[AgentFileRefConfig] = Field(default_factory=list)
|
||||
skills: list[AgentSkillRefConfig] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentSoulDifyToolCredentialRef(BaseModel):
|
||||
"""Reference to a stored Dify Plugin Tool credential.
|
||||
|
||||
@ -514,7 +509,6 @@ class AgentSoulConfig(BaseModel):
|
||||
|
||||
schema_version: int = 1
|
||||
prompt: AgentSoulPromptConfig = Field(default_factory=AgentSoulPromptConfig)
|
||||
skills_files: AgentSoulSkillsFilesConfig = Field(default_factory=AgentSoulSkillsFilesConfig)
|
||||
tools: AgentSoulToolsConfig = Field(default_factory=AgentSoulToolsConfig)
|
||||
knowledge: AgentSoulKnowledgeConfig = Field(default_factory=AgentSoulKnowledgeConfig)
|
||||
human: AgentSoulHumanConfig = Field(default_factory=AgentSoulHumanConfig)
|
||||
|
||||
@ -592,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
|
||||
|
||||
@ -602,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
|
||||
|
||||
@ -1608,7 +1622,7 @@ Inspect one drive-backed skill for slash-menu hover/detail UI
|
||||
| 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)
|
||||
Delete one drive file by key via drive commit-null semantics
|
||||
|
||||
#### Parameters
|
||||
|
||||
@ -1694,7 +1708,7 @@ Upload + standardize a Skill into the agent drive
|
||||
| 400 | Invalid skill package or no bound agent | |
|
||||
|
||||
### [DELETE] /apps/{app_id}/agent/skills/{slug}
|
||||
Delete a standardized skill: soul ref first, then the <slug>/ drive prefix (ENG-625 D5)
|
||||
Delete a standardized skill by removing its known drive keys via commit-null
|
||||
|
||||
#### Parameters
|
||||
|
||||
@ -8145,7 +8159,7 @@ Get all published workflows for a snippet
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Published workflows retrieved successfully | **application/json**: [WorkflowPaginationResponse](#workflowpaginationresponse)<br> |
|
||||
| 200 | Published workflows retrieved successfully | **application/json**: [SnippetWorkflowPaginationResponse](#snippetworkflowpaginationresponse)<br> |
|
||||
|
||||
### [GET] /snippets/{snippet_id}/workflows/default-workflow-block-configs
|
||||
**Get default block configurations for snippet workflow**
|
||||
@ -12143,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 |
|
||||
@ -12392,23 +12417,6 @@ Risk marker for CLI tool bootstrap commands.
|
||||
| provider_id | string | | No |
|
||||
| tools_count | integer | | No |
|
||||
|
||||
#### AgentComposerFileCandidateResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| drive_key | string | | No |
|
||||
| file_id | string | | No |
|
||||
| id | string | | No |
|
||||
| kind | string, <br>**Default:** file | | No |
|
||||
| name | string | | No |
|
||||
| reference | string | | No |
|
||||
| remote_url | string | | No |
|
||||
| tenant_id | string | | No |
|
||||
| transfer_method | string | | No |
|
||||
| type | string | | No |
|
||||
| upload_file_id | string | | No |
|
||||
| url | string | | No |
|
||||
|
||||
#### AgentComposerImpactBindingResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@ -12433,22 +12441,6 @@ Risk marker for CLI tool bootstrap commands.
|
||||
| human_contacts | [ [AgentHumanContactConfig](#agenthumancontactconfig) ] | | No |
|
||||
| previous_node_outputs | [ [WorkflowPreviousNodeOutputRef](#workflowpreviousnodeoutputref) ] | | No |
|
||||
|
||||
#### AgentComposerSkillCandidateResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| description | string | | No |
|
||||
| file_id | string | | No |
|
||||
| full_archive_file_id | string | | No |
|
||||
| full_archive_key | string | | No |
|
||||
| id | string | | No |
|
||||
| kind | string, <br>**Default:** skill | | No |
|
||||
| manifest_files | [ string ] | | No |
|
||||
| name | string | | No |
|
||||
| path | string | | No |
|
||||
| skill_md_file_id | string | | No |
|
||||
| skill_md_key | string | | No |
|
||||
|
||||
#### AgentComposerSoulCandidatesResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@ -12457,7 +12449,6 @@ Risk marker for CLI tool bootstrap commands.
|
||||
| dify_tools | [ [AgentComposerDifyToolCandidateResponse](#agentcomposerdifytoolcandidateresponse) ] | | No |
|
||||
| human_contacts | [ [AgentHumanContactConfig](#agenthumancontactconfig) ] | | No |
|
||||
| knowledge_datasets | [ [AgentKnowledgeDatasetConfig](#agentknowledgedatasetconfig) ] | | No |
|
||||
| skills_files | [ ] | | No |
|
||||
|
||||
#### AgentComposerSoulLockResponse
|
||||
|
||||
@ -12562,6 +12553,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 |
|
||||
@ -12572,7 +12569,6 @@ Audit operation recorded for Agent Soul version/revision changes.
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| config_version_id | string | | No |
|
||||
| removed_keys | [ string ] | | No |
|
||||
| result | string | | Yes |
|
||||
|
||||
@ -12586,7 +12582,6 @@ Audit operation recorded for Agent Soul version/revision changes.
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| config_version_id | string | | No |
|
||||
| file | [AgentDriveFileResponse](#agentdrivefileresponse) | | Yes |
|
||||
|
||||
#### AgentDriveFilePayload
|
||||
@ -13172,27 +13167,12 @@ Visibility and lifecycle scope of an Agent record.
|
||||
| enabled | boolean | | No |
|
||||
| type | string | | No |
|
||||
|
||||
#### AgentSkillRefConfig
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| description | string | | No |
|
||||
| file_id | string | | No |
|
||||
| full_archive_file_id | string | | No |
|
||||
| full_archive_key | string | | No |
|
||||
| id | string | | No |
|
||||
| manifest_files | [ string ] | | No |
|
||||
| name | string | | No |
|
||||
| path | string | | No |
|
||||
| skill_md_file_id | string | | No |
|
||||
| skill_md_key | string | | No |
|
||||
|
||||
#### AgentSkillUploadResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| manifest | [SkillManifest](#skillmanifest) | | Yes |
|
||||
| skill | [AgentSkillRefConfig](#agentskillrefconfig) | | Yes |
|
||||
| skill | [AgentUploadedSkillResponse](#agentuploadedskillresponse) | | Yes |
|
||||
|
||||
#### AgentSoulAppFeaturesConfig
|
||||
|
||||
@ -13221,7 +13201,6 @@ Visibility and lifecycle scope of an Agent record.
|
||||
| prompt | [AgentSoulPromptConfig](#agentsoulpromptconfig) | | No |
|
||||
| sandbox | [AgentSoulSandboxConfig](#agentsoulsandboxconfig) | | No |
|
||||
| schema_version | integer, <br>**Default:** 1 | | No |
|
||||
| skills_files | [AgentSoulSkillsFilesConfig](#agentsoulskillsfilesconfig) | | No |
|
||||
| tools | [AgentSoulToolsConfig](#agentsoultoolsconfig) | | No |
|
||||
|
||||
#### AgentSoulDifyToolConfig
|
||||
@ -13338,13 +13317,6 @@ Reference to model credentials resolved only at runtime.
|
||||
| config | [AgentSandboxProviderConfig](#agentsandboxproviderconfig) | | No |
|
||||
| provider | string | | No |
|
||||
|
||||
#### AgentSoulSkillsFilesConfig
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| files | [ [AgentFileRefConfig](#agentfilerefconfig) ] | | No |
|
||||
| skills | [ [AgentSkillRefConfig](#agentskillrefconfig) ] | | No |
|
||||
|
||||
#### AgentSoulToolsConfig
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@ -13476,6 +13448,16 @@ Soft lifecycle state for Agent records.
|
||||
| tool_output | object | | Yes |
|
||||
| tool_parameters | object | | Yes |
|
||||
|
||||
#### AgentUploadedSkillResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| archive_key | string | | No |
|
||||
| description | string | | Yes |
|
||||
| name | string | | Yes |
|
||||
| path | string | | Yes |
|
||||
| skill_md_key | string | | Yes |
|
||||
|
||||
#### AgentUserSatisfactionRateStatisticResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@ -19435,6 +19417,15 @@ Query parameters for listing snippet published workflows.
|
||||
| limit | integer, <br>**Default:** 10 | | No |
|
||||
| page | integer, <br>**Default:** 1 | | No |
|
||||
|
||||
#### SnippetWorkflowPaginationResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| has_more | boolean | | Yes |
|
||||
| items | [ [SnippetWorkflowResponse](#snippetworkflowresponse) ] | | Yes |
|
||||
| limit | integer | | Yes |
|
||||
| page | integer | | Yes |
|
||||
|
||||
#### SnippetWorkflowResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
|
||||
@ -80,7 +80,7 @@ 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 |
|
||||
| workspace_id | query | | Yes | string |
|
||||
@ -318,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 |
|
||||
|
||||
@ -592,12 +592,12 @@ Request body for POST /workspaces/<workspace_id>/apps/imports.
|
||||
|
||||
#### 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 |
|
||||
| workspace_id | string | | Yes |
|
||||
@ -922,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 |
|
||||
|
||||
@ -990,6 +990,24 @@ Pagination for GET /account/sessions. Strict (extra='forbid').
|
||||
| last_used_at | string | | No |
|
||||
| prefix | string | | Yes |
|
||||
|
||||
#### 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 |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| 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
|
||||
|
||||
200 body for POST /apps/<id>/tasks/<task_id>/stop. The handler always returns
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -137,9 +137,6 @@ def soul_candidates(
|
||||
soul = agent_soul or AgentSoulConfig()
|
||||
truncated = False
|
||||
|
||||
skills_files = [{"kind": "skill", **skill.model_dump(exclude_none=True)} for skill in soul.skills_files.skills]
|
||||
skills_files += [{"kind": "file", **file.model_dump(exclude_none=True)} for file in soul.skills_files.files]
|
||||
|
||||
cli_tools = [tool.model_dump(exclude_none=True) for tool in soul.tools.cli_tools if tool.enabled]
|
||||
|
||||
dataset_ids = [dataset.id for dataset in soul.knowledge.datasets if dataset.id]
|
||||
@ -162,7 +159,6 @@ def soul_candidates(
|
||||
dify_tools = workspace_tools_loader()
|
||||
|
||||
lists = {
|
||||
"skills_files": skills_files,
|
||||
"dify_tools": dify_tools,
|
||||
"cli_tools": cli_tools,
|
||||
"knowledge_datasets": knowledge_datasets,
|
||||
|
||||
@ -21,7 +21,6 @@ from models.agent import (
|
||||
WorkflowAgentNodeBinding,
|
||||
)
|
||||
from models.agent_config_entities import (
|
||||
AgentFileRefConfig,
|
||||
DeclaredOutputConfig,
|
||||
)
|
||||
from models.agent_config_entities import (
|
||||
@ -34,7 +33,6 @@ from services.agent.errors import (
|
||||
AgentNameConflictError,
|
||||
AgentNotFoundError,
|
||||
AgentVersionNotFoundError,
|
||||
InvalidComposerConfigError,
|
||||
)
|
||||
from services.entities.agent_entities import (
|
||||
AgentSoulConfig,
|
||||
@ -48,6 +46,13 @@ from services.entities.agent_entities import (
|
||||
# WorkflowAgentNodeBinding.workflow_version tag for the draft workflow row.
|
||||
# Mirrors Workflow.version when it is "draft" (see models/workflow.py).
|
||||
_DRAFT_WORKFLOW_VERSION = "draft"
|
||||
_PUBLISH_SAVE_STRATEGIES = frozenset(
|
||||
{
|
||||
ComposerSaveStrategy.SAVE_AS_NEW_VERSION,
|
||||
ComposerSaveStrategy.SAVE_AS_NEW_AGENT,
|
||||
ComposerSaveStrategy.SAVE_TO_ROSTER,
|
||||
}
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -73,6 +78,13 @@ def _backfill_cli_tool_ids(agent_soul: AgentSoulConfig | None) -> None:
|
||||
seen_ids.add(minted)
|
||||
|
||||
|
||||
def _validate_composer_payload_for_strategy(payload: ComposerSavePayload) -> None:
|
||||
if payload.save_strategy in _PUBLISH_SAVE_STRATEGIES:
|
||||
ComposerConfigValidator.validate_publish_payload(payload)
|
||||
return
|
||||
ComposerConfigValidator.validate_draft_save_payload(payload)
|
||||
|
||||
|
||||
class AgentComposerService:
|
||||
@classmethod
|
||||
def load_workflow_composer(cls, *, tenant_id: str, app_id: str, node_id: str) -> dict[str, Any]:
|
||||
@ -102,33 +114,10 @@ class AgentComposerService:
|
||||
raise ValueError("Workflow composer endpoint only accepts workflow variant")
|
||||
|
||||
_backfill_cli_tool_ids(payload.agent_soul)
|
||||
ComposerConfigValidator.validate_save_payload(payload)
|
||||
_validate_composer_payload_for_strategy(payload)
|
||||
workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id)
|
||||
binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id)
|
||||
|
||||
# ENG-623 §4.4: drive-backed refs must point at real drive rows before the
|
||||
# soul is persisted. Only strategies that write the soul onto an *existing*
|
||||
# agent are checked — new-agent strategies create a fresh (empty) drive, so
|
||||
# any carried drive key would be flagged on the next save instead.
|
||||
if (
|
||||
payload.agent_soul is not None
|
||||
and binding is not None
|
||||
and binding.agent_id
|
||||
and payload.save_strategy
|
||||
in (
|
||||
ComposerSaveStrategy.NODE_JOB_ONLY,
|
||||
ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION,
|
||||
ComposerSaveStrategy.SAVE_AS_NEW_VERSION,
|
||||
)
|
||||
and (
|
||||
payload.save_strategy != ComposerSaveStrategy.NODE_JOB_ONLY
|
||||
or binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT
|
||||
)
|
||||
):
|
||||
cls._require_drive_refs_resolved(
|
||||
tenant_id=tenant_id, agent_id=binding.agent_id, agent_soul=payload.agent_soul
|
||||
)
|
||||
|
||||
match payload.save_strategy:
|
||||
case ComposerSaveStrategy.NODE_JOB_ONLY:
|
||||
binding = cls._save_node_job_only(
|
||||
@ -176,7 +165,11 @@ class AgentComposerService:
|
||||
version_id=version_id,
|
||||
)
|
||||
state = cls._serialize_workflow_state(binding=binding, agent=agent, version=version)
|
||||
state["validation"] = cls.collect_validation_findings(tenant_id=tenant_id, payload=payload)
|
||||
state["validation"] = cls.collect_validation_findings(
|
||||
tenant_id=tenant_id,
|
||||
payload=payload,
|
||||
agent_id=binding.agent_id,
|
||||
)
|
||||
return state
|
||||
|
||||
@classmethod
|
||||
@ -215,7 +208,7 @@ class AgentComposerService:
|
||||
if payload.variant != ComposerVariant.AGENT_APP:
|
||||
raise ValueError("Agent App composer endpoint only accepts agent_app variant")
|
||||
_backfill_cli_tool_ids(payload.agent_soul)
|
||||
ComposerConfigValidator.validate_save_payload(payload)
|
||||
_validate_composer_payload_for_strategy(payload)
|
||||
if payload.agent_soul is None:
|
||||
raise ValueError("agent_soul is required")
|
||||
|
||||
@ -250,9 +243,6 @@ class AgentComposerService:
|
||||
db.session.rollback()
|
||||
raise AgentNameConflictError() from exc
|
||||
|
||||
# ENG-623 §4.4: dangling drive-backed refs are rejected before persisting.
|
||||
cls._require_drive_refs_resolved(tenant_id=tenant_id, agent_id=agent.id, agent_soul=payload.agent_soul)
|
||||
|
||||
if payload.save_strategy == ComposerSaveStrategy.SAVE_AS_NEW_VERSION or not agent.active_config_snapshot_id:
|
||||
version = cls._create_config_version(
|
||||
tenant_id=tenant_id,
|
||||
@ -281,7 +271,11 @@ class AgentComposerService:
|
||||
|
||||
db.session.commit()
|
||||
state = cls.load_agent_app_composer(tenant_id=tenant_id, app_id=app_id)
|
||||
state["validation"] = cls.collect_validation_findings(tenant_id=tenant_id, payload=payload)
|
||||
state["validation"] = cls.collect_validation_findings(
|
||||
tenant_id=tenant_id,
|
||||
payload=payload,
|
||||
agent_id=agent.id,
|
||||
)
|
||||
return state
|
||||
|
||||
@classmethod
|
||||
@ -292,11 +286,7 @@ class AgentComposerService:
|
||||
payload: ComposerSavePayload,
|
||||
agent_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""ENG-617 soft findings, with DB-backed dataset existence for placeholders.
|
||||
|
||||
With ``agent_id`` the drive-backed skill/file refs are also checked against
|
||||
the agent drive (ENG-623 §4.4) and dangling ones surface as warnings.
|
||||
"""
|
||||
"""ENG-617 soft findings, with DB-backed dataset and drive mention checks."""
|
||||
from services.agent.prompt_mentions import MentionKind, parse_prompt_mentions
|
||||
|
||||
mentioned_ids: set[str] = set()
|
||||
@ -312,136 +302,14 @@ class AgentComposerService:
|
||||
findings = ComposerConfigValidator.collect_soft_findings(payload, existing_dataset_ids=existing_dataset_ids)
|
||||
if agent_id and payload.agent_soul is not None:
|
||||
findings["warnings"].extend(
|
||||
cls._drive_ref_findings(tenant_id=tenant_id, agent_id=agent_id, agent_soul=payload.agent_soul)
|
||||
cls._drive_mention_findings(
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
prompt=payload.agent_soul.prompt.system_prompt,
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
@classmethod
|
||||
def remove_drive_refs(
|
||||
cls,
|
||||
*,
|
||||
tenant_id: str,
|
||||
agent_id: str,
|
||||
account_id: str,
|
||||
skill_slug: str | None = None,
|
||||
file_key: str | None = None,
|
||||
app_id: str | None = None,
|
||||
node_id: str | None = None,
|
||||
) -> str | None:
|
||||
"""Drop the soul refs backed by a drive skill/file before the drive rows go.
|
||||
|
||||
Soul-first ordering (ENG-625 D5): a mid-failure leaves harmless orphan KV
|
||||
rows that an idempotent DELETE retry cleans, instead of a soul ref that
|
||||
keeps failing dangling-ref validation. Returns the new config version id,
|
||||
or ``None`` when the soul held no matching ref (idempotent re-delete).
|
||||
"""
|
||||
if (skill_slug is None) == (file_key is None):
|
||||
raise ValueError("remove_drive_refs requires exactly one of skill_slug or file_key")
|
||||
agent = db.session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1))
|
||||
if agent is None or not agent.active_config_snapshot_id:
|
||||
return None
|
||||
current_snapshot = cls._require_version(
|
||||
tenant_id=tenant_id, agent_id=agent.id, version_id=agent.active_config_snapshot_id
|
||||
)
|
||||
agent_soul = AgentSoulConfig.model_validate(current_snapshot.config_snapshot_dict)
|
||||
|
||||
removed_display: str | None = None
|
||||
if skill_slug is not None:
|
||||
kept_skills = []
|
||||
for skill in agent_soul.skills_files.skills:
|
||||
slug = (skill.skill_md_key or "").split("/", 1)[0] or (skill.path or "").strip("/")
|
||||
if slug == skill_slug:
|
||||
removed_display = skill.name or skill.id or skill_slug
|
||||
continue
|
||||
kept_skills.append(skill)
|
||||
if removed_display is None:
|
||||
return None
|
||||
agent_soul.skills_files.skills = kept_skills
|
||||
note = f"Removed skill '{removed_display}' from the drive."
|
||||
else:
|
||||
kept_files = []
|
||||
for file in agent_soul.skills_files.files:
|
||||
if file.drive_key == file_key:
|
||||
removed_display = file.name or file.drive_key
|
||||
continue
|
||||
kept_files.append(file)
|
||||
if removed_display is None:
|
||||
return None
|
||||
agent_soul.skills_files.files = kept_files
|
||||
note = f"Removed file '{removed_display}' from the drive."
|
||||
|
||||
version = cls._update_current_version(
|
||||
current_snapshot=current_snapshot,
|
||||
account_id=account_id,
|
||||
agent_soul=agent_soul,
|
||||
operation=AgentConfigRevisionOperation.SAVE_CURRENT_VERSION,
|
||||
version_note=note,
|
||||
)
|
||||
agent.active_config_snapshot_id = version.id
|
||||
agent.updated_by = account_id
|
||||
cls._sync_draft_binding_snapshot(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
node_id=node_id,
|
||||
agent_id=agent_id,
|
||||
snapshot_id=version.id,
|
||||
account_id=account_id,
|
||||
)
|
||||
db.session.commit()
|
||||
return version.id
|
||||
|
||||
@classmethod
|
||||
def add_drive_file_ref(
|
||||
cls,
|
||||
*,
|
||||
tenant_id: str,
|
||||
agent_id: str,
|
||||
account_id: str,
|
||||
file_ref: AgentFileRefConfig,
|
||||
app_id: str | None = None,
|
||||
node_id: str | None = None,
|
||||
) -> str | None:
|
||||
"""Add or replace one drive-backed file ref in the active Agent Soul.
|
||||
|
||||
``POST /agent/files`` is an ADD FILE user action, not just a low-level
|
||||
drive commit. The committed file must be present in ``skills_files.files``
|
||||
because runtime ``dify.drive`` is built from the active Agent Soul.
|
||||
"""
|
||||
if not file_ref.drive_key:
|
||||
raise ValueError("file_ref.drive_key is required")
|
||||
agent = db.session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1))
|
||||
if agent is None or not agent.active_config_snapshot_id:
|
||||
return None
|
||||
current_snapshot = cls._require_version(
|
||||
tenant_id=tenant_id, agent_id=agent.id, version_id=agent.active_config_snapshot_id
|
||||
)
|
||||
agent_soul = AgentSoulConfig.model_validate(current_snapshot.config_snapshot_dict)
|
||||
kept_files = [item for item in agent_soul.skills_files.files if item.drive_key != file_ref.drive_key]
|
||||
kept_files.append(file_ref)
|
||||
agent_soul.skills_files.files = kept_files
|
||||
|
||||
display = file_ref.name or file_ref.drive_key
|
||||
version = cls._update_current_version(
|
||||
current_snapshot=current_snapshot,
|
||||
account_id=account_id,
|
||||
agent_soul=agent_soul,
|
||||
operation=AgentConfigRevisionOperation.SAVE_CURRENT_VERSION,
|
||||
version_note=f"Added file '{display}' to the drive.",
|
||||
)
|
||||
agent.active_config_snapshot_id = version.id
|
||||
agent.active_config_has_model = agent_soul_has_model(agent_soul)
|
||||
agent.updated_by = account_id
|
||||
cls._sync_draft_binding_snapshot(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
node_id=node_id,
|
||||
agent_id=agent_id,
|
||||
snapshot_id=version.id,
|
||||
account_id=account_id,
|
||||
)
|
||||
db.session.commit()
|
||||
return version.id
|
||||
|
||||
@classmethod
|
||||
def resolve_bound_agent_id(cls, *, tenant_id: str, app_id: str) -> str | None:
|
||||
"""The Agent App's bound roster agent id, if any (validate-endpoint context)."""
|
||||
@ -468,49 +336,25 @@ class AgentComposerService:
|
||||
return binding.agent_id if binding else None
|
||||
|
||||
@classmethod
|
||||
def _sync_draft_binding_snapshot(
|
||||
cls,
|
||||
*,
|
||||
tenant_id: str,
|
||||
app_id: str | None,
|
||||
node_id: str | None,
|
||||
agent_id: str,
|
||||
snapshot_id: str,
|
||||
account_id: str,
|
||||
) -> None:
|
||||
"""Keep workflow node bindings on the new active snapshot after direct drive edits."""
|
||||
if not app_id or not node_id:
|
||||
return
|
||||
try:
|
||||
workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id)
|
||||
except ValueError:
|
||||
return
|
||||
binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id)
|
||||
if binding is None or binding.agent_id != agent_id:
|
||||
return
|
||||
binding.current_snapshot_id = snapshot_id
|
||||
binding.updated_by = account_id
|
||||
|
||||
@classmethod
|
||||
def _drive_ref_findings(
|
||||
def _drive_mention_findings(
|
||||
cls,
|
||||
*,
|
||||
tenant_id: str,
|
||||
agent_id: str,
|
||||
agent_soul: AgentSoulConfig,
|
||||
prompt: str,
|
||||
) -> list[dict[str, str | None]]:
|
||||
"""Drive-backed refs whose keys have no row in the agent drive (ENG-623 §4.4).
|
||||
"""Soft warnings for missing drive-backed prompt mentions."""
|
||||
from services.agent.prompt_mentions import MentionKind, parse_prompt_mentions
|
||||
from services.agent_drive_service import decode_drive_mention_ref
|
||||
|
||||
Each finding message starts with its stable code token
|
||||
(``skill_ref_dangling`` / ``file_ref_dangling``) in the ENG-616/617 style.
|
||||
"""
|
||||
wanted_keys: dict[str, tuple[str, str]] = {}
|
||||
for skill in agent_soul.skills_files.skills:
|
||||
if skill.skill_md_key:
|
||||
wanted_keys[skill.skill_md_key] = ("skill_ref_dangling", skill.name or skill.id or "unknown")
|
||||
for file in agent_soul.skills_files.files:
|
||||
if file.drive_key:
|
||||
wanted_keys[file.drive_key] = ("file_ref_dangling", file.name or file.id or "unknown")
|
||||
for mention in parse_prompt_mentions(prompt):
|
||||
if mention.kind not in {MentionKind.SKILL, MentionKind.FILE}:
|
||||
continue
|
||||
decoded_key = decode_drive_mention_ref(mention.ref_id)
|
||||
if not decoded_key:
|
||||
continue
|
||||
wanted_keys[decoded_key] = (mention.kind.value, mention.label or decoded_key)
|
||||
if not wanted_keys:
|
||||
return []
|
||||
|
||||
@ -524,28 +368,20 @@ class AgentComposerService:
|
||||
)
|
||||
)
|
||||
findings: list[dict[str, str | None]] = []
|
||||
for key, (code, display) in wanted_keys.items():
|
||||
for key, (kind, display) in wanted_keys.items():
|
||||
if key in existing_keys:
|
||||
continue
|
||||
kind = "skill" if code == "skill_ref_dangling" else "file"
|
||||
findings.append(
|
||||
{
|
||||
"code": code,
|
||||
"code": "mention_target_missing",
|
||||
"surface": "agent_soul",
|
||||
"kind": kind,
|
||||
"id": key,
|
||||
"message": f"{code}: {kind} '{display}' has no drive entry for key '{key}'.",
|
||||
"message": f"{kind} '{display}' has no drive entry for key '{key}'.",
|
||||
}
|
||||
)
|
||||
return findings
|
||||
|
||||
@classmethod
|
||||
def _require_drive_refs_resolved(cls, *, tenant_id: str, agent_id: str, agent_soul: AgentSoulConfig) -> None:
|
||||
"""Hard save-time guard: dangling drive-backed refs are rejected (400)."""
|
||||
findings = cls._drive_ref_findings(tenant_id=tenant_id, agent_id=agent_id, agent_soul=agent_soul)
|
||||
if findings:
|
||||
raise InvalidComposerConfigError("; ".join(str(finding["message"]) for finding in findings))
|
||||
|
||||
@classmethod
|
||||
def get_workflow_candidates(cls, *, tenant_id: str, app_id: str, node_id: str, user_id: str) -> dict[str, Any]:
|
||||
"""Slash-menu data source for the workflow Agent node composer (ENG-615)."""
|
||||
|
||||
@ -50,7 +50,7 @@ _DANGEROUS_ACK_KEYS = (
|
||||
|
||||
class ComposerConfigValidator:
|
||||
@classmethod
|
||||
def validate_save_payload(cls, payload: ComposerSavePayload) -> None:
|
||||
def validate_draft_save_payload(cls, payload: ComposerSavePayload) -> None:
|
||||
if (
|
||||
payload.variant == ComposerVariant.WORKFLOW
|
||||
and payload.soul_lock.locked
|
||||
@ -59,6 +59,13 @@ class ComposerConfigValidator:
|
||||
):
|
||||
raise AgentSoulLockedError()
|
||||
|
||||
@classmethod
|
||||
def validate_save_payload(cls, payload: ComposerSavePayload) -> None:
|
||||
cls.validate_publish_payload(payload)
|
||||
|
||||
@classmethod
|
||||
def validate_publish_payload(cls, payload: ComposerSavePayload) -> None:
|
||||
cls.validate_draft_save_payload(payload)
|
||||
if payload.agent_soul is not None:
|
||||
cls.validate_agent_soul(payload.agent_soul)
|
||||
if payload.node_job is not None:
|
||||
@ -191,6 +198,8 @@ class ComposerConfigValidator:
|
||||
}
|
||||
)
|
||||
continue
|
||||
if mention.kind in {MentionKind.SKILL, MentionKind.FILE}:
|
||||
continue
|
||||
if resolved is None:
|
||||
warnings.append(
|
||||
{
|
||||
|
||||
@ -4,13 +4,14 @@ Slash-menu insertions are stored inline in the plain-string prompt as tokens:
|
||||
|
||||
[§<kind>:<id>[:<label>]§]
|
||||
|
||||
``kind`` is a fixed lowercase word; ``id`` points at an item in the Agent config
|
||||
lists (mentions are pointers — the entity itself lives in ``skills_files`` /
|
||||
``tools`` / ``knowledge.datasets`` / ``human.contacts`` /
|
||||
``previous_node_output_refs`` / ``declared_outputs``); ``label`` is an optional
|
||||
plain-text fallback only (the backend always re-resolves by id, so renames never
|
||||
break references). A single ``:`` separates all three fields; ``label`` is the
|
||||
trailing remainder and may itself contain ``:``.
|
||||
``kind`` is a fixed lowercase word; ``id`` points at an item in the Agent
|
||||
runtime context. For prompt-owned entities that means Agent Soul lists such as
|
||||
``tools`` / ``knowledge.datasets`` / ``human.contacts`` and workflow job lists
|
||||
such as ``previous_node_output_refs`` / ``declared_outputs``. For drive-backed
|
||||
``skill`` / ``file`` mentions the field stores a URL-encoded drive key and is
|
||||
resolved against ``agent_drive_files`` at runtime. ``label`` is an optional
|
||||
plain-text fallback only. A single ``:`` separates all three fields; ``label``
|
||||
is the trailing remainder and may itself contain ``:``.
|
||||
|
||||
The ``[§…§]`` wrapper uses the section sign ``§`` (U+00A7), which never appears
|
||||
in Dify template syntax (``{{var}}`` / ``{{#a.b#}}``) nor in normal prompt text,
|
||||
@ -55,7 +56,11 @@ MENTION_PATTERN = re.compile(
|
||||
_RESIDUAL_MENTION_PATTERN = re.compile(r"\[§([A-Za-z_][A-Za-z0-9_]*:[^§]*?)§\]")
|
||||
|
||||
MAX_MENTIONS_PER_PROMPT = 200
|
||||
MAX_MENTION_FIELD_LENGTH = 255
|
||||
# Drive keys are validated up to 512 Unicode code points before URL encoding.
|
||||
# Worst case, one code point becomes 4 UTF-8 bytes and each byte becomes a
|
||||
# 3-character ``%XX`` escape, so a valid encoded drive key can reach 6144 chars.
|
||||
MAX_MENTION_REF_ID_LENGTH = 6144
|
||||
MAX_MENTION_LABEL_LENGTH = 255
|
||||
|
||||
# Reserved ``tool`` mention id suffix: ``<provider>/*`` means "every tool of this
|
||||
# provider" (a provider hosts many tools, like an MCP server). Single tools use
|
||||
@ -102,7 +107,7 @@ def parse_prompt_mentions(prompt: str) -> list[PromptMention]:
|
||||
for match in MENTION_PATTERN.finditer(prompt or ""):
|
||||
ref_id = match.group(2)
|
||||
label = match.group(3)
|
||||
if len(ref_id) > MAX_MENTION_FIELD_LENGTH or (label is not None and len(label) > MAX_MENTION_FIELD_LENGTH):
|
||||
if len(ref_id) > MAX_MENTION_REF_ID_LENGTH or (label is not None and len(label) > MAX_MENTION_LABEL_LENGTH):
|
||||
continue
|
||||
mentions.append(
|
||||
PromptMention(
|
||||
@ -127,8 +132,8 @@ def expand_prompt_mentions(prompt: str, resolver: MentionResolver) -> str:
|
||||
def _replace(match: re.Match[str]) -> str:
|
||||
ref_id = match.group(2)
|
||||
label = match.group(3) or None
|
||||
fallback = (label or ref_id)[:MAX_MENTION_FIELD_LENGTH]
|
||||
if len(ref_id) > MAX_MENTION_FIELD_LENGTH or (label is not None and len(label) > MAX_MENTION_FIELD_LENGTH):
|
||||
fallback = (label or ref_id)[:MAX_MENTION_LABEL_LENGTH]
|
||||
if len(ref_id) > MAX_MENTION_REF_ID_LENGTH or (label is not None and len(label) > MAX_MENTION_LABEL_LENGTH):
|
||||
return fallback
|
||||
mention = PromptMention(
|
||||
kind=MentionKind(match.group(1)),
|
||||
@ -141,7 +146,7 @@ def expand_prompt_mentions(prompt: str, resolver: MentionResolver) -> str:
|
||||
resolved = resolver(mention)
|
||||
if resolved is None or not resolved.strip():
|
||||
return fallback
|
||||
return resolved[:MAX_MENTION_FIELD_LENGTH]
|
||||
return resolved[:MAX_MENTION_LABEL_LENGTH]
|
||||
|
||||
return scrub_mention_markers(MENTION_PATTERN.sub(_replace, prompt))
|
||||
|
||||
@ -163,27 +168,19 @@ def scrub_mention_markers(text: str) -> str:
|
||||
# inner is ``kind:id[:label]``; prefer the label, else the id.
|
||||
parts = match.group(1).split(":", 2)
|
||||
if len(parts) >= 3 and parts[2].strip():
|
||||
return parts[2].strip()[:MAX_MENTION_FIELD_LENGTH]
|
||||
return parts[2].strip()[:MAX_MENTION_LABEL_LENGTH]
|
||||
if len(parts) >= 2 and parts[1].strip():
|
||||
return parts[1].strip()[:MAX_MENTION_FIELD_LENGTH]
|
||||
return match.group(1)[:MAX_MENTION_FIELD_LENGTH]
|
||||
return parts[1].strip()[:MAX_MENTION_LABEL_LENGTH]
|
||||
return match.group(1)[:MAX_MENTION_LABEL_LENGTH]
|
||||
|
||||
return _RESIDUAL_MENTION_PATTERN.sub(_degrade, text)
|
||||
|
||||
|
||||
def build_soul_mention_resolver(agent_soul: AgentSoulConfig) -> MentionResolver:
|
||||
"""Resolve soul-surface mentions to canonical display names from the soul config."""
|
||||
"""Resolve non-drive soul-surface mentions to canonical display names."""
|
||||
|
||||
def _resolve(mention: PromptMention) -> str | None:
|
||||
match mention.kind:
|
||||
case MentionKind.SKILL:
|
||||
for skill in agent_soul.skills_files.skills:
|
||||
if mention.ref_id in (skill.id, skill.name):
|
||||
return skill.name or skill.id
|
||||
case MentionKind.FILE:
|
||||
for file in agent_soul.skills_files.files:
|
||||
if mention.ref_id in (file.id, file.name):
|
||||
return file.name or file.id
|
||||
case MentionKind.TOOL:
|
||||
for tool in agent_soul.tools.dify_tools:
|
||||
prefixes = {prefix for prefix in (tool.provider, tool.provider_id, tool.plugin_id) if prefix}
|
||||
@ -273,7 +270,8 @@ def _selector_from_ref(ref: WorkflowPreviousNodeOutputRef) -> tuple[str, str] |
|
||||
__all__ = [
|
||||
"ALL_PROVIDER_TOOLS_SUFFIX",
|
||||
"MAX_MENTIONS_PER_PROMPT",
|
||||
"MAX_MENTION_FIELD_LENGTH",
|
||||
"MAX_MENTION_LABEL_LENGTH",
|
||||
"MAX_MENTION_REF_ID_LENGTH",
|
||||
"MENTION_PATTERN",
|
||||
"NODE_JOB_PROMPT_ALLOWED_KINDS",
|
||||
"SOUL_PROMPT_ALLOWED_KINDS",
|
||||
|
||||
@ -492,6 +492,52 @@ class AgentRosterService:
|
||||
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]:
|
||||
@ -587,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,
|
||||
@ -598,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
|
||||
@ -608,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,
|
||||
|
||||
@ -4,11 +4,12 @@ A Skill is a ``.zip`` / ``.skill`` archive that must contain a ``SKILL.md`` entr
|
||||
file (Anthropic Skills convention: YAML frontmatter with ``name`` + ``description``,
|
||||
followed by markdown instructions). This service validates the archive (extension,
|
||||
size, zip integrity, zip-slip safety, SKILL.md presence/encoding/fields) and
|
||||
extracts a manifest the API can bind to an Agent config version's skill list.
|
||||
extracts a manifest consumed by drive standardization.
|
||||
|
||||
It does NOT execute or load the skill — the agent backend owns execution. It also
|
||||
does not (here) standardize the package into the agent drive; that is ENG-594 (S6),
|
||||
which consumes the manifest produced here.
|
||||
does not persist anything into Agent Soul or bind anything to config versions;
|
||||
``SkillStandardizeService`` consumes the manifest and commits the canonical drive
|
||||
rows instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@ -22,8 +23,6 @@ import zipfile
|
||||
import yaml
|
||||
from pydantic import BaseModel
|
||||
|
||||
from models.agent_config_entities import AgentSkillRefConfig
|
||||
|
||||
# Bounds — generous but finite so a hostile upload can't exhaust memory/disk.
|
||||
_MAX_ARCHIVE_BYTES = 50 * 1024 * 1024
|
||||
_MAX_UNCOMPRESSED_BYTES = 200 * 1024 * 1024
|
||||
@ -58,22 +57,6 @@ class SkillManifest(BaseModel):
|
||||
size: int # total uncompressed bytes
|
||||
hash: str # sha256 of the archive bytes
|
||||
|
||||
def to_skill_ref(self, *, file_id: str, path: str | None = None) -> AgentSkillRefConfig:
|
||||
"""Build a config skill ref. ``path`` is the stable drive path (set by S6)."""
|
||||
return AgentSkillRefConfig.model_validate(
|
||||
{
|
||||
"id": self.hash,
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"file_id": file_id,
|
||||
"path": path,
|
||||
"size": self.size,
|
||||
"hash": self.hash,
|
||||
"entry_path": self.entry_path,
|
||||
"manifest_files": self.files,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class SkillPackageService:
|
||||
"""Validate Skill archives and extract their manifest."""
|
||||
|
||||
@ -9,19 +9,21 @@ to the agent drive (Agent Files §5.4 / §4):
|
||||
|
||||
Both are stored as ``ToolFile`` records and bound via ``AgentDriveService.commit``
|
||||
with ``value_owned_by_drive=True`` (the drive owns their lifecycle). The returned
|
||||
skill ref records the stable drive paths + file ids (not just the raw upload id),
|
||||
so the Composer can reload the bound skill list. The console ``/skills/upload``
|
||||
endpoints delegate to this service so "upload" now always means drive-backed skill
|
||||
normalization.
|
||||
payload is the slim drive-derived skill DTO the UI needs to work with the drive
|
||||
catalog — ``name``, ``description``, ``path``, ``skill_md_key``, and
|
||||
``archive_key`` — plus the extracted manifest for upload feedback. The console
|
||||
``/skills/upload`` endpoints delegate to this service so "upload" now always means
|
||||
drive-backed skill normalization rather than Agent Soul binding.
|
||||
"""
|
||||
|
||||
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, DriveSkillMetadata
|
||||
|
||||
@ -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,
|
||||
@ -103,29 +130,24 @@ class SkillStandardizeService:
|
||||
file_ref=DriveFileRef(kind="tool_file", id=archive_tool_file.id),
|
||||
value_owned_by_drive=True,
|
||||
),
|
||||
*member_items,
|
||||
],
|
||||
)
|
||||
|
||||
skill_ref = AgentSkillRefConfig.model_validate(
|
||||
{
|
||||
"id": manifest.hash,
|
||||
"name": manifest.name,
|
||||
"description": manifest.description,
|
||||
"file_id": archive_tool_file.id,
|
||||
"path": slug,
|
||||
"size": manifest.size,
|
||||
"hash": manifest.hash,
|
||||
"entry_path": skill_md_key,
|
||||
"skill_md_file_id": md_tool_file.id,
|
||||
"skill_md_key": skill_md_key,
|
||||
"full_archive_file_id": archive_tool_file.id,
|
||||
"full_archive_key": archive_key,
|
||||
# ENG-371: zip member listing — strong signals (scripts/*.sh) for infer-tools.
|
||||
"manifest_files": manifest.files,
|
||||
}
|
||||
drive_skill = next(
|
||||
skill
|
||||
for skill in self._drive.list_skills(tenant_id=tenant_id, agent_id=agent_id)
|
||||
if skill["skill_md_key"] == skill_md_key
|
||||
)
|
||||
|
||||
return {
|
||||
"skill": skill_ref.model_dump(exclude_none=True),
|
||||
"skill": {
|
||||
"name": drive_skill["name"],
|
||||
"description": drive_skill["description"],
|
||||
"path": drive_skill["path"],
|
||||
"skill_md_key": drive_skill["skill_md_key"],
|
||||
"archive_key": drive_skill["archive_key"],
|
||||
},
|
||||
"manifest": manifest.model_dump(),
|
||||
}
|
||||
|
||||
|
||||
@ -19,15 +19,11 @@ from typing import Any
|
||||
|
||||
import json_repair
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
from sqlalchemy import select
|
||||
|
||||
from core.errors.error import ProviderTokenNotInitError
|
||||
from core.model_manager import ModelManager
|
||||
from extensions.ext_database import db
|
||||
from graphon.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage
|
||||
from graphon.model_runtime.entities.model_entities import ModelType
|
||||
from models.agent import Agent
|
||||
from models.agent_config_entities import AgentSoulConfig
|
||||
from services.agent_drive_service import AgentDriveError, AgentDriveService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -97,12 +93,8 @@ class SkillToolInferenceService:
|
||||
|
||||
def infer(self, *, tenant_id: str, agent_id: str, slug: str) -> dict[str, Any]:
|
||||
skill_md = self._load_skill_md(tenant_id=tenant_id, agent_id=agent_id, slug=slug)
|
||||
manifest_files = self._manifest_files_from_soul(tenant_id=tenant_id, agent_id=agent_id, slug=slug)
|
||||
|
||||
user_prompt = f"SKILL.md of skill '{slug}':\n\n{skill_md}"
|
||||
if manifest_files:
|
||||
listing = "\n".join(manifest_files[:200])
|
||||
user_prompt += f"\n\nFiles inside the skill package:\n{listing}"
|
||||
|
||||
raw = self._invoke(tenant_id=tenant_id, user_prompt=user_prompt)
|
||||
try:
|
||||
@ -138,37 +130,6 @@ class SkillToolInferenceService:
|
||||
)
|
||||
return str(preview["text"])
|
||||
|
||||
@staticmethod
|
||||
def _manifest_files_from_soul(*, tenant_id: str, agent_id: str, slug: str) -> list[str]:
|
||||
"""The zip path listing standardize persisted onto the ref, if present.
|
||||
|
||||
Degrades to an empty list (SKILL.md-only inference) for refs that
|
||||
predate ``manifest_files``.
|
||||
"""
|
||||
agent = db.session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1))
|
||||
if agent is None or not agent.active_config_snapshot_id:
|
||||
return []
|
||||
from models.agent import AgentConfigSnapshot
|
||||
|
||||
snapshot = db.session.scalar(
|
||||
select(AgentConfigSnapshot).where(
|
||||
AgentConfigSnapshot.tenant_id == tenant_id,
|
||||
AgentConfigSnapshot.agent_id == agent_id,
|
||||
AgentConfigSnapshot.id == agent.active_config_snapshot_id,
|
||||
)
|
||||
)
|
||||
if snapshot is None:
|
||||
return []
|
||||
soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict)
|
||||
for skill in soul.skills_files.skills:
|
||||
ref_slug = (skill.skill_md_key or "").split("/", 1)[0] or (skill.path or "").strip("/")
|
||||
if ref_slug != slug:
|
||||
continue
|
||||
files = skill.get("manifest_files")
|
||||
if isinstance(files, list):
|
||||
return [str(item) for item in files]
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _invoke(*, tenant_id: str, user_prompt: str) -> str:
|
||||
try:
|
||||
|
||||
@ -8,17 +8,25 @@ from pydantic import ValidationError
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidator
|
||||
from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidationError, WorkflowAgentNodeValidator
|
||||
from models.agent import (
|
||||
Agent,
|
||||
AgentConfigSnapshot,
|
||||
AgentDriveFile,
|
||||
AgentScope,
|
||||
AgentStatus,
|
||||
WorkflowAgentBindingType,
|
||||
WorkflowAgentNodeBinding,
|
||||
)
|
||||
from models.agent_config_entities import DeclaredOutputConfig, WorkflowNodeJobConfig
|
||||
from models.agent_config_entities import AgentSoulConfig, DeclaredOutputConfig, WorkflowNodeJobConfig
|
||||
from models.workflow import Workflow
|
||||
from services.agent.composer_validator import ComposerConfigValidator
|
||||
from services.entities.agent_entities import (
|
||||
ComposerSavePayload,
|
||||
ComposerSaveStrategy,
|
||||
ComposerSoulLockPayload,
|
||||
ComposerVariant,
|
||||
)
|
||||
|
||||
|
||||
class WorkflowAgentPublishService:
|
||||
@ -67,11 +75,131 @@ class WorkflowAgentPublishService:
|
||||
@classmethod
|
||||
def validate_agent_nodes_for_publish(cls, *, session: Session, draft_workflow: Workflow) -> None:
|
||||
WorkflowAgentNodeValidator.validate_published_workflow(session=session, workflow=draft_workflow)
|
||||
cls._validate_composer_configs_for_publish(session=session, draft_workflow=draft_workflow)
|
||||
|
||||
@classmethod
|
||||
def validate_agent_nodes_for_draft_sync(cls, *, session: Session, draft_workflow: Workflow) -> None:
|
||||
WorkflowAgentNodeValidator.validate_draft_workflow(session=session, workflow=draft_workflow)
|
||||
|
||||
@classmethod
|
||||
def _validate_composer_configs_for_publish(cls, *, session: Session, draft_workflow: Workflow) -> None:
|
||||
node_ids = {
|
||||
node_id for node_id, _node_data in WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict)
|
||||
}
|
||||
if not node_ids:
|
||||
return
|
||||
|
||||
bindings = session.scalars(
|
||||
select(WorkflowAgentNodeBinding).where(
|
||||
WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id,
|
||||
WorkflowAgentNodeBinding.app_id == draft_workflow.app_id,
|
||||
WorkflowAgentNodeBinding.workflow_id == draft_workflow.id,
|
||||
WorkflowAgentNodeBinding.workflow_version == draft_workflow.version,
|
||||
WorkflowAgentNodeBinding.node_id.in_(node_ids),
|
||||
)
|
||||
).all()
|
||||
for binding in bindings:
|
||||
cls._validate_binding_composer_config_for_publish(session=session, binding=binding)
|
||||
|
||||
@classmethod
|
||||
def _validate_binding_composer_config_for_publish(
|
||||
cls,
|
||||
*,
|
||||
session: Session,
|
||||
binding: WorkflowAgentNodeBinding,
|
||||
) -> None:
|
||||
if not binding.agent_id:
|
||||
return
|
||||
|
||||
agent = session.scalar(
|
||||
select(Agent)
|
||||
.where(
|
||||
Agent.tenant_id == binding.tenant_id,
|
||||
Agent.id == binding.agent_id,
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
if agent is None:
|
||||
return
|
||||
|
||||
snapshot_id = (
|
||||
agent.active_config_snapshot_id
|
||||
if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT
|
||||
else binding.current_snapshot_id
|
||||
)
|
||||
if snapshot_id is None:
|
||||
return
|
||||
|
||||
snapshot = session.scalar(
|
||||
select(AgentConfigSnapshot)
|
||||
.where(
|
||||
AgentConfigSnapshot.tenant_id == binding.tenant_id,
|
||||
AgentConfigSnapshot.agent_id == agent.id,
|
||||
AgentConfigSnapshot.id == snapshot_id,
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
if snapshot is None:
|
||||
return
|
||||
|
||||
agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict)
|
||||
node_job = WorkflowNodeJobConfig.model_validate(binding.node_job_config_dict)
|
||||
payload = ComposerSavePayload.model_construct(
|
||||
variant=ComposerVariant.WORKFLOW,
|
||||
save_strategy=ComposerSaveStrategy.NODE_JOB_ONLY,
|
||||
soul_lock=ComposerSoulLockPayload(locked=False),
|
||||
agent_soul=agent_soul,
|
||||
node_job=node_job,
|
||||
)
|
||||
ComposerConfigValidator.validate_publish_payload(payload)
|
||||
# ENG-623 §4.4: drive-backed refs must point at real drive rows before
|
||||
# publishing. This stays out of composer save so autosave/save-draft can
|
||||
# persist incomplete refs and surface them as non-blocking findings.
|
||||
cls._require_drive_refs_resolved_for_publish(session=session, binding=binding, agent_soul=agent_soul)
|
||||
|
||||
@classmethod
|
||||
def _require_drive_refs_resolved_for_publish(
|
||||
cls,
|
||||
*,
|
||||
session: Session,
|
||||
binding: WorkflowAgentNodeBinding,
|
||||
agent_soul: AgentSoulConfig,
|
||||
) -> None:
|
||||
from services.agent.prompt_mentions import MentionKind, parse_prompt_mentions
|
||||
from services.agent_drive_service import decode_drive_mention_ref
|
||||
|
||||
wanted_keys: dict[str, tuple[str, str]] = {}
|
||||
for mention in parse_prompt_mentions(agent_soul.prompt.system_prompt):
|
||||
if mention.kind not in {MentionKind.SKILL, MentionKind.FILE}:
|
||||
continue
|
||||
drive_key = decode_drive_mention_ref(mention.ref_id)
|
||||
if not drive_key:
|
||||
continue
|
||||
code = "skill_ref_dangling" if mention.kind == MentionKind.SKILL else "file_ref_dangling"
|
||||
wanted_keys[drive_key] = (code, mention.label or drive_key)
|
||||
if not wanted_keys or not binding.agent_id:
|
||||
return
|
||||
|
||||
existing_keys = set(
|
||||
session.scalars(
|
||||
select(AgentDriveFile.key).where(
|
||||
AgentDriveFile.tenant_id == binding.tenant_id,
|
||||
AgentDriveFile.agent_id == binding.agent_id,
|
||||
AgentDriveFile.key.in_(sorted(wanted_keys)),
|
||||
)
|
||||
).all()
|
||||
)
|
||||
messages: list[str] = []
|
||||
for key, (code, display) in wanted_keys.items():
|
||||
if key in existing_keys:
|
||||
continue
|
||||
kind = "skill" if code == "skill_ref_dangling" else "file"
|
||||
messages.append(f"{code}: {kind} '{display}' has no drive entry for key '{key}'.")
|
||||
if messages:
|
||||
raise WorkflowAgentNodeValidationError(
|
||||
f"Workflow Agent node {binding.node_id} has invalid Agent Soul drive refs: {'; '.join(messages)}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def sync_agent_bindings_for_draft(
|
||||
cls,
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
"""Agent 网盘 (agent drive) service — list/manifest + commit with lifecycle (ENG-591).
|
||||
"""Agent 网盘 (agent drive) service — manifest/catalog + commit lifecycle.
|
||||
|
||||
The agent drive is a per-agent path-like KV index over existing UploadFile /
|
||||
ToolFile records (see ``AgentDriveFile``). This service is the control plane:
|
||||
@ -8,11 +8,13 @@ ToolFile records (see ``AgentDriveFile``). This service is the control plane:
|
||||
``FileAccessScope`` (Agent Files §3.1.2). We reuse the standard
|
||||
``file_factory.build_from_mapping`` + ``resolve_file_url`` rebuild, which always
|
||||
filters by ``tenant_id`` in the builders, so omitting the scope is safe.
|
||||
* ``commit`` binds a batch of existing file refs to keys. Source ToolFiles must
|
||||
* ``commit`` is the single mutation entry point for writes and removals.
|
||||
``file_ref=None`` removes an exact key idempotently; otherwise the service
|
||||
binds the referenced UploadFile/ToolFile to the key. Source ToolFiles must
|
||||
belong to the current run user. Overwriting a key whose previous value is
|
||||
``value_owned_by_drive`` physically cleans the old value (storage + record),
|
||||
unless another drive entry still references it. Re-committing the same
|
||||
``key -> file_ref`` is idempotent.
|
||||
``key -> file_ref`` is idempotent and still refreshes skill metadata.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@ -30,7 +32,6 @@ from sqlalchemy.exc import DataError, SQLAlchemyError
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.file_access.controller import DatabaseFileAccessController
|
||||
from core.app.workflow.file_runtime import DifyWorkflowFileRuntime
|
||||
from core.db.session_factory import session_factory
|
||||
from extensions.ext_storage import storage
|
||||
from factories import file_factory
|
||||
@ -93,7 +94,7 @@ class DriveCommitItem(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
key: str
|
||||
file_ref: DriveFileRef
|
||||
file_ref: DriveFileRef | None = None
|
||||
# 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
|
||||
@ -385,6 +386,15 @@ class AgentDriveService:
|
||||
pending_storage_deletes: list[str],
|
||||
) -> dict[str, Any]:
|
||||
key = normalize_drive_key(item.key)
|
||||
if item.file_ref is None:
|
||||
return self._remove_one(
|
||||
session,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
key=key,
|
||||
pending_storage_deletes=pending_storage_deletes,
|
||||
)
|
||||
|
||||
skill_metadata = self._validate_skill_commit_fields(key=key, item=item)
|
||||
file_kind = AgentDriveFileKind(item.file_ref.kind)
|
||||
file_id = item.file_ref.id
|
||||
@ -447,6 +457,45 @@ class AgentDriveService:
|
||||
session.add(row)
|
||||
return self._row_dict(row)
|
||||
|
||||
def _remove_one(
|
||||
self,
|
||||
session: Session,
|
||||
*,
|
||||
tenant_id: str,
|
||||
agent_id: str,
|
||||
key: str,
|
||||
pending_storage_deletes: list[str],
|
||||
) -> dict[str, Any]:
|
||||
existing = session.scalar(
|
||||
select(AgentDriveFile).where(
|
||||
AgentDriveFile.tenant_id == tenant_id,
|
||||
AgentDriveFile.agent_id == agent_id,
|
||||
AgentDriveFile.key == key,
|
||||
)
|
||||
)
|
||||
if existing is None:
|
||||
return {"key": key, "removed": True, "noop": True}
|
||||
result = {
|
||||
"key": key,
|
||||
"removed": True,
|
||||
"file_kind": existing.file_kind.value,
|
||||
"file_id": existing.file_id,
|
||||
"value_owned_by_drive": existing.value_owned_by_drive,
|
||||
"is_skill": existing.is_skill,
|
||||
"skill_metadata": existing.skill_metadata,
|
||||
}
|
||||
if existing.value_owned_by_drive:
|
||||
self._cleanup_value(
|
||||
session,
|
||||
tenant_id=tenant_id,
|
||||
file_kind=existing.file_kind,
|
||||
file_id=existing.file_id,
|
||||
exclude_row_id=existing.id,
|
||||
pending_storage_deletes=pending_storage_deletes,
|
||||
)
|
||||
session.delete(existing)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _row_dict(row: AgentDriveFile) -> dict[str, Any]:
|
||||
return {
|
||||
@ -750,6 +799,11 @@ class AgentDriveService:
|
||||
else:
|
||||
mapping = {"transfer_method": "local_file", "upload_file_id": file_id}
|
||||
controller = DatabaseFileAccessController()
|
||||
# Keep workflow runtime wiring lazy: importing this service is part of
|
||||
# Agent v2 node bootstrap, while ``core.app.workflow`` re-exports the
|
||||
# node factory. A module-level import here would close that cycle.
|
||||
from core.app.workflow.file_runtime import DifyWorkflowFileRuntime
|
||||
|
||||
runtime = DifyWorkflowFileRuntime(file_access_controller=controller)
|
||||
try:
|
||||
if file_kind == AgentDriveFileKind.UPLOAD_FILE:
|
||||
|
||||
@ -770,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,
|
||||
)
|
||||
@ -1585,7 +1586,7 @@ class RBACService:
|
||||
account_id=member_account_id,
|
||||
roles=[
|
||||
RBACRole(
|
||||
id="",
|
||||
id=role,
|
||||
name=role,
|
||||
description="",
|
||||
is_builtin=True,
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
20
api/services/retention/workflow_run/tenant_prefix.py
Normal file
20
api/services/retention/workflow_run/tenant_prefix.py
Normal file
@ -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
|
||||
@ -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"
|
||||
|
||||
@ -2,17 +2,44 @@
|
||||
Testcontainers integration tests for archived workflow run deletion service.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import zipfile
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import MagicMock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from graphon.enums import WorkflowExecutionStatus
|
||||
from libs.archive_storage import ArchiveStorage
|
||||
from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
|
||||
from models.workflow import WorkflowArchiveLog, WorkflowRun
|
||||
from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_NAME, ARCHIVE_SCHEMA_VERSION
|
||||
from services.retention.workflow_run.delete_archived_workflow_run import ArchivedWorkflowRunDeletion
|
||||
|
||||
ARCHIVED_TABLES = [
|
||||
"workflow_runs",
|
||||
"workflow_app_logs",
|
||||
"workflow_node_executions",
|
||||
"workflow_node_execution_offload",
|
||||
"workflow_pauses",
|
||||
"workflow_pause_reasons",
|
||||
"workflow_trigger_logs",
|
||||
]
|
||||
|
||||
|
||||
class FakeArchiveStorage:
|
||||
def __init__(self, objects: dict[str, bytes]):
|
||||
self.objects = objects
|
||||
|
||||
def object_exists(self, key: str) -> bool:
|
||||
return key in self.objects
|
||||
|
||||
def get_object(self, key: str) -> bytes:
|
||||
return self.objects[key]
|
||||
|
||||
|
||||
class TestArchivedWorkflowRunDeletion:
|
||||
def _create_workflow_run(
|
||||
@ -47,7 +74,7 @@ class TestArchivedWorkflowRunDeletion:
|
||||
db_session_with_containers.commit()
|
||||
return run
|
||||
|
||||
def _create_archive_log(self, db_session_with_containers: Session, *, run: WorkflowRun) -> None:
|
||||
def _create_archive_log(self, db_session_with_containers: Session, *, run: WorkflowRun) -> WorkflowArchiveLog:
|
||||
archive_log = WorkflowArchiveLog(
|
||||
tenant_id=run.tenant_id,
|
||||
app_id=run.app_id,
|
||||
@ -72,6 +99,59 @@ class TestArchivedWorkflowRunDeletion:
|
||||
)
|
||||
db_session_with_containers.add(archive_log)
|
||||
db_session_with_containers.commit()
|
||||
return archive_log
|
||||
|
||||
def _archive_key(self, run: WorkflowRun) -> str:
|
||||
return (
|
||||
f"{run.tenant_id}/app_id={run.app_id}/year={run.created_at.strftime('%Y')}/"
|
||||
f"month={run.created_at.strftime('%m')}/workflow_run_id={run.id}/{ARCHIVE_BUNDLE_NAME}"
|
||||
)
|
||||
|
||||
def _archive_bundle(self, run: WorkflowRun, *, workflow_run_rows: int = 1) -> bytes:
|
||||
table_payloads: dict[str, bytes] = {}
|
||||
table_counts = {
|
||||
"workflow_runs": workflow_run_rows,
|
||||
"workflow_app_logs": 0,
|
||||
"workflow_node_executions": 0,
|
||||
"workflow_node_execution_offload": 0,
|
||||
"workflow_pauses": 0,
|
||||
"workflow_pause_reasons": 0,
|
||||
"workflow_trigger_logs": 0,
|
||||
}
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
records = [{"id": run.id}] if table_name == "workflow_runs" and workflow_run_rows else []
|
||||
table_payloads[table_name] = ArchiveStorage.serialize_to_jsonl(records)
|
||||
|
||||
manifest = {
|
||||
"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(),
|
||||
"archived_at": datetime.now(UTC).isoformat(),
|
||||
"tables": {
|
||||
table_name: {
|
||||
"row_count": table_counts[table_name],
|
||||
"checksum": ArchiveStorage.compute_checksum(payload),
|
||||
"size_bytes": len(payload),
|
||||
}
|
||||
for table_name, payload in table_payloads.items()
|
||||
},
|
||||
}
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
|
||||
archive.writestr("manifest.json", json.dumps(manifest).encode("utf-8"))
|
||||
for table_name, payload in table_payloads.items():
|
||||
archive.writestr(f"{table_name}.jsonl", payload)
|
||||
return buffer.getvalue()
|
||||
|
||||
def _patch_storage(self, run: WorkflowRun):
|
||||
storage = FakeArchiveStorage({self._archive_key(run): self._archive_bundle(run)})
|
||||
return patch(
|
||||
"services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
|
||||
return_value=storage,
|
||||
)
|
||||
|
||||
def test_delete_by_run_id_returns_error_when_run_missing(self, db_session_with_containers: Session):
|
||||
deleter = ArchivedWorkflowRunDeletion()
|
||||
@ -109,13 +189,23 @@ class TestArchivedWorkflowRunDeletion:
|
||||
self._create_archive_log(db_session_with_containers, run=run2)
|
||||
run_ids = [run1.id, run2.id]
|
||||
|
||||
deleter = ArchivedWorkflowRunDeletion()
|
||||
results = deleter.delete_batch(
|
||||
tenant_ids=[tenant_id],
|
||||
start_date=base_time - timedelta(minutes=1),
|
||||
end_date=base_time + timedelta(minutes=1),
|
||||
limit=2,
|
||||
storage = FakeArchiveStorage(
|
||||
{
|
||||
self._archive_key(run1): self._archive_bundle(run1),
|
||||
self._archive_key(run2): self._archive_bundle(run2),
|
||||
}
|
||||
)
|
||||
deleter = ArchivedWorkflowRunDeletion()
|
||||
with patch(
|
||||
"services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
|
||||
return_value=storage,
|
||||
):
|
||||
results = deleter.delete_batch(
|
||||
tenant_ids=[tenant_id],
|
||||
start_date=base_time - timedelta(minutes=1),
|
||||
end_date=base_time + timedelta(minutes=1),
|
||||
limit=2,
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
assert all(result.success for result in results)
|
||||
@ -133,9 +223,11 @@ class TestArchivedWorkflowRunDeletion:
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
run_id = run.id
|
||||
archive_log = self._create_archive_log(db_session_with_containers, run=run)
|
||||
deleter = ArchivedWorkflowRunDeletion()
|
||||
|
||||
result = deleter._delete_run(run)
|
||||
with self._patch_storage(run):
|
||||
result = deleter._delete_run(run, archive_log)
|
||||
|
||||
assert result.success is True
|
||||
assert result.deleted_counts["runs"] == 1
|
||||
@ -152,9 +244,11 @@ class TestArchivedWorkflowRunDeletion:
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
run_id = run.id
|
||||
archive_log = self._create_archive_log(db_session_with_containers, run=run)
|
||||
deleter = ArchivedWorkflowRunDeletion(dry_run=True)
|
||||
|
||||
result = deleter._delete_run(run)
|
||||
with self._patch_storage(run):
|
||||
result = deleter._delete_run(run, archive_log)
|
||||
|
||||
assert result.success is True
|
||||
assert result.run_id == run_id
|
||||
@ -164,22 +258,33 @@ class TestArchivedWorkflowRunDeletion:
|
||||
|
||||
def test_delete_run_exception_returns_error(self, db_session_with_containers: Session):
|
||||
"""Exception during deletion should return failure result."""
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
tenant_id = str(uuid4())
|
||||
run = self._create_workflow_run(
|
||||
db_session_with_containers,
|
||||
tenant_id=tenant_id,
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
archive_log = self._create_archive_log(db_session_with_containers, run=run)
|
||||
deleter = ArchivedWorkflowRunDeletion(dry_run=False)
|
||||
|
||||
expected_counts = {
|
||||
"runs": 1,
|
||||
"node_executions": 0,
|
||||
"offloads": 0,
|
||||
"app_logs": 0,
|
||||
"trigger_logs": 0,
|
||||
"pauses": 0,
|
||||
"pause_reasons": 0,
|
||||
}
|
||||
with patch.object(deleter, "_get_workflow_run_repo") as mock_get_repo:
|
||||
mock_repo = MagicMock()
|
||||
mock_get_repo.return_value = mock_repo
|
||||
mock_repo.get_archived_log_by_run_id.return_value = archive_log
|
||||
mock_repo.count_runs_with_related.return_value = expected_counts
|
||||
mock_repo.delete_runs_with_related.side_effect = Exception("Database error")
|
||||
|
||||
result = deleter._delete_run(run)
|
||||
with self._patch_storage(run):
|
||||
result = deleter._delete_run(run, archive_log)
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "Database error"
|
||||
@ -197,7 +302,8 @@ class TestArchivedWorkflowRunDeletion:
|
||||
run_id = run.id
|
||||
|
||||
deleter = ArchivedWorkflowRunDeletion()
|
||||
result = deleter.delete_by_run_id(run_id)
|
||||
with self._patch_storage(run):
|
||||
result = deleter.delete_by_run_id(run_id)
|
||||
|
||||
assert result.success is True
|
||||
db_session_with_containers.expunge_all()
|
||||
@ -212,3 +318,48 @@ class TestArchivedWorkflowRunDeletion:
|
||||
|
||||
assert repo1 is repo2
|
||||
assert deleter.workflow_run_repo is repo1
|
||||
|
||||
def test_delete_run_fails_when_archive_object_missing(self, db_session_with_containers: Session):
|
||||
tenant_id = str(uuid4())
|
||||
run = self._create_workflow_run(
|
||||
db_session_with_containers,
|
||||
tenant_id=tenant_id,
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
archive_log = self._create_archive_log(db_session_with_containers, run=run)
|
||||
deleter = ArchivedWorkflowRunDeletion()
|
||||
storage = FakeArchiveStorage({})
|
||||
|
||||
with patch(
|
||||
"services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
|
||||
return_value=storage,
|
||||
):
|
||||
result = deleter._delete_run(run, archive_log)
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == f"Archive bundle not found: {self._archive_key(run)}"
|
||||
db_session_with_containers.expire_all()
|
||||
assert db_session_with_containers.get(WorkflowRun, run.id) is not None
|
||||
|
||||
def test_delete_run_fails_when_manifest_count_differs_from_live_rows(self, db_session_with_containers: Session):
|
||||
tenant_id = str(uuid4())
|
||||
run = self._create_workflow_run(
|
||||
db_session_with_containers,
|
||||
tenant_id=tenant_id,
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
archive_log = self._create_archive_log(db_session_with_containers, run=run)
|
||||
bundle = self._archive_bundle(run, workflow_run_rows=0)
|
||||
storage = FakeArchiveStorage({self._archive_key(run): bundle})
|
||||
deleter = ArchivedWorkflowRunDeletion()
|
||||
|
||||
with patch(
|
||||
"services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
|
||||
return_value=storage,
|
||||
):
|
||||
result = deleter._delete_run(run, archive_log)
|
||||
|
||||
assert result.success is False
|
||||
assert "Archive row count mismatch before delete" in str(result.error)
|
||||
db_session_with_containers.expire_all()
|
||||
assert db_session_with_containers.get(WorkflowRun, run.id) is not None
|
||||
|
||||
@ -181,7 +181,9 @@ class TestHitTestingService:
|
||||
# ── Response formatting ────────────────────────────────────────────
|
||||
|
||||
@patch("core.rag.datasource.retrieval_service.RetrievalService.format_retrieval_documents")
|
||||
def test_compact_retrieve_response_should_format_correctly(self, mock_format: MagicMock) -> None:
|
||||
def test_compact_retrieve_response_should_format_correctly(
|
||||
self, mock_format: MagicMock, db_session_with_containers: Session
|
||||
) -> None:
|
||||
query = "test query"
|
||||
mock_doc = MagicMock(spec=Document)
|
||||
|
||||
@ -189,7 +191,9 @@ class TestHitTestingService:
|
||||
mock_record.model_dump.return_value = {"content": "formatted content"}
|
||||
mock_format.return_value = [mock_record]
|
||||
|
||||
response = _RetrieveResponse.model_validate(HitTestingService.compact_retrieve_response(query, [mock_doc]))
|
||||
response = _RetrieveResponse.model_validate(
|
||||
HitTestingService.compact_retrieve_response(db_session_with_containers, query, [mock_doc])
|
||||
)
|
||||
|
||||
assert response.query.content == query
|
||||
assert len(response.records) == 1
|
||||
@ -242,6 +246,7 @@ class TestHitTestingService:
|
||||
|
||||
response = _RetrieveResponse.model_validate(
|
||||
HitTestingService.external_retrieve(
|
||||
db_session_with_containers,
|
||||
dataset=dataset,
|
||||
query='test "query"',
|
||||
account=account,
|
||||
@ -269,7 +274,9 @@ class TestHitTestingService:
|
||||
dataset = _create_dataset(db_session_with_containers, provider="vendor")
|
||||
account = MagicMock()
|
||||
|
||||
response = _RetrieveResponse.model_validate(HitTestingService.external_retrieve(dataset, "test query", account))
|
||||
response = _RetrieveResponse.model_validate(
|
||||
HitTestingService.external_retrieve(db_session_with_containers, dataset, "test query", account)
|
||||
)
|
||||
|
||||
assert response.query.content == "test query"
|
||||
assert response.records == []
|
||||
@ -292,6 +299,7 @@ class TestHitTestingService:
|
||||
|
||||
response = _RetrieveResponse.model_validate(
|
||||
HitTestingService.retrieve(
|
||||
db_session_with_containers,
|
||||
dataset=dataset,
|
||||
query="test query",
|
||||
account=account,
|
||||
@ -320,7 +328,11 @@ class TestHitTestingService:
|
||||
|
||||
retrieval_model = {
|
||||
"search_method": "semantic_search",
|
||||
"metadata_filtering_conditions": {"some": "condition"},
|
||||
"metadata_filtering_conditions": {
|
||||
"conditions": [
|
||||
{"name": "category", "comparison_operator": "is", "value": "test"},
|
||||
],
|
||||
},
|
||||
"top_k": 5,
|
||||
"reranking_enable": False,
|
||||
"score_threshold_enabled": False,
|
||||
@ -330,6 +342,7 @@ class TestHitTestingService:
|
||||
mock_retrieve.return_value = retrieved_documents
|
||||
|
||||
HitTestingService.retrieve(
|
||||
db_session_with_containers,
|
||||
dataset=dataset,
|
||||
query="test query",
|
||||
account=account,
|
||||
@ -352,7 +365,11 @@ class TestHitTestingService:
|
||||
|
||||
retrieval_model = {
|
||||
"search_method": "semantic_search",
|
||||
"metadata_filtering_conditions": {"some": "condition"},
|
||||
"metadata_filtering_conditions": {
|
||||
"conditions": [
|
||||
{"name": "category", "comparison_operator": "is", "value": "test"},
|
||||
],
|
||||
},
|
||||
"top_k": 5,
|
||||
"reranking_enable": False,
|
||||
"score_threshold_enabled": False,
|
||||
@ -362,6 +379,7 @@ class TestHitTestingService:
|
||||
|
||||
response = _RetrieveResponse.model_validate(
|
||||
HitTestingService.retrieve(
|
||||
db_session_with_containers,
|
||||
dataset=dataset,
|
||||
query="test query",
|
||||
account=account,
|
||||
@ -393,6 +411,7 @@ class TestHitTestingService:
|
||||
mock_retrieve.return_value = retrieved_documents
|
||||
|
||||
HitTestingService.retrieve(
|
||||
db_session_with_containers,
|
||||
dataset=dataset,
|
||||
query="test query",
|
||||
account=account,
|
||||
@ -452,6 +471,7 @@ class TestHitTestingService:
|
||||
mock_retrieve.return_value = retrieved_documents
|
||||
|
||||
HitTestingService.retrieve(
|
||||
db_session_with_containers,
|
||||
dataset=dataset,
|
||||
query="test query",
|
||||
account=account,
|
||||
@ -477,11 +497,15 @@ class TestHitTestingService:
|
||||
"doc_metadata": {"source": "manual"},
|
||||
}
|
||||
|
||||
def test_dump_retrieval_records_returns_dumped_records_without_document_ids(self) -> None:
|
||||
def test_dump_retrieval_records_returns_dumped_records_without_document_ids(
|
||||
self, db_session_with_containers: Session
|
||||
) -> None:
|
||||
segment = _build_segment(document_id="")
|
||||
record = RetrievalSegments.model_validate({"segment": segment, "score": 0.95})
|
||||
|
||||
records = _DUMPED_RETRIEVAL_RECORDS.validate_python(HitTestingService._dump_retrieval_records([record]))
|
||||
records = _DUMPED_RETRIEVAL_RECORDS.validate_python(
|
||||
HitTestingService._dump_retrieval_records(db_session_with_containers, [record])
|
||||
)
|
||||
|
||||
assert len(records) == 1
|
||||
assert records[0].segment.id == segment.id
|
||||
@ -493,7 +517,9 @@ class TestHitTestingService:
|
||||
segment = _create_segment(db_session_with_containers, document=document)
|
||||
record = RetrievalSegments.model_validate({"segment": segment, "score": 0.9})
|
||||
|
||||
records = _DUMPED_RETRIEVAL_RECORDS.validate_python(HitTestingService._dump_retrieval_records([record]))
|
||||
records = _DUMPED_RETRIEVAL_RECORDS.validate_python(
|
||||
HitTestingService._dump_retrieval_records(db_session_with_containers, [record])
|
||||
)
|
||||
|
||||
assert len(records) == 1
|
||||
dumped_segment = records[0].segment
|
||||
@ -515,7 +541,7 @@ class TestHitTestingService:
|
||||
segment = _create_segment(db_session_with_containers)
|
||||
record = RetrievalSegments.model_validate({"segment": segment, "score": 0.95})
|
||||
|
||||
result = HitTestingService._dump_retrieval_records([record])
|
||||
result = HitTestingService._dump_retrieval_records(db_session_with_containers, [record])
|
||||
|
||||
assert result == []
|
||||
assert "Skipping hit-testing records with missing documents" in caplog.text
|
||||
|
||||
@ -14,6 +14,7 @@ from dify_agent.layers.dify_plugin import (
|
||||
DifyPluginToolConfig,
|
||||
DifyPluginToolsLayerConfig,
|
||||
)
|
||||
from dify_agent.layers.drive import DifyDriveLayerConfig
|
||||
from dify_agent.layers.execution_context import DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, DifyExecutionContextLayerConfig
|
||||
from dify_agent.layers.knowledge import DIFY_KNOWLEDGE_BASE_LAYER_TYPE_ID, DifyKnowledgeBaseLayerConfig
|
||||
from dify_agent.layers.output import DIFY_OUTPUT_LAYER_TYPE_ID
|
||||
@ -42,7 +43,7 @@ from clients.agent_backend import (
|
||||
extract_runtime_layer_specs,
|
||||
redact_for_agent_backend_log,
|
||||
)
|
||||
from clients.agent_backend.request_builder import DIFY_SHELL_LAYER_ID
|
||||
from clients.agent_backend.request_builder import DIFY_DRIVE_LAYER_ID, DIFY_SHELL_LAYER_ID
|
||||
|
||||
|
||||
def _run_input() -> AgentBackendWorkflowNodeRunInput:
|
||||
@ -331,6 +332,22 @@ def test_workflow_request_builder_adds_shell_layer_when_include_shell():
|
||||
assert shell_config.env[0].name == "PROJECT_NAME"
|
||||
|
||||
|
||||
def test_workflow_request_builder_binds_shell_to_drive_when_configured():
|
||||
run_input = _run_input()
|
||||
run_input.include_shell = True
|
||||
run_input.drive_config = DifyDriveLayerConfig(drive_ref="agent-agent-1")
|
||||
|
||||
request = AgentBackendRunRequestBuilder().build_for_workflow_node(run_input)
|
||||
layers = {layer.name: layer for layer in request.composition.layers}
|
||||
layer_names = [layer.name for layer in request.composition.layers]
|
||||
|
||||
assert layers[DIFY_SHELL_LAYER_ID].deps == {
|
||||
"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID,
|
||||
"drive": DIFY_DRIVE_LAYER_ID,
|
||||
}
|
||||
assert layer_names.index(DIFY_DRIVE_LAYER_ID) < layer_names.index(DIFY_SHELL_LAYER_ID)
|
||||
|
||||
|
||||
def test_agent_app_request_builder_omits_shell_layer_by_default():
|
||||
request = AgentBackendRunRequestBuilder().build_for_agent_app(_agent_app_input())
|
||||
assert DIFY_SHELL_LAYER_ID not in {layer.name for layer in request.composition.layers}
|
||||
@ -350,6 +367,21 @@ def test_agent_app_request_builder_adds_shell_layer_when_include_shell():
|
||||
assert shell_config.env[0].name == "APP_ENV"
|
||||
|
||||
|
||||
def test_agent_app_request_builder_binds_shell_to_drive_when_configured():
|
||||
run_input = _agent_app_input(include_shell=True)
|
||||
run_input.drive_config = DifyDriveLayerConfig(drive_ref="agent-agent-1")
|
||||
|
||||
request = AgentBackendRunRequestBuilder().build_for_agent_app(run_input)
|
||||
layers = {layer.name: layer for layer in request.composition.layers}
|
||||
layer_names = [layer.name for layer in request.composition.layers]
|
||||
|
||||
assert layers[DIFY_SHELL_LAYER_ID].deps == {
|
||||
"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID,
|
||||
"drive": DIFY_DRIVE_LAYER_ID,
|
||||
}
|
||||
assert layer_names.index(DIFY_DRIVE_LAYER_ID) < layer_names.index(DIFY_SHELL_LAYER_ID)
|
||||
|
||||
|
||||
def test_agent_app_request_builder_adds_knowledge_layer_when_configured():
|
||||
run_input = _agent_app_input()
|
||||
run_input.knowledge = DifyKnowledgeBaseLayerConfig.model_validate(
|
||||
|
||||
@ -27,6 +27,7 @@ from controllers.console.agent.roster import (
|
||||
AgentAppApi,
|
||||
AgentAppCopyApi,
|
||||
AgentAppListApi,
|
||||
AgentDebugConversationRefreshApi,
|
||||
AgentInviteOptionsApi,
|
||||
AgentLogMessagesApi,
|
||||
AgentLogsApi,
|
||||
@ -158,6 +159,7 @@ def test_agent_v2_console_routes_are_agent_id_first() -> None:
|
||||
"/agent/<uuid:agent_id>/api-enable",
|
||||
"/agent/<uuid:agent_id>/api-keys",
|
||||
"/agent/<uuid:agent_id>/api-keys/<uuid:api_key_id>",
|
||||
"/agent/<uuid:agent_id>/debug-conversation/refresh",
|
||||
"/agent/<uuid:agent_id>/chat-messages",
|
||||
"/agent/<uuid:agent_id>/chat-messages/<string:task_id>/stop",
|
||||
"/agent/<uuid:agent_id>/feedbacks",
|
||||
@ -462,6 +464,7 @@ def test_agent_app_copy_uses_agent_id_and_returns_agent_detail(
|
||||
json={
|
||||
"name": "Iris copy",
|
||||
"description": "Copied",
|
||||
"role": "Copied role",
|
||||
"icon_type": "emoji",
|
||||
"icon": "sparkles",
|
||||
"icon_background": "#fff",
|
||||
@ -477,12 +480,45 @@ def test_agent_app_copy_uses_agent_id_and_returns_agent_detail(
|
||||
"account": current_user,
|
||||
"name": "Iris copy",
|
||||
"description": "Copied",
|
||||
"role": "Copied role",
|
||||
"icon_type": "emoji",
|
||||
"icon": "sparkles",
|
||||
"icon_background": "#fff",
|
||||
}
|
||||
|
||||
|
||||
def test_agent_debug_conversation_refresh_uses_current_user(
|
||||
app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str
|
||||
) -> None:
|
||||
agent_id = "00000000-0000-0000-0000-000000000001"
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class FakeRosterService:
|
||||
def refresh_agent_app_debug_conversation_id(self, **kwargs: object) -> str:
|
||||
captured.update(kwargs)
|
||||
return "new-debug-conversation-id"
|
||||
|
||||
monkeypatch.setattr(roster_controller, "_agent_roster_service", lambda: FakeRosterService())
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/agent/00000000-0000-0000-0000-000000000001/debug-conversation/refresh",
|
||||
method="POST",
|
||||
):
|
||||
response = unwrap(AgentDebugConversationRefreshApi.post)(
|
||||
AgentDebugConversationRefreshApi(),
|
||||
"tenant-1",
|
||||
SimpleNamespace(id=account_id),
|
||||
agent_id,
|
||||
)
|
||||
|
||||
assert response == {"debug_conversation_id": "new-debug-conversation-id"}
|
||||
assert captured == {
|
||||
"tenant_id": "tenant-1",
|
||||
"agent_id": agent_id,
|
||||
"account_id": account_id,
|
||||
}
|
||||
|
||||
|
||||
def test_agent_api_access_uses_agent_id_and_returns_service_api_metadata(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
@ -936,7 +972,7 @@ def test_workflow_composer_get_put_validate_candidates_impact_and_save(
|
||||
"save_workflow_composer",
|
||||
lambda **kwargs: _workflow_composer_response(save_options=[kwargs["payload"].save_strategy.value]),
|
||||
)
|
||||
monkeypatch.setattr(composer_controller.ComposerConfigValidator, "validate_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(composer_controller.ComposerConfigValidator, "validate_publish_payload", lambda payload: None)
|
||||
monkeypatch.setattr(
|
||||
composer_controller.AgentComposerService, "resolve_workflow_node_agent_id", lambda **kwargs: None
|
||||
)
|
||||
@ -1031,7 +1067,7 @@ def test_agent_composer_routes_resolve_app_from_agent_id(
|
||||
"save_agent_app_composer",
|
||||
save_agent_app_composer,
|
||||
)
|
||||
monkeypatch.setattr(composer_controller.ComposerConfigValidator, "validate_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(composer_controller.ComposerConfigValidator, "validate_publish_payload", lambda payload: None)
|
||||
monkeypatch.setattr(
|
||||
composer_controller.AgentComposerService,
|
||||
"collect_validation_findings",
|
||||
|
||||
@ -12,6 +12,7 @@ import io
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from controllers.console.app.agent import (
|
||||
@ -152,27 +153,20 @@ def test_files_commit_validates_upload_and_returns_drive_ref():
|
||||
patch(f"{_MOD}.console_ns") as ns,
|
||||
patch(f"{_MOD}.db") as db_mock,
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
patch(f"{_MOD}.AgentComposerService") as composer,
|
||||
):
|
||||
ns.payload = {"upload_file_id": "0fa6f9bc-3416-4476-8857-a13129704dd9"}
|
||||
db_mock.session.scalar.return_value = upload
|
||||
drive.return_value.commit.return_value = [
|
||||
{"key": "files/sample qna.pdf", "size": 5, "mime_type": "application/pdf"}
|
||||
]
|
||||
composer.add_drive_file_ref.return_value = "ver-2"
|
||||
body, status = raw(AgentDriveFilesApi(), _USER, _APP)
|
||||
|
||||
assert status == 201
|
||||
assert body["file"]["drive_key"] == "files/sample qna.pdf"
|
||||
assert body["file"]["file_id"] == "uf-1"
|
||||
assert body["config_version_id"] == "ver-2"
|
||||
item = drive.return_value.commit.call_args.kwargs["items"][0]
|
||||
assert item.value_owned_by_drive is True
|
||||
assert item.file_ref.kind == "upload_file"
|
||||
file_ref = composer.add_drive_file_ref.call_args.kwargs["file_ref"]
|
||||
assert file_ref.drive_key == "files/sample qna.pdf"
|
||||
assert file_ref.name == "sample qna.pdf"
|
||||
assert composer.add_drive_file_ref.call_args.kwargs["app_id"] == "app-1"
|
||||
|
||||
|
||||
def test_files_by_agent_commit_uses_agent_route_and_ignores_node_id():
|
||||
@ -184,20 +178,16 @@ def test_files_by_agent_commit_uses_agent_route_and_ignores_node_id():
|
||||
patch(f"{_MOD}.console_ns") as ns,
|
||||
patch(f"{_MOD}.db") as db_mock,
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
patch(f"{_MOD}.AgentComposerService") as composer,
|
||||
):
|
||||
ns.payload = {"upload_file_id": "0fa6f9bc-3416-4476-8857-a13129704dd9"}
|
||||
db_mock.session.scalar.return_value = upload
|
||||
drive.return_value.commit.return_value = [
|
||||
{"key": "files/sample.pdf", "size": 5, "mime_type": "application/pdf"}
|
||||
]
|
||||
composer.add_drive_file_ref.return_value = "ver-2"
|
||||
body, status = raw(AgentDriveFilesByAgentApi(), "tenant-1", _USER, "agent-1")
|
||||
|
||||
assert status == 201
|
||||
assert body["config_version_id"] == "ver-2"
|
||||
resolve_app.assert_called_once_with(tenant_id="tenant-1", agent_id="agent-1")
|
||||
assert composer.add_drive_file_ref.call_args.kwargs["node_id"] is None
|
||||
|
||||
|
||||
def test_files_commit_404_when_upload_not_in_tenant():
|
||||
@ -234,13 +224,10 @@ def test_files_commit_resolves_workflow_node_agent():
|
||||
drive.return_value.commit.return_value = [
|
||||
{"key": "files/sample.pdf", "size": 5, "mime_type": "application/pdf"}
|
||||
]
|
||||
composer.add_drive_file_ref.return_value = "ver-2"
|
||||
body, status = raw(AgentDriveFilesApi(), _USER, _WORKFLOW_APP)
|
||||
|
||||
assert status == 201
|
||||
assert body["config_version_id"] == "ver-2"
|
||||
assert drive.return_value.commit.call_args.kwargs["agent_id"] == "wf-agent-1"
|
||||
assert composer.add_drive_file_ref.call_args.kwargs["node_id"] == "agent-node-1"
|
||||
|
||||
|
||||
def test_files_delete_updates_soul_then_drive():
|
||||
@ -250,17 +237,15 @@ def test_files_delete_updates_soul_then_drive():
|
||||
calls: list[str] = []
|
||||
with _json_ctx(method="DELETE", query_string="key=files/sample.pdf"):
|
||||
with (
|
||||
patch(f"{_MOD}.AgentComposerService") as composer,
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
):
|
||||
composer.remove_drive_refs.side_effect = lambda **kw: calls.append("soul") or "ver-2"
|
||||
drive.return_value.delete.side_effect = lambda **kw: calls.append("drive") or ["files/sample.pdf"]
|
||||
drive.return_value.commit.side_effect = lambda **kw: (
|
||||
calls.append("drive") or [{"key": "files/sample.pdf", "removed": True}]
|
||||
)
|
||||
body = raw(AgentDriveFilesApi(), _USER, _APP)
|
||||
|
||||
assert calls == ["soul", "drive"] # soul-first ordering
|
||||
assert body == {"result": "success", "removed_keys": ["files/sample.pdf"], "config_version_id": "ver-2"}
|
||||
assert composer.remove_drive_refs.call_args.kwargs["file_key"] == "files/sample.pdf"
|
||||
assert composer.remove_drive_refs.call_args.kwargs["app_id"] == "app-1"
|
||||
assert calls == ["drive"]
|
||||
assert body == {"result": "success", "removed_keys": ["files/sample.pdf"]}
|
||||
|
||||
|
||||
def test_files_by_agent_delete_uses_agent_route_and_ignores_node_id():
|
||||
@ -268,16 +253,13 @@ def test_files_by_agent_delete_uses_agent_route_and_ignores_node_id():
|
||||
with _json_ctx(method="DELETE", query_string="key=files/sample.pdf&node_id=ignored"):
|
||||
with (
|
||||
patch(f"{_MOD}.resolve_agent_app_model", return_value=_APP) as resolve_app,
|
||||
patch(f"{_MOD}.AgentComposerService") as composer,
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
):
|
||||
composer.remove_drive_refs.return_value = "ver-2"
|
||||
drive.return_value.delete.return_value = ["files/sample.pdf"]
|
||||
drive.return_value.commit.return_value = [{"key": "files/sample.pdf", "removed": True}]
|
||||
body = raw(AgentDriveFilesByAgentApi(), "tenant-1", _USER, "agent-1")
|
||||
|
||||
assert body["config_version_id"] == "ver-2"
|
||||
assert body == {"result": "success", "removed_keys": ["files/sample.pdf"]}
|
||||
resolve_app.assert_called_once_with(tenant_id="tenant-1", agent_id="agent-1")
|
||||
assert composer.remove_drive_refs.call_args.kwargs["node_id"] is None
|
||||
|
||||
|
||||
def test_files_delete_resolves_workflow_node_agent():
|
||||
@ -290,13 +272,11 @@ def test_files_delete_resolves_workflow_node_agent():
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
):
|
||||
composer.resolve_workflow_node_agent_id.return_value = "wf-agent-1"
|
||||
composer.remove_drive_refs.return_value = "ver-2"
|
||||
drive.return_value.delete.return_value = ["files/sample.pdf"]
|
||||
drive.return_value.commit.return_value = [{"key": "files/sample.pdf", "removed": True}]
|
||||
body = raw(AgentDriveFilesApi(), _USER, _WORKFLOW_APP)
|
||||
|
||||
assert body["config_version_id"] == "ver-2"
|
||||
assert drive.return_value.delete.call_args.kwargs["agent_id"] == "wf-agent-1"
|
||||
assert composer.remove_drive_refs.call_args.kwargs["node_id"] == "agent-node-1"
|
||||
assert body == {"result": "success", "removed_keys": ["files/sample.pdf"]}
|
||||
assert drive.return_value.commit.call_args.kwargs["agent_id"] == "wf-agent-1"
|
||||
|
||||
|
||||
def test_files_delete_survives_drive_failure():
|
||||
@ -305,14 +285,11 @@ def test_files_delete_survives_drive_failure():
|
||||
raw = _raw(AgentDriveFilesApi.delete)
|
||||
with _json_ctx(method="DELETE", query_string="key=files/sample.pdf"):
|
||||
with (
|
||||
patch(f"{_MOD}.AgentComposerService") as composer,
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
):
|
||||
composer.remove_drive_refs.return_value = "ver-2"
|
||||
drive.return_value.delete.side_effect = RuntimeError("storage down")
|
||||
body = raw(AgentDriveFilesApi(), _USER, _APP)
|
||||
# soul already updated; drive cleanup is best-effort and retryable
|
||||
assert body == {"result": "success", "removed_keys": [], "config_version_id": "ver-2"}
|
||||
drive.return_value.commit.side_effect = RuntimeError("storage down")
|
||||
with pytest.raises(RuntimeError, match="storage down"):
|
||||
raw(AgentDriveFilesApi(), _USER, _APP)
|
||||
|
||||
|
||||
def test_skill_delete_uses_slug_prefix_and_is_idempotent():
|
||||
@ -321,17 +298,18 @@ def test_skill_delete_uses_slug_prefix_and_is_idempotent():
|
||||
raw = _raw(AgentSkillApi.delete)
|
||||
with _json_ctx(method="DELETE"):
|
||||
with (
|
||||
patch(f"{_MOD}.AgentComposerService") as composer,
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
):
|
||||
composer.remove_drive_refs.return_value = None # ref already gone
|
||||
drive.return_value.delete.return_value = []
|
||||
drive.return_value.commit.return_value = [
|
||||
{"key": "tender-analyzer/SKILL.md", "removed": True},
|
||||
{"key": "tender-analyzer/.DIFY-SKILL-FULL.zip", "removed": True},
|
||||
]
|
||||
body = raw(AgentSkillApi(), _USER, _APP, "tender-analyzer")
|
||||
|
||||
assert body == {"result": "success", "removed_keys": [], "config_version_id": None}
|
||||
assert drive.return_value.delete.call_args.kwargs["prefix"] == "tender-analyzer/"
|
||||
assert composer.remove_drive_refs.call_args.kwargs["skill_slug"] == "tender-analyzer"
|
||||
assert composer.remove_drive_refs.call_args.kwargs["app_id"] == "app-1"
|
||||
assert body == {
|
||||
"result": "success",
|
||||
"removed_keys": ["tender-analyzer/SKILL.md", "tender-analyzer/.DIFY-SKILL-FULL.zip"],
|
||||
}
|
||||
|
||||
|
||||
def test_skill_delete_by_agent_uses_agent_route():
|
||||
@ -339,16 +317,13 @@ def test_skill_delete_by_agent_uses_agent_route():
|
||||
with _json_ctx(method="DELETE", query_string="node_id=ignored"):
|
||||
with (
|
||||
patch(f"{_MOD}.resolve_agent_app_model", return_value=_APP) as resolve_app,
|
||||
patch(f"{_MOD}.AgentComposerService") as composer,
|
||||
patch(f"{_MOD}.AgentDriveService") as drive,
|
||||
):
|
||||
composer.remove_drive_refs.return_value = "ver-2"
|
||||
drive.return_value.delete.return_value = ["tender-analyzer/SKILL.md"]
|
||||
drive.return_value.commit.return_value = [{"key": "tender-analyzer/SKILL.md", "removed": True}]
|
||||
body = raw(AgentSkillByAgentApi(), "tenant-1", _USER, "agent-1", "tender-analyzer")
|
||||
|
||||
assert body["config_version_id"] == "ver-2"
|
||||
assert body == {"result": "success", "removed_keys": ["tender-analyzer/SKILL.md"]}
|
||||
resolve_app.assert_called_once_with(tenant_id="tenant-1", agent_id="agent-1")
|
||||
assert composer.remove_drive_refs.call_args.kwargs["node_id"] is None
|
||||
|
||||
|
||||
def test_skill_delete_rejects_path_like_slug():
|
||||
|
||||
@ -8,12 +8,13 @@ controller's request parsing + error mapping, not auth (tested separately).
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from controllers.inner_api.plugin.agent_drive import AgentDriveCommitApi, AgentDriveManifestApi
|
||||
from controllers.inner_api.plugin.agent_drive import AgentDriveCommitApi, AgentDriveManifestApi, AgentDriveSkillsApi
|
||||
from services.agent_drive_service import AgentDriveError
|
||||
|
||||
_MOD = "controllers.inner_api.plugin.agent_drive"
|
||||
@ -52,6 +53,41 @@ def test_manifest_bad_drive_ref_is_400():
|
||||
assert body["code"] == "invalid_drive_ref"
|
||||
|
||||
|
||||
def test_skills_requires_tenant_id_and_returns_items():
|
||||
raw = _raw(AgentDriveSkillsApi.get)
|
||||
|
||||
with app.test_request_context("/"):
|
||||
body, status = raw(AgentDriveSkillsApi(), "agent-agent-1")
|
||||
assert status == 400
|
||||
assert body["code"] == "missing_tenant_id"
|
||||
|
||||
with app.test_request_context("/?tenant_id=tenant-1"):
|
||||
with patch(f"{_MOD}.AgentDriveService") as svc:
|
||||
svc.return_value.list_skills.return_value = [
|
||||
{
|
||||
"path": "tender-analyzer",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"archive_key": None,
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
}
|
||||
]
|
||||
result = raw(AgentDriveSkillsApi(), "agent-agent-1")
|
||||
|
||||
assert result == {
|
||||
"items": [
|
||||
{
|
||||
"path": "tender-analyzer",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"archive_key": None,
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
}
|
||||
]
|
||||
}
|
||||
assert svc.return_value.list_skills.call_args.kwargs == {"tenant_id": "tenant-1", "agent_id": "agent-1"}
|
||||
|
||||
|
||||
def test_commit_parses_body_and_returns_items():
|
||||
raw = _raw(AgentDriveCommitApi.post)
|
||||
payload = {
|
||||
@ -60,11 +96,35 @@ def test_commit_parses_body_and_returns_items():
|
||||
"items": [{"key": "a.txt", "file_ref": {"kind": "tool_file", "id": "tf-1"}}],
|
||||
}
|
||||
with app.test_request_context("/", method="POST", json=payload):
|
||||
with patch(f"{_MOD}.AgentDriveService") as svc:
|
||||
with (
|
||||
patch(f"{_MOD}.get_user", return_value=SimpleNamespace(id="user-1")) as get_user,
|
||||
patch(f"{_MOD}.AgentDriveService") as svc,
|
||||
):
|
||||
svc.return_value.commit.return_value = [{"key": "a.txt"}]
|
||||
result = raw(AgentDriveCommitApi(), "agent-agent-1")
|
||||
assert result == {"items": [{"key": "a.txt"}]}
|
||||
assert get_user.call_args.args == ("tenant-1", "user-1")
|
||||
assert svc.return_value.commit.call_args.kwargs["agent_id"] == "agent-1"
|
||||
assert svc.return_value.commit.call_args.kwargs["user_id"] == "user-1"
|
||||
|
||||
|
||||
def test_commit_canonicalizes_user_before_service_call():
|
||||
raw = _raw(AgentDriveCommitApi.post)
|
||||
payload = {
|
||||
"tenant_id": "tenant-1",
|
||||
"user_id": "session-1",
|
||||
"items": [{"key": "a.txt", "file_ref": {"kind": "tool_file", "id": "tf-1"}}],
|
||||
}
|
||||
with app.test_request_context("/", method="POST", json=payload):
|
||||
with (
|
||||
patch(f"{_MOD}.get_user", return_value=SimpleNamespace(id="end-user-1")),
|
||||
patch(f"{_MOD}.AgentDriveService") as svc,
|
||||
):
|
||||
svc.return_value.commit.return_value = [{"key": "a.txt"}]
|
||||
result = raw(AgentDriveCommitApi(), "agent-agent-1")
|
||||
|
||||
assert result == {"items": [{"key": "a.txt"}]}
|
||||
assert svc.return_value.commit.call_args.kwargs["user_id"] == "end-user-1"
|
||||
|
||||
|
||||
def test_commit_invalid_body_is_400():
|
||||
@ -83,13 +143,16 @@ def test_commit_maps_service_error():
|
||||
"items": [{"key": "a.txt", "file_ref": {"kind": "tool_file", "id": "tf-1"}}],
|
||||
}
|
||||
with app.test_request_context("/", method="POST", json=payload):
|
||||
with patch(f"{_MOD}.AgentDriveService") as svc:
|
||||
with (
|
||||
patch(f"{_MOD}.get_user", return_value=SimpleNamespace(id="user-1")),
|
||||
patch(f"{_MOD}.AgentDriveService") as svc,
|
||||
):
|
||||
svc.return_value.commit.side_effect = AgentDriveError("source_not_found", "nope", status_code=404)
|
||||
body, status = raw(AgentDriveCommitApi(), "agent-agent-1")
|
||||
assert status == 404
|
||||
assert body["code"] == "source_not_found"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api_cls", [AgentDriveManifestApi, AgentDriveCommitApi])
|
||||
@pytest.mark.parametrize("api_cls", [AgentDriveManifestApi, AgentDriveSkillsApi, AgentDriveCommitApi])
|
||||
def test_endpoints_have_handlers(api_cls):
|
||||
assert callable(getattr(api_cls(), "get", None) or getattr(api_cls(), "post", None))
|
||||
|
||||
@ -1,5 +1,7 @@
|
||||
"""Unit tests for the inner knowledge retrieval controller."""
|
||||
"""Unit tests for the plugin inner knowledge retrieval controller."""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@ -53,31 +55,38 @@ def _payload() -> dict[str, object]:
|
||||
}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _plugin_inner_auth() -> Iterator[None]:
|
||||
with (
|
||||
patch("configs.dify_config.PLUGIN_DAEMON_KEY", "plugin-daemon-key"),
|
||||
patch("configs.dify_config.INNER_API_KEY_FOR_PLUGIN", "inner-key"),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
class TestInnerKnowledgeRetrieveApi:
|
||||
def test_post_returns_401_when_api_key_missing(self, inner_api_app: Flask):
|
||||
with patch("configs.dify_config.INNER_API", True):
|
||||
def test_post_returns_404_when_api_key_missing(self, inner_api_app: Flask):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
headers=_headers(api_key=None),
|
||||
)
|
||||
|
||||
assert response.status_code == 401
|
||||
assert response.get_json()["code"] == "inner_api_unauthorized"
|
||||
assert response.status_code == 404
|
||||
|
||||
def test_post_returns_401_when_api_key_invalid(self, inner_api_app: Flask):
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
def test_post_returns_404_when_api_key_invalid(self, inner_api_app: Flask):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
headers=_headers(api_key="wrong-key"),
|
||||
)
|
||||
|
||||
assert response.status_code == 401
|
||||
assert response.get_json()["code"] == "inner_api_unauthorized"
|
||||
assert response.status_code == 404
|
||||
|
||||
def test_post_returns_400_for_invalid_body(self, inner_api_app: Flask):
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json={"caller": {"tenant_id": "tenant-1"}},
|
||||
@ -91,7 +100,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
def test_post_returns_404_for_service_not_found_error(self, mock_retrieve, inner_api_app: Flask):
|
||||
mock_retrieve.side_effect = InnerKnowledgeRetrieveAppNotFoundError("app missing")
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
@ -105,7 +114,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
def test_post_returns_403_for_service_forbidden_error(self, mock_retrieve, inner_api_app: Flask):
|
||||
mock_retrieve.side_effect = InnerKnowledgeRetrieveDatasetTenantMismatchError("wrong tenant")
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
@ -119,7 +128,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
def test_post_returns_422_for_retrieval_config_value_error(self, mock_retrieve, inner_api_app: Flask):
|
||||
mock_retrieve.side_effect = ValueError("invalid reranking config")
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
@ -133,7 +142,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
def test_post_returns_429_for_rate_limit_error(self, mock_retrieve, inner_api_app: Flask):
|
||||
mock_retrieve.side_effect = RateLimitExceededError("knowledge rate limited")
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
@ -147,7 +156,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
payload = _payload()
|
||||
payload["metadata_filtering"] = {"mode": "manual"}
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=payload,
|
||||
@ -161,7 +170,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
payload = _payload()
|
||||
payload["metadata_filtering"] = {"mode": "automatic"}
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=payload,
|
||||
@ -175,7 +184,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
def test_post_returns_502_for_external_knowledge_failure(self, mock_retrieve, inner_api_app: Flask):
|
||||
mock_retrieve.side_effect = ExternalKnowledgeRetrievalError("upstream failed")
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
@ -219,7 +228,7 @@ class TestInnerKnowledgeRetrieveApi:
|
||||
),
|
||||
)
|
||||
|
||||
with patch("configs.dify_config.INNER_API", True), patch("configs.dify_config.INNER_API_KEY", "inner-key"):
|
||||
with _plugin_inner_auth():
|
||||
response = inner_api_app.test_client().post(
|
||||
"/inner/api/knowledge/retrieve",
|
||||
json=_payload(),
|
||||
|
||||
10
api/tests/unit_tests/controllers/openapi/_mode_constants.py
Normal file
10
api/tests/unit_tests/controllers/openapi/_mode_constants.py
Normal file
@ -0,0 +1,10 @@
|
||||
"""Shared mode lists for the openapi app-list query tests.
|
||||
|
||||
Single source so adding/removing a listable app type is a one-line change
|
||||
across every query-validator test.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
LISTABLE_MODES = ["completion", "chat", "advanced-chat", "workflow", "agent-chat"]
|
||||
NON_LISTABLE_MODES = ["rag-pipeline", "channel", "agent"]
|
||||
@ -4,7 +4,7 @@ Runs against the model directly, not the HTTP layer. Pins:
|
||||
- defaults match the plan (page=1, limit=20).
|
||||
- workspace_id is required.
|
||||
- numeric bounds enforced (page >= 1, limit in [1, MAX_PAGE_LIMIT]).
|
||||
- mode validates against the AppMode enum.
|
||||
- mode validates against the SupportedAppType enum (listable app types only).
|
||||
- name has a length cap.
|
||||
"""
|
||||
|
||||
@ -16,10 +16,14 @@ from pydantic import ValidationError
|
||||
from controllers.openapi._models import MAX_PAGE_LIMIT
|
||||
from controllers.openapi.apps import AppListQuery
|
||||
|
||||
from ._mode_constants import LISTABLE_MODES, NON_LISTABLE_MODES
|
||||
|
||||
WS_ID = "00000000-0000-0000-0000-000000000001"
|
||||
|
||||
|
||||
def test_defaults():
|
||||
q = AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001"})
|
||||
assert q.workspace_id == "00000000-0000-0000-0000-000000000001"
|
||||
q = AppListQuery.model_validate({"workspace_id": WS_ID})
|
||||
assert q.workspace_id == WS_ID
|
||||
assert q.page == 1
|
||||
assert q.limit == 20
|
||||
assert q.mode is None
|
||||
@ -33,64 +37,71 @@ def test_workspace_id_required():
|
||||
|
||||
def test_page_must_be_positive():
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "page": 0})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "page": 0})
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "page": -1})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "page": -1})
|
||||
|
||||
|
||||
def test_page_rejects_non_integer_string():
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "page": "abc"})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "page": "abc"})
|
||||
|
||||
|
||||
def test_limit_must_be_positive():
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "limit": 0})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "limit": 0})
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "limit": -1})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "limit": -1})
|
||||
|
||||
|
||||
def test_limit_caps_at_max_page_limit():
|
||||
# Boundary accepts.
|
||||
q = AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "limit": MAX_PAGE_LIMIT})
|
||||
q = AppListQuery.model_validate({"workspace_id": WS_ID, "limit": MAX_PAGE_LIMIT})
|
||||
assert q.limit == MAX_PAGE_LIMIT
|
||||
|
||||
# Just over rejects.
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate(
|
||||
{"workspace_id": "00000000-0000-0000-0000-000000000001", "limit": MAX_PAGE_LIMIT + 1}
|
||||
)
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "limit": MAX_PAGE_LIMIT + 1})
|
||||
|
||||
|
||||
def test_mode_whitelisted_against_app_mode():
|
||||
# Valid mode passes.
|
||||
q = AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "mode": "chat"})
|
||||
@pytest.mark.parametrize("mode", LISTABLE_MODES)
|
||||
def test_mode_accepts_listable_app_types(mode: str):
|
||||
q = AppListQuery.model_validate({"workspace_id": WS_ID, "mode": mode})
|
||||
assert q.mode is not None
|
||||
assert q.mode.value == "chat"
|
||||
assert q.mode.value == mode
|
||||
|
||||
# Invalid mode rejects.
|
||||
|
||||
def test_mode_rejects_unknown_value():
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "mode": "not-a-mode"})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "mode": "not-a-mode"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", NON_LISTABLE_MODES)
|
||||
def test_mode_rejects_non_listable_app_modes(mode: str):
|
||||
"""rag-pipeline (a knowledge Pipeline), channel (unused) and agent (roster-owned)
|
||||
are AppMode members but not standalone listable apps — the `app` face rejects them."""
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "mode": mode})
|
||||
|
||||
|
||||
def test_name_length_capped():
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "name": "x" * 200})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "name": "x" * 200})
|
||||
with pytest.raises(ValidationError):
|
||||
AppListQuery.model_validate({"workspace_id": "00000000-0000-0000-0000-000000000001", "name": "x" * 201})
|
||||
AppListQuery.model_validate({"workspace_id": WS_ID, "name": "x" * 201})
|
||||
|
||||
|
||||
def test_all_fields_accept_valid_values():
|
||||
"""Pin the happy-path acceptance for every field in one place."""
|
||||
q = AppListQuery.model_validate(
|
||||
{
|
||||
"workspace_id": "00000000-0000-0000-0000-000000000001",
|
||||
"workspace_id": WS_ID,
|
||||
"page": 5,
|
||||
"limit": 50,
|
||||
"mode": "workflow",
|
||||
"name": "search",
|
||||
}
|
||||
)
|
||||
assert q.workspace_id == "00000000-0000-0000-0000-000000000001"
|
||||
assert q.workspace_id == WS_ID
|
||||
assert q.page == 5
|
||||
assert q.limit == 50
|
||||
assert q.mode is not None
|
||||
|
||||
@ -10,9 +10,11 @@ import pytest
|
||||
|
||||
from controllers.openapi.apps import ( # pyright: ignore[reportPrivateUsage]
|
||||
_EMPTY_PARAMETERS,
|
||||
_is_listable,
|
||||
parameters_payload,
|
||||
)
|
||||
from controllers.service_api.app.error import AppUnavailableError
|
||||
from models.model import AppMode
|
||||
|
||||
|
||||
def _fake_app(**overrides):
|
||||
@ -53,3 +55,16 @@ def test_empty_parameters_constant_matches_describe_fallback_shape():
|
||||
assert _EMPTY_PARAMETERS["opening_statement"] is None
|
||||
assert _EMPTY_PARAMETERS["file_upload"] is None
|
||||
assert _EMPTY_PARAMETERS["system_parameters"] == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode",
|
||||
[AppMode.COMPLETION, AppMode.CHAT, AppMode.ADVANCED_CHAT, AppMode.WORKFLOW, AppMode.AGENT_CHAT],
|
||||
)
|
||||
def test_is_listable_accepts_supported_app_types(mode):
|
||||
assert _is_listable(_fake_app(mode=mode)) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.CHANNEL, AppMode.RAG_PIPELINE])
|
||||
def test_is_listable_hides_non_app_modes(mode):
|
||||
assert _is_listable(_fake_app(mode=mode)) is False
|
||||
|
||||
@ -13,6 +13,8 @@ from pydantic import ValidationError
|
||||
|
||||
from controllers.openapi.apps_permitted_external import PermittedExternalAppsListQuery
|
||||
|
||||
from ._mode_constants import NON_LISTABLE_MODES
|
||||
|
||||
|
||||
def test_query_defaults_match_apps_list():
|
||||
q = PermittedExternalAppsListQuery.model_validate({})
|
||||
@ -36,11 +38,18 @@ def test_query_rejects_tag():
|
||||
PermittedExternalAppsListQuery.model_validate({"tag": "prod"})
|
||||
|
||||
|
||||
def test_query_validates_mode_against_app_mode():
|
||||
def test_query_validates_mode_against_supported_app_type():
|
||||
with pytest.raises(ValidationError):
|
||||
PermittedExternalAppsListQuery.model_validate({"mode": "not-a-mode"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", NON_LISTABLE_MODES)
|
||||
def test_query_rejects_non_listable_app_modes(mode: str):
|
||||
"""Non-app runtime modes and roster-owned agent are not listable here."""
|
||||
with pytest.raises(ValidationError):
|
||||
PermittedExternalAppsListQuery.model_validate({"mode": mode})
|
||||
|
||||
|
||||
def test_query_clamps_limit_at_max():
|
||||
with pytest.raises(ValidationError):
|
||||
PermittedExternalAppsListQuery.model_validate({"limit": 500})
|
||||
|
||||
@ -0,0 +1,24 @@
|
||||
"""Unit tests for SupportedAppType — the listable subset of AppMode that the
|
||||
openapi `app` face (`get app`) exposes and the CLI `--mode` whitelist derives from.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from controllers.openapi._models import SUPPORTED_APP_TYPES, SupportedAppType
|
||||
from models.model import AppMode
|
||||
|
||||
|
||||
def test_supported_app_type_is_the_listable_subset_of_app_mode():
|
||||
"""SupportedAppType (and the derived SUPPORTED_APP_TYPES tuple) is exactly the
|
||||
curated, listable subset of AppMode; non-app/runtime modes stay out."""
|
||||
assert {t.value for t in SupportedAppType} == {
|
||||
"completion",
|
||||
"chat",
|
||||
"advanced-chat",
|
||||
"workflow",
|
||||
"agent-chat",
|
||||
}
|
||||
assert set(SUPPORTED_APP_TYPES) <= set(AppMode)
|
||||
assert AppMode.AGENT not in SUPPORTED_APP_TYPES
|
||||
assert AppMode.RAG_PIPELINE not in SUPPORTED_APP_TYPES
|
||||
assert AppMode.CHANNEL not in SUPPORTED_APP_TYPES
|
||||
@ -226,19 +226,8 @@ class TestAgentAppRuntimeRequestBuilder:
|
||||
|
||||
|
||||
def _soul_with_model_and_skill() -> AgentSoulConfig:
|
||||
from models.agent_config_entities import AgentSkillRefConfig
|
||||
|
||||
soul = _soul_with_model()
|
||||
soul.skills_files.skills = [
|
||||
AgentSkillRefConfig.model_validate(
|
||||
{
|
||||
"id": "abc",
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
}
|
||||
)
|
||||
]
|
||||
soul.prompt.system_prompt = "Use [§skill:tender-analyzer%2FSKILL.md:Tender Analyzer§]"
|
||||
return soul
|
||||
|
||||
|
||||
@ -247,6 +236,28 @@ class TestAgentAppDriveLayer:
|
||||
monkeypatch.setattr(
|
||||
"core.app.apps.agent_app.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [
|
||||
{
|
||||
"path": "tender-analyzer",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"archive_key": None,
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
"size": 1,
|
||||
"mime_type": "text/markdown",
|
||||
"hash": None,
|
||||
"created_at": 1,
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [
|
||||
{"key": "tender-analyzer/SKILL.md", "is_skill": True}
|
||||
],
|
||||
)
|
||||
builder = AgentAppRuntimeRequestBuilder(
|
||||
credentials_provider=_FakeCredentialsProvider(),
|
||||
plugin_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type]
|
||||
@ -256,12 +267,42 @@ class TestAgentAppDriveLayer:
|
||||
|
||||
drive = next(layer for layer in result.request.composition.layers if layer.name == "drive")
|
||||
assert drive.type == "dify.drive"
|
||||
assert drive.deps == {"execution_context": "execution_context"}
|
||||
assert drive.config.drive_ref == "agent-agent-1"
|
||||
assert [skill.skill_md_key for skill in drive.config.skills] == ["tender-analyzer/SKILL.md"]
|
||||
assert drive.config.mentioned_skill_keys == ["tender-analyzer/SKILL.md"]
|
||||
# injected right after execution_context, mirroring the workflow surface
|
||||
names = [layer.name for layer in result.request.composition.layers]
|
||||
assert names.index("drive") == names.index("execution_context") + 1
|
||||
|
||||
def test_drive_layer_injected_with_empty_catalog_and_shell_depends_on_it(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"core.app.apps.agent_app.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr("core.app.apps.agent_app.runtime_request_builder.dify_config.AGENT_SHELL_ENABLED", True)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [],
|
||||
)
|
||||
builder = AgentAppRuntimeRequestBuilder(
|
||||
credentials_provider=_FakeCredentialsProvider(),
|
||||
plugin_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
result = builder.build(_ctx(_soul_with_model()))
|
||||
|
||||
layers = {layer.name: layer for layer in result.request.composition.layers}
|
||||
assert layers["drive"].config.drive_ref == "agent-agent-1"
|
||||
assert layers["drive"].config.skills == []
|
||||
assert layers[DIFY_SHELL_LAYER_ID].deps == {
|
||||
"execution_context": "execution_context",
|
||||
"drive": "drive",
|
||||
}
|
||||
|
||||
def test_no_drive_layer_when_flag_disabled(self):
|
||||
builder = AgentAppRuntimeRequestBuilder(
|
||||
credentials_provider=_FakeCredentialsProvider(),
|
||||
@ -269,3 +310,152 @@ class TestAgentAppDriveLayer:
|
||||
)
|
||||
result = builder.build(_ctx(_soul_with_model_and_skill()))
|
||||
assert all(layer.name != "drive" for layer in result.request.composition.layers)
|
||||
|
||||
def test_agent_app_runtime_expands_skill_and_file_mentions_in_agent_soul_prompt(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
"core.app.apps.agent_app.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [
|
||||
{
|
||||
"path": "tender-analyzer",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"archive_key": None,
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
"size": 1,
|
||||
"mime_type": "text/markdown",
|
||||
"hash": None,
|
||||
"created_at": 1,
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [
|
||||
{"key": "tender-analyzer/SKILL.md", "is_skill": True},
|
||||
{"key": "files/sample.pdf", "is_skill": False},
|
||||
],
|
||||
)
|
||||
soul = _soul_with_model()
|
||||
soul.prompt.system_prompt = (
|
||||
"Use [§skill:tender-analyzer%2FSKILL.md:Tender Analyzer§] and [§file:files%2Fsample.pdf:sample.pdf§]."
|
||||
)
|
||||
builder = AgentAppRuntimeRequestBuilder(
|
||||
credentials_provider=_FakeCredentialsProvider(),
|
||||
plugin_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
result = builder.build(_ctx(soul))
|
||||
|
||||
prompt_layer = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt")
|
||||
assert prompt_layer.config.prefix == "Use Tender Analyzer and sample.pdf."
|
||||
assert "[§" not in prompt_layer.config.prefix
|
||||
|
||||
def test_agent_app_runtime_missing_drive_mentions_fall_back_to_label_then_decoded_key(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
"core.app.apps.agent_app.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [],
|
||||
)
|
||||
soul = _soul_with_model()
|
||||
soul.prompt.system_prompt = (
|
||||
"Use [§skill:ghost%2FSKILL.md:Ghost Skill§], [§file:files%2Fghost.txt:Ghost File§], "
|
||||
"and [§file:files%2Fmissing.txt§]."
|
||||
)
|
||||
builder = AgentAppRuntimeRequestBuilder(
|
||||
credentials_provider=_FakeCredentialsProvider(),
|
||||
plugin_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
result = builder.build(_ctx(soul))
|
||||
|
||||
prompt_layer = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt")
|
||||
assert prompt_layer.config.prefix == "Use Ghost Skill, Ghost File, and files/missing.txt."
|
||||
assert "[§" not in prompt_layer.config.prefix
|
||||
|
||||
def test_agent_app_runtime_expands_drive_mentions_in_agent_soul_prompt(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"core.app.apps.agent_app.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [
|
||||
{
|
||||
"path": "tender-analyzer",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"archive_key": None,
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
"size": 1,
|
||||
"mime_type": "text/markdown",
|
||||
"hash": None,
|
||||
"created_at": 1,
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [
|
||||
{"key": "tender-analyzer/SKILL.md", "is_skill": True},
|
||||
{"key": "files/sample.pdf", "is_skill": False},
|
||||
],
|
||||
)
|
||||
soul = _soul_with_model()
|
||||
soul.prompt.system_prompt = (
|
||||
"Use [§skill:tender-analyzer%2FSKILL.md:Tender Analyzer§] and [§file:files%2Fsample.pdf:sample.pdf§]"
|
||||
)
|
||||
builder = AgentAppRuntimeRequestBuilder(
|
||||
credentials_provider=_FakeCredentialsProvider(),
|
||||
plugin_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
result = builder.build(_ctx(soul))
|
||||
|
||||
prompt_layer = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt")
|
||||
assert prompt_layer.config.prefix == "Use Tender Analyzer and sample.pdf"
|
||||
assert "[§" not in prompt_layer.config.prefix
|
||||
|
||||
def test_agent_app_runtime_missing_drive_mentions_fall_back_without_marker_leak(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
"core.app.apps.agent_app.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [],
|
||||
)
|
||||
soul = _soul_with_model()
|
||||
soul.prompt.system_prompt = (
|
||||
"Use [§skill:ghost%2FSKILL.md:Ghost Skill§], [§file:files%2Fghost.txt:Ghost File§], "
|
||||
"and [§file:files%2Fno-label.txt§]."
|
||||
)
|
||||
builder = AgentAppRuntimeRequestBuilder(
|
||||
credentials_provider=_FakeCredentialsProvider(),
|
||||
plugin_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
result = builder.build(_ctx(soul))
|
||||
|
||||
prompt_layer = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt")
|
||||
assert prompt_layer.config.prefix == "Use Ghost Skill, Ghost File, and files/no-label.txt."
|
||||
assert "[§" not in prompt_layer.config.prefix
|
||||
|
||||
@ -834,57 +834,116 @@ def test_mentions_expand_in_soul_and_job_prompts_without_token_leak():
|
||||
|
||||
def _soul_with_drive_skill() -> AgentSoulConfig:
|
||||
return AgentSoulConfig(
|
||||
prompt={"system_prompt": "You are careful."},
|
||||
model=AgentSoulModelConfig(plugin_id="langgenius/openai", model_provider="openai", model="gpt-test"),
|
||||
skills_files={
|
||||
"skills": [
|
||||
{
|
||||
"id": "abc123",
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"full_archive_key": "tender-analyzer/.DIFY-SKILL-FULL.zip",
|
||||
},
|
||||
{"id": "legacy", "name": "Legacy Skill"}, # pre-standardization: no drive key
|
||||
],
|
||||
"files": [
|
||||
{"name": "sample.pdf", "drive_key": "files/sample.pdf", "type": "application/pdf"},
|
||||
{"name": "plain-upload.pdf", "file_id": "upload-1"}, # not drive-backed
|
||||
],
|
||||
prompt={
|
||||
"system_prompt": (
|
||||
"You are careful. Use [§skill:tender-analyzer%2FSKILL.md:Tender Analyzer§] "
|
||||
"and [§file:files%2Fsample.pdf:sample.pdf§]."
|
||||
)
|
||||
},
|
||||
model=AgentSoulModelConfig(plugin_id="langgenius/openai", model_provider="openai", model="gpt-test"),
|
||||
)
|
||||
|
||||
|
||||
def test_build_drive_layer_config_catalogs_only_drive_backed_refs():
|
||||
def _mock_drive_catalog(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [
|
||||
{
|
||||
"path": "tender-analyzer",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"archive_key": "tender-analyzer/.DIFY-SKILL-FULL.zip",
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
"size": 123,
|
||||
"mime_type": "text/markdown",
|
||||
"hash": "hash-1",
|
||||
"created_at": 1,
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [
|
||||
{"key": "tender-analyzer/SKILL.md", "is_skill": True},
|
||||
{"key": "tender-analyzer/.DIFY-SKILL-FULL.zip", "is_skill": False},
|
||||
{"key": "files/sample.pdf", "is_skill": False},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _mock_empty_drive_catalog(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [],
|
||||
)
|
||||
|
||||
|
||||
def test_build_drive_layer_config_catalogs_drive_skills_and_mentions(monkeypatch: pytest.MonkeyPatch):
|
||||
from core.workflow.nodes.agent_v2.runtime_request_builder import build_drive_layer_config
|
||||
|
||||
config, warnings = build_drive_layer_config(_soul_with_drive_skill(), agent_id="agent-1")
|
||||
_mock_drive_catalog(monkeypatch)
|
||||
config, warnings = build_drive_layer_config(_soul_with_drive_skill(), tenant_id="tenant-1", agent_id="agent-1")
|
||||
|
||||
assert config is not None
|
||||
assert config.drive_ref == "agent-agent-1"
|
||||
assert [skill.skill_md_key for skill in config.skills] == ["tender-analyzer/SKILL.md"]
|
||||
assert config.skills[0].archive_key == "tender-analyzer/.DIFY-SKILL-FULL.zip"
|
||||
assert [file.key for file in config.files] == ["files/sample.pdf"]
|
||||
assert [w["code"] for w in warnings] == ["skill_ref_dangling"]
|
||||
assert "Legacy Skill" in warnings[0]["message"]
|
||||
assert config.mentioned_skill_keys == ["tender-analyzer/SKILL.md"]
|
||||
assert config.mentioned_file_keys == ["files/sample.pdf"]
|
||||
assert warnings == []
|
||||
|
||||
|
||||
def test_build_drive_layer_config_skips_when_nothing_configured():
|
||||
def test_build_drive_layer_config_emits_drive_ref_when_catalog_is_empty(monkeypatch: pytest.MonkeyPatch):
|
||||
from core.workflow.nodes.agent_v2.runtime_request_builder import build_drive_layer_config
|
||||
|
||||
_mock_empty_drive_catalog(monkeypatch)
|
||||
soul = AgentSoulConfig(
|
||||
model=AgentSoulModelConfig(plugin_id="langgenius/openai", model_provider="openai", model="gpt-test")
|
||||
)
|
||||
assert build_drive_layer_config(soul, agent_id="agent-1") == (None, [])
|
||||
config, warnings = build_drive_layer_config(soul, tenant_id="tenant-1", agent_id="agent-1")
|
||||
|
||||
assert config is not None
|
||||
assert config.drive_ref == "agent-agent-1"
|
||||
assert config.skills == []
|
||||
assert config.mentioned_skill_keys == []
|
||||
assert config.mentioned_file_keys == []
|
||||
assert warnings == []
|
||||
|
||||
|
||||
def test_workflow_run_request_contains_drive_layer_with_empty_catalog(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr("core.workflow.nodes.agent_v2.runtime_request_builder.dify_config.AGENT_SHELL_ENABLED", True)
|
||||
_mock_empty_drive_catalog(monkeypatch)
|
||||
|
||||
result = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()).build(_context())
|
||||
|
||||
dumped = result.request.model_dump(mode="json")
|
||||
layers = {layer["name"]: layer for layer in dumped["composition"]["layers"]}
|
||||
assert layers["drive"]["config"] == {
|
||||
"drive_ref": "agent-agent-1",
|
||||
"skills": [],
|
||||
"mentioned_skill_keys": [],
|
||||
"mentioned_file_keys": [],
|
||||
}
|
||||
assert layers[DIFY_SHELL_LAYER_ID]["deps"] == {
|
||||
"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID,
|
||||
"drive": "drive",
|
||||
}
|
||||
|
||||
|
||||
def test_build_drive_layer_config_requires_agent_identity():
|
||||
from core.workflow.nodes.agent_v2.runtime_request_builder import build_drive_layer_config
|
||||
|
||||
config, warnings = build_drive_layer_config(_soul_with_drive_skill(), agent_id=None)
|
||||
config, warnings = build_drive_layer_config(_soul_with_drive_skill(), tenant_id="tenant-1", agent_id=None)
|
||||
|
||||
assert config is None
|
||||
assert [w["code"] for w in warnings] == ["skill_ref_dangling"]
|
||||
assert [w["code"] for w in warnings] == ["drive_ref_dangling"]
|
||||
|
||||
|
||||
def test_workflow_run_request_contains_drive_layer_when_flag_enabled(monkeypatch: pytest.MonkeyPatch):
|
||||
@ -892,6 +951,7 @@ def test_workflow_run_request_contains_drive_layer_when_flag_enabled(monkeypatch
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
_mock_drive_catalog(monkeypatch)
|
||||
context = _context()
|
||||
context.snapshot.config_snapshot = _soul_with_drive_skill()
|
||||
|
||||
@ -904,21 +964,21 @@ def test_workflow_run_request_contains_drive_layer_when_flag_enabled(monkeypatch
|
||||
assert layer_names.index("drive") == layer_names.index("execution_context") + 1
|
||||
drive = next(layer for layer in dumped["composition"]["layers"] if layer["name"] == "drive")
|
||||
assert drive["type"] == "dify.drive"
|
||||
assert drive["deps"] == {"execution_context": "execution_context"}
|
||||
assert drive["config"]["drive_ref"] == "agent-agent-1"
|
||||
assert drive["config"]["skills"] == [
|
||||
{
|
||||
"path": "tender-analyzer",
|
||||
"name": "Tender Analyzer",
|
||||
"description": "Parses RFPs.",
|
||||
"skill_md_key": "tender-analyzer/SKILL.md",
|
||||
"archive_key": "tender-analyzer/.DIFY-SKILL-FULL.zip",
|
||||
}
|
||||
]
|
||||
assert drive["config"]["files"] == [
|
||||
{"name": "sample.pdf", "key": "files/sample.pdf", "size": None, "mime_type": "application/pdf"}
|
||||
]
|
||||
# the dangling legacy ref degraded to a warning instead of failing the run
|
||||
assert drive["config"]["mentioned_skill_keys"] == ["tender-analyzer/SKILL.md"]
|
||||
assert drive["config"]["mentioned_file_keys"] == ["files/sample.pdf"]
|
||||
warnings = result.metadata["runtime_support"]["unsupported_runtime_warnings"]
|
||||
assert any(w["code"] == "skill_ref_dangling" for w in warnings)
|
||||
assert warnings == []
|
||||
# the drive layer is non-sensitive and must survive into persistable specs
|
||||
from dify_agent.protocol import extract_runtime_layer_specs
|
||||
|
||||
@ -926,6 +986,51 @@ def test_workflow_run_request_contains_drive_layer_when_flag_enabled(monkeypatch
|
||||
assert any(spec.name == "drive" and spec.type == "dify.drive" for spec in specs)
|
||||
|
||||
|
||||
def test_workflow_runtime_expands_drive_mentions_in_agent_soul_prompt(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
_mock_drive_catalog(monkeypatch)
|
||||
context = _context()
|
||||
context.snapshot.config_snapshot = _soul_with_drive_skill()
|
||||
|
||||
result = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()).build(context)
|
||||
|
||||
soul_prompt = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt")
|
||||
assert soul_prompt.config.prefix == "You are careful. Use Tender Analyzer and sample.pdf."
|
||||
assert "[§" not in soul_prompt.config.prefix
|
||||
|
||||
|
||||
def test_workflow_runtime_missing_drive_mentions_fall_back_to_label_then_decoded_key(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.dify_config.AGENT_DRIVE_MANIFEST_ENABLED", True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.list_skills",
|
||||
lambda self, *, tenant_id, agent_id: [],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.workflow.nodes.agent_v2.runtime_request_builder.AgentDriveService.manifest",
|
||||
lambda self, *, tenant_id, agent_id, prefix="", include_download_url=False: [],
|
||||
)
|
||||
context = _context()
|
||||
context.snapshot.config_snapshot = AgentSoulConfig(
|
||||
prompt={
|
||||
"system_prompt": (
|
||||
"Use [§skill:ghost%2FSKILL.md:Ghost Skill§], [§file:files%2Fghost.txt:Ghost File§], "
|
||||
"and [§file:files%2Fno-label.txt§]."
|
||||
)
|
||||
},
|
||||
model=AgentSoulModelConfig(plugin_id="langgenius/openai", model_provider="openai", model="gpt-test"),
|
||||
)
|
||||
|
||||
result = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()).build(context)
|
||||
|
||||
soul_prompt = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt")
|
||||
assert soul_prompt.config.prefix == "Use Ghost Skill, Ghost File, and files/no-label.txt."
|
||||
assert "[§" not in soul_prompt.config.prefix
|
||||
|
||||
|
||||
def test_workflow_run_request_has_no_drive_layer_when_flag_disabled():
|
||||
context = _context()
|
||||
context.snapshot.config_snapshot = _soul_with_drive_skill()
|
||||
@ -934,20 +1039,20 @@ def test_workflow_run_request_has_no_drive_layer_when_flag_disabled():
|
||||
|
||||
dumped = result.request.model_dump(mode="json")
|
||||
assert all(layer["name"] != "drive" for layer in dumped["composition"]["layers"])
|
||||
warnings = result.metadata["runtime_support"]["unsupported_runtime_warnings"]
|
||||
assert any(w["code"] == "drive_manifest_disabled" for w in warnings)
|
||||
assert result.metadata["runtime_support"]["unsupported_runtime_warnings"] == []
|
||||
|
||||
|
||||
def test_build_drive_layer_config_all_refs_dangling_yields_no_config():
|
||||
def test_build_drive_layer_config_missing_mentions_warn_but_keep_skill_catalog(monkeypatch: pytest.MonkeyPatch):
|
||||
from core.workflow.nodes.agent_v2.runtime_request_builder import build_drive_layer_config
|
||||
|
||||
_mock_drive_catalog(monkeypatch)
|
||||
soul = AgentSoulConfig(
|
||||
model=AgentSoulModelConfig(plugin_id="langgenius/openai", model_provider="openai", model="gpt-test"),
|
||||
skills_files={"skills": [{"id": "legacy", "name": "Legacy"}], "files": [{"name": "u.pdf", "file_id": "u1"}]},
|
||||
prompt={"system_prompt": "Use [§skill:ghost%2FSKILL.md:Ghost§]"},
|
||||
)
|
||||
config, warnings = build_drive_layer_config(soul, agent_id="agent-1")
|
||||
assert config is None
|
||||
assert [w["code"] for w in warnings] == ["skill_ref_dangling"]
|
||||
config, warnings = build_drive_layer_config(soul, tenant_id="tenant-1", agent_id="agent-1")
|
||||
assert config is not None
|
||||
assert [w["code"] for w in warnings] == ["mention_target_missing"]
|
||||
|
||||
|
||||
# ── ENG-635: ask_human layer gating + feature manifest ───────────────────────
|
||||
|
||||
@ -17,7 +17,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.broadcast_channel.exc import BroadcastChannelError, SubscriptionClosedError
|
||||
from libs.broadcast_channel.exc import SubscriptionClosedError
|
||||
from libs.broadcast_channel.redis.pubsub_channel import (
|
||||
BroadcastChannel as RedisBroadcastChannel,
|
||||
)
|
||||
@ -395,11 +395,10 @@ class TestRedisSubscription:
|
||||
assert received_messages == test_messages
|
||||
|
||||
def test_message_iterator_when_closed(self, subscription: _RedisSubscription):
|
||||
"""Test that iterator raises error when subscription is closed."""
|
||||
"""Test that iterator stops when subscription is closed."""
|
||||
subscription.close()
|
||||
|
||||
with pytest.raises(BroadcastChannelError, match="The Redis regular subscription is closed"):
|
||||
iter(subscription)
|
||||
assert list(subscription) == []
|
||||
|
||||
# ==================== Message Enqueue Tests ====================
|
||||
|
||||
@ -616,8 +615,15 @@ class TestRedisSubscription:
|
||||
"""Test iterator behavior after close."""
|
||||
subscription.close()
|
||||
|
||||
with pytest.raises(SubscriptionClosedError, match="The Redis regular subscription is closed"):
|
||||
iter(subscription)
|
||||
assert list(subscription) == []
|
||||
|
||||
def test_close_does_not_join_unstarted_listener_thread(self, subscription: _RedisSubscription):
|
||||
"""close() should tolerate a listener object that has not been started yet."""
|
||||
subscription._listener_thread = threading.Thread(target=lambda: None)
|
||||
|
||||
subscription.close()
|
||||
|
||||
assert subscription._listener_thread is None
|
||||
|
||||
def test_start_after_close(self, subscription: _RedisSubscription):
|
||||
"""Test start attempts after close."""
|
||||
@ -818,11 +824,10 @@ class TestRedisShardedSubscription:
|
||||
assert received_messages == test_messages
|
||||
|
||||
def test_message_iterator_when_closed(self, sharded_subscription: _RedisShardedSubscription):
|
||||
"""Test that iterator raises error when sharded subscription is closed."""
|
||||
"""Test that iterator stops when sharded subscription is closed."""
|
||||
sharded_subscription.close()
|
||||
|
||||
with pytest.raises(SubscriptionClosedError, match="The Redis sharded subscription is closed"):
|
||||
iter(sharded_subscription)
|
||||
assert list(sharded_subscription) == []
|
||||
|
||||
# ==================== Message Enqueue Tests ====================
|
||||
|
||||
@ -1093,8 +1098,7 @@ class TestRedisShardedSubscription:
|
||||
"""Test iterator behavior after close for sharded subscription."""
|
||||
sharded_subscription.close()
|
||||
|
||||
with pytest.raises(SubscriptionClosedError, match="The Redis sharded subscription is closed"):
|
||||
iter(sharded_subscription)
|
||||
assert list(sharded_subscription) == []
|
||||
|
||||
def test_start_after_close(self, sharded_subscription: _RedisShardedSubscription):
|
||||
"""Test start attempts after close for sharded subscription."""
|
||||
@ -1312,12 +1316,10 @@ class TestRedisSubscriptionCommon:
|
||||
assert received_messages == test_messages
|
||||
|
||||
def test_message_iterator_when_closed(self, subscription, subscription_params):
|
||||
"""Test that iterator raises error when subscription is closed."""
|
||||
subscription_type, _ = subscription_params
|
||||
"""Test that iterator stops when subscription is closed."""
|
||||
subscription.close()
|
||||
|
||||
with pytest.raises(SubscriptionClosedError, match=f"The Redis {subscription_type} subscription is closed"):
|
||||
iter(subscription)
|
||||
assert list(subscription) == []
|
||||
|
||||
# ==================== Message Enqueue Tests ====================
|
||||
|
||||
@ -1390,11 +1392,9 @@ class TestRedisSubscriptionCommon:
|
||||
|
||||
def test_iterator_after_close(self, subscription, subscription_params):
|
||||
"""Test iterator behavior after close."""
|
||||
subscription_type, _ = subscription_params
|
||||
subscription.close()
|
||||
|
||||
with pytest.raises(SubscriptionClosedError, match=f"The Redis {subscription_type} subscription is closed"):
|
||||
iter(subscription)
|
||||
assert list(subscription) == []
|
||||
|
||||
def test_start_after_close(self, subscription, subscription_params):
|
||||
"""Test start attempts after close."""
|
||||
|
||||
@ -0,0 +1,122 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic.migration import MigrationContext
|
||||
from alembic.operations import Operations
|
||||
|
||||
_MIGRATION_PATH = (
|
||||
Path(__file__).resolve().parents[3]
|
||||
/ "migrations/versions/2026_06_18_2300-b2515f9d4c2a_agent_drive_skill_metadata_refactor.py"
|
||||
)
|
||||
|
||||
|
||||
def _load_migration_module():
|
||||
spec = importlib.util.spec_from_file_location("agent_drive_skill_metadata_refactor", _MIGRATION_PATH)
|
||||
if spec is None or spec.loader is None:
|
||||
raise RuntimeError("failed to load migration module")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _create_pre_upgrade_schema(engine: sa.Engine) -> None:
|
||||
metadata = sa.MetaData()
|
||||
sa.Table(
|
||||
"agent_drive_files",
|
||||
metadata,
|
||||
sa.Column("tenant_id", sa.String(36), nullable=False),
|
||||
sa.Column("agent_id", sa.String(36), nullable=False),
|
||||
sa.Column("key", sa.String(512), nullable=False),
|
||||
sa.Column("file_kind", sa.String(32), nullable=False),
|
||||
sa.Column("file_id", sa.String(36), nullable=False),
|
||||
sa.Column("value_owned_by_drive", sa.Boolean(), nullable=False, server_default=sa.text("false")),
|
||||
sa.Column("size", sa.BigInteger(), nullable=True),
|
||||
sa.Column("hash", sa.String(255), nullable=True),
|
||||
sa.Column("mime_type", sa.String(255), nullable=True),
|
||||
sa.Column("created_by", sa.String(36), nullable=True),
|
||||
sa.Column("id", sa.String(36), primary_key=True),
|
||||
sa.Column("created_at", sa.DateTime(), nullable=False),
|
||||
sa.Column("updated_at", sa.DateTime(), nullable=False),
|
||||
sa.UniqueConstraint("tenant_id", "agent_id", "key", name="agent_drive_file_scope_key_unique"),
|
||||
)
|
||||
sa.Table(
|
||||
"agent_config_snapshots",
|
||||
metadata,
|
||||
sa.Column("id", sa.String(36), primary_key=True),
|
||||
sa.Column("config_snapshot", sa.Text(), nullable=False),
|
||||
)
|
||||
metadata.create_all(engine)
|
||||
|
||||
|
||||
def _run_migration_step(module: object, engine: sa.Engine, step_name: str) -> None:
|
||||
with engine.begin() as connection:
|
||||
context = MigrationContext.configure(connection)
|
||||
operations = Operations(context)
|
||||
original_op = module.op
|
||||
module.op = operations
|
||||
try:
|
||||
getattr(module, step_name)()
|
||||
finally:
|
||||
module.op = original_op
|
||||
|
||||
|
||||
def test_upgrade_adds_skill_columns_and_index_and_strips_snapshot_data() -> None:
|
||||
engine = sa.create_engine("sqlite:///:memory:")
|
||||
_create_pre_upgrade_schema(engine)
|
||||
snapshot = {
|
||||
"prompt": {"system_prompt": "Use [§skill:legacy:Legacy§]"},
|
||||
"skills_files": {"skills": [{"name": "Legacy"}], "files": [{"name": "u.pdf"}]},
|
||||
}
|
||||
with engine.begin() as connection:
|
||||
connection.execute(
|
||||
sa.text("INSERT INTO agent_config_snapshots (id, config_snapshot) VALUES (:id, :config_snapshot)"),
|
||||
{"id": "snap-1", "config_snapshot": json.dumps(snapshot)},
|
||||
)
|
||||
|
||||
module = _load_migration_module()
|
||||
_run_migration_step(module, engine, "upgrade")
|
||||
|
||||
inspector = sa.inspect(engine)
|
||||
columns = {column["name"] for column in inspector.get_columns("agent_drive_files")}
|
||||
assert {"is_skill", "skill_metadata"}.issubset(columns)
|
||||
indexes = {index["name"] for index in inspector.get_indexes("agent_drive_files")}
|
||||
assert "agent_drive_files_tenant_agent_is_skill_key_idx" in indexes
|
||||
|
||||
with engine.begin() as connection:
|
||||
stored_snapshot = connection.execute(
|
||||
sa.text("SELECT config_snapshot FROM agent_config_snapshots WHERE id = :id"),
|
||||
{"id": "snap-1"},
|
||||
).scalar_one()
|
||||
assert "skills_files" not in json.loads(stored_snapshot)
|
||||
|
||||
|
||||
def test_downgrade_drops_skill_columns_and_index_without_reconstructing_legacy_data() -> None:
|
||||
engine = sa.create_engine("sqlite:///:memory:")
|
||||
_create_pre_upgrade_schema(engine)
|
||||
with engine.begin() as connection:
|
||||
connection.execute(
|
||||
sa.text("INSERT INTO agent_config_snapshots (id, config_snapshot) VALUES (:id, :config_snapshot)"),
|
||||
{"id": "snap-1", "config_snapshot": json.dumps({"prompt": {"system_prompt": "hello"}})},
|
||||
)
|
||||
|
||||
module = _load_migration_module()
|
||||
_run_migration_step(module, engine, "upgrade")
|
||||
_run_migration_step(module, engine, "downgrade")
|
||||
|
||||
inspector = sa.inspect(engine)
|
||||
columns = {column["name"] for column in inspector.get_columns("agent_drive_files")}
|
||||
assert "is_skill" not in columns
|
||||
assert "skill_metadata" not in columns
|
||||
indexes = {index["name"] for index in inspector.get_indexes("agent_drive_files")}
|
||||
assert "agent_drive_files_tenant_agent_is_skill_key_idx" not in indexes
|
||||
|
||||
with engine.begin() as connection:
|
||||
stored_snapshot = connection.execute(
|
||||
sa.text("SELECT config_snapshot FROM agent_config_snapshots WHERE id = :id"),
|
||||
{"id": "snap-1"},
|
||||
).scalar_one()
|
||||
assert "skills_files" not in json.loads(stored_snapshot)
|
||||
@ -3,7 +3,7 @@ import pytest
|
||||
from models.agent_config_entities import AgentKnowledgeQueryMode, AgentSoulModelConfig, DeclaredOutputType
|
||||
from services.agent.composer_service import AgentComposerService
|
||||
from services.agent.composer_validator import ComposerConfigValidator
|
||||
from services.agent.errors import AgentSoulLockedError, PlaintextSecretNotAllowedError
|
||||
from services.agent.errors import AgentSoulLockedError, InvalidComposerConfigError, PlaintextSecretNotAllowedError
|
||||
from services.entities.agent_entities import (
|
||||
AgentSoulConfig,
|
||||
ComposerSavePayload,
|
||||
@ -64,6 +64,24 @@ def test_locked_workflow_node_job_only_allows_inline_soul_payload():
|
||||
ComposerConfigValidator.validate_save_payload(payload)
|
||||
|
||||
|
||||
def test_draft_save_payload_skips_publish_only_agent_soul_validation():
|
||||
payload = ComposerSavePayload.model_validate(
|
||||
{
|
||||
"variant": ComposerVariant.AGENT_APP,
|
||||
"save_strategy": ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION,
|
||||
"agent_soul": {
|
||||
"prompt": {"system_prompt": "no human reference yet"},
|
||||
"human": {"contacts": [{"id": "human-1", "name": "Reviewer"}]},
|
||||
"env": {"variables": [{"name": "bad-name"}]},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
ComposerConfigValidator.validate_draft_save_payload(payload)
|
||||
with pytest.raises(InvalidComposerConfigError):
|
||||
ComposerConfigValidator.validate_publish_payload(payload)
|
||||
|
||||
|
||||
def test_agent_app_soul_allows_app_features_and_variables():
|
||||
payload = ComposerSavePayload.model_validate(
|
||||
{
|
||||
|
||||
@ -4,6 +4,7 @@ from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidationError
|
||||
from models.agent import (
|
||||
Agent,
|
||||
AgentConfigRevisionOperation,
|
||||
@ -17,7 +18,6 @@ from models.agent import (
|
||||
WorkflowAgentNodeBinding,
|
||||
)
|
||||
from models.agent_config_entities import (
|
||||
AgentFileRefConfig,
|
||||
DeclaredArrayItem,
|
||||
DeclaredOutputChildConfig,
|
||||
DeclaredOutputConfig,
|
||||
@ -169,7 +169,7 @@ def test_save_workflow_composer_dispatches_save_strategy(monkeypatch, strategy,
|
||||
calls = []
|
||||
|
||||
monkeypatch.setattr(composer_service.db, "session", fake_session)
|
||||
monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_draft_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: SimpleNamespace(id="workflow-1"))
|
||||
monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
@ -222,12 +222,52 @@ def test_save_workflow_composer_rejects_agent_app_variant():
|
||||
)
|
||||
|
||||
|
||||
def _duplicate_env_secret_payload(strategy: ComposerSaveStrategy) -> ComposerSavePayload:
|
||||
return ComposerSavePayload.model_validate(
|
||||
{
|
||||
"variant": ComposerVariant.AGENT_APP.value,
|
||||
"save_strategy": strategy.value,
|
||||
"agent_soul": {
|
||||
"prompt": {"system_prompt": "x"},
|
||||
"env": {
|
||||
"variables": [{"name": "TOKEN", "value": "plain"}],
|
||||
"secret_refs": [{"name": "TOKEN", "value": "credential-1"}],
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"strategy",
|
||||
[
|
||||
ComposerSaveStrategy.NODE_JOB_ONLY,
|
||||
ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION,
|
||||
],
|
||||
)
|
||||
def test_draft_save_strategies_skip_publish_validation(strategy: ComposerSaveStrategy):
|
||||
composer_service._validate_composer_payload_for_strategy(_duplicate_env_secret_payload(strategy))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"strategy",
|
||||
[
|
||||
ComposerSaveStrategy.SAVE_AS_NEW_VERSION,
|
||||
ComposerSaveStrategy.SAVE_AS_NEW_AGENT,
|
||||
ComposerSaveStrategy.SAVE_TO_ROSTER,
|
||||
],
|
||||
)
|
||||
def test_publish_save_strategies_run_publish_validation(strategy: ComposerSaveStrategy):
|
||||
with pytest.raises(InvalidComposerConfigError, match="duplicate env/secret name 'TOKEN'"):
|
||||
composer_service._validate_composer_payload_for_strategy(_duplicate_env_secret_payload(strategy))
|
||||
|
||||
|
||||
def test_save_agent_app_composer_creates_agent_when_missing(monkeypatch: pytest.MonkeyPatch):
|
||||
fake_session = FakeSession(scalar=[None])
|
||||
created_version = SimpleNamespace(id="version-1")
|
||||
|
||||
monkeypatch.setattr(composer_service.db, "session", fake_session)
|
||||
monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_draft_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(AgentComposerService, "_create_config_version", lambda **kwargs: created_version)
|
||||
monkeypatch.setattr(AgentComposerService, "load_agent_app_composer", lambda **kwargs: {"loaded": True})
|
||||
payload = ComposerSavePayload.model_validate(
|
||||
@ -257,7 +297,7 @@ def test_save_agent_app_composer_updates_current_version(monkeypatch: pytest.Mon
|
||||
updated = {}
|
||||
|
||||
monkeypatch.setattr(composer_service.db, "session", fake_session)
|
||||
monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_draft_save_payload", lambda payload: None)
|
||||
monkeypatch.setattr(AgentComposerService, "_require_version", lambda **kwargs: SimpleNamespace(id="version-1"))
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService,
|
||||
@ -1649,6 +1689,74 @@ class TestAgentAppBackingAgent:
|
||||
with pytest.raises(roster_service.AgentNotFoundError):
|
||||
service.get_agent_app_model(tenant_id="tenant-1", agent_id="agent-x")
|
||||
|
||||
def test_refresh_agent_app_debug_conversation_creates_mapping(self):
|
||||
agent = Agent(
|
||||
id="agent-1",
|
||||
tenant_id="tenant-1",
|
||||
name="Iris",
|
||||
description="",
|
||||
agent_kind=AgentKind.DIFY_AGENT,
|
||||
scope=AgentScope.ROSTER,
|
||||
source=AgentSource.AGENT_APP,
|
||||
status=AgentStatus.ACTIVE,
|
||||
app_id="app-1",
|
||||
)
|
||||
session = FakeSession(scalar=[agent, None])
|
||||
service = AgentRosterService(session)
|
||||
|
||||
conversation_id = service.refresh_agent_app_debug_conversation_id(
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
account_id="account-1",
|
||||
)
|
||||
|
||||
conversations = [a for a in session.added if isinstance(a, Conversation)]
|
||||
assert len(conversations) == 1
|
||||
assert conversations[0].id == conversation_id
|
||||
assert conversations[0].app_id == "app-1"
|
||||
assert conversations[0].from_source == ConversationFromSource.CONSOLE
|
||||
assert conversations[0].from_account_id == "account-1"
|
||||
mappings = [a for a in session.added if isinstance(a, AgentDebugConversation)]
|
||||
assert len(mappings) == 1
|
||||
assert mappings[0].tenant_id == "tenant-1"
|
||||
assert mappings[0].agent_id == "agent-1"
|
||||
assert mappings[0].app_id == "app-1"
|
||||
assert mappings[0].account_id == "account-1"
|
||||
assert mappings[0].conversation_id == conversation_id
|
||||
assert session.deleted == []
|
||||
assert session.commits == 1
|
||||
|
||||
def test_refresh_agent_app_debug_conversation_replaces_existing_mapping(self):
|
||||
agent = Agent(
|
||||
id="agent-1",
|
||||
tenant_id="tenant-1",
|
||||
name="Iris",
|
||||
description="",
|
||||
agent_kind=AgentKind.DIFY_AGENT,
|
||||
scope=AgentScope.ROSTER,
|
||||
source=AgentSource.AGENT_APP,
|
||||
status=AgentStatus.ACTIVE,
|
||||
app_id="app-1",
|
||||
)
|
||||
mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation")
|
||||
session = FakeSession(scalar=[agent, mapping])
|
||||
service = AgentRosterService(session)
|
||||
|
||||
conversation_id = service.refresh_agent_app_debug_conversation_id(
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
account_id="account-1",
|
||||
)
|
||||
|
||||
assert mapping.app_id == "app-1"
|
||||
assert mapping.conversation_id == conversation_id
|
||||
assert [a for a in session.added if isinstance(a, AgentDebugConversation)] == []
|
||||
conversations = [a for a in session.added if isinstance(a, Conversation)]
|
||||
assert len(conversations) == 1
|
||||
assert conversations[0].id == conversation_id
|
||||
assert session.deleted == []
|
||||
assert session.commits == 1
|
||||
|
||||
def test_duplicate_agent_app_copies_app_config_and_active_soul(self, monkeypatch: pytest.MonkeyPatch):
|
||||
source_config = SimpleNamespace(
|
||||
opening_statement="hello",
|
||||
@ -1815,8 +1923,11 @@ class TestAgentAppBackingAgent:
|
||||
monkeypatch.setattr(service, "_copy_agent_active_snapshot", lambda **_: None)
|
||||
monkeypatch.setattr(service, "_next_duplicate_agent_name", lambda **_: "Iris copy")
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class FakeAppService:
|
||||
def create_app(self, tenant_id: str, params, account: object) -> object:
|
||||
captured["params"] = params
|
||||
return target_app
|
||||
|
||||
access_mode_updates = []
|
||||
@ -1842,9 +1953,11 @@ class TestAgentAppBackingAgent:
|
||||
tenant_id="tenant-1",
|
||||
agent_id="source-agent",
|
||||
account=SimpleNamespace(id="account-1"),
|
||||
role="Custom Analyst",
|
||||
)
|
||||
|
||||
assert duplicated is target_app
|
||||
assert captured["params"].agent_role == "Custom Analyst"
|
||||
assert access_mode_updates == [("target-app", "private")]
|
||||
|
||||
def test_duplicate_agent_app_falls_back_to_public_access_mode(self, monkeypatch: pytest.MonkeyPatch):
|
||||
@ -1988,6 +2101,97 @@ class TestListWorkflowsReferencingAppAgent:
|
||||
|
||||
|
||||
class TestWorkflowAgentDraftBindingSync:
|
||||
def _agent_workflow(self) -> Workflow:
|
||||
return Workflow(
|
||||
id="workflow-1",
|
||||
tenant_id="tenant-1",
|
||||
app_id="app-1",
|
||||
version=Workflow.VERSION_DRAFT,
|
||||
graph=json.dumps(
|
||||
{
|
||||
"nodes": [{"id": "agent-node", "data": {"type": "agent", "version": "2"}}],
|
||||
"edges": [],
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def _agent_binding(self) -> WorkflowAgentNodeBinding:
|
||||
return WorkflowAgentNodeBinding(
|
||||
id="binding-1",
|
||||
tenant_id="tenant-1",
|
||||
app_id="app-1",
|
||||
workflow_id="workflow-1",
|
||||
workflow_version=Workflow.VERSION_DRAFT,
|
||||
node_id="agent-node",
|
||||
binding_type=WorkflowAgentBindingType.ROSTER_AGENT,
|
||||
agent_id="agent-1",
|
||||
current_snapshot_id="snapshot-1",
|
||||
node_job_config=WorkflowNodeJobConfig(),
|
||||
)
|
||||
|
||||
def _publish_agent(self) -> Agent:
|
||||
return Agent(
|
||||
id="agent-1",
|
||||
tenant_id="tenant-1",
|
||||
name="Iris",
|
||||
status=AgentStatus.ACTIVE,
|
||||
active_config_snapshot_id="snapshot-1",
|
||||
)
|
||||
|
||||
def _snapshot(self, agent_soul: AgentSoulConfig) -> AgentConfigSnapshot:
|
||||
return AgentConfigSnapshot(
|
||||
id="snapshot-1",
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
version=1,
|
||||
config_snapshot=agent_soul,
|
||||
)
|
||||
|
||||
def test_publish_validation_rejects_agent_soul_publish_only_errors(self):
|
||||
binding = self._agent_binding()
|
||||
agent_soul = AgentSoulConfig.model_validate(
|
||||
{
|
||||
"model": {
|
||||
"plugin_id": "langgenius/openai/openai",
|
||||
"model_provider": "openai",
|
||||
"model": "gpt-4o",
|
||||
},
|
||||
"prompt": {"system_prompt": "no human reference yet"},
|
||||
"human": {"contacts": [{"id": "human-1", "name": "Reviewer"}]},
|
||||
}
|
||||
)
|
||||
agent = self._publish_agent()
|
||||
snapshot = self._snapshot(agent_soul)
|
||||
session = FakeSession(scalar=[binding, agent, snapshot, agent, snapshot], scalars=[[binding]])
|
||||
|
||||
with pytest.raises(InvalidComposerConfigError, match="human_involvement_not_referenced"):
|
||||
WorkflowAgentPublishService.validate_agent_nodes_for_publish(
|
||||
session=session,
|
||||
draft_workflow=self._agent_workflow(),
|
||||
)
|
||||
|
||||
def test_publish_validation_rejects_dangling_agent_soul_drive_refs(self):
|
||||
binding = self._agent_binding()
|
||||
agent_soul = AgentSoulConfig.model_validate(
|
||||
{
|
||||
"model": {
|
||||
"plugin_id": "langgenius/openai/openai",
|
||||
"model_provider": "openai",
|
||||
"model": "gpt-4o",
|
||||
},
|
||||
"prompt": {"system_prompt": "Use [§skill:research%2FSKILL.md:Research§]."},
|
||||
}
|
||||
)
|
||||
agent = self._publish_agent()
|
||||
snapshot = self._snapshot(agent_soul)
|
||||
session = FakeSession(scalar=[binding, agent, snapshot, agent, snapshot], scalars=[[binding], []])
|
||||
|
||||
with pytest.raises(WorkflowAgentNodeValidationError, match="skill_ref_dangling"):
|
||||
WorkflowAgentPublishService.validate_agent_nodes_for_publish(
|
||||
session=session,
|
||||
draft_workflow=self._agent_workflow(),
|
||||
)
|
||||
|
||||
def test_projects_binding_declared_outputs_to_draft_graph_response(self):
|
||||
workflow = Workflow(
|
||||
id="workflow-1",
|
||||
@ -2576,18 +2780,17 @@ def test_workspace_dify_tools_returns_provider_and_tool_granularities(monkeypatc
|
||||
assert {entry["granularity"] for entry in entries[1:]} == {"tool"}
|
||||
|
||||
|
||||
# ── ENG-623 §4.4: drive-backed ref validation ────────────────────────────────
|
||||
# ── ENG-623 §4.4: drive-backed prompt mention validation ─────────────────────
|
||||
|
||||
|
||||
def _drive_soul(**overrides):
|
||||
from services.entities.agent_entities import AgentSoulConfig
|
||||
|
||||
base = {
|
||||
"skills_files": {
|
||||
"skills": [
|
||||
{"id": "sk-1", "name": "Tender Analyzer", "skill_md_key": "tender-analyzer/SKILL.md"},
|
||||
],
|
||||
"files": [{"name": "sample.pdf", "drive_key": "files/sample.pdf"}],
|
||||
"prompt": {
|
||||
"system_prompt": (
|
||||
"Use [§skill:tender-analyzer%2FSKILL.md:Tender Analyzer§] and [§file:files%2Fsample.pdf:sample.pdf§]."
|
||||
)
|
||||
},
|
||||
}
|
||||
base.update(overrides)
|
||||
@ -2607,47 +2810,47 @@ def _patch_drive_keys(monkeypatch, existing_keys):
|
||||
return captured
|
||||
|
||||
|
||||
def test_drive_ref_findings_reports_missing_keys(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_drive_mention_findings_reports_missing_keys(monkeypatch: pytest.MonkeyPatch):
|
||||
_patch_drive_keys(monkeypatch, existing_keys=["tender-analyzer/SKILL.md"])
|
||||
|
||||
findings = AgentComposerService._drive_ref_findings(
|
||||
tenant_id="tenant-1", agent_id="agent-1", agent_soul=_drive_soul()
|
||||
findings = AgentComposerService._drive_mention_findings(
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
prompt=_drive_soul().prompt.system_prompt,
|
||||
)
|
||||
|
||||
assert [(f["code"], f["id"]) for f in findings] == [("file_ref_dangling", "files/sample.pdf")]
|
||||
assert str(findings[0]["message"]).startswith("file_ref_dangling: ")
|
||||
assert [(f["code"], f["id"]) for f in findings] == [("mention_target_missing", "files/sample.pdf")]
|
||||
assert findings[0]["kind"] == "file"
|
||||
assert str(findings[0]["message"]).startswith("file 'sample.pdf' has no drive entry")
|
||||
|
||||
|
||||
def test_drive_ref_findings_clean_when_all_keys_exist(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_drive_mention_findings_clean_when_all_keys_exist(monkeypatch: pytest.MonkeyPatch):
|
||||
_patch_drive_keys(monkeypatch, existing_keys=["tender-analyzer/SKILL.md", "files/sample.pdf"])
|
||||
|
||||
assert (
|
||||
AgentComposerService._drive_ref_findings(tenant_id="tenant-1", agent_id="agent-1", agent_soul=_drive_soul())
|
||||
AgentComposerService._drive_mention_findings(
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
prompt=_drive_soul().prompt.system_prompt,
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
|
||||
def test_drive_ref_findings_skips_refs_without_drive_keys(monkeypatch: pytest.MonkeyPatch):
|
||||
# No drive-backed ref at all -> no DB roundtrip, no findings.
|
||||
soul = _drive_soul(
|
||||
skills_files={"skills": [{"id": "legacy", "name": "Legacy"}], "files": [{"name": "u.pdf", "file_id": "u-1"}]}
|
||||
def test_drive_mention_findings_skips_prompt_without_drive_mentions(monkeypatch: pytest.MonkeyPatch):
|
||||
# No drive-backed mention at all -> no DB roundtrip, no findings.
|
||||
soul = _drive_soul(prompt={"system_prompt": "Use [§knowledge:kb-1:Docs§]."})
|
||||
findings = AgentComposerService._drive_mention_findings(
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
prompt=soul.prompt.system_prompt,
|
||||
)
|
||||
findings = AgentComposerService._drive_ref_findings(tenant_id="tenant-1", agent_id="agent-1", agent_soul=soul)
|
||||
assert findings == []
|
||||
|
||||
|
||||
def test_require_drive_refs_resolved_raises_with_stable_code(monkeypatch: pytest.MonkeyPatch):
|
||||
from services.agent.errors import InvalidComposerConfigError
|
||||
|
||||
_patch_drive_keys(monkeypatch, existing_keys=[])
|
||||
|
||||
with pytest.raises(InvalidComposerConfigError, match="skill_ref_dangling"):
|
||||
AgentComposerService._require_drive_refs_resolved(
|
||||
tenant_id="tenant-1", agent_id="agent-1", agent_soul=_drive_soul()
|
||||
)
|
||||
|
||||
|
||||
def test_collect_validation_findings_appends_drive_findings_with_agent_context(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_collect_validation_findings_appends_drive_mention_findings_with_agent_context(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from services.entities.agent_entities import ComposerSavePayload
|
||||
|
||||
_patch_drive_keys(monkeypatch, existing_keys=[])
|
||||
@ -2664,149 +2867,14 @@ def test_collect_validation_findings_appends_drive_findings_with_agent_context(m
|
||||
)
|
||||
|
||||
codes = {w["code"] for w in findings["warnings"]}
|
||||
assert {"skill_ref_dangling", "file_ref_dangling"} <= codes
|
||||
assert codes >= {"mention_target_missing"}
|
||||
assert {w["id"] for w in findings["warnings"] if w["code"] == "mention_target_missing"} == {
|
||||
"tender-analyzer/SKILL.md",
|
||||
"files/sample.pdf",
|
||||
}
|
||||
# without agent context the drive check is skipped entirely
|
||||
findings_no_agent = AgentComposerService.collect_validation_findings(tenant_id="tenant-1", payload=payload)
|
||||
assert all(w["code"] not in {"skill_ref_dangling", "file_ref_dangling"} for w in findings_no_agent["warnings"])
|
||||
|
||||
|
||||
# ── ENG-625 D5: soul-first ref removal ───────────────────────────────────────
|
||||
|
||||
|
||||
def _patch_remove_drive_refs_env(monkeypatch: pytest.MonkeyPatch, *, soul_dict):
|
||||
"""Wire the classmethod's collaborators so soul editing + versioning is observable."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
import services.agent.composer_service as module
|
||||
|
||||
agent = SimpleNamespace(id="agent-1", active_config_snapshot_id="snap-1", updated_by=None)
|
||||
snapshot = SimpleNamespace(id="snap-1", tenant_id="tenant-1", agent_id="agent-1", config_snapshot_dict=soul_dict)
|
||||
committed: dict[str, object] = {}
|
||||
|
||||
fake_session = SimpleNamespace(scalar=lambda stmt: agent, commit=lambda: committed.setdefault("committed", True))
|
||||
monkeypatch.setattr(module.db, "session", fake_session)
|
||||
monkeypatch.setattr(AgentComposerService, "_require_version", classmethod(lambda cls, **kwargs: snapshot))
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_update(cls, *, current_snapshot, account_id, agent_soul, operation, version_note):
|
||||
captured["agent_soul"] = agent_soul
|
||||
captured["version_note"] = version_note
|
||||
return SimpleNamespace(id="snap-2")
|
||||
|
||||
monkeypatch.setattr(AgentComposerService, "_update_current_version", classmethod(fake_update))
|
||||
return agent, captured, committed
|
||||
|
||||
|
||||
def test_remove_drive_refs_drops_skill_by_slug_and_versions(monkeypatch: pytest.MonkeyPatch):
|
||||
soul_dict = {
|
||||
"skills_files": {
|
||||
"skills": [
|
||||
{"id": "sk-1", "name": "Tender Analyzer", "skill_md_key": "tender-analyzer/SKILL.md"},
|
||||
{"id": "sk-2", "name": "Other", "skill_md_key": "other-skill/SKILL.md"},
|
||||
],
|
||||
"files": [],
|
||||
}
|
||||
}
|
||||
agent, captured, committed = _patch_remove_drive_refs_env(monkeypatch, soul_dict=soul_dict)
|
||||
|
||||
version_id = AgentComposerService.remove_drive_refs(
|
||||
tenant_id="tenant-1", agent_id="agent-1", account_id="acc-1", skill_slug="tender-analyzer"
|
||||
)
|
||||
|
||||
assert version_id == "snap-2"
|
||||
assert agent.active_config_snapshot_id == "snap-2"
|
||||
kept = [s.skill_md_key for s in captured["agent_soul"].skills_files.skills]
|
||||
assert kept == ["other-skill/SKILL.md"]
|
||||
assert "Tender Analyzer" in str(captured["version_note"])
|
||||
assert committed.get("committed") is True
|
||||
|
||||
|
||||
def test_remove_drive_refs_is_noop_when_ref_absent(monkeypatch: pytest.MonkeyPatch):
|
||||
soul_dict = {"skills_files": {"skills": [], "files": []}}
|
||||
agent, captured, committed = _patch_remove_drive_refs_env(monkeypatch, soul_dict=soul_dict)
|
||||
|
||||
assert (
|
||||
AgentComposerService.remove_drive_refs(
|
||||
tenant_id="tenant-1", agent_id="agent-1", account_id="acc-1", file_key="files/none.pdf"
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert "agent_soul" not in captured
|
||||
assert committed == {}
|
||||
|
||||
|
||||
def test_remove_drive_refs_drops_file_by_key(monkeypatch: pytest.MonkeyPatch):
|
||||
soul_dict = {
|
||||
"skills_files": {
|
||||
"skills": [],
|
||||
"files": [
|
||||
{"name": "keep.pdf", "drive_key": "files/keep.pdf"},
|
||||
{"name": "drop.pdf", "drive_key": "files/drop.pdf"},
|
||||
],
|
||||
}
|
||||
}
|
||||
_, captured, _ = _patch_remove_drive_refs_env(monkeypatch, soul_dict=soul_dict)
|
||||
|
||||
version_id = AgentComposerService.remove_drive_refs(
|
||||
tenant_id="tenant-1", agent_id="agent-1", account_id="acc-1", file_key="files/drop.pdf"
|
||||
)
|
||||
|
||||
assert version_id == "snap-2"
|
||||
assert [f.drive_key for f in captured["agent_soul"].skills_files.files] == ["files/keep.pdf"]
|
||||
|
||||
|
||||
def test_add_drive_file_ref_adds_or_replaces_file_and_versions(monkeypatch: pytest.MonkeyPatch):
|
||||
soul_dict = {
|
||||
"skills_files": {
|
||||
"skills": [],
|
||||
"files": [
|
||||
{"name": "old.pdf", "drive_key": "files/old.pdf"},
|
||||
{"name": "stale.pdf", "drive_key": "files/new.pdf"},
|
||||
],
|
||||
}
|
||||
}
|
||||
agent, captured, committed = _patch_remove_drive_refs_env(monkeypatch, soul_dict=soul_dict)
|
||||
|
||||
version_id = AgentComposerService.add_drive_file_ref(
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
account_id="acc-1",
|
||||
file_ref=AgentFileRefConfig(name="new.pdf", file_id="uf-1", drive_key="files/new.pdf", type="application/pdf"),
|
||||
)
|
||||
|
||||
assert version_id == "snap-2"
|
||||
assert agent.active_config_snapshot_id == "snap-2"
|
||||
assert [f.drive_key for f in captured["agent_soul"].skills_files.files] == ["files/old.pdf", "files/new.pdf"]
|
||||
assert captured["agent_soul"].skills_files.files[-1].name == "new.pdf"
|
||||
assert "new.pdf" in str(captured["version_note"])
|
||||
assert committed.get("committed") is True
|
||||
|
||||
|
||||
def test_add_drive_file_ref_syncs_workflow_binding_snapshot(monkeypatch: pytest.MonkeyPatch):
|
||||
binding = SimpleNamespace(agent_id="agent-1", current_snapshot_id="snap-1", updated_by=None)
|
||||
_patch_remove_drive_refs_env(monkeypatch, soul_dict={"skills_files": {"skills": [], "files": []}})
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService, "_get_draft_workflow", classmethod(lambda cls, **kwargs: SimpleNamespace(id="wf-1"))
|
||||
)
|
||||
monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", classmethod(lambda cls, **kwargs: binding))
|
||||
|
||||
AgentComposerService.add_drive_file_ref(
|
||||
tenant_id="tenant-1",
|
||||
agent_id="agent-1",
|
||||
account_id="acc-1",
|
||||
file_ref=AgentFileRefConfig(name="new.pdf", file_id="uf-1", drive_key="files/new.pdf"),
|
||||
app_id="app-1",
|
||||
node_id="agent-node-1",
|
||||
)
|
||||
|
||||
assert binding.current_snapshot_id == "snap-2"
|
||||
assert binding.updated_by == "acc-1"
|
||||
|
||||
|
||||
def test_remove_drive_refs_requires_exactly_one_scope():
|
||||
with pytest.raises(ValueError):
|
||||
AgentComposerService.remove_drive_refs(tenant_id="t", agent_id="a", account_id="u")
|
||||
assert all(w["code"] != "mention_target_missing" for w in findings_no_agent["warnings"])
|
||||
|
||||
|
||||
# ── ENG-623/625: resolver helpers + save-path drive guard ────────────────────
|
||||
@ -2844,58 +2912,7 @@ def test_resolve_workflow_node_agent_id_degrades_without_workflow_or_binding(mon
|
||||
assert AgentComposerService.resolve_workflow_node_agent_id(tenant_id="t", app_id="a", node_id="n") == "agent-7"
|
||||
|
||||
|
||||
def test_remove_drive_refs_returns_none_without_agent_or_snapshot(monkeypatch: pytest.MonkeyPatch):
|
||||
from types import SimpleNamespace
|
||||
|
||||
import services.agent.composer_service as module
|
||||
|
||||
monkeypatch.setattr(module.db, "session", SimpleNamespace(scalar=lambda stmt: None))
|
||||
assert AgentComposerService.remove_drive_refs(tenant_id="t", agent_id="a", account_id="u", skill_slug="s") is None
|
||||
|
||||
agent_without_snapshot = SimpleNamespace(id="a", active_config_snapshot_id=None)
|
||||
monkeypatch.setattr(module.db, "session", SimpleNamespace(scalar=lambda stmt: agent_without_snapshot))
|
||||
assert AgentComposerService.remove_drive_refs(tenant_id="t", agent_id="a", account_id="u", skill_slug="s") is None
|
||||
|
||||
|
||||
def test_save_workflow_composer_guards_drive_refs_for_existing_agent_strategies(monkeypatch: pytest.MonkeyPatch):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from services.entities.agent_entities import ComposerSavePayload
|
||||
|
||||
payload = ComposerSavePayload.model_validate(
|
||||
{
|
||||
"variant": "workflow",
|
||||
"save_strategy": "save_to_current_version",
|
||||
"agent_soul": _drive_soul().model_dump(mode="json"),
|
||||
"soul_lock": {"locked": False},
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService, "_get_draft_workflow", classmethod(lambda cls, **kwargs: SimpleNamespace(id="wf-1"))
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService,
|
||||
"_get_workflow_binding",
|
||||
classmethod(lambda cls, **kwargs: SimpleNamespace(agent_id="agent-1")),
|
||||
)
|
||||
guarded: dict[str, str] = {}
|
||||
|
||||
def fake_guard(cls, *, tenant_id, agent_id, agent_soul):
|
||||
guarded["agent_id"] = agent_id
|
||||
raise InvalidComposerConfigError("skill_ref_dangling: boom")
|
||||
|
||||
from services.agent.errors import InvalidComposerConfigError
|
||||
|
||||
monkeypatch.setattr(AgentComposerService, "_require_drive_refs_resolved", classmethod(fake_guard))
|
||||
|
||||
with pytest.raises(InvalidComposerConfigError, match="skill_ref_dangling"):
|
||||
AgentComposerService.save_workflow_composer(
|
||||
tenant_id="t-1", app_id="app-1", node_id="n-1", account_id="acc-1", payload=payload
|
||||
)
|
||||
assert guarded["agent_id"] == "agent-1"
|
||||
|
||||
|
||||
def test_save_workflow_composer_guards_drive_refs_for_inline_node_job_only(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_save_workflow_composer_reports_drive_mentions_for_inline_node_job_only(monkeypatch: pytest.MonkeyPatch):
|
||||
payload = ComposerSavePayload.model_validate(
|
||||
{
|
||||
"variant": "workflow",
|
||||
@ -2933,26 +2950,27 @@ def test_save_workflow_composer_guards_drive_refs_for_inline_node_job_only(monke
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService, "_serialize_workflow_state", classmethod(lambda cls, **kwargs: {"state": "ok"})
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService, "collect_validation_findings", classmethod(lambda cls, **kwargs: {"warnings": []})
|
||||
)
|
||||
guarded: dict[str, str] = {}
|
||||
|
||||
def fake_guard(cls, *, tenant_id, agent_id, agent_soul):
|
||||
def fake_collect(cls, *, tenant_id, payload, agent_id=None):
|
||||
guarded["tenant_id"] = tenant_id
|
||||
guarded["agent_id"] = agent_id
|
||||
return {"warnings": [{"code": "mention_target_missing", "id": "files/sample.pdf"}]}
|
||||
|
||||
monkeypatch.setattr(AgentComposerService, "_require_drive_refs_resolved", classmethod(fake_guard))
|
||||
monkeypatch.setattr(AgentComposerService, "collect_validation_findings", classmethod(fake_collect))
|
||||
|
||||
result = AgentComposerService.save_workflow_composer(
|
||||
tenant_id="t-1", app_id="app-1", node_id="n-1", account_id="acc-1", payload=payload
|
||||
)
|
||||
|
||||
assert result == {"state": "ok", "validation": {"warnings": []}}
|
||||
assert result == {
|
||||
"state": "ok",
|
||||
"validation": {"warnings": [{"code": "mention_target_missing", "id": "files/sample.pdf"}]},
|
||||
}
|
||||
assert guarded == {"tenant_id": "t-1", "agent_id": "agent-1"}
|
||||
|
||||
|
||||
def test_save_workflow_composer_skips_drive_refs_for_roster_node_job_only(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_save_workflow_composer_reports_drive_mentions_for_roster_node_job_only(monkeypatch: pytest.MonkeyPatch):
|
||||
payload = ComposerSavePayload.model_validate(
|
||||
{
|
||||
"variant": "workflow",
|
||||
@ -2990,29 +3008,17 @@ def test_save_workflow_composer_skips_drive_refs_for_roster_node_job_only(monkey
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService, "_serialize_workflow_state", classmethod(lambda cls, **kwargs: {"state": "ok"})
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
AgentComposerService, "collect_validation_findings", classmethod(lambda cls, **kwargs: {"warnings": []})
|
||||
)
|
||||
captured: dict[str, str | None] = {}
|
||||
|
||||
def fail_guard(cls, *, tenant_id, agent_id, agent_soul):
|
||||
raise AssertionError("roster node-job-only saves must not validate agent drive refs")
|
||||
def fake_collect(cls, *, tenant_id, payload, agent_id=None):
|
||||
captured["agent_id"] = agent_id
|
||||
return {"warnings": []}
|
||||
|
||||
monkeypatch.setattr(AgentComposerService, "_require_drive_refs_resolved", classmethod(fail_guard))
|
||||
monkeypatch.setattr(AgentComposerService, "collect_validation_findings", classmethod(fake_collect))
|
||||
|
||||
result = AgentComposerService.save_workflow_composer(
|
||||
tenant_id="t-1", app_id="app-1", node_id="n-1", account_id="acc-1", payload=payload
|
||||
)
|
||||
|
||||
assert result == {"state": "ok", "validation": {"warnings": []}}
|
||||
|
||||
|
||||
def test_remove_drive_refs_noop_when_skill_slug_unmatched(monkeypatch: pytest.MonkeyPatch):
|
||||
soul_dict = {"skills_files": {"skills": [{"name": "Other", "skill_md_key": "other/SKILL.md"}], "files": []}}
|
||||
_, captured, committed = _patch_remove_drive_refs_env(monkeypatch, soul_dict=soul_dict)
|
||||
assert (
|
||||
AgentComposerService.remove_drive_refs(
|
||||
tenant_id="t-1", agent_id="agent-1", account_id="acc-1", skill_slug="ghost"
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert committed == {}
|
||||
assert captured["agent_id"] == "agent-1"
|
||||
|
||||
@ -118,10 +118,6 @@ def test_previous_outputs_capped_and_flagged():
|
||||
def _soul() -> AgentSoulConfig:
|
||||
return AgentSoulConfig.model_validate(
|
||||
{
|
||||
"skills_files": {
|
||||
"skills": [{"id": "sk-1", "name": "tender-analyzer"}],
|
||||
"files": [{"id": "f-1", "name": "qna_report.pdf"}],
|
||||
},
|
||||
"tools": {
|
||||
"cli_tools": [
|
||||
{"id": "ct-1", "name": "ffmpeg"},
|
||||
@ -144,7 +140,6 @@ def test_soul_candidates_lists_configured_items_only():
|
||||
)
|
||||
|
||||
assert truncated is False
|
||||
assert [item["kind"] for item in lists["skills_files"]] == ["skill", "file"]
|
||||
assert [item["name"] for item in lists["cli_tools"]] == ["ffmpeg"]
|
||||
# the stable mention id flows through so the frontend can mint [§cli_tool:<id>§]
|
||||
assert [item["id"] for item in lists["cli_tools"]] == ["ct-1"]
|
||||
@ -158,35 +153,19 @@ def test_soul_candidates_lists_configured_items_only():
|
||||
assert lists["dify_tools"][0]["id"] == "tavily/tavily_search"
|
||||
|
||||
|
||||
def test_candidates_response_preserves_skill_and_file_candidate_shapes():
|
||||
def test_candidates_response_omits_legacy_skill_file_candidates():
|
||||
response = AgentComposerCandidatesResponse.model_validate(
|
||||
{
|
||||
"variant": "agent_app",
|
||||
"allowed_node_job_candidates": {},
|
||||
"allowed_soul_candidates": {
|
||||
"skills_files": [
|
||||
{"kind": "skill", "id": "sk-1", "name": "tender-analyzer", "path": "skills/tender.md"},
|
||||
{
|
||||
"kind": "file",
|
||||
"id": "f-1",
|
||||
"name": "qna_report.pdf",
|
||||
"transfer_method": "local_file",
|
||||
"reference": "upload-1",
|
||||
"url": "https://files.example/qna_report.pdf",
|
||||
},
|
||||
]
|
||||
"cli_tools": [],
|
||||
},
|
||||
"capabilities": {"human_roster_available": False},
|
||||
}
|
||||
).model_dump(mode="json")
|
||||
|
||||
skill, file = response["allowed_soul_candidates"]["skills_files"]
|
||||
assert skill["kind"] == "skill"
|
||||
assert skill["path"] == "skills/tender.md"
|
||||
assert file["kind"] == "file"
|
||||
assert file["transfer_method"] == "local_file"
|
||||
assert file["reference"] == "upload-1"
|
||||
assert file["url"] == "https://files.example/qna_report.pdf"
|
||||
assert "skills_files" not in response["allowed_soul_candidates"]
|
||||
|
||||
|
||||
def test_soul_candidates_empty_config_yields_empty_lists():
|
||||
|
||||
@ -171,7 +171,6 @@ def test_configured_but_deleted_dataset_surfaces_as_placeholder():
|
||||
def test_unresolved_non_knowledge_mentions_warn_target_missing():
|
||||
findings = _findings(_soul_payload("use [§skill:nope:Ghost Skill§] and [§human:missing§]"))
|
||||
codes = [(w["code"], w["kind"]) for w in findings["warnings"]]
|
||||
assert ("mention_target_missing", "skill") in codes
|
||||
assert ("mention_target_missing", "human") in codes
|
||||
assert findings["knowledge_retrieval_placeholder"] == []
|
||||
|
||||
|
||||
@ -7,11 +7,13 @@ guarantees no mention-shaped marker survives to the model.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from urllib.parse import quote
|
||||
|
||||
import pytest
|
||||
|
||||
from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig
|
||||
from services.agent.prompt_mentions import (
|
||||
MAX_MENTION_FIELD_LENGTH,
|
||||
MAX_MENTION_REF_ID_LENGTH,
|
||||
NODE_JOB_PROMPT_ALLOWED_KINDS,
|
||||
SOUL_PROMPT_ALLOWED_KINDS,
|
||||
MentionKind,
|
||||
@ -26,11 +28,11 @@ from services.agent.prompt_mentions import (
|
||||
|
||||
|
||||
def test_parse_extracts_kind_id_and_optional_label():
|
||||
prompt = "Use [§skill:abc-1:tender-analyzer§] then ask [§human:c-1§]."
|
||||
prompt = "Use [§skill:tender-analyzer%2FSKILL.md:tender-analyzer§] then ask [§human:c-1§]."
|
||||
mentions = parse_prompt_mentions(prompt)
|
||||
|
||||
assert [(m.kind, m.ref_id, m.label) for m in mentions] == [
|
||||
(MentionKind.SKILL, "abc-1", "tender-analyzer"),
|
||||
(MentionKind.SKILL, "tender-analyzer%2FSKILL.md", "tender-analyzer"),
|
||||
(MentionKind.HUMAN, "c-1", None),
|
||||
]
|
||||
assert prompt[mentions[0].start : mentions[0].end] == mentions[0].raw
|
||||
@ -48,10 +50,16 @@ def test_parse_ignores_legacy_template_forms_and_unknown_kinds():
|
||||
|
||||
|
||||
def test_parse_skips_oversized_id_or_label():
|
||||
long_id = "x" * (MAX_MENTION_FIELD_LENGTH + 1)
|
||||
long_id = "x" * (MAX_MENTION_REF_ID_LENGTH + 1)
|
||||
assert parse_prompt_mentions(f"[§skill:{long_id}§]") == []
|
||||
|
||||
|
||||
def test_parse_accepts_long_unicode_encoded_drive_key_within_drive_limit():
|
||||
encoded_drive_key = quote("你" * 512)
|
||||
mentions = parse_prompt_mentions(f"[§skill:{encoded_drive_key}:Long Skill§]")
|
||||
assert [(mention.kind, mention.ref_id) for mention in mentions] == [(MentionKind.SKILL, encoded_drive_key)]
|
||||
|
||||
|
||||
# ── expand + scrub ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@ -88,10 +96,6 @@ def test_expand_empty_prompt_is_noop():
|
||||
def soul() -> AgentSoulConfig:
|
||||
return AgentSoulConfig.model_validate(
|
||||
{
|
||||
"skills_files": {
|
||||
"skills": [{"id": "sk-1", "name": "tender-analyzer"}],
|
||||
"files": [{"id": "f-1", "name": "qna_report.pdf"}],
|
||||
},
|
||||
"tools": {
|
||||
"dify_tools": [
|
||||
{
|
||||
@ -112,17 +116,13 @@ def soul() -> AgentSoulConfig:
|
||||
def test_soul_resolver_resolves_each_kind(soul: AgentSoulConfig):
|
||||
resolver = build_soul_mention_resolver(soul)
|
||||
prompt = (
|
||||
"Use [§skill:sk-1§] with [§file:f-1§], search via "
|
||||
"[§tool:tavily/tavily_search:tavily§], run [§cli_tool:ct-1:ffmpeg§], "
|
||||
"Use [§tool:tavily/tavily_search:tavily§], run [§cli_tool:ct-1:ffmpeg§], "
|
||||
"ground in [§knowledge:ds-1§], ask [§human:c-1§]."
|
||||
)
|
||||
|
||||
expanded = expand_prompt_mentions(prompt, resolver)
|
||||
|
||||
assert expanded == (
|
||||
"Use tender-analyzer with qna_report.pdf, search via tavily_search, "
|
||||
"run ffmpeg, ground in 产品手册, ask EMAIL · David Hayes."
|
||||
)
|
||||
assert expanded == ("Use tavily_search, run ffmpeg, ground in 产品手册, ask EMAIL · David Hayes.")
|
||||
|
||||
|
||||
def test_soul_resolver_unknown_ids_degrade(soul: AgentSoulConfig):
|
||||
|
||||
@ -121,16 +121,3 @@ def test_read_member_bytes_roundtrip_and_errors():
|
||||
with pytest.raises(SkillPackageError) as bad_zip:
|
||||
service.read_member_bytes(content=b"not a zip", member_path="SKILL.md")
|
||||
assert bad_zip.value.code == "invalid_archive"
|
||||
|
||||
|
||||
def test_to_skill_ref_carries_metadata():
|
||||
manifest = _extract({"SKILL.md": _SKILL_MD.encode()})
|
||||
ref = manifest.to_skill_ref(file_id="upload-1", path="pdf-toolkit/.DIFY-SKILL-FULL.zip")
|
||||
|
||||
assert ref.name == "PDF Toolkit"
|
||||
assert ref.file_id == "upload-1"
|
||||
assert ref.path == "pdf-toolkit/.DIFY-SKILL-FULL.zip"
|
||||
assert ref.id == manifest.hash
|
||||
dumped = ref.model_dump()
|
||||
assert dumped["hash"] == manifest.hash
|
||||
assert dumped["entry_path"] == "SKILL.md"
|
||||
|
||||
@ -32,16 +32,30 @@ def test_slugify_skill_name():
|
||||
assert slugify_skill_name("") == "skill"
|
||||
|
||||
|
||||
def test_standardize_creates_two_drive_owned_toolfiles_and_commits():
|
||||
def test_standardize_creates_drive_owned_toolfiles_and_commits_archive_members():
|
||||
content = _zip({"SKILL.md": _SKILL_MD, "scripts/run.py": b"print('x')\n"})
|
||||
|
||||
tool_files = MagicMock()
|
||||
tool_files.create_file_by_raw.side_effect = [
|
||||
SimpleNamespace(id="md-tool-file"),
|
||||
SimpleNamespace(id="zip-tool-file"),
|
||||
SimpleNamespace(id="script-tool-file"),
|
||||
]
|
||||
drive = MagicMock()
|
||||
drive.commit.return_value = []
|
||||
drive.list_skills.return_value = [
|
||||
{
|
||||
"path": "pdf-toolkit",
|
||||
"skill_md_key": "pdf-toolkit/SKILL.md",
|
||||
"archive_key": "pdf-toolkit/.DIFY-SKILL-FULL.zip",
|
||||
"name": "PDF Toolkit",
|
||||
"description": "Work with PDFs.",
|
||||
"size": len(_SKILL_MD),
|
||||
"mime_type": "text/markdown",
|
||||
"hash": None,
|
||||
"created_at": None,
|
||||
},
|
||||
]
|
||||
|
||||
service = SkillStandardizeService(tool_file_manager=tool_files, drive_service=drive)
|
||||
result = service.standardize(
|
||||
@ -52,33 +66,39 @@ def test_standardize_creates_two_drive_owned_toolfiles_and_commits():
|
||||
agent_id="agent-1",
|
||||
)
|
||||
|
||||
# Two ToolFiles: SKILL.md (markdown) + full archive (zip).
|
||||
assert tool_files.create_file_by_raw.call_count == 2
|
||||
md_call, zip_call = tool_files.create_file_by_raw.call_args_list
|
||||
# ToolFiles: SKILL.md, full archive, and each inspectable package member.
|
||||
assert tool_files.create_file_by_raw.call_count == 3
|
||||
md_call, zip_call, script_call = tool_files.create_file_by_raw.call_args_list
|
||||
assert md_call.kwargs["mimetype"] == "text/markdown"
|
||||
assert md_call.kwargs["file_binary"] == _SKILL_MD
|
||||
assert zip_call.kwargs["mimetype"] == "application/zip"
|
||||
assert zip_call.kwargs["file_binary"] == content
|
||||
assert script_call.kwargs["mimetype"] in {"text/x-python", "text/plain", "application/octet-stream"}
|
||||
assert script_call.kwargs["file_binary"] == b"print('x')\n"
|
||||
assert script_call.kwargs["filename"] == "run.py"
|
||||
|
||||
# Committed as drive-owned with the standardized keys.
|
||||
commit_kwargs = drive.commit.call_args.kwargs
|
||||
assert commit_kwargs["agent_id"] == "agent-1"
|
||||
items = commit_kwargs["items"]
|
||||
assert [item.key for item in items] == ["pdf-toolkit/SKILL.md", "pdf-toolkit/.DIFY-SKILL-FULL.zip"]
|
||||
assert [item.key for item in items] == [
|
||||
"pdf-toolkit/SKILL.md",
|
||||
"pdf-toolkit/.DIFY-SKILL-FULL.zip",
|
||||
"pdf-toolkit/scripts/run.py",
|
||||
]
|
||||
assert all(item.value_owned_by_drive for item in items)
|
||||
assert [item.file_ref.id for item in items] == ["md-tool-file", "zip-tool-file"]
|
||||
assert [item.file_ref.id for item in items] == ["md-tool-file", "zip-tool-file", "script-tool-file"]
|
||||
assert items[0].is_skill is True
|
||||
assert items[0].skill_metadata is not None
|
||||
assert items[0].skill_metadata.name == "PDF Toolkit"
|
||||
assert items[0].skill_metadata.manifest_files == ["SKILL.md", "scripts/run.py"]
|
||||
assert items[1].is_skill is False
|
||||
assert items[2].is_skill is False
|
||||
|
||||
# The returned skill ref carries stable drive paths + file ids.
|
||||
# The returned upload response carries only the drive-derived fields the UI needs.
|
||||
skill = result["skill"]
|
||||
assert skill["path"] == "pdf-toolkit"
|
||||
assert skill["name"] == "PDF Toolkit"
|
||||
assert skill["full_archive_file_id"] == "zip-tool-file"
|
||||
assert skill["skill_md_file_id"] == "md-tool-file"
|
||||
assert skill["archive_key"] == "pdf-toolkit/.DIFY-SKILL-FULL.zip"
|
||||
assert skill["skill_md_key"] == "pdf-toolkit/SKILL.md"
|
||||
# ENG-371: zip member listing persisted for infer-tools signals
|
||||
assert "SKILL.md" in skill["manifest_files"]
|
||||
assert "scripts/run.py" in skill["manifest_files"]
|
||||
assert result["manifest"]["files"] == ["SKILL.md", "scripts/run.py"]
|
||||
|
||||
@ -29,13 +29,8 @@ def _service(preview=_SKILL_MD_PREVIEW):
|
||||
return SkillToolInferenceService(drive_service=drive), drive
|
||||
|
||||
|
||||
def _patch_soul_files(monkeypatch, files):
|
||||
monkeypatch.setattr(SkillToolInferenceService, "_manifest_files_from_soul", staticmethod(lambda **kwargs: files))
|
||||
|
||||
|
||||
def test_infer_returns_suggestions_with_inferred_from(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_infer_returns_suggestions_with_inferred_from(monkeypatch):
|
||||
service, drive = _service()
|
||||
_patch_soul_files(monkeypatch, ["SKILL.md", "scripts/transcribe.sh"])
|
||||
raw = (
|
||||
'{"inferable": true, "reason": null, "cli_tools": [{"name": "ffmpeg",'
|
||||
' "description": "transcoding for step 2", "command": "ffmpeg",'
|
||||
@ -53,9 +48,8 @@ def test_infer_returns_suggestions_with_inferred_from(monkeypatch: pytest.Monkey
|
||||
drive.preview.assert_called_once_with(tenant_id="t-1", agent_id="a-1", key="audio-transcribe/SKILL.md")
|
||||
|
||||
|
||||
def test_infer_threads_manifest_files_into_the_prompt(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_infer_threads_skill_md_into_the_prompt(monkeypatch):
|
||||
service, _ = _service()
|
||||
_patch_soul_files(monkeypatch, ["scripts/run.sh"])
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
def fake_invoke(*, tenant_id, user_prompt):
|
||||
@ -65,22 +59,20 @@ def test_infer_threads_manifest_files_into_the_prompt(monkeypatch: pytest.Monkey
|
||||
with patch.object(SkillToolInferenceService, "_invoke", staticmethod(fake_invoke)):
|
||||
service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe")
|
||||
|
||||
assert "scripts/run.sh" in captured["prompt"]
|
||||
assert "Files inside the skill package" not in captured["prompt"]
|
||||
assert "ffmpeg" in captured["prompt"] # SKILL.md body present
|
||||
|
||||
|
||||
def test_infer_not_inferable_passes_reason_through(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_infer_not_inferable_passes_reason_through(monkeypatch):
|
||||
service, _ = _service()
|
||||
_patch_soul_files(monkeypatch, [])
|
||||
raw = '{"inferable": false, "cli_tools": [], "reason": "SKILL.md 未描述任何外部命令依赖"}'
|
||||
with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)):
|
||||
result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe")
|
||||
assert result == {"inferable": False, "cli_tools": [], "reason": "SKILL.md 未描述任何外部命令依赖"}
|
||||
|
||||
|
||||
def test_infer_retries_once_then_422(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_infer_retries_once_then_422(monkeypatch):
|
||||
service, _ = _service()
|
||||
_patch_soul_files(monkeypatch, [])
|
||||
calls: list[int] = []
|
||||
|
||||
def bad_invoke(**kwargs):
|
||||
@ -96,9 +88,8 @@ def test_infer_retries_once_then_422(monkeypatch: pytest.MonkeyPatch):
|
||||
assert exc_info.value.status_code == 422
|
||||
|
||||
|
||||
def test_infer_repairs_slightly_malformed_json(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_infer_repairs_slightly_malformed_json(monkeypatch):
|
||||
service, _ = _service()
|
||||
_patch_soul_files(monkeypatch, [])
|
||||
raw = 'Here you go: {"inferable": true, "cli_tools": [], "reason": null,}'
|
||||
with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)):
|
||||
result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe")
|
||||
@ -123,10 +114,10 @@ def test_binary_skill_md_maps_to_404():
|
||||
assert exc_info.value.code == "skill_not_found"
|
||||
|
||||
|
||||
# ── real-path coverage: _invoke / _manifest_files_from_soul / passthrough ────
|
||||
# ── real-path coverage: _invoke / passthrough ────────────────────────────────
|
||||
|
||||
|
||||
def test_invoke_maps_missing_default_model_to_400(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_invoke_maps_missing_default_model_to_400(monkeypatch):
|
||||
import services.agent.skill_tool_inference_service as module
|
||||
from core.errors.error import ProviderTokenNotInitError
|
||||
|
||||
@ -140,7 +131,7 @@ def test_invoke_maps_missing_default_model_to_400(monkeypatch: pytest.MonkeyPatc
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
||||
def test_invoke_maps_model_failure_to_422_and_success_returns_text(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_invoke_maps_model_failure_to_422_and_success_returns_text(monkeypatch):
|
||||
import services.agent.skill_tool_inference_service as module
|
||||
|
||||
fake_manager = MagicMock()
|
||||
@ -171,55 +162,3 @@ def test_load_skill_md_passes_through_non_missing_drive_errors():
|
||||
with pytest.raises(SkillToolInferenceError) as exc_info:
|
||||
service.infer(tenant_id="t-1", agent_id="a-1", slug="x")
|
||||
assert exc_info.value.code == "agent_not_found"
|
||||
|
||||
|
||||
def _patch_inference_db(monkeypatch: pytest.MonkeyPatch, *, agent, snapshot):
|
||||
from types import SimpleNamespace
|
||||
|
||||
import services.agent.skill_tool_inference_service as module
|
||||
|
||||
results = iter([agent, snapshot])
|
||||
monkeypatch.setattr(module.db, "session", SimpleNamespace(scalar=lambda stmt: next(results)))
|
||||
|
||||
|
||||
def test_manifest_files_from_soul_reads_active_snapshot(monkeypatch: pytest.MonkeyPatch):
|
||||
from types import SimpleNamespace
|
||||
|
||||
soul_dict = {
|
||||
"skills_files": {
|
||||
"skills": [
|
||||
{"name": "Other", "skill_md_key": "other/SKILL.md", "manifest_files": ["x.md"]},
|
||||
{"name": "Audio", "skill_md_key": "audio-transcribe/SKILL.md", "manifest_files": ["scripts/a.sh"]},
|
||||
]
|
||||
}
|
||||
}
|
||||
agent = SimpleNamespace(active_config_snapshot_id="snap-1")
|
||||
snapshot = SimpleNamespace(config_snapshot_dict=soul_dict)
|
||||
_patch_inference_db(monkeypatch, agent=agent, snapshot=snapshot)
|
||||
|
||||
files = SkillToolInferenceService._manifest_files_from_soul(
|
||||
tenant_id="t-1", agent_id="a-1", slug="audio-transcribe"
|
||||
)
|
||||
assert files == ["scripts/a.sh"]
|
||||
|
||||
|
||||
def test_manifest_files_from_soul_degrades_when_agent_or_snapshot_missing(monkeypatch: pytest.MonkeyPatch):
|
||||
_patch_inference_db(monkeypatch, agent=None, snapshot=None)
|
||||
assert SkillToolInferenceService._manifest_files_from_soul(tenant_id="t", agent_id="a", slug="s") == []
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
_patch_inference_db(monkeypatch, agent=SimpleNamespace(active_config_snapshot_id="snap-1"), snapshot=None)
|
||||
assert SkillToolInferenceService._manifest_files_from_soul(tenant_id="t", agent_id="a", slug="s") == []
|
||||
|
||||
|
||||
def test_manifest_files_from_soul_empty_when_slug_not_in_soul(monkeypatch: pytest.MonkeyPatch):
|
||||
from types import SimpleNamespace
|
||||
|
||||
soul_dict = {"skills_files": {"skills": [{"name": "Other", "skill_md_key": "other/SKILL.md"}]}}
|
||||
_patch_inference_db(
|
||||
monkeypatch,
|
||||
agent=SimpleNamespace(active_config_snapshot_id="snap-1"),
|
||||
snapshot=SimpleNamespace(config_snapshot_dict=soul_dict),
|
||||
)
|
||||
assert SkillToolInferenceService._manifest_files_from_soul(tenant_id="t", agent_id="a", slug="ghost") == []
|
||||
|
||||
@ -46,7 +46,7 @@ class TestCatalog:
|
||||
assert call.tenant_id == "tenant-1"
|
||||
assert call.account_id == "acct-1"
|
||||
assert call.json is None
|
||||
assert call.params is None
|
||||
assert call.params == {"billing_enabled": svc.dify_config.BILLING_ENABLED}
|
||||
assert len(out.groups) == 1
|
||||
assert out.groups[0].group_key == "workspace"
|
||||
|
||||
|
||||
@ -147,8 +147,7 @@ class TestHitTestingServiceRetrieve:
|
||||
Provides a mocked database session for testing database operations
|
||||
like adding and committing DatasetQuery records.
|
||||
"""
|
||||
with patch("services.hit_testing_service.db.session", autospec=True) as mock_db:
|
||||
yield mock_db
|
||||
return MagicMock()
|
||||
|
||||
def test_retrieve_success_with_default_retrieval_model(self, mock_db_session):
|
||||
"""
|
||||
@ -186,7 +185,9 @@ class TestHitTestingServiceRetrieve:
|
||||
mock_format.return_value = mock_records
|
||||
|
||||
# Act
|
||||
result = HitTestingService.retrieve(dataset, query, account, retrieval_model, external_retrieval_model)
|
||||
result = HitTestingService.retrieve(
|
||||
mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result["query"]["content"] == query
|
||||
@ -232,7 +233,9 @@ class TestHitTestingServiceRetrieve:
|
||||
mock_format.return_value = mock_records
|
||||
|
||||
# Act
|
||||
result = HitTestingService.retrieve(dataset, query, account, retrieval_model, external_retrieval_model)
|
||||
result = HitTestingService.retrieve(
|
||||
mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result["query"]["content"] == query
|
||||
@ -257,9 +260,11 @@ class TestHitTestingServiceRetrieve:
|
||||
retrieval_model = {
|
||||
"metadata_filtering_conditions": {
|
||||
"conditions": [
|
||||
{"field": "category", "operator": "is", "value": "test"},
|
||||
{"name": "category", "comparison_operator": "is", "value": "test"},
|
||||
],
|
||||
},
|
||||
"reranking_enable": False,
|
||||
"score_threshold_enabled": False,
|
||||
}
|
||||
external_retrieval_model = {}
|
||||
|
||||
@ -286,7 +291,9 @@ class TestHitTestingServiceRetrieve:
|
||||
mock_format.return_value = mock_records
|
||||
|
||||
# Act
|
||||
result = HitTestingService.retrieve(dataset, query, account, retrieval_model, external_retrieval_model)
|
||||
result = HitTestingService.retrieve(
|
||||
mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result["query"]["content"] == query
|
||||
@ -308,9 +315,11 @@ class TestHitTestingServiceRetrieve:
|
||||
retrieval_model = {
|
||||
"metadata_filtering_conditions": {
|
||||
"conditions": [
|
||||
{"field": "category", "operator": "is", "value": "test"},
|
||||
{"name": "category", "comparison_operator": "is", "value": "test"},
|
||||
],
|
||||
},
|
||||
"reranking_enable": False,
|
||||
"score_threshold_enabled": False,
|
||||
}
|
||||
external_retrieval_model = {}
|
||||
|
||||
@ -327,7 +336,9 @@ class TestHitTestingServiceRetrieve:
|
||||
mock_format.return_value = []
|
||||
|
||||
# Act
|
||||
result = HitTestingService.retrieve(dataset, query, account, retrieval_model, external_retrieval_model)
|
||||
result = HitTestingService.retrieve(
|
||||
mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result["query"]["content"] == query
|
||||
@ -344,6 +355,8 @@ class TestHitTestingServiceRetrieve:
|
||||
dataset_retrieval_model = {
|
||||
"search_method": RetrievalMethod.HYBRID_SEARCH,
|
||||
"top_k": 3,
|
||||
"reranking_enable": False,
|
||||
"score_threshold_enabled": False,
|
||||
}
|
||||
dataset = HitTestingTestDataFactory.create_dataset_mock(retrieval_model=dataset_retrieval_model)
|
||||
account = HitTestingTestDataFactory.create_user_mock()
|
||||
@ -366,7 +379,9 @@ class TestHitTestingServiceRetrieve:
|
||||
mock_format.return_value = mock_records
|
||||
|
||||
# Act
|
||||
result = HitTestingService.retrieve(dataset, query, account, retrieval_model, external_retrieval_model)
|
||||
result = HitTestingService.retrieve(
|
||||
mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result["query"]["content"] == query
|
||||
@ -391,8 +406,7 @@ class TestHitTestingServiceExternalRetrieve:
|
||||
Provides a mocked database session for testing database operations
|
||||
like adding and committing DatasetQuery records.
|
||||
"""
|
||||
with patch("services.hit_testing_service.db.session", autospec=True) as mock_db:
|
||||
yield mock_db
|
||||
return MagicMock()
|
||||
|
||||
def test_external_retrieve_success(self, mock_db_session):
|
||||
"""
|
||||
@ -424,7 +438,7 @@ class TestHitTestingServiceExternalRetrieve:
|
||||
|
||||
# Act
|
||||
result = HitTestingService.external_retrieve(
|
||||
dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
)
|
||||
|
||||
# Assert
|
||||
@ -455,7 +469,7 @@ class TestHitTestingServiceExternalRetrieve:
|
||||
|
||||
# Act
|
||||
result = HitTestingService.external_retrieve(
|
||||
dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
)
|
||||
|
||||
# Assert
|
||||
@ -490,7 +504,7 @@ class TestHitTestingServiceExternalRetrieve:
|
||||
|
||||
# Act
|
||||
result = HitTestingService.external_retrieve(
|
||||
dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
)
|
||||
|
||||
# Assert
|
||||
@ -524,7 +538,7 @@ class TestHitTestingServiceExternalRetrieve:
|
||||
|
||||
# Act
|
||||
result = HitTestingService.external_retrieve(
|
||||
dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions
|
||||
)
|
||||
|
||||
# Assert
|
||||
@ -565,7 +579,7 @@ class TestHitTestingServiceCompactRetrieveResponse:
|
||||
mock_format.return_value = mock_records
|
||||
|
||||
# Act
|
||||
result = HitTestingService.compact_retrieve_response(query, documents)
|
||||
result = HitTestingService.compact_retrieve_response(MagicMock(), query, documents)
|
||||
|
||||
# Assert
|
||||
assert result["query"]["content"] == query
|
||||
@ -591,7 +605,7 @@ class TestHitTestingServiceCompactRetrieveResponse:
|
||||
mock_format.return_value = []
|
||||
|
||||
# Act
|
||||
result = HitTestingService.compact_retrieve_response(query, documents)
|
||||
result = HitTestingService.compact_retrieve_response(MagicMock(), query, documents)
|
||||
|
||||
# Assert
|
||||
assert result["query"]["content"] == query
|
||||
@ -708,7 +722,7 @@ class TestHitTestingServiceHitTestingArgsCheck:
|
||||
args = {"query": ""}
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(ValueError, match="Query is required and cannot exceed 250 characters"):
|
||||
with pytest.raises(ValueError, match="Query or attachment_ids is required"):
|
||||
HitTestingService.hit_testing_args_check(args)
|
||||
|
||||
def test_hit_testing_args_check_none_query(self):
|
||||
@ -721,7 +735,7 @@ class TestHitTestingServiceHitTestingArgsCheck:
|
||||
args = {"query": None}
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(ValueError, match="Query is required and cannot exceed 250 characters"):
|
||||
with pytest.raises(ValueError, match="Query or attachment_ids is required"):
|
||||
HitTestingService.hit_testing_args_check(args)
|
||||
|
||||
def test_hit_testing_args_check_too_long_query(self):
|
||||
@ -734,7 +748,7 @@ class TestHitTestingServiceHitTestingArgsCheck:
|
||||
args = {"query": "a" * 251}
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(ValueError, match="Query is required and cannot exceed 250 characters"):
|
||||
with pytest.raises(ValueError, match="Query cannot exceed 250 characters"):
|
||||
HitTestingService.hit_testing_args_check(args)
|
||||
|
||||
def test_hit_testing_args_check_exactly_250_characters(self):
|
||||
|
||||
@ -0,0 +1,96 @@
|
||||
import io
|
||||
import json
|
||||
import zipfile
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.archive_storage import ArchiveStorage
|
||||
from services.retention.workflow_run.constants import ARCHIVE_SCHEMA_VERSION
|
||||
from services.retention.workflow_run.delete_archived_workflow_run import ArchivedWorkflowRunDeletion
|
||||
|
||||
ARCHIVED_TABLES = [
|
||||
"workflow_runs",
|
||||
"workflow_app_logs",
|
||||
"workflow_node_executions",
|
||||
"workflow_node_execution_offload",
|
||||
"workflow_pauses",
|
||||
"workflow_pause_reasons",
|
||||
"workflow_trigger_logs",
|
||||
]
|
||||
|
||||
|
||||
def _build_archive_bundle(
|
||||
*,
|
||||
run_id: str = "run-1",
|
||||
tenant_id: str = "tenant-1",
|
||||
app_id: str = "app-1",
|
||||
workflow_id: str = "workflow-1",
|
||||
corrupt_checksum_for: str | None = None,
|
||||
) -> bytes:
|
||||
table_payloads: dict[str, bytes] = {}
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
records = [{"id": run_id}] if table_name == "workflow_runs" else []
|
||||
table_payloads[table_name] = ArchiveStorage.serialize_to_jsonl(records)
|
||||
|
||||
manifest = {
|
||||
"schema_version": ARCHIVE_SCHEMA_VERSION,
|
||||
"workflow_run_id": run_id,
|
||||
"tenant_id": tenant_id,
|
||||
"app_id": app_id,
|
||||
"workflow_id": workflow_id,
|
||||
"created_at": datetime.now(UTC).isoformat(),
|
||||
"archived_at": datetime.now(UTC).isoformat(),
|
||||
"tables": {
|
||||
table_name: {
|
||||
"row_count": 1 if table_name == "workflow_runs" else 0,
|
||||
"checksum": ArchiveStorage.compute_checksum(payload),
|
||||
"size_bytes": len(payload),
|
||||
}
|
||||
for table_name, payload in table_payloads.items()
|
||||
},
|
||||
}
|
||||
if corrupt_checksum_for:
|
||||
manifest["tables"][corrupt_checksum_for]["checksum"] = "bad-checksum"
|
||||
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
|
||||
archive.writestr("manifest.json", json.dumps(manifest).encode("utf-8"))
|
||||
for table_name, payload in table_payloads.items():
|
||||
archive.writestr(f"{table_name}.jsonl", payload)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def test_validate_archive_bundle_accepts_valid_archive() -> None:
|
||||
manifest = ArchivedWorkflowRunDeletion._validate_archive_bundle(
|
||||
_build_archive_bundle(),
|
||||
run_id="run-1",
|
||||
tenant_id="tenant-1",
|
||||
app_id="app-1",
|
||||
workflow_id="workflow-1",
|
||||
)
|
||||
|
||||
assert manifest["schema_version"] == ARCHIVE_SCHEMA_VERSION
|
||||
assert manifest["tables"]["workflow_runs"]["row_count"] == 1
|
||||
|
||||
|
||||
def test_validate_archive_bundle_rejects_checksum_mismatch() -> None:
|
||||
with pytest.raises(ValueError, match="archive member checksum mismatch"):
|
||||
ArchivedWorkflowRunDeletion._validate_archive_bundle(
|
||||
_build_archive_bundle(corrupt_checksum_for="workflow_runs"),
|
||||
run_id="run-1",
|
||||
tenant_id="tenant-1",
|
||||
app_id="app-1",
|
||||
workflow_id="workflow-1",
|
||||
)
|
||||
|
||||
|
||||
def test_validate_archive_bundle_rejects_manifest_target_mismatch() -> None:
|
||||
with pytest.raises(ValueError, match="manifest tenant_id does not match delete target"):
|
||||
ArchivedWorkflowRunDeletion._validate_archive_bundle(
|
||||
_build_archive_bundle(),
|
||||
run_id="run-1",
|
||||
tenant_id="different-tenant",
|
||||
app_id="app-1",
|
||||
workflow_id="workflow-1",
|
||||
)
|
||||
@ -15,7 +15,7 @@ from sqlalchemy import delete, select
|
||||
|
||||
from core.db.session_factory import session_factory
|
||||
from extensions.storage.storage_type import StorageType
|
||||
from models.agent import Agent, AgentDriveFile, AgentScope, AgentSource
|
||||
from models.agent import Agent, AgentDriveFile, AgentDriveFileKind, AgentScope, AgentSource
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import UploadFile
|
||||
from models.tools import ToolFile
|
||||
@ -133,6 +133,123 @@ def test_commit_then_manifest_lists_the_entry():
|
||||
assert AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, prefix="other/") == []
|
||||
|
||||
|
||||
def test_commit_skill_row_persists_metadata_and_lists_catalog() -> None:
|
||||
tf = _seed_tool_file(name="SKILL.md")
|
||||
AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[
|
||||
DriveCommitItem(
|
||||
key="tender-analyzer/SKILL.md",
|
||||
file_ref={"kind": "tool_file", "id": tf},
|
||||
is_skill=True,
|
||||
skill_metadata=DriveSkillMetadata(name="Tender Analyzer", description="Parses RFPs."),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
with session_factory.create_session() as session:
|
||||
row = session.scalar(select(AgentDriveFile).where(AgentDriveFile.key == "tender-analyzer/SKILL.md"))
|
||||
assert row is not None
|
||||
assert row.is_skill is True
|
||||
assert row.skill_metadata == '{"description":"Parses RFPs.","name":"Tender Analyzer"}'
|
||||
|
||||
skills = AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT)
|
||||
assert len(skills) == 1
|
||||
assert skills[0]["path"] == "tender-analyzer"
|
||||
assert skills[0]["skill_md_key"] == "tender-analyzer/SKILL.md"
|
||||
assert skills[0]["archive_key"] is None
|
||||
assert skills[0]["name"] == "Tender Analyzer"
|
||||
assert skills[0]["description"] == "Parses RFPs."
|
||||
assert skills[0]["size"] == 5
|
||||
assert skills[0]["mime_type"] == "text/plain"
|
||||
|
||||
|
||||
def test_commit_rejects_skill_row_without_skill_metadata() -> None:
|
||||
tf = _seed_tool_file(name="SKILL.md")
|
||||
|
||||
with pytest.raises(AgentDriveError) as exc_info:
|
||||
AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[
|
||||
DriveCommitItem(
|
||||
key="tender-analyzer/SKILL.md",
|
||||
file_ref={"kind": "tool_file", "id": tf},
|
||||
is_skill=True,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
assert exc_info.value.code == "invalid_skill_metadata"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw_metadata", [None, '{"description":"oops"}'])
|
||||
def test_list_skills_raises_controlled_error_for_invalid_stored_metadata(raw_metadata: str | None) -> None:
|
||||
tf = _seed_tool_file(name="SKILL.md")
|
||||
|
||||
with session_factory.create_session() as session:
|
||||
session.add(
|
||||
AgentDriveFile(
|
||||
id="44444444-4444-4444-4444-444444444444",
|
||||
tenant_id=TENANT,
|
||||
agent_id=AGENT,
|
||||
key="broken-skill/SKILL.md",
|
||||
file_kind=AgentDriveFileKind.TOOL_FILE,
|
||||
file_id=tf,
|
||||
value_owned_by_drive=True,
|
||||
is_skill=True,
|
||||
skill_metadata=raw_metadata,
|
||||
size=5,
|
||||
mime_type="text/plain",
|
||||
created_by=USER,
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
|
||||
with pytest.raises(AgentDriveError) as exc_info:
|
||||
AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT)
|
||||
|
||||
assert exc_info.value.code == "invalid_skill_metadata"
|
||||
|
||||
|
||||
def test_commit_rejects_non_skill_row_with_skill_metadata() -> None:
|
||||
tf = _seed_tool_file()
|
||||
with pytest.raises(AgentDriveError, match="skill metadata"):
|
||||
AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[
|
||||
DriveCommitItem(
|
||||
key="files/report.txt",
|
||||
file_ref={"kind": "tool_file", "id": tf},
|
||||
skill_metadata=DriveSkillMetadata(name="Bad", description=""),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_commit_rejects_non_canonical_skill_key() -> None:
|
||||
tf = _seed_tool_file(name="README.md")
|
||||
with pytest.raises(AgentDriveError, match="canonical"):
|
||||
AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[
|
||||
DriveCommitItem(
|
||||
key="tender-analyzer/README.md",
|
||||
file_ref={"kind": "tool_file", "id": tf},
|
||||
is_skill=True,
|
||||
skill_metadata=DriveSkillMetadata(name="Tender Analyzer", description=""),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_commit_rejects_tool_file_not_owned_by_user():
|
||||
other = _seed_tool_file(user_id="99999999-9999-9999-9999-999999999999")
|
||||
with pytest.raises(AgentDriveError) as exc_info:
|
||||
@ -247,6 +364,49 @@ def test_recommit_same_value_is_idempotent_and_keeps_value():
|
||||
assert len(rows) == 1
|
||||
|
||||
|
||||
def test_recommit_same_skill_value_updates_metadata_without_cleaning_backing_file() -> None:
|
||||
tf = _seed_tool_file(name="SKILL.md")
|
||||
AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[
|
||||
DriveCommitItem(
|
||||
key="tender-analyzer/SKILL.md",
|
||||
file_ref={"kind": "tool_file", "id": tf},
|
||||
value_owned_by_drive=True,
|
||||
is_skill=True,
|
||||
skill_metadata=DriveSkillMetadata(name="Tender Analyzer", description="v1"),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
with patch("services.agent_drive_service.storage") as storage_mock:
|
||||
AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[
|
||||
DriveCommitItem(
|
||||
key="tender-analyzer/SKILL.md",
|
||||
file_ref={"kind": "tool_file", "id": tf},
|
||||
value_owned_by_drive=False,
|
||||
is_skill=True,
|
||||
skill_metadata=DriveSkillMetadata(name="Tender Analyzer v2", description="v2"),
|
||||
)
|
||||
],
|
||||
)
|
||||
storage_mock.delete.assert_not_called()
|
||||
|
||||
with session_factory.create_session() as session:
|
||||
row = session.scalar(select(AgentDriveFile).where(AgentDriveFile.key == "tender-analyzer/SKILL.md"))
|
||||
assert row is not None
|
||||
assert row.file_id == tf
|
||||
assert row.value_owned_by_drive is False
|
||||
assert row.skill_metadata == '{"description":"v2","name":"Tender Analyzer v2"}'
|
||||
assert session.scalar(select(ToolFile).where(ToolFile.id == tf)) is not None
|
||||
|
||||
|
||||
def _seed_upload_file(*, name: str = "u.txt") -> str:
|
||||
upload = UploadFile(
|
||||
tenant_id=TENANT,
|
||||
@ -319,7 +479,7 @@ def test_manifest_includes_internal_download_url():
|
||||
|
||||
with (
|
||||
patch("services.agent_drive_service.file_factory.build_from_mapping", return_value=object()),
|
||||
patch("services.agent_drive_service.DifyWorkflowFileRuntime") as runtime_cls,
|
||||
patch("core.app.workflow.file_runtime.DifyWorkflowFileRuntime") as runtime_cls,
|
||||
):
|
||||
runtime_cls.return_value.resolve_file_url.return_value = "http://internal/files/x?sign=1"
|
||||
items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, include_download_url=True)
|
||||
@ -349,16 +509,31 @@ def test_delete_by_key_cleans_drive_owned_value():
|
||||
_commit("files/doomed.txt", tf, owned=True)
|
||||
|
||||
with patch("services.agent_drive_service.storage") as storage_mock:
|
||||
removed = AgentDriveService().delete(tenant_id=TENANT, agent_id=AGENT, key="files/doomed.txt")
|
||||
removed = AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[DriveCommitItem(key="files/doomed.txt", file_ref=None)],
|
||||
)
|
||||
storage_mock.delete.assert_called_once()
|
||||
|
||||
assert removed == ["files/doomed.txt"]
|
||||
assert removed == [
|
||||
{
|
||||
"key": "files/doomed.txt",
|
||||
"file_kind": "tool_file",
|
||||
"file_id": tf,
|
||||
"value_owned_by_drive": True,
|
||||
"is_skill": False,
|
||||
"skill_metadata": None,
|
||||
"removed": True,
|
||||
}
|
||||
]
|
||||
with session_factory.create_session() as session:
|
||||
assert session.scalar(select(ToolFile).where(ToolFile.id == tf)) is None
|
||||
assert list(session.scalars(select(AgentDriveFile))) == []
|
||||
|
||||
|
||||
def test_delete_by_prefix_removes_all_skill_keys():
|
||||
def test_commit_null_batch_removes_multiple_skill_keys():
|
||||
md = _seed_tool_file(name="SKILL.md")
|
||||
zf = _seed_tool_file(name="full.zip")
|
||||
_commit("tender-analyzer/SKILL.md", md, owned=True)
|
||||
@ -367,9 +542,20 @@ def test_delete_by_prefix_removes_all_skill_keys():
|
||||
_commit("files/other.txt", other, owned=True)
|
||||
|
||||
with patch("services.agent_drive_service.storage"):
|
||||
removed = AgentDriveService().delete(tenant_id=TENANT, agent_id=AGENT, prefix="tender-analyzer/")
|
||||
removed = AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[
|
||||
DriveCommitItem(key="tender-analyzer/SKILL.md", file_ref=None),
|
||||
DriveCommitItem(key="tender-analyzer/.DIFY-SKILL-FULL.zip", file_ref=None),
|
||||
],
|
||||
)
|
||||
|
||||
assert sorted(removed) == ["tender-analyzer/.DIFY-SKILL-FULL.zip", "tender-analyzer/SKILL.md"]
|
||||
assert sorted(item["key"] for item in removed) == [
|
||||
"tender-analyzer/.DIFY-SKILL-FULL.zip",
|
||||
"tender-analyzer/SKILL.md",
|
||||
]
|
||||
with session_factory.create_session() as session:
|
||||
# both skill ToolFiles physically removed, the unrelated file untouched
|
||||
assert session.scalar(select(ToolFile).where(ToolFile.id == md)) is None
|
||||
@ -379,28 +565,30 @@ def test_delete_by_prefix_removes_all_skill_keys():
|
||||
assert keys == ["files/other.txt"]
|
||||
|
||||
|
||||
def test_delete_is_idempotent():
|
||||
assert AgentDriveService().delete(tenant_id=TENANT, agent_id=AGENT, key="files/never-there.txt") == []
|
||||
assert AgentDriveService().delete(tenant_id=TENANT, agent_id=AGENT, prefix="ghost-skill/") == []
|
||||
def test_commit_null_is_idempotent_for_missing_keys():
|
||||
removed = AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[DriveCommitItem(key="files/never-there.txt", file_ref=None)],
|
||||
)
|
||||
assert removed == [{"key": "files/never-there.txt", "removed": True, "noop": True}]
|
||||
|
||||
|
||||
def test_delete_requires_exactly_one_scope():
|
||||
with pytest.raises(AgentDriveError) as exc_info:
|
||||
AgentDriveService().delete(tenant_id=TENANT, agent_id=AGENT)
|
||||
assert exc_info.value.code == "invalid_delete_scope"
|
||||
with pytest.raises(AgentDriveError):
|
||||
AgentDriveService().delete(tenant_id=TENANT, agent_id=AGENT, prefix="a/", key="a/b")
|
||||
|
||||
|
||||
def test_delete_keeps_shared_value_records():
|
||||
def test_commit_null_keeps_shared_value_records():
|
||||
tf = _seed_tool_file(name="shared.txt")
|
||||
_commit("files/shared.txt", tf, owned=False)
|
||||
|
||||
with patch("services.agent_drive_service.storage") as storage_mock:
|
||||
removed = AgentDriveService().delete(tenant_id=TENANT, agent_id=AGENT, key="files/shared.txt")
|
||||
removed = AgentDriveService().commit(
|
||||
tenant_id=TENANT,
|
||||
user_id=USER,
|
||||
agent_id=AGENT,
|
||||
items=[DriveCommitItem(key="files/shared.txt", file_ref=None)],
|
||||
)
|
||||
storage_mock.delete.assert_not_called()
|
||||
|
||||
assert removed == ["files/shared.txt"]
|
||||
assert removed[0]["key"] == "files/shared.txt"
|
||||
with session_factory.create_session() as session:
|
||||
# only the KV row dropped; the shared ToolFile survives
|
||||
assert session.scalar(select(ToolFile).where(ToolFile.id == tf)) is not None
|
||||
@ -502,7 +690,7 @@ def test_upload_file_download_url_uses_attachment_filename():
|
||||
upload_file_id = _seed_upload_file(name="report.pdf")
|
||||
_commit_upload("files/report.pdf", upload_file_id)
|
||||
|
||||
with patch("services.agent_drive_service.DifyWorkflowFileRuntime") as runtime_cls:
|
||||
with patch("core.app.workflow.file_runtime.DifyWorkflowFileRuntime") as runtime_cls:
|
||||
runtime_cls.return_value.resolve_upload_file_url.return_value = "https://files.example/report.pdf"
|
||||
url = AgentDriveService().download_url(tenant_id=TENANT, agent_id=AGENT, key="files/report.pdf")
|
||||
|
||||
|
||||
@ -9,8 +9,6 @@ This module contains tests for:
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_NAME
|
||||
|
||||
|
||||
class TestWorkflowRunArchiver:
|
||||
"""Tests for the WorkflowRunArchiver class."""
|
||||
@ -37,18 +35,20 @@ class TestWorkflowRunArchiver:
|
||||
assert archiver.limit == 50
|
||||
assert archiver.dry_run is True
|
||||
|
||||
def test_get_archive_key(self):
|
||||
"""Test archive key generation."""
|
||||
def test_get_bundle_manifest_key(self):
|
||||
"""Test V2 bundle manifest key generation."""
|
||||
from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver
|
||||
|
||||
archiver = WorkflowRunArchiver.__new__(WorkflowRunArchiver)
|
||||
archiver = WorkflowRunArchiver(run_shard_index=1, run_shard_total=4)
|
||||
|
||||
mock_run = MagicMock()
|
||||
mock_run.tenant_id = "tenant-123"
|
||||
mock_run.app_id = "app-999"
|
||||
mock_run.tenant_id = "9enant-123"
|
||||
mock_run.id = "run-456"
|
||||
mock_run.created_at = datetime(2024, 1, 15, 12, 0, 0)
|
||||
|
||||
key = archiver._get_archive_key(mock_run)
|
||||
identity = archiver._build_bundle_identity([mock_run])
|
||||
key = archiver._get_manifest_object_key(identity)
|
||||
|
||||
assert key == f"tenant-123/app_id=app-999/year=2024/month=01/workflow_run_id=run-456/{ARCHIVE_BUNDLE_NAME}"
|
||||
assert key.endswith("/manifest.json")
|
||||
assert "workflow-runs/v2/tenant_prefix=9/tenant_id=9enant-123/year=2024/month=01" in key
|
||||
assert "/shard=01-of-04/" in key
|
||||
|
||||
2
api/uv.lock
generated
2
api/uv.lock
generated
@ -1304,7 +1304,7 @@ requires-dist = [
|
||||
{ name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=1.85.1,<2.0.0" },
|
||||
{ name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.12.0,<3.0.0" },
|
||||
{ name = "redis", marker = "extra == 'server'", specifier = ">=7.4.0,<8.0.0" },
|
||||
{ name = "shell-session-manager", marker = "extra == 'server'", specifier = "==2.2.0" },
|
||||
{ name = "shell-session-manager", marker = "extra == 'server'", specifier = "==2.2.1" },
|
||||
{ name = "typer", specifier = ">=0.16.1,<0.17" },
|
||||
{ name = "typing-extensions", specifier = ">=4.12.2,<5.0.0" },
|
||||
{ name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = "==0.46.0" },
|
||||
|
||||
@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@langgenius/difyctl",
|
||||
"type": "module",
|
||||
"version": "0.1.0-rc.1",
|
||||
"version": "0.1.0-alpha",
|
||||
"description": "Dify command-line interface",
|
||||
"difyctl": {
|
||||
"channel": "rc",
|
||||
"channel": "alpha",
|
||||
"compat": {
|
||||
"minDify": "1.14.0",
|
||||
"minDify": "1.15.0",
|
||||
"maxDify": "1.15.0"
|
||||
},
|
||||
"release": {
|
||||
|
||||
@ -4,7 +4,7 @@ import { readFileSync } from 'node:fs'
|
||||
import { dirname, resolve } from 'node:path'
|
||||
import { fileURLToPath } from 'node:url'
|
||||
|
||||
export const BUILD_CHANNELS = ['dev', 'rc', 'stable'] as const
|
||||
export const BUILD_CHANNELS = ['dev', 'alpha', 'rc', 'edge', 'stable'] as const
|
||||
export type BuildChannel = (typeof BUILD_CHANNELS)[number]
|
||||
|
||||
export type BuildInfo = {
|
||||
|
||||
@ -26,6 +26,7 @@ const SEMVER_CORE_LEN = 3
|
||||
// Add channels here: { name, prerelease, versionForm }.
|
||||
const CHANNELS = [
|
||||
{ name: 'stable', prerelease: false, versionForm: /^\d+\.\d+\.\d+(\+[0-9A-Z.-]+)?$/i },
|
||||
{ name: 'alpha', prerelease: true, versionForm: /^\d+\.\d+\.\d+-alpha(\.\d+)?$/ },
|
||||
{ name: 'rc', prerelease: true, versionForm: /^\d+\.\d+\.\d+-rc\.\d+$/ },
|
||||
{ name: 'edge', prerelease: true, versionForm: /^\d+\.\d+\.\d+-edge\.[0-9a-f]{7,40}$/ },
|
||||
]
|
||||
|
||||
@ -15,13 +15,13 @@ function run(args: string[]): { code: number, stdout: string, stderr: string } {
|
||||
}
|
||||
}
|
||||
|
||||
describe('release-naming compat-check (compat 1.14.0..1.15.0)', () => {
|
||||
describe('release-naming compat-check (compat 1.15.0..1.15.0)', () => {
|
||||
it('accepts a version inside the window', () => {
|
||||
expect(run(['compat-check', '1.14.7']).code).toBe(0)
|
||||
expect(run(['compat-check', '1.15.0']).code).toBe(0)
|
||||
})
|
||||
|
||||
it('accepts the inclusive lower bound', () => {
|
||||
expect(run(['compat-check', '1.14.0']).code).toBe(0)
|
||||
expect(run(['compat-check', '1.15.0']).code).toBe(0)
|
||||
})
|
||||
|
||||
it('accepts the inclusive upper bound', () => {
|
||||
@ -29,26 +29,22 @@ describe('release-naming compat-check (compat 1.14.0..1.15.0)', () => {
|
||||
})
|
||||
|
||||
it('accepts a v-prefixed tag', () => {
|
||||
expect(run(['compat-check', 'v1.14.2']).code).toBe(0)
|
||||
expect(run(['compat-check', 'v1.15.0']).code).toBe(0)
|
||||
})
|
||||
|
||||
it('rejects a version below the lower bound', () => {
|
||||
expect(run(['compat-check', '1.13.9']).code).not.toBe(0)
|
||||
expect(run(['compat-check', '1.14.9']).code).not.toBe(0)
|
||||
})
|
||||
|
||||
it('rejects a version above the upper bound', () => {
|
||||
expect(run(['compat-check', '1.15.1']).code).not.toBe(0)
|
||||
})
|
||||
|
||||
it('treats a prerelease of the upper bound as in range (1.15.0-rc1 <= 1.15.0)', () => {
|
||||
expect(run(['compat-check', '1.15.0-rc1']).code).toBe(0)
|
||||
it('treats a prerelease of the bound as below it (1.15.0-rc1 < 1.15.0)', () => {
|
||||
expect(run(['compat-check', '1.15.0-rc1']).code).not.toBe(0)
|
||||
})
|
||||
|
||||
it('treats a prerelease of the lower bound as below it (1.14.0-rc1 < 1.14.0)', () => {
|
||||
expect(run(['compat-check', '1.14.0-rc1']).code).not.toBe(0)
|
||||
})
|
||||
|
||||
it('ignores build metadata on the upper bound (1.15.0+build == 1.15.0)', () => {
|
||||
it('ignores build metadata on the bound (1.15.0+build == 1.15.0)', () => {
|
||||
expect(run(['compat-check', '1.15.0+build123']).code).toBe(0)
|
||||
})
|
||||
|
||||
@ -64,7 +60,7 @@ describe('release-naming compat-check (compat 1.14.0..1.15.0)', () => {
|
||||
describe('release-naming github-env', () => {
|
||||
it('emits difyctlTag = tagPrefix + version', () => {
|
||||
const { stdout } = run(['github-env'])
|
||||
expect(stdout).toMatch(/^difyctlTag=difyctl-v0\.1\.0-rc\.1$/m)
|
||||
expect(stdout).toMatch(/^difyctlTag=difyctl-v0\.1\.0-alpha$/m)
|
||||
})
|
||||
|
||||
it('still emits the existing trace fields', () => {
|
||||
@ -79,8 +75,8 @@ describe('release-naming edge channel', () => {
|
||||
expect(run(['channels']).stdout).toMatch(/^edge$/m)
|
||||
})
|
||||
|
||||
it('edge-version derives <pkgcore>-edge.<sha> stripping the rc prerelease', () => {
|
||||
// package.json version is 0.1.0-rc.1 -> core 0.1.0
|
||||
it('edge-version derives <pkgcore>-edge.<sha> stripping the alpha prerelease', () => {
|
||||
// package.json version is 0.1.0-alpha -> core 0.1.0
|
||||
expect(run(['edge-version', '2fd7b82']).stdout.trim()).toBe('0.1.0-edge.2fd7b82')
|
||||
})
|
||||
|
||||
|
||||
@ -81,7 +81,7 @@ describe('release-r2-edge manifest', () => {
|
||||
|
||||
it('carries the compat window from package.json', () => {
|
||||
const { json } = buildManifest()
|
||||
expect(json.compat).toEqual({ minDify: '1.14.0', maxDify: '1.15.0' })
|
||||
expect(json.compat).toEqual({ minDify: '1.15.0', maxDify: '1.15.0' })
|
||||
})
|
||||
|
||||
it('lists all 5 targets with asset name + sha256 from the checksums file', () => {
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
import type { AppDescribeResponse, AppListResponse, AppMode } from '@dify/contracts/api/openapi/types.gen'
|
||||
import type { AppDescribeResponse, AppListResponse, SupportedAppType } from '@dify/contracts/api/openapi/types.gen'
|
||||
import type { AppReader } from './app-reader'
|
||||
import type { OpenApiClient } from '@/http/orpc'
|
||||
import type { HttpClient } from '@/http/types'
|
||||
@ -8,12 +8,12 @@ export type ListQuery = {
|
||||
readonly workspaceId: string
|
||||
readonly page?: number
|
||||
readonly limit?: number
|
||||
readonly mode?: AppMode | ''
|
||||
readonly mode?: SupportedAppType | ''
|
||||
readonly name?: string
|
||||
}
|
||||
|
||||
// An absent or empty mode filter means "any mode" — collapse both to undefined for the query.
|
||||
export function normalizeMode(mode: AppMode | '' | undefined): AppMode | undefined {
|
||||
export function normalizeMode(mode: SupportedAppType | '' | undefined): SupportedAppType | undefined {
|
||||
return mode !== undefined && mode !== '' ? mode : undefined
|
||||
}
|
||||
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import type { AppMode } from '@dify/contracts/api/openapi/types.gen'
|
||||
import type { SupportedAppType } from '@dify/contracts/api/openapi/types.gen'
|
||||
import { zSupportedAppType } from '@dify/contracts/api/openapi/zod.gen'
|
||||
import { DifyCommand } from '@/commands/_shared/dify-command'
|
||||
import { httpRetryFlag } from '@/commands/_shared/global-flags'
|
||||
import { Args, Flags } from '@/framework/flags'
|
||||
@ -6,16 +7,9 @@ import { OutputFormat, table } from '@/framework/output'
|
||||
import { agentGuide } from './guide'
|
||||
import { runGetApp } from './run'
|
||||
|
||||
const APP_MODE_VALUES: readonly AppMode[] = [
|
||||
'advanced-chat',
|
||||
'agent',
|
||||
'agent-chat',
|
||||
'channel',
|
||||
'chat',
|
||||
'completion',
|
||||
'rag-pipeline',
|
||||
'workflow',
|
||||
]
|
||||
// Single source: derived from the backend's listable app types (openapi codegen).
|
||||
// Adding/removing a listable type is a backend-only change that flows here on regen.
|
||||
const APP_MODE_VALUES: readonly SupportedAppType[] = zSupportedAppType.options
|
||||
|
||||
export default class GetApp extends DifyCommand {
|
||||
static override description = 'List apps or describe one app\'s basic info'
|
||||
@ -56,7 +50,7 @@ export default class GetApp extends DifyCommand {
|
||||
allWorkspaces: flags['all-workspaces'],
|
||||
page: flags.page,
|
||||
limitRaw: flags.limit,
|
||||
mode: flags.mode as AppMode | undefined,
|
||||
mode: flags.mode as SupportedAppType | undefined,
|
||||
name: flags.name,
|
||||
format,
|
||||
}, { active: ctx.active, http: ctx.http, io: ctx.io })
|
||||
|
||||
13
cli/src/commands/get/app/mode-whitelist.test.ts
Normal file
13
cli/src/commands/get/app/mode-whitelist.test.ts
Normal file
@ -0,0 +1,13 @@
|
||||
import { zSupportedAppType } from '@dify/contracts/api/openapi/zod.gen'
|
||||
import { describe, expect, it } from 'vitest'
|
||||
|
||||
// The `get app --mode` whitelist is derived from this generated enum (see index.ts).
|
||||
// These pins guard the original bug: the CLI must not advertise modes the backend
|
||||
// rejects (rag-pipeline, channel) or modes that aren't listable here (agent).
|
||||
describe('get app --mode whitelist', () => {
|
||||
it('is exactly the listable app types', () => {
|
||||
expect([...zSupportedAppType.options].sort()).toEqual(
|
||||
['advanced-chat', 'agent-chat', 'chat', 'completion', 'workflow'],
|
||||
)
|
||||
})
|
||||
})
|
||||
@ -1,4 +1,4 @@
|
||||
import type { AppDescribeResponse, AppListResponse, AppMode } from '@dify/contracts/api/openapi/types.gen'
|
||||
import type { AppDescribeResponse, AppListResponse, AppMode, SupportedAppType } from '@dify/contracts/api/openapi/types.gen'
|
||||
import type { AppReader } from '@/api/app-reader'
|
||||
import type { ActiveContext } from '@/auth/hosts'
|
||||
import type { HttpClient } from '@/http/types'
|
||||
@ -20,7 +20,7 @@ export type GetAppOptions = {
|
||||
readonly allWorkspaces?: boolean
|
||||
readonly page?: number
|
||||
readonly limitRaw?: string
|
||||
readonly mode?: AppMode
|
||||
readonly mode?: SupportedAppType
|
||||
readonly name?: string
|
||||
readonly format?: string
|
||||
}
|
||||
|
||||
@ -158,6 +158,6 @@ describe('Version command', () => {
|
||||
if (output?.kind !== 'formatted')
|
||||
throw new Error('expected formatted output')
|
||||
|
||||
expect(output.data.text()).toContain('WARNING: This build is a release candidate')
|
||||
expect(output.data.text()).toContain('WARNING: This build is a rc release')
|
||||
})
|
||||
})
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
import { arch, platform } from '@/sys/index'
|
||||
import { compatString } from './compat'
|
||||
|
||||
export type Channel = 'dev' | 'edge' | 'rc' | 'stable'
|
||||
export type Channel = 'dev' | 'alpha' | 'edge' | 'rc' | 'stable'
|
||||
|
||||
export type VersionInfo = {
|
||||
version: string
|
||||
|
||||
@ -52,7 +52,7 @@ describe('renderVersionText', () => {
|
||||
expect(text).not.toContain('WARNING:')
|
||||
})
|
||||
|
||||
it('appends RC warning when channel is rc', () => {
|
||||
it('appends warning when channel is rc', () => {
|
||||
const report: VersionReport = {
|
||||
client: baseClient({ channel: 'rc' }),
|
||||
server: { endpoint: '', reachable: false },
|
||||
@ -60,8 +60,43 @@ describe('renderVersionText', () => {
|
||||
}
|
||||
const text = renderVersionText(report)
|
||||
|
||||
expect(text).toContain('WARNING: This build is a release candidate')
|
||||
expect(text).toContain('install the stable channel')
|
||||
expect(text).toContain('WARNING: This build is a rc release')
|
||||
expect(text).toContain('install or wait for the stable channel')
|
||||
})
|
||||
|
||||
it('appends warning when channel is alpha', () => {
|
||||
const report: VersionReport = {
|
||||
client: baseClient({ channel: 'alpha' }),
|
||||
server: { endpoint: '', reachable: false },
|
||||
compat: { ...compatible(), status: 'unknown', detail: 'server probe skipped' },
|
||||
}
|
||||
const text = renderVersionText(report)
|
||||
|
||||
expect(text).toContain('WARNING: This build is a alpha release')
|
||||
expect(text).toContain('install or wait for the stable channel')
|
||||
})
|
||||
|
||||
it('appends warning when channel is edge', () => {
|
||||
const report: VersionReport = {
|
||||
client: baseClient({ channel: 'edge' }),
|
||||
server: { endpoint: '', reachable: false },
|
||||
compat: { ...compatible(), status: 'unknown', detail: 'server probe skipped' },
|
||||
}
|
||||
const text = renderVersionText(report)
|
||||
|
||||
expect(text).toContain('WARNING: This build is a edge release')
|
||||
expect(text).toContain('install or wait for the stable channel')
|
||||
})
|
||||
|
||||
it('does not append warning when channel is stable', () => {
|
||||
const report: VersionReport = {
|
||||
client: baseClient({ channel: 'stable' }),
|
||||
server: { endpoint: '', reachable: false },
|
||||
compat: { ...compatible(), status: 'unknown', detail: 'server probe skipped' },
|
||||
}
|
||||
const text = renderVersionText(report)
|
||||
|
||||
expect(text).not.toContain('WARNING:')
|
||||
})
|
||||
|
||||
it('shows "(skipped …)" when server.endpoint is empty', () => {
|
||||
@ -105,7 +140,7 @@ describe('renderVersionText', () => {
|
||||
// RC warning) ran, yet the output is byte-clean.
|
||||
expect(plain).not.toMatch(ANSI_RE)
|
||||
expect(plain).toContain('Compatibility: incompatible')
|
||||
expect(plain).toContain('WARNING: This build is a release candidate')
|
||||
expect(plain).toContain('WARNING: This build is a rc release')
|
||||
})
|
||||
|
||||
describe('with picocolors stubbed to always emit ANSI', () => {
|
||||
@ -147,8 +182,8 @@ describe('renderVersionText', () => {
|
||||
const colored = render(report, { color: true })
|
||||
expect(colored).toMatch(ANSI_RE)
|
||||
expect(colored).toContain('Compatibility: incompatible')
|
||||
// RC warning lines also routed through yellow.
|
||||
expect(colored).toContain('release candidate')
|
||||
// prerelease warning lines also routed through yellow.
|
||||
expect(colored).toContain('WARNING: This build is a rc release')
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@ -1,10 +1,13 @@
|
||||
import type { Channel } from './info'
|
||||
import type { VersionReport } from './probe'
|
||||
import { colorScheme } from '@/sys/io/color'
|
||||
|
||||
const RC_WARNING_LINES = [
|
||||
'WARNING: This build is a release candidate. It is in beta test, not stable,',
|
||||
' and may have bugs. For production use, install the stable channel.',
|
||||
] as const
|
||||
function prereleaseWarning(channel: Channel): readonly string[] {
|
||||
return [
|
||||
`WARNING: This build is a ${channel} release. It is not stable`,
|
||||
' and may have bugs. For production use, install or wait for the stable channel.',
|
||||
]
|
||||
}
|
||||
|
||||
export type RenderOptions = {
|
||||
readonly color?: boolean
|
||||
@ -49,9 +52,9 @@ export function renderVersionText(report: VersionReport, opts: RenderOptions = {
|
||||
const verdictText = `Compatibility: ${COMPAT_LABEL[compat.status]} — ${compat.detail}`
|
||||
lines.push(compat.status === 'unsupported' ? c.yellow(verdictText) : verdictText)
|
||||
|
||||
if (client.channel === 'rc') {
|
||||
if (client.channel !== 'stable') {
|
||||
lines.push('')
|
||||
for (const line of RC_WARNING_LINES)
|
||||
for (const line of prereleaseWarning(client.channel))
|
||||
lines.push(c.yellow(line))
|
||||
}
|
||||
|
||||
|
||||
@ -177,6 +177,19 @@ describe('E2E / difyctl get app (list)', () => {
|
||||
expect(result.exitCode, '--mode chatbot should cause non-zero exit').not.toBe(0)
|
||||
})
|
||||
|
||||
// Regression: rag-pipeline (a knowledge Pipeline), channel (unused) and agent
|
||||
// (roster-owned) are AppMode members but not listable app types. The old CLI
|
||||
// whitelist advertised rag-pipeline/channel, so the CLI forwarded them and the
|
||||
// server replied 400. The whitelist now derives from SupportedAppType, so the
|
||||
// CLI rejects them before any HTTP call.
|
||||
it.each(['rag-pipeline', 'channel', 'agent'])(
|
||||
'[P0] non-listable mode %s is intercepted client-side',
|
||||
async (mode) => {
|
||||
const result = await fx.r(['get', 'app', '--mode', mode])
|
||||
expect(result.exitCode, `--mode ${mode} should be rejected client-side`).not.toBe(0)
|
||||
},
|
||||
)
|
||||
|
||||
// ── workspace override ────────────────────────────────────────────────────
|
||||
|
||||
it('[P0] -w overrides the default workspace', async () => {
|
||||
|
||||
59
dify-agent/docker/local-sandbox/Dockerfile
Normal file
59
dify-agent/docker/local-sandbox/Dockerfile
Normal file
@ -0,0 +1,59 @@
|
||||
# Local sandbox image for shellctl-managed Dify Agent workspaces.
|
||||
#
|
||||
# Build this from the dify-agent package root:
|
||||
# docker build -f docker/local-sandbox/Dockerfile -t dify-agent-local-sandbox:local .
|
||||
#
|
||||
# This image merges the former shellctl-only image with the sandbox-visible
|
||||
# Agent Stub client CLI. It runs shellctl by default, and shellctl-managed jobs
|
||||
# can call `dify-agent ...` without installing extra packages at runtime.
|
||||
|
||||
FROM python:3.12-slim-bookworm AS base
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1 \
|
||||
DIFY_AGENT_STUB_DRIVE_BASE=/mnt/drive
|
||||
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl \
|
||||
tmux \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
ENV UV_VERSION=0.8.9
|
||||
RUN python -m pip install --no-cache-dir "uv==${UV_VERSION}"
|
||||
|
||||
WORKDIR /opt/dify-agent
|
||||
|
||||
|
||||
FROM base AS packages
|
||||
|
||||
ENV SHELL_SESSION_MANAGER_VERSION=2.2.1
|
||||
|
||||
COPY pyproject.toml uv.lock README.md ./
|
||||
COPY src ./src
|
||||
|
||||
RUN uv sync --frozen --no-dev --no-editable --extra grpc \
|
||||
&& uv pip install --python .venv/bin/python "shell-session-manager==${SHELL_SESSION_MANAGER_VERSION}"
|
||||
|
||||
|
||||
FROM base AS production
|
||||
|
||||
ENV VIRTUAL_ENV=/opt/dify-agent/.venv
|
||||
ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
|
||||
|
||||
COPY --from=packages ${VIRTUAL_ENV} ${VIRTUAL_ENV}
|
||||
|
||||
RUN ln -s ${VIRTUAL_ENV}/bin/dify-agent /usr/local/bin/dify-agent \
|
||||
&& ln -s ${VIRTUAL_ENV}/bin/shellctl /usr/local/bin/shellctl \
|
||||
&& useradd --create-home --shell /bin/sh dify \
|
||||
&& mkdir -p /mnt/drive \
|
||||
&& chown -R dify:dify /home/dify /mnt/drive
|
||||
|
||||
USER dify
|
||||
WORKDIR /home/dify
|
||||
|
||||
EXPOSE 5004
|
||||
|
||||
CMD ["shellctl", "serve", "--listen", "0.0.0.0:5004"]
|
||||
@ -1,25 +0,0 @@
|
||||
FROM python:3.13-slim
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl \
|
||||
tmux \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN python -m pip install --no-cache-dir \
|
||||
shell-session-manager==2.2.0 \
|
||||
uv
|
||||
|
||||
RUN useradd --create-home --shell /bin/sh dify
|
||||
|
||||
USER dify
|
||||
WORKDIR /home/dify
|
||||
|
||||
EXPOSE 5004
|
||||
|
||||
CMD ["shellctl", "serve", "--listen", "0.0.0.0:5004"]
|
||||
@ -51,6 +51,9 @@ DIFY_AGENT_REDIS_PREFIX=dify-agent
|
||||
|
||||
DIFY_AGENT_PLUGIN_DAEMON_URL=http://localhost:5002
|
||||
DIFY_AGENT_PLUGIN_DAEMON_API_KEY=replace-with-plugin-daemon-server-key
|
||||
|
||||
DIFY_AGENT_INNER_API_URL=http://localhost:5001
|
||||
DIFY_AGENT_INNER_API_KEY=replace-with-dify-inner-api-key-for-plugin
|
||||
EOF
|
||||
```
|
||||
|
||||
@ -63,12 +66,16 @@ The minimum settings are:
|
||||
- `DIFY_AGENT_PLUGIN_DAEMON_API_KEY`: API key sent by the server to the plugin
|
||||
daemon. In a Dify Docker setup this is usually the value previously configured
|
||||
as `PLUGIN_DAEMON_KEY`.
|
||||
- `DIFY_AGENT_INNER_API_URL`: Dify API service root for `/inner/api/...` calls.
|
||||
- `DIFY_AGENT_INNER_API_KEY`: API key sent to Dify API inner plugin endpoints.
|
||||
In Docker this should match `PLUGIN_DIFY_INNER_API_KEY`, which maps to Dify
|
||||
API `INNER_API_KEY_FOR_PLUGIN`.
|
||||
|
||||
See `.example.env` for the full server settings template.
|
||||
|
||||
If you plan to run `dify.shell`, also configure `DIFY_AGENT_SHELLCTL_ENTRYPOINT`
|
||||
and, when shell jobs need to call back with the `dify-agent` command, set
|
||||
`DIFY_AGENT_STUB_URL` plus a 32-byte base64url
|
||||
`DIFY_AGENT_STUB_API_BASE_URL` plus a 32-byte base64url
|
||||
`DIFY_AGENT_SERVER_SECRET_KEY` as documented in `.example.env`.
|
||||
|
||||
## Start the Dify Agent server
|
||||
|
||||
@ -36,11 +36,13 @@ also reads `.env` and `dify-agent/.env` when present.
|
||||
| `DIFY_AGENT_RUN_RETENTION_SECONDS` | `259200` | Seconds to retain Redis run records and per-run event streams; defaults to 3 days. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_URL` | `http://localhost:5002` | Base URL for the Dify plugin daemon. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_API_KEY` | empty | API key sent to the Dify plugin daemon. |
|
||||
| `DIFY_AGENT_INNER_API_URL` | `http://localhost:5001` | Dify API service root used when dify-agent calls `/inner/api/...` endpoints. |
|
||||
| `DIFY_AGENT_INNER_API_KEY` | empty | API key sent to Dify API inner plugin endpoints. Set this to Dify API `INNER_API_KEY_FOR_PLUGIN` (Docker: `PLUGIN_DIFY_INNER_API_KEY`). |
|
||||
| `DIFY_AGENT_SHELLCTL_ENTRYPOINT` | empty | Base URL for the shellctl server used by `dify.shell`; required when runs include the shell layer. |
|
||||
| `DIFY_AGENT_SHELLCTL_AUTH_TOKEN` | empty | Optional bearer token sent to the shellctl server. |
|
||||
| `DIFY_AGENT_STUB_URL` | empty | Public Agent Stub URL reachable from shellctl-managed remote machines. Use `http(s)://.../agent-stub` for HTTP or `grpc://host:port` for gRPC; enables `DIFY_AGENT_STUB_*` env injection for user `shell.run` jobs. |
|
||||
| `DIFY_AGENT_STUB_GRPC_BIND_ADDRESS` | empty | Optional `host:port` bind override used only when `DIFY_AGENT_STUB_URL` uses `grpc://`. |
|
||||
| `DIFY_AGENT_SERVER_SECRET_KEY` | empty | Server-wide root secret used to derive Agent Stub JWE keys; required when `DIFY_AGENT_STUB_URL` is set and must be unpadded base64url for 32 bytes. |
|
||||
| `DIFY_AGENT_STUB_API_BASE_URL` | empty | Public Agent Stub API base URL reachable from shellctl-managed remote machines. HTTP may be the service root or `/agent-stub`; gRPC must be `grpc://host:port`. Enables `DIFY_AGENT_STUB_*` env injection for user `shell.run` jobs. |
|
||||
| `DIFY_AGENT_STUB_GRPC_BIND_ADDRESS` | empty | Optional `host:port` bind override used only when `DIFY_AGENT_STUB_API_BASE_URL` uses `grpc://`. |
|
||||
| `DIFY_AGENT_SERVER_SECRET_KEY` | empty | Server-wide root secret used to derive Agent Stub JWE keys; required when `DIFY_AGENT_STUB_API_BASE_URL` is set and must be unpadded base64url for 32 bytes. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_CONNECT_TIMEOUT` | `10` | Plugin-daemon HTTP connect timeout in seconds. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_READ_TIMEOUT` | `600` | Plugin-daemon HTTP read timeout in seconds. |
|
||||
| `DIFY_AGENT_PLUGIN_DAEMON_WRITE_TIMEOUT` | `30` | Plugin-daemon HTTP write timeout in seconds. |
|
||||
@ -58,10 +60,12 @@ DIFY_AGENT_SHUTDOWN_GRACE_SECONDS=30
|
||||
DIFY_AGENT_RUN_RETENTION_SECONDS=259200
|
||||
DIFY_AGENT_PLUGIN_DAEMON_URL=http://localhost:5002
|
||||
DIFY_AGENT_PLUGIN_DAEMON_API_KEY=replace-with-daemon-key
|
||||
DIFY_AGENT_INNER_API_URL=http://localhost:5001
|
||||
DIFY_AGENT_INNER_API_KEY=replace-with-dify-inner-api-key-for-plugin
|
||||
DIFY_AGENT_SHELLCTL_ENTRYPOINT=http://127.0.0.1:5004
|
||||
DIFY_AGENT_SHELLCTL_AUTH_TOKEN=replace-with-shellctl-token
|
||||
# Generate with: python -c 'import base64, secrets; print(base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode())'
|
||||
DIFY_AGENT_STUB_URL=https://agent.example.com/agent-stub
|
||||
DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=replace-with-base64url-32-byte-secret
|
||||
```
|
||||
|
||||
|
||||
@ -55,10 +55,14 @@ To let commands inside user-visible shell jobs call back to the Dify Agent serve
|
||||
with `dify-agent ...`, also enable the Agent Stub:
|
||||
|
||||
```env
|
||||
DIFY_AGENT_STUB_URL=https://agent.example.com/agent-stub
|
||||
DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub
|
||||
DIFY_AGENT_SERVER_SECRET_KEY=replace-with-base64url-32-byte-secret
|
||||
```
|
||||
|
||||
HTTP `DIFY_AGENT_STUB_API_BASE_URL` may be either the service root or the
|
||||
explicit `/agent-stub` API root; the server normalizes the service root to
|
||||
`/agent-stub`. Other HTTP paths are rejected at startup.
|
||||
|
||||
`DIFY_AGENT_SERVER_SECRET_KEY` must be unpadded base64url text for exactly 32
|
||||
decoded bytes. One way to generate it is:
|
||||
|
||||
@ -69,10 +73,14 @@ python -c 'import base64, secrets; print(base64.urlsafe_b64encode(secrets.token_
|
||||
## Client request shape
|
||||
|
||||
A client adds the shell layer as an ordinary composition layer. Basic shell jobs
|
||||
do not need dependencies. To inject `DIFY_AGENT_STUB_URL` and
|
||||
`DIFY_AGENT_STUB_AUTH_JWE` into user-visible `shell.run` jobs, declare the
|
||||
execution-context layer as the shell layer's `execution_context` dependency. A
|
||||
typical run still also includes:
|
||||
do not need dependencies. To inject `DIFY_AGENT_STUB_API_BASE_URL`,
|
||||
`DIFY_AGENT_STUB_AUTH_JWE`, and `DIFY_AGENT_STUB_DRIVE_BASE` into user-visible
|
||||
`shell.run` jobs, declare the execution-context layer as the shell layer's
|
||||
`execution_context` dependency. When the run also includes `dify.drive`, declare
|
||||
it as the shell layer's `drive` dependency; the injected drive base is then
|
||||
computed from the fixed Agent Stub drive mount and the drive reference, for
|
||||
example `/mnt/drive/agent-123`. Without a drive dependency, the CLI keeps the
|
||||
historical `/mnt/drive` fallback. A typical run still also includes:
|
||||
|
||||
- a prompt layer that supplies the task;
|
||||
- an execution-context layer carrying tenant/user context;
|
||||
@ -194,33 +202,34 @@ Here is the analysis of the sales dataset:
|
||||
* **SHA-256 Hash:** `e86521a0d759037a09b059cb3cb2419f0a3f06e674db8151ccf2f93811dac0b8`
|
||||
````
|
||||
|
||||
## Running shellctl in Docker
|
||||
## Running the local sandbox in Docker
|
||||
|
||||
Build the shellctl image from the Dify Agent package root:
|
||||
Build the local sandbox image from the Dify Agent package root:
|
||||
|
||||
```bash
|
||||
docker build -f docker/shellctl/Dockerfile -t dify-agent-shellctl:local .
|
||||
docker build -f docker/local-sandbox/Dockerfile -t dify-agent-local-sandbox:local .
|
||||
```
|
||||
|
||||
Run it with a bearer token and publish the API on localhost:
|
||||
|
||||
```bash
|
||||
docker run --rm --name dify-agent-shellctl \
|
||||
docker run --rm --name dify-agent-local-sandbox \
|
||||
-e SHELLCTL_AUTH_TOKEN=replace-with-a-token \
|
||||
-p 127.0.0.1:5004:5004 \
|
||||
dify-agent-shellctl:local
|
||||
dify-agent-local-sandbox:local
|
||||
```
|
||||
|
||||
The image starts `shellctl serve --listen 0.0.0.0:5004` as the non-root
|
||||
`dify` user and leaves shellctl state/runtime directories at their package
|
||||
defaults.
|
||||
`dify` user. It also sets the fallback `DIFY_AGENT_STUB_DRIVE_BASE=/mnt/drive`
|
||||
and pre-creates that directory with write access for the same user.
|
||||
|
||||
## Docker image contents
|
||||
|
||||
The provided `docker/shellctl/Dockerfile` installs:
|
||||
The provided `docker/local-sandbox/Dockerfile` installs:
|
||||
|
||||
- `tmux`, required by `shellctl` to manage shell jobs;
|
||||
- `shell-session-manager==2.2.0`, which provides the `shellctl` CLI/server;
|
||||
- `shell-session-manager==2.2.1`, which provides the `shellctl` CLI/server;
|
||||
- `uv`, so uv shebang scripts with PEP 723 metadata can run inside the shell
|
||||
workspace;
|
||||
- the `dify-agent` Agent Stub client CLI, including its gRPC transport extra;
|
||||
- a non-root default user named `dify`.
|
||||
|
||||
@ -26,7 +26,7 @@ server = [
|
||||
"pydantic-ai-slim[anthropic,google,openai]>=1.85.1,<2.0.0",
|
||||
"pydantic-settings>=2.12.0,<3.0.0",
|
||||
"redis>=7.4.0,<8.0.0",
|
||||
"shell-session-manager==2.2.0",
|
||||
"shell-session-manager==2.2.1",
|
||||
"uvicorn[standard]==0.46.0",
|
||||
]
|
||||
|
||||
|
||||
@ -11,6 +11,7 @@ ToolFile ids back into the drive.
|
||||
from __future__ import annotations
|
||||
|
||||
import stat
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path, PurePosixPath
|
||||
from tempfile import TemporaryDirectory
|
||||
@ -32,6 +33,7 @@ from dify_agent.agent_stub.protocol.agent_stub import (
|
||||
AgentStubDriveFileRef,
|
||||
AgentStubDriveItem,
|
||||
AgentStubDriveManifestResponse,
|
||||
DEFAULT_AGENT_STUB_DRIVE_BASE,
|
||||
)
|
||||
|
||||
_SKILL_MD_FILENAME = "SKILL.md"
|
||||
@ -81,11 +83,15 @@ def list_drive_from_environment(prefix: str, json_output: bool) -> str | AgentSt
|
||||
return _format_manifest(response)
|
||||
|
||||
|
||||
def pull_drive_from_environment(prefix: str, drive_base: str = "/mnt/drive") -> list[Path]:
|
||||
def pull_drive_from_environment(
|
||||
targets: list[str] | None = None,
|
||||
drive_base: str = DEFAULT_AGENT_STUB_DRIVE_BASE,
|
||||
) -> list[Path]:
|
||||
"""Pull drive files into one local drive base via signed download URLs.
|
||||
|
||||
Args:
|
||||
prefix: Optional drive-key prefix forwarded to the manifest request.
|
||||
targets: Optional drive-key targets or prefixes. An empty list preserves
|
||||
the historical whole-drive pull by using ``[""]``.
|
||||
drive_base: Local base directory that receives downloaded drive files.
|
||||
|
||||
Returns:
|
||||
@ -117,16 +123,24 @@ def pull_drive_from_environment(prefix: str, drive_base: str = "/mnt/drive") ->
|
||||
"""
|
||||
|
||||
environment = read_agent_stub_environment()
|
||||
response = request_agent_stub_drive_manifest_sync(
|
||||
url=environment.url,
|
||||
auth_jwe=environment.auth_jwe,
|
||||
prefix=prefix,
|
||||
include_download_url=True,
|
||||
)
|
||||
manifest_targets = targets or [""]
|
||||
with ThreadPoolExecutor(max_workers=min(len(manifest_targets), 4)) as executor:
|
||||
responses = list(
|
||||
executor.map(
|
||||
lambda target: request_agent_stub_drive_manifest_sync(
|
||||
url=environment.url,
|
||||
auth_jwe=environment.auth_jwe,
|
||||
prefix=target,
|
||||
include_download_url=True,
|
||||
),
|
||||
manifest_targets,
|
||||
)
|
||||
)
|
||||
base_path = Path(drive_base).expanduser().resolve()
|
||||
base_path.mkdir(parents=True, exist_ok=True)
|
||||
written_paths: list[Path] = []
|
||||
for item in response.items:
|
||||
deduplicated_items = {item.key: item for response in responses for item in response.items}
|
||||
for item in [deduplicated_items[key] for key in sorted(deduplicated_items)]:
|
||||
download_url = item.download_url
|
||||
if not isinstance(download_url, str) or not download_url:
|
||||
raise AgentStubValidationError(f"drive manifest item is missing download_url: {item.key}")
|
||||
|
||||
@ -8,8 +8,10 @@ import os
|
||||
|
||||
from dify_agent.agent_stub.protocol.agent_stub import (
|
||||
AGENT_STUB_AUTH_JWE_ENV_VAR,
|
||||
AGENT_STUB_URL_ENV_VAR,
|
||||
normalize_agent_stub_url,
|
||||
AGENT_STUB_DRIVE_BASE_ENV_VAR,
|
||||
AGENT_STUB_API_BASE_URL_ENV_VAR,
|
||||
DEFAULT_AGENT_STUB_DRIVE_BASE,
|
||||
normalize_agent_stub_api_base_url,
|
||||
)
|
||||
|
||||
|
||||
@ -28,32 +30,44 @@ class AgentStubEnvironment:
|
||||
def has_agent_stub_environment(env: Mapping[str, str] | None = None) -> bool:
|
||||
"""Return whether both required Agent Stub environment variables exist."""
|
||||
values = env or os.environ
|
||||
return bool(values.get(AGENT_STUB_URL_ENV_VAR) and values.get(AGENT_STUB_AUTH_JWE_ENV_VAR))
|
||||
return bool(values.get(AGENT_STUB_API_BASE_URL_ENV_VAR) and values.get(AGENT_STUB_AUTH_JWE_ENV_VAR))
|
||||
|
||||
|
||||
def read_agent_stub_environment(env: Mapping[str, str] | None = None) -> AgentStubEnvironment:
|
||||
"""Read and validate the Agent Stub environment variables."""
|
||||
values = env or os.environ
|
||||
url = (values.get(AGENT_STUB_URL_ENV_VAR) or "").strip()
|
||||
url = (values.get(AGENT_STUB_API_BASE_URL_ENV_VAR) or "").strip()
|
||||
auth_jwe = (values.get(AGENT_STUB_AUTH_JWE_ENV_VAR) or "").strip()
|
||||
missing: list[str] = []
|
||||
if not url:
|
||||
missing.append(AGENT_STUB_URL_ENV_VAR)
|
||||
missing.append(AGENT_STUB_API_BASE_URL_ENV_VAR)
|
||||
if not auth_jwe:
|
||||
missing.append(AGENT_STUB_AUTH_JWE_ENV_VAR)
|
||||
if missing:
|
||||
names = ", ".join(missing)
|
||||
raise MissingAgentStubEnvironmentError(f"missing required Agent Stub environment variables: {names}")
|
||||
try:
|
||||
normalized_url = normalize_agent_stub_url(url)
|
||||
normalized_url = normalize_agent_stub_api_base_url(url)
|
||||
except ValueError as exc:
|
||||
raise MissingAgentStubEnvironmentError(f"invalid {AGENT_STUB_URL_ENV_VAR}: {exc}") from exc
|
||||
raise MissingAgentStubEnvironmentError(f"invalid {AGENT_STUB_API_BASE_URL_ENV_VAR}: {exc}") from exc
|
||||
return AgentStubEnvironment(url=normalized_url, auth_jwe=auth_jwe)
|
||||
|
||||
|
||||
def read_agent_stub_drive_base(env: Mapping[str, str] | None = None) -> str:
|
||||
"""Read the sandbox-local drive base used by ``dify-agent drive pull``.
|
||||
|
||||
The variable is optional because older Agent Stub environments only injected
|
||||
URL/auth values. Blank values keep the historical ``/mnt/drive`` fallback.
|
||||
"""
|
||||
values = env or os.environ
|
||||
configured_drive_base = (values.get(AGENT_STUB_DRIVE_BASE_ENV_VAR) or "").strip()
|
||||
return configured_drive_base or DEFAULT_AGENT_STUB_DRIVE_BASE
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AgentStubEnvironment",
|
||||
"MissingAgentStubEnvironmentError",
|
||||
"has_agent_stub_environment",
|
||||
"read_agent_stub_drive_base",
|
||||
"read_agent_stub_environment",
|
||||
]
|
||||
|
||||
@ -21,9 +21,14 @@ from dify_agent.agent_stub.cli._drive import (
|
||||
pull_drive_from_environment,
|
||||
push_drive_from_environment,
|
||||
)
|
||||
from dify_agent.agent_stub.cli._env import MissingAgentStubEnvironmentError, has_agent_stub_environment
|
||||
from dify_agent.agent_stub.cli._env import (
|
||||
MissingAgentStubEnvironmentError,
|
||||
has_agent_stub_environment,
|
||||
read_agent_stub_drive_base,
|
||||
)
|
||||
from dify_agent.agent_stub.cli._files import download_file_from_environment, upload_file_from_environment
|
||||
from dify_agent.agent_stub.client._errors import AgentStubClientError
|
||||
from dify_agent.agent_stub.protocol.agent_stub import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE
|
||||
|
||||
|
||||
app = typer.Typer(
|
||||
@ -78,11 +83,22 @@ def drive_list(
|
||||
|
||||
@drive_app.command("pull")
|
||||
def drive_pull(
|
||||
path_prefix: str = typer.Argument("", metavar="PATH_PREFIX"),
|
||||
drive_base: str = typer.Option("/mnt/drive", "--drive-base", help="Local base directory for pulled drive files."),
|
||||
targets: list[str] = typer.Argument(None, metavar="TARGET"),
|
||||
drive_base: str | None = typer.Option(
|
||||
None,
|
||||
"--drive-base",
|
||||
help=(
|
||||
f"Local base directory for pulled drive files. Defaults to ${AGENT_STUB_DRIVE_BASE_ENV_VAR} "
|
||||
f"or {DEFAULT_AGENT_STUB_DRIVE_BASE}."
|
||||
),
|
||||
),
|
||||
) -> None:
|
||||
"""Pull drive files into one local directory tree."""
|
||||
_run_drive_pull(path_prefix=path_prefix, drive_base=drive_base)
|
||||
"""Pull one or more drive keys/prefixes into one local directory tree.
|
||||
|
||||
Passing no ``TARGET`` preserves the historical whole-drive behavior by
|
||||
pulling from the empty prefix.
|
||||
"""
|
||||
_run_drive_pull(targets=targets, drive_base=drive_base)
|
||||
|
||||
|
||||
@drive_app.command("push")
|
||||
@ -207,9 +223,9 @@ def _run_drive_list(*, path_prefix: str, json_output: bool) -> None:
|
||||
typer.echo(response)
|
||||
|
||||
|
||||
def _run_drive_pull(*, path_prefix: str, drive_base: str) -> None:
|
||||
def _run_drive_pull(*, targets: list[str] | None, drive_base: str | None) -> None:
|
||||
try:
|
||||
response = pull_drive_from_environment(prefix=path_prefix, drive_base=drive_base)
|
||||
response = pull_drive_from_environment(targets=targets, drive_base=drive_base or read_agent_stub_drive_base())
|
||||
except MissingAgentStubEnvironmentError as exc:
|
||||
typer.echo(str(exc), err=True)
|
||||
raise SystemExit(2) from exc
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user