Compare commits
49
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5c6372d2f7 | ||
|
|
5fd06fafe0 | ||
|
|
af4b65b295 | ||
|
|
6e5fc1081b | ||
|
|
0e84ae7338 | ||
|
|
5a792945f5 | ||
|
|
2f3c785f27 | ||
|
|
38045078dc | ||
|
|
83f40b85f4 | ||
|
|
aff8a49bbf | ||
|
|
b0329cd53c | ||
|
|
2a008423d8 | ||
|
|
f962c9e47a | ||
|
|
af7e59de7c | ||
|
|
61b07ab17d | ||
|
|
a7aff83d52 | ||
|
|
9de7e0fe44 | ||
|
|
50341357b3 | ||
|
|
96e34e7b24 | ||
|
|
1aea3460af | ||
|
|
48e536ba39 | ||
|
|
33fe0dfd60 | ||
|
|
a8ffc93a5a | ||
|
|
14ac8dfbf6 | ||
|
|
b820ccf086 | ||
|
|
415c0db22e | ||
|
|
6ad1a66a50 | ||
|
|
2f2d7c26d4 | ||
|
|
18f4640ad1 | ||
|
|
f4ad6ce978 | ||
|
|
e9c0e9de9e | ||
|
|
8a33161080 | ||
|
|
752af8a270 | ||
|
|
678ce2ab37 | ||
|
|
b737833e2a | ||
|
|
872b6906f2 | ||
|
|
d876f6dba5 | ||
|
|
cd98193234 | ||
|
|
bfffcb7d0f | ||
|
|
d66de8a47b | ||
|
|
1a2ba2237d | ||
|
|
8ff92c9ef8 | ||
|
|
511dbb9974 | ||
|
|
4381ec8fee | ||
|
|
cb41fb3e76 | ||
|
|
7dfd84472f | ||
|
|
62bdbc8628 | ||
|
|
85cc183501 | ||
|
|
5af8f6af4d |
@@ -9,4 +9,3 @@
|
||||
# Codegen output must stay byte-identical across platforms so
|
||||
# `pnpm tree:check` in CI does not trip on CRLF rewrites.
|
||||
*.generated.ts text eol=lf
|
||||
*.gen.ts text eol=lf
|
||||
|
||||
@@ -72,6 +72,7 @@ jobs:
|
||||
- 'docker/volumes/sandbox/conf/**'
|
||||
cli:
|
||||
- 'cli/**'
|
||||
- 'packages/contracts/**'
|
||||
- 'packages/tsconfig/**'
|
||||
- 'package.json'
|
||||
- 'pnpm-lock.yaml'
|
||||
@@ -105,6 +106,7 @@ jobs:
|
||||
- 'docker/docker-compose.middleware.yaml'
|
||||
- 'docker/envs/middleware.env.example'
|
||||
- '.github/workflows/web-e2e.yml'
|
||||
- '.github/workflows/main-ci.yml'
|
||||
- '.github/actions/setup-web/**'
|
||||
vdb:
|
||||
- 'api/core/rag/datasource/**'
|
||||
|
||||
@@ -19,6 +19,7 @@ jobs:
|
||||
test:
|
||||
name: Web Full-Stack E2E
|
||||
runs-on: depot-ubuntu-24.04-4
|
||||
timeout-minutes: 120
|
||||
defaults:
|
||||
run:
|
||||
shell: bash
|
||||
@@ -62,7 +63,52 @@ jobs:
|
||||
E2E_INIT_PASSWORD: E2eInit12345
|
||||
run: vp run e2e:full
|
||||
|
||||
- name: Run external runtime E2E tests
|
||||
- name: Preserve Chromium E2E report and logs
|
||||
if: ${{ !cancelled() }}
|
||||
run: |
|
||||
if [[ -d e2e/cucumber-report ]]; then
|
||||
mv e2e/cucumber-report e2e/cucumber-report-non-external
|
||||
fi
|
||||
if [[ -d e2e/.logs ]]; then
|
||||
mv e2e/.logs e2e/.logs-non-external
|
||||
fi
|
||||
|
||||
- name: Run WebKit keyboard and browser smoke tests
|
||||
working-directory: ./e2e
|
||||
env:
|
||||
E2E_ADMIN_EMAIL: e2e-admin@example.com
|
||||
E2E_ADMIN_NAME: E2E Admin
|
||||
E2E_ADMIN_PASSWORD: E2eAdmin12345
|
||||
E2E_BROWSER: webkit
|
||||
E2E_INIT_PASSWORD: E2eInit12345
|
||||
run: |
|
||||
teardown_webkit_smoke() {
|
||||
local run_status=$?
|
||||
trap - EXIT
|
||||
if ! vp run e2e:middleware:down; then
|
||||
echo "::error title=E2E teardown failed::WebKit smoke middleware did not shut down cleanly."
|
||||
if [[ "$run_status" -eq 0 ]]; then
|
||||
run_status=1
|
||||
fi
|
||||
fi
|
||||
exit "$run_status"
|
||||
}
|
||||
|
||||
trap teardown_webkit_smoke EXIT
|
||||
vp run e2e:middleware:up
|
||||
vp run e2e -- --tags '@browser-smoke'
|
||||
|
||||
- name: Preserve WebKit E2E report and logs
|
||||
if: ${{ !cancelled() }}
|
||||
run: |
|
||||
if [[ -d e2e/cucumber-report ]]; then
|
||||
mv e2e/cucumber-report e2e/cucumber-report-webkit
|
||||
fi
|
||||
if [[ -d e2e/.logs ]]; then
|
||||
mv e2e/.logs e2e/.logs-webkit
|
||||
fi
|
||||
|
||||
- name: Run prepared and external runtime E2E tests
|
||||
if: ${{ inputs.run-external-runtime }}
|
||||
working-directory: ./e2e
|
||||
env:
|
||||
@@ -72,13 +118,10 @@ jobs:
|
||||
E2E_AGENT_DECISION_MODEL_NAME: ${{ vars.E2E_AGENT_DECISION_MODEL_NAME || 'gpt-5.5' }}
|
||||
E2E_AGENT_DECISION_MODEL_PROVIDER: ${{ vars.E2E_AGENT_DECISION_MODEL_PROVIDER || 'openai' }}
|
||||
E2E_AGENT_DECISION_MODEL_TYPE: ${{ vars.E2E_AGENT_DECISION_MODEL_TYPE || 'llm' }}
|
||||
E2E_EXTERNAL_RUNTIME_SEED_SPECS: ${{ vars.E2E_EXTERNAL_RUNTIME_SEED_SPECS }}
|
||||
E2E_EXTERNAL_RUNTIME_TAGS: ${{ vars.E2E_EXTERNAL_RUNTIME_TAGS }}
|
||||
E2E_FORCE_WEB_BUILD: "1"
|
||||
E2E_INIT_PASSWORD: E2eInit12345
|
||||
E2E_MARKETPLACE_API_URL: ${{ vars.E2E_MARKETPLACE_API_URL }}
|
||||
E2E_MARKETPLACE_PLUGIN_IDS: ${{ vars.E2E_MARKETPLACE_PLUGIN_IDS }}
|
||||
E2E_MARKETPLACE_PLUGIN_UNIQUE_IDENTIFIERS: ${{ vars.E2E_MARKETPLACE_PLUGIN_UNIQUE_IDENTIFIERS }}
|
||||
E2E_MODEL_PROVIDER_CREDENTIALS_JSON: ${{ secrets.E2E_MODEL_PROVIDER_CREDENTIALS_JSON }}
|
||||
E2E_SPEECH_TO_TEXT_MODEL_NAME: ${{ vars.E2E_SPEECH_TO_TEXT_MODEL_NAME || 'gpt-4o-mini-transcribe' }}
|
||||
E2E_SPEECH_TO_TEXT_MODEL_PROVIDER: ${{ vars.E2E_SPEECH_TO_TEXT_MODEL_PROVIDER || 'openai' }}
|
||||
@@ -102,10 +145,22 @@ jobs:
|
||||
mv .logs .logs-non-external
|
||||
fi
|
||||
|
||||
trap 'vp run e2e:middleware:down' EXIT
|
||||
teardown_external_runtime() {
|
||||
local run_status=$?
|
||||
trap - EXIT
|
||||
if ! vp run e2e:middleware:down; then
|
||||
echo "::error title=E2E teardown failed::External runtime middleware did not shut down cleanly."
|
||||
if [[ "$run_status" -eq 0 ]]; then
|
||||
run_status=1
|
||||
fi
|
||||
fi
|
||||
exit "$run_status"
|
||||
}
|
||||
|
||||
trap teardown_external_runtime EXIT
|
||||
vp run e2e:middleware:up
|
||||
vp run e2e:external:prepare
|
||||
vp run e2e:external
|
||||
vp run e2e:post-merge:prepare
|
||||
vp run e2e:post-merge
|
||||
|
||||
- name: Upload Cucumber report
|
||||
if: ${{ !cancelled() }}
|
||||
@@ -115,6 +170,7 @@ jobs:
|
||||
path: |
|
||||
e2e/cucumber-report
|
||||
e2e/cucumber-report-non-external
|
||||
e2e/cucumber-report-webkit
|
||||
retention-days: 7
|
||||
|
||||
- name: Upload E2E logs
|
||||
@@ -125,5 +181,15 @@ jobs:
|
||||
path: |
|
||||
e2e/.logs/*.log
|
||||
e2e/.logs-non-external/*.log
|
||||
e2e/.logs-webkit/*.log
|
||||
include-hidden-files: true
|
||||
retention-days: 7
|
||||
|
||||
- name: Upload E2E seed report
|
||||
if: ${{ !cancelled() && inputs.run-external-runtime }}
|
||||
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
|
||||
with:
|
||||
name: e2e-seed-report
|
||||
path: e2e/seed-report
|
||||
if-no-files-found: ignore
|
||||
retention-days: 7
|
||||
|
||||
Vendored
+4
-18
@@ -5,7 +5,9 @@
|
||||
"name": "Python: API (gevent)",
|
||||
"type": "debugpy",
|
||||
"request": "launch",
|
||||
"program": "${workspaceFolder}/api/app.py",
|
||||
"module": "gevent.monkey",
|
||||
"args": ["--module", "app"],
|
||||
"gevent": true,
|
||||
"jinja": true,
|
||||
"justMyCode": true,
|
||||
"cwd": "${workspaceFolder}/api",
|
||||
@@ -33,22 +35,6 @@
|
||||
"justMyCode": false,
|
||||
"cwd": "${workspaceFolder}/api",
|
||||
"python": "${workspaceFolder}/api/.venv/bin/python"
|
||||
},
|
||||
{
|
||||
"name": "Next.js: debug full stack",
|
||||
"type": "node",
|
||||
"request": "launch",
|
||||
"program": "${workspaceFolder}/web/node_modules/next/dist/bin/next",
|
||||
"runtimeArgs": ["--inspect"],
|
||||
"skipFiles": ["<node_internals>/**"],
|
||||
"serverReadyAction": {
|
||||
"action": "debugWithChrome",
|
||||
"killOnServerStop": true,
|
||||
"pattern": "- Local:.+(https?://.+)",
|
||||
"uriFormat": "%s",
|
||||
"webRoot": "${workspaceFolder}/web"
|
||||
},
|
||||
"cwd": "${workspaceFolder}/web"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
<a href="https://discord.gg/FngNHpbcY7" target="_blank">
|
||||
<img src="https://img.shields.io/discord/1082486657678311454?logo=discord&labelColor=%20%235462eb&logoColor=%20%23f5f5f5&color=%20%235462eb"
|
||||
alt="chat on Discord"></a>
|
||||
<a href="https://reddit.com/r/difyai" target="_blank">
|
||||
<a href="https://reddit.com/r/difyai" target="_blank">
|
||||
<img src="https://img.shields.io/reddit/subreddit-subscribers/difyai?style=plastic&logo=reddit&label=r%2Fdifyai&labelColor=white"
|
||||
alt="join Reddit"></a>
|
||||
<a href="https://twitter.com/intent/follow?screen_name=dify_ai" target="_blank">
|
||||
|
||||
+11
-7
@@ -561,6 +561,8 @@ WORKFLOW_MAX_EXECUTION_STEPS=500
|
||||
WORKFLOW_MAX_EXECUTION_TIME=1200
|
||||
WORKFLOW_CALL_MAX_DEPTH=5
|
||||
MAX_VARIABLE_SIZE=204800
|
||||
# Maximum concurrent node-builder LLM calls per workflow generation request
|
||||
WORKFLOW_GENERATOR_NODE_BUILDER_MAX_WORKERS=6
|
||||
|
||||
# GraphEngine Worker Pool Configuration
|
||||
# Minimum number of workers per GraphEngine instance (default: 1)
|
||||
@@ -665,17 +667,19 @@ PLUGIN_REMOTE_INSTALL_HOST=localhost
|
||||
PLUGIN_MAX_PACKAGE_SIZE=15728640
|
||||
PLUGIN_MODEL_SCHEMA_CACHE_TTL=3600
|
||||
PLUGIN_MODEL_PROVIDERS_CACHE_TTL=86400
|
||||
# Comma-separated marketplace plugin IDs whose latest versions are installed for newly registered users.
|
||||
# Example: langgenius/openai,langgenius/gemini
|
||||
NEW_USER_DEFAULT_PLUGIN_IDS=
|
||||
# Comma-separated model_type:provider:model entries assigned after default plugins finish installing.
|
||||
# Example: llm:langgenius/openai/openai:gpt-4o-mini,text-embedding:langgenius/openai/openai:text-embedding-3-small
|
||||
NEW_USER_DEFAULT_MODELS=
|
||||
INNER_API_KEY_FOR_PLUGIN=QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1
|
||||
|
||||
# Dify Agent backend
|
||||
AGENT_BACKEND_BASE_URL=http://localhost:5050
|
||||
|
||||
# KnowledgeFS (Dataset 2.0)
|
||||
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
|
||||
AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS=30
|
||||
AGENT_BACKEND_STREAM_MAX_RECONNECTS=3
|
||||
AGENT_BACKEND_RUN_TIMEOUT_SECONDS=1200
|
||||
|
||||
# Marketplace configuration
|
||||
MARKETPLACE_ENABLED=true
|
||||
|
||||
+19
@@ -1,5 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
# ``python -m app`` (docker DEBUG=true, or IDE debugging) serves through the
|
||||
# gevent pywsgi server at the bottom of this file, so the stdlib must be
|
||||
# monkey-patched BEFORE any other import pulls in sockets or locks. Without
|
||||
# this, every request runs as a greenlet on one OS thread while blocking
|
||||
# calls (LLM invokes, ``Future.result`` waits, DB I/O) pin that thread — the
|
||||
# whole process freezes until the call returns. Gunicorn and Celery apply
|
||||
# their own patching (see gunicorn.conf.py / celery_entrypoint.py), and
|
||||
# ``flask run`` uses real Werkzeug threads, so both skip this branch.
|
||||
if __name__ == "__main__":
|
||||
from gevent import monkey
|
||||
|
||||
monkey.patch_all()
|
||||
|
||||
import psycogreen.gevent as psycogreen_gevent
|
||||
from grpc.experimental import gevent as grpc_gevent
|
||||
|
||||
grpc_gevent.init_gevent()
|
||||
psycogreen_gevent.patch_psycopg()
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
@@ -8,7 +8,7 @@ creating another wire contract.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Protocol
|
||||
|
||||
from dify_agent.client import (
|
||||
@@ -45,7 +45,13 @@ class AgentBackendRunClient(Protocol):
|
||||
def cancel_run(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse:
|
||||
"""Request explicit cancellation for one Agent backend run."""
|
||||
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
"""Yield public ``dify-agent`` run events in stream order."""
|
||||
|
||||
def wait_run(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse:
|
||||
@@ -61,7 +67,15 @@ class _DifyAgentSyncClient(Protocol):
|
||||
def cancel_run_sync(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse:
|
||||
"""Cancel one run synchronously."""
|
||||
|
||||
def stream_events_sync(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
def stream_events_sync(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
max_reconnects: int | None = None,
|
||||
timeout_seconds: float | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
"""Stream run events synchronously."""
|
||||
|
||||
def wait_run_sync(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse:
|
||||
@@ -73,8 +87,16 @@ class DifyAgentBackendRunClient:
|
||||
|
||||
client: _DifyAgentSyncClient
|
||||
|
||||
def __init__(self, client: _DifyAgentSyncClient) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
client: _DifyAgentSyncClient,
|
||||
*,
|
||||
stream_max_reconnects: int = 3,
|
||||
stream_timeout_seconds: float = 1200,
|
||||
) -> None:
|
||||
self.client = client
|
||||
self._stream_max_reconnects = stream_max_reconnects
|
||||
self._stream_timeout_seconds = stream_timeout_seconds
|
||||
|
||||
def create_run(self, request: CreateRunRequest) -> CreateRunResponse:
|
||||
"""Create one run through ``POST /runs`` and normalize client exceptions."""
|
||||
@@ -90,10 +112,22 @@ class DifyAgentBackendRunClient:
|
||||
except Exception as exc:
|
||||
raise _normalize_dify_agent_error(exc) from exc
|
||||
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
"""Stream run events from ``/events/sse`` with the wrapped client's reconnect policy."""
|
||||
try:
|
||||
yield from self.client.stream_events_sync(run_id, after=after)
|
||||
yield from self.client.stream_events_sync(
|
||||
run_id,
|
||||
after=after,
|
||||
max_reconnects=self._stream_max_reconnects,
|
||||
timeout_seconds=self._stream_timeout_seconds,
|
||||
should_stop=should_stop,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise _normalize_dify_agent_error(exc) from exc
|
||||
|
||||
|
||||
@@ -13,10 +13,17 @@ def create_agent_backend_run_client(
|
||||
base_url: str | None = None,
|
||||
use_fake: bool = False,
|
||||
fake_scenario: str | FakeAgentBackendScenario = FakeAgentBackendScenario.SUCCESS,
|
||||
stream_read_timeout_seconds: float = 30,
|
||||
stream_max_reconnects: int = 3,
|
||||
stream_run_timeout_seconds: float = 1200,
|
||||
) -> AgentBackendRunClient:
|
||||
"""Create the API-side run client without hiding the ``dify-agent`` protocol."""
|
||||
if use_fake:
|
||||
return FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario(fake_scenario))
|
||||
if base_url is None:
|
||||
raise ValueError("base_url is required when creating a real Agent backend client")
|
||||
return DifyAgentBackendRunClient(Client(base_url=base_url))
|
||||
return DifyAgentBackendRunClient(
|
||||
Client(base_url=base_url, stream_timeout=stream_read_timeout_seconds),
|
||||
stream_max_reconnects=stream_max_reconnects,
|
||||
stream_timeout_seconds=stream_run_timeout_seconds,
|
||||
)
|
||||
|
||||
@@ -7,7 +7,7 @@ separate ``agent-backend.v1`` event stream.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Callable, Iterator
|
||||
from datetime import UTC, datetime
|
||||
from enum import StrEnum
|
||||
|
||||
@@ -69,9 +69,17 @@ class FakeAgentBackendRunClient:
|
||||
del request
|
||||
return CancelRunResponse(run_id=run_id, status="cancelled")
|
||||
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
"""Yield the deterministic public ``RunEvent`` sequence for ``run_id``."""
|
||||
for event in self._events(run_id):
|
||||
if should_stop is not None and should_stop():
|
||||
return
|
||||
if after is not None and event.id is not None and event.id <= after:
|
||||
continue
|
||||
yield event
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
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
|
||||
|
||||
@@ -9,7 +8,6 @@ class ExtraServiceConfig(
|
||||
# place the configs in alphabet order
|
||||
AgentBackendConfig,
|
||||
ArchiveStorageConfig,
|
||||
KnowledgeFSConfig,
|
||||
NotionConfig,
|
||||
SentryConfig,
|
||||
):
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from pydantic import Field, NonNegativeFloat
|
||||
from pydantic import Field, NonNegativeFloat, NonNegativeInt, PositiveFloat
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
@@ -22,6 +22,21 @@ class AgentBackendConfig(BaseSettings):
|
||||
default="success",
|
||||
)
|
||||
|
||||
AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS: PositiveFloat = Field(
|
||||
description="Read timeout for one Agent backend SSE connection.",
|
||||
default=30,
|
||||
)
|
||||
|
||||
AGENT_BACKEND_STREAM_MAX_RECONNECTS: NonNegativeInt = Field(
|
||||
description="Maximum Agent backend SSE reconnects before failing the run.",
|
||||
default=3,
|
||||
)
|
||||
|
||||
AGENT_BACKEND_RUN_TIMEOUT_SECONDS: PositiveFloat = Field(
|
||||
description="Total deadline for one Agent backend run event stream.",
|
||||
default=1200,
|
||||
)
|
||||
|
||||
AGENT_SHELL_ENABLED: bool = Field(
|
||||
description=(
|
||||
"Inject the dify.shell layer (sandboxed bash workspace) into Agent runs. "
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
"""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_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")
|
||||
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_complete_connection(self) -> "KnowledgeFSConfig":
|
||||
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")
|
||||
return self
|
||||
@@ -275,6 +275,42 @@ class PluginConfig(BaseSettings):
|
||||
default=50 * 1024 * 1024,
|
||||
)
|
||||
|
||||
NEW_USER_DEFAULT_PLUGIN_IDS: str = Field(
|
||||
description="Comma-separated marketplace plugin IDs whose latest versions are installed for new users",
|
||||
default="",
|
||||
)
|
||||
|
||||
@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()]
|
||||
|
||||
NEW_USER_DEFAULT_MODELS: str = Field(
|
||||
description=("Comma-separated default models for new users in 'model_type:provider:model' format"),
|
||||
default="",
|
||||
)
|
||||
|
||||
@property
|
||||
def NEW_USER_DEFAULT_MODEL_LIST(self) -> list[tuple[str, str, str]]:
|
||||
default_models: list[tuple[str, str, str]] = []
|
||||
configured_model_types: set[str] = set()
|
||||
|
||||
for item in self.NEW_USER_DEFAULT_MODELS.split(","):
|
||||
if not item.strip():
|
||||
continue
|
||||
|
||||
parts = tuple(part.strip() for part in item.split(":", 2))
|
||||
if len(parts) != 3 or not all(parts):
|
||||
raise ValueError("NEW_USER_DEFAULT_MODELS entries must use 'model_type:provider:model' format")
|
||||
|
||||
model_type, provider, model = parts
|
||||
if model_type in configured_model_types:
|
||||
raise ValueError(f"NEW_USER_DEFAULT_MODELS contains duplicate model type: {model_type}")
|
||||
|
||||
configured_model_types.add(model_type)
|
||||
default_models.append((model_type, provider, model))
|
||||
|
||||
return default_models
|
||||
|
||||
|
||||
class MarketplaceConfig(BaseSettings):
|
||||
"""
|
||||
@@ -784,6 +820,11 @@ class WorkflowConfig(BaseSettings):
|
||||
default=500,
|
||||
)
|
||||
|
||||
WORKFLOW_GENERATOR_NODE_BUILDER_MAX_WORKERS: PositiveInt = Field(
|
||||
description="Maximum concurrent node-builder LLM calls per workflow generation request",
|
||||
default=6,
|
||||
)
|
||||
|
||||
WORKFLOW_MAX_EXECUTION_TIME: PositiveInt = Field(
|
||||
description="Maximum execution time in seconds for a single workflow",
|
||||
default=1200,
|
||||
|
||||
@@ -38,7 +38,6 @@ from . import (
|
||||
feature,
|
||||
human_input_form,
|
||||
init_validate,
|
||||
knowledge_fs_proxy,
|
||||
notification,
|
||||
ping,
|
||||
setup,
|
||||
@@ -196,7 +195,6 @@ __all__ = [
|
||||
"human_input_form",
|
||||
"init_validate",
|
||||
"installed_app",
|
||||
"knowledge_fs_proxy",
|
||||
"load_balancing_config",
|
||||
"login",
|
||||
"mcp_server",
|
||||
|
||||
@@ -58,6 +58,7 @@ from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from factories import file_factory, variable_factory
|
||||
from fields.base import ResponseModel
|
||||
from fields.conversation_variable_fields import WorkflowConversationVariableResponse
|
||||
from fields.member_fields import SimpleAccount
|
||||
from fields.workflow_run_fields import WorkflowRunNodeExecutionResponse
|
||||
from graphon.enums import NodeType
|
||||
@@ -238,21 +239,6 @@ class WorkflowOnlineUsersPayload(BaseModel):
|
||||
return list(dict.fromkeys(app_id.strip() for app_id in app_ids if app_id.strip()))
|
||||
|
||||
|
||||
class WorkflowConversationVariableResponse(ResponseModel):
|
||||
id: str
|
||||
name: str
|
||||
value_type: str
|
||||
value: Any
|
||||
description: str
|
||||
|
||||
@field_validator("value_type", mode="before")
|
||||
@classmethod
|
||||
def _serialize_value_type(cls, value: Any) -> str:
|
||||
if hasattr(value, "exposed_type"):
|
||||
return str(value.exposed_type())
|
||||
return str(value)
|
||||
|
||||
|
||||
class PipelineVariableResponse(ResponseModel):
|
||||
label: str
|
||||
variable: str
|
||||
|
||||
@@ -16,7 +16,6 @@ from controllers.console.auth.error import (
|
||||
)
|
||||
from controllers.console.error import AccountNotFound, EmailSendIpLimitError
|
||||
from controllers.console.wraps import email_password_login_enabled, setup_required
|
||||
from events.tenant_event import tenant_was_created
|
||||
from extensions.ext_database import db
|
||||
from libs.helper import EmailStr, extract_remote_ip
|
||||
from libs.password import hash_password
|
||||
@@ -201,7 +200,4 @@ class ForgotPasswordResetApi(Resource):
|
||||
not TenantService.get_join_tenants(account, session=db.session())
|
||||
and FeatureService.get_system_features().is_allow_create_workspace
|
||||
):
|
||||
tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
|
||||
TenantService.create_tenant_member(tenant, account, db.session(), role="owner")
|
||||
account.set_current_tenant_with_session(tenant, session=db.session())
|
||||
tenant_was_created.send(tenant)
|
||||
TenantService.create_owner_tenant(account, session=db.session())
|
||||
|
||||
@@ -42,7 +42,6 @@ from controllers.console.wraps import (
|
||||
setup_required,
|
||||
with_current_user,
|
||||
)
|
||||
from events.tenant_event import tenant_was_created
|
||||
from extensions.ext_database import db
|
||||
from libs.helper import EmailStr, extract_remote_ip
|
||||
from libs.helper import timezone as validate_timezone_string
|
||||
@@ -317,10 +316,7 @@ class EmailCodeLoginApi(Resource):
|
||||
if not FeatureService.get_system_features().is_allow_create_workspace:
|
||||
raise NotAllowedCreateWorkspace()
|
||||
else:
|
||||
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
|
||||
TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner")
|
||||
account.set_current_tenant_with_session(new_tenant, session=db.session())
|
||||
tenant_was_created.send(new_tenant)
|
||||
TenantService.create_owner_tenant(account, session=db.session())
|
||||
|
||||
if account is None:
|
||||
try:
|
||||
|
||||
@@ -11,7 +11,6 @@ from configs import dify_config
|
||||
from constants.languages import languages
|
||||
from controllers.common.fields import RedirectResponse
|
||||
from controllers.common.schema import query_params_from_model, register_response_schema_model, register_schema_models
|
||||
from events.tenant_event import tenant_was_created
|
||||
from extensions.ext_database import db
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from libs.helper import extract_remote_ip
|
||||
@@ -282,10 +281,7 @@ def _generate_account(
|
||||
if not FeatureService.get_system_features().is_allow_create_workspace:
|
||||
raise WorkSpaceNotAllowedCreateError()
|
||||
else:
|
||||
new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session())
|
||||
TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner")
|
||||
account.set_current_tenant_with_session(new_tenant, session=db.session())
|
||||
tenant_was_created.send(new_tenant)
|
||||
TenantService.create_owner_tenant(account, session=db.session())
|
||||
|
||||
if not account:
|
||||
normalized_email = user_info.email.lower()
|
||||
|
||||
@@ -60,6 +60,7 @@ from core.errors.error import (
|
||||
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.message_fields import SuggestedQuestionsResponse
|
||||
from graphon.graph_engine.manager import GraphEngineManager
|
||||
from graphon.model_runtime.errors.invoke import InvokeError
|
||||
@@ -373,7 +374,7 @@ class TrialWorkflowResponse(ResponseModel):
|
||||
updated_at: int | None = None
|
||||
tool_published: bool | None = None
|
||||
environment_variables: list[JsonObject] = Field(default_factory=list)
|
||||
conversation_variables: list[JsonObject] = Field(default_factory=list)
|
||||
conversation_variables: list[WorkflowConversationVariableResponse] = Field(default_factory=list)
|
||||
rag_pipeline_variables: list[JsonObject] = Field(default_factory=list)
|
||||
|
||||
@field_validator("created_at", "updated_at", mode="before")
|
||||
|
||||
@@ -1,280 +0,0 @@
|
||||
"""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
|
||||
the exact operation templates generated from the pinned KnowledgeFS contract.
|
||||
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,
|
||||
MethodNotAllowed,
|
||||
NotFound,
|
||||
RequestEntityTooLarge,
|
||||
ServiceUnavailable,
|
||||
)
|
||||
|
||||
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,
|
||||
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 _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 _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.
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@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"])
|
||||
@_console_api_errors
|
||||
@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 MethodNotAllowed(valid_methods=["GET"])
|
||||
return _proxy_request("GET", upstream_path)
|
||||
|
||||
|
||||
@bp.route("/knowledge-fs/<path:upstream_path>", methods=["DELETE", "PATCH", "POST", "PUT"])
|
||||
@_console_api_errors
|
||||
@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)
|
||||
try:
|
||||
get_knowledge_fs_operation(method, upstream_path)
|
||||
except KnowledgeFSRouteNotAllowedError as exc:
|
||||
raise NotFound() from exc
|
||||
return _proxy_knowledge_fs_non_get(method, upstream_path)
|
||||
@@ -88,6 +88,7 @@ class TenantInfoResponse(ResponseModel):
|
||||
custom_config: WorkspaceCustomConfigResponse | None = None
|
||||
trial_credits: int | None = None
|
||||
trial_credits_used: int | None = None
|
||||
trial_credits_exhausted_at: int | None = None
|
||||
next_credit_reset_date: int | None = None
|
||||
|
||||
@field_validator("plan", "status", "trial_end_reason", mode="before")
|
||||
|
||||
@@ -47,10 +47,12 @@ class EnterpriseWorkspace(Resource):
|
||||
if account is None:
|
||||
return {"message": "owner account not found."}, 404
|
||||
|
||||
tenant = TenantService.create_tenant(args.name, is_from_dashboard=True, session=db.session())
|
||||
TenantService.create_tenant_member(tenant, account, db.session(), role="owner")
|
||||
|
||||
tenant_was_created.send(tenant)
|
||||
tenant = TenantService.create_owner_tenant(
|
||||
account,
|
||||
name=args.name,
|
||||
is_from_dashboard=True,
|
||||
session=db.session(),
|
||||
)
|
||||
|
||||
resp = {
|
||||
"id": tenant.id,
|
||||
|
||||
@@ -242,6 +242,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
workflow_triggered_from = WorkflowRunTriggeredFrom.APP_RUN
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=workflow_triggered_from,
|
||||
@@ -249,6 +250,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
# Create workflow node execution repository
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
@@ -375,6 +377,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
# Create workflow execution(aka workflow run) repository
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.DEBUGGING,
|
||||
@@ -382,6 +385,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
# Create workflow node execution repository
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
@@ -466,6 +470,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
# Create workflow execution(aka workflow run) repository
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.DEBUGGING,
|
||||
@@ -473,6 +478,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
|
||||
# Create workflow node execution repository
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
|
||||
@@ -540,6 +540,9 @@ class AgentAppGenerator(MessageBasedAppGenerator):
|
||||
base_url=dify_config.AGENT_BACKEND_BASE_URL,
|
||||
use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
|
||||
fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
|
||||
stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS,
|
||||
stream_max_reconnects=dify_config.AGENT_BACKEND_STREAM_MAX_RECONNECTS,
|
||||
stream_run_timeout_seconds=dify_config.AGENT_BACKEND_RUN_TIMEOUT_SECONDS,
|
||||
),
|
||||
event_adapter=AgentBackendRunEventAdapter(),
|
||||
session_store=AgentAppRuntimeSessionStore(),
|
||||
|
||||
@@ -941,48 +941,64 @@ class AgentAppRunner:
|
||||
if pending_text:
|
||||
persist_answer_text(pending_text)
|
||||
|
||||
for public_event in self._agent_backend_client.stream_events(run_id):
|
||||
if queue_manager.is_stopped():
|
||||
flush_pending_agent_message_text()
|
||||
self._cancel_run(run_id)
|
||||
raise GenerateTaskStoppedError()
|
||||
for internal_event in self._event_adapter.adapt(public_event):
|
||||
try:
|
||||
public_events = self._agent_backend_client.stream_events(
|
||||
run_id,
|
||||
should_stop=queue_manager.is_stopped,
|
||||
)
|
||||
for public_event in public_events:
|
||||
if queue_manager.is_stopped():
|
||||
flush_pending_agent_message_text()
|
||||
self._cancel_run(run_id)
|
||||
raise GenerateTaskStoppedError()
|
||||
if internal_event.type in (
|
||||
AgentBackendInternalEventType.RUN_STARTED,
|
||||
AgentBackendInternalEventType.STREAM_EVENT,
|
||||
AgentBackendInternalEventType.AGENT_MESSAGE_DELTA,
|
||||
):
|
||||
if isinstance(internal_event, AgentBackendAgentMessageDeltaInternalEvent):
|
||||
debounced_delta = text_delta_debouncer.push(internal_event.delta)
|
||||
if debounced_delta:
|
||||
persist_answer_text(debounced_delta)
|
||||
continue
|
||||
|
||||
if isinstance(internal_event, AgentBackendStreamInternalEvent):
|
||||
for internal_event in self._event_adapter.adapt(public_event):
|
||||
if queue_manager.is_stopped():
|
||||
flush_pending_agent_message_text()
|
||||
try:
|
||||
process_recorder.handle_stream_event(internal_event)
|
||||
except Exception:
|
||||
db.session.rollback()
|
||||
logger.warning(
|
||||
"Failed to persist Agent App process event: run_id=%s message_id=%s event_kind=%s",
|
||||
run_id,
|
||||
message_id,
|
||||
internal_event.event_kind,
|
||||
exc_info=True,
|
||||
)
|
||||
self._cancel_run(run_id)
|
||||
raise GenerateTaskStoppedError()
|
||||
if internal_event.type in (
|
||||
AgentBackendInternalEventType.RUN_STARTED,
|
||||
AgentBackendInternalEventType.STREAM_EVENT,
|
||||
AgentBackendInternalEventType.AGENT_MESSAGE_DELTA,
|
||||
):
|
||||
if isinstance(internal_event, AgentBackendAgentMessageDeltaInternalEvent):
|
||||
debounced_delta = text_delta_debouncer.push(internal_event.delta)
|
||||
if debounced_delta:
|
||||
persist_answer_text(debounced_delta)
|
||||
continue
|
||||
|
||||
if isinstance(internal_event, AgentBackendStreamInternalEvent):
|
||||
flush_pending_agent_message_text()
|
||||
try:
|
||||
process_recorder.handle_stream_event(internal_event)
|
||||
except Exception:
|
||||
db.session.rollback()
|
||||
logger.warning(
|
||||
"Failed to persist Agent App process event: run_id=%s message_id=%s event_kind=%s",
|
||||
run_id,
|
||||
message_id,
|
||||
internal_event.event_kind,
|
||||
exc_info=True,
|
||||
)
|
||||
continue
|
||||
continue
|
||||
continue
|
||||
flush_pending_agent_message_text()
|
||||
terminal = internal_event
|
||||
break
|
||||
if terminal is not None:
|
||||
break
|
||||
flush_pending_agent_message_text()
|
||||
terminal = internal_event
|
||||
break
|
||||
if terminal is not None:
|
||||
break
|
||||
except GenerateTaskStoppedError:
|
||||
raise
|
||||
except Exception as error:
|
||||
flush_pending_agent_message_text()
|
||||
self._cancel_run(run_id)
|
||||
if queue_manager.is_stopped():
|
||||
raise GenerateTaskStoppedError() from error
|
||||
raise
|
||||
flush_pending_agent_message_text()
|
||||
if queue_manager.is_stopped():
|
||||
self._cancel_run(run_id)
|
||||
raise GenerateTaskStoppedError()
|
||||
return terminal, process_recorder
|
||||
|
||||
def _cancel_run(self, run_id: str) -> None:
|
||||
|
||||
@@ -21,6 +21,7 @@ from core.app.entities.queue_entities import (
|
||||
WorkflowQueueMessage,
|
||||
)
|
||||
from extensions.ext_redis import redis_client
|
||||
from graphon.graph_engine.manager import GraphEngineManager
|
||||
from graphon.runtime import GraphRuntimeState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -51,6 +52,9 @@ class AppQueueManager(ABC):
|
||||
self._graph_runtime_state: GraphRuntimeState | None = None
|
||||
self._stopped_cache: TTLCache[tuple, bool] = TTLCache(maxsize=1, ttl=1)
|
||||
self._cache_lock = threading.Lock()
|
||||
self._execution_terminal = threading.Event()
|
||||
self._abort_sent = threading.Event()
|
||||
self._lifecycle_lock = threading.Lock()
|
||||
|
||||
def listen(self):
|
||||
"""
|
||||
@@ -59,7 +63,7 @@ class AppQueueManager(ABC):
|
||||
"""
|
||||
# wait for APP_MAX_EXECUTION_TIME seconds to stop listen
|
||||
listen_timeout = dify_config.APP_MAX_EXECUTION_TIME
|
||||
start_time = time.time()
|
||||
start_time = time.monotonic()
|
||||
last_ping_time: int | float = 0
|
||||
try:
|
||||
while True:
|
||||
@@ -72,8 +76,14 @@ class AppQueueManager(ABC):
|
||||
except queue.Empty:
|
||||
continue
|
||||
finally:
|
||||
elapsed_time = time.time() - start_time
|
||||
if elapsed_time >= listen_timeout or self._is_stopped():
|
||||
elapsed_time = time.monotonic() - start_time
|
||||
timed_out = elapsed_time >= listen_timeout
|
||||
manually_stopped = self._is_stopped()
|
||||
if not self._execution_terminal.is_set() and (timed_out or manually_stopped):
|
||||
reason = (
|
||||
f"App execution exceeded {listen_timeout} seconds" if timed_out else "App task was stopped"
|
||||
)
|
||||
self._abort_execution(reason)
|
||||
# publish two messages to make sure the client can receive the stop signal
|
||||
# and stop listening after the stop signal processed
|
||||
self.publish(
|
||||
@@ -84,16 +94,33 @@ class AppQueueManager(ABC):
|
||||
self.publish(QueuePingEvent(), PublishFrom.TASK_PIPELINE)
|
||||
last_ping_time = elapsed_time // 10
|
||||
finally:
|
||||
if not self._execution_terminal.is_set():
|
||||
self._abort_execution("Client response stream closed before app execution completed")
|
||||
self._graph_runtime_state = None # Release reference once consumers finish or close the generator.
|
||||
|
||||
def stop_listen(self):
|
||||
def stop_listen(self, *, execution_terminal: bool = False):
|
||||
"""
|
||||
Stop listen to queue
|
||||
:return:
|
||||
"""
|
||||
if execution_terminal:
|
||||
self._execution_terminal.set()
|
||||
self._clear_task_belong_cache()
|
||||
self._q.put(None)
|
||||
|
||||
def _abort_execution(self, reason: str) -> None:
|
||||
"""Propagate response timeout/disconnect to legacy and GraphEngine runners."""
|
||||
with self._lifecycle_lock:
|
||||
if self._execution_terminal.is_set() or self._abort_sent.is_set():
|
||||
return
|
||||
self._abort_sent.set()
|
||||
|
||||
try:
|
||||
self.set_stop_flag_no_user_check(self._task_id)
|
||||
GraphEngineManager(redis_client).send_stop_command(self._task_id, reason=reason)
|
||||
except Exception:
|
||||
logger.exception("Failed to abort app execution for task %s", self._task_id)
|
||||
|
||||
def _clear_task_belong_cache(self) -> None:
|
||||
"""
|
||||
Remove the task belong cache key once listening is finished.
|
||||
|
||||
@@ -45,7 +45,7 @@ class MessageBasedAppQueueManager(AppQueueManager):
|
||||
if isinstance(
|
||||
event, QueueStopEvent | QueueErrorEvent | QueueMessageEndEvent | QueueAdvancedChatMessageEndEvent
|
||||
):
|
||||
self.stop_listen()
|
||||
self.stop_listen(execution_terminal=True)
|
||||
|
||||
if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped():
|
||||
if self._app_mode == AppMode.ADVANCED_CHAT.value:
|
||||
|
||||
@@ -216,6 +216,7 @@ class PipelineGenerator(BaseAppGenerator):
|
||||
session_factory = sessionmaker(bind=db.engine, expire_on_commit=False)
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=workflow_triggered_from,
|
||||
@@ -223,6 +224,7 @@ class PipelineGenerator(BaseAppGenerator):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN,
|
||||
@@ -425,6 +427,7 @@ class PipelineGenerator(BaseAppGenerator):
|
||||
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING,
|
||||
@@ -432,6 +435,7 @@ class PipelineGenerator(BaseAppGenerator):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
@@ -524,6 +528,7 @@ class PipelineGenerator(BaseAppGenerator):
|
||||
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING,
|
||||
@@ -531,6 +536,7 @@ class PipelineGenerator(BaseAppGenerator):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
|
||||
@@ -42,7 +42,7 @@ class PipelineQueueManager(AppQueueManager):
|
||||
| QueueWorkflowFailedEvent
|
||||
| QueueWorkflowPartialSuccessEvent,
|
||||
):
|
||||
self.stop_listen()
|
||||
self.stop_listen(execution_terminal=True)
|
||||
|
||||
if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped():
|
||||
raise GenerateTaskStoppedError()
|
||||
|
||||
@@ -243,6 +243,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
workflow_triggered_from = WorkflowRunTriggeredFrom.APP_RUN
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=workflow_triggered_from,
|
||||
@@ -250,6 +251,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
# Create workflow node execution repository
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
@@ -470,6 +472,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
# Create workflow execution(aka workflow run) repository
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.DEBUGGING,
|
||||
@@ -477,6 +480,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
# Create workflow node execution repository
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
@@ -560,6 +564,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
# Create workflow execution(aka workflow run) repository
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.DEBUGGING,
|
||||
@@ -567,6 +572,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
|
||||
# Create workflow node execution repository
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=application_generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
|
||||
@@ -41,4 +41,4 @@ class WorkflowAppQueueManager(AppQueueManager):
|
||||
| QueueWorkflowFailedEvent
|
||||
| QueueWorkflowPartialSuccessEvent,
|
||||
):
|
||||
self.stop_listen()
|
||||
self.stop_listen(execution_terminal=True)
|
||||
|
||||
@@ -26,6 +26,7 @@ from core.app.entities.queue_entities import (
|
||||
QueueNodeSucceededEvent,
|
||||
QueueReasoningChunkEvent,
|
||||
QueueRetrieverResourcesEvent,
|
||||
QueueStopEvent,
|
||||
QueueTextChunkEvent,
|
||||
QueueWorkflowFailedEvent,
|
||||
QueueWorkflowPartialSuccessEvent,
|
||||
@@ -424,7 +425,12 @@ class WorkflowBasedAppRunner:
|
||||
QueueWorkflowFailedEvent(error=event.error, exceptions_count=event.exceptions_count)
|
||||
)
|
||||
case GraphRunAbortedEvent():
|
||||
self._publish_event(QueueWorkflowFailedEvent(error=event.reason or "Unknown error", exceptions_count=0))
|
||||
self._publish_event(
|
||||
QueueStopEvent(
|
||||
stopped_by=QueueStopEvent.StopBy.USER_MANUAL,
|
||||
reason=event.reason or "Workflow execution aborted",
|
||||
)
|
||||
)
|
||||
case GraphRunPausedEvent():
|
||||
runtime_state = workflow_entry.graph_engine.graph_runtime_state
|
||||
paused_nodes = runtime_state.get_paused_nodes()
|
||||
|
||||
@@ -500,11 +500,15 @@ class QueueStopEvent(AppQueueEvent):
|
||||
|
||||
event: QueueEvent = QueueEvent.STOP
|
||||
stopped_by: StopBy
|
||||
reason: str | None = None
|
||||
|
||||
def get_stop_reason(self) -> str:
|
||||
"""
|
||||
To stop reason
|
||||
"""
|
||||
if self.reason:
|
||||
return self.reason
|
||||
|
||||
reason_mapping = {
|
||||
QueueStopEvent.StopBy.USER_MANUAL: "Stopped by user.",
|
||||
QueueStopEvent.StopBy.ANNOTATION_REPLY: "Stopped by annotation reply.",
|
||||
|
||||
@@ -47,24 +47,6 @@ 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
|
||||
|
||||
@@ -160,31 +142,7 @@ def _inject_trace_headers(headers: Headers | None) -> Headers:
|
||||
return headers
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
def make_request(method: str, url: str, max_retries: int = SSRF_DEFAULT_MAX_RETRIES, **kwargs: Any) -> httpx.Response:
|
||||
# Convert requests-style allow_redirects to httpx-style follow_redirects
|
||||
if "allow_redirects" in kwargs:
|
||||
allow_redirects = kwargs.pop("allow_redirects")
|
||||
@@ -217,11 +175,6 @@ def make_request(
|
||||
# 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:
|
||||
@@ -232,11 +185,7 @@ def make_request(
|
||||
if user_provided_host is not None:
|
||||
headers["host"] = user_provided_host
|
||||
kwargs["headers"] = headers
|
||||
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)
|
||||
response = client.request(method=method, url=url, **kwargs)
|
||||
|
||||
# Check for SSRF protection by Squid proxy
|
||||
if response.status_code in (401, 403):
|
||||
@@ -246,7 +195,6 @@ def make_request(
|
||||
|
||||
# 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. "
|
||||
@@ -260,7 +208,6 @@ def make_request(
|
||||
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)
|
||||
@@ -273,42 +220,6 @@ def make_request(
|
||||
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)
|
||||
|
||||
|
||||
@@ -403,55 +403,6 @@ class LLMGenerator:
|
||||
return ""
|
||||
return "\n\n".join(sections) + "\n\n"
|
||||
|
||||
@classmethod
|
||||
def classify_workflow_mode(
|
||||
cls,
|
||||
tenant_id: str,
|
||||
instruction: str,
|
||||
model_config: ModelConfig,
|
||||
) -> Literal["workflow", "advanced-chat"]:
|
||||
"""Classify a free-text instruction into a concrete app mode.
|
||||
|
||||
One tiny LLM call using the model the user already picked (so no extra
|
||||
provider setup is needed). Parsed leniently; defaults to
|
||||
``advanced-chat`` on anything unexpected or any error, so a
|
||||
``mode="auto"`` request never blocks generation. NEVER raises.
|
||||
"""
|
||||
default_mode: Literal["workflow", "advanced-chat"] = "advanced-chat"
|
||||
try:
|
||||
model_instance = ModelManager.for_tenant(tenant_id=tenant_id).get_model_instance(
|
||||
tenant_id=tenant_id,
|
||||
model_type=ModelType.LLM,
|
||||
provider=model_config.provider,
|
||||
model=model_config.name,
|
||||
)
|
||||
prompt_messages: list[PromptMessage] = [
|
||||
UserPromptMessage(
|
||||
content=(
|
||||
"Reply with exactly one word: 'workflow' (one-shot automation, no chat) "
|
||||
"or 'advanced-chat' (conversational multi-turn). "
|
||||
f"Instruction: {instruction.strip()}"
|
||||
)
|
||||
),
|
||||
]
|
||||
response: LLMResult = model_instance.invoke_llm(
|
||||
prompt_messages=prompt_messages,
|
||||
model_parameters={"max_tokens": 4, "temperature": 0},
|
||||
stream=False,
|
||||
)
|
||||
text = (response.message.get_text_content() or "").strip().lower()
|
||||
except Exception:
|
||||
logger.info("Workflow mode classification failed; defaulting to %s", default_mode, exc_info=True)
|
||||
return default_mode
|
||||
|
||||
# Lenient parse: an affirmative "workflow" wins; everything else
|
||||
# (including a truncated / empty / garbled reply) falls back to the
|
||||
# conversational default. "advanced-chat" needs no positive match
|
||||
# because it IS the default.
|
||||
if "workflow" in text:
|
||||
return "workflow"
|
||||
return default_mode
|
||||
|
||||
@classmethod
|
||||
def generate_rule_config(cls, tenant_id: str, args: RuleGeneratePayload):
|
||||
output_parser = RuleConfigGeneratorOutputParser()
|
||||
|
||||
@@ -13,7 +13,6 @@ from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from core.repositories.factory import WorkflowExecutionRepository
|
||||
from graphon.entities import WorkflowExecution
|
||||
from libs.helper import extract_tenant_id
|
||||
from models import Account, CreatorUserRole, EndUser
|
||||
from models.enums import WorkflowRunTriggeredFrom
|
||||
from tasks.workflow_execution_tasks import (
|
||||
@@ -47,6 +46,7 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
def __init__(
|
||||
self,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str | None,
|
||||
triggered_from: WorkflowRunTriggeredFrom | None,
|
||||
@@ -56,7 +56,8 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine for fallback operations
|
||||
user: Account or EndUser object containing tenant_id, user ID, and role information
|
||||
tenant_id: Tenant that owns the workflow execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: App ID for filtering by application (can be None)
|
||||
triggered_from: Source of the execution trigger (DEBUGGING or APP_RUN)
|
||||
"""
|
||||
@@ -71,10 +72,8 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine"
|
||||
)
|
||||
|
||||
# Extract tenant_id from user
|
||||
tenant_id = extract_tenant_id(user)
|
||||
if not tenant_id:
|
||||
raise ValueError("User must have a tenant_id or current_tenant_id")
|
||||
raise ValueError("tenant_id is required")
|
||||
self._tenant_id = tenant_id
|
||||
|
||||
# Store app context
|
||||
|
||||
@@ -17,7 +17,6 @@ from core.repositories.factory import (
|
||||
WorkflowNodeExecutionRepository,
|
||||
)
|
||||
from graphon.entities import WorkflowNodeExecution
|
||||
from libs.helper import extract_tenant_id
|
||||
from models import Account, CreatorUserRole, EndUser
|
||||
from models.workflow import WorkflowNodeExecutionTriggeredFrom
|
||||
from tasks.workflow_node_execution_tasks import (
|
||||
@@ -54,6 +53,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
|
||||
def __init__(
|
||||
self,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str | None,
|
||||
triggered_from: WorkflowNodeExecutionTriggeredFrom | None,
|
||||
@@ -63,7 +63,8 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine for fallback operations
|
||||
user: Account or EndUser object containing tenant_id, user ID, and role information
|
||||
tenant_id: Tenant that owns the workflow node execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: App ID for filtering by application (can be None)
|
||||
triggered_from: Source of the execution trigger (SINGLE_STEP or WORKFLOW_RUN)
|
||||
"""
|
||||
@@ -78,10 +79,8 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
|
||||
f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine"
|
||||
)
|
||||
|
||||
# Extract tenant_id from user
|
||||
tenant_id = extract_tenant_id(user)
|
||||
if not tenant_id:
|
||||
raise ValueError("User must have a tenant_id or current_tenant_id")
|
||||
raise ValueError("tenant_id is required")
|
||||
self._tenant_id = tenant_id
|
||||
|
||||
# Store app context
|
||||
|
||||
@@ -62,6 +62,7 @@ class DifyCoreRepositoryFactory:
|
||||
def create_workflow_execution_repository(
|
||||
cls,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str,
|
||||
triggered_from: WorkflowRunTriggeredFrom,
|
||||
@@ -71,7 +72,8 @@ class DifyCoreRepositoryFactory:
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine
|
||||
user: Account or EndUser object
|
||||
tenant_id: Tenant that owns the workflow execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: Application ID
|
||||
triggered_from: Source of the execution trigger
|
||||
|
||||
@@ -87,6 +89,7 @@ class DifyCoreRepositoryFactory:
|
||||
repository_class = import_string(class_path)
|
||||
return repository_class(
|
||||
session_factory=session_factory,
|
||||
tenant_id=tenant_id,
|
||||
user=user,
|
||||
app_id=app_id,
|
||||
triggered_from=triggered_from,
|
||||
@@ -98,6 +101,7 @@ class DifyCoreRepositoryFactory:
|
||||
def create_workflow_node_execution_repository(
|
||||
cls,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str,
|
||||
triggered_from: WorkflowNodeExecutionTriggeredFrom,
|
||||
@@ -107,7 +111,8 @@ class DifyCoreRepositoryFactory:
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine
|
||||
user: Account or EndUser object
|
||||
tenant_id: Tenant that owns the workflow node execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: Application ID
|
||||
triggered_from: Source of the execution trigger
|
||||
|
||||
@@ -123,6 +128,7 @@ class DifyCoreRepositoryFactory:
|
||||
repository_class = import_string(class_path)
|
||||
return repository_class(
|
||||
session_factory=session_factory,
|
||||
tenant_id=tenant_id,
|
||||
user=user,
|
||||
app_id=app_id,
|
||||
triggered_from=triggered_from,
|
||||
|
||||
@@ -13,7 +13,6 @@ from core.repositories.factory import WorkflowExecutionRepository
|
||||
from graphon.entities import WorkflowExecution
|
||||
from graphon.enums import WorkflowExecutionStatus, WorkflowType
|
||||
from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter
|
||||
from libs.helper import extract_tenant_id
|
||||
from models import (
|
||||
Account,
|
||||
CreatorUserRole,
|
||||
@@ -40,6 +39,7 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
def __init__(
|
||||
self,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str | None,
|
||||
triggered_from: WorkflowRunTriggeredFrom | None,
|
||||
@@ -49,7 +49,8 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine for creating sessions
|
||||
user: Account or EndUser object containing tenant_id, user ID, and role information
|
||||
tenant_id: Tenant that owns the workflow execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: App ID for filtering by application (can be None)
|
||||
triggered_from: Source of the execution trigger (DEBUGGING or APP_RUN)
|
||||
"""
|
||||
@@ -64,10 +65,8 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine"
|
||||
)
|
||||
|
||||
# Extract tenant_id from user
|
||||
tenant_id = extract_tenant_id(user)
|
||||
if not tenant_id:
|
||||
raise ValueError("User must have a tenant_id or current_tenant_id")
|
||||
raise ValueError("tenant_id is required")
|
||||
self._tenant_id = tenant_id
|
||||
|
||||
# Store app context
|
||||
|
||||
@@ -23,7 +23,6 @@ from graphon.entities import WorkflowNodeExecution
|
||||
from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
|
||||
from graphon.model_runtime.utils.encoders import jsonable_encoder
|
||||
from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter
|
||||
from libs.helper import extract_tenant_id
|
||||
from libs.uuid_utils import uuidv7
|
||||
from models import (
|
||||
Account,
|
||||
@@ -63,6 +62,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
|
||||
def __init__(
|
||||
self,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str | None,
|
||||
triggered_from: WorkflowNodeExecutionTriggeredFrom | None,
|
||||
@@ -72,7 +72,8 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine for creating sessions
|
||||
user: Account or EndUser object containing tenant_id, user ID, and role information
|
||||
tenant_id: Tenant that owns the workflow node execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: App ID for filtering by application (can be None)
|
||||
triggered_from: Source of the execution trigger (SINGLE_STEP or WORKFLOW_RUN)
|
||||
"""
|
||||
@@ -87,10 +88,8 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
|
||||
f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine"
|
||||
)
|
||||
|
||||
# Extract tenant_id from user
|
||||
tenant_id = extract_tenant_id(user)
|
||||
if not tenant_id:
|
||||
raise ValueError("User must have a tenant_id or current_tenant_id")
|
||||
raise ValueError("tenant_id is required")
|
||||
self._tenant_id = tenant_id
|
||||
|
||||
# Store app context
|
||||
@@ -299,6 +298,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
|
||||
content=value_json.encode("utf-8"),
|
||||
mimetype="application/json",
|
||||
user=self._user,
|
||||
tenant_id=self._tenant_id,
|
||||
)
|
||||
offload = WorkflowNodeExecutionOffload(
|
||||
id=uuidv7(),
|
||||
|
||||
@@ -4,12 +4,12 @@ Workflow generator package.
|
||||
Generates a Dify workflow graph (nodes, edges, viewport) from a natural-language
|
||||
instruction. Intended for the cmd+k `/create` slash command's preview/apply flow.
|
||||
|
||||
Pipeline (slim, single-shot variant):
|
||||
Pipeline:
|
||||
|
||||
runner.WorkflowGenerator.generate_workflow_graph(...)
|
||||
├── planner_prompts: short LLM call → high-level node plan
|
||||
└── builder_prompts: structured-output LLM call → full graph JSON
|
||||
└── postprocess: fill defaults, auto-layout viewport, sanity-check edges
|
||||
├── planner_prompts: short LLM call → node and edge plan
|
||||
├── node_builder_prompts: bounded parallel calls → semantic node configs
|
||||
└── postprocess: assemble wrappers, auto-layout, validate graph
|
||||
|
||||
The runner is pure domain logic; ``WorkflowGeneratorService`` (in ``services/``)
|
||||
owns the model-manager dependency and is what controllers call.
|
||||
|
||||
@@ -1,19 +1,4 @@
|
||||
"""
|
||||
Builder prompts.
|
||||
|
||||
The builder is the second step of the slim planner→builder pipeline. It takes
|
||||
the planner's high-level node list and emits the *full* graph JSON consumed by
|
||||
``WorkflowService.sync_draft_workflow``.
|
||||
|
||||
The builder owns: node configuration (prompts, code, headers, etc.), edge wiring,
|
||||
handle ids ("source"/"target"), positions, and the viewport. It is the only
|
||||
prompt that needs to know the concrete shape of each node type — keep its
|
||||
examples accurate or the LLM will invent fields.
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
"""Compact semantic configuration references for workflow node builders."""
|
||||
|
||||
# Per-node-type configuration cheatsheet.
|
||||
#
|
||||
@@ -23,50 +8,9 @@ from typing import Any
|
||||
# both ``WorkflowService.sync_draft_workflow``'s structural checks and the
|
||||
# runtime entity validation each node performs when the workflow runs.
|
||||
#
|
||||
# The cheatsheet is assembled DYNAMICALLY per request: the planner decides
|
||||
# which node types the workflow needs, and ``build_node_config_cheatsheet``
|
||||
# stitches together only the snippets for those types (plus the always-needed
|
||||
# wrapper / shared-field / edge-handle preamble, and the containers section
|
||||
# when an iteration / loop is planned). This keeps the builder prompt tight —
|
||||
# a 3-node summariser no longer carries the schema for 12 unrelated node
|
||||
# types — and lets each snippet document its FULL schema (e.g. a "file" start
|
||||
# variable's required ``allowed_file_types``) without bloating every prompt.
|
||||
#
|
||||
# The postprocessor in ``runner.py`` fills missing wrapper fields (``type``,
|
||||
# ``positionAbsolute``, ``width``, ``height``, ``sourcePosition`` /
|
||||
# ``targetPosition``, edge ``data.sourceType`` / ``data.targetType``), so the
|
||||
# LLM only needs to emit semantically meaningful fields.
|
||||
|
||||
# Always-included preamble: the node/edge wrapper shape and the shared
|
||||
# ``data`` fields that apply to every node type, plus the "## Per type" header
|
||||
# the per-type snippets slot under.
|
||||
_CHEATSHEET_PREAMBLE = """\
|
||||
## Node wrapper (every node, top-level)
|
||||
|
||||
{"id": "node1" (digits + letters only — see "Node IDs" below),
|
||||
"type": "custom", # ReactFlow renderer key. Iteration/loop
|
||||
# *start* children use special types
|
||||
# (see Containers below).
|
||||
"position": {"x": <number>, "y": <number>},
|
||||
"data": { ... per-type fields ... }}
|
||||
|
||||
Children of iteration / loop containers additionally need
|
||||
``parentId``, ``zIndex: 1002`` and ``extent: "parent"`` — see Containers.
|
||||
|
||||
## Shared "data" fields (every node)
|
||||
|
||||
{"type": "<node-type>", # e.g. "llm", "start", "if-else"
|
||||
"title": "<short label>",
|
||||
"desc": "<one-liner>",
|
||||
"selected": false}
|
||||
|
||||
## Per type — additional "data" fields (only the node types in your plan are shown)"""
|
||||
|
||||
|
||||
# node_type → its per-type schema snippet. Keyed by the exact ``node_type``
|
||||
# string the planner emits so ``build_node_config_cheatsheet`` can look each
|
||||
# one up directly. Iteration / loop are documented in the Containers section
|
||||
# (they are subgraphs, not leaf nodes) rather than here.
|
||||
# Each snippet mirrors the production node default closely enough for one
|
||||
# model call to emit only meaningful ``data`` fields. The runner owns wrappers,
|
||||
# topology, container metadata, layout, and validation.
|
||||
_NODE_SNIPPETS: dict[str, str] = {
|
||||
"start": """\
|
||||
- start:
|
||||
@@ -248,484 +192,30 @@ _NODE_SNIPPETS: dict[str, str] = {
|
||||
Enable only the sub-features you need; ``conditions`` reuse the if-else
|
||||
condition shape (key / comparison_operator / value). Outputs: ``result``
|
||||
(the processed array), ``first_record``, ``last_record``.""",
|
||||
"assigner": """\
|
||||
- assigner (write to an existing conversation / loop variable):
|
||||
{"version": "2",
|
||||
"items": [{"variable_selector": ["<target-node>", "<target-var>"],
|
||||
"input_type": "variable",
|
||||
"operation": "over-write",
|
||||
"value": ["<source-node>", "<source-var>"]}]}
|
||||
``input_type`` is "variable" (value is a selector) or "constant".
|
||||
Operations: over-write | clear | append | extend | set | += | -= | *= |
|
||||
/= | remove-first | remove-last.""",
|
||||
"human-input": """\
|
||||
- human-input (pause for a person; use webapp delivery by default):
|
||||
{"delivery_methods": [{"id": "webapp", "type": "webapp", "enabled": true}],
|
||||
"form_content": "<short review / approval instructions>",
|
||||
"inputs": [{"type": "paragraph", "output_variable_name": "comment",
|
||||
"default": {"type": "constant", "selector": [], "value": ""}}],
|
||||
"user_actions": [{"id": "approve", "title": "Approve",
|
||||
"button_style": "primary"}],
|
||||
"timeout": 3, "timeout_unit": "day"}
|
||||
Each ``inputs[].output_variable_name`` is an output variable. Outgoing
|
||||
edges use the matching user-action id as ``sourceHandle``.""",
|
||||
}
|
||||
|
||||
|
||||
# Pulled into the cheatsheet only when an iteration / loop appears in the plan.
|
||||
_CONTAINERS_SECTION = """\
|
||||
## Containers — iteration / loop
|
||||
|
||||
These are SUBGRAPH nodes. To use one you MUST emit, in order:
|
||||
|
||||
1. The container node itself, e.g. for iteration:
|
||||
id: "nodeK"
|
||||
type: "custom"
|
||||
data: {"type": "iteration",
|
||||
"title": "<label>",
|
||||
"desc": "",
|
||||
"selected": false,
|
||||
"start_node_id": "nodeKstart",
|
||||
"iterator_selector": ["<src>", "<list-var>"],
|
||||
"output_selector": ["<inner-last-node>", "<out-var>"],
|
||||
"is_parallel": false,
|
||||
"parallel_nums": 10,
|
||||
"error_handle_mode": "terminated",
|
||||
"flatten_output": true}
|
||||
width: 808
|
||||
height: 204
|
||||
zIndex: 1
|
||||
|
||||
For loop, swap "iteration" → "loop" and use:
|
||||
data: {"type": "loop", "title": "...", "desc": "",
|
||||
"selected": false, "start_node_id": "nodeKstart",
|
||||
"break_conditions": [], "loop_count": 10,
|
||||
"logical_operator": "and"}
|
||||
|
||||
2. The auto-start child (one per container):
|
||||
id: "nodeKstart"
|
||||
type: "custom-iteration-start" # loop → "custom-loop-start"
|
||||
parentId: "nodeK"
|
||||
extent: "parent"
|
||||
draggable: false
|
||||
selectable: false
|
||||
zIndex: 1002
|
||||
position: {"x": 60, "y": 78} # relative to parent
|
||||
data: {"type": "iteration-start", # loop → "loop-start"
|
||||
"title": "", "desc": "",
|
||||
"isInIteration": true, # loop → "isInLoop": true
|
||||
"selected": false}
|
||||
|
||||
3. Each inner-pipeline node (any node type, follows normal data rules) MUST add:
|
||||
parentId: "nodeK"
|
||||
extent: "parent"
|
||||
zIndex: 1002
|
||||
position: {x, y} # relative to parent
|
||||
data: {..., "isInIteration": true, # loop → "isInLoop": true
|
||||
"iteration_id": "nodeK"} # loop → "loop_id"
|
||||
|
||||
4. Edges INSIDE a container must add to ``data``:
|
||||
"isInIteration": true # loop → "isInLoop": true
|
||||
"iteration_id": "nodeK" # loop → "loop_id"
|
||||
and use ``zIndex: 1002``. Edges OUTSIDE containers use the default
|
||||
``isInIteration: false`` / ``isInLoop: false``.
|
||||
|
||||
5. The container's incoming/outgoing edges connect to the container's id
|
||||
(``nodeK``), NOT to inner nodes. The first inner edge connects from
|
||||
``nodeKstart``."""
|
||||
|
||||
|
||||
# Always-included trailer: edge handle conventions for every graph.
|
||||
_EDGE_HANDLES_SECTION = """\
|
||||
## Edge handles
|
||||
|
||||
- Most nodes: sourceHandle "source", targetHandle "target".
|
||||
- if-else cases: sourceHandle is the case_id ("true" / "false" / ...).
|
||||
- question-classifier: sourceHandle is the class_id ("1" / "2" / ...).
|
||||
- iteration-start / sourceHandle "source"; the edge from the *start node
|
||||
loop-start: is what kicks off the first inner step."""
|
||||
|
||||
|
||||
# Container node types are described in ``_CONTAINERS_SECTION`` rather than as
|
||||
# leaf snippets; their presence in a plan pulls that section in.
|
||||
_CONTAINER_NODE_TYPES = frozenset({"iteration", "loop"})
|
||||
|
||||
|
||||
def build_node_config_cheatsheet(node_types: Iterable[str] | None = None) -> str:
|
||||
"""
|
||||
Assemble the builder cheatsheet for exactly the node types in the plan.
|
||||
|
||||
``node_types`` is the set of ``node_type`` strings the planner chose. We
|
||||
emit the always-on preamble (wrapper / shared fields), then only the
|
||||
per-type snippets for the requested types (``start`` is always included —
|
||||
every graph has one), the Containers section when an iteration / loop is
|
||||
planned, and the edge-handles trailer. Unknown / unrecognised type strings
|
||||
are ignored (the runtime / structural validator catches genuinely bogus
|
||||
types).
|
||||
|
||||
``None`` returns the FULL cheatsheet (every snippet + containers) — used to
|
||||
build the static back-compat constants below and as a safe fallback.
|
||||
"""
|
||||
if node_types is None:
|
||||
requested: set[str] = set(_NODE_SNIPPETS) | set(_CONTAINER_NODE_TYPES)
|
||||
else:
|
||||
requested = {str(t).strip() for t in node_types if str(t).strip()}
|
||||
requested.add("start") # every workflow has exactly one start node
|
||||
|
||||
parts: list[str] = [_CHEATSHEET_PREAMBLE]
|
||||
# Iterate _NODE_SNIPPETS (not ``requested``) to keep a stable, readable order.
|
||||
parts.extend(snippet for node_type, snippet in _NODE_SNIPPETS.items() if node_type in requested)
|
||||
if requested & _CONTAINER_NODE_TYPES:
|
||||
parts.append(_CONTAINERS_SECTION)
|
||||
parts.append(_EDGE_HANDLES_SECTION)
|
||||
return "\n\n".join(parts) + "\n"
|
||||
|
||||
|
||||
# Full cheatsheet (all node types) — retained as a module constant so callers
|
||||
# and tests that want the complete reference can import it directly. The
|
||||
# dynamic per-request prompt is built by ``get_builder_system_prompt``.
|
||||
NODE_CONFIG_CHEATSHEET = build_node_config_cheatsheet()
|
||||
|
||||
|
||||
_BASE_SYSTEM_PROMPT_HEAD = """You are a Dify workflow builder.
|
||||
|
||||
You are given:
|
||||
1. A user instruction (what the workflow should do).
|
||||
2. A node plan from the planner (which nodes to use, in execution order).
|
||||
|
||||
Your job: emit a complete Dify workflow graph as JSON. The graph will be written
|
||||
directly into a Studio draft, so it must be syntactically valid and structurally
|
||||
correct.
|
||||
|
||||
# Hard rules
|
||||
|
||||
1. The output is a single JSON object — no prose, no Markdown, no code fences.
|
||||
2. NODE IDs MUST USE ONLY ALPHANUMERICS + UNDERSCORES — never hyphens.
|
||||
Dify's run-time placeholder regex (see ``variable_pool.VARIABLE_PATTERN``)
|
||||
is ``\\{\\{#([a-zA-Z0-9_]{1,50}(?:\\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){1,10})#\\}\\}``,
|
||||
so any placeholder pointing at a hyphenated id (e.g. ``{{#node-1.text#}}``)
|
||||
silently fails to match at run time and the literal string survives into
|
||||
the prompt — the user then sees ``{{#node-1.text#}}`` in their output.
|
||||
Use the EXACT ids from the plan, formatted as ``node1``, ``node2``, ... in
|
||||
plan order. Edge ``source`` / ``target`` must reference these ids.
|
||||
3. Every node has top-level fields: id, type, position, data.
|
||||
- "type" is always "custom" (ReactFlow node renderer).
|
||||
- "data.type" is the actual node type ("llm", "start", etc.).
|
||||
4. Every edge has top-level fields: id, source, target, type, sourceHandle, targetHandle.
|
||||
- "type" is always "custom".
|
||||
- "sourceHandle"/"targetHandle" follow the cheatsheet (default: "source"/"target").
|
||||
- Edge id format: "<source>-<sourceHandle>-<target>-<targetHandle>".
|
||||
5. Use the model from the planner context for ALL "llm" / "question-classifier" /
|
||||
"parameter-extractor" nodes (provider, name, mode, completion_params).
|
||||
6. Reference upstream outputs with the literal placeholder syntax
|
||||
``{{#<node-id>.<output-var>#}}`` — that's DOUBLE curly braces with ``#``
|
||||
markers inside (matching Dify's runtime placeholder regex
|
||||
``\\{\\{#[^#]+#\\}\\}``). NEVER emit single-brace ``{#…#}`` — Dify will
|
||||
not interpolate it, so the LLM at run time would see the literal
|
||||
placeholder string in its prompt and echo it back as output. Use
|
||||
``["<node-id>", "<output-var>"]`` for ``value_selector`` /
|
||||
``query_variable_selector`` / etc.
|
||||
7. The "start" node owns input variables; downstream nodes reference them as
|
||||
``["<start-node-id>", "<var-name>"]`` for selectors or
|
||||
``{{#<start-node-id>.<var-name>#}}`` inside prompt strings.
|
||||
8. NEVER emit "code" or "http-request" nodes if a tool from the "Available tools"
|
||||
section below covers the same task — replace them with a "tool" node referencing
|
||||
the exact provider/tool identifier from the catalogue. "code" / "http-request"
|
||||
are last-resort escape hatches for arbitrary transformations and APIs that no
|
||||
installed tool can express.
|
||||
9. EVERY variable reference MUST resolve to a real, declared variable on the
|
||||
source node — never invent a variable name. Specifically:
|
||||
- ``{{#<node-id>.<var>#}}`` inside a prompt / ``answer`` / ``template-transform``
|
||||
template (DOUBLE braces — single ``{#…#}`` is NOT a Dify placeholder
|
||||
and will NOT be substituted), AND ``["<node-id>", "<var>"]`` inside a
|
||||
``value_selector`` /
|
||||
``query_variable_selector`` / ``iterator_selector`` / ``output_selector`` /
|
||||
``tool_parameters[*].value`` (when ``type: "variable"``), MUST point at a
|
||||
value that the source node actually exposes:
|
||||
* ``start`` → one of the ``data.variables[*].variable`` entries you
|
||||
declared on the start node. Add an entry if you need a new input.
|
||||
* ``llm`` → ``text`` (the default LLM output) or, when structured
|
||||
output is enabled, a key from its schema.
|
||||
* ``code`` → a key in ``data.outputs``.
|
||||
* ``knowledge-retrieval`` → ``result`` (the standard array output).
|
||||
* ``parameter-extractor`` → one of the ``data.parameters[*].name``.
|
||||
* ``document-extractor`` → ``text`` (extracted file text; an array of
|
||||
strings when ``is_array_file`` is true).
|
||||
* ``variable-aggregator`` → ``output``.
|
||||
* ``list-operator`` → ``result`` (array), ``first_record``,
|
||||
``last_record``.
|
||||
* ``tool`` → any parameter declared by the tool — the run time
|
||||
validates these, so you can name them freely, but pick from the
|
||||
documented provider/tool.
|
||||
If the planner's "Start inputs" list (see user prompt) is non-empty,
|
||||
copy each entry verbatim into ``start.data.variables`` so the
|
||||
downstream references resolve.
|
||||
- In Advanced-Chat mode you may also reference ``sys.query`` and
|
||||
``sys.files`` without declaring them. In selector fields, spell these as
|
||||
exactly ``["sys", "query"]`` and ``["sys", "files"]`` — never as a
|
||||
one-item array such as ``["sys,query"]`` or ``["sys.query"]``.
|
||||
10. MULTIPLE KNOWLEDGE-RETRIEVAL INPUTS TO ONE LLM require a template fan-in.
|
||||
``context.variable_selector accepts only one selector`` and therefore
|
||||
cannot carry two retrieval outputs. When an LLM must synthesize two or
|
||||
more retrieval results:
|
||||
- Run the retrieval nodes as parallel siblings from the same query input.
|
||||
- Add one ``template-transform`` node after them. Give it one variable per
|
||||
retrieval, such as ``value_selector: ["node2", "result"]`` and
|
||||
``value_selector: ["node3", "result"]``, and render every source's
|
||||
content into one labelled text output.
|
||||
- Add an edge from EACH retrieval node to the template, then one edge from
|
||||
the template to the LLM. The LLM must NOT receive direct retrieval edges.
|
||||
- Enable the LLM's context, set ``context.variable_selector`` to the
|
||||
template's ``["<template-node-id>", "output"]``, and put
|
||||
``{{#context#}}`` in its prompt.
|
||||
- Do not use ``variable-aggregator`` for this: it selects the first value
|
||||
produced by mutually exclusive branches; it does not concatenate two
|
||||
retrieval results that both ran.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
_BASE_SYSTEM_PROMPT_TAIL = """\
|
||||
|
||||
# Layout
|
||||
|
||||
- Place nodes left-to-right with x=80 + 320 * index, y=280.
|
||||
- Viewport: {"x": 0, "y": 0, "zoom": 0.7}.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
_BASE_SYSTEM_PROMPT_FOOTER = """
|
||||
|
||||
# Output schema
|
||||
|
||||
{
|
||||
"nodes": [...],
|
||||
"edges": [...],
|
||||
"viewport": {"x": 0, "y": 0, "zoom": 0.7}
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
_WORKFLOW_MODE_RULES = """# Mode-specific rules — Workflow
|
||||
|
||||
- The graph MUST start with exactly one "start" node and end with exactly one "end" node.
|
||||
- Do NOT use "answer" nodes (those are for Advanced Chat only).
|
||||
- The "end" node's outputs[].value_selector must point at a real upstream output.
|
||||
"""
|
||||
|
||||
|
||||
_ADVANCED_CHAT_MODE_RULES = """# Mode-specific rules — Advanced Chat (Chatflow)
|
||||
|
||||
- The graph MUST start with exactly one "start" node and end with exactly one "answer" node.
|
||||
- Do NOT use "end" nodes (those are for plain Workflow apps).
|
||||
- The "start" node should expose "sys.query" / "sys.files" automatically; user-defined
|
||||
variables go in start.data.variables.
|
||||
- The "answer" node's "answer" field references upstream outputs as
|
||||
{{#<node-id>.<var>#}} and is what the user sees in chat.
|
||||
"""
|
||||
|
||||
|
||||
def _assemble_builder_system_prompt(mode: str, node_types: Iterable[str] | None) -> str:
|
||||
"""Stitch the builder system prompt for ``mode`` around a cheatsheet built
|
||||
for ``node_types`` (``None`` → full cheatsheet)."""
|
||||
mode_rules = _ADVANCED_CHAT_MODE_RULES if mode == "advanced-chat" else _WORKFLOW_MODE_RULES
|
||||
return (
|
||||
_BASE_SYSTEM_PROMPT_HEAD
|
||||
+ mode_rules
|
||||
+ _BASE_SYSTEM_PROMPT_TAIL
|
||||
+ build_node_config_cheatsheet(node_types)
|
||||
+ _BASE_SYSTEM_PROMPT_FOOTER
|
||||
)
|
||||
|
||||
|
||||
# Static full-cheatsheet prompts — the back-compat default returned by
|
||||
# ``get_builder_system_prompt`` when the caller doesn't pin a node-type set.
|
||||
BUILDER_SYSTEM_PROMPT_WORKFLOW = _assemble_builder_system_prompt("workflow", None)
|
||||
|
||||
BUILDER_SYSTEM_PROMPT_ADVANCED_CHAT = _assemble_builder_system_prompt("advanced-chat", None)
|
||||
|
||||
|
||||
BUILDER_USER_PROMPT = """# User instruction
|
||||
|
||||
{instruction}
|
||||
|
||||
{ideal_output_section}\
|
||||
{existing_graph_section}\
|
||||
# Selected model (use for all LLM-based nodes)
|
||||
|
||||
provider={provider}, name={name}, mode={mode_label}
|
||||
|
||||
{tool_catalogue_section}\
|
||||
{start_inputs_section}\
|
||||
# Node plan (from planner — use these labels and node_types in this order)
|
||||
|
||||
{plan_block}
|
||||
|
||||
Now emit the complete workflow graph JSON.
|
||||
"""
|
||||
|
||||
|
||||
# Node wrapper fields that carry no meaning the builder needs: pure canvas /
|
||||
# selection state, plus geometry the runner's postprocess recomputes anyway.
|
||||
# Stripping them out of the refine prompt cuts its size roughly in half on
|
||||
# hand-edited graphs — fewer tokens in, and (because the builder echoes
|
||||
# untouched nodes verbatim) far fewer tokens out, which is where the latency
|
||||
# lives.
|
||||
_PRUNED_NODE_KEYS = frozenset(
|
||||
{
|
||||
"positionAbsolute",
|
||||
"sourcePosition",
|
||||
"targetPosition",
|
||||
"selected",
|
||||
"dragging",
|
||||
"measured",
|
||||
}
|
||||
)
|
||||
|
||||
# Additionally pruned from TOP-LEVEL nodes only: the layered auto-layout
|
||||
# recomputes their position and size defaults, so the builder never needs to
|
||||
# reproduce them. Container children keep ``position`` (relative to the
|
||||
# parent, which we cannot recompute) and containers keep ``width`` /
|
||||
# ``height`` (their canvas size is real config, not a default).
|
||||
_PRUNED_TOP_LEVEL_NODE_KEYS = _PRUNED_NODE_KEYS | {"position", "width", "height"}
|
||||
|
||||
_CONTAINER_DATA_TYPES = frozenset({"iteration", "loop"})
|
||||
|
||||
# Edge fields the builder must echo; everything else (ids, zIndex,
|
||||
# sourceType / targetType, isInIteration / isInLoop markers) is recomputed
|
||||
# by the runner's postprocess from the node topology.
|
||||
_KEPT_EDGE_KEYS = ("source", "target", "sourceHandle", "targetHandle")
|
||||
|
||||
|
||||
def compact_graph_for_builder(current_graph: dict) -> dict:
|
||||
"""
|
||||
Strip canvas noise out of a draft graph before prompt injection.
|
||||
|
||||
Keeps everything semantically meaningful — ids, wrapper ``type``,
|
||||
``parentId``, the full ``data`` config, child positions, container
|
||||
sizes — and drops geometry / selection state the postprocess pass
|
||||
recomputes. The builder echoes untouched nodes verbatim, so every byte
|
||||
removed here is removed twice (prompt AND completion).
|
||||
"""
|
||||
nodes_out: list[dict] = []
|
||||
for node in current_graph.get("nodes") or []:
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
is_child = bool(node.get("parentId"))
|
||||
is_container = isinstance(node.get("data"), dict) and node["data"].get("type") in _CONTAINER_DATA_TYPES
|
||||
pruned = _PRUNED_NODE_KEYS if (is_child or is_container) else _PRUNED_TOP_LEVEL_NODE_KEYS
|
||||
compact = {k: v for k, v in node.items() if k not in pruned}
|
||||
if is_container:
|
||||
# Container position is still recomputed by the layout pass.
|
||||
compact.pop("position", None)
|
||||
nodes_out.append(compact)
|
||||
edges_out = [
|
||||
{k: edge[k] for k in _KEPT_EDGE_KEYS if k in edge}
|
||||
for edge in (current_graph.get("edges") or [])
|
||||
if isinstance(edge, dict)
|
||||
]
|
||||
return {"nodes": nodes_out, "edges": edges_out}
|
||||
|
||||
|
||||
def format_builder_existing_graph_section(current_graph: dict | None) -> str:
|
||||
"""
|
||||
Refine mode: give the builder the existing graph JSON so it can keep
|
||||
every node and edge the user's change does not touch byte-for-byte — same
|
||||
ids, same config, same prompt templates. Without the full config the
|
||||
builder would regenerate untouched nodes from scratch and silently drop
|
||||
the user's hand-tuned settings. Canvas-only fields are stripped first
|
||||
(see ``compact_graph_for_builder``) — they're recomputed in postprocess,
|
||||
so carrying them only slows the call down.
|
||||
|
||||
Returns an empty string in create mode (no ``current_graph``); the builder
|
||||
then behaves exactly as before, constructing the graph purely from the
|
||||
planner's node plan.
|
||||
"""
|
||||
if not current_graph:
|
||||
return ""
|
||||
graph_json = json.dumps(compact_graph_for_builder(current_graph), ensure_ascii=False, separators=(",", ":"))
|
||||
return (
|
||||
"# Existing graph to refine (JSON)\n\n"
|
||||
"You are REFINING this existing graph, NOT building from scratch. Apply "
|
||||
"ONLY the change the user instruction describes. Every node and edge the "
|
||||
"change does not affect MUST be preserved verbatim — keep the same node "
|
||||
"ids, the same `data` config, and the same prompt templates. The node "
|
||||
"plan below is the target node set after your change; use the existing "
|
||||
"graph as the source of truth for the config of nodes that carry over.\n\n"
|
||||
f"```json\n{graph_json}\n```\n\n"
|
||||
)
|
||||
|
||||
|
||||
def format_start_inputs_section(start_inputs: list[dict[str, Any]]) -> str:
|
||||
"""
|
||||
Surface the planner's ``start_inputs`` list to the builder so it can
|
||||
populate ``start.data.variables`` with the exact set of inputs every
|
||||
downstream variable reference will need. Empty list → empty section,
|
||||
because the builder may then declare no input variables (e.g. an
|
||||
Advanced-Chat workflow that only consumes ``sys.query``).
|
||||
"""
|
||||
if not start_inputs:
|
||||
return ""
|
||||
lines = ["# Start inputs (copy each entry verbatim into start.data.variables)"]
|
||||
lines.append("")
|
||||
for inp in start_inputs:
|
||||
variable = str(inp.get("variable") or "").strip()
|
||||
label = str(inp.get("label") or "").strip()
|
||||
type_ = str(inp.get("type") or "paragraph").strip()
|
||||
if not variable:
|
||||
continue
|
||||
lines.append(f"- variable={variable!r} label={label!r} type={type_!r}")
|
||||
lines.append("")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def format_builder_tool_catalogue_section(catalogue_text: str) -> str:
|
||||
"""
|
||||
Builder-facing catalogue block. The builder needs the same identifiers
|
||||
the planner saw, plus a stern reminder that ``tool`` nodes MUST set
|
||||
``provider_id`` / ``provider_name`` / ``tool_name`` to entries that
|
||||
actually exist in this list — hallucinated tools fail at draft sync.
|
||||
"""
|
||||
if not catalogue_text.strip():
|
||||
return ""
|
||||
return (
|
||||
"# Available tools (use these exact provider/tool identifiers — "
|
||||
"for each 'tool' node, set provider_id and provider_name to the "
|
||||
"provider portion and tool_name to the tool portion)\n\n"
|
||||
f"{catalogue_text}\n\n"
|
||||
)
|
||||
|
||||
|
||||
def format_plan_block(plan_nodes: list[dict[str, Any]]) -> str:
|
||||
"""
|
||||
Render the planner output as a numbered list the builder can quote.
|
||||
|
||||
Node IDs use no separator (``node1``, ``node2``, ...) because Dify's
|
||||
run-time placeholder regex requires ``[a-zA-Z0-9_]`` in the node-id
|
||||
slot — a hyphenated id like ``node-1`` would silently fail to match
|
||||
at run time and the literal ``{{#node-1.var#}}`` survives into the
|
||||
LLM prompt.
|
||||
|
||||
For container children (planner emitted a ``"parent": "<label>"`` key),
|
||||
we resolve the parent label to its ``nodeN`` id and surface it on the
|
||||
same line so the builder knows to set ``parentId`` and the
|
||||
``isInIteration`` / ``isInLoop`` markers on inner nodes.
|
||||
"""
|
||||
# First pass: label → node-id so we can resolve "parent" hints.
|
||||
label_to_id: dict[str, str] = {}
|
||||
for idx, node in enumerate(plan_nodes, start=1):
|
||||
label = str(node.get("label") or "")
|
||||
if label and label not in label_to_id:
|
||||
label_to_id[label] = f"node{idx}"
|
||||
|
||||
lines = []
|
||||
for idx, node in enumerate(plan_nodes, start=1):
|
||||
node_id = f"node{idx}"
|
||||
label = node.get("label", "")
|
||||
node_type = node.get("node_type", "")
|
||||
purpose = node.get("purpose", "")
|
||||
parent_label = str(node.get("parent") or "")
|
||||
parent_clause = ""
|
||||
if parent_label:
|
||||
parent_id = label_to_id.get(parent_label, "")
|
||||
if parent_id:
|
||||
parent_clause = f" parent={parent_id}"
|
||||
else:
|
||||
parent_clause = f" parent={parent_label!r}"
|
||||
lines.append(f"{idx}. id={node_id} type={node_type} label={label!r}{parent_clause}\n purpose: {purpose}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def get_builder_system_prompt(mode: str, node_types: Iterable[str] | None = None) -> str:
|
||||
"""
|
||||
Build the builder system prompt for ``mode``, with a cheatsheet scoped to
|
||||
``node_types`` (the planner's chosen node types).
|
||||
|
||||
When ``node_types`` is ``None`` we return the cached full-cheatsheet
|
||||
constant (back-compat default). When the runner passes the plan's node-type
|
||||
set we assemble a fresh prompt carrying only the relevant per-type schemas,
|
||||
so the builder isn't handed config for node types the workflow never uses.
|
||||
"""
|
||||
if node_types is None:
|
||||
return BUILDER_SYSTEM_PROMPT_ADVANCED_CHAT if mode == "advanced-chat" else BUILDER_SYSTEM_PROMPT_WORKFLOW
|
||||
return _assemble_builder_system_prompt(mode, node_types)
|
||||
def get_node_config_snippet(node_type: str) -> str:
|
||||
"""Return the semantic config reference for one leaf node type."""
|
||||
return _NODE_SNIPPETS.get(node_type, "")
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Compact prompts for parallel, per-node workflow configuration.
|
||||
|
||||
Each call produces only the semantic ``data`` fields for one planned node.
|
||||
Canvas wrappers, shared labels, topology, layout, and edge defaults are owned
|
||||
by ``WorkflowGenerator`` so completion length scales with node configuration
|
||||
rather than with the full ReactFlow graph.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from core.workflow.generator.prompts.builder_prompts import get_node_config_snippet
|
||||
|
||||
_CONTAINER_CONFIG_SNIPPETS = {
|
||||
"iteration": """- iteration:
|
||||
{"iterator_selector": ["<src>", "<list-var>"],
|
||||
"output_selector": ["<last-child>", "<out-var>"],
|
||||
"is_parallel": false, "parallel_nums": 10,
|
||||
"error_handle_mode": "terminated", "flatten_output": true}
|
||||
The runner supplies start_node_id, child wrappers, and the synthetic start node.""",
|
||||
"loop": """- loop:
|
||||
{"break_conditions": [{"id": "c1",
|
||||
"variable_selector": ["<child>", "<var>"],
|
||||
"comparison_operator": "is",
|
||||
"value": "<value>"}],
|
||||
"loop_count": 10, "logical_operator": "and"}
|
||||
The runner supplies start_node_id, child wrappers, and the synthetic start node.""",
|
||||
}
|
||||
|
||||
_NODE_BUILDER_HEAD = """You configure exactly ONE node in a Dify workflow.
|
||||
|
||||
Return one JSON object with exactly this shape: {"config": {...}}.
|
||||
``config`` contains only node-type-specific ``data`` fields. Do NOT repeat id,
|
||||
type, title, desc, selected, position, wrapper fields, edges, or viewport.
|
||||
|
||||
Rules:
|
||||
- Use only ids from the supplied normalized plan.
|
||||
- Placeholder strings use ``{{#node_id.variable#}}``; selector fields use
|
||||
``["node_id", "variable"]``. Never invent an upstream output.
|
||||
- Use the selected model verbatim for llm, question-classifier, and
|
||||
parameter-extractor nodes.
|
||||
- Keep prompts/code concise but complete for the user's requested behavior.
|
||||
- Emit strict JSON only: no prose, Markdown, comments, or trailing commas.
|
||||
|
||||
# Target node schema
|
||||
|
||||
"""
|
||||
|
||||
|
||||
NODE_BUILDER_USER_PROMPT = """# Target node
|
||||
|
||||
id={node_id}, type={node_type}, label={label!r}
|
||||
purpose={purpose}
|
||||
|
||||
# User instruction
|
||||
|
||||
{instruction}
|
||||
|
||||
{ideal_output_section}{mode_section}{model_section}{tool_catalogue_section}{start_inputs_section}{existing_config_section}\
|
||||
# Normalized plan and topology
|
||||
|
||||
{plan_json}
|
||||
|
||||
Return {{"config": {{...}}}} for target node {node_id} now.
|
||||
"""
|
||||
|
||||
|
||||
def get_node_builder_system_prompt(node_type: str) -> str:
|
||||
"""Build a one-node prompt containing only that node's semantic schema."""
|
||||
snippet = _CONTAINER_CONFIG_SNIPPETS.get(node_type) or get_node_config_snippet(node_type)
|
||||
return _NODE_BUILDER_HEAD + (snippet or f"- {node_type}: emit the minimum valid config fields.")
|
||||
|
||||
|
||||
def format_parallel_plan(
|
||||
plan_nodes: list[dict[str, Any]],
|
||||
plan_edges: list[dict[str, Any]],
|
||||
start_inputs: list[dict[str, Any]] | None = None,
|
||||
) -> str:
|
||||
"""Serialize the shared plan compactly so every node call has graph context.
|
||||
|
||||
``start_inputs`` rides along so downstream builders reference the declared
|
||||
``{{#<start-id>.<variable>#}}`` names instead of guessing them from prose
|
||||
— a guessed name gets auto-injected as a spurious form input later.
|
||||
"""
|
||||
payload: dict[str, Any] = {"nodes": plan_nodes, "edges": plan_edges}
|
||||
if start_inputs:
|
||||
payload["start_inputs"] = start_inputs
|
||||
return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
|
||||
|
||||
|
||||
def format_mode_section(mode: str) -> str:
|
||||
"""Tell each builder which app mode it is configuring for.
|
||||
|
||||
Matters most in advanced-chat, where ``sys.query`` / ``sys.files`` are the
|
||||
sanctioned way to reference the user's message — without this the model
|
||||
invents start-node variables that postprocess then materializes as
|
||||
spurious form inputs.
|
||||
"""
|
||||
if mode == "advanced-chat":
|
||||
return (
|
||||
"# App mode\n\n"
|
||||
"advanced-chat: the user's chat message is available as sys.query and uploaded files "
|
||||
'as sys.files — placeholder {{#sys.query#}}, selector ["sys", "query"]. Reference them '
|
||||
"directly; do NOT invent start-node variables for the chat message.\n\n"
|
||||
)
|
||||
return (
|
||||
"# App mode\n\n"
|
||||
"workflow: there are NO automatic system variables; reference user input only through "
|
||||
"the start node's declared variables.\n\n"
|
||||
)
|
||||
|
||||
|
||||
def format_start_inputs_section(start_inputs: list[dict[str, Any]]) -> str:
|
||||
"""Render planner-declared inputs for the start-node builder only."""
|
||||
if not start_inputs:
|
||||
return ""
|
||||
lines = ["# Start inputs (copy each entry verbatim into start.data.variables)", ""]
|
||||
for input_ in start_inputs:
|
||||
variable = str(input_.get("variable") or "").strip()
|
||||
if not variable:
|
||||
continue
|
||||
label = str(input_.get("label") or "").strip()
|
||||
type_ = str(input_.get("type") or "paragraph").strip()
|
||||
lines.append(f"- variable={variable!r} label={label!r} type={type_!r}")
|
||||
lines.append("")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def format_tool_catalogue_section(catalogue_text: str) -> str:
|
||||
"""Render exact tool identifiers for a tool-node builder only."""
|
||||
if not catalogue_text.strip():
|
||||
return ""
|
||||
return (
|
||||
"# Available tools (use these exact provider/tool identifiers — "
|
||||
"set provider_id and provider_name to the provider portion and "
|
||||
"tool_name to the tool portion)\n\n"
|
||||
f"{catalogue_text}\n\n"
|
||||
)
|
||||
@@ -1,13 +1,14 @@
|
||||
"""
|
||||
Planner prompts.
|
||||
|
||||
The planner is the lightweight first step in the slim planner→builder pipeline.
|
||||
The planner is the lightweight first step in the slim planner→node-builders pipeline.
|
||||
It receives the user's natural-language instruction and emits a high-level
|
||||
node plan in JSON. The builder later turns that plan into the final graph.
|
||||
node and edge plan in JSON. Node builders later produce configs that the runner
|
||||
assembles into the final graph.
|
||||
|
||||
We keep the planner deliberately short — the heavy lifting (config schemas,
|
||||
edge wiring, default values) belongs in the builder. The planner only commits
|
||||
to the *which-node-types* decision so the builder gets a tight scaffold.
|
||||
default values) belongs in the builders. The planner commits to the minimum
|
||||
topology and node types so every builder gets a tight scaffold.
|
||||
"""
|
||||
|
||||
PLANNER_SYSTEM_PROMPT = """You are a Dify workflow planner.
|
||||
@@ -39,6 +40,8 @@ minimum set of Dify workflow nodes needed to fulfil it, in execution order.
|
||||
mutually-exclusive paths before "end" / "answer".
|
||||
- "list-operator" — filter / sort / slice an array variable (e.g. the items
|
||||
fed into or produced by an "iteration").
|
||||
- "assigner" — update an existing conversation or loop variable.
|
||||
- "human-input" — pause for a person to review, approve, or enter data.
|
||||
|
||||
# Rules
|
||||
|
||||
@@ -97,22 +100,45 @@ minimum set of Dify workflow nodes needed to fulfil it, in execution order.
|
||||
variables are automatic — downstream nodes may reference them without
|
||||
a ``start_inputs`` entry. In Workflow mode there is NO automatic
|
||||
variable; everything the user supplies must be in ``start_inputs``.
|
||||
11. Output strictly the JSON object — no prose, no Markdown, no code fences.
|
||||
11. Give every node a unique runtime-safe ``id`` using only letters, digits,
|
||||
and underscores. In create mode use ``node1``, ``node2``, ... in node-list
|
||||
order. In refine mode preserve the existing id for every retained node.
|
||||
12. Emit the target graph's edges in ``edges``. Each edge is
|
||||
``{"source": "<id>", "target": "<id>"}``; add ``source_handle`` only
|
||||
for branch nodes: if-else case id, question-classifier class id, or
|
||||
human-input action id. Container children reference the container id in
|
||||
their ``parent`` field; do not emit the synthetic iteration/loop start node.
|
||||
13. In refine mode add ``action`` to every retained target node:
|
||||
``"keep"`` when its data config is unchanged, ``"update"`` when the user
|
||||
asked to change its config, and ``"add"`` for a new node. Removed nodes are
|
||||
omitted. Edge-only rewiring does not require changing a node's action.
|
||||
14. Output strictly the JSON object — no prose, no Markdown, no code fences.
|
||||
15. Echo the app mode in the ``mode`` output field — exactly "workflow" or
|
||||
"advanced-chat". When the ``# Mode`` section says auto, YOU decide:
|
||||
"workflow" for one-shot automations (run once with form inputs, return a
|
||||
result), "advanced-chat" for conversational multi-turn assistants. The
|
||||
terminal node must match the chosen mode (rule 2): "end" for workflow,
|
||||
"answer" for advanced-chat.
|
||||
|
||||
# Output schema
|
||||
|
||||
{
|
||||
"title": "<≤ 40-char title of the workflow>",
|
||||
"description": "<one-sentence summary>",
|
||||
"mode": "workflow | advanced-chat",
|
||||
"app_name": "<≤ 30-char product-style name, e.g. 'URL Summarizer'>",
|
||||
"icon": "<single emoji that captures the workflow's purpose, e.g. '📰'>",
|
||||
"start_inputs": [
|
||||
{"variable": "url", "label": "URL", "type": "text-input"}
|
||||
],
|
||||
"nodes": [
|
||||
{"label": "Start", "node_type": "start", "purpose": "..."},
|
||||
{"label": "Summarize", "node_type": "llm", "purpose": "..."},
|
||||
{"label": "End", "node_type": "end", "purpose": "..."}
|
||||
{"id": "node1", "label": "Start", "node_type": "start", "purpose": "..."},
|
||||
{"id": "node2", "label": "Summarize", "node_type": "llm", "purpose": "..."},
|
||||
{"id": "node3", "label": "End", "node_type": "end", "purpose": "..."}
|
||||
],
|
||||
"edges": [
|
||||
{"source": "node1", "target": "node2"},
|
||||
{"source": "node2", "target": "node3"}
|
||||
]
|
||||
}
|
||||
"""
|
||||
@@ -140,7 +166,8 @@ def format_existing_graph_section(current_graph: dict | None) -> str:
|
||||
|
||||
We pass only ids / node-types / titles + edge endpoints here — the planner
|
||||
decides *which nodes* exist, so it needs the shape, not the per-node config.
|
||||
The builder gets the full graph JSON to preserve untouched node config.
|
||||
Node builders receive only the config of a node marked ``update``;
|
||||
configs marked ``keep`` are reused directly.
|
||||
"""
|
||||
if not current_graph:
|
||||
return ""
|
||||
@@ -156,7 +183,13 @@ def format_existing_graph_section(current_graph: dict | None) -> str:
|
||||
for edge in edges:
|
||||
if not isinstance(edge, dict):
|
||||
continue
|
||||
edge_lines.append(f"- {edge.get('source', '')} -> {edge.get('target', '')}")
|
||||
# Branch wiring (if-else case ids, classifier class ids, human-input
|
||||
# action ids) lives in ``sourceHandle``. The planner is the only
|
||||
# source of edges for the rebuilt graph, so the real handle must be
|
||||
# surfaced here or refine silently rewires branches.
|
||||
handle = str(edge.get("sourceHandle") or "")
|
||||
handle_suffix = f" (source_handle={handle!r})" if handle and handle != "source" else ""
|
||||
edge_lines.append(f"- {edge.get('source', '')} -> {edge.get('target', '')}{handle_suffix}")
|
||||
nodes_block = "\n".join(node_lines) or "(none)"
|
||||
edges_block = "\n".join(edge_lines) or "(none)"
|
||||
return (
|
||||
@@ -166,7 +199,9 @@ def format_existing_graph_section(current_graph: dict | None) -> str:
|
||||
"node list to reflect that change while keeping everything the "
|
||||
"instruction does not mention — preserve existing nodes, their order, "
|
||||
"and their labels wherever the change leaves them untouched. Only add, "
|
||||
"remove, or rename nodes the requested change actually requires.\n\n"
|
||||
"remove, or rename nodes the requested change actually requires. "
|
||||
"For every retained edge, copy its source_handle verbatim from the "
|
||||
"list below — branch wiring must survive the refine unchanged.\n\n"
|
||||
f"Current nodes:\n{nodes_block}\n\n"
|
||||
f"Current edges:\n{edges_block}\n\n"
|
||||
)
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
"""
|
||||
Workflow generator runner.
|
||||
|
||||
Slim planner→builder pipeline. Pure domain logic; the model instance is
|
||||
Slim planner→parallel-node-builder pipeline. Pure domain logic; the model instance is
|
||||
injected by ``WorkflowGeneratorService`` so this module stays cleanly
|
||||
separated from the infrastructure layer.
|
||||
|
||||
Pipeline:
|
||||
|
||||
1. PLANNER — short LLM call producing a high-level node list.
|
||||
2. BUILDER — structured-output LLM call producing the full graph JSON.
|
||||
2. BUILDERS — bounded concurrent LLM calls producing compact node configs.
|
||||
3. POSTPROC — fill safe defaults, lay nodes out left-to-right, dedupe
|
||||
edge ids, and run a final structural sanity check.
|
||||
|
||||
@@ -20,7 +20,8 @@ Intentionally NOT here (deferred to a future iteration):
|
||||
- Tool / model catalogue filtering
|
||||
|
||||
If quality regresses below product threshold we add those back; for now the
|
||||
single planner+builder pair shipped behind cmd+k `/create` is enough.
|
||||
planner and bounded parallel node builders shipped behind cmd+k `/create` are
|
||||
enough.
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -28,17 +29,22 @@ import logging
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from copy import deepcopy
|
||||
from typing import Any, ClassVar, cast
|
||||
|
||||
import json_repair
|
||||
|
||||
from core.workflow.generator.prompts.builder_prompts import (
|
||||
BUILDER_USER_PROMPT,
|
||||
format_builder_existing_graph_section,
|
||||
format_builder_tool_catalogue_section,
|
||||
format_plan_block,
|
||||
from configs import dify_config
|
||||
from core.workflow.generator.prompts.node_builder_prompts import (
|
||||
NODE_BUILDER_USER_PROMPT,
|
||||
format_mode_section,
|
||||
format_parallel_plan,
|
||||
format_start_inputs_section,
|
||||
get_builder_system_prompt,
|
||||
get_node_builder_system_prompt,
|
||||
)
|
||||
from core.workflow.generator.prompts.node_builder_prompts import (
|
||||
format_tool_catalogue_section as format_node_tool_catalogue_section,
|
||||
)
|
||||
from core.workflow.generator.prompts.planner_prompts import (
|
||||
PLANNER_SYSTEM_PROMPT,
|
||||
@@ -55,6 +61,7 @@ from core.workflow.generator.types import (
|
||||
WorkflowGenerateErrorDict,
|
||||
WorkflowGenerateResultDict,
|
||||
WorkflowGenerationMode,
|
||||
WorkflowGenerationModeRequest,
|
||||
)
|
||||
from graphon.enums import BuiltinNodeTypes
|
||||
from graphon.model_runtime.entities.llm_entities import LLMResult
|
||||
@@ -96,11 +103,31 @@ _DEFAULT_FILE_UPLOAD_METHODS = ("local_file", "remote_url")
|
||||
|
||||
# Token ceiling for the planner call when the caller didn't pin one. The plan
|
||||
# is a short JSON node list (a handful of nodes with labels/purposes), so this
|
||||
# is generous headroom while still bounding a runaway response. The builder is
|
||||
# left on the caller's budget — it emits the full graph and genuinely needs it.
|
||||
# is generous headroom while still bounding a runaway response. Builder calls
|
||||
# keep the caller's budget so complex node configs are not truncated.
|
||||
_PLANNER_DEFAULT_MAX_TOKENS = 4096
|
||||
|
||||
|
||||
# Per-node calls trade a larger request count for a shorter critical path.
|
||||
# The cap comes from ``WORKFLOW_GENERATOR_NODE_BUILDER_MAX_WORKERS`` (default 6,
|
||||
# enough to run the planner's recommended 3–6-node plans as a single wave);
|
||||
# provider rate-limit bursts are absorbed by ``_invoke_with_retry``'s bounded
|
||||
# backoff, and operators can dial the env var back down if their provider is
|
||||
# stricter. Read at call time so tests (and live config reloads) can adjust it
|
||||
# without re-importing.
|
||||
def _node_builder_max_workers() -> int:
|
||||
return dify_config.WORKFLOW_GENERATOR_NODE_BUILDER_MAX_WORKERS
|
||||
|
||||
|
||||
_MODEL_NODE_TYPES = frozenset(
|
||||
{
|
||||
BuiltinNodeTypes.LLM,
|
||||
BuiltinNodeTypes.QUESTION_CLASSIFIER,
|
||||
BuiltinNodeTypes.PARAMETER_EXTRACTOR,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# Appended as a trailing user message on the SECOND (and only) attempt when
|
||||
# the first response wasn't parseable as JSON. Keep this terse — the model
|
||||
# already has its full instructions in the original system message; this is
|
||||
@@ -110,6 +137,12 @@ _JSON_RETRY_HINT = (
|
||||
"Do not include any prose, markdown code fences, comments, or trailing commas."
|
||||
)
|
||||
|
||||
_PLANNER_SCHEMA_RETRY_HINT = (
|
||||
"Your plan did not match the required topology schema. Return the complete plan again with "
|
||||
"a unique non-empty id on every node and a non-empty edges array whose source and target "
|
||||
"reference those ids. Return ONLY the JSON object."
|
||||
)
|
||||
|
||||
|
||||
# Provider hiccups we retry: a dropped connection, a 5xx, or a rate-limit are
|
||||
# all transient — the same request usually succeeds moments later. We do NOT
|
||||
@@ -189,14 +222,58 @@ def _result_with_errors(
|
||||
def _with_mode(result: WorkflowGenerateResultDict, mode: WorkflowGenerationMode) -> WorkflowGenerateResultDict:
|
||||
"""Stamp the resolved concrete ``mode`` onto a result envelope.
|
||||
|
||||
``mode="auto"`` requests are resolved to a concrete mode before planning;
|
||||
echoing it back lets the frontend pick the right app type to create. It's
|
||||
present for explicit modes too so the response shape stays uniform.
|
||||
``mode="auto"`` requests are resolved to a concrete mode from the planner
|
||||
output; echoing it back lets the frontend pick the right app type to
|
||||
create. It's present for explicit modes too so the response shape stays
|
||||
uniform.
|
||||
"""
|
||||
result["mode"] = mode
|
||||
return result
|
||||
|
||||
|
||||
def _fallback_mode(mode: WorkflowGenerationModeRequest) -> WorkflowGenerationMode:
|
||||
"""Concrete mode for envelopes emitted before the planner resolved one.
|
||||
|
||||
``auto`` maps to the conversational default — the same never-fail fallback
|
||||
the old standalone classifier used — so ``result.mode`` never leaks the
|
||||
``auto`` sentinel to the frontend.
|
||||
"""
|
||||
return "advanced-chat" if mode == "auto" else mode
|
||||
|
||||
|
||||
def _planner_prompt_mode(mode: WorkflowGenerationModeRequest) -> str:
|
||||
"""Mode string interpolated into the planner user prompt.
|
||||
|
||||
For ``auto`` the value is self-describing so the planner knows the choice
|
||||
is delegated to it (system-prompt rule 15).
|
||||
"""
|
||||
return "auto (choose workflow or advanced-chat)" if mode == "auto" else mode
|
||||
|
||||
|
||||
def _resolve_generation_mode(
|
||||
requested: WorkflowGenerationModeRequest, plan: PlannerResultDict
|
||||
) -> WorkflowGenerationMode:
|
||||
"""Resolve the request mode into the concrete generation mode.
|
||||
|
||||
An explicit request always wins — a contradictory planner ``mode`` field is
|
||||
ignored. For ``auto``: trust the planner's echoed ``mode``, else infer from
|
||||
the plan's terminal node type (the structural source of truth the graph is
|
||||
validated against), else fall back to the conversational default. Lenient
|
||||
on purpose — a bad ``mode`` value must never fail the plan.
|
||||
"""
|
||||
if requested != "auto":
|
||||
return requested
|
||||
planner_mode = str(plan.get("mode") or "").strip().lower()
|
||||
if planner_mode in ("workflow", "advanced-chat"):
|
||||
return cast(WorkflowGenerationMode, planner_mode)
|
||||
node_types = {str(node.get("node_type") or "") for node in plan.get("nodes") or [] if isinstance(node, dict)}
|
||||
if BuiltinNodeTypes.ANSWER in node_types:
|
||||
return "advanced-chat"
|
||||
if BuiltinNodeTypes.END in node_types:
|
||||
return "workflow"
|
||||
return "advanced-chat"
|
||||
|
||||
|
||||
def _build_plan_event(
|
||||
*,
|
||||
plan: PlannerResultDict,
|
||||
@@ -255,7 +332,7 @@ class WorkflowGenerator:
|
||||
provider: str,
|
||||
model_name: str,
|
||||
model_mode: str,
|
||||
mode: WorkflowGenerationMode,
|
||||
mode: WorkflowGenerationModeRequest,
|
||||
instruction: str,
|
||||
ideal_output: str = "",
|
||||
tool_catalogue_text: str = "",
|
||||
@@ -263,20 +340,26 @@ class WorkflowGenerator:
|
||||
current_graph: dict[str, Any] | None = None,
|
||||
) -> WorkflowGenerateResultDict:
|
||||
"""
|
||||
Run planner → builder → postprocess and return a graph payload.
|
||||
Run planner → node builders → postprocess and return a graph payload.
|
||||
|
||||
``mode`` accepts the ``"auto"`` sentinel — the planner then chooses the
|
||||
concrete mode itself (echoed in its ``mode`` output field) so no extra
|
||||
classification call is needed; the resolution is stamped onto the
|
||||
result envelope.
|
||||
|
||||
``current_graph`` switches the pipeline from create mode to REFINE
|
||||
mode: the existing draft graph is injected into both the planner
|
||||
(compact node/edge summary) and the builder (full JSON) so the LLM
|
||||
amends the graph the user is editing instead of inventing a new one.
|
||||
``None`` (the default) is plain create-from-scratch behaviour.
|
||||
mode: the existing draft graph is summarized for the planner. Node
|
||||
builders receive only the config of the node they update, while configs
|
||||
marked ``keep`` are reused without an LLM call. ``None`` (the default)
|
||||
is plain create-from-scratch behaviour.
|
||||
|
||||
``tool_catalogue_text`` is the formatted list of installed tools for
|
||||
the calling tenant (see ``tool_catalogue.build_tool_catalogue`` /
|
||||
``format_tool_catalogue``). It's injected into both the planner and
|
||||
builder prompts so the LLM can pick concrete ``provider/tool``
|
||||
identifiers instead of inventing names; an empty string skips the
|
||||
section entirely (useful for unit tests).
|
||||
identifiers instead of inventing names; node builders receive it
|
||||
only for tool nodes. An empty string skips the section entirely (useful
|
||||
for unit tests).
|
||||
|
||||
``installed_tools`` is the structural sibling — a set of
|
||||
``(provider_name, tool_name)`` pairs the validator consults to reject
|
||||
@@ -316,7 +399,7 @@ class WorkflowGenerator:
|
||||
# The event generator always emits exactly one result envelope; this
|
||||
# fallback only guards against a future refactor that forgets to.
|
||||
if result is None:
|
||||
result = _with_mode(_empty_result(), mode)
|
||||
result = _with_mode(_empty_result(), _fallback_mode(mode))
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
@@ -328,7 +411,7 @@ class WorkflowGenerator:
|
||||
provider: str,
|
||||
model_name: str,
|
||||
model_mode: str,
|
||||
mode: WorkflowGenerationMode,
|
||||
mode: WorkflowGenerationModeRequest,
|
||||
instruction: str,
|
||||
ideal_output: str = "",
|
||||
tool_catalogue_text: str = "",
|
||||
@@ -368,7 +451,7 @@ class WorkflowGenerator:
|
||||
provider: str,
|
||||
model_name: str,
|
||||
model_mode: str,
|
||||
mode: WorkflowGenerationMode,
|
||||
mode: WorkflowGenerationModeRequest,
|
||||
instruction: str,
|
||||
ideal_output: str = "",
|
||||
tool_catalogue_text: str = "",
|
||||
@@ -376,7 +459,7 @@ class WorkflowGenerator:
|
||||
current_graph: dict[str, Any] | None = None,
|
||||
) -> Iterator[tuple[str, dict[str, Any]]]:
|
||||
"""
|
||||
Drive planner → builder → postprocess and yield generation events.
|
||||
Drive planner → node builders → postprocess and yield generation events.
|
||||
|
||||
Shared core for both ``generate_workflow_graph`` (keeps only the final
|
||||
``result``) and ``generate_workflow_graph_stream`` (streams every
|
||||
@@ -402,11 +485,16 @@ class WorkflowGenerator:
|
||||
),
|
||||
)
|
||||
if plan_err is not None:
|
||||
yield "result", cast(dict[str, Any], _with_mode(_result_with_errors(_empty_result(), [plan_err]), mode))
|
||||
failed = _with_mode(_result_with_errors(_empty_result(), [plan_err]), _fallback_mode(mode))
|
||||
yield "result", cast(dict[str, Any], failed)
|
||||
return
|
||||
|
||||
# The lambda return is non-None when no error fired — narrow it for type-checkers.
|
||||
plan = cast(PlannerResultDict, plan)
|
||||
# ``auto`` requests resolve here — the planner echoed its mode choice
|
||||
# (or we infer it from the plan's terminal node). Explicit modes pass
|
||||
# through unchanged. Everything downstream uses the concrete mode.
|
||||
resolved_mode = _resolve_generation_mode(mode, plan)
|
||||
plan_nodes: list[dict[str, Any]] = cast(list[dict[str, Any]], plan.get("nodes", []))
|
||||
if not plan_nodes:
|
||||
empty_plan = _with_mode(
|
||||
@@ -414,15 +502,12 @@ class WorkflowGenerator:
|
||||
_empty_result(),
|
||||
[_err(WorkflowGenerateErrorCode.EMPTY_PLAN, "Planner returned no nodes")],
|
||||
),
|
||||
mode,
|
||||
resolved_mode,
|
||||
)
|
||||
yield "result", cast(dict[str, Any], empty_plan)
|
||||
return
|
||||
|
||||
# A single LLM cannot select multiple retrieval outputs as context.
|
||||
# Make the required template fan-in explicit in the plan so the
|
||||
# builder receives its schema and assigns stable sequential ids.
|
||||
cls._insert_multi_retrieval_template_plan(plan_nodes)
|
||||
plan_edges = [cast(dict[str, Any], edge) for edge in (plan.get("edges") or []) if isinstance(edge, dict)]
|
||||
|
||||
# Planner-supplied user-input declarations. The builder uses these to
|
||||
# populate ``start.data.variables`` so downstream ``{#start.<var>#}``
|
||||
@@ -436,34 +521,48 @@ class WorkflowGenerator:
|
||||
|
||||
# First event the stream sees: the high-level plan, before the slower
|
||||
# builder call. Non-streaming callers ignore it.
|
||||
yield "plan", _build_plan_event(plan=plan, plan_nodes=plan_nodes, start_inputs=start_inputs, mode=mode)
|
||||
yield "plan", _build_plan_event(plan=plan, plan_nodes=plan_nodes, start_inputs=start_inputs, mode=resolved_mode)
|
||||
|
||||
# ── 2. BUILDER ────────────────────────────────────────────────────
|
||||
graph, build_err = cls._run_stage(
|
||||
stage="Builder",
|
||||
failure_fallback_message="Failed to build workflow graph",
|
||||
run=lambda: cls._run_builder(
|
||||
builder_started_at = time.monotonic()
|
||||
|
||||
def build_graph() -> GraphDict:
|
||||
return cls._run_parallel_node_builders(
|
||||
model_instance=model_instance,
|
||||
model_parameters=model_parameters,
|
||||
provider=provider,
|
||||
model_name=model_name,
|
||||
model_mode=model_mode,
|
||||
mode=mode,
|
||||
mode=resolved_mode,
|
||||
instruction=instruction,
|
||||
ideal_output=ideal_output,
|
||||
plan_nodes=plan_nodes,
|
||||
plan_edges=plan_edges,
|
||||
tool_catalogue_text=tool_catalogue_text,
|
||||
start_inputs=start_inputs,
|
||||
current_graph=current_graph,
|
||||
),
|
||||
)
|
||||
|
||||
graph, build_err = cls._run_stage(
|
||||
stage="Builder",
|
||||
failure_fallback_message="Failed to build workflow graph",
|
||||
run=build_graph,
|
||||
)
|
||||
logger.info(
|
||||
"Workflow generator: node builders completed nodes=%s elapsed_ms=%.1f",
|
||||
len(plan_nodes),
|
||||
(time.monotonic() - builder_started_at) * 1000,
|
||||
)
|
||||
if build_err is not None:
|
||||
yield "result", cast(dict[str, Any], _with_mode(_result_with_errors(_empty_result(), [build_err]), mode))
|
||||
yield (
|
||||
"result",
|
||||
cast(dict[str, Any], _with_mode(_result_with_errors(_empty_result(), [build_err]), resolved_mode)),
|
||||
)
|
||||
return
|
||||
graph = cast(GraphDict, graph)
|
||||
|
||||
# ── 3. POSTPROC + VALIDATE ────────────────────────────────────────
|
||||
graph = cls._postprocess_graph(graph=graph, mode=mode)
|
||||
graph = cls._postprocess_graph(graph=graph, mode=resolved_mode)
|
||||
|
||||
# ``app_name`` / ``icon`` are planner display metadata; both default
|
||||
# to "" when the LLM omits them — the FE owns the fallback.
|
||||
@@ -475,13 +574,13 @@ class WorkflowGenerator:
|
||||
"error": "",
|
||||
"errors": [],
|
||||
}
|
||||
_with_mode(result, mode)
|
||||
_with_mode(result, resolved_mode)
|
||||
|
||||
# Final structural sanity check — fail closed if start/end shape is
|
||||
# wrong, container topology is broken, a tool was hallucinated, or a
|
||||
# variable reference points at a node that won't expose it. We still
|
||||
# return the partial graph so the caller can debug or salvage it.
|
||||
structural_errors = cls._validate_structure(graph=graph, mode=mode, installed_tools=installed_tools)
|
||||
structural_errors = cls._validate_structure(graph=graph, mode=resolved_mode, installed_tools=installed_tools)
|
||||
if structural_errors:
|
||||
logger.warning("Workflow generator: structural validation failed: %s", structural_errors)
|
||||
yield "result", cast(dict[str, Any], _result_with_errors(result, structural_errors))
|
||||
@@ -622,14 +721,14 @@ class WorkflowGenerator:
|
||||
*,
|
||||
model_instance,
|
||||
model_parameters: dict[str, Any],
|
||||
mode: WorkflowGenerationMode,
|
||||
mode: WorkflowGenerationModeRequest,
|
||||
instruction: str,
|
||||
ideal_output: str,
|
||||
tool_catalogue_text: str,
|
||||
current_graph: dict[str, Any] | None = None,
|
||||
) -> PlannerResultDict:
|
||||
user_prompt = PLANNER_USER_PROMPT.format(
|
||||
mode=mode,
|
||||
mode=_planner_prompt_mode(mode),
|
||||
instruction=instruction.strip(),
|
||||
existing_graph_section=format_existing_graph_section(current_graph),
|
||||
ideal_output_section=format_ideal_output_section(ideal_output),
|
||||
@@ -639,59 +738,65 @@ class WorkflowGenerator:
|
||||
SystemPromptMessage(content=PLANNER_SYSTEM_PROMPT),
|
||||
UserPromptMessage(content=user_prompt),
|
||||
]
|
||||
clamped_parameters = _clamp_for_planner(model_parameters)
|
||||
parsed = cls._invoke_and_parse_json(
|
||||
model_instance=model_instance,
|
||||
messages=messages,
|
||||
model_parameters=_clamp_for_planner(model_parameters),
|
||||
model_parameters=clamped_parameters,
|
||||
stage="Planner",
|
||||
)
|
||||
try:
|
||||
return cls._validate_planner_schema(parsed)
|
||||
except _StageSchemaError:
|
||||
logger.info("Workflow generator: planner schema invalid; retrying once")
|
||||
parsed = cls._invoke_and_parse_json(
|
||||
model_instance=model_instance,
|
||||
messages=[*messages, UserPromptMessage(content=_PLANNER_SCHEMA_RETRY_HINT)],
|
||||
model_parameters=clamped_parameters,
|
||||
stage="Planner",
|
||||
)
|
||||
return cls._validate_planner_schema(parsed)
|
||||
|
||||
@staticmethod
|
||||
def _validate_planner_schema(parsed: dict[str, Any]) -> PlannerResultDict:
|
||||
"""Require the single planner contract consumed by node builders."""
|
||||
nodes = parsed.get("nodes")
|
||||
if not isinstance(nodes, list):
|
||||
raise _StageSchemaError("Planner", "missing 'nodes' array")
|
||||
if not nodes:
|
||||
return cast(PlannerResultDict, parsed)
|
||||
|
||||
node_ids: set[str] = set()
|
||||
for node in nodes:
|
||||
if not isinstance(node, dict) or "node_type" not in node:
|
||||
if not isinstance(node, dict) or not node.get("node_type"):
|
||||
raise _StageSchemaError("Planner", f"malformed node entry: {node!r}")
|
||||
node_id = node.get("id")
|
||||
if not isinstance(node_id, str) or not node_id.strip():
|
||||
raise _StageSchemaError("Planner", f"node missing non-empty id: {node!r}")
|
||||
if node_id in node_ids:
|
||||
raise _StageSchemaError("Planner", f"duplicate node id: {node_id!r}")
|
||||
node_ids.add(node_id)
|
||||
|
||||
edges = parsed.get("edges")
|
||||
if not isinstance(edges, list) or not edges:
|
||||
raise _StageSchemaError("Planner", "missing non-empty 'edges' array")
|
||||
for edge in edges:
|
||||
if not isinstance(edge, dict):
|
||||
raise _StageSchemaError("Planner", f"malformed edge entry: {edge!r}")
|
||||
source = edge.get("source")
|
||||
target = edge.get("target")
|
||||
if not isinstance(source, str) or not isinstance(target, str):
|
||||
raise _StageSchemaError("Planner", f"edge missing source or target: {edge!r}")
|
||||
if source not in node_ids or target not in node_ids:
|
||||
raise _StageSchemaError("Planner", f"edge references unknown node: {edge!r}")
|
||||
|
||||
return cast(PlannerResultDict, parsed)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Plan normalization
|
||||
# ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _insert_multi_retrieval_template_plan(plan_nodes: list[dict[str, Any]]) -> None:
|
||||
"""Insert the unambiguous multi-retrieval template step when omitted.
|
||||
|
||||
Multiple LLM nodes make ownership ambiguous, so that case remains in
|
||||
the planner's hands. With exactly one LLM, every independent retrieval
|
||||
result can safely fan into one template immediately before that LLM.
|
||||
"""
|
||||
node_types = [str(node.get("node_type") or "") for node in plan_nodes]
|
||||
if node_types.count(BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL) < 2:
|
||||
return
|
||||
if node_types.count(BuiltinNodeTypes.LLM) != 1:
|
||||
return
|
||||
if BuiltinNodeTypes.TEMPLATE_TRANSFORM in node_types:
|
||||
return
|
||||
|
||||
llm_index = node_types.index(BuiltinNodeTypes.LLM)
|
||||
retrievals_before_llm = node_types[:llm_index].count(BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL)
|
||||
if retrievals_before_llm < 2:
|
||||
return
|
||||
plan_nodes.insert(
|
||||
llm_index,
|
||||
{
|
||||
"label": "Combine Knowledge",
|
||||
"node_type": BuiltinNodeTypes.TEMPLATE_TRANSFORM,
|
||||
"purpose": "Combine every knowledge retrieval result into one labelled context for the LLM.",
|
||||
},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Builder
|
||||
# ------------------------------------------------------------------
|
||||
@classmethod
|
||||
def _run_builder(
|
||||
def _run_parallel_node_builders(
|
||||
cls,
|
||||
*,
|
||||
model_instance,
|
||||
@@ -703,53 +808,275 @@ class WorkflowGenerator:
|
||||
instruction: str,
|
||||
ideal_output: str,
|
||||
plan_nodes: list[dict[str, Any]],
|
||||
plan_edges: list[dict[str, Any]],
|
||||
tool_catalogue_text: str,
|
||||
start_inputs: list[dict[str, Any]] | None = None,
|
||||
current_graph: dict[str, Any] | None = None,
|
||||
start_inputs: list[dict[str, Any]],
|
||||
current_graph: dict[str, Any] | None,
|
||||
) -> GraphDict:
|
||||
user_prompt = BUILDER_USER_PROMPT.format(
|
||||
"""Build changed node configs concurrently and expand them into a graph.
|
||||
|
||||
Refine plans can mark existing nodes as ``keep``; those nodes bypass
|
||||
the model entirely and retain their full data config. Every other node
|
||||
gets one compact call, with at most ``_node_builder_max_workers()``
|
||||
calls in flight. Any fragment failure aborts the graph, preserving the
|
||||
generator's existing fail-closed contract.
|
||||
"""
|
||||
existing_by_id = {
|
||||
str(node.get("id")): node
|
||||
for node in ((current_graph or {}).get("nodes") or [])
|
||||
if isinstance(node, dict) and node.get("id")
|
||||
}
|
||||
existing_edges = [edge for edge in ((current_graph or {}).get("edges") or []) if isinstance(edge, dict)]
|
||||
nodes_to_build = [
|
||||
node for node in plan_nodes if not (node.get("action") == "keep" and str(node.get("id")) in existing_by_id)
|
||||
]
|
||||
|
||||
# Shared across every builder call in this request — compute once.
|
||||
plan_json = format_parallel_plan(plan_nodes, plan_edges, start_inputs)
|
||||
mode_section = format_mode_section(mode)
|
||||
|
||||
configs_by_id: dict[str, dict[str, Any]] = {}
|
||||
if nodes_to_build:
|
||||
max_workers = min(_node_builder_max_workers(), len(nodes_to_build))
|
||||
with ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="workflow-node-builder") as executor:
|
||||
futures = {
|
||||
executor.submit(
|
||||
cls._run_node_builder,
|
||||
model_instance=model_instance,
|
||||
model_parameters=model_parameters,
|
||||
provider=provider,
|
||||
model_name=model_name,
|
||||
model_mode=model_mode,
|
||||
mode_section=mode_section,
|
||||
instruction=instruction,
|
||||
ideal_output=ideal_output,
|
||||
target_node=node,
|
||||
plan_json=plan_json,
|
||||
tool_catalogue_text=tool_catalogue_text,
|
||||
start_inputs=start_inputs,
|
||||
existing_node=existing_by_id.get(str(node.get("id"))),
|
||||
): str(node.get("id"))
|
||||
for node in nodes_to_build
|
||||
}
|
||||
try:
|
||||
for future in as_completed(futures):
|
||||
node_id = futures[future]
|
||||
configs_by_id[node_id] = future.result()
|
||||
except BaseException:
|
||||
# Fail fast: one failed fragment aborts the whole graph, so
|
||||
# queued builder calls would only burn quota and delay the
|
||||
# error envelope. In-flight calls cannot be interrupted;
|
||||
# they finish while the pool shuts down.
|
||||
for pending in futures:
|
||||
pending.cancel()
|
||||
raise
|
||||
|
||||
return cls._assemble_parallel_graph(
|
||||
plan_nodes=plan_nodes,
|
||||
plan_edges=plan_edges,
|
||||
configs_by_id=configs_by_id,
|
||||
existing_by_id=existing_by_id,
|
||||
existing_edges=existing_edges,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _run_node_builder(
|
||||
cls,
|
||||
*,
|
||||
model_instance,
|
||||
model_parameters: dict[str, Any],
|
||||
provider: str,
|
||||
model_name: str,
|
||||
model_mode: str,
|
||||
mode_section: str,
|
||||
instruction: str,
|
||||
ideal_output: str,
|
||||
target_node: dict[str, Any],
|
||||
plan_json: str,
|
||||
tool_catalogue_text: str,
|
||||
start_inputs: list[dict[str, Any]],
|
||||
existing_node: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Generate only the semantic config for one normalized plan node."""
|
||||
node_id = str(target_node.get("id") or "")
|
||||
node_type = str(target_node.get("node_type") or "")
|
||||
model_section = ""
|
||||
if node_type in _MODEL_NODE_TYPES:
|
||||
model_section = (
|
||||
f"# Selected model (copy verbatim)\n\nprovider={provider}, name={model_name}, mode={model_mode}\n\n"
|
||||
)
|
||||
existing_config_section = ""
|
||||
if existing_node:
|
||||
existing_data = existing_node.get("data") if isinstance(existing_node.get("data"), dict) else {}
|
||||
existing_config_section = (
|
||||
"# Existing config to preserve unless the instruction changes it\n\n"
|
||||
f"{json.dumps(existing_data, ensure_ascii=False, separators=(',', ':'))}\n\n"
|
||||
)
|
||||
user_prompt = NODE_BUILDER_USER_PROMPT.format(
|
||||
node_id=node_id,
|
||||
node_type=node_type,
|
||||
label=str(target_node.get("label") or ""),
|
||||
purpose=str(target_node.get("purpose") or ""),
|
||||
instruction=instruction.strip(),
|
||||
ideal_output_section=format_ideal_output_section(ideal_output),
|
||||
existing_graph_section=format_builder_existing_graph_section(current_graph),
|
||||
provider=provider,
|
||||
name=model_name,
|
||||
mode_label=model_mode,
|
||||
plan_block=format_plan_block(plan_nodes),
|
||||
tool_catalogue_section=format_builder_tool_catalogue_section(tool_catalogue_text),
|
||||
start_inputs_section=format_start_inputs_section(start_inputs or []),
|
||||
mode_section=mode_section,
|
||||
model_section=model_section,
|
||||
tool_catalogue_section=(
|
||||
format_node_tool_catalogue_section(tool_catalogue_text) if node_type == BuiltinNodeTypes.TOOL else ""
|
||||
),
|
||||
start_inputs_section=(
|
||||
format_start_inputs_section(start_inputs) if node_type == BuiltinNodeTypes.START else ""
|
||||
),
|
||||
existing_config_section=existing_config_section,
|
||||
plan_json=plan_json,
|
||||
)
|
||||
# Scope the builder cheatsheet to exactly the node types the planner
|
||||
# chose, so the prompt carries each type's FULL schema (e.g. a file
|
||||
# start variable's required ``allowed_file_types``) without dragging in
|
||||
# config for unrelated node types.
|
||||
plan_node_types = {
|
||||
str(node.get("node_type") or "").strip() for node in plan_nodes if str(node.get("node_type") or "").strip()
|
||||
}
|
||||
messages = [
|
||||
SystemPromptMessage(content=get_builder_system_prompt(mode, plan_node_types)),
|
||||
UserPromptMessage(content=user_prompt),
|
||||
]
|
||||
parsed = cls._invoke_and_parse_json(
|
||||
model_instance=model_instance,
|
||||
messages=messages,
|
||||
messages=[
|
||||
SystemPromptMessage(content=get_node_builder_system_prompt(node_type)),
|
||||
UserPromptMessage(content=user_prompt),
|
||||
],
|
||||
model_parameters=model_parameters,
|
||||
stage="Builder",
|
||||
stage=f"Builder {node_id}",
|
||||
)
|
||||
config = parsed.get("config")
|
||||
if not isinstance(config, dict):
|
||||
raise _StageSchemaError(f"Builder {node_id}", "missing 'config' object")
|
||||
return cast(dict[str, Any], config)
|
||||
|
||||
nodes = parsed.get("nodes")
|
||||
edges = parsed.get("edges")
|
||||
if not isinstance(nodes, list) or not isinstance(edges, list):
|
||||
raise _StageSchemaError("Builder", "graph missing 'nodes' or 'edges' arrays")
|
||||
@classmethod
|
||||
def _assemble_parallel_graph(
|
||||
cls,
|
||||
*,
|
||||
plan_nodes: list[dict[str, Any]],
|
||||
plan_edges: list[dict[str, Any]],
|
||||
configs_by_id: dict[str, dict[str, Any]],
|
||||
existing_by_id: dict[str, dict[str, Any]],
|
||||
existing_edges: list[dict[str, Any]] | None = None,
|
||||
) -> GraphDict:
|
||||
"""Expand compact node configs and planner topology into graph JSON.
|
||||
|
||||
viewport = parsed.get("viewport") or _DEFAULT_VIEWPORT
|
||||
return cast(
|
||||
GraphDict,
|
||||
{
|
||||
"nodes": nodes,
|
||||
"edges": edges,
|
||||
"viewport": viewport,
|
||||
},
|
||||
)
|
||||
``existing_edges`` (refine only) preserves wiring the planner cannot
|
||||
express: the synthetic ``<container>start`` entry edge keeps its
|
||||
existing target instead of being re-pointed at whichever child the
|
||||
planner happened to list first.
|
||||
"""
|
||||
label_to_id = {
|
||||
str(node.get("label")): str(node.get("id")) for node in plan_nodes if node.get("label") and node.get("id")
|
||||
}
|
||||
type_by_id = {str(node.get("id")): str(node.get("node_type") or "") for node in plan_nodes}
|
||||
children_by_parent: dict[str, list[str]] = {}
|
||||
nodes: list[dict[str, Any]] = []
|
||||
|
||||
for planned in plan_nodes:
|
||||
node_id = str(planned.get("id") or "")
|
||||
node_type = str(planned.get("node_type") or "")
|
||||
existing = existing_by_id.get(node_id)
|
||||
node: dict[str, Any]
|
||||
if planned.get("action") == "keep" and existing is not None:
|
||||
node = deepcopy(existing)
|
||||
else:
|
||||
config = dict(configs_by_id.get(node_id) or {})
|
||||
for shared_key in ("type", "title", "desc", "selected"):
|
||||
config.pop(shared_key, None)
|
||||
data: dict[str, Any] = {
|
||||
"type": node_type,
|
||||
"title": str(planned.get("label") or node_id),
|
||||
"desc": str(planned.get("purpose") or ""),
|
||||
**config,
|
||||
}
|
||||
node = deepcopy(existing) if existing is not None else {"id": node_id}
|
||||
node["id"] = node_id
|
||||
node["data"] = data
|
||||
|
||||
parent_ref = str(planned.get("parent") or "")
|
||||
parent_id = label_to_id.get(parent_ref, parent_ref)
|
||||
if not parent_id and str(node.get("parentId") or "") in type_by_id:
|
||||
# Kept nodes rarely re-state containment — recover the parent
|
||||
# from the deepcopied wrapper so entry-edge synthesis still
|
||||
# counts this child.
|
||||
parent_id = str(node["parentId"])
|
||||
if parent_id:
|
||||
child_index = len(children_by_parent.get(parent_id, []))
|
||||
node["parentId"] = parent_id
|
||||
node.setdefault("position", {"x": 240 + 260 * child_index, "y": 60})
|
||||
node.setdefault("data", {})
|
||||
parent_type = type_by_id.get(parent_id)
|
||||
if parent_type == BuiltinNodeTypes.ITERATION:
|
||||
node["data"].setdefault("isInIteration", True)
|
||||
node["data"].setdefault("iteration_id", parent_id)
|
||||
elif parent_type == BuiltinNodeTypes.LOOP:
|
||||
node["data"].setdefault("isInLoop", True)
|
||||
node["data"].setdefault("loop_id", parent_id)
|
||||
children_by_parent.setdefault(parent_id, []).append(node_id)
|
||||
elif node.get("parentId"):
|
||||
# The container was dropped from the plan: strip the stale
|
||||
# containment markers so the kept node rejoins the top level
|
||||
# (and its auto-layout) instead of pointing at a deleted parent.
|
||||
for wrapper_key in ("parentId", "extent", "zIndex", "position", "positionAbsolute"):
|
||||
node.pop(wrapper_key, None)
|
||||
if isinstance(node.get("data"), dict):
|
||||
for marker_key in ("isInIteration", "iteration_id", "isInLoop", "loop_id"):
|
||||
node["data"].pop(marker_key, None)
|
||||
|
||||
nodes.append(node)
|
||||
if node_type in cls._CONTAINER_TYPES:
|
||||
start_id = f"{node_id}start"
|
||||
node.setdefault("data", {})["start_node_id"] = start_id
|
||||
node.setdefault("width", 808)
|
||||
node.setdefault("height", 204)
|
||||
node.setdefault("zIndex", 1)
|
||||
is_iteration = node_type == BuiltinNodeTypes.ITERATION
|
||||
nodes.append(
|
||||
{
|
||||
"id": start_id,
|
||||
"type": "custom-iteration-start" if is_iteration else "custom-loop-start",
|
||||
"parentId": node_id,
|
||||
"extent": "parent",
|
||||
"draggable": False,
|
||||
"selectable": False,
|
||||
"zIndex": 1002,
|
||||
"position": {"x": 60, "y": 78},
|
||||
"data": {
|
||||
"type": "iteration-start" if is_iteration else "loop-start",
|
||||
"title": "",
|
||||
"desc": "",
|
||||
"selected": False,
|
||||
"isInIteration" if is_iteration else "isInLoop": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
edges: list[dict[str, Any]] = []
|
||||
for planned_edge in plan_edges:
|
||||
edge: dict[str, Any] = {
|
||||
"source": str(planned_edge.get("source") or ""),
|
||||
"target": str(planned_edge.get("target") or ""),
|
||||
}
|
||||
source_handle = planned_edge.get("source_handle") or planned_edge.get("sourceHandle")
|
||||
target_handle = planned_edge.get("target_handle") or planned_edge.get("targetHandle")
|
||||
if source_handle:
|
||||
edge["sourceHandle"] = str(source_handle)
|
||||
if target_handle:
|
||||
edge["targetHandle"] = str(target_handle)
|
||||
edges.append(edge)
|
||||
|
||||
# Synthesize each container's entry edge. Refine keeps the existing
|
||||
# entry target when it is still a child — the planner's node listing
|
||||
# order says nothing about execution order inside a kept container.
|
||||
planned_sources = {str(edge.get("source") or "") for edge in edges}
|
||||
existing_entry_targets = {
|
||||
str(edge.get("source") or ""): str(edge.get("target") or "") for edge in (existing_edges or [])
|
||||
}
|
||||
for parent_id, child_ids in children_by_parent.items():
|
||||
start_id = f"{parent_id}start"
|
||||
if not child_ids or start_id in planned_sources:
|
||||
continue
|
||||
preferred = existing_entry_targets.get(start_id, "")
|
||||
entry_target = preferred if preferred in child_ids else child_ids[0]
|
||||
edges.append({"source": start_id, "target": entry_target})
|
||||
|
||||
return cast(GraphDict, {"nodes": nodes, "edges": edges, "viewport": _DEFAULT_VIEWPORT})
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Postprocessing
|
||||
@@ -1225,6 +1552,11 @@ class WorkflowGenerator:
|
||||
return var == "output"
|
||||
if node_type == BuiltinNodeTypes.LIST_OPERATOR:
|
||||
return var in {"result", "first_record", "last_record"}
|
||||
if node_type == BuiltinNodeTypes.HUMAN_INPUT:
|
||||
return any(
|
||||
isinstance(item, dict) and item.get("output_variable_name") == var
|
||||
for item in (data.get("inputs") or [])
|
||||
)
|
||||
# Other node types (if-else, iteration-start, loop-start, ...) don't
|
||||
# produce outputs of their own.
|
||||
return False
|
||||
@@ -1247,6 +1579,15 @@ class WorkflowGenerator:
|
||||
if isinstance(parameter, dict) and isinstance(parameter.get("name"), str)
|
||||
]
|
||||
return parameters[0] if len(parameters) == 1 else None
|
||||
if node_type == BuiltinNodeTypes.HUMAN_INPUT:
|
||||
human_outputs: list[str] = []
|
||||
for item in data.get("inputs") or []:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
output_name = item.get("output_variable_name")
|
||||
if isinstance(output_name, str):
|
||||
human_outputs.append(output_name)
|
||||
return human_outputs[0] if len(human_outputs) == 1 else None
|
||||
if not isinstance(node_type, str):
|
||||
return None
|
||||
single_output_by_type: dict[str, str] = {
|
||||
@@ -1533,6 +1874,12 @@ class WorkflowGenerator:
|
||||
for klass in (data.get("classes") or [])
|
||||
if isinstance(klass, dict) and klass.get("id")
|
||||
]
|
||||
elif node_type == BuiltinNodeTypes.HUMAN_INPUT:
|
||||
branch_handles = [
|
||||
str(action["id"])
|
||||
for action in (data.get("user_actions") or [])
|
||||
if isinstance(action, dict) and action.get("id")
|
||||
]
|
||||
else:
|
||||
continue
|
||||
|
||||
@@ -1941,8 +2288,8 @@ class WorkflowGenerator:
|
||||
"""
|
||||
Validate iteration / loop topology:
|
||||
|
||||
* every container has at least one child whose ``parentId``
|
||||
points at it;
|
||||
* every container has at least one executable child whose
|
||||
``parentId`` points at it;
|
||||
* every non-container node with a ``parentId`` points at a real
|
||||
container, not at a non-container node;
|
||||
* no cycles in the parent chain (a node cannot be its own
|
||||
@@ -1959,7 +2306,9 @@ class WorkflowGenerator:
|
||||
if not isinstance(parent, str) or not parent:
|
||||
continue
|
||||
if parent in container_ids:
|
||||
children_by_parent.setdefault(parent, []).append(n.get("id", ""))
|
||||
node_type = (n.get("data") or {}).get("type")
|
||||
if node_type not in {"iteration-start", "loop-start"}:
|
||||
children_by_parent.setdefault(parent, []).append(n.get("id", ""))
|
||||
elif parent in by_id:
|
||||
# Parent exists but isn't a container — that's a topology bug.
|
||||
out.append(
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
"""
|
||||
Typed payloads for workflow generation.
|
||||
|
||||
These TypedDicts describe the shape that the planner and builder LLM calls are
|
||||
required to return after ``json_repair`` parsing. They mirror the runtime
|
||||
``graph`` shape consumed by ``WorkflowService.sync_draft_workflow`` so the output
|
||||
can be written straight into a draft workflow without further translation.
|
||||
These TypedDicts describe the planner payload and the runtime graph assembled
|
||||
from builder LLM responses after ``json_repair`` parsing. The graph types mirror
|
||||
the shape consumed by ``WorkflowService.sync_draft_workflow`` so the output can
|
||||
be written straight into a draft workflow.
|
||||
"""
|
||||
|
||||
from enum import StrEnum
|
||||
@@ -12,11 +12,10 @@ from typing import Literal, NotRequired, TypedDict
|
||||
|
||||
WorkflowGenerationMode = Literal["workflow", "advanced-chat"]
|
||||
|
||||
# The mode accepted at the API boundary. ``auto`` is a sentinel that asks the
|
||||
# service to classify the instruction into a concrete ``WorkflowGenerationMode``
|
||||
# (one tiny LLM call) BEFORE planning — see
|
||||
# ``WorkflowGeneratorService._resolve_mode`` and
|
||||
# ``LLMGenerator.classify_workflow_mode``.
|
||||
# The mode accepted at the API boundary. ``auto`` is a sentinel that delegates
|
||||
# the choice to the planner: it echoes a concrete mode in its ``mode`` output
|
||||
# field (falling back to terminal-node inference, then ``advanced-chat``) —
|
||||
# see ``runner._resolve_generation_mode``. No extra LLM call is involved.
|
||||
WorkflowGenerationModeRequest = Literal["workflow", "advanced-chat", "auto"]
|
||||
|
||||
|
||||
@@ -58,9 +57,21 @@ class WorkflowGenerateErrorDict(TypedDict):
|
||||
class PlannerNodeDict(TypedDict):
|
||||
"""One node from the planner's high-level plan."""
|
||||
|
||||
id: NotRequired[str]
|
||||
label: str
|
||||
node_type: str
|
||||
purpose: str
|
||||
parent: NotRequired[str]
|
||||
action: NotRequired[Literal["keep", "update", "add"]]
|
||||
|
||||
|
||||
class PlannerEdgeDict(TypedDict):
|
||||
"""Compact topology emitted by the planner for parallel node building."""
|
||||
|
||||
source: str
|
||||
target: str
|
||||
source_handle: NotRequired[str]
|
||||
target_handle: NotRequired[str]
|
||||
|
||||
|
||||
class PlannerStartInputDict(TypedDict):
|
||||
@@ -82,10 +93,15 @@ class PlannerResultDict(TypedDict):
|
||||
|
||||
title: str
|
||||
description: str
|
||||
# Concrete mode the planner chose ("workflow" / "advanced-chat"). Parsed
|
||||
# leniently — an ``auto`` request infers the mode from the terminal node
|
||||
# when this is missing or invalid, so a bad value never fails the plan.
|
||||
mode: NotRequired[str]
|
||||
app_name: NotRequired[str]
|
||||
icon: NotRequired[str]
|
||||
start_inputs: NotRequired[list[PlannerStartInputDict]]
|
||||
nodes: list[PlannerNodeDict]
|
||||
edges: NotRequired[list[PlannerEdgeDict]]
|
||||
|
||||
|
||||
class GraphNodePositionDict(TypedDict):
|
||||
|
||||
@@ -499,6 +499,9 @@ class DifyNodeFactory(NodeFactory):
|
||||
base_url=dify_config.AGENT_BACKEND_BASE_URL,
|
||||
use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
|
||||
fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
|
||||
stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS,
|
||||
stream_max_reconnects=dify_config.AGENT_BACKEND_STREAM_MAX_RECONNECTS,
|
||||
stream_run_timeout_seconds=dify_config.AGENT_BACKEND_RUN_TIMEOUT_SECONDS,
|
||||
),
|
||||
"event_adapter": AgentBackendRunEventAdapter(),
|
||||
# Agent Files §4.6: reback file outputs from the ToolFile row so
|
||||
|
||||
@@ -5,6 +5,7 @@ from collections.abc import Generator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, override
|
||||
|
||||
from agenton.compositor import CompositorSessionSnapshot
|
||||
from dify_agent.protocol import CancelRunRequest
|
||||
|
||||
from clients.agent_backend import (
|
||||
AgentBackendAgentMessageDeltaInternalEvent,
|
||||
@@ -473,7 +474,10 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
|
||||
"""
|
||||
stream_event_count = 0
|
||||
try:
|
||||
for public_event in self._agent_backend_client.stream_events(run_id):
|
||||
for public_event in self._agent_backend_client.stream_events(
|
||||
run_id,
|
||||
should_stop=self._is_graph_aborted,
|
||||
):
|
||||
stream_event_count += 1
|
||||
for internal_event in self._event_adapter.adapt(public_event):
|
||||
if internal_event.type == AgentBackendInternalEventType.RUN_STARTED:
|
||||
@@ -501,6 +505,7 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
|
||||
| AgentBackendDeferredToolCallInternalEvent,
|
||||
):
|
||||
return internal_event, None
|
||||
self._cancel_backend_run(run_id, reason="unexpected_event")
|
||||
return None, self._failure_event(
|
||||
inputs={},
|
||||
process_data={},
|
||||
@@ -509,6 +514,7 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
|
||||
error_type="agent_backend_stream_error",
|
||||
)
|
||||
except AgentBackendError as error:
|
||||
self._cancel_backend_run(run_id, reason=self._stream_stop_reason())
|
||||
return None, self._failure_event(
|
||||
inputs={},
|
||||
process_data={},
|
||||
@@ -517,6 +523,7 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
|
||||
error_type=self._agent_backend_error_type(error),
|
||||
)
|
||||
except Exception as error:
|
||||
self._cancel_backend_run(run_id, reason=self._stream_stop_reason())
|
||||
return None, self._failure_event(
|
||||
inputs={},
|
||||
process_data={},
|
||||
@@ -525,8 +532,28 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
|
||||
error_type="agent_backend_stream_error",
|
||||
)
|
||||
|
||||
self._cancel_backend_run(run_id, reason="stream_ended_without_terminal_event")
|
||||
return None, None
|
||||
|
||||
def _is_graph_aborted(self) -> bool:
|
||||
"""Let Agent SSE consumption observe GraphEngine's cooperative abort state."""
|
||||
try:
|
||||
return self.graph_runtime_state.graph_execution.aborted
|
||||
except (AttributeError, RuntimeError):
|
||||
return False
|
||||
|
||||
def _stream_stop_reason(self) -> str:
|
||||
return "workflow_graph_aborted" if self._is_graph_aborted() else "event_stream_failed"
|
||||
|
||||
def _cancel_backend_run(self, run_id: str, *, reason: str) -> None:
|
||||
try:
|
||||
self._agent_backend_client.cancel_run(
|
||||
run_id,
|
||||
CancelRunRequest(reason=reason, message="Workflow Agent event consumption stopped"),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to cancel Workflow Agent backend run: run_id=%s", run_id, exc_info=True)
|
||||
|
||||
@staticmethod
|
||||
def _record_type_check_metadata(metadata: dict[str, Any], outcome: OutputTypeCheckOutcome) -> None:
|
||||
# Surface enough detail in metadata for Inspector / debug logs without
|
||||
|
||||
@@ -61,16 +61,33 @@ class WorkflowAgentNodeValidator:
|
||||
|
||||
@classmethod
|
||||
def validate_draft_workflow(cls, *, session: Session, workflow: Workflow) -> None:
|
||||
cls._validate_workflow(session=session, workflow=workflow, require_binding=False)
|
||||
cls._validate_workflow(
|
||||
session=session,
|
||||
workflow=workflow,
|
||||
require_binding=False,
|
||||
validate_previous_node_topology=False,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def validate_published_workflow(cls, *, session: Session, workflow: Workflow) -> None:
|
||||
cls._validate_workflow(session=session, workflow=workflow, require_binding=True)
|
||||
cls._validate_workflow(
|
||||
session=session,
|
||||
workflow=workflow,
|
||||
require_binding=True,
|
||||
validate_previous_node_topology=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _validate_workflow(cls, *, session: Session, workflow: Workflow, require_binding: bool) -> None:
|
||||
def _validate_workflow(
|
||||
cls,
|
||||
*,
|
||||
session: Session,
|
||||
workflow: Workflow,
|
||||
require_binding: bool,
|
||||
validate_previous_node_topology: bool,
|
||||
) -> None:
|
||||
graph = workflow.graph_dict
|
||||
topology = _WorkflowGraphTopology.from_graph(graph)
|
||||
topology = _WorkflowGraphTopology.from_graph(graph) if validate_previous_node_topology else None
|
||||
for node_id, node_data in cls.iter_agent_v2_nodes(graph):
|
||||
cls._validate_node_schema(node_id=node_id, node_data=node_data)
|
||||
binding = cls._find_binding(
|
||||
@@ -185,12 +202,12 @@ class WorkflowAgentNodeValidator:
|
||||
raise WorkflowAgentNodeValidationError(
|
||||
f"Workflow Agent node {binding.node_id} has invalid previous node output ref."
|
||||
)
|
||||
if topology is None:
|
||||
continue
|
||||
if len(selector) < 2:
|
||||
raise WorkflowAgentNodeValidationError(
|
||||
f"Workflow Agent node {binding.node_id} has incomplete previous node output ref."
|
||||
)
|
||||
if topology is None:
|
||||
continue
|
||||
source_node_id = selector[0]
|
||||
if not topology.has_node(source_node_id):
|
||||
raise WorkflowAgentNodeValidationError(
|
||||
|
||||
@@ -7,6 +7,9 @@ from .delete_tool_parameters_cache_when_sync_draft_workflow import (
|
||||
handle as handle_delete_tool_parameters_cache_when_sync_draft_workflow,
|
||||
)
|
||||
from .queue_credential_sync_when_tenant_created import handle as handle_queue_credential_sync_when_tenant_created
|
||||
from .queue_default_plugin_install_when_tenant_created import (
|
||||
handle as handle_queue_default_plugin_install_when_tenant_created,
|
||||
)
|
||||
from .sync_plugin_trigger_when_app_created import handle as handle_sync_plugin_trigger_when_app_created
|
||||
from .sync_webhook_when_app_created import handle as handle_sync_webhook_when_app_created
|
||||
from .sync_workflow_schedule_when_app_published import handle as handle_sync_workflow_schedule_when_app_published
|
||||
@@ -32,6 +35,7 @@ __all__ = [
|
||||
"handle_create_site_record_when_app_created",
|
||||
"handle_delete_tool_parameters_cache_when_sync_draft_workflow",
|
||||
"handle_queue_credential_sync_when_tenant_created",
|
||||
"handle_queue_default_plugin_install_when_tenant_created",
|
||||
"handle_sync_plugin_trigger_when_app_created",
|
||||
"handle_sync_webhook_when_app_created",
|
||||
"handle_sync_workflow_schedule_when_app_published",
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Queue default marketplace plugin installation after tenant creation."""
|
||||
|
||||
import logging
|
||||
|
||||
from configs import dify_config
|
||||
from events.tenant_event import tenant_was_created
|
||||
from tasks.install_default_plugins_task import install_default_plugins_task
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@tenant_was_created.connect
|
||||
def handle(sender, **kwargs) -> None:
|
||||
"""Keep tenant creation non-blocking while installing configured plugins asynchronously."""
|
||||
plugin_ids = dify_config.NEW_USER_DEFAULT_PLUGIN_ID_LIST
|
||||
if not plugin_ids:
|
||||
return
|
||||
|
||||
try:
|
||||
install_default_plugins_task.delay(sender.id, plugin_ids)
|
||||
except Exception:
|
||||
logger.exception("Failed to queue default plugin installation for tenant %s", sender.id)
|
||||
@@ -156,6 +156,7 @@ def init_app(app: DifyApp) -> Celery:
|
||||
"tasks.generate_summary_index_task", # summary index generation
|
||||
"tasks.regenerate_summary_index_task", # summary index regeneration
|
||||
"tasks.initialize_created_app_rbac_access_task", # app access initialization
|
||||
"tasks.install_default_plugins_task", # tenant default plugin installation
|
||||
"tasks.app_generate.resume_agent_app_task", # ENG-635: Agent v2 chat ask_human resume
|
||||
"tasks.workflow_run_archive_download_tasks", # workflow-run archive download preparation
|
||||
]
|
||||
|
||||
@@ -12,7 +12,6 @@ from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchem
|
||||
from extensions.logstore.aliyun_logstore import AliyunLogStore
|
||||
from graphon.entities import WorkflowExecution
|
||||
from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter
|
||||
from libs.helper import extract_tenant_id
|
||||
from models import (
|
||||
Account,
|
||||
CreatorUserRole,
|
||||
@@ -27,6 +26,7 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
def __init__(
|
||||
self,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str | None,
|
||||
triggered_from: WorkflowRunTriggeredFrom | None,
|
||||
@@ -36,7 +36,8 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine for creating sessions
|
||||
user: Account or EndUser object containing tenant_id, user ID, and role information
|
||||
tenant_id: Tenant that owns the workflow execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: App ID for filtering by application (can be None)
|
||||
triggered_from: Source of the execution trigger (DEBUGGING or APP_RUN)
|
||||
"""
|
||||
@@ -47,10 +48,8 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
# Note: Project/logstore/index initialization is done at app startup via ext_logstore
|
||||
self.logstore_client = AliyunLogStore()
|
||||
|
||||
# Extract tenant_id from user
|
||||
tenant_id = extract_tenant_id(user)
|
||||
if not tenant_id:
|
||||
raise ValueError("User must have a tenant_id or current_tenant_id")
|
||||
raise ValueError("tenant_id is required")
|
||||
self._tenant_id = tenant_id
|
||||
|
||||
# Store app context
|
||||
@@ -64,7 +63,13 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository):
|
||||
self._creator_user_role = CreatorUserRole.ACCOUNT if isinstance(user, Account) else CreatorUserRole.END_USER
|
||||
|
||||
# Initialize SQL repository for dual-write support
|
||||
self.sql_repository = SQLAlchemyWorkflowExecutionRepository(session_factory, user, app_id, triggered_from)
|
||||
self.sql_repository = SQLAlchemyWorkflowExecutionRepository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=tenant_id,
|
||||
user=user,
|
||||
app_id=app_id,
|
||||
triggered_from=triggered_from,
|
||||
)
|
||||
|
||||
# Control flag for dual-write (write to both LogStore and SQL database)
|
||||
# Set to True to enable dual-write for safe migration, False to use LogStore only
|
||||
|
||||
+11
-6
@@ -26,7 +26,6 @@ from graphon.entities import WorkflowNodeExecution
|
||||
from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
|
||||
from graphon.model_runtime.utils.encoders import jsonable_encoder
|
||||
from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter
|
||||
from libs.helper import extract_tenant_id
|
||||
from models import (
|
||||
Account,
|
||||
CreatorUserRole,
|
||||
@@ -109,6 +108,7 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
|
||||
def __init__(
|
||||
self,
|
||||
session_factory: sessionmaker | Engine,
|
||||
tenant_id: str,
|
||||
user: Account | EndUser,
|
||||
app_id: str | None,
|
||||
triggered_from: WorkflowNodeExecutionTriggeredFrom | None,
|
||||
@@ -118,7 +118,8 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
|
||||
|
||||
Args:
|
||||
session_factory: SQLAlchemy sessionmaker or engine for creating sessions
|
||||
user: Account or EndUser object containing tenant_id, user ID, and role information
|
||||
tenant_id: Tenant that owns the workflow node execution
|
||||
user: Account or EndUser used for creator attribution
|
||||
app_id: App ID for filtering by application (can be None)
|
||||
triggered_from: Source of the execution trigger (SINGLE_STEP or WORKFLOW_RUN)
|
||||
"""
|
||||
@@ -128,10 +129,8 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
|
||||
# Initialize LogStore client
|
||||
self.logstore_client = AliyunLogStore()
|
||||
|
||||
# Extract tenant_id from user
|
||||
tenant_id = extract_tenant_id(user)
|
||||
if not tenant_id:
|
||||
raise ValueError("User must have a tenant_id or current_tenant_id")
|
||||
raise ValueError("tenant_id is required")
|
||||
self._tenant_id = tenant_id
|
||||
|
||||
# Store app context
|
||||
@@ -145,7 +144,13 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
|
||||
self._creator_user_role = CreatorUserRole.ACCOUNT if isinstance(user, Account) else CreatorUserRole.END_USER
|
||||
|
||||
# Initialize SQL repository for dual-write support
|
||||
self.sql_repository = SQLAlchemyWorkflowNodeExecutionRepository(session_factory, user, app_id, triggered_from)
|
||||
self.sql_repository = SQLAlchemyWorkflowNodeExecutionRepository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=tenant_id,
|
||||
user=user,
|
||||
app_id=app_id,
|
||||
triggered_from=triggered_from,
|
||||
)
|
||||
|
||||
# Control flag for dual-write (write to both LogStore and SQL database)
|
||||
# Set to True to enable dual-write for safe migration, False to use LogStore only
|
||||
|
||||
@@ -60,6 +60,21 @@ class ConversationVariableResponse(ResponseModel):
|
||||
return to_timestamp(value)
|
||||
|
||||
|
||||
class WorkflowConversationVariableResponse(ResponseModel):
|
||||
id: str
|
||||
name: str
|
||||
value_type: str
|
||||
value: Any
|
||||
description: str
|
||||
|
||||
@field_validator("value_type", mode="before")
|
||||
@classmethod
|
||||
def _serialize_value_type(cls, value: Any) -> str:
|
||||
if hasattr(value, "exposed_type"):
|
||||
return str(value.exposed_type())
|
||||
return str(value)
|
||||
|
||||
|
||||
class PaginatedConversationVariableResponse(ResponseModel):
|
||||
page: int
|
||||
limit: int
|
||||
|
||||
@@ -188,6 +188,14 @@ def validate_config_skill_name(name: str) -> str:
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalize_legacy_missing_asset_file_id(value: Any) -> Any:
|
||||
"""Canonicalize the null placeholder emitted by early portable Agent DSLs."""
|
||||
|
||||
if isinstance(value, dict) and value.get("is_missing") is True and value.get("file_id") is None:
|
||||
return {**value, "file_id": ""}
|
||||
return value
|
||||
|
||||
|
||||
class AgentConfigFileRefConfig(BaseModel):
|
||||
"""Stable Agent Soul reference to one config file payload."""
|
||||
|
||||
@@ -201,6 +209,11 @@ class AgentConfigFileRefConfig(BaseModel):
|
||||
hash: str | None = None
|
||||
mime_type: str | None = None
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _normalize_legacy_file_id(cls, value: Any) -> Any:
|
||||
return _normalize_legacy_missing_asset_file_id(value)
|
||||
|
||||
@field_validator("name")
|
||||
@classmethod
|
||||
def _validate_name(cls, value: str) -> str:
|
||||
@@ -231,6 +244,11 @@ class AgentConfigSkillRefConfig(BaseModel):
|
||||
hash: str | None = None
|
||||
mime_type: str | None = "application/zip"
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _normalize_legacy_file_id(cls, value: Any) -> Any:
|
||||
return _normalize_legacy_missing_asset_file_id(value)
|
||||
|
||||
@field_validator("name")
|
||||
@classmethod
|
||||
def _validate_name(cls, value: str) -> str:
|
||||
|
||||
@@ -22008,6 +22008,7 @@ Tag type
|
||||
| role | string | | No |
|
||||
| status | string | | No |
|
||||
| trial_credits | integer | | No |
|
||||
| trial_credits_exhausted_at | integer | | No |
|
||||
| trial_credits_used | integer | | No |
|
||||
| trial_end_reason | string | | No |
|
||||
|
||||
@@ -22557,7 +22558,7 @@ Enum class for tool provider
|
||||
|
||||
| Name | Type | Description | Required |
|
||||
| ---- | ---- | ----------- | -------- |
|
||||
| conversation_variables | [ [JsonObject](#jsonobject) ] | | No |
|
||||
| conversation_variables | [ [WorkflowConversationVariableResponse](#workflowconversationvariableresponse) ] | | No |
|
||||
| created_at | integer | | No |
|
||||
| created_by | [TrialSimpleAccount](#trialsimpleaccount) | | No |
|
||||
| environment_variables | [ [JsonObject](#jsonobject) ] | | No |
|
||||
|
||||
@@ -296,6 +296,7 @@ class AliyunDataTrace(BaseTraceInstance):
|
||||
session_factory = sessionmaker(bind=db.engine)
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=trace_info.tenant_id,
|
||||
user=service_account,
|
||||
app_id=app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
+1
@@ -845,6 +845,7 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=trace_info.tenant_id,
|
||||
user=service_account,
|
||||
app_id=app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
@@ -186,6 +186,7 @@ class LangFuseDataTrace(BaseTraceInstance):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=trace_info.tenant_id,
|
||||
user=service_account,
|
||||
app_id=app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
@@ -154,6 +154,7 @@ class LangSmithDataTrace(BaseTraceInstance):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=trace_info.tenant_id,
|
||||
user=service_account,
|
||||
app_id=app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
@@ -174,6 +174,7 @@ class OpikDataTrace(BaseTraceInstance):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=trace_info.tenant_id,
|
||||
user=service_account,
|
||||
app_id=app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
@@ -252,18 +252,9 @@ class TencentDataTrace(BaseTraceInstance):
|
||||
if not service_account:
|
||||
raise ValueError(f"Creator account not found for app {app_id}")
|
||||
|
||||
current_tenant = session.scalar(
|
||||
select(TenantAccountJoin)
|
||||
.where(TenantAccountJoin.account_id == service_account.id, TenantAccountJoin.current.is_(True))
|
||||
.limit(1)
|
||||
)
|
||||
if not current_tenant:
|
||||
raise ValueError(f"Current tenant not found for account {service_account.id}")
|
||||
|
||||
service_account.set_tenant_id_with_session(current_tenant.tenant_id, session=session)
|
||||
|
||||
repository = SQLAlchemyWorkflowNodeExecutionRepository(
|
||||
session_factory=session_maker,
|
||||
tenant_id=app.tenant_id,
|
||||
user=service_account,
|
||||
app_id=trace_info.metadata.get("app_id"),
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
+4
-6
@@ -18,7 +18,7 @@ from core.ops.entities.trace_entity import (
|
||||
)
|
||||
from graphon.entities import WorkflowNodeExecution
|
||||
from graphon.enums import BuiltinNodeTypes
|
||||
from models import Account, App, TenantAccountJoin
|
||||
from models import Account, App
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -420,20 +420,18 @@ class TestTencentDataTrace:
|
||||
app = MagicMock(spec=App)
|
||||
app.id = "app-1"
|
||||
app.created_by = "user-1"
|
||||
app.tenant_id = "tenant-1"
|
||||
|
||||
account = MagicMock(spec=Account)
|
||||
account.id = "user-1"
|
||||
|
||||
tenant_join = MagicMock(spec=TenantAccountJoin)
|
||||
tenant_join.tenant_id = "tenant-1"
|
||||
|
||||
mock_executions = [MagicMock()]
|
||||
|
||||
with patch("dify_trace_tencent.tencent_trace.db") as mock_db:
|
||||
mock_db.engine = "engine"
|
||||
with patch("dify_trace_tencent.tencent_trace.Session") as mock_session_ctx:
|
||||
session = mock_session_ctx.return_value.__enter__.return_value
|
||||
session.scalar.side_effect = [app, account, tenant_join]
|
||||
session.scalar.side_effect = [app, account]
|
||||
|
||||
with patch("dify_trace_tencent.tencent_trace.SQLAlchemyWorkflowNodeExecutionRepository") as mock_repo:
|
||||
mock_repo.return_value.get_by_workflow_execution.return_value = mock_executions
|
||||
@@ -441,7 +439,7 @@ class TestTencentDataTrace:
|
||||
results = tencent_data_trace._get_workflow_node_executions(trace_info)
|
||||
|
||||
assert results == mock_executions
|
||||
account.set_tenant_id_with_session.assert_called_once_with("tenant-1", session=session)
|
||||
assert mock_repo.call_args.kwargs["tenant_id"] == "tenant-1"
|
||||
|
||||
def test_get_workflow_node_executions_no_app_id(self, tencent_data_trace, caplog: pytest.LogCaptureFixture):
|
||||
trace_info = MagicMock(spec=WorkflowTraceInfo)
|
||||
|
||||
@@ -159,6 +159,7 @@ class WeaveDataTrace(BaseTraceInstance):
|
||||
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=trace_info.tenant_id,
|
||||
user=service_account,
|
||||
app_id=app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "dify-api"
|
||||
version = "1.16.0-rc1"
|
||||
version = "1.16.0"
|
||||
requires-python = "~=3.12.0"
|
||||
|
||||
dependencies = [
|
||||
|
||||
@@ -1297,7 +1297,7 @@ class TenantService:
|
||||
def create_owner_tenant_if_not_exist(
|
||||
account: Account, name: str | None = None, is_setup: bool | None = False, *, session: Session
|
||||
):
|
||||
"""Check if user have a workspace or not"""
|
||||
"""Create an owner workspace only when the account has no membership."""
|
||||
available_ta = session.scalar(
|
||||
select(TenantAccountJoin)
|
||||
.where(TenantAccountJoin.account_id == account.id)
|
||||
@@ -1308,8 +1308,28 @@ class TenantService:
|
||||
if available_ta:
|
||||
return
|
||||
|
||||
"""Create owner tenant if not exist"""
|
||||
if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup:
|
||||
TenantService.create_owner_tenant(account, name=name, is_setup=is_setup, session=session)
|
||||
|
||||
@staticmethod
|
||||
def create_owner_tenant(
|
||||
account: Account,
|
||||
name: str | None = None,
|
||||
is_setup: bool | None = False,
|
||||
is_from_dashboard: bool | None = False,
|
||||
*,
|
||||
session: Session,
|
||||
) -> Tenant:
|
||||
"""Create an owner workspace and bind its owner RBAC role when enabled.
|
||||
|
||||
This is the single write path for a newly created workspace with an
|
||||
owner. It persists the legacy membership before creating the matching
|
||||
RBAC role binding, then makes the workspace current for the account.
|
||||
"""
|
||||
if (
|
||||
not FeatureService.get_system_features().is_allow_create_workspace
|
||||
and not is_setup
|
||||
and not is_from_dashboard
|
||||
):
|
||||
raise WorkSpaceNotAllowedCreateError()
|
||||
|
||||
workspaces = FeatureService.get_system_features().license.workspaces
|
||||
@@ -1317,9 +1337,19 @@ class TenantService:
|
||||
raise WorkspacesLimitExceededError()
|
||||
|
||||
if name:
|
||||
tenant = TenantService.create_tenant(name=name, is_setup=is_setup, session=session)
|
||||
tenant = TenantService.create_tenant(
|
||||
name=name,
|
||||
is_setup=is_setup,
|
||||
is_from_dashboard=is_from_dashboard,
|
||||
session=session,
|
||||
)
|
||||
else:
|
||||
tenant = TenantService.create_tenant(name=f"{account.name}'s Workspace", is_setup=is_setup, session=session)
|
||||
tenant = TenantService.create_tenant(
|
||||
name=f"{account.name}'s Workspace",
|
||||
is_setup=is_setup,
|
||||
is_from_dashboard=is_from_dashboard,
|
||||
session=session,
|
||||
)
|
||||
TenantService.create_tenant_member(tenant, account, session, role="owner")
|
||||
if dify_config.RBAC_ENABLED:
|
||||
owner_role_id = AccountService._resolve_legacy_role_id(str(tenant.id), account.id, TenantAccountRole.OWNER)
|
||||
@@ -1333,6 +1363,7 @@ class TenantService:
|
||||
account.set_current_tenant_with_session(tenant, session=session)
|
||||
session.commit()
|
||||
tenant_was_created.send(tenant)
|
||||
return tenant
|
||||
|
||||
@staticmethod
|
||||
def create_tenant_member(
|
||||
@@ -1984,10 +2015,7 @@ class RegisterService:
|
||||
and FeatureService.get_system_features().license.workspaces.is_available()
|
||||
):
|
||||
try:
|
||||
tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=session)
|
||||
TenantService.create_tenant_member(tenant, account, session, role="owner")
|
||||
account.set_current_tenant_with_session(tenant, session=session)
|
||||
tenant_was_created.send(tenant)
|
||||
TenantService.create_owner_tenant(account, session=session)
|
||||
except Exception:
|
||||
_try_join_enterprise_default_workspace(str(account.id))
|
||||
raise
|
||||
|
||||
@@ -11,17 +11,14 @@ from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from sqlalchemy import event, func, select
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from constants.model_template import default_app_templates
|
||||
from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidator
|
||||
from events.app_event import app_was_created
|
||||
from graphon.enums import BuiltinNodeTypes
|
||||
from models import Account
|
||||
from models.agent import (
|
||||
@@ -41,7 +38,7 @@ from models.agent import (
|
||||
WorkflowAgentNodeBinding,
|
||||
)
|
||||
from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig
|
||||
from models.model import App, AppMode, AppModelConfig, IconType
|
||||
from models.model import App, AppModelConfig
|
||||
from models.workflow import Workflow
|
||||
from services.agent.agent_soul_state import agent_soul_has_model
|
||||
from services.agent.dsl_entities import (
|
||||
@@ -57,8 +54,6 @@ from services.agent.roster_service import AgentRosterService
|
||||
from services.entities.dsl_entities import DslImportWarning
|
||||
from services.plugin.dependencies_analysis import DependenciesAnalysisService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AgentPackageImportResult(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
@@ -249,7 +244,7 @@ class AgentDslService:
|
||||
raw_packages: Mapping[str, Any],
|
||||
account: Account,
|
||||
) -> tuple[dict[str, Any], list[DslImportWarning]]:
|
||||
"""Materialize packages and bindings for a Workflow or Snippet draft."""
|
||||
"""Materialize every packaged Agent as a node-owned inline Agent."""
|
||||
|
||||
graph = copy.deepcopy(dict(portable_graph))
|
||||
packages = {key: AgentPackage.model_validate(value) for key, value in raw_packages.items()}
|
||||
@@ -264,7 +259,6 @@ class AgentDslService:
|
||||
for binding in previous_bindings:
|
||||
self.session.delete(binding)
|
||||
self.session.flush()
|
||||
imported_roster: dict[str, AgentPackageImportResult] = {}
|
||||
warnings: list[DslImportWarning] = []
|
||||
|
||||
for node_id, raw_node_data in WorkflowAgentNodeValidator.iter_agent_v2_nodes(graph):
|
||||
@@ -280,28 +274,17 @@ class AgentDslService:
|
||||
raise ValueError(f"Workflow Agent node {node_id} references unknown package {package_ref!r}.")
|
||||
|
||||
try:
|
||||
binding_type = WorkflowAgentBindingType(str(raw_binding.get("binding_type")))
|
||||
WorkflowAgentBindingType(str(raw_binding.get("binding_type")))
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"Workflow Agent node {node_id} has an invalid binding type.") from exc
|
||||
|
||||
if binding_type == WorkflowAgentBindingType.ROSTER_AGENT:
|
||||
imported = imported_roster.get(package_ref)
|
||||
if imported is None:
|
||||
imported = self._create_imported_roster_agent_app(
|
||||
tenant_id=workflow.tenant_id,
|
||||
account=account,
|
||||
package=package,
|
||||
package_path=f"agent_packages.{package_ref}",
|
||||
)
|
||||
imported_roster[package_ref] = imported
|
||||
else:
|
||||
imported = self._create_imported_inline_agent(
|
||||
workflow=workflow,
|
||||
node_id=node_id,
|
||||
account=account,
|
||||
package=package,
|
||||
package_path=f"agent_packages.{package_ref}",
|
||||
)
|
||||
imported = self._create_imported_inline_agent(
|
||||
workflow=workflow,
|
||||
node_id=node_id,
|
||||
account=account,
|
||||
package=package,
|
||||
package_path=f"agent_packages.{package_ref}",
|
||||
)
|
||||
|
||||
node_job = WorkflowNodeJobConfig.model_validate(node_data.get(AGENT_NODE_JOB_DSL_KEY) or {})
|
||||
self.session.add(
|
||||
@@ -311,7 +294,7 @@ class AgentDslService:
|
||||
workflow_id=workflow.id,
|
||||
workflow_version=workflow.version,
|
||||
node_id=node_id,
|
||||
binding_type=binding_type,
|
||||
binding_type=WorkflowAgentBindingType.INLINE_AGENT,
|
||||
agent_id=imported.agent.id,
|
||||
current_snapshot_id=imported.snapshot.id,
|
||||
node_job_config=node_job,
|
||||
@@ -320,7 +303,7 @@ class AgentDslService:
|
||||
)
|
||||
)
|
||||
node_data["agent_binding"] = {
|
||||
"binding_type": binding_type.value,
|
||||
"binding_type": WorkflowAgentBindingType.INLINE_AGENT.value,
|
||||
"agent_id": imported.agent.id,
|
||||
"current_snapshot_id": imported.snapshot.id,
|
||||
}
|
||||
@@ -402,43 +385,6 @@ class AgentDslService:
|
||||
)
|
||||
return dependencies
|
||||
|
||||
def _create_imported_roster_agent_app(
|
||||
self,
|
||||
*,
|
||||
tenant_id: str,
|
||||
account: Account,
|
||||
package: AgentPackage,
|
||||
package_path: str,
|
||||
) -> AgentPackageImportResult:
|
||||
metadata = package.metadata
|
||||
app_template = dict(default_app_templates[AppMode.AGENT]["app"])
|
||||
app = App(**app_template)
|
||||
app.name = metadata.name
|
||||
app.description = metadata.description
|
||||
app.mode = AppMode.AGENT
|
||||
app.icon_type = self._app_icon_type(metadata.icon_type)
|
||||
app.icon = metadata.icon
|
||||
app.icon_background = metadata.icon_background
|
||||
app.tenant_id = tenant_id
|
||||
app.enable_site = True
|
||||
app.enable_api = True
|
||||
app.created_by = account.id
|
||||
app.maintainer = account.id
|
||||
app.updated_by = account.id
|
||||
self.session.add(app)
|
||||
self.session.flush()
|
||||
app_was_created.send(app, account=account, session=self.session)
|
||||
self._configure_visible_agent_app_after_commit(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app.id,
|
||||
account_id=account.id,
|
||||
)
|
||||
result = self.import_agent_app_package(app=app, account=account, package=package)
|
||||
result.warnings = [
|
||||
warning.model_copy(update={"path": f"{package_path}.{warning.path}"}) for warning in result.warnings
|
||||
]
|
||||
return result
|
||||
|
||||
def _create_imported_inline_agent(
|
||||
self,
|
||||
*,
|
||||
@@ -654,28 +600,6 @@ class AgentDslService:
|
||||
)
|
||||
return next(candidate for candidate in candidates if candidate not in existing)
|
||||
|
||||
def _configure_visible_agent_app_after_commit(self, *, tenant_id: str, app_id: str, account_id: str) -> None:
|
||||
"""Apply external RBAC and web-app visibility only after the DB transaction commits."""
|
||||
|
||||
def configure(_session: Session) -> None:
|
||||
try:
|
||||
from services.enterprise import rbac_service as enterprise_rbac_service
|
||||
from services.enterprise.enterprise_service import EnterpriseService
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
enterprise_rbac_service.try_sync_creator_access_policy_member_bindings(
|
||||
tenant_id,
|
||||
account_id,
|
||||
enterprise_rbac_service.RBACResourceType.APP,
|
||||
app_id,
|
||||
)
|
||||
if FeatureService.get_system_features().webapp_auth.enabled:
|
||||
EnterpriseService.WebAppAuth.update_app_access_mode(app_id, "private")
|
||||
except Exception:
|
||||
logger.exception("Failed to configure imported Agent App %s after commit", app_id)
|
||||
|
||||
event.listen(self.session, "after_commit", configure, once=True)
|
||||
|
||||
def _require_agent(self, *, tenant_id: str, agent_id: str) -> Agent:
|
||||
agent = self.session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1))
|
||||
if agent is None:
|
||||
@@ -702,10 +626,6 @@ class AgentDslService:
|
||||
def _agent_icon_type(value: str | None) -> AgentIconType | None:
|
||||
return AgentIconType(value) if value else None
|
||||
|
||||
@staticmethod
|
||||
def _app_icon_type(value: str | None) -> IconType:
|
||||
return IconType(value) if value else IconType.EMOJI
|
||||
|
||||
|
||||
def is_agent_v2_graph(graph: Mapping[str, Any]) -> bool:
|
||||
return any(
|
||||
|
||||
@@ -55,6 +55,7 @@ class QuotaBalanceResult(TypedDict):
|
||||
reserved: int
|
||||
quota: int
|
||||
usage: int
|
||||
exhausted_at: NotRequired[int]
|
||||
|
||||
|
||||
class QuotaConsumeCappedResult(TypedDict):
|
||||
|
||||
@@ -32,6 +32,7 @@ class CreditPoolBalance:
|
||||
pool_type: str
|
||||
quota_limit: int
|
||||
quota_used: int
|
||||
exhausted_at: int | None = None
|
||||
|
||||
@property
|
||||
def remaining_credits(self) -> int:
|
||||
@@ -133,6 +134,7 @@ class CreditPoolService:
|
||||
pool_type=normalized_pool_type,
|
||||
quota_limit=balance["quota"],
|
||||
quota_used=balance["usage"],
|
||||
exhausted_at=balance.get("exhausted_at"),
|
||||
)
|
||||
|
||||
session = cls._require_session(session)
|
||||
|
||||
@@ -53,6 +53,7 @@ class FileService:
|
||||
content: bytes,
|
||||
mimetype: str,
|
||||
user: Account | EndUser,
|
||||
tenant_id: str | None = None,
|
||||
source: Literal["datasets"] | None = None,
|
||||
source_url: str = "",
|
||||
) -> UploadFile:
|
||||
@@ -84,16 +85,16 @@ class FileService:
|
||||
# generate file key
|
||||
file_uuid = str(uuid.uuid4())
|
||||
|
||||
current_tenant_id = extract_tenant_id(user)
|
||||
resource_tenant_id = tenant_id if tenant_id is not None else extract_tenant_id(user)
|
||||
|
||||
file_key = "upload_files/" + (current_tenant_id or "") + "/" + file_uuid + "." + extension
|
||||
file_key = "upload_files/" + (resource_tenant_id or "") + "/" + file_uuid + "." + extension
|
||||
|
||||
# save file to storage
|
||||
storage.save(file_key, content)
|
||||
|
||||
# save file to db
|
||||
upload_file = UploadFile(
|
||||
tenant_id=current_tenant_id or "",
|
||||
tenant_id=resource_tenant_id or "",
|
||||
storage_type=StorageType(dify_config.STORAGE_TYPE),
|
||||
key=file_key,
|
||||
name=filename,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,369 +0,0 @@
|
||||
"""Transport-only forwarding for the allowlisted KnowledgeFS Console routes.
|
||||
|
||||
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. The dedicated
|
||||
request path uses Dify's shared SSRF policy, accepts only exact contract
|
||||
operations, never follows redirects, bounds buffered identity 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 Literal, NamedTuple, Protocol, cast
|
||||
|
||||
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
|
||||
from services.knowledge_fs_contract_routes import KNOWLEDGE_FS_CONTRACT_OPERATIONS
|
||||
|
||||
type KnowledgeFSMethod = Literal["DELETE", "GET", "PATCH", "POST", "PUT"]
|
||||
type KnowledgeFSResponseKind = Literal["binary", "buffered", "stream"]
|
||||
type KnowledgeFSAccess = Literal["read", "write"]
|
||||
type KnowledgeFSRBACPermission = Literal[
|
||||
"dataset_access_config",
|
||||
"dataset_api_key_manage",
|
||||
"dataset_create_and_management",
|
||||
"dataset_document_download",
|
||||
"dataset_edit",
|
||||
"dataset_external_connect",
|
||||
"dataset_readonly",
|
||||
]
|
||||
|
||||
_JWT_AUDIENCE = "knowledge-fs"
|
||||
_JWT_ISSUER = "dify"
|
||||
_JWT_TTL_SECONDS = 60
|
||||
_MAX_BUFFERED_RESPONSE_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
class KnowledgeFSOperation(NamedTuple):
|
||||
method: KnowledgeFSMethod
|
||||
path: str
|
||||
response_kind: KnowledgeFSResponseKind
|
||||
access: KnowledgeFSAccess
|
||||
rbac_permission: KnowledgeFSRBACPermission
|
||||
max_response_bytes: int
|
||||
request_headers: tuple[str, ...]
|
||||
response_headers: tuple[str, ...]
|
||||
response_media_types: tuple[str, ...]
|
||||
|
||||
|
||||
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,
|
||||
method: KnowledgeFSMethod,
|
||||
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.
|
||||
method: Allowlisted upstream HTTP method.
|
||||
operation: Generated KnowledgeFS operation metadata.
|
||||
|
||||
Raises:
|
||||
KnowledgeFSAccessDeniedError: The account lacks a required legacy or enterprise permission.
|
||||
"""
|
||||
permission = RBACPermission(operation.rbac_permission)
|
||||
if permission == RBACPermission.DATASET_ACCESS_CONFIG and not account.is_admin_or_owner:
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS access-policy changes require a workspace administrator")
|
||||
if permission == RBACPermission.DATASET_API_KEY_MANAGE and (
|
||||
(method == "DELETE" and not account.is_admin_or_owner)
|
||||
or (method != "DELETE" and not account.has_edit_permission)
|
||||
):
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS API-key changes require elevated workspace access")
|
||||
if permission == RBACPermission.DATASET_EXTERNAL_CONNECT and not account.has_edit_permission:
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS connections require workspace edit access")
|
||||
if operation.access == "write" and not account.is_dataset_editor:
|
||||
raise KnowledgeFSAccessDeniedError("KnowledgeFS mutations require dataset edit access")
|
||||
if not RBACService.CheckAccess.check(
|
||||
tenant_id,
|
||||
account.id,
|
||||
scene=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,
|
||||
method=method,
|
||||
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": [f"knowledge-spaces:{operation.access}"],
|
||||
"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 (
|
||||
allowed_method,
|
||||
template,
|
||||
response_kind,
|
||||
access,
|
||||
rbac_permission,
|
||||
max_response_bytes,
|
||||
request_headers,
|
||||
response_headers,
|
||||
response_media_types,
|
||||
) in KNOWLEDGE_FS_CONTRACT_OPERATIONS:
|
||||
if method == allowed_method and _matches_route_template(template, path):
|
||||
return KnowledgeFSOperation(
|
||||
method,
|
||||
path,
|
||||
response_kind=cast(KnowledgeFSResponseKind, response_kind),
|
||||
access=cast(KnowledgeFSAccess, access),
|
||||
rbac_permission=cast(KnowledgeFSRBACPermission, rbac_permission),
|
||||
max_response_bytes=max_response_bytes,
|
||||
request_headers=request_headers,
|
||||
response_headers=response_headers,
|
||||
response_media_types=response_media_types,
|
||||
)
|
||||
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
|
||||
@@ -586,6 +586,7 @@ class RagPipelineService:
|
||||
|
||||
repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=db.engine,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=account,
|
||||
app_id=pipeline.id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
@@ -1377,6 +1378,7 @@ class RagPipelineService:
|
||||
# Create repository and save the node execution
|
||||
repository = SQLAlchemyWorkflowNodeExecutionRepository(
|
||||
session_factory=db.engine,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=current_user,
|
||||
app_id=pipeline.id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
|
||||
@@ -16,13 +16,11 @@ from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
from core.app.app_config.entities import ModelConfig
|
||||
from core.llm_generator.llm_generator import LLMGenerator
|
||||
from core.model_manager import ModelInstance, ModelManager
|
||||
from core.workflow.generator import WorkflowGenerator
|
||||
from core.workflow.generator.tool_catalogue import build_tool_catalogue, format_tool_catalogue, installed_tool_keys
|
||||
from core.workflow.generator.types import (
|
||||
WorkflowGenerateResultDict,
|
||||
WorkflowGenerationMode,
|
||||
WorkflowGenerationModeRequest,
|
||||
)
|
||||
from graphon.model_runtime.entities.model_entities import ModelType
|
||||
@@ -52,11 +50,9 @@ class WorkflowGeneratorService:
|
||||
"""
|
||||
Resolve a model instance for the tenant and run the generator.
|
||||
|
||||
``mode`` accepts the ``"auto"`` sentinel — when set, the instruction is
|
||||
classified into a concrete ``workflow`` / ``advanced-chat`` mode (one
|
||||
tiny LLM call) before planning so the rest of the pipeline runs against
|
||||
a concrete mode. The resolved mode is echoed back under the result's
|
||||
``mode`` key.
|
||||
``mode`` accepts the ``"auto"`` sentinel — the planner itself picks the
|
||||
concrete ``workflow`` / ``advanced-chat`` mode (no extra LLM call) and
|
||||
the resolution is echoed back under the result's ``mode`` key.
|
||||
|
||||
``current_graph`` is the existing draft graph for the cmd+k `/refine`
|
||||
flow — when present the generator refines it instead of creating a new
|
||||
@@ -66,9 +62,6 @@ class WorkflowGeneratorService:
|
||||
controller can map them to existing HTTP error envelopes (same
|
||||
envelope as ``/rule-generate``).
|
||||
"""
|
||||
resolved_mode = cls._resolve_mode(
|
||||
tenant_id=tenant_id, mode=mode, instruction=instruction, model_config=model_config
|
||||
)
|
||||
model_instance, model_parameters, tool_catalogue_text, installed_tools = cls._resolve_generation_context(
|
||||
tenant_id=tenant_id, model_config=model_config
|
||||
)
|
||||
@@ -79,7 +72,7 @@ class WorkflowGeneratorService:
|
||||
provider=model_config.provider,
|
||||
model_name=model_config.name,
|
||||
model_mode=model_config.mode.value,
|
||||
mode=resolved_mode,
|
||||
mode=mode,
|
||||
instruction=instruction,
|
||||
ideal_output=ideal_output,
|
||||
tool_catalogue_text=tool_catalogue_text,
|
||||
@@ -101,16 +94,13 @@ class WorkflowGeneratorService:
|
||||
"""
|
||||
Streaming sibling of ``generate_workflow_graph``.
|
||||
|
||||
Resolves the same model instance / tool catalogue / concrete mode, then
|
||||
delegates to ``WorkflowGenerator.generate_workflow_graph_stream`` and
|
||||
yields its ``(event_name, payload)`` tuples through to the controller's
|
||||
SSE writer. Provider-init / invoke errors raised while resolving the
|
||||
model instance propagate to the caller (the controller emits them as a
|
||||
Resolves the same model instance / tool catalogue, then delegates to
|
||||
``WorkflowGenerator.generate_workflow_graph_stream`` and yields its
|
||||
``(event_name, payload)`` tuples through to the controller's SSE
|
||||
writer. Provider-init / invoke errors raised while resolving the model
|
||||
instance propagate to the caller (the controller emits them as a
|
||||
single ``result`` SSE event).
|
||||
"""
|
||||
resolved_mode = cls._resolve_mode(
|
||||
tenant_id=tenant_id, mode=mode, instruction=instruction, model_config=model_config
|
||||
)
|
||||
model_instance, model_parameters, tool_catalogue_text, installed_tools = cls._resolve_generation_context(
|
||||
tenant_id=tenant_id, model_config=model_config
|
||||
)
|
||||
@@ -121,7 +111,7 @@ class WorkflowGeneratorService:
|
||||
provider=model_config.provider,
|
||||
model_name=model_config.name,
|
||||
model_mode=model_config.mode.value,
|
||||
mode=resolved_mode,
|
||||
mode=mode,
|
||||
instruction=instruction,
|
||||
ideal_output=ideal_output,
|
||||
tool_catalogue_text=tool_catalogue_text,
|
||||
@@ -129,28 +119,6 @@ class WorkflowGeneratorService:
|
||||
current_graph=current_graph,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _resolve_mode(
|
||||
cls,
|
||||
*,
|
||||
tenant_id: str,
|
||||
mode: WorkflowGenerationModeRequest,
|
||||
instruction: str,
|
||||
model_config: ModelConfig,
|
||||
) -> WorkflowGenerationMode:
|
||||
"""Resolve the request mode into a concrete generation mode.
|
||||
|
||||
``"auto"`` triggers a one-word LLM classification using the model the
|
||||
user already picked; everything else passes through unchanged. The
|
||||
classifier never raises (defaults to ``advanced-chat``), so ``auto``
|
||||
never blocks generation.
|
||||
"""
|
||||
if mode == "auto":
|
||||
return LLMGenerator.classify_workflow_mode(
|
||||
tenant_id=tenant_id, instruction=instruction, model_config=model_config
|
||||
)
|
||||
return mode
|
||||
|
||||
@classmethod
|
||||
def _resolve_generation_context(
|
||||
cls,
|
||||
|
||||
@@ -1039,6 +1039,7 @@ class WorkflowService:
|
||||
# Create repository and save the node execution
|
||||
repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=db.engine,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=account,
|
||||
app_id=app_model.id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
|
||||
@@ -9,6 +9,15 @@ from services.account_service import TenantService
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
|
||||
def _set_credit_pool_info(
|
||||
tenant_info: dict[str, object], *, quota_limit: int, quota_used: int, exhausted_at: int | None = None
|
||||
) -> None:
|
||||
tenant_info["trial_credits"] = quota_limit
|
||||
tenant_info["trial_credits_used"] = quota_used
|
||||
if isinstance(exhausted_at, int) and exhausted_at > 0 and quota_limit > 0 and quota_used >= quota_limit:
|
||||
tenant_info["trial_credits_exhausted_at"] = exhausted_at
|
||||
|
||||
|
||||
class WorkspaceService:
|
||||
@classmethod
|
||||
def get_tenant_info(cls, tenant: Tenant, session: Session):
|
||||
@@ -54,7 +63,7 @@ class WorkspaceService:
|
||||
if dify_config.EDITION == "CLOUD":
|
||||
tenant_info["next_credit_reset_date"] = feature.next_credit_reset_date
|
||||
|
||||
from services.credit_pool_service import CreditPoolService
|
||||
from services.credit_pool_service import CreditPoolBalance, CreditPoolService
|
||||
|
||||
paid_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="paid", session=session)
|
||||
# if the tenant is not on the sandbox plan and the paid pool is not full, use the paid pool
|
||||
@@ -63,12 +72,22 @@ class WorkspaceService:
|
||||
and paid_pool is not None
|
||||
and (paid_pool.quota_limit == -1 or paid_pool.quota_limit > paid_pool.quota_used)
|
||||
):
|
||||
tenant_info["trial_credits"] = paid_pool.quota_limit
|
||||
tenant_info["trial_credits_used"] = paid_pool.quota_used
|
||||
exhausted_at = paid_pool.exhausted_at if isinstance(paid_pool, CreditPoolBalance) else None
|
||||
_set_credit_pool_info(
|
||||
tenant_info,
|
||||
quota_limit=paid_pool.quota_limit,
|
||||
quota_used=paid_pool.quota_used,
|
||||
exhausted_at=exhausted_at,
|
||||
)
|
||||
else:
|
||||
trial_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="trial", session=session)
|
||||
if trial_pool:
|
||||
tenant_info["trial_credits"] = trial_pool.quota_limit
|
||||
tenant_info["trial_credits_used"] = trial_pool.quota_used
|
||||
exhausted_at = trial_pool.exhausted_at if isinstance(trial_pool, CreditPoolBalance) else None
|
||||
_set_credit_pool_info(
|
||||
tenant_info,
|
||||
quota_limit=trial_pool.quota_limit,
|
||||
quota_used=trial_pool.quota_used,
|
||||
exhausted_at=exhausted_at,
|
||||
)
|
||||
|
||||
return tenant_info
|
||||
|
||||
@@ -24,6 +24,9 @@ def _create_agent_backend_client():
|
||||
base_url=dify_config.AGENT_BACKEND_BASE_URL,
|
||||
use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
|
||||
fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
|
||||
stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS,
|
||||
stream_max_reconnects=dify_config.AGENT_BACKEND_STREAM_MAX_RECONNECTS,
|
||||
stream_run_timeout_seconds=dify_config.AGENT_BACKEND_RUN_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -248,7 +248,6 @@ class _AppRunner:
|
||||
case _Account():
|
||||
with self._session() as session:
|
||||
user: Account = session.get(Account, user_params.user_id)
|
||||
user.set_tenant_id_with_session(self._exec_params.tenant_id, session=session)
|
||||
return user
|
||||
case _:
|
||||
raise AssertionError(f"user should only be _Account or _EndUser, got {type(user_params)}")
|
||||
@@ -257,10 +256,7 @@ class _AppRunner:
|
||||
def _resolve_user_for_run(session: Session, workflow_run: WorkflowRun) -> Account | EndUser | None:
|
||||
role = CreatorUserRole(workflow_run.created_by_role)
|
||||
if role == CreatorUserRole.ACCOUNT:
|
||||
user = session.get(Account, workflow_run.created_by)
|
||||
if user:
|
||||
user.set_tenant_id_with_session(workflow_run.tenant_id, session=session)
|
||||
return user
|
||||
return session.get(Account, workflow_run.created_by)
|
||||
|
||||
return session.get(EndUser, workflow_run.created_by)
|
||||
|
||||
@@ -617,12 +613,14 @@ def _resume_advanced_chat(
|
||||
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=app_model.id,
|
||||
triggered_from=triggered_from,
|
||||
)
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=app_model.id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
@@ -684,12 +682,14 @@ def _resume_workflow(
|
||||
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=app_model.id,
|
||||
triggered_from=triggered_from,
|
||||
)
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=app_model.id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
@@ -247,12 +247,14 @@ def resume_workflow_execution(task_data_dict: dict[str, Any]) -> None:
|
||||
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom(workflow_run.triggered_from),
|
||||
)
|
||||
workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=app_model.tenant_id,
|
||||
user=user,
|
||||
app_id=generate_entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Install configured marketplace plugins for a newly created tenant."""
|
||||
|
||||
import logging
|
||||
|
||||
from celery import shared_task
|
||||
|
||||
from configs import dify_config
|
||||
from core.helper import marketplace
|
||||
from core.plugin.entities.plugin_daemon import PluginInstallTaskStatus
|
||||
from core.plugin.plugin_service import PluginService
|
||||
from services.model_provider_service import ModelProviderService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@shared_task(queue="plugin", bind=True, max_retries=60, default_retry_delay=5)
|
||||
def configure_default_models_task(self, tenant_id: str, plugin_install_task_id: str | None) -> None:
|
||||
"""Set explicitly configured default models after default plugins finish installing."""
|
||||
if not dify_config.NEW_USER_DEFAULT_MODELS:
|
||||
return
|
||||
|
||||
plugin_install_failed = False
|
||||
if plugin_install_task_id:
|
||||
try:
|
||||
install_task = PluginService.fetch_install_task(tenant_id, plugin_install_task_id)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to fetch default plugin installation task for tenant %s; retrying",
|
||||
tenant_id,
|
||||
)
|
||||
raise self.retry(exc=exc)
|
||||
|
||||
if install_task.status in (PluginInstallTaskStatus.Pending, PluginInstallTaskStatus.Running):
|
||||
raise self.retry()
|
||||
|
||||
plugin_install_failed = install_task.status == PluginInstallTaskStatus.Failed
|
||||
if plugin_install_failed:
|
||||
failed_plugin_ids = [
|
||||
plugin.plugin_id for plugin in install_task.plugins if plugin.status == PluginInstallTaskStatus.Failed
|
||||
]
|
||||
logger.error(
|
||||
"Default plugin installation failed for tenant %s: %s",
|
||||
tenant_id,
|
||||
", ".join(failed_plugin_ids),
|
||||
)
|
||||
|
||||
model_provider_service = ModelProviderService()
|
||||
failed_model_types: list[str] = []
|
||||
for model_type, provider, model in dify_config.NEW_USER_DEFAULT_MODEL_LIST:
|
||||
try:
|
||||
model_provider_service.update_default_model_of_model_type(
|
||||
tenant_id=tenant_id,
|
||||
model_type=model_type,
|
||||
provider=provider,
|
||||
model=model,
|
||||
)
|
||||
except Exception:
|
||||
failed_model_types.append(model_type)
|
||||
logger.exception(
|
||||
"Failed to configure default model for tenant %s: model_type=%s provider=%s model=%s",
|
||||
tenant_id,
|
||||
model_type,
|
||||
provider,
|
||||
model,
|
||||
)
|
||||
|
||||
if plugin_install_failed or failed_model_types:
|
||||
raise RuntimeError(
|
||||
f"Failed to initialize defaults for tenant {tenant_id}; "
|
||||
f"model types: {', '.join(failed_model_types) or 'none'}"
|
||||
)
|
||||
|
||||
|
||||
@shared_task(queue="plugin")
|
||||
def install_default_plugins_task(tenant_id: str, plugin_ids: list[str]) -> None:
|
||||
"""Install the latest marketplace versions of the configured plugins."""
|
||||
if not plugin_ids:
|
||||
return
|
||||
|
||||
try:
|
||||
manifests = {manifest.plugin_id: manifest for manifest in marketplace.batch_fetch_plugin_manifests(plugin_ids)}
|
||||
plugin_identifiers = [
|
||||
manifests[plugin_id].latest_package_identifier for plugin_id in plugin_ids if plugin_id in manifests
|
||||
]
|
||||
missing_plugin_ids = [plugin_id for plugin_id in plugin_ids if plugin_id not in manifests]
|
||||
if missing_plugin_ids:
|
||||
logger.warning("Default plugins not found in marketplace: %s", ", ".join(missing_plugin_ids))
|
||||
if not plugin_identifiers:
|
||||
return
|
||||
|
||||
response = PluginService.install_from_marketplace_pkg(tenant_id, plugin_identifiers)
|
||||
if dify_config.NEW_USER_DEFAULT_MODELS:
|
||||
configure_default_models_task.delay(
|
||||
tenant_id,
|
||||
None if response.all_installed else response.task_id,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to install default plugins for tenant %s", tenant_id)
|
||||
raise
|
||||
@@ -146,6 +146,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any],
|
||||
session_factory = sessionmaker(bind=db.engine, expire_on_commit=False)
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=account,
|
||||
app_id=entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN,
|
||||
@@ -154,6 +155,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any],
|
||||
workflow_node_execution_repository = (
|
||||
DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=account,
|
||||
app_id=entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN,
|
||||
|
||||
@@ -160,6 +160,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any],
|
||||
session_factory = sessionmaker(bind=db.engine, expire_on_commit=False)
|
||||
workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=account,
|
||||
app_id=entity.app_config.app_id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN,
|
||||
@@ -168,6 +169,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any],
|
||||
workflow_node_execution_repository = (
|
||||
DifyCoreRepositoryFactory.create_workflow_node_execution_repository(
|
||||
session_factory=session_factory,
|
||||
tenant_id=pipeline.tenant_id,
|
||||
user=account,
|
||||
app_id=entity.app_config.app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN,
|
||||
|
||||
@@ -656,10 +656,8 @@ class TestAccountGeneration:
|
||||
@patch("controllers.console.auth.oauth.TenantService")
|
||||
@patch("controllers.console.auth.oauth.FeatureService")
|
||||
@patch("controllers.console.auth.oauth.AccountService")
|
||||
@patch("controllers.console.auth.oauth.tenant_was_created")
|
||||
def test_should_create_workspace_for_account_without_tenant(
|
||||
self,
|
||||
mock_event: MagicMock,
|
||||
mock_account_service: MagicMock,
|
||||
mock_feature_service: MagicMock,
|
||||
mock_tenant_service: MagicMock,
|
||||
@@ -672,16 +670,9 @@ class TestAccountGeneration:
|
||||
mock_tenant_service.get_join_tenants.return_value = []
|
||||
mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
|
||||
|
||||
mock_new_tenant = MagicMock()
|
||||
mock_tenant_service.create_tenant.return_value = mock_new_tenant
|
||||
|
||||
with app.test_request_context(headers={"Accept-Language": "en-US,en;q=0.9"}):
|
||||
result, oauth_new_user = _generate_account("github", user_info)
|
||||
|
||||
assert result == mock_account
|
||||
assert oauth_new_user is False
|
||||
mock_tenant_service.create_tenant.assert_called_once_with("Test User's Workspace", session=ANY)
|
||||
mock_tenant_service.create_tenant_member.assert_called_once_with(
|
||||
mock_new_tenant, mock_account, ANY, role="owner"
|
||||
)
|
||||
mock_event.send.assert_called_once_with(mock_new_tenant)
|
||||
mock_tenant_service.create_owner_tenant.assert_called_once_with(mock_account, session=ANY)
|
||||
|
||||
+2
@@ -267,12 +267,14 @@ class TestHumanInputResumeNodeExecutionIntegration:
|
||||
)
|
||||
execution_repo = SQLAlchemyWorkflowExecutionRepository(
|
||||
session_factory=self.session.get_bind(),
|
||||
tenant_id=self.tenant.id,
|
||||
user=self.account,
|
||||
app_id=self.app.id,
|
||||
triggered_from=WorkflowRunTriggeredFrom.DEBUGGING,
|
||||
)
|
||||
node_execution_repo = SQLAlchemyWorkflowNodeExecutionRepository(
|
||||
session_factory=self.session.get_bind(),
|
||||
tenant_id=self.tenant.id,
|
||||
user=self.account,
|
||||
app_id=self.app.id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
+3
@@ -40,8 +40,11 @@ def _create_account_with_tenant(session: Session) -> Account:
|
||||
def _make_repo(session: Session, account: Account, app_id: str) -> SQLAlchemyWorkflowNodeExecutionRepository:
|
||||
engine = session.get_bind()
|
||||
assert isinstance(engine, Engine)
|
||||
tenant_id = account.current_tenant_id
|
||||
assert tenant_id is not None
|
||||
return SQLAlchemyWorkflowNodeExecutionRepository(
|
||||
session_factory=sessionmaker(bind=engine, expire_on_commit=False),
|
||||
tenant_id=tenant_id,
|
||||
user=account,
|
||||
app_id=app_id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
|
||||
@@ -11,7 +11,7 @@ import pytest
|
||||
import yaml
|
||||
from faker import Faker
|
||||
from flask import Flask
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.trigger.constants import (
|
||||
@@ -22,11 +22,24 @@ from core.trigger.constants import (
|
||||
from extensions.ext_redis import redis_client
|
||||
from graphon.enums import BuiltinNodeTypes
|
||||
from models import Account, App, AppMode
|
||||
from models.agent import Agent, AgentConfigDraft, AgentConfigDraftType, AgentConfigSnapshot, AgentScope, AgentSource
|
||||
from models.agent import (
|
||||
Agent,
|
||||
AgentConfigDraft,
|
||||
AgentConfigDraftType,
|
||||
AgentConfigSnapshot,
|
||||
AgentScope,
|
||||
AgentSource,
|
||||
AgentStatus,
|
||||
WorkflowAgentBindingType,
|
||||
WorkflowAgentNodeBinding,
|
||||
)
|
||||
from models.agent_config_entities import AgentSoulConfig
|
||||
from models.model import AppModelConfig, IconType
|
||||
from models.workflow import Workflow, WorkflowType
|
||||
from services import app_dsl_service
|
||||
from services.account_service import AccountService, TenantService
|
||||
from services.agent.dsl_entities import AGENT_PACKAGE_REF_KEY, make_portable_agent_package
|
||||
from services.agent.dsl_service import AgentDslService
|
||||
from services.app_dsl_service import (
|
||||
CHECK_DEPENDENCIES_REDIS_KEY_PREFIX,
|
||||
CURRENT_DSL_VERSION,
|
||||
@@ -952,6 +965,126 @@ class TestAppDslService:
|
||||
assert "model_config" in exported_data
|
||||
assert "dependencies" in exported_data
|
||||
|
||||
def test_workflow_package_import_materializes_all_agent_bindings_as_inline(
|
||||
self, db_session_with_containers: Session, mock_external_service_dependencies
|
||||
):
|
||||
app, account = self._create_test_app_and_account(db_session_with_containers, mock_external_service_dependencies)
|
||||
app.mode = AppMode.WORKFLOW
|
||||
workflow = Workflow.new(
|
||||
tenant_id=app.tenant_id,
|
||||
app_id=app.id,
|
||||
type=WorkflowType.WORKFLOW.value,
|
||||
version=Workflow.VERSION_DRAFT,
|
||||
graph=json.dumps({"nodes": [], "edges": []}),
|
||||
features=json.dumps({}),
|
||||
created_by=account.id,
|
||||
environment_variables=[],
|
||||
conversation_variables=[],
|
||||
rag_pipeline_variables=[],
|
||||
)
|
||||
db_session_with_containers.add(workflow)
|
||||
db_session_with_containers.flush()
|
||||
|
||||
source_agent = Agent(
|
||||
tenant_id=app.tenant_id,
|
||||
name="Portable Agent",
|
||||
description="Imported into each node",
|
||||
role="researcher",
|
||||
scope=AgentScope.ROSTER,
|
||||
source=AgentSource.AGENT_APP,
|
||||
status=AgentStatus.ACTIVE,
|
||||
created_by=account.id,
|
||||
updated_by=account.id,
|
||||
)
|
||||
package = make_portable_agent_package(source_agent, AgentSoulConfig(config_note="portable"))
|
||||
graph = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": "roster-node",
|
||||
"data": {
|
||||
"type": BuiltinNodeTypes.AGENT,
|
||||
"version": "2",
|
||||
"agent_binding": {
|
||||
"binding_type": WorkflowAgentBindingType.ROSTER_AGENT.value,
|
||||
AGENT_PACKAGE_REF_KEY: "agent_1",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"id": "inline-node",
|
||||
"data": {
|
||||
"type": BuiltinNodeTypes.AGENT,
|
||||
"version": "2",
|
||||
"agent_binding": {
|
||||
"binding_type": WorkflowAgentBindingType.INLINE_AGENT.value,
|
||||
AGENT_PACKAGE_REF_KEY: "agent_1",
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
"edges": [],
|
||||
}
|
||||
imported_roster_count_before = db_session_with_containers.scalar(
|
||||
select(func.count())
|
||||
.select_from(Agent)
|
||||
.where(
|
||||
Agent.tenant_id == app.tenant_id,
|
||||
Agent.scope == AgentScope.ROSTER,
|
||||
Agent.source == AgentSource.IMPORTED,
|
||||
)
|
||||
)
|
||||
|
||||
imported_graph, warnings = AgentDslService(db_session_with_containers).import_workflow_packages(
|
||||
workflow=workflow,
|
||||
portable_graph=graph,
|
||||
raw_packages={"agent_1": package.model_dump(mode="json")},
|
||||
account=account,
|
||||
)
|
||||
db_session_with_containers.commit()
|
||||
|
||||
assert warnings == []
|
||||
graph_bindings = [node["data"]["agent_binding"] for node in imported_graph["nodes"]]
|
||||
assert all(binding["binding_type"] == WorkflowAgentBindingType.INLINE_AGENT.value for binding in graph_bindings)
|
||||
assert len({binding["agent_id"] for binding in graph_bindings}) == 2
|
||||
|
||||
bindings = db_session_with_containers.scalars(
|
||||
select(WorkflowAgentNodeBinding).where(
|
||||
WorkflowAgentNodeBinding.tenant_id == app.tenant_id,
|
||||
WorkflowAgentNodeBinding.workflow_id == workflow.id,
|
||||
WorkflowAgentNodeBinding.workflow_version == Workflow.VERSION_DRAFT,
|
||||
)
|
||||
).all()
|
||||
assert len(bindings) == 2
|
||||
assert all(binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT for binding in bindings)
|
||||
|
||||
imported_agents = db_session_with_containers.scalars(
|
||||
select(Agent).where(Agent.id.in_({binding.agent_id for binding in bindings if binding.agent_id}))
|
||||
).all()
|
||||
assert len(imported_agents) == 2
|
||||
assert all(agent.scope == AgentScope.WORKFLOW_ONLY for agent in imported_agents)
|
||||
assert all(agent.source == AgentSource.IMPORTED for agent in imported_agents)
|
||||
assert all(agent.app_id == app.id and agent.workflow_id == workflow.id for agent in imported_agents)
|
||||
assert {agent.workflow_node_id for agent in imported_agents} == {"roster-node", "inline-node"}
|
||||
assert all(agent.backing_app_id for agent in imported_agents)
|
||||
|
||||
backing_apps = db_session_with_containers.scalars(
|
||||
select(App).where(App.id.in_({agent.backing_app_id for agent in imported_agents if agent.backing_app_id}))
|
||||
).all()
|
||||
assert len(backing_apps) == 2
|
||||
assert all(backing_app.mode == AppMode.AGENT for backing_app in backing_apps)
|
||||
assert all(backing_app.enable_site is False and backing_app.enable_api is False for backing_app in backing_apps)
|
||||
|
||||
imported_roster_count_after = db_session_with_containers.scalar(
|
||||
select(func.count())
|
||||
.select_from(Agent)
|
||||
.where(
|
||||
Agent.tenant_id == app.tenant_id,
|
||||
Agent.scope == AgentScope.ROSTER,
|
||||
Agent.source == AgentSource.IMPORTED,
|
||||
)
|
||||
)
|
||||
assert imported_roster_count_after == imported_roster_count_before
|
||||
|
||||
def test_agent_app_dsl_round_trip_creates_unpublished_imported_agent(
|
||||
self, db_session_with_containers: Session, mock_external_service_dependencies
|
||||
):
|
||||
|
||||
@@ -7,6 +7,7 @@ from faker import Faker
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
|
||||
from services.credit_pool_service import CreditPoolBalance
|
||||
from services.workspace_service import WorkspaceService
|
||||
|
||||
|
||||
@@ -720,7 +721,13 @@ class TestWorkspaceService:
|
||||
mock_external_service_dependencies["tenant_service"].has_roles.return_value = False
|
||||
|
||||
paid_pool = MagicMock(quota_limit=500, quota_used=500)
|
||||
trial_pool = MagicMock(quota_limit=100, quota_used=10)
|
||||
trial_pool = CreditPoolBalance(
|
||||
tenant_id=tenant.id,
|
||||
pool_type="trial",
|
||||
quota_limit=100,
|
||||
quota_used=100,
|
||||
exhausted_at=1748908800,
|
||||
)
|
||||
|
||||
with (
|
||||
patch("services.workspace_service.current_user", account),
|
||||
@@ -730,7 +737,8 @@ class TestWorkspaceService:
|
||||
|
||||
assert result is not None
|
||||
assert result["trial_credits"] == 100
|
||||
assert result["trial_credits_used"] == 10
|
||||
assert result["trial_credits_used"] == 100
|
||||
assert result["trial_credits_exhausted_at"] == 1748908800
|
||||
|
||||
def test_get_tenant_info_cloud_fall_back_to_trial_when_paid_none(
|
||||
self, db_session_with_containers: Session, mock_external_service_dependencies
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import override
|
||||
|
||||
import pytest
|
||||
@@ -47,6 +47,8 @@ def _request():
|
||||
|
||||
|
||||
class _SuccessfulClient:
|
||||
stream_options: tuple[int | None, float | None, Callable[[], bool] | None] | None = None
|
||||
|
||||
def create_run_sync(self, request: CreateRunRequest) -> CreateRunResponse:
|
||||
assert isinstance(request, CreateRunRequest)
|
||||
return CreateRunResponse(run_id="run-1", status="running")
|
||||
@@ -55,8 +57,17 @@ class _SuccessfulClient:
|
||||
del request
|
||||
return CancelRunResponse(run_id=run_id, status="cancelled")
|
||||
|
||||
def stream_events_sync(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
def stream_events_sync(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
max_reconnects: int | None = None,
|
||||
timeout_seconds: float | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after
|
||||
self.stream_options = (max_reconnects, timeout_seconds, should_stop)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id)
|
||||
|
||||
def wait_run_sync(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse:
|
||||
@@ -72,17 +83,22 @@ class _SuccessfulClient:
|
||||
|
||||
|
||||
def test_dify_agent_backend_run_client_delegates_sync_methods():
|
||||
client = DifyAgentBackendRunClient(_SuccessfulClient())
|
||||
wrapped = _SuccessfulClient()
|
||||
client = DifyAgentBackendRunClient(wrapped, stream_max_reconnects=2, stream_timeout_seconds=45)
|
||||
|
||||
def should_stop() -> bool:
|
||||
return False
|
||||
|
||||
created = client.create_run(_request())
|
||||
cancelled = client.cancel_run(created.run_id)
|
||||
events = list(client.stream_events(created.run_id))
|
||||
events = list(client.stream_events(created.run_id, should_stop=should_stop))
|
||||
status = client.wait_run(created.run_id)
|
||||
|
||||
assert created.run_id == "run-1"
|
||||
assert cancelled.status == "cancelled"
|
||||
assert events[0].type == "run_started"
|
||||
assert status.status == "succeeded"
|
||||
assert wrapped.stream_options == (2, 45, should_stop)
|
||||
|
||||
|
||||
def test_dify_agent_backend_run_client_maps_validation_error():
|
||||
@@ -125,7 +141,16 @@ def test_dify_agent_backend_run_client_maps_timeout_error():
|
||||
def test_dify_agent_backend_run_client_maps_stream_error():
|
||||
class StreamClient(_SuccessfulClient):
|
||||
@override
|
||||
def stream_events_sync(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
def stream_events_sync(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
max_reconnects: int | None = None,
|
||||
timeout_seconds: float | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del run_id, after, max_reconnects, timeout_seconds, should_stop
|
||||
raise DifyAgentStreamError("bad stream")
|
||||
yield
|
||||
|
||||
|
||||
@@ -89,6 +89,54 @@ def test_dify_config(monkeypatch: pytest.MonkeyPatch):
|
||||
assert Version(config.project.version) >= Version("1.0.0")
|
||||
|
||||
|
||||
def test_new_user_default_plugin_ids_are_parsed_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_set_basic_config_env(monkeypatch)
|
||||
monkeypatch.setenv(
|
||||
"NEW_USER_DEFAULT_PLUGIN_IDS",
|
||||
"langgenius/openai, langgenius/gemini",
|
||||
)
|
||||
|
||||
config = DifyConfig(_env_file=None)
|
||||
|
||||
assert config.NEW_USER_DEFAULT_PLUGIN_ID_LIST == [
|
||||
"langgenius/openai",
|
||||
"langgenius/gemini",
|
||||
]
|
||||
|
||||
|
||||
def test_new_user_default_models_are_parsed_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_set_basic_config_env(monkeypatch)
|
||||
monkeypatch.setenv(
|
||||
"NEW_USER_DEFAULT_MODELS",
|
||||
(
|
||||
"llm:langgenius/openai/openai:gpt-4o-mini, "
|
||||
"text-embedding:langgenius/openai/openai:text-embedding-3-small, "
|
||||
"rerank:langgenius/ollama/ollama:reranker:latest"
|
||||
),
|
||||
)
|
||||
|
||||
config = DifyConfig(_env_file=None)
|
||||
|
||||
assert config.NEW_USER_DEFAULT_MODEL_LIST == [
|
||||
("llm", "langgenius/openai/openai", "gpt-4o-mini"),
|
||||
("text-embedding", "langgenius/openai/openai", "text-embedding-3-small"),
|
||||
("rerank", "langgenius/ollama/ollama", "reranker:latest"),
|
||||
]
|
||||
|
||||
|
||||
def test_new_user_default_models_reject_duplicate_model_types(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_set_basic_config_env(monkeypatch)
|
||||
monkeypatch.setenv(
|
||||
"NEW_USER_DEFAULT_MODELS",
|
||||
"llm:langgenius/openai/openai:gpt-4o-mini,llm:langgenius/anthropic/anthropic:claude-sonnet-4",
|
||||
)
|
||||
|
||||
config = DifyConfig(_env_file=None)
|
||||
|
||||
with pytest.raises(ValueError, match="duplicate model type: llm"):
|
||||
_ = config.NEW_USER_DEFAULT_MODEL_LIST
|
||||
|
||||
|
||||
def test_http_timeout_defaults(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that HTTP timeout defaults are correctly set"""
|
||||
# clear system environment variables
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
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_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_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
|
||||
|
||||
|
||||
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_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://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)
|
||||
@@ -33,6 +33,7 @@ from core.errors.error import (
|
||||
QuotaExceededError,
|
||||
)
|
||||
from graphon.model_runtime.errors.invoke import InvokeError
|
||||
from graphon.variables import StringVariable
|
||||
from models import Account
|
||||
from models.account import TenantStatus
|
||||
from models.model import AppMode
|
||||
@@ -1146,7 +1147,14 @@ class TestAppWorkflowApi:
|
||||
created_at=datetime(2024, 1, 1, tzinfo=UTC),
|
||||
updated_at=datetime(2024, 1, 2, tzinfo=UTC),
|
||||
environment_variables=[],
|
||||
conversation_variables=[],
|
||||
conversation_variables=[
|
||||
StringVariable(
|
||||
id="conversation-variable-1",
|
||||
name="topic",
|
||||
value="sqlite",
|
||||
selector=["conversation", "topic"],
|
||||
)
|
||||
],
|
||||
rag_pipeline_variables=[],
|
||||
get_created_by_account=MagicMock(return_value=created_by),
|
||||
get_updated_by_account=MagicMock(return_value=None),
|
||||
@@ -1174,7 +1182,15 @@ class TestAppWorkflowApi:
|
||||
"updated_at": 1704153600,
|
||||
"tool_published": True,
|
||||
"environment_variables": [],
|
||||
"conversation_variables": [],
|
||||
"conversation_variables": [
|
||||
{
|
||||
"id": "conversation-variable-1",
|
||||
"name": "topic",
|
||||
"value_type": "string",
|
||||
"value": "sqlite",
|
||||
"description": "",
|
||||
}
|
||||
],
|
||||
"rag_pipeline_variables": [],
|
||||
}
|
||||
app_model.workflow_with_session.assert_called_once_with(session=session)
|
||||
|
||||
@@ -1,748 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
from collections.abc import Iterator
|
||||
from inspect import unwrap
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from flask import Flask, Response
|
||||
from werkzeug.exceptions import (
|
||||
BadGateway,
|
||||
Forbidden,
|
||||
MethodNotAllowed,
|
||||
NotFound,
|
||||
RequestEntityTooLarge,
|
||||
ServiceUnavailable,
|
||||
)
|
||||
|
||||
from controllers.console import bp
|
||||
from controllers.console.knowledge_fs_proxy import (
|
||||
_console_api_errors,
|
||||
_proxy_knowledge_fs_non_get,
|
||||
_proxy_response,
|
||||
proxy_knowledge_fs_get,
|
||||
proxy_knowledge_fs_write,
|
||||
)
|
||||
from controllers.console.wraps import RBACPermission
|
||||
from core.helper import ssrf_proxy
|
||||
from services.knowledge_fs_proxy import (
|
||||
KnowledgeFSAccessDeniedError,
|
||||
KnowledgeFSConfigurationError,
|
||||
KnowledgeFSOperation,
|
||||
KnowledgeFSResponseKind,
|
||||
KnowledgeFSRouteNotAllowedError,
|
||||
KnowledgeFSUpstreamResponse,
|
||||
get_knowledge_fs_operation,
|
||||
)
|
||||
|
||||
|
||||
def _upstream(
|
||||
response: httpx.Response,
|
||||
kind: KnowledgeFSResponseKind = "buffered",
|
||||
*,
|
||||
max_response_bytes: int | None = None,
|
||||
) -> KnowledgeFSUpstreamResponse:
|
||||
operation = KnowledgeFSOperation(
|
||||
method="GET",
|
||||
path="test",
|
||||
response_kind=kind,
|
||||
access="read",
|
||||
rbac_permission="dataset_readonly",
|
||||
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)
|
||||
|
||||
|
||||
class _EventStream(httpx.SyncByteStream):
|
||||
def __init__(self, content: bytes) -> None:
|
||||
self._content = content
|
||||
|
||||
def __iter__(self) -> Iterator[bytes]:
|
||||
yield self._content
|
||||
|
||||
|
||||
class _FailingEventStream(httpx.SyncByteStream):
|
||||
def __iter__(self) -> Iterator[bytes]:
|
||||
yield b"event: delta\ndata: first\n\n"
|
||||
request = httpx.Request("POST", "http://knowledge-fs.test/queries")
|
||||
raise httpx.ReadTimeout("stream timed out", request=request)
|
||||
|
||||
|
||||
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"}
|
||||
|
||||
|
||||
def test_generic_get_rejects_flasks_implicit_head_route(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> 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="HEAD",
|
||||
):
|
||||
with pytest.raises(MethodNotAllowed):
|
||||
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",
|
||||
"queries",
|
||||
RBACPermission.DATASET_READONLY,
|
||||
),
|
||||
(
|
||||
proxy_knowledge_fs_write,
|
||||
"POST",
|
||||
"knowledge-spaces",
|
||||
RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
|
||||
),
|
||||
(proxy_knowledge_fs_write, "DELETE", "knowledge-spaces/space-1", RBACPermission.DATASET_EDIT),
|
||||
(
|
||||
proxy_knowledge_fs_write,
|
||||
"PATCH",
|
||||
"knowledge-spaces/space-1/access-policy",
|
||||
RBACPermission.DATASET_ACCESS_CONFIG,
|
||||
),
|
||||
(
|
||||
proxy_knowledge_fs_write,
|
||||
"POST",
|
||||
"knowledge-spaces/space-1/api-keys",
|
||||
RBACPermission.DATASET_API_KEY_MANAGE,
|
||||
),
|
||||
(
|
||||
proxy_knowledge_fs_get,
|
||||
"GET",
|
||||
"knowledge-spaces/space-1/documents/document-1/multimodal/item-1/asset",
|
||||
RBACPermission.DATASET_DOCUMENT_DOWNLOAD,
|
||||
),
|
||||
(
|
||||
proxy_knowledge_fs_write,
|
||||
"POST",
|
||||
"knowledge-spaces/space-1/source-connections",
|
||||
RBACPermission.DATASET_EXTERNAL_CONNECT,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_generic_routes_delegate_to_the_authorized_service_use_case(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route,
|
||||
method: str,
|
||||
path: str,
|
||||
permission: RBACPermission,
|
||||
) -> None:
|
||||
response_kind: KnowledgeFSResponseKind = (
|
||||
"binary" if permission == RBACPermission.DATASET_DOCUMENT_DOWNLOAD else "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.value
|
||||
|
||||
|
||||
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)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.wraps.current_account_with_tenant",
|
||||
lambda: (MagicMock(id="account-1"), "tenant-1"),
|
||||
)
|
||||
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/queries", method="POST"):
|
||||
response = _proxy_knowledge_fs_non_get("POST", "queries")
|
||||
|
||||
assert isinstance(response, Response)
|
||||
zadd.assert_called_once()
|
||||
proxy.assert_called_once_with("POST", "queries")
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "path"),
|
||||
[
|
||||
("DELETE", "knowledge-spaces/space-1"),
|
||||
("PATCH", "knowledge-spaces/space-1"),
|
||||
("POST", "knowledge-spaces"),
|
||||
("PUT", "knowledge-spaces/space-1/retrieval-profile"),
|
||||
],
|
||||
)
|
||||
def test_generic_write_forwards_path_raw_body_and_current_tenant(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
method: str,
|
||||
path: str,
|
||||
) -> None:
|
||||
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'{"challenge":"delete-space","expectedRevision":1}'
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/knowledge-spaces/space-1",
|
||||
method="DELETE",
|
||||
data=body,
|
||||
content_type="application/json",
|
||||
headers={"Idempotency-Key": "delete-space-1"},
|
||||
):
|
||||
response = route("knowledge-spaces/space-1")
|
||||
|
||||
assert isinstance(response, Response)
|
||||
assert response.status_code == 202
|
||||
assert forward.call_args.kwargs["request_headers"].get("Idempotency-Key") == "delete-space-1"
|
||||
|
||||
|
||||
def test_processing_events_forwards_last_event_id(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
stream=_EventStream(b"event: done\ndata: {}\n\n"),
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
forward = MagicMock(return_value=_upstream(upstream, "stream"))
|
||||
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request", forward)
|
||||
_set_current_workspace(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_get)
|
||||
path = "knowledge-spaces/space-1/documents/document-1/processing-tasks/task-1/events"
|
||||
|
||||
with app.test_request_context(
|
||||
f"/console/api/knowledge-fs/{path}",
|
||||
headers={"Last-Event-ID": "42"},
|
||||
):
|
||||
response = route(path)
|
||||
assert isinstance(response, Response)
|
||||
response.close()
|
||||
|
||||
assert forward.call_args.kwargs["request_headers"].get("Last-Event-ID") == "42"
|
||||
|
||||
|
||||
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_read_post_allows_non_editor_and_streams_sse(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
stream=_EventStream(b"event: delta\ndata: first\n\nevent: done\ndata: {}\n\n"),
|
||||
headers={
|
||||
"Cache-Control": "no-store",
|
||||
"Content-Type": "text/event-stream",
|
||||
"X-Query-Run-Id": "query-run-1",
|
||||
"X-Session-Id": "session-1",
|
||||
},
|
||||
)
|
||||
forward = MagicMock(return_value=_upstream(upstream, "stream"))
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
forward,
|
||||
)
|
||||
_set_current_workspace(monkeypatch, editor=False)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context(
|
||||
"/console/api/knowledge-fs/queries",
|
||||
method="POST",
|
||||
data=b'{"knowledgeSpaceId":"space-1","query":"hello"}',
|
||||
content_type="application/json",
|
||||
):
|
||||
response = route("queries")
|
||||
assert isinstance(response, Response)
|
||||
assert response.get_data() == b"event: delta\ndata: first\n\nevent: done\ndata: {}\n\n"
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["Content-Type"].startswith("text/event-stream")
|
||||
assert response.headers["Cache-Control"] == "no-store"
|
||||
assert response.headers["X-Query-Run-Id"] == "query-run-1"
|
||||
assert response.headers["X-Session-Id"] == "session-1"
|
||||
assert upstream.is_closed
|
||||
forward.assert_called_once()
|
||||
|
||||
|
||||
def test_unconsumed_sse_response_closes_upstream(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
stream=_EventStream(b"event: done\ndata: {}\n\n"),
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(return_value=_upstream(upstream, "stream")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch, editor=False)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/queries", method="POST", data=b"{}"):
|
||||
response = route("queries")
|
||||
assert isinstance(response, Response)
|
||||
response.close()
|
||||
|
||||
assert upstream.is_closed
|
||||
|
||||
|
||||
def test_sse_response_uses_the_generated_operation_limit(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
stream=_EventStream(b"event: done\ndata: {}\n\n"),
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(return_value=_upstream(upstream, "stream", max_response_bytes=1)),
|
||||
)
|
||||
_set_current_workspace(monkeypatch, editor=False)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/queries", method="POST", data=b"{}"):
|
||||
response = route("queries")
|
||||
assert isinstance(response, Response)
|
||||
with pytest.raises(ssrf_proxy.ResponseTooLargeError, match="response exceeded 1 bytes"):
|
||||
response.get_data()
|
||||
|
||||
assert upstream.is_closed
|
||||
|
||||
|
||||
def test_sse_transport_failure_is_visible_to_the_client(
|
||||
app: Flask,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
upstream = httpx.Response(
|
||||
200,
|
||||
stream=_FailingEventStream(),
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
|
||||
MagicMock(return_value=_upstream(upstream, "stream")),
|
||||
)
|
||||
_set_current_workspace(monkeypatch, editor=False)
|
||||
_bypass_policy_wrappers(monkeypatch)
|
||||
route = unwrap(proxy_knowledge_fs_write)
|
||||
|
||||
with app.test_request_context("/console/api/knowledge-fs/queries", method="POST", data=b"{}"):
|
||||
response = route("queries")
|
||||
assert isinstance(response, Response)
|
||||
with pytest.raises(httpx.ReadTimeout, match="stream timed out"):
|
||||
response.get_data()
|
||||
|
||||
assert upstream.is_closed
|
||||
|
||||
|
||||
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: str,
|
||||
) -> 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}
|
||||
@@ -6,18 +6,22 @@ import builtins
|
||||
import importlib
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from inspect import unwrap
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
from types import ModuleType
|
||||
from unittest.mock import MagicMock, patch
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from flask.views import MethodView
|
||||
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
||||
|
||||
from core.tools.entities.api_entities import ToolProviderApiEntity as CoreToolProviderApiEntity
|
||||
from core.tools.entities.common_entities import I18nObject
|
||||
from core.tools.entities.tool_entities import ToolParameter
|
||||
from models import Account
|
||||
from models import Account, BuiltinToolProvider, Tenant, TenantAccountJoin
|
||||
from models.account import TenantAccountRole
|
||||
from models.credential_permission import CredentialPermission
|
||||
from models.enums import PermissionEnum
|
||||
|
||||
if not hasattr(builtins, "MethodView"):
|
||||
builtins.MethodView = MethodView # type: ignore[attr-defined]
|
||||
@@ -27,13 +31,6 @@ _CONTROLLER_MODULE: ModuleType | None = None
|
||||
_WRAPS_MODULE: ModuleType | None = None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _mock_db():
|
||||
mock_session = SimpleNamespace(scalar=lambda *args, **kwargs: True)
|
||||
with patch("extensions.ext_database.db.session", mock_session):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app() -> Flask:
|
||||
flask_app = Flask(__name__)
|
||||
@@ -69,8 +66,7 @@ def controller_module(monkeypatch: pytest.MonkeyPatch):
|
||||
with ExitStack() as stack:
|
||||
for target, value in patch_targets:
|
||||
stack.enter_context(patch(target, value))
|
||||
with _mock_db():
|
||||
_CONTROLLER_MODULE = importlib.import_module(module_name)
|
||||
_CONTROLLER_MODULE = importlib.import_module(module_name)
|
||||
|
||||
module = _CONTROLLER_MODULE
|
||||
|
||||
@@ -88,11 +84,79 @@ def controller_module(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
def _mock_account(user_id: str = "user-123") -> Account:
|
||||
user = Account(name="Test User", email=f"{user_id}@example.com")
|
||||
user.id = user_id
|
||||
user.id = _stable_uuid(f"account:{user_id}")
|
||||
user.role = TenantAccountRole.NORMAL
|
||||
return user
|
||||
|
||||
|
||||
def _stable_uuid(value: str) -> str:
|
||||
return str(uuid5(NAMESPACE_URL, value))
|
||||
|
||||
|
||||
def _persist_workspace(session: Session, user: Account, tenant_name: str) -> Tenant:
|
||||
tenant = Tenant(name=tenant_name)
|
||||
tenant.id = _stable_uuid(f"tenant:{tenant_name}")
|
||||
membership = TenantAccountJoin(
|
||||
tenant_id=tenant.id,
|
||||
account_id=user.id,
|
||||
current=True,
|
||||
role=TenantAccountRole.NORMAL,
|
||||
)
|
||||
session.add_all([user, tenant, membership])
|
||||
session.commit()
|
||||
return tenant
|
||||
|
||||
|
||||
def _provider_credential(
|
||||
*,
|
||||
tenant_id: str,
|
||||
user_id: str,
|
||||
credential_name: str,
|
||||
visibility: PermissionEnum = PermissionEnum.ALL_TEAM,
|
||||
) -> BuiltinToolProvider:
|
||||
provider = BuiltinToolProvider(
|
||||
name=credential_name,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
provider="demo",
|
||||
encrypted_credentials='{"api_key": "sk-secret"}',
|
||||
visibility=visibility,
|
||||
)
|
||||
provider.id = _stable_uuid(f"credential:{tenant_id}:{credential_name}")
|
||||
return provider
|
||||
|
||||
|
||||
@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
|
||||
finally:
|
||||
database_session.remove()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _mock_credential_encryption(controller_module: ModuleType):
|
||||
encrypter = MagicMock()
|
||||
encrypter.decrypt.side_effect = lambda credentials: credentials
|
||||
encrypter.mask_plugin_credentials.return_value = {"api_key": "[__HIDDEN__]"}
|
||||
with (
|
||||
patch(
|
||||
"services.tools.builtin_tools_manage_service.ToolManager.get_builtin_provider",
|
||||
return_value=MagicMock(),
|
||||
),
|
||||
patch.object(
|
||||
controller_module.BuiltinToolManageService,
|
||||
"create_tool_encrypter",
|
||||
return_value=(encrypter, MagicMock()),
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _set_current_account(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
controller_module: ModuleType,
|
||||
@@ -190,23 +254,6 @@ def _provider_list_item(
|
||||
return service_payload, provider.model_dump(mode="json", exclude_unset=True)
|
||||
|
||||
|
||||
def _credential_response(controller_module: ModuleType, credential_id: str = "cred-1") -> tuple[dict, dict]:
|
||||
expected = {
|
||||
"id": credential_id,
|
||||
"name": "Credential",
|
||||
"provider": "demo",
|
||||
"credential_type": controller_module.CredentialType.API_KEY,
|
||||
"is_default": False,
|
||||
"credentials": {},
|
||||
"visibility": "all_team_members",
|
||||
"created_by": "",
|
||||
"partial_member_list": [],
|
||||
"from_other_member": False,
|
||||
}
|
||||
credential = controller_module.ToolProviderCredentialApiEntity.model_validate(expected)
|
||||
return credential.model_dump(mode="json"), credential.model_dump(mode="json")
|
||||
|
||||
|
||||
def _provider_config_response(controller_module: ModuleType) -> tuple[dict, dict]:
|
||||
expected = {
|
||||
"type": "secret-input",
|
||||
@@ -308,7 +355,7 @@ def test_builtin_provider_add_passes_payload(
|
||||
|
||||
assert response == {"result": "success"}
|
||||
service_mock.assert_called_once_with(
|
||||
user_id="user-123",
|
||||
user_id=user.id,
|
||||
tenant_id="tenant-456",
|
||||
provider="openai",
|
||||
credentials={"api_key": "sk-test"},
|
||||
@@ -397,54 +444,96 @@ def test_builtin_provider_info_uses_core_to_dict_tool_projection(
|
||||
assert "original_credentials" not in resp
|
||||
|
||||
|
||||
def test_builtin_provider_credentials_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch):
|
||||
user = _mock_account("user-tenant-cred")
|
||||
_set_current_account(monkeypatch, controller_module, user, "tenant-cred")
|
||||
service_payload, expected_response = _credential_response(controller_module)
|
||||
service_mock = MagicMock(return_value=[service_payload])
|
||||
monkeypatch.setattr(
|
||||
controller_module.BuiltinToolManageService,
|
||||
"get_builtin_tool_provider_credentials",
|
||||
service_mock,
|
||||
)
|
||||
|
||||
with app.test_request_context("/creds", method="GET"):
|
||||
resp = controller_module.ToolBuiltinProviderGetCredentialsApi().get(provider="demo")
|
||||
|
||||
assert resp == [expected_response]
|
||||
service_mock.assert_called_once_with(
|
||||
tenant_id="tenant-cred",
|
||||
provider_name="demo",
|
||||
session=ANY,
|
||||
user=user,
|
||||
include_credential_ids=None,
|
||||
)
|
||||
|
||||
|
||||
def test_builtin_provider_credentials_get_reads_repeated_include_ids(
|
||||
app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, BuiltinToolProvider, CredentialPermission)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_builtin_provider_credentials_get(
|
||||
app: Flask,
|
||||
controller_module: ModuleType,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sqlite_session: Session,
|
||||
):
|
||||
user = _mock_account("user-tenant-cred")
|
||||
credential_payload, expected = _credential_response(controller_module)
|
||||
service_mock = MagicMock(return_value=[credential_payload])
|
||||
monkeypatch.setattr(
|
||||
controller_module.BuiltinToolManageService,
|
||||
"get_builtin_tool_provider_credentials",
|
||||
service_mock,
|
||||
tenant = _persist_workspace(sqlite_session, user, "tenant-cred")
|
||||
credential = _provider_credential(
|
||||
tenant_id=tenant.id,
|
||||
user_id=user.id,
|
||||
credential_name="Credential",
|
||||
)
|
||||
other_tenant = Tenant(name="other-tenant")
|
||||
other_tenant.id = _stable_uuid("tenant:other-tenant")
|
||||
other_credential = _provider_credential(
|
||||
tenant_id=other_tenant.id,
|
||||
user_id=user.id,
|
||||
credential_name="Other Tenant Credential",
|
||||
)
|
||||
sqlite_session.add_all([credential, other_tenant, other_credential])
|
||||
sqlite_session.commit()
|
||||
_set_current_account(monkeypatch, controller_module, user, tenant.id)
|
||||
|
||||
with app.test_request_context("/creds?include_credential_ids=cred-1&include_credential_ids=cred-2", method="GET"):
|
||||
with (
|
||||
_bind_database_session(sqlite_session),
|
||||
_mock_credential_encryption(controller_module),
|
||||
app.test_request_context("/creds", method="GET"),
|
||||
):
|
||||
response = controller_module.ToolBuiltinProviderGetCredentialsApi().get(provider="demo")
|
||||
|
||||
assert [item["id"] for item in response] == [credential.id]
|
||||
assert response[0]["name"] == "Credential"
|
||||
assert response[0]["credentials"] == {"api_key": "[__HIDDEN__]"}
|
||||
assert response[0]["created_by"] == user.id
|
||||
assert other_credential.id not in {item["id"] for item in response}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sqlite_session",
|
||||
[(Account, Tenant, TenantAccountJoin, BuiltinToolProvider, CredentialPermission)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_builtin_provider_credentials_get_reads_repeated_include_ids(
|
||||
app: Flask,
|
||||
controller_module: ModuleType,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sqlite_session: Session,
|
||||
):
|
||||
user = _mock_account("user-tenant-cred")
|
||||
tenant = _persist_workspace(sqlite_session, user, "tenant-cred")
|
||||
visible_credential = _provider_credential(
|
||||
tenant_id=tenant.id,
|
||||
user_id=user.id,
|
||||
credential_name="Visible Credential",
|
||||
)
|
||||
other_user = _mock_account("other-user")
|
||||
borrowed_credential = _provider_credential(
|
||||
tenant_id=tenant.id,
|
||||
user_id=other_user.id,
|
||||
credential_name="Borrowed Credential",
|
||||
visibility=PermissionEnum.ONLY_ME,
|
||||
)
|
||||
other_membership = TenantAccountJoin(
|
||||
tenant_id=tenant.id,
|
||||
account_id=other_user.id,
|
||||
role=TenantAccountRole.NORMAL,
|
||||
)
|
||||
sqlite_session.add_all([visible_credential, other_user, other_membership, borrowed_credential])
|
||||
sqlite_session.commit()
|
||||
|
||||
request_path = (
|
||||
f"/creds?include_credential_ids={visible_credential.id}&include_credential_ids={borrowed_credential.id}"
|
||||
)
|
||||
with (
|
||||
_bind_database_session(sqlite_session),
|
||||
_mock_credential_encryption(controller_module),
|
||||
app.test_request_context(request_path, method="GET"),
|
||||
):
|
||||
api = controller_module.ToolBuiltinProviderGetCredentialsApi()
|
||||
resp = unwrap(api.get)(api, "tenant-cred", user, provider="demo")
|
||||
response = unwrap(api.get)(api, tenant.id, user, provider="demo")
|
||||
|
||||
assert resp == [expected]
|
||||
service_mock.assert_called_once_with(
|
||||
tenant_id="tenant-cred",
|
||||
provider_name="demo",
|
||||
session=ANY,
|
||||
user=user,
|
||||
include_credential_ids=["cred-1", "cred-2"],
|
||||
)
|
||||
assert [item["id"] for item in response] == [visible_credential.id, borrowed_credential.id]
|
||||
assert response[0]["from_other_member"] is False
|
||||
assert response[1]["from_other_member"] is True
|
||||
|
||||
|
||||
def test_api_provider_remote_schema_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
@@ -84,10 +84,9 @@ class TestEnterpriseWorkspace:
|
||||
assert hasattr(api_instance, "post")
|
||||
assert callable(api_instance.post)
|
||||
|
||||
@patch("controllers.inner_api.workspace.workspace.tenant_was_created")
|
||||
@patch("controllers.inner_api.workspace.workspace.TenantService")
|
||||
@patch("controllers.inner_api.workspace.workspace.db")
|
||||
def test_post_creates_workspace_with_owner(self, mock_db, mock_tenant_svc, mock_event, api_instance, app: Flask):
|
||||
def test_post_creates_workspace_with_owner(self, mock_db, mock_tenant_svc, api_instance, app: Flask):
|
||||
"""Test that post() creates a workspace and assigns the owner account"""
|
||||
# Arrange
|
||||
mock_account = MagicMock()
|
||||
@@ -102,7 +101,7 @@ class TestEnterpriseWorkspace:
|
||||
mock_tenant.status = TenantStatus.NORMAL
|
||||
mock_tenant.created_at = now
|
||||
mock_tenant.updated_at = now
|
||||
mock_tenant_svc.create_tenant.return_value = mock_tenant
|
||||
mock_tenant_svc.create_owner_tenant.return_value = mock_tenant
|
||||
|
||||
# Act — unwrap to bypass auth/setup decorators (tested in test_auth_wraps.py)
|
||||
unwrapped_post = inspect.unwrap(api_instance.post)
|
||||
@@ -115,11 +114,12 @@ class TestEnterpriseWorkspace:
|
||||
assert result["message"] == "enterprise workspace created."
|
||||
assert result["tenant"]["id"] == "tenant-id"
|
||||
assert result["tenant"]["name"] == "My Workspace"
|
||||
mock_tenant_svc.create_tenant.assert_called_once_with("My Workspace", is_from_dashboard=True, session=ANY)
|
||||
mock_tenant_svc.create_tenant_member.assert_called_once_with(
|
||||
mock_tenant, mock_account, mock_db.session(), role="owner"
|
||||
mock_tenant_svc.create_owner_tenant.assert_called_once_with(
|
||||
mock_account,
|
||||
name="My Workspace",
|
||||
is_from_dashboard=True,
|
||||
session=ANY,
|
||||
)
|
||||
mock_event.send.assert_called_once_with(mock_tenant)
|
||||
|
||||
@patch("controllers.inner_api.workspace.workspace.db")
|
||||
def test_post_returns_404_when_owner_not_found(self, mock_db, api_instance, app: Flask):
|
||||
|
||||
@@ -4,11 +4,13 @@ from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
import controllers.web.human_input_form as human_input_module
|
||||
@@ -16,13 +18,21 @@ import controllers.web.site as site_module
|
||||
from controllers.web.error import WebFormRateLimitExceededError
|
||||
from core.workflow.nodes.human_input.entities import ParagraphInputConfig, SelectInputConfig, StringListSource
|
||||
from core.workflow.nodes.human_input.enums import ValueSourceType
|
||||
from models import Tenant
|
||||
from models.enums import CustomizeTokenStrategy
|
||||
from models.human_input import RecipientType
|
||||
from models.model import App, AppMode, IconType, Site
|
||||
from services.feature_service import FeatureModel
|
||||
from services.human_input_service import FormExpiredError
|
||||
|
||||
HumanInputFormApi = human_input_module.HumanInputFormApi
|
||||
HumanInputFormUploadTokenApi = human_input_module.HumanInputFormUploadTokenApi
|
||||
TenantStatus = human_input_module.TenantStatus
|
||||
|
||||
SQLITE_MODELS = (Tenant, App, Site)
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures("sqlite_session"),
|
||||
pytest.mark.parametrize("sqlite_session", [SQLITE_MODELS], indirect=True),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -34,36 +44,65 @@ def app() -> Flask:
|
||||
return app
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
"""Simple stand-in for db.session that returns pre-seeded objects."""
|
||||
@pytest.fixture
|
||||
def database_session(
|
||||
sqlite_session: Session,
|
||||
sqlite_engine: Engine,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> Session:
|
||||
"""Bind model/controller database access to the shared SQLite session."""
|
||||
|
||||
def __init__(self, mapping: dict[str, Any]):
|
||||
self._mapping = mapping
|
||||
|
||||
def get(self, model, ident):
|
||||
return self._mapping.get(model.__name__)
|
||||
|
||||
def scalar(self, stmt):
|
||||
# Extract the model name from the select statement's column_descriptions
|
||||
try:
|
||||
name = stmt.column_descriptions[0]["entity"].__name__
|
||||
except (AttributeError, IndexError, KeyError):
|
||||
return None
|
||||
return self._mapping.get(name)
|
||||
database = SimpleNamespace(engine=sqlite_engine, session=sqlite_session)
|
||||
monkeypatch.setattr(human_input_module, "db", database)
|
||||
monkeypatch.setattr("models.model.db", database)
|
||||
return sqlite_session
|
||||
|
||||
|
||||
class _FakeDB:
|
||||
"""Minimal db stub exposing engine and session."""
|
||||
|
||||
def __init__(self, session: _FakeSession):
|
||||
self.session = session
|
||||
self.engine = object()
|
||||
def _persist_app_site(session: Session, *, include_site: bool = True) -> tuple[Tenant, App, Site | None]:
|
||||
tenant = Tenant(name="Tenant", plan="basic")
|
||||
tenant.custom_config_dict = {"remove_webapp_brand": True, "replace_webapp_logo": None}
|
||||
app_model = App(
|
||||
id=str(uuid4()),
|
||||
tenant_id=tenant.id,
|
||||
name="Human Input App",
|
||||
mode=AppMode.CHAT,
|
||||
icon_type=IconType.EMOJI,
|
||||
icon="robot",
|
||||
icon_background="#fff",
|
||||
enable_site=True,
|
||||
enable_api=False,
|
||||
)
|
||||
site = None
|
||||
models: list[object] = [tenant, app_model]
|
||||
if include_site:
|
||||
site = Site(
|
||||
app_id=app_model.id,
|
||||
title="My Site",
|
||||
default_language="en",
|
||||
customize_token_strategy=CustomizeTokenStrategy.UUID,
|
||||
icon_type=IconType.EMOJI,
|
||||
icon="robot",
|
||||
icon_background="#fff",
|
||||
description="desc",
|
||||
input_placeholder="Ask the app",
|
||||
chat_color_theme="light",
|
||||
chat_color_theme_inverted=False,
|
||||
custom_disclaimer="",
|
||||
prompt_public=False,
|
||||
show_workflow_steps=True,
|
||||
use_icon_as_answer_icon=False,
|
||||
)
|
||||
models.append(site)
|
||||
session.add_all(models)
|
||||
session.commit()
|
||||
return tenant, app_model, site
|
||||
|
||||
|
||||
def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask, database_session: Session):
|
||||
"""GET returns form definition merged with site payload."""
|
||||
|
||||
expiration_time = datetime(2099, 1, 1, tzinfo=UTC)
|
||||
_, app_model, _ = _persist_app_site(database_session)
|
||||
|
||||
class _FakeDefinition:
|
||||
def model_dump(self, mode: str | None = None):
|
||||
@@ -77,9 +116,9 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
|
||||
class _FakeForm:
|
||||
def __init__(self, expiration: datetime):
|
||||
self.workflow_run_id = "workflow-1"
|
||||
self.app_id = "app-1"
|
||||
self.tenant_id = "tenant-1"
|
||||
self.workflow_run_id = None
|
||||
self.app_id = app_model.id
|
||||
self.tenant_id = app_model.tenant_id
|
||||
self.expiration_time = expiration
|
||||
self.recipient_type = RecipientType.BACKSTAGE
|
||||
|
||||
@@ -92,32 +131,6 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
monkeypatch.setattr(human_input_module, "_FORM_ACCESS_RATE_LIMITER", limiter_mock)
|
||||
monkeypatch.setattr(human_input_module, "extract_remote_ip", lambda req: "203.0.113.10")
|
||||
|
||||
tenant = SimpleNamespace(
|
||||
id="tenant-1",
|
||||
status=TenantStatus.NORMAL,
|
||||
plan="basic",
|
||||
custom_config_dict={"remove_webapp_brand": True, "replace_webapp_logo": False},
|
||||
)
|
||||
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", tenant=tenant, enable_site=True)
|
||||
workflow_run = SimpleNamespace(app_id="app-1")
|
||||
site_model = SimpleNamespace(
|
||||
title="My Site",
|
||||
icon_type="emoji",
|
||||
icon="robot",
|
||||
icon_background="#fff",
|
||||
description="desc",
|
||||
input_placeholder="Ask the app",
|
||||
default_language="en",
|
||||
chat_color_theme="light",
|
||||
chat_color_theme_inverted=False,
|
||||
copyright=None,
|
||||
privacy_policy=None,
|
||||
custom_disclaimer="",
|
||||
prompt_public=False,
|
||||
show_workflow_steps=True,
|
||||
use_icon_as_answer_icon=False,
|
||||
)
|
||||
|
||||
# Patch service to return fake form.
|
||||
service_mock = MagicMock()
|
||||
service_mock.get_form_by_token.return_value = form
|
||||
@@ -125,10 +138,6 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
service_mock.resolve_form_inputs.return_value = [resolved_input]
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
|
||||
# Patch db session.
|
||||
db_stub = _FakeDB(_FakeSession({"WorkflowRun": workflow_run, "App": app_model, "Site": site_model}))
|
||||
monkeypatch.setattr(human_input_module, "db", db_stub)
|
||||
|
||||
monkeypatch.setattr(
|
||||
site_module.FeatureService,
|
||||
"get_features",
|
||||
@@ -153,7 +162,7 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
assert body["user_actions"] == [{"id": "approve", "title": "Approve", "button_style": "default"}]
|
||||
assert body["expiration_time"] == int(expiration_time.timestamp())
|
||||
assert body["site"] == {
|
||||
"app_id": "app-1",
|
||||
"app_id": app_model.id,
|
||||
"end_user_id": None,
|
||||
"enable_site": True,
|
||||
"site": {
|
||||
@@ -187,10 +196,11 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
limiter_mock.increment_rate_limit.assert_called_once_with("203.0.113.10")
|
||||
|
||||
|
||||
def test_get_form_uses_runtime_select_options(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
def test_get_form_uses_runtime_select_options(monkeypatch: pytest.MonkeyPatch, app: Flask, database_session: Session):
|
||||
"""GET returns variable-backed select options resolved from runtime state."""
|
||||
|
||||
expiration_time = datetime(2099, 1, 1, tzinfo=UTC)
|
||||
_, app_model, _ = _persist_app_site(database_session)
|
||||
configured_inputs = [
|
||||
{
|
||||
"type": "select",
|
||||
@@ -225,9 +235,9 @@ def test_get_form_uses_runtime_select_options(monkeypatch: pytest.MonkeyPatch, a
|
||||
|
||||
class _FakeForm:
|
||||
def __init__(self, expiration: datetime):
|
||||
self.workflow_run_id = "workflow-1"
|
||||
self.app_id = "app-1"
|
||||
self.tenant_id = "tenant-1"
|
||||
self.workflow_run_id = None
|
||||
self.app_id = app_model.id
|
||||
self.tenant_id = app_model.tenant_id
|
||||
self.recipient_type = RecipientType.STANDALONE_WEB_APP
|
||||
self.expiration_time = expiration
|
||||
|
||||
@@ -239,37 +249,11 @@ def test_get_form_uses_runtime_select_options(monkeypatch: pytest.MonkeyPatch, a
|
||||
monkeypatch.setattr(human_input_module, "_FORM_ACCESS_RATE_LIMITER", limiter_mock)
|
||||
monkeypatch.setattr(human_input_module, "extract_remote_ip", lambda req: "203.0.113.10")
|
||||
|
||||
tenant = SimpleNamespace(
|
||||
id="tenant-1",
|
||||
status=TenantStatus.NORMAL,
|
||||
plan="basic",
|
||||
custom_config_dict={},
|
||||
)
|
||||
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", tenant=tenant, enable_site=True)
|
||||
site_model = SimpleNamespace(
|
||||
title="My Site",
|
||||
icon_type="emoji",
|
||||
icon="robot",
|
||||
icon_background="#fff",
|
||||
description="desc",
|
||||
input_placeholder="Ask the app",
|
||||
default_language="en",
|
||||
chat_color_theme="light",
|
||||
chat_color_theme_inverted=False,
|
||||
copyright=None,
|
||||
privacy_policy=None,
|
||||
custom_disclaimer="",
|
||||
prompt_public=False,
|
||||
show_workflow_steps=True,
|
||||
use_icon_as_answer_icon=False,
|
||||
)
|
||||
|
||||
form = _FakeForm(expiration_time)
|
||||
service_mock = MagicMock()
|
||||
service_mock.get_form_by_token.return_value = form
|
||||
service_mock.resolve_form_inputs.return_value = runtime_inputs
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
monkeypatch.setattr(human_input_module, "db", _FakeDB(_FakeSession({"App": app_model, "Site": site_model})))
|
||||
|
||||
def mock_get_features(tenant_id: str, exclude_vector_space: bool = False):
|
||||
return FeatureModel(can_replace_logo=True)
|
||||
@@ -284,7 +268,9 @@ def test_get_form_uses_runtime_select_options(monkeypatch: pytest.MonkeyPatch, a
|
||||
service_mock.resolve_form_inputs.assert_called_once_with(form)
|
||||
|
||||
|
||||
def test_create_upload_token_returns_token_and_form_expiration(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
def test_create_upload_token_returns_token_and_form_expiration(
|
||||
monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine
|
||||
):
|
||||
"""POST returns a HITL upload token for an active form token."""
|
||||
|
||||
expiration_time = datetime(2099, 1, 1, tzinfo=UTC)
|
||||
@@ -312,7 +298,7 @@ def test_create_upload_token_returns_token_and_form_expiration(monkeypatch: pyte
|
||||
"HumanInputFileUploadService",
|
||||
_service_factory,
|
||||
)
|
||||
monkeypatch.setattr(human_input_module, "db", SimpleNamespace(engine=object()))
|
||||
monkeypatch.setattr(human_input_module, "db", SimpleNamespace(engine=sqlite_engine))
|
||||
|
||||
limiter_mock = MagicMock()
|
||||
limiter_mock.is_rate_limited.return_value = False
|
||||
@@ -329,14 +315,18 @@ def test_create_upload_token_returns_token_and_form_expiration(monkeypatch: pyte
|
||||
}
|
||||
repo_factory.assert_called_once()
|
||||
assert captured["workflow_run_repository"] is workflow_run_repository
|
||||
session_factory = captured["session_factory"]
|
||||
assert isinstance(session_factory, sessionmaker)
|
||||
assert session_factory.kw["bind"] is sqlite_engine
|
||||
service_mock.issue_upload_token.assert_called_once_with("token-1")
|
||||
limiter_mock.increment_rate_limit.assert_called_once_with("203.0.113.10")
|
||||
|
||||
|
||||
def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: Flask, database_session: Session):
|
||||
"""GET returns form payload for backstage token."""
|
||||
|
||||
expiration_time = datetime(2099, 1, 2, tzinfo=UTC)
|
||||
_, app_model, _ = _persist_app_site(database_session)
|
||||
|
||||
class _FakeDefinition:
|
||||
def model_dump(self, mode: str | None = None):
|
||||
@@ -350,9 +340,9 @@ def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: F
|
||||
|
||||
class _FakeForm:
|
||||
def __init__(self, expiration: datetime):
|
||||
self.workflow_run_id = "workflow-1"
|
||||
self.app_id = "app-1"
|
||||
self.tenant_id = "tenant-1"
|
||||
self.workflow_run_id = None
|
||||
self.app_id = app_model.id
|
||||
self.tenant_id = app_model.tenant_id
|
||||
self.expiration_time = expiration
|
||||
|
||||
def get_definition(self):
|
||||
@@ -363,40 +353,11 @@ def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: F
|
||||
limiter_mock.is_rate_limited.return_value = False
|
||||
monkeypatch.setattr(human_input_module, "_FORM_ACCESS_RATE_LIMITER", limiter_mock)
|
||||
monkeypatch.setattr(human_input_module, "extract_remote_ip", lambda req: "203.0.113.10")
|
||||
tenant = SimpleNamespace(
|
||||
id="tenant-1",
|
||||
status=TenantStatus.NORMAL,
|
||||
plan="basic",
|
||||
custom_config_dict={"remove_webapp_brand": True, "replace_webapp_logo": False},
|
||||
)
|
||||
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", tenant=tenant, enable_site=True)
|
||||
workflow_run = SimpleNamespace(app_id="app-1")
|
||||
site_model = SimpleNamespace(
|
||||
title="My Site",
|
||||
icon_type="emoji",
|
||||
icon="robot",
|
||||
icon_background="#fff",
|
||||
description="desc",
|
||||
input_placeholder="Ask the app",
|
||||
default_language="en",
|
||||
chat_color_theme="light",
|
||||
chat_color_theme_inverted=False,
|
||||
copyright=None,
|
||||
privacy_policy=None,
|
||||
custom_disclaimer="",
|
||||
prompt_public=False,
|
||||
show_workflow_steps=True,
|
||||
use_icon_as_answer_icon=False,
|
||||
)
|
||||
|
||||
service_mock = MagicMock()
|
||||
service_mock.get_form_by_token.return_value = form
|
||||
service_mock.resolve_form_inputs.return_value = []
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
|
||||
db_stub = _FakeDB(_FakeSession({"WorkflowRun": workflow_run, "App": app_model, "Site": site_model}))
|
||||
monkeypatch.setattr(human_input_module, "db", db_stub)
|
||||
|
||||
monkeypatch.setattr(
|
||||
site_module.FeatureService,
|
||||
"get_features",
|
||||
@@ -421,7 +382,7 @@ def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: F
|
||||
assert body["user_actions"] == []
|
||||
assert body["expiration_time"] == int(expiration_time.timestamp())
|
||||
assert body["site"] == {
|
||||
"app_id": "app-1",
|
||||
"app_id": app_model.id,
|
||||
"end_user_id": None,
|
||||
"enable_site": True,
|
||||
"site": {
|
||||
@@ -455,10 +416,13 @@ def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: F
|
||||
limiter_mock.increment_rate_limit.assert_called_once_with("203.0.113.10")
|
||||
|
||||
|
||||
def test_get_form_raises_forbidden_when_site_missing(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
def test_get_form_raises_forbidden_when_site_missing(
|
||||
monkeypatch: pytest.MonkeyPatch, app: Flask, database_session: Session
|
||||
):
|
||||
"""GET raises Forbidden if site cannot be resolved."""
|
||||
|
||||
expiration_time = datetime(2099, 1, 3, tzinfo=UTC)
|
||||
_, app_model, _ = _persist_app_site(database_session, include_site=False)
|
||||
|
||||
class _FakeDefinition:
|
||||
def model_dump(self, mode: str | None = None):
|
||||
@@ -472,9 +436,9 @@ def test_get_form_raises_forbidden_when_site_missing(monkeypatch: pytest.MonkeyP
|
||||
|
||||
class _FakeForm:
|
||||
def __init__(self, expiration: datetime):
|
||||
self.workflow_run_id = "workflow-1"
|
||||
self.app_id = "app-1"
|
||||
self.tenant_id = "tenant-1"
|
||||
self.workflow_run_id = None
|
||||
self.app_id = app_model.id
|
||||
self.tenant_id = app_model.tenant_id
|
||||
self.expiration_time = expiration
|
||||
|
||||
def get_definition(self):
|
||||
@@ -485,17 +449,10 @@ def test_get_form_raises_forbidden_when_site_missing(monkeypatch: pytest.MonkeyP
|
||||
limiter_mock.is_rate_limited.return_value = False
|
||||
monkeypatch.setattr(human_input_module, "_FORM_ACCESS_RATE_LIMITER", limiter_mock)
|
||||
monkeypatch.setattr(human_input_module, "extract_remote_ip", lambda req: "203.0.113.10")
|
||||
tenant = SimpleNamespace(status=TenantStatus.NORMAL)
|
||||
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", tenant=tenant)
|
||||
workflow_run = SimpleNamespace(app_id="app-1")
|
||||
|
||||
service_mock = MagicMock()
|
||||
service_mock.get_form_by_token.return_value = form
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
|
||||
db_stub = _FakeDB(_FakeSession({"WorkflowRun": workflow_run, "App": app_model, "Site": None}))
|
||||
monkeypatch.setattr(human_input_module, "db", db_stub)
|
||||
|
||||
with app.test_request_context("/api/form/human_input/token-1", method="GET"):
|
||||
with pytest.raises(Forbidden):
|
||||
HumanInputFormApi().get("token-1")
|
||||
@@ -503,7 +460,7 @@ def test_get_form_raises_forbidden_when_site_missing(monkeypatch: pytest.MonkeyP
|
||||
limiter_mock.increment_rate_limit.assert_called_once_with("203.0.113.10")
|
||||
|
||||
|
||||
def test_submit_form_accepts_backstage_token(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
def test_submit_form_accepts_backstage_token(monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine):
|
||||
"""POST forwards backstage submissions to the service."""
|
||||
|
||||
class _FakeForm:
|
||||
@@ -517,7 +474,7 @@ def test_submit_form_accepts_backstage_token(monkeypatch: pytest.MonkeyPatch, ap
|
||||
service_mock = MagicMock()
|
||||
service_mock.get_form_by_token.return_value = form
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
monkeypatch.setattr(human_input_module, "db", _FakeDB(_FakeSession({})))
|
||||
monkeypatch.setattr(human_input_module, "db", SimpleNamespace(engine=sqlite_engine))
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/form/human_input/token-1",
|
||||
@@ -550,7 +507,6 @@ def test_submit_form_rate_limited(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
service_mock = MagicMock()
|
||||
service_mock.get_form_by_token.return_value = None
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
monkeypatch.setattr(human_input_module, "db", _FakeDB(_FakeSession({})))
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/form/human_input/token-1",
|
||||
@@ -576,7 +532,6 @@ def test_get_form_rate_limited(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
service_mock = MagicMock()
|
||||
service_mock.get_form_by_token.return_value = None
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
monkeypatch.setattr(human_input_module, "db", _FakeDB(_FakeSession({})))
|
||||
|
||||
with app.test_request_context("/api/form/human_input/token-1", method="GET"):
|
||||
with pytest.raises(WebFormRateLimitExceededError):
|
||||
@@ -587,7 +542,7 @@ def test_get_form_rate_limited(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
service_mock.get_form_by_token.assert_not_called()
|
||||
|
||||
|
||||
def test_get_form_raises_expired(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
def test_get_form_raises_expired(monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine):
|
||||
class _FakeForm:
|
||||
pass
|
||||
|
||||
@@ -600,7 +555,7 @@ def test_get_form_raises_expired(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
service_mock.get_form_by_token.return_value = form
|
||||
service_mock.ensure_form_active.side_effect = FormExpiredError("form-id")
|
||||
monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock)
|
||||
monkeypatch.setattr(human_input_module, "db", _FakeDB(_FakeSession({})))
|
||||
monkeypatch.setattr(human_input_module, "db", SimpleNamespace(engine=sqlite_engine))
|
||||
|
||||
with app.test_request_context("/api/form/human_input/token-1", method="GET"):
|
||||
with pytest.raises(FormExpiredError):
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
from uuid import NAMESPACE_URL, UUID, uuid5
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
from werkzeug.exceptions import NotFound, Unauthorized
|
||||
|
||||
from controllers.web.error import WebAppAuthRequiredError
|
||||
@@ -12,9 +15,41 @@ from controllers.web.passport import (
|
||||
exchange_token_for_existing_web_user,
|
||||
generate_session_id,
|
||||
)
|
||||
from models.enums import CustomizeTokenStrategy, EndUserType
|
||||
from models.model import App, AppMode, EndUser, IconType, Site
|
||||
from services.webapp_auth_service import WebAppAuthType
|
||||
|
||||
|
||||
def _stable_uuid(value: str) -> str:
|
||||
return str(uuid5(NAMESPACE_URL, value))
|
||||
|
||||
|
||||
def _persist_webapp(session: Session, *, app_code: str = "code") -> tuple[App, Site]:
|
||||
tenant_id = _stable_uuid(f"tenant:{app_code}")
|
||||
app_model = App(
|
||||
id=_stable_uuid(f"app:{app_code}"),
|
||||
tenant_id=tenant_id,
|
||||
name="Web App",
|
||||
mode=AppMode.CHAT,
|
||||
icon_type=IconType.EMOJI,
|
||||
icon="chat",
|
||||
icon_background="#FFFFFF",
|
||||
enable_site=True,
|
||||
enable_api=False,
|
||||
)
|
||||
site = Site(
|
||||
id=_stable_uuid(f"site:{app_code}"),
|
||||
app_id=app_model.id,
|
||||
title="Web App Site",
|
||||
default_language="en-US",
|
||||
customize_token_strategy=CustomizeTokenStrategy.UUID,
|
||||
code=app_code,
|
||||
)
|
||||
session.add_all([app_model, site])
|
||||
session.commit()
|
||||
return app_model, site
|
||||
|
||||
|
||||
def test_decode_enterprise_webapp_user_id_none() -> None:
|
||||
assert decode_enterprise_webapp_user_id(None) is None
|
||||
|
||||
@@ -31,70 +66,68 @@ def test_decode_enterprise_webapp_user_id_valid(monkeypatch: pytest.MonkeyPatch)
|
||||
assert decode_enterprise_webapp_user_id("token") == decoded
|
||||
|
||||
|
||||
def test_exchange_token_public_flow(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
site = SimpleNamespace(id="s1", app_id="a1", code="code", status="normal")
|
||||
app_model = SimpleNamespace(id="a1", status="normal", enable_site=True)
|
||||
call_state = {"calls": 0}
|
||||
|
||||
def _scalar_side_effect(*_args, **_kwargs):
|
||||
call_state["calls"] += 1
|
||||
return site if call_state["calls"] == 1 else app_model
|
||||
|
||||
db_session = SimpleNamespace(scalar=_scalar_side_effect)
|
||||
monkeypatch.setattr("controllers.web.passport.db", SimpleNamespace(session=db_session))
|
||||
monkeypatch.setattr("controllers.web.passport._exchange_for_public_app_token", lambda *_args, **_kwargs: "resp")
|
||||
@pytest.mark.parametrize("sqlite_session", [(App, Site)], indirect=True)
|
||||
def test_exchange_token_public_flow(sqlite_session: Session) -> None:
|
||||
app_model, site = _persist_webapp(sqlite_session)
|
||||
|
||||
decoded = {"auth_type": "public"}
|
||||
result = exchange_token_for_existing_web_user("code", decoded, WebAppAuthType.PUBLIC)
|
||||
with (
|
||||
patch("controllers.web.passport.db.session", sqlite_session),
|
||||
patch("controllers.web.passport._exchange_for_public_app_token", return_value="resp") as exchange_mock,
|
||||
):
|
||||
result = exchange_token_for_existing_web_user("code", decoded, WebAppAuthType.PUBLIC)
|
||||
|
||||
assert result == "resp"
|
||||
exchange_mock.assert_called_once_with(app_model, site, decoded)
|
||||
|
||||
|
||||
def test_exchange_token_requires_external(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
site = SimpleNamespace(id="s1", app_id="a1", code="code", status="normal")
|
||||
app_model = SimpleNamespace(id="a1", status="normal", enable_site=True)
|
||||
call_state = {"calls": 0}
|
||||
|
||||
def _scalar_side_effect(*_args, **_kwargs):
|
||||
call_state["calls"] += 1
|
||||
return site if call_state["calls"] == 1 else app_model
|
||||
|
||||
db_session = SimpleNamespace(scalar=_scalar_side_effect)
|
||||
monkeypatch.setattr("controllers.web.passport.db", SimpleNamespace(session=db_session))
|
||||
@pytest.mark.parametrize("sqlite_session", [(App, Site)], indirect=True)
|
||||
def test_exchange_token_requires_external(sqlite_session: Session) -> None:
|
||||
_persist_webapp(sqlite_session)
|
||||
|
||||
decoded = {"auth_type": "internal"}
|
||||
with pytest.raises(WebAppAuthRequiredError):
|
||||
with (
|
||||
patch("controllers.web.passport.db.session", sqlite_session),
|
||||
pytest.raises(WebAppAuthRequiredError),
|
||||
):
|
||||
exchange_token_for_existing_web_user("code", decoded, WebAppAuthType.EXTERNAL)
|
||||
|
||||
|
||||
def test_exchange_token_missing_session_id(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
site = SimpleNamespace(id="s1", app_id="a1", code="code", status="normal")
|
||||
app_model = SimpleNamespace(id="a1", status="normal", enable_site=True, tenant_id="t1")
|
||||
call_state = {"calls": 0}
|
||||
|
||||
def _scalar_side_effect(*_args, **_kwargs):
|
||||
call_state["calls"] += 1
|
||||
if call_state["calls"] == 1:
|
||||
return site
|
||||
if call_state["calls"] == 2:
|
||||
return app_model
|
||||
return None
|
||||
|
||||
db_session = SimpleNamespace(scalar=_scalar_side_effect, add=lambda *_a, **_k: None, commit=lambda: None)
|
||||
monkeypatch.setattr("controllers.web.passport.db", SimpleNamespace(session=db_session))
|
||||
@pytest.mark.parametrize("sqlite_session", [(App, Site, EndUser)], indirect=True)
|
||||
def test_exchange_token_missing_session_id(sqlite_session: Session) -> None:
|
||||
_persist_webapp(sqlite_session)
|
||||
|
||||
decoded = {"auth_type": "internal"}
|
||||
with pytest.raises(NotFound):
|
||||
with (
|
||||
patch("controllers.web.passport.db.session", sqlite_session),
|
||||
pytest.raises(NotFound),
|
||||
):
|
||||
exchange_token_for_existing_web_user("code", decoded, WebAppAuthType.INTERNAL)
|
||||
assert sqlite_session.scalars(select(EndUser)).all() == []
|
||||
|
||||
|
||||
def test_generate_session_id(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
counts = [1, 0]
|
||||
@pytest.mark.parametrize("sqlite_session", [(EndUser,)], indirect=True)
|
||||
def test_generate_session_id(sqlite_session: Session) -> None:
|
||||
collision_id = _stable_uuid("session:collision")
|
||||
generated_id = _stable_uuid("session:generated")
|
||||
sqlite_session.add(
|
||||
EndUser(
|
||||
id=_stable_uuid("end-user:collision"),
|
||||
tenant_id=_stable_uuid("tenant:collision"),
|
||||
type=EndUserType.BROWSER,
|
||||
name="Existing User",
|
||||
session_id=collision_id,
|
||||
)
|
||||
)
|
||||
sqlite_session.commit()
|
||||
|
||||
def _scalar(*_args, **_kwargs):
|
||||
return counts.pop(0)
|
||||
with (
|
||||
patch("controllers.web.passport.db.session", sqlite_session),
|
||||
patch(
|
||||
"controllers.web.passport.uuid.uuid4",
|
||||
side_effect=[UUID(collision_id), UUID(generated_id)],
|
||||
),
|
||||
):
|
||||
session_id = generate_session_id()
|
||||
|
||||
db_session = SimpleNamespace(scalar=_scalar)
|
||||
monkeypatch.setattr("controllers.web.passport.db", SimpleNamespace(session=db_session))
|
||||
|
||||
session_id = generate_session_id()
|
||||
assert session_id
|
||||
assert session_id == generated_id
|
||||
|
||||
@@ -4,7 +4,7 @@ saved, using the deterministic fake backend client (no live stack)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Callable, Iterator
|
||||
from datetime import UTC, datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, override
|
||||
@@ -106,8 +106,14 @@ class _RecordingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
|
||||
class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield PydanticAIStreamRunEvent(
|
||||
@@ -138,8 +144,14 @@ class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
|
||||
class _StreamingRecordingFakeAgentBackendRunClient(_RecordingFakeAgentBackendRunClient):
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield PydanticAIStreamRunEvent(
|
||||
@@ -173,8 +185,14 @@ class _StreamingStopAfterFirstDeltaFakeAgentBackendRunClient(_RecordingFakeAgent
|
||||
self._queue_manager = queue_manager
|
||||
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield PydanticAIStreamRunEvent(
|
||||
@@ -196,8 +214,14 @@ class _StreamingStopAfterFirstDeltaFakeAgentBackendRunClient(_RecordingFakeAgent
|
||||
|
||||
class _StreamingSingleAgentMessageDeltaFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield PydanticAIStreamRunEvent(
|
||||
@@ -220,8 +244,14 @@ class _StreamingSingleAgentMessageDeltaFakeAgentBackendRunClient(FakeAgentBacken
|
||||
|
||||
class _NullOutputFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield RunSucceededEvent(
|
||||
@@ -237,8 +267,14 @@ class _NullOutputFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
|
||||
class _StreamingTextNullOutputFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield PydanticAIStreamRunEvent(
|
||||
@@ -261,8 +297,14 @@ class _StreamingTextNullOutputFakeAgentBackendRunClient(FakeAgentBackendRunClien
|
||||
|
||||
class _AgentAnswerStreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield PydanticAIStreamRunEvent(
|
||||
@@ -292,8 +334,14 @@ class _AgentAnswerStreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
|
||||
class _ProcessStreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
|
||||
@override
|
||||
def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]:
|
||||
del after
|
||||
def stream_events(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
after: str | None = None,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
) -> Iterator[RunEvent]:
|
||||
del after, should_stop
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
|
||||
yield PydanticAIStreamRunEvent(
|
||||
|
||||
@@ -61,16 +61,37 @@ class TestBaseAppQueueManager:
|
||||
manager._check_for_sqlalchemy_models(bad)
|
||||
|
||||
def test_stop_listen_defers_graph_runtime_state_cleanup_until_listener_exits(self):
|
||||
with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis:
|
||||
with (
|
||||
patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis,
|
||||
patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager,
|
||||
):
|
||||
mock_redis.setex.return_value = True
|
||||
mock_redis.get.return_value = None
|
||||
manager = DummyQueueManager(task_id="t1", user_id="u1", invoke_from=InvokeFrom.SERVICE_API)
|
||||
runtime_state = SimpleNamespace(name="runtime-state")
|
||||
manager.graph_runtime_state = runtime_state
|
||||
|
||||
runtime_state = SimpleNamespace(name="runtime-state")
|
||||
manager.graph_runtime_state = runtime_state
|
||||
manager.stop_listen()
|
||||
|
||||
manager.stop_listen()
|
||||
assert manager.graph_runtime_state is runtime_state
|
||||
assert list(manager.listen()) == []
|
||||
assert manager.graph_runtime_state is None
|
||||
graph_engine_manager.return_value.send_stop_command.assert_called_once_with(
|
||||
"t1",
|
||||
reason="Client response stream closed before app execution completed",
|
||||
)
|
||||
|
||||
assert manager.graph_runtime_state is runtime_state
|
||||
assert list(manager.listen()) == []
|
||||
assert manager.graph_runtime_state is None
|
||||
def test_abort_execution_is_idempotent_when_graph_stop_command_fails(self, caplog):
|
||||
with (
|
||||
patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis,
|
||||
patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager,
|
||||
):
|
||||
mock_redis.setex.return_value = True
|
||||
graph_engine_manager.return_value.send_stop_command.side_effect = RuntimeError("redis unavailable")
|
||||
manager = DummyQueueManager(task_id="t1", user_id="u1", invoke_from=InvokeFrom.SERVICE_API)
|
||||
|
||||
manager._abort_execution("stream closed")
|
||||
manager._abort_execution("duplicate")
|
||||
|
||||
graph_engine_manager.return_value.send_stop_command.assert_called_once_with("t1", reason="stream closed")
|
||||
assert "Failed to abort app execution for task t1" in caplog.text
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user