Merge feat/openapi-app-build-loop into feat/difyctl-app-build-loop

This commit is contained in:
GareArc
2026-10-15 15:17:14 -07:00
489 changed files with 5143 additions and 51219 deletions
-8
View File
@@ -733,14 +733,6 @@ AGENT_BACKEND_STREAM_MAX_RECONNECTS=3
# Client deadline for converting a Binding file to a ToolFile through the Agent backend.
AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS=240
# KnowledgeFS (Dataset 2.0)
KNOWLEDGE_FS_ENABLED=false
KNOWLEDGE_FS_BASE_URL=
# Shared with KnowledgeFS; use at least 32 random characters.
KNOWLEDGE_FS_JWT_SECRET=
KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS=300
KNOWLEDGE_FS_TIMEOUT_SECONDS=10
# Marketplace configuration
MARKETPLACE_ENABLED=true
MARKETPLACE_API_URL=https://marketplace.dify.ai
-3
View File
@@ -59,7 +59,6 @@ select = [
ignore = [
"E402", # module-import-not-at-top-of-file
"E712", # true-false-comparison
"F821", # undefined-name
"F841", # unused-variable
"UP007", # non-pep604-annotation
@@ -71,8 +70,6 @@ ignore = [
"N815", # mixed-case-variable-in-class-scope
"PT011", # pytest-raises-too-broad
"SIM102", # collapsible-if
"SIM105", # suppressible-exception
"SIM108", # if-else-block-instead-of-if-exp
"SIM117", # multiple-with-statements
]
+3 -3
View File
@@ -252,8 +252,8 @@ def migrate_knowledge_vector_database():
select(DatasetDocument).where(
DatasetDocument.dataset_id == dataset.id,
DatasetDocument.indexing_status == IndexingStatus.COMPLETED,
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
).all()
@@ -264,7 +264,7 @@ def migrate_knowledge_vector_database():
select(DocumentSegment).where(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.status == SegmentStatus.COMPLETED,
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
)
).all()
-2
View File
@@ -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.logstore_config import LogStoreConfig
from configs.extra.notion_config import NotionConfig
from configs.extra.sentry_config import SentryConfig
@@ -11,7 +10,6 @@ class ExtraServiceConfig(
# place the configs in alphabet order
AgentBackendConfig,
ArchiveStorageConfig,
KnowledgeFSConfig,
LogStoreConfig,
NotionConfig,
SentryConfig,
-64
View File
@@ -1,64 +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_ENABLED: bool = Field(
default=False,
description="Enable the private KnowledgeFS Console bridge.",
)
KNOWLEDGE_FS_BASE_URL: str | None = Field(default=None, description="KnowledgeFS gateway base URL.")
KNOWLEDGE_FS_JWT_SECRET: SecretStr | None = Field(
default=None,
min_length=32,
description="Shared secret used to sign short-lived KnowledgeFS service JWTs.",
)
KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS: PositiveFloat = Field(default=300.0, le=3600.0, allow_inf_nan=False)
KNOWLEDGE_FS_TIMEOUT_SECONDS: PositiveFloat = Field(default=10.0, le=60.0, allow_inf_nan=False)
@field_validator(
"KNOWLEDGE_FS_BASE_URL",
"KNOWLEDGE_FS_JWT_SECRET",
mode="before",
)
@classmethod
def normalize_optional_string(cls, value: object) -> object:
if isinstance(value, SecretStr):
normalized = value.get_secret_value().strip()
return SecretStr(normalized) if normalized else None
if isinstance(value, str):
normalized = value.strip()
return normalized or None
return value
@field_validator("KNOWLEDGE_FS_BASE_URL")
@classmethod
def validate_base_url(cls, value: str | None) -> str | None:
if value is None:
return None
parsed = urlsplit(value)
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
raise ValueError("KNOWLEDGE_FS_BASE_URL must be an absolute HTTP(S) URL")
try:
_ = parsed.port
except ValueError as exc:
raise ValueError("KNOWLEDGE_FS_BASE_URL must include a valid port") from exc
if parsed.username or parsed.password or parsed.query or parsed.fragment:
raise ValueError("KNOWLEDGE_FS_BASE_URL must not include credentials, query, or fragment")
return value.rstrip("/")
@model_validator(mode="after")
def validate_enabled_connection(self) -> "KnowledgeFSConfig":
if not self.KNOWLEDGE_FS_ENABLED:
return self
if bool(self.KNOWLEDGE_FS_BASE_URL) != bool(self.KNOWLEDGE_FS_JWT_SECRET):
raise ValueError("KNOWLEDGE_FS_BASE_URL and KNOWLEDGE_FS_JWT_SECRET must be configured together")
if not self.KNOWLEDGE_FS_BASE_URL:
raise ValueError("KnowledgeFS connection settings are required when the integration is enabled")
return self
+1 -4
View File
@@ -71,9 +71,6 @@ def resolve_agent_access_filter(
if has_agent_list_permission(override.permission_keys)
}
if whitelist_scope.unrestricted:
accessible_agent_ids = permission_agent_ids
else:
accessible_agent_ids = set(whitelist_scope.resource_ids)
accessible_agent_ids = permission_agent_ids if whitelist_scope.unrestricted else set(whitelist_scope.resource_ids)
return AgentAccessFilter(accessible_agent_ids=accessible_agent_ids)
+1 -4
View File
@@ -3,10 +3,7 @@ import urllib.parse
def decode_remote_url(url: str, query_string: bytes | str = b"") -> str:
decoded_url = urllib.parse.unquote(url)
if isinstance(query_string, bytes):
raw_query = query_string.decode()
else:
raw_query = query_string
raw_query = query_string.decode() if isinstance(query_string, bytes) else query_string
if not raw_query:
return decoded_url
-2
View File
@@ -136,7 +136,6 @@ from . import (
feature,
human_input_form,
init_validate,
knowledge_fs_proxy,
notification,
onboarding,
resource_access_token,
@@ -294,7 +293,6 @@ __all__ = [
"human_input_form",
"init_validate",
"installed_app",
"knowledge_fs_proxy",
"load_balancing_config",
"login",
"mcp_server",
+2 -3
View File
@@ -1,4 +1,3 @@
from typing import Any
from uuid import UUID
from flask_restx import Resource
@@ -55,7 +54,7 @@ class TraceProviderQuery(BaseModel):
class TraceConfigPayload(BaseModel):
tracing_provider: str = Field(..., description="Tracing provider name")
tracing_config: dict[str, Any] = Field(
tracing_config: dict[str, str | None] = Field(
...,
description="Tracing configuration data",
)
@@ -68,7 +67,7 @@ class TraceAppConfigResponse(ResponseModel):
id: str | None = None
app_id: str | None = None
tracing_provider: str | None = None
tracing_config: dict[str, Any] | None = Field(default=None)
tracing_config: dict[str, str | None] | None = Field(default=None)
is_active: bool | None = None
created_at: str | None = None
updated_at: str | None = None
+1 -4
View File
@@ -117,10 +117,7 @@ def get_app_model[**P, R](
app_mode = AppMode.value_of(app_model.mode)
if mode is not None:
if isinstance(mode, list):
modes = mode
else:
modes = [mode]
modes = mode if isinstance(mode, list) else [mode]
if app_mode not in modes:
mode_values = {m.value for m in modes}
@@ -1,419 +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
only the operations explicitly enabled by Dify's product registry. The registry
can be validated explicitly against the pinned KnowledgeFS contract during development.
Console auth and contract-specific dataset RBAC run before forwarding. Request
bodies are capped at 64 MiB, JSON and binary responses have separate bounds,
SSE responses remain streaming with a bounded idle read timeout, and only safe
response headers are exposed. Operation-specific upstream error mappings are
applied before Console JSON error handling; the default maps 401 to 502 so it
cannot trigger browser-session recovery and preserves resource-level 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,
HTTPException,
NotFound,
RequestEntityTooLarge,
ServiceUnavailable,
default_exceptions,
)
from configs import dify_config
from controllers.console import api, bp
from controllers.console.wraps import (
account_initialization_required,
cloud_edition_billing_rate_limit_check,
setup_required,
)
from core.helper import ssrf_proxy
from libs.login import current_account_with_tenant, login_required
from services.knowledge_fs_operations import KnowledgeFSMethod
from services.knowledge_fs_proxy import (
KnowledgeFSAccessDeniedError,
KnowledgeFSAuthorization,
KnowledgeFSConfigurationError,
KnowledgeFSRouteNotAllowedError,
KnowledgeFSTimeoutError,
KnowledgeFSTransportError,
KnowledgeFSUpstreamResponse,
authorize_knowledge_fs_request,
get_knowledge_fs_operation,
proxy_authorized_knowledge_fs_request,
proxy_knowledge_fs_request,
)
logger = logging.getLogger(__name__)
type _KnowledgeFSRequestForwarder = Callable[
[str | None, str | None, bytes | None, bytes | None],
KnowledgeFSUpstreamResponse,
]
_MAX_PROXY_BODY_BYTES = 64 * 1024 * 1024
_RESPONSE_HEADER_ALLOWLIST = (
"Cache-Control",
"Content-Disposition",
"Content-Type",
"Retry-After",
"X-Trace-Id",
)
_RESPONSE_HEADER_DENYLIST = frozenset(
{
"authorization",
"connection",
"cookie",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"set-cookie",
"te",
"trailer",
"transfer-encoding",
"upgrade",
}
)
def _console_api_errors[**P](
view: Callable[P, ResponseReturnValue],
) -> Callable[P, ResponseReturnValue]:
"""Route raw Blueprint exceptions through the Console API JSON handlers."""
@wraps(view)
def decorated(*args: P.args, **kwargs: P.kwargs) -> ResponseReturnValue:
try:
return view(*args, **kwargs)
except Exception as exc:
return api.handle_error(exc)
return decorated
def _knowledge_fs_enabled[**P](
view: Callable[P, ResponseReturnValue],
) -> Callable[P, ResponseReturnValue]:
"""Hide the complete KnowledgeFS route surface while the bridge is disabled."""
@wraps(view)
def decorated(*args: P.args, **kwargs: P.kwargs) -> ResponseReturnValue:
if not dify_config.KNOWLEDGE_FS_ENABLED:
raise NotFound()
return view(*args, **kwargs)
return decorated
def _translate_proxy_error(exc: Exception, *, tenant_id: str) -> NoReturn:
"""Map forwarding failures to the stable Console HTTP error surface."""
if isinstance(exc, KnowledgeFSRouteNotAllowedError):
raise NotFound() from exc
if isinstance(exc, KnowledgeFSAccessDeniedError):
raise Forbidden() from exc
if isinstance(exc, KnowledgeFSConfigurationError):
logger.error("KnowledgeFS request was blocked by invalid configuration for tenant_id=%s", tenant_id)
raise ServiceUnavailable("KnowledgeFS integration is misconfigured") from exc
if isinstance(exc, KnowledgeFSTimeoutError):
raise GatewayTimeout("KnowledgeFS request timed out") from exc
if isinstance(exc, KnowledgeFSTransportError):
logger.warning("KnowledgeFS transport request failed for tenant_id=%s", tenant_id)
raise BadGateway("KnowledgeFS is unavailable") from exc
raise exc
def _knowledge_fs_operation_access_required(
view: Callable[[KnowledgeFSAuthorization], ResponseReturnValue],
) -> Callable[[KnowledgeFSMethod, str], ResponseReturnValue]:
"""Authorize one declared operation before billing and request-body work."""
@wraps(view)
def decorated(method: KnowledgeFSMethod, upstream_path: str) -> ResponseReturnValue:
current_user, tenant_id = current_account_with_tenant()
try:
authorization = authorize_knowledge_fs_request(
account=current_user,
tenant_id=tenant_id,
method=method,
path=upstream_path,
)
except KnowledgeFSRouteNotAllowedError as exc:
raise NotFound() from exc
except KnowledgeFSAccessDeniedError as exc:
_translate_proxy_error(exc, tenant_id=tenant_id)
return view(authorization)
return decorated
def _request_body() -> bytes:
"""Read the raw body up to the proxy limit or raise RequestEntityTooLarge."""
body = request.stream.read(_MAX_PROXY_BODY_BYTES + 1)
if len(body) > _MAX_PROXY_BODY_BYTES:
raise RequestEntityTooLarge("KnowledgeFS proxy request body is too large")
return body
def _stream_response_body(
upstream: httpx.Response,
*,
tenant_id: str,
max_response_bytes: int,
) -> Iterator[bytes]:
"""Yield one bounded SSE response and always release its pooled connection."""
total_bytes = 0
try:
for chunk in upstream.iter_bytes():
total_bytes += len(chunk)
if total_bytes > max_response_bytes:
logger.warning("KnowledgeFS stream exceeded the proxy limit for tenant_id=%s", tenant_id)
raise ssrf_proxy.ResponseTooLargeError(f"response exceeded {max_response_bytes} bytes")
yield chunk
finally:
upstream.close()
def _proxy_response(
upstream_result: KnowledgeFSUpstreamResponse,
*,
tenant_id: str,
contract_response_headers: tuple[str, ...],
max_response_bytes: int,
) -> Response:
"""Expose raw content, status, and allowlisted headers from KnowledgeFS.
Raises:
HTTPException: KnowledgeFS returns a status normalized by the operation contract.
"""
upstream = upstream_result.response
mapped_status = dict(upstream_result.operation.error_status_map).get(upstream.status_code)
if mapped_status is not None:
upstream.close()
description = "KnowledgeFS upstream request failed"
if upstream.status_code == HTTPStatus.UNAUTHORIZED:
description = "KnowledgeFS authentication failed"
logger.error(
"KnowledgeFS rejected the Dify server credential with HTTP %s for tenant_id=%s",
upstream.status_code,
tenant_id,
)
exception_type = default_exceptions.get(mapped_status)
if exception_type is None:
exception = HTTPException(description)
exception.code = mapped_status
raise exception
raise exception_type(description)
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_current_request(
*,
method: KnowledgeFSMethod,
tenant_id: str,
forward: _KnowledgeFSRequestForwarder,
) -> Response:
"""Forward the current raw request through one preconfigured service entry."""
if not dify_config.KNOWLEDGE_FS_ENABLED:
raise NotFound()
try:
proxy_result = forward(
request.headers.get("Accept"),
request.content_type,
request.query_string or None,
_request_body() if method != "GET" else None,
)
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,
)
def _proxy_request(
method: KnowledgeFSMethod,
upstream_path: str,
) -> Response:
"""Authorize and forward the current request through the combined service use case."""
if not dify_config.KNOWLEDGE_FS_ENABLED:
raise NotFound()
current_user, tenant_id = current_account_with_tenant()
def forward(
accept: str | None,
content_type: str | None,
query: bytes | None,
body: bytes | None,
) -> KnowledgeFSUpstreamResponse:
return proxy_knowledge_fs_request(
account=current_user,
method=method,
path=upstream_path,
tenant_id=tenant_id,
accept=accept,
content_type=content_type,
query=query,
body=body,
request_headers=request.headers,
)
return _proxy_current_request(method=method, tenant_id=tenant_id, forward=forward)
def _proxy_authorized_request(authorization: KnowledgeFSAuthorization) -> Response:
"""Forward the current request using one previously authorized operation capability.
Args:
authorization: Request-scoped capability produced before billing and body parsing.
Returns:
The filtered response returned by KnowledgeFS.
Raises:
HTTPException: The integration is disabled or forwarding fails.
"""
operation = authorization.operation
tenant_id = authorization.tenant_id
def forward(
accept: str | None,
content_type: str | None,
query: bytes | None,
body: bytes | None,
) -> KnowledgeFSUpstreamResponse:
return proxy_authorized_knowledge_fs_request(
authorization=authorization,
accept=accept,
content_type=content_type,
query=query,
body=body,
request_headers=request.headers,
)
return _proxy_current_request(method=operation.method, tenant_id=tenant_id, forward=forward)
@_knowledge_fs_enabled
@_knowledge_fs_operation_access_required
@cloud_edition_billing_rate_limit_check("knowledge")
def _proxy_knowledge_fs_non_get(
authorization: KnowledgeFSAuthorization,
) -> ResponseReturnValue:
"""Apply knowledge billing checks to one allowlisted non-GET operation."""
return _proxy_authorized_request(authorization)
@bp.route(
"/knowledge-fs/<path:upstream_path>",
methods=["OPTIONS"],
provide_automatic_options=False,
)
@_console_api_errors
@_knowledge_fs_enabled
def proxy_knowledge_fs_options(upstream_path: str) -> ResponseReturnValue:
"""Complete a CORS preflight only for an enabled Console operation."""
requested_method = cast(KnowledgeFSMethod, request.headers.get("Access-Control-Request-Method", "").upper())
try:
get_knowledge_fs_operation(requested_method, upstream_path)
except KnowledgeFSRouteNotAllowedError as exc:
raise NotFound() from exc
return Response(status=HTTPStatus.NO_CONTENT)
@bp.route(
"/knowledge-fs/<path:upstream_path>",
methods=["GET"],
provide_automatic_options=False,
)
@_console_api_errors
@_knowledge_fs_enabled
@setup_required
@login_required
@account_initialization_required
def proxy_knowledge_fs_get(upstream_path: str) -> ResponseReturnValue:
"""Forward one authenticated, dataset-readable GET request.
Args:
upstream_path: Relative KFS path captured after the Console proxy prefix.
Returns:
The filtered raw KnowledgeFS response or a Console JSON error response.
"""
if request.method != "GET":
raise NotFound()
return _proxy_request("GET", upstream_path)
@bp.route(
"/knowledge-fs/<path:upstream_path>",
methods=["DELETE", "PATCH", "POST", "PUT"],
provide_automatic_options=False,
)
@_console_api_errors
@_knowledge_fs_enabled
@setup_required
@login_required
@account_initialization_required
def proxy_knowledge_fs_write(upstream_path: str) -> ResponseReturnValue:
"""Forward one authenticated non-GET request under its contract access policy.
Args:
upstream_path: Relative KFS path captured after the Console proxy prefix.
Returns:
The filtered raw KnowledgeFS response or a Console JSON error response.
"""
method = cast(KnowledgeFSMethod, request.method)
return _proxy_knowledge_fs_non_get(method, upstream_path)
+2 -3
View File
@@ -23,6 +23,7 @@ The transport-generic enum members, ``_CODE_BY_STATUS`` and the
promote them to ``libs`` if a second surface adopts ``ErrorBody``.
"""
import contextlib
import logging
from enum import StrEnum
from typing import Any
@@ -206,10 +207,8 @@ class OpenApiErrorFormatter:
# flask-restx Api.handle_error does `data = getattr(e, "data", default_data)`
# AFTER our handler returns, so a pre-existing e.data (flask_restx.abort,
# BaseHTTPException) would override the canonical body. Rewrite it.
try:
with contextlib.suppress(AttributeError):
e.data = wire # type: ignore[attr-defined]
except AttributeError:
pass
return wire
def _resolve_code(self, e: Exception, status_code: int) -> str:
+1 -4
View File
@@ -128,10 +128,7 @@ class OpenApiWorkflowEventsApi(Resource):
else:
msg_generator = MessageGenerator()
generator: BaseAppGenerator
if app_mode == AppMode.ADVANCED_CHAT:
generator = AdvancedChatAppGenerator()
else:
generator = WorkflowAppGenerator()
generator = AdvancedChatAppGenerator() if app_mode == AppMode.ADVANCED_CHAT else WorkflowAppGenerator()
include_state_snapshot = query.include_state_snapshot
continue_on_pause = query.continue_on_pause
+2 -3
View File
@@ -1,3 +1,4 @@
import contextlib
import json
import logging
from abc import ABC, abstractmethod
@@ -333,10 +334,8 @@ class CotAgentRunner(BaseAgentRunner, ABC):
return answer, ToolInvokeMeta.error_instance(answer)
if isinstance(tool_call_args, str):
try:
with contextlib.suppress(json.JSONDecodeError):
tool_call_args = json.loads(tool_call_args)
except json.JSONDecodeError:
pass
# invoke tool
with use_credit_usage_metadata({"app_type": CreditUsageAppType.AGENT}):
@@ -1,3 +1,4 @@
import contextlib
import json
import re
from typing import Any
@@ -42,11 +43,9 @@ class WorkflowVariablesConfigManager:
"""
json_schema = variable.get("json_schema")
if isinstance(json_schema, str):
try:
# Leave as-is; Pydantic validation will surface the error.
with contextlib.suppress(json.JSONDecodeError, TypeError):
variable["json_schema"] = json.loads(json_schema)
except (json.JSONDecodeError, TypeError):
# Leave as-is; Pydantic validation will surface the error.
pass
@classmethod
def convert_rag_pipeline_variable(cls, workflow: Workflow, start_node_id: str) -> list[RagPipelineVariableEntity]:
+1 -4
View File
@@ -143,10 +143,7 @@ def _parse_signed_file_path(path: str) -> _SignedFileUrl | None:
upload_match = _UPLOAD_FILE_PATH_PATTERN.match(path)
if upload_match:
preview_kind: Literal["file-preview", "image-preview"]
if upload_match.group("preview_kind") == "image-preview":
preview_kind = "image-preview"
else:
preview_kind = "file-preview"
preview_kind = "image-preview" if upload_match.group("preview_kind") == "image-preview" else "file-preview"
return _SignedFileUrl(
file_id=upload_match.group("file_id"),
+2 -8
View File
@@ -75,10 +75,7 @@ def _normalize_completion_params(completion_params: dict[str, object]) -> tuple[
"""
normalized_parameters = dict(completion_params)
stop_value = normalized_parameters.pop("stop", [])
if isinstance(stop_value, list) and all(isinstance(item, str) for item in stop_value):
stop = stop_value
else:
stop = []
stop = stop_value if isinstance(stop_value, list) and all(isinstance(item, str) for item in stop_value) else []
for token_limit_key in ("max_tokens", "max_output_tokens"):
token_limit = normalized_parameters.get(token_limit_key)
@@ -276,10 +273,7 @@ class LLMGenerator:
answer = query
else:
output = result_dict.get("Your Output")
if isinstance(output, str) and output.strip():
answer = output.strip()
else:
answer = query
answer = output.strip() if isinstance(output, str) and output.strip() else query
name = answer.strip()
+2 -4
View File
@@ -2,7 +2,7 @@ import logging
import queue
from collections.abc import Generator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from contextlib import contextmanager, suppress
from typing import Any, final
from urllib.parse import urljoin, urlparse
@@ -310,10 +310,8 @@ def sse_client(
finally:
# Close the SSE connection to unblock the reader thread
if transport.event_source is not None:
try:
with suppress(RuntimeError):
transport.event_source.response.close()
except RuntimeError:
pass
# Clean up queues
if read_queue:
+24 -2
View File
@@ -16,6 +16,10 @@ from core.mcp.types import CallToolResult, Tool
logger = logging.getLogger(__name__)
_MAX_DISCOVERED_TOOLS = 10_000
# Bound round trips independently of inventory size, including empty continuation pages.
_MAX_TOOL_DISCOVERY_REQUESTS = 100
class MCPClient:
def __init__(
@@ -111,11 +115,29 @@ class MCPClient:
self._session.initialize()
def list_tools(self) -> list[Tool]:
"""List available tools from the MCP server"""
"""List the complete tool inventory within tool and request limits, or raise."""
if not self._session:
raise ValueError("Session not initialized.")
response = self._session.list_tools()
return response.tools
requests_sent = 1
tools: list[Tool] = []
seen_cursors: set[str] = set()
while True:
tool_count = len(tools) + len(response.tools)
if tool_count > _MAX_DISCOVERED_TOOLS:
raise MCPConnectionError("MCP tools/list exceeded the 10,000-tool discovery limit.")
if tool_count == _MAX_DISCOVERED_TOOLS and response.nextCursor is not None:
raise MCPConnectionError("MCP tools/list reached the 10,000-tool discovery limit with more pages.")
tools.extend(response.tools)
if response.nextCursor is None:
return tools
if response.nextCursor in seen_cursors:
raise MCPConnectionError("MCP tools/list returned a repeated pagination cursor.")
if requests_sent >= _MAX_TOOL_DISCOVERY_REQUESTS:
raise MCPConnectionError("MCP tools/list reached the 100-request discovery limit with more pages.")
seen_cursors.add(response.nextCursor)
response = self._session.list_tools(cursor=response.nextCursor)
requests_sent += 1
def invoke_tool(self, tool_name: str, tool_args: dict[str, Any]) -> CallToolResult:
"""Call a tool"""
+3 -2
View File
@@ -307,12 +307,13 @@ class ClientSession(
types.CompleteResult,
)
def list_tools(self) -> types.ListToolsResult:
"""Send a tools/list request."""
def list_tools(self, cursor: str | None = None) -> types.ListToolsResult:
"""Send a tools/list request for one page."""
return self.send_request(
types.ClientRequest(
types.ListToolsRequest(
method="tools/list",
params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None,
)
),
types.ListToolsResult,
+1 -4
View File
@@ -217,10 +217,7 @@ class TokenBufferMemory:
# fetch limited messages, and return reversed
stmt = select(Message).where(Message.conversation_id == conversation.id).order_by(Message.created_at.desc())
if message_limit and message_limit > 0:
message_limit = min(message_limit, 500)
else:
message_limit = 500
message_limit = min(message_limit, 500) if message_limit and message_limit > 0 else 500
msg_limit_stmt = stmt.limit(message_limit)
+1 -4
View File
@@ -55,10 +55,7 @@ class OutputModeration(BaseModel):
if not result or not result.flagged:
return completion, False
if result.action == ModerationAction.DIRECT_OUTPUT:
final_output = result.preset_response
else:
final_output = result.text
final_output = result.preset_response if result.action == ModerationAction.DIRECT_OUTPUT else result.text
if public_event:
self.queue_manager.publish(
+1 -4
View File
@@ -187,10 +187,7 @@ class BasePluginClient:
data if isinstance(data, (bytes, str, dict)) or data is None else None
)
if isinstance(data, dict):
if prepared_headers.get("Content-Type") == "application/json":
prepared_data = json.dumps(data)
else:
prepared_data = data
prepared_data = json.dumps(data) if prepared_headers.get("Content-Type") == "application/json" else data
return str(url), prepared_headers, prepared_data, params, files
+4 -8
View File
@@ -460,7 +460,7 @@ def _attach_active_credentials(
def _load_provider_model_cache_entries(tenant_id: str) -> list[_ProviderModelCacheEntry]:
with session_factory.create_session() as session:
stmt = select(ProviderModel).where(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid == True)
stmt = select(ProviderModel).where(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid.is_(True))
provider_models = list(session.scalars(stmt))
_attach_active_credentials(
session=session,
@@ -1002,7 +1002,7 @@ class ProviderManager:
def _get_all_providers(tenant_id: str) -> dict[str, list[Provider]]:
provider_name_to_provider_records_dict = defaultdict(list)
with session_factory.create_session() as session:
stmt = select(Provider).where(Provider.tenant_id == tenant_id, Provider.is_valid == True)
stmt = select(Provider).where(Provider.tenant_id == tenant_id, Provider.is_valid.is_(True))
providers = list(session.scalars(stmt))
_attach_active_credentials(
session=session,
@@ -1661,13 +1661,11 @@ class ProviderManager:
for variable in provider_credential_secret_variables:
if variable in provider_credentials:
try:
with contextlib.suppress(ValueError):
provider_credentials[variable] = encrypter.decrypt_token_with_decoding(
provider_credentials.get(variable, ""),
decoding_context,
)
except ValueError:
pass
current_using_credentials = provider_credentials or {}
@@ -1803,13 +1801,11 @@ class ProviderManager:
for variable in model_credential_secret_variables:
if variable in provider_model_credentials:
try:
with contextlib.suppress(ValueError):
provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding(
provider_model_credentials.get(variable) or "",
decoding_context,
)
except ValueError:
pass
# cache provider model credentials
provider_model_credentials_cache.set(credentials=provider_model_credentials)
+4 -7
View File
@@ -603,7 +603,7 @@ class RetrievalService:
if index_node_ids:
document_segment_stmt = select(DocumentSegment).where(
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
DocumentSegment.status == "completed",
DocumentSegment.index_node_id.in_(index_node_ids),
)
@@ -614,7 +614,7 @@ class RetrievalService:
if segment_ids:
document_segment_stmt = select(DocumentSegment).where(
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
DocumentSegment.status == "completed",
DocumentSegment.id.in_(segment_ids),
)
@@ -627,7 +627,7 @@ class RetrievalService:
if summary_segment_ids:
summary_segment_ids_list = list(summary_segment_ids)
summary_segment_stmt = select(DocumentSegment).where(
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
DocumentSegment.status == "completed",
DocumentSegment.id.in_(summary_segment_ids_list),
)
@@ -671,10 +671,7 @@ class RetrievalService:
child_chunk_details: list[ChildChunkDetail] = []
for child_chunk in child_chunks:
child_document: Document | None = doc_to_document_map.get(child_chunk.index_node_id)
if child_document:
child_score = child_document.metadata.get("score", 0.0)
else:
child_score = 0.0
child_score = child_document.metadata.get("score", 0.0) if child_document else 0.0
child_chunk_detail: ChildChunkDetail = {
"id": child_chunk.id,
"content": child_chunk.content,
+1 -4
View File
@@ -109,10 +109,7 @@ class Blob(BaseModel):
Returns:
Blob instance
"""
if mime_type is None and guess_type:
_mimetype = mimetypes.guess_type(path)[0]
else:
_mimetype = mime_type
_mimetype = mimetypes.guess_type(path)[0] if mime_type is None and guess_type else mime_type
# We do not load the data immediately, instead we treat the blob as a
# reference to the underlying data.
return cls(data=None, mimetype=_mimetype, encoding=encoding, path=path)
+1 -4
View File
@@ -96,10 +96,7 @@ class ExtractProcessor:
if filename_match:
filename = unquote(filename_match.group(1))
match = re.search(r"\.(\w+)$", filename)
if match:
suffix = "." + match.group(1)
else:
suffix = ""
suffix = "." + match.group(1) if match else ""
# https://stackoverflow.com/questions/26541416/generate-temporary-file-names-without-creating-actual-file-in-python#comment90414256_26541521
# Generate a temporary filename under the created temp_dir and ensure the directory exists
file_path = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}" # type: ignore
+1 -4
View File
@@ -127,10 +127,7 @@ class NotionExtractor(BaseExtractor):
# join them all so no part of the value is dropped.
value = "".join(segment.get("plain_text", "") for segment in property_value[type])
elif type in {"select", "status"}:
if property_value[type]:
value = property_value[type]["name"]
else:
value = ""
value = property_value[type]["name"] if property_value[type] else ""
else:
value = property_value[type]
data[property_name] = value
@@ -1,5 +1,6 @@
"""Paragraph index processor."""
import contextlib
import logging
import re
import uuid
@@ -459,11 +460,9 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
# This is especially important for image-only chunks where text is empty or minimal
if is_default_prompt:
language_for_prompt = document_language or "the same language as the input content"
try:
# If default prompt doesn't have {language} placeholder, use it as-is
with contextlib.suppress(KeyError):
summary_prompt = summary_prompt.format(language=language_for_prompt)
except KeyError:
# If default prompt doesn't have {language} placeholder, use it as-is
pass
model_manager = ModelManager.for_tenant(tenant_id=tenant_id)
model_instance = model_manager.get_model_instance(
+8 -14
View File
@@ -470,10 +470,7 @@ class DatasetRetrieval:
planning_strategy = PlanningStrategy.ROUTER
available_datasets = self._get_available_datasets(tenant_id, dataset_ids)
if inputs:
inputs = {key: str(value) for key, value in inputs.items()}
else:
inputs = {}
inputs = {key: str(value) for key, value in inputs.items()} if inputs else {}
available_datasets_ids = [dataset.id for dataset in available_datasets]
metadata_filter_document_ids, metadata_condition = self.get_metadata_filter_condition(
session,
@@ -566,10 +563,7 @@ class DatasetRetrieval:
segment_content = sign_segment_content(segment, session=session)
# If summary exists, prepend it to the content
if record.summary:
final_content = f"{record.summary}\n{segment_content}"
else:
final_content = segment_content
final_content = f"{record.summary}\n{segment_content}" if record.summary else segment_content
document_context_list.append(
DocumentContext(
@@ -601,8 +595,8 @@ class DatasetRetrieval:
document_ids = [record.segment.document_id for record in authorized_records]
dataset_document_stmt = select(DatasetDocument).where(
DatasetDocument.id.in_(document_ids),
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
documents = session.execute(dataset_document_stmt).scalars().all() # type: ignore
dataset_stmt = select(Dataset).where(
@@ -1527,8 +1521,8 @@ class DatasetRetrieval:
document_query = select(DatasetDocument).where(
DatasetDocument.dataset_id.in_(dataset_ids),
DatasetDocument.indexing_status == "completed",
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
filters = [] # type: ignore
metadata_condition = None
@@ -2081,8 +2075,8 @@ class DatasetRetrieval:
select(DocumentModel.dataset_id, func.count(DocumentModel.id).label("available_document_count"))
.where(
DocumentModel.indexing_status == "completed",
DocumentModel.enabled == True,
DocumentModel.archived == False,
DocumentModel.enabled.is_(True),
DocumentModel.archived.is_(False),
DocumentModel.dataset_id.in_(dataset_ids),
)
.group_by(DocumentModel.dataset_id)
+3 -12
View File
@@ -41,10 +41,7 @@ class FixedRecursiveCharacterTextSplitter(EnhanceRecursiveCharacterTextSplitter)
@override
def split_text(self, text: str) -> list[str]:
"""Split incoming text and return chunks."""
if self._fixed_separator:
chunks = text.split(self._fixed_separator)
else:
chunks = [text]
chunks = text.split(self._fixed_separator) if self._fixed_separator else [text]
final_chunks = []
chunks_lengths = self._length_function(chunks)
@@ -74,18 +71,12 @@ class FixedRecursiveCharacterTextSplitter(EnhanceRecursiveCharacterTextSplitter)
# Now that we have the separator, split the text
if separator:
if separator == " ":
splits = re.split(r" +", text)
else:
splits = text.split(separator)
splits = re.split(r" +", text) if separator == " " else text.split(separator)
if self._keep_separator:
splits = [s + separator for s in splits[:-1]] + splits[-1:]
else:
splits = list(text)
if separator == "\n":
splits = [s for s in splits if s != ""]
else:
splits = [s for s in splits if (s not in {"", "\n"})]
splits = [s for s in splits if s != ""] if separator == "\n" else [s for s in splits if s not in {"", "\n"}]
_good_splits = []
_good_splits_lengths = [] # cache the lengths of the splits
_separator = "" if self._keep_separator else separator
+1 -1
View File
@@ -47,7 +47,7 @@ class SummaryIndex:
DocumentSegment.dataset_id == dataset_id,
DocumentSegment.document_id == document_id,
DocumentSegment.status == "completed",
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
)
).all()
segment_ids = [segment.id for segment in segments]
@@ -94,7 +94,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
DocumentSegment.dataset_id.in_(self.dataset_ids),
DocumentSegment.completed_at.isnot(None),
DocumentSegment.status == "completed",
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
DocumentSegment.index_node_id.in_(index_node_ids),
)
segments = session.scalars(document_segment_stmt).all()
@@ -117,8 +117,8 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
dataset = session.get(Dataset, segment.dataset_id)
document_stmt = select(Document).where(
Document.id == segment.document_id,
Document.enabled == True,
Document.archived == False,
Document.enabled.is_(True),
Document.archived.is_(False),
)
document = session.scalar(document_stmt)
if dataset and document:
@@ -78,10 +78,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
self.retrieve_config.metadata_filtering_conditions,
self.inputs,
)
if metadata_filter_document_ids:
document_ids_filter = metadata_filter_document_ids.get(dataset.id, [])
else:
document_ids_filter = None
document_ids_filter = metadata_filter_document_ids.get(dataset.id, []) if metadata_filter_document_ids else None
if dataset.provider == "external":
results: list[RetrievalDocument] = []
external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval(
@@ -183,10 +180,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
segment_content = sign_segment_content(segment, session=session)
# If summary exists, prepend it to the content
if record.summary:
final_content = f"{record.summary}\n{segment_content}"
else:
final_content = segment_content
final_content = f"{record.summary}\n{segment_content}" if record.summary else segment_content
document_context_list.append(
DocumentContext(
@@ -201,8 +195,8 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
dataset = session.get(Dataset, segment.dataset_id)
dataset_document_stmt = select(DatasetDocument).where(
DatasetDocument.id == segment.document_id,
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
document = session.scalar(dataset_document_stmt)
if dataset and document:
+1 -4
View File
@@ -70,10 +70,7 @@ def get_url(url: str, user_agent: str | None = None) -> str:
# Detect encoding using charset_normalizer
detected_encoding = charset_normalizer.from_bytes(response.content).best()
if detected_encoding:
encoding = detected_encoding.encoding
else:
encoding = "utf-8"
encoding = detected_encoding.encoding if detected_encoding else "utf-8"
if encoding:
try:
content = response.content.decode(encoding)
+2 -3
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import contextlib
import json
import logging
from collections.abc import Generator, Mapping, Sequence
@@ -170,10 +171,8 @@ class WorkflowTool(Tool):
usage_metadata: dict[str, Any] = {}
if total_tokens is not None:
try:
with contextlib.suppress(TypeError, ValueError):
usage_metadata["total_tokens"] = int(str(total_tokens))
except (TypeError, ValueError):
pass
if total_price is not None:
usage_metadata["total_price"] = str(total_price)
currency = data.get("currency")
-385
View File
@@ -1,385 +0,0 @@
"""Validate Dify Console KnowledgeFS declarations against a pinned OpenAPI document.
The OpenAPI document is exported only during explicit development validation. Runtime declarations live with Dify
product policy; this module validates their transport metadata without generating a complete operation catalog.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import subprocess
import sys
import tempfile
from copy import deepcopy
from pathlib import Path
from typing import Any, Literal, TypedDict
API_ROOT = Path(__file__).resolve().parents[1]
if str(API_ROOT) not in sys.path:
sys.path.insert(0, str(API_ROOT))
WORKSPACE_ROOT = API_ROOT.parent
LOCK_PATH = API_ROOT / "knowledge-fs-contract.lock.json"
DEFAULT_REPOSITORY = WORKSPACE_ROOT.parent / "knowledge-fs"
OPENAPI_METHODS = ("delete", "get", "head", "options", "patch", "post", "put", "trace")
PROXY_METHODS = frozenset({"delete", "get", "patch", "post", "put"})
CONSOLE_PROXY_ERROR_SCHEMA_NAME = "ConsoleProxyError"
CONSOLE_PROXY_ERROR_SCHEMA: dict[str, Any] = {
"type": "object",
"required": ["code", "message", "status"],
"properties": {
"code": {"type": "string"},
"message": {"type": "string"},
"status": {"type": "integer"},
},
}
class ContractDeclaration(TypedDict):
"""KnowledgeFS transport contract declared by one Dify Console registry entry."""
operation_id: str
method: str
path: str
required_scope: str | None
response_kind: str
max_response_bytes: int
request_headers: tuple[str, ...]
response_headers: tuple[str, ...]
response_media_types: tuple[str, ...]
error_status_map: tuple[tuple[int, int], ...]
type DeclarationField = Literal[
"method",
"path",
"required_scope",
"response_kind",
"max_response_bytes",
"request_headers",
"response_headers",
"response_media_types",
]
DECLARATION_FIELDS: tuple[DeclarationField, ...] = (
"method",
"path",
"required_scope",
"response_kind",
"max_response_bytes",
"request_headers",
"response_headers",
"response_media_types",
)
def main() -> None:
"""Update or verify the pin and validate Console declarations against its OpenAPI document."""
parser = argparse.ArgumentParser()
mode = parser.add_mutually_exclusive_group()
mode.add_argument("--check", action="store_true")
mode.add_argument("--update-lock", action="store_true")
parser.add_argument("--repository", type=Path, default=DEFAULT_REPOSITORY)
parser.add_argument("--output-openapi", type=Path)
args = parser.parse_args()
repository = args.repository.resolve()
lock = json.loads(LOCK_PATH.read_text())
tracked_changes = run("git", "status", "--porcelain", "--untracked-files=no", cwd=repository).strip()
if tracked_changes:
raise RuntimeError("KnowledgeFS checkout must not contain tracked changes during contract export")
commit = run("git", "rev-parse", "HEAD", cwd=repository).strip()
if not args.update_lock and commit != lock["commit"]:
raise RuntimeError(
f"KnowledgeFS checkout mismatch: expected {lock['commit']}, received {commit}. "
"Use the pinned commit or pass --update-lock intentionally."
)
with tempfile.TemporaryDirectory(prefix="dify-knowledge-fs-contract-") as directory:
openapi_path = Path(directory) / "knowledge-fs.openapi.json"
subprocess.run(
["pnpm", "openapi:export", "--", "--output", str(openapi_path)],
cwd=repository,
check=True,
)
openapi_content = openapi_path.read_bytes()
openapi_sha256 = sha256(openapi_content)
if not args.update_lock and openapi_sha256 != lock["openapiSha256"]:
raise RuntimeError(
f"KnowledgeFS OpenAPI hash mismatch: expected {lock['openapiSha256']}, received {openapi_sha256}"
)
document: dict[str, Any] = json.loads(openapi_content)
declarations = console_contract_declarations()
validate_declarations(document, declarations)
if args.output_openapi:
filtered_document = filter_openapi_document(document, declarations)
filtered_document["x-dify-source-openapi-sha256"] = openapi_sha256
filtered_document["x-dify-console-declarations-sha256"] = contract_declarations_sha256(declarations)
args.output_openapi.parent.mkdir(parents=True, exist_ok=True)
args.output_openapi.write_text(json.dumps(filtered_document, indent=2) + "\n")
if args.update_lock:
LOCK_PATH.write_text(
json.dumps(
{
"commit": commit,
"openapiSha256": openapi_sha256,
"repository": lock["repository"],
},
indent=2,
)
+ "\n"
)
def validate_declarations(document: dict[str, Any], declarations: tuple[ContractDeclaration, ...]) -> None:
"""Validate Dify Console declarations against matching pinned OpenAPI operations."""
operations_by_id: dict[str, list[tuple[str, str, dict[str, Any], dict[str, Any]]]] = {}
for path, path_item in document.get("paths", {}).items():
for method in OPENAPI_METHODS:
operation = path_item.get(method)
if operation is None:
continue
operation_id = operation.get("operationId")
if isinstance(operation_id, str) and operation_id:
operations_by_id.setdefault(operation_id, []).append((method, path, path_item, operation))
declared_ids: set[str] = set()
for declaration in declarations:
operation_id = declaration["operation_id"]
if operation_id in declared_ids:
raise ValueError(f"Dify Console registry has duplicate operationId: {operation_id}")
declared_ids.add(operation_id)
matches = operations_by_id.get(operation_id, [])
if not matches:
raise ValueError(f"KnowledgeFS OpenAPI has no operationId: {operation_id}")
if len(matches) > 1:
raise ValueError(f"KnowledgeFS OpenAPI has duplicate operationId: {operation_id}")
method, path, path_item, operation = matches[0]
if not path.startswith("/"):
raise ValueError(f"KnowledgeFS OpenAPI path must be absolute: {path}")
if method not in PROXY_METHODS:
raise ValueError(f"KnowledgeFS proxy does not support {method.upper()} {path}")
expected: dict[DeclarationField, object] = {
"method": method.upper(),
"path": path[1:],
"required_scope": required_scope(operation),
"response_kind": response_kind(operation),
"max_response_bytes": required_max_response_bytes(operation),
"request_headers": request_header_names(path_item, operation),
"response_headers": response_header_names(operation),
"response_media_types": response_media_types(operation),
}
for field in DECLARATION_FIELDS:
expected_value = expected[field]
received_value = declaration[field]
if received_value != expected_value:
raise ValueError(
f"KnowledgeFS operation {operation_id} field {field} drifted: "
f"expected {expected_value!r}, received {received_value!r}"
)
validate_error_status_map(operation_id, declaration["error_status_map"])
def filter_openapi_document(
document: dict[str, Any],
declarations: tuple[ContractDeclaration, ...],
) -> dict[str, Any]:
"""Return a code-generation document containing only Console-allowlisted operations."""
filtered_document: dict[str, Any] = {
key: value for key, value in document.items() if key not in {"components", "paths"}
}
source_paths = document.get("paths", {})
filtered_paths: dict[str, Any] = {}
for declaration in declarations:
path = f"/{declaration['path']}"
method = declaration["method"].lower()
source_path_item = source_paths[path]
path_metadata = {key: value for key, value in source_path_item.items() if key not in OPENAPI_METHODS}
filtered_path_item = filtered_paths.setdefault(path, path_metadata)
filtered_operation = deepcopy(source_path_item[method])
_rewrite_proxy_error_responses(filtered_operation, declaration["error_status_map"])
filtered_path_item[method] = filtered_operation
filtered_document["paths"] = filtered_paths
source_components = document.get("components", {})
filtered_components = {key: value for key, value in source_components.items() if key != "schemas"}
source_schemas = source_components.get("schemas", {})
available_schemas = {**source_schemas, CONSOLE_PROXY_ERROR_SCHEMA_NAME: CONSOLE_PROXY_ERROR_SCHEMA}
schema_names = _referenced_schema_names(filtered_paths, available_schemas)
filtered_components["schemas"] = {
name: schema for name, schema in available_schemas.items() if name in schema_names
}
filtered_document["components"] = filtered_components
return filtered_document
def validate_error_status_map(operation_id: str, error_status_map: tuple[tuple[int, int], ...]) -> None:
"""Validate the status normalization advertised by one Console operation."""
upstream_statuses: set[int] = set()
for upstream_status, console_status in error_status_map:
if upstream_status in upstream_statuses:
raise ValueError(f"KnowledgeFS operation {operation_id} has duplicate error status: {upstream_status}")
if not 400 <= upstream_status <= 599 or not 400 <= console_status <= 599:
raise ValueError(f"KnowledgeFS operation {operation_id} has invalid error status mapping")
upstream_statuses.add(upstream_status)
def _rewrite_proxy_error_responses(
operation: dict[str, Any],
error_status_map: tuple[tuple[int, int], ...],
) -> None:
responses = operation.setdefault("responses", {})
proxy_error_response = {
"description": "Error normalized by the Dify Console KnowledgeFS proxy.",
"content": {
"application/json": {"schema": {"$ref": f"#/components/schemas/{CONSOLE_PROXY_ERROR_SCHEMA_NAME}"}}
},
}
for upstream_status, console_status in error_status_map:
existing_target = responses.get(str(console_status)) if upstream_status != console_status else None
responses.pop(str(upstream_status), None)
normalized_response: dict[str, Any] = deepcopy(proxy_error_response)
existing_schema = (
existing_target.get("content", {}).get("application/json", {}).get("schema")
if isinstance(existing_target, dict)
else None
)
if existing_schema is not None:
normalized_response["content"]["application/json"]["schema"] = {
"oneOf": [
deepcopy(existing_schema),
{"$ref": f"#/components/schemas/{CONSOLE_PROXY_ERROR_SCHEMA_NAME}"},
]
}
responses[str(console_status)] = normalized_response
def _referenced_schema_names(value: Any, schemas: dict[str, Any]) -> set[str]:
reference_prefix = "#/components/schemas/"
selected: set[str] = set()
pending: list[Any] = [value]
while pending:
current = pending.pop()
if isinstance(current, list):
pending.extend(current)
continue
if not isinstance(current, dict):
continue
reference = current.get("$ref")
if isinstance(reference, str) and reference.startswith(reference_prefix):
name = reference.removeprefix(reference_prefix)
if name not in selected:
if name not in schemas:
raise ValueError(f"KnowledgeFS OpenAPI references missing schema: {name}")
selected.add(name)
pending.append(schemas[name])
pending.extend(current.values())
return selected
def console_contract_declarations() -> tuple[ContractDeclaration, ...]:
"""Return transport declarations from the runtime Console operation registry."""
from services.knowledge_fs_operations import KNOWLEDGE_FS_CONSOLE_OPERATIONS
return tuple(
{
"operation_id": operation.operation_id,
"method": operation.method,
"path": operation.path,
"required_scope": operation.required_scope,
"response_kind": operation.response_kind,
"max_response_bytes": operation.max_response_bytes,
"request_headers": operation.request_headers,
"response_headers": operation.response_headers,
"response_media_types": operation.response_media_types,
"error_status_map": operation.error_status_map,
}
for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS
)
def contract_declarations_sha256(declarations: tuple[ContractDeclaration, ...]) -> str:
"""Return a stable digest for the runtime Console operation declarations."""
content = json.dumps(declarations, separators=(",", ":"), sort_keys=True).encode()
return sha256(content)
def response_kind(operation: dict[str, Any]) -> str:
media_types = response_media_types(operation)
if "text/event-stream" in media_types:
return "stream"
if "application/octet-stream" in media_types:
return "binary"
return "buffered"
def response_media_types(operation: dict[str, Any]) -> tuple[str, ...]:
media_types: set[str] = set()
for status, response in operation.get("responses", {}).items():
if status == "2XX" or (len(status) == 3 and status.startswith("2") and status.isdigit()):
media_types.update(response.get("content", {}))
return tuple(sorted(media_types))
def required_scope(operation: dict[str, Any]) -> str | None:
scope = operation.get("x-knowledge-fs-required-scope")
if scope in ("knowledge-spaces:read", "knowledge-spaces:write"):
return scope
if operation.get("security") == []:
return None
raise ValueError(f"KnowledgeFS operation has no supported required scope: {scope}")
def required_max_response_bytes(operation: dict[str, Any]) -> int:
value = operation.get("x-knowledge-fs-max-response-bytes")
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
raise ValueError(f"KnowledgeFS operation has no valid response byte limit: {value}")
return value
def request_header_names(path_item: dict[str, Any], operation: dict[str, Any]) -> tuple[str, ...]:
names: set[str] = set()
for parameter in [*path_item.get("parameters", []), *operation.get("parameters", [])]:
if "$ref" in parameter:
raise ValueError(f"KnowledgeFS request header references are not supported: {parameter['$ref']}")
if parameter.get("in") == "header":
names.add(parameter["name"].lower())
return tuple(sorted(names))
def response_header_names(operation: dict[str, Any]) -> tuple[str, ...]:
return tuple(
sorted(
{
name.lower()
for response in operation.get("responses", {}).values()
for name in response.get("headers", {})
}
)
)
def sha256(content: bytes) -> str:
return hashlib.sha256(content).hexdigest()
def run(*command: str, cwd: Path) -> str:
return subprocess.run(command, cwd=cwd, check=True, capture_output=True, text=True).stdout
if __name__ == "__main__":
main()
+1 -4
View File
@@ -61,10 +61,7 @@ def _parse_otlp_headers(raw: str) -> dict[str, str]:
def _datetime_to_ns(dt: datetime) -> int:
"""Convert a datetime to nanoseconds since epoch (OTEL convention)."""
# Ensure we always interpret naive datetimes as UTC instead of local time.
if dt.tzinfo is None:
dt = dt.replace(tzinfo=UTC)
else:
dt = dt.astimezone(UTC)
dt = dt.replace(tzinfo=UTC) if dt.tzinfo is None else dt.astimezone(UTC)
return int(dt.timestamp() * 1_000_000_000)
+16 -14
View File
@@ -196,22 +196,23 @@ def build_account_services(
sessions=RedisAccountSessionGateway(redis=redis),
now=naive_utc_now,
)
activation = AccountActivationService(
tokens=invitation_tokens,
accounts=activation_accounts,
workspace_policy=DeploymentWorkspaceInvitePolicy(),
eligibility=BillingAccountActivationEligibility(
enabled=deployment_edition == DeploymentEdition.CLOUD,
),
membership_cache=BillingWorkspaceMembershipCache(
enabled=deployment_edition == DeploymentEdition.CLOUD,
),
member_access_sync=RBACWorkspaceMemberAccessSync(
enabled=dify_config.RBAC_ENABLED,
),
)
return AccountServices(
lifecycle=lifecycle,
activation=AccountActivationService(
tokens=invitation_tokens,
accounts=activation_accounts,
workspace_policy=DeploymentWorkspaceInvitePolicy(),
eligibility=BillingAccountActivationEligibility(
enabled=deployment_edition == DeploymentEdition.CLOUD,
),
membership_cache=BillingWorkspaceMembershipCache(
enabled=deployment_edition == DeploymentEdition.CLOUD,
),
member_access_sync=RBACWorkspaceMemberAccessSync(
enabled=dify_config.RBAC_ENABLED,
),
),
activation=activation,
identity=AccountIdentityGateway(accounts=accounts, redis=redis),
access=AccountAccessService(
accounts=accounts,
@@ -226,6 +227,7 @@ def build_account_services(
invitations=AccountActivationConsoleAuthInvitationGateway(
tokens=invitation_tokens,
accounts=activation_accounts,
activation=activation,
),
policies=DeploymentConsoleAuthPolicyGateway(
billing_enabled=deployment_edition == DeploymentEdition.CLOUD,
+3
View File
@@ -149,6 +149,9 @@ class RedisClientWrapper:
def ttl(self, name: str | bytes) -> Any:
return self._require_client().ttl(_serialize_redis_name_arg(name, self._get_prefix()))
def pttl(self, name: str | bytes) -> Any:
return self._require_client().pttl(_serialize_redis_name_arg(name, self._get_prefix()))
def getdel(self, name: str | bytes) -> Any:
return self._require_client().getdel(_serialize_redis_name_arg(name, self._get_prefix()))
+3 -3
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import contextlib
import logging
import os
import socket
@@ -247,10 +248,9 @@ class AliyunLogStore:
finally:
# Ensure socket is properly closed
if sock:
try:
# Ignore errors during cleanup
with contextlib.suppress(Exception):
sock.close()
except Exception: # noqa: S110
pass # Ignore errors during cleanup
@property
def supports_pg_protocol(self) -> bool:
@@ -220,10 +220,7 @@ class LogstoreAPIWorkflowNodeExecutionRepository(DifyAPIWorkflowNodeExecutionRep
# For each id, select the row with max log_version
deduplicated_results = []
for rows in id_to_results.values():
if len(rows) > 1:
max_row = max(rows, key=lambda x: int(x.get("log_version", 0)))
else:
max_row = rows[0]
max_row = max(rows, key=lambda x: int(x.get("log_version", 0))) if len(rows) > 1 else rows[0]
deduplicated_results.append(max_row)
else:
# For PG mode, results are already deduplicated by the SQL query
@@ -325,10 +322,7 @@ class LogstoreAPIWorkflowNodeExecutionRepository(DifyAPIWorkflowNodeExecutionRep
# For each id, select the row with max log_version
for rows in id_to_results.values():
if len(rows) > 1:
max_row = max(rows, key=lambda x: int(x.get("log_version", 0)))
else:
max_row = rows[0]
max_row = max(rows, key=lambda x: int(x.get("log_version", 0))) if len(rows) > 1 else rows[0]
model = _dict_to_workflow_node_execution_model(max_row)
if model and model.id: # Ensure model is valid
+2 -3
View File
@@ -10,6 +10,7 @@ by opentelemetry.instrumentation.sqlcommenter_utils._add_framework_tags() when t
SQLAlchemy instrumentor appends comments to SQL statements.
"""
import contextlib
import logging
from typing import Any, TypedDict
@@ -105,10 +106,8 @@ def _on_task_postrun(*args: object, **kwargs: object) -> None:
except Exception:
logger.debug("Failed to detach SQL commenter context", exc_info=True)
finally:
try:
with contextlib.suppress(AttributeError):
delattr(task, _TOKEN_ATTR)
except AttributeError:
pass
def setup_celery_sqlcommenter() -> None:
@@ -191,11 +191,8 @@ class ClickZettaVolumeStorage(BaseStorage):
# Only use TABLE VOLUME for actual dataset-specific paths
# like "dataset_12345/file.pdf" or paths with dataset_ prefix
if dataset_id:
table_name = f"{self._config.table_prefix}{dataset_id}"
else:
# Default table name for generic operations
table_name = "default_dataset"
# Default table name for generic operations
table_name = f"{self._config.table_prefix}{dataset_id}" if dataset_id else "default_dataset"
return f"TABLE VOLUME {table_name}"
elif self._config.volume_type == "external":
return f"VOLUME {self._config.volume_name}"
@@ -517,10 +514,7 @@ class ClickZettaVolumeStorage(BaseStorage):
else:
sql = f"LIST {volume_prefix} SUBDIRECTORY '{self._config.dify_prefix}'"
else:
if path:
sql = f"LIST {volume_prefix} SUBDIRECTORY '{path}'"
else:
sql = f"LIST {volume_prefix}"
sql = f"LIST {volume_prefix} SUBDIRECTORY '{path}'" if path else f"LIST {volume_prefix}"
rows = self._execute_sql(sql, fetch=True)
-5
View File
@@ -1,5 +0,0 @@
{
"commit": "a0f50470612cc0b3656f89e4f2435aaf412e6e3b",
"openapiSha256": "f18910e9c45a64f0855e0643a7a626fb2889021b4f943458de86c6bd2469facb",
"repository": "https://github.com/langgenius/knowledge-fs"
}
@@ -1,3 +1,4 @@
import contextlib
import logging
import queue
import threading
@@ -246,14 +247,10 @@ class RedisSubscriptionBase(Subscription):
try:
self._queue.put_nowait(SIG_CLOSE)
except queue.Full:
try:
with contextlib.suppress(queue.Empty):
self._queue.get_nowait()
except queue.Empty:
pass
try:
with contextlib.suppress(queue.Full):
self._queue.put_nowait(SIG_CLOSE)
except queue.Full:
pass
# Abstract methods to be implemented by subclasses
def _get_subscription_type(self) -> str:
@@ -202,10 +202,7 @@ class _StreamsSubscription(Subscription):
self._start_if_needed()
try:
if timeout is None:
item = self._queue.get()
else:
item = self._queue.get(timeout=timeout)
item = self._queue.get() if timeout is None else self._queue.get(timeout=timeout)
except queue.Empty:
return None
+3 -4
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import contextlib
import json
import re
import uuid
@@ -1368,11 +1369,9 @@ class Conversation(Base):
if not workflow_run:
continue
try:
# Handle invalid status values gracefully
with contextlib.suppress(ValueError, KeyError):
status_counts[WorkflowExecutionStatus(workflow_run.status)] += 1
except (ValueError, KeyError):
# Handle invalid status values gracefully
pass
return {
"success": status_counts[WorkflowExecutionStatus.SUCCEEDED],
-1
View File
@@ -26683,7 +26683,6 @@ Non-sensitive bootstrap snapshot exposed before Console or Web authentication.
| enable_step_by_step_tour | boolean | | Yes |
| is_allow_register | boolean | | Yes |
| is_email_setup | boolean | | Yes |
| knowledge_fs_enabled | boolean | | Yes |
| license | [LicenseStatusModel](#licensestatusmodel) | | Yes |
| plugin_installation_permission | [PluginInstallationPermissionModel](#plugininstallationpermissionmodel) | | Yes |
| rbac_enabled | boolean | | Yes |
-1
View File
@@ -1957,7 +1957,6 @@ Non-sensitive bootstrap snapshot exposed before Console or Web authentication.
| enable_step_by_step_tour | boolean | | Yes |
| is_allow_register | boolean | | Yes |
| is_email_setup | boolean | | Yes |
| knowledge_fs_enabled | boolean | | Yes |
| license | [LicenseStatusModel](#licensestatusmodel) | | Yes |
| plugin_installation_permission | [PluginInstallationPermissionModel](#plugininstallationpermissionmodel) | | Yes |
| rbac_enabled | boolean | | Yes |
@@ -557,7 +557,7 @@ def test_record_methods_handle_exceptions(
def test_metrics_initializes_grpc_metric_exporter() -> None:
client = _build_client()
metric_reader = metric_reader_instances[-1]
exporter = cast(DummyGrpcMetricExporter, metric_reader.exporter)
exporter = metric_reader.exporter
assert isinstance(exporter, DummyGrpcMetricExporter)
assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000
@@ -570,7 +570,7 @@ def test_metrics_initializes_http_protobuf_metric_exporter(monkeypatch: pytest.M
monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf")
client = _build_client()
metric_reader = metric_reader_instances[-1]
exporter = cast(DummyHttpMetricExporter, metric_reader.exporter)
exporter = metric_reader.exporter
assert isinstance(exporter, DummyHttpMetricExporter)
assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000
@@ -582,7 +582,7 @@ def test_metrics_initializes_http_json_metric_exporter(monkeypatch: pytest.Monke
monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json")
client = _build_client()
metric_reader = metric_reader_instances[-1]
exporter = cast(DummyJsonMetricExporter, metric_reader.exporter)
exporter = metric_reader.exporter
assert isinstance(exporter, DummyJsonMetricExporter)
assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000
@@ -597,7 +597,7 @@ def test_metrics_http_json_metric_exporter_falls_back_without_temporality(monkey
monkeypatch.setattr(exporter_module, "OTLPMetricExporter", DummyJsonMetricExporterNoTemporality)
_ = _build_client()
metric_reader = metric_reader_instances[-1]
exporter = cast(DummyJsonMetricExporterNoTemporality, metric_reader.exporter)
exporter = metric_reader.exporter
assert isinstance(exporter, DummyJsonMetricExporterNoTemporality)
assert "preferred_temporality" not in exporter.kwargs
@@ -262,13 +262,10 @@ class AlibabaCloudMySQLVector(BaseVector):
try:
distance = float(record["distance"])
# Convert distance to similarity score
if self.distance_function == "cosine":
# For cosine distance: similarity = 1 - distance
similarity = 1.0 - distance
else:
# For euclidean distance: use inverse relationship
# similarity = 1 / (1 + distance)
similarity = 1.0 / (1.0 + distance)
# For cosine distance: similarity = 1 - distance
# For euclidean distance: use inverse relationship
# similarity = 1 / (1 + distance)
similarity = 1.0 - distance if self.distance_function == "cosine" else 1.0 / (1.0 + distance)
metadata = parse_metadata_json(record["meta"])
metadata["score"] = similarity
@@ -10,7 +10,7 @@ import json
import logging
import threading
import uuid
from contextlib import contextmanager
from contextlib import contextmanager, suppress
from typing import TYPE_CHECKING, Any, override
from sqlalchemy.orm import Session
@@ -104,10 +104,8 @@ class IrisConnectionPool:
is_valid = True
except (OSError, RuntimeError) as e:
logger.debug("Connection validation failed: %s", e)
try:
with suppress(OSError, RuntimeError):
conn.close()
except (OSError, RuntimeError):
pass
with self._lock:
self._pool.append(conn if is_valid else self._create_connection())
@@ -168,10 +166,8 @@ class IrisConnectionPool:
"""Close all connections (application shutdown only)."""
with self._lock:
for conn in self._pool:
try:
with suppress(OSError, RuntimeError):
conn.close()
except (OSError, RuntimeError):
pass
self._pool.clear()
self._in_use = 0
self._schemas_initialized.clear()
@@ -457,7 +457,7 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory):
else:
idle_tidb_auth_binding = db.session.scalar(
select(TidbAuthBinding)
.where(TidbAuthBinding.active == False, TidbAuthBinding.status == "ACTIVE")
.where(TidbAuthBinding.active.is_(False), TidbAuthBinding.status == "ACTIVE")
.limit(1)
)
if idle_tidb_auth_binding:
@@ -531,10 +531,7 @@ class WeaviateVector(BaseVector):
for obj in res.objects:
properties = dict(obj.properties or {})
text = properties.pop(Field.TEXT_KEY.value, "")
if obj.metadata and obj.metadata.distance is not None:
distance = obj.metadata.distance
else:
distance = 1.0
distance = obj.metadata.distance if obj.metadata and obj.metadata.distance is not None else 1.0
score = 1.0 - distance
if score > score_threshold:
@@ -58,6 +58,7 @@ class SQLAlchemyAccountActivationRepository(AccountActivationRepository):
*,
role: str,
setup: AccountSetup | None,
create_membership: bool = True,
) -> ActivationPersistenceResult | None:
with self._session_factory.begin() as session:
tenant_id = session.scalar(
@@ -83,6 +84,8 @@ class SQLAlchemyAccountActivationRepository(AccountActivationRepository):
TenantAccountJoin.account_id == account.id,
)
)
if membership is None and not create_membership:
return None
membership_created = membership is None
if membership is None:
membership = TenantAccountJoin(
+2 -5
View File
@@ -378,7 +378,7 @@ class ConsoleAppRepository(ConsoleApps, DiscoveryApps, RelatedApps):
def _build_app_list_filters(
user_id: str, tenant_id: str, params: AppListBaseParams, session: Session
) -> list[sa.ColumnElement[bool]]:
filters = [App.tenant_id == tenant_id, App.is_universal == False]
filters = [App.tenant_id == tenant_id, App.is_universal.is_(False)]
if params.mode == "workflow":
filters.append(App.mode == AppMode.WORKFLOW)
@@ -794,10 +794,7 @@ class ConsoleAppRepository(ConsoleApps, DiscoveryApps, RelatedApps):
app.name = args["name"]
app.description = args["description"]
icon_type = args.get("icon_type")
if icon_type is None:
resolved_icon_type = app.icon_type
else:
resolved_icon_type = IconType(icon_type)
resolved_icon_type = app.icon_type if icon_type is None else IconType(icon_type)
app.icon_type = resolved_icon_type
app.icon = args["icon"]
@@ -3,7 +3,7 @@
from typing import cast, override
from pydantic import JsonValue
from sqlalchemy import select
from sqlalchemy import and_, or_, select
from sqlalchemy.orm import Session, sessionmaker
from models.enums import ConversationFromSource, FeedbackFromSource, FeedbackRating
@@ -65,12 +65,19 @@ class SQLAlchemyInstalledAppMessageRepository(InstalledAppMessageStore):
f"The first_id cursor {first_id} does not belong to conversation {conversation_id}."
)
statement = statement.where(
Message.created_at < first_message.created_at, Message.id != first_message.id
or_(
Message.created_at < first_message.created_at,
and_(
Message.created_at == first_message.created_at,
Message.id < first_message.id,
),
)
)
messages = list(session.scalars(statement.order_by(Message.created_at.desc()).limit(limit + 1)).all())
messages = list(
session.scalars(statement.order_by(Message.created_at.desc(), Message.id.desc()).limit(limit + 1)).all()
)
has_more = len(messages) > limit
# Preserve ascending presentation of each newest-first page and the
# existing strict timestamp cursor, including tied timestamps.
# Preserve ascending presentation of each newest-first page.
messages = messages[:limit]
messages.reverse()
return MessagePage(
@@ -322,8 +322,8 @@ def get_dataset_available_document_count(dataset: Dataset, *, session: Session)
Document.dataset_id == dataset.id,
Document.tenant_id == dataset.tenant_id,
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
Document.enabled.is_(True),
Document.archived.is_(False),
)
)
or 0
@@ -337,7 +337,7 @@ def get_dataset_available_segment_count(dataset: Dataset, *, session: Session) -
DocumentSegment.dataset_id == dataset.id,
DocumentSegment.tenant_id == dataset.tenant_id,
DocumentSegment.status == "completed",
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
)
)
or 0
@@ -629,8 +629,8 @@ def build_dataset_detail_prefetch(datasets: Sequence[Dataset], *, session: Sessi
(
sa.and_(
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
Document.enabled.is_(True),
Document.archived.is_(False),
),
1,
),
+10 -3
View File
@@ -1,6 +1,6 @@
from typing import cast, override
from sqlalchemy import select
from sqlalchemy import and_, or_, select
from sqlalchemy.orm import Session, sessionmaker
from models.enums import CreatorUserRole
@@ -50,9 +50,16 @@ class SQLAlchemySavedMessageRepository(SavedMessageStore):
last_message = session.scalar(stmt.where(Message.id == last_id).limit(1))
if last_message is None:
raise LastMessageNotExistsError("The last_id cursor does not belong to the current saved messages.")
stmt = stmt.where(Message.created_at < last_message.created_at, Message.id != last_message.id)
stmt = stmt.where(
or_(
Message.created_at < last_message.created_at,
and_(Message.created_at == last_message.created_at, Message.id < last_message.id),
)
)
messages = list(session.scalars(stmt.order_by(Message.created_at.desc()).limit(fetch_limit)).all())
messages = list(
session.scalars(stmt.order_by(Message.created_at.desc(), Message.id.desc()).limit(fetch_limit)).all()
)
has_more = len(messages) > limit
if has_more:
messages = messages[:limit]
@@ -125,10 +125,8 @@ class SQLAlchemyExecutionExtraContentRepository(ExecutionExtraContentRepository)
submitted = form.submitted_at is not None or form.status == HumanInputFormStatus.SUBMITTED
# Only a waiting form can still be answered; a timed-out or expired one gets no token.
if not submitted and form.status == HumanInputFormStatus.WAITING:
form_token = self._resolve_form_token(recipients_by_form_id.get(form.id, []))
else:
form_token = None
answerable = not submitted and form.status == HumanInputFormStatus.WAITING
form_token = self._resolve_form_token(recipients_by_form_id.get(form.id, [])) if answerable else None
form_definition_domain_model = HumanInputFormDefinition(
form_id=form.id,
node_id=form.node_id,
+6 -6
View File
@@ -56,8 +56,8 @@ def clean_unused_datasets_task():
select(Document.dataset_id, func.count(Document.id).label("document_count"))
.where(
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
Document.enabled.is_(True),
Document.archived.is_(False),
Document.updated_at > clean_day,
)
.group_by(Document.dataset_id)
@@ -69,8 +69,8 @@ def clean_unused_datasets_task():
select(Document.dataset_id, func.count(Document.id).label("document_count"))
.where(
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
Document.enabled.is_(True),
Document.archived.is_(False),
Document.updated_at < clean_day,
)
.group_by(Document.dataset_id)
@@ -127,8 +127,8 @@ def clean_unused_datasets_task():
documents = session.scalars(
select(Document).where(
Document.dataset_id == dataset.id,
Document.enabled == True,
Document.archived == False,
Document.enabled.is_(True),
Document.archived.is_(False),
)
).all()
for document in documents:
+1 -1
View File
@@ -22,7 +22,7 @@ def create_tidb_serverless_task():
try:
# check the number of idle tidb serverless
idle_tidb_serverless_number = (
db.session.scalar(select(func.count(TidbAuthBinding.id)).where(TidbAuthBinding.active == False)) or 0
db.session.scalar(select(func.count(TidbAuthBinding.id)).where(TidbAuthBinding.active.is_(False))) or 0
)
if idle_tidb_serverless_number >= tidb_serverless_number:
break
@@ -20,7 +20,7 @@ def update_tidb_serverless_status_task():
# check the number of idle tidb serverless
tidb_serverless_list = db.session.scalars(
select(TidbAuthBinding).where(
TidbAuthBinding.active == False,
TidbAuthBinding.active.is_(False),
TidbAuthBinding.status == TidbAuthBindingStatus.CREATING,
)
).all()
+16
View File
@@ -185,6 +185,18 @@ class RedisInvitationTokenStore(InvitationTokenStore):
return None
return _invitation_token_adapter.validate_json(data)
@override
def consume_membership(self, invitation: InvitationLookup) -> bool:
key = self._invitation_token_key(invitation.token)
ttl = self._redis.pttl(key)
if ttl <= 0:
return False
return bool(self._redis.set(self._membership_used_key(invitation.token), "1", nx=True, px=ttl))
@override
def is_membership_consumed(self, invitation: InvitationLookup) -> bool:
return self._redis.get(self._membership_used_key(invitation.token)) == b"1"
@override
def revoke(self, invitation: InvitationLookup) -> None:
if invitation.workspace_id is not None and invitation.email is not None:
@@ -196,6 +208,10 @@ class RedisInvitationTokenStore(InvitationTokenStore):
def _invitation_token_key(token: str) -> str:
return f"member_invite:token:{token}"
@classmethod
def _membership_used_key(cls, token: str) -> str:
return f"{cls._invitation_token_key(token)}:membership_used"
@staticmethod
def _workspace_invitation_key(invitation: InvitationLookup) -> str:
assert invitation.workspace_id is not None
+18 -3
View File
@@ -27,7 +27,11 @@ from services.account.login_service import (
HumanVerificationGateway,
ResetPasswordEmailGateway,
)
from services.account_activation_service import AccountActivationRepository, InvitationTokenStore
from services.account_activation_service import (
AccountActivationRepository,
AccountActivationService,
InvitationTokenStore,
)
from services.account_security_gateway import RedisAccountEmailSecurityGateway
from services.billing_service import BillingService
from services.email_code_login_challenge import EmailCodeLoginChallengeStore, EmailCodeLoginChallengeUnavailableError
@@ -62,11 +66,18 @@ _RESET_PASSWORD_RATE_LIMIT_SECONDS = 60
class AccountActivationConsoleAuthInvitationGateway(ConsoleAuthInvitationGateway):
"""Resolve login invitations through the shared activation contracts."""
"""Resolve login invitations and restore invited membership through activation."""
def __init__(self, *, tokens: InvitationTokenStore, accounts: AccountActivationRepository) -> None:
def __init__(
self,
*,
tokens: InvitationTokenStore,
accounts: AccountActivationRepository,
activation: AccountActivationService,
) -> None:
self._tokens = tokens
self._accounts = accounts
self._activation = activation
@override
def resolve(self, *, email: str, token: str) -> LoginInvitation | None:
@@ -82,6 +93,10 @@ class AccountActivationConsoleAuthInvitationGateway(ConsoleAuthInvitationGateway
return LoginInvitation(email=invitation.account_email)
return None
@override
def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool:
return self._activation.ensure_membership(account_id=account_id, email=email, token=token)
class DeploymentConsoleAuthPolicyGateway(ConsoleAuthPolicyGateway):
def __init__(self, *, billing_enabled: bool) -> None:
+8 -3
View File
@@ -41,6 +41,8 @@ class ConsoleAuthWorkspaceQuery(Protocol):
class ConsoleAuthInvitationGateway(Protocol):
def resolve(self, *, email: str, token: str) -> LoginInvitation | None: ...
def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool: ...
class ConsoleAuthPolicyGateway(Protocol):
def get_email_freeze_type(self, email: str) -> str | None: ...
@@ -200,9 +202,12 @@ class ConsoleAuthenticationService:
self._invalid_credentials(normalized_email, command.ip_address)
if not self._workspaces.has_active_for_account(account.id):
if self._policies.is_workspace_creation_allowed() and not self._policies.has_workspace_capacity():
raise account_errors.LoginWorkspaceLimitError
return PasswordLoginResult(tokens=None, workspace_found=False)
if invitation_token is not None:
self._invitations.ensure_membership(account_id=account.id, email=command.email, token=invitation_token)
if not self._workspaces.has_active_for_account(account.id):
if self._policies.is_workspace_creation_allowed() and not self._policies.has_workspace_capacity():
raise account_errors.LoginWorkspaceLimitError
return PasswordLoginResult(tokens=None, workspace_found=False)
tokens = self._issue_session(
account_id=account.id,
+31 -3
View File
@@ -25,6 +25,10 @@ _PENDING_ACCOUNT_STATUS = "pending"
class InvitationTokenStore(Protocol):
def consume_membership(self, invitation: InvitationLookup) -> bool: ...
def is_membership_consumed(self, invitation: InvitationLookup) -> bool: ...
def find(self, invitation: InvitationLookup) -> InvitationToken | None: ...
def revoke(self, invitation: InvitationLookup) -> None: ...
@@ -39,6 +43,7 @@ class AccountActivationRepository(Protocol):
*,
role: str,
setup: AccountSetup | None,
create_membership: bool = True,
) -> ActivationPersistenceResult | None: ...
@@ -108,8 +113,9 @@ class AccountActivationService:
raise FrozenAccountError
setup = self._resolve_setup(invitation, command)
raw_role = invitation.role
role = raw_role if raw_role is not None and raw_role in _NON_OWNER_ROLES else _DEFAULT_ROLE
role = self._membership_role(invitation.role)
# An invitation that already restored membership at sign-in must not grant it again after removal.
membership_consumed = self._tokens.is_membership_consumed(command.invitation)
normalized_email = command.invitation.email.lower() if command.invitation.email else None
self._tokens.revoke(
@@ -119,13 +125,29 @@ class AccountActivationService:
token=command.invitation.token,
)
)
result = self._accounts.activate(invitation, role=role, setup=setup)
result = self._accounts.activate(invitation, role=role, setup=setup, create_membership=not membership_consumed)
if result is None:
raise InvalidInvitationError
if result.membership_created:
self._membership_cache.invalidate(invitation.workspace_id)
self._member_access_sync.sync(invitation.workspace_id, invitation.account_id)
def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool:
lookup = InvitationLookup(workspace_id=None, email=email, token=token)
invitation = self._resolve(lookup)
if invitation is None or invitation.account_id != account_id:
return False
if not self._tokens.consume_membership(lookup):
return False
result = self._accounts.activate(invitation, role=self._membership_role(invitation.role), setup=None)
if result is None:
return False
if result.membership_created:
self._membership_cache.invalidate(invitation.workspace_id)
self._member_access_sync.sync(invitation.workspace_id, invitation.account_id)
return True
def _resolve(self, invitation: InvitationLookup) -> AccountInvitation | None:
token = self._tokens.find(invitation)
resolved = self._accounts.resolve(token) if token is not None else None
@@ -146,6 +168,12 @@ class AccountActivationService:
return None
return self._accounts.resolve(token)
@staticmethod
def _membership_role(role: str | None) -> str:
if role is not None and role in _NON_OWNER_ROLES:
return role
return _DEFAULT_ROLE
@staticmethod
def _requires_setup(invitation: AccountInvitation) -> bool:
if invitation.requires_setup is not None:
+1 -4
View File
@@ -750,10 +750,7 @@ class AppDslService:
current_draft_workflow = locked_draft or workflow_service.get_draft_workflow(
app_model=app, session=self._session
)
if current_draft_workflow:
unique_hash = current_draft_workflow.unique_hash
else:
unique_hash = None
unique_hash = current_draft_workflow.unique_hash if current_draft_workflow else None
graph = workflow_data.get("graph", {})
if not isinstance(graph, dict):
raise ValueError("Workflow graph must be a mapping")
+2 -5
View File
@@ -223,7 +223,7 @@ class AppService:
def _build_app_list_filters(
user_id: str, tenant_id: str, params: AppListBaseParams, session: Session
) -> list[sa.ColumnElement[bool]]:
filters = [App.tenant_id == tenant_id, App.is_universal == False]
filters = [App.tenant_id == tenant_id, App.is_universal.is_(False)]
if params.mode == "workflow":
filters.append(App.mode == AppMode.WORKFLOW)
@@ -700,10 +700,7 @@ class AppService:
app.name = args["name"]
app.description = args["description"]
icon_type = args.get("icon_type")
if icon_type is None:
resolved_icon_type = app.icon_type
else:
resolved_icon_type = IconType(icon_type)
resolved_icon_type = app.icon_type if icon_type is None else IconType(icon_type)
app.icon_type = resolved_icon_type
app.icon = args["icon"]
+2 -2
View File
@@ -49,7 +49,7 @@ class ConversationService:
return InfiniteScrollPagination(data=[], limit=limit, has_more=False)
stmt = select(Conversation).where(
Conversation.is_deleted == False,
Conversation.is_deleted.is_(False),
Conversation.app_id == app_model.id,
Conversation.from_source == ("api" if isinstance(user, EndUser) else "console"),
Conversation.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
@@ -183,7 +183,7 @@ class ConversationService:
Conversation.from_source == ("api" if isinstance(user, EndUser) else "console"),
Conversation.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
Conversation.from_account_id == (user.id if isinstance(user, Account) else None),
Conversation.is_deleted == False,
Conversation.is_deleted.is_(False),
)
.limit(1)
)
+1 -1
View File
@@ -453,7 +453,7 @@ class DatasourceProviderService:
DatasourceOauthTenantParamConfig.tenant_id == tenant_id,
DatasourceOauthTenantParamConfig.provider == datasource_provider_id.provider_name,
DatasourceOauthTenantParamConfig.plugin_id == datasource_provider_id.plugin_id,
DatasourceOauthTenantParamConfig.enabled == True,
DatasourceOauthTenantParamConfig.enabled.is_(True),
)
)
or 0
@@ -199,4 +199,3 @@ class SystemFeatureModel(FeatureResponseModel):
enable_learn_app: bool = True
enable_step_by_step_tour: bool = False
rbac_enabled: bool = False
knowledge_fs_enabled: bool = False
+11 -14
View File
@@ -1834,9 +1834,9 @@ class DocumentService:
Document.tenant_id == dataset_ref.tenant_id,
Document.dataset_id == dataset_ref.dataset_id,
Document.id.in_(document_ids),
Document.enabled == True,
Document.enabled.is_(True),
Document.indexing_status == IndexingStatus.COMPLETED,
Document.archived == False,
Document.archived.is_(False),
)
).all()
return documents
@@ -1846,7 +1846,7 @@ class DocumentService:
documents = session.scalars(
select(Document).where(
Document.dataset_id == dataset_id,
Document.enabled == True,
Document.enabled.is_(True),
)
).all()
@@ -1857,9 +1857,9 @@ class DocumentService:
documents = session.scalars(
select(Document).where(
Document.dataset_id == dataset_id,
Document.enabled == True,
Document.enabled.is_(True),
Document.indexing_status == IndexingStatus.COMPLETED,
Document.archived == False,
Document.archived.is_(False),
)
).all()
@@ -2401,7 +2401,7 @@ class DocumentService:
Document.dataset_id == dataset.id,
Document.tenant_id == account.current_tenant_id,
Document.data_source_type == DataSourceType.UPLOAD_FILE,
Document.enabled == True,
Document.enabled.is_(True),
Document.name.in_(file_names),
)
).all()
@@ -2456,7 +2456,7 @@ class DocumentService:
Document.dataset_id == dataset.id,
Document.tenant_id == account.current_tenant_id,
Document.data_source_type == DataSourceType.NOTION_IMPORT,
Document.enabled == True,
Document.enabled.is_(True),
)
).all()
)
@@ -2515,10 +2515,7 @@ class DocumentService:
"only_main_content": website_info.only_main_content,
"mode": "crawl",
}
if len(url) > 255:
document_name = url[:200] + "..."
else:
document_name = url
document_name = url[:200] + "..." if len(url) > 255 else url
document = DocumentService.build_document(
dataset,
dataset_process_rule.id,
@@ -2902,8 +2899,8 @@ class DocumentService:
session.scalar(
select(func.count(Document.id)).where(
Document.completed_at.isnot(None),
Document.enabled == True,
Document.archived == False,
Document.enabled.is_(True),
Document.archived.is_(False),
Document.tenant_id == current_user.current_tenant_id,
)
)
@@ -3003,7 +3000,7 @@ class DocumentService:
sa.and_(
DataSourceOauthBinding.tenant_id == account.current_tenant_id,
DataSourceOauthBinding.provider == "notion",
DataSourceOauthBinding.disabled == False,
DataSourceOauthBinding.disabled.is_(False),
DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
)
)
+1 -4
View File
@@ -261,10 +261,7 @@ class ExternalDatasetService:
@staticmethod
def assembling_headers(authorization: Authorization, headers: dict[str, Any] | None = None) -> dict[str, Any]:
authorization = deepcopy(authorization)
if headers:
headers = deepcopy(headers)
else:
headers = {}
headers = deepcopy(headers) if headers else {}
if authorization.type == "api-key":
if authorization.config is None:
raise ValueError("authorization config is required")
-500
View File
@@ -1,500 +0,0 @@
"""Product-facing KnowledgeFS operation and authorization declarations.
This registry is Dify's explicit Console surface. Transport concerns live in
knowledge_fs_proxy so contract review does not require reading proxy mechanics.
"""
from __future__ import annotations
from typing import Final, Literal, NamedTuple
from core.rbac import RBACPermission
type KnowledgeFSMethod = Literal["DELETE", "GET", "PATCH", "POST", "PUT"]
type KnowledgeFSResponseKind = Literal["binary", "buffered", "stream"]
type KnowledgeFSRequiredScope = Literal["knowledge-spaces:read", "knowledge-spaces:write"]
type KnowledgeFSLegacyRole = Literal["reader", "dataset_editor", "admin"]
type KnowledgeFSErrorStatusMap = tuple[tuple[int, int], ...]
class KnowledgeFSOperation(NamedTuple):
operation_id: str
method: KnowledgeFSMethod
path: str
response_kind: KnowledgeFSResponseKind
required_scope: KnowledgeFSRequiredScope
rbac_permission: RBACPermission
legacy_role: KnowledgeFSLegacyRole
max_response_bytes: int
request_headers: tuple[str, ...]
response_headers: tuple[str, ...]
response_media_types: tuple[str, ...]
error_status_map: KnowledgeFSErrorStatusMap
def _console_operation(
operation_id: str,
method: KnowledgeFSMethod,
path: str,
*,
rbac_permission: RBACPermission,
legacy_role: KnowledgeFSLegacyRole,
max_response_bytes: int = 1_048_576,
request_headers: tuple[str, ...] = ("x-trace-id",),
response_kind: KnowledgeFSResponseKind = "buffered",
response_media_types: tuple[str, ...] = ("application/json",),
error_status_map: KnowledgeFSErrorStatusMap = ((401, 502), (403, 403)),
) -> KnowledgeFSOperation:
"""Declare one contract-pinned operation with an explicit Dify authorization policy."""
is_read = method == "GET"
return KnowledgeFSOperation(
operation_id=operation_id,
method=method,
path=path,
response_kind=response_kind,
required_scope="knowledge-spaces:read" if is_read else "knowledge-spaces:write",
rbac_permission=rbac_permission,
legacy_role=legacy_role,
max_response_bytes=max_response_bytes,
request_headers=request_headers,
response_headers=("x-trace-id",),
response_media_types=response_media_types,
error_status_map=error_status_map,
)
def _dataset_read_operation(operation_id: str, path: str) -> KnowledgeFSOperation:
"""Declare a dataset-readable buffered JSON operation."""
return _console_operation(
operation_id,
"GET",
path,
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
)
def _dataset_edit_operation(
operation_id: str,
method: KnowledgeFSMethod,
path: str,
*,
request_headers: tuple[str, ...] = ("x-trace-id",),
) -> KnowledgeFSOperation:
"""Declare a dataset-editable buffered JSON operation."""
return _console_operation(
operation_id,
method,
path,
rbac_permission=RBACPermission.DATASET_EDIT,
legacy_role="dataset_editor",
request_headers=request_headers,
)
def _external_source_operation(
operation_id: str,
method: KnowledgeFSMethod,
path: str,
*,
request_headers: tuple[str, ...] = ("x-trace-id",),
) -> KnowledgeFSOperation:
"""Declare an external-source operation restricted to dataset editors."""
return _console_operation(
operation_id,
method,
path,
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
request_headers=request_headers,
)
KNOWLEDGE_FS_CONSOLE_OPERATIONS: Final[tuple[KnowledgeFSOperation, ...]] = (
_console_operation(
operation_id="listKnowledgeSpaces",
method="GET",
path="knowledge-spaces",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_console_operation(
operation_id="createKnowledgeSpace",
method="POST",
path="knowledge-spaces",
rbac_permission=RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="getKnowledgeSpacesById",
method="GET",
path="knowledge-spaces/{id}",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_dataset_edit_operation("patchKnowledgeSpacesById", "PATCH", "knowledge-spaces/{id}"),
_dataset_edit_operation(
"deleteKnowledgeSpacesById",
"DELETE",
"knowledge-spaces/{id}",
request_headers=("idempotency-key", "x-trace-id"),
),
_dataset_read_operation("getKnowledgeSpacesByIdStats", "knowledge-spaces/{id}/stats"),
_console_operation(
operation_id="getKnowledgeSpacesByIdAccessPolicy",
method="GET",
path="knowledge-spaces/{id}/access-policy",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_console_operation(
operation_id="patchKnowledgeSpacesByIdAccessPolicy",
method="PATCH",
path="knowledge-spaces/{id}/access-policy",
rbac_permission=RBACPermission.DATASET_ACCESS_CONFIG,
legacy_role="admin",
),
_console_operation(
operation_id="getSourceProviders",
method="GET",
path="source-providers",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdSourceConnections",
method="GET",
path="knowledge-spaces/{id}/source-connections",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="postKnowledgeSpacesByIdSourceConnections",
method="POST",
path="knowledge-spaces/{id}/source-connections",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_external_source_operation(
"postKnowledgeSpacesByIdSourceConnectionsOauth",
"POST",
"knowledge-spaces/{id}/source-connections/oauth",
),
_external_source_operation("postSourceOauthCallback", "POST", "source-oauth/callback"),
_external_source_operation(
"getKnowledgeSpacesByIdSourceConnectionsByConnectionId",
"GET",
"knowledge-spaces/{id}/source-connections/{connectionId}",
),
_external_source_operation(
"deleteKnowledgeSpacesByIdSourceConnectionsByConnectionId",
"DELETE",
"knowledge-spaces/{id}/source-connections/{connectionId}",
),
_console_operation(
operation_id="postKnowledgeSpacesByIdSourceConnectionsByConnectionIdRefresh",
method="POST",
path="knowledge-spaces/{id}/source-connections/{connectionId}/refresh",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdSources",
method="GET",
path="knowledge-spaces/{id}/sources",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_console_operation(
operation_id="postKnowledgeSpacesByIdSources",
method="POST",
path="knowledge-spaces/{id}/sources",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_external_source_operation(
"getKnowledgeSpacesByIdSourcesBySourceId",
"GET",
"knowledge-spaces/{id}/sources/{sourceId}",
),
_external_source_operation(
"patchKnowledgeSpacesByIdSourcesBySourceId",
"PATCH",
"knowledge-spaces/{id}/sources/{sourceId}",
),
_external_source_operation(
"deleteKnowledgeSpacesByIdSourcesBySourceId",
"DELETE",
"knowledge-spaces/{id}/sources/{sourceId}",
request_headers=("idempotency-key", "x-trace-id"),
),
_external_source_operation(
"putKnowledgeSpacesByIdSourcesBySourceIdCredentials",
"PUT",
"knowledge-spaces/{id}/sources/{sourceId}/credentials",
),
_external_source_operation(
"deleteKnowledgeSpacesByIdSourcesBySourceIdCredentials",
"DELETE",
"knowledge-spaces/{id}/sources/{sourceId}/credentials",
),
_external_source_operation(
"postKnowledgeSpacesByIdSourcesBySourceIdSync",
"POST",
"knowledge-spaces/{id}/sources/{sourceId}/sync",
request_headers=("idempotency-key", "x-trace-id"),
),
_console_operation(
operation_id="postKnowledgeSpacesByIdSourcesBySourceIdCrawlPreview",
method="POST",
path="knowledge-spaces/{id}/sources/{sourceId}/crawl-preview",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
request_headers=("idempotency-key", "x-trace-id"),
),
_external_source_operation(
"postKnowledgeSpacesByIdSourcesBySourceIdWorkflowImports",
"POST",
"knowledge-spaces/{id}/sources/{sourceId}/workflow-imports",
request_headers=("idempotency-key", "x-trace-id"),
),
_external_source_operation(
"getKnowledgeSpacesByIdSourcesBySourceIdPages",
"GET",
"knowledge-spaces/{id}/sources/{sourceId}/pages",
),
_external_source_operation(
"getKnowledgeSpacesByIdSourcesBySourceIdFiles",
"GET",
"knowledge-spaces/{id}/sources/{sourceId}/files",
),
_external_source_operation(
"postKnowledgeSpacesByIdSourcesBySourceIdCrawl",
"POST",
"knowledge-spaces/{id}/sources/{sourceId}/crawl",
),
_external_source_operation(
"postKnowledgeSpacesByIdSourcesBySourceIdImport",
"POST",
"knowledge-spaces/{id}/sources/{sourceId}/import",
),
_external_source_operation(
"postKnowledgeSpacesByIdSourcesBySourceIdTest",
"POST",
"knowledge-spaces/{id}/sources/{sourceId}/test",
),
_external_source_operation(
"postKnowledgeSpacesByIdSourcesBySourceIdImportFiles",
"POST",
"knowledge-spaces/{id}/sources/{sourceId}/import-files",
),
_external_source_operation(
"postKnowledgeSpacesByIdSourcesBulk",
"POST",
"knowledge-spaces/{id}/sources/bulk",
request_headers=("idempotency-key", "x-trace-id"),
),
_external_source_operation(
"getKnowledgeSpacesByIdSourceWorkflows",
"GET",
"knowledge-spaces/{id}/source-workflows",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdSourceWorkflowsByRunId",
method="GET",
path="knowledge-spaces/{id}/source-workflows/{runId}",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_external_source_operation(
"getKnowledgeSpacesByIdSourceWorkflowsByRunIdBulkItems",
"GET",
"knowledge-spaces/{id}/source-workflows/{runId}/bulk-items",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdSourceWorkflowsByRunIdPages",
method="GET",
path="knowledge-spaces/{id}/source-workflows/{runId}/pages",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="postKnowledgeSpacesByIdSourceWorkflowsByRunIdCancel",
method="POST",
path="knowledge-spaces/{id}/source-workflows/{runId}/cancel",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="postKnowledgeSpacesByIdSourceWorkflowsByRunIdRetry",
method="POST",
path="knowledge-spaces/{id}/source-workflows/{runId}/retry",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="postKnowledgeSpacesByIdSourceWorkflowsByRunIdSelection",
method="POST",
path="knowledge-spaces/{id}/source-workflows/{runId}/selection",
rbac_permission=RBACPermission.DATASET_EXTERNAL_CONNECT,
legacy_role="dataset_editor",
request_headers=("idempotency-key", "x-trace-id"),
),
_console_operation(
operation_id="getKnowledgeSpacesByIdSourcesBySourceIdSyncPolicy",
method="GET",
path="knowledge-spaces/{id}/sources/{sourceId}/sync-policy",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_external_source_operation(
"putKnowledgeSpacesByIdSourcesBySourceIdSyncPolicy",
"PUT",
"knowledge-spaces/{id}/sources/{sourceId}/sync-policy",
),
_dataset_read_operation("getKnowledgeSpacesByIdDocuments", "knowledge-spaces/{id}/documents"),
_dataset_edit_operation("postKnowledgeSpacesByIdDocuments", "POST", "knowledge-spaces/{id}/documents"),
_dataset_edit_operation(
"deleteKnowledgeSpacesByIdDocumentsBulk",
"DELETE",
"knowledge-spaces/{id}/documents/bulk",
request_headers=("idempotency-key", "x-trace-id"),
),
_dataset_edit_operation(
"postKnowledgeSpacesByIdDocumentsBulk",
"POST",
"knowledge-spaces/{id}/documents/bulk",
),
_dataset_edit_operation(
"postKnowledgeSpacesByIdDocumentsBulkReindex",
"POST",
"knowledge-spaces/{id}/documents/bulk/reindex",
),
_dataset_read_operation(
"getKnowledgeSpacesByIdDocumentsByDocumentId",
"knowledge-spaces/{id}/documents/{documentId}",
),
_dataset_edit_operation(
"deleteKnowledgeSpacesByIdDocumentsByDocumentId",
"DELETE",
"knowledge-spaces/{id}/documents/{documentId}",
request_headers=("idempotency-key", "x-trace-id"),
),
_console_operation(
operation_id="getKnowledgeSpacesByIdLogicalDocuments",
method="GET",
path="knowledge-spaces/{id}/logical-documents",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_dataset_edit_operation(
"deleteKnowledgeSpacesByIdLogicalDocumentsByDocumentId",
"DELETE",
"knowledge-spaces/{id}/logical-documents/{documentId}",
request_headers=("idempotency-key", "x-trace-id"),
),
_dataset_read_operation(
"getKnowledgeSpacesByIdDocumentsByDocumentIdOutline",
"knowledge-spaces/{id}/documents/{documentId}/outline",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdLogicalDocumentsByDocumentId",
method="GET",
path="knowledge-spaces/{id}/logical-documents/{documentId}",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdDocumentsByDocumentIdRevisions",
method="GET",
path="knowledge-spaces/{id}/documents/{documentId}/revisions",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_dataset_edit_operation(
"postKnowledgeSpacesByIdDocumentsByDocumentIdRevisionsByRevisionRollback",
"POST",
"knowledge-spaces/{id}/documents/{documentId}/revisions/{revision}/rollback",
),
_dataset_edit_operation(
"patchKnowledgeSpacesByIdDocumentsByDocumentIdMetadata",
"PATCH",
"knowledge-spaces/{id}/documents/{documentId}/metadata",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdDocumentsByDocumentIdRevisionsByRevisionChunks",
method="GET",
path="knowledge-spaces/{id}/documents/{documentId}/revisions/{revision}/chunks",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_dataset_read_operation(
"getKnowledgeSpacesByIdDocumentsByDocumentIdRevisionsByRevisionChunksByChunkId",
"knowledge-spaces/{id}/documents/{documentId}/revisions/{revision}/chunks/{chunkId}",
),
_dataset_edit_operation(
"postKnowledgeSpacesByIdDocumentsByDocumentIdRevisionsByRevisionChunksByChunkIdState",
"POST",
"knowledge-spaces/{id}/documents/{documentId}/revisions/{revision}/chunks/{chunkId}/state",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdProcessingTasks",
method="GET",
path="knowledge-spaces/{id}/processing-tasks",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
),
_dataset_read_operation(
"getKnowledgeSpacesByIdDocumentsByDocumentIdProcessingTasks",
"knowledge-spaces/{id}/documents/{documentId}/processing-tasks",
),
_dataset_read_operation(
"getKnowledgeSpacesByIdDocumentsByDocumentIdProcessingTasksByTaskId",
"knowledge-spaces/{id}/documents/{documentId}/processing-tasks/{taskId}",
),
_console_operation(
operation_id="getKnowledgeSpacesByIdDocumentsByDocumentIdProcessingTasksByTaskIdEvents",
method="GET",
path="knowledge-spaces/{id}/documents/{documentId}/processing-tasks/{taskId}/events",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
max_response_bytes=67_108_864,
request_headers=("last-event-id", "x-trace-id"),
response_kind="stream",
response_media_types=("text/event-stream",),
),
_console_operation(
operation_id="deleteKnowledgeSpacesByIdDocumentsByDocumentIdProcessingTasksByTaskId",
method="DELETE",
path="knowledge-spaces/{id}/documents/{documentId}/processing-tasks/{taskId}",
rbac_permission=RBACPermission.DATASET_EDIT,
legacy_role="dataset_editor",
),
_console_operation(
operation_id="postKnowledgeSpacesByIdDocumentsByDocumentIdProcessingTasksByTaskIdRetry",
method="POST",
path="knowledge-spaces/{id}/documents/{documentId}/processing-tasks/{taskId}/retry",
rbac_permission=RBACPermission.DATASET_EDIT,
legacy_role="dataset_editor",
),
_dataset_read_operation(
"getKnowledgeSpacesByIdDocumentsByDocumentIdSettings",
"knowledge-spaces/{id}/documents/{documentId}/settings",
),
_dataset_edit_operation(
"putKnowledgeSpacesByIdDocumentsByDocumentIdSettings",
"PUT",
"knowledge-spaces/{id}/documents/{documentId}/settings",
),
_dataset_read_operation("getJobsById", "jobs/{id}"),
_dataset_edit_operation("deleteJobsById", "DELETE", "jobs/{id}"),
_dataset_edit_operation("postJobsByIdRetry", "POST", "jobs/{id}/retry"),
_dataset_read_operation("getDeletionJobsByJobId", "deletion-jobs/{jobId}"),
_dataset_edit_operation(
"postDeletionJobsByJobIdRetry",
"POST",
"deletion-jobs/{jobId}/retry",
request_headers=("idempotency-key", "x-trace-id"),
),
_dataset_read_operation("getBulkJobsById", "bulk-jobs/{id}"),
)
-408
View File
@@ -1,408 +0,0 @@
"""Authorize and forward the explicitly enabled KnowledgeFS Console operations.
The dedicated request path uses Dify's shared SSRF policy, never follows redirects,
bounds buffered responses, and rejects compressed streaming responses.
"""
from __future__ import annotations
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from http import HTTPStatus
from typing import NamedTuple, Protocol
import httpx
import jwt
from configs import dify_config
from core.helper import ssrf_proxy
from core.rbac import RBACResourceScope
from core.tools.errors import ToolSSRFError
from models import Account
from services.enterprise.rbac_service import RBACService
from services.knowledge_fs_operations import (
KNOWLEDGE_FS_CONSOLE_OPERATIONS,
KnowledgeFSMethod,
KnowledgeFSOperation,
KnowledgeFSResponseKind,
)
_JWT_AUDIENCE = "knowledge-fs"
_JWT_ISSUER = "dify"
_JWT_TTL_SECONDS = 60
_MAX_BUFFERED_RESPONSE_BYTES = 1024 * 1024
class KnowledgeFSUpstreamResponse(NamedTuple):
response: httpx.Response
response_kind: KnowledgeFSResponseKind
operation: KnowledgeFSOperation
_AUTHORIZATION_MARKER = object()
@dataclass(eq=False, frozen=True, init=False, slots=True)
class KnowledgeFSAuthorization:
"""Single-use forwarding capability created after Dify workspace policy checks.
Callers obtain this value from :func:`authorize_knowledge_fs_request`. Direct
construction and repeated forwarding are rejected before outbound I/O.
"""
account_id: str
tenant_id: str
operation: KnowledgeFSOperation
_used: bool
def __init__(
self,
account_id: str,
tenant_id: str,
operation: KnowledgeFSOperation,
*,
_marker: object | None = None,
) -> None:
if _marker is not _AUTHORIZATION_MARKER:
raise KnowledgeFSAccessDeniedError("KnowledgeFS authorization must be created by workspace authorization")
object.__setattr__(self, "account_id", account_id)
object.__setattr__(self, "tenant_id", tenant_id)
object.__setattr__(self, "operation", operation)
object.__setattr__(self, "_used", False)
def consume(self) -> tuple[str, str, KnowledgeFSOperation]:
"""Return the authorized principals and canonical operation exactly once.
Raises:
KnowledgeFSAccessDeniedError: The capability was already consumed.
"""
if self._used:
raise KnowledgeFSAccessDeniedError("KnowledgeFS authorization has already been used")
object.__setattr__(self, "_used", True)
return self.account_id, self.tenant_id, self.operation
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,
path: str,
) -> KnowledgeFSAuthorization:
"""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: Requested upstream HTTP method.
path: Requested relative KnowledgeFS path.
Raises:
KnowledgeFSRouteNotAllowedError: The method and path do not resolve to a declared operation.
KnowledgeFSAccessDeniedError: The account lacks a required legacy or enterprise permission.
Returns:
A request-scoped capability binding the authorized account, workspace, and operation.
"""
operation = get_knowledge_fs_operation(method, path)
if operation.legacy_role == "dataset_editor" and not account.is_dataset_editor:
raise KnowledgeFSAccessDeniedError("KnowledgeFS operation requires dataset edit access")
if operation.legacy_role == "admin" and not account.is_admin_or_owner:
raise KnowledgeFSAccessDeniedError("KnowledgeFS operation requires workspace administration access")
if not RBACService.CheckAccess.check(
tenant_id,
account.id,
scene=operation.rbac_permission.value,
resource_type=RBACResourceScope.DATASET.value,
):
raise KnowledgeFSAccessDeniedError("KnowledgeFS operation is denied by workspace RBAC")
return KnowledgeFSAuthorization(account.id, tenant_id, operation, _marker=_AUTHORIZATION_MARKER)
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."""
authorization = authorize_knowledge_fs_request(
account=account,
tenant_id=tenant_id,
method=method,
path=path,
)
return proxy_authorized_knowledge_fs_request(
authorization=authorization,
accept=accept,
content_type=content_type,
query=query,
body=body,
request_headers=request_headers,
)
def proxy_authorized_knowledge_fs_request(
*,
authorization: KnowledgeFSAuthorization,
accept: str | None = None,
content_type: str | None = None,
query: bytes | None = None,
body: bytes | None = None,
request_headers: _RequestHeaders | None = None,
) -> KnowledgeFSUpstreamResponse:
"""Forward one request whose operation and workspace policy were already authorized.
This performs one outbound KnowledgeFS request and does not repeat Dify RBAC checks.
Args:
authorization: Request-scoped capability returned by :func:`authorize_knowledge_fs_request`.
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: Incoming headers; only names declared by the operation are forwarded.
Returns:
The bounded KnowledgeFS response together with its transport metadata.
Raises:
KnowledgeFSConfigurationError: The connection is incomplete or blocked by outbound policy.
KnowledgeFSRouteNotAllowedError: A forwarded request header is outside the operation contract.
KnowledgeFSTimeoutError: KnowledgeFS exceeds the configured timeout.
KnowledgeFSTransportError: The request fails or its response violates transport bounds.
"""
account_id, tenant_id, operation = authorization.consume()
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=operation.method,
path=operation.path,
tenant_id=tenant_id,
accept=accept,
content_type=content_type,
query=query,
body=body,
request_headers=contract_request_headers,
)
def _forward_knowledge_fs_request(
*,
account_id: str,
method: KnowledgeFSMethod,
path: str,
tenant_id: str,
accept: str | None = None,
content_type: str | None = None,
query: bytes | None = None,
body: bytes | None = None,
request_headers: Mapping[str, str] | None = None,
) -> KnowledgeFSUpstreamResponse:
"""Forward one fixed-route request without parsing its KnowledgeFS payload.
Args:
account_id: Current Dify account used as the KFS member identity.
method: Allowlisted upstream HTTP method.
path: Relative KnowledgeFS path under an allowlisted product surface.
tenant_id: Current Dify workspace used as the KFS tenant identity.
accept: Original Accept header, when present.
content_type: Original request Content-Type header, when present.
query: Original encoded query string from the Console request.
body: Original request body, when present.
request_headers: Contract-declared request headers forwarded by the Console adapter.
Returns:
The KnowledgeFS response and its actual transport kind. Non-success responses are buffered.
Raises:
KnowledgeFSConfigurationError: The connection is incomplete or blocked by outbound policy.
KnowledgeFSRouteNotAllowedError: The path is outside the allowlisted product surface.
KnowledgeFSTimeoutError: KnowledgeFS exceeds the configured timeout.
KnowledgeFSTransportError: The request fails or its response cannot be safely bounded.
Each request is bound to stable Dify account and workspace principals with a short expiration.
"""
operation = get_knowledge_fs_operation(method, path)
base_url = dify_config.KNOWLEDGE_FS_BASE_URL
jwt_secret = dify_config.KNOWLEDGE_FS_JWT_SECRET
if base_url is None or jwt_secret is None:
raise KnowledgeFSConfigurationError("KnowledgeFS connection configuration is incomplete")
now = datetime.now(UTC)
token = jwt.encode(
{
"aud": _JWT_AUDIENCE,
"caller_kind": "interactive",
"dify_account_id": f"dify-account:{account_id}",
"exp": now + timedelta(seconds=_JWT_TTL_SECONDS),
"iat": now,
"iss": _JWT_ISSUER,
"scopes": [operation.required_scope],
"sub": f"dify-workspace:{tenant_id}",
"tenant_id": tenant_id,
},
jwt_secret.get_secret_value(),
algorithm="HS256",
)
headers = {
"Accept": accept or "application/json",
"Accept-Encoding": "identity",
"Authorization": f"Bearer {token}",
}
if body is not None:
headers["Content-Type"] = content_type or "application/json"
allowed_request_headers = set(operation.request_headers)
for name, value in (request_headers or {}).items():
normalized_name = name.lower()
if normalized_name not in allowed_request_headers:
raise KnowledgeFSRouteNotAllowedError("KnowledgeFS request header is not allowed")
headers[normalized_name] = value
try:
upstream_url = httpx.URL(f"{base_url}/").join(operation.path)
response = ssrf_proxy.make_request(
method=operation.method,
url=str(upstream_url),
params=query,
content=body,
headers=headers,
timeout=dify_config.KNOWLEDGE_FS_TIMEOUT_SECONDS,
follow_redirects=False,
max_retries=0,
stream_response=True,
)
response_kind = _classify_response(operation, response)
if response_kind == "stream":
content_encoding = response.headers.get("content-encoding", "identity").strip().lower()
if content_encoding not in {"", "identity"}:
response.close()
raise KnowledgeFSTransportError("KnowledgeFS streaming response used an unsupported encoding")
_set_response_read_timeout(response, dify_config.KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS)
return KnowledgeFSUpstreamResponse(response, response_kind, operation)
max_response_bytes = (
operation.max_response_bytes
if HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES
else _MAX_BUFFERED_RESPONSE_BYTES
)
buffered_response = ssrf_proxy.buffer_response(response, max_response_bytes=max_response_bytes)
if buffered_response.content and not buffered_response.headers.get("content-type", "").strip():
buffered_response.close()
raise KnowledgeFSTransportError("KnowledgeFS buffered response used an unsupported media type")
return KnowledgeFSUpstreamResponse(buffered_response, response_kind, operation)
except ssrf_proxy.ResponseLimitError as exc:
raise KnowledgeFSTransportError("KnowledgeFS response violated the proxy limit") from exc
except ToolSSRFError as exc:
raise KnowledgeFSConfigurationError("KnowledgeFS origin was blocked by outbound policy") from exc
except httpx.TimeoutException as exc:
raise KnowledgeFSTimeoutError("KnowledgeFS request timed out") from exc
except httpx.RequestError as exc:
raise KnowledgeFSTransportError("KnowledgeFS transport request failed") from exc
def get_knowledge_fs_operation(method: KnowledgeFSMethod, path: str) -> KnowledgeFSOperation:
"""Resolve an exact operation and its transport/access contract metadata."""
for operation in KNOWLEDGE_FS_CONSOLE_OPERATIONS:
if method == operation.method and _matches_route_template(operation.path, path):
return operation._replace(path=path)
raise KnowledgeFSRouteNotAllowedError("KnowledgeFS route is not allowed")
def _classify_response(operation: KnowledgeFSOperation, response: httpx.Response) -> KnowledgeFSResponseKind:
"""Resolve the actual response kind from status and Content-Type before reading its body."""
content_type = response.headers.get("content-type", "").partition(";")[0].strip().lower()
is_success = HTTPStatus.OK <= response.status_code < HTTPStatus.MULTIPLE_CHOICES
if not is_success:
if content_type and not _is_json_content_type(content_type):
response.close()
raise KnowledgeFSTransportError("KnowledgeFS error response used an unsupported media type")
return "buffered"
if operation.response_kind == "stream":
if content_type != "text/event-stream":
response.close()
raise KnowledgeFSTransportError("KnowledgeFS stream response used an unsupported media type")
return "stream"
if operation.response_kind == "binary":
if content_type not in operation.response_media_types:
response.close()
raise KnowledgeFSTransportError("KnowledgeFS binary response used an unsupported media type")
return "binary"
if content_type and not _is_json_content_type(content_type):
response.close()
raise KnowledgeFSTransportError("KnowledgeFS buffered response used an unsupported media type")
return "buffered"
def _is_json_content_type(content_type: str) -> bool:
return content_type == "application/json" or content_type.endswith("+json")
def _set_response_read_timeout(response: httpx.Response, timeout_seconds: float | None) -> None:
"""Set the body-read timeout after headers identify a valid SSE response."""
try:
request = response.request
except RuntimeError:
return
timeout = request.extensions.get("timeout")
if isinstance(timeout, dict):
timeout["read"] = timeout_seconds
def _matches_route_template(template: str, path: str) -> bool:
"""Match path parameters without permitting encoded or traversal-like segments."""
template_segments = template.split("/")
path_segments = path.split("/")
if len(template_segments) != len(path_segments):
return False
for template_segment, path_segment in zip(template_segments, path_segments, strict=True):
if template_segment.startswith("{") and template_segment.endswith("}"):
if (
not path_segment
or path_segment in {".", ".."}
or "\\" in path_segment
or "%" in path_segment
or "?" in path_segment
or "#" in path_segment
):
return False
continue
if template_segment != path_segment:
return False
return True
+2 -3
View File
@@ -1,3 +1,4 @@
import contextlib
import re
from collections.abc import Mapping
from typing import Any
@@ -70,10 +71,8 @@ class DependenciesAnalysisService:
provider = tool.get("provider_id") or tool.get("provider_name") or tool.get("provider")
if not isinstance(provider, str):
continue
try:
with contextlib.suppress(ValueError, NotFound):
dependencies.append(cls.analyze_tool_provider_reference(provider))
except (ValueError, NotFound):
pass
return dependencies
@classmethod
+2 -5
View File
@@ -1467,7 +1467,7 @@ class RagPipelineService:
def get_recommended_plugins(self, type: str, current_user: Account, current_tenant_id: str) -> dict[str, Any]:
# Query active recommended plugins
stmt = select(PipelineRecommendedPlugin).where(PipelineRecommendedPlugin.active == True)
stmt = select(PipelineRecommendedPlugin).where(PipelineRecommendedPlugin.active.is_(True))
if type and type != "all":
stmt = stmt.where(PipelineRecommendedPlugin.type == type)
@@ -1595,10 +1595,7 @@ class RagPipelineService:
continue
variables = workflow.rag_pipeline_variables
if variables:
variables_map = {item["variable"]: item for item in variables}
else:
variables_map = {}
variables_map = {item["variable"]: item for item in variables} if variables else {}
datasource_parameters = datasource_node_data.get("datasource_parameters", {})
user_input_variables_keys = []
+1 -4
View File
@@ -118,10 +118,7 @@ class RemoteFileService:
):
raise FileTooLargeError()
if response.request.method == "GET":
content = response.content
else:
content = self._fetch_content(url=url)
content = response.content if response.request.method == "GET" else self._fetch_content(url=url)
upload_file = self._files.upload_file(
filename=file_info.filename,
@@ -325,10 +325,7 @@ class WorkflowRunArchiver:
# Initialize archive storage (will raise if not configured)
try:
if not self.dry_run:
storage = get_archive_storage()
else:
storage = None
storage = get_archive_storage() if not self.dry_run else None
except ArchiveStorageNotConfiguredError as e:
click.echo(click.style(f"Archive storage not configured: {e}", fg="red"))
return summary
@@ -32,6 +32,8 @@ class ArchiveManifest(TypedDict, total=False):
_manifest_adapter: TypeAdapter[ArchiveManifest] = TypeAdapter(ArchiveManifest)
import contextlib
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
from extensions.ext_database import db
@@ -338,10 +340,8 @@ class WorkflowRunRestore:
if isinstance(column.type, DateTime):
value = result.get(column.key)
if isinstance(value, str):
try:
with contextlib.suppress(ValueError):
result[column.key] = datetime.fromisoformat(value)
except ValueError:
pass
return result
+2 -4
View File
@@ -22,7 +22,7 @@ import xml.etree.ElementTree as ET
import zipfile
from base64 import b64encode
from collections.abc import Generator
from contextlib import contextmanager
from contextlib import contextmanager, suppress
from dataclasses import dataclass
from datetime import datetime
from enum import StrEnum
@@ -4708,10 +4708,8 @@ class SkillManagementService:
continue
if source[index : index + 2] == "\\'":
try:
with suppress(ValueError):
text_parts.append(bytes.fromhex(source[index + 2 : index + 4]).decode("latin-1"))
except ValueError:
pass
index += 4
continue
if index + 1 < len(source):
-1
View File
@@ -153,7 +153,6 @@ class SystemFeatureService:
system_features.enable_learn_app = dify_config.ENABLE_LEARN_APP
system_features.webapp_auth.allow_public_access = dify_config.WEBAPP_PUBLIC_ACCESS_ENABLED
system_features.enable_step_by_step_tour = dify_config.ENABLE_STEP_BY_STEP_TOUR
system_features.knowledge_fs_enabled = dify_config.KNOWLEDGE_FS_ENABLED
@classmethod
def _fulfill_params_from_enterprise(cls, features: feature_entities.SystemFeatureModel) -> None:
@@ -1023,10 +1023,7 @@ class DraftVariableSaver:
self._node_type,
)
continue
if isinstance(value, Segment):
value_seg = value
else:
value_seg = _build_segment_for_serialized_values(value)
value_seg = value if isinstance(value, Segment) else _build_segment_for_serialized_values(value)
draft_vars.append(
self._create_draft_variable(
name=name,
@@ -1,5 +1,6 @@
from __future__ import annotations
import contextlib
import json
import logging
import queue
@@ -692,10 +693,8 @@ def _start_buffering(subscription) -> BufferState:
buffer_state.queue.put_nowait(event)
except queue.Full:
dropped_count += 1
try:
with contextlib.suppress(queue.Empty):
buffer_state.queue.get_nowait()
except queue.Empty:
pass
try:
buffer_state.queue.put_nowait(event)
except queue.Full:
+6 -6
View File
@@ -40,8 +40,8 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
select(DatasetDocument).where(
DatasetDocument.dataset_id == dataset_id,
DatasetDocument.indexing_status == "completed",
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
).all()
@@ -61,7 +61,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
select(DocumentSegment)
.where(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
)
.order_by(DocumentSegment.position.asc())
).all()
@@ -107,8 +107,8 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
select(DatasetDocument).where(
DatasetDocument.dataset_id == dataset_id,
DatasetDocument.indexing_status == "completed",
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
).all()
# add new index
@@ -134,7 +134,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
select(DocumentSegment)
.where(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
)
.order_by(DocumentSegment.position.asc())
).all()
+6 -6
View File
@@ -44,8 +44,8 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
select(DatasetDocument).where(
DatasetDocument.dataset_id == dataset_id,
DatasetDocument.indexing_status == "completed",
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
).all()
@@ -65,7 +65,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
select(DocumentSegment)
.where(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
)
.order_by(DocumentSegment.position.asc())
).all()
@@ -103,8 +103,8 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
select(DatasetDocument).where(
DatasetDocument.dataset_id == dataset_id,
DatasetDocument.indexing_status == "completed",
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
).all()
# add new index
@@ -130,7 +130,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
select(DocumentSegment)
.where(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
)
.order_by(DocumentSegment.position.asc())
).all()
+6 -6
View File
@@ -101,7 +101,7 @@ def regenerate_summary_index_task(
.where(
DocumentSegment.dataset_id == dataset_id,
DocumentSegment.status == "completed", # Segment must be completed
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
DocumentSegmentSummary.dataset_id == dataset_id,
DocumentSegmentSummary.summary_content.isnot(None), # Must have summary content
# Include completed summaries or error summaries (with content)
@@ -109,8 +109,8 @@ def regenerate_summary_index_task(
DocumentSegmentSummary.status == SummaryStatus.COMPLETED,
DocumentSegmentSummary.status == SummaryStatus.ERROR,
),
DatasetDocument.enabled == True, # Document must be enabled
DatasetDocument.archived == False, # Document must not be archived
DatasetDocument.enabled.is_(True), # Document must be enabled
DatasetDocument.archived.is_(False), # Document must not be archived
DatasetDocument.doc_form != IndexStructureType.QA_INDEX, # Skip qa_model documents
)
.order_by(DocumentSegment.document_id.asc(), DocumentSegment.position.asc())
@@ -177,8 +177,8 @@ def regenerate_summary_index_task(
select(DatasetDocument).where(
DatasetDocument.dataset_id == dataset_id,
DatasetDocument.indexing_status == "completed",
DatasetDocument.enabled == True,
DatasetDocument.archived == False,
DatasetDocument.enabled.is_(True),
DatasetDocument.archived.is_(False),
)
).all()
@@ -214,7 +214,7 @@ def regenerate_summary_index_task(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.dataset_id == dataset_id,
DocumentSegment.status == "completed",
DocumentSegment.enabled == True,
DocumentSegment.enabled.is_(True),
DocumentSegmentSummary.dataset_id == dataset_id,
)
.order_by(DocumentSegment.position.asc())
@@ -251,7 +251,7 @@ class TestRagPipelineRunTasks:
assert call_kwargs["workflow_id"] == workflow.id
assert call_kwargs["user"].id == account.id
assert call_kwargs["invoke_from"] == InvokeFrom.PUBLISHED_PIPELINE
assert call_kwargs["streaming"] == False
assert call_kwargs["streaming"] is False
assert isinstance(call_kwargs["application_generate_entity"], RagPipelineGenerateEntity)
def test_rag_pipeline_run_task_success(
@@ -296,7 +296,7 @@ class TestRagPipelineRunTasks:
assert call_kwargs["workflow_id"] == workflow.id
assert call_kwargs["user"].id == account.id
assert call_kwargs["invoke_from"] == InvokeFrom.PUBLISHED_PIPELINE
assert call_kwargs["streaming"] == False
assert call_kwargs["streaming"] is False
assert isinstance(call_kwargs["application_generate_entity"], RagPipelineGenerateEntity)
def test_priority_rag_pipeline_run_task_with_waiting_tasks(
@@ -757,7 +757,7 @@ class TestRagPipelineRunTasks:
assert call_kwargs["workflow_id"] == workflow.id
assert call_kwargs["user"].id == account.id
assert call_kwargs["invoke_from"] == InvokeFrom.PUBLISHED_PIPELINE
assert call_kwargs["streaming"] == False
assert call_kwargs["streaming"] is False
assert isinstance(call_kwargs["application_generate_entity"], RagPipelineGenerateEntity)
def test_run_single_rag_pipeline_task_entity_validation_error(
@@ -1,142 +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_ENABLED",
"KNOWLEDGE_FS_BASE_URL",
"KNOWLEDGE_FS_JWT_SECRET",
"KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS",
"KNOWLEDGE_FS_TIMEOUT_SECONDS",
)
def test_knowledge_fs_config_normalizes_complete_connection() -> None:
config = KnowledgeFSConfig(
KNOWLEDGE_FS_ENABLED=True,
KNOWLEDGE_FS_BASE_URL=" https://knowledge-fs.test/ ",
KNOWLEDGE_FS_JWT_SECRET=" production-secret-with-at-least-32-bytes ",
)
assert config.KNOWLEDGE_FS_BASE_URL == "https://knowledge-fs.test"
assert isinstance(config.KNOWLEDGE_FS_JWT_SECRET, SecretStr)
assert config.KNOWLEDGE_FS_JWT_SECRET.get_secret_value() == "production-secret-with-at-least-32-bytes"
assert "production-secret" not in repr(config)
assert "production-secret" not in config.model_dump_json()
assert config.KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS == 300.0
def test_knowledge_fs_config_treats_blank_connection_as_disabled() -> None:
config = KnowledgeFSConfig(
KNOWLEDGE_FS_BASE_URL=" ",
KNOWLEDGE_FS_JWT_SECRET="",
)
assert config.KNOWLEDGE_FS_BASE_URL is None
assert config.KNOWLEDGE_FS_JWT_SECRET is None
assert config.KNOWLEDGE_FS_ENABLED is False
def test_knowledge_fs_config_requires_connection_when_enabled() -> None:
with pytest.raises(ValidationError, match="connection settings are required"):
KnowledgeFSConfig(KNOWLEDGE_FS_ENABLED=True)
@pytest.mark.parametrize(
("base_url", "jwt_secret"),
[
("https://knowledge-fs.test", None),
(None, "production-secret-with-at-least-32-bytes"),
],
)
def test_disabled_knowledge_fs_config_allows_partial_connection(base_url: str | None, jwt_secret: str | None) -> None:
config = KnowledgeFSConfig(
KNOWLEDGE_FS_ENABLED=False,
KNOWLEDGE_FS_BASE_URL=base_url,
KNOWLEDGE_FS_JWT_SECRET=jwt_secret,
)
assert config.KNOWLEDGE_FS_ENABLED is False
def test_knowledge_fs_docker_config_is_not_shadowed_by_root_env() -> None:
root_env_example = (_REPOSITORY_ROOT / "docker/.env.example").read_text(encoding="utf-8")
api_env_example = (_REPOSITORY_ROOT / "docker/envs/core-services/api.env.example").read_text(encoding="utf-8")
for variable in _KNOWLEDGE_FS_DOCKER_VARIABLES:
assert f"{variable}=" not in root_env_example
assert f"{variable}=" in api_env_example
@pytest.mark.parametrize(
("base_url", "jwt_secret"),
[
("https://knowledge-fs.test", None),
(None, "production-secret-with-at-least-32-bytes"),
],
)
def test_knowledge_fs_config_rejects_partial_connection(base_url: str | None, jwt_secret: str | None) -> None:
with pytest.raises(ValidationError, match="must be configured together"):
KnowledgeFSConfig(
KNOWLEDGE_FS_ENABLED=True,
KNOWLEDGE_FS_BASE_URL=base_url,
KNOWLEDGE_FS_JWT_SECRET=jwt_secret,
)
@pytest.mark.parametrize("base_url", ["knowledge-fs.test", "ftp://knowledge-fs.test", "http:///missing-host"])
def test_knowledge_fs_config_rejects_non_http_absolute_urls(base_url: str) -> None:
with pytest.raises(ValidationError, match="absolute HTTP\\(S\\) URL"):
KnowledgeFSConfig(
KNOWLEDGE_FS_BASE_URL=base_url,
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
)
@pytest.mark.parametrize(
"base_url",
[
"https://knowledge-fs.test:notaport",
"https://knowledge-fs.test:65536",
],
)
def test_knowledge_fs_config_rejects_invalid_ports(base_url: str) -> None:
with pytest.raises(ValidationError, match="valid port"):
KnowledgeFSConfig(
KNOWLEDGE_FS_BASE_URL=base_url,
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
)
@pytest.mark.parametrize(
"base_url",
[
"https://user:password@knowledge-fs.test",
"https://knowledge-fs.test?region=us",
"https://knowledge-fs.test#gateway",
],
)
def test_knowledge_fs_config_rejects_unsafe_base_url_components(base_url: str) -> None:
with pytest.raises(ValidationError, match="must not include credentials, query, or fragment"):
KnowledgeFSConfig(
KNOWLEDGE_FS_BASE_URL=base_url,
KNOWLEDGE_FS_JWT_SECRET="production-secret-with-at-least-32-bytes",
)
@pytest.mark.parametrize("timeout_seconds", [float("inf"), float("nan"), 60.0001])
def test_knowledge_fs_config_rejects_unbounded_timeouts(timeout_seconds: float) -> None:
with pytest.raises(ValidationError):
KnowledgeFSConfig(KNOWLEDGE_FS_TIMEOUT_SECONDS=timeout_seconds)
@pytest.mark.parametrize("timeout_seconds", [float("inf"), float("nan"), 3600.0001])
def test_knowledge_fs_config_rejects_unbounded_sse_read_timeouts(timeout_seconds: float) -> None:
with pytest.raises(ValidationError):
KnowledgeFSConfig(KNOWLEDGE_FS_SSE_READ_TIMEOUT_SECONDS=timeout_seconds)
@@ -256,24 +256,42 @@ def test_trace_app_config_get_empty_returns_exact_legacy_body(
)
@pytest.mark.parametrize(
("provider", "config"),
[
("langfuse", {"public_key": "pk", "secret_key": "******"}),
(
"mlflow",
{
"tracking_uri": "http://localhost:5000",
"experiment_id": "0",
"username": None,
"password": "******",
"project_url": "http://localhost:5000",
},
),
],
)
def test_trace_app_config_get_configured_returns_exact_legacy_body(
app: Flask,
tracing_configs: MagicMock,
provider: str,
config: dict[str, str | None],
) -> None:
tracing_configs.get.return_value = AppTracingConfigRecord(
id="trace-config-1",
app_id=APP_ID,
tracing_provider=PROVIDER,
tracing_config={"public_key": "pk", "secret_key": "******"},
tracing_provider=provider,
tracing_config=config,
is_active=True,
created_at=datetime(2026, 1, 2, 3, 4, 5),
updated_at=datetime(2026, 1, 3, 4, 5, 6),
)
with app.test_request_context("/?tracing_provider=langfuse"):
with app.test_request_context(f"/?tracing_provider={provider}"):
result = _original(ops_trace_module.TraceAppConfigApi.get)(
ops_trace_module.TraceAppConfigApi(),
ops_trace_module.TraceProviderQuery(tracing_provider=PROVIDER),
ops_trace_module.TraceProviderQuery(tracing_provider=provider),
_request_context(),
UUID(APP_ID),
)
@@ -281,8 +299,8 @@ def test_trace_app_config_get_configured_returns_exact_legacy_body(
assert result == {
"id": "trace-config-1",
"app_id": APP_ID,
"tracing_provider": PROVIDER,
"tracing_config": {"public_key": "pk", "secret_key": "******"},
"tracing_provider": provider,
"tracing_config": config,
"is_active": True,
"created_at": "2026-01-02 03:04:05",
"updated_at": "2026-01-03 04:05:06",
@@ -296,15 +314,24 @@ def test_trace_app_config_get_configured_returns_exact_legacy_body(
pytest.param("patch", {"result": "success"}, id="update"),
],
)
@pytest.mark.parametrize(
("provider", "config"),
[
("langfuse", {"public_key": "pk", "secret_key": "sk"}),
("mlflow", {"tracking_uri": "", "experiment_id": "0", "username": None, "password": "******"}),
],
)
def test_trace_app_config_write_returns_expected_response(
app: Flask,
tracing_configs: MagicMock,
method_name: str,
expected_result: object,
provider: str,
config: dict[str, str | None],
) -> None:
payload = ops_trace_module.TraceConfigPayload(
tracing_provider=PROVIDER,
tracing_config={"public_key": "pk", "secret_key": "sk"},
tracing_provider=provider,
tracing_config=config,
)
with app.test_request_context("/", method=method_name.upper()):
@@ -319,8 +346,8 @@ def test_trace_app_config_write_returns_expected_response(
_service_method(tracing_configs, method_name).assert_called_once_with(
context=_request_context(),
app_id=APP_ID,
tracing_provider=PROVIDER,
tracing_config={"public_key": "pk", "secret_key": "sk"},
tracing_provider=provider,
tracing_config=config,
)
@@ -1,807 +0,0 @@
from __future__ import annotations
import gzip
from collections.abc import Callable
from inspect import unwrap
from unittest.mock import MagicMock
import httpx
import pytest
from flask import Flask, Response
from pydantic import SecretStr
from werkzeug.exceptions import (
BadGateway,
Forbidden,
HTTPException,
NotFound,
RequestEntityTooLarge,
ServiceUnavailable,
)
from controllers.console import bp
from controllers.console.knowledge_fs_proxy import (
_console_api_errors,
_proxy_knowledge_fs_non_get,
_proxy_request,
_proxy_response,
proxy_knowledge_fs_get,
proxy_knowledge_fs_options,
proxy_knowledge_fs_write,
)
from controllers.console.wraps import RBACPermission
from models.account import Account, TenantAccountRole
from services.knowledge_fs_operations import (
KnowledgeFSMethod,
KnowledgeFSOperation,
KnowledgeFSResponseKind,
)
from services.knowledge_fs_proxy import (
KnowledgeFSAccessDeniedError,
KnowledgeFSConfigurationError,
KnowledgeFSRouteNotAllowedError,
KnowledgeFSUpstreamResponse,
get_knowledge_fs_operation,
)
@pytest.fixture(autouse=True)
def enable_knowledge_fs(config_overrides: Callable[..., None]) -> None:
config_overrides(KNOWLEDGE_FS_ENABLED=True)
def _upstream(
response: httpx.Response,
kind: KnowledgeFSResponseKind = "buffered",
*,
error_status_map: tuple[tuple[int, int], ...] = ((401, 502), (403, 403)),
max_response_bytes: int | None = None,
) -> KnowledgeFSUpstreamResponse:
operation = KnowledgeFSOperation(
operation_id="testOperation",
method="GET",
path="test",
response_kind=kind,
required_scope="knowledge-spaces:read",
rbac_permission=RBACPermission.DATASET_READONLY,
legacy_role="reader",
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=(),
error_status_map=error_status_map,
)
return KnowledgeFSUpstreamResponse(response, kind, operation)
def _set_current_workspace(
monkeypatch: pytest.MonkeyPatch,
*,
editor: bool = True,
has_edit_permission: bool = True,
admin_or_owner: bool = True,
) -> None:
if admin_or_owner:
role = TenantAccountRole.ADMIN
elif editor and has_edit_permission:
role = TenantAccountRole.EDITOR
elif editor:
role = TenantAccountRole.DATASET_OPERATOR
else:
role = TenantAccountRole.NORMAL
account = Account(name="Knowledge User", email="knowledge@example.com")
account.id = "account-1"
account.role = role
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",
lambda method, path: _proxy_request(method, path),
)
def test_console_blueprint_registers_generic_knowledge_fs_routes() -> None:
app = Flask("knowledge-fs-route-registration")
app.register_blueprint(bp)
adapter = app.url_map.bind("localhost")
get_endpoint, get_values = adapter.match(
"/console/api/knowledge-fs/knowledge-spaces",
method="GET",
)
assert get_endpoint.endswith("proxy_knowledge_fs_get")
assert get_values == {"upstream_path": "knowledge-spaces"}
for method in ("DELETE", "PATCH", "POST", "PUT"):
write_endpoint, write_values = adapter.match(
"/console/api/knowledge-fs/knowledge-spaces/space-1",
method=method,
)
assert write_endpoint.endswith("proxy_knowledge_fs_write")
assert write_values == {"upstream_path": "knowledge-spaces/space-1"}
options_endpoint, options_values = adapter.match(
"/console/api/knowledge-fs/knowledge-spaces",
method="OPTIONS",
)
assert options_endpoint.endswith("proxy_knowledge_fs_options")
assert options_values == {"upstream_path": "knowledge-spaces"}
def test_proxy_options_does_not_require_an_authenticated_account(app: Flask) -> None:
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
method="OPTIONS",
headers={"Access-Control-Request-Method": "GET"},
):
response = app.make_response(proxy_knowledge_fs_options("knowledge-spaces"))
assert response.status_code == 204
def test_proxy_options_is_hidden_when_knowledge_fs_is_disabled(
app: Flask,
config_overrides: Callable[..., None],
) -> None:
config_overrides(KNOWLEDGE_FS_ENABLED=False)
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
method="OPTIONS",
headers={"Access-Control-Request-Method": "GET"},
):
response = app.make_response(proxy_knowledge_fs_options("knowledge-spaces"))
assert response.status_code == 404
@pytest.mark.parametrize(
("upstream_path", "requested_method"),
[
("unregistered", "GET"),
("knowledge-spaces", "DELETE"),
("knowledge-spaces", ""),
],
)
def test_proxy_options_hides_unregistered_operations(
app: Flask,
upstream_path: str,
requested_method: str,
) -> None:
headers = {"Access-Control-Request-Method": requested_method} if requested_method else None
with app.test_request_context(
f"/console/api/knowledge-fs/{upstream_path}",
method="OPTIONS",
headers=headers,
):
response = app.make_response(proxy_knowledge_fs_options(upstream_path))
assert response.status_code == 404
def test_proxy_is_hidden_when_knowledge_fs_is_disabled(app: Flask, config_overrides: Callable[..., None]) -> None:
config_overrides(KNOWLEDGE_FS_ENABLED=False)
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method="GET"):
with pytest.raises(NotFound):
_proxy_request("GET", "knowledge-spaces")
@pytest.mark.parametrize(
("route", "method"),
[
(proxy_knowledge_fs_get, "GET"),
(proxy_knowledge_fs_write, "POST"),
],
)
def test_proxy_routes_are_hidden_before_downstream_work_when_disabled(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
config_overrides: Callable[..., None],
route,
method: KnowledgeFSMethod,
) -> None:
config_overrides(KNOWLEDGE_FS_ENABLED=False)
proxy_request = MagicMock()
proxy_non_get = MagicMock()
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy_request)
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_knowledge_fs_non_get", proxy_non_get)
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method=method):
response = app.make_response(route("knowledge-spaces"))
assert response.status_code == 404
proxy_request.assert_not_called()
proxy_non_get.assert_not_called()
@pytest.mark.parametrize("method", ["HEAD", "OPTIONS"])
def test_generic_get_hides_unregistered_methods(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
method: str,
) -> None:
proxy = MagicMock(return_value=Response(status=200))
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._proxy_request", proxy)
route = unwrap(proxy_knowledge_fs_get)
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
method=method,
):
with pytest.raises(NotFound):
route("knowledge-spaces")
proxy.assert_not_called()
@pytest.mark.parametrize(
("route", "method", "path", "permission"),
[
(proxy_knowledge_fs_get, "GET", "knowledge-spaces", RBACPermission.DATASET_READONLY),
(
proxy_knowledge_fs_write,
"POST",
"knowledge-spaces",
RBACPermission.DATASET_CREATE_AND_MANAGEMENT,
),
],
)
def test_generic_routes_delegate_to_the_authorized_service_use_case(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
route,
method: KnowledgeFSMethod,
path: str,
permission: RBACPermission,
) -> None:
response_kind: KnowledgeFSResponseKind = "buffered"
upstream = httpx.Response(
200,
content=b"asset" if response_kind == "binary" else b"{}",
headers={"Content-Type": "application/octet-stream" if response_kind == "binary" else "application/json"},
)
proxy = MagicMock(return_value=_upstream(upstream, response_kind))
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
proxy,
)
_set_current_workspace(monkeypatch)
_bypass_policy_wrappers(monkeypatch)
raw_route = unwrap(route)
with app.test_request_context(f"/console/api/knowledge-fs/{path}", method=method, data=b"{}"):
response = raw_route(path)
assert isinstance(response, Response)
assert proxy.call_args.kwargs["account"].id == "account-1"
assert proxy.call_args.kwargs["tenant_id"] == "tenant-1"
assert proxy.call_args.kwargs["method"] == method
assert proxy.call_args.kwargs["path"] == path
assert get_knowledge_fs_operation(method, path).rbac_permission == permission
def test_read_post_applies_knowledge_rate_limit_once(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
config_overrides: Callable[..., None],
) -> None:
config_overrides(RBAC_ENABLED=False)
account = Account(name="Knowledge User", email="knowledge@example.com")
account.id = "account-1"
account.role = TenantAccountRole.DATASET_OPERATOR
def current_workspace() -> tuple[Account, str]:
return account, "tenant-1"
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.current_account_with_tenant", current_workspace)
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", current_workspace)
check_access = MagicMock(return_value=True)
monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access)
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_authorized_request", proxy)
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces", method="POST"):
response = _proxy_knowledge_fs_non_get("POST", "knowledge-spaces")
assert isinstance(response, Response)
zadd.assert_called_once()
proxy.assert_called_once()
authorization = proxy.call_args.args[0]
assert authorization.account_id == "account-1"
assert authorization.tenant_id == "tenant-1"
assert authorization.operation.operation_id == "createKnowledgeSpace"
check_access.assert_called_once()
def test_denied_write_does_not_consume_the_workspace_rate_limit(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
) -> None:
account = Account(name="Knowledge Viewer", email="viewer@example.com")
account.id = "account-1"
account.role = TenantAccountRole.NORMAL
def current_workspace() -> tuple[Account, str]:
return account, "tenant-1"
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.current_account_with_tenant",
current_workspace,
)
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", current_workspace)
monkeypatch.setattr(
"controllers.console.wraps.FeatureService.get_knowledge_rate_limit",
MagicMock(return_value=MagicMock(enabled=True, limit=10)),
)
zadd = MagicMock()
monkeypatch.setattr("controllers.console.wraps.redis_client.zadd", zadd)
upstream_request = MagicMock()
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", upstream_request)
route = unwrap(proxy_knowledge_fs_write)
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
method="POST",
data=b"{}",
content_type="application/json",
):
with pytest.raises(Forbidden):
route("knowledge-spaces")
zadd.assert_not_called()
upstream_request.assert_not_called()
def test_generic_get_forwards_path_query_and_raw_response(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
) -> None:
forward = MagicMock(
return_value=_upstream(
httpx.Response(
200,
content=gzip.compress(b'{"items":[],"nextCursor":null}'),
headers={
"Cache-Control": "no-store",
"Content-Encoding": "gzip",
"Content-Disposition": 'attachment; filename="result.json"',
"Content-Type": "application/json",
"Retry-After": "3",
"Set-Cookie": "kfs=secret",
"X-Trace-Id": "trace-1",
},
)
)
)
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
forward,
)
_set_current_workspace(monkeypatch)
route = unwrap(proxy_knowledge_fs_get)
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
query_string=[("limit", "20"), ("cursor", "first"), ("cursor", "second")],
):
response = route("knowledge-spaces")
request = forward.call_args.kwargs
assert request["account"].id == "account-1"
assert request["method"] == "GET"
assert request["path"] == "knowledge-spaces"
assert request["tenant_id"] == "tenant-1"
assert request["accept"] is None
assert request["content_type"] is None
assert request["query"] == b"limit=20&cursor=first&cursor=second"
assert request["body"] is None
assert isinstance(response, Response)
assert response.status_code == 200
assert response.get_json() == {"items": [], "nextCursor": None}
assert response.headers["Cache-Control"] == "no-store"
assert response.headers["Content-Disposition"] == 'attachment; filename="result.json"'
assert response.headers["Retry-After"] == "3"
assert response.headers["X-Trace-Id"] == "trace-1"
assert "Content-Encoding" not in response.headers
assert "Set-Cookie" not in response.headers
def test_generic_write_forwards_path_raw_body_and_current_tenant(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
) -> None:
method = "POST"
path = "knowledge-spaces"
forward = MagicMock(
return_value=_upstream(
httpx.Response(
201,
content=b'{"id":"space-1","tenantId":"tenant-1"}',
headers={"Content-Type": "application/json"},
)
)
)
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
forward,
)
_set_current_workspace(monkeypatch)
_bypass_policy_wrappers(monkeypatch)
route = unwrap(proxy_knowledge_fs_write)
body = b'{"idempotencyKey":"create-product-docs","name":"Product docs"}'
with app.test_request_context(
f"/console/api/knowledge-fs/{path}",
method=method,
data=body,
content_type="application/json",
):
response = route(path)
request = forward.call_args.kwargs
assert request["account"].id == "account-1"
assert request["method"] == method
assert request["path"] == path
assert request["tenant_id"] == "tenant-1"
assert request["accept"] is None
assert request["content_type"] == "application/json"
assert request["query"] is None
assert request["body"] == body
assert isinstance(response, Response)
assert response.status_code == 201
assert response.get_json()["tenantId"] == "tenant-1"
def test_generic_write_forwards_through_the_authorized_production_path(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
config_overrides: Callable[..., None],
) -> None:
config_overrides(
KNOWLEDGE_FS_BASE_URL="http://knowledge-fs.test",
KNOWLEDGE_FS_JWT_SECRET=SecretStr("production-secret-with-at-least-32-bytes"),
)
account = Account(name="Knowledge User", email="knowledge@example.com")
account.id = "account-1"
account.role = TenantAccountRole.DATASET_OPERATOR
def current_workspace() -> tuple[Account, str]:
return account, "tenant-1"
monkeypatch.setattr("controllers.console.knowledge_fs_proxy.current_account_with_tenant", current_workspace)
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", current_workspace)
check_access = MagicMock(return_value=True)
monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access)
monkeypatch.setattr(
"controllers.console.wraps.FeatureService.get_knowledge_rate_limit",
MagicMock(return_value=MagicMock(enabled=False)),
)
upstream_request = MagicMock(
return_value=httpx.Response(
201,
content=b'{"id":"space-1","tenantId":"tenant-1"}',
headers={"Content-Type": "application/json"},
)
)
monkeypatch.setattr("services.knowledge_fs_proxy.ssrf_proxy.make_request", upstream_request)
route = unwrap(proxy_knowledge_fs_write)
body = b'{"idempotencyKey":"create-product-docs","name":"Product docs"}'
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
method="POST",
query_string={"source": "console"},
data=body,
content_type="application/json",
headers={"X-Trace-Id": "trace-1"},
):
response = route("knowledge-spaces")
assert isinstance(response, Response)
assert response.status_code == 201
assert response.get_json() == {"id": "space-1", "tenantId": "tenant-1"}
check_access.assert_called_once()
assert upstream_request.call_args.kwargs["method"] == "POST"
assert upstream_request.call_args.kwargs["url"] == "http://knowledge-fs.test/knowledge-spaces"
assert upstream_request.call_args.kwargs["params"] == b"source=console"
assert upstream_request.call_args.kwargs["content"] == body
assert upstream_request.call_args.kwargs["headers"]["x-trace-id"] == "trace-1"
def test_generic_write_forwards_contract_declared_request_headers(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
) -> None:
forward = MagicMock(
return_value=_upstream(
httpx.Response(202, content=b'{"status":"accepted"}', headers={"Content-Type": "application/json"})
)
)
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
forward,
)
_set_current_workspace(monkeypatch)
_bypass_policy_wrappers(monkeypatch)
route = unwrap(proxy_knowledge_fs_write)
body = b'{"name":"Product docs"}'
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
method="POST",
data=body,
content_type="application/json",
headers={"X-Trace-Id": "trace-1"},
):
response = route("knowledge-spaces")
assert isinstance(response, Response)
assert response.status_code == 202
assert forward.call_args.kwargs["request_headers"].get("X-Trace-Id") == "trace-1"
def test_contract_response_headers_cannot_bypass_the_proxy_denylist() -> None:
denied_headers = (
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailer",
"transfer-encoding",
"upgrade",
)
upstream = httpx.Response(
200,
content=b"{}",
headers=dict.fromkeys(denied_headers, "blocked"),
)
response = _proxy_response(
_upstream(upstream),
tenant_id="tenant-1",
contract_response_headers=denied_headers,
max_response_bytes=1024 * 1024,
)
for name in denied_headers:
assert name not in response.headers
def test_contract_response_headers_forward_binary_hardening_headers() -> None:
upstream = httpx.Response(
200,
content=b"asset",
headers={
"Content-Security-Policy": "sandbox; default-src 'none'",
"Content-Type": "image/png",
"X-Content-Type-Options": "nosniff",
},
)
response = _proxy_response(
_upstream(upstream, "binary"),
tenant_id="tenant-1",
contract_response_headers=("content-security-policy", "x-content-type-options"),
max_response_bytes=25 * 1024 * 1024,
)
assert response.headers["Content-Security-Policy"] == "sandbox; default-src 'none'"
assert response.headers["X-Content-Type-Options"] == "nosniff"
def test_authorized_service_denial_is_exposed_as_forbidden(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
MagicMock(side_effect=KnowledgeFSAccessDeniedError("workspace access denied")),
)
_set_current_workspace(monkeypatch)
route = unwrap(proxy_knowledge_fs_get)
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
with pytest.raises(Forbidden):
route("knowledge-spaces")
def test_generic_post_rejects_oversized_body(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
forward = MagicMock()
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
forward,
)
_set_current_workspace(monkeypatch)
_bypass_policy_wrappers(monkeypatch)
route = unwrap(proxy_knowledge_fs_write)
monkeypatch.setattr("controllers.console.knowledge_fs_proxy._MAX_PROXY_BODY_BYTES", 8)
with app.test_request_context(
"/console/api/knowledge-fs/knowledge-spaces",
method="POST",
data=b"x" * 9,
content_type="application/json",
):
with pytest.raises(RequestEntityTooLarge):
route("knowledge-spaces")
forward.assert_not_called()
def test_server_credential_rejection_is_not_exposed_as_browser_auth_failure(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
MagicMock(
return_value=_upstream(
httpx.Response(
401,
content=b'{"error":"invalid server credential"}',
headers={"Content-Type": "application/json", "WWW-Authenticate": "Bearer"},
)
)
),
)
_set_current_workspace(monkeypatch)
route = unwrap(proxy_knowledge_fs_get)
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
with pytest.raises(BadGateway, match="authentication failed"):
route("knowledge-spaces")
def test_resource_authorization_rejection_is_exposed_as_forbidden(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
MagicMock(
return_value=_upstream(
httpx.Response(
403,
content=b'{"error":"resource access denied"}',
headers={"Content-Type": "application/json"},
)
)
),
)
_set_current_workspace(monkeypatch)
route = unwrap(proxy_knowledge_fs_get)
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
with pytest.raises(Forbidden):
route("knowledge-spaces")
def test_proxy_response_applies_operation_specific_error_status_mapping() -> None:
upstream = httpx.Response(
429,
content=b'{"error":"rate limited"}',
headers={"Content-Type": "application/json"},
)
with pytest.raises(ServiceUnavailable):
_proxy_response(
_upstream(upstream, error_status_map=((429, 503),)),
tenant_id="tenant-1",
contract_response_headers=(),
max_response_bytes=1024 * 1024,
)
assert upstream.is_closed
def test_proxy_response_preserves_nonstandard_mapped_error_status() -> None:
upstream = httpx.Response(
429,
content=b'{"error":"rate limited"}',
headers={"Content-Type": "application/json"},
)
with pytest.raises(HTTPException) as exc_info:
_proxy_response(
_upstream(upstream, error_status_map=((429, 499),)),
tenant_id="tenant-1",
contract_response_headers=(),
max_response_bytes=1024 * 1024,
)
assert exc_info.value.code == 499
assert upstream.is_closed
def test_contract_response_headers_are_deduplicated_case_insensitively() -> None:
upstream = httpx.Response(
200,
content=b"asset",
headers={
"Cache-Control": "private",
"Content-Disposition": 'inline; filename="asset.png"',
"Content-Type": "image/png",
},
)
response = _proxy_response(
_upstream(upstream, "binary"),
tenant_id="tenant-1",
contract_response_headers=("cache-control", "content-disposition"),
max_response_bytes=25 * 1024 * 1024,
)
assert response.headers.getlist("Cache-Control") == ["private"]
assert response.headers.getlist("Content-Disposition") == ['inline; filename="asset.png"']
def test_configuration_error_is_reported_as_unavailable(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
MagicMock(side_effect=KnowledgeFSConfigurationError("missing token")),
)
_set_current_workspace(monkeypatch)
route = unwrap(proxy_knowledge_fs_get)
with app.test_request_context("/console/api/knowledge-fs/knowledge-spaces"):
with pytest.raises(ServiceUnavailable, match="misconfigured"):
route("knowledge-spaces")
def test_disallowed_kfs_route_is_hidden_as_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(
"controllers.console.knowledge_fs_proxy.proxy_knowledge_fs_request",
MagicMock(side_effect=KnowledgeFSRouteNotAllowedError("blocked")),
)
_set_current_workspace(monkeypatch)
route = unwrap(proxy_knowledge_fs_get)
with app.test_request_context("/console/api/knowledge-fs/openapi.json"):
with pytest.raises(NotFound):
route("openapi.json")
@pytest.mark.parametrize("method", ["PATCH", "POST"])
def test_disallowed_non_get_route_is_hidden_as_not_found(
app: Flask,
monkeypatch: pytest.MonkeyPatch,
method: KnowledgeFSMethod,
) -> None:
_set_current_workspace(monkeypatch)
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}
@@ -177,15 +177,15 @@ def test_sse_transport_endpoint_validation():
# Valid endpoint (same origin)
valid_endpoint = "http://example.com/messages/session123"
assert transport._validate_endpoint_url(valid_endpoint) == True
assert transport._validate_endpoint_url(valid_endpoint) is True
# Invalid endpoint (different origin)
invalid_endpoint = "http://malicious.com/messages/session123"
assert transport._validate_endpoint_url(invalid_endpoint) == False
assert transport._validate_endpoint_url(invalid_endpoint) is False
# Invalid endpoint (different scheme)
invalid_scheme = "https://example.com/messages/session123"
assert transport._validate_endpoint_url(invalid_scheme) == False
assert transport._validate_endpoint_url(invalid_scheme) is False
def test_sse_transport_message_parsing():
@@ -1,3 +1,4 @@
import contextlib
import logging
import queue
import time
@@ -466,10 +467,8 @@ def test_check_receiver_status_fail(streams):
future = executor.submit(raise_err)
session._receiver_future = future
try:
with contextlib.suppress(BaseException):
future.result()
except BaseException:
pass
with pytest.raises(RuntimeError, match="Receiver failed"):
session.check_receiver_status()
@@ -0,0 +1,280 @@
"""Offline pagination regressions using the real client and in-memory JSON-RPC."""
from collections.abc import Generator, Mapping
from contextlib import contextmanager
from datetime import timedelta
from queue import Queue
from threading import Thread
import pytest
from core.mcp.auth_client import MCPClientWithAuthRetry
from core.mcp.error import MCPAuthError, MCPConnectionError
from core.mcp.mcp_client import MCPClient
from core.mcp.session.client_session import ClientSession
from core.mcp.types import (
ErrorData,
JSONRPCError,
JSONRPCMessage,
JSONRPCRequest,
JSONRPCResponse,
ListToolsRequest,
ListToolsResult,
SessionMessage,
Tool,
)
def tool(name: str) -> Tool:
return Tool(name=name, description=name, inputSchema={"type": "object", "properties": {}})
@contextmanager
def client_for_pages(
pages: Mapping[str | None, ListToolsResult | ErrorData], client: MCPClient | None = None
) -> Generator[tuple[MCPClient, list[ListToolsRequest]]]:
incoming: Queue[SessionMessage | Exception | None] = Queue()
outgoing: Queue[SessionMessage | None] = Queue()
requests: list[ListToolsRequest] = []
errors: list[Exception] = []
def serve_pages() -> None:
while (message := outgoing.get()) is not None:
request = message.message.root
assert isinstance(request, JSONRPCRequest)
try:
payload = ListToolsRequest.model_validate({"method": request.method, "params": request.params})
requests.append(payload)
cursor = payload.params.cursor if payload.params is not None else None
page = pages[cursor]
response: JSONRPCResponse | JSONRPCError
if isinstance(page, ErrorData):
response = JSONRPCError(jsonrpc="2.0", id=request.id, error=page)
else:
response = JSONRPCResponse(
jsonrpc="2.0",
id=request.id,
result=page.model_dump(mode="json", by_alias=True, exclude_none=True),
)
except Exception as error:
errors.append(error)
response = JSONRPCError(
jsonrpc="2.0", id=request.id, error=ErrorData(code=500, message="Unexpected pagination request")
)
incoming.put(SessionMessage(message=JSONRPCMessage(response)))
thread = Thread(target=serve_pages, daemon=True)
thread.start()
client = client if client is not None else MCPClient("http://127.0.0.1:1/mcp")
try:
with client._exit_stack:
client._session = client._exit_stack.enter_context(
ClientSession(incoming, outgoing, read_timeout_seconds=timedelta(seconds=2))
)
yield client, requests
finally:
outgoing.put(None)
thread.join(timeout=2)
assert not thread.is_alive()
assert not errors, errors
def test_single_page_control() -> None:
with client_for_pages({None: ListToolsResult(tools=[tool("alpha")])}) as (client, requests):
assert [x.name for x in client.list_tools()] == ["alpha"]
assert len(requests) == 1
@pytest.mark.parametrize("first_page", [[tool("alpha")], []], ids=["populated-first-page", "empty-first-page"])
def test_all_pages_are_available(first_page: list[Tool]) -> None:
cursor = "opaque:+/=?page2"
with client_for_pages(
{
None: ListToolsResult(tools=first_page, nextCursor=cursor),
cursor: ListToolsResult(tools=[tool("beta")]),
}
) as (client, requests):
expected = [t.name for t in first_page] + ["beta"]
assert [x.name for x in client.list_tools()] == expected
assert requests[1].params is not None
assert requests[1].params.cursor == cursor
def test_no_session_control() -> None:
with pytest.raises(ValueError, match="Session not initialized"):
MCPClient("http://127.0.0.1:1/mcp").list_tools()
@pytest.mark.parametrize("cursors", [("opaque-A", "opaque-A"), ("opaque-A", "opaque-B", "opaque-A")])
def test_repeated_cursor_is_a_bounded_error(cursors: tuple[str, ...]) -> None:
pages = {
None if index == 0 else cursors[index - 1]: ListToolsResult(tools=[tool(f"page-{index}")], nextCursor=cursor)
for index, cursor in enumerate(cursors)
}
with client_for_pages(pages) as (client, requests):
with pytest.raises(MCPConnectionError, match="repeated pagination cursor"):
client.list_tools()
assert len(requests) == len(cursors)
def test_later_page_auth_error_restarts_clean_inventory(monkeypatch: pytest.MonkeyPatch) -> None:
client = MCPClientWithAuthRetry("http://127.0.0.1:1/mcp")
old_pages = {
None: ListToolsResult(tools=[tool("stale-alpha")], nextCursor="opaque-old"),
"opaque-old": ErrorData(code=401, message="controlled later-page authentication failure"),
}
new_pages = {
None: ListToolsResult(tools=[tool("fresh-alpha")], nextCursor="opaque-new"),
"opaque-new": ListToolsResult(tools=[tool("beta")]),
}
auth_errors: list[MCPAuthError] = []
cleanup: list[str] = []
reconnects: list[str] = []
with (
client_for_pages(old_pages, client) as (_, old_requests),
client_for_pages(new_pages) as (
reconnected_client,
new_requests,
),
):
old_session = client._session
assert old_session is not None
client._initialized = True
client._exit_stack.callback(lambda: cleanup.append("closed"))
def reconnect() -> None:
reconnects.append("reconnected")
client._session = reconnected_client._session
# Patch only the OAuth exchange and transport reconnect, not session requests.
monkeypatch.setattr(client, "_handle_auth_error", auth_errors.append)
monkeypatch.setattr(client, "_initialize", reconnect)
assert [t.name for t in client.list_tools()] == ["fresh-alpha", "beta"]
assert [p.params.cursor if p.params is not None else None for p in old_requests] == [None, "opaque-old"]
assert [p.params.cursor if p.params is not None else None for p in new_requests] == [None, "opaque-new"]
assert len(auth_errors) == 1
assert str(auth_errors[0]) == "controlled later-page authentication failure"
assert cleanup == ["closed"]
assert reconnects == ["reconnected"]
assert old_session._receiver_future is not None
assert old_session._receiver_future.done()
assert client._has_retried is False
@pytest.mark.parametrize("page_sizes", [(10_000,), (9_999, 1)], ids=["single-page", "two-pages"])
def test_exact_tool_limit_on_terminal_page(page_sizes: tuple[int, ...]) -> None:
pages: dict[str | None, ListToolsResult] = {}
expected: list[str] = []
for index, size in enumerate(page_sizes):
names = [f"page-{index}-tool-{number}" for number in range(size)]
expected.extend(names)
cursor = None if index == 0 else f"page-{index}"
next_cursor = f"page-{index + 1}" if index + 1 < len(page_sizes) else None
pages[cursor] = ListToolsResult(tools=[tool(name) for name in names], nextCursor=next_cursor)
with client_for_pages(pages) as (client, requests):
assert [result.name for result in client.list_tools()] == expected
assert len(requests) == len(page_sizes)
@pytest.mark.parametrize("next_cursor", [None, "more"])
@pytest.mark.parametrize("page_sizes", [(10_001,), (9_999, 2)], ids=["oversized-page", "accumulated-overflow"])
def test_tool_limit_overflow_never_returns_partial_inventory(
page_sizes: tuple[int, ...], next_cursor: str | None
) -> None:
pages: dict[str | None, ListToolsResult] = {}
for index, size in enumerate(page_sizes):
cursor = None if index == 0 else f"page-{index}"
continuation = f"page-{index + 1}" if index + 1 < len(page_sizes) else next_cursor
pages[cursor] = ListToolsResult(
tools=[tool(f"page-{index}-tool-{number}") for number in range(size)], nextCursor=continuation
)
with client_for_pages(pages) as (client, requests):
with pytest.raises(MCPConnectionError, match="10,000-tool discovery limit"):
client.list_tools()
assert len(requests) == len(page_sizes)
@pytest.mark.parametrize("next_cursor", ["more", ""])
@pytest.mark.parametrize("first_page_size", [10_000, 9_999], ids=["single-page", "two-pages"])
def test_exact_tool_limit_with_continuation_fails_without_fetching_again(
first_page_size: int, next_cursor: str
) -> None:
pages: dict[str | None, ListToolsResult] = {
None: ListToolsResult(
tools=[tool(f"tool-{number}") for number in range(first_page_size)],
nextCursor=next_cursor if first_page_size == 10_000 else "last-tool",
)
}
if first_page_size < 10_000:
pages["last-tool"] = ListToolsResult(tools=[tool("tool-9999")], nextCursor=next_cursor)
with client_for_pages(pages) as (client, requests):
with pytest.raises(MCPConnectionError, match="10,000-tool discovery limit"):
client.list_tools()
assert len(requests) == (1 if first_page_size == 10_000 else 2)
@pytest.mark.parametrize("auth_retry", [False, True], ids=["base-client", "auth-retry-client"])
def test_empty_pages_exhaust_request_budget_without_extra_request_or_auth_retry(auth_retry: bool) -> None:
client = MCPClientWithAuthRetry("http://127.0.0.1:1/mcp") if auth_retry else MCPClient("http://127.0.0.1:1/mcp")
pages = {
None if index == 0 else f"fresh-cursor-{index - 1}": ListToolsResult(
tools=[], nextCursor=f"fresh-cursor-{index}"
)
for index in range(100)
}
with client_for_pages(pages, client) as (_, requests):
with pytest.raises(MCPConnectionError, match="100-request discovery limit") as exc_info:
client.list_tools()
assert not isinstance(exc_info.value, MCPAuthError)
assert len(requests) == 100
@pytest.mark.parametrize(
"last_page_tools", [[], [tool("last-page-tool")]], ids=["empty-terminal", "populated-terminal"]
)
def test_terminal_page_on_last_allowed_request_succeeds(last_page_tools: list[Tool]) -> None:
pages = {
None if index == 0 else f"fresh-cursor-{index - 1}": ListToolsResult(
tools=[], nextCursor=f"fresh-cursor-{index}"
)
for index in range(99)
}
pages["fresh-cursor-98"] = ListToolsResult(tools=last_page_tools)
with client_for_pages(pages) as (client, requests):
assert client.list_tools() == last_page_tools
assert len(requests) == 100
@pytest.mark.parametrize("error_on_first_page", [True, False], ids=["first-page", "later-page"])
def test_session_error_propagates_unchanged(error_on_first_page: bool, monkeypatch: pytest.MonkeyPatch) -> None:
error = RuntimeError("controlled session error")
calls: list[str | None] = []
with client_for_pages({None: ListToolsResult(tools=[tool("alpha")], nextCursor="more")}) as (client, _):
session = client._session
assert session is not None
list_tools = session.list_tools
def fail_on_page(cursor: str | None = None) -> ListToolsResult:
calls.append(cursor)
if error_on_first_page or cursor is not None:
raise error
return list_tools(cursor=cursor)
monkeypatch.setattr(session, "list_tools", fail_on_page)
with pytest.raises(RuntimeError) as exc_info:
client.list_tools()
assert exc_info.value is error
assert calls == ([None] if error_on_first_page else [None, "more"])
def test_empty_string_is_an_opaque_continuation_cursor() -> None:
with client_for_pages(
{
None: ListToolsResult(tools=[tool("alpha")], nextCursor=""),
"": ListToolsResult(tools=[tool("beta")]),
}
) as (client, requests):
assert [result.name for result in client.list_tools()] == ["alpha", "beta"]
assert len(requests) == 2
assert requests[1].params is not None
assert requests[1].params.cursor == ""
@@ -3,6 +3,7 @@
from __future__ import annotations
import base64
import contextlib
from unittest.mock import patch
import pytest
@@ -288,8 +289,6 @@ def test_invoke_skips_forwarding_outside_enterprise_edition(config_overrides):
with patch("core.entities.mcp_provider.MCPProviderEntity.decrypt_server_url", return_value="u"):
with patch("core.entities.mcp_provider.MCPProviderEntity.decrypt_headers", return_value={}):
# Should not raise; should not call enterprise.
try:
with contextlib.suppress(Exception):
tool.invoke_remote_mcp_tool({}, user_id=None, app_id=None)
except Exception:
pass
inject.assert_not_called()
@@ -758,7 +758,7 @@ class TestSegmentTypeValidationIntegration:
result_first = segment_type.is_valid(valid_value, ArrayValidation.FIRST)
result_all = segment_type.is_valid(valid_value, ArrayValidation.ALL)
assert result_none == result_first == result_all == True, (
assert result_none == result_first == result_all is True, (
f"{segment_type} should ignore array_validation parameter"
)
@@ -2,7 +2,7 @@
import contextvars
import threading
from contextlib import contextmanager
from contextlib import contextmanager, suppress
from typing import Any
from unittest.mock import MagicMock
@@ -176,10 +176,8 @@ class TestExecutionContext:
return
except Exception as exc:
errors.append(exc)
try:
with suppress(Exception):
barrier.abort()
except Exception:
pass
t1 = threading.Thread(target=worker)
t2 = threading.Thread(target=worker)
@@ -28,39 +28,39 @@ def test_declares_variable():
# BuiltinNodeTypes is not imported directly, we need to mock or just use the generator method
# LLM node
assert WorkflowGenerator._declares_variable({"data": {"type": "llm"}}, "text") == True
assert WorkflowGenerator._declares_variable({"data": {"type": "llm"}}, "text") is True
llm_so = {"data": {"type": "llm", "structured_output": {"schema": {"properties": {"json_var": {}}}}}}
assert WorkflowGenerator._declares_variable(llm_so, "json_var") == True
assert WorkflowGenerator._declares_variable(llm_so, "other_var") == False
assert WorkflowGenerator._declares_variable(llm_so, "json_var") is True
assert WorkflowGenerator._declares_variable(llm_so, "other_var") is False
# Code node
assert (
WorkflowGenerator._declares_variable({"data": {"type": "code", "outputs": {"code_var": "str"}}}, "code_var")
== True
is True
)
# Knowledge retrieval
assert WorkflowGenerator._declares_variable({"data": {"type": "knowledge-retrieval"}}, "result") == True
assert WorkflowGenerator._declares_variable({"data": {"type": "knowledge-retrieval"}}, "result") is True
# Parameter extractor
assert (
WorkflowGenerator._declares_variable(
{"data": {"type": "parameter-extractor", "parameters": [{"name": "param1"}]}}, "param1"
)
== True
is True
)
# HTTP request
assert WorkflowGenerator._declares_variable({"data": {"type": "http-request"}}, "body") == True
assert WorkflowGenerator._declares_variable({"data": {"type": "http-request"}}, "body") is True
# Template transform
assert WorkflowGenerator._declares_variable({"data": {"type": "template-transform"}}, "output") == True
assert WorkflowGenerator._declares_variable({"data": {"type": "template-transform"}}, "output") is True
# Tool
assert WorkflowGenerator._declares_variable({"data": {"type": "tool"}}, "anything") == True
assert WorkflowGenerator._declares_variable({"data": {"type": "tool"}}, "anything") is True
# Other node (default false)
assert WorkflowGenerator._declares_variable({"data": {"type": "unknown"}}, "anything") == False
assert WorkflowGenerator._declares_variable({"data": {"type": "unknown"}}, "anything") is False
def test_collect_container_errors():

Some files were not shown because too many files have changed in this diff Show More