Compare commits
72
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b95d313128 | ||
|
|
45ac70132c | ||
|
|
6e772259bb | ||
|
|
af465fb527 | ||
|
|
62bbc0dbeb | ||
|
|
ed1289281d | ||
|
|
e37f32e4a0 | ||
|
|
150613cc05 | ||
|
|
a5194ff17b | ||
|
|
cdaf5d4562 | ||
|
|
145b0e622e | ||
|
|
4bab62d458 | ||
|
|
b99f8941bb | ||
|
|
32d2594877 | ||
|
|
7b188287da | ||
|
|
f6db23444d | ||
|
|
a5c989d82b | ||
|
|
2622b1ce0a | ||
|
|
1e672530c4 | ||
|
|
852169d99c | ||
|
|
d7860571aa | ||
|
|
3683831738 | ||
|
|
f5024f5440 | ||
|
|
4f8477427e | ||
|
|
8d406aec31 | ||
|
|
f8754286ec | ||
|
|
7da6b2ce36 | ||
|
|
aa718add24 | ||
|
|
fcc1d39f5e | ||
|
|
861831f176 | ||
|
|
d333aa31ea | ||
|
|
97e6a14c13 | ||
|
|
6a338654fb | ||
|
|
8f075c82bc | ||
|
|
28d17da3b1 | ||
|
|
891b2dc537 | ||
|
|
75f7069541 | ||
|
|
b7193d1cba | ||
|
|
ef0115d340 | ||
|
|
3afb9b3230 | ||
|
|
0862641533 | ||
|
|
cbd7872520 | ||
|
|
23882a704e | ||
|
|
1bd5254ea7 | ||
|
|
e13069dbe3 | ||
|
|
edd3d4a218 | ||
|
|
79751f7622 | ||
|
|
cb64446fa3 | ||
|
|
e77ee82526 | ||
|
|
b3f163fb0f | ||
|
|
873ef2b592 | ||
|
|
43d725a1ce | ||
|
|
c8ccfba960 | ||
|
|
eee8d6cf7b | ||
|
|
11016563b1 | ||
|
|
b7edc5ef0c | ||
|
|
f93a5b95c6 | ||
|
|
6022939bf0 | ||
|
|
30df1433cc | ||
|
|
aa7b65c602 | ||
|
|
ba24745c07 | ||
|
|
38ea0de907 | ||
|
|
bbc31a3403 | ||
|
|
605c0fe02b | ||
|
|
15674a9db6 | ||
|
|
9a2d039e93 | ||
|
|
58d1dd8873 | ||
|
|
5ea884f799 | ||
|
|
2ec34b2cfb | ||
|
|
fa65f45bff | ||
|
|
63ca2b94b5 | ||
|
|
a7ef41a8f5 |
@@ -9,12 +9,11 @@ Use this skill for Vitest work under `web/` and `packages/dify-ui/`. Do not use
|
||||
|
||||
## Required Source
|
||||
|
||||
Before writing, changing, or reviewing frontend tests, read `web/docs/test.md` completely. It is the single source of truth. This skill defines the execution workflow and must not add requirements that conflict with or duplicate that guide.
|
||||
Before writing, changing, or reviewing frontend tests, read `web/docs/test.md` completely. It is the single source of truth. This skill provides an execution checklist and must not redefine or extend that policy.
|
||||
|
||||
## Workflow
|
||||
|
||||
1. Read the source, its behavior owner, nearby specs, and relevant public dependencies.
|
||||
1. Identify whether the contract belongs in `web/`, Dify UI Browser Mode, or a styled Storybook test.
|
||||
1. Apply the canonical guide to decide whether a test is needed and choose its boundary.
|
||||
1. For a behavior change or bug fix, write or identify the failing scenario first when practical.
|
||||
1. Implement one coherent scenario at a time and run the focused spec before expanding scope.
|
||||
@@ -33,4 +32,4 @@ vp test run path/to/spec-or-directory
|
||||
vp test run --project unit src/path/to/spec
|
||||
```
|
||||
|
||||
For styled Dify UI behavior, run `vp test --project storybook --run`. Run broader checks only after the focused behavior passes.
|
||||
Run Dify UI Storybook tests with `vp test --project storybook --run`. Run broader checks only after the focused behavior passes.
|
||||
|
||||
@@ -47,9 +47,6 @@ jobs:
|
||||
- name: Install dependencies
|
||||
run: uv sync --project api --dev
|
||||
|
||||
- name: Run dify config tests
|
||||
run: uv run --project api pytest api/tests/unit_tests/configs/test_env_consistency.py
|
||||
|
||||
- name: Run Unit Tests
|
||||
run: |
|
||||
uv run --project api pytest \
|
||||
|
||||
@@ -27,7 +27,7 @@ jobs:
|
||||
steps:
|
||||
- id: skip_check
|
||||
continue-on-error: true
|
||||
uses: fkirc/skip-duplicate-actions@f75f66ce1886f00957d99748a42c724f4330bdcf # v5.3.1
|
||||
uses: fkirc/skip-duplicate-actions@b974a9395958c231af965b70070979a577efa578 # v5.3.2
|
||||
with:
|
||||
cancel_others: 'true'
|
||||
concurrent_skipping: same_content_newer
|
||||
@@ -81,7 +81,6 @@ jobs:
|
||||
- '.npmrc'
|
||||
- '.nvmrc'
|
||||
- '.github/workflows/cli-tests.yml'
|
||||
- '.github/workflows/cli-docker-build.yml'
|
||||
- '.github/actions/setup-web/**'
|
||||
web:
|
||||
- 'web/**'
|
||||
|
||||
@@ -27,7 +27,7 @@ jobs:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0
|
||||
uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0
|
||||
with:
|
||||
go-version-file: dify-agent-runtime/go.mod
|
||||
cache-dependency-path: dify-agent-runtime/go.sum
|
||||
@@ -51,13 +51,13 @@ jobs:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0
|
||||
uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0
|
||||
with:
|
||||
go-version-file: dify-agent-runtime/go.mod
|
||||
cache-dependency-path: dify-agent-runtime/go.sum
|
||||
|
||||
- name: Run golangci-lint
|
||||
uses: golangci/golangci-lint-action@4afd733a84b1f43292c63897423277bb7f4313a9 # v6.5.0
|
||||
uses: golangci/golangci-lint-action@ba0d7d2ec06a0ea1cb5fa41b2e4a3ab91d21278a # v6.5.0
|
||||
with:
|
||||
working-directory: dify-agent-runtime
|
||||
version: latest
|
||||
@@ -78,7 +78,7 @@ jobs:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0
|
||||
uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0
|
||||
with:
|
||||
go-version-file: dify-agent-runtime/go.mod
|
||||
cache-dependency-path: dify-agent-runtime/go.sum
|
||||
|
||||
@@ -29,7 +29,7 @@ jobs:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Use Node.js
|
||||
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
|
||||
uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
|
||||
with:
|
||||
node-version: 22
|
||||
cache: ''
|
||||
|
||||
@@ -158,7 +158,7 @@ jobs:
|
||||
|
||||
- name: Run Claude Code for Translation Sync
|
||||
if: steps.context.outputs.CHANGED_FILES != ''
|
||||
uses: anthropics/claude-code-action@e90deca47693f9457b72f2b53c17d7c445a87342 # v1.0.171
|
||||
uses: anthropics/claude-code-action@af0559ee4f514d1ef21826982bed13f7edc3c35e # v1.0.178
|
||||
with:
|
||||
anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
github_token: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
@@ -135,16 +135,6 @@ jobs:
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ -d cucumber-report ]]; then
|
||||
rm -rf cucumber-report-non-external
|
||||
mv cucumber-report cucumber-report-non-external
|
||||
fi
|
||||
|
||||
if [[ -d .logs ]]; then
|
||||
rm -rf .logs-non-external
|
||||
mv .logs .logs-non-external
|
||||
fi
|
||||
|
||||
teardown_external_runtime() {
|
||||
local run_status=$?
|
||||
trap - EXIT
|
||||
|
||||
@@ -17,8 +17,6 @@ jobs:
|
||||
test:
|
||||
name: Web Tests (${{ matrix.shardIndex }}/${{ matrix.shardTotal }})
|
||||
runs-on: depot-ubuntu-24.04-4
|
||||
env:
|
||||
VITEST_COVERAGE_SCOPE: app-components
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
|
||||
@@ -107,6 +107,7 @@ test:
|
||||
echo "Target: $(TARGET_TESTS)"; \
|
||||
uv run --project api --dev pytest $(TARGET_TESTS); \
|
||||
else \
|
||||
set -e; \
|
||||
echo "Running backend unit tests"; \
|
||||
uv run --project api --dev pytest -p no:benchmark --timeout "$${PYTEST_TIMEOUT:-20}" -n auto \
|
||||
api/tests/unit_tests \
|
||||
@@ -124,6 +125,7 @@ test-all:
|
||||
echo "Target: $(TARGET_TESTS)"; \
|
||||
uv run --project api --dev pytest $(TARGET_TESTS); \
|
||||
else \
|
||||
set -e; \
|
||||
echo "Running backend unit tests"; \
|
||||
uv run --project api --dev pytest -p no:benchmark --timeout "$${PYTEST_TIMEOUT:-20}" -n auto \
|
||||
api/tests/unit_tests \
|
||||
|
||||
@@ -681,6 +681,14 @@ AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS=30
|
||||
AGENT_BACKEND_STREAM_MAX_RECONNECTS=3
|
||||
AGENT_BACKEND_RUN_TIMEOUT_SECONDS=1200
|
||||
|
||||
# KnowledgeFS (Dataset 2.0)
|
||||
KNOWLEDGE_FS_ENABLED=false
|
||||
KNOWLEDGE_FS_BASE_URL=
|
||||
# Shared with KnowledgeFS; use at least 32 random characters.
|
||||
KNOWLEDGE_FS_JWT_SECRET=
|
||||
KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS=300
|
||||
KNOWLEDGE_FS_TIMEOUT_SECONDS=10
|
||||
|
||||
# Marketplace configuration
|
||||
MARKETPLACE_ENABLED=true
|
||||
MARKETPLACE_API_URL=https://marketplace.dify.ai
|
||||
|
||||
+209
-63
@@ -1,6 +1,8 @@
|
||||
import datetime
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from typing import TypedDict
|
||||
|
||||
@@ -21,6 +23,7 @@ from tasks.remove_app_and_related_data_task import delete_draft_variables_batch
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_HEX_PREFIXES = tuple("0123456789abcdef")
|
||||
_TARGET_MONTH_PATTERN = re.compile(r"^\d{4}-(0[1-9]|1[0-2])$")
|
||||
|
||||
|
||||
class WorkflowRunArchivePlanRow(TypedDict):
|
||||
@@ -66,6 +69,7 @@ def _parse_tenant_prefixes(prefixes: str | None) -> list[str]:
|
||||
|
||||
|
||||
def _parse_comma_separated_ids(raw_ids: str | None, *, param_name: str) -> list[str] | None:
|
||||
"""Keep an omitted scope unset while rejecting an explicitly empty scope."""
|
||||
if raw_ids is None:
|
||||
return None
|
||||
parsed = sorted({raw_id.strip() for raw_id in raw_ids.split(",") if raw_id.strip()})
|
||||
@@ -74,6 +78,27 @@ def _parse_comma_separated_ids(raw_ids: str | None, *, param_name: str) -> list[
|
||||
return parsed
|
||||
|
||||
|
||||
def _parse_archive_target_month(target_month: str) -> tuple[int, int]:
|
||||
"""Validate the V2 catalog month selector and return its numeric components."""
|
||||
if not _TARGET_MONTH_PATTERN.fullmatch(target_month):
|
||||
raise click.BadParameter("target-month must use YYYY-MM format", param_hint="--target-month")
|
||||
year_text, month_text = target_month.split("-", maxsplit=1)
|
||||
return int(year_text), int(month_text)
|
||||
|
||||
|
||||
def _parse_archive_catalog_cursor(after_catalog_id: str | None) -> str | None:
|
||||
"""Normalize the exclusive V2 catalog keyset cursor when one is provided."""
|
||||
if after_catalog_id is None:
|
||||
return None
|
||||
try:
|
||||
return str(uuid.UUID(after_catalog_id))
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter(
|
||||
"after-catalog-id must be a UUID returned by the same V2 operation and scope",
|
||||
param_hint="--after-catalog-id",
|
||||
) from exc
|
||||
|
||||
|
||||
def _get_archive_candidate_tenant_ids_by_prefix(
|
||||
session: Session,
|
||||
prefix: str,
|
||||
@@ -810,9 +835,11 @@ def backfill_workflow_run_archive_bundles(
|
||||
click.echo(click.style(f" ... and {len(summary.errors) - 10} more failures", fg="red"))
|
||||
|
||||
|
||||
def _echo_bundle_archive_operation_summary(summary) -> None:
|
||||
def _echo_bundle_archive_operation_summary(summary, *, dry_run: bool) -> None:
|
||||
status = "completed successfully" if summary.bundles_failed == 0 else "completed with failures"
|
||||
fg = "green" if summary.bundles_failed == 0 else "red"
|
||||
cursor_label = "preview_next_catalog_id" if dry_run else "next_catalog_id"
|
||||
cursor_value = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id
|
||||
click.echo(
|
||||
click.style(
|
||||
f"{summary.operation} {status}. "
|
||||
@@ -821,10 +848,12 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
|
||||
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}",
|
||||
f"bytes_per_second={summary.bytes_per_second:.2f} {cursor_label}={cursor_value or 'none'}",
|
||||
fg=fg,
|
||||
)
|
||||
)
|
||||
if dry_run:
|
||||
click.echo(click.style("Dry-run cursor is preview-only; do not persist it for a destructive run.", fg="yellow"))
|
||||
click.echo(click.style("table,row_count", fg="white"))
|
||||
for table_name in [
|
||||
"workflow_runs",
|
||||
@@ -842,7 +871,8 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
|
||||
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",
|
||||
f"catalog_id={result.catalog_id} time={result.elapsed_time:.2f}s "
|
||||
f"validation={result.validation_time:.2f}s",
|
||||
fg="white",
|
||||
)
|
||||
)
|
||||
@@ -850,7 +880,7 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
|
||||
click.echo(
|
||||
click.style(
|
||||
f" failed bundle={result.bundle_id} tenant={result.tenant_id} "
|
||||
f"object_prefix={result.object_prefix} error={result.error}",
|
||||
f"catalog_id={result.catalog_id} object_prefix={result.object_prefix} error={result.error}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
@@ -867,25 +897,24 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
|
||||
)
|
||||
@click.option("--run-id", required=False, help="Workflow run ID to restore.")
|
||||
@click.option(
|
||||
"--start-from",
|
||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||
"--target-month",
|
||||
metavar="YYYY-MM",
|
||||
default=None,
|
||||
help="Optional lower bound (inclusive) for created_at; must be paired with --end-before.",
|
||||
help="V2 catalog month to restore; required unless --run-id is used.",
|
||||
)
|
||||
@click.option(
|
||||
"--end-before",
|
||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||
"--after-catalog-id",
|
||||
default=None,
|
||||
help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.",
|
||||
help="Exclusive V2 cursor from the same restore month and tenant scope.",
|
||||
)
|
||||
@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("--limit", type=click.IntRange(min=1), default=100, show_default=True, help="Maximum V2 catalog rows.")
|
||||
@click.option("--dry-run", is_flag=True, help="Preview without restoring.")
|
||||
def restore_workflow_runs(
|
||||
tenant_ids: str | None,
|
||||
run_id: str | None,
|
||||
start_from: datetime.datetime | None,
|
||||
end_before: datetime.datetime | None,
|
||||
target_month: str | None,
|
||||
after_catalog_id: str | None,
|
||||
workers: int,
|
||||
limit: int,
|
||||
dry_run: bool,
|
||||
@@ -905,23 +934,20 @@ def restore_workflow_runs(
|
||||
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
|
||||
if tenant_ids:
|
||||
parsed_tenant_ids = [tid.strip() for tid in tenant_ids.split(",") if tid.strip()]
|
||||
if not parsed_tenant_ids:
|
||||
raise click.BadParameter("tenant-ids must not be empty")
|
||||
parsed_tenant_ids = _parse_comma_separated_ids(tenant_ids, param_name="tenant-ids")
|
||||
|
||||
if (start_from is None) ^ (end_before is None):
|
||||
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 restore.")
|
||||
if workers < 1:
|
||||
raise click.BadParameter("workers must be at least 1")
|
||||
if run_id is not None and (target_month is not None or after_catalog_id is not None):
|
||||
raise click.UsageError("--target-month and --after-catalog-id are only valid for V2 batch restore.")
|
||||
if run_id is None and target_month is None:
|
||||
raise click.UsageError("--target-month is required for V2 batch restore.")
|
||||
|
||||
start_time = datetime.datetime.now(datetime.UTC)
|
||||
target_desc = f"workflow run {run_id}" if run_id else f"workflow archive catalog month {target_month}"
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Starting restore of workflow run {run_id} at {start_time.isoformat()}.",
|
||||
f"Starting restore of {target_desc} at {start_time.isoformat()}.",
|
||||
fg="white",
|
||||
)
|
||||
)
|
||||
@@ -955,17 +981,20 @@ def restore_workflow_runs(
|
||||
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
|
||||
assert target_month is not None
|
||||
target_year, target_month_number = _parse_archive_target_month(target_month)
|
||||
catalog_cursor = _parse_archive_catalog_cursor(after_catalog_id)
|
||||
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,
|
||||
target_year=target_year,
|
||||
target_month=target_month_number,
|
||||
after_catalog_id=catalog_cursor,
|
||||
limit=limit,
|
||||
)
|
||||
_echo_bundle_archive_operation_summary(summary)
|
||||
return
|
||||
_echo_bundle_archive_operation_summary(summary, dry_run=dry_run)
|
||||
if summary.bundles_failed:
|
||||
raise click.exceptions.Exit(1)
|
||||
|
||||
|
||||
@click.command(
|
||||
@@ -979,23 +1008,41 @@ def restore_workflow_runs(
|
||||
)
|
||||
@click.option("--run-id", required=False, help="Workflow run ID to delete.")
|
||||
@click.option(
|
||||
"--start-from",
|
||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||
"--target-month",
|
||||
metavar="YYYY-MM",
|
||||
default=None,
|
||||
help="Optional lower bound (inclusive) for created_at; must be paired with --end-before.",
|
||||
help="V2 catalog month to delete; required unless --run-id is used.",
|
||||
)
|
||||
@click.option(
|
||||
"--end-before",
|
||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||
"--after-catalog-id",
|
||||
default=None,
|
||||
help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.",
|
||||
help="Exclusive V2 cursor from the same delete month and tenant scope.",
|
||||
)
|
||||
@click.option(
|
||||
"--run-shard-index",
|
||||
default=None,
|
||||
type=click.IntRange(min=0),
|
||||
help="Zero-based archive shard index. Must be paired with --run-shard-total.",
|
||||
)
|
||||
@click.option(
|
||||
"--run-shard-total",
|
||||
default=None,
|
||||
type=click.IntRange(min=1, max=16),
|
||||
help="Total archive shard count. Must be paired with --run-shard-index.",
|
||||
)
|
||||
@click.option("--all-pages", is_flag=True, help="Process catalog pages until an empty page is reached.")
|
||||
@click.option(
|
||||
"--limit",
|
||||
type=click.IntRange(min=1),
|
||||
default=100,
|
||||
show_default=True,
|
||||
help="Maximum V2 catalog rows per page.",
|
||||
)
|
||||
@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.",
|
||||
help="V1 --run-id only: continue when one archive object fails validation.",
|
||||
)
|
||||
@click.option(
|
||||
"--restore-sample-interval",
|
||||
@@ -1007,8 +1054,11 @@ def restore_workflow_runs(
|
||||
def delete_archived_workflow_runs(
|
||||
tenant_ids: str | None,
|
||||
run_id: str | None,
|
||||
start_from: datetime.datetime | None,
|
||||
end_before: datetime.datetime | None,
|
||||
target_month: str | None,
|
||||
after_catalog_id: str | None,
|
||||
run_shard_index: int | None,
|
||||
run_shard_total: int | None,
|
||||
all_pages: bool,
|
||||
limit: int,
|
||||
dry_run: bool,
|
||||
skip_bad_archives: bool,
|
||||
@@ -1018,26 +1068,38 @@ def delete_archived_workflow_runs(
|
||||
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.
|
||||
counts, and source/archive content checksums before deleting source rows. Parallel workers may select one exact
|
||||
archive shard; all-pages mode keeps only the current bounded page in memory. `--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
|
||||
if tenant_ids:
|
||||
parsed_tenant_ids = [tid.strip() for tid in tenant_ids.split(",") if tid.strip()]
|
||||
if not parsed_tenant_ids:
|
||||
raise click.BadParameter("tenant-ids must not be empty")
|
||||
parsed_tenant_ids = _parse_comma_separated_ids(tenant_ids, param_name="tenant-ids")
|
||||
|
||||
if (start_from is None) ^ (end_before is None):
|
||||
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")
|
||||
if run_id is not None and (
|
||||
target_month is not None
|
||||
or after_catalog_id is not None
|
||||
or run_shard_index is not None
|
||||
or run_shard_total is not None
|
||||
or all_pages
|
||||
):
|
||||
raise click.UsageError(
|
||||
"--target-month, --after-catalog-id, --run-shard-index, --run-shard-total, and --all-pages "
|
||||
"are only valid for V2 batch delete."
|
||||
)
|
||||
if run_id is None and target_month is None:
|
||||
raise click.UsageError("--target-month is required for V2 batch delete.")
|
||||
if run_id is None and skip_bad_archives:
|
||||
raise click.UsageError("--skip-bad-archives is not supported for V2 catalog batches; they fail fast.")
|
||||
if (run_shard_index is None) ^ (run_shard_total is None):
|
||||
raise click.UsageError("--run-shard-index and --run-shard-total must be provided together.")
|
||||
if run_shard_index is not None and run_shard_total is not None and run_shard_index >= run_shard_total:
|
||||
raise click.UsageError("--run-shard-index must be less than --run-shard-total.")
|
||||
|
||||
start_time = datetime.datetime.now(datetime.UTC)
|
||||
target_desc = f"workflow run {run_id}" if run_id else "workflow runs"
|
||||
target_desc = f"workflow run {run_id}" if run_id else f"workflow archive catalog month {target_month}"
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Starting delete of {target_desc} at {start_time.isoformat()}.",
|
||||
@@ -1110,20 +1172,104 @@ def delete_archived_workflow_runs(
|
||||
|
||||
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,
|
||||
assert target_month is not None
|
||||
target_year, target_month_number = _parse_archive_target_month(target_month)
|
||||
catalog_cursor = _parse_archive_catalog_cursor(after_catalog_id)
|
||||
shard = (
|
||||
f"{run_shard_index:02d}-of-{run_shard_total:02d}"
|
||||
if run_shard_index is not None and run_shard_total is not None
|
||||
else None
|
||||
)
|
||||
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)
|
||||
bundle_deleter = WorkflowRunBundleArchiveMaintenance(dry_run=dry_run, strict_content_validation=True)
|
||||
if run_shard_total is not None:
|
||||
try:
|
||||
bundle_deleter.validate_catalog_shards(
|
||||
target_year=target_year,
|
||||
target_month=target_month_number,
|
||||
shard_total=run_shard_total,
|
||||
tenant_ids=parsed_tenant_ids,
|
||||
)
|
||||
except ValueError as exc:
|
||||
logger.exception(
|
||||
"Archive catalog shard preflight failed: target_month=%s shard=%s",
|
||||
target_month,
|
||||
shard,
|
||||
)
|
||||
raise click.ClickException(
|
||||
f"Archive catalog shard preflight failed for target_month={target_month} shard={shard}: {exc}"
|
||||
) from exc
|
||||
|
||||
initial_catalog_cursor = catalog_cursor
|
||||
pages_processed = 0
|
||||
bundles_succeeded = 0
|
||||
runs_processed = 0
|
||||
rows_processed = 0
|
||||
archive_bytes = 0
|
||||
while True:
|
||||
summary = bundle_deleter.delete_batch(
|
||||
tenant_ids=parsed_tenant_ids,
|
||||
target_year=target_year,
|
||||
target_month=target_month_number,
|
||||
after_catalog_id=catalog_cursor,
|
||||
limit=limit,
|
||||
shard=shard,
|
||||
)
|
||||
_echo_bundle_archive_operation_summary(summary, dry_run=dry_run)
|
||||
if summary.bundles_failed:
|
||||
failed_result = next((result for result in summary.results if not result.success), None)
|
||||
failed_catalog_id = failed_result.catalog_id if failed_result is not None else "unknown"
|
||||
page_resume_cursor = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id
|
||||
resume_cursor = page_resume_cursor or catalog_cursor
|
||||
if dry_run:
|
||||
cursor_details = (
|
||||
f"preview_after_catalog_id={resume_cursor or 'none'} "
|
||||
f"destructive_retry_after_catalog_id={initial_catalog_cursor or 'none'}"
|
||||
)
|
||||
else:
|
||||
cursor_details = f"resume_after_catalog_id={resume_cursor or 'none'}"
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Delete stopped: target_month={target_month} shard={shard or 'all'} "
|
||||
f"failed_catalog_id={failed_catalog_id} "
|
||||
f"{cursor_details}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
raise click.exceptions.Exit(1)
|
||||
|
||||
if not all_pages:
|
||||
break
|
||||
if summary.bundles_processed == 0:
|
||||
break
|
||||
|
||||
pages_processed += 1
|
||||
bundles_succeeded += summary.bundles_succeeded
|
||||
runs_processed += summary.runs_processed
|
||||
rows_processed += summary.rows_processed
|
||||
archive_bytes += summary.archive_bytes
|
||||
next_catalog_id = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id
|
||||
if next_catalog_id is None or (catalog_cursor is not None and next_catalog_id <= catalog_cursor):
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Delete cursor did not advance: target_month={target_month} shard={shard or 'all'} "
|
||||
f"after_catalog_id={catalog_cursor or 'none'} next_catalog_id={next_catalog_id or 'none'}",
|
||||
fg="red",
|
||||
)
|
||||
)
|
||||
raise click.exceptions.Exit(1)
|
||||
catalog_cursor = next_catalog_id
|
||||
|
||||
if all_pages:
|
||||
final_cursor_label = "preview_final_catalog_id" if dry_run else "final_catalog_id"
|
||||
click.echo(
|
||||
click.style(
|
||||
f"Delete all-pages completed successfully. target_month={target_month} shard={shard or 'all'} "
|
||||
f"pages={pages_processed} bundles_success={bundles_succeeded} runs={runs_processed} "
|
||||
f"rows={rows_processed} archive_bytes={archive_bytes} "
|
||||
f"{final_cursor_label}={catalog_cursor or 'none'}",
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _find_orphaned_draft_variables(batch_size: int = 1000) -> list[str]:
|
||||
|
||||
@@ -14,6 +14,12 @@ class EnterpriseFeatureConfig(BaseSettings):
|
||||
default=False,
|
||||
)
|
||||
|
||||
WEBAPP_PUBLIC_ACCESS_ENABLED: bool = Field(
|
||||
description="Whether admins are allowed to set a webapp's access mode to public (anyone with the link, "
|
||||
"no auth). Disable in security-sensitive on-prem deployments.",
|
||||
default=True,
|
||||
)
|
||||
|
||||
CAN_REPLACE_LOGO: bool = Field(
|
||||
description="Allow customization of the enterprise logo.",
|
||||
default=False,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from configs.extra.agent_backend_config import AgentBackendConfig
|
||||
from configs.extra.archive_config import ArchiveStorageConfig
|
||||
from configs.extra.knowledge_fs_config import KnowledgeFSConfig
|
||||
from configs.extra.notion_config import NotionConfig
|
||||
from configs.extra.sentry_config import SentryConfig
|
||||
|
||||
@@ -8,6 +9,7 @@ class ExtraServiceConfig(
|
||||
# place the configs in alphabet order
|
||||
AgentBackendConfig,
|
||||
ArchiveStorageConfig,
|
||||
KnowledgeFSConfig,
|
||||
NotionConfig,
|
||||
SentryConfig,
|
||||
):
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Configuration for the optional KnowledgeFS Console bridge."""
|
||||
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from pydantic import Field, PositiveFloat, SecretStr, field_validator, model_validator
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class KnowledgeFSConfig(BaseSettings):
|
||||
"""Server-only settings for the KnowledgeFS production connection."""
|
||||
|
||||
KNOWLEDGE_FS_ENABLED: bool = Field(
|
||||
default=False,
|
||||
description="Enable the private KnowledgeFS Console bridge.",
|
||||
)
|
||||
KNOWLEDGE_FS_BASE_URL: str | None = Field(default=None, description="KnowledgeFS gateway base URL.")
|
||||
KNOWLEDGE_FS_JWT_SECRET: SecretStr | None = Field(
|
||||
default=None,
|
||||
min_length=32,
|
||||
description="Shared secret used to sign short-lived KnowledgeFS service JWTs.",
|
||||
)
|
||||
KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS: PositiveFloat = Field(default=300.0, le=3600.0, allow_inf_nan=False)
|
||||
KNOWLEDGE_FS_TIMEOUT_SECONDS: PositiveFloat = Field(default=10.0, le=60.0, allow_inf_nan=False)
|
||||
|
||||
@field_validator(
|
||||
"KNOWLEDGE_FS_BASE_URL",
|
||||
"KNOWLEDGE_FS_JWT_SECRET",
|
||||
mode="before",
|
||||
)
|
||||
@classmethod
|
||||
def normalize_optional_string(cls, value: object) -> object:
|
||||
if isinstance(value, SecretStr):
|
||||
normalized = value.get_secret_value().strip()
|
||||
return SecretStr(normalized) if normalized else None
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip()
|
||||
return normalized or None
|
||||
return value
|
||||
|
||||
@field_validator("KNOWLEDGE_FS_BASE_URL")
|
||||
@classmethod
|
||||
def validate_base_url(cls, value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
parsed = urlsplit(value)
|
||||
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL must be an absolute HTTP(S) URL")
|
||||
try:
|
||||
_ = parsed.port
|
||||
except ValueError as exc:
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL must include a valid port") from exc
|
||||
if parsed.username or parsed.password or parsed.query or parsed.fragment:
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL must not include credentials, query, or fragment")
|
||||
return value.rstrip("/")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_enabled_connection(self) -> "KnowledgeFSConfig":
|
||||
if not self.KNOWLEDGE_FS_ENABLED:
|
||||
return self
|
||||
if bool(self.KNOWLEDGE_FS_BASE_URL) != bool(self.KNOWLEDGE_FS_JWT_SECRET):
|
||||
raise ValueError("KNOWLEDGE_FS_BASE_URL and KNOWLEDGE_FS_JWT_SECRET must be configured together")
|
||||
if not self.KNOWLEDGE_FS_BASE_URL:
|
||||
raise ValueError("KnowledgeFS connection settings are required when the integration is enabled")
|
||||
return self
|
||||
@@ -1,4 +1,4 @@
|
||||
from datetime import timedelta
|
||||
from datetime import datetime, timedelta
|
||||
from enum import StrEnum
|
||||
from typing import Literal
|
||||
|
||||
@@ -11,6 +11,7 @@ from pydantic import (
|
||||
PositiveFloat,
|
||||
PositiveInt,
|
||||
computed_field,
|
||||
field_validator,
|
||||
)
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
@@ -280,6 +281,27 @@ class PluginConfig(BaseSettings):
|
||||
default="",
|
||||
)
|
||||
|
||||
@field_validator("PLUGIN_REMOTE_INSTALL_PORT", mode="before")
|
||||
@classmethod
|
||||
def _reject_host_port_shaped_plugin_remote_install_port(cls, v):
|
||||
"""Reject ``host:port``-shaped values with an actionable hint.
|
||||
|
||||
``EXPOSE_PLUGIN_DEBUGGING_PORT`` is overloaded: it feeds both the
|
||||
plugin_daemon ``ports:`` mapping (where ``127.0.0.1:5003`` is valid
|
||||
compose syntax) and this integer app setting advertised in the console.
|
||||
Without this guard a loopback bind spec crashloops the api container
|
||||
with an opaque ``int_parsing`` traceback. See issue #39323.
|
||||
"""
|
||||
if isinstance(v, str) and ":" in v.strip():
|
||||
raise ValueError(
|
||||
"PLUGIN_REMOTE_INSTALL_PORT must be a bare port number, got "
|
||||
f"{v!r}. A 'host:port' value usually means "
|
||||
"EXPOSE_PLUGIN_DEBUGGING_PORT was set to a compose publish spec "
|
||||
"like '127.0.0.1:5003'; bind loopback via a "
|
||||
"docker-compose.override.yaml instead of overloading this var."
|
||||
)
|
||||
return v
|
||||
|
||||
@property
|
||||
def NEW_USER_DEFAULT_PLUGIN_ID_LIST(self) -> list[str]:
|
||||
return [item.strip() for item in self.NEW_USER_DEFAULT_PLUGIN_IDS.split(",") if item.strip()]
|
||||
@@ -1138,6 +1160,16 @@ class HomepageConfig(BaseSettings):
|
||||
default=True,
|
||||
)
|
||||
|
||||
ENABLE_STEP_BY_STEP_TOUR: bool = Field(
|
||||
description="Enable account-level Step-by-step Tour eligibility checks",
|
||||
default=False,
|
||||
)
|
||||
|
||||
STEP_BY_STEP_TOUR_ROLLOUT_STARTED_AT: datetime | None = Field(
|
||||
description="UTC timestamp after which newly initialized accounts are eligible for Step-by-step Tour",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
class RagEtlConfig(BaseSettings):
|
||||
"""
|
||||
|
||||
@@ -38,7 +38,9 @@ from . import (
|
||||
feature,
|
||||
human_input_form,
|
||||
init_validate,
|
||||
knowledge_fs_proxy,
|
||||
notification,
|
||||
onboarding,
|
||||
ping,
|
||||
setup,
|
||||
spec,
|
||||
@@ -195,6 +197,7 @@ __all__ = [
|
||||
"human_input_form",
|
||||
"init_validate",
|
||||
"installed_app",
|
||||
"knowledge_fs_proxy",
|
||||
"load_balancing_config",
|
||||
"login",
|
||||
"mcp_server",
|
||||
@@ -207,6 +210,7 @@ __all__ = [
|
||||
"notification",
|
||||
"oauth",
|
||||
"oauth_server",
|
||||
"onboarding",
|
||||
"ops_trace",
|
||||
"parameter",
|
||||
"ping",
|
||||
|
||||
@@ -105,6 +105,7 @@ class SyncDraftWorkflowPayload(BaseModel):
|
||||
graph: dict[str, Any]
|
||||
features: dict[str, Any]
|
||||
hash: str | None = None
|
||||
is_collaborative: bool = Field(default=False, alias="_is_collaborative")
|
||||
environment_variables: list[dict[str, Any]] = Field(
|
||||
default_factory=list,
|
||||
)
|
||||
@@ -610,6 +611,7 @@ class DraftWorkflowApi(Resource):
|
||||
environment_variables=environment_variables,
|
||||
conversation_variables=conversation_variables,
|
||||
session=db.session(),
|
||||
graph_only=args["is_collaborative"],
|
||||
)
|
||||
except WorkflowHashNotEqualError:
|
||||
raise DraftWorkflowNotSync()
|
||||
|
||||
@@ -607,6 +607,13 @@ class DatasetListApi(Resource):
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_created_app_rbac_access_task.delay(current_tenant_id, current_user.id, dataset_id=dataset.id)
|
||||
else:
|
||||
enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist(
|
||||
current_tenant_id,
|
||||
current_user.id,
|
||||
dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.SPECIFIC),
|
||||
)
|
||||
|
||||
permission_keys_map = enterprise_rbac_service.RBACService.DatasetPermissions.batch_get(
|
||||
current_tenant_id,
|
||||
@@ -875,7 +882,7 @@ class DatasetIndexingEstimateApi(Resource):
|
||||
file_details = session.scalars(
|
||||
select(UploadFile).where(UploadFile.tenant_id == current_tenant_id, UploadFile.id.in_(file_ids))
|
||||
).all()
|
||||
if file_details is None:
|
||||
if not file_details:
|
||||
raise NotFound("File not found.")
|
||||
|
||||
if file_details:
|
||||
|
||||
@@ -47,7 +47,9 @@ from controllers.console.explore.error import (
|
||||
NotWorkflowAppError,
|
||||
)
|
||||
from controllers.console.explore.wraps import TrialAppResource, trial_feature_enable
|
||||
from controllers.console.wraps import with_current_user
|
||||
from controllers.console.files import FILE_UPLOAD_PARAMS, upload_file_from_request
|
||||
from controllers.console.remote_files import RemoteFileUploadPayload, upload_remote_file_from_request
|
||||
from controllers.console.wraps import cloud_edition_billing_resource_check, with_current_user
|
||||
from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError
|
||||
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager
|
||||
@@ -61,12 +63,13 @@ from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from fields.base import ResponseModel
|
||||
from fields.conversation_variable_fields import WorkflowConversationVariableResponse
|
||||
from fields.file_fields import FileResponse, FileWithSignedUrl
|
||||
from fields.message_fields import SuggestedQuestionsResponse
|
||||
from graphon.graph_engine.manager import GraphEngineManager
|
||||
from graphon.model_runtime.errors.invoke import InvokeError
|
||||
from libs import helper
|
||||
from libs.helper import dump_response, to_timestamp, uuid_value
|
||||
from models import Account
|
||||
from models import Account, App
|
||||
from models.account import TenantStatus
|
||||
from models.model import AppMode, Site, load_annotation_reply_config
|
||||
from models.workflow import Workflow
|
||||
@@ -428,6 +431,36 @@ register_response_schema_models(
|
||||
simple_account_model = console_ns.models[TrialSimpleAccount.__name__]
|
||||
|
||||
|
||||
class TrialAppFileUploadApi(TrialAppResource):
|
||||
@trial_feature_enable
|
||||
@cloud_edition_billing_resource_check("documents")
|
||||
@console_ns.doc(consumes=["multipart/form-data"], params=FILE_UPLOAD_PARAMS)
|
||||
@console_ns.response(201, "File uploaded successfully", console_ns.models[FileResponse.__name__])
|
||||
@with_current_user
|
||||
def post(self, current_user: Account, app_model: App):
|
||||
"""Upload a file into the tenant that owns the trial app."""
|
||||
upload_file = upload_file_from_request(
|
||||
current_user=current_user,
|
||||
resource_tenant_id=app_model.tenant_id,
|
||||
)
|
||||
return dump_response(FileResponse, upload_file), 201
|
||||
|
||||
|
||||
class TrialAppRemoteFileUploadApi(TrialAppResource):
|
||||
@trial_feature_enable
|
||||
@cloud_edition_billing_resource_check("documents")
|
||||
@console_ns.expect(console_ns.models[RemoteFileUploadPayload.__name__])
|
||||
@console_ns.response(201, "File uploaded successfully", console_ns.models[FileWithSignedUrl.__name__])
|
||||
@with_current_user
|
||||
def post(self, current_user: Account, app_model: App):
|
||||
"""Upload a remote file into the tenant that owns the trial app."""
|
||||
remote_file = upload_remote_file_from_request(
|
||||
current_user=current_user,
|
||||
resource_tenant_id=app_model.tenant_id,
|
||||
)
|
||||
return remote_file.model_dump(mode="json"), 201
|
||||
|
||||
|
||||
class TrialAppWorkflowRunApi(TrialAppResource):
|
||||
@trial_feature_enable
|
||||
@console_ns.expect(console_ns.models[WorkflowRunRequest.__name__])
|
||||
@@ -612,7 +645,7 @@ class TrialChatAudioApi(TrialAppResource):
|
||||
def post(self, current_user: Account, trial_app):
|
||||
app_model = trial_app
|
||||
|
||||
file = request.files["file"]
|
||||
file = request.files.get("file")
|
||||
|
||||
try:
|
||||
# Get IDs before they might be detached from session
|
||||
@@ -887,6 +920,18 @@ class DatasetListApi(Resource):
|
||||
|
||||
console_ns.add_resource(TrialChatApi, "/trial-apps/<uuid:app_id>/chat-messages", endpoint="trial_app_chat_completion")
|
||||
|
||||
console_ns.add_resource(
|
||||
TrialAppFileUploadApi,
|
||||
"/trial-apps/<uuid:app_id>/files/upload",
|
||||
endpoint="trial_app_file_upload",
|
||||
)
|
||||
|
||||
console_ns.add_resource(
|
||||
TrialAppRemoteFileUploadApi,
|
||||
"/trial-apps/<uuid:app_id>/remote-files/upload",
|
||||
endpoint="trial_app_remote_file_upload",
|
||||
)
|
||||
|
||||
console_ns.add_resource(
|
||||
TrialMessageSuggestedQuestionApi,
|
||||
"/trial-apps/<uuid:app_id>/messages/<uuid:message_id>/suggested-questions",
|
||||
|
||||
@@ -29,7 +29,7 @@ from extensions.ext_database import db
|
||||
from fields.file_fields import FileResponse, UploadConfig
|
||||
from libs.helper import dump_response
|
||||
from libs.login import login_required
|
||||
from models import Account
|
||||
from models import Account, UploadFile
|
||||
from services.file_service import FileService
|
||||
|
||||
from . import console_ns
|
||||
@@ -39,7 +39,7 @@ register_response_schema_models(console_ns, AllowedExtensionsResponse, TextConte
|
||||
|
||||
PREVIEW_WORDS_LIMIT = 3000
|
||||
|
||||
_FILE_UPLOAD_PARAMS = {
|
||||
FILE_UPLOAD_PARAMS = {
|
||||
"file": {
|
||||
"description": "File to upload",
|
||||
"in": "formData",
|
||||
@@ -56,6 +56,43 @@ _FILE_UPLOAD_PARAMS = {
|
||||
}
|
||||
|
||||
|
||||
def upload_file_from_request(*, current_user: Account, resource_tenant_id: str | None = None) -> UploadFile:
|
||||
"""Validate the multipart request and persist the file under the requested resource tenant."""
|
||||
source_str = request.form.get("source")
|
||||
source: Literal["datasets"] | None = "datasets" if source_str == "datasets" else None
|
||||
|
||||
if "file" not in request.files:
|
||||
raise NoFileUploadedError()
|
||||
|
||||
if len(request.files) > 1:
|
||||
raise TooManyFilesError()
|
||||
file = request.files["file"]
|
||||
|
||||
if not file.filename:
|
||||
raise FilenameNotExistsError
|
||||
if source == "datasets" and not current_user.is_dataset_editor:
|
||||
raise Forbidden()
|
||||
|
||||
if source not in ("datasets", None):
|
||||
source = None
|
||||
|
||||
try:
|
||||
return FileService(db.engine).upload_file(
|
||||
filename=file.filename,
|
||||
content=file.stream.read(),
|
||||
mimetype=file.mimetype,
|
||||
user=current_user,
|
||||
tenant_id=resource_tenant_id,
|
||||
source=source,
|
||||
)
|
||||
except services.errors.file.FileTooLargeError as file_too_large_error:
|
||||
raise FileTooLargeError(file_too_large_error.description)
|
||||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
except services.errors.file.BlockedFileExtensionError as blocked_extension_error:
|
||||
raise BlockedFileExtensionError(blocked_extension_error.description)
|
||||
|
||||
|
||||
@console_ns.route("/files/upload")
|
||||
class FileApi(Resource):
|
||||
@setup_required
|
||||
@@ -81,42 +118,11 @@ class FileApi(Resource):
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@cloud_edition_billing_resource_check("documents")
|
||||
@console_ns.doc(consumes=["multipart/form-data"], params=_FILE_UPLOAD_PARAMS)
|
||||
@console_ns.doc(consumes=["multipart/form-data"], params=FILE_UPLOAD_PARAMS)
|
||||
@console_ns.response(201, "File uploaded successfully", console_ns.models[FileResponse.__name__])
|
||||
@with_current_user
|
||||
def post(self, current_user: Account):
|
||||
source_str = request.form.get("source")
|
||||
source: Literal["datasets"] | None = "datasets" if source_str == "datasets" else None
|
||||
|
||||
if "file" not in request.files:
|
||||
raise NoFileUploadedError()
|
||||
|
||||
if len(request.files) > 1:
|
||||
raise TooManyFilesError()
|
||||
file = request.files["file"]
|
||||
|
||||
if not file.filename:
|
||||
raise FilenameNotExistsError
|
||||
if source == "datasets" and not current_user.is_dataset_editor:
|
||||
raise Forbidden()
|
||||
|
||||
if source not in ("datasets", None):
|
||||
source = None
|
||||
|
||||
try:
|
||||
upload_file = FileService(db.engine).upload_file(
|
||||
filename=file.filename,
|
||||
content=file.stream.read(),
|
||||
mimetype=file.mimetype,
|
||||
user=current_user,
|
||||
source=source,
|
||||
)
|
||||
except services.errors.file.FileTooLargeError as file_too_large_error:
|
||||
raise FileTooLargeError(file_too_large_error.description)
|
||||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
except services.errors.file.BlockedFileExtensionError as blocked_extension_error:
|
||||
raise BlockedFileExtensionError(blocked_extension_error.description)
|
||||
upload_file = upload_file_from_request(current_user=current_user)
|
||||
|
||||
return dump_response(FileResponse, upload_file), 201
|
||||
|
||||
|
||||
@@ -0,0 +1,332 @@
|
||||
"""Authenticated transport adapter for the Console-to-KnowledgeFS proxy.
|
||||
|
||||
These raw Blueprint routes deliberately stay outside Dify's OpenAPI surface:
|
||||
KnowledgeFS owns the wire contract consumed by the frontend. The catch-all path
|
||||
avoids resource-specific Dify controllers, while the forwarding module consumes
|
||||
only the operations explicitly enabled by Dify's product registry. The registry
|
||||
can be validated explicitly against the pinned KnowledgeFS contract during development.
|
||||
Console auth and contract-specific dataset RBAC run before forwarding. Request
|
||||
bodies are capped at 64 MiB, JSON and binary responses have separate bounds,
|
||||
SSE responses remain streaming with a bounded idle read timeout, and only safe
|
||||
response headers are exposed. Upstream 401 responses become 502 so they cannot
|
||||
trigger Dify browser-session recovery; resource-level 403 responses remain 403.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable, Iterator
|
||||
from functools import wraps
|
||||
from http import HTTPStatus
|
||||
from typing import NoReturn, cast
|
||||
|
||||
import httpx
|
||||
from flask import Response, request, stream_with_context
|
||||
from flask.typing import ResponseReturnValue
|
||||
from werkzeug.exceptions import (
|
||||
BadGateway,
|
||||
Forbidden,
|
||||
GatewayTimeout,
|
||||
NotFound,
|
||||
RequestEntityTooLarge,
|
||||
ServiceUnavailable,
|
||||
)
|
||||
|
||||
from configs import dify_config
|
||||
from controllers.console import api, bp
|
||||
from controllers.console.wraps import (
|
||||
account_initialization_required,
|
||||
cloud_edition_billing_rate_limit_check,
|
||||
setup_required,
|
||||
)
|
||||
from core.helper import ssrf_proxy
|
||||
from libs.login import current_account_with_tenant, login_required
|
||||
from services.knowledge_fs_proxy import (
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSMethod,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSTimeoutError,
|
||||
KnowledgeFSTransportError,
|
||||
KnowledgeFSUpstreamResponse,
|
||||
authorize_knowledge_fs_request,
|
||||
get_knowledge_fs_operation,
|
||||
proxy_knowledge_fs_request,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MAX_PROXY_BODY_BYTES = 64 * 1024 * 1024
|
||||
_RESPONSE_HEADER_ALLOWLIST = (
|
||||
"Cache-Control",
|
||||
"Content-Disposition",
|
||||
"Content-Type",
|
||||
"Retry-After",
|
||||
"X-Trace-Id",
|
||||
)
|
||||
_RESPONSE_HEADER_DENYLIST = frozenset(
|
||||
{
|
||||
"authorization",
|
||||
"connection",
|
||||
"cookie",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
"set-cookie",
|
||||
"te",
|
||||
"trailer",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _console_api_errors[**P](
|
||||
view: Callable[P, ResponseReturnValue],
|
||||
) -> Callable[P, ResponseReturnValue]:
|
||||
"""Route raw Blueprint exceptions through the Console API JSON handlers."""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs) -> ResponseReturnValue:
|
||||
try:
|
||||
return view(*args, **kwargs)
|
||||
except Exception as exc:
|
||||
return api.handle_error(exc)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def _knowledge_fs_enabled[**P](
|
||||
view: Callable[P, ResponseReturnValue],
|
||||
) -> Callable[P, ResponseReturnValue]:
|
||||
"""Hide the complete KnowledgeFS route surface while the bridge is disabled."""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs) -> ResponseReturnValue:
|
||||
if not dify_config.KNOWLEDGE_FS_ENABLED:
|
||||
raise NotFound()
|
||||
return view(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def _translate_proxy_error(exc: Exception, *, tenant_id: str) -> NoReturn:
|
||||
"""Map forwarding failures to the stable Console HTTP error surface."""
|
||||
if isinstance(exc, KnowledgeFSRouteNotAllowedError):
|
||||
raise NotFound() from exc
|
||||
if isinstance(exc, KnowledgeFSAccessDeniedError):
|
||||
raise Forbidden() from exc
|
||||
if isinstance(exc, KnowledgeFSConfigurationError):
|
||||
logger.error("KnowledgeFS request was blocked by invalid configuration for tenant_id=%s", tenant_id)
|
||||
raise ServiceUnavailable("KnowledgeFS integration is misconfigured") from exc
|
||||
if isinstance(exc, KnowledgeFSTimeoutError):
|
||||
raise GatewayTimeout("KnowledgeFS request timed out") from exc
|
||||
if isinstance(exc, KnowledgeFSTransportError):
|
||||
logger.warning("KnowledgeFS transport request failed for tenant_id=%s", tenant_id)
|
||||
raise BadGateway("KnowledgeFS is unavailable") from exc
|
||||
raise exc
|
||||
|
||||
|
||||
def _knowledge_fs_operation_access_required(
|
||||
view: Callable[[KnowledgeFSMethod, str], ResponseReturnValue],
|
||||
) -> Callable[[KnowledgeFSMethod, str], ResponseReturnValue]:
|
||||
"""Authorize one declared operation before billing and request-body work."""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(method: KnowledgeFSMethod, upstream_path: str) -> ResponseReturnValue:
|
||||
try:
|
||||
operation = get_knowledge_fs_operation(method, upstream_path)
|
||||
except KnowledgeFSRouteNotAllowedError as exc:
|
||||
raise NotFound() from exc
|
||||
|
||||
current_user, tenant_id = current_account_with_tenant()
|
||||
try:
|
||||
authorize_knowledge_fs_request(
|
||||
account=current_user,
|
||||
tenant_id=tenant_id,
|
||||
operation=operation,
|
||||
)
|
||||
except KnowledgeFSAccessDeniedError as exc:
|
||||
_translate_proxy_error(exc, tenant_id=tenant_id)
|
||||
return view(method, upstream_path)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def _request_body() -> bytes:
|
||||
"""Read the raw body up to the proxy limit or raise RequestEntityTooLarge."""
|
||||
body = request.stream.read(_MAX_PROXY_BODY_BYTES + 1)
|
||||
if len(body) > _MAX_PROXY_BODY_BYTES:
|
||||
raise RequestEntityTooLarge("KnowledgeFS proxy request body is too large")
|
||||
return body
|
||||
|
||||
|
||||
def _stream_response_body(
|
||||
upstream: httpx.Response,
|
||||
*,
|
||||
tenant_id: str,
|
||||
max_response_bytes: int,
|
||||
) -> Iterator[bytes]:
|
||||
"""Yield one bounded SSE response and always release its pooled connection."""
|
||||
total_bytes = 0
|
||||
try:
|
||||
for chunk in upstream.iter_bytes():
|
||||
total_bytes += len(chunk)
|
||||
if total_bytes > max_response_bytes:
|
||||
logger.warning("KnowledgeFS stream exceeded the proxy limit for tenant_id=%s", tenant_id)
|
||||
raise ssrf_proxy.ResponseTooLargeError(f"response exceeded {max_response_bytes} bytes")
|
||||
yield chunk
|
||||
finally:
|
||||
upstream.close()
|
||||
|
||||
|
||||
def _proxy_response(
|
||||
upstream_result: KnowledgeFSUpstreamResponse,
|
||||
*,
|
||||
tenant_id: str,
|
||||
contract_response_headers: tuple[str, ...],
|
||||
max_response_bytes: int,
|
||||
) -> Response:
|
||||
"""Expose raw content, status, and allowlisted headers from KnowledgeFS.
|
||||
|
||||
Raises:
|
||||
BadGateway: KnowledgeFS rejects the configured server credential.
|
||||
Forbidden: KnowledgeFS denies the account access to the requested resource.
|
||||
"""
|
||||
upstream = upstream_result.response
|
||||
if upstream.status_code == HTTPStatus.UNAUTHORIZED:
|
||||
upstream.close()
|
||||
logger.error(
|
||||
"KnowledgeFS rejected the Dify server credential with HTTP %s for tenant_id=%s",
|
||||
upstream.status_code,
|
||||
tenant_id,
|
||||
)
|
||||
raise BadGateway("KnowledgeFS authentication failed")
|
||||
if upstream.status_code == HTTPStatus.FORBIDDEN:
|
||||
upstream.close()
|
||||
raise Forbidden()
|
||||
|
||||
allowed_header_names = dict.fromkeys(
|
||||
name.lower() for name in (*_RESPONSE_HEADER_ALLOWLIST, *contract_response_headers)
|
||||
)
|
||||
headers = {
|
||||
name: value
|
||||
for name in allowed_header_names
|
||||
if name not in _RESPONSE_HEADER_DENYLIST
|
||||
if (value := upstream.headers.get(name)) is not None
|
||||
}
|
||||
if upstream_result.response_kind == "stream":
|
||||
response = Response(
|
||||
stream_with_context( # pyrefly: ignore[no-matching-overload]
|
||||
_stream_response_body(
|
||||
upstream,
|
||||
tenant_id=tenant_id,
|
||||
max_response_bytes=max_response_bytes,
|
||||
)
|
||||
),
|
||||
status=upstream.status_code,
|
||||
headers=headers,
|
||||
)
|
||||
response.call_on_close(upstream.close)
|
||||
return response
|
||||
|
||||
try:
|
||||
content = upstream.content
|
||||
finally:
|
||||
upstream.close()
|
||||
return Response(content, status=upstream.status_code, headers=headers)
|
||||
|
||||
|
||||
def _proxy_request(method: KnowledgeFSMethod, upstream_path: str) -> Response:
|
||||
"""Forward the current raw request and return its filtered upstream response.
|
||||
|
||||
The call performs one outbound KnowledgeFS request. Integration failures are
|
||||
converted to Console HTTP exceptions for the outer JSON error adapter.
|
||||
"""
|
||||
if not dify_config.KNOWLEDGE_FS_ENABLED:
|
||||
raise NotFound()
|
||||
current_user, tenant_id = current_account_with_tenant()
|
||||
try:
|
||||
proxy_result = proxy_knowledge_fs_request(
|
||||
account=current_user,
|
||||
method=method,
|
||||
path=upstream_path,
|
||||
tenant_id=tenant_id,
|
||||
accept=request.headers.get("Accept"),
|
||||
content_type=request.content_type,
|
||||
query=request.query_string or None,
|
||||
body=_request_body() if method != "GET" else None,
|
||||
request_headers=request.headers,
|
||||
)
|
||||
except (
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSTimeoutError,
|
||||
KnowledgeFSTransportError,
|
||||
) as exc:
|
||||
_translate_proxy_error(exc, tenant_id=tenant_id)
|
||||
return _proxy_response(
|
||||
proxy_result,
|
||||
tenant_id=tenant_id,
|
||||
contract_response_headers=proxy_result.operation.response_headers,
|
||||
max_response_bytes=proxy_result.operation.max_response_bytes,
|
||||
)
|
||||
|
||||
|
||||
@_knowledge_fs_enabled
|
||||
@_knowledge_fs_operation_access_required
|
||||
@cloud_edition_billing_rate_limit_check("knowledge")
|
||||
def _proxy_knowledge_fs_non_get(
|
||||
method: KnowledgeFSMethod,
|
||||
upstream_path: str,
|
||||
) -> ResponseReturnValue:
|
||||
"""Apply knowledge billing checks to one allowlisted non-GET operation."""
|
||||
return _proxy_request(method, upstream_path)
|
||||
|
||||
|
||||
@bp.route(
|
||||
"/knowledge-fs/<path:upstream_path>",
|
||||
methods=["GET", "OPTIONS"],
|
||||
provide_automatic_options=False,
|
||||
)
|
||||
@_console_api_errors
|
||||
@_knowledge_fs_enabled
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def proxy_knowledge_fs_get(upstream_path: str) -> ResponseReturnValue:
|
||||
"""Forward one authenticated, dataset-readable GET request.
|
||||
|
||||
Args:
|
||||
upstream_path: Relative KFS path captured after the Console proxy prefix.
|
||||
|
||||
Returns:
|
||||
The filtered raw KnowledgeFS response or a Console JSON error response.
|
||||
"""
|
||||
if request.method != "GET":
|
||||
raise NotFound()
|
||||
return _proxy_request("GET", upstream_path)
|
||||
|
||||
|
||||
@bp.route(
|
||||
"/knowledge-fs/<path:upstream_path>",
|
||||
methods=["DELETE", "PATCH", "POST", "PUT"],
|
||||
provide_automatic_options=False,
|
||||
)
|
||||
@_console_api_errors
|
||||
@_knowledge_fs_enabled
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def proxy_knowledge_fs_write(upstream_path: str) -> ResponseReturnValue:
|
||||
"""Forward one authenticated non-GET request under its contract access policy.
|
||||
|
||||
Args:
|
||||
upstream_path: Relative KFS path captured after the Console proxy prefix.
|
||||
|
||||
Returns:
|
||||
The filtered raw KnowledgeFS response or a Console JSON error response.
|
||||
"""
|
||||
method = cast(KnowledgeFSMethod, request.method)
|
||||
return _proxy_knowledge_fs_non_get(method, upstream_path)
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Console onboarding APIs.
|
||||
|
||||
This module keeps Step-by-step Tour persistence account-scoped. Workspace IDs
|
||||
are accepted only as presentation overrides; UI-only state such as minimized
|
||||
panels or the currently active task stays on the frontend. PATCH requests are
|
||||
action-based so callers do not replace server-side arrays with stale snapshots.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Literal, cast
|
||||
|
||||
from flask_restx import Resource
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
from controllers.common.schema import register_response_schema_models, register_schema_models
|
||||
from extensions.ext_database import db
|
||||
from fields.base import ResponseModel
|
||||
from libs.helper import dump_response
|
||||
from libs.login import login_required
|
||||
from models import Account
|
||||
from services.step_by_step_tour_service import StepByStepTourPatch, StepByStepTourService
|
||||
|
||||
from . import console_ns
|
||||
from .wraps import account_initialization_required, setup_required, with_current_tenant_id, with_current_user
|
||||
|
||||
StepByStepTourAction = Literal[
|
||||
"skip",
|
||||
"complete_task",
|
||||
"uncomplete_task",
|
||||
"enable_current_workspace",
|
||||
"disable_current_workspace",
|
||||
]
|
||||
StepByStepTourTaskId = Literal["home", "studio", "knowledge", "integration"]
|
||||
|
||||
|
||||
class StepByStepTourStatePatchPayload(BaseModel):
|
||||
action: StepByStepTourAction = Field(description="State update action")
|
||||
task_id: StepByStepTourTaskId | None = Field(default=None, description="Task ID for task actions")
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_patch_shape(self) -> "StepByStepTourStatePatchPayload":
|
||||
task_actions = {"complete_task", "uncomplete_task"}
|
||||
if self.action in task_actions and self.task_id is None:
|
||||
raise ValueError("task_id is required for task actions")
|
||||
if self.action not in task_actions and self.task_id is not None:
|
||||
raise ValueError("task_id is only supported for task actions")
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class StepByStepTourStateResponse(ResponseModel):
|
||||
first_workspace_id: str | None = None
|
||||
skipped: bool = False
|
||||
completed_task_ids: list[StepByStepTourTaskId] = Field(default_factory=list)
|
||||
manually_enabled_workspace_ids: list[str] = Field(default_factory=list)
|
||||
manually_disabled_workspace_ids: list[str] = Field(default_factory=list)
|
||||
updated_at: datetime | None = None
|
||||
|
||||
|
||||
register_schema_models(console_ns, StepByStepTourStatePatchPayload)
|
||||
register_response_schema_models(console_ns, StepByStepTourStateResponse)
|
||||
|
||||
|
||||
@console_ns.route("/onboarding/step-by-step-tour/state")
|
||||
class StepByStepTourStateApi(Resource):
|
||||
@console_ns.doc("get_step_by_step_tour_state")
|
||||
@console_ns.doc(description="Get account-level Step-by-step Tour state")
|
||||
@console_ns.response(200, "Success", console_ns.models[StepByStepTourStateResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def get(self, current_tenant_id: str, current_user: Account):
|
||||
return dump_response(
|
||||
StepByStepTourStateResponse,
|
||||
StepByStepTourService.get_state(
|
||||
account=current_user,
|
||||
current_tenant_id=current_tenant_id,
|
||||
session=db.session,
|
||||
),
|
||||
)
|
||||
|
||||
@console_ns.doc("patch_step_by_step_tour_state")
|
||||
@console_ns.doc(description="Update account-level Step-by-step Tour state")
|
||||
@console_ns.expect(console_ns.models[StepByStepTourStatePatchPayload.__name__])
|
||||
@console_ns.response(200, "Success", console_ns.models[StepByStepTourStateResponse.__name__])
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@with_current_user
|
||||
@with_current_tenant_id
|
||||
def patch(self, current_tenant_id: str, current_user: Account):
|
||||
payload = StepByStepTourStatePatchPayload.model_validate(console_ns.payload or {})
|
||||
patch = cast(StepByStepTourPatch, payload.model_dump(exclude_unset=True, exclude_none=True))
|
||||
return dump_response(
|
||||
StepByStepTourStateResponse,
|
||||
StepByStepTourService.patch_state(
|
||||
account=current_user,
|
||||
current_tenant_id=current_tenant_id,
|
||||
patch=patch,
|
||||
session=db.session,
|
||||
),
|
||||
)
|
||||
@@ -46,6 +46,61 @@ class GetRemoteFileInfo(Resource):
|
||||
).model_dump(mode="json")
|
||||
|
||||
|
||||
def upload_remote_file_from_request(
|
||||
*,
|
||||
current_user: Account,
|
||||
resource_tenant_id: str | None = None,
|
||||
) -> FileWithSignedUrl:
|
||||
"""Validate the JSON request, fetch its remote file, and persist it under the requested tenant."""
|
||||
payload = RemoteFileUploadPayload.model_validate(console_ns.payload)
|
||||
url = payload.url
|
||||
|
||||
# Try to fetch remote file metadata/content first
|
||||
try:
|
||||
resp = remote_fetcher.make_request("HEAD", url=url)
|
||||
if resp.status_code != httpx.codes.OK:
|
||||
resp = remote_fetcher.make_request("GET", url=url, timeout=3, follow_redirects=True)
|
||||
if resp.status_code != httpx.codes.OK:
|
||||
# Normalize into a user-friendly error message expected by tests
|
||||
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}")
|
||||
except httpx.RequestError as e:
|
||||
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}")
|
||||
|
||||
file_info = helpers.guess_file_info_from_response(resp)
|
||||
|
||||
# Enforce file size limit with 400 (Bad Request) per tests' expectation
|
||||
if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):
|
||||
raise FileTooLargeError()
|
||||
|
||||
# Load content if needed
|
||||
content = resp.content if resp.request.method == "GET" else remote_fetcher.make_request("GET", url).content
|
||||
|
||||
try:
|
||||
upload_file = FileService(db.engine).upload_file(
|
||||
filename=file_info.filename,
|
||||
content=content,
|
||||
mimetype=file_info.mimetype,
|
||||
user=current_user,
|
||||
tenant_id=resource_tenant_id,
|
||||
source_url=url,
|
||||
)
|
||||
except services.errors.file.FileTooLargeError as file_too_large_error:
|
||||
raise FileTooLargeError(file_too_large_error.description)
|
||||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
return FileWithSignedUrl(
|
||||
id=upload_file.id,
|
||||
name=upload_file.name,
|
||||
size=upload_file.size,
|
||||
extension=upload_file.extension,
|
||||
url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
|
||||
mime_type=upload_file.mime_type,
|
||||
created_by=upload_file.created_by,
|
||||
created_at=int(upload_file.created_at.timestamp()),
|
||||
)
|
||||
|
||||
|
||||
@console_ns.route("/remote-files/upload")
|
||||
class RemoteFileUpload(Resource):
|
||||
@console_ns.expect(console_ns.models[RemoteFileUploadPayload.__name__])
|
||||
@@ -53,53 +108,8 @@ class RemoteFileUpload(Resource):
|
||||
@login_required
|
||||
@with_current_user
|
||||
def post(self, current_user: Account):
|
||||
payload = RemoteFileUploadPayload.model_validate(console_ns.payload)
|
||||
url = payload.url
|
||||
|
||||
# Try to fetch remote file metadata/content first
|
||||
try:
|
||||
resp = remote_fetcher.make_request("HEAD", url=url)
|
||||
if resp.status_code != httpx.codes.OK:
|
||||
resp = remote_fetcher.make_request("GET", url=url, timeout=3, follow_redirects=True)
|
||||
if resp.status_code != httpx.codes.OK:
|
||||
# Normalize into a user-friendly error message expected by tests
|
||||
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}")
|
||||
except httpx.RequestError as e:
|
||||
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}")
|
||||
|
||||
file_info = helpers.guess_file_info_from_response(resp)
|
||||
|
||||
# Enforce file size limit with 400 (Bad Request) per tests' expectation
|
||||
if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):
|
||||
raise FileTooLargeError()
|
||||
|
||||
# Load content if needed
|
||||
content = resp.content if resp.request.method == "GET" else remote_fetcher.make_request("GET", url).content
|
||||
|
||||
try:
|
||||
upload_file = FileService(db.engine).upload_file(
|
||||
filename=file_info.filename,
|
||||
content=content,
|
||||
mimetype=file_info.mimetype,
|
||||
user=current_user,
|
||||
source_url=url,
|
||||
)
|
||||
except services.errors.file.FileTooLargeError as file_too_large_error:
|
||||
raise FileTooLargeError(file_too_large_error.description)
|
||||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
# Success: return created resource with 201 status
|
||||
remote_file = upload_remote_file_from_request(current_user=current_user)
|
||||
return (
|
||||
FileWithSignedUrl(
|
||||
id=upload_file.id,
|
||||
name=upload_file.name,
|
||||
size=upload_file.size,
|
||||
extension=upload_file.extension,
|
||||
url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
|
||||
mime_type=upload_file.mime_type,
|
||||
created_by=upload_file.created_by,
|
||||
created_at=int(upload_file.created_at.timestamp()),
|
||||
).model_dump(mode="json"),
|
||||
remote_file.model_dump(mode="json"),
|
||||
201,
|
||||
)
|
||||
|
||||
@@ -98,6 +98,7 @@ def handle_collaboration_event(sid, data):
|
||||
6. workflow_update
|
||||
7. comments_update
|
||||
8. node_panel_presence
|
||||
9. graph_view_state (session reports tab visibility; drives leader election)
|
||||
"""
|
||||
return collaboration_service.relay_collaboration_event(sid, data)
|
||||
|
||||
|
||||
@@ -4,20 +4,16 @@ from http import HTTPStatus
|
||||
from flask import redirect
|
||||
from flask_restx import Resource
|
||||
from pydantic import BaseModel, Field
|
||||
from werkzeug.exceptions import Conflict, NotFound
|
||||
from werkzeug.exceptions import Conflict, Forbidden, NotFound
|
||||
|
||||
from controllers.common.fields import RedirectResponse
|
||||
from controllers.common.schema import register_response_schema_models, register_schema_models
|
||||
from controllers.console import console_ns
|
||||
from controllers.console.wraps import (
|
||||
RBACPermission,
|
||||
RBACResourceScope,
|
||||
account_initialization_required,
|
||||
cloud_edition_billing_enabled,
|
||||
cloud_edition_billing_paid_plan_required,
|
||||
is_admin_or_owner_required,
|
||||
only_edition_cloud,
|
||||
rbac_permission_required,
|
||||
setup_required,
|
||||
)
|
||||
from extensions.ext_database import db
|
||||
@@ -25,6 +21,7 @@ from fields.base import ResponseModel
|
||||
from libs.archive_storage import get_export_storage
|
||||
from libs.helper import dump_response
|
||||
from libs.login import current_account_with_tenant, login_required
|
||||
from models import TenantAccountRole
|
||||
from services.retention.workflow_run.archive_download_preparation import ARCHIVE_DOWNLOAD_MIME_TYPE
|
||||
from services.retention.workflow_run.archive_download_task_cache import (
|
||||
WorkflowRunArchiveDownloadStatus,
|
||||
@@ -98,11 +95,13 @@ register_response_schema_models(
|
||||
)
|
||||
|
||||
|
||||
def _current_ids() -> tuple[str, str]:
|
||||
"""Return current `(tenant_id, account_id)` or raise when no workspace is selected."""
|
||||
def _current_owner_or_admin_ids() -> tuple[str, str]:
|
||||
"""Return current Cloud workspace IDs for an owner or admin, independently of enterprise RBAC."""
|
||||
current_user, current_tenant_id = current_account_with_tenant()
|
||||
if not current_tenant_id:
|
||||
raise NotFound("Current workspace not found")
|
||||
if not TenantAccountRole.is_privileged_role(current_user.current_role):
|
||||
raise Forbidden()
|
||||
return current_tenant_id, current_user.id
|
||||
|
||||
|
||||
@@ -124,12 +123,8 @@ class WorkflowRunArchivesApi(Resource):
|
||||
@only_edition_cloud
|
||||
@cloud_edition_billing_enabled
|
||||
@cloud_edition_billing_paid_plan_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(
|
||||
RBACResourceScope.WORKSPACE, RBACPermission.WORKSPACE_ROLE_MANAGE, resource_required=False
|
||||
)
|
||||
def get(self):
|
||||
tenant_id, _ = _current_ids()
|
||||
tenant_id, _ = _current_owner_or_admin_ids()
|
||||
return dump_response(WorkflowRunArchiveListResponse, list_workflow_run_archives(db.session(), tenant_id))
|
||||
|
||||
|
||||
@@ -149,12 +144,8 @@ class WorkflowRunArchiveDownloadsApi(Resource):
|
||||
@only_edition_cloud
|
||||
@cloud_edition_billing_enabled
|
||||
@cloud_edition_billing_paid_plan_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(
|
||||
RBACResourceScope.WORKSPACE, RBACPermission.WORKSPACE_ROLE_MANAGE, resource_required=False
|
||||
)
|
||||
def post(self):
|
||||
tenant_id, account_id = _current_ids()
|
||||
tenant_id, account_id = _current_owner_or_admin_ids()
|
||||
payload = WorkflowRunArchiveDownloadPayload.model_validate(console_ns.payload or {})
|
||||
try:
|
||||
task = create_workflow_run_archive_download_task(
|
||||
@@ -180,12 +171,8 @@ class WorkflowRunArchiveDownloadApi(Resource):
|
||||
@only_edition_cloud
|
||||
@cloud_edition_billing_enabled
|
||||
@cloud_edition_billing_paid_plan_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(
|
||||
RBACResourceScope.WORKSPACE, RBACPermission.WORKSPACE_ROLE_MANAGE, resource_required=False
|
||||
)
|
||||
def get(self, download_id: str):
|
||||
tenant_id, _ = _current_ids()
|
||||
tenant_id, _ = _current_owner_or_admin_ids()
|
||||
try:
|
||||
task = get_workflow_run_archive_download_task(tenant_id=tenant_id, download_id=download_id)
|
||||
except WorkflowRunArchiveDownloadTaskNotFoundError as exc:
|
||||
@@ -209,12 +196,8 @@ class WorkflowRunArchiveDownloadFileApi(Resource):
|
||||
@only_edition_cloud
|
||||
@cloud_edition_billing_enabled
|
||||
@cloud_edition_billing_paid_plan_required
|
||||
@is_admin_or_owner_required
|
||||
@rbac_permission_required(
|
||||
RBACResourceScope.WORKSPACE, RBACPermission.WORKSPACE_ROLE_MANAGE, resource_required=False
|
||||
)
|
||||
def get(self, download_id: str):
|
||||
tenant_id, _ = _current_ids()
|
||||
tenant_id, _ = _current_owner_or_admin_ids()
|
||||
try:
|
||||
task = get_ready_workflow_run_archive_download_task(tenant_id=tenant_id, download_id=download_id)
|
||||
except WorkflowRunArchiveDownloadTaskNotFoundError as exc:
|
||||
|
||||
@@ -101,7 +101,7 @@ class AudioApi(Resource):
|
||||
|
||||
Accepts an audio file upload and returns the transcribed text.
|
||||
"""
|
||||
file = request.files["file"]
|
||||
file = request.files.get("file")
|
||||
|
||||
try:
|
||||
response = AudioService.transcript_asr(
|
||||
|
||||
@@ -531,14 +531,22 @@ class DatasetListApi(DatasetApiResource):
|
||||
except services.errors.dataset.DatasetNameDuplicateError:
|
||||
raise DatasetNameDuplicateError()
|
||||
|
||||
if payload.permission == DatasetPermissionEnum.ALL_TEAM and dify_config.RBAC_ENABLED:
|
||||
RBACService.DatasetAccess.replace_whitelist(
|
||||
tenant_id,
|
||||
current_user.id,
|
||||
dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_created_app_rbac_access_task.delay(tenant_id, current_user.id, dataset_id=dataset.id)
|
||||
if dify_config.RBAC_ENABLED:
|
||||
if payload.permission == DatasetPermissionEnum.ALL_TEAM:
|
||||
RBACService.DatasetAccess.replace_whitelist(
|
||||
tenant_id,
|
||||
current_user.id,
|
||||
dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL),
|
||||
)
|
||||
initialize_created_app_rbac_access_task.delay(tenant_id, current_user.id, dataset_id=dataset.id)
|
||||
else:
|
||||
RBACService.DatasetAccess.replace_whitelist(
|
||||
tenant_id,
|
||||
current_user.id,
|
||||
dataset.id,
|
||||
ReplaceMemberBindings(scope=RBACResourceWhitelistScope.SPECIFIC),
|
||||
)
|
||||
|
||||
return _dump_service_dataset_detail(dataset, session=session), 200
|
||||
|
||||
|
||||
@@ -76,7 +76,7 @@ class AudioApi(WebApiResource):
|
||||
@web_ns.response(200, "Success", web_ns.models[AudioToTextResponse.__name__])
|
||||
def post(self, app_model: App, end_user: EndUser):
|
||||
"""Convert audio to text"""
|
||||
file = request.files["file"]
|
||||
file = request.files.get("file")
|
||||
|
||||
try:
|
||||
response = AudioService.transcript_asr(
|
||||
|
||||
@@ -625,7 +625,11 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
message=message_snapshot,
|
||||
user=user,
|
||||
stream=stream,
|
||||
draft_var_saver_factory=self._get_draft_var_saver_factory(invoke_from, account=user),
|
||||
draft_var_saver_factory=self._get_draft_var_saver_factory(
|
||||
invoke_from,
|
||||
account=user,
|
||||
tenant_id=application_generate_entity.app_config.tenant_id,
|
||||
),
|
||||
)
|
||||
|
||||
return AdvancedChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
|
||||
|
||||
@@ -32,6 +32,7 @@ class _DebuggerDraftVariableSaver:
|
||||
self,
|
||||
*,
|
||||
account: Account,
|
||||
tenant_id: str,
|
||||
app_id: str,
|
||||
node_id: str,
|
||||
node_type: NodeType,
|
||||
@@ -39,6 +40,7 @@ class _DebuggerDraftVariableSaver:
|
||||
enclosing_node_id: str | None = None,
|
||||
) -> None:
|
||||
self._account = account
|
||||
self._tenant_id = tenant_id
|
||||
self._app_id = app_id
|
||||
self._node_id = node_id
|
||||
self._node_type = node_type
|
||||
@@ -49,6 +51,7 @@ class _DebuggerDraftVariableSaver:
|
||||
with Session(db.engine) as session, session.begin():
|
||||
DraftVariableSaverImpl(
|
||||
session=session,
|
||||
tenant_id=self._tenant_id,
|
||||
app_id=self._app_id,
|
||||
node_id=self._node_id,
|
||||
node_type=self._node_type,
|
||||
@@ -287,7 +290,12 @@ class BaseAppGenerator:
|
||||
|
||||
@final
|
||||
@staticmethod
|
||||
def _get_draft_var_saver_factory(invoke_from: InvokeFrom, account: Account | EndUser) -> DraftVariableSaverFactory:
|
||||
def _get_draft_var_saver_factory(
|
||||
invoke_from: InvokeFrom,
|
||||
account: Account | EndUser,
|
||||
*,
|
||||
tenant_id: str,
|
||||
) -> DraftVariableSaverFactory:
|
||||
if invoke_from == InvokeFrom.DEBUGGER:
|
||||
assert isinstance(account, Account)
|
||||
|
||||
@@ -300,6 +308,7 @@ class BaseAppGenerator:
|
||||
) -> DraftVariableSaver:
|
||||
return _DebuggerDraftVariableSaver(
|
||||
account=account,
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
node_id=node_id,
|
||||
node_type=node_type,
|
||||
|
||||
@@ -349,6 +349,7 @@ class PipelineGenerator(BaseAppGenerator):
|
||||
draft_var_saver_factory = self._get_draft_var_saver_factory(
|
||||
invoke_from,
|
||||
user,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
)
|
||||
# return response or stream generator
|
||||
response = self._handle_response(
|
||||
|
||||
@@ -399,7 +399,11 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
|
||||
worker_thread.start()
|
||||
|
||||
draft_var_saver_factory = self._get_draft_var_saver_factory(invoke_from, user)
|
||||
draft_var_saver_factory = self._get_draft_var_saver_factory(
|
||||
invoke_from,
|
||||
user,
|
||||
tenant_id=app_model.tenant_id,
|
||||
)
|
||||
|
||||
# return response or stream generator
|
||||
response = self._handle_response(
|
||||
|
||||
@@ -47,6 +47,24 @@ class MaxRetriesExceededError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class ResponseLimitError(ValueError):
|
||||
"""Base error for responses that cannot be safely bounded."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ResponseTooLargeError(ResponseLimitError):
|
||||
"""Raised when an identity response exceeds the configured byte limit."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class UnsupportedResponseEncodingError(ResponseLimitError):
|
||||
"""Raised when response encoding prevents safe decoded-size enforcement."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
request_error = httpx.RequestError
|
||||
max_retries_exceeded_error = MaxRetriesExceededError
|
||||
|
||||
@@ -142,7 +160,31 @@ def _inject_trace_headers(headers: Headers | None) -> Headers:
|
||||
return headers
|
||||
|
||||
|
||||
def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETRIES, **kwargs: Any) -> httpx.Response:
|
||||
def make_request(
|
||||
method: str,
|
||||
url: str,
|
||||
max_retries: int = SSRF_DEFAULT_MAX_RETRIES,
|
||||
stream_response: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> httpx.Response:
|
||||
"""Send one SSRF-protected request with optional streaming.
|
||||
|
||||
Args:
|
||||
method: HTTP method sent through the configured SSRF client.
|
||||
url: Absolute request URL.
|
||||
max_retries: Number of retry attempts after the initial request.
|
||||
stream_response: Return an open streaming response that the caller must close.
|
||||
**kwargs: Additional keyword arguments forwarded to ``httpx.Client``.
|
||||
|
||||
Returns:
|
||||
A buffered response, or an open response when ``stream_response`` is true.
|
||||
|
||||
Raises:
|
||||
ToolSSRFError: The configured SSRF proxy rejects the destination.
|
||||
MaxRetriesExceededError: All configured request attempts fail.
|
||||
httpx.RequestError: A request fails while retries are disabled.
|
||||
ValueError: The SSL verification option or request headers are invalid.
|
||||
"""
|
||||
# Convert requests-style allow_redirects to httpx-style follow_redirects
|
||||
if "allow_redirects" in kwargs:
|
||||
allow_redirects = kwargs.pop("allow_redirects")
|
||||
@@ -175,6 +217,11 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
# When using a forward proxy, httpx may override the Host header based on the URL.
|
||||
# We extract and preserve any explicitly set Host header to support virtual hosting.
|
||||
user_provided_host = _get_user_provided_host_header(headers)
|
||||
send_kwargs: dict[str, Any] = {}
|
||||
if "auth" in kwargs:
|
||||
send_kwargs["auth"] = kwargs.pop("auth")
|
||||
if "follow_redirects" in kwargs:
|
||||
send_kwargs["follow_redirects"] = kwargs.pop("follow_redirects")
|
||||
|
||||
retries = 0
|
||||
while retries <= max_retries:
|
||||
@@ -185,7 +232,11 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
if user_provided_host is not None:
|
||||
headers["host"] = user_provided_host
|
||||
kwargs["headers"] = headers
|
||||
response = client.request(method=method, url=url, **kwargs)
|
||||
request = client.build_request(method=method, url=url, **kwargs)
|
||||
if stream_response:
|
||||
response = client.send(request, stream=True, **send_kwargs)
|
||||
else:
|
||||
response = client.send(request, **send_kwargs)
|
||||
|
||||
# Check for SSRF protection by Squid proxy
|
||||
if response.status_code in (401, 403):
|
||||
@@ -195,6 +246,7 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
|
||||
# Squid typically identifies itself in Server or Via headers
|
||||
if "squid" in server_header or "squid" in via_header:
|
||||
response.close()
|
||||
raise ToolSSRFError(
|
||||
f"Access to '{url}' was blocked by SSRF protection. "
|
||||
f"The URL may point to a private or local network address. "
|
||||
@@ -208,6 +260,7 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
response.status_code,
|
||||
url,
|
||||
)
|
||||
response.close()
|
||||
|
||||
except httpx.RequestError as e:
|
||||
logger.warning("Request to URL %s failed on attempt %s: %s", url, retries + 1, e)
|
||||
@@ -220,6 +273,42 @@ def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETR
|
||||
raise MaxRetriesExceededError(f"Reached maximum retries ({max_retries}) for URL {url}")
|
||||
|
||||
|
||||
def buffer_response(response: httpx.Response, *, max_response_bytes: int) -> httpx.Response:
|
||||
"""Consume one open identity response under a decoded byte limit and close its stream."""
|
||||
if max_response_bytes <= 0:
|
||||
raise ValueError("max_response_bytes must be positive")
|
||||
|
||||
try:
|
||||
content_encoding = response.headers.get("content-encoding", "identity").strip().lower()
|
||||
if content_encoding not in {"", "identity"}:
|
||||
raise UnsupportedResponseEncodingError(f"content encoding {content_encoding} cannot be safely bounded")
|
||||
content = bytearray()
|
||||
for chunk in response.iter_bytes():
|
||||
if len(content) + len(chunk) > max_response_bytes:
|
||||
raise ResponseTooLargeError(f"response exceeded {max_response_bytes} bytes")
|
||||
content.extend(chunk)
|
||||
decoded_headers = {
|
||||
name: value
|
||||
for name, value in response.headers.items()
|
||||
if name.lower() not in {"content-encoding", "content-length", "transfer-encoding"}
|
||||
}
|
||||
try:
|
||||
request = response.request
|
||||
except RuntimeError:
|
||||
request = None
|
||||
return httpx.Response(
|
||||
response.status_code,
|
||||
headers=decoded_headers,
|
||||
content=bytes(content),
|
||||
request=request,
|
||||
extensions=response.extensions,
|
||||
history=response.history,
|
||||
default_encoding=response.default_encoding,
|
||||
)
|
||||
finally:
|
||||
response.close()
|
||||
|
||||
|
||||
def get(url: str, max_retries: int = SSRF_DEFAULT_MAX_RETRIES, **kwargs: Any) -> httpx.Response:
|
||||
return make_request("GET", url, max_retries=max_retries, **kwargs)
|
||||
|
||||
|
||||
@@ -131,7 +131,10 @@ class WaterCrawlAPIClient(BaseAPIClient):
|
||||
content_type = response.headers.get("Content-Type", "")
|
||||
media_type = content_type.split(";", 1)[0].strip().lower()
|
||||
if media_type == "application/json":
|
||||
return response.json() or {}
|
||||
try:
|
||||
return response.json() or {}
|
||||
except ValueError as exc:
|
||||
raise ValueError("Invalid JSON response from WaterCrawl") from exc
|
||||
|
||||
if media_type == "application/octet-stream":
|
||||
return response.content
|
||||
|
||||
@@ -1,5 +1,15 @@
|
||||
"""WaterCrawl domain exceptions.
|
||||
|
||||
These exceptions are constructed from upstream HTTP responses, which may be
|
||||
JSON API errors or plain text/HTML proxy errors. Keep the exception type stable
|
||||
even when the body is not JSON so callers can handle WaterCrawl failures by
|
||||
domain type instead of low-level parser errors.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import override
|
||||
from typing import Any, override
|
||||
|
||||
from httpx import Response
|
||||
|
||||
|
||||
class WaterCrawlError(Exception):
|
||||
@@ -7,11 +17,16 @@ class WaterCrawlError(Exception):
|
||||
|
||||
|
||||
class WaterCrawlBadRequestError(WaterCrawlError):
|
||||
def __init__(self, response):
|
||||
def __init__(self, response: Response):
|
||||
self.status_code = response.status_code
|
||||
self.response = response
|
||||
data = response.json()
|
||||
self.message = data.get("message", "Unknown error occurred")
|
||||
try:
|
||||
data: Any = response.json()
|
||||
except ValueError:
|
||||
data = {}
|
||||
if not isinstance(data, dict):
|
||||
data = {}
|
||||
self.message = data.get("message") or response.text or "Unknown error occurred"
|
||||
self.errors = data.get("errors", {})
|
||||
super().__init__(self.message)
|
||||
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
"""Validate Dify Console KnowledgeFS declarations against a pinned OpenAPI document.
|
||||
|
||||
The OpenAPI document is exported only during explicit development validation. Runtime declarations live with Dify
|
||||
product policy; this module validates their transport metadata without generating a complete operation catalog.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, TypedDict
|
||||
|
||||
API_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(API_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(API_ROOT))
|
||||
|
||||
WORKSPACE_ROOT = API_ROOT.parent
|
||||
LOCK_PATH = API_ROOT / "knowledge-fs-contract.lock.json"
|
||||
DEFAULT_REPOSITORY = WORKSPACE_ROOT.parent / "knowledge-fs"
|
||||
OPENAPI_METHODS = ("delete", "get", "head", "options", "patch", "post", "put", "trace")
|
||||
PROXY_METHODS = frozenset({"delete", "get", "patch", "post", "put"})
|
||||
|
||||
|
||||
class ContractDeclaration(TypedDict):
|
||||
"""KnowledgeFS transport contract declared by one Dify Console registry entry."""
|
||||
|
||||
operation_id: str
|
||||
method: str
|
||||
path: str
|
||||
required_scope: str | None
|
||||
response_kind: str
|
||||
max_response_bytes: int
|
||||
request_headers: tuple[str, ...]
|
||||
response_headers: tuple[str, ...]
|
||||
response_media_types: tuple[str, ...]
|
||||
|
||||
|
||||
type DeclarationField = Literal[
|
||||
"method",
|
||||
"path",
|
||||
"required_scope",
|
||||
"response_kind",
|
||||
"max_response_bytes",
|
||||
"request_headers",
|
||||
"response_headers",
|
||||
"response_media_types",
|
||||
]
|
||||
|
||||
DECLARATION_FIELDS: tuple[DeclarationField, ...] = (
|
||||
"method",
|
||||
"path",
|
||||
"required_scope",
|
||||
"response_kind",
|
||||
"max_response_bytes",
|
||||
"request_headers",
|
||||
"response_headers",
|
||||
"response_media_types",
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Update or verify the pin and validate Console declarations against its OpenAPI document."""
|
||||
parser = argparse.ArgumentParser()
|
||||
mode = parser.add_mutually_exclusive_group()
|
||||
mode.add_argument("--check", action="store_true")
|
||||
mode.add_argument("--update-lock", action="store_true")
|
||||
parser.add_argument("--repository", type=Path, default=DEFAULT_REPOSITORY)
|
||||
args = parser.parse_args()
|
||||
|
||||
repository = args.repository.resolve()
|
||||
lock = json.loads(LOCK_PATH.read_text())
|
||||
tracked_changes = run("git", "status", "--porcelain", "--untracked-files=no", cwd=repository).strip()
|
||||
if tracked_changes:
|
||||
raise RuntimeError("KnowledgeFS checkout must not contain tracked changes during contract export")
|
||||
|
||||
commit = run("git", "rev-parse", "HEAD", cwd=repository).strip()
|
||||
if not args.update_lock and commit != lock["commit"]:
|
||||
raise RuntimeError(
|
||||
f"KnowledgeFS checkout mismatch: expected {lock['commit']}, received {commit}. "
|
||||
"Use the pinned commit or pass --update-lock intentionally."
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="dify-knowledge-fs-contract-") as directory:
|
||||
openapi_path = Path(directory) / "knowledge-fs.openapi.json"
|
||||
subprocess.run(
|
||||
["pnpm", "openapi:export", "--", "--output", str(openapi_path)],
|
||||
cwd=repository,
|
||||
check=True,
|
||||
)
|
||||
openapi_content = openapi_path.read_bytes()
|
||||
|
||||
openapi_sha256 = sha256(openapi_content)
|
||||
if not args.update_lock and openapi_sha256 != lock["openapiSha256"]:
|
||||
raise RuntimeError(
|
||||
f"KnowledgeFS OpenAPI hash mismatch: expected {lock['openapiSha256']}, received {openapi_sha256}"
|
||||
)
|
||||
|
||||
document: dict[str, Any] = json.loads(openapi_content)
|
||||
validate_declarations(document, console_contract_declarations())
|
||||
|
||||
if args.update_lock:
|
||||
LOCK_PATH.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"commit": commit,
|
||||
"openapiSha256": openapi_sha256,
|
||||
"repository": lock["repository"],
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
|
||||
def validate_declarations(document: dict[str, Any], declarations: tuple[ContractDeclaration, ...]) -> None:
|
||||
"""Validate Dify Console declarations against matching pinned OpenAPI operations."""
|
||||
operations_by_id: dict[str, list[tuple[str, str, dict[str, Any], dict[str, Any]]]] = {}
|
||||
for path, path_item in document.get("paths", {}).items():
|
||||
for method in OPENAPI_METHODS:
|
||||
operation = path_item.get(method)
|
||||
if operation is None:
|
||||
continue
|
||||
operation_id = operation.get("operationId")
|
||||
if isinstance(operation_id, str) and operation_id:
|
||||
operations_by_id.setdefault(operation_id, []).append((method, path, path_item, operation))
|
||||
|
||||
declared_ids: set[str] = set()
|
||||
for declaration in declarations:
|
||||
operation_id = declaration["operation_id"]
|
||||
if operation_id in declared_ids:
|
||||
raise ValueError(f"Dify Console registry has duplicate operationId: {operation_id}")
|
||||
declared_ids.add(operation_id)
|
||||
|
||||
matches = operations_by_id.get(operation_id, [])
|
||||
if not matches:
|
||||
raise ValueError(f"KnowledgeFS OpenAPI has no operationId: {operation_id}")
|
||||
if len(matches) > 1:
|
||||
raise ValueError(f"KnowledgeFS OpenAPI has duplicate operationId: {operation_id}")
|
||||
|
||||
method, path, path_item, operation = matches[0]
|
||||
if not path.startswith("/"):
|
||||
raise ValueError(f"KnowledgeFS OpenAPI path must be absolute: {path}")
|
||||
if method not in PROXY_METHODS:
|
||||
raise ValueError(f"KnowledgeFS proxy does not support {method.upper()} {path}")
|
||||
expected: ContractDeclaration = {
|
||||
"operation_id": operation_id,
|
||||
"method": method.upper(),
|
||||
"path": path[1:],
|
||||
"required_scope": required_scope(operation),
|
||||
"response_kind": response_kind(operation),
|
||||
"max_response_bytes": required_max_response_bytes(operation),
|
||||
"request_headers": request_header_names(path_item, operation),
|
||||
"response_headers": response_header_names(operation),
|
||||
"response_media_types": response_media_types(operation),
|
||||
}
|
||||
for field in DECLARATION_FIELDS:
|
||||
expected_value = expected[field]
|
||||
received_value = declaration[field]
|
||||
if received_value != expected_value:
|
||||
raise ValueError(
|
||||
f"KnowledgeFS operation {operation_id} field {field} drifted: "
|
||||
f"expected {expected_value!r}, received {received_value!r}"
|
||||
)
|
||||
|
||||
|
||||
def console_contract_declarations() -> tuple[ContractDeclaration, ...]:
|
||||
"""Return transport declarations from the runtime Console operation registry."""
|
||||
from services.knowledge_fs_proxy import KNOWLEDGE_FS_CONSOLE_OPERATIONS
|
||||
|
||||
return tuple(
|
||||
{
|
||||
"operation_id": operation.operation_id,
|
||||
"method": operation.method,
|
||||
"path": operation.path,
|
||||
"required_scope": operation.required_scope,
|
||||
"response_kind": operation.response_kind,
|
||||
"max_response_bytes": operation.max_response_bytes,
|
||||
"request_headers": operation.request_headers,
|
||||
"response_headers": operation.response_headers,
|
||||
"response_media_types": operation.response_media_types,
|
||||
}
|
||||
for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS
|
||||
)
|
||||
|
||||
|
||||
def response_kind(operation: dict[str, Any]) -> str:
|
||||
media_types = response_media_types(operation)
|
||||
if "text/event-stream" in media_types:
|
||||
return "stream"
|
||||
if "application/octet-stream" in media_types:
|
||||
return "binary"
|
||||
return "buffered"
|
||||
|
||||
|
||||
def response_media_types(operation: dict[str, Any]) -> tuple[str, ...]:
|
||||
media_types: set[str] = set()
|
||||
for status, response in operation.get("responses", {}).items():
|
||||
if status == "2XX" or (len(status) == 3 and status.startswith("2") and status.isdigit()):
|
||||
media_types.update(response.get("content", {}))
|
||||
return tuple(sorted(media_types))
|
||||
|
||||
|
||||
def required_scope(operation: dict[str, Any]) -> str | None:
|
||||
scope = operation.get("x-knowledge-fs-required-scope")
|
||||
if scope in ("knowledge-spaces:read", "knowledge-spaces:write"):
|
||||
return scope
|
||||
if operation.get("security") == []:
|
||||
return None
|
||||
raise ValueError(f"KnowledgeFS operation has no supported required scope: {scope}")
|
||||
|
||||
|
||||
def required_max_response_bytes(operation: dict[str, Any]) -> int:
|
||||
value = operation.get("x-knowledge-fs-max-response-bytes")
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
||||
raise ValueError(f"KnowledgeFS operation has no valid response byte limit: {value}")
|
||||
return value
|
||||
|
||||
|
||||
def request_header_names(path_item: dict[str, Any], operation: dict[str, Any]) -> tuple[str, ...]:
|
||||
names: set[str] = set()
|
||||
for parameter in [*path_item.get("parameters", []), *operation.get("parameters", [])]:
|
||||
if "$ref" in parameter:
|
||||
raise ValueError(f"KnowledgeFS request header references are not supported: {parameter['$ref']}")
|
||||
if parameter.get("in") == "header":
|
||||
names.add(parameter["name"].lower())
|
||||
return tuple(sorted(names))
|
||||
|
||||
|
||||
def response_header_names(operation: dict[str, Any]) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
sorted(
|
||||
{
|
||||
name.lower()
|
||||
for response in operation.get("responses", {}).values()
|
||||
for name in response.get("headers", {})
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def sha256(content: bytes) -> str:
|
||||
return hashlib.sha256(content).hexdigest()
|
||||
|
||||
|
||||
def run(*command: str, cwd: Path) -> str:
|
||||
return subprocess.run(command, cwd=cwd, check=True, capture_output=True, text=True).stdout
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"commit": "4310e2d582d25e7de58183f27720afab01e123cf",
|
||||
"openapiSha256": "5827ca930ce38462bfd1b2bef387efbf37eb7ffcaedde4558af2fbaeccbfbc4b",
|
||||
"repository": "https://github.com/langgenius/knowledge-fs"
|
||||
}
|
||||
+33
-12
@@ -16,12 +16,19 @@ from urllib.parse import quote
|
||||
import boto3
|
||||
import orjson
|
||||
from botocore.client import Config
|
||||
from botocore.exceptions import ClientError
|
||||
from botocore.exceptions import BotoCoreError, ClientError
|
||||
|
||||
from configs import dify_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_OBJECT_NOT_FOUND_ERROR_CODES = frozenset({"404", "NoSuchKey", "NotFound"})
|
||||
|
||||
|
||||
def _is_object_not_found_error(error: ClientError) -> bool:
|
||||
error_code = str(error.response.get("Error", {}).get("Code", ""))
|
||||
return error_code in _OBJECT_NOT_FOUND_ERROR_CODES
|
||||
|
||||
|
||||
class ArchiveStorageError(Exception):
|
||||
"""Base exception for archive storage operations."""
|
||||
@@ -138,10 +145,11 @@ class ArchiveStorage:
|
||||
response = self.client.get_object(Bucket=self.bucket, Key=key)
|
||||
return response["Body"].read()
|
||||
except ClientError as e:
|
||||
error_code = e.response.get("Error", {}).get("Code")
|
||||
if error_code == "NoSuchKey":
|
||||
raise FileNotFoundError(f"Archive object not found: {key}")
|
||||
raise ArchiveStorageError(f"Failed to download object '{key}': {e}")
|
||||
if _is_object_not_found_error(e):
|
||||
raise FileNotFoundError(f"Archive object not found: {key}") from e
|
||||
raise ArchiveStorageError(f"Failed to download object '{key}': {e}") from e
|
||||
except BotoCoreError as e:
|
||||
raise ArchiveStorageError(f"Failed to download object '{key}': {e}") from e
|
||||
|
||||
def get_object_stream(self, key: str) -> Generator[bytes, None, None]:
|
||||
"""
|
||||
@@ -161,10 +169,11 @@ class ArchiveStorage:
|
||||
response = self.client.get_object(Bucket=self.bucket, Key=key)
|
||||
yield from response["Body"].iter_chunks()
|
||||
except ClientError as e:
|
||||
error_code = e.response.get("Error", {}).get("Code")
|
||||
if error_code == "NoSuchKey":
|
||||
raise FileNotFoundError(f"Archive object not found: {key}")
|
||||
raise ArchiveStorageError(f"Failed to stream object '{key}': {e}")
|
||||
if _is_object_not_found_error(e):
|
||||
raise FileNotFoundError(f"Archive object not found: {key}") from e
|
||||
raise ArchiveStorageError(f"Failed to stream object '{key}': {e}") from e
|
||||
except BotoCoreError as e:
|
||||
raise ArchiveStorageError(f"Failed to stream object '{key}': {e}") from e
|
||||
|
||||
def object_exists(self, key: str) -> bool:
|
||||
"""
|
||||
@@ -175,12 +184,19 @@ class ArchiveStorage:
|
||||
|
||||
Returns:
|
||||
True if object exists, False otherwise
|
||||
|
||||
Raises:
|
||||
ArchiveStorageError: If storage cannot authoritatively determine object existence
|
||||
"""
|
||||
try:
|
||||
self.client.head_object(Bucket=self.bucket, Key=key)
|
||||
return True
|
||||
except ClientError:
|
||||
return False
|
||||
except ClientError as e:
|
||||
if _is_object_not_found_error(e):
|
||||
return False
|
||||
raise ArchiveStorageError(f"Failed to check archive object '{key}': {e}") from e
|
||||
except BotoCoreError as e:
|
||||
raise ArchiveStorageError(f"Failed to check archive object '{key}': {e}") from e
|
||||
|
||||
def delete_object(self, key: str) -> None:
|
||||
"""
|
||||
@@ -196,7 +212,12 @@ class ArchiveStorage:
|
||||
self.client.delete_object(Bucket=self.bucket, Key=key)
|
||||
logger.debug("Deleted object: %s", key)
|
||||
except ClientError as e:
|
||||
raise ArchiveStorageError(f"Failed to delete object '{key}': {e}")
|
||||
if _is_object_not_found_error(e):
|
||||
logger.debug("Archive object was already absent: %s", key)
|
||||
return
|
||||
raise ArchiveStorageError(f"Failed to delete object '{key}': {e}") from e
|
||||
except BotoCoreError as e:
|
||||
raise ArchiveStorageError(f"Failed to delete object '{key}': {e}") from e
|
||||
|
||||
def generate_presigned_url(
|
||||
self,
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
"""add step by step tour state
|
||||
|
||||
Revision ID: b8c9d0e1f2a3
|
||||
Revises: 3c9f8e2a1d7b
|
||||
Create Date: 2026-06-29 12:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
import models
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b8c9d0e1f2a3"
|
||||
down_revision = "3c9f8e2a1d7b"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
op.create_table(
|
||||
"account_step_by_step_tour_states",
|
||||
sa.Column("id", models.types.StringUUID(), nullable=False),
|
||||
sa.Column("account_id", models.types.StringUUID(), nullable=False),
|
||||
sa.Column("first_workspace_id", models.types.StringUUID(), nullable=True),
|
||||
sa.Column("skipped", sa.Boolean(), server_default=sa.text("false"), nullable=False),
|
||||
sa.Column("completed_task_ids", models.types.AdjustedJSON(), nullable=False),
|
||||
sa.Column("manually_enabled_workspace_ids", models.types.AdjustedJSON(), nullable=False),
|
||||
sa.Column("manually_disabled_workspace_ids", models.types.AdjustedJSON(), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False),
|
||||
sa.Column("updated_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id", name="account_step_by_step_tour_state_pkey"),
|
||||
sa.UniqueConstraint("account_id", name="account_step_by_step_tour_state_account_id_key"),
|
||||
)
|
||||
|
||||
|
||||
def downgrade():
|
||||
op.drop_table("account_step_by_step_tour_states")
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
"""add workflow-run archive bundle cursor indexes
|
||||
|
||||
Revision ID: 3c9f8e2a1d7b
|
||||
Revises: 7a1c2d9e4b60
|
||||
Create Date: 2026-07-15 16:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "3c9f8e2a1d7b"
|
||||
down_revision = "7a1c2d9e4b60"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_TABLE_NAME = "workflow_run_archive_bundles"
|
||||
_INDEXES = (
|
||||
("workflow_run_archive_bundle_month_id_idx", ("year", "month", "id")),
|
||||
("workflow_run_archive_bundle_month_shard_id_idx", ("year", "month", "shard", "id")),
|
||||
)
|
||||
|
||||
|
||||
def _is_postgresql() -> bool:
|
||||
return op.get_bind().dialect.name == "postgresql"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if _is_postgresql():
|
||||
with op.get_context().autocommit_block():
|
||||
for index_name, columns in _INDEXES:
|
||||
op.create_index(index_name, _TABLE_NAME, columns, postgresql_concurrently=True)
|
||||
return
|
||||
for index_name, columns in _INDEXES:
|
||||
op.create_index(index_name, _TABLE_NAME, columns)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _is_postgresql():
|
||||
with op.get_context().autocommit_block():
|
||||
for index_name, _ in reversed(_INDEXES):
|
||||
op.drop_index(index_name, table_name=_TABLE_NAME, postgresql_concurrently=True)
|
||||
return
|
||||
for index_name, _ in reversed(_INDEXES):
|
||||
op.drop_index(index_name, table_name=_TABLE_NAME)
|
||||
@@ -101,6 +101,7 @@ from .model import (
|
||||
UploadFile,
|
||||
)
|
||||
from .oauth import DatasourceOauthParamConfig, DatasourceProvider, OAuthAccessToken
|
||||
from .onboarding import AccountStepByStepTourState
|
||||
from .provider import (
|
||||
LoadBalancingModelConfig,
|
||||
Provider,
|
||||
@@ -155,6 +156,7 @@ __all__ = [
|
||||
"Account",
|
||||
"AccountIntegrate",
|
||||
"AccountStatus",
|
||||
"AccountStepByStepTourState",
|
||||
"AccountTrialAppRecord",
|
||||
"Agent",
|
||||
"AgentConfigDraft",
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Account-level onboarding state models."""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import DateTime, func
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from .base import TypeBase, gen_uuidv7_string
|
||||
from .types import AdjustedJSON, StringUUID
|
||||
|
||||
|
||||
class AccountStepByStepTourState(TypeBase):
|
||||
"""Persistent account-level Step-by-step Tour state.
|
||||
|
||||
The tour is account-owned, with workspace IDs stored only as presentation
|
||||
overrides. The first workspace is the workspace context where an eligible
|
||||
account first asks for tour state; subsequent workspaces are opt-in only.
|
||||
"""
|
||||
|
||||
__tablename__ = "account_step_by_step_tour_states"
|
||||
__table_args__ = (
|
||||
sa.PrimaryKeyConstraint("id", name="account_step_by_step_tour_state_pkey"),
|
||||
sa.UniqueConstraint("account_id", name="account_step_by_step_tour_state_account_id_key"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(
|
||||
StringUUID,
|
||||
insert_default=gen_uuidv7_string,
|
||||
default_factory=gen_uuidv7_string,
|
||||
init=False,
|
||||
)
|
||||
account_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
first_workspace_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
|
||||
skipped: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"), default=False)
|
||||
completed_task_ids: Mapped[list[str]] = mapped_column(AdjustedJSON, nullable=False, default_factory=list)
|
||||
manually_enabled_workspace_ids: Mapped[list[str]] = mapped_column(
|
||||
AdjustedJSON,
|
||||
nullable=False,
|
||||
default_factory=list,
|
||||
)
|
||||
manually_disabled_workspace_ids: Mapped[list[str]] = mapped_column(
|
||||
AdjustedJSON,
|
||||
nullable=False,
|
||||
default_factory=list,
|
||||
)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime,
|
||||
server_default=func.current_timestamp(),
|
||||
nullable=False,
|
||||
init=False,
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime,
|
||||
server_default=func.current_timestamp(),
|
||||
nullable=False,
|
||||
init=False,
|
||||
onupdate=func.current_timestamp(),
|
||||
)
|
||||
@@ -1476,6 +1476,8 @@ class WorkflowRunArchiveBundle(DefaultFieldsDCMixin, TypeBase):
|
||||
name="workflow_run_archive_bundle_identity_uq",
|
||||
),
|
||||
sa.Index("workflow_run_archive_bundle_tenant_month_idx", "tenant_id", "year", "month"),
|
||||
sa.Index("workflow_run_archive_bundle_month_id_idx", "year", "month", "id"),
|
||||
sa.Index("workflow_run_archive_bundle_month_shard_id_idx", "year", "month", "shard", "id"),
|
||||
)
|
||||
|
||||
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
|
||||
@@ -7787,6 +7787,30 @@ Initiate OAuth login process
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Success | **application/json**: [OAuthProviderTokenResponse](#oauthprovidertokenresponse)<br> |
|
||||
|
||||
### [GET] /onboarding/step-by-step-tour/state
|
||||
Get account-level Step-by-step Tour state
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Success | **application/json**: [StepByStepTourStateResponse](#stepbysteptourstateresponse)<br> |
|
||||
|
||||
### [PATCH] /onboarding/step-by-step-tour/state
|
||||
Update account-level Step-by-step Tour state
|
||||
|
||||
#### Request Body
|
||||
|
||||
| Required | Schema |
|
||||
| -------- | ------ |
|
||||
| Yes | **application/json**: [StepByStepTourStatePatchPayload](#stepbysteptourstatepatchpayload)<br> |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Success | **application/json**: [StepByStepTourStateResponse](#stepbysteptourstateresponse)<br> |
|
||||
|
||||
### [DELETE] /rag/pipeline/customized/templates/{template_id}
|
||||
#### Parameters
|
||||
|
||||
@@ -9639,6 +9663,27 @@ Bedrock retrieval test (internal use only)
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Success | **application/json**: [TrialDatasetListResponse](#trialdatasetlistresponse)<br> |
|
||||
|
||||
### [POST] /trial-apps/{app_id}/files/upload
|
||||
**Upload a file into the tenant that owns the trial app**
|
||||
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| app_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Request Body
|
||||
|
||||
| Required | Schema |
|
||||
| -------- | ------ |
|
||||
| Yes | **multipart/form-data**: { **"file"**: binary, **"source"**: string, <br>**Available values:** "datasets" }<br> |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 201 | File uploaded successfully | **application/json**: [FileResponse](#fileresponse)<br> |
|
||||
|
||||
### [GET] /trial-apps/{app_id}/messages/{message_id}/suggested-questions
|
||||
#### Parameters
|
||||
|
||||
@@ -9668,6 +9713,27 @@ Bedrock retrieval test (internal use only)
|
||||
| ---- | ----------- | ------ |
|
||||
| 200 | Success | **application/json**: [Parameters](#parameters)<br> |
|
||||
|
||||
### [POST] /trial-apps/{app_id}/remote-files/upload
|
||||
**Upload a remote file into the tenant that owns the trial app**
|
||||
|
||||
#### Parameters
|
||||
|
||||
| Name | Located in | Description | Required | Schema |
|
||||
| ---- | ---------- | ----------- | -------- | ------ |
|
||||
| app_id | path | | Yes | string (uuid) |
|
||||
|
||||
#### Request Body
|
||||
|
||||
| Required | Schema |
|
||||
| -------- | ------ |
|
||||
| Yes | **application/json**: [RemoteFileUploadPayload](#remotefileuploadpayload)<br> |
|
||||
|
||||
#### Responses
|
||||
|
||||
| Code | Description | Schema |
|
||||
| ---- | ----------- | ------ |
|
||||
| 201 | File uploaded successfully | **application/json**: [FileWithSignedUrl](#filewithsignedurl)<br> |
|
||||
|
||||
### [GET] /trial-apps/{app_id}/site
|
||||
**Retrieve app site info**
|
||||
|
||||
@@ -21746,6 +21812,24 @@ Query parameters for listing snippet published workflows.
|
||||
| paused | integer | | Yes |
|
||||
| success | integer | | Yes |
|
||||
|
||||
#### StepByStepTourStatePatchPayload
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| action | string, <br>**Available values:** "complete_task", "disable_current_workspace", "enable_current_workspace", "skip", "uncomplete_task" | State update action<br>*Enum:* `"complete_task"`, `"disable_current_workspace"`, `"enable_current_workspace"`, `"skip"`, `"uncomplete_task"` | Yes |
|
||||
| task_id | string | Task ID for task actions | No |
|
||||
|
||||
#### StepByStepTourStateResponse
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| completed_task_ids | [ string, <br>**Available values:** "home", "integration", "knowledge", "studio" ] | | No |
|
||||
| first_workspace_id | string | | No |
|
||||
| manually_disabled_workspace_ids | [ string ] | | No |
|
||||
| manually_enabled_workspace_ids | [ string ] | | No |
|
||||
| skipped | boolean | | No |
|
||||
| updated_at | string | | No |
|
||||
|
||||
#### Storage
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
@@ -21859,6 +21943,7 @@ The subscription constructor of the trigger provider
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| _is_collaborative | boolean | | No |
|
||||
| conversation_variables | [ object ] | | No |
|
||||
| environment_variables | [ object ] | | No |
|
||||
| features | object | | Yes |
|
||||
@@ -21898,6 +21983,7 @@ Model class for provider system configuration response.
|
||||
| enable_learn_app | boolean, <br>**Default:** true | | Yes |
|
||||
| enable_marketplace | boolean | | Yes |
|
||||
| enable_social_oauth_login | boolean | | Yes |
|
||||
| enable_step_by_step_tour | boolean | | Yes |
|
||||
| enable_trial_app | boolean | | Yes |
|
||||
| is_allow_create_workspace | boolean | | Yes |
|
||||
| is_allow_register | boolean | | Yes |
|
||||
@@ -22854,6 +22940,7 @@ in form definiton, or a variable while the workflow is running.
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| allow_email_code_login | boolean | | Yes |
|
||||
| allow_email_password_login | boolean | | Yes |
|
||||
| allow_public_access | boolean, <br>**Default:** true | | Yes |
|
||||
| allow_sso | boolean | | Yes |
|
||||
| enabled | boolean | | Yes |
|
||||
| sso_config | [WebAppAuthSSOModel](#webappauthssomodel) | | Yes |
|
||||
|
||||
@@ -1577,6 +1577,7 @@ Default configuration for form inputs.
|
||||
| enable_learn_app | boolean, <br>**Default:** true | | Yes |
|
||||
| enable_marketplace | boolean | | Yes |
|
||||
| enable_social_oauth_login | boolean | | Yes |
|
||||
| enable_step_by_step_tour | boolean | | Yes |
|
||||
| enable_trial_app | boolean | | Yes |
|
||||
| is_allow_create_workspace | boolean | | Yes |
|
||||
| is_allow_register | boolean | | Yes |
|
||||
@@ -1642,6 +1643,7 @@ in form definiton, or a variable while the workflow is running.
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| allow_email_code_login | boolean | | Yes |
|
||||
| allow_email_password_login | boolean | | Yes |
|
||||
| allow_public_access | boolean, <br>**Default:** true | | Yes |
|
||||
| allow_sso | boolean | | Yes |
|
||||
| enabled | boolean | | Yes |
|
||||
| sso_config | [WebAppAuthSSOModel](#webappauthssomodel) | | Yes |
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Shared database fixtures for provider tests.
|
||||
|
||||
Provider tests live outside ``tests/unit_tests`` and therefore cannot use that
|
||||
suite's SQLite fixtures. Keep this fixture scoped to ``providers`` so each
|
||||
provider package can exercise real SQLAlchemy queries without a service
|
||||
database or mocked sessions.
|
||||
"""
|
||||
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from models.base import TypeBase
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sqlite3_session(request: pytest.FixtureRequest) -> Iterator[Session]:
|
||||
"""Yield an isolated SQLite session with the parametrized model tables.
|
||||
|
||||
Pass the required ORM classes through indirect parametrization. The engine
|
||||
is per-test so committed rows and identity-map state cannot leak between
|
||||
provider packages.
|
||||
"""
|
||||
|
||||
models: tuple[type[TypeBase], ...] = request.param
|
||||
engine = create_engine("sqlite:///:memory:")
|
||||
tables = [model.metadata.tables[model.__tablename__] for model in models]
|
||||
TypeBase.metadata.create_all(engine, tables=tables)
|
||||
try:
|
||||
with Session(engine, expire_on_commit=False) as session:
|
||||
yield session
|
||||
finally:
|
||||
engine.dispose()
|
||||
@@ -3,7 +3,10 @@ from __future__ import annotations
|
||||
import json
|
||||
from typing import NotRequired, TypedDict, override
|
||||
|
||||
from redis.lock import Lock
|
||||
|
||||
from extensions.ext_redis import redis_client
|
||||
from extensions.redis_names import serialize_redis_name
|
||||
|
||||
SESSION_STATE_TTL_SECONDS = 3600
|
||||
SERVER_HEARTBEAT_TTL_SECONDS = 90
|
||||
@@ -12,6 +15,31 @@ WORKFLOW_LEADER_PREFIX = "workflow_leader:"
|
||||
WS_SID_MAP_PREFIX = "ws_sid_map:"
|
||||
WS_SERVER_HEARTBEAT_PREFIX = "ws_server_heartbeat:"
|
||||
WS_SERVER_SESSIONS_PREFIX = "ws_server_sessions:"
|
||||
GRAPH_VIEW_STATE_LOCK_PREFIX = "workflow_graph_view_state_lock:"
|
||||
GRAPH_VIEW_STATE_LOCK_TIMEOUT_SECONDS = 15
|
||||
|
||||
_UPDATE_SESSION_GRAPH_ACTIVE_LUA = """
|
||||
local raw = redis.call('HGET', KEYS[1], ARGV[1])
|
||||
if not raw then
|
||||
return 0
|
||||
end
|
||||
|
||||
local decoded_ok, session_info = pcall(cjson.decode, raw)
|
||||
if not decoded_ok or type(session_info) ~= 'table' then
|
||||
return 0
|
||||
end
|
||||
|
||||
local incoming_sequence = tonumber(ARGV[3])
|
||||
local current_sequence = tonumber(session_info.graph_active_sequence)
|
||||
if current_sequence and incoming_sequence <= current_sequence then
|
||||
return 0
|
||||
end
|
||||
|
||||
session_info.graph_active = ARGV[2] == '1'
|
||||
session_info.graph_active_sequence = incoming_sequence
|
||||
redis.call('HSET', KEYS[1], ARGV[1], cjson.encode(session_info))
|
||||
return 1
|
||||
"""
|
||||
|
||||
|
||||
class WorkflowSessionInfo(TypedDict):
|
||||
@@ -156,6 +184,55 @@ class WorkflowCollaborationRepository:
|
||||
|
||||
return users
|
||||
|
||||
def get_session_info(self, workflow_id: str, sid: str) -> WorkflowSessionInfo | None:
|
||||
raw = self._redis.hget(self.workflow_key(workflow_id), sid)
|
||||
value = self._decode(raw)
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
session_info = json.loads(value)
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
return None
|
||||
|
||||
if not isinstance(session_info, dict):
|
||||
return None
|
||||
if "user_id" not in session_info or "username" not in session_info or "sid" not in session_info:
|
||||
return None
|
||||
|
||||
user: WorkflowSessionInfo = {
|
||||
"user_id": str(session_info["user_id"]),
|
||||
"username": str(session_info["username"]),
|
||||
"avatar": session_info.get("avatar"),
|
||||
"sid": str(session_info["sid"]),
|
||||
"connected_at": int(session_info.get("connected_at") or 0),
|
||||
}
|
||||
if isinstance(session_info.get("server_id"), str):
|
||||
user["server_id"] = session_info["server_id"]
|
||||
if isinstance(session_info.get("graph_active"), bool):
|
||||
user["graph_active"] = session_info["graph_active"]
|
||||
return user
|
||||
|
||||
def update_session_graph_active(self, workflow_id: str, sid: str, active: bool, sequence: int) -> bool:
|
||||
"""Atomically apply a graph visibility update when its client sequence is newer."""
|
||||
# RedisClientWrapper prefixes regular hash calls, but eval is delegated to the raw client.
|
||||
workflow_key = serialize_redis_name(self.workflow_key(workflow_id))
|
||||
result = self._redis.eval(
|
||||
_UPDATE_SESSION_GRAPH_ACTIVE_LUA,
|
||||
1,
|
||||
workflow_key,
|
||||
sid,
|
||||
"1" if active else "0",
|
||||
sequence,
|
||||
)
|
||||
return bool(result)
|
||||
|
||||
def graph_view_state_lock(self, workflow_id: str) -> Lock:
|
||||
"""Serialize visibility state changes and their leader-election side effects."""
|
||||
return self._redis.lock(
|
||||
f"{GRAPH_VIEW_STATE_LOCK_PREFIX}{workflow_id}",
|
||||
timeout=GRAPH_VIEW_STATE_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
def refresh_server_heartbeat(self, server_id: str) -> None:
|
||||
self._redis.set(self.server_key(server_id), "1", ex=SERVER_HEARTBEAT_TTL_SECONDS)
|
||||
|
||||
|
||||
@@ -9,12 +9,13 @@ import sqlalchemy as sa
|
||||
from sqlalchemy import and_, func, or_, select
|
||||
from sqlalchemy.orm import aliased
|
||||
|
||||
from configs import dify_config
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from libs.helper import convert_datetime_to_date, escape_like_pattern, to_timestamp
|
||||
from models.agent import WorkflowAgentNodeBinding
|
||||
from models.enums import MessageStatus
|
||||
from models.enums import CreatorUserRole, MessageStatus
|
||||
from models.model import App, Conversation, Message
|
||||
from models.workflow import WorkflowNodeExecutionModel, WorkflowRun
|
||||
from models.workflow import WorkflowNodeExecutionModel, WorkflowRun, WorkflowType
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -580,9 +581,26 @@ class AgentObservabilityService:
|
||||
|
||||
def _load_daily_statistics(
|
||||
self, *, app: App, agent_id: str, params: AgentStatisticsQueryParams, source_filter: AgentSourceFilter
|
||||
) -> list[dict[str, Any]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
if source_filter.kind in {"all", "webapp"}:
|
||||
rows.extend(self._load_webapp_daily_statistics(app=app, params=params, source_filter=source_filter))
|
||||
if source_filter.kind in {"all", "workflow"}:
|
||||
rows.extend(
|
||||
self._load_workflow_daily_statistics(
|
||||
app=app,
|
||||
agent_id=agent_id,
|
||||
params=params,
|
||||
source_filter=source_filter,
|
||||
)
|
||||
)
|
||||
return self._merge_daily_statistics(rows)
|
||||
|
||||
def _load_webapp_daily_statistics(
|
||||
self, *, app: App, params: AgentStatisticsQueryParams, source_filter: AgentSourceFilter
|
||||
) -> list[dict[str, Any]]:
|
||||
converted_created_at = convert_datetime_to_date("m.created_at")
|
||||
message_scope = self._statistics_message_scope_sql(source_filter)
|
||||
message_scope = self._statistics_webapp_message_scope_sql(source_filter)
|
||||
sql_query = f"""SELECT
|
||||
{converted_created_at} AS date,
|
||||
COUNT(m.id) AS message_count,
|
||||
@@ -602,12 +620,154 @@ WHERE
|
||||
args: dict[str, Any] = {
|
||||
"tz": params.timezone,
|
||||
"app_id": app.id,
|
||||
"tenant_id": app.tenant_id,
|
||||
"agent_id": agent_id,
|
||||
"debugger": InvokeFrom.DEBUGGER,
|
||||
}
|
||||
if source_filter.invoke_from is not None:
|
||||
args["source"] = source_filter.invoke_from
|
||||
if params.start:
|
||||
sql_query += " AND m.created_at >= :start"
|
||||
args["start"] = params.start
|
||||
if params.end:
|
||||
sql_query += " AND m.created_at < :end"
|
||||
args["end"] = params.end
|
||||
sql_query += " GROUP BY date ORDER BY date"
|
||||
|
||||
return [dict(row._mapping) for row in self._session.execute(sa.text(sql_query), args).all()]
|
||||
|
||||
@staticmethod
|
||||
def _statistics_webapp_message_scope_sql(source_filter: AgentSourceFilter) -> str:
|
||||
app_scope = "m.app_id = :app_id"
|
||||
if source_filter.invoke_from is not None:
|
||||
app_scope += " AND m.invoke_from = :source"
|
||||
return app_scope
|
||||
|
||||
def _load_workflow_daily_statistics(
|
||||
self,
|
||||
*,
|
||||
app: App,
|
||||
agent_id: str,
|
||||
params: AgentStatisticsQueryParams,
|
||||
source_filter: AgentSourceFilter,
|
||||
) -> list[dict[str, Any]]:
|
||||
converted_run_created_at = convert_datetime_to_date("aru.created_at")
|
||||
total_tokens = self._workflow_execution_metadata_numeric_sql(("total_tokens",), "BIGINT")
|
||||
nested_total_tokens = self._workflow_execution_metadata_numeric_sql(
|
||||
("agent_log", "agent_backend", "usage", "total_tokens"), "BIGINT"
|
||||
)
|
||||
total_price = self._workflow_execution_metadata_numeric_sql(("total_price",), "DECIMAL(65, 30)")
|
||||
nested_total_price = self._workflow_execution_metadata_numeric_sql(
|
||||
("agent_log", "agent_backend", "usage", "total_price"), "DECIMAL(65, 30)"
|
||||
)
|
||||
completion_tokens = self._workflow_execution_metadata_numeric_sql(
|
||||
("agent_log", "agent_backend", "usage", "completion_tokens"), "BIGINT"
|
||||
)
|
||||
binding_filters = self._statistics_workflow_binding_filters_sql(source_filter)
|
||||
run_date_filters = ""
|
||||
args: dict[str, Any] = {
|
||||
"tz": params.timezone,
|
||||
"tenant_id": app.tenant_id,
|
||||
"agent_id": agent_id,
|
||||
"chat_workflow_type": WorkflowType.CHAT,
|
||||
"end_user_role": CreatorUserRole.END_USER,
|
||||
}
|
||||
if source_filter.app_id:
|
||||
args["source_app_id"] = source_filter.app_id
|
||||
if source_filter.workflow_id:
|
||||
args["workflow_id"] = source_filter.workflow_id
|
||||
if source_filter.workflow_version:
|
||||
args["workflow_version"] = source_filter.workflow_version
|
||||
if source_filter.node_id:
|
||||
args["node_id"] = source_filter.node_id
|
||||
if params.start:
|
||||
run_date_filters += " AND wr.created_at >= :start"
|
||||
args["start"] = params.start
|
||||
if params.end:
|
||||
run_date_filters += " AND wr.created_at < :end"
|
||||
args["end"] = params.end
|
||||
|
||||
run_query = f"""WITH agent_run_usage AS (
|
||||
SELECT
|
||||
wr.id,
|
||||
wr.created_by_role,
|
||||
wr.created_by,
|
||||
wr.created_at,
|
||||
COALESCE(SUM(COALESCE({total_tokens}, {nested_total_tokens}, 0)), 0) AS token_count,
|
||||
COALESCE(SUM(COALESCE({total_price}, {nested_total_price}, 0)), 0) AS total_price,
|
||||
COALESCE(SUM(COALESCE(wne.elapsed_time, 0)), 0) AS latency,
|
||||
COALESCE(SUM(COALESCE({completion_tokens}, 0)), 0) AS answer_tokens
|
||||
FROM workflow_runs wr
|
||||
JOIN workflow_agent_node_bindings wanb
|
||||
ON wanb.tenant_id = :tenant_id
|
||||
AND wanb.agent_id = :agent_id
|
||||
AND wanb.app_id = wr.app_id
|
||||
AND wanb.workflow_id = wr.workflow_id
|
||||
AND wanb.workflow_version = wr.version
|
||||
{binding_filters}
|
||||
JOIN workflow_node_executions wne
|
||||
ON wne.workflow_run_id = wr.id
|
||||
AND wne.node_id = wanb.node_id
|
||||
WHERE wr.type != :chat_workflow_type{run_date_filters}
|
||||
GROUP BY wr.id, wr.created_by_role, wr.created_by, wr.created_at
|
||||
)
|
||||
SELECT
|
||||
{converted_run_created_at} AS date,
|
||||
COUNT(aru.id) AS message_count,
|
||||
COUNT(aru.id) AS conversation_count,
|
||||
COUNT(DISTINCT CASE
|
||||
WHEN aru.created_by_role = :end_user_role THEN aru.created_by
|
||||
ELSE NULL
|
||||
END) AS end_user_count,
|
||||
COALESCE(SUM(aru.token_count), 0) AS token_count,
|
||||
COALESCE(SUM(aru.total_price), 0) AS total_price,
|
||||
COALESCE(AVG(aru.latency), 0) AS avg_latency,
|
||||
COALESCE(SUM(aru.latency), 0) AS latency_sum,
|
||||
COALESCE(SUM(aru.answer_tokens), 0) AS answer_tokens,
|
||||
0 AS like_count
|
||||
FROM agent_run_usage aru
|
||||
GROUP BY date
|
||||
ORDER BY date"""
|
||||
rows = [dict(row._mapping) for row in self._session.execute(sa.text(run_query), args).all()]
|
||||
rows.extend(
|
||||
self._load_workflow_chat_daily_context(
|
||||
app=app,
|
||||
agent_id=agent_id,
|
||||
params=params,
|
||||
source_filter=source_filter,
|
||||
)
|
||||
)
|
||||
return self._merge_daily_statistics(rows)
|
||||
|
||||
def _load_workflow_chat_daily_context(
|
||||
self,
|
||||
*,
|
||||
app: App,
|
||||
agent_id: str,
|
||||
params: AgentStatisticsQueryParams,
|
||||
source_filter: AgentSourceFilter,
|
||||
) -> list[dict[str, Any]]:
|
||||
converted_created_at = convert_datetime_to_date("m.created_at")
|
||||
workflow_scope = self._statistics_workflow_message_scope_sql(source_filter)
|
||||
sql_query = f"""SELECT
|
||||
{converted_created_at} AS date,
|
||||
COUNT(m.id) AS message_count,
|
||||
COUNT(DISTINCT m.conversation_id) AS conversation_count,
|
||||
COUNT(DISTINCT m.from_end_user_id) AS end_user_count,
|
||||
COALESCE(SUM(COALESCE(m.message_tokens, 0) + COALESCE(m.answer_tokens, 0)), 0) AS token_count,
|
||||
COALESCE(SUM(COALESCE(m.total_price, 0)), 0) AS total_price,
|
||||
COALESCE(AVG(m.provider_response_latency), 0) AS avg_latency,
|
||||
COALESCE(SUM(m.provider_response_latency), 0) AS latency_sum,
|
||||
COALESCE(SUM(m.answer_tokens), 0) AS answer_tokens,
|
||||
COUNT(mf.id) AS like_count
|
||||
FROM messages m
|
||||
LEFT JOIN message_feedbacks mf
|
||||
ON mf.message_id = m.id AND mf.rating = 'like'
|
||||
WHERE
|
||||
{workflow_scope}"""
|
||||
args: dict[str, Any] = {
|
||||
"tz": params.timezone,
|
||||
"tenant_id": app.tenant_id,
|
||||
"agent_id": agent_id,
|
||||
"chat_workflow_type": WorkflowType.CHAT,
|
||||
}
|
||||
if source_filter.app_id:
|
||||
args["source_app_id"] = source_filter.app_id
|
||||
if source_filter.workflow_id:
|
||||
@@ -627,10 +787,7 @@ WHERE
|
||||
return [dict(row._mapping) for row in self._session.execute(sa.text(sql_query), args).all()]
|
||||
|
||||
@staticmethod
|
||||
def _statistics_message_scope_sql(source_filter: AgentSourceFilter) -> str:
|
||||
app_scope = "m.app_id = :app_id"
|
||||
if source_filter.invoke_from is not None:
|
||||
app_scope += " AND m.invoke_from = :source"
|
||||
def _statistics_workflow_binding_filters_sql(source_filter: AgentSourceFilter) -> str:
|
||||
workflow_binding_filters = []
|
||||
if source_filter.app_id:
|
||||
workflow_binding_filters.append("wanb.app_id = :source_app_id")
|
||||
@@ -640,8 +797,12 @@ WHERE
|
||||
workflow_binding_filters.append("wanb.workflow_version = :workflow_version")
|
||||
if source_filter.node_id:
|
||||
workflow_binding_filters.append("wanb.node_id = :node_id")
|
||||
extra_workflow_filters = f"AND {' AND '.join(workflow_binding_filters)}" if workflow_binding_filters else ""
|
||||
workflow_scope = f"""m.workflow_run_id IS NOT NULL
|
||||
return f"AND {' AND '.join(workflow_binding_filters)}" if workflow_binding_filters else ""
|
||||
|
||||
@classmethod
|
||||
def _statistics_workflow_message_scope_sql(cls, source_filter: AgentSourceFilter) -> str:
|
||||
binding_filters = cls._statistics_workflow_binding_filters_sql(source_filter)
|
||||
return f"""m.workflow_run_id IS NOT NULL
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM workflow_runs wr
|
||||
@@ -651,17 +812,65 @@ WHERE
|
||||
AND wanb.app_id = wr.app_id
|
||||
AND wanb.workflow_id = wr.workflow_id
|
||||
AND wanb.workflow_version = wr.version
|
||||
{extra_workflow_filters}
|
||||
{binding_filters}
|
||||
JOIN workflow_node_executions wne
|
||||
ON wne.workflow_run_id = wr.id
|
||||
AND wne.node_id = wanb.node_id
|
||||
WHERE wr.id = m.workflow_run_id
|
||||
AND wr.type = :chat_workflow_type
|
||||
)"""
|
||||
if source_filter.kind == "webapp":
|
||||
return app_scope
|
||||
if source_filter.kind == "workflow":
|
||||
return workflow_scope
|
||||
return f"(({app_scope}) OR ({workflow_scope}))"
|
||||
|
||||
@staticmethod
|
||||
def _workflow_execution_metadata_numeric_sql(path: tuple[str, ...], numeric_type: str) -> str:
|
||||
if dify_config.DB_TYPE == "postgresql":
|
||||
json_path = ",".join(path)
|
||||
value = f"CAST(wne.execution_metadata AS JSONB) #>> '{{{json_path}}}'"
|
||||
return f"CAST(NULLIF({value}, '') AS {numeric_type})"
|
||||
if dify_config.DB_TYPE in {"mysql", "oceanbase", "seekdb"}:
|
||||
json_path = "$." + ".".join(path)
|
||||
mysql_numeric_type = "UNSIGNED" if numeric_type == "BIGINT" else numeric_type
|
||||
value = f"JSON_UNQUOTE(JSON_EXTRACT(wne.execution_metadata, '{json_path}'))"
|
||||
return f"CAST(NULLIF(NULLIF({value}, ''), 'null') AS {mysql_numeric_type})"
|
||||
raise NotImplementedError(f"Unsupported database type: {dify_config.DB_TYPE}")
|
||||
|
||||
@staticmethod
|
||||
def _merge_daily_statistics(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
merged: dict[Any, dict[str, Any]] = {}
|
||||
weighted_latency: dict[Any, float] = {}
|
||||
for row in rows:
|
||||
date = row["date"]
|
||||
target = merged.setdefault(
|
||||
date,
|
||||
{
|
||||
"date": date,
|
||||
"message_count": 0,
|
||||
"conversation_count": 0,
|
||||
"end_user_count": 0,
|
||||
"token_count": 0,
|
||||
"total_price": Decimal(0),
|
||||
"avg_latency": 0.0,
|
||||
"latency_sum": 0.0,
|
||||
"answer_tokens": 0,
|
||||
"like_count": 0,
|
||||
},
|
||||
)
|
||||
message_count = int(row.get("message_count") or 0)
|
||||
target["message_count"] += message_count
|
||||
target["conversation_count"] += int(row.get("conversation_count") or 0)
|
||||
target["end_user_count"] += int(row.get("end_user_count") or 0)
|
||||
target["token_count"] += int(row.get("token_count") or 0)
|
||||
target["total_price"] += Decimal(str(row.get("total_price") or 0))
|
||||
target["latency_sum"] += float(row.get("latency_sum") or 0)
|
||||
target["answer_tokens"] += int(row.get("answer_tokens") or 0)
|
||||
target["like_count"] += int(row.get("like_count") or 0)
|
||||
weighted_latency[date] = (
|
||||
weighted_latency.get(date, 0.0) + float(row.get("avg_latency") or 0) * message_count
|
||||
)
|
||||
|
||||
for date, row in merged.items():
|
||||
message_count = int(row["message_count"])
|
||||
row["avg_latency"] = weighted_latency[date] / message_count if message_count else 0.0
|
||||
return sorted(merged.values(), key=lambda row: str(row["date"]))
|
||||
|
||||
@staticmethod
|
||||
def _build_charts(rows: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]:
|
||||
|
||||
+24
-12
@@ -408,6 +408,7 @@ class AppService:
|
||||
default_model_config = app_template.get("model_config")
|
||||
default_model_config = default_model_config.copy() if default_model_config else None
|
||||
if default_model_config and "model" in default_model_config:
|
||||
default_model_dict = default_model_config["model"]
|
||||
# get model provider
|
||||
model_manager = ModelManager.for_tenant(tenant_id=account.current_tenant_id or "")
|
||||
|
||||
@@ -422,7 +423,7 @@ class AppService:
|
||||
logger.exception("Get default model instance failed, tenant_id: %s", tenant_id)
|
||||
model_instance = None
|
||||
|
||||
if model_instance:
|
||||
if model_instance is not None:
|
||||
if (
|
||||
model_instance.model_name == default_model_config["model"]["name"]
|
||||
and model_instance.provider == default_model_config["model"]["provider"]
|
||||
@@ -430,17 +431,28 @@ class AppService:
|
||||
default_model_dict = default_model_config["model"]
|
||||
else:
|
||||
llm_model = cast(LargeLanguageModel, model_instance.model_type_instance)
|
||||
model_schema = llm_model.get_model_schema(model_instance.model_name, model_instance.credentials)
|
||||
if model_schema is None:
|
||||
raise ValueError(f"model schema not found for model {model_instance.model_name}")
|
||||
|
||||
default_model_dict = {
|
||||
"provider": model_instance.provider,
|
||||
"name": model_instance.model_name,
|
||||
"mode": model_schema.model_properties.get(ModelPropertyKey.MODE),
|
||||
"completion_params": {},
|
||||
}
|
||||
else:
|
||||
try:
|
||||
model_schema = llm_model.get_model_schema(model_instance.model_name, model_instance.credentials)
|
||||
if model_schema is None:
|
||||
raise ValueError(f"model schema not found for model {model_instance.model_name}")
|
||||
except Exception:
|
||||
# A removed provider model must not prevent creating an app.
|
||||
logger.warning(
|
||||
"Default model schema is unavailable, tenant_id: %s, provider: %s, model: %s",
|
||||
tenant_id,
|
||||
model_instance.provider,
|
||||
model_instance.model_name,
|
||||
exc_info=True,
|
||||
)
|
||||
model_instance = None
|
||||
else:
|
||||
default_model_dict = {
|
||||
"provider": model_instance.provider,
|
||||
"name": model_instance.model_name,
|
||||
"mode": model_schema.model_properties.get(ModelPropertyKey.MODE),
|
||||
"completion_params": {},
|
||||
}
|
||||
if model_instance is None:
|
||||
try:
|
||||
provider, model = model_manager.get_default_provider_model_name(
|
||||
tenant_id=account.current_tenant_id or "", model_type=ModelType.LLM
|
||||
|
||||
@@ -739,8 +739,13 @@ class DatasetService:
|
||||
dataset.id, external_knowledge_id, external_knowledge_api_id, session
|
||||
)
|
||||
|
||||
# Commit changes to database
|
||||
session.commit()
|
||||
# Flush changes to the database without closing the caller-managed
|
||||
# transaction. This helper receives a session opened by the caller
|
||||
# (`with Session(...) as session`); calling commit() here closed that
|
||||
# context manager early and raised
|
||||
# sqlalchemy.exc.InvalidRequestError: Can't operate on closed transaction
|
||||
# (#39191).
|
||||
session.flush()
|
||||
|
||||
return dataset
|
||||
|
||||
@@ -810,9 +815,15 @@ class DatasetService:
|
||||
if data.get("icon_info"):
|
||||
filtered_data["icon_info"] = data.get("icon_info")
|
||||
|
||||
# Update dataset in database
|
||||
# Update dataset in database. Use flush() rather than commit() so the
|
||||
# caller-managed transaction (opened with `with Session(...) as session`)
|
||||
# stays open for subsequent operations — _update_pipeline_knowledge_base
|
||||
# node data and any caller follow-ups run on the same session. Calling
|
||||
# commit() here closed the context manager early and raised
|
||||
# sqlalchemy.exc.InvalidRequestError: Can't operate on closed transaction
|
||||
# (#39191).
|
||||
session.execute(update(Dataset).where(Dataset.id == dataset.id).values(**filtered_data))
|
||||
session.commit()
|
||||
session.flush()
|
||||
|
||||
# Reload dataset to get updated values
|
||||
session.refresh(dataset)
|
||||
|
||||
@@ -317,9 +317,6 @@ _LEGACY_WORKSPACE_OWNER_KEYS: list[str] = [
|
||||
"credential.use",
|
||||
"credential.create",
|
||||
"credential.manage",
|
||||
"billing.view",
|
||||
"billing.subscription.manage",
|
||||
"billing.manage",
|
||||
"app.acl.preview",
|
||||
"app_library.access",
|
||||
"app.create_and_management",
|
||||
@@ -349,9 +346,6 @@ _LEGACY_WORKSPACE_ADMIN_KEYS: list[str] = [
|
||||
"credential.use",
|
||||
"credential.create",
|
||||
"credential.manage",
|
||||
"billing.view",
|
||||
"billing.subscription.manage",
|
||||
"billing.manage",
|
||||
"app_library.access",
|
||||
"app.create_and_management",
|
||||
"app.tag.manage",
|
||||
@@ -377,9 +371,6 @@ _LEGACY_WORKSPACE_EDITOR_KEYS: list[str] = [
|
||||
"dataset.external.connect",
|
||||
"snippets.create_and_modify",
|
||||
"tool.manage",
|
||||
"billing.view",
|
||||
"billing.subscription.manage",
|
||||
"billing.manage",
|
||||
]
|
||||
|
||||
_LEGACY_WORKSPACE_NORMAL_KEYS: list[str] = [
|
||||
@@ -387,9 +378,6 @@ _LEGACY_WORKSPACE_NORMAL_KEYS: list[str] = [
|
||||
"plugin.install",
|
||||
"credential.use",
|
||||
"app_library.access",
|
||||
"billing.view",
|
||||
"billing.subscription.manage",
|
||||
"billing.manage",
|
||||
]
|
||||
|
||||
_LEGACY_WORKSPACE_DATASET_OPERATOR_KEYS: list[str] = [
|
||||
@@ -805,7 +793,6 @@ 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,
|
||||
)
|
||||
|
||||
@@ -100,6 +100,7 @@ class WebAppAuthModel(FeatureResponseModel):
|
||||
sso_config: WebAppAuthSSOModel = WebAppAuthSSOModel()
|
||||
allow_email_code_login: bool = False
|
||||
allow_email_password_login: bool = False
|
||||
allow_public_access: bool = True
|
||||
|
||||
|
||||
class KnowledgePipeline(FeatureResponseModel):
|
||||
@@ -183,6 +184,7 @@ class SystemFeatureModel(FeatureResponseModel):
|
||||
enable_trial_app: bool = False
|
||||
enable_explore_banner: bool = False
|
||||
enable_learn_app: bool = True
|
||||
enable_step_by_step_tour: bool = False
|
||||
rbac_enabled: bool = False
|
||||
|
||||
|
||||
@@ -285,6 +287,8 @@ class FeatureService:
|
||||
system_features.enable_trial_app = dify_config.ENABLE_TRIAL_APP
|
||||
system_features.enable_explore_banner = dify_config.ENABLE_EXPLORE_BANNER
|
||||
system_features.enable_learn_app = dify_config.ENABLE_LEARN_APP
|
||||
system_features.webapp_auth.allow_public_access = dify_config.WEBAPP_PUBLIC_ACCESS_ENABLED
|
||||
system_features.enable_step_by_step_tour = dify_config.ENABLE_STEP_BY_STEP_TOUR
|
||||
|
||||
@classmethod
|
||||
def _fulfill_trial_models_from_env(cls) -> list[str]:
|
||||
|
||||
@@ -0,0 +1,358 @@
|
||||
"""Transport-only forwarding for the explicitly enabled KnowledgeFS Console operations.
|
||||
|
||||
KnowledgeFS owns the request and response contract. This module binds short-lived
|
||||
account and workspace identities, enforces Dify's coarse workspace policy, and
|
||||
normalizes transport failures. Dify deliberately maintains a small product-facing
|
||||
operation registry instead of exposing the full upstream OpenAPI surface. The
|
||||
dedicated request path uses Dify's shared SSRF policy, never follows redirects,
|
||||
bounds buffered responses, and rejects compressed responses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from http import HTTPStatus
|
||||
from typing import Final, Literal, NamedTuple, Protocol
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
|
||||
from configs import dify_config
|
||||
from core.helper import ssrf_proxy
|
||||
from core.rbac import RBACPermission, RBACResourceScope
|
||||
from core.tools.errors import ToolSSRFError
|
||||
from models import Account
|
||||
from services.enterprise.rbac_service import RBACService
|
||||
|
||||
type KnowledgeFSMethod = Literal["DELETE", "GET", "PATCH", "POST", "PUT"]
|
||||
type KnowledgeFSResponseKind = Literal["binary", "buffered", "stream"]
|
||||
type KnowledgeFSRequiredScope = Literal["knowledge-spaces:read", "knowledge-spaces:write"]
|
||||
|
||||
_JWT_AUDIENCE = "knowledge-fs"
|
||||
_JWT_ISSUER = "dify"
|
||||
_JWT_TTL_SECONDS = 60
|
||||
_MAX_BUFFERED_RESPONSE_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
class KnowledgeFSOperation(NamedTuple):
|
||||
operation_id: str
|
||||
method: KnowledgeFSMethod
|
||||
path: str
|
||||
response_kind: KnowledgeFSResponseKind
|
||||
required_scope: KnowledgeFSRequiredScope
|
||||
rbac_permission: RBACPermission
|
||||
requires_dataset_editor: bool
|
||||
max_response_bytes: int
|
||||
request_headers: tuple[str, ...]
|
||||
response_headers: tuple[str, ...]
|
||||
response_media_types: tuple[str, ...]
|
||||
|
||||
|
||||
KNOWLEDGE_FS_CONSOLE_OPERATIONS: Final[tuple[KnowledgeFSOperation, ...]] = (
|
||||
KnowledgeFSOperation(
|
||||
operation_id="listKnowledgeSpaces",
|
||||
method="GET",
|
||||
path="knowledge-spaces",
|
||||
response_kind="buffered",
|
||||
required_scope="knowledge-spaces:read",
|
||||
rbac_permission=RBACPermission.DATASET_READONLY,
|
||||
requires_dataset_editor=False,
|
||||
max_response_bytes=1_048_576,
|
||||
request_headers=("x-trace-id",),
|
||||
response_headers=("x-trace-id",),
|
||||
response_media_types=("application/json",),
|
||||
),
|
||||
KnowledgeFSOperation(
|
||||
operation_id="createKnowledgeSpace",
|
||||
method="POST",
|
||||
path="knowledge-spaces",
|
||||
response_kind="buffered",
|
||||
required_scope="knowledge-spaces:write",
|
||||
rbac_permission=RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
|
||||
requires_dataset_editor=True,
|
||||
max_response_bytes=1_048_576,
|
||||
request_headers=("x-trace-id",),
|
||||
response_headers=("x-trace-id",),
|
||||
response_media_types=("application/json",),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class KnowledgeFSUpstreamResponse(NamedTuple):
|
||||
response: httpx.Response
|
||||
response_kind: KnowledgeFSResponseKind
|
||||
operation: KnowledgeFSOperation
|
||||
|
||||
|
||||
class _RequestHeaders(Protocol):
|
||||
def items(self) -> Iterable[tuple[str, str]]: ...
|
||||
|
||||
|
||||
class KnowledgeFSConfigurationError(RuntimeError):
|
||||
"""KnowledgeFS is incompletely configured or blocked by outbound policy."""
|
||||
|
||||
|
||||
class KnowledgeFSTimeoutError(RuntimeError):
|
||||
"""KnowledgeFS exceeded the configured request timeout."""
|
||||
|
||||
|
||||
class KnowledgeFSTransportError(RuntimeError):
|
||||
"""KnowledgeFS could not be reached or returned a response outside safety bounds."""
|
||||
|
||||
|
||||
class KnowledgeFSRouteNotAllowedError(RuntimeError):
|
||||
"""The requested path is outside the Console-visible KnowledgeFS surface."""
|
||||
|
||||
|
||||
class KnowledgeFSAccessDeniedError(RuntimeError):
|
||||
"""The Dify account lacks the workspace permission required by the operation."""
|
||||
|
||||
|
||||
def authorize_knowledge_fs_request(
|
||||
*,
|
||||
account: Account,
|
||||
tenant_id: str,
|
||||
operation: KnowledgeFSOperation,
|
||||
) -> None:
|
||||
"""Enforce Dify's workspace policy before KFS performs resource authorization.
|
||||
|
||||
Args:
|
||||
account: Authenticated Dify account with its current workspace role.
|
||||
tenant_id: Current Dify workspace identifier.
|
||||
operation: Dify-maintained KnowledgeFS operation and policy metadata.
|
||||
|
||||
Raises:
|
||||
KnowledgeFSAccessDeniedError: The account lacks a required legacy or enterprise permission.
|
||||
"""
|
||||
if operation.requires_dataset_editor and not account.is_dataset_editor:
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS mutations require dataset edit access")
|
||||
if not RBACService.CheckAccess.check(
|
||||
tenant_id,
|
||||
account.id,
|
||||
scene=operation.rbac_permission.value,
|
||||
resource_type=RBACResourceScope.DATASET.value,
|
||||
):
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS operation is denied by workspace RBAC")
|
||||
|
||||
|
||||
def proxy_knowledge_fs_request(
|
||||
*,
|
||||
account: Account,
|
||||
method: KnowledgeFSMethod,
|
||||
path: str,
|
||||
tenant_id: str,
|
||||
accept: str | None = None,
|
||||
content_type: str | None = None,
|
||||
query: bytes | None = None,
|
||||
body: bytes | None = None,
|
||||
request_headers: _RequestHeaders | None = None,
|
||||
) -> KnowledgeFSUpstreamResponse:
|
||||
"""Authorize and forward one allowlisted KnowledgeFS request as a single use case."""
|
||||
operation = get_knowledge_fs_operation(method, path)
|
||||
authorize_knowledge_fs_request(
|
||||
account=account,
|
||||
tenant_id=tenant_id,
|
||||
operation=operation,
|
||||
)
|
||||
incoming_request_headers = {name.lower(): value for name, value in (request_headers or {}).items()}
|
||||
contract_request_headers = {
|
||||
name: incoming_request_headers[name] for name in operation.request_headers if name in incoming_request_headers
|
||||
}
|
||||
return _forward_knowledge_fs_request(
|
||||
account_id=account.id,
|
||||
method=method,
|
||||
path=path,
|
||||
tenant_id=tenant_id,
|
||||
accept=accept,
|
||||
content_type=content_type,
|
||||
query=query,
|
||||
body=body,
|
||||
request_headers=contract_request_headers,
|
||||
)
|
||||
|
||||
|
||||
def _forward_knowledge_fs_request(
|
||||
*,
|
||||
account_id: str,
|
||||
method: KnowledgeFSMethod,
|
||||
path: str,
|
||||
tenant_id: str,
|
||||
accept: str | None = None,
|
||||
content_type: str | None = None,
|
||||
query: bytes | None = None,
|
||||
body: bytes | None = None,
|
||||
request_headers: Mapping[str, str] | None = None,
|
||||
) -> KnowledgeFSUpstreamResponse:
|
||||
"""Forward one fixed-route request without parsing its KnowledgeFS payload.
|
||||
|
||||
Args:
|
||||
account_id: Current Dify account used as the KFS member identity.
|
||||
method: Allowlisted upstream HTTP method.
|
||||
path: Relative KnowledgeFS path under an allowlisted product surface.
|
||||
tenant_id: Current Dify workspace used as the KFS tenant identity.
|
||||
accept: Original Accept header, when present.
|
||||
content_type: Original request Content-Type header, when present.
|
||||
query: Original encoded query string from the Console request.
|
||||
body: Original request body, when present.
|
||||
request_headers: Contract-declared request headers forwarded by the Console adapter.
|
||||
|
||||
Returns:
|
||||
The KnowledgeFS response and its actual transport kind. Non-success responses are buffered.
|
||||
|
||||
Raises:
|
||||
KnowledgeFSConfigurationError: The connection is incomplete or blocked by outbound policy.
|
||||
KnowledgeFSRouteNotAllowedError: The path is outside the allowlisted product surface.
|
||||
KnowledgeFSTimeoutError: KnowledgeFS exceeds the configured timeout.
|
||||
KnowledgeFSTransportError: The request fails or its response cannot be safely bounded.
|
||||
|
||||
Each request is bound to stable Dify account and workspace principals with a short expiration.
|
||||
"""
|
||||
operation = get_knowledge_fs_operation(method, path)
|
||||
base_url = dify_config.KNOWLEDGE_FS_BASE_URL
|
||||
jwt_secret = dify_config.KNOWLEDGE_FS_JWT_SECRET
|
||||
if base_url is None or jwt_secret is None:
|
||||
raise KnowledgeFSConfigurationError("KnowledgeFS connection configuration is incomplete")
|
||||
now = datetime.now(UTC)
|
||||
token = jwt.encode(
|
||||
{
|
||||
"aud": _JWT_AUDIENCE,
|
||||
"caller_kind": "interactive",
|
||||
"dify_account_id": f"dify-account:{account_id}",
|
||||
"exp": now + timedelta(seconds=_JWT_TTL_SECONDS),
|
||||
"iat": now,
|
||||
"iss": _JWT_ISSUER,
|
||||
"scopes": [operation.required_scope],
|
||||
"sub": f"dify-workspace:{tenant_id}",
|
||||
"tenant_id": tenant_id,
|
||||
},
|
||||
jwt_secret.get_secret_value(),
|
||||
algorithm="HS256",
|
||||
)
|
||||
headers = {
|
||||
"Accept": accept or "application/json",
|
||||
"Accept-Encoding": "identity",
|
||||
"Authorization": f"Bearer {token}",
|
||||
}
|
||||
if body is not None:
|
||||
headers["Content-Type"] = content_type or "application/json"
|
||||
allowed_request_headers = set(operation.request_headers)
|
||||
for name, value in (request_headers or {}).items():
|
||||
normalized_name = name.lower()
|
||||
if normalized_name not in allowed_request_headers:
|
||||
raise KnowledgeFSRouteNotAllowedError("KnowledgeFS request header is not allowed")
|
||||
headers[normalized_name] = value
|
||||
|
||||
try:
|
||||
upstream_url = httpx.URL(f"{base_url}/").join(operation.path)
|
||||
response = ssrf_proxy.make_request(
|
||||
method=operation.method,
|
||||
url=str(upstream_url),
|
||||
params=query,
|
||||
content=body,
|
||||
headers=headers,
|
||||
timeout=dify_config.KNOWLEDGE_FS_TIMEOUT_SECONDS,
|
||||
follow_redirects=False,
|
||||
max_retries=0,
|
||||
stream_response=True,
|
||||
)
|
||||
response_kind = _classify_response(operation, response)
|
||||
if response_kind == "stream":
|
||||
content_encoding = response.headers.get("content-encoding", "identity").strip().lower()
|
||||
if content_encoding not in {"", "identity"}:
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS streaming response used an unsupported encoding")
|
||||
_set_response_read_timeout(response, dify_config.KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS)
|
||||
return KnowledgeFSUpstreamResponse(response, response_kind, operation)
|
||||
|
||||
max_response_bytes = (
|
||||
operation.max_response_bytes
|
||||
if HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES
|
||||
else _MAX_BUFFERED_RESPONSE_BYTES
|
||||
)
|
||||
buffered_response = ssrf_proxy.buffer_response(response, max_response_bytes=max_response_bytes)
|
||||
if buffered_response.content and not buffered_response.headers.get("content-type", "").strip():
|
||||
buffered_response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS buffered response used an unsupported media type")
|
||||
return KnowledgeFSUpstreamResponse(buffered_response, response_kind, operation)
|
||||
except ssrf_proxy.ResponseLimitError as exc:
|
||||
raise KnowledgeFSTransportError("KnowledgeFS response violated the proxy limit") from exc
|
||||
except ToolSSRFError as exc:
|
||||
raise KnowledgeFSConfigurationError("KnowledgeFS origin was blocked by outbound policy") from exc
|
||||
except httpx.TimeoutException as exc:
|
||||
raise KnowledgeFSTimeoutError("KnowledgeFS request timed out") from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise KnowledgeFSTransportError("KnowledgeFS transport request failed") from exc
|
||||
|
||||
|
||||
def get_knowledge_fs_operation(method: KnowledgeFSMethod, path: str) -> KnowledgeFSOperation:
|
||||
"""Resolve an exact operation and its transport/access contract metadata."""
|
||||
for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS:
|
||||
if method == operation.method and _matches_route_template(operation.path, path):
|
||||
return operation._replace(path=path)
|
||||
raise KnowledgeFSRouteNotAllowedError("KnowledgeFS route is not allowed")
|
||||
|
||||
|
||||
def _classify_response(operation: KnowledgeFSOperation, response: httpx.Response) -> KnowledgeFSResponseKind:
|
||||
"""Resolve the actual response kind from status and Content-Type before reading its body."""
|
||||
content_type = response.headers.get("content-type", "").partition(";")[0].strip().lower()
|
||||
is_success = HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES
|
||||
if not is_success:
|
||||
if content_type and not _is_json_content_type(content_type):
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS error response used an unsupported media type")
|
||||
return "buffered"
|
||||
|
||||
if operation.response_kind == "stream":
|
||||
if content_type != "text/event-stream":
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS stream response used an unsupported media type")
|
||||
return "stream"
|
||||
if operation.response_kind == "binary":
|
||||
if content_type not in operation.response_media_types:
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS binary response used an unsupported media type")
|
||||
return "binary"
|
||||
if content_type and not _is_json_content_type(content_type):
|
||||
response.close()
|
||||
raise KnowledgeFSTransportError("KnowledgeFS buffered response used an unsupported media type")
|
||||
return "buffered"
|
||||
|
||||
|
||||
def _is_json_content_type(content_type: str) -> bool:
|
||||
return content_type == "application/json" or content_type.endswith("+json")
|
||||
|
||||
|
||||
def _set_response_read_timeout(response: httpx.Response, timeout_seconds: float | None) -> None:
|
||||
"""Set the body-read timeout after headers identify a valid SSE response."""
|
||||
try:
|
||||
request = response.request
|
||||
except RuntimeError:
|
||||
return
|
||||
timeout = request.extensions.get("timeout")
|
||||
if isinstance(timeout, dict):
|
||||
timeout["read"] = timeout_seconds
|
||||
|
||||
|
||||
def _matches_route_template(template: str, path: str) -> bool:
|
||||
"""Match path parameters without permitting encoded or traversal-like segments."""
|
||||
template_segments = template.split("/")
|
||||
path_segments = path.split("/")
|
||||
if len(template_segments) != len(path_segments):
|
||||
return False
|
||||
|
||||
for template_segment, path_segment in zip(template_segments, path_segments, strict=True):
|
||||
if template_segment.startswith("{") and template_segment.endswith("}"):
|
||||
if (
|
||||
not path_segment
|
||||
or path_segment in {".", ".."}
|
||||
or "\\" in path_segment
|
||||
or "%" in path_segment
|
||||
or "?" in path_segment
|
||||
or "#" in path_segment
|
||||
):
|
||||
return False
|
||||
continue
|
||||
if template_segment != path_segment:
|
||||
return False
|
||||
return True
|
||||
@@ -601,6 +601,7 @@ class RagPipelineService:
|
||||
with sessionmaker(bind=db.engine).begin() as session:
|
||||
draft_var_saver = DraftVariableSaver(
|
||||
session=session,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
app_id=pipeline.id,
|
||||
node_id=workflow_node_execution.node_id,
|
||||
node_type=workflow_node_execution.node_type,
|
||||
@@ -1391,6 +1392,7 @@ class RagPipelineService:
|
||||
with sessionmaker(bind=db.engine).begin() as session:
|
||||
draft_var_saver = DraftVariableSaver(
|
||||
session=session,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
app_id=pipeline.id,
|
||||
node_id=workflow_node_execution_db_model.node_id,
|
||||
node_type=workflow_node_execution_db_model.node_type,
|
||||
|
||||
@@ -6,7 +6,10 @@ This service archives workflow run logs for paid plan users older than the confi
|
||||
|
||||
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 as the recoverable source of truth. Completed bundles are also
|
||||
mirrored into a small database index so console listing and download jobs do not list object storage online.
|
||||
published into a small database catalog so console listing, download, and maintenance jobs do not list object storage
|
||||
online. Archive success requires that catalog publication to commit. A retry reconciles only its known manifest key
|
||||
against an already-published shard index; a missing shard index fails closed rather than rebuilding it by scanning
|
||||
historical manifests.
|
||||
|
||||
Archive campaigns should use fixed absolute UTC windows for every tenant-prefix/shard execution. Relative windows are
|
||||
evaluated at process start and are not safe for multi-day rollout because each command would scan a different window.
|
||||
@@ -639,7 +642,6 @@ class WorkflowRunArchiver:
|
||||
if storage is None:
|
||||
raise ArchiveStorageNotConfiguredError("Archive storage not configured")
|
||||
if storage.object_exists(self._get_manifest_object_key(identity)):
|
||||
self._write_bundle_index(storage, identity)
|
||||
self._sync_existing_bundle_index(session, storage, identity)
|
||||
result.success = True
|
||||
result.skipped = True
|
||||
@@ -651,6 +653,8 @@ class WorkflowRunArchiver:
|
||||
runs = [run for run in runs if run.id not in archived_run_ids]
|
||||
result.skipped_run_count = original_run_count - len(runs)
|
||||
if not runs:
|
||||
# Historical catalog rows are a rollout precondition. New bundles commit their catalog row before
|
||||
# publishing the shard index, so this idempotency path never needs to rescan prior manifests.
|
||||
result.run_count = 0
|
||||
result.success = True
|
||||
result.skipped = True
|
||||
@@ -663,7 +667,6 @@ class WorkflowRunArchiver:
|
||||
result.object_prefix = identity.object_prefix
|
||||
result.run_count = len(runs)
|
||||
if storage.object_exists(self._get_manifest_object_key(identity)):
|
||||
self._write_bundle_index(storage, identity)
|
||||
self._sync_existing_bundle_index(session, storage, identity)
|
||||
result.success = True
|
||||
result.skipped = True
|
||||
@@ -695,10 +698,10 @@ class WorkflowRunArchiver:
|
||||
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)
|
||||
self._merge_bundle_manifest_into_index(storage, identity, [run.id for run in runs])
|
||||
manifest = decode_archive_bundle_manifest(manifest_data)
|
||||
upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data))
|
||||
session.commit()
|
||||
self._merge_bundle_manifest_into_index(storage, identity, [run.id for run in runs])
|
||||
|
||||
logger.info(
|
||||
"Archived workflow run bundle %s: tenant=%s runs=%s tables=%s object_prefix=%s",
|
||||
@@ -730,16 +733,18 @@ class WorkflowRunArchiver:
|
||||
storage: ArchiveStorage,
|
||||
identity: ArchiveBundleIdentity,
|
||||
) -> None:
|
||||
"""Best-effort DB index sync for a bundle whose manifest already exists in archive storage."""
|
||||
"""Publish a known manifest to the DB catalog and reconcile its existing shard index."""
|
||||
manifest_key = self._get_manifest_object_key(identity)
|
||||
try:
|
||||
manifest_data = storage.get_object(manifest_key)
|
||||
manifest = decode_archive_bundle_manifest(manifest_data)
|
||||
upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data))
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
logger.warning("Failed to sync workflow archive bundle index for %s", manifest_key, exc_info=True)
|
||||
manifest_data = storage.get_object(manifest_key)
|
||||
manifest = decode_archive_bundle_manifest(manifest_data)
|
||||
upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data))
|
||||
session.commit()
|
||||
self._merge_bundle_manifest_into_index(
|
||||
storage,
|
||||
identity,
|
||||
manifest["run_ids"],
|
||||
require_existing_index=True,
|
||||
)
|
||||
|
||||
def _lock_runs_for_archive(
|
||||
self,
|
||||
@@ -1029,10 +1034,20 @@ class WorkflowRunArchiver:
|
||||
storage: ArchiveStorage,
|
||||
identity: ArchiveBundleIdentity,
|
||||
run_ids: Sequence[str],
|
||||
*,
|
||||
require_existing_index: bool = False,
|
||||
) -> ArchiveBundleIndexDict:
|
||||
"""
|
||||
Merge one bundle into its shard index.
|
||||
|
||||
Retries for a known manifest set ``require_existing_index`` so they never create a partial index by scanning
|
||||
or overwriting a shard whose historical entries cannot be proven from this one manifest.
|
||||
"""
|
||||
index_key = self._get_index_object_key(identity)
|
||||
if storage.object_exists(index_key):
|
||||
index = self._load_bundle_index(storage, identity)
|
||||
elif require_existing_index:
|
||||
raise RuntimeError(f"archive shard index missing while reconciling known manifest: {index_key}")
|
||||
else:
|
||||
index = self._build_bundle_index(storage, identity)
|
||||
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
"""
|
||||
Maintain V2 workflow-run archive bundles.
|
||||
|
||||
Archive V2 keeps object-store manifests as the recoverable bundle source of truth. This maintenance module still
|
||||
discovers delete/restore targets by listing `manifest.json` objects and uses object-store marker files for
|
||||
delete/restore state. The separate database bundle index is intended for console listing and download jobs, not as the
|
||||
source of truth for destructive maintenance.
|
||||
Archive V2 keeps object-store manifests as the recoverable bundle source of truth. Delete and restore discover a
|
||||
bounded page of candidates from `workflow_run_archive_bundles`, optionally restricted to one exact archive shard, then
|
||||
construct each immutable manifest key from the catalog identity. They never list the object-store namespace.
|
||||
Object-store marker files keep delete/restore idempotent, while the caller persists a non-dry-run cursor only after a
|
||||
candidate succeeds.
|
||||
|
||||
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.
|
||||
not-updated case. Restore never skips a bundle with a missing deleted marker when deletion started or source rows have
|
||||
drifted, so an external cursor cannot pass an interrupted delete.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
@@ -38,6 +40,7 @@ from models.workflow import (
|
||||
WorkflowPause,
|
||||
WorkflowPauseReason,
|
||||
WorkflowRun,
|
||||
WorkflowRunArchiveBundle,
|
||||
)
|
||||
from services.retention.workflow_run.constants import (
|
||||
ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME,
|
||||
@@ -51,7 +54,6 @@ from services.retention.workflow_run.constants import (
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ARCHIVE_ROOT_PREFIX = "workflow-runs/v2/"
|
||||
_CHUNK_SIZE = 5_000
|
||||
|
||||
|
||||
@@ -84,11 +86,28 @@ class BundleManifest(TypedDict):
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BundleReference:
|
||||
"""Object-store reference for one V2 archive bundle."""
|
||||
class ArchiveBundleCatalogEntry:
|
||||
"""Immutable catalog identity and manifest-derived metrics for one V2 archive bundle."""
|
||||
|
||||
catalog_id: str
|
||||
tenant_id: str
|
||||
year: int
|
||||
month: int
|
||||
shard: str
|
||||
bundle_id: str
|
||||
workflow_run_count: int
|
||||
row_count: int
|
||||
archive_bytes: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BundleReference:
|
||||
"""Verified object-store reference for one catalog candidate."""
|
||||
|
||||
catalog: ArchiveBundleCatalogEntry
|
||||
object_prefix: str
|
||||
manifest_key: str
|
||||
manifest_size_bytes: int
|
||||
manifest: BundleManifest
|
||||
|
||||
|
||||
@@ -96,6 +115,7 @@ class BundleReference:
|
||||
class BundleOperationResult:
|
||||
"""Result for one V2 bundle delete or restore operation."""
|
||||
|
||||
catalog_id: str
|
||||
bundle_id: str
|
||||
tenant_id: str
|
||||
object_prefix: str
|
||||
@@ -128,6 +148,8 @@ class BundleOperationSummary:
|
||||
archive_bytes: int = 0
|
||||
elapsed_time: float = 0.0
|
||||
validation_time: float = 0.0
|
||||
next_catalog_id: str | None = None
|
||||
preview_next_catalog_id: str | None = None
|
||||
table_counts: dict[str, int] = field(default_factory=dict)
|
||||
results: list[BundleOperationResult] = field(default_factory=list)
|
||||
|
||||
@@ -180,162 +202,334 @@ RESTORE_ORDER = [
|
||||
"workflow_trigger_logs",
|
||||
]
|
||||
|
||||
DELETE_ORDER = [
|
||||
"workflow_pause_reasons",
|
||||
"workflow_node_execution_offload",
|
||||
"workflow_trigger_logs",
|
||||
"workflow_app_logs",
|
||||
"workflow_node_executions",
|
||||
"workflow_pauses",
|
||||
"workflow_runs",
|
||||
]
|
||||
|
||||
|
||||
class WorkflowRunBundleArchiveMaintenance:
|
||||
"""
|
||||
Delete and restore V2 workflow-run archive bundles.
|
||||
|
||||
Delete accepts already-missing source rows only when every remaining row in the bundle scope is an unchanged
|
||||
archive subset. It then removes only those verified primary keys, so unrelated or unarchived rows fail closed.
|
||||
Non-dry-run delete and restore serialize on the existing archive catalog row before checking markers or changing
|
||||
source rows.
|
||||
|
||||
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.
|
||||
strict_content_validation: Compare restored source-table content checksums against Parquet content. Delete
|
||||
always validates that every remaining live row belongs to and matches the archive before removing it.
|
||||
storage: Optional archive storage implementation. Tests may provide an in-memory implementation.
|
||||
session_factory: Optional session factory. Each candidate is processed in its own transaction.
|
||||
|
||||
Batches stop at the first error so a returned cursor cannot pass an unhandled candidate.
|
||||
"""
|
||||
|
||||
dry_run: bool
|
||||
strict_content_validation: bool
|
||||
stop_on_error: bool
|
||||
storage: ArchiveStorage | None
|
||||
session_factory: sessionmaker[Session]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dry_run: bool = False,
|
||||
strict_content_validation: bool = True,
|
||||
stop_on_error: bool = True,
|
||||
storage: ArchiveStorage | None = None,
|
||||
session_factory: sessionmaker[Session] | None = None,
|
||||
) -> None:
|
||||
self.dry_run = dry_run
|
||||
self.strict_content_validation = strict_content_validation
|
||||
self.stop_on_error = stop_on_error
|
||||
self.storage = storage
|
||||
self.session_factory = session_factory or sessionmaker(bind=db.engine, expire_on_commit=False)
|
||||
|
||||
def delete_batch(
|
||||
self,
|
||||
*,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
target_year: int,
|
||||
target_month: int,
|
||||
after_catalog_id: str | None,
|
||||
limit: int,
|
||||
shard: str | None = None,
|
||||
) -> BundleOperationSummary:
|
||||
"""Validate and delete source rows for archived V2 bundles in the requested created_at window."""
|
||||
"""Validate and delete one keyset page, optionally scoped to an exact archive shard."""
|
||||
return self._process_batch(
|
||||
operation="delete",
|
||||
tenant_ids=tenant_ids,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
target_year=target_year,
|
||||
target_month=target_month,
|
||||
after_catalog_id=after_catalog_id,
|
||||
limit=limit,
|
||||
shard=shard,
|
||||
)
|
||||
|
||||
def restore_batch(
|
||||
self,
|
||||
*,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
target_year: int,
|
||||
target_month: int,
|
||||
after_catalog_id: str | None,
|
||||
limit: int,
|
||||
) -> BundleOperationSummary:
|
||||
"""Restore source rows for deleted V2 bundles in the requested created_at window."""
|
||||
"""Restore source rows for a keyset page of deleted V2 bundles in one calendar month."""
|
||||
return self._process_batch(
|
||||
operation="restore",
|
||||
tenant_ids=tenant_ids,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
target_year=target_year,
|
||||
target_month=target_month,
|
||||
after_catalog_id=after_catalog_id,
|
||||
limit=limit,
|
||||
shard=None,
|
||||
)
|
||||
|
||||
def validate_catalog_shards(
|
||||
self,
|
||||
*,
|
||||
target_year: int,
|
||||
target_month: int,
|
||||
shard_total: int,
|
||||
tenant_ids: Sequence[str] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Fail before a parallel delete when the requested closed-month scope contains a different shard layout.
|
||||
|
||||
A subset of the expected shards is valid because an archive shard may legitimately contain no bundles. Any
|
||||
other shard name indicates a historical or mixed-layout month that must be handled by the serial delete path.
|
||||
"""
|
||||
if not 1 <= shard_total <= 16:
|
||||
raise ValueError("shard_total must be between 1 and 16")
|
||||
expected_shards = tuple(f"{index:02d}-of-{shard_total:02d}" for index in range(shard_total))
|
||||
conditions = [
|
||||
WorkflowRunArchiveBundle.year == target_year,
|
||||
WorkflowRunArchiveBundle.month == target_month,
|
||||
]
|
||||
if tenant_ids is not None:
|
||||
conditions.append(WorkflowRunArchiveBundle.tenant_id.in_(tenant_ids))
|
||||
statement = (
|
||||
select(WorkflowRunArchiveBundle.shard)
|
||||
.where(
|
||||
*conditions,
|
||||
WorkflowRunArchiveBundle.shard.not_in(expected_shards),
|
||||
)
|
||||
.distinct()
|
||||
.order_by(WorkflowRunArchiveBundle.shard.asc())
|
||||
)
|
||||
with self.session_factory() as session:
|
||||
unexpected_shards = list(session.scalars(statement))
|
||||
if unexpected_shards:
|
||||
raise ValueError(
|
||||
"archive catalog month contains unexpected shards for "
|
||||
f"{shard_total}-way delete: {', '.join(unexpected_shards)}"
|
||||
)
|
||||
|
||||
def _process_batch(
|
||||
self,
|
||||
*,
|
||||
operation: str,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
target_year: int,
|
||||
target_month: int,
|
||||
after_catalog_id: str | None,
|
||||
limit: int,
|
||||
shard: str | None,
|
||||
) -> 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,
|
||||
storage = self.storage or self._get_archive_storage()
|
||||
catalog_entries = self._list_catalog_entries(
|
||||
tenant_ids=tenant_ids,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
target_year=target_year,
|
||||
target_month=target_month,
|
||||
after_catalog_id=after_catalog_id,
|
||||
limit=limit,
|
||||
shard=shard,
|
||||
)
|
||||
|
||||
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}")
|
||||
logger.info(
|
||||
"Found %s V2 archive catalog candidates for %s: year=%s month=%s shard=%s after_catalog_id=%s",
|
||||
len(catalog_entries),
|
||||
operation,
|
||||
target_year,
|
||||
target_month,
|
||||
shard,
|
||||
after_catalog_id,
|
||||
)
|
||||
for catalog_entry in catalog_entries:
|
||||
try:
|
||||
bundle_ref = self._build_bundle_reference(storage, catalog_entry)
|
||||
with self.session_factory() 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}")
|
||||
except Exception as exc:
|
||||
result = self._new_result_from_catalog_entry(catalog_entry)
|
||||
result.error = str(exc)
|
||||
logger.exception(
|
||||
"Failed to prepare V2 archive bundle %s from catalog %s",
|
||||
catalog_entry.bundle_id,
|
||||
catalog_entry.catalog_id,
|
||||
)
|
||||
|
||||
self._merge_result(summary, result)
|
||||
if not result.success and self.stop_on_error:
|
||||
if result.success:
|
||||
if self.dry_run:
|
||||
summary.preview_next_catalog_id = catalog_entry.catalog_id
|
||||
else:
|
||||
summary.next_catalog_id = catalog_entry.catalog_id
|
||||
else:
|
||||
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(
|
||||
def _list_catalog_entries(
|
||||
self,
|
||||
storage: ArchiveStorage,
|
||||
*,
|
||||
operation: str,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
start_date: datetime.datetime,
|
||||
end_date: datetime.datetime,
|
||||
target_year: int,
|
||||
target_month: int,
|
||||
after_catalog_id: str | None,
|
||||
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))
|
||||
shard: str | None = None,
|
||||
) -> list[ArchiveBundleCatalogEntry]:
|
||||
"""Read one bounded, stable-order candidate page from the database catalog."""
|
||||
conditions = [
|
||||
WorkflowRunArchiveBundle.year == target_year,
|
||||
WorkflowRunArchiveBundle.month == target_month,
|
||||
]
|
||||
if tenant_ids:
|
||||
conditions.append(WorkflowRunArchiveBundle.tenant_id.in_(tenant_ids))
|
||||
if shard is not None:
|
||||
conditions.append(WorkflowRunArchiveBundle.shard == shard)
|
||||
if after_catalog_id:
|
||||
conditions.append(WorkflowRunArchiveBundle.id > after_catalog_id)
|
||||
|
||||
refs.sort(
|
||||
key=lambda ref: (
|
||||
self._parse_manifest_datetime(ref.manifest["min_created_at"]),
|
||||
ref.manifest["tenant_id"],
|
||||
ref.manifest["bundle_id"],
|
||||
)
|
||||
statement = (
|
||||
select(WorkflowRunArchiveBundle).where(*conditions).order_by(WorkflowRunArchiveBundle.id.asc()).limit(limit)
|
||||
)
|
||||
return refs[:limit]
|
||||
with self.session_factory() as session:
|
||||
if after_catalog_id:
|
||||
cursor_bundle = session.get(WorkflowRunArchiveBundle, after_catalog_id)
|
||||
self._validate_catalog_cursor_scope(
|
||||
cursor_bundle,
|
||||
tenant_ids=tenant_ids,
|
||||
target_year=target_year,
|
||||
target_month=target_month,
|
||||
shard=shard,
|
||||
)
|
||||
bundles = list(session.scalars(statement))
|
||||
return [
|
||||
ArchiveBundleCatalogEntry(
|
||||
catalog_id=bundle.id,
|
||||
tenant_id=bundle.tenant_id,
|
||||
year=bundle.year,
|
||||
month=bundle.month,
|
||||
shard=bundle.shard,
|
||||
bundle_id=bundle.bundle_id,
|
||||
workflow_run_count=bundle.workflow_run_count,
|
||||
row_count=bundle.row_count,
|
||||
archive_bytes=bundle.archive_bytes,
|
||||
)
|
||||
for bundle in bundles
|
||||
]
|
||||
|
||||
@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 _validate_catalog_cursor_scope(
|
||||
cursor_bundle: WorkflowRunArchiveBundle | None,
|
||||
*,
|
||||
tenant_ids: Sequence[str] | None,
|
||||
target_year: int,
|
||||
target_month: int,
|
||||
shard: str | None,
|
||||
) -> None:
|
||||
"""Reject a keyset cursor that cannot safely represent the requested catalog scope."""
|
||||
if cursor_bundle is None:
|
||||
raise ValueError("after_catalog_id does not exist in the workflow run archive bundle catalog")
|
||||
if cursor_bundle.year != target_year or cursor_bundle.month != target_month:
|
||||
raise ValueError("after_catalog_id is outside the requested archive month")
|
||||
if tenant_ids is not None and cursor_bundle.tenant_id not in tenant_ids:
|
||||
raise ValueError("after_catalog_id is outside the requested tenant scope")
|
||||
if shard is not None and cursor_bundle.shard != shard:
|
||||
raise ValueError("after_catalog_id is outside the requested archive shard")
|
||||
|
||||
def _build_bundle_reference(
|
||||
self,
|
||||
storage: ArchiveStorage,
|
||||
catalog_entry: ArchiveBundleCatalogEntry,
|
||||
) -> BundleReference:
|
||||
"""Load and bind one manifest to the catalog row that selected it."""
|
||||
object_prefix = self._catalog_object_prefix(catalog_entry)
|
||||
manifest_key = f"{object_prefix}/{ARCHIVE_BUNDLE_MANIFEST_NAME}"
|
||||
manifest_data = storage.get_object(manifest_key)
|
||||
manifest = self._load_and_validate_manifest(manifest_data, object_prefix=object_prefix)
|
||||
self._validate_manifest_catalog_identity(manifest, catalog_entry)
|
||||
return BundleReference(
|
||||
catalog=catalog_entry,
|
||||
object_prefix=object_prefix,
|
||||
manifest_key=manifest_key,
|
||||
manifest_size_bytes=len(manifest_data),
|
||||
manifest=manifest,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _catalog_object_prefix(catalog_entry: ArchiveBundleCatalogEntry) -> str:
|
||||
"""Construct the immutable V2 bundle prefix from its database catalog identity."""
|
||||
if not catalog_entry.tenant_id:
|
||||
raise ValueError("archive catalog tenant_id must not be empty")
|
||||
if not 1 <= catalog_entry.month <= 12:
|
||||
raise ValueError(f"archive catalog month is invalid: {catalog_entry.month}")
|
||||
return (
|
||||
f"workflow-runs/v2/tenant_prefix={catalog_entry.tenant_id[0].lower()}/"
|
||||
f"tenant_id={catalog_entry.tenant_id}/year={catalog_entry.year:04d}/"
|
||||
f"month={catalog_entry.month:02d}/shard={catalog_entry.shard}/bundle={catalog_entry.bundle_id}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_manifest_catalog_identity(
|
||||
manifest: BundleManifest,
|
||||
catalog_entry: ArchiveBundleCatalogEntry,
|
||||
) -> None:
|
||||
"""Fail closed when the catalog locator and manifest identify different immutable bundles."""
|
||||
expected_identity = (
|
||||
catalog_entry.tenant_id,
|
||||
catalog_entry.year,
|
||||
catalog_entry.month,
|
||||
catalog_entry.shard,
|
||||
catalog_entry.bundle_id,
|
||||
)
|
||||
manifest_identity = (
|
||||
manifest["tenant_id"],
|
||||
manifest["year"],
|
||||
manifest["month"],
|
||||
manifest["shard"],
|
||||
manifest["bundle_id"],
|
||||
)
|
||||
if manifest_identity != expected_identity:
|
||||
raise ValueError(
|
||||
f"archive manifest identity does not match catalog: expected={expected_identity}, "
|
||||
f"actual={manifest_identity}"
|
||||
)
|
||||
if manifest["workflow_run_count"] != catalog_entry.workflow_run_count:
|
||||
raise ValueError("archive manifest workflow_run_count does not match catalog")
|
||||
manifest_row_count = sum(table["row_count"] for table in manifest["tables"].values())
|
||||
if manifest_row_count != catalog_entry.row_count:
|
||||
raise ValueError("archive manifest row_count does not match catalog")
|
||||
|
||||
def _delete_bundle(
|
||||
self,
|
||||
@@ -344,37 +538,61 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
bundle_ref: BundleReference,
|
||||
) -> BundleOperationResult:
|
||||
start_time = time.time()
|
||||
result = self._new_result(bundle_ref.manifest)
|
||||
result = self._new_result(bundle_ref.manifest, bundle_ref.catalog.catalog_id)
|
||||
try:
|
||||
validation_start = time.time()
|
||||
if not self.dry_run:
|
||||
self._lock_catalog_entry(session, bundle_ref.catalog)
|
||||
if self._is_restore_started(storage, bundle_ref.object_prefix):
|
||||
raise ValueError("restore started marker exists; reconcile restore before delete")
|
||||
|
||||
deleted_marker_exists = self._is_deleted(storage, bundle_ref.object_prefix)
|
||||
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
|
||||
):
|
||||
live_records = self._load_live_bundle_records(
|
||||
session,
|
||||
manifest,
|
||||
table_records,
|
||||
lock=not self.dry_run,
|
||||
)
|
||||
if deleted_marker_exists:
|
||||
live_counts = {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES}
|
||||
if any(live_counts.values()):
|
||||
raise ValueError(f"Live rows exist for bundle with deleted marker: {live_counts}")
|
||||
if not self.dry_run:
|
||||
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.validation_time = time.time() - validation_start
|
||||
result.success = True
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
self._validate_live_archive_subset(manifest, table_records, live_records)
|
||||
result.validation_time = time.time() - validation_start
|
||||
|
||||
if not any(live_records[table_name] for table_name in ARCHIVED_TABLES):
|
||||
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)
|
||||
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
|
||||
result.success = True
|
||||
result.elapsed_time = time.time() - start_time
|
||||
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:
|
||||
expected_deleted_counts = {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES}
|
||||
deleted_counts = self._delete_bundle_rows(session, live_records)
|
||||
if deleted_counts != expected_deleted_counts:
|
||||
raise ValueError(
|
||||
f"Deleted row count mismatch: expected={result.table_counts}, actual={deleted_counts}"
|
||||
f"Deleted row count mismatch: expected={expected_deleted_counts}, actual={deleted_counts}"
|
||||
)
|
||||
self._validate_live_counts(session, manifest, expected_present=False)
|
||||
remaining_records = self._load_live_bundle_records(session, manifest, table_records, lock=True)
|
||||
remaining_counts = {table_name: len(remaining_records[table_name]) for table_name in ARCHIVED_TABLES}
|
||||
if any(remaining_counts.values()):
|
||||
raise ValueError(f"Live rows remain after bundle delete: {remaining_counts}")
|
||||
session.commit()
|
||||
self._mark_deleted(storage, bundle_ref.object_prefix)
|
||||
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
|
||||
@@ -394,9 +612,25 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
bundle_ref: BundleReference,
|
||||
) -> BundleOperationResult:
|
||||
start_time = time.time()
|
||||
result = self._new_result(bundle_ref.manifest)
|
||||
result = self._new_result(bundle_ref.manifest, bundle_ref.catalog.catalog_id)
|
||||
try:
|
||||
validation_start = time.time()
|
||||
if not self.dry_run:
|
||||
self._lock_catalog_entry(session, bundle_ref.catalog)
|
||||
if not self._is_deleted(storage, bundle_ref.object_prefix):
|
||||
# A committed delete may be interrupted before `.deleted` is written. Do not let restore advance its
|
||||
# cursor over that state: retry delete to reconcile it, or investigate source-row drift first.
|
||||
if self._is_delete_started(storage, bundle_ref.object_prefix):
|
||||
raise ValueError("delete started marker exists without a deleted marker; reconcile delete first")
|
||||
restore_started = self._is_restore_started(storage, bundle_ref.object_prefix)
|
||||
self._validate_live_counts(session, bundle_ref.manifest, expected_present=True)
|
||||
result.validation_time = time.time() - validation_start
|
||||
if restore_started and not self.dry_run:
|
||||
self._mark_restored(storage, bundle_ref.object_prefix)
|
||||
result.success = True
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
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
|
||||
@@ -408,6 +642,7 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
if not self.dry_run:
|
||||
self._mark_restored(storage, bundle_ref.object_prefix)
|
||||
result.success = True
|
||||
result.elapsed_time = time.time() - start_time
|
||||
return result
|
||||
|
||||
self._validate_live_counts(session, manifest, expected_present=False)
|
||||
@@ -432,13 +667,40 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _new_result(manifest: BundleManifest) -> BundleOperationResult:
|
||||
def _new_result(manifest: BundleManifest, catalog_id: str) -> BundleOperationResult:
|
||||
return BundleOperationResult(
|
||||
catalog_id=catalog_id,
|
||||
bundle_id=manifest["bundle_id"],
|
||||
tenant_id=manifest["tenant_id"],
|
||||
object_prefix=manifest["object_prefix"],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _new_result_from_catalog_entry(catalog_entry: ArchiveBundleCatalogEntry) -> BundleOperationResult:
|
||||
try:
|
||||
object_prefix = WorkflowRunBundleArchiveMaintenance._catalog_object_prefix(catalog_entry)
|
||||
except ValueError:
|
||||
object_prefix = "<invalid archive catalog identity>"
|
||||
return BundleOperationResult(
|
||||
catalog_id=catalog_entry.catalog_id,
|
||||
bundle_id=catalog_entry.bundle_id,
|
||||
tenant_id=catalog_entry.tenant_id,
|
||||
object_prefix=object_prefix,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _lock_catalog_entry(session: Session, catalog_entry: ArchiveBundleCatalogEntry) -> None:
|
||||
locked_catalog_id = session.scalar(
|
||||
select(WorkflowRunArchiveBundle.id)
|
||||
.where(
|
||||
WorkflowRunArchiveBundle.id == catalog_entry.catalog_id,
|
||||
WorkflowRunArchiveBundle.tenant_id == catalog_entry.tenant_id,
|
||||
)
|
||||
.with_for_update()
|
||||
)
|
||||
if locked_catalog_id is None:
|
||||
raise ValueError("archive catalog row disappeared before bundle maintenance")
|
||||
|
||||
def _validate_archive_object(
|
||||
self,
|
||||
storage: ArchiveStorage,
|
||||
@@ -446,7 +708,7 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
) -> 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))
|
||||
total_size = bundle_ref.manifest_size_bytes
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
info = manifest["tables"][table_name]
|
||||
payload = self._get_checked_object(storage, info["object_key"])
|
||||
@@ -469,12 +731,14 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
f"expected={info['row_count']}, actual={len(records)}"
|
||||
)
|
||||
table_records[table_name] = records
|
||||
if total_size != bundle_ref.catalog.archive_bytes:
|
||||
raise ValueError(
|
||||
f"Archive object total size mismatch: expected={bundle_ref.catalog.archive_bytes}, actual={total_size}"
|
||||
)
|
||||
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
|
||||
@@ -539,6 +803,110 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
table = pq.read_table(io.BytesIO(payload))
|
||||
return table.to_pylist()
|
||||
|
||||
def _load_live_bundle_records(
|
||||
self,
|
||||
session: Session,
|
||||
manifest: BundleManifest,
|
||||
table_records: dict[str, list[dict[str, Any]]],
|
||||
*,
|
||||
lock: bool,
|
||||
) -> dict[str, list[dict[str, Any]]]:
|
||||
"""Load the complete live scope for a bundle, including archived rows whose relationship fields drifted."""
|
||||
run_ids = manifest["run_ids"]
|
||||
archive_ids = {
|
||||
table_name: [str(record["id"]) for record in table_records[table_name]] for table_name in ARCHIVED_TABLES
|
||||
}
|
||||
live_node_ids = self._select_ids_by_run_ids(session, WorkflowNodeExecutionModel, run_ids)
|
||||
live_pause_ids = self._select_ids_by_run_ids(session, WorkflowPause, run_ids)
|
||||
node_ids = sorted(set(archive_ids["workflow_node_executions"]) | set(live_node_ids))
|
||||
pause_ids = sorted(set(archive_ids["workflow_pauses"]) | set(live_pause_ids))
|
||||
|
||||
def load_scope(
|
||||
table_name: str,
|
||||
model: Any,
|
||||
scope_column: Any,
|
||||
scope_ids: Sequence[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
return self._merge_records_by_id(
|
||||
self._load_records_by_column(session, model, scope_column, scope_ids, lock=lock),
|
||||
self._load_records_by_column(session, model, model.id, archive_ids[table_name], lock=lock),
|
||||
)
|
||||
|
||||
return {
|
||||
"workflow_pause_reasons": load_scope(
|
||||
"workflow_pause_reasons", WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids
|
||||
),
|
||||
"workflow_node_execution_offload": load_scope(
|
||||
"workflow_node_execution_offload",
|
||||
WorkflowNodeExecutionOffload,
|
||||
WorkflowNodeExecutionOffload.node_execution_id,
|
||||
node_ids,
|
||||
),
|
||||
"workflow_trigger_logs": load_scope(
|
||||
"workflow_trigger_logs", WorkflowTriggerLog, WorkflowTriggerLog.workflow_run_id, run_ids
|
||||
),
|
||||
"workflow_app_logs": load_scope(
|
||||
"workflow_app_logs", WorkflowAppLog, WorkflowAppLog.workflow_run_id, run_ids
|
||||
),
|
||||
"workflow_node_executions": load_scope(
|
||||
"workflow_node_executions",
|
||||
WorkflowNodeExecutionModel,
|
||||
WorkflowNodeExecutionModel.workflow_run_id,
|
||||
run_ids,
|
||||
),
|
||||
"workflow_pauses": load_scope("workflow_pauses", WorkflowPause, WorkflowPause.workflow_run_id, run_ids),
|
||||
"workflow_runs": self._load_records_by_column(
|
||||
session,
|
||||
WorkflowRun,
|
||||
WorkflowRun.id,
|
||||
sorted(set(run_ids) | set(archive_ids["workflow_runs"])),
|
||||
lock=lock,
|
||||
),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _validate_live_archive_subset(
|
||||
cls,
|
||||
manifest: BundleManifest,
|
||||
table_records: dict[str, list[dict[str, Any]]],
|
||||
live_records: dict[str, list[dict[str, Any]]],
|
||||
) -> None:
|
||||
"""Require every live row in the bundle scope to exist unchanged in the validated archive."""
|
||||
manifest_run_ids = {str(run_id) for run_id in manifest["run_ids"]}
|
||||
if len(manifest_run_ids) != len(manifest["run_ids"]):
|
||||
raise ValueError("archive manifest contains duplicate workflow run IDs")
|
||||
|
||||
archive_records_by_id: dict[str, dict[str, dict[str, Any]]] = {}
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
records_by_id = {str(record["id"]): record for record in table_records[table_name]}
|
||||
if len(records_by_id) != len(table_records[table_name]):
|
||||
raise ValueError(f"archive contains duplicate row IDs for {table_name}")
|
||||
archive_records_by_id[table_name] = records_by_id
|
||||
|
||||
if set(archive_records_by_id["workflow_runs"]) != manifest_run_ids:
|
||||
raise ValueError("archive workflow run IDs do not match manifest run_ids")
|
||||
|
||||
for table_name in ARCHIVED_TABLES:
|
||||
live_ids = [str(record["id"]) for record in live_records[table_name]]
|
||||
if len(set(live_ids)) != len(live_ids):
|
||||
raise ValueError(f"live scope contains duplicate row IDs for {table_name}")
|
||||
|
||||
archive_by_id = archive_records_by_id[table_name]
|
||||
extra_ids = sorted(set(live_ids) - set(archive_by_id))
|
||||
if extra_ids:
|
||||
raise ValueError(
|
||||
f"Live bundle scope contains rows missing from archive for {table_name}: {extra_ids[:10]}"
|
||||
)
|
||||
|
||||
archive_subset = [archive_by_id[row_id] for row_id in live_ids]
|
||||
live_checksum = cls._records_checksum(live_records[table_name])
|
||||
archive_checksum = cls._records_checksum(archive_subset)
|
||||
if live_checksum != archive_checksum:
|
||||
raise ValueError(
|
||||
f"Live/archive subset content checksum mismatch for {table_name}: "
|
||||
f"expected={archive_checksum}, actual={live_checksum}"
|
||||
)
|
||||
|
||||
def _validate_live_counts(
|
||||
self,
|
||||
session: Session,
|
||||
@@ -617,26 +985,13 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
def _delete_bundle_rows(
|
||||
self,
|
||||
session: Session,
|
||||
table_records: dict[str, list[dict[str, Any]]],
|
||||
live_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)
|
||||
for table_name in DELETE_ORDER:
|
||||
model = TABLE_MODELS[table_name]
|
||||
row_ids = [str(record["id"]) for record in live_records[table_name]]
|
||||
deleted_counts[table_name] = self._delete_by_column(session, model, model.id, row_ids)
|
||||
return deleted_counts
|
||||
|
||||
def _restore_bundle_rows(
|
||||
@@ -708,11 +1063,6 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
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,
|
||||
@@ -757,8 +1107,16 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
session: Session,
|
||||
model: Any,
|
||||
run_ids: Sequence[str],
|
||||
*,
|
||||
lock: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
return self._load_records_by_column(session, model, self._run_id_column(model), run_ids)
|
||||
return self._load_records_by_column(
|
||||
session,
|
||||
model,
|
||||
self._run_id_column(model),
|
||||
run_ids,
|
||||
lock=lock,
|
||||
)
|
||||
|
||||
def _load_records_by_column(
|
||||
self,
|
||||
@@ -766,23 +1124,23 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
model: Any,
|
||||
column: Any,
|
||||
values: Sequence[str],
|
||||
*,
|
||||
lock: bool = False,
|
||||
) -> 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))))
|
||||
statement = select(model).where(column.in_(chunk)).order_by(model.id.asc())
|
||||
if lock:
|
||||
statement = statement.with_for_update()
|
||||
rows.extend(session.scalars(statement))
|
||||
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
|
||||
)
|
||||
def _merge_records_by_id(*record_groups: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
records_by_id = {str(record["id"]): record for record_group in record_groups for record in record_group}
|
||||
return [records_by_id[row_id] for row_id in sorted(records_by_id)]
|
||||
|
||||
@staticmethod
|
||||
def _run_id_column(model: Any) -> Any:
|
||||
@@ -813,6 +1171,10 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
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 _is_restore_started(storage: ArchiveStorage, object_prefix: str) -> bool:
|
||||
return storage.object_exists(f"{object_prefix}/{ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME}")
|
||||
|
||||
@staticmethod
|
||||
def _mark_deleted(storage: ArchiveStorage, object_prefix: str) -> None:
|
||||
WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME)
|
||||
@@ -820,10 +1182,13 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
@staticmethod
|
||||
def _mark_restored(storage: ArchiveStorage, object_prefix: str) -> None:
|
||||
WorkflowRunBundleArchiveMaintenance._delete_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME)
|
||||
WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
|
||||
WorkflowRunBundleArchiveMaintenance._delete_marker(
|
||||
storage, object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_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:
|
||||
@@ -836,16 +1201,6 @@ class WorkflowRunBundleArchiveMaintenance:
|
||||
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)]
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
"""Account-level Step-by-step Tour persistence."""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import NotRequired, TypedDict
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.orm import Session, scoped_session
|
||||
|
||||
from configs import dify_config
|
||||
from libs.datetime_utils import ensure_naive_utc
|
||||
from models.account import Account
|
||||
from models.onboarding import AccountStepByStepTourState
|
||||
|
||||
STEP_BY_STEP_TOUR_TASK_IDS = frozenset(("home", "studio", "knowledge", "integration"))
|
||||
|
||||
|
||||
class StepByStepTourStateResponse(TypedDict):
|
||||
first_workspace_id: str | None
|
||||
skipped: bool
|
||||
completed_task_ids: list[str]
|
||||
manually_enabled_workspace_ids: list[str]
|
||||
manually_disabled_workspace_ids: list[str]
|
||||
updated_at: datetime | None
|
||||
|
||||
|
||||
class StepByStepTourPatch(TypedDict):
|
||||
action: str
|
||||
task_id: NotRequired[str | None]
|
||||
|
||||
|
||||
class StepByStepTourService:
|
||||
"""Coordinate persisted tour state with account eligibility rules."""
|
||||
|
||||
@classmethod
|
||||
def get_state(
|
||||
cls,
|
||||
*,
|
||||
account: Account,
|
||||
current_tenant_id: str,
|
||||
session: Session | scoped_session,
|
||||
) -> StepByStepTourStateResponse:
|
||||
eligible = cls.is_eligible(account)
|
||||
state = cls._get_state(account.id, session=session)
|
||||
|
||||
if eligible:
|
||||
state = cls._ensure_state(account.id, session=session, state=state)
|
||||
if state.first_workspace_id is None:
|
||||
state.first_workspace_id = current_tenant_id
|
||||
session.commit()
|
||||
session.refresh(state)
|
||||
|
||||
return cls._build_response(state=state)
|
||||
|
||||
@classmethod
|
||||
def patch_state(
|
||||
cls,
|
||||
*,
|
||||
account: Account,
|
||||
current_tenant_id: str,
|
||||
patch: StepByStepTourPatch,
|
||||
session: Session | scoped_session,
|
||||
) -> StepByStepTourStateResponse:
|
||||
state = cls._ensure_state(account.id, session=session, state=None)
|
||||
cls._apply_action(
|
||||
state=state,
|
||||
action=patch["action"],
|
||||
task_id=patch.get("task_id"),
|
||||
current_tenant_id=current_tenant_id,
|
||||
)
|
||||
|
||||
session.commit()
|
||||
session.refresh(state)
|
||||
return cls._build_response(state=state)
|
||||
|
||||
@classmethod
|
||||
def is_eligible(cls, account: Account) -> bool:
|
||||
if not dify_config.ENABLE_STEP_BY_STEP_TOUR:
|
||||
return False
|
||||
|
||||
rollout_started_at = dify_config.STEP_BY_STEP_TOUR_ROLLOUT_STARTED_AT
|
||||
if rollout_started_at is None:
|
||||
return False
|
||||
|
||||
account_started_at = account.initialized_at or account.created_at
|
||||
if account_started_at is None:
|
||||
return False
|
||||
|
||||
return ensure_naive_utc(account_started_at) >= ensure_naive_utc(rollout_started_at)
|
||||
|
||||
@classmethod
|
||||
def _get_state(
|
||||
cls,
|
||||
account_id: str,
|
||||
*,
|
||||
session: Session | scoped_session,
|
||||
) -> AccountStepByStepTourState | None:
|
||||
stmt = select(AccountStepByStepTourState).where(AccountStepByStepTourState.account_id == account_id).limit(1)
|
||||
return session.execute(stmt).scalar_one_or_none()
|
||||
|
||||
@classmethod
|
||||
def _ensure_state(
|
||||
cls,
|
||||
account_id: str,
|
||||
*,
|
||||
session: Session | scoped_session,
|
||||
state: AccountStepByStepTourState | None,
|
||||
) -> AccountStepByStepTourState:
|
||||
if state is None:
|
||||
state = cls._get_state(account_id, session=session)
|
||||
if state is not None:
|
||||
return state
|
||||
|
||||
state = AccountStepByStepTourState(account_id=account_id)
|
||||
session.add(state)
|
||||
try:
|
||||
session.flush()
|
||||
except IntegrityError:
|
||||
# Another tab/device can create the account row between our read and insert.
|
||||
session.rollback()
|
||||
state = cls._get_state(account_id, session=session)
|
||||
if state is None:
|
||||
raise
|
||||
return state
|
||||
|
||||
@classmethod
|
||||
def _apply_action(
|
||||
cls,
|
||||
*,
|
||||
state: AccountStepByStepTourState,
|
||||
action: str,
|
||||
task_id: str | None,
|
||||
current_tenant_id: str,
|
||||
) -> None:
|
||||
match action:
|
||||
case "skip":
|
||||
state.skipped = True
|
||||
state.manually_enabled_workspace_ids = cls._remove_id(
|
||||
state.manually_enabled_workspace_ids,
|
||||
current_tenant_id,
|
||||
)
|
||||
case "complete_task":
|
||||
if task_id is None:
|
||||
raise ValueError("task_id is required")
|
||||
cls._validate_task_id(task_id)
|
||||
state.completed_task_ids = cls._add_id(state.completed_task_ids, task_id)
|
||||
case "uncomplete_task":
|
||||
if task_id is None:
|
||||
raise ValueError("task_id is required")
|
||||
cls._validate_task_id(task_id)
|
||||
state.completed_task_ids = cls._remove_id(state.completed_task_ids, task_id)
|
||||
case "enable_current_workspace":
|
||||
state.skipped = False
|
||||
state.manually_enabled_workspace_ids = cls._add_id(
|
||||
state.manually_enabled_workspace_ids,
|
||||
current_tenant_id,
|
||||
)
|
||||
state.manually_disabled_workspace_ids = cls._remove_id(
|
||||
state.manually_disabled_workspace_ids,
|
||||
current_tenant_id,
|
||||
)
|
||||
case "disable_current_workspace":
|
||||
state.manually_enabled_workspace_ids = cls._remove_id(
|
||||
state.manually_enabled_workspace_ids,
|
||||
current_tenant_id,
|
||||
)
|
||||
state.manually_disabled_workspace_ids = cls._add_id(
|
||||
state.manually_disabled_workspace_ids,
|
||||
current_tenant_id,
|
||||
)
|
||||
case _:
|
||||
raise ValueError(f"Unsupported action: {action}")
|
||||
|
||||
@classmethod
|
||||
def _build_response(
|
||||
cls,
|
||||
*,
|
||||
state: AccountStepByStepTourState | None,
|
||||
) -> StepByStepTourStateResponse:
|
||||
if state is None:
|
||||
return {
|
||||
"first_workspace_id": None,
|
||||
"skipped": False,
|
||||
"completed_task_ids": [],
|
||||
"manually_enabled_workspace_ids": [],
|
||||
"manually_disabled_workspace_ids": [],
|
||||
"updated_at": None,
|
||||
}
|
||||
|
||||
return {
|
||||
"first_workspace_id": state.first_workspace_id,
|
||||
"skipped": state.skipped,
|
||||
"completed_task_ids": cls._normalize_ids(state.completed_task_ids),
|
||||
"manually_enabled_workspace_ids": cls._normalize_ids(state.manually_enabled_workspace_ids),
|
||||
"manually_disabled_workspace_ids": cls._normalize_ids(state.manually_disabled_workspace_ids),
|
||||
"updated_at": state.updated_at,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _validate_task_id(task_id: str) -> None:
|
||||
if task_id not in STEP_BY_STEP_TOUR_TASK_IDS:
|
||||
raise ValueError(f"Unsupported task_id: {task_id}")
|
||||
|
||||
@classmethod
|
||||
def _add_id(cls, values: list[str], value: str) -> list[str]:
|
||||
normalized = cls._normalize_ids(values)
|
||||
if value in normalized:
|
||||
return normalized
|
||||
return [*normalized, value]
|
||||
|
||||
@classmethod
|
||||
def _remove_id(cls, values: list[str], value: str) -> list[str]:
|
||||
return [item for item in cls._normalize_ids(values) if item != value]
|
||||
|
||||
@staticmethod
|
||||
def _normalize_ids(values: list[str]) -> list[str]:
|
||||
normalized: list[str] = []
|
||||
for value in values:
|
||||
if value not in normalized:
|
||||
normalized.append(value)
|
||||
return normalized
|
||||
@@ -8,6 +8,7 @@ import uuid
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, override
|
||||
|
||||
from socketio.exceptions import TimeoutError as SocketIOTimeoutError # type: ignore[reportMissingTypeStubs]
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -18,6 +19,7 @@ from repositories.workflow_collaboration_repository import WorkflowCollaboration
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SERVER_HEARTBEAT_INTERVAL_SECONDS = 30
|
||||
SYNC_REQUEST_TIMEOUT_SECONDS = 15
|
||||
_PROCESS_SERVER_ID: str | None = None
|
||||
_PROCESS_SERVER_ID_PID: int | None = None
|
||||
|
||||
@@ -38,7 +40,8 @@ class WorkflowCollaborationService:
|
||||
|
||||
Socket.IO rooms are process-local unless backed by a message queue, while online users and leader election live in
|
||||
Redis. Each websocket worker writes a small heartbeat keyed by `server_id`; session rows store their owner so a
|
||||
worker can distinguish a live remote sid from a stale sid left behind by a dead worker.
|
||||
worker can distinguish a live remote sid from a stale sid left behind by a dead worker. Visibility events are
|
||||
ordered per socket, and sync requests wait for the selected saver to acknowledge its draft write.
|
||||
"""
|
||||
|
||||
_heartbeat_started: bool
|
||||
@@ -128,6 +131,9 @@ class WorkflowCollaborationService:
|
||||
"sid": sid,
|
||||
"connected_at": int(time.time()),
|
||||
"server_id": self.server_id,
|
||||
# Joins are assumed visible; hidden tabs re-report via a graph_view_state
|
||||
# event right after receiving the post-join "status" emit.
|
||||
"graph_active": True,
|
||||
}
|
||||
|
||||
self._repository.set_session_info(workflow_id, session_info)
|
||||
@@ -158,7 +164,8 @@ class WorkflowCollaborationService:
|
||||
self.handle_leader_disconnect(workflow_id, sid)
|
||||
self.broadcast_online_users(workflow_id)
|
||||
|
||||
def relay_collaboration_event(self, sid: str, data: Mapping[str, object]) -> tuple[dict[str, str], int]:
|
||||
def relay_collaboration_event(self, sid: str, data: Mapping[str, object]) -> tuple[dict[str, object], int]:
|
||||
"""Route collaboration control events, directing save and graph-resync requests to the active leader."""
|
||||
mapping = self._repository.get_sid_mapping(sid)
|
||||
if not mapping:
|
||||
return {"msg": "unauthorized"}, 401
|
||||
@@ -174,11 +181,54 @@ class WorkflowCollaborationService:
|
||||
if not event_type:
|
||||
return {"msg": "invalid event type"}, 400
|
||||
|
||||
if event_type == "graph_view_state":
|
||||
if not isinstance(event_data, Mapping):
|
||||
return {"msg": "invalid graph_view_state"}, 400
|
||||
|
||||
graph_active = event_data.get("graphActive")
|
||||
sequence = event_data.get("sequence")
|
||||
if not isinstance(graph_active, bool):
|
||||
return {"msg": "invalid graph_view_state"}, 400
|
||||
if not isinstance(sequence, int) or isinstance(sequence, bool) or sequence < 0:
|
||||
return {"msg": "invalid graph_view_state"}, 400
|
||||
|
||||
# Sequence the visibility write together with its leader side effect. Otherwise a newer
|
||||
# visible event can land after the Lua write but before an older hidden handler demotes.
|
||||
with self._repository.graph_view_state_lock(workflow_id):
|
||||
applied = self._repository.update_session_graph_active(workflow_id, sid, graph_active, sequence)
|
||||
if not applied:
|
||||
return {"msg": "graph_view_state_ignored"}, 200
|
||||
|
||||
# Write the flag before re-electing so tier-1 selection already excludes the
|
||||
# session that just went hidden (including one _ensure_leader just promoted).
|
||||
if not graph_active:
|
||||
self._demote_leader_if_hidden(workflow_id, sid)
|
||||
return {"msg": "graph_view_state_updated"}, 200
|
||||
|
||||
if event_type == "sync_request":
|
||||
if not isinstance(event_data, Mapping):
|
||||
return {"msg": "invalid sync_request"}, 400
|
||||
request_id = event_data.get("requestId")
|
||||
if not isinstance(request_id, str) or not request_id.strip():
|
||||
return {"msg": "invalid sync_request"}, 400
|
||||
|
||||
leader_sid = self._repository.get_current_leader(workflow_id)
|
||||
target_sid: str | None
|
||||
if leader_sid and self.is_session_active(workflow_id, leader_sid):
|
||||
target_sid = leader_sid
|
||||
if self._is_session_graph_active(workflow_id, leader_sid):
|
||||
target_sid = leader_sid
|
||||
else:
|
||||
# The leader is connected but its tab is hidden: its canvas is frozen
|
||||
# (rAF paused), so saving through it would persist stale data. Hand
|
||||
# leadership to a visible session — or, if every tab is hidden, to the
|
||||
# requester, whose canvas at least contains its own edits.
|
||||
replacement = self._select_graph_leader(workflow_id, preferred_sid=sid)
|
||||
if replacement and replacement != leader_sid:
|
||||
self._repository.set_leader(workflow_id, replacement)
|
||||
self.broadcast_leader_change(workflow_id, replacement)
|
||||
target_sid = replacement
|
||||
else:
|
||||
target_sid = leader_sid
|
||||
else:
|
||||
if leader_sid:
|
||||
self._repository.delete_leader(workflow_id)
|
||||
@@ -188,15 +238,73 @@ class WorkflowCollaborationService:
|
||||
self.broadcast_leader_change(workflow_id, target_sid)
|
||||
|
||||
if not target_sid:
|
||||
return {"msg": "no_active_leader"}, 200
|
||||
return {"msg": "no_active_leader", "requestId": request_id}, 503
|
||||
|
||||
target_data = dict(event_data)
|
||||
target_data["requestId"] = request_id
|
||||
try:
|
||||
result = self._socketio.call(
|
||||
"collaboration_update",
|
||||
{"type": event_type, "userId": user_id, "data": target_data, "timestamp": timestamp},
|
||||
to=target_sid,
|
||||
timeout=SYNC_REQUEST_TIMEOUT_SECONDS,
|
||||
)
|
||||
except SocketIOTimeoutError:
|
||||
logger.warning(
|
||||
"Workflow collaboration sync request timed out: workflow_id=%s requester_sid=%s target_sid=%s",
|
||||
workflow_id,
|
||||
sid,
|
||||
target_sid,
|
||||
)
|
||||
return {"msg": "sync_request_timeout", "requestId": request_id}, 504
|
||||
|
||||
if not isinstance(result, Mapping) or result.get("success") is not True:
|
||||
response: dict[str, object] = {
|
||||
"msg": "workflow_sync_failed",
|
||||
"requestId": request_id,
|
||||
"success": False,
|
||||
}
|
||||
if isinstance(result, Mapping) and isinstance(result.get("error"), str):
|
||||
response["error"] = result["error"]
|
||||
return response, 502
|
||||
|
||||
workflow_hash = result.get("hash")
|
||||
updated_at = result.get("updatedAt")
|
||||
if not isinstance(workflow_hash, str) or not workflow_hash:
|
||||
return {"msg": "invalid_sync_response", "requestId": request_id}, 502
|
||||
if not isinstance(updated_at, int) or isinstance(updated_at, bool):
|
||||
return {"msg": "invalid_sync_response", "requestId": request_id}, 502
|
||||
|
||||
return {
|
||||
"msg": "workflow_synced",
|
||||
"requestId": request_id,
|
||||
"success": True,
|
||||
"hash": workflow_hash,
|
||||
"updatedAt": updated_at,
|
||||
}, 200
|
||||
|
||||
if event_type == "graph_resync_request":
|
||||
leader_sid = self._repository.get_current_leader(workflow_id)
|
||||
resync_target_sid: str | None
|
||||
if leader_sid and self.is_session_active(workflow_id, leader_sid):
|
||||
resync_target_sid = leader_sid
|
||||
else:
|
||||
if leader_sid:
|
||||
self._repository.delete_leader(workflow_id)
|
||||
resync_target_sid = self._select_graph_leader(workflow_id, preferred_sid=sid)
|
||||
if resync_target_sid:
|
||||
self._repository.set_leader(workflow_id, resync_target_sid)
|
||||
self.broadcast_leader_change(workflow_id, resync_target_sid)
|
||||
|
||||
if not resync_target_sid:
|
||||
return {"msg": "no_active_leader"}, 503
|
||||
|
||||
self._socketio.emit(
|
||||
"collaboration_update",
|
||||
{"type": event_type, "userId": user_id, "data": event_data, "timestamp": timestamp},
|
||||
room=target_sid,
|
||||
to=resync_target_sid,
|
||||
)
|
||||
|
||||
return {"msg": "sync_request_forwarded"}, 200
|
||||
return {"msg": "graph_resync_request_forwarded"}, 200
|
||||
|
||||
self._socketio.emit(
|
||||
"collaboration_update",
|
||||
@@ -329,17 +437,56 @@ class WorkflowCollaborationService:
|
||||
self._repository.set_leader(workflow_id, sid)
|
||||
self.broadcast_leader_change(workflow_id, sid)
|
||||
|
||||
def _select_graph_leader(self, workflow_id: str, preferred_sid: str | None = None) -> str | None:
|
||||
session_sids = [
|
||||
session["sid"]
|
||||
def _select_graph_leader(
|
||||
self,
|
||||
workflow_id: str,
|
||||
preferred_sid: str | None = None,
|
||||
*,
|
||||
require_graph_active: bool = False,
|
||||
) -> str | None:
|
||||
"""Pick a leader, preferring sessions whose canvas tab is visible.
|
||||
|
||||
Hidden tabs freeze rAF-driven CRDT->canvas application, so a visible session is
|
||||
always the freshest saver. When every tab is hidden the room must still keep a
|
||||
leader (followers drop sync_requests unless they believe they are the leader),
|
||||
so tier 2 falls back to any active session unless require_graph_active is set.
|
||||
"""
|
||||
active_sessions = [
|
||||
session
|
||||
for session in self._repository.list_sessions(workflow_id)
|
||||
if session.get("graph_active", True) and self.is_session_active(workflow_id, session["sid"])
|
||||
if self.is_session_active(workflow_id, session["sid"])
|
||||
]
|
||||
if not session_sids:
|
||||
visible_sids = [session["sid"] for session in active_sessions if session.get("graph_active", True)]
|
||||
|
||||
candidate_sids = visible_sids
|
||||
if not candidate_sids and not require_graph_active:
|
||||
candidate_sids = [session["sid"] for session in active_sessions]
|
||||
|
||||
if not candidate_sids:
|
||||
return None
|
||||
if preferred_sid and preferred_sid in session_sids:
|
||||
if preferred_sid and preferred_sid in candidate_sids:
|
||||
return preferred_sid
|
||||
return session_sids[0]
|
||||
return candidate_sids[0]
|
||||
|
||||
def _is_session_graph_active(self, workflow_id: str, sid: str) -> bool:
|
||||
"""Default to True on missing/unreadable session data so read failures never churn leadership."""
|
||||
session_info = self._repository.get_session_info(workflow_id, sid)
|
||||
if session_info is None:
|
||||
return True
|
||||
return bool(session_info.get("graph_active", True))
|
||||
|
||||
def _demote_leader_if_hidden(self, workflow_id: str, hidden_sid: str) -> None:
|
||||
current_leader = self._repository.get_current_leader(workflow_id)
|
||||
if current_leader != hidden_sid:
|
||||
return
|
||||
|
||||
new_leader = self._select_graph_leader(workflow_id, require_graph_active=True)
|
||||
# No visible session: keep the hidden leader rather than leaving the room leaderless.
|
||||
if not new_leader or new_leader == hidden_sid:
|
||||
return
|
||||
|
||||
self._repository.set_leader(workflow_id, new_leader)
|
||||
self.broadcast_leader_change(workflow_id, new_leader)
|
||||
|
||||
def is_session_active(self, workflow_id: str, sid: str) -> bool:
|
||||
if not sid:
|
||||
|
||||
@@ -823,6 +823,8 @@ _FILENAME_TRANS_TABLE = _make_filename_trans_table()
|
||||
|
||||
|
||||
class DraftVariableSaver:
|
||||
"""Persist draft outputs under the tenant that owns the app or pipeline."""
|
||||
|
||||
# _DUMMY_OUTPUT_IDENTITY is a placeholder output for workflow nodes.
|
||||
# Its sole possible value is `None`.
|
||||
#
|
||||
@@ -842,6 +844,10 @@ class DraftVariableSaver:
|
||||
# Database session used for persisting draft variables.
|
||||
_session: Session
|
||||
|
||||
# Resource owner tenant. An account's current tenant may be unset or point elsewhere
|
||||
# when draft variables are persisted by an asynchronous workflow execution.
|
||||
_tenant_id: str
|
||||
|
||||
# The application ID associated with the draft variables.
|
||||
# This should match the `Workflow.app_id` of the workflow to which the current node belongs.
|
||||
_app_id: str
|
||||
@@ -867,6 +873,7 @@ class DraftVariableSaver:
|
||||
def __init__(
|
||||
self,
|
||||
session: Session,
|
||||
tenant_id: str,
|
||||
app_id: str,
|
||||
node_id: str,
|
||||
node_type: NodeType,
|
||||
@@ -878,6 +885,7 @@ class DraftVariableSaver:
|
||||
# WorkflowNodeExecutionModel/WorkflowNodeExecution, not their `node_execution_id`
|
||||
# field. These are distinct database fields with different purposes.
|
||||
self._session = session
|
||||
self._tenant_id = tenant_id
|
||||
self._app_id = app_id
|
||||
self._node_id = node_id
|
||||
self._node_type = node_type
|
||||
@@ -885,12 +893,6 @@ class DraftVariableSaver:
|
||||
self._user = user
|
||||
self._enclosing_node_id = enclosing_node_id
|
||||
|
||||
def _resolve_app_tenant_id(self) -> str:
|
||||
tenant_id = self._session.scalar(select(App.tenant_id).where(App.id == self._app_id))
|
||||
if not tenant_id:
|
||||
raise ValueError(f"Unable to resolve tenant_id for app {self._app_id}")
|
||||
return tenant_id
|
||||
|
||||
def _create_dummy_output_variable(self):
|
||||
return WorkflowDraftVariable.new_node_variable(
|
||||
app_id=self._app_id,
|
||||
@@ -949,11 +951,10 @@ class DraftVariableSaver:
|
||||
if name == SystemVariableKey.FILES:
|
||||
# Here we know the type of variable must be `array[file]`, we
|
||||
# just rebuild files from the serialized payload.
|
||||
tenant_id = self._resolve_app_tenant_id()
|
||||
files = [
|
||||
build_file_from_stored_mapping(
|
||||
file_mapping=v,
|
||||
tenant_id=tenant_id,
|
||||
tenant_id=self._tenant_id,
|
||||
)
|
||||
for v in value
|
||||
]
|
||||
@@ -1096,8 +1097,8 @@ class DraftVariableSaver:
|
||||
content=original_content_serialized.encode(),
|
||||
mimetype=content_type,
|
||||
user=self._user,
|
||||
tenant_id=self._tenant_id,
|
||||
)
|
||||
assert self._user.current_tenant_id
|
||||
# Create WorkflowDraftVariableFile record
|
||||
variable_file = WorkflowDraftVariableFile(
|
||||
upload_file_id=upload_file.id,
|
||||
@@ -1105,7 +1106,7 @@ class DraftVariableSaver:
|
||||
length=original_length,
|
||||
value_type=value_seg.value_type,
|
||||
app_id=self._app_id,
|
||||
tenant_id=self._user.current_tenant_id,
|
||||
tenant_id=self._tenant_id,
|
||||
user_id=self._user.id,
|
||||
)
|
||||
variable_file.id = str(uuidv7())
|
||||
|
||||
@@ -13,6 +13,7 @@ from sqlalchemy import desc, select
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from core.app.apps.message_generator import MessageGenerator
|
||||
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity
|
||||
from core.app.entities.task_entities import (
|
||||
HumanInputRequiredResponse,
|
||||
MessageReplaceStreamResponse,
|
||||
@@ -84,9 +85,6 @@ def build_workflow_event_stream(
|
||||
topic = MessageGenerator.get_response_topic(app_mode, workflow_run.id)
|
||||
workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
|
||||
node_execution_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(session_maker)
|
||||
message_context = (
|
||||
_get_message_context(session_maker, workflow_run.id) if app_mode == AppMode.ADVANCED_CHAT else None
|
||||
)
|
||||
|
||||
pause_entity: WorkflowPauseEntity | None = None
|
||||
if workflow_run.status == WorkflowExecutionStatus.PAUSED:
|
||||
@@ -97,6 +95,38 @@ def build_workflow_event_stream(
|
||||
pause_entity = None
|
||||
|
||||
resumption_context = _load_resumption_context(pause_entity)
|
||||
message_context: MessageContext | None = None
|
||||
if app_mode == AppMode.ADVANCED_CHAT:
|
||||
if workflow_run.status == WorkflowExecutionStatus.PAUSED:
|
||||
if resumption_context is None:
|
||||
raise AssertionError(
|
||||
"WorkflowResumptionContext is required for advanced-chat snapshot replay, "
|
||||
f"workflow_run_id={workflow_run.id}"
|
||||
)
|
||||
generate_entity = resumption_context.get_generate_entity()
|
||||
if not isinstance(generate_entity, AdvancedChatAppGenerateEntity):
|
||||
raise AssertionError(
|
||||
"AdvancedChatAppGenerateEntity is required for advanced-chat snapshot replay, "
|
||||
f"workflow_run_id={workflow_run.id}, generate_entity_type={type(generate_entity).__name__}"
|
||||
)
|
||||
if not generate_entity.conversation_id:
|
||||
raise AssertionError(
|
||||
f"conversation_id is required for advanced-chat snapshot replay, workflow_run_id={workflow_run.id}"
|
||||
)
|
||||
message_context = _get_message_context_by_conversation(
|
||||
session_maker,
|
||||
conversation_id=generate_entity.conversation_id,
|
||||
workflow_run_id=workflow_run.id,
|
||||
)
|
||||
else:
|
||||
# Compatibility fallback for non-suspended snapshot requests. This app-scoped lookup is not optimal;
|
||||
# a dedicated index or stronger lookup key would be preferable.
|
||||
message_context = _get_message_context_by_app(
|
||||
session_maker,
|
||||
app_id=app_id,
|
||||
workflow_run_id=workflow_run.id,
|
||||
)
|
||||
|
||||
node_snapshots = node_execution_repo.get_execution_snapshots_by_workflow_run(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
@@ -175,19 +205,68 @@ def build_workflow_event_stream(
|
||||
return _generate()
|
||||
|
||||
|
||||
def _get_message_context(session_maker: sessionmaker[Session], workflow_run_id: str) -> MessageContext | None:
|
||||
def _get_message_context_by_conversation(
|
||||
session_maker: sessionmaker[Session],
|
||||
*,
|
||||
conversation_id: str,
|
||||
workflow_run_id: str,
|
||||
) -> MessageContext | None:
|
||||
"""Look up a paused or suspended Advanced Chat snapshot message by conversation and workflow run.
|
||||
|
||||
Use this exact lookup after recovering ``conversation_id`` from persisted resumption context. Its predicates match
|
||||
``message_workflow_run_id_idx``.
|
||||
"""
|
||||
with session_maker() as session:
|
||||
stmt = select(Message).where(Message.workflow_run_id == workflow_run_id).order_by(desc(Message.created_at))
|
||||
stmt = (
|
||||
select(Message)
|
||||
.where(
|
||||
Message.conversation_id == conversation_id,
|
||||
Message.workflow_run_id == workflow_run_id,
|
||||
)
|
||||
.order_by(desc(Message.created_at))
|
||||
.limit(1)
|
||||
)
|
||||
message = session.scalar(stmt)
|
||||
if message is None:
|
||||
return None
|
||||
created_at = int(message.created_at.timestamp()) if message.created_at else 0
|
||||
return MessageContext(
|
||||
conversation_id=message.conversation_id,
|
||||
message_id=message.id,
|
||||
created_at=created_at,
|
||||
answer=message.answer,
|
||||
return _to_message_context(message)
|
||||
|
||||
|
||||
def _get_message_context_by_app(
|
||||
session_maker: sessionmaker[Session],
|
||||
*,
|
||||
app_id: str,
|
||||
workflow_run_id: str,
|
||||
) -> MessageContext | None:
|
||||
"""Look up a non-suspended or running Advanced Chat reconnect snapshot by app and workflow run.
|
||||
|
||||
This compatibility path applies only when no resumption context is expected. The app-scoped query is not optimal;
|
||||
a dedicated index or stronger lookup key would be preferable.
|
||||
"""
|
||||
with session_maker() as session:
|
||||
stmt = (
|
||||
select(Message)
|
||||
.where(
|
||||
Message.app_id == app_id,
|
||||
Message.workflow_run_id == workflow_run_id,
|
||||
)
|
||||
.order_by(desc(Message.created_at))
|
||||
.limit(1)
|
||||
)
|
||||
message = session.scalar(stmt)
|
||||
if message is None:
|
||||
return None
|
||||
return _to_message_context(message)
|
||||
|
||||
|
||||
def _to_message_context(message: Message) -> MessageContext:
|
||||
created_at = int(message.created_at.timestamp()) if message.created_at else 0
|
||||
return MessageContext(
|
||||
conversation_id=message.conversation_id,
|
||||
message_id=message.id,
|
||||
created_at=created_at,
|
||||
answer=message.answer,
|
||||
)
|
||||
|
||||
|
||||
def _load_resumption_context(pause_entity: WorkflowPauseEntity | None) -> WorkflowResumptionContext | None:
|
||||
|
||||
@@ -322,6 +322,7 @@ class WorkflowService:
|
||||
session: Session,
|
||||
commit: bool = True,
|
||||
sync_agent_bindings: bool = True,
|
||||
graph_only: bool = False,
|
||||
) -> Workflow:
|
||||
"""
|
||||
Sync draft workflow.
|
||||
@@ -338,8 +339,10 @@ class WorkflowService:
|
||||
if workflow and workflow.unique_hash != unique_hash:
|
||||
raise WorkflowHashNotEqualError()
|
||||
|
||||
# validate features structure
|
||||
self.validate_features_structure(app_model=app_model, features=features)
|
||||
# Collaboration persists features and variables through dedicated endpoints. A graph save
|
||||
# must not overwrite those newer database values with another collaborator's stale cache.
|
||||
if not graph_only or not workflow:
|
||||
self.validate_features_structure(app_model=app_model, features=features)
|
||||
|
||||
# validate graph structure
|
||||
self.validate_graph_structure(graph=graph)
|
||||
@@ -361,11 +364,12 @@ class WorkflowService:
|
||||
# update draft workflow if found
|
||||
else:
|
||||
workflow.graph = json.dumps(graph)
|
||||
workflow.features = json.dumps(features)
|
||||
workflow.updated_by = account.id
|
||||
workflow.updated_at = naive_utc_now()
|
||||
workflow.environment_variables = environment_variables
|
||||
workflow.conversation_variables = conversation_variables
|
||||
if not graph_only:
|
||||
workflow.features = json.dumps(features)
|
||||
workflow.environment_variables = environment_variables
|
||||
workflow.conversation_variables = conversation_variables
|
||||
|
||||
from services.agent.workflow_publish_service import WorkflowAgentPublishService
|
||||
|
||||
@@ -1056,6 +1060,7 @@ class WorkflowService:
|
||||
with sessionmaker(bind=db.engine).begin() as session:
|
||||
draft_var_saver = DraftVariableSaver(
|
||||
session=session,
|
||||
tenant_id=app_model.tenant_id,
|
||||
app_id=app_model.id,
|
||||
node_id=workflow_node_execution.node_id,
|
||||
node_type=workflow_node_execution.node_type,
|
||||
@@ -1206,6 +1211,7 @@ class WorkflowService:
|
||||
with sessionmaker(bind=db.engine).begin() as session:
|
||||
draft_var_saver = DraftVariableSaver(
|
||||
session=session,
|
||||
tenant_id=app_model.tenant_id,
|
||||
app_id=app_model.id,
|
||||
node_id=node_id,
|
||||
node_type=BuiltinNodeTypes.HUMAN_INPUT,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
FLASK_APP=app.py
|
||||
FLASK_DEBUG=0
|
||||
ENABLE_COLLABORATION_MODE=false
|
||||
SECRET_KEY='uhySf6a3aZuvRNfAlcr47paOw9TRYBY6j8ZHXpVw1yx5RP27Yj3w2uvI'
|
||||
|
||||
CONSOLE_API_URL=http://127.0.0.1:5001
|
||||
|
||||
@@ -513,11 +513,94 @@ class TestArchiveRunIdempotency:
|
||||
storage = MagicMock()
|
||||
storage.object_exists.return_value = True
|
||||
|
||||
result = archiver._archive_bundle(MagicMock(), storage, [run])
|
||||
with patch.object(archiver, "_sync_existing_bundle_index") as sync_existing_bundle_index:
|
||||
result = archiver._archive_bundle(MagicMock(), storage, [run])
|
||||
|
||||
assert result.success is True
|
||||
assert result.skipped is True
|
||||
assert result.error == "bundle already archived"
|
||||
sync_existing_bundle_index.assert_called_once()
|
||||
|
||||
def test_existing_bundle_catalog_publication_failure_is_not_success(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = _run()
|
||||
session = MagicMock()
|
||||
storage = MagicMock()
|
||||
storage.object_exists.return_value = True
|
||||
|
||||
with patch.object(archiver, "_sync_existing_bundle_index", side_effect=RuntimeError("catalog unavailable")):
|
||||
result = archiver._archive_bundle(session, storage, [run])
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "catalog unavailable"
|
||||
session.rollback.assert_called_once()
|
||||
|
||||
def test_retry_repairs_index_after_catalog_commit_then_index_write_failure(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = _run()
|
||||
identity = archiver._build_bundle_identity([run])
|
||||
index_key = archiver._get_index_object_key(identity)
|
||||
manifest_key = archiver._get_manifest_object_key(identity)
|
||||
storage = FakeArchiveStorage()
|
||||
original_put_object = storage.put_object
|
||||
index_write_count = 0
|
||||
|
||||
def put_object(key: str, data: bytes) -> str:
|
||||
nonlocal index_write_count
|
||||
if key == index_key:
|
||||
index_write_count += 1
|
||||
if index_write_count == 2:
|
||||
raise RuntimeError("index write failed")
|
||||
return original_put_object(key, data)
|
||||
|
||||
storage.put_object = MagicMock(side_effect=put_object)
|
||||
first_session = MagicMock()
|
||||
first_session.scalar.return_value = None
|
||||
table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}
|
||||
|
||||
with (
|
||||
patch.object(archiver, "_lock_runs_for_archive", return_value=[run]),
|
||||
patch.object(archiver, "_extract_bundle_data", return_value=table_data),
|
||||
):
|
||||
first_result = archiver._archive_bundle(first_session, storage, [run])
|
||||
|
||||
assert first_result.success is False
|
||||
assert first_result.error == "index write failed"
|
||||
assert manifest_key in storage.objects
|
||||
assert json.loads(storage.objects[index_key])["run_ids"] == []
|
||||
|
||||
storage.list_objects = MagicMock(wraps=storage.list_objects)
|
||||
retry_session = MagicMock()
|
||||
retry_session.scalar.return_value = None
|
||||
|
||||
retry_result = archiver._archive_bundle(retry_session, storage, [run])
|
||||
|
||||
assert retry_result.success is True
|
||||
assert retry_result.skipped is True
|
||||
assert json.loads(storage.objects[index_key])["manifest_keys"] == [manifest_key]
|
||||
assert json.loads(storage.objects[index_key])["run_ids"] == [run.id]
|
||||
storage.list_objects.assert_not_called()
|
||||
|
||||
def test_existing_manifest_with_missing_index_fails_without_partial_rebuild(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = _run()
|
||||
identity = archiver._build_bundle_identity([run])
|
||||
_, _, manifest_data = archiver._build_archive_payload(
|
||||
identity,
|
||||
[run],
|
||||
{"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]},
|
||||
)
|
||||
manifest_key = archiver._get_manifest_object_key(identity)
|
||||
index_key = archiver._get_index_object_key(identity)
|
||||
storage = FakeArchiveStorage({manifest_key: manifest_data})
|
||||
storage.list_objects = MagicMock(wraps=storage.list_objects)
|
||||
|
||||
result = archiver._archive_bundle(MagicMock(), storage, [run])
|
||||
|
||||
assert result.success is False
|
||||
assert "archive shard index missing" in (result.error or "")
|
||||
assert index_key not in storage.objects
|
||||
storage.list_objects.assert_not_called()
|
||||
|
||||
def test_successful_bundle_persists_archive_index(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
@@ -550,6 +633,28 @@ class TestArchiveRunIdempotency:
|
||||
assert archived_bundle.row_count == 2
|
||||
session.commit.assert_called_once()
|
||||
|
||||
def test_new_bundle_catalog_commit_failure_is_not_success(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = _run(str(uuid.uuid4()))
|
||||
run.tenant_id = str(uuid.uuid4())
|
||||
session = MagicMock()
|
||||
session.scalar.return_value = None
|
||||
session.commit.side_effect = RuntimeError("catalog commit failed")
|
||||
storage = MagicMock()
|
||||
storage.object_exists.return_value = False
|
||||
storage.list_objects.return_value = []
|
||||
table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}
|
||||
|
||||
with (
|
||||
patch.object(archiver, "_lock_runs_for_archive", return_value=[run]),
|
||||
patch.object(archiver, "_extract_bundle_data", return_value=table_data),
|
||||
):
|
||||
result = archiver._archive_bundle(session, storage, [run])
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "catalog commit failed"
|
||||
session.rollback.assert_called_once()
|
||||
|
||||
def test_index_skips_all_already_archived_runs(self):
|
||||
archiver = WorkflowRunArchiver(days=90)
|
||||
run = MagicMock()
|
||||
|
||||
@@ -311,6 +311,7 @@ class TestDraftVariableLoader(unittest.TestCase):
|
||||
# Use DraftVariableSaver to create offloaded variable (this mimics production)
|
||||
saver = DraftVariableSaver(
|
||||
session=session,
|
||||
tenant_id=self._test_tenant_id,
|
||||
app_id=self._test_app_id,
|
||||
node_id="test_offload_node",
|
||||
node_type=BuiltinNodeTypes.LLM, # Use a real node type
|
||||
|
||||
-179
@@ -1,179 +0,0 @@
|
||||
"""Testcontainers integration tests for controllers.console.app.app_import endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from inspect import unwrap
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from controllers.console.app import app_import as app_import_module
|
||||
from services.app_dsl_service import ImportStatus
|
||||
|
||||
|
||||
class _Result:
|
||||
def __init__(self, status: ImportStatus, app_id: str | None = "app-1"):
|
||||
self.status = status
|
||||
self.app_id = app_id
|
||||
|
||||
def model_dump(self, mode: str = "json"):
|
||||
return {"status": self.status, "app_id": self.app_id}
|
||||
|
||||
|
||||
def _install_features(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
|
||||
features = SimpleNamespace(webapp_auth=SimpleNamespace(enabled=enabled))
|
||||
monkeypatch.setattr(app_import_module.FeatureService, "get_system_features", lambda: features)
|
||||
|
||||
|
||||
class TestAppImportApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask):
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_import_post_returns_failed_status(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = app_import_module.AppImportApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED, app_id=None),
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
|
||||
assert status == 400
|
||||
assert response["status"] == ImportStatus.FAILED
|
||||
|
||||
def test_import_post_returns_pending_status(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = app_import_module.AppImportApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.PENDING),
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
|
||||
assert status == 202
|
||||
assert response["status"] == ImportStatus.PENDING
|
||||
|
||||
def test_import_post_updates_webapp_auth_when_enabled(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = app_import_module.AppImportApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=True)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.COMPLETED, app_id="app-123"),
|
||||
)
|
||||
update_access = MagicMock()
|
||||
monkeypatch.setattr(app_import_module.EnterpriseService.WebAppAuth, "update_app_access_mode", update_access)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
|
||||
update_access.assert_called_once_with("app-123", "private")
|
||||
assert status == 200
|
||||
assert response["status"] == ImportStatus.COMPLETED
|
||||
|
||||
def test_import_post_commits_session_on_success(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = app_import_module.AppImportApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.COMPLETED, app_id="app-123"),
|
||||
)
|
||||
|
||||
fake_session = MagicMock()
|
||||
fake_session.__enter__.return_value = fake_session
|
||||
fake_session.__exit__.return_value = None
|
||||
monkeypatch.setattr(app_import_module, "Session", lambda *_args, **_kwargs: fake_session)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
|
||||
fake_session.commit.assert_called_once_with()
|
||||
fake_session.rollback.assert_not_called()
|
||||
assert status == 200
|
||||
assert response["status"] == ImportStatus.COMPLETED
|
||||
|
||||
def test_import_post_rolls_back_session_on_failure(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = app_import_module.AppImportApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED, app_id=None),
|
||||
)
|
||||
|
||||
fake_session = MagicMock()
|
||||
fake_session.__enter__.return_value = fake_session
|
||||
fake_session.__exit__.return_value = None
|
||||
monkeypatch.setattr(app_import_module, "Session", lambda *_args, **_kwargs: fake_session)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
|
||||
fake_session.rollback.assert_called_once_with()
|
||||
fake_session.commit.assert_not_called()
|
||||
assert status == 400
|
||||
assert response["status"] == ImportStatus.FAILED
|
||||
|
||||
|
||||
class TestAppImportConfirmApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask):
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_import_confirm_returns_failed_status(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = app_import_module.AppImportConfirmApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"confirm_import",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED),
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports/import-1/confirm", method="POST"):
|
||||
response, status = method(api, SimpleNamespace(id="u1"), import_id="import-1")
|
||||
|
||||
assert status == 400
|
||||
assert response["status"] == ImportStatus.FAILED
|
||||
|
||||
|
||||
class TestAppImportCheckDependenciesApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask):
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_import_check_dependencies_returns_result(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = app_import_module.AppImportCheckDependenciesApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"check_dependencies",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(model_dump=lambda mode="json": {"leaked_dependencies": []}),
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports/app-1/check-dependencies", method="GET"):
|
||||
response, status = method(api, app_model=SimpleNamespace(id="app-1"))
|
||||
|
||||
assert status == 200
|
||||
assert response["leaked_dependencies"] == []
|
||||
@@ -1,3 +1,5 @@
|
||||
from decimal import Decimal
|
||||
|
||||
import pytest
|
||||
from agenton.compositor import CompositorSessionSnapshot
|
||||
from dify_agent.protocol import (
|
||||
@@ -83,7 +85,19 @@ def test_event_adapter_maps_run_succeeded_to_final_output():
|
||||
data=RunSucceededEventData(
|
||||
output={"summary": "done"},
|
||||
session_snapshot=snapshot,
|
||||
usage=AgentRunUsage(prompt_tokens=2, completion_tokens=3),
|
||||
usage=AgentRunUsage(
|
||||
prompt_tokens=2,
|
||||
prompt_unit_price=Decimal(5),
|
||||
prompt_price_unit=Decimal("0.000001"),
|
||||
prompt_price=Decimal("0.000010"),
|
||||
completion_tokens=3,
|
||||
completion_unit_price=Decimal(30),
|
||||
completion_price_unit=Decimal("0.000001"),
|
||||
completion_price=Decimal("0.000090"),
|
||||
total_price=Decimal("0.000100"),
|
||||
currency="USD",
|
||||
latency=0.4,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
@@ -94,7 +108,22 @@ def test_event_adapter_maps_run_succeeded_to_final_output():
|
||||
source_event_id="3-0",
|
||||
output={"summary": "done"},
|
||||
session_snapshot=snapshot,
|
||||
usage={"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
|
||||
usage={
|
||||
"prompt_tokens": 2,
|
||||
"prompt_unit_price": "5",
|
||||
"prompt_price_unit": "0.000001",
|
||||
"prompt_price": "0.000010",
|
||||
"completion_tokens": 3,
|
||||
"completion_unit_price": "30",
|
||||
"completion_price_unit": "0.000001",
|
||||
"completion_price": "0.000090",
|
||||
"total_tokens": 5,
|
||||
"total_price": "0.000100",
|
||||
"currency": "USD",
|
||||
"latency": 0.4,
|
||||
"time_to_first_token": None,
|
||||
"time_to_generate": None,
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@@ -3,9 +3,19 @@ from unittest.mock import MagicMock
|
||||
|
||||
import click
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
from sqlalchemy.exc import OperationalError
|
||||
|
||||
from commands import retention
|
||||
from services.retention.workflow_run import bundle_archive_maintenance
|
||||
from services.retention.workflow_run.bundle_archive_maintenance import (
|
||||
BundleOperationResult,
|
||||
BundleOperationSummary,
|
||||
)
|
||||
|
||||
_CURSOR_0 = "00000000-0000-0000-0000-000000000000"
|
||||
_CURSOR_1 = "00000000-0000-0000-0000-000000000001"
|
||||
_CURSOR_2 = "00000000-0000-0000-0000-000000000002"
|
||||
|
||||
|
||||
def _db_disconnect_error() -> OperationalError:
|
||||
@@ -24,6 +34,58 @@ def _session_context(session):
|
||||
return context
|
||||
|
||||
|
||||
def _delete_summary(
|
||||
*,
|
||||
processed: int,
|
||||
succeeded: int = 0,
|
||||
failed: int = 0,
|
||||
next_catalog_id: str | None = None,
|
||||
preview_next_catalog_id: str | None = None,
|
||||
results: list[BundleOperationResult] | None = None,
|
||||
) -> BundleOperationSummary:
|
||||
return BundleOperationSummary(
|
||||
operation="delete",
|
||||
bundles_processed=processed,
|
||||
bundles_succeeded=succeeded,
|
||||
bundles_failed=failed,
|
||||
next_catalog_id=next_catalog_id,
|
||||
preview_next_catalog_id=preview_next_catalog_id,
|
||||
results=results or [],
|
||||
)
|
||||
|
||||
|
||||
def _patch_bundle_deleter(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
summaries: list[BundleOperationSummary],
|
||||
) -> MagicMock:
|
||||
deleter = MagicMock()
|
||||
deleter.delete_batch.side_effect = summaries
|
||||
monkeypatch.setattr(
|
||||
bundle_archive_maintenance,
|
||||
"WorkflowRunBundleArchiveMaintenance",
|
||||
MagicMock(return_value=deleter),
|
||||
)
|
||||
return deleter
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"command",
|
||||
[retention.restore_workflow_runs, retention.delete_archived_workflow_runs],
|
||||
)
|
||||
def test_v2_archive_maintenance_rejects_explicitly_empty_tenant_ids(command):
|
||||
result = CliRunner().invoke(
|
||||
command,
|
||||
["--tenant-ids", "", "--target-month", "2025-03"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 2
|
||||
assert "tenant-ids must not be empty" in result.output
|
||||
|
||||
|
||||
def test_archive_tenant_id_parser_keeps_omitted_scope_unset():
|
||||
assert retention._parse_comma_separated_ids(None, param_name="tenant-ids") is None
|
||||
|
||||
|
||||
def test_resolve_archive_tenant_ids_from_plan_uses_explicit_sessions(monkeypatch):
|
||||
end_before = datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC)
|
||||
sessions = [MagicMock(name="session-a"), MagicMock(name="session-b")]
|
||||
@@ -137,3 +199,265 @@ def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails(mon
|
||||
dry_run=True,
|
||||
delete_after_archive=False,
|
||||
)
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_keeps_single_page_behavior_without_all_pages(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(
|
||||
monkeypatch,
|
||||
[_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1)],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
["--target-month", "2025-03", "--limit", "2"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
deleter.delete_batch.assert_called_once()
|
||||
assert deleter.delete_batch.call_args.kwargs["target_year"] == 2025
|
||||
assert deleter.delete_batch.call_args.kwargs["target_month"] == 3
|
||||
assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] is None
|
||||
assert deleter.delete_batch.call_args.kwargs["limit"] == 2
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_all_pages_continues_until_empty_page(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(
|
||||
monkeypatch,
|
||||
[
|
||||
_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1),
|
||||
_delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_2),
|
||||
_delete_summary(processed=0),
|
||||
],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
["--target-month", "2025-03", "--all-pages", "--limit", "2"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [
|
||||
None,
|
||||
_CURSOR_1,
|
||||
_CURSOR_2,
|
||||
]
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_all_pages_fetches_empty_page_after_exact_full_page(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(
|
||||
monkeypatch,
|
||||
[
|
||||
_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1),
|
||||
_delete_summary(processed=0),
|
||||
],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
["--target-month", "2025-03", "--all-pages", "--limit", "2"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert deleter.delete_batch.call_count == 2
|
||||
assert deleter.delete_batch.call_args_list[1].kwargs["after_catalog_id"] == _CURSOR_1
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_all_pages_stops_at_first_failed_page(monkeypatch):
|
||||
failed_result = BundleOperationResult(
|
||||
catalog_id=_CURSOR_2,
|
||||
bundle_id="bundle-failed",
|
||||
tenant_id="tenant-1",
|
||||
object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed",
|
||||
error="archive checksum mismatch",
|
||||
)
|
||||
deleter = _patch_bundle_deleter(
|
||||
monkeypatch,
|
||||
[
|
||||
_delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_1),
|
||||
_delete_summary(processed=1, failed=1, results=[failed_result]),
|
||||
_delete_summary(processed=0),
|
||||
],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
["--target-month", "2025-03", "--all-pages"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert deleter.delete_batch.call_count == 2
|
||||
assert "target_month=2025-03" in result.output
|
||||
assert f"failed_catalog_id={_CURSOR_2}" in result.output
|
||||
assert f"resume_after_catalog_id={_CURSOR_1}" in result.output
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_all_pages_fails_when_cursor_does_not_advance(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(
|
||||
monkeypatch,
|
||||
[_delete_summary(processed=1, succeeded=1, next_catalog_id=None)],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
["--target-month", "2025-03", "--all-pages"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
deleter.delete_batch.assert_called_once()
|
||||
assert "cursor did not advance" in result.output.lower()
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_all_pages_uses_preview_cursor_for_dry_run(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(
|
||||
monkeypatch,
|
||||
[
|
||||
_delete_summary(processed=1, succeeded=1, preview_next_catalog_id=_CURSOR_1),
|
||||
_delete_summary(processed=0),
|
||||
],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
["--target-month", "2025-03", "--all-pages", "--dry-run"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [
|
||||
None,
|
||||
_CURSOR_1,
|
||||
]
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_dry_run_failure_separates_preview_and_destructive_cursors(monkeypatch):
|
||||
failed_result = BundleOperationResult(
|
||||
catalog_id=_CURSOR_2,
|
||||
bundle_id="bundle-failed",
|
||||
tenant_id="tenant-1",
|
||||
object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed",
|
||||
error="archive checksum mismatch",
|
||||
)
|
||||
deleter = _patch_bundle_deleter(
|
||||
monkeypatch,
|
||||
[
|
||||
_delete_summary(
|
||||
processed=1,
|
||||
succeeded=1,
|
||||
preview_next_catalog_id=_CURSOR_1,
|
||||
),
|
||||
_delete_summary(
|
||||
processed=1,
|
||||
failed=1,
|
||||
results=[failed_result],
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
[
|
||||
"--target-month",
|
||||
"2025-03",
|
||||
"--after-catalog-id",
|
||||
_CURSOR_0,
|
||||
"--all-pages",
|
||||
"--dry-run",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert deleter.delete_batch.call_count == 2
|
||||
assert f"failed_catalog_id={_CURSOR_2}" in result.output
|
||||
assert f"preview_after_catalog_id={_CURSOR_1}" in result.output
|
||||
assert f"destructive_retry_after_catalog_id={_CURSOR_0}" in result.output
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_all_pages_starts_after_explicit_cursor(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
[
|
||||
"--target-month",
|
||||
"2025-03",
|
||||
"--after-catalog-id",
|
||||
_CURSOR_0,
|
||||
"--all-pages",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
deleter.delete_batch.assert_called_once()
|
||||
assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] == _CURSOR_0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shard_args",
|
||||
[
|
||||
["--run-shard-index", "0"],
|
||||
["--run-shard-total", "16"],
|
||||
["--run-shard-index", "16", "--run-shard-total", "16"],
|
||||
["--run-shard-index", "-1", "--run-shard-total", "16"],
|
||||
["--run-shard-index", "0", "--run-shard-total", "0"],
|
||||
["--run-shard-index", "0", "--run-shard-total", "17"],
|
||||
],
|
||||
)
|
||||
def test_delete_archived_workflow_runs_rejects_invalid_run_shard_options(monkeypatch, shard_args):
|
||||
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
["--target-month", "2025-03", *shard_args],
|
||||
)
|
||||
|
||||
assert result.exit_code == 2
|
||||
deleter.delete_batch.assert_not_called()
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_passes_formatted_run_shard_to_service(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
[
|
||||
"--target-month",
|
||||
"2025-03",
|
||||
"--tenant-ids",
|
||||
"tenant-1",
|
||||
"--run-shard-index",
|
||||
"3",
|
||||
"--run-shard-total",
|
||||
"16",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
deleter.validate_catalog_shards.assert_called_once_with(
|
||||
target_year=2025,
|
||||
target_month=3,
|
||||
shard_total=16,
|
||||
tenant_ids=["tenant-1"],
|
||||
)
|
||||
deleter.delete_batch.assert_called_once()
|
||||
assert deleter.delete_batch.call_args.kwargs["shard"] == "03-of-16"
|
||||
|
||||
|
||||
def test_delete_archived_workflow_runs_rejects_mixed_catalog_shards_before_delete(monkeypatch):
|
||||
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
|
||||
deleter.validate_catalog_shards.side_effect = ValueError("unexpected shards: 00-of-01")
|
||||
|
||||
result = CliRunner().invoke(
|
||||
retention.delete_archived_workflow_runs,
|
||||
[
|
||||
"--target-month",
|
||||
"2025-03",
|
||||
"--run-shard-index",
|
||||
"3",
|
||||
"--run-shard-total",
|
||||
"16",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "shard preflight failed" in result.output.lower()
|
||||
assert "00-of-01" in result.output
|
||||
deleter.delete_batch.assert_not_called()
|
||||
|
||||
@@ -8,8 +8,13 @@ from yarl import URL
|
||||
from configs.app_config import DifyConfig
|
||||
|
||||
|
||||
def _clear_environment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in tuple(os.environ):
|
||||
monkeypatch.delenv(name)
|
||||
|
||||
|
||||
def _set_basic_config_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
|
||||
monkeypatch.setenv("DB_TYPE", "postgresql")
|
||||
@@ -51,7 +56,7 @@ def test_dify_config_preserves_explicit_secret_key(
|
||||
|
||||
def test_dify_config(monkeypatch: pytest.MonkeyPatch):
|
||||
# clear system environment variables
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
# Set environment variables using monkeypatch
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
@@ -104,6 +109,24 @@ def test_new_user_default_plugin_ids_are_parsed_from_env(monkeypatch: pytest.Mon
|
||||
]
|
||||
|
||||
|
||||
def test_plugin_remote_install_port_rejects_host_port_spec(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A 'host:port' compose publish spec must produce an actionable error, not an opaque int_parsing traceback."""
|
||||
_set_basic_config_env(monkeypatch)
|
||||
monkeypatch.setenv("PLUGIN_REMOTE_INSTALL_PORT", "127.0.0.1:5003")
|
||||
|
||||
with pytest.raises(ValueError, match="must be a bare port number"):
|
||||
DifyConfig(_env_file=None)
|
||||
|
||||
|
||||
def test_plugin_remote_install_port_accepts_bare_port(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_set_basic_config_env(monkeypatch)
|
||||
monkeypatch.setenv("PLUGIN_REMOTE_INSTALL_PORT", "5003")
|
||||
|
||||
config = DifyConfig(_env_file=None)
|
||||
|
||||
assert config.PLUGIN_REMOTE_INSTALL_PORT == 5003
|
||||
|
||||
|
||||
def test_new_user_default_models_are_parsed_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_set_basic_config_env(monkeypatch)
|
||||
monkeypatch.setenv(
|
||||
@@ -140,7 +163,7 @@ def test_new_user_default_models_reject_duplicate_model_types(monkeypatch: pytes
|
||||
def test_http_timeout_defaults(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that HTTP timeout defaults are correctly set"""
|
||||
# clear system environment variables
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
# Set minimal required env vars
|
||||
monkeypatch.setenv("DB_TYPE", "postgresql")
|
||||
@@ -160,7 +183,7 @@ def test_http_timeout_defaults(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
|
||||
def test_internal_files_url_falls_back_to_server_console_api_url(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
monkeypatch.setenv("SERVER_CONSOLE_API_URL", "http://api:5001")
|
||||
|
||||
config = DifyConfig(_env_file=None)
|
||||
@@ -169,7 +192,7 @@ def test_internal_files_url_falls_back_to_server_console_api_url(monkeypatch: py
|
||||
|
||||
|
||||
def test_internal_files_url_prefers_explicit_value(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
monkeypatch.setenv("INTERNAL_FILES_URL", "http://files-internal:5001")
|
||||
monkeypatch.setenv("SERVER_CONSOLE_API_URL", "http://api:5001")
|
||||
|
||||
@@ -183,7 +206,7 @@ def test_internal_files_url_prefers_explicit_value(monkeypatch: pytest.MonkeyPat
|
||||
def test_flask_configs(monkeypatch: pytest.MonkeyPatch):
|
||||
flask_app = Flask("app")
|
||||
# clear system environment variables
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
# Set environment variables using monkeypatch
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
@@ -290,7 +313,7 @@ def test_db_session_timezone_override_can_disable_app_level_timezone_injection(m
|
||||
|
||||
|
||||
def test_pubsub_redis_url_default(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
|
||||
@@ -313,7 +336,7 @@ def test_pubsub_redis_url_default(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
|
||||
def test_pubsub_redis_url_override(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
|
||||
@@ -330,7 +353,7 @@ def test_pubsub_redis_url_override(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
|
||||
def test_pubsub_redis_url_required_when_default_unavailable(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
|
||||
@@ -346,7 +369,7 @@ def test_pubsub_redis_url_required_when_default_unavailable(monkeypatch: pytest.
|
||||
|
||||
|
||||
def test_dify_config_exposes_redis_key_prefix_default(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
|
||||
@@ -363,7 +386,7 @@ def test_dify_config_exposes_redis_key_prefix_default(monkeypatch: pytest.Monkey
|
||||
|
||||
|
||||
def test_dify_config_reads_redis_key_prefix_from_env(monkeypatch: pytest.MonkeyPatch):
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
|
||||
@@ -415,7 +438,7 @@ def test_celery_broker_url_with_special_chars_password(
|
||||
from kombu.utils.url import parse_url
|
||||
|
||||
# clear system environment variables
|
||||
os.environ.clear()
|
||||
_clear_environment(monkeypatch)
|
||||
|
||||
# Set up basic required environment variables (following existing pattern)
|
||||
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic import SecretStr, ValidationError
|
||||
|
||||
from configs.extra.knowledge_fs_config import KnowledgeFSConfig
|
||||
|
||||
_REPOSITORY_ROOT = Path(__file__).resolve().parents[4]
|
||||
_KNOWLEDGE_FS_DOCKER_VARIABLES = (
|
||||
"KNOWLEDGE_FS_ENABLED",
|
||||
"KNOWLEDGE_FS_BASE_URL",
|
||||
"KNOWLEDGE_FS_JWT_SECRET",
|
||||
"KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS",
|
||||
"KNOWLEDGE_FS_TIMEOUT_SECONDS",
|
||||
)
|
||||
|
||||
|
||||
def test_knowledge_fs_config_normalizes_complete_connection() -> None:
|
||||
config = KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_ENABLED=True,
|
||||
KNOWLEDGE_FS_BASE_URL=" https://knowledge-fs.test/ ",
|
||||
KNOWLEDGE_FS_JWT_SECRET=" production-secret-with-at-least-32-bytes ",
|
||||
)
|
||||
|
||||
assert config.KNOWLEDGE_FS_BASE_URL == "https://knowledge-fs.test"
|
||||
assert isinstance(config.KNOWLEDGE_FS_JWT_SECRET, SecretStr)
|
||||
assert config.KNOWLEDGE_FS_JWT_SECRET.get_secret_value() == "production-secret-with-at-least-32-bytes"
|
||||
assert "production-secret" not in repr(config)
|
||||
assert "production-secret" not in config.model_dump_json()
|
||||
assert config.KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS == 300.0
|
||||
|
||||
|
||||
def test_knowledge_fs_config_treats_blank_connection_as_disabled() -> None:
|
||||
config = KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=" ",
|
||||
KNOWLEDGE_FS_JWT_SECRET="",
|
||||
)
|
||||
|
||||
assert config.KNOWLEDGE_FS_BASE_URL is None
|
||||
assert config.KNOWLEDGE_FS_JWT_SECRET is None
|
||||
assert config.KNOWLEDGE_FS_ENABLED is False
|
||||
|
||||
|
||||
def test_knowledge_fs_config_requires_connection_when_enabled() -> None:
|
||||
with pytest.raises(ValidationError, match="connection settings are required"):
|
||||
KnowledgeFSConfig(KNOWLEDGE_FS_ENABLED=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_url", "jwt_secret"),
|
||||
[
|
||||
("https://knowledge-fs.test", None),
|
||||
(None, "production-secret-with-at-least-32-bytes"),
|
||||
],
|
||||
)
|
||||
def test_disabled_knowledge_fs_config_allows_partial_connection(base_url: str | None, jwt_secret: str | None) -> None:
|
||||
config = KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_ENABLED=False,
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET=jwt_secret,
|
||||
)
|
||||
|
||||
assert config.KNOWLEDGE_FS_ENABLED is False
|
||||
|
||||
|
||||
def test_knowledge_fs_docker_config_is_not_shadowed_by_root_env() -> None:
|
||||
root_env_example = (_REPOSITORY_ROOT / "docker/.env.example").read_text(encoding="utf-8")
|
||||
api_env_example = (_REPOSITORY_ROOT / "docker/envs/core-services/api.env.example").read_text(encoding="utf-8")
|
||||
|
||||
for variable in _KNOWLEDGE_FS_DOCKER_VARIABLES:
|
||||
assert f"{variable}=" not in root_env_example
|
||||
assert f"{variable}=" in api_env_example
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_url", "jwt_secret"),
|
||||
[
|
||||
("https://knowledge-fs.test", None),
|
||||
(None, "production-secret-with-at-least-32-bytes"),
|
||||
],
|
||||
)
|
||||
def test_knowledge_fs_config_rejects_partial_connection(base_url: str | None, jwt_secret: str | None) -> None:
|
||||
with pytest.raises(ValidationError, match="must be configured together"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_ENABLED=True,
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET=jwt_secret,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("base_url", ["knowledge-fs.test", "ftp://knowledge-fs.test", "http:///missing-host"])
|
||||
def test_knowledge_fs_config_rejects_non_http_absolute_urls(base_url: str) -> None:
|
||||
with pytest.raises(ValidationError, match="absolute HTTP\\(S\\) URL"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base_url",
|
||||
[
|
||||
"https://knowledge-fs.test:notaport",
|
||||
"https://knowledge-fs.test:65536",
|
||||
],
|
||||
)
|
||||
def test_knowledge_fs_config_rejects_invalid_ports(base_url: str) -> None:
|
||||
with pytest.raises(ValidationError, match="valid port"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base_url",
|
||||
[
|
||||
"https://user:password@knowledge-fs.test",
|
||||
"https://knowledge-fs.test?region=us",
|
||||
"https://knowledge-fs.test#gateway",
|
||||
],
|
||||
)
|
||||
def test_knowledge_fs_config_rejects_unsafe_base_url_components(base_url: str) -> None:
|
||||
with pytest.raises(ValidationError, match="must not include credentials, query, or fragment"):
|
||||
KnowledgeFSConfig(
|
||||
KNOWLEDGE_FS_BASE_URL=base_url,
|
||||
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("timeout_seconds", [float("inf"), float("nan"), 60.0001])
|
||||
def test_knowledge_fs_config_rejects_unbounded_timeouts(timeout_seconds: float) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
KnowledgeFSConfig(KNOWLEDGE_FS_TIMEOUT_SECONDS=timeout_seconds)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("timeout_seconds", [float("inf"), float("nan"), 3600.0001])
|
||||
def test_knowledge_fs_config_rejects_unbounded_sse_read_timeouts(timeout_seconds: float) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
KnowledgeFSConfig(KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS=timeout_seconds)
|
||||
@@ -1,23 +1,82 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
from typing import Any
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from controllers.common.agent_app_parameters import get_published_agent_app_feature_dict_and_user_input_form
|
||||
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
|
||||
from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError
|
||||
from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource, AgentStatus
|
||||
from models.model import AppAnnotationSetting
|
||||
|
||||
|
||||
def _app_model(*, bound_agent_id: str | None, app_model_config=None):
|
||||
def _stable_uuid(value: str) -> str:
|
||||
return str(uuid5(NAMESPACE_URL, value))
|
||||
|
||||
|
||||
def _app_model(*, tenant_id: str, bound_agent_id: str | None, app_model_config: object | None = None):
|
||||
return SimpleNamespace(
|
||||
id="app-1",
|
||||
tenant_id="tenant-1",
|
||||
id=_stable_uuid(f"app:{tenant_id}"),
|
||||
tenant_id=tenant_id,
|
||||
bound_agent_id=bound_agent_id,
|
||||
app_model_config_with_session=lambda *, session: app_model_config,
|
||||
)
|
||||
|
||||
|
||||
def test_published_agent_app_parameters_use_soul_file_upload():
|
||||
def _persist_agent(
|
||||
session: Session,
|
||||
*,
|
||||
tenant_id: str,
|
||||
agent_id: str,
|
||||
active_config_snapshot_id: str | None,
|
||||
active_config_is_published: bool,
|
||||
) -> Agent:
|
||||
agent = Agent(
|
||||
id=agent_id,
|
||||
tenant_id=tenant_id,
|
||||
name="Agent",
|
||||
scope=AgentScope.ROSTER,
|
||||
source=AgentSource.AGENT_APP,
|
||||
status=AgentStatus.ACTIVE,
|
||||
active_config_snapshot_id=active_config_snapshot_id,
|
||||
active_config_is_published=active_config_is_published,
|
||||
)
|
||||
session.add(agent)
|
||||
session.commit()
|
||||
return agent
|
||||
|
||||
|
||||
def _persist_snapshot(
|
||||
session: Session,
|
||||
*,
|
||||
snapshot_id: str,
|
||||
tenant_id: str,
|
||||
agent_id: str,
|
||||
config_snapshot: dict[str, Any],
|
||||
) -> AgentConfigSnapshot:
|
||||
snapshot = AgentConfigSnapshot(
|
||||
id=snapshot_id,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
version=1,
|
||||
config_snapshot=config_snapshot,
|
||||
)
|
||||
session.add(snapshot)
|
||||
session.commit()
|
||||
return snapshot
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Agent, AgentConfigSnapshot, AppAnnotationSetting)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_published_agent_app_parameters_use_soul_file_upload(sqlite_session: Session):
|
||||
tenant_id = _stable_uuid("tenant:one")
|
||||
agent_id = _stable_uuid("agent:one")
|
||||
snapshot_id = _stable_uuid("snapshot:one")
|
||||
app_model_config = SimpleNamespace(
|
||||
to_dict=lambda **_kwargs: {
|
||||
"opening_statement": "Hi from legacy presentation config",
|
||||
@@ -27,14 +86,24 @@ def test_published_agent_app_parameters_use_soul_file_upload():
|
||||
},
|
||||
}
|
||||
)
|
||||
app_model = _app_model(bound_agent_id="agent-1", app_model_config=app_model_config)
|
||||
agent = SimpleNamespace(
|
||||
id="agent-1",
|
||||
active_config_snapshot_id="snapshot-1",
|
||||
app_model = _app_model(
|
||||
tenant_id=tenant_id,
|
||||
bound_agent_id=agent_id,
|
||||
app_model_config=app_model_config,
|
||||
)
|
||||
_persist_agent(
|
||||
sqlite_session,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
active_config_snapshot_id=snapshot_id,
|
||||
active_config_is_published=True,
|
||||
)
|
||||
snapshot = SimpleNamespace(
|
||||
config_snapshot_dict={
|
||||
_persist_snapshot(
|
||||
sqlite_session,
|
||||
snapshot_id=snapshot_id,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
config_snapshot={
|
||||
"app_features": {
|
||||
"file_upload": {
|
||||
"enabled": True,
|
||||
@@ -46,14 +115,12 @@ def test_published_agent_app_parameters_use_soul_file_upload():
|
||||
}
|
||||
},
|
||||
"app_variables": [{"name": "topic", "type": "string", "required": True}],
|
||||
}
|
||||
},
|
||||
)
|
||||
session = MagicMock()
|
||||
session.scalar.side_effect = [agent, snapshot, None]
|
||||
|
||||
features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(
|
||||
app_model,
|
||||
session=session,
|
||||
session=sqlite_session,
|
||||
)
|
||||
parameters = get_parameters_from_feature_dict(features_dict=features_dict, user_input_form=user_input_form)
|
||||
|
||||
@@ -69,20 +136,30 @@ def test_published_agent_app_parameters_use_soul_file_upload():
|
||||
assert parameters["user_input_form"] == [{"text-input": {"label": "topic", "variable": "topic", "required": True}}]
|
||||
|
||||
|
||||
def test_published_agent_app_parameters_requires_bound_agent():
|
||||
app_model = _app_model(bound_agent_id=None)
|
||||
@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot)], indirect=True)
|
||||
def test_published_agent_app_parameters_requires_bound_agent(sqlite_session: Session):
|
||||
tenant_id = _stable_uuid("tenant:unbound")
|
||||
app_model = _app_model(tenant_id=tenant_id, bound_agent_id=None)
|
||||
|
||||
with pytest.raises(AgentAppGeneratorError, match="no bound Agent"):
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=MagicMock())
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=sqlite_session)
|
||||
|
||||
|
||||
def test_published_agent_app_parameters_requires_existing_active_agent():
|
||||
app_model = _app_model(bound_agent_id="agent-1")
|
||||
session = MagicMock()
|
||||
session.scalar.return_value = None
|
||||
@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot)], indirect=True)
|
||||
def test_published_agent_app_parameters_requires_existing_active_agent(sqlite_session: Session):
|
||||
requested_tenant_id = _stable_uuid("tenant:requested")
|
||||
agent_id = _stable_uuid("agent:cross-tenant")
|
||||
app_model = _app_model(tenant_id=requested_tenant_id, bound_agent_id=agent_id)
|
||||
_persist_agent(
|
||||
sqlite_session,
|
||||
tenant_id=_stable_uuid("tenant:other"),
|
||||
agent_id=agent_id,
|
||||
active_config_snapshot_id=None,
|
||||
active_config_is_published=False,
|
||||
)
|
||||
|
||||
with pytest.raises(AgentAppGeneratorError, match="no bound Agent"):
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=session)
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=sqlite_session)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -92,68 +169,96 @@ def test_published_agent_app_parameters_requires_existing_active_agent():
|
||||
False,
|
||||
],
|
||||
)
|
||||
def test_published_agent_app_parameters_requires_published_agent(active_config_is_published):
|
||||
app_model = _app_model(bound_agent_id="agent-1")
|
||||
agent = SimpleNamespace(
|
||||
id="agent-1",
|
||||
@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot)], indirect=True)
|
||||
def test_published_agent_app_parameters_requires_published_agent(
|
||||
active_config_is_published: bool, sqlite_session: Session
|
||||
):
|
||||
tenant_id = _stable_uuid(f"tenant:published:{active_config_is_published}")
|
||||
agent_id = _stable_uuid(f"agent:published:{active_config_is_published}")
|
||||
app_model = _app_model(tenant_id=tenant_id, bound_agent_id=agent_id)
|
||||
_persist_agent(
|
||||
sqlite_session,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
active_config_snapshot_id=None,
|
||||
active_config_is_published=active_config_is_published,
|
||||
)
|
||||
session = MagicMock()
|
||||
session.scalar.return_value = agent
|
||||
|
||||
with pytest.raises(AgentAppNotPublishedError, match="not been published"):
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=session)
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=sqlite_session)
|
||||
|
||||
|
||||
def test_published_agent_app_parameters_allows_unpublished_draft_with_active_snapshot():
|
||||
app_model = _app_model(bound_agent_id="agent-1")
|
||||
agent = SimpleNamespace(
|
||||
id="agent-1",
|
||||
active_config_snapshot_id="snapshot-1",
|
||||
@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot)], indirect=True)
|
||||
def test_published_agent_app_parameters_allows_unpublished_draft_with_active_snapshot(sqlite_session: Session):
|
||||
tenant_id = _stable_uuid("tenant:unpublished-draft")
|
||||
agent_id = _stable_uuid("agent:unpublished-draft")
|
||||
snapshot_id = _stable_uuid("snapshot:unpublished-draft")
|
||||
app_model = _app_model(tenant_id=tenant_id, bound_agent_id=agent_id)
|
||||
_persist_agent(
|
||||
sqlite_session,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
active_config_snapshot_id=snapshot_id,
|
||||
active_config_is_published=False,
|
||||
)
|
||||
snapshot = SimpleNamespace(config_snapshot_dict={})
|
||||
session = MagicMock()
|
||||
session.scalar.side_effect = [agent, snapshot]
|
||||
_persist_snapshot(
|
||||
sqlite_session,
|
||||
snapshot_id=snapshot_id,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
config_snapshot={},
|
||||
)
|
||||
|
||||
features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(
|
||||
app_model,
|
||||
session=session,
|
||||
session=sqlite_session,
|
||||
)
|
||||
|
||||
assert features_dict["file_upload"]["enabled"] is True
|
||||
assert user_input_form == []
|
||||
|
||||
|
||||
def test_published_agent_app_parameters_requires_published_snapshot():
|
||||
app_model = _app_model(bound_agent_id="agent-1")
|
||||
agent = SimpleNamespace(
|
||||
id="agent-1",
|
||||
active_config_snapshot_id="snapshot-1",
|
||||
@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot)], indirect=True)
|
||||
def test_published_agent_app_parameters_requires_published_snapshot(sqlite_session: Session):
|
||||
tenant_id = _stable_uuid("tenant:missing-snapshot")
|
||||
agent_id = _stable_uuid("agent:missing-snapshot")
|
||||
app_model = _app_model(tenant_id=tenant_id, bound_agent_id=agent_id)
|
||||
_persist_agent(
|
||||
sqlite_session,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
active_config_snapshot_id=_stable_uuid("snapshot:missing"),
|
||||
active_config_is_published=True,
|
||||
)
|
||||
session = MagicMock()
|
||||
session.scalar.side_effect = [agent, None]
|
||||
|
||||
with pytest.raises(AgentAppGeneratorError, match="published version not found"):
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=session)
|
||||
get_published_agent_app_feature_dict_and_user_input_form(app_model, session=sqlite_session)
|
||||
|
||||
|
||||
def test_published_agent_app_parameters_allows_missing_legacy_app_model_config():
|
||||
app_model = _app_model(bound_agent_id="agent-1")
|
||||
agent = SimpleNamespace(
|
||||
id="agent-1",
|
||||
active_config_snapshot_id="snapshot-1",
|
||||
@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot)], indirect=True)
|
||||
def test_published_agent_app_parameters_allows_missing_legacy_app_model_config(sqlite_session: Session):
|
||||
tenant_id = _stable_uuid("tenant:no-legacy-config")
|
||||
agent_id = _stable_uuid("agent:no-legacy-config")
|
||||
snapshot_id = _stable_uuid("snapshot:no-legacy-config")
|
||||
app_model = _app_model(tenant_id=tenant_id, bound_agent_id=agent_id)
|
||||
_persist_agent(
|
||||
sqlite_session,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
active_config_snapshot_id=snapshot_id,
|
||||
active_config_is_published=True,
|
||||
)
|
||||
snapshot = SimpleNamespace(config_snapshot_dict={})
|
||||
session = MagicMock()
|
||||
session.scalar.side_effect = [agent, snapshot]
|
||||
_persist_snapshot(
|
||||
sqlite_session,
|
||||
snapshot_id=snapshot_id,
|
||||
tenant_id=tenant_id,
|
||||
agent_id=agent_id,
|
||||
config_snapshot={},
|
||||
)
|
||||
|
||||
features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(
|
||||
app_model,
|
||||
session=session,
|
||||
session=sqlite_session,
|
||||
)
|
||||
|
||||
assert features_dict["file_upload"] == {
|
||||
|
||||
@@ -398,6 +398,12 @@ class TestWorkflowEndpoints:
|
||||
def test_workflow_copy_payload(self):
|
||||
payload = SyncDraftWorkflowPayload(graph={}, features={})
|
||||
assert payload.graph == {}
|
||||
assert payload.model_dump()["is_collaborative"] is False
|
||||
|
||||
def test_workflow_sync_payload_accepts_collaboration_marker(self):
|
||||
payload = SyncDraftWorkflowPayload.model_validate({"graph": {}, "features": {}, "_is_collaborative": True})
|
||||
assert payload.is_collaborative is True
|
||||
assert payload.model_dump()["is_collaborative"] is True
|
||||
|
||||
def test_workflow_mode_query(self):
|
||||
payload = AdvancedChatWorkflowRunPayload(inputs={}, query="hi")
|
||||
|
||||
@@ -2,15 +2,23 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from inspect import unwrap
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from controllers.console.app import app_import as app_import_module
|
||||
from models.account import Account
|
||||
from models.engine import db
|
||||
from models.model import App
|
||||
from services.app_dsl_service import ImportStatus
|
||||
from services.entities.dsl_entities import CheckDependenciesResult
|
||||
from services.feature_service import SystemFeatureModel, WebAppAuthModel
|
||||
|
||||
|
||||
def _unwrap(func):
|
||||
@@ -38,17 +46,58 @@ class _Result:
|
||||
|
||||
|
||||
def _install_features(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
|
||||
features = SimpleNamespace(webapp_auth=SimpleNamespace(enabled=enabled))
|
||||
features = SystemFeatureModel(webapp_auth=WebAppAuthModel(enabled=enabled))
|
||||
monkeypatch.setattr(app_import_module.FeatureService, "get_system_features", lambda: features)
|
||||
|
||||
|
||||
def _mock_session(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
|
||||
fake_session = MagicMock()
|
||||
fake_session.__enter__.return_value = fake_session
|
||||
fake_session.__exit__.return_value = None
|
||||
monkeypatch.setattr(app_import_module, "db", SimpleNamespace(engine=object()))
|
||||
monkeypatch.setattr(app_import_module, "Session", lambda *_args, **_kwargs: fake_session)
|
||||
return fake_session
|
||||
def _make_account(account_id: str = "u1") -> Account:
|
||||
account = Account(name="Test User", email="test@example.com")
|
||||
account.id = account_id
|
||||
return account
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app() -> Iterator[Flask]:
|
||||
app = Flask(__name__)
|
||||
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
|
||||
db.init_app(app)
|
||||
|
||||
with app.app_context():
|
||||
yield app
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransactionEvents:
|
||||
commits: int = 0
|
||||
rollbacks: int = 0
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def transaction_events() -> TransactionEvents:
|
||||
"""Observe transaction decisions while keeping the controller on a real SQLAlchemy session."""
|
||||
|
||||
observed = TransactionEvents()
|
||||
|
||||
def record_commit(_session: Session) -> None:
|
||||
observed.commits += 1
|
||||
|
||||
def record_rollback(_session: Session) -> None:
|
||||
observed.rollbacks += 1
|
||||
|
||||
event.listen(Session, "after_commit", record_commit)
|
||||
event.listen(Session, "after_rollback", record_rollback)
|
||||
try:
|
||||
yield observed
|
||||
finally:
|
||||
event.remove(Session, "after_commit", record_commit)
|
||||
event.remove(Session, "after_rollback", record_rollback)
|
||||
|
||||
|
||||
def _failed_result_after_starting_transaction(
|
||||
service: app_import_module.AppDslService, *, app_id: str | None = None
|
||||
) -> _Result:
|
||||
service._session.begin()
|
||||
return _Result(ImportStatus.FAILED, app_id=app_id)
|
||||
|
||||
|
||||
class TestAppImportApi:
|
||||
@@ -57,33 +106,39 @@ class TestAppImportApi:
|
||||
return app_import_module.AppImportApi()
|
||||
|
||||
def test_import_post_returns_failed_status_and_rolls_back(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED, app_id=None),
|
||||
lambda service, *_args, **_kwargs: _failed_result_after_starting_transaction(service, app_id=None),
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
response, status = method(api, _make_account())
|
||||
|
||||
session.rollback.assert_called_once_with()
|
||||
session.commit.assert_not_called()
|
||||
assert transaction_events.rollbacks == 1
|
||||
assert transaction_events.commits == 0
|
||||
assert status == 400
|
||||
assert response["status"] == ImportStatus.FAILED
|
||||
|
||||
def test_import_post_returns_pending_status_and_commits(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
@@ -91,20 +146,23 @@ class TestAppImportApi:
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
response, status = method(api, _make_account())
|
||||
|
||||
session.commit.assert_called_once_with()
|
||||
session.rollback.assert_not_called()
|
||||
assert transaction_events.commits == 1
|
||||
assert transaction_events.rollbacks == 0
|
||||
assert status == 202
|
||||
assert response["status"] == ImportStatus.PENDING
|
||||
|
||||
def test_import_post_updates_webapp_auth_when_enabled(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=True)
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"import_app",
|
||||
@@ -114,25 +172,28 @@ class TestAppImportApi:
|
||||
monkeypatch.setattr(app_import_module.EnterpriseService.WebAppAuth, "update_app_access_mode", update_access)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method(api, SimpleNamespace(id="u1"))
|
||||
response, status = method(api, _make_account())
|
||||
|
||||
session.commit.assert_called_once_with()
|
||||
session.rollback.assert_not_called()
|
||||
assert transaction_events.commits == 1
|
||||
assert transaction_events.rollbacks == 0
|
||||
update_access.assert_called_once_with("app-123", "private")
|
||||
assert status == 200
|
||||
assert response["status"] == ImportStatus.COMPLETED
|
||||
|
||||
def test_import_post_attaches_permission_keys_when_creating_new_app_and_rbac_enabled(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = _unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module,
|
||||
"current_account_with_tenant",
|
||||
lambda: (SimpleNamespace(id="u1"), "tenant-1"),
|
||||
lambda: (_make_account(), "tenant-1"),
|
||||
)
|
||||
monkeypatch.setattr(app_import_module.dify_config, "RBAC_ENABLED", True)
|
||||
monkeypatch.setattr(
|
||||
@@ -149,21 +210,24 @@ class TestAppImportApi:
|
||||
with app.test_request_context("/console/api/apps/imports", method="POST", json={"mode": "yaml-content"}):
|
||||
response, status = method()
|
||||
|
||||
session.commit.assert_called_once_with()
|
||||
assert transaction_events.commits == 1
|
||||
assert status == 200
|
||||
assert response["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"]
|
||||
|
||||
def test_import_post_does_not_attach_permission_keys_when_overwriting_existing_app(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = _unwrap(api.post)
|
||||
|
||||
_install_features(monkeypatch, enabled=False)
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module,
|
||||
"current_account_with_tenant",
|
||||
lambda: (SimpleNamespace(id="u1"), "tenant-1"),
|
||||
lambda: (_make_account(), "tenant-1"),
|
||||
)
|
||||
monkeypatch.setattr(app_import_module.dify_config, "RBAC_ENABLED", True)
|
||||
monkeypatch.setattr(
|
||||
@@ -184,7 +248,7 @@ class TestAppImportApi:
|
||||
):
|
||||
response, status = method()
|
||||
|
||||
session.commit.assert_called_once_with()
|
||||
assert transaction_events.commits == 1
|
||||
assert status == 200
|
||||
assert response["permission_keys"] == []
|
||||
|
||||
@@ -195,35 +259,41 @@ class TestAppImportConfirmApi:
|
||||
return app_import_module.AppImportConfirmApi()
|
||||
|
||||
def test_import_confirm_returns_failed_status_and_rolls_back(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = unwrap(api.post)
|
||||
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"confirm_import",
|
||||
lambda *_args, **_kwargs: _Result(ImportStatus.FAILED),
|
||||
lambda service, *_args, **_kwargs: _failed_result_after_starting_transaction(service),
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports/import-1/confirm", method="POST"):
|
||||
response, status = method(api, SimpleNamespace(id="u1"), import_id="import-1")
|
||||
response, status = method(api, _make_account(), import_id="import-1")
|
||||
|
||||
session.rollback.assert_called_once_with()
|
||||
session.commit.assert_not_called()
|
||||
assert transaction_events.rollbacks == 1
|
||||
assert transaction_events.commits == 0
|
||||
assert status == 400
|
||||
assert response["status"] == ImportStatus.FAILED
|
||||
|
||||
def test_import_confirm_attaches_permission_keys_when_creating_new_app_and_rbac_enabled(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = _unwrap(api.post)
|
||||
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module,
|
||||
"current_account_with_tenant",
|
||||
lambda: (SimpleNamespace(id="u1"), "tenant-1"),
|
||||
lambda: (_make_account(), "tenant-1"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.redis_client,
|
||||
@@ -248,20 +318,23 @@ class TestAppImportConfirmApi:
|
||||
with app.test_request_context("/console/api/apps/imports/import-1/confirm", method="POST"):
|
||||
response, status = method(import_id="import-1")
|
||||
|
||||
session.commit.assert_called_once_with()
|
||||
assert transaction_events.commits == 1
|
||||
assert status == 200
|
||||
assert response["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"]
|
||||
|
||||
def test_import_confirm_does_not_attach_permission_keys_when_overwriting_existing_app(
|
||||
self, api, app: Flask, monkeypatch: pytest.MonkeyPatch
|
||||
self,
|
||||
api,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
transaction_events: TransactionEvents,
|
||||
) -> None:
|
||||
method = _unwrap(api.post)
|
||||
|
||||
session = _mock_session(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
app_import_module,
|
||||
"current_account_with_tenant",
|
||||
lambda: (SimpleNamespace(id="u1"), "tenant-1"),
|
||||
lambda: (_make_account(), "tenant-1"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.redis_client,
|
||||
@@ -286,6 +359,27 @@ class TestAppImportConfirmApi:
|
||||
with app.test_request_context("/console/api/apps/imports/import-1/confirm", method="POST"):
|
||||
response, status = method(import_id="import-1")
|
||||
|
||||
session.commit.assert_called_once_with()
|
||||
assert transaction_events.commits == 1
|
||||
assert status == 200
|
||||
assert response["permission_keys"] == []
|
||||
|
||||
|
||||
class TestAppImportCheckDependenciesApi:
|
||||
def test_import_check_dependencies_returns_result(
|
||||
self,
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
api = app_import_module.AppImportCheckDependenciesApi()
|
||||
method = unwrap(api.get)
|
||||
monkeypatch.setattr(
|
||||
app_import_module.AppDslService,
|
||||
"check_dependencies",
|
||||
lambda *_args, **_kwargs: CheckDependenciesResult(leaked_dependencies=[]),
|
||||
)
|
||||
|
||||
with app.test_request_context("/console/api/apps/imports/app-1/check-dependencies", method="GET"):
|
||||
response, status = method(api, app_model=App(id="app-1"))
|
||||
|
||||
assert status == 200
|
||||
assert response["leaked_dependencies"] == []
|
||||
|
||||
@@ -7,10 +7,28 @@ from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from controllers.console.app import generator as generator_module
|
||||
from controllers.console.app.error import ProviderNotInitializeError
|
||||
from core.errors.error import ProviderTokenNotInitError
|
||||
from models.model import App, AppMode
|
||||
|
||||
|
||||
def _persist_app(session: Session, *, tenant_id: str = "t1") -> App:
|
||||
app_model = App(
|
||||
id="app-1",
|
||||
tenant_id=tenant_id,
|
||||
name="Workflow App",
|
||||
description="",
|
||||
mode=AppMode.WORKFLOW,
|
||||
enable_site=False,
|
||||
enable_api=False,
|
||||
max_active_requests=None,
|
||||
)
|
||||
session.add(app_model)
|
||||
session.commit()
|
||||
return app_model
|
||||
|
||||
|
||||
def _model_config_payload():
|
||||
@@ -66,12 +84,11 @@ def test_rule_code_generate_maps_token_error(app: Flask, monkeypatch: pytest.Mon
|
||||
method(api, "t1")
|
||||
|
||||
|
||||
def test_instruction_generate_app_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
|
||||
def test_instruction_generate_app_not_found(app: Flask, sqlite_session: Session) -> None:
|
||||
api = generator_module.InstructionGenerateApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
session = MagicMock()
|
||||
session.scalar.return_value = None
|
||||
_persist_app(sqlite_session, tenant_id="other-tenant")
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/instruction-generate",
|
||||
@@ -83,25 +100,21 @@ def test_instruction_generate_app_not_found(app: Flask, monkeypatch: pytest.Monk
|
||||
"model_config": _model_config_payload(),
|
||||
},
|
||||
):
|
||||
response, status = method(api, session, "t1")
|
||||
response, status = method(api, sqlite_session, "t1")
|
||||
|
||||
assert status == 400
|
||||
assert response["error"] == "app app-1 not found"
|
||||
stmt = session.scalar.call_args.args[0]
|
||||
compiled = stmt.compile()
|
||||
statement = str(compiled)
|
||||
assert "apps.id" in statement
|
||||
assert "apps.tenant_id" in statement
|
||||
assert "app-1" in compiled.params.values()
|
||||
assert "t1" in compiled.params.values()
|
||||
assert sqlite_session.get(App, "app-1") is not None
|
||||
|
||||
|
||||
def test_instruction_generate_workflow_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
|
||||
def test_instruction_generate_workflow_not_found(
|
||||
app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
|
||||
) -> None:
|
||||
api = generator_module.InstructionGenerateApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
app_model = SimpleNamespace(id="app-1")
|
||||
session = SimpleNamespace(scalar=lambda *_args, **_kwargs: app_model)
|
||||
app_model = _persist_app(sqlite_session)
|
||||
_install_workflow_service(monkeypatch, workflow=None)
|
||||
|
||||
with app.test_request_context(
|
||||
@@ -114,18 +127,20 @@ def test_instruction_generate_workflow_not_found(app: Flask, monkeypatch: pytest
|
||||
"model_config": _model_config_payload(),
|
||||
},
|
||||
):
|
||||
response, status = method(api, session, "t1")
|
||||
response, status = method(api, sqlite_session, "t1")
|
||||
|
||||
assert status == 400
|
||||
assert response["error"] == "workflow app-1 not found"
|
||||
|
||||
|
||||
def test_instruction_generate_node_missing(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
|
||||
def test_instruction_generate_node_missing(
|
||||
app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
|
||||
) -> None:
|
||||
api = generator_module.InstructionGenerateApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
app_model = SimpleNamespace(id="app-1")
|
||||
session = SimpleNamespace(scalar=lambda *_args, **_kwargs: app_model)
|
||||
app_model = _persist_app(sqlite_session)
|
||||
|
||||
workflow = SimpleNamespace(graph_dict={"nodes": []})
|
||||
_install_workflow_service(monkeypatch, workflow=workflow)
|
||||
@@ -140,18 +155,18 @@ def test_instruction_generate_node_missing(app: Flask, monkeypatch: pytest.Monke
|
||||
"model_config": _model_config_payload(),
|
||||
},
|
||||
):
|
||||
response, status = method(api, session, "t1")
|
||||
response, status = method(api, sqlite_session, "t1")
|
||||
|
||||
assert status == 400
|
||||
assert response["error"] == "node node-1 not found"
|
||||
|
||||
|
||||
def test_instruction_generate_code_node(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
|
||||
def test_instruction_generate_code_node(app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None:
|
||||
api = generator_module.InstructionGenerateApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
app_model = SimpleNamespace(id="app-1")
|
||||
session = SimpleNamespace(scalar=lambda *_args, **_kwargs: app_model)
|
||||
app_model = _persist_app(sqlite_session)
|
||||
|
||||
workflow = SimpleNamespace(
|
||||
graph_dict={
|
||||
@@ -173,18 +188,19 @@ def test_instruction_generate_code_node(app: Flask, monkeypatch: pytest.MonkeyPa
|
||||
"model_config": _model_config_payload(),
|
||||
},
|
||||
):
|
||||
response = method(api, session, "t1")
|
||||
response = method(api, sqlite_session, "t1")
|
||||
|
||||
assert response == {"code": "x"}
|
||||
assert workflow_service.app_model is app_model
|
||||
assert workflow_service.session is session
|
||||
assert workflow_service.session is sqlite_session
|
||||
|
||||
|
||||
def test_instruction_generate_legacy_modify(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
|
||||
def test_instruction_generate_legacy_modify(
|
||||
app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
|
||||
) -> None:
|
||||
api = generator_module.InstructionGenerateApi()
|
||||
method = unwrap(api.post)
|
||||
session = SimpleNamespace()
|
||||
|
||||
monkeypatch.setattr(
|
||||
generator_module.LLMGenerator,
|
||||
"instruction_modify_legacy",
|
||||
@@ -202,16 +218,15 @@ def test_instruction_generate_legacy_modify(app: Flask, monkeypatch: pytest.Monk
|
||||
"model_config": _model_config_payload(),
|
||||
},
|
||||
):
|
||||
response = method(api, session, "t1")
|
||||
response = method(api, sqlite_session, "t1")
|
||||
|
||||
assert response == {"instruction": "ok"}
|
||||
|
||||
|
||||
def test_instruction_generate_incompatible_params(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
|
||||
def test_instruction_generate_incompatible_params(app: Flask, sqlite_session: Session) -> None:
|
||||
api = generator_module.InstructionGenerateApi()
|
||||
method = unwrap(api.post)
|
||||
session = SimpleNamespace()
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/instruction-generate",
|
||||
method="POST",
|
||||
@@ -223,7 +238,7 @@ def test_instruction_generate_incompatible_params(app: Flask, monkeypatch: pytes
|
||||
"model_config": _model_config_payload(),
|
||||
},
|
||||
):
|
||||
response, status = method(api, session, "t1")
|
||||
response, status = method(api, sqlite_session, "t1")
|
||||
|
||||
assert status == 400
|
||||
assert response["error"] == "incompatible parameters"
|
||||
|
||||
+17
-13
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from inspect import unwrap
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import PropertyMock, patch
|
||||
|
||||
import pytest
|
||||
@@ -275,9 +276,9 @@ class TestCustomizedPipelineTemplateApi:
|
||||
assert (response, status) == ("", 204)
|
||||
assert deleted_templates == [("template-1", tenant_id)]
|
||||
|
||||
@pytest.mark.parametrize("sqlite_session", [(PipelineCustomizedTemplate,)], indirect=True)
|
||||
def test_post_exports_yaml_from_orm_template(
|
||||
self,
|
||||
database_app: Flask,
|
||||
self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session
|
||||
) -> None:
|
||||
api = CustomizedPipelineTemplateApi()
|
||||
method = unwrap(api.post)
|
||||
@@ -293,26 +294,29 @@ class TestCustomizedPipelineTemplateApi:
|
||||
language="en-US",
|
||||
created_by="00000000-0000-0000-0000-000000000002",
|
||||
)
|
||||
db.session.add(template)
|
||||
db.session.commit()
|
||||
template.id = "template-1"
|
||||
sqlite_session.add(template)
|
||||
sqlite_session.commit()
|
||||
monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine))
|
||||
|
||||
with database_app.test_request_context("/rag/pipeline/customized/templates/template-1", method="POST"):
|
||||
response, status = method(api, template.id)
|
||||
with app.test_request_context("/rag/pipeline/customized/templates/template-1", method="POST"):
|
||||
response, status = method(api, "template-1")
|
||||
|
||||
assert status == 200
|
||||
assert response == {"data": "dsl: value"}
|
||||
|
||||
@pytest.mark.parametrize("sqlite_session", [(PipelineCustomizedTemplate,)], indirect=True)
|
||||
def test_post_raises_when_template_is_missing(
|
||||
self,
|
||||
database_app: Flask,
|
||||
self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session
|
||||
) -> None:
|
||||
api = CustomizedPipelineTemplateApi()
|
||||
method = unwrap(api.post)
|
||||
with (
|
||||
database_app.test_request_context("/rag/pipeline/customized/templates/missing", method="POST"),
|
||||
pytest.raises(ValueError, match="Customized pipeline template not found"),
|
||||
):
|
||||
method(api, "44444444-4444-4444-4444-444444444444")
|
||||
assert sqlite_session.get(PipelineCustomizedTemplate, "missing") is None
|
||||
monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine))
|
||||
|
||||
with app.test_request_context("/rag/pipeline/customized/templates/missing", method="POST"):
|
||||
with pytest.raises(ValueError, match="Customized pipeline template not found"):
|
||||
method(api, "missing")
|
||||
|
||||
|
||||
class TestPublishCustomizedPipelineTemplateApi:
|
||||
|
||||
@@ -189,6 +189,44 @@ def test_trial_app_detail_serializes_with_explicit_session(app: Flask, monkeypat
|
||||
module.TrialAppDetailResponse.model_validate.assert_called_once_with(response_view, from_attributes=True)
|
||||
|
||||
|
||||
class TestTrialAppFileUploadApi:
|
||||
def test_upload_uses_trial_app_tenant(self, app: Flask, account: Account) -> None:
|
||||
api = module.TrialAppFileUploadApi()
|
||||
method = unwrap(api.post)
|
||||
app_model = SimpleNamespace(tenant_id="app-tenant-id")
|
||||
upload_file = MagicMock()
|
||||
|
||||
with (
|
||||
app.test_request_context("/", method="POST"),
|
||||
patch.object(module, "upload_file_from_request", return_value=upload_file) as upload,
|
||||
patch.object(module, "dump_response", return_value={"id": "upload-file-id"}),
|
||||
):
|
||||
response, status = method(api, account, app_model)
|
||||
|
||||
assert status == 201
|
||||
assert response == {"id": "upload-file-id"}
|
||||
upload.assert_called_once_with(current_user=account, resource_tenant_id="app-tenant-id")
|
||||
|
||||
|
||||
class TestTrialAppRemoteFileUploadApi:
|
||||
def test_upload_uses_trial_app_tenant(self, app: Flask, account: Account) -> None:
|
||||
api = module.TrialAppRemoteFileUploadApi()
|
||||
method = unwrap(api.post)
|
||||
app_model = SimpleNamespace(tenant_id="app-tenant-id")
|
||||
remote_file = MagicMock()
|
||||
remote_file.model_dump.return_value = {"id": "upload-file-id"}
|
||||
|
||||
with (
|
||||
app.test_request_context("/", method="POST", json={"url": "https://example.com/file.txt"}),
|
||||
patch.object(module, "upload_remote_file_from_request", return_value=remote_file) as upload,
|
||||
):
|
||||
response, status = method(api, account, app_model)
|
||||
|
||||
assert status == 201
|
||||
assert response == {"id": "upload-file-id"}
|
||||
upload.assert_called_once_with(current_user=account, resource_tenant_id="app-tenant-id")
|
||||
|
||||
|
||||
class TestTrialAppWorkflowRunApi:
|
||||
def test_not_workflow_app(self, app: Flask, account: Account) -> None:
|
||||
api = module.TrialAppWorkflowRunApi()
|
||||
@@ -800,6 +838,29 @@ class TestTrialChatAudioApi:
|
||||
with pytest.raises(module.NoAudioUploadedError):
|
||||
method(api, account, trial_app_chat)
|
||||
|
||||
def test_missing_file_field_returns_400(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None:
|
||||
"""A multipart POST with no `file` field must surface as 400, not 500.
|
||||
|
||||
Verifies the controller passes file=None to AudioService.transcript_asr
|
||||
instead of raising a KeyError that would yield HTTP 500.
|
||||
"""
|
||||
|
||||
def fake_asr(*args, **kwargs):
|
||||
assert kwargs["file"] is None
|
||||
raise module.services.errors.audio.NoAudioUploadedServiceError()
|
||||
|
||||
api = module.TrialChatAudioApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
with (
|
||||
app.test_request_context("/", method="POST", data={}, content_type="multipart/form-data"),
|
||||
patch.object(module.AudioService, "transcript_asr", side_effect=fake_asr),
|
||||
):
|
||||
with pytest.raises(module.NoAudioUploadedError) as exc_info:
|
||||
method(api, account, trial_app_chat)
|
||||
|
||||
assert exc_info.value.code == 400
|
||||
|
||||
def test_audio_too_large(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None:
|
||||
api = module.TrialChatAudioApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
@@ -18,6 +18,7 @@ from controllers.console.files import (
|
||||
FileApi,
|
||||
FilePreviewApi,
|
||||
FileSupportTypeApi,
|
||||
upload_file_from_request,
|
||||
)
|
||||
from models import Account
|
||||
from models.account import AccountStatus, TenantAccountRole
|
||||
@@ -181,6 +182,22 @@ class TestFileApiPost:
|
||||
assert response["id"] == "file-id-123"
|
||||
assert response["name"] == "test.txt"
|
||||
|
||||
def test_upload_with_resource_tenant(self, app: Flask, mock_account_context, mock_file_service):
|
||||
upload_file = MagicMock()
|
||||
mock_file_service.upload_file.return_value = upload_file
|
||||
|
||||
with app.test_request_context(
|
||||
method="POST",
|
||||
data={"file": (io.BytesIO(b"hello"), "test.txt")},
|
||||
):
|
||||
result = upload_file_from_request(
|
||||
current_user=mock_account_context,
|
||||
resource_tenant_id="app-tenant-id",
|
||||
)
|
||||
|
||||
assert result is upload_file
|
||||
assert mock_file_service.upload_file.call_args.kwargs["tenant_id"] == "app-tenant-id"
|
||||
|
||||
def test_upload_with_invalid_source(self, app: Flask, mock_account_context, mock_file_service):
|
||||
"""Test that invalid source parameter gets normalized to None"""
|
||||
api = FileApi()
|
||||
|
||||
@@ -0,0 +1,639 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
from inspect import unwrap
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from flask import Flask, Response
|
||||
from werkzeug.exceptions import (
|
||||
BadGateway,
|
||||
Forbidden,
|
||||
NotFound,
|
||||
RequestEntityTooLarge,
|
||||
ServiceUnavailable,
|
||||
)
|
||||
|
||||
from controllers.console import bp
|
||||
from controllers.console.knowledge_fs_proxy import (
|
||||
_console_api_errors,
|
||||
_proxy_knowledge_fs_non_get,
|
||||
_proxy_request,
|
||||
_proxy_response,
|
||||
proxy_knowledge_fs_get,
|
||||
proxy_knowledge_fs_write,
|
||||
)
|
||||
from controllers.console.wraps import RBACPermission
|
||||
from services.knowledge_fs_proxy import (
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSMethod,
|
||||
KnowledgeFSOperation,
|
||||
KnowledgeFSResponseKind,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSUpstreamResponse,
|
||||
get_knowledge_fs_operation,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_knowledge_fs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.dify_config.KNOWLEDGE_FS_ENABLED", True)
|
||||
|
||||
|
||||
def _upstream(
|
||||
response: httpx.Response,
|
||||
kind: KnowledgeFSResponseKind = "buffered",
|
||||
*,
|
||||
max_response_bytes: int | None = None,
|
||||
) -> KnowledgeFSUpstreamResponse:
|
||||
operation = KnowledgeFSOperation(
|
||||
operation_id="testOperation",
|
||||
method="GET",
|
||||
path="test",
|
||||
response_kind=kind,
|
||||
required_scope="knowledge-spaces:read",
|
||||
rbac_permission=RBACPermission.DATASET_READONLY,
|
||||
requires_dataset_editor=False,
|
||||
max_response_bytes=max_response_bytes
|
||||
or (64 * 1024 * 1024 if kind == "stream" else 25 * 1024 * 1024 if kind == "binary" else 1024 * 1024),
|
||||
request_headers=(),
|
||||
response_headers=(
|
||||
"content-security-policy",
|
||||
"x-content-type-options",
|
||||
"x-query-run-id",
|
||||
"x-session-id",
|
||||
),
|
||||
response_media_types=(),
|
||||
)
|
||||
return KnowledgeFSUpstreamResponse(response, kind, operation)
|
||||
|
||||
|
||||
def _set_current_workspace(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
editor: bool = True,
|
||||
has_edit_permission: bool = True,
|
||||
admin_or_owner: bool = True,
|
||||
) -> None:
|
||||
account = MagicMock(
|
||||
id="account-1",
|
||||
has_edit_permission=has_edit_permission,
|
||||
is_admin_or_owner=admin_or_owner,
|
||||
is_dataset_editor=editor,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.current_account_with_tenant",
|
||||
lambda: (account, "tenant-1"),
|
||||
)
|
||||
|
||||
|
||||
def _bypass_policy_wrappers(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy._proxy_knowledge_fs_non_get",
|
||||
unwrap(_proxy_knowledge_fs_non_get),
|
||||
)
|
||||
|
||||
|
||||
def test_console_blueprint_registers_generic_knowledge_fs_routes() -> None:
|
||||
app = Flask("knowledge-fs-route-registration")
|
||||
app.register_blueprint(bp)
|
||||
adapter = app.url_map.bind("localhost")
|
||||
|
||||
get_endpoint, get_values = adapter.match(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="GET",
|
||||
)
|
||||
assert get_endpoint.endswith("proxy_knowledge_fs_get")
|
||||
assert get_values == {"upstream_path": "knowledge-spaces"}
|
||||
for method in ("DELETE", "PATCH", "POST", "PUT"):
|
||||
write_endpoint, write_values = adapter.match(
|
||||
"/console/api/knowledge-fs/knowledge-spaces/space-1",
|
||||
method=method,
|
||||
)
|
||||
assert write_endpoint.endswith("proxy_knowledge_fs_write")
|
||||
assert write_values == {"upstream_path": "knowledge-spaces/space-1"}
|
||||
options_endpoint, options_values = adapter.match(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="OPTIONS",
|
||||
)
|
||||
assert options_endpoint.endswith("proxy_knowledge_fs_get")
|
||||
assert options_values == {"upstream_path": "knowledge-spaces"}
|
||||
|
||||
|
||||
def test_proxy_is_hidden_when_knowledge_fs_is_disabled(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.dify_config.KNOWLEDGE_FS_ENABLED", False)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method="GET"):
|
||||
with pytest.raises(NotFound):
|
||||
_proxy_request("GET", "knowledge-spaces")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "method"),
|
||||
[
|
||||
(proxy_knowledge_fs_get, "GET"),
|
||||
(proxy_knowledge_fs_write, "POST"),
|
||||
],
|
||||
)
|
||||
def test_proxy_routes_are_hidden_before_downstream_work_when_disabled(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route,
|
||||
method: KnowledgeFSMethod,
|
||||
) -> None:
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.dify_config.KNOWLEDGE_FS_ENABLED", False)
|
||||
proxy_request = MagicMock()
|
||||
proxy_non_get = MagicMock()
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy_request)
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_knowledge_fs_non_get", proxy_non_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method=method):
|
||||
response = app.make_response(route("knowledge-spaces"))
|
||||
|
||||
assert response.status_code == 404
|
||||
proxy_request.assert_not_called()
|
||||
proxy_non_get.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["HEAD", "OPTIONS"])
|
||||
def test_generic_get_hides_unregistered_methods(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method: str,
|
||||
) -> None:
|
||||
proxy = MagicMock(return_value=Response(status=200))
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method=method,
|
||||
):
|
||||
with pytest.raises(NotFound):
|
||||
route("knowledge-spaces")
|
||||
|
||||
proxy.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "method", "path", "permission"),
|
||||
[
|
||||
(proxy_knowledge_fs_get, "GET", "knowledge-spaces", RBACPermission.DATASET_READONLY),
|
||||
(
|
||||
proxy_knowledge_fs_write,
|
||||
"POST",
|
||||
"knowledge-spaces",
|
||||
RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_generic_routes_delegate_to_the_authorized_service_use_case(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route,
|
||||
method: KnowledgeFSMethod,
|
||||
path: str,
|
||||
permission: RBACPermission,
|
||||
) -> None:
|
||||
response_kind: KnowledgeFSResponseKind = "buffered"
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"asset" if response_kind == "binary" else b"{}",
|
||||
headers={"Content-Type": "application/octet-stream" if response_kind == "binary" else "application/json"},
|
||||
)
|
||||
proxy = MagicMock(return_value=_upstream(upstream, response_kind))
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
proxy,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
raw_route = unwrap(route)
|
||||
|
||||
with app.test_request_context(f"/console/api/knowledge-fs/{path}", method=method, data=b"{}"):
|
||||
response = raw_route(path)
|
||||
|
||||
assert isinstance(response, Response)
|
||||
assert proxy.call_args.kwargs["account"].id == "account-1"
|
||||
assert proxy.call_args.kwargs["tenant_id"] == "tenant-1"
|
||||
assert proxy.call_args.kwargs["method"] == method
|
||||
assert proxy.call_args.kwargs["path"] == path
|
||||
assert get_knowledge_fs_operation(method, path).rbac_permission == permission
|
||||
|
||||
|
||||
def test_read_post_applies_knowledge_rate_limit_once(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr("controllers.common.wraps.dify_config.RBAC_ENABLED", False)
|
||||
account = MagicMock(id="account-1", is_dataset_editor=True)
|
||||
|
||||
def current_workspace() -> tuple[MagicMock, str]:
|
||||
return account, "tenant-1"
|
||||
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.current_account_with_tenant", current_workspace)
|
||||
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", current_workspace)
|
||||
monkeypatch.setattr(
|
||||
"services.knowledge_fs_proxy.RBACService.CheckAccess.check",
|
||||
MagicMock(return_value=True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.wraps.FeatureService.get_knowledge_rate_limit",
|
||||
MagicMock(return_value=MagicMock(enabled=True, limit=10)),
|
||||
)
|
||||
zadd = MagicMock()
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zadd", zadd)
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zremrangebyscore", MagicMock())
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zcard", MagicMock(return_value=1))
|
||||
proxy = MagicMock(return_value=Response(status=200))
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method="POST"):
|
||||
response = _proxy_knowledge_fs_non_get("POST", "knowledge-spaces")
|
||||
|
||||
assert isinstance(response, Response)
|
||||
zadd.assert_called_once()
|
||||
proxy.assert_called_once_with("POST", "knowledge-spaces")
|
||||
|
||||
|
||||
def test_denied_write_does_not_consume_the_workspace_rate_limit(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
account = MagicMock(id="account-1", is_dataset_editor=False)
|
||||
|
||||
def current_workspace() -> tuple[MagicMock, str]:
|
||||
return account, "tenant-1"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.current_account_with_tenant",
|
||||
current_workspace,
|
||||
)
|
||||
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", current_workspace)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.wraps.FeatureService.get_knowledge_rate_limit",
|
||||
MagicMock(return_value=MagicMock(enabled=True, limit=10)),
|
||||
)
|
||||
zadd = MagicMock()
|
||||
monkeypatch.setattr("controllers.console.wraps.redis_client.zadd", zadd)
|
||||
upstream_request = MagicMock()
|
||||
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", upstream_request)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="POST",
|
||||
data=b"{}",
|
||||
content_type="application/json",
|
||||
):
|
||||
with pytest.raises(Forbidden):
|
||||
route("knowledge-spaces")
|
||||
|
||||
zadd.assert_not_called()
|
||||
upstream_request.assert_not_called()
|
||||
|
||||
|
||||
def test_generic_get_forwards_path_query_and_raw_response(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
forward = MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
200,
|
||||
content=gzip.compress(b'{"items":[],"nextCursor":null}'),
|
||||
headers={
|
||||
"Cache-Control": "no-store",
|
||||
"Content-Encoding": "gzip",
|
||||
"Content-Disposition": 'attachment; filename="result.json"',
|
||||
"Content-Type": "application/json",
|
||||
"Retry-After": "3",
|
||||
"Set-Cookie": "kfs=secret",
|
||||
"X-Trace-Id": "trace-1",
|
||||
},
|
||||
)
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
query_string=[("limit", "20"), ("cursor", "first"), ("cursor", "second")],
|
||||
):
|
||||
response = route("knowledge-spaces")
|
||||
|
||||
request = forward.call_args.kwargs
|
||||
assert request["account"].id == "account-1"
|
||||
assert request["method"] == "GET"
|
||||
assert request["path"] == "knowledge-spaces"
|
||||
assert request["tenant_id"] == "tenant-1"
|
||||
assert request["accept"] is None
|
||||
assert request["content_type"] is None
|
||||
assert request["query"] == b"limit=20&cursor=first&cursor=second"
|
||||
assert request["body"] is None
|
||||
assert isinstance(response, Response)
|
||||
assert response.status_code == 200
|
||||
assert response.get_json() == {"items": [], "nextCursor": None}
|
||||
assert response.headers["Cache-Control"] == "no-store"
|
||||
assert response.headers["Content-Disposition"] == 'attachment; filename="result.json"'
|
||||
assert response.headers["Retry-After"] == "3"
|
||||
assert response.headers["X-Trace-Id"] == "trace-1"
|
||||
assert "Content-Encoding" not in response.headers
|
||||
assert "Set-Cookie" not in response.headers
|
||||
|
||||
|
||||
def test_generic_write_forwards_path_raw_body_and_current_tenant(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
method = "POST"
|
||||
path = "knowledge-spaces"
|
||||
forward = MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
201,
|
||||
content=b'{"id":"space-1","tenantId":"tenant-1"}',
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
body = b'{"idempotencyKey":"create-product-docs","name":"Product docs"}'
|
||||
|
||||
with app.test_request_context(
|
||||
f"/console/api/knowledge-fs/{path}",
|
||||
method=method,
|
||||
data=body,
|
||||
content_type="application/json",
|
||||
):
|
||||
response = route(path)
|
||||
|
||||
request = forward.call_args.kwargs
|
||||
assert request["account"].id == "account-1"
|
||||
assert request["method"] == method
|
||||
assert request["path"] == path
|
||||
assert request["tenant_id"] == "tenant-1"
|
||||
assert request["accept"] is None
|
||||
assert request["content_type"] == "application/json"
|
||||
assert request["query"] is None
|
||||
assert request["body"] == body
|
||||
assert isinstance(response, Response)
|
||||
assert response.status_code == 201
|
||||
assert response.get_json()["tenantId"] == "tenant-1"
|
||||
|
||||
|
||||
def test_generic_write_forwards_contract_declared_request_headers(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
forward = MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(202, content=b'{"status":"accepted"}', headers={"Content-Type": "application/json"})
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
body = b'{"name":"Product docs"}'
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="POST",
|
||||
data=body,
|
||||
content_type="application/json",
|
||||
headers={"X-Trace-Id": "trace-1"},
|
||||
):
|
||||
response = route("knowledge-spaces")
|
||||
|
||||
assert isinstance(response, Response)
|
||||
assert response.status_code == 202
|
||||
assert forward.call_args.kwargs["request_headers"].get("X-Trace-Id") == "trace-1"
|
||||
|
||||
|
||||
def test_contract_response_headers_cannot_bypass_the_proxy_denylist() -> None:
|
||||
denied_headers = (
|
||||
"connection",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
"te",
|
||||
"trailer",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
)
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"{}",
|
||||
headers=dict.fromkeys(denied_headers, "blocked"),
|
||||
)
|
||||
|
||||
response = _proxy_response(
|
||||
_upstream(upstream),
|
||||
tenant_id="tenant-1",
|
||||
contract_response_headers=denied_headers,
|
||||
max_response_bytes=1024 * 1024,
|
||||
)
|
||||
|
||||
for name in denied_headers:
|
||||
assert name not in response.headers
|
||||
|
||||
|
||||
def test_contract_response_headers_forward_binary_hardening_headers() -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"asset",
|
||||
headers={
|
||||
"Content-Security-Policy": "sandbox; default-src 'none'",
|
||||
"Content-Type": "image/png",
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
},
|
||||
)
|
||||
|
||||
response = _proxy_response(
|
||||
_upstream(upstream, "binary"),
|
||||
tenant_id="tenant-1",
|
||||
contract_response_headers=("content-security-policy", "x-content-type-options"),
|
||||
max_response_bytes=25 * 1024 * 1024,
|
||||
)
|
||||
|
||||
assert response.headers["Content-Security-Policy"] == "sandbox; default-src 'none'"
|
||||
assert response.headers["X-Content-Type-Options"] == "nosniff"
|
||||
|
||||
|
||||
def test_authorized_service_denial_is_exposed_as_forbidden(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(side_effect=KnowledgeFSAccessDeniedError("workspace access denied")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(Forbidden):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_generic_post_rejects_oversized_body(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
forward = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._MAX_PROXY_BODY_BYTES", 8)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces",
|
||||
method="POST",
|
||||
data=b"x" * 9,
|
||||
content_type="application/json",
|
||||
):
|
||||
with pytest.raises(RequestEntityTooLarge):
|
||||
route("knowledge-spaces")
|
||||
|
||||
forward.assert_not_called()
|
||||
|
||||
|
||||
def test_server_credential_rejection_is_not_exposed_as_browser_auth_failure(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
401,
|
||||
content=b'{"error":"invalid server credential"}',
|
||||
headers={"Content-Type": "application/json", "WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
)
|
||||
),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(BadGateway, match="authentication failed"):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_resource_authorization_rejection_is_exposed_as_forbidden(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(
|
||||
return_value=_upstream(
|
||||
httpx.Response(
|
||||
403,
|
||||
content=b'{"error":"resource access denied"}',
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
)
|
||||
),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(Forbidden):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_contract_response_headers_are_deduplicated_case_insensitively() -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
content=b"asset",
|
||||
headers={
|
||||
"Cache-Control": "private",
|
||||
"Content-Disposition": 'inline; filename="asset.png"',
|
||||
"Content-Type": "image/png",
|
||||
},
|
||||
)
|
||||
|
||||
response = _proxy_response(
|
||||
_upstream(upstream, "binary"),
|
||||
tenant_id="tenant-1",
|
||||
contract_response_headers=("cache-control", "content-disposition"),
|
||||
max_response_bytes=25 * 1024 * 1024,
|
||||
)
|
||||
|
||||
assert response.headers.getlist("Cache-Control") == ["private"]
|
||||
assert response.headers.getlist("Content-Disposition") == ['inline; filename="asset.png"']
|
||||
|
||||
|
||||
def test_configuration_error_is_reported_as_unavailable(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(side_effect=KnowledgeFSConfigurationError("missing token")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
with pytest.raises(ServiceUnavailable, match="misconfigured"):
|
||||
route("knowledge-spaces")
|
||||
|
||||
|
||||
def test_disallowed_kfs_route_is_hidden_as_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(side_effect=KnowledgeFSRouteNotAllowedError("blocked")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/openapi.json"):
|
||||
with pytest.raises(NotFound):
|
||||
route("openapi.json")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["PATCH", "POST"])
|
||||
def test_disallowed_non_get_route_is_hidden_as_not_found(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method: KnowledgeFSMethod,
|
||||
) -> None:
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/not-a-route", method=method):
|
||||
with pytest.raises(NotFound):
|
||||
route("not-a-route")
|
||||
|
||||
|
||||
def test_raw_route_uses_console_json_error_handler(app: Flask) -> None:
|
||||
def forbidden(_upstream_path: str) -> Response:
|
||||
raise Forbidden("blocked")
|
||||
|
||||
route = _console_api_errors(forbidden)
|
||||
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
|
||||
response = app.make_response(route("knowledge-spaces"))
|
||||
|
||||
assert response.status_code == 403
|
||||
assert response.is_json
|
||||
assert response.get_json() == {"code": "forbidden", "message": "blocked", "status": 403}
|
||||
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from inspect import unwrap
|
||||
from unittest.mock import Mock, PropertyMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from pydantic import ValidationError
|
||||
|
||||
from controllers.console import console_ns
|
||||
from controllers.console.onboarding import (
|
||||
StepByStepTourStateApi,
|
||||
StepByStepTourStatePatchPayload,
|
||||
)
|
||||
from extensions.ext_database import db
|
||||
from models.account import Account, AccountStatus
|
||||
from services.step_by_step_tour_service import StepByStepTourService
|
||||
|
||||
|
||||
def _account() -> Account:
|
||||
account = Account(name="User", email="user@example.com", status=AccountStatus.ACTIVE)
|
||||
account.id = "account-1"
|
||||
return account
|
||||
|
||||
|
||||
def _state_response() -> dict[str, object]:
|
||||
return {
|
||||
"first_workspace_id": "workspace-1",
|
||||
"skipped": False,
|
||||
"completed_task_ids": ["home"],
|
||||
"manually_enabled_workspace_ids": [],
|
||||
"manually_disabled_workspace_ids": [],
|
||||
"updated_at": datetime(2026, 6, 28, tzinfo=UTC),
|
||||
}
|
||||
|
||||
|
||||
def test_get_step_by_step_tour_state(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
get_state = Mock(return_value=_state_response())
|
||||
monkeypatch.setattr(StepByStepTourService, "get_state", get_state)
|
||||
|
||||
api = StepByStepTourStateApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
with app.test_request_context("/console/api/onboarding/step-by-step-tour/state", method="GET"):
|
||||
result = method(api, "workspace-1", _account())
|
||||
|
||||
assert result == {
|
||||
"first_workspace_id": "workspace-1",
|
||||
"skipped": False,
|
||||
"completed_task_ids": ["home"],
|
||||
"manually_enabled_workspace_ids": [],
|
||||
"manually_disabled_workspace_ids": [],
|
||||
"updated_at": "2026-06-28T00:00:00Z",
|
||||
}
|
||||
get_state.assert_called_once()
|
||||
assert get_state.call_args.kwargs["current_tenant_id"] == "workspace-1"
|
||||
assert get_state.call_args.kwargs["session"] is db.session
|
||||
|
||||
|
||||
def test_patch_step_by_step_tour_state_passes_action_payload(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
patch_state = Mock(return_value=_state_response())
|
||||
monkeypatch.setattr(StepByStepTourService, "patch_state", patch_state)
|
||||
|
||||
api = StepByStepTourStateApi()
|
||||
method = unwrap(api.patch)
|
||||
payload = {"action": "complete_task", "task_id": "studio"}
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/onboarding/step-by-step-tour/state",
|
||||
method="PATCH",
|
||||
json=payload,
|
||||
):
|
||||
with patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload):
|
||||
result = method(api, "workspace-1", _account())
|
||||
|
||||
assert result["completed_task_ids"] == ["home"]
|
||||
patch_state.assert_called_once()
|
||||
assert patch_state.call_args.kwargs["current_tenant_id"] == "workspace-1"
|
||||
assert patch_state.call_args.kwargs["patch"] == payload
|
||||
assert patch_state.call_args.kwargs["session"] is db.session
|
||||
|
||||
|
||||
def test_patch_payload_rejects_non_action_fields() -> None:
|
||||
with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
|
||||
StepByStepTourStatePatchPayload.model_validate({"action": "skip", "skipped": True})
|
||||
|
||||
|
||||
def test_patch_payload_rejects_task_id_without_task_action() -> None:
|
||||
with pytest.raises(ValidationError, match="task_id is only supported for task actions"):
|
||||
StepByStepTourStatePatchPayload.model_validate({"action": "skip", "task_id": "home"})
|
||||
|
||||
|
||||
def test_patch_payload_requires_action() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
StepByStepTourStatePatchPayload.model_validate({"task_id": "home"})
|
||||
@@ -186,6 +186,39 @@ def test_remote_file_upload_success_when_fetch_falls_back_to_get(app: Flask, mon
|
||||
content=b"fallback-content",
|
||||
mimetype="text/plain",
|
||||
user=current_user,
|
||||
tenant_id=None,
|
||||
source_url=url,
|
||||
)
|
||||
|
||||
|
||||
def test_remote_file_upload_assigns_resource_tenant(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
url = "https://example.com/report.txt"
|
||||
response = _FakeResponse(status_code=200, method="GET", content=b"content")
|
||||
monkeypatch.setattr(remote_files_module.remote_fetcher, "make_request", MagicMock(return_value=response))
|
||||
|
||||
file_service_cls, current_user = _mock_upload_dependencies(monkeypatch)
|
||||
file_service_cls.return_value.upload_file.return_value = SimpleNamespace(
|
||||
id="file-1",
|
||||
name="report.txt",
|
||||
size=7,
|
||||
extension=".txt",
|
||||
mime_type="text/plain",
|
||||
created_by="u1",
|
||||
created_at=datetime(2024, 1, 1, tzinfo=UTC),
|
||||
)
|
||||
|
||||
with app.test_request_context(method="POST", json={"url": url}):
|
||||
remote_files_module.upload_remote_file_from_request(
|
||||
current_user=current_user,
|
||||
resource_tenant_id="app-tenant-id",
|
||||
)
|
||||
|
||||
file_service_cls.return_value.upload_file.assert_called_once_with(
|
||||
filename="report.txt",
|
||||
content=b"content",
|
||||
mimetype="text/plain",
|
||||
user=current_user,
|
||||
tenant_id="app-tenant-id",
|
||||
source_url=url,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,11 +1,77 @@
|
||||
import pytest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
from configs import dify_config
|
||||
from controllers.console import workflow_run_archive
|
||||
from controllers.console.workflow_run_archive import (
|
||||
WorkflowRunArchiveDownloadApi,
|
||||
WorkflowRunArchiveDownloadFileApi,
|
||||
WorkflowRunArchiveDownloadsApi,
|
||||
WorkflowRunArchivesApi,
|
||||
)
|
||||
from models import TenantAccountRole
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
[TenantAccountRole.EDITOR, TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR],
|
||||
)
|
||||
@pytest.mark.parametrize("rbac_enabled", [False, True])
|
||||
def test_current_owner_or_admin_ids_rejects_non_manager(
|
||||
monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole, rbac_enabled: bool
|
||||
) -> None:
|
||||
current_user = SimpleNamespace(id="account-1", current_role=role)
|
||||
monkeypatch.setattr(dify_config, "RBAC_ENABLED", rbac_enabled)
|
||||
monkeypatch.setattr(
|
||||
workflow_run_archive,
|
||||
"current_account_with_tenant",
|
||||
lambda: (current_user, "tenant-1"),
|
||||
)
|
||||
|
||||
with pytest.raises(Forbidden):
|
||||
workflow_run_archive._current_owner_or_admin_ids()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("role", [TenantAccountRole.OWNER, TenantAccountRole.ADMIN])
|
||||
@pytest.mark.parametrize("rbac_enabled", [False, True])
|
||||
def test_current_owner_or_admin_ids_returns_current_ids_for_manager(
|
||||
monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole, rbac_enabled: bool
|
||||
) -> None:
|
||||
current_user = SimpleNamespace(id="account-1", current_role=role)
|
||||
monkeypatch.setattr(dify_config, "RBAC_ENABLED", rbac_enabled)
|
||||
monkeypatch.setattr(
|
||||
workflow_run_archive,
|
||||
"current_account_with_tenant",
|
||||
lambda: (current_user, "tenant-1"),
|
||||
)
|
||||
|
||||
assert workflow_run_archive._current_owner_or_admin_ids() == ("tenant-1", "account-1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "args"),
|
||||
[
|
||||
(WorkflowRunArchivesApi.get, ()),
|
||||
(WorkflowRunArchiveDownloadsApi.post, ()),
|
||||
(WorkflowRunArchiveDownloadApi.get, ("download-1",)),
|
||||
(WorkflowRunArchiveDownloadFileApi.get, ("download-1",)),
|
||||
],
|
||||
)
|
||||
def test_workflow_run_archive_endpoints_enforce_fixed_workspace_roles(
|
||||
monkeypatch: pytest.MonkeyPatch, method, args: tuple[str, ...]
|
||||
) -> None:
|
||||
while hasattr(method, "__wrapped__"):
|
||||
method = method.__wrapped__
|
||||
|
||||
def reject_non_manager() -> tuple[str, str]:
|
||||
raise Forbidden()
|
||||
|
||||
monkeypatch.setattr(workflow_run_archive, "_current_owner_or_admin_ids", reject_non_manager)
|
||||
|
||||
with pytest.raises(Forbidden):
|
||||
method(None, *args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -28,3 +94,4 @@ def test_workflow_run_archive_endpoints_require_cloud_paid_plan(method) -> None:
|
||||
"cloud_edition_billing_enabled",
|
||||
"cloud_edition_billing_paid_plan_required",
|
||||
} <= decorator_names
|
||||
assert "rbac_permission_required" not in decorator_names
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
import inspect
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
||||
|
||||
from controllers.console.workspace.account import (
|
||||
AccountDeleteUpdateFeedbackApi,
|
||||
@@ -12,7 +15,8 @@ from controllers.console.workspace.account import (
|
||||
ChangeEmailSendEmailApi,
|
||||
CheckEmailUnique,
|
||||
)
|
||||
from models import Account, AccountStatus, Tenant
|
||||
from models import Account, AccountIntegrate, AccountStatus, Tenant, TenantAccountJoin
|
||||
from models.account import TenantAccountRole
|
||||
from services.account_service import AccountService
|
||||
from services.entities.auth_entities import (
|
||||
ChangeEmailNewEmailToken,
|
||||
@@ -45,6 +49,38 @@ def _build_account(email: str, account_id: str = "acc", tenant: Tenant | None =
|
||||
return account
|
||||
|
||||
|
||||
def _stable_uuid(value: str) -> str:
|
||||
return str(uuid5(NAMESPACE_URL, value))
|
||||
|
||||
|
||||
def _persist_account_with_tenant(session: Session, email: str, account_name: str = "account") -> tuple[Account, Tenant]:
|
||||
tenant = Tenant(name=f"{account_name} tenant")
|
||||
tenant.id = _stable_uuid(f"tenant:{account_name}")
|
||||
account = Account(name=account_name, email=email, status=AccountStatus.ACTIVE)
|
||||
account.id = _stable_uuid(f"account:{account_name}")
|
||||
membership = TenantAccountJoin(
|
||||
tenant_id=tenant.id,
|
||||
account_id=account.id,
|
||||
current=True,
|
||||
role=TenantAccountRole.OWNER,
|
||||
)
|
||||
session.add_all([account, tenant, membership])
|
||||
session.commit()
|
||||
account._current_tenant = tenant
|
||||
account.role = TenantAccountRole.OWNER
|
||||
return account, tenant
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _bind_database_session(session: Session):
|
||||
database_session = scoped_session(sessionmaker(bind=session.get_bind(), expire_on_commit=False))
|
||||
try:
|
||||
with patch("extensions.ext_database.db.session", database_session):
|
||||
yield database_session
|
||||
finally:
|
||||
database_session.remove()
|
||||
|
||||
|
||||
def _build_change_email_token(
|
||||
phase: str,
|
||||
*,
|
||||
@@ -387,47 +423,64 @@ class TestChangeEmailValidity:
|
||||
|
||||
class TestChangeEmailReset:
|
||||
@patch("controllers.console.workspace.account.AccountService.send_change_email_completed_notify_email")
|
||||
@patch("controllers.console.workspace.account.AccountService.update_account_email")
|
||||
@patch("controllers.console.workspace.account.AccountService.revoke_change_email_token")
|
||||
@patch("controllers.console.workspace.account.AccountService.get_change_email_data")
|
||||
@patch("controllers.console.workspace.account.AccountService.check_email_unique")
|
||||
@patch("controllers.console.workspace.account.AccountService.is_account_in_freeze")
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, AccountIntegrate)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_should_normalize_new_email_before_update(
|
||||
self,
|
||||
mock_is_freeze: MagicMock,
|
||||
mock_check_unique: MagicMock,
|
||||
mock_get_data: MagicMock,
|
||||
mock_revoke_token: MagicMock,
|
||||
mock_update_account: MagicMock,
|
||||
mock_send_notify: MagicMock,
|
||||
app: Flask,
|
||||
sqlite_session: Session,
|
||||
):
|
||||
current_user = _build_account("old@example.com", "acc3")
|
||||
mock_is_freeze.return_value = False
|
||||
mock_check_unique.return_value = True
|
||||
mock_get_data.return_value = _build_change_email_token(
|
||||
AccountService.CHANGE_EMAIL_PHASE_NEW_VERIFIED,
|
||||
account_id="acc3",
|
||||
email="new@example.com",
|
||||
old_email="OLD@example.com",
|
||||
)
|
||||
mock_account_after_update = _build_account("new@example.com", "acc3-updated")
|
||||
mock_update_account.return_value = mock_account_after_update
|
||||
|
||||
with app.test_request_context(
|
||||
"/account/change-email/reset",
|
||||
method="POST",
|
||||
json={"new_email": "New@Example.com", "token": "token-123"},
|
||||
):
|
||||
api = ChangeEmailResetApi()
|
||||
method = inspect.unwrap(api.post)
|
||||
method(api, current_user)
|
||||
with _bind_database_session(sqlite_session) as database_session:
|
||||
current_user, _ = _persist_account_with_tenant(
|
||||
database_session(),
|
||||
"old@example.com",
|
||||
"email-reset-account",
|
||||
)
|
||||
account_integration = AccountIntegrate(
|
||||
account_id=current_user.id,
|
||||
provider="google",
|
||||
open_id="google-user",
|
||||
encrypted_token="encrypted-token",
|
||||
)
|
||||
database_session.add(account_integration)
|
||||
database_session.commit()
|
||||
mock_get_data.return_value = _build_change_email_token(
|
||||
AccountService.CHANGE_EMAIL_PHASE_NEW_VERIFIED,
|
||||
account_id=current_user.id,
|
||||
email="new@example.com",
|
||||
old_email="OLD@example.com",
|
||||
)
|
||||
|
||||
mock_is_freeze.assert_called_once_with("new@example.com")
|
||||
mock_check_unique.assert_called_once_with("new@example.com", session=ANY)
|
||||
mock_revoke_token.assert_called_once_with("token-123")
|
||||
mock_update_account.assert_called_once_with(current_user, email="new@example.com", session=ANY)
|
||||
mock_send_notify.assert_called_once_with(email="new@example.com")
|
||||
with app.test_request_context(
|
||||
"/account/change-email/reset",
|
||||
method="POST",
|
||||
json={"new_email": "New@Example.com", "token": "token-123"},
|
||||
):
|
||||
api = ChangeEmailResetApi()
|
||||
method = inspect.unwrap(api.post)
|
||||
response = method(api, current_user)
|
||||
|
||||
sqlite_session.expire_all()
|
||||
persisted_account = sqlite_session.get(Account, current_user.id)
|
||||
assert response["email"] == "new@example.com"
|
||||
assert persisted_account is not None
|
||||
assert persisted_account.email == "new@example.com"
|
||||
assert sqlite_session.get(AccountIntegrate, account_integration.id) is None
|
||||
mock_is_freeze.assert_called_once_with("new@example.com")
|
||||
mock_revoke_token.assert_called_once_with("token-123")
|
||||
mock_send_notify.assert_called_once_with(email="new@example.com")
|
||||
|
||||
@patch("controllers.console.workspace.account.AccountService.send_change_email_completed_notify_email")
|
||||
@patch("controllers.console.workspace.account.AccountService.update_account_email")
|
||||
@@ -663,36 +716,47 @@ class TestAccountDeletionFeedback:
|
||||
|
||||
|
||||
class TestCheckEmailUnique:
|
||||
@patch("controllers.console.workspace.account.AccountService.check_email_unique")
|
||||
@patch("controllers.console.workspace.account.AccountService.is_account_in_freeze")
|
||||
def test_should_normalize_email(self, mock_is_freeze: MagicMock, mock_check_unique: MagicMock, app: Flask):
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_should_normalize_email(
|
||||
self,
|
||||
mock_is_freeze: MagicMock,
|
||||
app: Flask,
|
||||
sqlite_session: Session,
|
||||
):
|
||||
mock_is_freeze.return_value = False
|
||||
mock_check_unique.return_value = True
|
||||
|
||||
with app.test_request_context(
|
||||
"/account/change-email/check-email-unique",
|
||||
method="POST",
|
||||
json={"email": "Case@Test.com"},
|
||||
):
|
||||
api = CheckEmailUnique()
|
||||
method = inspect.unwrap(api.post)
|
||||
response = method(api)
|
||||
with _bind_database_session(sqlite_session) as database_session:
|
||||
_persist_account_with_tenant(database_session(), "different@test.com", "uniqueness-account")
|
||||
with app.test_request_context(
|
||||
"/account/change-email/check-email-unique",
|
||||
method="POST",
|
||||
json={"email": "Case@Test.com"},
|
||||
):
|
||||
api = CheckEmailUnique()
|
||||
method = inspect.unwrap(api.post)
|
||||
response = method(api)
|
||||
|
||||
assert response == {"result": "success"}
|
||||
mock_is_freeze.assert_called_once_with("case@test.com")
|
||||
mock_check_unique.assert_called_once_with("case@test.com", session=ANY)
|
||||
|
||||
|
||||
def test_get_account_by_email_with_case_fallback_uses_lowercase_lookup():
|
||||
mock_session = MagicMock()
|
||||
first = MagicMock()
|
||||
first.scalar_one_or_none.return_value = None
|
||||
second = MagicMock()
|
||||
expected_account = MagicMock()
|
||||
second.scalar_one_or_none.return_value = expected_account
|
||||
mock_session.execute.side_effect = [first, second]
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_get_account_by_email_with_case_fallback_uses_lowercase_lookup(sqlite_session: Session):
|
||||
expected_account, _ = _persist_account_with_tenant(
|
||||
sqlite_session,
|
||||
"mixed@test.com",
|
||||
"case-fallback-account",
|
||||
)
|
||||
|
||||
result = AccountService.get_account_by_email_with_case_fallback("Mixed@Test.com", session=mock_session)
|
||||
result = AccountService.get_account_by_email_with_case_fallback("Mixed@Test.com", session=sqlite_session)
|
||||
|
||||
assert result is expected_account
|
||||
assert mock_session.execute.call_count == 2
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from types import SimpleNamespace
|
||||
from typing import override
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -5,6 +6,8 @@ import pytest
|
||||
from flask import Flask, request
|
||||
from flask_login import LoginManager, UserMixin
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
from werkzeug.exceptions import HTTPException
|
||||
|
||||
from controllers.common.wraps import _extract_resource_id
|
||||
@@ -34,6 +37,7 @@ from controllers.console.wraps import (
|
||||
)
|
||||
from models import Account
|
||||
from models.account import AccountStatus, TenantAccountRole
|
||||
from models.dataset import RateLimitLog
|
||||
from services.feature_service import LicenseStatus
|
||||
|
||||
|
||||
@@ -657,8 +661,7 @@ class TestRateLimiting:
|
||||
"""Test rate limiting decorator"""
|
||||
|
||||
@patch("controllers.console.wraps.redis_client")
|
||||
@patch("controllers.console.wraps.db")
|
||||
def test_should_allow_requests_within_rate_limit(self, mock_db: MagicMock, mock_redis: MagicMock):
|
||||
def test_should_allow_requests_within_rate_limit(self, mock_redis: MagicMock):
|
||||
"""Test that requests within rate limit are allowed"""
|
||||
# Arrange
|
||||
mock_rate_limit = MagicMock()
|
||||
@@ -685,8 +688,13 @@ class TestRateLimiting:
|
||||
mock_redis.zremrangebyscore.assert_called_once()
|
||||
|
||||
@patch("controllers.console.wraps.redis_client")
|
||||
@patch("controllers.console.wraps.db")
|
||||
def test_should_reject_requests_over_rate_limit(self, mock_db: MagicMock, mock_redis: MagicMock):
|
||||
@pytest.mark.parametrize("sqlite_session", [(RateLimitLog,)], indirect=True)
|
||||
def test_should_reject_requests_over_rate_limit(
|
||||
self,
|
||||
mock_redis: MagicMock,
|
||||
sqlite_session: Session,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Test that requests over rate limit are rejected and logged"""
|
||||
# Arrange
|
||||
app = create_app_with_login()
|
||||
@@ -696,8 +704,7 @@ class TestRateLimiting:
|
||||
mock_rate_limit.subscription_plan = "pro"
|
||||
mock_redis.zcard.return_value = 11 # Over limit
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_db.session = mock_session
|
||||
monkeypatch.setattr("controllers.console.wraps.db", SimpleNamespace(session=sqlite_session))
|
||||
|
||||
@cloud_edition_billing_rate_limit_check("knowledge")
|
||||
def knowledge_request():
|
||||
@@ -719,9 +726,11 @@ class TestRateLimiting:
|
||||
assert exc_info.value.code == 403
|
||||
assert "rate limit" in str(exc_info.value.description)
|
||||
|
||||
# Verify rate limit log was created
|
||||
mock_session.add.assert_called_once()
|
||||
mock_session.commit.assert_called_once()
|
||||
rate_limit_log = sqlite_session.scalar(select(RateLimitLog))
|
||||
assert rate_limit_log is not None
|
||||
assert rate_limit_log.tenant_id == "tenant123"
|
||||
assert rate_limit_log.subscription_plan == "pro"
|
||||
assert rate_limit_log.operation == "knowledge"
|
||||
|
||||
|
||||
class TestCloudUtmRecord:
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
import inspect
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import MagicMock, PropertyMock, patch
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy.orm import Session
|
||||
from werkzeug.exceptions import NotFound
|
||||
|
||||
from controllers.console import console_ns
|
||||
@@ -32,15 +35,17 @@ from controllers.console.workspace.error import (
|
||||
CurrentPasswordIncorrectError,
|
||||
InvalidAccountDeletionCodeError,
|
||||
)
|
||||
from models import Account
|
||||
from models.account import AccountStatus
|
||||
from extensions.storage.storage_type import StorageType
|
||||
from models import Account, AccountIntegrate, InvitationCode, Tenant, TenantAccountJoin
|
||||
from models.account import AccountStatus, InvitationCodeStatus, TenantAccountRole
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import UploadFile
|
||||
from services.errors.account import CurrentPasswordIncorrectError as ServicePwdError
|
||||
|
||||
|
||||
def make_account(account_id: str = "u1", *, status: AccountStatus = AccountStatus.ACTIVE) -> Account:
|
||||
account = Account(name="John", email=f"{account_id}@test.com", status=status)
|
||||
account.id = account_id
|
||||
account.id = str(uuid5(NAMESPACE_URL, f"account:{account_id}"))
|
||||
account.avatar = "avatar.png"
|
||||
account.interface_language = "en-US"
|
||||
account.interface_theme = "light"
|
||||
@@ -49,12 +54,62 @@ def make_account(account_id: str = "u1", *, status: AccountStatus = AccountStatu
|
||||
return account
|
||||
|
||||
|
||||
def persist_account_with_tenant(
|
||||
session: Session,
|
||||
account_id: str = "u1",
|
||||
*,
|
||||
status: AccountStatus = AccountStatus.ACTIVE,
|
||||
tenant_id: str = "tenant-1",
|
||||
) -> tuple[Account, Tenant]:
|
||||
account = make_account(account_id, status=status)
|
||||
tenant = Tenant(name=tenant_id)
|
||||
tenant.id = str(uuid5(NAMESPACE_URL, f"tenant:{tenant_id}"))
|
||||
membership = TenantAccountJoin(
|
||||
tenant_id=tenant.id,
|
||||
account_id=account.id,
|
||||
current=True,
|
||||
role=TenantAccountRole.OWNER,
|
||||
)
|
||||
session.add_all([account, tenant, membership])
|
||||
session.commit()
|
||||
account._current_tenant = tenant
|
||||
account.role = TenantAccountRole.OWNER
|
||||
return account, tenant
|
||||
|
||||
|
||||
def make_upload_file(*, tenant_id: str, created_by: str) -> UploadFile:
|
||||
return UploadFile(
|
||||
tenant_id=tenant_id,
|
||||
storage_type=StorageType.LOCAL,
|
||||
key="avatar.png",
|
||||
name="avatar.png",
|
||||
size=128,
|
||||
extension="png",
|
||||
mime_type="image/png",
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=created_by,
|
||||
created_at=datetime.now(UTC).replace(tzinfo=None),
|
||||
used=False,
|
||||
)
|
||||
|
||||
|
||||
class TestAccountInitApi:
|
||||
def test_init_success(self, app: Flask):
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, InvitationCode)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_init_success(self, app: Flask, sqlite_session: Session):
|
||||
api = AccountInitApi()
|
||||
method = inspect.unwrap(api.post)
|
||||
|
||||
account = make_account(status=AccountStatus.UNINITIALIZED)
|
||||
account, tenant = persist_account_with_tenant(
|
||||
sqlite_session,
|
||||
status=AccountStatus.UNINITIALIZED,
|
||||
)
|
||||
invitation_code = InvitationCode(batch="batch-1", code="code123")
|
||||
sqlite_session.add(invitation_code)
|
||||
sqlite_session.commit()
|
||||
payload = {
|
||||
"interface_language": "en-US",
|
||||
"timezone": "UTC",
|
||||
@@ -63,14 +118,24 @@ class TestAccountInitApi:
|
||||
|
||||
with (
|
||||
app.test_request_context("/account/init", json=payload),
|
||||
patch("controllers.console.workspace.account.db.session.commit", return_value=None),
|
||||
patch("controllers.console.workspace.account.dify_config.EDITION", "CLOUD"),
|
||||
patch("controllers.console.workspace.account.db.session.scalar") as scalar_mock,
|
||||
patch("controllers.console.workspace.account.db.session", sqlite_session),
|
||||
):
|
||||
scalar_mock.return_value = MagicMock(status="unused")
|
||||
resp = method(api, account)
|
||||
|
||||
assert resp["result"] == "success"
|
||||
sqlite_session.expire_all()
|
||||
persisted_account = sqlite_session.get(Account, account.id)
|
||||
persisted_invitation = sqlite_session.get(InvitationCode, invitation_code.id)
|
||||
assert persisted_account is not None
|
||||
assert persisted_account.status == AccountStatus.ACTIVE
|
||||
assert persisted_account.interface_language == "en-US"
|
||||
assert persisted_account.timezone == "UTC"
|
||||
assert persisted_account.initialized_at is not None
|
||||
assert persisted_invitation is not None
|
||||
assert persisted_invitation.status == InvitationCodeStatus.USED
|
||||
assert persisted_invitation.used_by_account_id == account.id
|
||||
assert persisted_invitation.used_by_tenant_id == tenant.id
|
||||
|
||||
def test_init_already_initialized(self, app: Flask):
|
||||
api = AccountInitApi()
|
||||
@@ -93,7 +158,7 @@ class TestAccountProfileApi:
|
||||
with app.test_request_context("/account/profile"):
|
||||
result = method(api, user)
|
||||
|
||||
assert result["id"] == "u1"
|
||||
assert result["id"] == user.id
|
||||
|
||||
|
||||
class TestAccountUpdateApis:
|
||||
@@ -119,29 +184,33 @@ class TestAccountUpdateApis:
|
||||
):
|
||||
result = method(api, user)
|
||||
|
||||
assert result["id"] == "u1"
|
||||
assert result["id"] == user.id
|
||||
|
||||
|
||||
class TestAccountAvatarApiGet:
|
||||
"""GET /account/avatar must not sign arbitrary upload_file IDs (IDOR)."""
|
||||
|
||||
def test_get_avatar_signed_url_when_upload_owned_by_current_account(self, app: Flask):
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, UploadFile)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_get_avatar_signed_url_when_upload_owned_by_current_account(self, app: Flask, sqlite_session: Session):
|
||||
api = AccountAvatarApi()
|
||||
method = inspect.unwrap(api.get)
|
||||
|
||||
user = make_account("acc-owner")
|
||||
tenant_id = "tenant-1"
|
||||
user, tenant = persist_account_with_tenant(sqlite_session, "acc-owner")
|
||||
tenant_id = tenant.id
|
||||
file_id = "550e8400-e29b-41d4-a716-446655440000"
|
||||
|
||||
upload_file = MagicMock()
|
||||
upload_file = make_upload_file(tenant_id=tenant_id, created_by=user.id)
|
||||
upload_file.id = file_id
|
||||
upload_file.tenant_id = tenant_id
|
||||
upload_file.created_by = user.id
|
||||
upload_file.created_by_role = CreatorUserRole.ACCOUNT
|
||||
sqlite_session.add(upload_file)
|
||||
sqlite_session.commit()
|
||||
|
||||
with (
|
||||
app.test_request_context(f"/account/avatar?avatar={file_id}"),
|
||||
patch("controllers.console.workspace.account.db.session.scalar", return_value=upload_file),
|
||||
patch("controllers.console.workspace.account.db.session", sqlite_session),
|
||||
patch(
|
||||
"controllers.console.workspace.account.file_helpers.get_signed_file_url",
|
||||
return_value="https://signed/example",
|
||||
@@ -152,23 +221,35 @@ class TestAccountAvatarApiGet:
|
||||
assert result == {"avatar_url": "https://signed/example"}
|
||||
sign_mock.assert_called_once_with(upload_file_id=file_id)
|
||||
|
||||
def test_get_avatar_not_found_when_upload_created_by_other_account_same_tenant(self, app: Flask):
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, UploadFile)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_get_avatar_not_found_when_upload_created_by_other_account_same_tenant(
|
||||
self, app: Flask, sqlite_session: Session
|
||||
):
|
||||
api = AccountAvatarApi()
|
||||
method = inspect.unwrap(api.get)
|
||||
|
||||
user = make_account("acc-a")
|
||||
tenant_id = "tenant-1"
|
||||
user, tenant = persist_account_with_tenant(sqlite_session, "acc-a")
|
||||
tenant_id = tenant.id
|
||||
file_id = "550e8400-e29b-41d4-a716-446655440001"
|
||||
|
||||
upload_file = MagicMock()
|
||||
other_account = make_account("acc-b")
|
||||
other_membership = TenantAccountJoin(
|
||||
tenant_id=tenant_id,
|
||||
account_id=other_account.id,
|
||||
role=TenantAccountRole.NORMAL,
|
||||
)
|
||||
upload_file = make_upload_file(tenant_id=tenant_id, created_by=other_account.id)
|
||||
upload_file.id = file_id
|
||||
upload_file.tenant_id = tenant_id
|
||||
upload_file.created_by = "acc-b"
|
||||
upload_file.created_by_role = CreatorUserRole.ACCOUNT
|
||||
sqlite_session.add_all([other_account, other_membership, upload_file])
|
||||
sqlite_session.commit()
|
||||
|
||||
with (
|
||||
app.test_request_context(f"/account/avatar?avatar={file_id}"),
|
||||
patch("controllers.console.workspace.account.db.session.scalar", return_value=upload_file),
|
||||
patch("controllers.console.workspace.account.db.session", sqlite_session),
|
||||
patch(
|
||||
"controllers.console.workspace.account.file_helpers.get_signed_file_url",
|
||||
return_value="https://signed/leak",
|
||||
@@ -179,23 +260,29 @@ class TestAccountAvatarApiGet:
|
||||
|
||||
sign_mock.assert_not_called()
|
||||
|
||||
def test_get_avatar_not_found_when_upload_belongs_to_other_tenant(self, app: Flask):
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, UploadFile)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_get_avatar_not_found_when_upload_belongs_to_other_tenant(self, app: Flask, sqlite_session: Session):
|
||||
api = AccountAvatarApi()
|
||||
method = inspect.unwrap(api.get)
|
||||
|
||||
user = make_account("acc-owner")
|
||||
tenant_id = "tenant-1"
|
||||
user, tenant = persist_account_with_tenant(sqlite_session, "acc-owner")
|
||||
tenant_id = tenant.id
|
||||
file_id = "550e8400-e29b-41d4-a716-446655440002"
|
||||
|
||||
upload_file = MagicMock()
|
||||
other_tenant = Tenant(name="tenant-other")
|
||||
other_tenant.id = str(uuid5(NAMESPACE_URL, "tenant:tenant-other"))
|
||||
upload_file = make_upload_file(tenant_id=other_tenant.id, created_by=user.id)
|
||||
upload_file.id = file_id
|
||||
upload_file.tenant_id = "tenant-other"
|
||||
upload_file.created_by = user.id
|
||||
upload_file.created_by_role = CreatorUserRole.ACCOUNT
|
||||
sqlite_session.add_all([other_tenant, upload_file])
|
||||
sqlite_session.commit()
|
||||
|
||||
with (
|
||||
app.test_request_context(f"/account/avatar?avatar={file_id}"),
|
||||
patch("controllers.console.workspace.account.db.session.scalar", return_value=upload_file),
|
||||
patch("controllers.console.workspace.account.db.session", sqlite_session),
|
||||
patch(
|
||||
"controllers.console.workspace.account.file_helpers.get_signed_file_url",
|
||||
return_value="https://signed/leak",
|
||||
@@ -246,7 +333,7 @@ class TestAccountPasswordApi:
|
||||
):
|
||||
result = method(api, user)
|
||||
|
||||
assert result["id"] == "u1"
|
||||
assert result["id"] == user.id
|
||||
|
||||
def test_password_wrong_current(self, app: Flask):
|
||||
api = AccountPasswordApi()
|
||||
@@ -271,21 +358,38 @@ class TestAccountPasswordApi:
|
||||
|
||||
|
||||
class TestAccountIntegrateApi:
|
||||
def test_get_integrates(self, app: Flask):
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, AccountIntegrate)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_get_integrates(self, app: Flask, sqlite_session: Session):
|
||||
api = AccountIntegrateApi()
|
||||
method = inspect.unwrap(api.get)
|
||||
|
||||
account = make_account("acc1")
|
||||
account, _ = persist_account_with_tenant(sqlite_session, "acc1")
|
||||
sqlite_session.add(
|
||||
AccountIntegrate(
|
||||
account_id=account.id,
|
||||
provider="github",
|
||||
open_id="github-user",
|
||||
encrypted_token="encrypted-token",
|
||||
)
|
||||
)
|
||||
sqlite_session.commit()
|
||||
|
||||
with (
|
||||
app.test_request_context("/"),
|
||||
patch("controllers.console.workspace.account.db.session.scalars") as scalars_mock,
|
||||
patch("controllers.console.workspace.account.db.session", sqlite_session),
|
||||
):
|
||||
scalars_mock.return_value.all.return_value = []
|
||||
result = method(api, account)
|
||||
|
||||
assert "data" in result
|
||||
assert len(result["data"]) == 2
|
||||
assert result["data"][0]["provider"] == "github"
|
||||
assert result["data"][0]["is_bound"] is True
|
||||
assert result["data"][0]["link"] is None
|
||||
assert result["data"][1]["provider"] == "google"
|
||||
assert result["data"][1]["is_bound"] is False
|
||||
assert result["data"][1]["link"].endswith("/console/api/oauth/login/google")
|
||||
|
||||
|
||||
class TestAccountDeleteApi:
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import logging
|
||||
from collections.abc import Iterator
|
||||
from http import HTTPStatus
|
||||
from inspect import unwrap
|
||||
from io import BytesIO
|
||||
@@ -6,6 +7,8 @@ from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import Engine, event
|
||||
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
||||
from werkzeug.datastructures import FileStorage
|
||||
from werkzeug.exceptions import Unauthorized
|
||||
|
||||
@@ -36,6 +39,17 @@ from libs.datetime_utils import naive_utc_now
|
||||
from models.account import Account, Tenant, TenantCustomConfigDict, TenantStatus
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def workspace_session(sqlite_engine: Engine) -> Iterator[scoped_session[Session]]:
|
||||
"""Provide the callable scoped session expected by Flask-SQLAlchemy controllers."""
|
||||
Tenant.metadata.create_all(sqlite_engine, tables=[Tenant.__table__])
|
||||
session = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False))
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
session.remove()
|
||||
|
||||
|
||||
def make_account(account_id: str = "u1") -> Account:
|
||||
account = Account(name="Test User", email=f"{account_id}@example.com")
|
||||
account.id = account_id
|
||||
@@ -370,11 +384,13 @@ class TestTenantInfoResponse:
|
||||
|
||||
|
||||
class TestSwitchWorkspaceApi:
|
||||
def test_switch_success(self, app: Flask):
|
||||
def test_switch_success(self, app: Flask, workspace_session: scoped_session[Session]):
|
||||
api = SwitchWorkspaceApi()
|
||||
method = unwrap(api.post)
|
||||
payload = {"tenant_id": "t2"}
|
||||
tenant = make_tenant("t2")
|
||||
workspace_session.add(tenant)
|
||||
workspace_session.commit()
|
||||
user = make_account()
|
||||
with (
|
||||
app.test_request_context("/workspaces/switch", json=payload),
|
||||
@@ -383,11 +399,10 @@ class TestSwitchWorkspaceApi:
|
||||
"controllers.console.workspace.workspace.WorkspaceService.get_tenant_info", return_value={"id": "t2"}
|
||||
),
|
||||
):
|
||||
session = MagicMock()
|
||||
session.get.return_value = tenant
|
||||
result = method(api, session, user)
|
||||
result = method(api, workspace_session, user)
|
||||
|
||||
assert result["result"] == "success"
|
||||
switch_tenant.assert_called_once_with(user, "t2", session=session)
|
||||
switch_tenant.assert_called_once_with(user, "t2", session=workspace_session)
|
||||
|
||||
def test_switch_not_linked(self, app: Flask):
|
||||
api = SwitchWorkspaceApi()
|
||||
@@ -401,7 +416,7 @@ class TestSwitchWorkspaceApi:
|
||||
with pytest.raises(AccountNotLinkTenantError):
|
||||
method(api, MagicMock(), user)
|
||||
|
||||
def test_switch_tenant_not_found(self, app: Flask):
|
||||
def test_switch_tenant_not_found(self, app: Flask, workspace_session: scoped_session[Session]):
|
||||
api = SwitchWorkspaceApi()
|
||||
method = unwrap(api.post)
|
||||
payload = {"tenant_id": "missing"}
|
||||
@@ -410,19 +425,21 @@ class TestSwitchWorkspaceApi:
|
||||
app.test_request_context("/workspaces/switch", json=payload),
|
||||
patch("controllers.console.workspace.workspace.TenantService.switch_tenant"),
|
||||
):
|
||||
session = MagicMock()
|
||||
session.get.return_value = None
|
||||
with pytest.raises(ValueError):
|
||||
method(api, session, user)
|
||||
method(api, workspace_session, user)
|
||||
|
||||
|
||||
class TestCustomConfigWorkspaceApi:
|
||||
def test_post_success(self, app: Flask):
|
||||
def test_post_success(self, app: Flask, workspace_session: scoped_session[Session]):
|
||||
api = CustomConfigWorkspaceApi()
|
||||
method = unwrap(api.post)
|
||||
tenant = make_tenant(custom_config={})
|
||||
workspace_session.add(tenant)
|
||||
workspace_session.commit()
|
||||
|
||||
payload = {"remove_webapp_brand": True}
|
||||
events = []
|
||||
event.listen(workspace_session, "after_commit", lambda _: events.append("commit"))
|
||||
with (
|
||||
app.test_request_context("/workspaces/custom-config", json=payload),
|
||||
patch(
|
||||
@@ -430,27 +447,29 @@ class TestCustomConfigWorkspaceApi:
|
||||
side_effect=lambda *args, **kwargs: events.append("get_tenant_info") or {"id": "t1"},
|
||||
),
|
||||
):
|
||||
session = MagicMock()
|
||||
session.get.return_value = tenant
|
||||
session.commit.side_effect = lambda: events.append("commit")
|
||||
result = method(api, session, "t1")
|
||||
result = method(api, workspace_session, "t1")
|
||||
assert result["result"] == "success"
|
||||
assert events == ["commit", "get_tenant_info"]
|
||||
|
||||
def test_logo_fallback(self, app: Flask):
|
||||
def test_logo_fallback(self, app: Flask, workspace_session: scoped_session[Session]):
|
||||
api = CustomConfigWorkspaceApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
tenant = make_tenant(custom_config={"replace_webapp_logo": "old-logo"})
|
||||
workspace_session.add(tenant)
|
||||
workspace_session.commit()
|
||||
|
||||
payload = {"remove_webapp_brand": False}
|
||||
|
||||
with (
|
||||
app.test_request_context("/workspaces/custom-config", json=payload),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.WorkspaceService.get_tenant_info", return_value={"id": "t1"}
|
||||
"controllers.console.workspace.workspace.WorkspaceService.get_tenant_info",
|
||||
return_value={"id": "t1"},
|
||||
),
|
||||
):
|
||||
session = MagicMock()
|
||||
session.get.return_value = tenant
|
||||
result = method(api, session, "t1")
|
||||
result = method(api, workspace_session, "t1")
|
||||
|
||||
assert tenant.custom_config_dict["replace_webapp_logo"] == "old-logo"
|
||||
assert result["result"] == "success"
|
||||
|
||||
@@ -541,14 +560,19 @@ class TestWebappLogoWorkspaceApi:
|
||||
|
||||
|
||||
class TestWorkspaceInfoApi:
|
||||
def test_post_success(self, app: Flask):
|
||||
def test_post_success(self, app: Flask, workspace_session: scoped_session[Session]):
|
||||
api = WorkspaceInfoApi()
|
||||
method = unwrap(api.post)
|
||||
tenant = make_tenant()
|
||||
workspace_session.add(tenant)
|
||||
workspace_session.commit()
|
||||
|
||||
payload = {"name": "New Name"}
|
||||
events = []
|
||||
with (
|
||||
app.test_request_context("/workspaces/info", json=payload),
|
||||
patch("controllers.console.workspace.workspace.db.get_or_404", return_value=tenant),
|
||||
patch("controllers.console.workspace.workspace.db.session", workspace_session),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.WorkspaceService.get_tenant_info",
|
||||
side_effect=lambda *args, **kwargs: (
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.agent.cot_agent_runner import CotAgentRunner
|
||||
from core.agent.entities import AgentScratchpadUnit
|
||||
@@ -26,7 +29,7 @@ class DummyRunner(CotAgentRunner):
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner(mocker: MockerFixture):
|
||||
def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[DummyRunner]:
|
||||
# Prevent BaseAgentRunner __init__ from hitting database
|
||||
mocker.patch(
|
||||
"core.agent.base_agent_runner.BaseAgentRunner.organize_agent_history",
|
||||
@@ -81,9 +84,12 @@ def runner(mocker: MockerFixture):
|
||||
runner.agent_callback = None
|
||||
runner.memory = None
|
||||
runner.history_prompt_messages = []
|
||||
runner.session = MagicMock()
|
||||
runner.session = Session(sqlite_engine)
|
||||
|
||||
return runner
|
||||
try:
|
||||
yield runner
|
||||
finally:
|
||||
runner.session.close()
|
||||
|
||||
|
||||
class TestFillInputs:
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.agent.errors import AgentMaxIterationError
|
||||
from core.agent.fc_agent_runner import FunctionCallAgentRunner
|
||||
@@ -69,7 +72,7 @@ class DummyResult:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner(mocker: MockerFixture):
|
||||
def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCallAgentRunner]:
|
||||
# Completely bypass BaseAgentRunner __init__ to avoid DB / Flask context
|
||||
mocker.patch(
|
||||
"core.agent.base_agent_runner.BaseAgentRunner.__init__",
|
||||
@@ -131,7 +134,7 @@ def runner(mocker: MockerFixture):
|
||||
runner._current_thoughts = []
|
||||
runner.files = []
|
||||
runner.agent_callback = MagicMock()
|
||||
runner.session = MagicMock()
|
||||
runner.session = Session(sqlite_engine)
|
||||
|
||||
runner._init_prompt_tools = MagicMock(return_value=({}, []))
|
||||
runner.create_agent_thought = MagicMock(return_value="thought1")
|
||||
@@ -139,7 +142,10 @@ def runner(mocker: MockerFixture):
|
||||
runner.recalc_llm_max_tokens = MagicMock()
|
||||
runner.update_prompt_message_tool = MagicMock()
|
||||
|
||||
return runner
|
||||
try:
|
||||
yield runner
|
||||
finally:
|
||||
runner.session.close()
|
||||
|
||||
|
||||
# ==============================
|
||||
|
||||
@@ -6,6 +6,7 @@ from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Iterator
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, override
|
||||
from unittest.mock import MagicMock
|
||||
@@ -137,7 +138,19 @@ class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
data=RunSucceededEventData(
|
||||
output={"text": "hello agent"},
|
||||
session_snapshot=CompositorSessionSnapshot(layers=[]),
|
||||
usage=AgentRunUsage(prompt_tokens=3, completion_tokens=5),
|
||||
usage=AgentRunUsage(
|
||||
prompt_tokens=3,
|
||||
prompt_unit_price=Decimal(5),
|
||||
prompt_price_unit=Decimal("0.000001"),
|
||||
prompt_price=Decimal("0.000015"),
|
||||
completion_tokens=5,
|
||||
completion_unit_price=Decimal(30),
|
||||
completion_price_unit=Decimal("0.000001"),
|
||||
completion_price=Decimal("0.000150"),
|
||||
total_price=Decimal("0.000165"),
|
||||
currency="USD",
|
||||
latency=0.5,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
@@ -714,6 +727,10 @@ def test_successful_turn_routes_stream_text_to_agent_message_and_uses_terminal_o
|
||||
assert end_events[0].llm_result.usage.prompt_tokens == 3
|
||||
assert end_events[0].llm_result.usage.completion_tokens == 5
|
||||
assert end_events[0].llm_result.usage.total_tokens == 8
|
||||
assert end_events[0].llm_result.usage.prompt_price == Decimal("0.000015")
|
||||
assert end_events[0].llm_result.usage.completion_price == Decimal("0.000150")
|
||||
assert end_events[0].llm_result.usage.total_price == Decimal("0.000165")
|
||||
assert end_events[0].llm_result.usage.currency == "USD"
|
||||
rows = sorted(fake_session.rows.values(), key=lambda row: row.position)
|
||||
assert rows == []
|
||||
assert store.saved
|
||||
|
||||
@@ -485,7 +485,11 @@ class TestBaseAppGeneratorExtras:
|
||||
account.id = "account-id"
|
||||
account.tenant_id = "tenant-id"
|
||||
|
||||
factory = base_app_generator._get_draft_var_saver_factory(InvokeFrom.DEBUGGER, account)
|
||||
factory = base_app_generator._get_draft_var_saver_factory(
|
||||
InvokeFrom.DEBUGGER,
|
||||
account,
|
||||
tenant_id="tenant-id",
|
||||
)
|
||||
saver = factory(
|
||||
app_id="app-id",
|
||||
node_id="node-id",
|
||||
|
||||
@@ -185,7 +185,11 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture):
|
||||
return_value="converted",
|
||||
)
|
||||
mocker.patch.object(WorkflowAppGenerator, "_handle_response", return_value="response")
|
||||
mocker.patch.object(WorkflowAppGenerator, "_get_draft_var_saver_factory", return_value=MagicMock())
|
||||
draft_saver_factory = mocker.patch.object(
|
||||
WorkflowAppGenerator,
|
||||
"_get_draft_var_saver_factory",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
|
||||
pause_layer = MagicMock(name="pause-layer")
|
||||
mocker.patch(
|
||||
@@ -209,7 +213,7 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture):
|
||||
|
||||
mocker.patch("core.app.apps.workflow.app_generator.threading.Thread", DummyThread)
|
||||
|
||||
app_model = SimpleNamespace(mode="workflow")
|
||||
app_model = SimpleNamespace(mode="workflow", tenant_id="tenant")
|
||||
app_config = SimpleNamespace(app_id="app", tenant_id="tenant", workflow_id="wf")
|
||||
application_generate_entity = SimpleNamespace(
|
||||
task_id="task",
|
||||
@@ -240,6 +244,7 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture):
|
||||
assert result == "converted"
|
||||
assert worker_kwargs["kwargs"]["graph_engine_layers"] == ("base-layer", pause_layer)
|
||||
assert worker_kwargs["kwargs"]["graph_runtime_state"] is graph_runtime_state
|
||||
assert draft_saver_factory.call_args.kwargs["tenant_id"] == app_model.tenant_id
|
||||
|
||||
|
||||
def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
|
||||
@@ -301,7 +306,7 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture):
|
||||
|
||||
pause_config = SimpleNamespace(session_factory=MagicMock(), state_owner_user_id="owner")
|
||||
|
||||
app_model = SimpleNamespace(mode="workflow")
|
||||
app_model = SimpleNamespace(mode="workflow", tenant_id="tenant")
|
||||
app_config = SimpleNamespace(app_id="app", tenant_id="tenant", workflow_id="workflow")
|
||||
application_generate_entity = SimpleNamespace(
|
||||
task_id="task",
|
||||
|
||||
@@ -4,6 +4,8 @@ from time import time
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from core.app.app_config.entities import WorkflowUIBasedAppConfig
|
||||
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity
|
||||
@@ -32,6 +34,13 @@ from models.model import AppMode
|
||||
from repositories.factory import DifyAPIRepositoryFactory
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sqlite_session_factory(sqlite_engine: Engine) -> sessionmaker[Session]:
|
||||
"""Provide the real session factory injected into the persistence layer."""
|
||||
|
||||
return sessionmaker(sqlite_engine, expire_on_commit=False)
|
||||
|
||||
|
||||
def _create_initialized_response_stream_filter() -> ResponseStreamFilter:
|
||||
"""Build a `ResponseStreamFilter` that has already run `initialize()`.
|
||||
|
||||
@@ -211,27 +220,25 @@ class TestPauseStatePersistenceLayer:
|
||||
workflow_execution_id=workflow_execution_id,
|
||||
)
|
||||
|
||||
def test_init_with_dependency_injection(self):
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_init_with_dependency_injection(self, sqlite_session_factory: sessionmaker[Session]):
|
||||
state_owner_user_id = "user-123"
|
||||
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id=state_owner_user_id,
|
||||
generate_entity=self._create_generate_entity(),
|
||||
response_stream_filter=ResponseStreamFilter(),
|
||||
)
|
||||
|
||||
assert layer._session_maker is session_factory
|
||||
assert layer._session_maker is sqlite_session_factory
|
||||
assert layer._state_owner_user_id == state_owner_user_id
|
||||
with pytest.raises(GraphEngineLayerNotInitializedError):
|
||||
_ = layer.graph_runtime_state
|
||||
assert layer.command_channel is None
|
||||
|
||||
def test_initialize_sets_dependencies(self):
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_initialize_sets_dependencies(self, sqlite_session_factory: sessionmaker[Session]):
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id="owner",
|
||||
generate_entity=self._create_generate_entity(),
|
||||
response_stream_filter=ResponseStreamFilter(),
|
||||
@@ -245,11 +252,12 @@ class TestPauseStatePersistenceLayer:
|
||||
assert layer.graph_runtime_state is graph_runtime_state
|
||||
assert layer.command_channel is command_channel
|
||||
|
||||
def test_on_event_with_graph_run_paused_event(self, monkeypatch: pytest.MonkeyPatch):
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_on_event_with_graph_run_paused_event(
|
||||
self, monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]
|
||||
):
|
||||
generate_entity = self._create_generate_entity(workflow_execution_id="run-123")
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id="owner-123",
|
||||
generate_entity=generate_entity,
|
||||
response_stream_filter=_create_initialized_response_stream_filter(),
|
||||
@@ -272,7 +280,7 @@ class TestPauseStatePersistenceLayer:
|
||||
|
||||
layer.on_event(event)
|
||||
|
||||
mock_factory.assert_called_once_with(session_factory)
|
||||
mock_factory.assert_called_once_with(sqlite_session_factory)
|
||||
assert mock_repo.create_workflow_pause.call_count == 1
|
||||
call_kwargs = mock_repo.create_workflow_pause.call_args.kwargs
|
||||
assert call_kwargs["workflow_run_id"] == "run-123"
|
||||
@@ -285,11 +293,12 @@ class TestPauseStatePersistenceLayer:
|
||||
|
||||
assert isinstance(pause_reasons, list)
|
||||
|
||||
def test_on_event_enriches_hitl_pause_reasons_before_persisting(self, monkeypatch: pytest.MonkeyPatch):
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_on_event_enriches_hitl_pause_reasons_before_persisting(
|
||||
self, monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]
|
||||
):
|
||||
generate_entity = self._create_generate_entity(workflow_execution_id="run-123")
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id="owner-123",
|
||||
generate_entity=generate_entity,
|
||||
response_stream_filter=_create_initialized_response_stream_filter(),
|
||||
@@ -343,10 +352,11 @@ class TestPauseStatePersistenceLayer:
|
||||
)
|
||||
assert mock_repo.create_workflow_pause.call_args.kwargs["pause_reasons"] == [enriched_reason]
|
||||
|
||||
def test_on_event_ignores_non_paused_events(self, monkeypatch: pytest.MonkeyPatch):
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_on_event_ignores_non_paused_events(
|
||||
self, monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]
|
||||
):
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id="owner-123",
|
||||
generate_entity=self._create_generate_entity(),
|
||||
response_stream_filter=ResponseStreamFilter(),
|
||||
@@ -372,10 +382,11 @@ class TestPauseStatePersistenceLayer:
|
||||
mock_factory.assert_not_called()
|
||||
mock_repo.create_workflow_pause.assert_not_called()
|
||||
|
||||
def test_on_event_raises_when_graph_runtime_state_is_uninitialized(self):
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_on_event_raises_when_graph_runtime_state_is_uninitialized(
|
||||
self, sqlite_session_factory: sessionmaker[Session]
|
||||
):
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id="owner-123",
|
||||
generate_entity=self._create_generate_entity(),
|
||||
response_stream_filter=ResponseStreamFilter(),
|
||||
@@ -386,10 +397,11 @@ class TestPauseStatePersistenceLayer:
|
||||
with pytest.raises(GraphEngineLayerNotInitializedError):
|
||||
layer.on_event(event)
|
||||
|
||||
def test_on_event_asserts_when_workflow_execution_id_missing(self, monkeypatch: pytest.MonkeyPatch):
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_on_event_asserts_when_workflow_execution_id_missing(
|
||||
self, monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]
|
||||
):
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id="owner-123",
|
||||
generate_entity=self._create_generate_entity(),
|
||||
response_stream_filter=_create_initialized_response_stream_filter(),
|
||||
@@ -494,12 +506,13 @@ def test_workflow_resumption_context_dumps_loads_roundtrip(state: WorkflowResump
|
||||
assert restored_entity.extras["trace_session_id"] == "session-1"
|
||||
|
||||
|
||||
def test_on_event_persists_response_stream_filter_dump(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
session_factory = Mock(name="session_factory")
|
||||
def test_on_event_persists_response_stream_filter_dump(
|
||||
monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]
|
||||
) -> None:
|
||||
generate_entity = TestPauseStatePersistenceLayer._create_generate_entity(workflow_execution_id="run-123")
|
||||
response_stream_filter = _create_initialized_response_stream_filter()
|
||||
layer = PauseStatePersistenceLayer(
|
||||
session_factory=session_factory,
|
||||
session_factory=sqlite_session_factory,
|
||||
state_owner_user_id="owner-123",
|
||||
generate_entity=generate_entity,
|
||||
response_stream_filter=response_stream_filter,
|
||||
|
||||
@@ -1,15 +1,52 @@
|
||||
from decimal import Decimal
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from clients.agent_backend.errors import AgentBackendRunFailedError
|
||||
from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.app.entities.queue_entities import QueueErrorEvent
|
||||
from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
|
||||
from core.errors.error import QuotaExceededError
|
||||
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError, InvokeRateLimitError
|
||||
from models.enums import MessageStatus
|
||||
from models.enums import ConversationFromSource, MessageStatus
|
||||
from models.model import AppMode, Message
|
||||
|
||||
|
||||
def _persist_message(session: Session, *, message_id: str) -> Message:
|
||||
message = Message(
|
||||
id=message_id,
|
||||
app_id="app-1",
|
||||
model_provider=None,
|
||||
model_id=None,
|
||||
override_model_configs=None,
|
||||
conversation_id="conversation-1",
|
||||
inputs={},
|
||||
query="query",
|
||||
message={},
|
||||
message_unit_price=Decimal(0),
|
||||
answer="",
|
||||
answer_unit_price=Decimal(0),
|
||||
parent_message_id=None,
|
||||
total_price=None,
|
||||
currency="USD",
|
||||
status=MessageStatus.NORMAL,
|
||||
error=None,
|
||||
message_metadata=None,
|
||||
invoke_from=InvokeFrom.WEB_APP,
|
||||
from_source=ConversationFromSource.CONSOLE,
|
||||
from_end_user_id=None,
|
||||
from_account_id="account-1",
|
||||
workflow_run_id=None,
|
||||
app_mode=AppMode.COMPLETION,
|
||||
)
|
||||
session.add(message)
|
||||
session.commit()
|
||||
session.expunge_all()
|
||||
return message
|
||||
|
||||
|
||||
class TestBasedGenerateTaskPipeline:
|
||||
@@ -58,26 +95,33 @@ class TestBasedGenerateTaskPipeline:
|
||||
assert "Knowledge retrieval failed" in str(err)
|
||||
assert "agent_run_id=run-1" in str(err)
|
||||
|
||||
def test_handle_error_updates_message_when_found(self, pipeline):
|
||||
@pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True)
|
||||
def test_handle_error_updates_message_when_found(self, pipeline, sqlite_session: Session):
|
||||
event = QueueErrorEvent(error=ValueError("oops"))
|
||||
message = SimpleNamespace(status=MessageStatus.NORMAL, error=None)
|
||||
session = Mock()
|
||||
session.scalar.return_value = message
|
||||
_persist_message(sqlite_session, message_id="msg-1")
|
||||
|
||||
err = pipeline.handle_error(event=event, session=session, message_id="msg-1")
|
||||
err = pipeline.handle_error(event=event, session=sqlite_session, message_id="msg-1")
|
||||
|
||||
assert err is event.error
|
||||
assert message.status == MessageStatus.ERROR
|
||||
assert message.error == "oops"
|
||||
sqlite_session.flush()
|
||||
sqlite_session.expire_all()
|
||||
updated_message = sqlite_session.get(Message, "msg-1")
|
||||
assert updated_message is not None
|
||||
assert updated_message.status == MessageStatus.ERROR
|
||||
assert updated_message.error == "oops"
|
||||
|
||||
def test_handle_error_returns_err_when_message_missing(self, pipeline):
|
||||
@pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True)
|
||||
def test_handle_error_returns_err_when_message_missing(self, pipeline, sqlite_session: Session):
|
||||
event = QueueErrorEvent(error=ValueError("oops"))
|
||||
session = Mock()
|
||||
session.scalar.return_value = None
|
||||
_persist_message(sqlite_session, message_id="other-message")
|
||||
|
||||
err = pipeline.handle_error(event=event, session=session, message_id="msg-1")
|
||||
err = pipeline.handle_error(event=event, session=sqlite_session, message_id="msg-1")
|
||||
|
||||
assert err is event.error
|
||||
untouched_message = sqlite_session.get(Message, "other-message")
|
||||
assert untouched_message is not None
|
||||
assert untouched_message.status == MessageStatus.NORMAL
|
||||
assert untouched_message.error is None
|
||||
|
||||
def test_error_to_stream_response_and_ping(self, pipeline):
|
||||
error_response = pipeline.error_to_stream_response(ValueError("boom"))
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
import gzip
|
||||
from typing import override
|
||||
from unittest.mock import ANY, MagicMock, call, patch
|
||||
|
||||
import httpx
|
||||
@@ -5,10 +7,13 @@ import pytest
|
||||
|
||||
from core.helper.ssrf_proxy import (
|
||||
SSRF_DEFAULT_MAX_RETRIES,
|
||||
ResponseTooLargeError,
|
||||
SSRFProxy,
|
||||
UnsupportedResponseEncodingError,
|
||||
_build_ssrf_client,
|
||||
_get_user_provided_host_header,
|
||||
_to_graphon_http_response,
|
||||
buffer_response,
|
||||
graphon_ssrf_proxy,
|
||||
make_request,
|
||||
max_retries_exceeded_error,
|
||||
@@ -21,12 +26,91 @@ def test_successful_request(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
response = make_request("GET", "http://example.com")
|
||||
assert response.status_code == 200
|
||||
mock_client.request.assert_called_once()
|
||||
mock_client.build_request.assert_called_once()
|
||||
mock_client.send.assert_called_once()
|
||||
|
||||
|
||||
def test_buffer_response_rejects_encoded_response_before_decoding() -> None:
|
||||
payload = b"x" * (8 * 1024 * 1024)
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
content=gzip.compress(payload),
|
||||
headers={"Content-Encoding": "gzip"},
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
with httpx.Client(transport=transport) as client:
|
||||
response = client.send(client.build_request("GET", "http://example.com"), stream=True)
|
||||
with pytest.raises(UnsupportedResponseEncodingError, match="content encoding gzip"):
|
||||
buffer_response(response, max_response_bytes=1024 * 1024)
|
||||
|
||||
|
||||
def test_buffer_response_returns_response_within_decoded_byte_limit() -> None:
|
||||
payload = b"response-body"
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
content=payload,
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
with httpx.Client(transport=transport) as client:
|
||||
streaming_response = client.send(client.build_request("GET", "http://example.com"), stream=True)
|
||||
response = buffer_response(streaming_response, max_response_bytes=32)
|
||||
|
||||
assert response.content == payload
|
||||
assert str(response.request.url) == "http://example.com"
|
||||
|
||||
|
||||
def test_buffer_response_rejects_identity_response_exceeding_byte_limit() -> None:
|
||||
payload = b"response-body"
|
||||
transport = httpx.MockTransport(lambda request: httpx.Response(200, content=payload, request=request))
|
||||
|
||||
with httpx.Client(transport=transport) as client:
|
||||
response = client.send(client.build_request("GET", "http://example.com"), stream=True)
|
||||
with pytest.raises(ResponseTooLargeError, match="response exceeded 8 bytes"):
|
||||
buffer_response(response, max_response_bytes=8)
|
||||
|
||||
|
||||
def test_request_can_return_an_open_stream_the_caller_closes() -> None:
|
||||
class EventStream(httpx.SyncByteStream):
|
||||
@override
|
||||
def __iter__(self):
|
||||
yield b"event: delta\ndata: first\n\n"
|
||||
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
stream=EventStream(),
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
httpx.Client(transport=transport) as client,
|
||||
patch("core.helper.ssrf_proxy._get_ssrf_client", return_value=client),
|
||||
):
|
||||
response = make_request(
|
||||
"GET",
|
||||
"http://example.com/events",
|
||||
max_retries=0,
|
||||
stream_response=True,
|
||||
)
|
||||
|
||||
assert response.is_stream_consumed is False
|
||||
assert response.is_closed is False
|
||||
assert b"".join(response.iter_bytes()) == b"event: delta\ndata: first\n\n"
|
||||
response.close()
|
||||
|
||||
assert response.is_closed
|
||||
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
@@ -34,7 +118,7 @@ def test_retry_exceed_max_retries(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with pytest.raises(Exception) as e:
|
||||
@@ -47,13 +131,13 @@ def test_force_list_response_returns_when_retries_disabled(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
response = make_request("GET", "http://example.com", max_retries=0)
|
||||
|
||||
assert response is mock_response
|
||||
mock_client.request.assert_called_once()
|
||||
mock_client.send.assert_called_once()
|
||||
|
||||
|
||||
def test_build_ssrf_client_passes_ssl_verify_to_proxy_mount_transports():
|
||||
@@ -126,15 +210,15 @@ def test_host_header_preservation_with_user_header(mock_get_client):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
custom_host = "custom.example.com:8080"
|
||||
response = make_request("GET", "http://example.com", headers={"Host": custom_host})
|
||||
|
||||
assert response.status_code == 200
|
||||
# Verify client.request was called with the host header preserved (lowercase)
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
# Verify the request was built with the host header preserved (lowercase)
|
||||
call_kwargs = mock_client.build_request.call_args.kwargs
|
||||
assert call_kwargs["headers"]["host"] == custom_host
|
||||
|
||||
|
||||
@@ -145,36 +229,36 @@ def test_host_header_preservation_case_insensitive(mock_get_client, host_key):
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
response = make_request("GET", "http://example.com", headers={host_key: "api.example.com"})
|
||||
|
||||
assert response.status_code == 200
|
||||
# Host header should be normalized to lowercase "host"
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.build_request.call_args.kwargs
|
||||
assert call_kwargs["headers"]["host"] == "api.example.com"
|
||||
|
||||
|
||||
class TestFollowRedirectsParameter:
|
||||
"""Tests for follow_redirects parameter handling.
|
||||
|
||||
These tests verify that follow_redirects is correctly passed to client.request().
|
||||
These tests verify that follow_redirects is correctly passed to client.send().
|
||||
"""
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
def test_follow_redirects_passed_to_request(self, mock_get_client):
|
||||
"""Verify follow_redirects IS passed to client.request()."""
|
||||
"""Verify follow_redirects IS passed to client.send()."""
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
make_request("GET", "http://example.com", follow_redirects=True)
|
||||
|
||||
# Verify follow_redirects was passed to request
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
# Verify follow_redirects was passed to send
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert call_kwargs.get("follow_redirects") is True
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
@@ -183,14 +267,14 @@ class TestFollowRedirectsParameter:
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
# Use allow_redirects (requests-style parameter)
|
||||
make_request("GET", "http://example.com", allow_redirects=True)
|
||||
|
||||
# Verify it was converted to follow_redirects
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert call_kwargs.get("follow_redirects") is True
|
||||
assert "allow_redirects" not in call_kwargs
|
||||
|
||||
@@ -200,13 +284,13 @@ class TestFollowRedirectsParameter:
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
make_request("GET", "http://example.com")
|
||||
|
||||
# follow_redirects should not be in kwargs, letting httpx use its default
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert "follow_redirects" not in call_kwargs
|
||||
|
||||
@patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True)
|
||||
@@ -215,13 +299,13 @@ class TestFollowRedirectsParameter:
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_client.request.return_value = mock_response
|
||||
mock_client.send.return_value = mock_response
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
# Both specified - follow_redirects should take precedence
|
||||
make_request("GET", "http://example.com", allow_redirects=False, follow_redirects=True)
|
||||
|
||||
call_kwargs = mock_client.request.call_args.kwargs
|
||||
call_kwargs = mock_client.send.call_args.kwargs
|
||||
assert call_kwargs.get("follow_redirects") is True
|
||||
|
||||
|
||||
|
||||
@@ -1,11 +1,51 @@
|
||||
"""Unit tests for Notion extraction, HTTP parsing, and persisted metadata updates."""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from sqlalchemy import Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from core.rag.extractor import notion_extractor
|
||||
from core.rag.index_processor.constant.index_type import IndexStructureType
|
||||
from models.base import TypeBase
|
||||
from models.dataset import Document as DocumentModel
|
||||
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def persisted_document(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sqlite_engine: Engine,
|
||||
) -> Iterator[tuple[sessionmaker[Session], DocumentModel]]:
|
||||
"""Persist a Notion document and bind the extractor's ``db.session`` to SQLite."""
|
||||
TypeBase.metadata.create_all(sqlite_engine, tables=[DocumentModel.__table__])
|
||||
session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
|
||||
document = DocumentModel(
|
||||
id=str(uuid.uuid4()),
|
||||
tenant_id=str(uuid.uuid4()),
|
||||
dataset_id=str(uuid.uuid4()),
|
||||
position=1,
|
||||
data_source_type=DataSourceType.NOTION_IMPORT,
|
||||
data_source_info=json.dumps({"source": "notion", "last_edited_time": "2025-01-01T00:00:00.000Z"}),
|
||||
batch="batch",
|
||||
name="Notion page",
|
||||
created_from=DocumentCreatedFrom.API,
|
||||
created_by=str(uuid.uuid4()),
|
||||
indexing_status=IndexingStatus.COMPLETED,
|
||||
doc_form=IndexStructureType.PARAGRAPH_INDEX,
|
||||
)
|
||||
with session_maker() as sqlite_session:
|
||||
sqlite_session.add(document)
|
||||
sqlite_session.commit()
|
||||
monkeypatch.setattr(notion_extractor, "db", SimpleNamespace(session=sqlite_session))
|
||||
yield session_maker, document
|
||||
|
||||
|
||||
def _mock_response(data, status_code: int = 200, text: str = ""):
|
||||
@@ -394,7 +434,11 @@ class TestNotionMetadataAndCredentialMethods:
|
||||
|
||||
assert extractor.update_last_edited_time(None) is None
|
||||
|
||||
def test_update_last_edited_time_updates_document_and_commits(self, monkeypatch: pytest.MonkeyPatch):
|
||||
def test_update_last_edited_time_updates_document_and_commits(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
persisted_document: tuple[sessionmaker[Session], DocumentModel],
|
||||
):
|
||||
extractor = notion_extractor.NotionExtractor(
|
||||
notion_workspace_id="ws",
|
||||
notion_obj_id="obj",
|
||||
@@ -402,40 +446,20 @@ class TestNotionMetadataAndCredentialMethods:
|
||||
tenant_id="tenant",
|
||||
notion_access_token="token",
|
||||
)
|
||||
|
||||
class FakeDocumentModel:
|
||||
data_source_info = "data_source_info"
|
||||
id = "id"
|
||||
|
||||
execute_calls = []
|
||||
|
||||
class FakeUpdateStmt:
|
||||
def where(self, *args):
|
||||
return self
|
||||
|
||||
def values(self, **kwargs):
|
||||
return self
|
||||
|
||||
class FakeSession:
|
||||
committed = False
|
||||
|
||||
def execute(self, stmt):
|
||||
execute_calls.append(stmt)
|
||||
|
||||
def commit(self):
|
||||
self.committed = True
|
||||
|
||||
fake_db = SimpleNamespace(session=FakeSession())
|
||||
monkeypatch.setattr(notion_extractor, "DocumentModel", FakeDocumentModel)
|
||||
monkeypatch.setattr(notion_extractor, "update", lambda model: FakeUpdateStmt())
|
||||
monkeypatch.setattr(notion_extractor, "db", fake_db)
|
||||
monkeypatch.setattr(extractor, "get_notion_last_edited_time", lambda: "2026-01-01T00:00:00.000Z")
|
||||
session_maker, document = persisted_document
|
||||
|
||||
doc_model = SimpleNamespace(id="doc-1", data_source_info_dict={"source": "notion"})
|
||||
extractor.update_last_edited_time(doc_model)
|
||||
extractor.update_last_edited_time(document)
|
||||
|
||||
assert execute_calls
|
||||
assert fake_db.session.committed is True
|
||||
# Closing the writer session rolls back an uncommitted update before the independent read below.
|
||||
notion_extractor.db.session.close()
|
||||
with session_maker() as verification_session:
|
||||
stored_document = verification_session.get(DocumentModel, document.id)
|
||||
assert stored_document is not None
|
||||
assert stored_document.data_source_info_dict == {
|
||||
"source": "notion",
|
||||
"last_edited_time": "2026-01-01T00:00:00.000Z",
|
||||
}
|
||||
|
||||
def test_get_notion_last_edited_time_uses_page_and_database_urls(self, mocker: MockerFixture):
|
||||
extractor_page = notion_extractor.NotionExtractor(
|
||||
|
||||
@@ -176,6 +176,22 @@ class TestWaterCrawlAPIClient:
|
||||
with pytest.raises(expected_exception):
|
||||
client.process_response(_response(status, {"message": "bad", "errors": {"url": ["x"]}}))
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status", "expected_exception"),
|
||||
[
|
||||
(401, WaterCrawlAuthenticationError),
|
||||
(403, WaterCrawlPermissionError),
|
||||
(422, WaterCrawlBadRequestError),
|
||||
],
|
||||
)
|
||||
def test_process_response_error_statuses_with_non_json_body(self, status: int, expected_exception: type[Exception]):
|
||||
client = WaterCrawlAPIClient(api_key="k")
|
||||
response = _response(status, text="<html>upstream error</html>")
|
||||
response.json.side_effect = json.JSONDecodeError("Expecting value", response.text, 0)
|
||||
|
||||
with pytest.raises(expected_exception):
|
||||
client.process_response(response)
|
||||
|
||||
def test_process_response_204_returns_none(self):
|
||||
client = WaterCrawlAPIClient(api_key="k")
|
||||
assert client.process_response(_response(204, None)) is None
|
||||
@@ -185,6 +201,14 @@ class TestWaterCrawlAPIClient:
|
||||
assert client.process_response(_response(200, {"ok": True})) == {"ok": True}
|
||||
assert client.process_response(_response(200, None)) == {}
|
||||
|
||||
def test_process_response_json_payload_with_invalid_body_raises_clear_error(self):
|
||||
client = WaterCrawlAPIClient(api_key="k")
|
||||
response = _response(200, text="<html>upstream error</html>")
|
||||
response.json.side_effect = json.JSONDecodeError("Expecting value", response.text, 0)
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid JSON response from WaterCrawl"):
|
||||
client.process_response(response)
|
||||
|
||||
def test_process_response_accepts_json_content_type_parameters(self):
|
||||
client = WaterCrawlAPIClient(api_key="k")
|
||||
|
||||
|
||||
@@ -3,7 +3,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.app_config.entities import DatasetRetrieveConfigEntity
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
@@ -14,13 +17,14 @@ def _retrieve_config() -> DatasetRetrieveConfigEntity:
|
||||
return DatasetRetrieveConfigEntity(retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE)
|
||||
|
||||
|
||||
def test_get_dataset_tools_returns_empty_for_empty_dataset_ids() -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_get_dataset_tools_returns_empty_for_empty_dataset_ids(sqlite_session: Session) -> None:
|
||||
# Arrange
|
||||
retrieve_config = _retrieve_config()
|
||||
|
||||
# Act
|
||||
tools = DatasetRetrieverTool.get_dataset_tools(
|
||||
session=MagicMock(),
|
||||
session=sqlite_session,
|
||||
tenant_id="tenant",
|
||||
dataset_ids=[],
|
||||
retrieve_config=retrieve_config,
|
||||
@@ -35,13 +39,14 @@ def test_get_dataset_tools_returns_empty_for_empty_dataset_ids() -> None:
|
||||
assert tools == []
|
||||
|
||||
|
||||
def test_get_dataset_tools_returns_empty_for_missing_retrieve_config() -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_get_dataset_tools_returns_empty_for_missing_retrieve_config(sqlite_session: Session) -> None:
|
||||
# Arrange
|
||||
dataset_ids = ["d1"]
|
||||
|
||||
# Act
|
||||
tools = DatasetRetrieverTool.get_dataset_tools(
|
||||
session=MagicMock(),
|
||||
session=sqlite_session,
|
||||
tenant_id="tenant",
|
||||
dataset_ids=dataset_ids,
|
||||
retrieve_config=None, # type: ignore[arg-type]
|
||||
@@ -56,7 +61,8 @@ def test_get_dataset_tools_returns_empty_for_missing_retrieve_config() -> None:
|
||||
assert tools == []
|
||||
|
||||
|
||||
def test_get_dataset_tools_builds_tool_and_restores_strategy() -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_get_dataset_tools_builds_tool_and_restores_strategy(sqlite_session: Session) -> None:
|
||||
# Arrange
|
||||
retrieve_config = _retrieve_config()
|
||||
retrieval_tool = SimpleNamespace(name="dataset_tool", description="desc", run=lambda query: f"result:{query}")
|
||||
@@ -66,7 +72,7 @@ def test_get_dataset_tools_builds_tool_and_restores_strategy() -> None:
|
||||
# Act
|
||||
with patch("core.tools.utils.dataset_retriever_tool.DatasetRetrieval", return_value=feature):
|
||||
tools = DatasetRetrieverTool.get_dataset_tools(
|
||||
session=MagicMock(),
|
||||
session=sqlite_session,
|
||||
tenant_id="tenant",
|
||||
dataset_ids=["d1"],
|
||||
retrieve_config=retrieve_config,
|
||||
@@ -83,7 +89,7 @@ def test_get_dataset_tools_builds_tool_and_restores_strategy() -> None:
|
||||
assert retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE
|
||||
|
||||
|
||||
def _build_dataset_tool() -> tuple[DatasetRetrieverTool, SimpleNamespace]:
|
||||
def _build_dataset_tool(sqlite_session: Session) -> tuple[DatasetRetrieverTool, SimpleNamespace]:
|
||||
retrieval_tool = SimpleNamespace(
|
||||
name="dataset_tool",
|
||||
description="desc",
|
||||
@@ -93,7 +99,7 @@ def _build_dataset_tool() -> tuple[DatasetRetrieverTool, SimpleNamespace]:
|
||||
feature.to_dataset_retriever_tool.return_value = [retrieval_tool]
|
||||
with patch("core.tools.utils.dataset_retriever_tool.DatasetRetrieval", return_value=feature):
|
||||
tools = DatasetRetrieverTool.get_dataset_tools(
|
||||
session=MagicMock(),
|
||||
session=sqlite_session,
|
||||
tenant_id="tenant",
|
||||
dataset_ids=["d1"],
|
||||
retrieve_config=_retrieve_config(),
|
||||
@@ -106,9 +112,10 @@ def _build_dataset_tool() -> tuple[DatasetRetrieverTool, SimpleNamespace]:
|
||||
return tools[0], retrieval_tool
|
||||
|
||||
|
||||
def test_runtime_parameters_shape() -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_runtime_parameters_shape(sqlite_session: Session) -> None:
|
||||
# Arrange
|
||||
tool, _ = _build_dataset_tool()
|
||||
tool, _ = _build_dataset_tool(sqlite_session)
|
||||
|
||||
# Act
|
||||
params = tool.get_runtime_parameters()
|
||||
@@ -118,33 +125,36 @@ def test_runtime_parameters_shape() -> None:
|
||||
assert params[0].name == "query"
|
||||
|
||||
|
||||
def test_empty_query_behavior() -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_empty_query_behavior(sqlite_session: Session) -> None:
|
||||
# Arrange
|
||||
tool, _ = _build_dataset_tool()
|
||||
tool, _ = _build_dataset_tool(sqlite_session)
|
||||
|
||||
# Act
|
||||
empty_query = list(tool.invoke(session=MagicMock(), user_id="u", tool_parameters={}))
|
||||
empty_query = list(tool.invoke(session=sqlite_session, user_id="u", tool_parameters={}))
|
||||
|
||||
# Assert
|
||||
assert len(empty_query) == 1
|
||||
assert empty_query[0].message.text == "please input query"
|
||||
|
||||
|
||||
def test_query_invocation_result() -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_query_invocation_result(sqlite_session: Session) -> None:
|
||||
# Arrange
|
||||
tool, _ = _build_dataset_tool()
|
||||
tool, _ = _build_dataset_tool(sqlite_session)
|
||||
|
||||
# Act
|
||||
result = list(tool.invoke(session=MagicMock(), user_id="u", tool_parameters={"query": "hello"}))
|
||||
result = list(tool.invoke(session=sqlite_session, user_id="u", tool_parameters={"query": "hello"}))
|
||||
|
||||
# Assert
|
||||
assert len(result) == 1
|
||||
assert result[0].message.text == "result:hello"
|
||||
|
||||
|
||||
def test_validate_credentials() -> None:
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_validate_credentials(sqlite_session: Session) -> None:
|
||||
# Arrange
|
||||
tool, _ = _build_dataset_tool()
|
||||
tool, _ = _build_dataset_tool(sqlite_session)
|
||||
|
||||
# Act
|
||||
result = tool.validate_credentials(credentials={}, parameters={}, format_only=False)
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
"""MCP tool tests using real SQLite sessions for the ORM invocation contract."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.mcp.types import (
|
||||
@@ -97,7 +100,8 @@ def test_mcp_tool_usage_extraction_helpers():
|
||||
assert derived.total_tokens == 0
|
||||
|
||||
|
||||
def test_mcp_tool_invoke_handles_content_types_and_structured_output():
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_mcp_tool_invoke_handles_content_types_and_structured_output(sqlite_session: Session):
|
||||
tool = _build_mcp_tool()
|
||||
img_data = base64.b64encode(b"img").decode()
|
||||
blob_data = base64.b64encode(b"blob").decode()
|
||||
@@ -123,7 +127,7 @@ def test_mcp_tool_invoke_handles_content_types_and_structured_output():
|
||||
)
|
||||
|
||||
with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result):
|
||||
messages = list(tool.invoke(session=MagicMock(), user_id="user-1", tool_parameters={"a": 1}))
|
||||
messages = list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={"a": 1}))
|
||||
|
||||
types = [m.type for m in messages]
|
||||
assert ToolInvokeMessage.MessageType.JSON in types
|
||||
@@ -133,7 +137,8 @@ def test_mcp_tool_invoke_handles_content_types_and_structured_output():
|
||||
assert tool.latest_usage.total_tokens == 5
|
||||
|
||||
|
||||
def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource():
|
||||
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
|
||||
def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource(sqlite_session: Session):
|
||||
tool = _build_mcp_tool()
|
||||
# Use model_construct to bypass pydantic validation and force unsupported resource path.
|
||||
bad_resource = EmbeddedResource.model_construct(type="resource", resource=object())
|
||||
@@ -141,7 +146,7 @@ def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource():
|
||||
|
||||
with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result):
|
||||
with pytest.raises(ToolInvokeError, match="Unsupported embedded resource type"):
|
||||
list(tool.invoke(session=MagicMock(), user_id="user-1", tool_parameters={}))
|
||||
list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={}))
|
||||
|
||||
|
||||
def test_mcp_tool_handle_none_parameter_filters_empty_values():
|
||||
|
||||
@@ -1,27 +1,61 @@
|
||||
"""Unit tests for `ToolFileManager` behavior.
|
||||
"""Unit tests for ``ToolFileManager`` behavior.
|
||||
|
||||
Covers signing, file persistence flows, and retrieval APIs with mocked
|
||||
storage/session boundaries (httpx, SimpleNamespace, Mock/patch) to avoid real
|
||||
IO.
|
||||
File metadata is persisted through real SQLite-backed sessions. Storage and
|
||||
remote HTTP remain mocked because they are external I/O boundaries.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
from collections.abc import Iterator
|
||||
from unittest.mock import Mock, patch
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
import core.tools.tool_file_manager as tool_file_manager_module
|
||||
from core.tools.tool_file_manager import ToolFileManager
|
||||
from graphon.file import FileTransferMethod
|
||||
from graphon.file import FileTransferMethod, FileType
|
||||
from models.base import TypeBase
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import MessageFile
|
||||
from models.tools import ToolFile
|
||||
|
||||
|
||||
def _patch_session_factory(session: Mock):
|
||||
session_cm = MagicMock()
|
||||
session_cm.__enter__.return_value = session
|
||||
session_cm.__exit__.return_value = False
|
||||
return patch("core.tools.tool_file_manager.session_factory.create_session", return_value=session_cm)
|
||||
@pytest.fixture
|
||||
def sqlite_tool_file_session(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> Iterator[Session]:
|
||||
"""Bind manager-owned sessions to SQLite and expose a setup/assertion session."""
|
||||
TypeBase.metadata.create_all(sqlite_engine, tables=[ToolFile.__table__, MessageFile.__table__])
|
||||
factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
|
||||
monkeypatch.setattr(tool_file_manager_module.session_factory, "create_session", factory)
|
||||
with factory() as session:
|
||||
yield session
|
||||
|
||||
|
||||
def _tool_file(*, file_key: str = "k1", mimetype: str = "text/plain", name: str = "file.txt") -> ToolFile:
|
||||
return ToolFile(
|
||||
user_id=str(uuid4()),
|
||||
tenant_id=str(uuid4()),
|
||||
conversation_id=str(uuid4()),
|
||||
file_key=file_key,
|
||||
mimetype=mimetype,
|
||||
original_url=None,
|
||||
name=name,
|
||||
size=12,
|
||||
)
|
||||
|
||||
|
||||
def _message_file(*, url: str | None) -> MessageFile:
|
||||
return MessageFile(
|
||||
message_id=str(uuid4()),
|
||||
type=FileType.IMAGE,
|
||||
transfer_method=FileTransferMethod.TOOL_FILE,
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=str(uuid4()),
|
||||
url=url,
|
||||
)
|
||||
|
||||
|
||||
def test_tool_file_manager_sign_file_builds_url() -> None:
|
||||
@@ -29,120 +63,105 @@ def test_tool_file_manager_sign_file_builds_url() -> None:
|
||||
assert "/files/tools/tf-1.png" in url
|
||||
|
||||
|
||||
def test_create_file_by_raw_stores_file_and_persists_record() -> None:
|
||||
def test_create_file_by_raw_stores_file_and_persists_record(sqlite_tool_file_session: Session) -> None:
|
||||
manager = ToolFileManager()
|
||||
session = Mock()
|
||||
session.refresh.side_effect = lambda model: setattr(model, "id", "tf-1")
|
||||
|
||||
def tool_file_factory(**kwargs):
|
||||
return SimpleNamespace(**kwargs)
|
||||
user_id = str(uuid4())
|
||||
tenant_id = str(uuid4())
|
||||
conversation_id = str(uuid4())
|
||||
|
||||
with (
|
||||
patch("core.tools.tool_file_manager.storage") as storage,
|
||||
patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory),
|
||||
patch("core.tools.tool_file_manager.guess_extension", return_value=".txt"),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="abc")),
|
||||
_patch_session_factory(session),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=UUID(int=0xABC)),
|
||||
):
|
||||
file_model = manager.create_file_by_raw(
|
||||
user_id="u1",
|
||||
tenant_id="t1",
|
||||
conversation_id="c1",
|
||||
user_id=user_id,
|
||||
tenant_id=tenant_id,
|
||||
conversation_id=conversation_id,
|
||||
file_binary=b"hello",
|
||||
mimetype="text/plain",
|
||||
filename="readme",
|
||||
)
|
||||
|
||||
assert file_model.name.endswith(".txt")
|
||||
storage.save.assert_called_once()
|
||||
session.add.assert_called_once()
|
||||
session.commit.assert_called_once()
|
||||
session.refresh.assert_called_once_with(file_model)
|
||||
persisted = sqlite_tool_file_session.get(ToolFile, file_model.id)
|
||||
assert persisted is not None
|
||||
assert persisted.name == "readme.txt"
|
||||
assert persisted.file_key == f"tools/{tenant_id}/{UUID(int=0xABC).hex}.txt"
|
||||
storage.save.assert_called_once_with(persisted.file_key, b"hello")
|
||||
|
||||
|
||||
def test_create_file_by_raw_prefers_filename_extension_over_mimetype() -> None:
|
||||
def test_create_file_by_raw_prefers_filename_extension_over_mimetype(
|
||||
sqlite_tool_file_session: Session,
|
||||
) -> None:
|
||||
manager = ToolFileManager()
|
||||
session = Mock()
|
||||
session.refresh.side_effect = lambda model: setattr(model, "id", "tf-docx")
|
||||
|
||||
def tool_file_factory(**kwargs):
|
||||
return SimpleNamespace(**kwargs)
|
||||
tenant_id = str(uuid4())
|
||||
|
||||
with (
|
||||
patch("core.tools.tool_file_manager.storage") as storage,
|
||||
patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="abc")),
|
||||
_patch_session_factory(session),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=UUID(int=0xABC)),
|
||||
):
|
||||
file_model = manager.create_file_by_raw(
|
||||
user_id="u1",
|
||||
tenant_id="t1",
|
||||
conversation_id="c1",
|
||||
user_id=str(uuid4()),
|
||||
tenant_id=tenant_id,
|
||||
conversation_id=str(uuid4()),
|
||||
file_binary=b"docx",
|
||||
mimetype="application/octet-stream",
|
||||
filename="report.docx",
|
||||
)
|
||||
|
||||
assert file_model.name == "report.docx"
|
||||
assert file_model.file_key == "tools/t1/abc.docx"
|
||||
storage.save.assert_called_once_with("tools/t1/abc.docx", b"docx")
|
||||
session.add.assert_called_once_with(file_model)
|
||||
session.commit.assert_called_once()
|
||||
session.refresh.assert_called_once_with(file_model)
|
||||
persisted = sqlite_tool_file_session.get(ToolFile, file_model.id)
|
||||
assert persisted is not None
|
||||
assert persisted.name == "report.docx"
|
||||
assert persisted.file_key == f"tools/{tenant_id}/{UUID(int=0xABC).hex}.docx"
|
||||
storage.save.assert_called_once_with(persisted.file_key, b"docx")
|
||||
|
||||
|
||||
def test_create_file_by_url_downloads_and_persists_record() -> None:
|
||||
def test_create_file_by_url_downloads_and_persists_record(sqlite_tool_file_session: Session) -> None:
|
||||
manager = ToolFileManager()
|
||||
tenant_id = str(uuid4())
|
||||
response = Mock()
|
||||
response.content = b"binary"
|
||||
response.headers = {"Content-Type": "application/octet-stream"}
|
||||
response.raise_for_status.return_value = None
|
||||
session = Mock()
|
||||
|
||||
def tool_file_factory(**kwargs):
|
||||
return SimpleNamespace(**kwargs)
|
||||
|
||||
session.refresh.side_effect = lambda model: setattr(model, "id", "tf-2")
|
||||
with (
|
||||
patch("core.tools.tool_file_manager.storage") as storage,
|
||||
patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="def")),
|
||||
_patch_session_factory(session),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=UUID(int=0xDEF)),
|
||||
patch("core.tools.tool_file_manager.remote_fetcher.make_request", return_value=response),
|
||||
):
|
||||
file_model = manager.create_file_by_url("u1", "t1", "https://example.com/f.bin", "c1")
|
||||
file_model = manager.create_file_by_url(str(uuid4()), tenant_id, "https://example.com/f.bin", str(uuid4()))
|
||||
|
||||
assert file_model.file_key.startswith("tools/t1/")
|
||||
storage.save.assert_called_once()
|
||||
session.add.assert_called_once_with(file_model)
|
||||
session.commit.assert_called_once()
|
||||
session.refresh.assert_called_once_with(file_model)
|
||||
persisted = sqlite_tool_file_session.get(ToolFile, file_model.id)
|
||||
assert persisted is not None
|
||||
assert persisted.file_key == f"tools/{tenant_id}/{UUID(int=0xDEF).hex}.bin"
|
||||
assert persisted.original_url == "https://example.com/f.bin"
|
||||
storage.save.assert_called_once_with(persisted.file_key, b"binary")
|
||||
|
||||
|
||||
def test_create_file_by_url_prefers_url_extension_over_mimetype() -> None:
|
||||
def test_create_file_by_url_prefers_url_extension_over_mimetype(
|
||||
sqlite_tool_file_session: Session,
|
||||
) -> None:
|
||||
manager = ToolFileManager()
|
||||
tenant_id = str(uuid4())
|
||||
response = Mock()
|
||||
response.content = b"docx"
|
||||
response.headers = {"Content-Type": "application/octet-stream"}
|
||||
response.raise_for_status.return_value = None
|
||||
session = Mock()
|
||||
|
||||
def tool_file_factory(**kwargs):
|
||||
return SimpleNamespace(**kwargs)
|
||||
|
||||
session.refresh.side_effect = lambda model: setattr(model, "id", "tf-docx")
|
||||
with (
|
||||
patch("core.tools.tool_file_manager.storage") as storage,
|
||||
patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="urlabc")),
|
||||
_patch_session_factory(session),
|
||||
patch("core.tools.tool_file_manager.uuid4", return_value=UUID(int=0xABC)),
|
||||
patch("core.tools.tool_file_manager.remote_fetcher.make_request", return_value=response),
|
||||
):
|
||||
file_model = manager.create_file_by_url("u1", "t1", "https://example.com/report.docx?download=1", "c1")
|
||||
file_model = manager.create_file_by_url(
|
||||
str(uuid4()), tenant_id, "https://example.com/report.docx?download=1", str(uuid4())
|
||||
)
|
||||
|
||||
assert file_model.file_key == "tools/t1/urlabc.docx"
|
||||
assert file_model.name == "urlabc.docx"
|
||||
storage.save.assert_called_once_with("tools/t1/urlabc.docx", b"docx")
|
||||
persisted = sqlite_tool_file_session.get(ToolFile, file_model.id)
|
||||
assert persisted is not None
|
||||
assert persisted.file_key == f"tools/{tenant_id}/{UUID(int=0xABC).hex}.docx"
|
||||
assert persisted.name == f"{UUID(int=0xABC).hex}.docx"
|
||||
storage.save.assert_called_once_with(persisted.file_key, b"docx")
|
||||
|
||||
|
||||
def test_create_file_by_url_raises_on_timeout() -> None:
|
||||
@@ -156,121 +175,72 @@ def test_create_file_by_url_raises_on_timeout() -> None:
|
||||
manager.create_file_by_url("u1", "t1", "https://example.com/f.bin", "c1")
|
||||
|
||||
|
||||
def test_get_file_binary_returns_none_when_not_found() -> None:
|
||||
# Arrange
|
||||
manager = ToolFileManager()
|
||||
session = Mock()
|
||||
session.scalar.return_value = None
|
||||
|
||||
# Act
|
||||
with _patch_session_factory(session):
|
||||
result = manager.get_file_binary("missing")
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
def test_get_file_binary_returns_none_when_not_found(sqlite_tool_file_session: Session) -> None:
|
||||
assert ToolFileManager().get_file_binary(str(uuid4())) is None
|
||||
|
||||
|
||||
def test_get_file_binary_returns_bytes_when_found() -> None:
|
||||
# Arrange
|
||||
manager = ToolFileManager()
|
||||
tool_file = SimpleNamespace(file_key="k1", mimetype="text/plain")
|
||||
session = Mock()
|
||||
session.scalar.return_value = tool_file
|
||||
def test_get_file_binary_returns_bytes_when_found(sqlite_tool_file_session: Session) -> None:
|
||||
tool_file = _tool_file()
|
||||
sqlite_tool_file_session.add(tool_file)
|
||||
sqlite_tool_file_session.commit()
|
||||
|
||||
# Act
|
||||
with patch("core.tools.tool_file_manager.storage") as storage:
|
||||
storage.load_once.return_value = b"hello"
|
||||
with _patch_session_factory(session):
|
||||
result = manager.get_file_binary("id1")
|
||||
result = ToolFileManager().get_file_binary(tool_file.id)
|
||||
|
||||
# Assert
|
||||
assert result == (b"hello", "text/plain")
|
||||
storage.load_once.assert_called_once_with("k1")
|
||||
|
||||
|
||||
def test_get_file_binary_by_message_file_id_when_messagefile_missing() -> None:
|
||||
# Arrange
|
||||
manager = ToolFileManager()
|
||||
session = Mock()
|
||||
session.scalar.side_effect = [None, None]
|
||||
|
||||
# Act
|
||||
with _patch_session_factory(session):
|
||||
result = manager.get_file_binary_by_message_file_id("mf-1")
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
def test_get_file_binary_by_message_file_id_when_messagefile_missing(
|
||||
sqlite_tool_file_session: Session,
|
||||
) -> None:
|
||||
assert ToolFileManager().get_file_binary_by_message_file_id(str(uuid4())) is None
|
||||
|
||||
|
||||
def test_get_file_binary_by_message_file_id_when_url_is_none() -> None:
|
||||
# Arrange
|
||||
manager = ToolFileManager()
|
||||
message_file = SimpleNamespace(url=None)
|
||||
session = Mock()
|
||||
session.scalar.side_effect = [message_file, None]
|
||||
def test_get_file_binary_by_message_file_id_when_url_is_none(sqlite_tool_file_session: Session) -> None:
|
||||
message_file = _message_file(url=None)
|
||||
sqlite_tool_file_session.add(message_file)
|
||||
sqlite_tool_file_session.commit()
|
||||
|
||||
# Act
|
||||
with _patch_session_factory(session):
|
||||
result = manager.get_file_binary_by_message_file_id("mf-1")
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
assert ToolFileManager().get_file_binary_by_message_file_id(message_file.id) is None
|
||||
|
||||
|
||||
def test_get_file_binary_by_message_file_id_returns_bytes_when_found() -> None:
|
||||
# Arrange
|
||||
manager = ToolFileManager()
|
||||
message_file = SimpleNamespace(url="https://x/files/tools/tool123.png")
|
||||
tool_file = SimpleNamespace(file_key="k2", mimetype="image/png")
|
||||
session = Mock()
|
||||
session.scalar.side_effect = [message_file, tool_file]
|
||||
def test_get_file_binary_by_message_file_id_returns_bytes_when_found(
|
||||
sqlite_tool_file_session: Session,
|
||||
) -> None:
|
||||
tool_file = _tool_file(file_key="k2", mimetype="image/png", name="image.png")
|
||||
message_file = _message_file(url=f"https://x/files/tools/{tool_file.id}.png")
|
||||
sqlite_tool_file_session.add_all([tool_file, message_file])
|
||||
sqlite_tool_file_session.commit()
|
||||
|
||||
# Act
|
||||
with patch("core.tools.tool_file_manager.storage") as storage:
|
||||
storage.load_once.return_value = b"img"
|
||||
with _patch_session_factory(session):
|
||||
result = manager.get_file_binary_by_message_file_id("mf-1")
|
||||
result = ToolFileManager().get_file_binary_by_message_file_id(message_file.id)
|
||||
|
||||
# Assert
|
||||
assert result == (b"img", "image/png")
|
||||
storage.load_once.assert_called_once_with("k2")
|
||||
|
||||
|
||||
def test_get_file_generator_returns_none_when_toolfile_missing() -> None:
|
||||
# Arrange
|
||||
manager = ToolFileManager()
|
||||
session = Mock()
|
||||
session.scalar.return_value = None
|
||||
def test_get_file_generator_returns_none_when_toolfile_missing(sqlite_tool_file_session: Session) -> None:
|
||||
stream, tool_file = ToolFileManager().get_file_generator_by_tool_file_id(str(uuid4()))
|
||||
|
||||
# Act
|
||||
with _patch_session_factory(session):
|
||||
stream, tool_file = manager.get_file_generator_by_tool_file_id("tool123")
|
||||
|
||||
# Assert
|
||||
assert stream is None
|
||||
assert tool_file is None
|
||||
|
||||
|
||||
def test_get_file_generator_returns_stream_when_found() -> None:
|
||||
# Arrange
|
||||
manager = ToolFileManager()
|
||||
tool_file = SimpleNamespace(
|
||||
id="tool123",
|
||||
file_key="k2",
|
||||
mimetype="image/png",
|
||||
original_url=None,
|
||||
name="image.png",
|
||||
size=12,
|
||||
)
|
||||
session = Mock()
|
||||
session.scalar.return_value = tool_file
|
||||
def test_get_file_generator_returns_stream_when_found(sqlite_tool_file_session: Session) -> None:
|
||||
tool_file = _tool_file(file_key="k2", mimetype="image/png", name="image.png")
|
||||
sqlite_tool_file_session.add(tool_file)
|
||||
sqlite_tool_file_session.commit()
|
||||
|
||||
# Act
|
||||
with patch("core.tools.tool_file_manager.storage") as storage:
|
||||
stream = iter([b"a", b"b"])
|
||||
storage.load_stream.return_value = stream
|
||||
with _patch_session_factory(session):
|
||||
result_stream, result_file = manager.get_file_generator_by_tool_file_id("tool123")
|
||||
assert list(result_stream) == [b"a", b"b"]
|
||||
assert result_file is not None
|
||||
assert result_file.related_id == "tool123"
|
||||
assert result_file.mime_type == "image/png"
|
||||
assert result_file.transfer_method == FileTransferMethod.TOOL_FILE
|
||||
storage.load_stream.return_value = iter([b"a", b"b"])
|
||||
result_stream, result_file = ToolFileManager().get_file_generator_by_tool_file_id(tool_file.id)
|
||||
|
||||
assert result_stream is not None
|
||||
assert list(result_stream) == [b"a", b"b"]
|
||||
assert result_file is not None
|
||||
assert result_file.related_id == tool_file.id
|
||||
assert result_file.mime_type == "image/png"
|
||||
assert result_file.transfer_method == FileTransferMethod.TOOL_FILE
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from decimal import Decimal
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@@ -523,24 +524,36 @@ def test_success_output_adapter_maps_backend_usage_to_llm_usage_and_metadata():
|
||||
source_event_id="2-0",
|
||||
output={"summary": "ok"},
|
||||
session_snapshot=CompositorSessionSnapshot(layers=[]),
|
||||
usage={
|
||||
"prompt_tokens": 10,
|
||||
"prompt_unit_price": "5",
|
||||
"prompt_price_unit": "0.000001",
|
||||
"prompt_price": "0.000050",
|
||||
"completion_tokens": 5,
|
||||
"completion_unit_price": "30",
|
||||
"completion_price_unit": "0.000001",
|
||||
"completion_price": "0.000150",
|
||||
"total_tokens": 15,
|
||||
"total_price": "0.000200",
|
||||
"currency": "USD",
|
||||
"latency": 0.5,
|
||||
},
|
||||
),
|
||||
inputs={},
|
||||
process_data={},
|
||||
metadata={
|
||||
"agent_backend": {
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
}
|
||||
}
|
||||
},
|
||||
metadata={},
|
||||
)
|
||||
|
||||
assert result.llm_usage.prompt_tokens == 10
|
||||
assert result.llm_usage.completion_tokens == 5
|
||||
assert result.llm_usage.total_tokens == 15
|
||||
assert result.llm_usage.prompt_price == Decimal("0.000050")
|
||||
assert result.llm_usage.completion_price == Decimal("0.000150")
|
||||
assert result.llm_usage.total_price == Decimal("0.000200")
|
||||
assert result.llm_usage.currency == "USD"
|
||||
assert result.metadata[WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS] == 15
|
||||
assert result.metadata[WorkflowNodeExecutionMetadataKey.TOTAL_PRICE] == Decimal("0.000200")
|
||||
assert result.metadata[WorkflowNodeExecutionMetadataKey.CURRENCY] == "USD"
|
||||
|
||||
|
||||
def test_failure_output_adapter_maps_cancelled_to_failure_code():
|
||||
|
||||
@@ -0,0 +1,356 @@
|
||||
"""Tests for pinned KnowledgeFS declaration validation."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import pytest
|
||||
|
||||
from dev import generate_knowledge_fs_contract as contract_validator
|
||||
from dev.generate_knowledge_fs_contract import ContractDeclaration, validate_declarations
|
||||
from services.knowledge_fs_proxy import KNOWLEDGE_FS_CONSOLE_OPERATIONS, KnowledgeFSOperation
|
||||
|
||||
|
||||
def test_contract_cli_updates_checks_and_detects_openapi_drift(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
repository = tmp_path / "knowledge-fs"
|
||||
repository.mkdir()
|
||||
subprocess.run(["git", "init", "--quiet"], cwd=repository, check=True)
|
||||
subprocess.run(
|
||||
[
|
||||
"git",
|
||||
"-c",
|
||||
"user.email=contract-test@example.com",
|
||||
"-c",
|
||||
"user.name=Contract Test",
|
||||
"commit",
|
||||
"--allow-empty",
|
||||
"--quiet",
|
||||
"-m",
|
||||
"fixture",
|
||||
],
|
||||
cwd=repository,
|
||||
check=True,
|
||||
)
|
||||
commit = subprocess.run(
|
||||
["git", "rev-parse", "HEAD"], cwd=repository, check=True, capture_output=True, text=True
|
||||
).stdout.strip()
|
||||
|
||||
document = console_registry_document()
|
||||
executable_directory = tmp_path / "bin"
|
||||
executable_directory.mkdir()
|
||||
fake_pnpm = executable_directory / "pnpm"
|
||||
write_fake_pnpm(fake_pnpm, document)
|
||||
|
||||
lock_path = tmp_path / "knowledge-fs-contract.lock.json"
|
||||
lock_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"commit": "",
|
||||
"openapiSha256": "",
|
||||
"repository": "https://github.com/langgenius/knowledge-fs",
|
||||
}
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(contract_validator, "LOCK_PATH", lock_path)
|
||||
monkeypatch.setenv("PATH", f"{executable_directory}{os.pathsep}{os.environ['PATH']}")
|
||||
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
["generate_knowledge_fs_contract.py", "--repository", str(repository), "--update-lock"],
|
||||
)
|
||||
contract_validator.main()
|
||||
|
||||
updated_lock = json.loads(lock_path.read_text())
|
||||
assert updated_lock["commit"] == commit
|
||||
assert set(updated_lock) == {"commit", "openapiSha256", "repository"}
|
||||
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
["generate_knowledge_fs_contract.py", "--repository", str(repository), "--check"],
|
||||
)
|
||||
contract_validator.main()
|
||||
|
||||
with monkeypatch.context() as registry_drift:
|
||||
registry_drift.setattr(
|
||||
contract_validator,
|
||||
"console_contract_declarations",
|
||||
lambda: (declaration(method="DELETE"),),
|
||||
raising=False,
|
||||
)
|
||||
with pytest.raises(ValueError, match="listKnowledgeSpaces.*method.*expected.*received"):
|
||||
contract_validator.main()
|
||||
|
||||
write_fake_pnpm(fake_pnpm, {"paths": {}})
|
||||
with pytest.raises(RuntimeError, match="OpenAPI hash mismatch"):
|
||||
contract_validator.main()
|
||||
|
||||
|
||||
def test_contract_script_loads_runtime_registry_outside_api_directory(tmp_path: Path) -> None:
|
||||
script_path = Path(contract_validator.__file__).resolve()
|
||||
command = (
|
||||
"import runpy; "
|
||||
f"namespace = runpy.run_path({str(script_path)!r}); "
|
||||
"print(len(namespace['console_contract_declarations']()))"
|
||||
)
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", command],
|
||||
cwd=tmp_path,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
assert result.stdout.strip() == "2"
|
||||
|
||||
|
||||
def test_validate_declarations_accepts_matching_contract() -> None:
|
||||
route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
route["parameters"] = [{"in": "header", "name": "X-Trace-Id"}]
|
||||
route["responses"] = {
|
||||
"200": {
|
||||
"content": {"application/json": {}},
|
||||
"headers": {"X-Trace-Id": {}},
|
||||
}
|
||||
}
|
||||
document = {"paths": {"/knowledge-spaces": {"get": route}}}
|
||||
|
||||
validate_declarations(
|
||||
document,
|
||||
(
|
||||
declaration(
|
||||
request_headers=("x-trace-id",),
|
||||
response_headers=("x-trace-id",),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_console_operation_registry_matches_contract() -> None:
|
||||
list_route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
create_route = operation("knowledge-spaces:write", "createKnowledgeSpace")
|
||||
for route in (list_route, create_route):
|
||||
route["parameters"] = [{"in": "header", "name": "X-Trace-Id"}]
|
||||
route["responses"] = {
|
||||
"200": {
|
||||
"content": {"application/json": {}},
|
||||
"headers": {"X-Trace-Id": {}},
|
||||
}
|
||||
}
|
||||
|
||||
validate_declarations(
|
||||
{
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": list_route,
|
||||
"post": create_route,
|
||||
}
|
||||
}
|
||||
},
|
||||
tuple(_contract_declaration(operation) for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS),
|
||||
)
|
||||
|
||||
|
||||
def console_registry_document() -> dict[str, object]:
|
||||
list_route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
create_route = operation("knowledge-spaces:write", "createKnowledgeSpace")
|
||||
for route in (list_route, create_route):
|
||||
route["parameters"] = [{"in": "header", "name": "X-Trace-Id"}]
|
||||
route["responses"] = {
|
||||
"200": {
|
||||
"content": {"application/json": {}},
|
||||
"headers": {"X-Trace-Id": {}},
|
||||
}
|
||||
}
|
||||
return {"paths": {"/knowledge-spaces": {"get": list_route, "post": create_route}}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("method", "POST"),
|
||||
("path", "spaces"),
|
||||
("required_scope", "knowledge-spaces:write"),
|
||||
("response_kind", "stream"),
|
||||
("max_response_bytes", 2_097_152),
|
||||
("request_headers", ("authorization",)),
|
||||
("response_headers", ("cache-control",)),
|
||||
("response_media_types", ("text/event-stream",)),
|
||||
],
|
||||
)
|
||||
def test_validate_declarations_reports_contract_field_drift(field: str, value: object) -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match=rf"listKnowledgeSpaces.*{field}.*expected.*received"):
|
||||
validate_declarations(document, (declaration(**{field: value}),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_unknown_operation_id() -> None:
|
||||
with pytest.raises(ValueError, match="no operationId: listKnowledgeSpaces"):
|
||||
validate_declarations({"paths": {}}, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_duplicate_declared_operation_ids() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="registry has duplicate operationId: listKnowledgeSpaces"):
|
||||
validate_declarations(document, (declaration(), declaration()))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_duplicate_upstream_operation_ids() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
},
|
||||
"/spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="OpenAPI has duplicate operationId: listKnowledgeSpaces"):
|
||||
validate_declarations(document, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_ignores_undeclared_operations() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
},
|
||||
"/internal-maintenance": {
|
||||
"head": {"responses": {"200": {}}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
validate_declarations(document, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_preserves_public_operation_scope() -> None:
|
||||
document = {"paths": {"/health": {"get": operation(None, "getHealth", security=[])}}}
|
||||
|
||||
validate_declarations(
|
||||
document,
|
||||
(
|
||||
declaration(
|
||||
operation_id="getHealth",
|
||||
path="health",
|
||||
required_scope=None,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_unsupported_declared_method() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"/knowledge-spaces": {
|
||||
"head": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="does not support HEAD /knowledge-spaces"):
|
||||
validate_declarations(document, (declaration(method="HEAD"),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_non_absolute_upstream_path() -> None:
|
||||
document = {
|
||||
"paths": {
|
||||
"knowledge-spaces": {
|
||||
"get": operation("knowledge-spaces:read", "listKnowledgeSpaces"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="path must be absolute: knowledge-spaces"):
|
||||
validate_declarations(document, (declaration(),))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [None, True, 0, "1048576"])
|
||||
def test_validate_declarations_rejects_invalid_response_byte_limits(value: object) -> None:
|
||||
route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
route["x-knowledge-fs-max-response-bytes"] = value
|
||||
|
||||
with pytest.raises(ValueError, match="no valid response byte limit"):
|
||||
validate_declarations({"paths": {"/knowledge-spaces": {"get": route}}}, (declaration(),))
|
||||
|
||||
|
||||
def test_validate_declarations_rejects_request_header_references() -> None:
|
||||
route = operation("knowledge-spaces:read", "listKnowledgeSpaces")
|
||||
route["parameters"] = [{"$ref": "#/components/parameters/TraceId"}]
|
||||
|
||||
with pytest.raises(ValueError, match="request header references are not supported"):
|
||||
validate_declarations({"paths": {"/knowledge-spaces": {"get": route}}}, (declaration(),))
|
||||
|
||||
|
||||
def operation(scope: str | None, operation_id: str, **overrides: object) -> dict[str, object]:
|
||||
value: dict[str, object] = {
|
||||
"operationId": operation_id,
|
||||
"responses": {"200": {"content": {"application/json": {}}}},
|
||||
"x-knowledge-fs-max-response-bytes": 1_048_576,
|
||||
}
|
||||
if scope is not None:
|
||||
value["x-knowledge-fs-required-scope"] = scope
|
||||
value.update(overrides)
|
||||
return value
|
||||
|
||||
|
||||
def declaration(**overrides: object) -> ContractDeclaration:
|
||||
value: dict[str, object] = {
|
||||
"operation_id": "listKnowledgeSpaces",
|
||||
"method": "GET",
|
||||
"path": "knowledge-spaces",
|
||||
"required_scope": "knowledge-spaces:read",
|
||||
"response_kind": "buffered",
|
||||
"max_response_bytes": 1_048_576,
|
||||
"request_headers": (),
|
||||
"response_headers": (),
|
||||
"response_media_types": ("application/json",),
|
||||
}
|
||||
value.update(overrides)
|
||||
return cast(ContractDeclaration, value)
|
||||
|
||||
|
||||
def _contract_declaration(operation: KnowledgeFSOperation) -> ContractDeclaration:
|
||||
return {
|
||||
"operation_id": operation.operation_id,
|
||||
"method": operation.method,
|
||||
"path": operation.path,
|
||||
"required_scope": operation.required_scope,
|
||||
"response_kind": operation.response_kind,
|
||||
"max_response_bytes": operation.max_response_bytes,
|
||||
"request_headers": operation.request_headers,
|
||||
"response_headers": operation.response_headers,
|
||||
"response_media_types": operation.response_media_types,
|
||||
}
|
||||
|
||||
|
||||
def write_fake_pnpm(path: Path, document: dict[str, object]) -> None:
|
||||
path.write_text(
|
||||
"#!/usr/bin/env python3\n"
|
||||
"import sys\n"
|
||||
"from pathlib import Path\n"
|
||||
"output = Path(sys.argv[sys.argv.index('--output') + 1])\n"
|
||||
f"output.write_text({json.dumps(document)!r})\n"
|
||||
)
|
||||
path.chmod(0o755)
|
||||
@@ -1,13 +1,12 @@
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from collections.abc import Iterator
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import ANY, patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import create_engine, select
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from core.app.entities.app_invoke_entities import ChatAppGenerateEntity
|
||||
from core.entities.provider_entities import ProviderQuotaType, QuotaUnit
|
||||
@@ -16,40 +15,29 @@ from models import TenantCreditPool
|
||||
from models.provider import ProviderType
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_credit_pool_session_factory(engine: Engine) -> Generator[None, None, None]:
|
||||
session_maker = sessionmaker(bind=engine, expire_on_commit=False)
|
||||
sessions = []
|
||||
|
||||
def _session():
|
||||
session = session_maker()
|
||||
sessions.append(session)
|
||||
return session
|
||||
|
||||
with patch("events.event_handlers.update_provider_when_message_created.db", SimpleNamespace(session=_session)):
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for session in sessions:
|
||||
session.close()
|
||||
@pytest.fixture
|
||||
def credit_pool_session_factory(sqlite_engine: Engine) -> Iterator[sessionmaker[Session]]:
|
||||
"""Bind message-created accounting to fixture-owned SQLite sessions."""
|
||||
TenantCreditPool.__table__.create(sqlite_engine)
|
||||
session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
|
||||
with patch("events.event_handlers.update_provider_when_message_created.db.session", session_factory):
|
||||
yield session_factory
|
||||
|
||||
|
||||
def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_insufficient() -> None:
|
||||
engine = create_engine("sqlite:///:memory:")
|
||||
TenantCreditPool.__table__.create(engine)
|
||||
def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_insufficient(
|
||||
credit_pool_session_factory: sessionmaker[Session],
|
||||
) -> None:
|
||||
tenant_id = str(uuid4())
|
||||
pool_id = str(uuid4())
|
||||
with engine.begin() as connection:
|
||||
connection.execute(
|
||||
TenantCreditPool.__table__.insert(),
|
||||
{
|
||||
"id": pool_id,
|
||||
"tenant_id": tenant_id,
|
||||
"pool_type": ProviderQuotaType.TRIAL,
|
||||
"quota_limit": 10,
|
||||
"quota_used": 9,
|
||||
},
|
||||
)
|
||||
pool = TenantCreditPool(
|
||||
tenant_id=tenant_id,
|
||||
pool_type=ProviderQuotaType.TRIAL,
|
||||
quota_limit=10,
|
||||
quota_used=9,
|
||||
)
|
||||
pool.id = pool_id
|
||||
with credit_pool_session_factory.begin() as session:
|
||||
session.add(pool)
|
||||
|
||||
system_configuration = SimpleNamespace(
|
||||
current_quota_type=ProviderQuotaType.TRIAL,
|
||||
@@ -77,7 +65,6 @@ def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_
|
||||
message = SimpleNamespace(message_tokens=2, answer_tokens=1)
|
||||
|
||||
with (
|
||||
_patched_credit_pool_session_factory(engine),
|
||||
patch.object(update_provider_when_message_created, "_execute_provider_updates"),
|
||||
):
|
||||
update_provider_when_message_created.handle(
|
||||
@@ -85,8 +72,8 @@ def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_
|
||||
application_generate_entity=application_generate_entity,
|
||||
)
|
||||
|
||||
with engine.connect() as connection:
|
||||
quota_used = connection.scalar(select(TenantCreditPool.quota_used).where(TenantCreditPool.id == pool_id))
|
||||
with credit_pool_session_factory() as session:
|
||||
quota_used = session.scalar(select(TenantCreditPool.quota_used).where(TenantCreditPool.id == pool_id))
|
||||
|
||||
assert quota_used == 10
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user