Compare commits

...
Author SHA1 Message Date
yyhandGitHub b95d313128 refactor(web): improve workflow block selector (#39377) 2026-07-21 14:18:38 +00:00
github-actions[bot]GitHubclaude[bot] <41898282+claude[bot]@users.noreply.github.com>yyh
45ac70132c chore(i18n): sync translations with en-US (#39362)
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: yyh <92089059+lyzno1@users.noreply.github.com>
2026-07-21 11:54:20 +00:00
JoelandGitHub 6e772259bb fix: update the agent roster documentation link (#39368) 2026-07-21 09:56:21 +00:00
QuantumGhostandGitHub af465fb527 fix(api): scope HITL snapshot message lookup to utilize db index (#39351) 2026-07-21 09:47:17 +00:00
Xiyuan ChenGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
62bbc0dbeb feat(webapp-access): add WEBAPP_PUBLIC_ACCESS_ENABLED toggle (#38567)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 08:46:28 +00:00
zyssyz123andGitHub ed1289281d fix(agent): include workflow runs in monitoring stats (#39354) 2026-07-21 08:33:23 +00:00
Asuka MinatoandGitHub e37f32e4a0 test: add SQLite session fixture for providers (#38982) 2026-07-21 08:29:44 +00:00
Asuka MinatoandGitHub 150613cc05 test: use sqlite3 session in test_build_from_mapping (#38769) 2026-07-21 08:21:36 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
a5194ff17b test: move human input delivery coverage to unit tests (#38926)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 08:18:32 +00:00
Asuka MinatoandGitHub cdaf5d4562 test: use sqlite3 session in test_tool_file_manager (#38771) 2026-07-21 08:10:23 +00:00
Asuka MinatoandGitHub 145b0e622e test: use sqlite3 session in test_workspace_account (#38757) 2026-07-21 08:07:57 +00:00
Asuka MinatoandGitHub 4bab62d458 test: use sqlite3 session in test_agent_app_parameters (#38759) 2026-07-21 07:59:00 +00:00
Asuka MinatoandGitHub b99f8941bb test: use sqlite3 session in test_database_retrieval (#38743) 2026-07-21 07:56:34 +00:00
Asuka MinatoandGitHub 32d2594877 test: move app import API coverage to unit tests (#38920) 2026-07-21 07:54:32 +00:00
Asuka MinatoandGitHub 7b188287da test: use sqlite3 session in test_accounts (#38752) 2026-07-21 07:50:05 +00:00
Asuka MinatoandGitHub f6db23444d test: move API token service coverage to unit tests (#38918) 2026-07-21 07:48:10 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
a5c989d82b test: use sqlite3 session in test_conversation_service (#38774)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 07:38:53 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2622b1ce0a test: use sqlite3 session in test_based_generate_task_pipeline (#38738)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 07:37:20 +00:00
Asuka MinatoandGitHub 1e672530c4 test: use sqlite3 session in test_built_in_retrieval (#38735) 2026-07-21 07:36:42 +00:00
Asuka MinatoandGitHub 852169d99c test: use SQLite sessions in unit misc (#39118) 2026-07-21 07:27:13 +00:00
Asuka MinatoandGitHub d7860571aa test: use SQLite sessions in services workflow (#39115) 2026-07-21 07:24:52 +00:00
Asuka MinatoandGitHub 3683831738 test: use sqlite3 session in test_account_service (#38751) 2026-07-21 07:22:05 +00:00
非法操作andGitHub f5024f5440 fix(workflow): avoid false draft save error during Strict Mode initia… (#39353) 2026-07-21 07:11:19 +00:00
JoelGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
4f8477427e fix: not support jump to detail in marketplace (#39347)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 07:01:00 +00:00
非法操作andGitHub 8d406aec31 fix(dev): use gevent WebSocket server for local API (#39349) 2026-07-21 06:42:55 +00:00
落尘andGitHub f8754286ec fix(watercrawl): handle non-json error responses (#37514) 2026-07-21 06:41:06 +00:00
yyhandGitHub 7da6b2ce36 refactor(web): remove permission catalog service hooks (#39342) 2026-07-21 06:01:00 +00:00
aa718add24 fix(api): return 400 instead of 500 when 'file' is missing from trial-app audio-to-text (#39305)
Co-authored-by: Harsh Kashyap <Harsh23Kashyap@users.noreply.github.com>
2026-07-21 05:59:48 +00:00
yyhandGitHub fcc1d39f5e fix(web): restore billing access boundary (#39343) 2026-07-21 05:35:36 +00:00
yyhandGitHub 861831f176 fix(api): keep workflow log archives owner-admin only (#39338) 2026-07-21 04:31:14 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
d333aa31ea test: use SQLite sessions in services workflow (#39116)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 03:40:02 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
97e6a14c13 test: use SQLite sessions in unit misc (#39119)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 03:20:59 +00:00
shaktiandGitHub 6a338654fb fix: return 400 instead of 500 when file field is missing in audio-to-text endpoints (#39303) (#39322) 2026-07-21 02:19:24 +00:00
ChecoandGitHub 8f075c82bc fix(configs): reject host:port-shaped PLUGIN_REMOTE_INSTALL_PORT with actionable hint (#39329)
Signed-off-by: sergioperezcheco <checo520@outlook.com>
2026-07-21 02:18:29 +00:00
林玮 (Jade Lin)andGitHub 28d17da3b1 fix(api): use resource tenant for draft variable files (#39307) 2026-07-21 02:07:47 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
891b2dc537 test: use SQLite sessions in unit misc (#39120)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-21 01:51:24 +00:00
WH-2099andGitHub 75f7069541 ci: simplify test workflow execution (#39318) 2026-07-21 00:18:59 +00:00
yyhandGitHub b7193d1cba refactor(web): remove basePath from hydration boundary (#39308) 2026-07-20 14:46:37 +00:00
wangxiaoleiGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
ef0115d340 chore: when dataset permission is not all team, update rbac config to… (#39227)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-20 12:59:29 +00:00
3afb9b3230 fix(web): stop doubling basePath in auth refresh redirects (#39273)
Co-authored-by: yyh <yuanyouhuilyz@gmail.com>
2026-07-20 12:22:40 +00:00
JoelandGitHub 0862641533 feat: support export agent dsl in sidebar (#39299) 2026-07-20 10:41:26 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
cbd7872520 test: use SQLite sessions in unit misc (#39121)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-20 10:38:52 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
23882a704e test: use SQLite sessions in unit misc (#39122)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-20 10:38:49 +00:00
Asuka MinatoGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
1bd5254ea7 test: use SQLite sessions in unit misc (#39123)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-20 10:38:46 +00:00
yyhandGitHub e13069dbe3 revert(ci): restore default Depot runners (#39297) 2026-07-20 10:18:52 +00:00
github-actions[bot]GitHubclaude[bot] <41898282+claude[bot]@users.noreply.github.com>yyh
edd3d4a218 chore(i18n): sync translations with en-US (#39296)
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: yyh <92089059+lyzno1@users.noreply.github.com>
2026-07-20 10:15:37 +00:00
JoelandGitHub 79751f7622 chore: display non-LLM settings in integration tool details (#39295) 2026-07-20 09:52:54 +00:00
JingyiGitHubhjlarryautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>yyhyyh
cb64446fa3 feat(web): add step-by-step tour shell (#38785)
Co-authored-by: hjlarry <hjlarry@163.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: yyh <yuanyouhuilyz@gmail.com>
Co-authored-by: yyh <92089059+lyzno1@users.noreply.github.com>
2026-07-20 09:23:25 +00:00
JoelandGitHub e77ee82526 fix: show persistent error when workflow draft save fails (#39293) 2026-07-20 08:48:38 +00:00
yyhandGitHub b3f163fb0f chore(deps): update pnpm and workspace dependencies (#39292) 2026-07-20 08:43:21 +00:00
yyhandGitHub 873ef2b592 docs: refine frontend testing guidance (#39290) 2026-07-20 08:18:26 +00:00
Stephen ZhouGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
43d725a1ce feat(dataset): proxy KnowledgeFS Console requests (#39158)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-20 08:10:06 +00:00
JoelandGitHub c8ccfba960 fix: the input of agent chat would be replaced if conversation history updated (#39289) 2026-07-20 08:07:12 +00:00
eee8d6cf7b feat(web): add webapp access control to agent access point (#39285)
Co-authored-by: yyh <yuanyouhuilyz@gmail.com>
2026-07-20 08:05:29 +00:00
zyssyz123andGitHub 11016563b1 fix(agent): preserve complete model usage pricing (#39201) 2026-07-20 08:03:13 +00:00
Ingram ZandGitHub b7edc5ef0c refactor: use catalog for workflow archive maintenance (#39001) 2026-07-20 07:59:29 +00:00
林玮 (Jade Lin)GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
f93a5b95c6 fix(api, web): scope trial app file uploads to the app tenant (#39149)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-20 07:00:45 +00:00
Kevin9703GitHubhjlarryautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>yyhyyh
6022939bf0 fix: prevent hidden-tab collaboration leader from saving stale drafts (#38997)
Co-authored-by: hjlarry <hjlarry@163.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: yyh <yuanyouhuilyz@gmail.com>
Co-authored-by: yyh <92089059+lyzno1@users.noreply.github.com>
2026-07-20 06:58:01 +00:00
非法操作andGitHub 30df1433cc fix: can't create chat app when get tenant default model schema error (#39266) 2026-07-20 06:47:16 +00:00
dependabot[bot]GitHubdependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
aa7b65c602 chore: bump the github-actions-dependencies group with 5 updates (#39247)
Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-07-20 06:47:01 +00:00
yyhandGitHub ba24745c07 fix(web): restore permission selector semantics (#39269)
Signed-off-by: yyh <yuanyouhuilyz@gmail.com>
2026-07-20 06:30:09 +00:00
38ea0de907 refactor(web): avoid direct setState in effect in jina reader crawler (#25194) (#39216)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Wu Tianwei <30284043+WTW0313@users.noreply.github.com>
2026-07-20 06:28:08 +00:00
JoelandGitHub bbc31a3403 fix: text generation web app keeps loading after timeout (#39268) 2026-07-20 06:27:32 +00:00
605c0fe02b refactor(web): avoid setState in website crawler effects (#39237)
Co-authored-by: Wu Tianwei <30284043+WTW0313@users.noreply.github.com>
2026-07-20 06:26:50 +00:00
yyhandGitHub 15674a9db6 test(web): align test state boundaries (#39239) 2026-07-20 06:06:17 +00:00
非法操作GitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
9a2d039e93 chore: the node should not draggable when preview a workflow (#39131)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-20 05:59:14 +00:00
yyhandGitHub 58d1dd8873 ci: route default Depot jobs to 4-core runners (#39277) 2026-07-20 05:42:56 +00:00
ChecoandGitHub 5ea884f799 fix(dataset_service): use flush() instead of commit() to preserve caller transaction (#39223) 2026-07-19 07:00:50 +00:00
yyhGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2ec34b2cfb feat(workflow): improve block selector navigation and previews (#39212)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-18 09:56:36 +00:00
fa65f45bff refactor(web): align integrations sidebar components (#39202)
Co-authored-by: Jingyi <jingyi.qi@dify.ai>
2026-07-17 22:52:47 +00:00
yyhGitHubautofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
63ca2b94b5 refactor(web): use generated OAuth query options (#39204)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-17 14:50:56 +00:00
yyhandGitHub a7ef41a8f5 refactor(web): use generated dataset card queries (#39203) 2026-07-17 14:50:40 +00:00
941 changed files with 38019 additions and 20178 deletions
+2 -3
View File
@@ -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.
-3
View File
@@ -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 \
+1 -2
View File
@@ -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/**'
+4 -4
View File
@@ -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
+1 -1
View File
@@ -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: ''
+1 -1
View File
@@ -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 }}
-10
View File
@@ -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
-2
View File
@@ -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:
+2
View File
@@ -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 \
+8
View File
@@ -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
View File
@@ -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]:
+6
View File
@@ -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,
+2
View File
@@ -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,
):
+64
View File
@@ -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
+33 -1
View File
@@ -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):
"""
+4
View File
@@ -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",
+2
View File
@@ -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()
+8 -1
View File
@@ -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:
+48 -3
View File
@@ -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",
+41 -35
View File
@@ -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)
+106
View File
@@ -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,
),
)
+57 -47
View File
@@ -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)
+10 -27
View File
@@ -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:
+1 -1
View File
@@ -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(
+16 -8
View File
@@ -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
+1 -1
View File
@@ -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)
+10 -1
View File
@@ -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(
+5 -1
View File
@@ -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(
+91 -2
View File
@@ -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)
+4 -1
View File
@@ -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)
+255
View File
@@ -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()
+5
View File
@@ -0,0 +1,5 @@
{
"commit": "4310e2d582d25e7de58183f27720afab01e123cf",
"openapiSha256": "5827ca930ce38462bfd1b2bef387efbf37eb7ffcaedde4558af2fbaeccbfbc4b",
"repository": "https://github.com/langgenius/knowledge-fs"
}
+33 -12
View File
@@ -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")
@@ -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)
+2
View File
@@ -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",
+59
View File
@@ -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(),
)
+2
View File
@@ -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)
+87
View File
@@ -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 |
+2
View File
@@ -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 |
+35
View File
@@ -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)
+227 -18
View File
@@ -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
View File
@@ -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
+15 -4
View File
@@ -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)
-13
View File
@@ -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,
)
+4
View File
@@ -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]:
+358
View File
@@ -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)]
+221
View File
@@ -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
+161 -14
View File
@@ -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:
+11 -10
View File
@@ -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())
+90 -11
View File
@@ -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:
+11 -5
View File
@@ -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
View File
@@ -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
@@ -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"
@@ -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